107 lines
4.3 KiB
Python
107 lines
4.3 KiB
Python
"""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._pick import pick_for_question
|
|
from api.analyses.base import Analysis
|
|
from api.analyses.types import Finding
|
|
from api.composer import compose
|
|
from api.llm import chat
|
|
from api.prompts import load, render
|
|
from api.recon import load_recon
|
|
from api.runtime import events
|
|
from api.tools.execute_sql import execute_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."},
|
|
"allowed_metrics": {"type": "array", "items": "string", "description": "Metrics the planner judged relevant (optional; defaults to all)."},
|
|
}
|
|
|
|
async def run(self, args: dict[str, Any], question: str) -> Finding:
|
|
sub_q = args.get("question") or question
|
|
allowed_metrics = args.get("allowed_metrics")
|
|
|
|
with lf.span("analysis.direct_answer", input={"question": sub_q}) as span:
|
|
try:
|
|
recon = load_recon()
|
|
|
|
await events.publish_current(events.tool_call_start(
|
|
"pick_for_question", input={"question": sub_q},
|
|
))
|
|
outcome = pick_for_question(
|
|
sub_q, recon=recon, allowed_metrics=allowed_metrics,
|
|
span_name="direct_answer.pick",
|
|
)
|
|
await events.publish_current(events.tool_call_end(
|
|
"pick_for_question", output={"pick": outcome.pick.to_dict()},
|
|
))
|
|
|
|
await events.publish_current(events.tool_call_start(
|
|
"compose_sql", input={"pick": outcome.pick.to_dict()},
|
|
))
|
|
composed = compose(outcome.pick, recon)
|
|
await events.publish_current(events.tool_call_end(
|
|
"compose_sql",
|
|
output={"sql": composed.sql, "tables": composed.used_tables},
|
|
))
|
|
|
|
await events.publish_current(events.tool_call_start("execute_sql", input={"sql": composed.sql}))
|
|
result = execute_sql(composed.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=composed.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,
|
|
span_name="direct_answer.interpret",
|
|
)
|
|
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=[composed.sql],
|
|
metadata={
|
|
"row_count": result.row_count,
|
|
"truncated": result.truncated,
|
|
"metric": outcome.pick.metric,
|
|
},
|
|
)
|