eval/chat_sim/run_chat.py

"""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

# --- Windows: psycopg3 async needs the selector loop (mirrors run.py). Set BEFORE
# anything touches asyncio.
if sys.platform == "win32":
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

# Windows consoles default to cp1252 and choke on the box-drawing glyphs below.
for _stream in (sys.stdout, sys.stderr):
    try:
        _stream.reconfigure(encoding="utf-8")  # type: ignore[union-attr]
    except Exception:
        pass

from langchain_core.callbacks import BaseCallbackHandler  # noqa: E402
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage  # noqa: E402

from src.agents.chat_handler import ChatHandler  # noqa: E402

# This user's catalog (verified in the playground DB):
#   tabular  source 9b565bc8-… = Titanic-Dataset.csv (891 rows)
#   schema   source aaa0a4c6-… = "dummy" postgres (orders/customers/products/…)
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 scripted conversation. Chosen to exercise every router intent against the
# real Titanic columns (Survived, Sex, Pclass, Age, Fare, Embarked).
DEFAULT_TURNS = [
    "What can you help me do in this analysis?",                                   # -> help
    "What data do I have available here?",                                         # -> check
    "What was the overall passenger survival rate, and how did it differ "
    "between male and female passengers?",                                         # -> structured_flow
    "Did higher passenger class (Pclass) come with a higher average fare and "
    "a higher survival rate?",                                                     # -> structured_flow
]

# ANSI (Windows Terminal / VS Code support it). Disable with --plain.
_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


# ───────────────────────── transparency capture ──────────────────────────────


@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):  # type: ignore[override]
        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):  # type: ignore[override]
        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):  # type: ignore[override]
        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:  # type: ignore[override]
        return ScriptTracer(self.sink)


# ───────────────────────────── pretty printing ───────────────────────────────


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])           # last call = final author
        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}"))
            # Local tool over your own data → show output regardless of the masked
            # flag (masking only matters for Langfuse Cloud). Note when it's a
            # cloud-masked call or has no text (structured / tool-call output).
            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"))


# ───────────────────────────── setup / turns ─────────────────────────────────


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

            # Scope structured_flow by seeding the analysis-scope catalog (B): the
            # user's catalog restricted to Titanic. structured_flow reads this row via
            # CatalogStore.get_by_analysis (the data_sources binding table was removed;
            # in production Go materializes B from analyses.data_bind).
            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:  # noqa: BLE001 — fail-open to whole catalog
            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}")


# ──────────────────────────────── main ───────────────────────────────────────


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()