Aniruddha7/QueryLens-Text2SQL_DocVQA-V2

1import os2import asyncio3import io4import sys5import contextlib6from fastapi import FastAPI, HTTPException7from pydantic import BaseModel8from typing import Optional9from fastapi import File, UploadFile10import base6411import uuid12import uvicorn13 14# Worker that imports the agent and exposes a simple HTTP API15app = FastAPI(title="Agent Worker")16 17class QueryReq(BaseModel):18    question: str19    doc_id: Optional[str] = None20    mode: Optional[str] = None  # 'text2sql' or 'ocr_qa'21 22 23class UploadResp(BaseModel):24    doc_id: str25    text_preview: str26    image_url: str27 28class OcrQAReq(BaseModel):29    question: str30    doc_id: str31 32class OcrQAResp(BaseModel):33    doc_id: str34    answer: str35    used_llm: bool36 37# Lazy import agent module (do at startup so heavy init happens in worker)38agent = None39# Lock to prevent concurrent capture of stdout/stderr which would interleave logs40capture_lock = asyncio.Lock()41 42@app.on_event("startup")43async def startup_event():44    global agent45    # Import here so the LLM and DB init occur inside the worker process46    # Use package-relative import so we load Agent.agentic_workflow (not a top-level file)47    try:48        from . import agentic_workflow as agentic_workflow49    except Exception:50        # fallback to absolute package import51        import Agent.agentic_workflow as agentic_workflow52    agent = agentic_workflow53    # Initialize DB (best-effort)54    db_url = os.environ.get("DB_CONNECTION_URL")55    if db_url:56        try:57            await agent.initialize_database(uri=db_url)58        except Exception as e:59            # Log but continue; agentic_workflow has fallback behavior60            print(f"Worker: initialize_database failed: {e}")61    else:62        print("Worker: DB_CONNECTION_URL not set; skipping DB init")63 64@app.get("/health")65async def health():66    return {"status": "ok"}67 68 69@app.post("/query")70async def query(req: QueryReq):71    global agent72    if agent is None:73        raise HTTPException(status_code=503, detail="Agent not initialized")74    q = req.question75    doc_id = getattr(req, 'doc_id', None)76    mode = getattr(req, 'mode', None)77    # Capture stdout/stderr produced by the agent for this request78    buf_out = io.StringIO()79    buf_err = io.StringIO()80    async with capture_lock:81        try:82            with contextlib.redirect_stdout(buf_out), contextlib.redirect_stderr(buf_err):83                # Robust mode switch: if mode is 'ocr_qa' or doc_id is present, use OCR agent; else use SQL logic84                if (mode == 'ocr_qa') or (doc_id and (mode is None or mode == 'auto')):85                    # Use OCR Q&A agent86                    # ocr_agent_qa may be sync, so run in executor87                    import asyncio88                    answer = await asyncio.get_event_loop().run_in_executor(89                        None, lambda: agent.ocr_agent_qa(q, doc_id)90                    )91                    result = answer92                    sql_out = None  # No SQL for OCR Q&A93                else:94                    # Use standard Text-to-SQL agentic workflow95                    result = await agent.agentic_query_process(q, None)96                    try:97                        sql_out = getattr(agent, 'LAST_GENERATED_SQL', None)98                    except Exception:99                        sql_out = None100        except ValueError as ve:101            # Treat known generation failures as normal answers so UI doesn't show nested JSON error102            err_msg = str(ve)103            if "Unable to generate SQL query" in err_msg:104                result = err_msg  # Provide user-visible plain text105                sql_out = None106            else:107                import traceback108                traceback.print_exc(file=buf_err)109                raise HTTPException(status_code=500, detail=err_msg)110        except Exception as e:111            import traceback112            traceback.print_exc(file=buf_err)113            raise HTTPException(status_code=500, detail=str(e))114    logs = buf_out.getvalue()115    errlogs = buf_err.getvalue()116    if errlogs:117        logs = logs + "\n" + errlogs118    return {"answer": result, "logs": logs, "sql": sql_out}119 120 121@app.post("/ocr-qa", response_model=OcrQAResp)122async def ocr_qa(req: OcrQAReq):123    """Answer a question about a previously uploaded document using its doc_id.124 125    This uses the agent's ocr_agent_qa helper which pulls stored OCR text and126    consults the LLM (or returns a low-memory fallback excerpt).127    """128    global agent129    if agent is None:130        raise HTTPException(status_code=503, detail="Agent not initialized")131    question = req.question132    doc_id = req.doc_id133    # Run potentially blocking OCR QA in thread pool if it's sync134    try:135        answer = await asyncio.get_event_loop().run_in_executor(136            None, lambda: agent.ocr_agent_qa(question, doc_id)137        )138    except Exception as e:139        raise HTTPException(status_code=500, detail=f"ocr_qa failed: {e}")140    used_llm = not answer.startswith("[Low memory]") and not answer.startswith("No OCR text") and not answer.startswith("Error while")141    return OcrQAResp(doc_id=doc_id, answer=answer, used_llm=used_llm)142 143 144@app.post("/upload-image", response_model=UploadResp)145async def upload_image(file: UploadFile = File(...)):146    """Accept an image, save locally, call MCP document_scanner.process eagerly,147    persist OCR to private chat store, and return a doc_id plus preview.148    """149    global agent150    if agent is None:151        raise HTTPException(status_code=503, detail="Agent not initialized")152 153    # Ensure uploads directory exists154    uploads_dir = "uploads"155    os.makedirs(uploads_dir, exist_ok=True)156 157    # Save uploaded file158    filename = f"{uuid.uuid4()}_{file.filename}"159    path = os.path.join(uploads_dir, filename)160    try:161        contents = await file.read()162        with open(path, "wb") as f:163            f.write(contents)164    except Exception as e:165        raise HTTPException(status_code=500, detail=f"Failed to save upload: {e}")166 167    # Call MCP document_scanner tool eagerly with base64 payload to avoid exposing file server168    b64 = base64.b64encode(contents).decode('utf-8')169    try:170        from .mcp_client import mcp_call_tool171        ctx = {"tool_call": {"name": "document_scanner.process", "args": {"image_bytes": b64}}}172        resp = mcp_call_tool(prompt="run document scanner", timeout=60.0, context=ctx)173        # Expect tool_result in response174        tool_result = resp.get("tool_result") if isinstance(resp, dict) else None175    except Exception:176        # Fall back to local direct call if MCP unavailable or tool failed177        tool_result = None178        try:179            if agent and hasattr(agent, 'tools') and hasattr(agent.tools, 'document_scanner'):180                try:181                    tool_result = agent.tools.document_scanner.process_image(image_bytes=b64)182                except Exception:183                    tool_result = None184        except Exception:185            tool_result = None186        if not tool_result:187            tool_result = {"doc_id": str(uuid.uuid4()), "text": "(ocr unavailable)", "metadata": {}}188 189    # Persist OCR to private chat store file for retrieval by doc_id190    doc_id = tool_result.get("doc_id") if isinstance(tool_result, dict) and tool_result.get("doc_id") else str(uuid.uuid4())191    ocr_text = tool_result.get("text", "") if isinstance(tool_result, dict) else str(tool_result)192 193    # Save to docs/ and also to chat_store_private for history194    docs_dir = os.path.join("chat_store", "docs")195    os.makedirs(docs_dir, exist_ok=True)196    doc_path = os.path.join(docs_dir, f"{doc_id}.txt")197    try:198        with open(doc_path, "w", encoding="utf-8") as df:199            df.write(ocr_text)200    except Exception as e:201        print(f"Failed to persist OCR text: {e}")202 203    # Also persist the original image path so Granite Vision can do direct visual Q&A204    img_path_file = os.path.join(docs_dir, f"{doc_id}.img_path")205    try:206        abs_img_path = os.path.abspath(path)207        with open(img_path_file, "w", encoding="utf-8") as ipf:208            ipf.write(abs_img_path)209        print(f"[upload] Saved image path for doc_id={doc_id}: {abs_img_path}")210    except Exception as e:211        print(f"[upload] Failed to persist image path: {e}")212 213    # Also add a short entry to chat_store_private for traceability214    try:215        try:216            from .agentic_workflow import chat_store_private, ChatMessage, MessageRole217        except ImportError:218            from Agent.agentic_workflow import chat_store_private, ChatMessage, MessageRole219        # Create minimal messages220        user_msg = ChatMessage(role=MessageRole.USER, content=f"Uploaded document {doc_id}")221        assistant_msg = ChatMessage(role=MessageRole.ASSISTANT, content=f"OCR stored: {ocr_text[:400]}")222        chat_store_private.add_message(key="conversation", message=user_msg)223        chat_store_private.add_message(key="conversation", message=assistant_msg)224        chat_store_private.persist(str("chat_store/chat_store_private.json"))225    except Exception:226        pass227 228    # Return doc id and small preview229    preview = ocr_text[:400]230    image_url = path231    return UploadResp(doc_id=doc_id, text_preview=preview, image_url=image_url)232 233if __name__ == '__main__':234    port = int(os.environ.get("WORKER_PORT", 8700))235    uvicorn.run(app, host="127.0.0.1", port=port)236