src/observability/langfuse/tracing.py

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