| """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)}, |
| ) |
|
|
| def end(self, out: Any) -> None: |
| with contextlib.suppress(Exception): |
| 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, |
| ) |
| return cls(trace) |
| except Exception as e: |
| 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: |
| |
| 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 |
|
|
