[KM-631][AI] Langfuse tracing: tokens + latency across the chat pipeline

Wire Langfuse (was scaffolded-only) so each chat request emits one trace with
per-call token counts + latency, plus tool spans. Gated and PII-aware.

- New src/observability/langfuse/tracing.py: RequestTracer (one trace/request),
NullTracer (no-op), TracingToolInvoker (metadata-only span per tool call),
_redact MaskFunction. Best-effort: never raises; pipeline unchanged if Langfuse
is down or disabled.
- Each LLM call (Orchestrator, Planner incl. retries, Assembler, Chatbot) gained an
optional `callbacks` param, passed to the LangChain call only when present (so
existing fakes/tests are untouched). Coordinator threads planner/assembler
callbacks through.
- ChatHandler: enable_tracing flag (default OFF so tests never hit Langfuse), creates
one trace per request, wraps the tool invoker for spans, ends the trace.
- PII policy for Cloud: UNMASKED = Orchestrator + Planner (PII-safe inputs);
MASKED (tokens+latency only) = Assembler + Chatbot (see real rows / doc chunks);
tool spans carry counts + status only, never rows.
- chatbot: stream_options include_usage so the streamed answer reports tokens too.
- api/v1/chat.py: activate tracing (ChatHandler(enable_tracing=True)).

Zero added LLM token cost (only reads usage the API returns); negligible latency
(SDK sends in a background thread). Verified live end-to-end against Langfuse Cloud.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

src/agents/chat_handler.py

@@ -70,12 +70,16 @@ class ChatHandler:
             Callable[[str], SlowPathCoordinator] | None
         ) = None,
         analysis_store: AnalysisStore | None = None,
+        enable_tracing: bool = False,
     ) -> None:
         self._intent_router = intent_router
         self._answer_agent = answer_agent
         self._catalog_reader = catalog_reader
         self._query_service = query_service
         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). OFF by default:
         # gated until the lead's real BusinessContext lands. When True, `structured`
         # intents route here instead of the single-query QueryService path. The
@@ -130,9 +134,13 @@ class ChatHandler:
         user_id: str,
         history: list[BaseMessage] | None = None,
     ) -> AsyncIterator[dict[str, Any]]:
+        tracer = self._make_tracer(user_id, message)
+
         # ---- 1. Classify intent --------------------------------------
         try:
-            decision = await self._get_intent_router().classify(message, history)
+            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:
             logger.error("intent classification failed", error=str(e))
             yield {"event": "error", "data": f"Could not classify message: {e}"}
@@ -150,7 +158,9 @@ class ChatHandler:
             try:
                 catalog = await self._get_catalog_reader().read(user_id, "structured")
                 if self._enable_slow_path:
-                    async for event in self._run_slow_path(user_id, rewritten, catalog):
+                    async for event in self._run_slow_path(
+                        user_id, rewritten, catalog, tracer
+                    ):
                         yield event
                     return
                 query_result = await self._get_query_service().run(
@@ -193,12 +203,16 @@ class ChatHandler:
         yield {"event": "sources", "data": json.dumps(sources)}
 
         # ---- 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:
@@ -206,19 +220,33 @@ class ChatHandler:
             yield {"event": "error", "data": f"Answer generation failed: {e}"}
             return
 
+        tracer.end()
         yield {"event": "done", "data": ""}
 
     # ------------------------------------------------------------------
     # Slow analytical path (gated, off by default)
     # ------------------------------------------------------------------
 
-    def _get_slow_path_coordinator(self, user_id: str) -> SlowPathCoordinator:
+    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
+    ) -> 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.
+        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)
@@ -231,10 +259,14 @@ class ChatHandler:
         from .slow_path.coordinator import SlowPathCoordinator
         from .slow_path.task_runner import TaskRunner
 
-        invoker = CompositeToolInvoker(
+        invoker: Any = CompositeToolInvoker(
             DataAccessToolInvoker(user_id, 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)
         registry = default_registry()
         return SlowPathCoordinator(
             PlannerService(), TaskRunner(invoker, registry), Assembler(), registry
@@ -252,6 +284,7 @@ class ChatHandler:
         user_id: str,
         query: str,
         catalog: Any,
+        tracer: Any = None,
     ) -> AsyncIterator[dict[str, Any]]:
         """Run the slow path and stream its assembled answer as SSE events.
 
@@ -263,10 +296,22 @@ class ChatHandler:
         from .planner.business_context import get_business_context
         from .planner.inputs import Constraints
 
-        coordinator = self._get_slow_path_coordinator(user_id)
+        if tracer is None:
+            from ..observability.langfuse.tracing import NullTracer
+
+            tracer = NullTracer()
+
+        coordinator = self._get_slow_path_coordinator(user_id, tracer)
         context = await get_business_context(user_id)
+        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
         try:
-            result = await coordinator.run(context, catalog, query, Constraints())
+            result = await coordinator.run(context, catalog, query, Constraints(), **run_kw)
         except Exception as e:
             logger.error("slow path failed", user_id=user_id, error=str(e))
             yield {"event": "error", "data": f"Analysis failed: {e}"}
@@ -278,6 +323,7 @@ class ChatHandler:
             await self._get_analysis_store().save(result.analysis_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)
         yield {"event": "done", "data": ""}
 
 

src/agents/chatbot.py

@@ -119,6 +119,9 @@ def _build_default_chain() -> Runnable:
         azure_endpoint=settings.azureai_endpoint_url_4o,
         api_key=settings.azureai_api_key_4o,
         temperature=0.3,
+        # Emit token usage on the final streamed chunk (this agent only streams), so
+        # the fast-path answer reports tokens to Langfuse like the non-streaming calls.
+        model_kwargs={"stream_options": {"include_usage": True}},
     )
     prompt = ChatPromptTemplate.from_messages(
         [
@@ -153,6 +156,7 @@ class ChatbotAgent:
         history: list[BaseMessage] | None = None,
         query_result: QueryResult | None = None,
         chunks: list[DocumentChunk] | None = None,
+        callbacks: list | None = None,
     ) -> AsyncIterator[str]:
         """Stream tokens of the final answer.
 
@@ -165,5 +169,9 @@ class ChatbotAgent:
             "history": history or [],
             "context": _build_context_block(query_result, chunks),
         }
-        async for token in chain.astream(payload):
-            yield token
+        if callbacks:
+            async for token in chain.astream(payload, config={"callbacks": callbacks}):
+                yield token
+        else:
+            async for token in chain.astream(payload):
+                yield token

src/agents/orchestration.py

@@ -96,11 +96,16 @@ class OrchestratorAgent:
         self,
         message: str,
         history: list[BaseMessage] | None = None,
+        callbacks: list | None = None,
     ) -> IntentRouterDecision:
         chain = self._ensure_chain()
-        decision: IntentRouterDecision = await chain.ainvoke(
-            {"message": message, "history": history or []}
-        )
+        payload = {"message": message, "history": history or []}
+        if callbacks:
+            decision: IntentRouterDecision = await chain.ainvoke(
+                payload, config={"callbacks": callbacks}
+            )
+        else:
+            decision = await chain.ainvoke(payload)
         logger.info(
             "intent classified",
             source_hint=decision.source_hint,

src/agents/planner/service.py

@@ -101,6 +101,7 @@ class PlannerService:
         tools: ToolRegistry,
         query: str,
         constraints: Constraints,
+        callbacks: list | None = None,
     ) -> TaskList:
         summary = CatalogSummary.from_catalog(catalog)
         chain = self._ensure_chain()
@@ -110,7 +111,14 @@ class PlannerService:
             human_content = build_planner_prompt(
                 context, summary, tools, query, constraints, previous_error
             )
-            task_list: TaskList = await chain.ainvoke({"human_content": human_content})
+            # All retry attempts share `callbacks`, so each shows up under the same
+            # trace — that is how retry token cost becomes visible.
+            if callbacks:
+                task_list: TaskList = await chain.ainvoke(
+                    {"human_content": human_content}, config={"callbacks": callbacks}
+                )
+            else:
+                task_list = await chain.ainvoke({"human_content": human_content})
             try:
                 self._validator.validate(task_list, tools, catalog, constraints)
             except PlannerValidationError as e:

src/agents/slow_path/assembler.py

@@ -92,13 +92,17 @@ class Assembler:
         run_state: RunState,
         context: BusinessContext,
         question: str | None = None,
+        callbacks: list | None = None,
     ) -> AssembledOutput:
         chain = self._ensure_chain()
         human_content = build_assembler_prompt(run_state, context, question)
         try:
-            narrative: AssemblerNarrative = await chain.ainvoke(
-                {"human_content": human_content}
-            )
+            if callbacks:
+                narrative: AssemblerNarrative = await chain.ainvoke(
+                    {"human_content": human_content}, config={"callbacks": callbacks}
+                )
+            else:
+                narrative = await chain.ainvoke({"human_content": human_content})
         except Exception as exc:  # surface as a typed error for the caller
             raise AssemblerError(f"assembler call failed: {exc}") from exc
 

src/agents/slow_path/coordinator.py

@@ -38,11 +38,17 @@ class SlowPathCoordinator:
         catalog: Catalog,
         query: str,
         constraints: Constraints,
+        planner_callbacks: list | None = None,
+        assembler_callbacks: list | None = None,
     ) -> AssembledOutput:
+        plan_kw = {"callbacks": planner_callbacks} if planner_callbacks else {}
         task_list = await self._planner.plan(
-            context, catalog, self._registry, query, constraints
+            context, catalog, self._registry, query, constraints, **plan_kw
         )
         run_state = await self._task_runner.run(
             task_list, business_context_id=context.project_id
         )
-        return await self._assembler.assemble(run_state, context, question=query)
+        asm_kw = {"callbacks": assembler_callbacks} if assembler_callbacks else {}
+        return await self._assembler.assemble(
+            run_state, context, question=query, **asm_kw
+        )

src/api/v1/chat.py

@@ -169,7 +169,7 @@ async def chat_stream(request: ChatRequest, db: AsyncSession = Depends(get_db)):
             return EventSourceResponse(stream_direct())
 
         history = await load_history(db, request.room_id, limit=10)
-        handler = ChatHandler()
+        handler = ChatHandler(enable_tracing=True)
 
         async def stream_response():
             logger.info("stream_response started", room_id=request.room_id, user_id=request.user_id)

src/observability/langfuse/tracing.py

@@ -0,0 +1,169 @@
+"""Langfuse request tracing — tokens + latency for the chat pipeline.
+
+One Langfuse trace per chat request. LangChain LLM calls attach a `CallbackHandler`
+(auto-captures prompt/completion tokens + latency); deterministic tool calls are
+recorded as metadata-only spans.
+
+PII policy for Langfuse **Cloud** (data leaves to Langfuse's servers):
+  - UNMASKED (full input/output): **Orchestrator + Planner** — their inputs are the
+    user question and a PII-safe `CatalogSummary` (sample values stripped by design).
+  - MASKED (tokens + latency only; input/output redacted): **Assembler + Chatbot** —
+    their inputs carry real query rows / document chunks that may contain PII.
+  - Tool spans carry only metadata (tool name, output kind, row COUNT, status) —
+    never the rows themselves.
+
+Everything here is best-effort and **never raises**: if Langfuse is unreachable or
+disabled, the chat pipeline runs unchanged. Tracing is created only when the caller
+opts in (ChatHandler(enable_tracing=True)); otherwise a `NullTracer` is used.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import functools
+import time
+from typing import Any
+
+from src.config.settings import settings
+from src.middlewares.logging import get_logger
+
+logger = get_logger("tracing")
+
+
+def _redact(*, data: Any) -> Any:
+    """Langfuse MaskFunction: drop the value entirely (used for PII-bearing calls)."""
+    return "<redacted: omitted from Langfuse (may contain user data)>"
+
+
+@functools.cache
+def _client() -> Any:
+    from langfuse import Langfuse
+
+    return Langfuse(
+        public_key=settings.LANGFUSE_PUBLIC_KEY,
+        secret_key=settings.LANGFUSE_SECRET_KEY,
+        host=settings.LANGFUSE_HOST,
+    )
+
+
+class _NullSpan:
+    def end(self, _out: Any) -> None: ...
+
+
+class NullTracer:
+    """No-op tracer (tracing disabled). Same surface as RequestTracer."""
+
+    active = False
+
+    def callbacks(self, *, masked: bool = False) -> list:
+        return []
+
+    def tool_span(self, tool: str, args: dict) -> Any:
+        return _NullSpan()
+
+    def end(self, *, output: Any = None) -> None: ...
+
+
+class _ToolSpan:
+    """A metadata-only span around one tool call. Never records row data."""
+
+    def __init__(self, trace: Any, tool: str, args: dict) -> None:
+        self._t0 = time.perf_counter()
+        self._span = trace.span(
+            name=f"tool:{tool}",
+            metadata={"tool": tool, "arg_keys": sorted(args)},  # keys only, no values
+        )
+
+    def end(self, out: Any) -> None:
+        with contextlib.suppress(Exception):  # never let a span break the run
+            kind = getattr(out, "kind", None)
+            is_err = kind == "error"
+            meta: dict[str, Any] = {
+                "kind": kind,
+                "elapsed_ms": round((time.perf_counter() - self._t0) * 1000),
+            }
+            if kind == "table":
+                meta["rows"] = len(getattr(out, "rows", None) or [])
+            err_msg = (getattr(out, "error", None) or "")[:300] if is_err else None
+            if err_msg:
+                meta["error"] = err_msg
+            self._span.end(
+                metadata=meta,
+                level="ERROR" if is_err else "DEFAULT",
+                status_message=err_msg,
+            )
+
+
+class RequestTracer:
+    """One Langfuse trace per chat request; hands out callbacks + tool spans."""
+
+    active = True
+
+    def __init__(self, trace: Any) -> None:
+        self._trace = trace
+
+    @classmethod
+    def start(
+        cls,
+        *,
+        user_id: str,
+        question: str | None = None,
+        session_id: str | None = None,
+    ) -> RequestTracer | NullTracer:
+        try:
+            trace = _client().trace(
+                name="chat_request",
+                user_id=user_id,
+                session_id=session_id,
+                input=question,  # the user's question (same exposure as Planner prompt)
+            )
+            return cls(trace)
+        except Exception as e:  # never let tracing break the request
+            logger.warning("tracing disabled (init failed)", error=str(e))
+            return NullTracer()
+
+    def callbacks(self, *, masked: bool = False) -> list:
+        """A LangChain callback nested under this trace. `masked=True` redacts the
+        call's input/output (tokens + latency are still captured)."""
+        try:
+            from langfuse.callback import CallbackHandler
+
+            return [
+                CallbackHandler(
+                    stateful_client=self._trace,
+                    mask=_redact if masked else None,
+                )
+            ]
+        except Exception as e:
+            logger.warning("tracing handler unavailable", error=str(e))
+            return []
+
+    def tool_span(self, tool: str, args: dict) -> Any:
+        try:
+            return _ToolSpan(self._trace, tool, args)
+        except Exception:
+            return _NullSpan()
+
+    def end(self, *, output: Any = None) -> None:
+        # Note: callers pass output=None on PII-bearing paths so no answer text is sent.
+        with contextlib.suppress(Exception):
+            if output is not None:
+                self._trace.update(output=output)
+
+
+class TracingToolInvoker:
+    """Wraps a ToolInvoker to record a metadata-only span per tool call.
+
+    Implements the ToolInvoker protocol; created at the composition root (ChatHandler)
+    so the slow-path agent code stays tool-agnostic and tracing-agnostic.
+    """
+
+    def __init__(self, inner: Any, tracer: RequestTracer) -> None:
+        self._inner = inner
+        self._tracer = tracer
+
+    async def invoke(self, tool_name: str, args: dict[str, Any]) -> Any:
+        span = self._tracer.tool_span(tool_name, args)
+        out = await self._inner.invoke(tool_name, args)
+        span.end(out)
+        return out