File size: 6,659 Bytes
de36225 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | 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")
@app.get("/")
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
@app.get('/health')
async def health():
return {"status": "ok"}
@app.post('/api/query')
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()})
@app.post('/api/ocr-qa')
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()})
@app.on_event('startup')
async def startup_event():
global client
if client is None:
client = httpx.AsyncClient(timeout=WORKER_TIMEOUT)
@app.on_event('shutdown')
async def shutdown_event():
global client
if client is not None:
await client.aclose()
@app.post('/api/test')
async def api_test(req: QueryReq):
# Simple echo endpoint to verify POST handling without contacting the worker
return {"ok": True, "echo": req.question}
@app.post('/upload-image')
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))
@app.get('/api/test')
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)
|