Agent/worker.py

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)