verbose live UI + tool-level SSE events + Groq default + regression tests

2026-06-03 05:04:29 -03:00
parent 131f4d9b86
commit e124a8a7d9
69 changed files with 3030 additions and 137 deletions
--- a/api/analyses/direct_answer.py
+++ b/api/analyses/direct_answer.py
@@ -0,0 +1,81 @@
+"""direct_answer Analysis — CoT, one-shot.
+
+Pattern: text-to-SQL → execute → ask the LLM to interpret the rows. No loop.
+For questions that resolve to a single SQL query and a short interpretation.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from api import langfuse_client as lf
+from api.analyses.base import Analysis
+from api.analyses.types import Finding
+from api.llm import chat
+from api.prompts import load, render
+from api.runtime import events
+from api.tools.execute_sql import execute_sql
+from api.tools.text_to_sql import text_to_sql
+
+logger = logging.getLogger("nvi.analyses.direct_answer")
+
+
+class DirectAnswer(Analysis):
+    name = "direct_answer"
+    description = (
+        "Answer a question that resolves to a single SQL query. Best for "
+        "lookup / aggregation questions with a single, well-defined metric."
+    )
+    args_schema = {
+        "question": {"type": "string", "description": "Refined question this Analysis should answer."},
+        "hint_tables": {"type": "array", "items": "string", "description": "Tables the planner thinks are relevant (optional)."},
+    }
+
+    async def run(self, args: dict[str, Any], question: str) -> Finding:
+        sub_q = args.get("question") or question
+        hint_tables = args.get("hint_tables")
+
+        with lf.span("analysis.direct_answer", input={"question": sub_q}) as span:
+            try:
+                await events.publish_current(events.tool_call_start("text_to_sql", input={"question": sub_q}))
+                t2s = text_to_sql(sub_q, hint_tables=hint_tables)
+                await events.publish_current(events.tool_call_end(
+                    "text_to_sql", output={"sql": t2s.sql, "tables": t2s.used_tables},
+                ))
+
+                await events.publish_current(events.tool_call_start("execute_sql", input={"sql": t2s.sql}))
+                result = execute_sql(t2s.sql)
+                await events.publish_current(events.tool_call_end(
+                    "execute_sql",
+                    output={"row_count": result.row_count, "truncated": result.truncated,
+                            "preview": result.as_dicts()[:5]},
+                ))
+
+                interpret_system = load("direct_answer.interpret")
+                interpret_user = render(
+                    "direct_answer.interpret.user",
+                    question=sub_q,
+                    sql=t2s.sql,
+                    rows=repr(result.as_dicts()[:20]),
+                )
+                await events.publish_current(events.llm_call(
+                    "direct_answer.interpret",
+                    system_len=len(interpret_system),
+                    user_len=len(interpret_user),
+                ))
+                summary = chat(system=interpret_system, user=interpret_user, max_tokens=512)
+            except Exception as e:
+                logger.exception("direct_answer failed")
+                await events.publish_current(events.tool_call_end("direct_answer", error=str(e)))
+                span.update(output={"error": str(e)})
+                return Finding(analysis=self.name, summary="", error=str(e))
+
+            span.update(output={"summary": summary, "rows": result.row_count})
+            return Finding(
+                analysis=self.name,
+                summary=summary,
+                rows=result.as_dicts()[:20],
+                sql=[t2s.sql],
+                metadata={"row_count": result.row_count, "truncated": result.truncated},
+            )