439 lines
16 KiB
Python
439 lines
16 KiB
Python
# pyright: reportPrivateUsage=false, reportUnannotatedClassAttribute=false, reportUnusedFunction=false, reportUnknownLambdaType=false, reportUnknownArgumentType=false, reportArgumentType=false, reportIndexIssue=false
|
|
"""Tests for the Langfuse projection over the captured LLM stream.
|
|
|
|
Exercises `LangfuseExporter` with a fake Langfuse client so we can assert the
|
|
reconstructed trace tree (trace ids, parent linkage, names, usage, trace-level
|
|
user attributes, session-as-metadata) without a real Langfuse backend.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from src.config import settings
|
|
from src.llm.backend import CompletionResult, ToolCallResult
|
|
from src.llm.capture import build_captured_call
|
|
from src.llm.types import LLMTelemetryContext
|
|
from src.telemetry import langfuse_session
|
|
from src.telemetry.langfuse_exporter import LangfuseExporter
|
|
|
|
|
|
class FakeOtelSpan:
|
|
def __init__(self) -> None:
|
|
self.attributes: dict[str, object] = {}
|
|
|
|
def set_attribute(self, key: str, value: object) -> None:
|
|
self.attributes[key] = value
|
|
|
|
|
|
class FakeObs:
|
|
_counter = 0
|
|
|
|
def __init__(self, **kwargs: object) -> None:
|
|
FakeObs._counter += 1
|
|
self.id = f"obs-{FakeObs._counter}"
|
|
self.kwargs = kwargs
|
|
self._otel_span = FakeOtelSpan()
|
|
self.ended = False
|
|
|
|
def end(self) -> None:
|
|
self.ended = True
|
|
|
|
|
|
class FakeClient:
|
|
def __init__(self) -> None:
|
|
self.observations: list[FakeObs] = []
|
|
|
|
def create_trace_id(self, *, seed: str | None = None) -> str:
|
|
return f"lf-{seed}"
|
|
|
|
def start_observation(self, **kwargs: object) -> FakeObs:
|
|
obs = FakeObs(**kwargs)
|
|
self.observations.append(obs)
|
|
return obs
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _exporter_env(monkeypatch: pytest.MonkeyPatch):
|
|
"""Enable the exporter and install a fake langfuse client + clean registry."""
|
|
monkeypatch.setattr(settings, "LANGFUSE_PUBLIC_KEY", "pk-test")
|
|
monkeypatch.setattr(settings, "LANGFUSE_EXPORTER_MODE", "exporter")
|
|
monkeypatch.setattr(settings, "NAMESPACE", "tenant1")
|
|
client = FakeClient()
|
|
import langfuse
|
|
|
|
monkeypatch.setattr(langfuse, "get_client", lambda: client)
|
|
langfuse_session.reset()
|
|
FakeObs._counter = 0
|
|
yield client
|
|
langfuse_session.reset()
|
|
|
|
|
|
def _call(
|
|
*,
|
|
run_id: str | None,
|
|
trace_id: str,
|
|
iteration: int | None = None,
|
|
step_seq: int = 0,
|
|
attempt: int = 1,
|
|
session_id: str | None = None,
|
|
track_name: str | None = None,
|
|
agent_type: str = "dialectic",
|
|
parent_category: str = "dialectic",
|
|
tool_names: list[str] | None = None,
|
|
finish_reason: str = "stop",
|
|
content: str = "answer",
|
|
):
|
|
telemetry = LLMTelemetryContext(
|
|
workspace_name="ws",
|
|
call_purpose="dialectic.answer",
|
|
parent_category=parent_category,
|
|
agent_type=agent_type,
|
|
run_id=run_id,
|
|
trace_id=trace_id,
|
|
span_id=trace_id,
|
|
session_id=session_id,
|
|
track_name=track_name,
|
|
iteration=iteration,
|
|
step_seq=step_seq,
|
|
)
|
|
result = CompletionResult(
|
|
content=content,
|
|
input_tokens=10,
|
|
output_tokens=5,
|
|
cache_read_input_tokens=2,
|
|
finish_reason=finish_reason,
|
|
tool_calls=[
|
|
ToolCallResult(id=f"tc-{i}", name=name, input={"q": name})
|
|
for i, name in enumerate(tool_names or [])
|
|
],
|
|
)
|
|
return build_captured_call(
|
|
telemetry=telemetry,
|
|
transport="anthropic",
|
|
provider_label=None,
|
|
model="claude-x",
|
|
messages=[{"role": "user", "content": "q"}],
|
|
tools=None,
|
|
tool_choice=None,
|
|
result=result,
|
|
attempt=attempt,
|
|
was_fallback=False,
|
|
was_stream=False,
|
|
finish_reason=finish_reason,
|
|
)
|
|
|
|
|
|
def test_single_shot_generation_is_trace_root(_exporter_env: FakeClient):
|
|
# Deriver/summarizer style: run_id None → no run/step span, generation is root.
|
|
client = _exporter_env
|
|
LangfuseExporter().export(
|
|
_call(run_id=None, trace_id="t1", track_name="Minimal Deriver")
|
|
)
|
|
|
|
assert len(client.observations) == 1
|
|
gen = client.observations[0]
|
|
assert gen.kwargs["as_type"] == "generation"
|
|
assert gen.kwargs["trace_context"] == {"trace_id": "lf-t1"}
|
|
assert gen.kwargs["model"] == "claude-x"
|
|
assert gen.kwargs["usage_details"] == {
|
|
"input": 10,
|
|
"output": 5,
|
|
"cache_read_input_tokens": 2,
|
|
"cache_creation_input_tokens": 0,
|
|
}
|
|
# Trace attrs stamped on the root generation; no session (session_id None).
|
|
assert gen._otel_span.attributes.get("user.id") == "tenant1"
|
|
assert "session.id" not in gen._otel_span.attributes
|
|
|
|
|
|
def test_agentic_run_builds_run_step_generation(_exporter_env: FakeClient):
|
|
client = _exporter_env
|
|
LangfuseExporter().export(
|
|
_call(
|
|
run_id="r1",
|
|
trace_id="r1",
|
|
iteration=1,
|
|
session_id="sess_abc",
|
|
track_name="Dialectic Agent",
|
|
)
|
|
)
|
|
|
|
by_type: dict[str, list[FakeObs]] = {}
|
|
for obs in client.observations:
|
|
by_type.setdefault(str(obs.kwargs["as_type"]), []).append(obs)
|
|
assert len(by_type["span"]) == 2 # run span + step span
|
|
assert len(by_type["generation"]) == 1
|
|
|
|
run_span, step_span = by_type["span"]
|
|
gen = by_type["generation"][0]
|
|
assert run_span.kwargs["trace_context"] == {"trace_id": "lf-r1"}
|
|
assert step_span.kwargs["trace_context"] == {
|
|
"trace_id": "lf-r1",
|
|
"parent_span_id": run_span.id,
|
|
}
|
|
assert gen.kwargs["trace_context"] == {
|
|
"trace_id": "lf-r1",
|
|
"parent_span_id": step_span.id,
|
|
}
|
|
# Trace attrs stamped once, on the run span (the root). The Honcho session is
|
|
# NOT a Langfuse session (one-shot queries aren't a conversation thread) — it
|
|
# rides in metadata as a correlation key instead.
|
|
assert "session.id" not in run_span._otel_span.attributes
|
|
assert run_span._otel_span.attributes["user.id"] == "tenant1"
|
|
assert run_span._otel_span.attributes["langfuse.trace.name"] == "Dialectic Agent"
|
|
assert run_span.kwargs["metadata"]["honcho_session"] == "sess_abc"
|
|
|
|
|
|
def test_run_span_created_once_across_iterations(_exporter_env: FakeClient):
|
|
client = _exporter_env
|
|
exporter = LangfuseExporter()
|
|
exporter.export(_call(run_id="r1", trace_id="r1", iteration=1, session_id="s"))
|
|
exporter.export(_call(run_id="r1", trace_id="r1", iteration=2, session_id="s"))
|
|
|
|
spans = [o for o in client.observations if o.kwargs["as_type"] == "span"]
|
|
gens = [o for o in client.observations if o.kwargs["as_type"] == "generation"]
|
|
# One run span shared, one step span per iteration, one generation per call.
|
|
assert len(gens) == 2
|
|
assert len(spans) == 3 # 1 run + 2 step
|
|
# Trace attrs (user/name) stamped exactly once across the whole run.
|
|
stamped = [o for o in client.observations if "user.id" in o._otel_span.attributes]
|
|
assert len(stamped) == 1
|
|
|
|
|
|
def test_langfuse_session_lru_evicts_least_recently_used(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
):
|
|
"""Past _MAX_TRACES the least-recently-touched trace is evicted (not refused),
|
|
so an active trace keeps its remembered span ids no matter the run volume."""
|
|
langfuse_session.reset()
|
|
monkeypatch.setattr(langfuse_session, "_MAX_TRACES", 2)
|
|
|
|
langfuse_session.ensure_run_span("t1", "b", lambda _s: "t1-span")
|
|
langfuse_session.ensure_run_span("t2", "b", lambda _s: "t2-span")
|
|
# Touch t1 so t2 becomes the least-recently-used trace.
|
|
assert (
|
|
langfuse_session.ensure_run_span("t1", "b", lambda _s: "ignored") == "t1-span"
|
|
)
|
|
# A third trace evicts the LRU trace (t2), keeping t1.
|
|
langfuse_session.ensure_run_span("t3", "b", lambda _s: "t3-span")
|
|
|
|
created: list[str] = []
|
|
# t1 still tracked → remembered span returned, create NOT re-invoked.
|
|
assert (
|
|
langfuse_session.ensure_run_span(
|
|
"t1", "b", lambda _s: created.append("t1") or "new"
|
|
)
|
|
== "t1-span"
|
|
)
|
|
assert created == []
|
|
# t2 was evicted → fresh state, create IS re-invoked.
|
|
assert (
|
|
langfuse_session.ensure_run_span(
|
|
"t2", "b", lambda _s: created.append("t2") or "t2-span2"
|
|
)
|
|
== "t2-span2"
|
|
)
|
|
assert created == ["t2"]
|
|
|
|
|
|
def test_error_finish_marks_generation_level(_exporter_env: FakeClient):
|
|
client = _exporter_env
|
|
LangfuseExporter().export(
|
|
_call(run_id=None, trace_id="t1", finish_reason="error", content="")
|
|
)
|
|
gen = client.observations[0]
|
|
assert gen.kwargs["level"] == "ERROR"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("attr", "value"),
|
|
[
|
|
("LANGFUSE_EXPORTER_MODE", "inline"), # exporter off in inline mode
|
|
("LANGFUSE_PUBLIC_KEY", None), # exporter off without a public key
|
|
],
|
|
)
|
|
def test_exporter_disabled_emits_nothing(
|
|
_exporter_env: FakeClient,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
attr: str,
|
|
value: object,
|
|
):
|
|
client = _exporter_env
|
|
monkeypatch.setattr(settings, attr, value)
|
|
LangfuseExporter().export(_call(run_id="r1", trace_id="r1", iteration=1))
|
|
assert client.observations == []
|
|
|
|
|
|
def test_generation_name_uses_generation_suffix(_exporter_env: FakeClient):
|
|
client = _exporter_env
|
|
LangfuseExporter().export(
|
|
_call(run_id="r1", trace_id="r1", iteration=1, track_name="Dialectic Agent")
|
|
)
|
|
gen = [o for o in client.observations if o.kwargs["as_type"] == "generation"][0]
|
|
assert gen.kwargs["name"] == "Dialectic Agent generation"
|
|
step = [o for o in client.observations if o.kwargs["as_type"] == "span"][1]
|
|
assert step.kwargs["name"] == "Dialectic Agent step"
|
|
|
|
|
|
def test_tool_calls_become_spans_under_the_step(_exporter_env: FakeClient):
|
|
client = _exporter_env
|
|
LangfuseExporter().export(
|
|
_call(
|
|
run_id="r1",
|
|
trace_id="r1",
|
|
iteration=1,
|
|
track_name="Dialectic Agent",
|
|
tool_names=["search_memory", "search_messages"],
|
|
)
|
|
)
|
|
|
|
spans = [o for o in client.observations if o.kwargs["as_type"] == "span"]
|
|
gen = [o for o in client.observations if o.kwargs["as_type"] == "generation"][0]
|
|
tools = [o for o in client.observations if o.kwargs["as_type"] == "tool"]
|
|
step_span = spans[1] # run span, then step span
|
|
|
|
assert [t.kwargs["name"] for t in tools] == ["search_memory", "search_messages"]
|
|
# Tool spans are siblings of the generation: same parent (the step span).
|
|
for t in tools:
|
|
assert t.kwargs["trace_context"]["parent_span_id"] == step_span.id
|
|
assert gen.kwargs["trace_context"]["parent_span_id"] == step_span.id
|
|
# The model's requested input args ride on the tool span.
|
|
assert tools[0].kwargs["input"] == {"q": "search_memory"}
|
|
|
|
|
|
def test_only_the_root_span_keeps_as_root(_exporter_env: FakeClient):
|
|
# The SDK stamps AS_ROOT on every trace_context span; the exporter must
|
|
# demote children so exactly one root survives — otherwise Langfuse races to
|
|
# pick the trace name/root and names the trace after a child span.
|
|
from langfuse import LangfuseOtelSpanAttributes as Attr
|
|
|
|
client = _exporter_env
|
|
LangfuseExporter().export(
|
|
_call(
|
|
run_id="r1",
|
|
trace_id="r1",
|
|
iteration=1,
|
|
track_name="Dialectic Agent",
|
|
tool_names=["search_memory"],
|
|
)
|
|
)
|
|
|
|
def is_demoted(obs: FakeObs) -> bool:
|
|
return obs._otel_span.attributes.get(Attr.AS_ROOT) is False
|
|
|
|
spans = [o for o in client.observations if o.kwargs["as_type"] == "span"]
|
|
run_span, step_span = spans[0], spans[1]
|
|
gen = [o for o in client.observations if o.kwargs["as_type"] == "generation"][0]
|
|
tools = [o for o in client.observations if o.kwargs["as_type"] == "tool"]
|
|
|
|
# Exactly one root: the run span is never demoted; everything with a real
|
|
# parent is.
|
|
assert not is_demoted(run_span)
|
|
assert is_demoted(step_span)
|
|
assert is_demoted(gen)
|
|
assert all(is_demoted(t) for t in tools)
|
|
demoted = [o for o in client.observations if is_demoted(o)]
|
|
assert len(demoted) == len(client.observations) - 1
|
|
|
|
|
|
def test_single_shot_generation_keeps_as_root(_exporter_env: FakeClient):
|
|
# No parent → the generation is the trace root and must not be demoted.
|
|
from langfuse import LangfuseOtelSpanAttributes as Attr
|
|
|
|
client = _exporter_env
|
|
LangfuseExporter().export(
|
|
_call(run_id=None, trace_id="t1", track_name="Minimal Deriver")
|
|
)
|
|
gen = client.observations[0]
|
|
assert gen._otel_span.attributes.get(Attr.AS_ROOT) is not False
|
|
|
|
|
|
def test_single_shot_tool_calls_are_skipped(_exporter_env: FakeClient):
|
|
# No step span to anchor to (deriver-style); tools don't orphan to the root.
|
|
client = _exporter_env
|
|
LangfuseExporter().export(
|
|
_call(run_id=None, trace_id="t1", tool_names=["search_memory"])
|
|
)
|
|
assert [o.kwargs["as_type"] for o in client.observations] == ["generation"]
|
|
|
|
|
|
def test_dreamer_specialists_nest_under_one_dream_root(_exporter_env: FakeClient):
|
|
# Both specialists share ONE dream trace (run_id) and both start at
|
|
# iteration 1. They must nest under a single synthetic "Dream" root (so the
|
|
# trace has one root, not one per specialist) while staying distinct
|
|
# sub-trees (no step-span collision).
|
|
from langfuse import LangfuseOtelSpanAttributes as Attr
|
|
|
|
client = _exporter_env
|
|
exporter = LangfuseExporter()
|
|
for agent_type in ("deduction", "induction"):
|
|
exporter.export(
|
|
_call(
|
|
run_id="dream1",
|
|
trace_id="dream1",
|
|
iteration=1,
|
|
agent_type=agent_type,
|
|
parent_category="dream",
|
|
track_name=f"Dreamer/{agent_type}",
|
|
)
|
|
)
|
|
|
|
by_name: dict[str, list[FakeObs]] = {}
|
|
for o in client.observations:
|
|
by_name.setdefault(str(o.kwargs["name"]), []).append(o)
|
|
gens = [o for o in client.observations if o.kwargs["as_type"] == "generation"]
|
|
|
|
def is_demoted(o: FakeObs) -> bool:
|
|
return o._otel_span.attributes.get(Attr.AS_ROOT) is False
|
|
|
|
# Exactly one trace root: the synthetic "Dream" span — no parent, not demoted.
|
|
roots = [
|
|
o
|
|
for o in client.observations
|
|
if o.kwargs["trace_context"] == {"trace_id": "lf-dream1"}
|
|
]
|
|
assert len(roots) == 1
|
|
dream_root = roots[0]
|
|
assert dream_root.kwargs["name"] == "Dream"
|
|
assert dream_root.kwargs["as_type"] == "span"
|
|
assert not is_demoted(dream_root)
|
|
|
|
# Both specialist run spans hang off the Dream root and are demoted.
|
|
run_dd = by_name["Dreamer/deduction"][0]
|
|
run_in = by_name["Dreamer/induction"][0]
|
|
assert len(by_name["Dreamer/deduction"]) == 1
|
|
assert len(by_name["Dreamer/induction"]) == 1
|
|
for rs in (run_dd, run_in):
|
|
assert rs.kwargs["trace_context"] == {
|
|
"trace_id": "lf-dream1",
|
|
"parent_span_id": dream_root.id,
|
|
}
|
|
assert is_demoted(rs)
|
|
|
|
# One step span per specialist, parented to its own run span; no collapsing.
|
|
assert len(by_name["Dreamer/deduction step"]) == 1
|
|
assert len(by_name["Dreamer/induction step"]) == 1
|
|
assert (
|
|
by_name["Dreamer/deduction step"][0].kwargs["trace_context"]["parent_span_id"]
|
|
== run_dd.id
|
|
)
|
|
assert (
|
|
by_name["Dreamer/induction step"][0].kwargs["trace_context"]["parent_span_id"]
|
|
== run_in.id
|
|
)
|
|
|
|
# Each generation nests under its OWN specialist's step.
|
|
assert len({g.kwargs["trace_context"]["parent_span_id"] for g in gens}) == 2
|
|
|
|
# Trace name is the branch-agnostic "Dream", stamped exactly once — on the
|
|
# Dream root, not on a specialist's run span.
|
|
named = [
|
|
o
|
|
for o in client.observations
|
|
if o._otel_span.attributes.get("langfuse.trace.name")
|
|
]
|
|
assert len(named) == 1
|
|
assert named[0] is dream_root
|
|
assert named[0]._otel_span.attributes["langfuse.trace.name"] == "Dream"
|