import os import asyncio import io import sys import contextlib from fastapi import FastAPI, HTTPException from pydantic import BaseModel from typing import Optional from fastapi import File, UploadFile import base64 import uuid import uvicorn # Worker that imports the agent and exposes a simple HTTP API app = FastAPI(title="Agent Worker") class QueryReq(BaseModel): question: str doc_id: Optional[str] = None mode: Optional[str] = None # 'text2sql' or 'ocr_qa' class UploadResp(BaseModel): doc_id: str text_preview: str image_url: str class OcrQAReq(BaseModel): question: str doc_id: str class OcrQAResp(BaseModel): doc_id: str answer: str used_llm: bool # Lazy import agent module (do at startup so heavy init happens in worker) agent = None # Lock to prevent concurrent capture of stdout/stderr which would interleave logs capture_lock = asyncio.Lock() @app.on_event("startup") async def startup_event(): global agent # Import here so the LLM and DB init occur inside the worker process # Use package-relative import so we load Agent.agentic_workflow (not a top-level file) try: from . import agentic_workflow as agentic_workflow except Exception: # fallback to absolute package import import Agent.agentic_workflow as agentic_workflow agent = agentic_workflow # Initialize DB (best-effort) db_url = os.environ.get("DB_CONNECTION_URL") if db_url: try: await agent.initialize_database(uri=db_url) except Exception as e: # Log but continue; agentic_workflow has fallback behavior print(f"Worker: initialize_database failed: {e}") else: print("Worker: DB_CONNECTION_URL not set; skipping DB init") @app.get("/health") async def health(): return {"status": "ok"} @app.post("/query") async def query(req: QueryReq): global agent if agent is None: raise HTTPException(status_code=503, detail="Agent not initialized") q = req.question doc_id = getattr(req, 'doc_id', None) mode = getattr(req, 'mode', None) # Capture stdout/stderr produced by the agent for this request buf_out = io.StringIO() buf_err = io.StringIO() async with capture_lock: try: with contextlib.redirect_stdout(buf_out), contextlib.redirect_stderr(buf_err): # Robust mode switch: if mode is 'ocr_qa' or doc_id is present, use OCR agent; else use SQL logic if (mode == 'ocr_qa') or (doc_id and (mode is None or mode == 'auto')): # Use OCR Q&A agent # ocr_agent_qa may be sync, so run in executor import asyncio answer = await asyncio.get_event_loop().run_in_executor( None, lambda: agent.ocr_agent_qa(q, doc_id) ) result = answer sql_out = None # No SQL for OCR Q&A else: # Use standard Text-to-SQL agentic workflow result = await agent.agentic_query_process(q, None) try: sql_out = getattr(agent, 'LAST_GENERATED_SQL', None) except Exception: sql_out = None except ValueError as ve: # Treat known generation failures as normal answers so UI doesn't show nested JSON error err_msg = str(ve) if "Unable to generate SQL query" in err_msg: result = err_msg # Provide user-visible plain text sql_out = None else: import traceback traceback.print_exc(file=buf_err) raise HTTPException(status_code=500, detail=err_msg) except Exception as e: import traceback traceback.print_exc(file=buf_err) raise HTTPException(status_code=500, detail=str(e)) logs = buf_out.getvalue() errlogs = buf_err.getvalue() if errlogs: logs = logs + "\n" + errlogs return {"answer": result, "logs": logs, "sql": sql_out} @app.post("/ocr-qa", response_model=OcrQAResp) async def ocr_qa(req: OcrQAReq): """Answer a question about a previously uploaded document using its doc_id. This uses the agent's ocr_agent_qa helper which pulls stored OCR text and consults the LLM (or returns a low-memory fallback excerpt). """ global agent if agent is None: raise HTTPException(status_code=503, detail="Agent not initialized") question = req.question doc_id = req.doc_id # Run potentially blocking OCR QA in thread pool if it's sync try: answer = await asyncio.get_event_loop().run_in_executor( None, lambda: agent.ocr_agent_qa(question, doc_id) ) except Exception as e: raise HTTPException(status_code=500, detail=f"ocr_qa failed: {e}") used_llm = not answer.startswith("[Low memory]") and not answer.startswith("No OCR text") and not answer.startswith("Error while") return OcrQAResp(doc_id=doc_id, answer=answer, used_llm=used_llm) @app.post("/upload-image", response_model=UploadResp) async def upload_image(file: UploadFile = File(...)): """Accept an image, save locally, call MCP document_scanner.process eagerly, persist OCR to private chat store, and return a doc_id plus preview. """ global agent if agent is None: raise HTTPException(status_code=503, detail="Agent not initialized") # Ensure uploads directory exists uploads_dir = "uploads" os.makedirs(uploads_dir, exist_ok=True) # Save uploaded file filename = f"{uuid.uuid4()}_{file.filename}" path = os.path.join(uploads_dir, filename) try: contents = await file.read() with open(path, "wb") as f: f.write(contents) except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to save upload: {e}") # Call MCP document_scanner tool eagerly with base64 payload to avoid exposing file server b64 = base64.b64encode(contents).decode('utf-8') try: from .mcp_client import mcp_call_tool ctx = {"tool_call": {"name": "document_scanner.process", "args": {"image_bytes": b64}}} resp = mcp_call_tool(prompt="run document scanner", timeout=60.0, context=ctx) # Expect tool_result in response tool_result = resp.get("tool_result") if isinstance(resp, dict) else None except Exception: # Fall back to local direct call if MCP unavailable or tool failed tool_result = None try: if agent and hasattr(agent, 'tools') and hasattr(agent.tools, 'document_scanner'): try: tool_result = agent.tools.document_scanner.process_image(image_bytes=b64) except Exception: tool_result = None except Exception: tool_result = None if not tool_result: tool_result = {"doc_id": str(uuid.uuid4()), "text": "(ocr unavailable)", "metadata": {}} # Persist OCR to private chat store file for retrieval by doc_id doc_id = tool_result.get("doc_id") if isinstance(tool_result, dict) and tool_result.get("doc_id") else str(uuid.uuid4()) ocr_text = tool_result.get("text", "") if isinstance(tool_result, dict) else str(tool_result) # Save to docs/ and also to chat_store_private for history docs_dir = os.path.join("chat_store", "docs") os.makedirs(docs_dir, exist_ok=True) doc_path = os.path.join(docs_dir, f"{doc_id}.txt") try: with open(doc_path, "w", encoding="utf-8") as df: df.write(ocr_text) except Exception as e: print(f"Failed to persist OCR text: {e}") # Also persist the original image path so Granite Vision can do direct visual Q&A img_path_file = os.path.join(docs_dir, f"{doc_id}.img_path") try: abs_img_path = os.path.abspath(path) with open(img_path_file, "w", encoding="utf-8") as ipf: ipf.write(abs_img_path) print(f"[upload] Saved image path for doc_id={doc_id}: {abs_img_path}") except Exception as e: print(f"[upload] Failed to persist image path: {e}") # Also add a short entry to chat_store_private for traceability try: try: from .agentic_workflow import chat_store_private, ChatMessage, MessageRole except ImportError: from Agent.agentic_workflow import chat_store_private, ChatMessage, MessageRole # Create minimal messages user_msg = ChatMessage(role=MessageRole.USER, content=f"Uploaded document {doc_id}") assistant_msg = ChatMessage(role=MessageRole.ASSISTANT, content=f"OCR stored: {ocr_text[:400]}") chat_store_private.add_message(key="conversation", message=user_msg) chat_store_private.add_message(key="conversation", message=assistant_msg) chat_store_private.persist(str("chat_store/chat_store_private.json")) except Exception: pass # Return doc id and small preview preview = ocr_text[:400] image_url = path return UploadResp(doc_id=doc_id, text_preview=preview, image_url=image_url) if __name__ == '__main__': port = int(os.environ.get("WORKER_PORT", 8700)) uvicorn.run(app, host="127.0.0.1", port=port)