diff --git a/py/src/braintrust/integrations/langchain/callbacks.py b/py/src/braintrust/integrations/langchain/callbacks.py index dbd413b8..2c50ffb2 100644 --- a/py/src/braintrust/integrations/langchain/callbacks.py +++ b/py/src/braintrust/integrations/langchain/callbacks.py @@ -72,6 +72,13 @@ def _provider_from_serialized(serialized: Mapping[str, Any] | None) -> str | Non return None +def _error_for_span(error: BaseException) -> str | None: + """``GeneratorExit`` means the consumer stopped a stream early, which is not a failure.""" + if isinstance(error, GeneratorExit): + return None + return str(error) + + def _resolve_name(name: str | None, serialized: Mapping[str, Any] | None, default: str) -> str: return name or (serialized or {}).get("name") or last_item((serialized or {}).get("id") or []) or default @@ -258,7 +265,7 @@ def on_llm_error( parent_run_id: UUID | None = None, **kwargs: Any, ) -> None: - self._end_span(run_id, error=str(error)) + self._end_span(run_id, error=_error_for_span(error)) self._start_times.pop(run_id, None) self._first_token_times.pop(run_id, None) self._ttft_ms.pop(run_id, None) @@ -271,7 +278,7 @@ def on_chain_error( parent_run_id: UUID | None = None, **kwargs: Any, ) -> None: - self._end_span(run_id, error=str(error)) + self._end_span(run_id, error=_error_for_span(error)) def on_tool_error( self, @@ -281,7 +288,7 @@ def on_tool_error( parent_run_id: UUID | None = None, **kwargs: Any, ) -> None: - self._end_span(run_id, error=str(error)) + self._end_span(run_id, error=_error_for_span(error)) def on_retriever_error( self, @@ -291,7 +298,7 @@ def on_retriever_error( parent_run_id: UUID | None = None, **kwargs: Any, ) -> None: - self._end_span(run_id, error=str(error)) + self._end_span(run_id, error=_error_for_span(error)) def on_agent_action( self, diff --git a/py/src/braintrust/integrations/langchain/test_callbacks.py b/py/src/braintrust/integrations/langchain/test_callbacks.py index 1cdac475..362c89d2 100644 --- a/py/src/braintrust/integrations/langchain/test_callbacks.py +++ b/py/src/braintrust/integrations/langchain/test_callbacks.py @@ -1,5 +1,6 @@ # pyright: reportTypedDictNotRequiredAccess=none import base64 +import contextlib import time import uuid from pathlib import Path @@ -1060,3 +1061,53 @@ def calculate(operation: str, a: float, b: float) -> float: ] ], ) + + +@pytest.mark.asyncio +async def test_langgraph_astream_early_exit_is_not_an_error(logger_memory_logger): + """Leaving an ``astream()`` loop early closes the stream with ``GeneratorExit``. + + LangGraph reports that via ``on_chain_error``; it is intentional cleanup, not a + failure, so the spans must not be marked as errored. + """ + test_logger, memory_logger = logger_memory_logger + assert not memory_logger.pop() + + try: + from langgraph.graph import END, START, StateGraph + except ImportError: + pytest.skip("langgraph not installed") + + handler = BraintrustCallbackHandler(logger=test_logger) + + def step_one(state: dict[str, str]): + return {"step": "one"} + + def step_two(state: dict[str, str]): + return {"step": "two"} + + graph = ( + StateGraph(state_schema=dict[str, str]) + .add_node("stepOne", step_one) + .add_node("stepTwo", step_two) + .add_edge(START, "stepOne") + .add_edge("stepOne", "stepTwo") + .add_edge("stepTwo", END) + .compile() + ) + + # ``aclosing`` closes the stream as soon as we stop iterating, rather than whenever + # the event loop finalizes the abandoned async generator. + async with contextlib.aclosing( + graph.astream({}, config={"callbacks": [handler]}, stream_mode="updates") + ) as stream: + async for chunk in stream: + assert "stepOne" in chunk + break + flush() + + spans = memory_logger.pop() + langgraph_spans = find_spans_by_attributes(spans, name="LangGraph") + assert len(langgraph_spans) == 1 + assert langgraph_spans[0].get("metrics", {}).get("end") is not None + assert [span.get("error") for span in spans] == [None] * len(spans)