204 lines
7.2 KiB
Python
204 lines
7.2 KiB
Python
"""langgraph wiring — the one place the framework appears.
|
|
|
|
Graph: START → plan → execute → synthesize → END
|
|
|
|
`execute` iterates Plan steps and dispatches each to its Analysis. A step
|
|
can declare an optional `fallback` Analysis that runs when the primary
|
|
errors or returns an empty finding.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
import logging
|
|
from typing import Any
|
|
|
|
from langgraph.graph import END, START, StateGraph
|
|
|
|
from api import langfuse_client as lf
|
|
from api.analyses.registry import REGISTRY as ANALYSES
|
|
from api.analyses.types import Finding
|
|
from api.llm import chat
|
|
from api.plan.planner import plan as run_planner
|
|
from api.prompts import load, render
|
|
from api.runtime import events
|
|
from api.runtime.context import current_analysis, current_run_id
|
|
from api.runtime.state import RunState
|
|
|
|
logger = logging.getLogger("nvi.runtime")
|
|
|
|
|
|
def _plan_node(state: RunState) -> dict[str, Any]:
|
|
plan_obj = run_planner(state["question"])
|
|
return {
|
|
"plan_rationale": plan_obj.rationale,
|
|
"plan_steps": [s.to_dict() for s in plan_obj.steps],
|
|
}
|
|
|
|
|
|
def _step_id(idx: int, name: str, suffix: str | None = None) -> str:
|
|
base = f"{name}#{idx}"
|
|
return f"{base}.{suffix}" if suffix else base
|
|
|
|
|
|
async def _invoke_analysis(name: str, args: dict[str, Any], question: str,
|
|
*, step_id: str, run_id: str, why: str) -> Finding:
|
|
"""Run a single Analysis with contextvar set, events around it, and
|
|
error→Finding catch. Returns the Finding (success or failure)."""
|
|
analysis = ANALYSES.get(name)
|
|
await events.publish(run_id, events.analysis_start(name, args, why, step_id=step_id))
|
|
if analysis is None:
|
|
f = Finding(analysis=name, summary="", error=f"unknown analysis {name!r}")
|
|
else:
|
|
token = current_analysis.set(step_id)
|
|
try:
|
|
f = await analysis.run(args, question)
|
|
except Exception as e:
|
|
logger.exception("analysis %s raised", name)
|
|
f = Finding(analysis=name, summary="", error=str(e))
|
|
finally:
|
|
current_analysis.reset(token)
|
|
await events.publish(run_id, events.analysis_end(name, f, step_id=step_id))
|
|
return f
|
|
|
|
|
|
def _finding_is_empty_or_errored(f: Finding) -> bool:
|
|
if f.error:
|
|
return True
|
|
if not f.summary and not f.rows:
|
|
return True
|
|
return False
|
|
|
|
|
|
async def _execute_node(state: RunState) -> dict[str, Any]:
|
|
findings: list[dict[str, Any]] = []
|
|
run_id = state["run_id"]
|
|
for idx, step in enumerate(state.get("plan_steps", [])):
|
|
primary_name = step["analysis"]
|
|
primary_args = step.get("args", {})
|
|
primary_why = step.get("why", "")
|
|
primary_step_id = _step_id(idx, primary_name)
|
|
|
|
primary = await _invoke_analysis(
|
|
primary_name, primary_args, state["question"],
|
|
step_id=primary_step_id, run_id=run_id, why=primary_why,
|
|
)
|
|
primary_dict = dataclasses.asdict(primary)
|
|
primary_dict["step_id"] = primary_step_id
|
|
findings.append(primary_dict)
|
|
|
|
fb = step.get("fallback")
|
|
if fb and _finding_is_empty_or_errored(primary):
|
|
fb_name = fb["analysis"]
|
|
fb_step_id = _step_id(idx, fb_name, suffix="fallback")
|
|
reason = primary.error or "primary returned empty finding"
|
|
await events.publish(run_id, events.analysis_fallback(
|
|
primary=primary_name, fallback=fb_name, reason=reason,
|
|
step_id=fb_step_id,
|
|
))
|
|
fb_finding = await _invoke_analysis(
|
|
fb_name, fb.get("args", {}), state["question"],
|
|
step_id=fb_step_id, run_id=run_id, why=fb.get("why", ""),
|
|
)
|
|
fb_dict = dataclasses.asdict(fb_finding)
|
|
fb_dict["step_id"] = fb_step_id
|
|
fb_dict["is_fallback_for"] = primary_step_id
|
|
findings.append(fb_dict)
|
|
return {"findings": findings}
|
|
|
|
|
|
async def _synthesize_node(state: RunState) -> dict[str, Any]:
|
|
await events.publish(state["run_id"], events.synth_start())
|
|
|
|
findings = state.get("findings", [])
|
|
# Effective findings = primary OR its fallback if the primary was empty.
|
|
# Pass everything through; the LLM weighs which one to cite.
|
|
usable = [f for f in findings if not _finding_is_empty_or_errored(_dict_to_finding(f))]
|
|
if not usable:
|
|
return {
|
|
"answer": "I couldn't answer this question — see the per-analysis errors above.",
|
|
"error": "all_analyses_failed",
|
|
}
|
|
|
|
findings_block = "\n\n".join(
|
|
f"[{f['analysis']}{ ' (fallback)' if f.get('is_fallback_for') else '' }]\n"
|
|
f"summary: {f.get('summary') or '(none)'}\n"
|
|
f"error: {f.get('error') or '(none)'}\n"
|
|
f"sql: {f.get('sql')}"
|
|
for f in findings
|
|
)
|
|
|
|
with lf.span("synthesize", as_type="generation", input={"findings": len(findings)}) as span:
|
|
answer = chat(
|
|
system=load("synthesize.system"),
|
|
user=render(
|
|
"synthesize.user",
|
|
question=state["question"],
|
|
findings_block=findings_block,
|
|
),
|
|
max_tokens=512,
|
|
)
|
|
span.update(output={"answer": answer})
|
|
return {"answer": answer}
|
|
|
|
|
|
def _dict_to_finding(d: dict[str, Any]) -> Finding:
|
|
return Finding(
|
|
analysis=d.get("analysis", ""),
|
|
summary=d.get("summary", ""),
|
|
rows=d.get("rows") or [],
|
|
sql=d.get("sql") or [],
|
|
error=d.get("error"),
|
|
metadata=d.get("metadata") or {},
|
|
)
|
|
|
|
|
|
def build_graph():
|
|
g: StateGraph = StateGraph(RunState)
|
|
g.add_node("plan", _plan_node)
|
|
g.add_node("execute", _execute_node)
|
|
g.add_node("synthesize", _synthesize_node)
|
|
g.add_edge(START, "plan")
|
|
g.add_edge("plan", "execute")
|
|
g.add_edge("execute", "synthesize")
|
|
g.add_edge("synthesize", END)
|
|
return g.compile()
|
|
|
|
|
|
GRAPH = build_graph()
|
|
|
|
|
|
async def run(run_id: str, question: str) -> RunState:
|
|
"""Drive the graph end-to-end, publishing SSE events along the way."""
|
|
initial: RunState = {"run_id": run_id, "question": question}
|
|
final: RunState = initial
|
|
token = current_run_id.set(run_id)
|
|
try:
|
|
await events.publish(run_id, events.run_start(question))
|
|
async for event in GRAPH.astream(initial, stream_mode="updates"):
|
|
for node, update in event.items():
|
|
await events.publish(run_id, events.node_update(node, update))
|
|
if isinstance(update, dict):
|
|
final = {**final, **update}
|
|
if node == "plan" and isinstance(update, dict):
|
|
await events.publish(
|
|
run_id,
|
|
events.plan_ready(
|
|
update.get("plan_rationale", ""),
|
|
update.get("plan_steps", []),
|
|
),
|
|
)
|
|
await events.publish(
|
|
run_id,
|
|
events.run_end(answer=final.get("answer", ""), error=final.get("error")),
|
|
)
|
|
except Exception as e:
|
|
logger.exception("run %s failed", run_id)
|
|
await events.publish(run_id, events.run_end(error=str(e)))
|
|
final = {**final, "error": str(e)}
|
|
finally:
|
|
current_run_id.reset(token)
|
|
await events.close(run_id)
|
|
lf.flush()
|
|
return final
|