File size: 5,766 Bytes
81e5fe7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """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
|