"""Chanakya ASGI server — A2A + REST query + auth + conversation endpoints.""" import asyncio import contextlib import json import os import time as _time from datetime import datetime from pathlib import Path from typing import Any import socketio import logfire from fastapi import FastAPI, Form from fastapi.responses import HTMLResponse from fastcore.xml import to_xml from fasthtml.components import Body, Button, Form as FHForm, H2, Html, Input, Option, P, Select, Strong from loguru import logger as _log from pylogue.shell import app_factory as create_core_app from starlette.middleware.cors import CORSMiddleware from starlette.requests import Request from starlette.responses import JSONResponse, StreamingResponse from backend.chanakya.a2a_context import RequestContext, make_metadata_aware_app, _IMPERSONATION_FILE from backend.chanakya.auth.utils import get_username_from_token from backend.chanakya.auth import ( close_auth_db, get_current_user as auth_me_handler, init_auth_db, login as auth_login_handler, upsert_user as auth_upsert_handler, verify_token as auth_verify_handler, ) from backend.chanakya.config import AGENT_REGISTRY from backend.chanakya.conversation import ( create_conversation_handler, delete_conversation_handler, get_conversation_handler, get_conversations_handler, get_messages_handler, get_user_bookmarks_handler, append_to_history_cache, load_conversation_history, save_message, save_message_handler, serialize_history, toggle_bookmark_handler, update_conversation_title_handler, ) from backend.chanakya.conversation.database import ensure_pool as _ensure_db_pool from backend.chanakya.rt_agent.agent import chanakya from backend.chanakya.schemas import AgentQueryResponse from data.mock_data.users import USERS # --------------------------------------------------------------------------- # Landing metrics queries — ts4_* tables # --------------------------------------------------------------------------- _LM_Q_CONTRACTS = "SELECT COUNT(*) AS total_contracts, SUM(initial_contract_value) AS total_initial_budget FROM ts4_contracts" _LM_Q_ACCEPTED = ( "SELECT COALESCE(SUM(change_to_prices), 0) AS accepted_revenue," " COALESCE(SUM(change_to_days), 0) AS accepted_days" " FROM ts4_quotations" " INNER JOIN ts4_contracts AS c USING (contract_number)" " WHERE status IN ('Accepted_No_CBS', 'Accepted', 'Acc_with_Instruction')" ) _LM_Q_PENDING = ( "SELECT COUNT(*) AS total_quotations, COALESCE(SUM(change_to_prices), 0) AS revenue," " COALESCE(SUM(change_to_days), 0) AS accepted_days, COUNT(DISTINCT contract_number) AS active_contracts" " FROM ts4_quotations" " WHERE status IN ('Submitted', 'In_Review', 'Assessment', 'Internal_Review', 'Revised_Quotation')" ) _LM_Q_OPEN_EW = "SELECT COUNT(*) AS total_ew FROM ts4_early_warnings WHERE status = 'Open'" _LM_Q_OPEN_NCE = "SELECT COUNT(*) AS total_nce FROM ts4_nce WHERE reply_date IS NULL AND quotation_number IS NULL" # Decisions — Combined impact (both budget + timeline) _LM_PENDING_STATUSES = "('Submitted', 'In_Review', 'Assessment', 'Internal_Review', 'Revised_Quotation')" _LM_Q_CO_COUNT = ( "SELECT COUNT(*) AS total_quotations," " COALESCE(SUM(q.change_to_prices), 0) AS ce_revenue," " COALESCE(SUM(q.change_to_days), 0) AS ce_days" " FROM ts4_quotations AS q" " WHERE q.change_to_prices IS NOT NULL AND q.change_to_prices != 0" " AND q.change_to_days IS NOT NULL AND q.change_to_days != 0" f" AND q.status IN {_LM_PENDING_STATUSES}" ) _LM_Q_CO_CARDS = ( "SELECT q.area, q.title, q.change_to_prices, q.change_to_days, q.quotation_date, q.contract_number" " FROM ts4_quotations AS q" " WHERE q.change_to_prices IS NOT NULL AND q.change_to_prices != 0" " AND q.change_to_days IS NOT NULL AND q.change_to_days != 0" f" AND q.status IN {_LM_PENDING_STATUSES}" " ORDER BY q.change_to_prices DESC, q.change_to_days DESC LIMIT 3" ) # Decisions — Timeline only _LM_Q_TL_COUNT = ( "SELECT COUNT(*) AS total_quotations," " COALESCE(SUM(q.change_to_days), 0) AS ce_days" " FROM ts4_quotations AS q" " WHERE q.change_to_prices IS NULL AND q.change_to_days IS NOT NULL" f" AND q.status IN {_LM_PENDING_STATUSES}" ) _LM_Q_TL_CARDS = ( "SELECT q.area, q.title, q.change_to_days, q.quotation_date, q.contract_number" " FROM ts4_quotations AS q" " WHERE q.change_to_days IS NOT NULL AND q.change_to_days != 0" " AND (q.change_to_prices IS NULL OR q.change_to_prices = 0)" f" AND q.status IN {_LM_PENDING_STATUSES}" " ORDER BY q.change_to_days DESC LIMIT 3" ) # Decisions — Budget only _LM_Q_BD_COUNT = ( "SELECT COUNT(*) AS total_quotations," " COALESCE(SUM(q.change_to_prices), 0) AS ce_revenue" " FROM ts4_quotations AS q" " WHERE q.change_to_prices IS NOT NULL AND q.change_to_prices != 0" " AND (q.change_to_days IS NULL OR q.change_to_days = 0)" f" AND q.status IN {_LM_PENDING_STATUSES}" ) _LM_Q_BD_CARDS = ( "SELECT q.area, q.title, q.change_to_prices, q.quotation_date, q.contract_number" " FROM ts4_quotations AS q" " WHERE q.change_to_prices IS NOT NULL AND q.change_to_prices != 0" " AND (q.change_to_days IS NULL OR q.change_to_days = 0)" f" AND q.status IN {_LM_PENDING_STATUSES}" " ORDER BY q.change_to_prices DESC LIMIT 3" ) # Contractor lookup queries — reuse the CARDS query as a derived table to avoid duplicating conditions _LM_Q_CO_CONTRACTORS = ( "SELECT c.contract_number, c.vendor, c.initial_contract_value" " FROM ts4_contracts c" " WHERE c.contract_number IN (" " SELECT q.contract_number FROM (" + _LM_Q_CO_CARDS + ") AS q" " )" ) _LM_Q_TL_CONTRACTORS = ( "SELECT c.contract_number, c.vendor, c.initial_contract_value" " FROM ts4_contracts c" " WHERE c.contract_number IN (" " SELECT q.contract_number FROM (" + _LM_Q_TL_CARDS + ") AS q" " )" ) _LM_Q_BD_CONTRACTORS = ( "SELECT c.contract_number, c.vendor, c.initial_contract_value" " FROM ts4_contracts c" " WHERE c.contract_number IN (" " SELECT q.contract_number FROM (" + _LM_Q_BD_CARDS + ") AS q" " )" ) def _fmt_date(d: object) -> str: if d is None: return "" if hasattr(d, "day"): return f"{d.day} {d.strftime('%b %Y')}" # type: ignore[union-attr] return str(d) def _fmt_budget_gbp(v: object) -> str | None: if v is None: return None m = float(v) / 1_000_000 return f"£{m:.1f}M" _cfg = AGENT_REGISTRY["chanakya"] # --------------------------------------------------------------------------- # ASGI app # --------------------------------------------------------------------------- app = make_metadata_aware_app( chanakya, name=_cfg.name, description=_cfg.description, ) app.add_middleware( CORSMiddleware, allow_origins=[os.environ.get("CORS_ALLOW_ORIGINS", "*")], allow_methods=["*"], allow_headers=["*"], ) sio = socketio.AsyncServer(async_mode="asgi", cors_allowed_origins="*") # Maps Socket.IO session ID → running agent task so we can cancel on stop/disconnect. _active_tasks: dict[str, asyncio.Task] = {} # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _make_run_context( conversation_id: int | None = None, user_email: str | None = None, persona: str | None = None, progress_queue: asyncio.Queue | None = None, ) -> RequestContext: """Build a ``RequestContext`` pre-loaded with typed runtime channels. Centralises all per-request deps setup so both the sync ``/query`` and the streaming ``/query/stream`` handlers share identical initialisation logic. """ ctx = RequestContext( user_email=user_email or "", role=persona or "PM", conversation_id=conversation_id, ) if progress_queue is not None: ctx.progress_queue = progress_queue return ctx # --------------------------------------------------------------------------- # Flow log — one structured logfire record per question capturing full pipeline # --------------------------------------------------------------------------- def _emit_flow_log( *, question: str, history_count: int, deps: RequestContext, start_time: float, endpoint: str, error: str | None = None, ) -> None: """Emit a single structured logfire record capturing the complete question flow. This log is emitted once per question — after the agent run completes — and contains everything that happened: which tools were called, what tatasteel returned, how many widgets were generated, and total duration. Feed this record to an LLM to diagnose pipeline inconsistencies. """ duration_ms = round((_time.monotonic() - start_time) * 1000) tatasteel_rows = 0 tatasteel_business_summary = "" tatasteel_sql_assumptions = "" tatasteel_sql_query = None tatasteel_rows_sample = None tatasteel_rows_all = None if deps.tatasteel_result: ts = deps.tatasteel_result[0] tatasteel_rows = len(ts.get("rows", [])) tatasteel_business_summary = ts.get("explanation", "") tatasteel_sql_assumptions = ts.get("sql_explanation", "") tatasteel_sql_query = ts.get("generated_sql") tatasteel_rows_sample = ts.get("rows_sample") tatasteel_rows_all = ts.get("rows", []) ui_widget_types: list[str] = [] ui_widgets_sample = None ui_widgets_all = None if deps.ui_result: ui_widget_types = [w.get("type", "unknown") for w in (deps.ui_result.get("data") or [])] ui_widgets_sample = (deps.ui_result.get("data") or [])[:3] ui_widgets_all = deps.ui_result.get("data") or [] # Find market_context and storylines from flow_log market_context = None storylines = None for step in deps.flow_log: if step["step"] == "ask_ui_agent": market_context = step.get("market_context") storylines = step.get("storylines") break tools_called = [step["step"] for step in deps.flow_log] final_viz_type = _extract_viz_type(ui_widgets_all) if ui_widgets_all else "" with logfire.span( "pipeline_summary | {question}", question=question[:200], history_messages=history_count, endpoint=endpoint, user_email=deps.user_email or "(anonymous)", conversation_id=deps.conversation_id, tools_called=tools_called, flow_steps=deps.flow_log, tatasteel_rows=tatasteel_rows, tatasteel_business_summary=tatasteel_business_summary, tatasteel_sql_assumptions=tatasteel_sql_assumptions, tatasteel_sql_query=tatasteel_sql_query, tatasteel_rows_sample=tatasteel_rows_sample, tatasteel_rows_all=tatasteel_rows_all, ui_widget_types=ui_widget_types, ui_widgets_sample=ui_widgets_sample, ui_widgets_all=ui_widgets_all, market_context=market_context, storylines=storylines, final_viz_type=final_viz_type or None, final_response=ui_widgets_all if ui_widgets_all else ("error: " + error if error else "no response"), error=error, total_duration_ms=duration_ms, ): pass # --------------------------------------------------------------------------- # Auth routes # --------------------------------------------------------------------------- app.add_route("/auth/login", auth_login_handler, methods=["POST"]) app.add_route("/auth/verify-token", auth_verify_handler, methods=["POST"]) app.add_route("/auth/me", auth_me_handler, methods=["GET"]) app.add_route("/auth/upsert-user", auth_upsert_handler, methods=["POST"]) # --------------------------------------------------------------------------- # Tatasteel landing metrics endpoint # --------------------------------------------------------------------------- async def _landing_metrics_handler(request: Request) -> JSONResponse: """GET /tatasteel/landing-metrics — returns live aggregate metrics for the landing screen.""" try: pool = await _ensure_db_pool() except Exception as exc: _log.error("landing-metrics: DB connect failed: {}", exc) return JSONResponse({"error": "Database unavailable"}, status_code=503) try: ( contracts, accepted, pending, open_ew, open_nce, co_count, co_cards, tl_count, tl_cards, bd_count, bd_cards, co_contractors, tl_contractors, bd_contractors, ) = await asyncio.gather( pool.fetchrow(_LM_Q_CONTRACTS), pool.fetchrow(_LM_Q_ACCEPTED), pool.fetchrow(_LM_Q_PENDING), pool.fetchrow(_LM_Q_OPEN_EW), pool.fetchrow(_LM_Q_OPEN_NCE), pool.fetchrow(_LM_Q_CO_COUNT), pool.fetch(_LM_Q_CO_CARDS), pool.fetchrow(_LM_Q_TL_COUNT), pool.fetch(_LM_Q_TL_CARDS), pool.fetchrow(_LM_Q_BD_COUNT), pool.fetch(_LM_Q_BD_CARDS), pool.fetch(_LM_Q_CO_CONTRACTORS), pool.fetch(_LM_Q_TL_CONTRACTORS), pool.fetch(_LM_Q_BD_CONTRACTORS), ) except Exception as exc: _log.error("landing-metrics: query failed: {}", exc) return JSONResponse({"error": str(exc)}, status_code=500) def _make_lookup(rows) -> dict[str, dict]: return { r["contract_number"]: { "vendor": r["vendor"], "initial_contract_value": r["initial_contract_value"], } for r in rows } co_lookup = _make_lookup(co_contractors) tl_lookup = _make_lookup(tl_contractors) bd_lookup = _make_lookup(bd_contractors) initial_budget = float(contracts["total_initial_budget"] or 0) return JSONResponse({ "overview": { "initialBudget": initial_budget, "acceptedQuotationRevenue": float(accepted["accepted_revenue"] or 0), "acceptedQuotationDays": int(accepted["accepted_days"] or 0), "pendingQuotationRevenue": float(pending["revenue"] or 0), "pendingQuotationDays": int(pending["accepted_days"] or 0), "activeContracts": int(pending["active_contracts"] or 0), }, "pipeline": { "totalContracts": int(contracts["total_contracts"] or 0), "totalInitialBudget": initial_budget, "openEarlyWarning": int(open_ew["total_ew"] or 0), "openNces": int(open_nce["total_nce"] or 0), "totalPendingQuotations": int(pending["total_quotations"] or 0), "pendingQuotationsRevenue": float(pending["revenue"] or 0), }, "decisions": { "co": { "count": int(co_count["total_quotations"] or 0), "revenueGbp": float(co_count["ce_revenue"] or 0), "days": int(co_count["ce_days"] or 0), "cards": [ { "title": row["title"] or "", "area": row["area"] or "", "dateRaised": _fmt_date(row["quotation_date"]), "timelineDays": int(row["change_to_days"] or 0), "budgetGbp": float(row["change_to_prices"] or 0), "contractorName": co_lookup.get(row["contract_number"], {}).get("vendor"), "contractInitialBudget": _fmt_budget_gbp(co_lookup.get(row["contract_number"], {}).get("initial_contract_value")), } for row in co_cards ], }, "tl": { "count": int(tl_count["total_quotations"] or 0), "revenueGbp": 0, "days": int(tl_count["ce_days"] or 0), "cards": [ { "title": row["title"] or "", "area": row["area"] or "", "dateRaised": _fmt_date(row["quotation_date"]), "timelineDays": int(row["change_to_days"] or 0), "contractorName": tl_lookup.get(row["contract_number"], {}).get("vendor"), "contractInitialBudget": _fmt_budget_gbp(tl_lookup.get(row["contract_number"], {}).get("initial_contract_value")), } for row in tl_cards ], }, "bd": { "count": int(bd_count["total_quotations"] or 0), "revenueGbp": float(bd_count["ce_revenue"] or 0), "days": 0, "cards": [ { "title": row["title"] or "", "area": row["area"] or "", "dateRaised": _fmt_date(row["quotation_date"]), "budgetGbp": float(row["change_to_prices"] or 0), "contractorName": bd_lookup.get(row["contract_number"], {}).get("vendor"), "contractInitialBudget": _fmt_budget_gbp(bd_lookup.get(row["contract_number"], {}).get("initial_contract_value")), } for row in bd_cards ], }, }, }) app.add_route("/tatasteel/landing-metrics", _landing_metrics_handler, methods=["GET"]) # --------------------------------------------------------------------------- # Conversation history routes # --------------------------------------------------------------------------- app.add_route("/conversations", create_conversation_handler, methods=["POST"]) app.add_route("/conversations", get_conversations_handler, methods=["GET"]) app.add_route("/conversations/{conversation_id:int}", get_conversation_handler, methods=["GET"]) app.add_route("/conversations/{conversation_id:int}/title", update_conversation_title_handler, methods=["PUT"]) app.add_route("/conversations/{conversation_id:int}", delete_conversation_handler, methods=["DELETE"]) app.add_route("/conversations/{conversation_id:int}/messages", save_message_handler, methods=["POST"]) app.add_route("/conversations/{conversation_id:int}/messages", get_messages_handler, methods=["GET"]) app.add_route("/conversations/{conversation_id:int}/messages/{message_id:int}/bookmark", toggle_bookmark_handler, methods=["PATCH"]) app.add_route("/users/me/bookmarks", get_user_bookmarks_handler, methods=["GET"]) # --------------------------------------------------------------------------- # Startup / shutdown # --------------------------------------------------------------------------- @app.on_event("startup") async def _startup_banner() -> None: """Initialize auth database and display startup banner.""" port = os.environ.get("CHANAKYA_PORT", "8010") try: _log.info("Initializing authentication database...") await init_auth_db() _log.info("✅ Auth database initialized successfully") except Exception as e: _log.error("❌ Failed to initialize auth database: {}", e) _log.warning("Authentication endpoints will not work until database is configured!") _log.info("\n Chanakya ready:") _log.info(" A2A → http://0.0.0.0:{}/", port) _log.info(" Chat → http://0.0.0.0:{}/chat", port) _log.info(" Query → http://0.0.0.0:{}/query", port) _log.info(" Stream→ http://0.0.0.0:{}/query/stream", port) _log.info(" SIO → http://0.0.0.0:{}/socket.io/ (Socket.IO)", port) _log.info(" Card → http://0.0.0.0:{}/.well-known/agent-card.json", port) _log.info(" Auth → http://0.0.0.0:{}/auth/login (POST)", port) @app.on_event("shutdown") async def _shutdown() -> None: """Close auth database connection pool on shutdown.""" await close_auth_db() _log.info("Auth database connection pool closed") # --------------------------------------------------------------------------- # Pylogue chat UI # --------------------------------------------------------------------------- _PROJECT_ROOT = Path(__file__).parents[3] _chat_db_path = os.environ.get( "PYLOGUE_DB_PATH", str(_PROJECT_ROOT / "data" / "conversation-histories" / "chanakya" / "chat_history.db"), ) Path(_chat_db_path).parent.mkdir(parents=True, exist_ok=True) app.mount( "/chat", create_core_app( responder_factory=lambda: chanakya, hero_title="Chanakya", hero_subtitle="Ask questions about Tata Steel's construction project data.", db_path=_chat_db_path, ), ) # --------------------------------------------------------------------------- # Dev impersonation admin UI # --------------------------------------------------------------------------- def _admin_html(message: str = "") -> str: active = _IMPERSONATION_FILE.read_text().strip() if _IMPERSONATION_FILE.exists() else "" options = [ Option( f'{u["name"]} ({u["role"]})', value=u["email"], selected=u["email"] == active, ) for u in USERS ] page = Html( Body( H2("Enterprise Brain — Dev Impersonation"), P(Strong(message), style="color:green") if message else "", P(f"Active: {active or '(none)'}"), FHForm( Select(*options, name="email", style="font-size:1rem;padding:.3rem"), Button("Set", type="submit", style="margin-left:.5rem;padding:.3rem .8rem"), method="post", ), FHForm( Input(type="hidden", name="email", value=""), Button("Clear", type="submit", style="margin-top:.5rem;padding:.3rem .8rem"), method="post", ), style="font-family:monospace;padding:2rem", ) ) return to_xml(page) _admin_app = FastAPI() @_admin_app.get("/", response_class=HTMLResponse) async def admin_get() -> HTMLResponse: return HTMLResponse(_admin_html()) @_admin_app.post("/", response_class=HTMLResponse) async def admin_post(email: str = Form(default="")) -> HTMLResponse: if email: _IMPERSONATION_FILE.write_text(email) msg = f"Impersonating {email} — all agents will use this identity" else: _IMPERSONATION_FILE.unlink(missing_ok=True) msg = "Cleared — agents will use real RequestContext.user_email" _log.info("Dev impersonation set to: {!r}", email) return HTMLResponse(_admin_html(msg)) app.mount("/admin", _admin_app) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _extract_viz_type(widgets: list[dict]) -> str: """Return the chart type from the first visualization widget, or empty string.""" for w in widgets: if w.get("type") == "visualization": return w.get("content", {}).get("type", "") return "" def _stamp(question: str) -> str: """Prepend current date/time to the user question so the system prompt stays static.""" now = datetime.now().strftime("%d %B %Y, %H:%M") return f"[Today: {now}]\n{question}" # --------------------------------------------------------------------------- # REST endpoints # --------------------------------------------------------------------------- async def _query_handler(request: Request) -> JSONResponse: """POST /query — run Chanakya; forwards to UI agent if LLM called ask_ui_agent.""" try: body = await request.json() question: str = body["question"] user_email: str | None = body.get("user_email") conversation_id: int | None = body.get("conversation_id") persona: str | None = body.get("persona") except Exception as exc: return JSONResponse( AgentQueryResponse.failure(message=f"Invalid request: {exc}", code="PIPELINE_ERROR").model_dump(), status_code=400, ) # Load history FIRST so the current question is not included in context. message_history = await load_conversation_history(conversation_id) # Then persist the user turn — single source of truth for user messages. # Saved AFTER load so 'question' does not appear in the LLM history for this turn. if conversation_id: try: await save_message( conversation_id=conversation_id, message_author="user", message=question, message_type="request", ) except Exception as exc: _log.error("[QUERY] Failed to save user message: {}", exc) append_to_history_cache(conversation_id, "user", question) deps_obj = _make_run_context(conversation_id=conversation_id, user_email=user_email, persona=persona) deps_obj.message_history_wire = serialize_history(message_history) _flow_start = _time.monotonic() try: with logfire.span("{q}", q=question[:120]): result = await chanakya.agent.run(_stamp(question), message_history=message_history, deps=deps_obj) except Exception as exc: _log.error("chanakya /query agent error: {}", exc) # Emit OUTSIDE the agent span — appears as a top-level row in logfire. _emit_flow_log( question=question, history_count=len(message_history), deps=deps_obj, start_time=_flow_start, endpoint="/query", error=str(exc), ) return JSONResponse( AgentQueryResponse.failure(message=str(exc), code="PIPELINE_ERROR").model_dump(), status_code=500, ) # Emit OUTSIDE the agent span — appears as a top-level row in logfire. _emit_flow_log( question=question, history_count=len(message_history), deps=deps_obj, start_time=_flow_start, endpoint="/query", ) ui_result = deps_obj.ui_result if ui_result: # ui_result is the done event from ask_ui_agent (sync handler has no progress_queue) widgets: list[dict] = ui_result.get("data") or [] text_content = next( (w["content"] for w in widgets if w.get("type") == "text"), "" ) viz_type = _extract_viz_type(widgets) message_text = text_content or "(no text response)" if viz_type: message_text += f"\n[Visualisation shown: visualization({viz_type})]" ts_cache: dict = {} if deps_obj.tatasteel_result: ts = deps_obj.tatasteel_result[0] ts_cache = { "rows": ts.get("rows", []), "explanation": ts.get("explanation", ""), "question": ts.get("question", ""), } if conversation_id: await save_message( conversation_id=conversation_id, message_author="assistant", message=message_text, message_type="response", metadata={ "status": ui_result.get("status", "success"), "data": widgets, "error": ui_result.get("error"), "_tatasteel_cache": ts_cache, }, ) append_to_history_cache(conversation_id, "assistant", message_text, tatasteel_cache=ts_cache) return JSONResponse(content={"status": ui_result.get("status", "success"), "data": widgets, "error": ui_result.get("error")}) text_widget = {"type": "text", "content": result.output} response = AgentQueryResponse.success(widgets=[text_widget]) if conversation_id: await save_message( conversation_id=conversation_id, message_author="assistant", message=result.output, message_type="response", metadata=response.model_dump(), ) append_to_history_cache(conversation_id, "assistant", result.output) return JSONResponse(response.model_dump()) app.add_route("/query", _query_handler, methods=["POST"]) async def _stream_query_handler(request: Request) -> StreamingResponse | JSONResponse: """POST /query/stream — Chanakya runs; if tatasteel was consulted, streams via UI agent.""" try: body = await request.json() question: str = body["question"] user_email: str | None = body.get("user_email") conversation_id: int | None = body.get("conversation_id") persona: str | None = body.get("persona") except Exception as exc: return JSONResponse( AgentQueryResponse.failure(message=f"Invalid request: {exc}", code="PIPELINE_ERROR").model_dump(), status_code=400, ) async def sse_events(): yield f"data: {json.dumps({'type': 'progress', 'step': 'thinking', 'message': 'Processing your question...'})}\n\n" # Load history FIRST so the current question is not in the LLM context for this turn. message_history = await load_conversation_history(conversation_id) _log.debug("[QUERY/STREAM] Loaded {} history messages for conv={}", len(message_history), conversation_id) # Then persist current user turn — backend is the single source of truth. if conversation_id: try: await save_message( conversation_id=conversation_id, message_author="user", message=question, message_type="request", ) except Exception as exc: _log.error("[QUERY/STREAM] Failed to save user message: {}", exc) append_to_history_cache(conversation_id, "user", question) progress_queue: asyncio.Queue = asyncio.Queue() deps = _make_run_context( conversation_id=conversation_id, user_email=user_email, persona=persona, progress_queue=progress_queue, ) deps.message_history_wire = serialize_history(message_history) # Run agent as background task so we can drain progress events concurrently. # CancelledError is caught inside the span so logfire doesn't record it as an # error — client disconnect is expected, not a bug. _flow_start = _time.monotonic() async def _run_with_span() -> Any: _cancelled = False with logfire.span("{q}", q=question[:120]): try: return await chanakya.agent.run(_stamp(question), message_history=message_history, deps=deps) except asyncio.CancelledError: _cancelled = True if _cancelled: raise asyncio.CancelledError() agent_task: asyncio.Task = asyncio.create_task(_run_with_span()) # Drain progress events while agent works (polling every 100 ms). # NOTE: do NOT use asyncio.shield here — it creates orphaned queue waiters that # silently consume events without yielding them. try: while not agent_task.done(): try: event = await asyncio.wait_for(progress_queue.get(), timeout=0.1) yield f"data: {json.dumps(event)}\n\n" except asyncio.TimeoutError: pass except (asyncio.CancelledError, GeneratorExit): agent_task.cancel() with contextlib.suppress(asyncio.CancelledError, Exception): await agent_task raise # Drain any events queued during the final agent step while not progress_queue.empty(): yield f"data: {json.dumps(progress_queue.get_nowait())}\n\n" try: result = agent_task.result() except Exception as exc: _log.error("[QUERY/STREAM] agent error: {}", exc) # Emit OUTSIDE the agent span — appears as a top-level row in logfire. _emit_flow_log( question=question, history_count=len(message_history), deps=deps, start_time=_flow_start, endpoint="/query/stream", error=str(exc), ) yield f"data: {json.dumps({'type': 'error', 'message': str(exc), 'code': 'PIPELINE_ERROR'})}\n\n" return # Emit OUTSIDE the agent span — appears as a top-level row in logfire. _emit_flow_log( question=question, history_count=len(message_history), deps=deps, start_time=_flow_start, endpoint="/query/stream", ) # Events were already forwarded live to frontend via progress_queue drain loop. # We only need ui_result (done event) here for DB save. done_event = deps.ui_result if done_event: _log.debug("[QUERY/STREAM] UI agent done, saving to DB") if conversation_id: widgets_data: list[dict] = done_event.get("data") or [] text_content = next( (w["content"] for w in widgets_data if w.get("type") == "text"), "" ) viz_type = _extract_viz_type(widgets_data) message_text = text_content or "(no text response)" if viz_type: message_text += f"\n[Visualisation shown: visualization({viz_type})]" try: ts_cache: dict = {} if deps.tatasteel_result: ts = deps.tatasteel_result[0] ts_cache = { "rows": ts.get("rows", []), "explanation": ts.get("explanation", ""), "question": ts.get("question", ""), } saved = await save_message( conversation_id=conversation_id, message_author="assistant", message=message_text, message_type="response", metadata={ "status": done_event.get("status", "success"), "data": widgets_data, "error": done_event.get("error"), "_tatasteel_cache": ts_cache, }, ) _log.info("[QUERY/STREAM] ✅ Saved assistant message for conv={}", conversation_id) append_to_history_cache(conversation_id, "assistant", message_text, tatasteel_cache=ts_cache) yield f"data: {json.dumps({'type': 'message_saved', 'message_id': saved['id']})}\n\n" except Exception as exc: _log.error("[QUERY/STREAM] ❌ Failed to save assistant message: {}", exc) return # General / conversational — stream chanakya's text output word by word _log.debug("[QUERY/STREAM] Direct text response") text = result.output for word in text.split(" "): yield f"data: {json.dumps({'type': 'text_chunk', 'content': word + ' '})}\n\n" text_widget = {"type": "text", "content": text} done_payload = {"type": "done", **AgentQueryResponse.success(widgets=[text_widget]).model_dump()} # Save direct text response to DB before emitting done. saved_msg_id: int | None = None if conversation_id: try: saved = await save_message( conversation_id=conversation_id, message_author="assistant", message=text, message_type="response", metadata={ "status": "success", "data": [text_widget], "error": None, }, ) saved_msg_id = saved["id"] _log.info("[QUERY/STREAM] ✅ Saved direct text message for conv={}", conversation_id) append_to_history_cache(conversation_id, "assistant", text) except Exception as exc: _log.error("[QUERY/STREAM] ❌ Failed to save direct text message: {}", exc) if saved_msg_id is not None: done_payload["message_id"] = saved_msg_id yield f"data: {json.dumps(done_payload)}\n\n" return StreamingResponse( sse_events(), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, ) app.add_route("/query/stream", _stream_query_handler, methods=["POST"]) # --------------------------------------------------------------------------- # Socket.IO server — handles real-time query streaming # --------------------------------------------------------------------------- async def _run_sio_query( sid: str, question: str, user_email: str | None, conversation_id: int | None, ) -> None: """Execute one query turn over Socket.IO, emitting the same event types as the SSE stream (progress, text_chunk, widget, done, error). """ await sio.emit("progress", {"type": "progress", "step": "thinking", "message": "Processing your question..."}, to=sid) # Load history FIRST so the current question is not included in context. message_history = await load_conversation_history(conversation_id) _log.debug("[SIO] Loaded {} history messages for conv={}", len(message_history), conversation_id) # Persist user turn — backend is the single source of truth. if conversation_id: try: await save_message( conversation_id=conversation_id, message_author="user", message=question, message_type="request", ) except Exception as exc: _log.error("[SIO] Failed to save user message: {}", exc) append_to_history_cache(conversation_id, "user", question) progress_queue: asyncio.Queue = asyncio.Queue() deps = _make_run_context( conversation_id=conversation_id, user_email=user_email, progress_queue=progress_queue, ) deps.message_history_wire = serialize_history(message_history) _flow_start = _time.monotonic() async def _run_with_span() -> Any: _cancelled = False with logfire.span("{q}", q=question[:120]): try: return await chanakya.agent.run(_stamp(question), message_history=message_history, deps=deps) except asyncio.CancelledError: _cancelled = True if _cancelled: raise asyncio.CancelledError() agent_task: asyncio.Task = asyncio.create_task(_run_with_span()) _active_tasks[sid] = agent_task try: # Drain progress events while agent works (polling every 100 ms). while not agent_task.done(): try: event = await asyncio.wait_for(progress_queue.get(), timeout=0.1) await sio.emit(event["type"], event, to=sid) except asyncio.TimeoutError: pass # Drain any events queued during the final agent step. while not progress_queue.empty(): event = progress_queue.get_nowait() await sio.emit(event["type"], event, to=sid) except asyncio.CancelledError: agent_task.cancel() with contextlib.suppress(asyncio.CancelledError, Exception): await agent_task _active_tasks.pop(sid, None) raise finally: _active_tasks.pop(sid, None) # Emit OUTSIDE the agent span — appears as a top-level row in logfire. _emit_flow_log( question=question, history_count=len(message_history), deps=deps, start_time=_flow_start, endpoint="/sio/query", ) # Check for UI-agent result first — this is the primary path and must be # evaluated before calling agent_task.result(), because an agent that sets # ui_result and then raises will cause agent_task.result() to throw, which # previously returned early and skipped the DB save entirely. done_event = deps.ui_result if done_event: _log.debug("[SIO] UI agent done, saving to DB") if conversation_id: widgets_data: list[dict] = done_event.get("data") or [] text_content = next( (w["content"] for w in widgets_data if w.get("type") == "text"), "" ) viz_type = _extract_viz_type(widgets_data) message_text = text_content or "(no text response)" if viz_type: message_text += f"\n[Visualisation shown: visualization({viz_type})]" try: ts_cache: dict = {} if deps.tatasteel_result: ts = deps.tatasteel_result[0] ts_cache = { "rows": ts.get("rows", []), "explanation": ts.get("explanation", ""), "question": ts.get("question", ""), } saved = await save_message( conversation_id=conversation_id, message_author="assistant", message=message_text, message_type="response", metadata={ "status": done_event.get("status", "success"), "data": widgets_data, "error": done_event.get("error"), "_tatasteel_cache": ts_cache, }, ) _log.info("[SIO] Saved assistant message for conv={}", conversation_id) append_to_history_cache(conversation_id, "assistant", message_text, tatasteel_cache=ts_cache) await sio.emit("message_saved", {"type": "message_saved", "message_id": saved["id"]}, to=sid) except Exception as exc: _log.error("[SIO] Failed to save assistant message: {}", exc) return # General / conversational — retrieve agent result (raises if agent errored). if agent_task.cancelled(): _log.info("[SIO] Agent task cancelled for sid={}", sid) return try: result = agent_task.result() except Exception as exc: _log.error("[SIO] agent error: {}", exc) await sio.emit( "error", {"type": "error", "message": str(exc), "code": "PIPELINE_ERROR"}, to=sid, ) return # Stream text word by word. _log.debug("[SIO] Direct text response") text = result.output for word in text.split(" "): await sio.emit("text_chunk", {"type": "text_chunk", "content": word + " "}, to=sid) text_widget = {"type": "text", "content": text} done_payload = {"type": "done", **AgentQueryResponse.success(widgets=[text_widget]).model_dump()} saved_msg_id: int | None = None if conversation_id: try: saved = await save_message( conversation_id=conversation_id, message_author="assistant", message=text, message_type="response", metadata={"status": "success", "data": [text_widget], "error": None}, ) saved_msg_id = saved["id"] _log.info("[SIO] Saved direct text message for conv={}", conversation_id) append_to_history_cache(conversation_id, "assistant", text) except Exception as exc: _log.error("[SIO] Failed to save direct text message: {}", exc) if saved_msg_id is not None: done_payload["message_id"] = saved_msg_id await sio.emit("done", done_payload, to=sid) @sio.on("connect") async def _sio_connect(sid: str, environ: dict, auth: dict | None = None) -> None: """Resolve user identity from JWT auth token passed in Socket.IO handshake auth.""" token: str | None = (auth or {}).get("token") user_email: str | None = get_username_from_token(token) if token else None await sio.save_session(sid, {"user_email": user_email}) _log.info("[SIO] Client connected sid={}, user={}", sid, user_email or "(unauthenticated)") @sio.on("disconnect") async def _sio_disconnect(sid: str) -> None: task = _active_tasks.pop(sid, None) if task and not task.done(): task.cancel() _log.info("[SIO] Cancelled agent task on disconnect sid={}", sid) _log.info("[SIO] Client disconnected sid={}", sid) @sio.on("stop") async def _sio_on_stop(sid: str) -> None: task = _active_tasks.get(sid) if task and not task.done(): task.cancel() _log.info("[SIO] Stop requested — cancelled agent task sid={}", sid) @sio.on("query") async def _sio_on_query(sid: str, data: dict) -> None: """Handle 'query' event — execute one question turn for the connected client.""" session = await sio.get_session(sid) user_email: str | None = session.get("user_email") if session else None question: str = (data.get("question") or "").strip() conversation_id: int | None = data.get("conversation_id") if not question: await sio.emit("error", {"type": "error", "message": "Missing question", "code": "INVALID_REQUEST"}, to=sid) return await _run_sio_query(sid, question, user_email, conversation_id) await sio.disconnect(sid) # Wrap the FastAPI app with the Socket.IO ASGI app. # Socket.IO requests (/socket.io/*) are handled by `sio`; # all other HTTP requests are forwarded to the inner FastAPI app. app = socketio.ASGIApp(sio, app) if __name__ == "__main__": import uvicorn uvicorn.run("backend.chanakya.rt_agent.server:app", host="0.0.0.0", port=_cfg.port, reload=True)