"""Infinithesim Postgres data agent — A2A server using pydantic-ai. Connects to the Infinithesim database using credentials from environment variables. The agent reads a schema definition to understand the DB structure, generates accurate SQL queries, and executes them safely (SELECT-only) to answer user questions. Required env vars: INFINITHESIM_DB_HOST — database host (default: localhost) INFINITHESIM_DB_PORT — database port (default: 5432) INFINITHESIM_DB_NAME — database name INFINITHESIM_DB_USER — database user INFINITHESIM_DB_PASSWORD — database password INFINITHESIM_SCHEMA_PATH — path to the schema file (default: backend/chanakya/infinithesim/schema.sql) """ import asyncio import json import os from pathlib import Path from loguru import logger as _log from starlette.middleware.cors import CORSMiddleware from starlette.requests import Request from starlette.responses import JSONResponse, StreamingResponse from pydantic_ai.messages import ModelMessage, ModelRequest, ModelResponse, UserPromptPart, TextPart from backend.chanakya.a2a_context import make_metadata_aware_app from backend.chanakya.base import EnterpriseAgent from backend.chanakya.config import AGENT_REGISTRY from backend.chanakya.nl2sql_core import NL2SQLDeps, register_nl2sql_tools from backend.chanakya.infinithesim.prompt import make_infinithesim_instructions from backend.chanakya.schemas import ( AgentQueryRequest, AgentQueryResponse, ) def _reconstruct_history(history: list[dict] | None) -> list[ModelMessage]: """Convert serialized [{role, content}] dicts from Chanakya hub into pydantic-ai ModelMessages.""" if not history: return [] msgs: list[ModelMessage] = [] for msg in history: role = msg.get("role", "") content = msg.get("content", "") if role == "user": msgs.append(ModelRequest(parts=[UserPromptPart(content=content)])) elif role in ("assistant", "agent"): msgs.append(ModelResponse(parts=[TextPart(content=content)])) return msgs from pylogue.shell import app_factory as create_core_app _cfg = AGENT_REGISTRY["infinithesim"] _PROJECT_ROOT = Path(__file__).parents[3] _DEFAULT_SCHEMA_PATH = Path(__file__).parent / "schema.sql" # --------------------------------------------------------------------------- # Agent # --------------------------------------------------------------------------- infinithesim_agent = EnterpriseAgent( deps_type=NL2SQLDeps, api_base=os.environ["LITELLM_PROVIDER_BASE_URL"], instructions=make_infinithesim_instructions(), logfire_env="chanakya-alpha", service_name="infinithesim", ) register_nl2sql_tools( infinithesim_agent, schema_path_env="INFINITHESIM_SCHEMA_PATH", db_env_prefix="INFINITHESIM_DB", default_schema_path=_DEFAULT_SCHEMA_PATH, auto_limit=False, ) # --------------------------------------------------------------------------- # ASGI app # --------------------------------------------------------------------------- app = make_metadata_aware_app( infinithesim_agent, name=_cfg.name, description=_cfg.description, ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) # --------------------------------------------------------------------------- # REST endpoint — structured JSON # --------------------------------------------------------------------------- async def _query_handler(request: Request) -> JSONResponse: """POST /query — returns a widget list for frontend rendering. The agent calls the ``format_response`` tool during its run, which populates ``deps.formatted_widgets``. The handler returns those widgets directly. The ``/chat`` endpoint remains available for prose output. """ try: body = await request.json() req = AgentQueryRequest(**body) except Exception as exc: return JSONResponse( AgentQueryResponse.failure(message=f"Invalid request: {exc}", code="PIPELINE_ERROR").model_dump(), status_code=400, ) deps = NL2SQLDeps(user_email=req.user_email or "", widget_mode=True) message_history = _reconstruct_history(req.history) _log.debug("infinithesim /query: using {} history messages", len(message_history)) try: result = await infinithesim_agent.agent.run(req.question, deps=deps, message_history=message_history) except Exception as exc: _log.error("infinithesim query endpoint error: {}", exc) return JSONResponse( AgentQueryResponse.failure(message=str(exc), code="PIPELINE_ERROR").model_dump() ) # Primary path: agent called format_response tool — prepend explanation as text widget if deps.formatted_widgets: text_widget = {"type": "text", "content": result.output} widgets = [text_widget] + deps.formatted_widgets _log.debug("infinithesim query endpoint widgets={} rows={}", len(widgets), len(deps.query_result or [])) return JSONResponse(AgentQueryResponse.success(widgets=widgets).model_dump()) # Fallback: agent didn't call format_response _log.warning("infinithesim query endpoint: format_response was not called — building fallback widgets") fallback: list[dict] = [{"type": "text", "content": result.output}] if deps.query_result: fallback.append({ "type": "table", "content": { "columns": list(deps.query_result[0].keys()), "rows": deps.query_result, }, }) return JSONResponse(AgentQueryResponse.success(widgets=fallback).model_dump()) app.add_route("/query", _query_handler, methods=["POST"]) async def _stream_query_handler(request: Request) -> StreamingResponse | JSONResponse: """POST /query/stream — SSE streaming via pydantic-ai run_stream(). The agent runs its tools silently (get_schema → run_query → format_response), then writes the explanation as its final text response which is streamed token-by-token via stream_text(delta=True). SSE events (server → client, in order): data: {"type": "progress", "step": "schema"|"query"|"formatting", "message": str} data: {"type": "widget_loading", "widget_type": "chart"|"table"} data: {"type": "text_chunk", "content": str} ← real LLM tokens data: {"type": "widget", "widget": {...}} data: {"type": "done", "success": true, "widgets": [text, data]} data: {"type": "error", "message": str, "code": str} """ try: body = await request.json() req = AgentQueryRequest(**body) except Exception as exc: return JSONResponse( AgentQueryResponse.failure(message=f"Invalid request: {exc}", code="PIPELINE_ERROR").model_dump(), status_code=400, ) progress_queue: asyncio.Queue = asyncio.Queue() deps = NL2SQLDeps(user_email=req.user_email or "", widget_mode=True, progress_queue=progress_queue) message_history = _reconstruct_history(req.history) _log.debug("infinithesim /query/stream: using {} history messages", len(message_history)) async def sse_events(): full_text: list[str] = [] merged: asyncio.Queue = asyncio.Queue() agent_done = asyncio.Event() had_error = False async def _drain_progress(): """Forward progress events from tools into merged queue as they arrive.""" while not agent_done.is_set(): try: event = await asyncio.wait_for(progress_queue.get(), timeout=0.05) await merged.put(event) except asyncio.TimeoutError: continue # Drain any remaining events after agent finishes while not progress_queue.empty(): await merged.put(progress_queue.get_nowait()) await merged.put(None) # sentinel async def _run_agent(): nonlocal had_error try: async with infinithesim_agent.agent.run_stream(req.question, deps=deps, message_history=message_history) as stream: async for chunk in stream.stream_text(delta=True): # LiteLLM may buffer the entire response into one chunk. # Split on spaces and emit word-by-word for a typing effect. if len(chunk) > 20: words = chunk.split(" ") for i, word in enumerate(words): token = word + (" " if i < len(words) - 1 else "") full_text.append(token) await merged.put({"type": "text_chunk", "content": token}) await asyncio.sleep(0.03) else: full_text.append(chunk) await merged.put({"type": "text_chunk", "content": chunk}) except Exception as exc: _log.error("infinithesim stream query error: {}", exc) had_error = True await merged.put({"type": "error", "message": str(exc), "code": "PIPELINE_ERROR"}) finally: agent_done.set() drain_task = asyncio.create_task(_drain_progress()) agent_task = asyncio.create_task(_run_agent()) while True: item = await merged.get() if item is None: break yield f"data: {json.dumps(item)}\n\n" await asyncio.gather(drain_task, agent_task, return_exceptions=True) if had_error: return text_widget = {"type": "text", "content": "".join(full_text)} all_widgets = [text_widget] + deps.formatted_widgets for w in deps.formatted_widgets: yield f"data: {json.dumps({'type': 'widget', 'widget': w})}\n\n" yield f"data: {json.dumps({'type': 'done', **AgentQueryResponse.success(widgets=all_widgets).model_dump()})}\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"]) # --------------------------------------------------------------------------- # Startup banner + chat UI # --------------------------------------------------------------------------- @app.on_event("startup") async def _startup_banner() -> None: port = os.environ.get("INFINITHESIM_AGENT_PORT", "8006") db = os.environ.get("INFINITHESIM_DB_NAME", "(not set)") host = os.environ.get("INFINITHESIM_DB_HOST", "localhost") _log.info("\n Infinithesim Agent 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(" DB → postgres://{}:{}/{}", host, os.environ.get("INFINITHESIM_DB_PORT", "5432"), db) _log.info(" Schema→ {}", _DEFAULT_SCHEMA_PATH) _INFINITHESIM_DB_PATH = os.environ.get( "PYLOGUE_DB_PATH", str(_PROJECT_ROOT / "data" / "conversation-histories" / "infinithesim" / "chat_history.db"), ) Path(_INFINITHESIM_DB_PATH).parent.mkdir(parents=True, exist_ok=True) app.mount( "/chat", create_core_app( responder_factory=lambda: infinithesim_agent, hero_title="Infinithesim Agent", hero_subtitle="Query the Infinithesim database with natural language.", db_path=_INFINITHESIM_DB_PATH, ), name="infinithesim-chat", )