| """End-to-end chat simulator with full step transparency (in-process). |
| |
| Simulates a user chatting inside ONE analysis session, from creation onward, and |
| prints what each step of the pipeline does: |
| |
| - the ROUTER decision (intent / rewritten query / confidence) |
| - slow-path STATUS pings (Planning… / Running N steps…) |
| - every TOOL call the slow path makes (check_data / retrieve_data / analyze_* — |
| with result kind, row count, latency, error) |
| - every LLM call (router / planner / assembler / chatbot / help) with |
| input-token / output-token / latency, and a snippet of the raw model output |
| - the streamed ANSWER + its SOURCES |
| - a per-turn timing + token summary |
| - finally, a REPORT generated from the slow-path report_inputs the run produced |
| |
| It calls `ChatHandler.handle()` IN-PROCESS (no server) so it can see the internal |
| LLM outputs the SSE endpoint hides. Transparency is captured by injecting a custom |
| tracer (`ScriptTracer`) into the exact seam the handler already threads its Langfuse |
| callbacks + tool spans through — no source changes. |
| |
| Run as a module from the repo root (so `src` imports resolve): |
| |
| uv run python -m eval.chat_sim.run_chat # predefined Titanic convo + report |
| uv run python -m eval.chat_sim.run_chat --interactive # you type the messages |
| uv run python -m eval.chat_sim.run_chat --no-report # skip the report capstone |
| uv run python -m eval.chat_sim.run_chat --no-bind # don't scope to Titanic (whole catalog) |
| |
| Needs a populated `.env` (Azure OpenAI + Postgres + Azure Blob for the Titanic |
| Parquet). Writes to the DB the `.env` points at (analysis state + report_inputs + |
| report) — point it at the playground DB. ENABLE_SLOW_PATH is forced on here. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import asyncio |
| import json |
| import sys |
| import time |
| import uuid |
| from dataclasses import dataclass, field |
| from typing import Any |
|
|
| |
| |
| if sys.platform == "win32": |
| asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) |
|
|
| |
| for _stream in (sys.stdout, sys.stderr): |
| try: |
| _stream.reconfigure(encoding="utf-8") |
| except Exception: |
| pass |
|
|
| from langchain_core.callbacks import BaseCallbackHandler |
| from langchain_core.messages import AIMessage, BaseMessage, HumanMessage |
|
|
| from src.agents.chat_handler import ChatHandler |
|
|
| |
| |
| |
| DEFAULT_USER_ID = "4b5d1bac-7211-490f-9a3d-66fed0168d5a" |
| TITANIC_SOURCE_ID = "9b565bc8-ccc4-4d10-9382-0bad416a091b" |
| TITANIC_NAME = "Titanic-Dataset.csv" |
|
|
| OBJECTIVE = "Understand what drove passenger survival on the Titanic — by sex, class, and fare." |
|
|
| |
| |
| DEFAULT_TURNS = [ |
| "What can you help me do in this analysis?", |
| "What data do I have available here?", |
| "What was the overall passenger survival rate, and how did it differ " |
| "between male and female passengers?", |
| "Did higher passenger class (Pclass) come with a higher average fare and " |
| "a higher survival rate?", |
| ] |
|
|
| |
| _C = { |
| "h": "\033[1;36m", "u": "\033[1;33m", "ai": "\033[1;32m", |
| "dim": "\033[2m", "warn": "\033[1;31m", "r": "\033[0m", |
| } |
|
|
|
|
| def c(key: str, text: str) -> str: |
| return f"{_C.get(key, '')}{text}{_C['r']}" if _C.get("_on", True) else text |
|
|
|
|
| |
|
|
|
|
| @dataclass |
| class LlmCall: |
| idx: int |
| ms: int | None |
| tin: int |
| tout: int |
| ttot: int |
| prompt_preview: str |
| output_preview: str |
| masked: bool |
|
|
|
|
| @dataclass |
| class ToolCall: |
| tool: str |
| arg_keys: list[str] |
| kind: str | None |
| rows: int | None |
| error: str | None |
| ms: int |
|
|
|
|
| @dataclass |
| class Sink: |
| """Per-turn collector shared by all StepLoggers + spans of that turn.""" |
| llm: list[LlmCall] = field(default_factory=list) |
| tools: list[ToolCall] = field(default_factory=list) |
|
|
|
|
| def _usage(response: Any) -> tuple[int, int, int]: |
| """Sum token usage off an LLMResult (usage_metadata, legacy fallback).""" |
| tin = tout = ttot = 0 |
| for gens in getattr(response, "generations", []) or []: |
| for g in gens: |
| msg = getattr(g, "message", None) |
| um = getattr(msg, "usage_metadata", None) if msg else None |
| if um: |
| tin += um.get("input_tokens", 0) |
| tout += um.get("output_tokens", 0) |
| ttot += um.get("total_tokens", 0) |
| if ttot == 0 and getattr(response, "llm_output", None): |
| u = response.llm_output.get("token_usage") or {} |
| tin += u.get("prompt_tokens", 0) |
| tout += u.get("completion_tokens", 0) |
| ttot += u.get("total_tokens", 0) |
| return tin, tout, ttot |
|
|
|
|
| def _out_text(response: Any) -> str: |
| try: |
| gens = response.generations |
| g = gens[0][0] |
| msg = getattr(g, "message", None) |
| return (getattr(msg, "content", None) or getattr(g, "text", "") or "").strip() |
| except Exception: |
| return "" |
|
|
|
|
| def _preview(text: str, n: int = 240) -> str: |
| text = " ".join(str(text).split()) |
| return text if len(text) <= n else text[: n - 1] + "…" |
|
|
|
|
| class StepLogger(BaseCallbackHandler): |
| """One per `tracer.callbacks()` call; all share the turn's Sink. |
| |
| Captures each LLM call's latency + tokens + a snippet of prompt/output. Matches |
| start->end by run_id so concurrent/streamed calls don't cross wires. |
| """ |
|
|
| def __init__(self, sink: Sink, masked: bool = False) -> None: |
| self.sink = sink |
| self.masked = masked |
| self._t0: dict[Any, float] = {} |
| self._prompt: dict[Any, str] = {} |
|
|
| def on_chat_model_start(self, serialized, messages, *, run_id=None, **kw): |
| self._t0[run_id] = time.perf_counter() |
| try: |
| flat = [m for grp in messages for m in grp] |
| self._prompt[run_id] = _preview( |
| next((getattr(m, "content", "") for m in flat |
| if m.__class__.__name__.startswith("System")), "" |
| ) or (flat[-1].content if flat else ""), 120 |
| ) |
| except Exception: |
| self._prompt[run_id] = "" |
|
|
| def on_llm_start(self, serialized, prompts, *, run_id=None, **kw): |
| self._t0[run_id] = time.perf_counter() |
| self._prompt[run_id] = _preview(prompts[0] if prompts else "", 120) |
|
|
| def on_llm_end(self, response, *, run_id=None, **kw): |
| t0 = self._t0.pop(run_id, None) |
| ms = round((time.perf_counter() - t0) * 1000) if t0 else None |
| tin, tout, ttot = _usage(response) |
| self.sink.llm.append(LlmCall( |
| idx=len(self.sink.llm) + 1, ms=ms, tin=tin, tout=tout, ttot=ttot, |
| prompt_preview=self._prompt.pop(run_id, ""), |
| output_preview=_preview(_out_text(response)), |
| masked=self.masked, |
| )) |
|
|
|
|
| class ScriptSpan: |
| """Mirrors tracing._ToolSpan: a metadata-only span around one slow-path tool call.""" |
|
|
| def __init__(self, sink: Sink, tool: str, args: dict) -> None: |
| self.sink = sink |
| self.tool = tool |
| self.args = args |
| self.t0 = time.perf_counter() |
|
|
| def end(self, out: Any) -> None: |
| kind = getattr(out, "kind", None) |
| rows = len(getattr(out, "rows", None) or []) if kind == "table" else None |
| err = getattr(out, "error", None) |
| self.sink.tools.append(ToolCall( |
| tool=self.tool, |
| arg_keys=sorted(self.args) if isinstance(self.args, dict) else [], |
| kind=kind, rows=rows, |
| error=_preview(err, 160) if err else None, |
| ms=round((time.perf_counter() - self.t0) * 1000), |
| )) |
|
|
|
|
| class ScriptTracer: |
| """Drop-in for RequestTracer/NullTracer. active=True so the slow path wraps its |
| ToolInvoker in TracingToolInvoker and routes tool spans here.""" |
|
|
| active = True |
|
|
| def __init__(self, sink: Sink) -> None: |
| self.sink = sink |
|
|
| def callbacks(self, *, masked: bool = False) -> list: |
| return [StepLogger(self.sink, masked)] |
|
|
| def tool_span(self, tool: str, args: dict) -> Any: |
| return ScriptSpan(self.sink, tool, args) |
|
|
| def end(self, *, output: Any = None) -> None: |
| return None |
|
|
|
|
| class InstrumentedChatHandler(ChatHandler): |
| """ChatHandler that emits our ScriptTracer instead of Langfuse/Null, so every |
| LLM + tool step of a turn lands in `self.sink`.""" |
|
|
| def __init__(self, *a, **k) -> None: |
| super().__init__(*a, **k) |
| self.sink = Sink() |
|
|
| def _make_tracer(self, user_id: str, question: str) -> Any: |
| return ScriptTracer(self.sink) |
|
|
|
|
| |
|
|
|
|
| def banner(text: str, ch: str = "═") -> None: |
| print(f"\n{c('h', ch * 78)}\n{c('h', text)}\n{c('h', ch * 78)}") |
|
|
|
|
| def _llm_labels(intent: str | None, n: int) -> list[str]: |
| """Best-effort name per LLM call, by the path's known call order.""" |
| seq = { |
| "structured_flow": ["router", "planner", "assembler"], |
| "help": ["router", "help"], |
| "unstructured_flow": ["router", "chatbot"], |
| "chat": ["router", "chatbot"], |
| "check": ["router"], |
| }.get(intent or "", ["router"]) |
| out = [] |
| for i in range(n): |
| if i < len(seq) - 1: |
| out.append(seq[i]) |
| elif i == n - 1: |
| out.append(seq[-1]) |
| else: |
| out.append(f"{seq[1] if len(seq) > 1 else 'llm'}·retry") |
| return out |
|
|
|
|
| def print_turn_steps(sink: Sink, intent: str | None, total_ms: int) -> None: |
| if sink.tools: |
| print(c("dim", "\n tool calls (slow path):")) |
| for t in sink.tools: |
| tag = c("warn", "ERROR") if t.error else (t.kind or "ok") |
| extra = f" rows={t.rows}" if t.rows is not None else "" |
| print(f" • {t.tool:<18} {tag:<7}{extra:<10} {t.ms:>5}ms" |
| f" args={t.arg_keys}") |
| if t.error: |
| print(c("warn", f" ↳ {t.error}")) |
|
|
| if sink.llm: |
| labels = _llm_labels(intent, len(sink.llm)) |
| print(c("dim", "\n llm calls (output / tokens / latency):")) |
| print(c("dim", f" {'#':<2} {'step':<14} {'in':>6} {'out':>6} {'tot':>6} {'ms':>6}")) |
| for call, label in zip(sink.llm, labels): |
| ms = f"{call.ms}" if call.ms is not None else "?" |
| print(f" {call.idx:<2} {label:<14} {call.tin:>6} {call.tout:>6} " |
| f"{call.ttot:>6} {ms:>6}") |
| print(c("dim", f" prompt: {call.prompt_preview}")) |
| |
| |
| |
| out = call.output_preview or "<no text content — structured / tool-call output>" |
| tag = " (masked→cloud)" if call.masked else "" |
| print(c("dim", f" output{tag}: {out}")) |
|
|
| tin = sum(c_.tin for c_ in sink.llm) |
| tout = sum(c_.tout for c_ in sink.llm) |
| print(c("dim", f"\n ── turn: {total_ms}ms · {len(sink.llm)} llm call(s) · " |
| f"{len(sink.tools)} tool call(s) · {tin}+{tout} tokens")) |
|
|
|
|
| |
|
|
|
|
| async def setup_analysis(user_id: str, bind_titanic: bool) -> str: |
| """Create a fresh analysis session (state row) + optionally bind it to Titanic. |
| |
| Mirrors what `/analysis/create` does: a state row carrying the goal, plus an |
| analysis-scope `data_catalog` row (B) restricting the analysis to one source, so |
| structured_flow is scoped deterministically. Returns the analysis_id (== room_id). |
| """ |
| from src.agents.state_store import AnalysisStateStore |
|
|
| analysis_id = str(uuid.uuid4()) |
| await AnalysisStateStore().create( |
| analysis_id=analysis_id, |
| user_id=user_id, |
| analysis_title="Titanic survival analysis (sim)", |
| objective=OBJECTIVE, |
| ) |
| print(f" created analysis {c('h', analysis_id)}") |
| print(f" objective: {OBJECTIVE}") |
|
|
| if bind_titanic: |
| try: |
| from datetime import UTC, datetime |
|
|
| from src.catalog.models import Catalog as CatalogModel |
| from src.catalog.store import CatalogStore |
| from src.db.postgres.connection import AsyncSessionLocal |
| from src.db.postgres.models import Catalog as CatalogRow |
|
|
| |
| |
| |
| |
| user_cat = await CatalogStore().get(user_id) |
| titanic = [ |
| s for s in (user_cat.sources if user_cat else []) |
| if s.source_id == TITANIC_SOURCE_ID |
| ] |
| if not titanic: |
| print(c("warn", f" Titanic source {TITANIC_SOURCE_ID} not in user " |
| "catalog — running unscoped (whole catalog)")) |
| else: |
| scoped = CatalogModel( |
| user_id=user_id, generated_at=datetime.now(UTC), sources=titanic, |
| ) |
| async with AsyncSessionLocal() as s: |
| s.add(CatalogRow( |
| scope_type="analysis", user_id=user_id, analysis_id=analysis_id, |
| catalog_payload=scoped.model_dump(mode="json"), |
| )) |
| await s.commit() |
| print(f" bound source: {TITANIC_NAME} ({TITANIC_SOURCE_ID}) " |
| f"{c('dim', '→ structured_flow scoped to Titanic (analysis catalog)')}") |
| except Exception as e: |
| print(c("warn", f" binding skipped ({type(e).__name__}: {e}) — " |
| f"fail-open to whole catalog")) |
| else: |
| print(c("dim", " no binding → structured_flow sees the whole catalog")) |
| return analysis_id |
|
|
|
|
| async def run_turn( |
| handler: InstrumentedChatHandler, |
| user_id: str, |
| analysis_id: str, |
| message: str, |
| history: list[BaseMessage], |
| ) -> None: |
| handler.sink = Sink() |
| banner(f"USER ▸ {message}", "─") |
|
|
| answer = "" |
| sources: list[dict] = [] |
| intent: str | None = None |
| t0 = time.perf_counter() |
|
|
| async for ev in handler.handle(message, user_id, history, analysis_id=analysis_id): |
| kind, data = ev["event"], ev["data"] |
| if kind == "intent": |
| try: |
| d = json.loads(data) |
| intent = d.get("intent") |
| print(f" {c('h', 'ROUTER')} → intent={c('h', intent)} " |
| f"confidence={d.get('confidence')}") |
| rq = d.get("rewritten_query") |
| if rq and rq != message: |
| print(c("dim", f" rewritten: {rq}")) |
| except Exception: |
| pass |
| elif kind == "status": |
| print(c("dim", f" · {data}")) |
| elif kind == "sources": |
| try: |
| sources = json.loads(data) or [] |
| except Exception: |
| sources = [] |
| elif kind == "chunk": |
| answer += data |
| elif kind == "error": |
| print(c("warn", f" ERROR: {data}")) |
|
|
| total_ms = round((time.perf_counter() - t0) * 1000) |
| print(f"\n {c('ai', 'ANSWER')} ▾") |
| for line in (answer or "(empty)").splitlines() or ["(empty)"]: |
| print(f" {line}") |
| if sources: |
| print(c("dim", f"\n sources ({len(sources)}): " |
| + ", ".join(s.get("filename") or s.get("document_id", "?") |
| for s in sources))) |
| print_turn_steps(handler.sink, intent, total_ms) |
|
|
| history.append(HumanMessage(content=message)) |
| history.append(AIMessage(content=answer)) |
|
|
|
|
| async def generate_report(user_id: str, analysis_id: str) -> None: |
| """Mirror POST /report: floor check → ReportGenerator → ReportStore → print.""" |
| banner("REPORT ▸ generating from accumulated report_inputs") |
| from src.agents.gate import stub_analysis_state |
| from src.agents.report.generator import ReportGenerator |
| from src.agents.report.readiness import report_floor |
| from src.agents.report.schemas import ProblemStatement |
| from src.agents.report.store import ReportStore |
| from src.agents.state_store import AnalysisStateStore |
|
|
| state = await AnalysisStateStore().get(analysis_id) |
| missing, _ = await report_floor( |
| analysis_id, state or stub_analysis_state(problem_validated=False) |
| ) |
| if missing: |
| print(c("warn", f" floor not met (409 in the API): {', '.join(missing)}")) |
| print(c("dim", " → need ≥1 successful slow-path analysis first " |
| "(did the structured turns run analyze_* tools?)")) |
| return |
|
|
| objective = (getattr(state, "objective", "") or |
| getattr(state, "problem_statement", "") or "") |
| ps = ProblemStatement( |
| objective=objective, |
| business_questions=list(getattr(state, "business_questions", []) or []), |
| ) |
| t0 = time.perf_counter() |
| report = await ReportGenerator().generate( |
| analysis_id, user_id, problem_statement=ps, user_name=None |
| ) |
| saved = await ReportStore().save(report) |
| print(f" generated v{saved.version} in {round((time.perf_counter()-t0)*1000)}ms " |
| f"· report_id={saved.report_id} · built from {len(saved.record_ids)} record(s)\n") |
| print(c("dim", " ── rendered markdown ──")) |
| for line in saved.rendered_markdown.splitlines(): |
| print(f" {line}") |
|
|
|
|
| |
|
|
|
|
| async def amain(args: argparse.Namespace) -> None: |
| if args.plain: |
| _C["_on"] = False |
|
|
| banner("DATA EYOND — end-to-end chat simulator (in-process)") |
| print(f" user_id: {args.user_id}") |
| print(f" slow_path: ON tracing→terminal: ON db: (from .env)") |
|
|
| handler = InstrumentedChatHandler( |
| enable_tracing=False, enable_gate=False |
| ) |
| analysis_id = await setup_analysis(args.user_id, bind_titanic=not args.no_bind) |
| history: list[BaseMessage] = [] |
|
|
| if args.interactive: |
| print(c("dim", "\n interactive mode — type a message, 'report' to generate, " |
| "'exit' to quit.\n")) |
| loop = asyncio.get_event_loop() |
| while True: |
| try: |
| msg = (await loop.run_in_executor(None, input, "you ▸ ")).strip() |
| except (EOFError, KeyboardInterrupt): |
| break |
| if not msg: |
| continue |
| if msg.lower() in {"exit", "quit"}: |
| break |
| if msg.lower() == "report": |
| await generate_report(args.user_id, analysis_id) |
| continue |
| await run_turn(handler, args.user_id, analysis_id, msg, history) |
| else: |
| turns = DEFAULT_TURNS[: args.max_turns] if args.max_turns else DEFAULT_TURNS |
| for msg in turns: |
| await run_turn(handler, args.user_id, analysis_id, msg, history) |
|
|
| if not args.no_report: |
| await generate_report(args.user_id, analysis_id) |
|
|
| banner("DONE") |
| print(f" analysis_id (== room_id): {analysis_id}") |
| print(c("dim", " state, report_inputs, and report were written to the .env DB.")) |
|
|
|
|
| def main() -> None: |
| p = argparse.ArgumentParser(description="End-to-end chat simulator with step transparency") |
| p.add_argument("--user-id", default=DEFAULT_USER_ID) |
| p.add_argument("--interactive", action="store_true", help="type messages yourself") |
| p.add_argument("--no-report", action="store_true", help="skip the report capstone") |
| p.add_argument("--no-bind", action="store_true", |
| help="don't scope to Titanic (planner sees the whole catalog)") |
| p.add_argument("--max-turns", type=int, default=0, |
| help="run only the first N scripted turns (cheap smoke test)") |
| p.add_argument("--plain", action="store_true", help="disable ANSI colors") |
| asyncio.run(amain(p.parse_args())) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
