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, Form, UploadFile import base64 import uuid import uvicorn # Worker that imports the agent and exposes a simple HTTP API app = FastAPI(title="Agent Worker") # --------------------------------------------------------------------------- # Schema profile registry (dynamic multi-DB support) # --------------------------------------------------------------------------- try: from .schema_profile import ProfileRegistry, _redact_db_url except Exception: import Agent.schema_profile as _sp ProfileRegistry = _sp.ProfileRegistry _redact_db_url = _sp._redact_db_url SCHEMAS_DIR = os.path.join("uploads", "schemas") PROFILE_STORE_PATH = os.path.join("chat_store", "schema_profiles.json") profile_registry = ProfileRegistry(PROFILE_STORE_PATH) def _safe_profile(entry): """Copy of a registry entry safe for public API responses (password masked).""" if not isinstance(entry, dict): return entry safe = dict(entry) safe["db_url"] = _redact_db_url(safe.get("db_url")) return safe SUPPORTED_DB_SUFFIXES = {".sqlite", ".db", ".sqlite3", ".csv"} class AddSchemaJsonReq(BaseModel): name: str connection_url: str activate: Optional[bool] = True notes: Optional[str] = None def _detect_demo_kind(name: str, tblset: Optional[set] = None) -> str: """Classify a profile as the built-in employee demo or a user schema.""" n = (name or "").lower() if "demo_employee" in n: return "demo" if tblset and all( t in tblset for t in ["employee", "department", "salary", "dept_emp"] ): return "demo" # Connection URL pointing at employee-ish database if "employee" in (name or ""): return "demo" return "user" 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 # ------------------------------------------------------------------ # Built-in demo employee database: ALWAYS available as a profile. # Sourced from the authentic employees sample dataset bundled in the repo # (employee-sample-database/sqlite/dataset_small/employee.db). Used as the # active database out of the box when no DB_CONNECTION_URL is provided. # ------------------------------------------------------------------ demo_db_url = None demo_db_path = os.path.join( os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "employee-sample-database", "sqlite", "dataset_small", "employee.db", ) if os.path.isfile(demo_db_path): demo_db_url = f"sqlite:///{demo_db_path.replace(os.sep, '/')}" try: has_demo = any((p or {}).get("kind") == "demo" for p in profile_registry.profiles.values()) if not has_demo: profile_registry.add(name="demo_employee", db_url=demo_db_url, kind="demo", activate=False) print("[WORKER] Registered built-in demo profile 'demo_employee'") except Exception as reg_err: print(f"[WORKER] demo profile registration skipped: {reg_err}") else: print(f"[WORKER] Bundled demo employee DB not found at {demo_db_path}") # ------------------------------------------------------------------ # Decide which database to initialize against: # 1. DB_CONNECTION_URL (env/compose wins) 2. last active profile 3. demo # ------------------------------------------------------------------ db_url = os.environ.get("DB_CONNECTION_URL") active = profile_registry.active() # Boot priority: 1) the user's last ACTIVATED profile (persisted across # restarts; never the auto-registered demo placeholder), 2) the env # DB_CONNECTION_URL, 3) the bundled demo file. Previously the env secret # ALWAYS won, so every restart silently yanked the engine back to the # owner's default database even when the registry said the user was # working in (e.g.) their car database — engine and registry disagreed # and queries ran against the wrong database. _active_kind = (active or {}).get("kind") target_uri = ( (active or {}).get("db_url") if active and _active_kind != "demo" else None ) or db_url or demo_db_url if target_uri: try: # Re-attach owner schema notes stored on the matching profile, so a # restarted server carries the domain hints with no re-registration. _pre = next( (p for p in profile_registry.list() if ((p.get("db_url") or "").replace("\\", "/") == target_uri.replace("\\", "/"))), None, ) await agent.initialize_database(uri=target_uri, schema_notes=(_pre or {}).get("notes")) # Keep the registry's active entry in sync with what we connected to. # IMPORTANT: force the booted database to OWN the active slot — the # demo registration above grabs "active" on a fresh store (ProfileRegistry # .add() activates when no active exists), which left Space deployments # showing the demo DB even though the env DB_CONNECTION_URL booted. try: entries = profile_registry.list() match = next( (p for p in entries if ((p.get("db_url") or "").replace("\\", "/") == target_uri.replace("\\", "/"))), None, ) if match is None: prof = getattr(agent, "ACTIVE_SCHEMA_PROFILE", None) _from_env = bool(db_url and target_uri == db_url) _kind = "user" if _from_env else ( "demo" if (prof and getattr(prof, "is_employee_db", False)) else "user") _name = (getattr(prof, "name", None) if prof else None) or ( "Employee_DB" if _from_env else "demo_employee") pid = f"boot_{uuid.uuid4().hex[:8]}" profile_registry.profiles[pid] = { "name": _name, "db_url": target_uri, "kind": _kind, "created_at": __import__("time").time(), } profile_registry.active_id = pid profile_registry.save() print(f"[WORKER] Registered ACTIVE schema profile '{_name}' (kind={_kind}) for env DB") else: if profile_registry.active_id != match["id"]: profile_registry.switch(match["id"]) print(f"[WORKER] Active profile switched to '{match['name']}'") else: print(f"[WORKER] Active profile already '{match['name']}'") except Exception as reg_err: print(f"[WORKER] profile bootstrap skipped: {reg_err}") except Exception as e: # Log but continue; agentic_workflow has fallback behavior print(f"Worker: initialize_database failed: {e}") # If the persisted profile's DB became unreachable (rotated # credentials, deleted file), retry once with the env DB so the # worker still boots with a usable database. if db_url and target_uri != db_url: try: print("[WORKER] Falling back to DB_CONNECTION_URL after failed profile boot") await agent.initialize_database(uri=db_url) _env_match = next( (p for p in profile_registry.list() if ((p.get("db_url") or "").replace("\\", "/") == db_url.replace("\\", "/"))), None, ) if _env_match and profile_registry.active_id != _env_match["id"]: profile_registry.switch(_env_match["id"]) print(f"[WORKER] Active profile switched to '{_env_match['name']}' (env fallback)") except Exception as e2: print(f"Worker: env fallback failed: {e2}") else: print("Worker: no database available (no env DB, no profiles, no demo bundle)") # Optional Groq Vision health check in cloud mode try: from .tools import groq_vision_health hc = groq_vision_health.health_check() if not hc.get("ok", True): print(f"[WORKER] Groq Vision health warning: {hc.get('error')} (model={hc.get('model')})") else: print(f"[WORKER] Groq Vision health: ok (model={hc.get('model')})") except Exception as e: print(f"[WORKER] Groq Vision health check skipped: {e}") @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 # Chart contract: strip the embedded chart marker from the answer and # expose it as a first-class field for the web UI (ECharts). chart_out = None try: from .chart_contract import extract_chart_marker clean_answer, chart_out = extract_chart_marker(result) if chart_out is not None: result = clean_answer except Exception: chart_out = None return {"answer": result, "logs": logs, "sql": sql_out, "chart": chart_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) # --------------------------------------------------------------------------- # Schema profile endpoints — dynamic multi-database support # --------------------------------------------------------------------------- @app.get("/schemas") async def schemas_list(): return {"profiles": [_safe_profile(p) for p in profile_registry.list()], "active_id": profile_registry.active_id} @app.post("/schemas/add") async def schemas_add_json(req: AddSchemaJsonReq): """Register a database via connection URL (postgres/mysql/sqlite).""" global agent name = (req.name or "").strip()[:80] if not name or not req.connection_url: raise HTTPException(status_code=400, detail="name and connection_url are required") entry = profile_registry.add( name=name, db_url=req.connection_url, kind=_detect_demo_kind(name), activate=bool(req.activate), notes=req.notes, ) if bool(req.activate) and agent is not None: try: await agent.initialize_database(uri=req.connection_url, schema_notes=req.notes) except Exception as e: raise HTTPException(status_code=400, detail=f"Profile saved, but DB init failed: {e}") # Return full schema digest preview return {"profile": _safe_profile(entry), "schema": agent.ACTIVE_SCHEMA_PROFILE.summarize() if agent else None} @app.post("/schemas/add-file") async def schemas_add_file(file: UploadFile = File(...), name: Optional[str] = Form(None), activate: bool = Form(True), notes: Optional[str] = Form(None)): """Register a database via file upload (.sqlite/.db/.sqlite3 or .csv).""" global agent fname = file.filename or "upload" ext = os.path.splitext(fname)[-1].lower() if ext not in SUPPORTED_DB_SUFFIXES: raise HTTPException( status_code=400, detail=f"Unsupported file type '{ext}'. Supported: {sorted(SUPPORTED_DB_SUFFIXES)}", ) display_name = (name or os.path.splitext(fname)[0]).strip()[:80] or "user_schema" os.makedirs(SCHEMAS_DIR, exist_ok=True) if ext == ".csv": # Convert CSV into a single-table SQLite database import pandas as pd import tempfile sqlite_path = os.path.join(SCHEMAS_DIR, f"{uuid.uuid4().hex}.db") try: contents = await file.read() df = pd.read_csv(io.BytesIO(contents)) except Exception as e: raise HTTPException(status_code=400, detail=f"CSV parse failed: {e}") table_name = "".join(ch for ch in display_name if ch.isalnum() or ch == "_") or "data" try: engine_path = os.path.abspath(sqlite_path).replace("\\", "/") df.to_sql(table_name, f"sqlite:///{engine_path}", index=False, if_exists="replace") except Exception as e: raise HTTPException(status_code=500, detail=f"CSV→SQLite conversion failed: {e}") db_url = f"sqlite:///{engine_path}" kind = "user" else: # Copy uploaded SQLite file into schemas dir sqlite_path = os.path.join(SCHEMAS_DIR, f"{uuid.uuid4().hex}{ext}") try: contents = await file.read() with open(sqlite_path, "wb") as fh: fh.write(contents) except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to save upload: {e}") engine_path = sqlite_path.replace("\\", "/") db_url = f"sqlite:///{engine_path}" # Poke the file to validate it is really a SQLite database try: import sqlite3 conn = sqlite3.connect(sqlite_path) tblset = {r[0] for r in conn.execute( "SELECT name FROM sqlite_master WHERE type IN ('table','view')" ).fetchall()} conn.close() except Exception as e: raise HTTPException(status_code=400, detail=f"Invalid SQLite file: {e}") kind = _detect_demo_kind(display_name, tblset) entry = profile_registry.add(name=display_name, db_url=db_url, kind=kind, activate=activate, notes=notes) if activate and agent is not None: try: await agent.initialize_database(uri=db_url, schema_notes=notes) except Exception as e: raise HTTPException(status_code=400, detail=f"Profile saved, but DB init failed: {e}") return {"profile": _safe_profile(entry), "schema": agent.ACTIVE_SCHEMA_PROFILE.summarize() if agent else None} @app.post("/schemas/switch/{profile_id}") async def schemas_switch(profile_id: str): """Switch the active database profile and re-initialize the agent against it.""" global agent if agent is None: raise HTTPException(status_code=503, detail="Agent not initialized") entry = profile_registry.switch(profile_id) if not entry: raise HTTPException(status_code=404, detail="Profile not found") # Dispose previous async engine best-effort try: if getattr(agent, "async_engine", None) is not None: await agent.async_engine.dispose() except Exception: pass try: await agent.initialize_database(uri=entry["db_url"], schema_notes=entry.get("notes")) except Exception as e: raise HTTPException(status_code=400, detail=f"Switch failed: {e}") return {"profile": _safe_profile(entry), "schema": agent.ACTIVE_SCHEMA_PROFILE.summarize()} class SetSchemaNotesReq(BaseModel): notes: str = "" @app.post("/schemas/{profile_id}/notes") async def schemas_set_notes(profile_id: str, req: SetSchemaNotesReq): """Persist owner-supplied schema notes (domain hints) for a profile. Notes are injected verbatim into every agent prompt for that database, supplying the semantics that live introspection cannot infer. """ notes = (req.notes or "").strip()[:4000] if not profile_registry.set_notes(profile_id, notes): raise HTTPException(status_code=404, detail="Profile not found") # Live effect on the running agent when the profile is currently active # (no re-initialization needed). if agent is not None and profile_id == profile_registry.active_id: try: if getattr(agent, "ACTIVE_SCHEMA_PROFILE", None) is not None: agent.ACTIVE_SCHEMA_PROFILE.notes = notes except Exception: pass return {"id": profile_id, "notes": notes} @app.delete("/schemas/{profile_id}") async def schemas_delete(profile_id: str): entry = profile_registry.get(profile_id) if not entry: raise HTTPException(status_code=404, detail="Profile not found") # Remove uploaded backing file for file-based profiles db_url = entry.get("db_url") or "" try: if db_url.startswith("sqlite:///"): p = db_url.replace("sqlite:///", "", 1) if os.path.exists(p) and str(SCHEMAS_DIR) in p: os.remove(p) except Exception: pass profile_registry.remove(profile_id) return {"deleted": profile_id} @app.get("/schema/current") async def schema_current(): """Current active schema summary (tables/columns/samples) for the UI.""" global agent if agent is None: raise HTTPException(status_code=503, detail="Agent not initialized") return agent.ACTIVE_SCHEMA_PROFILE.summarize() if __name__ == '__main__': port = int(os.environ.get("WORKER_PORT", 8700)) uvicorn.run(app, host="127.0.0.1", port=port)