"""UI Agent — visualisation renderer for Enterprise Brain. Receives raw query rows and a plain-text explanation from a data agent (e.g. Tatasteel) and turns them into a frontend-ready widget stream. Endpoints: POST /query — sync: returns widget list JSON POST /query/stream — SSE: streams text tokens then emits the widget """ from __future__ import annotations import asyncio import json import os import re from dataclasses import dataclass, field from pathlib import Path from typing import Any import pandas as pd import logfire # --------------------------------------------------------------------------- # Utilities # --------------------------------------------------------------------------- _MERMAID_RE = re.compile(r'```mermaid\s*\n(.*?)\n```', re.DOTALL | re.IGNORECASE) _BR_TAG_RE = re.compile(r'', re.IGNORECASE) def _split_explanation(text: str) -> tuple[str, list[str]]: """Strip
tags and extract mermaid code blocks from *text*. Returns: (cleaned_text, list_of_mermaid_diagrams) Each extracted diagram is emitted as a ``{"type": "mermaid"}`` widget so the frontend can render it natively instead of parsing raw markdown. """ diagrams: list[str] = [] def _pull(m: re.Match) -> str: # type: ignore[type-arg] diagrams.append(m.group(1).strip()) return "" cleaned = _MERMAID_RE.sub(_pull, text) cleaned = _BR_TAG_RE.sub("", cleaned).strip() return cleaned, diagrams from loguru import logger as _log from pydantic_ai import RunContext from starlette.middleware.cors import CORSMiddleware from starlette.requests import Request from starlette.responses import JSONResponse, StreamingResponse 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.conversation._history import deserialize_history, load_conversation_history from backend.chanakya.nl2sql_core.chart_config import ( VALID_CHART_KEYS, validate_chart_content, validate_highlight_content, ) from backend.chanakya.schemas import AgentQueryResponse, UIRenderRequest from backend.chanakya.ui_agent.prompt import make_ui_agent_instructions _cfg = AGENT_REGISTRY["ui_agent"] # --------------------------------------------------------------------------- # Deps # --------------------------------------------------------------------------- @dataclass class UIAgentDeps: """Per-request context for the UI agent.""" rows: list[dict[str, Any]] = field(default_factory=list) question: str = "" explanation: str = "" # Market enrichment from Chanakya's web_search step. market_context: str = "" # Suggested follow-up enquiry lines from Chanakya's intent analysis. storyline_hints: list[str] = field(default_factory=list) # Additional supporting datasets gathered for action/decision analysis. decision_support_context: list[dict[str, Any]] = field(default_factory=list) # Decision Context mapping explanation passed through from the routing agent. dc_explanation: str = "" # Ordered list of all widgets produced this turn. # Both tools append here so visualization → highlights → visualization # ordering is preserved exactly as the LLM called the tools. widgets: list[dict[str, Any]] = field(default_factory=list) # Slot index reserved by the first (failed) build_visualization call so that # a successful retry fills the original position instead of appending at the end. pending_visualization_slot: int | None = field(default=None) progress_queue: asyncio.Queue | None = field(default=None) currency_code: str = "GBP" # --------------------------------------------------------------------------- # Agent # --------------------------------------------------------------------------- ui_agent = EnterpriseAgent( deps_type=UIAgentDeps, api_base=os.environ["LITELLM_PROVIDER_BASE_URL"], instructions=make_ui_agent_instructions(), logfire_env="chanakya-alpha", service_name="ui-agent", retries=3, ) @ui_agent.tool async def build_visualization( ctx: RunContext[UIAgentDeps], chart_type: str, content: dict[str, Any], ) -> str: """Store a frontend-ready visualisation widget. The agent is responsible for constructing ``content`` according to the exact ``BaseVisualizationConfig`` schema for the chosen ``chart_type`` (schemas are listed in the system prompt). This tool simply validates the type key and records the widget so the HTTP handler can return it. Args: chart_type: One of the supported chart keys (e.g. ``contract-bars``, ``status-arc``, ``variation-split``, …). Used only for the ``widget_loading`` progress event; the ``"type"`` field inside ``content`` is the authoritative discriminant. content: Fully-formed ``BaseVisualizationConfig`` dict — constructed by the agent by mapping SQL row columns to the schema fields shown in the system prompt. """ ct = chart_type.strip().lower() if ct not in VALID_CHART_KEYS: ct = "multi-metric-constellation-chart" q = ctx.deps.progress_queue if q is not None: await q.put({"type": "widget_loading", "widget_type": ct}) if not content: # Agent signalled null content (no SQL rows) — still emit a widget so the # frontend always has something to render rather than silently showing nothing. ctx.deps.widgets.append({ "type": "visualization", "content": {"type": ct, "_empty": True}, }) _log.debug("build_visualization empty content, chart_type={}", ct) return "No data to visualise." # Guard: the LLM sometimes wraps the real payload in an extra {"content": {...}} # layer, e.g. {"type": "ew-category", "content": {"categories": [...]}}. # Unwrap when the dict has only "type" + "content" and the inner value is a dict. if ( isinstance(content.get("content"), dict) and set(content.keys()) <= {"type", "content"} ): inner = content["content"] content = {"type": content.get("type", ct), **inner} _log.warning("build_visualization: unwrapped double-nested content for chart_type={}", ct) # Guarantee the "type" discriminant is present and correct. # The LLM sometimes omits it from content even though it passes chart_type # correctly as a separate argument. Use the validated `ct` as source of truth. if content.get("type") != ct: _log.warning( "build_visualization: content missing/wrong type ({!r}), injecting {!r}", content.get("type"), ct, ) content = {"type": ct, **{k: v for k, v in content.items() if k != "type"}} # Reserve a slot on first attempt so a retry fills the original position # rather than appending after widgets added between the two calls. if ctx.deps.pending_visualization_slot is None: ctx.deps.widgets.append({"type": "visualization", "_pending": True}) ctx.deps.pending_visualization_slot = len(ctx.deps.widgets) - 1 slot = ctx.deps.pending_visualization_slot try: validated = validate_chart_content(content) content = validated.model_dump() except Exception as exc: # Slot stays reserved; LLM will retry and fill it. return ( f"Schema validation failed: {exc}\n" "Only the fields listed above are wrong — keep all other fields exactly as they were. " "Fix only the failing fields and call build_visualization again with the corrected content." ) ctx.deps.widgets[slot] = {"type": "visualization", "content": content} ctx.deps.pending_visualization_slot = None _log.debug("build_visualization chart_type={} slot={} total_widgets={}", ct, slot, len(ctx.deps.widgets)) return "Widget built." @ui_agent.tool async def build_text( ctx: RunContext[UIAgentDeps], content: str, ) -> str: """Insert a narrative text block at this exact position in the widget sequence. Call this instead of writing natural response text so that narrative appears interleaved with visualisations and highlights — not forced to the top. Args: content: Plain-English narrative for this position in the story. """ if not content or not content.strip(): return "Empty text — skipped." ctx.deps.widgets.append({"type": "text", "content": content.strip()}) _log.debug("build_text total_widgets={}", len(ctx.deps.widgets)) return "Text block added." @ui_agent.tool async def build_primary_highlights( ctx: RunContext[UIAgentDeps], block: dict[str, Any], ) -> str: """Insert a primary highlights widget surfacing key numbers. Use when key numbers need to be prominently visible. Can be called at any position — not required to be first. Args: block: A stats block, e.g. ``{"type": "stats", "items": [{"value": "93", "label": "Total NCEs"}, {"value": "£35.5M", "label": "CE Exposure"}]}``. """ if not block or not isinstance(block, dict): _log.warning("build_primary_highlights: empty or non-dict block, ignoring") return "No primary highlights to render." if block.get("type") != "stats": _log.warning("build_primary_highlights: block type must be 'stats', got {}", block.get("type")) block["type"] = "stats" items = block.get("items") or [] if not items: _log.warning("build_primary_highlights: stats block has no items") return "Primary highlights block has no items — skipped." # Hard cap at 2 items regardless of what the LLM produced block["items"] = items[:2] try: validated = validate_highlight_content(block) block = validated.model_dump() except Exception as exc: return ( f"Schema validation failed: {exc}\n" "Only the fields listed above are wrong — keep all other fields exactly as they were. " "Fix only the failing fields and call build_primary_highlights again." ) ctx.deps.widgets.append({"type": "primary_highlights", "content": block}) _log.debug("build_primary_highlights total_widgets={}", len(ctx.deps.widgets)) return "Primary highlights built." @ui_agent.tool async def build_key_highlights( ctx: RunContext[UIAgentDeps], block: dict[str, Any], ) -> str: """Store a key-highlights block at this position in the widget sequence. Can be called multiple times. The ``block`` must match one of the ``KeyHighlightBlock`` shapes in the system prompt — always include ``"type"`` as the first field. Args: block: A fully-formed ``KeyHighlightBlock`` dict, e.g. ``{"type": "ranked", "items": [...]}``. """ if not block or not isinstance(block, dict): _log.warning("build_key_highlights: empty or non-dict block, ignoring") return "No key highlights to render." if "type" not in block: _log.warning("build_key_highlights: missing 'type' discriminant in block") return "Key highlights block missing 'type' field — skipped." try: validated = validate_highlight_content(block) block = validated.model_dump() except Exception as exc: return ( f"Schema validation failed: {exc}\n" "Only the fields listed above are wrong — keep all other fields exactly as they were. " "Fix only the failing fields and call build_key_highlights again." ) ctx.deps.widgets.append({"type": "key_highlights", "content": block}) _log.debug("build_key_highlights type={} total_widgets={}", block.get("type"), len(ctx.deps.widgets)) return "Key highlights built." @ui_agent.tool async def build_table( ctx: RunContext[UIAgentDeps], columns: list[str], column_keys: list[str], currency_keys: list[str] = [], ) -> str: """Render a data table using ALL rows from the query result. Call this when the user explicitly asks for a table, or when the data is record-level detail that no chart encodes better (e.g. a list of specific items, a multi-column detail view, raw lookup results). **Do NOT pass rows yourself.** This tool maps every row from the query result automatically using ``column_keys``. You only declare which columns to show and which SQL key each column maps to. Args: columns: Ordered list of human-readable column header strings. Write what a business reader would see — never raw SQL column names. e.g. ``["Contractor", "Cost Impact (£)", "Status"]``. column_keys: Ordered list of SQL column names from the raw query result, corresponding 1-to-1 with ``columns``. e.g. ``["contractor", "change_to_prices", "status"]``. Must have the same length as ``columns``. currency_keys: Subset of ``column_keys`` whose values are monetary amounts (GBP or converted currency). The frontend uses this list to apply compact currency formatting (1K / 1M) to those columns. e.g. ``["change_to_prices", "base_contract"]``. Omit or pass ``[]`` for columns that are counts, days, or text. """ if not columns: _log.warning("build_table: empty columns list, skipping") return "Table skipped — no columns provided." if not column_keys: _log.warning("build_table: empty column_keys list, skipping") return "Table skipped — no column_keys provided." if len(columns) != len(column_keys): _log.warning( "build_table: columns/column_keys length mismatch ({} vs {})", len(columns), len(column_keys), ) return ( f"Table skipped — columns ({len(columns)}) and column_keys " f"({len(column_keys)}) must have the same length." ) raw_rows = ctx.deps.rows if not raw_rows: _log.warning("build_table: no rows in deps, skipping") return "Table skipped — no query data available." available_keys = set(raw_rows[0].keys()) matching = [k for k in column_keys if k in available_keys] if not matching: missing = column_keys _log.warning( "build_table: no column_keys match primary rows — keys={} available={}, skipping", missing, list(available_keys), ) return ( f"Table skipped — none of the requested column_keys {missing} exist in the " f"primary query result. Available columns: {sorted(available_keys)}. " "Use only columns from the primary rows dataset." ) unmatched = [k for k in column_keys if k not in available_keys] if unmatched: _log.warning("build_table: some column_keys not in primary rows — missing={}", unmatched) def _fmt(val: Any) -> str: if val is None: return "—" return str(val) mapped_rows = [ {col: _fmt(row.get(key)) for col, key in zip(columns, column_keys)} for row in raw_rows ] # Resolve currency_keys → column headers so the frontend can format by header name currency_key_set = set(currency_keys) currency_columns = [ col for col, key in zip(columns, column_keys) if key in currency_key_set ] # Deduplicate: skip if an identical table already exists in the widget list for existing in ctx.deps.widgets: if ( existing.get("type") == "table" and existing.get("content", {}).get("columns") == columns and existing.get("content", {}).get("rows") == mapped_rows ): _log.warning("build_table: duplicate table detected, skipping") return "Table already built — duplicate skipped." ctx.deps.widgets.append({ "type": "table", "content": { "type": "table", "columns": columns, "rows": mapped_rows, "currency_columns": currency_columns, }, }) _log.debug( "build_table cols={} rows={} total_widgets={}", len(columns), len(mapped_rows), len(ctx.deps.widgets), ) return f"Table built — {len(columns)} columns, {len(mapped_rows)} rows." @ui_agent.tool async def compute_aggregate( ctx: RunContext[UIAgentDeps], column: str, operation: str, ) -> str: """Compute an exact aggregate over a numeric column in the row data. ALWAYS call this tool instead of computing totals, averages, counts, medians, or any other arithmetic mentally from the raw rows. Supported operations: sum, mean, median, min, max, count, std, var. For grouped aggregates (e.g. sum per category), call this once per group filtering is not needed — pandas handles mixed types gracefully. Args: column: Column name from the row data to aggregate. operation: Aggregation function — one of: sum, mean, median, min, max, count, std, var. """ if not ctx.deps.rows: return "No row data available." df = pd.DataFrame(ctx.deps.rows) if column not in df.columns: return f"Column '{column}' not found. Available columns: {list(df.columns)}" series = pd.to_numeric(df[column], errors="coerce").dropna() if series.empty: return f"Column '{column}' has no numeric values." try: result = series.agg(operation) except Exception as exc: return f"Operation '{operation}' failed: {exc}. Supported: sum, mean, median, min, max, count, std, var." return f"{operation}({column}) = {result}" # --------------------------------------------------------------------------- # Prompt builder # --------------------------------------------------------------------------- def _build_ui_prompt(req: "UIRenderRequest") -> str: """Build the full prompt for the UI agent.""" parts = [ f"Question: {req.question}", f"Currency: {req.currency_code}", f"Data agent explanation: {req.explanation}", ] if req.market_context: parts.append( "Market & Industry Context (sourced from live web search — " "weave benchmarks and comparisons into your narrative):\n" + req.market_context ) if req.decision_support_context: parts.append( "Decision Support Context (additional Tata Steel analysis datasets gathered after the main answer query):\n" + json.dumps(req.decision_support_context) ) if req.dc_explanation: parts.append( "Decision Context Mapping (for your background reasoning only — " "DO NOT emit a build_text widget for this content; it is automatically " "rendered in a separate hidden panel):\n" + req.dc_explanation ) if req.what_to_act_on: parts.append( "Action items (highest-priority items the user should act on — surface the top item " "as a build_primary_highlights block and list the rest via build_key_highlights):\n" + "\n".join(f"- {a}" for a in req.what_to_act_on) ) if req.urgency: parts.append( "Urgency signal (use this to write the build_text narrative — emphasise the deadline " "or financial risk described here):\n" + req.urgency ) if req.cost_of_inaction: parts.append( "Cost of inaction (weave into the build_text narrative):\n" + req.cost_of_inaction ) parts.append(f"Rows ({len(req.rows)} total): {json.dumps(req.rows)}") return "\n\n".join(parts) # --------------------------------------------------------------------------- # Request handlers # --------------------------------------------------------------------------- async def _render_handler(request: Request) -> JSONResponse: """POST /query — sync render: returns widget list.""" try: body = await request.json() req = UIRenderRequest(**body) except Exception as exc: return JSONResponse( AgentQueryResponse.failure( message=f"Invalid request: {exc}", code="PIPELINE_ERROR" ).model_dump(), status_code=400, ) if req.history is not None: message_history = deserialize_history(req.history) else: message_history = await load_conversation_history(req.conversation_id) deps = UIAgentDeps( rows=req.rows, question=req.question, explanation=req.explanation, market_context=req.market_context or "", storyline_hints=req.storyline_hints or [], decision_support_context=req.decision_support_context or [], dc_explanation=req.dc_explanation or "", currency_code=req.currency_code, ) prompt = _build_ui_prompt(req) try: with logfire.span("{q}", q=req.question[:120]): result = await ui_agent.agent.run(prompt, deps=deps, message_history=message_history) except Exception as exc: _log.error("ui_agent render error: {}", exc) return JSONResponse( AgentQueryResponse.failure(message=str(exc), code="PIPELINE_ERROR").model_dump() ) clean_text, diagrams = _split_explanation(req.explanation) # Drop any placeholder slots that were never filled (validation kept failing). filled_widgets = [w for w in deps.widgets if not w.get("_pending")] has_text_widget = any(w.get("type") == "text" for w in filled_widgets) widgets: list[dict[str, Any]] = [] if not has_text_widget and clean_text: widgets.append({"type": "text", "content": clean_text}) widgets += [{"type": "mermaid", "content": d} for d in diagrams] widgets += filled_widgets # all text + visualizations + highlights in call order if req.storyline_hints: widgets.append({"type": "followup_questions", "content": req.storyline_hints}) if req.dc_explanation: widgets.append({"type": "followup_explanation", "content": req.dc_explanation}) return JSONResponse(AgentQueryResponse.success(widgets=widgets).model_dump()) async def _render_stream_handler(request: Request) -> StreamingResponse | JSONResponse: """POST /query/stream — SSE streaming render. SSE events (server → client, in order): data: {"type": "widget_loading", "widget_type": str} data: {"type": "text_chunk", "content": str} data: {"type": "widget", "widget": {...}} data: {"type": "done", "success": true, "widgets": [...]} data: {"type": "error", "message": str, "code": str} """ try: body = await request.json() req = UIRenderRequest(**body) except Exception as exc: return JSONResponse( AgentQueryResponse.failure( message=f"Invalid request: {exc}", code="PIPELINE_ERROR" ).model_dump(), status_code=400, ) if req.history is not None: message_history = deserialize_history(req.history) else: message_history = await load_conversation_history(req.conversation_id) progress_queue: asyncio.Queue = asyncio.Queue() deps = UIAgentDeps( rows=req.rows, question=req.question, explanation=req.explanation, market_context=req.market_context or "", storyline_hints=req.storyline_hints or [], decision_support_context=req.decision_support_context or [], dc_explanation=req.dc_explanation or "", progress_queue=progress_queue, currency_code=req.currency_code, ) prompt = _build_ui_prompt(req) async def sse_events(): full_text: list[str] = [] merged: asyncio.Queue = asyncio.Queue() agent_done = asyncio.Event() had_error = False async def _drain_progress(): 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 while not progress_queue.empty(): await merged.put(progress_queue.get_nowait()) await merged.put(None) async def _run_agent(): nonlocal had_error try: with logfire.span("{q}", q=req.question[:120]): async with ui_agent.agent.run_stream(prompt, deps=deps, message_history=message_history) as stream: async for chunk in stream.stream_text(delta=True): full_text.append(chunk) await merged.put({"type": "text_chunk", "content": chunk}) except Exception as exc: _log.error("ui_agent stream render 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 agent_text = "".join(full_text) clean_text, diagrams = _split_explanation(agent_text) # Drop any placeholder slots that were never filled (validation kept failing). filled_widgets = [w for w in deps.widgets if not w.get("_pending")] has_text_widget = any(w.get("type") == "text" for w in filled_widgets) widgets: list[dict[str, Any]] = [] if not has_text_widget and clean_text: widgets.append({"type": "text", "content": clean_text}) widgets += [{"type": "mermaid", "content": d} for d in diagrams] # Emit every widget (text / visualization / key_highlights) in call order for w in filled_widgets: yield f"data: {json.dumps({'type': 'widget', 'widget': w})}\n\n" widgets.append(w) # Append follow-up questions directly — bypasses the LLM so hints are # always present and rendered by the frontend's SuggestedFollowups component. if req.storyline_hints: fq = {"type": "followup_questions", "content": req.storyline_hints} yield f"data: {json.dumps({'type': 'widget', 'widget': fq})}\n\n" widgets.append(fq) # Append DC mapping explanation directly — explains how the question was # mapped to a Decision Context and how each follow-up was derived. if req.dc_explanation: fe = {"type": "followup_explanation", "content": req.dc_explanation} yield f"data: {json.dumps({'type': 'widget', 'widget': fe})}\n\n" widgets.append(fe) yield f"data: {json.dumps({'type': 'done', **AgentQueryResponse.success(widgets=widgets).model_dump()})}\n\n" return StreamingResponse( sse_events(), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, ) # --------------------------------------------------------------------------- # ASGI app # --------------------------------------------------------------------------- app = make_metadata_aware_app( ui_agent, name=_cfg.name, description=_cfg.description, ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) app.add_route("/query", _render_handler, methods=["POST"]) app.add_route("/query/stream", _render_stream_handler, methods=["POST"])