Rifqi Hafizuddin
Scope structured_flow + report to analysis catalog; drop dead data_sources binding
c333ec8 | """ChatHandler — top-level Phase 2 chat orchestrator. | |
| End-to-end flow per user message: | |
| 1. `OrchestratorAgent.classify` → RouterDecision (one of six intents). | |
| 2. Route by intent: | |
| - `chat` → no context. Pass straight to ChatbotAgent. | |
| - `structured_flow` → CatalogReader → slow analytical path | |
| (Planner → TaskRunner → Assembler). | |
| - `unstructured_flow` → DocumentRetriever (RAG over PGVector) → | |
| list[DocumentChunk]. | |
| - `check` → check_data / check_knowledge tool → rendered table. | |
| - `help` → Help skill: analysis state + history → streamed guidance. | |
| (`problem_statement` was removed 2026-06-24 — the goal is now user-entered | |
| `objective` + `business_questions` captured at onboarding, with no agent skill.) | |
| 3. `ChatbotAgent.astream` → yield text tokens. | |
| 4. Wrap each step into an SSE-style event dict so the API endpoint can | |
| stream them as Server-Sent Events. | |
| The chat endpoint (`src/api/v1/chat.py`) calls `ChatHandler.handle(...)` per | |
| request, behind two endpoint-level pre-filters: a greeting/farewell | |
| short-circuit and a Redis response cache (both skip the LLM on a hit). | |
| All dependencies are injectable for tests. Default constructors lazy-build | |
| production deps (no `Settings()` triggered at import time as long as you | |
| inject mocks). | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import json | |
| from collections.abc import AsyncIterator, Callable | |
| from typing import TYPE_CHECKING, Any | |
| from langchain_core.messages import BaseMessage | |
| from src.middlewares.logging import get_logger | |
| from src.retrieval.base import RetrievalResult | |
| from src.traceability import TraceabilityScratchpad, TraceabilityToolInvoker | |
| from .chatbot import ChatbotAgent, DocumentChunk | |
| from .guard import InputGuard | |
| from .handlers.check import run_check | |
| from .handlers.help import HelpAgent | |
| # `run_problem_statement` unwired 2026-06-24 (problem_statement removed from the router). | |
| # `ProblemStatementAgent` kept — still referenced by the constructor + _get_ps_agent. | |
| from .handlers.problem_statement import ProblemStatementAgent | |
| from .language import detect_reply_language | |
| from .orchestration import OrchestratorAgent | |
| from .refusals import blocked_message, out_of_scope_message | |
| if TYPE_CHECKING: | |
| from ..catalog.reader import CatalogReader | |
| from ..retrieval.router import RetrievalRouter | |
| from ..traceability.store import TraceabilityStore | |
| from .gate import AnalysisState | |
| from .slow_path.coordinator import SlowPathCoordinator | |
| from .slow_path.store import ReportInputStore | |
| logger = get_logger("chat_handler") | |
| def _is_content_filter_error(err: Exception) -> bool: | |
| """True when an exception is Azure's content-filter / jailbreak rejection. | |
| Azure OpenAI returns a 400 (`code='content_filter'`, `jailbreak.detected=True`) | |
| when a prompt trips its Responsible-AI policy. LangChain surfaces it as a raised | |
| exception; we string-match rather than import the concrete openai error type so | |
| the check survives SDK/version changes. | |
| """ | |
| s = str(err).lower() | |
| return ( | |
| "content_filter" in s | |
| or "responsibleai" in s | |
| or "jailbreak" in s | |
| or "content management policy" in s | |
| ) | |
| class ChatHandler: | |
| """Top-level chat orchestrator. | |
| Returns an `AsyncIterator[dict]` of SSE-style events with shape | |
| `{"event": <name>, "data": <str>}`. Event types: | |
| - `intent` — emitted once after classification (JSON-encoded decision) | |
| - `sources` — JSON array of source refs (one per structured table, or | |
| per (document_id, page_label) for unstructured) | |
| - `chunk` — text fragment of the streaming answer (one per token) | |
| - `done` — end of stream (data is empty string) | |
| - `error` — failure; data is a user-facing message | |
| """ | |
| def __init__( | |
| self, | |
| intent_router: OrchestratorAgent | None = None, | |
| answer_agent: ChatbotAgent | None = None, | |
| catalog_reader: CatalogReader | None = None, | |
| document_retriever: RetrievalRouter | None = None, | |
| *, | |
| slow_path_coordinator_factory: ( | |
| Callable[[str], SlowPathCoordinator] | None | |
| ) = None, | |
| analysis_store: ReportInputStore | None = None, | |
| traceability_store: TraceabilityStore | None = None, | |
| check_invoker_factory: Callable[[str], Any] | None = None, | |
| ps_agent: ProblemStatementAgent | None = None, | |
| help_agent: HelpAgent | None = None, | |
| state_store: Any | None = None, | |
| input_guard: InputGuard | None = None, | |
| enable_gate: bool = False, | |
| enable_tracing: bool = False, | |
| ) -> None: | |
| self._intent_router = intent_router | |
| self._answer_agent = answer_agent | |
| self._catalog_reader = catalog_reader | |
| self._document_retriever = document_retriever | |
| # Langfuse tracing (tokens + latency). OFF by default so tests never hit | |
| # Langfuse; the live endpoint opts in with ChatHandler(enable_tracing=True). | |
| self._enable_tracing = enable_tracing | |
| # Slow analytical path (Planner -> TaskRunner -> Assembler): the only route for | |
| # `structured_flow` now (the ENABLE_SLOW_PATH flag was removed 2026-07-02). The | |
| # factory + store are injectable for tests. | |
| self._slow_path_factory = slow_path_coordinator_factory | |
| self._analysis_store = analysis_store | |
| # Traceability (KM-691): user-facing per-turn provenance store. Injectable for | |
| # tests; lazily built (Postgres) in production. Distinct from Langfuse tracing. | |
| self._traceability_store = traceability_store | |
| # `check` skill: builds the data-access invoker (check_data/check_knowledge) | |
| # per request with the authenticated user_id. Injectable for tests. | |
| self._check_invoker_factory = check_invoker_factory | |
| # `problem_statement` skill: LLM drafter + the Analysis State store it writes | |
| # `problem_validated` to. Both injectable for tests. | |
| self._ps_agent = ps_agent | |
| # `help` skill: LLM guide that reads the Analysis State + chat history. | |
| self._help_agent = help_agent | |
| self._state_store = state_store | |
| # Input guard: screens each message for prompt-injection / secret-extraction / | |
| # abuse BEFORE the router. Injectable for tests; lazily built in production. | |
| self._input_guard = input_guard | |
| # Deterministic gate — DEPRECATED 2026-06-24 (problem_validated gate removed). | |
| # Unused flag; the gate call site in handle() is commented out. | |
| self._enable_gate = enable_gate | |
| # ------------------------------------------------------------------ | |
| # Lazy default-dep builders | |
| # ------------------------------------------------------------------ | |
| def _get_intent_router(self) -> OrchestratorAgent: | |
| if self._intent_router is None: | |
| self._intent_router = OrchestratorAgent() | |
| return self._intent_router | |
| def _get_input_guard(self) -> InputGuard: | |
| if self._input_guard is None: | |
| self._input_guard = InputGuard() | |
| return self._input_guard | |
| def _get_answer_agent(self) -> ChatbotAgent: | |
| if self._answer_agent is None: | |
| self._answer_agent = ChatbotAgent() | |
| return self._answer_agent | |
| def _get_catalog_reader(self) -> CatalogReader: | |
| if self._catalog_reader is None: | |
| from ..catalog.reader import CatalogReader | |
| from ..catalog.store import CatalogStore | |
| self._catalog_reader = CatalogReader(CatalogStore()) | |
| return self._catalog_reader | |
| def _get_document_retriever(self) -> RetrievalRouter: | |
| if self._document_retriever is None: | |
| from ..retrieval.router import RetrievalRouter | |
| self._document_retriever = RetrievalRouter() | |
| return self._document_retriever | |
| def _get_check_invoker(self, user_id: str, catalog_reader: Any = None) -> Any: | |
| """Build the per-request data-access invoker for the `check` skill. | |
| `catalog_reader` lets the caller scope the read (e.g. to the analysis | |
| catalog); defaults to the user-scope reader. | |
| """ | |
| if self._check_invoker_factory is not None: | |
| return self._check_invoker_factory(user_id) | |
| from ..tools.data_access import DataAccessToolInvoker | |
| return DataAccessToolInvoker(user_id, catalog_reader or self._get_catalog_reader()) | |
| def _get_ps_agent(self) -> ProblemStatementAgent: | |
| if self._ps_agent is None: | |
| self._ps_agent = ProblemStatementAgent() | |
| return self._ps_agent | |
| def _get_help_agent(self) -> HelpAgent: | |
| if self._help_agent is None: | |
| self._help_agent = HelpAgent() | |
| return self._help_agent | |
| def _get_state_store(self) -> Any: | |
| if self._state_store is None: | |
| from .state_store import AnalysisStateStore | |
| self._state_store = AnalysisStateStore() | |
| return self._state_store | |
| async def _load_analysis_state(self, analysis_id: str | None) -> AnalysisState: | |
| """Load Analysis State for the Help skill; fail closed to a not-validated stub. | |
| Mirrors the gate's never-throw fallback so Help degrades gracefully on a | |
| missing row, a read error, or a legacy room with no `analysis_id`. | |
| """ | |
| from .gate import stub_analysis_state | |
| if not analysis_id: | |
| return stub_analysis_state() | |
| try: | |
| state = await self._get_state_store().get(analysis_id) | |
| except Exception as e: | |
| logger.warning("help state read failed — not-validated", error=str(e)) | |
| state = None | |
| return state if state is not None else stub_analysis_state() | |
| # ------------------------------------------------------------------ | |
| # Public entry | |
| # ------------------------------------------------------------------ | |
| async def stream_help( | |
| self, | |
| user_id: str, | |
| analysis_id: str | None, | |
| history: list[BaseMessage] | None = None, | |
| message: str | None = None, | |
| message_id: str | None = None, | |
| ) -> AsyncIterator[dict[str, Any]]: | |
| """Deterministic `help` dispatch for the dedicated `/api/v1/tools/help` endpoint. | |
| Bypasses the intent router — the slash command IS the intent, so there is no | |
| classify round-trip and no misclassification risk. Streams the same guidance as | |
| the `help` branch of `handle()`, reusing the warm HelpAgent + state store. | |
| Emits SSE-style events: `sources` (always `[]` — help never references | |
| documents), `chunk`*, then `done` (data left empty; the endpoint stamps the | |
| `message_id`). On failure, yields a terminal `error` event. | |
| """ | |
| # Traceability (KM-691): a help turn has no planning/tools/sources — an | |
| # empty payload stamped `help`, flushed before `done`. | |
| pad = TraceabilityScratchpad() | |
| pad.message_id = message_id | |
| pad.set_intent("help") | |
| # Load (or lazily create) the analysis state; fail closed to a not-validated | |
| # stub so help degrades gracefully on a missing row / read error / legacy id. | |
| state: AnalysisState | None = None | |
| if analysis_id: | |
| try: | |
| state = await self._get_state_store().ensure(analysis_id, user_id) | |
| except Exception as e: # noqa: BLE001 — never block help on a state read | |
| logger.warning("help state ensure failed", analysis_id=analysis_id, error=str(e)) | |
| if state is None: | |
| state = await self._load_analysis_state(analysis_id) | |
| # report_ready (seam #5): deterministic, never-throws (fails closed to | |
| # not-ready) — the HelpAgent guard only offers generate_report when ready. | |
| from .report.readiness import is_report_ready | |
| report_ready = await is_report_ready(analysis_id, state) | |
| yield {"event": "sources", "data": json.dumps([])} | |
| try: | |
| async for token in self._get_help_agent().astream( | |
| state, | |
| history=history, | |
| message=message, | |
| report_ready=report_ready, | |
| ): | |
| yield {"event": "chunk", "data": token} | |
| except Exception as e: # noqa: BLE001 | |
| logger.error("help streaming failed", user_id=user_id, error=str(e)) | |
| yield {"event": "error", "data": f"Help generation failed: {e}"} | |
| return | |
| await self._flush_trace(pad, analysis_id, user_id) | |
| yield {"event": "done", "data": ""} | |
| async def handle( | |
| self, | |
| message: str, | |
| user_id: str, | |
| history: list[BaseMessage] | None = None, | |
| analysis_id: str | None = None, | |
| message_id: str | None = None, | |
| ) -> AsyncIterator[dict[str, Any]]: | |
| tracer = self._make_tracer(user_id, message) | |
| # Traceability (KM-691): per-request accumulator, flushed before EVERY `done` | |
| # (§5 matrix). Default intent `chat` until the router classifies; the two | |
| # refusal branches below stamp `blocked` explicitly. | |
| pad = TraceabilityScratchpad() | |
| pad.message_id = message_id | |
| # ---- 0. Input guard ------------------------------------------ | |
| # Deliberate input-filtering layer BEFORE the router: screen for prompt- | |
| # injection / secret-extraction / abuse. Fail-open on a guard *error* (never | |
| # take chat down); fail-closed on a positive detection → canned refusal, no | |
| # router, no answer. Benign off-topic messages pass here and are refused at | |
| # the `out_of_scope` branch below instead. | |
| gc = tracer.callbacks() # PII-safe, full capture (same policy as the router) | |
| gkw = {"callbacks": gc} if gc else {} | |
| verdict = await self._get_input_guard().screen(message, **gkw) | |
| if not verdict.allow: | |
| logger.info( | |
| "input guard blocked", user_id=user_id, category=verdict.category | |
| ) | |
| yield {"event": "sources", "data": json.dumps([])} | |
| yield {"event": "chunk", "data": blocked_message(message)} | |
| tracer.end() | |
| pad.set_intent("blocked") | |
| await self._flush_trace(pad, analysis_id, user_id) | |
| yield {"event": "done", "data": ""} | |
| return | |
| # ---- 1. Classify intent -------------------------------------- | |
| try: | |
| oc = tracer.callbacks() # orchestrator: PII-safe, full capture | |
| ckw = {"callbacks": oc} if oc else {} | |
| decision = await self._get_intent_router().classify(message, history, **ckw) | |
| except Exception as e: | |
| # Azure's own content filter (jailbreak detection) surfaces here as a 400. | |
| # Return a clean refusal instead of leaking the raw Azure error blob. | |
| if _is_content_filter_error(e): | |
| logger.info("router blocked by content filter", user_id=user_id) | |
| yield {"event": "sources", "data": json.dumps([])} | |
| yield {"event": "chunk", "data": blocked_message(message)} | |
| tracer.end() | |
| pad.set_intent("blocked") | |
| await self._flush_trace(pad, analysis_id, user_id) | |
| yield {"event": "done", "data": ""} | |
| return | |
| logger.error("intent classification failed", error=repr(e)) | |
| yield { | |
| "event": "error", | |
| "data": "Sorry, I couldn't process that message. Please try rephrasing.", | |
| } | |
| return | |
| intent = decision.intent | |
| pad.set_intent(intent) # traceability: chat | check | *_flow | out_of_scope | help | |
| # ---- 1a. Ensure session state row (T-A) ---------------------- | |
| # Rooms created via /room/create have no `analysis` row. Without one, Help and | |
| # the report_id write-back silently no-op. Lazily get-or-create it (idempotent). | |
| analysis_state: AnalysisState | None = None | |
| if analysis_id: | |
| try: | |
| analysis_state = await self._get_state_store().ensure(analysis_id, user_id) | |
| except Exception as e: | |
| logger.warning( | |
| "analysis state ensure failed", analysis_id=analysis_id, error=str(e) | |
| ) | |
| # ---- 1b. Gate (REMOVED 2026-06-24) --------------------------- | |
| # The problem_validated gate was dropped: structured_flow is no longer | |
| # redirected to problem_statement (the goal is now user-entered objective + | |
| # business_questions, no agent validation). `gate()` is neutered to a no-op; the | |
| # call site is left commented for restorability. | |
| # if self._enable_gate and analysis_id: | |
| # from .gate import gate, stub_analysis_state | |
| # | |
| # intent = gate( | |
| # intent, | |
| # analysis_state | |
| # if analysis_state is not None | |
| # else stub_analysis_state(), | |
| # ) | |
| # The `intent` event is consumed by the endpoint (it gates response caching | |
| # on the effective intent) and is NOT forwarded to the frontend. We emit the | |
| # post-gate intent so the cache keys on what actually ran. | |
| event_data = decision.model_dump() | |
| event_data["intent"] = intent | |
| yield {"event": "intent", "data": json.dumps(event_data)} | |
| rewritten = decision.rewritten_query or message | |
| query_result = None | |
| chunks: list[DocumentChunk] | None = None | |
| raw_chunks: Any = None | |
| # ---- 2. Route ------------------------------------------------ | |
| if intent == "out_of_scope": | |
| # Off-topic or manipulation the router flagged: canned refusal, no LLM, | |
| # no data lookup. (Malicious injections are usually stopped earlier by the | |
| # input guard; this catches benign off-topic + anything the guard let by.) | |
| yield {"event": "sources", "data": json.dumps([])} | |
| yield {"event": "chunk", "data": out_of_scope_message(message)} | |
| tracer.end() | |
| await self._flush_trace(pad, analysis_id, user_id) | |
| yield {"event": "done", "data": ""} | |
| return | |
| if intent == "structured_flow": | |
| try: | |
| # One memoizing reader per request: the same catalog is otherwise | |
| # re-fetched from the catalog DB 4-5x across the slow-path run. This | |
| # collapses those to one round-trip per source_hint and pins a single | |
| # consistent snapshot for plan + execution. | |
| from ..catalog.reader import ( | |
| AnalysisScopedCatalogReader, | |
| MemoizingCatalogReader, | |
| ) | |
| # Scope every catalog read — the Planner's AND the data-access tools' | |
| # own re-reads — to the analysis-scope catalog: Go materializes it with | |
| # exactly this analysis's bound db + file sources under their real | |
| # names. Falls back to the user-scope catalog when no analysis row | |
| # exists. Memoized so plan + execution share one snapshot. | |
| scoped = AnalysisScopedCatalogReader( | |
| self._get_catalog_reader(), analysis_id | |
| ) | |
| reader = MemoizingCatalogReader(scoped) | |
| catalog = await reader.read(user_id, "structured") | |
| # structured_flow always runs the slow analytical path (the | |
| # ENABLE_SLOW_PATH flag was removed 2026-07-02). | |
| # Detect reply language from the ORIGINAL message (not `rewritten` — the | |
| # router's rewritten_query is often normalized to English, which would | |
| # make the assembled answer English for an Indonesian question). | |
| reply_language = detect_reply_language(history, message=message) | |
| async for event in self._run_slow_path( | |
| user_id, rewritten, catalog, tracer, reader, analysis_id, | |
| reply_language, pad, | |
| ): | |
| yield event | |
| return | |
| except Exception as e: | |
| logger.error( | |
| "structured route failed", | |
| user_id=user_id, | |
| error=str(e), | |
| ) | |
| yield {"event": "error", "data": f"Structured query failed: {e}"} | |
| return | |
| elif intent == "unstructured_flow": | |
| try: | |
| raw_chunks = await self._get_document_retriever().retrieve( | |
| rewritten, user_id | |
| ) | |
| chunks = _normalize_chunks(raw_chunks) | |
| # Traceability (KM-691): retrieval bypasses the tool invoker, so synth | |
| # the retrieve_knowledge call (input = rewritten query) + document sources. | |
| pad.record_tool_call( | |
| "retrieve_knowledge", | |
| {"query": rewritten}, | |
| {"kind": "documents", "row_count": len(raw_chunks or [])}, | |
| ) | |
| pad.add_document_sources(raw_chunks, rewritten) | |
| except Exception as e: | |
| logger.error( | |
| "unstructured route failed", user_id=user_id, error=str(e) | |
| ) | |
| yield {"event": "error", "data": f"Document retrieval failed: {e}"} | |
| return | |
| elif intent == "check": | |
| try: | |
| # Scope check to the analysis catalog: it holds only this room's | |
| # bound sources and their real names (a DB shows as "xl test", not | |
| # the user-scope `postgres_<hash>` placeholder). Falls back to the | |
| # user-scope reader when the analysis has no catalog row. | |
| from ..catalog.reader import AnalysisScopedCatalogReader | |
| scoped_reader = AnalysisScopedCatalogReader( | |
| self._get_catalog_reader(), analysis_id | |
| ) | |
| # Wrap the check invoker so its check_* tool calls land in the trace. | |
| invoker = TraceabilityToolInvoker( | |
| self._get_check_invoker(user_id, scoped_reader), pad | |
| ) | |
| # Detect from the ORIGINAL message (not `rewritten`, which the | |
| # router normalizes to English) so the deterministic check reply | |
| # matches the user's language like the other paths. | |
| reply_language = detect_reply_language(history, message=message) | |
| text = await run_check(rewritten, invoker, reply_language) | |
| except Exception as e: | |
| logger.error("check route failed", user_id=user_id, error=str(e)) | |
| yield {"event": "error", "data": f"Lookup failed: {e}"} | |
| return | |
| yield {"event": "chunk", "data": text} | |
| await self._flush_trace(pad, analysis_id, user_id) | |
| yield {"event": "done", "data": ""} | |
| return | |
| # problem_statement dispatch removed 2026-06-24 (skill unwired; intent no longer | |
| # emitted by the router). Branch kept commented for restorability. | |
| # elif intent == "problem_statement": | |
| # try: | |
| # text = await run_problem_statement( | |
| # message, | |
| # analysis_id, | |
| # agent=self._get_ps_agent(), | |
| # store=self._get_state_store(), | |
| # history=history, | |
| # ) | |
| # except Exception as e: | |
| # logger.error("problem_statement route failed", user_id=user_id, error=str(e)) | |
| # yield {"event": "error", "data": f"Problem statement failed: {e}"} | |
| # return | |
| # yield {"event": "chunk", "data": text} | |
| # yield {"event": "done", "data": ""} | |
| # return | |
| elif intent == "help": | |
| try: | |
| state = analysis_state or await self._load_analysis_state(analysis_id) | |
| except Exception as e: | |
| logger.error("help route failed", user_id=user_id, error=str(e)) | |
| yield {"event": "error", "data": f"Help failed: {e}"} | |
| return | |
| # report_ready (seam #5): deterministic — validated goal + ≥1 recorded | |
| # analysis (mirrors the report API's own 409 gate). Never-throws (fails | |
| # closed to not-ready), so Help degrades safely. The consistency guard in | |
| # HelpAgent only offers `generate_report` when this says ready. | |
| from .report.readiness import is_report_ready | |
| report_ready = await is_report_ready(analysis_id, state) | |
| # The prompt sees chat history -> masked. | |
| hc = tracer.callbacks(masked=True) | |
| hkw = {"callbacks": hc} if hc else {} | |
| try: | |
| async for token in self._get_help_agent().astream( | |
| state, | |
| history=history, | |
| message=message, | |
| report_ready=report_ready, | |
| **hkw, | |
| ): | |
| yield {"event": "chunk", "data": token} | |
| except Exception as e: | |
| logger.error("help streaming failed", user_id=user_id, error=str(e)) | |
| yield {"event": "error", "data": f"Help generation failed: {e}"} | |
| return | |
| tracer.end() | |
| await self._flush_trace(pad, analysis_id, user_id) | |
| yield {"event": "done", "data": ""} | |
| return | |
| # else: chat path — no context | |
| # ---- 2b. Emit sources --------------------------------------- | |
| # Sources moved to traceability (KM-691): the stream stays text-only. The FE | |
| # reads sources from GET /api/v1/traceability (richer + intent-consistent; the | |
| # document sources for this turn are captured on the pad above). The empty | |
| # `sources` event is kept for SSE backward-compat. | |
| yield {"event": "sources", "data": json.dumps([])} | |
| # ---- 3. Stream answer ---------------------------------------- | |
| # masked: the answer call sees real query rows / doc chunks (possible PII). | |
| mc = tracer.callbacks(masked=True) | |
| akw = {"callbacks": mc} if mc else {} | |
| try: | |
| async for token in self._get_answer_agent().astream( | |
| message, | |
| history=history, | |
| query_result=query_result, | |
| chunks=chunks, | |
| **akw, | |
| ): | |
| yield {"event": "chunk", "data": token} | |
| except Exception as e: | |
| logger.error("answer streaming failed", user_id=user_id, error=str(e)) | |
| yield {"event": "error", "data": f"Answer generation failed: {e}"} | |
| return | |
| tracer.end() | |
| # chat: empty payload; unstructured_flow: synth retrieve_knowledge + doc sources. | |
| await self._flush_trace(pad, analysis_id, user_id) | |
| yield {"event": "done", "data": ""} | |
| # ------------------------------------------------------------------ | |
| # Slow analytical path (gated, off by default) | |
| # ------------------------------------------------------------------ | |
| def _make_tracer(self, user_id: str, question: str) -> Any: | |
| """One Langfuse trace per request (or a NullTracer when disabled).""" | |
| if not self._enable_tracing: | |
| from ..observability.langfuse.tracing import NullTracer | |
| return NullTracer() | |
| from ..observability.langfuse.tracing import RequestTracer | |
| return RequestTracer.start(user_id=user_id, question=question) | |
| def _get_slow_path_coordinator( | |
| self, | |
| user_id: str, | |
| tracer: Any = None, | |
| catalog_reader: CatalogReader | None = None, | |
| pad: TraceabilityScratchpad | None = None, | |
| ) -> SlowPathCoordinator: | |
| """Build the per-request slow-path coordinator (composition root). | |
| The data-access tools need the authenticated `user_id` + `CatalogReader`, | |
| so the `CompositeToolInvoker` is constructed per request. The slow-path | |
| agent code stays tool-agnostic (INV-7) — only here, the composition root, | |
| do we name concrete tool implementations. When tracing is active the invoker | |
| is wrapped so each tool call records a metadata-only span. | |
| """ | |
| if self._slow_path_factory is not None: | |
| return self._slow_path_factory(user_id) | |
| from ..tools.data_access import DataAccessToolInvoker | |
| from ..tools.invoker import AnalyticsToolInvoker, CompositeToolInvoker | |
| from .planner.registry import default_registry | |
| from .planner.service import PlannerService | |
| from .slow_path.assembler import Assembler | |
| from .slow_path.coordinator import SlowPathCoordinator | |
| from .slow_path.task_runner import TaskRunner | |
| invoker: Any = CompositeToolInvoker( | |
| DataAccessToolInvoker(user_id, catalog_reader or self._get_catalog_reader()), | |
| AnalyticsToolInvoker(), | |
| ) | |
| if tracer is not None and getattr(tracer, "active", False): | |
| from ..observability.langfuse.tracing import TracingToolInvoker | |
| invoker = TracingToolInvoker(invoker, tracer) | |
| # Traceability outermost: records the SAME real I/O the tools return (both | |
| # wrappers see it; order is immaterial). KM-691. | |
| if pad is not None: | |
| invoker = TraceabilityToolInvoker(invoker, pad) | |
| registry = default_registry() | |
| return SlowPathCoordinator( | |
| PlannerService(), TaskRunner(invoker, registry), Assembler(), registry | |
| ) | |
| def _get_analysis_store(self) -> ReportInputStore: | |
| if self._analysis_store is None: | |
| from .slow_path.store import PostgresReportInputStore | |
| self._analysis_store = PostgresReportInputStore() | |
| return self._analysis_store | |
| def _get_traceability_store(self) -> TraceabilityStore: | |
| if self._traceability_store is None: | |
| from ..traceability import PostgresTraceabilityStore | |
| self._traceability_store = PostgresTraceabilityStore() | |
| return self._traceability_store | |
| async def _flush_trace( | |
| self, pad: TraceabilityScratchpad, analysis_id: str | None, user_id: str | |
| ) -> None: | |
| """Persist the turn's traceability row right before `done` (KM-691). | |
| No-op without a `message_id` (tests / unwired callers never flush). The | |
| store save is itself never-throw, but we also guard here: a trace failure | |
| must never break — or delay past `done` — the user's answer. | |
| """ | |
| if pad.message_id is None: | |
| return | |
| try: | |
| payload = pad.build(analysis_id or "", user_id, pad.message_id) | |
| await self._get_traceability_store().save(payload) | |
| except Exception as e: # noqa: BLE001 — never break the answer on a trace slip | |
| logger.warning("traceability flush failed", error=str(e)) | |
| async def _run_slow_path( | |
| self, | |
| user_id: str, | |
| query: str, | |
| catalog: Any, | |
| tracer: Any = None, | |
| catalog_reader: CatalogReader | None = None, | |
| analysis_id: str | None = None, | |
| reply_language: str | None = None, | |
| pad: TraceabilityScratchpad | None = None, | |
| ) -> AsyncIterator[dict[str, Any]]: | |
| """Run the slow path and stream its assembled answer as SSE events. | |
| Context comes from the `get_business_context` seam (a stub today); the | |
| `analysis_record` is persisted via the `ReportInputStore` seam (PostgresReportInputStore), | |
| stamped with the request's user_id + analysis_id so the report can group it. | |
| `chat_answer` is emitted as a single `chunk` (the Assembler returns the whole | |
| object — true token streaming is a later step). | |
| """ | |
| from .planner.business_context import get_business_context | |
| from .planner.inputs import Constraints | |
| if tracer is None: | |
| from ..observability.langfuse.tracing import NullTracer | |
| tracer = NullTracer() | |
| coordinator = self._get_slow_path_coordinator(user_id, tracer, catalog_reader, pad) | |
| context = await get_business_context(user_id) | |
| # DB3: warm the user's DB connection in parallel with planning so the | |
| # handshake overlaps the ~4s Planner call. Default path only — an injected | |
| # coordinator factory (tests / custom) may not use the real DbExecutor. | |
| if self._slow_path_factory is None: | |
| from ..query.executor.db import DbExecutor | |
| asyncio.create_task(DbExecutor.prewarm(catalog, user_id)) # noqa: RUF006 | |
| pc = tracer.callbacks() # planner: PII-safe, full capture | |
| ac = tracer.callbacks(masked=True) # assembler: sees real rows -> masked | |
| run_kw: dict[str, Any] = {} | |
| if pc: | |
| run_kw["planner_callbacks"] = pc | |
| if ac: | |
| run_kw["assembler_callbacks"] = ac | |
| # R4: bridge the coordinator's per-stage progress callback to SSE `status` | |
| # events so the stream isn't silent for ~12s (and proxies don't drop the | |
| # idle connection). Status events only appear if the coordinator calls back. | |
| progress_q: asyncio.Queue[str] = asyncio.Queue() | |
| async def _progress(stage: str) -> None: | |
| await progress_q.put(stage) | |
| run_task = asyncio.create_task( | |
| coordinator.run( | |
| context, catalog, query, Constraints(), | |
| progress=_progress, reply_language=reply_language, **run_kw | |
| ) | |
| ) | |
| getter: asyncio.Task = asyncio.create_task(progress_q.get()) | |
| pending: set[asyncio.Task] = {run_task, getter} | |
| while True: | |
| done, pending = await asyncio.wait( | |
| pending, return_when=asyncio.FIRST_COMPLETED | |
| ) | |
| if getter in done: | |
| yield {"event": "status", "data": getter.result()} | |
| getter = asyncio.create_task(progress_q.get()) | |
| pending = pending | {getter} | |
| if run_task in done: | |
| getter.cancel() | |
| while not progress_q.empty(): | |
| yield {"event": "status", "data": progress_q.get_nowait()} | |
| break | |
| try: | |
| result = run_task.result() | |
| except Exception as e: | |
| logger.error("slow path failed", user_id=user_id, error=str(e)) | |
| yield {"event": "error", "data": f"Analysis failed: {e}"} | |
| return | |
| # Sources live in traceability now (KM-691), derived from the run's | |
| # retrieve_data calls; the stream stays text-only. | |
| yield {"event": "sources", "data": json.dumps([])} | |
| yield {"event": "chunk", "data": result.chat_answer} | |
| try: | |
| # Stamp identity from the request scope: owner + the shared session id | |
| # (analysis_id == room_id). Without analysis_id the record is orphaned — | |
| # list_for_analysis can't find it, so the report + is_report_ready go | |
| # blind. The store is never-throw. | |
| record = result.analysis_record.model_copy( | |
| update={"user_id": user_id, "analysis_id": analysis_id} | |
| ) | |
| await self._get_analysis_store().save(record) | |
| if pad is not None: | |
| # Traceability planning = goal_restated + tasks_run from the record; | |
| # tool_calls were already recorded by the wrapped invoker. | |
| pad.set_planning_from_record(record) | |
| except Exception as e: # persistence must never break the user's answer | |
| logger.error("analysis_record persist failed", user_id=user_id, error=str(e)) | |
| tracer.end() # output omitted (chat_answer may contain PII on Cloud) | |
| if pad is not None: | |
| await self._flush_trace(pad, analysis_id, user_id) | |
| yield {"event": "done", "data": ""} | |
| def _normalize_chunks(raw: Any) -> list[DocumentChunk]: | |
| """Convert whatever the retriever returns into list[DocumentChunk]. | |
| The Phase 2 `DocumentRetriever.retrieve` interface is a stub today; | |
| when TAB owner ships it, it should return `list[DocumentChunk]` | |
| directly so this normalizer becomes a no-op. Until then we coerce | |
| common shapes (dict-with-content, plain string) defensively. | |
| """ | |
| if not raw: | |
| return [] | |
| if isinstance(raw, list) and all(isinstance(c, DocumentChunk) for c in raw): | |
| return raw | |
| chunks: list[DocumentChunk] = [] | |
| for item in raw: | |
| if isinstance(item, DocumentChunk): | |
| chunks.append(item) | |
| elif isinstance(item, dict): | |
| chunks.append( | |
| DocumentChunk( | |
| content=str(item.get("content", "")), | |
| filename=item.get("filename"), | |
| page_label=item.get("page_label"), | |
| ) | |
| ) | |
| elif isinstance(item, RetrievalResult): | |
| data = item.metadata.get("data", {}) | |
| page = data.get("page_label") | |
| chunks.append(DocumentChunk( | |
| content=item.content, | |
| filename=data.get("filename"), | |
| page_label=str(page) if page is not None else None, | |
| )) | |
| elif isinstance(item, str): | |
| chunks.append(DocumentChunk(content=item)) | |
| return chunks | |