Spaces:
Running on Zero
Running on Zero
| import os | |
| import httpx | |
| import logging | |
| import traceback | |
| from typing import Optional | |
| from fastapi import FastAPI, HTTPException, Request | |
| from pydantic import BaseModel | |
| from fastapi.staticfiles import StaticFiles | |
| from fastapi.responses import FileResponse, JSONResponse | |
| from fastapi.middleware.cors import CORSMiddleware | |
| WORKER_URL = os.environ.get("WORKER_URL", "http://localhost:8700") | |
| # Timeout (seconds) when calling the worker; re-execute and heavy queries can take longer. | |
| # Make configurable via environment variable WORKER_TIMEOUT (default 300s). | |
| WORKER_TIMEOUT = float(os.environ.get("WORKER_TIMEOUT", "450")) # Increased to 450s to handle complex queries | |
| # Reusable AsyncClient to avoid per-request connection setup overhead | |
| client: httpx.AsyncClient | None = None | |
| app = FastAPI(title="Text2SQL Web API") | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], # adjust in production | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| # Serve static frontend under /static and provide a root route | |
| static_dir = os.path.join(os.path.dirname(__file__), 'static') | |
| if os.path.isdir(static_dir): | |
| app.mount("/static", StaticFiles(directory=static_dir, html=True), name="static") | |
| async def root_index(): | |
| index_path = os.path.join(static_dir, 'index.html') | |
| return FileResponse(index_path) | |
| class QueryReq(BaseModel): | |
| question: str | |
| # Optional doc_id produced by a prior /upload-image call | |
| doc_id: Optional[str] = None | |
| # Explicit mode hint from client: 'text2sql' or 'ocr_qa'. When omitted, worker may auto-detect. | |
| mode: Optional[str] = None | |
| class OcrQAReq(BaseModel): | |
| question: str | |
| doc_id: str | |
| class OcrQAResp(BaseModel): | |
| doc_id: str | |
| answer: str | |
| used_llm: bool | |
| async def health(): | |
| return {"status": "ok"} | |
| async def api_query(req: QueryReq): | |
| # Use a shared AsyncClient created at startup to avoid TCP/TLS/HTTP handshakes per request. | |
| global client | |
| if client is None: | |
| # fallback to create a temporary client (shouldn't happen if startup ran) | |
| client = httpx.AsyncClient(timeout=WORKER_TIMEOUT) | |
| try: | |
| payload = {"question": req.question} | |
| # Forward explicit mode if provided so the worker doesn't infer from lingering doc_id | |
| if getattr(req, 'mode', None): | |
| payload['mode'] = req.mode | |
| if getattr(req, 'doc_id', None): | |
| payload['doc_id'] = req.doc_id | |
| r = await client.post(f"{WORKER_URL}/query", json=payload) | |
| r.raise_for_status() | |
| return r.json() | |
| except httpx.HTTPStatusError as e: | |
| # Log the full traceback and response body for debugging | |
| logging.error("Worker returned non-2xx status: %s", e.response.status_code) | |
| logging.error("Response text: %s", getattr(e.response, 'text', str(e))) | |
| logging.error(traceback.format_exc()) | |
| raise HTTPException(status_code=e.response.status_code, detail=getattr(e.response, 'text', str(e))) | |
| except httpx.ReadTimeout as e: | |
| # Worker took too long. Return a 504 gateway timeout with a helpful message. | |
| logging.error("Timeout contacting worker after %ss: %s", WORKER_TIMEOUT, str(e)) | |
| logging.error(traceback.format_exc()) | |
| return JSONResponse(status_code=504, content={ | |
| "error": "worker_timeout", | |
| "message": f"Worker did not respond within {WORKER_TIMEOUT} seconds. The query may be long-running (re-execute) or the worker is busy. Try increasing WORKER_TIMEOUT or running the worker with more resources.", | |
| }) | |
| except Exception as e: | |
| logging.error("Error contacting worker: %s", str(e)) | |
| logging.error(traceback.format_exc()) | |
| # Return a JSON response containing the error detail to help debugging from the client | |
| return JSONResponse(status_code=500, content={"error": str(e), "trace": traceback.format_exc()}) | |
| async def api_ocr_qa(req: OcrQAReq): | |
| global client | |
| if client is None: | |
| client = httpx.AsyncClient(timeout=WORKER_TIMEOUT) | |
| try: | |
| payload = {"question": req.question, "doc_id": req.doc_id} | |
| r = await client.post(f"{WORKER_URL}/ocr-qa", json=payload) | |
| r.raise_for_status() | |
| return r.json() | |
| except httpx.HTTPStatusError as e: | |
| logging.error("Worker OCR QA non-2xx: %s", e.response.status_code) | |
| logging.error("Response text: %s", getattr(e.response, 'text', str(e))) | |
| return JSONResponse(status_code=e.response.status_code, content={"error": getattr(e.response, 'text', str(e))}) | |
| except Exception as e: | |
| logging.error("Error contacting worker OCR QA: %s", str(e)) | |
| return JSONResponse(status_code=500, content={"error": str(e), "trace": traceback.format_exc()}) | |
| async def startup_event(): | |
| global client | |
| if client is None: | |
| client = httpx.AsyncClient(timeout=WORKER_TIMEOUT) | |
| async def shutdown_event(): | |
| global client | |
| if client is not None: | |
| await client.aclose() | |
| async def api_test(req: QueryReq): | |
| # Simple echo endpoint to verify POST handling without contacting the worker | |
| return {"ok": True, "echo": req.question} | |
| async def upload_image(file: bytes = None, request: Request = None): | |
| """Proxy an image upload to the worker's /upload-image endpoint. | |
| Accepts multipart/form-data file field 'file' from the browser. | |
| Returns the worker JSON (doc_id, text_preview, image_url) on success. | |
| """ | |
| # Use the request.form() to access the uploaded file | |
| try: | |
| form = await request.form() | |
| f = form.get('file') | |
| if f is None: | |
| raise HTTPException(status_code=400, detail='file field is required') | |
| # forward the file to the worker | |
| files = {'file': (f.filename, await f.read(), f.content_type)} | |
| async with httpx.AsyncClient(timeout=120) as c: | |
| r = await c.post(f"{WORKER_URL}/upload-image", files=files) | |
| r.raise_for_status() | |
| return JSONResponse(status_code=200, content=r.json()) | |
| except HTTPException: | |
| raise | |
| except Exception as e: | |
| logging.error('upload-image proxy failed: %s', str(e)) | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| # --------------------------------------------------------------------------- | |
| # Schema profile proxy endpoints — dynamic multi-database support | |
| # --------------------------------------------------------------------------- | |
| async def api_schemas(): | |
| global client | |
| try: | |
| r = await client.get(f"{WORKER_URL}/schemas") | |
| r.raise_for_status() | |
| return r.json() | |
| except httpx.HTTPStatusError as e: | |
| return JSONResponse(status_code=e.response.status_code, | |
| content={"error": getattr(e.response, 'text', str(e))}) | |
| except Exception as e: | |
| logging.error('schemas list proxy failed: %s', str(e)) | |
| return JSONResponse(status_code=500, content={"error": str(e)}) | |
| class AddSchemaReq(BaseModel): | |
| name: str | |
| connection_url: str | |
| activate: Optional[bool] = True | |
| notes: Optional[str] = None | |
| async def api_schemas_add(req: AddSchemaReq): | |
| global client | |
| try: | |
| r = await client.post(f"{WORKER_URL}/schemas/add", json=req.model_dump(), | |
| timeout=WORKER_TIMEOUT) | |
| r.raise_for_status() | |
| return r.json() | |
| except httpx.HTTPStatusError as e: | |
| return JSONResponse(status_code=e.response.status_code, | |
| content={"error": getattr(e.response, 'text', str(e))}) | |
| except Exception as e: | |
| logging.error('schemas add proxy failed: %s', str(e)) | |
| return JSONResponse(status_code=500, content={"error": str(e)}) | |
| class SetSchemaNotesReq(BaseModel): | |
| notes: str = "" | |
| async def api_schemas_set_notes(profile_id: str, req: SetSchemaNotesReq): | |
| """Proxy per-database schema notes (domain hints) updates to the worker.""" | |
| global client | |
| try: | |
| r = await client.post(f"{WORKER_URL}/schemas/{profile_id}/notes", | |
| json={"notes": req.notes}, timeout=WORKER_TIMEOUT) | |
| r.raise_for_status() | |
| return r.json() | |
| except httpx.HTTPStatusError as e: | |
| return JSONResponse(status_code=e.response.status_code, | |
| content={"error": getattr(e.response, 'text', str(e))}) | |
| except Exception as e: | |
| logging.error('schemas notes proxy failed: %s', str(e)) | |
| return JSONResponse(status_code=500, content={"error": str(e)}) | |
| async def api_schemas_add_file(request: Request): | |
| """Proxy a schema file upload (.sqlite/.db/.csv) to the worker.""" | |
| global client | |
| try: | |
| form = await request.form() | |
| f = form.get('file') | |
| if f is None: | |
| raise HTTPException(status_code=400, detail='file field is required') | |
| data = {'name': form.get('name') or '', 'activate': form.get('activate', 'true'), | |
| 'notes': form.get('notes') or ''} | |
| files = {'file': (f.filename, await f.read(), f.content_type)} | |
| # DB init can be slow on first connect; generous timeout | |
| async with httpx.AsyncClient(timeout=300) as c: | |
| r = await c.post(f"{WORKER_URL}/schemas/add-file", files=files, data=data) | |
| r.raise_for_status() | |
| return r.json() | |
| except HTTPException: | |
| raise | |
| except httpx.HTTPStatusError as e: | |
| return JSONResponse(status_code=e.response.status_code, | |
| content={"error": getattr(e.response, 'text', str(e))}) | |
| except Exception as e: | |
| logging.error('schemas add-file proxy failed: %s', str(e)) | |
| return JSONResponse(status_code=500, content={"error": str(e)}) | |
| async def api_schemas_switch(profile_id: str): | |
| global client | |
| try: | |
| r = await client.post(f"{WORKER_URL}/schemas/switch/{profile_id}", | |
| timeout=WORKER_TIMEOUT) | |
| r.raise_for_status() | |
| return r.json() | |
| except httpx.HTTPStatusError as e: | |
| return JSONResponse(status_code=e.response.status_code, | |
| content={"error": getattr(e.response, 'text', str(e))}) | |
| except Exception as e: | |
| logging.error('schemas switch proxy failed: %s', str(e)) | |
| return JSONResponse(status_code=500, content={"error": str(e)}) | |
| async def api_schemas_delete(profile_id: str): | |
| global client | |
| try: | |
| r = await client.delete(f"{WORKER_URL}/schemas/{profile_id}") | |
| r.raise_for_status() | |
| return r.json() | |
| except httpx.HTTPStatusError as e: | |
| return JSONResponse(status_code=e.response.status_code, | |
| content={"error": getattr(e.response, 'text', str(e))}) | |
| except Exception as e: | |
| logging.error('schemas delete proxy failed: %s', str(e)) | |
| return JSONResponse(status_code=500, content={"error": str(e)}) | |
| async def api_schema_current(): | |
| global client | |
| try: | |
| r = await client.get(f"{WORKER_URL}/schema/current") | |
| r.raise_for_status() | |
| return r.json() | |
| except httpx.HTTPStatusError as e: | |
| return JSONResponse(status_code=e.response.status_code, | |
| content={"error": getattr(e.response, 'text', str(e))}) | |
| except Exception as e: | |
| logging.error('schema current proxy failed: %s', str(e)) | |
| return JSONResponse(status_code=500, content={"error": str(e)}) | |
| async def api_test_get(): | |
| return {"ok": True, "msg": "POST to /api/test with JSON {question: ...} to echo"} | |
| if __name__ == '__main__': | |
| import uvicorn | |
| port = int(os.environ.get('WEB_PORT', 8000)) | |
| uvicorn.run(app, host='0.0.0.0', port=port) | |