wehub-resource-sync
283 lines
10 KiB
Python
283 lines
10 KiB
Python
"""Unit tests for ``with next_*_span(...)`` support exercised through
|
|
LangGraph ``StateGraph`` execution.
|
|
|
|
LangGraph reuses the LangChain ``CallbackHandler`` (one shared
|
|
codepath), so the underlying ``pop_pending_for(...)`` +
|
|
``apply_pending_to_span(...)`` plumbing is the same as in
|
|
``test_langchain/test_next_span.py``. What's distinct here is the
|
|
LangGraph orchestration surface: nodes scheduled across asyncio tasks,
|
|
multi-node graphs that fire the LLM callback more than once per
|
|
``ainvoke``, and the conditional-edge / multi-step flow where the
|
|
"first LLM span only" one-shot rule is the surprising behavior users
|
|
need a regression guard for.
|
|
"""
|
|
|
|
from typing import List
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
from langchain_core.language_models.fake import FakeListLLM
|
|
from langgraph.graph import END, START, StateGraph
|
|
from typing_extensions import TypedDict
|
|
|
|
from deepeval.integrations.langchain import CallbackHandler
|
|
from deepeval.metrics import BaseMetric
|
|
from deepeval.tracing import (
|
|
next_llm_span,
|
|
next_span,
|
|
next_tool_span,
|
|
trace_manager,
|
|
)
|
|
from deepeval.tracing.types import LlmSpan, ToolSpan
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _RecordingCallbackHandler(CallbackHandler):
|
|
"""Capture span object refs at start so tests can assert against
|
|
them after ``graph.ainvoke(...)`` (the trace ends and
|
|
``trace_manager.active_spans`` clears, but span objects stay
|
|
attached to the trace tree)."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.llm_spans: List[LlmSpan] = []
|
|
self.tool_spans: List[ToolSpan] = []
|
|
|
|
def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
|
|
res = super().on_chat_model_start(
|
|
serialized, messages, run_id=run_id, **kwargs
|
|
)
|
|
span = trace_manager.get_span_by_uuid(str(run_id))
|
|
if span is not None:
|
|
self.llm_spans.append(span)
|
|
return res
|
|
|
|
def on_llm_start(self, serialized, prompts, *, run_id, **kwargs):
|
|
res = super().on_llm_start(serialized, prompts, run_id=run_id, **kwargs)
|
|
span = trace_manager.get_span_by_uuid(str(run_id))
|
|
if span is not None:
|
|
self.llm_spans.append(span)
|
|
return res
|
|
|
|
|
|
class _State(TypedDict, total=False):
|
|
prompt: str
|
|
output: str
|
|
|
|
|
|
def _fake_metric(name: str = "fake") -> BaseMetric:
|
|
metric = MagicMock(spec=BaseMetric)
|
|
metric.__name__ = name
|
|
return metric
|
|
|
|
|
|
def _build_single_llm_graph(llm: FakeListLLM):
|
|
"""Smallest meaningful graph: START → llm node → END. The node
|
|
invokes ``llm`` so the handler sees one chain call + one LLM call
|
|
per ``graph.ainvoke``."""
|
|
|
|
async def node(state: _State, config=None) -> dict:
|
|
out = await llm.ainvoke(state["prompt"], config=config)
|
|
return {"output": out}
|
|
|
|
builder = StateGraph(_State)
|
|
builder.add_node("llm", node)
|
|
builder.add_edge(START, "llm")
|
|
builder.add_edge("llm", END)
|
|
return builder.compile()
|
|
|
|
|
|
def _build_two_llm_graph(llm: FakeListLLM):
|
|
"""Two LLM nodes back-to-back so we can pin down the "first LLM
|
|
span only" one-shot semantics that bites ``create_agent`` /
|
|
multi-step graphs in real workloads."""
|
|
|
|
async def first(state: _State, config=None) -> dict:
|
|
out = await llm.ainvoke(state["prompt"], config=config)
|
|
return {"output": out}
|
|
|
|
async def second(state: _State, config=None) -> dict:
|
|
out = await llm.ainvoke(state["output"], config=config)
|
|
return {"output": out}
|
|
|
|
builder = StateGraph(_State)
|
|
builder.add_node("first", first)
|
|
builder.add_node("second", second)
|
|
builder.add_edge(START, "first")
|
|
builder.add_edge("first", "second")
|
|
builder.add_edge("second", END)
|
|
return builder.compile()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# next_llm_span via StateGraph nodes
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.filterwarnings(
|
|
"ignore:The 'config' parameter should be typed as 'RunnableConfig' or 'RunnableConfig \\| None'"
|
|
)
|
|
class TestNextLlmSpanInStateGraph:
|
|
async def test_metric_collection_lands_on_llm_span(self):
|
|
callback = _RecordingCallbackHandler()
|
|
llm = FakeListLLM(responses=["pong"])
|
|
graph = _build_single_llm_graph(llm)
|
|
|
|
with next_llm_span(metric_collection="graph_llm_v1"):
|
|
await graph.ainvoke(
|
|
{"prompt": "ping"}, config={"callbacks": [callback]}
|
|
)
|
|
|
|
assert len(callback.llm_spans) == 1
|
|
assert callback.llm_spans[0].metric_collection == "graph_llm_v1"
|
|
|
|
async def test_metrics_lands_on_llm_span(self):
|
|
callback = _RecordingCallbackHandler()
|
|
llm = FakeListLLM(responses=["pong"])
|
|
graph = _build_single_llm_graph(llm)
|
|
metric = _fake_metric()
|
|
|
|
with next_llm_span(metrics=[metric]):
|
|
await graph.ainvoke(
|
|
{"prompt": "ping"}, config={"callbacks": [callback]}
|
|
)
|
|
|
|
assert callback.llm_spans[0].metrics == [metric]
|
|
|
|
async def test_metadata_lands_on_llm_span(self):
|
|
callback = _RecordingCallbackHandler()
|
|
llm = FakeListLLM(responses=["pong"])
|
|
graph = _build_single_llm_graph(llm)
|
|
|
|
with next_llm_span(metadata={"node": "llm"}):
|
|
await graph.ainvoke(
|
|
{"prompt": "ping"}, config={"callbacks": [callback]}
|
|
)
|
|
|
|
assert callback.llm_spans[0].metadata == {"node": "llm"}
|
|
|
|
async def test_only_first_llm_span_in_multi_node_graph(self):
|
|
"""The "create_agent gotcha" — a graph that opens two LLM spans
|
|
in one ``ainvoke`` only stamps the FIRST one. This is what the
|
|
docs caution-block warns about for ``StateGraph`` /
|
|
``create_agent`` loops; pin it down so a future change to drain
|
|
order doesn't silently flip the contract."""
|
|
callback = _RecordingCallbackHandler()
|
|
llm = FakeListLLM(responses=["pong-1", "pong-2"])
|
|
graph = _build_two_llm_graph(llm)
|
|
|
|
with next_llm_span(metric_collection="only-first-node"):
|
|
await graph.ainvoke(
|
|
{"prompt": "ping"}, config={"callbacks": [callback]}
|
|
)
|
|
|
|
assert len(callback.llm_spans) == 2
|
|
assert callback.llm_spans[0].metric_collection == "only-first-node"
|
|
assert callback.llm_spans[1].metric_collection is None
|
|
|
|
async def test_unconsumed_payload_does_not_leak_across_invocations(
|
|
self,
|
|
):
|
|
"""Token-based reset: a ``with`` that never opens an LLM span
|
|
(because we don't invoke the graph) doesn't pollute the next
|
|
graph invocation."""
|
|
callback = _RecordingCallbackHandler()
|
|
llm = FakeListLLM(responses=["pong"])
|
|
graph = _build_single_llm_graph(llm)
|
|
|
|
with next_llm_span(metric_collection="leaked"):
|
|
pass # no ainvoke → nothing pops
|
|
|
|
with next_llm_span(metric_collection="fresh"):
|
|
await graph.ainvoke(
|
|
{"prompt": "ping"}, config={"callbacks": [callback]}
|
|
)
|
|
|
|
assert callback.llm_spans[0].metric_collection == "fresh"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Cross-type isolation in graph context
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.filterwarnings(
|
|
"ignore:The 'config' parameter should be typed as 'RunnableConfig' or 'RunnableConfig \\| None'"
|
|
)
|
|
async def test_next_tool_span_does_not_leak_to_llm_span_in_graph():
|
|
"""The handler pops only the slot matching the span type it's
|
|
opening; staging a tool default and then opening an LLM span
|
|
leaves the LLM span clean."""
|
|
callback = _RecordingCallbackHandler()
|
|
llm = FakeListLLM(responses=["pong"])
|
|
graph = _build_single_llm_graph(llm)
|
|
|
|
with next_tool_span(metric_collection="tool-only"):
|
|
await graph.ainvoke(
|
|
{"prompt": "ping"}, config={"callbacks": [callback]}
|
|
)
|
|
|
|
assert callback.llm_spans[0].metric_collection is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Base ``next_span`` slot via StateGraph
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.filterwarnings(
|
|
"ignore:The 'config' parameter should be typed as 'RunnableConfig' or 'RunnableConfig \\| None'"
|
|
)
|
|
async def test_base_next_span_lands_on_first_llm_span_in_graph():
|
|
"""``next_span(...)`` is "next of any type" — base slot also
|
|
plumbs through the handler's ``pop_pending_for(...)`` merge for
|
|
LLM spans inside a ``StateGraph`` node."""
|
|
callback = _RecordingCallbackHandler()
|
|
llm = FakeListLLM(responses=["pong"])
|
|
graph = _build_single_llm_graph(llm)
|
|
|
|
with next_span(metric_collection="from_base_in_graph"):
|
|
await graph.ainvoke(
|
|
{"prompt": "ping"}, config={"callbacks": [callback]}
|
|
)
|
|
|
|
assert callback.llm_spans[0].metric_collection == "from_base_in_graph"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Sync StateGraph: typically users go async, but the same wiring must
|
|
# hold under ``graph.invoke(...)`` since the handler is the same code
|
|
# path.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.filterwarnings(
|
|
"ignore:The 'config' parameter should be typed as 'RunnableConfig' or 'RunnableConfig \\| None'"
|
|
)
|
|
def test_next_llm_span_in_sync_state_graph():
|
|
callback = _RecordingCallbackHandler()
|
|
llm = FakeListLLM(responses=["pong"])
|
|
|
|
def node(state: _State, config=None) -> dict:
|
|
out = llm.invoke(state["prompt"], config=config)
|
|
return {"output": out}
|
|
|
|
builder = StateGraph(_State)
|
|
builder.add_node("llm", node)
|
|
builder.add_edge(START, "llm")
|
|
builder.add_edge("llm", END)
|
|
graph = builder.compile()
|
|
|
|
with next_llm_span(metric_collection="sync_graph_v1"):
|
|
graph.invoke({"prompt": "ping"}, config={"callbacks": [callback]})
|
|
|
|
assert len(callback.llm_spans) == 1
|
|
assert callback.llm_spans[0].metric_collection == "sync_graph_v1"
|