diff --git a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py index efe8163..e007ed4 100644 --- a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py +++ b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py @@ -13,21 +13,29 @@ AiConfigRep, LDContext, ProviderHandler, + SpanMessage, + SpanMessagePart, config, create_handler, + create_run_usage, + end_span_once, + lang_chain_span_messages, + lang_chain_span_usage, parse_template, - set_ld_span_attributes, - set_openllmetry_completion, - set_openllmetry_prompt, + set_input_content_attributes, + set_output_content_attributes, ) -try: - from opentelemetry import trace - from opentelemetry.trace import StatusCode as SpanStatusCode - - _HAS_OTEL = True -except ImportError: - _HAS_OTEL = False +from .spans import ( + build_span_callbacks, + fail_span, + finish_root_span, + mark_ok, + parent_context_of, + start_root_span, + succeed_span, + to_tool_definitions, +) def _build_agent_tools( @@ -142,9 +150,36 @@ def _make_default_chat_model(config: AiConfigRep) -> Any: return lc_openai.ChatOpenAI(model=model_name or "gpt-4o") -def create_langchain_agents_handler(llm: Any = None) -> ProviderHandler: - """Creates a ``ProviderHandler`` for LangChain via ``create_react_agent``.""" - tracer_name = "@launchdarkly/ai-langchain-agents" +def _run_usage_from_messages(messages: list[Any]) -> Any: + """Sums ``usage_metadata`` over a run's messages, the same set of numbers the callbacks see + from the other side. + + Only ``AIMessage`` carries usage. Summing here rather than trusting the callbacks' own total + matters when there is a real ``result``/stepped state to read: it is the same path the TypeScript + handler takes. + + The two sides do not read the same fields, though. This one sees ``usage_metadata`` only, and the + callbacks also fall back to ``llm_output.token_usage``. The caller reconciles them, because a + provider that reports only in ``llm_output`` would otherwise give a successful run a root that + says zero and chat spans that say otherwise. + """ + run_usage = create_run_usage() + for msg in messages: + usage = getattr(msg, "usage_metadata", None) + if usage: + run_usage.add(lang_chain_span_usage(usage)) + return run_usage + + +def create_langchain_agents_handler( + llm: Any = None, *, capture_content: bool = False +) -> ProviderHandler: + """Creates a ``ProviderHandler`` for LangChain via ``create_react_agent``. + + Set *capture_content* to put prompts, model output, tool arguments and tool results on the + emitted spans. It defaults to off. Conversation content is PII, so a run emits only metadata, + meaning models, token counts, timings and tool names, until a caller asks for more. + """ async def _call_impl( config: AiConfigRep, @@ -158,19 +193,8 @@ async def _call_impl( th = tool_handlers or {} vs = variables or {} - if _HAS_OTEL: - span = trace.get_tracer(tracer_name).start_span("langchain.agent") - span.set_attribute("gen_ai.operation.name", "chat") - span.set_attribute( - "gen_ai.system", - config.get("provider", {}).get("name", "langchain").lower(), - ) - span.set_attribute( - "gen_ai.request.model", config.get("model", {}).get("name", "") - ) - set_ld_span_attributes(span, vs) - else: - span = None + span = start_root_span(config, vs) + parent = parent_context_of(span) system_prompt = _extract_system_prompt(config, vs, history) if config.get("outputFormat"): @@ -181,34 +205,24 @@ async def _call_impl( initial_messages = _build_initial_messages(config, user_input, vs) - if span: - prompt_text = "\n".join( - [ - *(["system: " + system_prompt] if system_prompt else []), - *[ - f"{getattr(m, 'type', type(m).__name__)}: {m.content}" - for m in initial_messages - ], - ] - ) - span.add_event("gen_ai.content.prompt", {"gen_ai.prompt": prompt_text}) - prompt_msgs: list[dict[str, str]] = [] - if system_prompt: - prompt_msgs.append({"role": "system", "content": system_prompt}) - prompt_msgs.extend( - [ - { - "role": getattr(m, "type", type(m).__name__), - "content": m.content - if isinstance(m.content, str) - else str(m.content), - } - for m in initial_messages - ] - ) - set_openllmetry_prompt(span, prompt_msgs) + span_callbacks = build_span_callbacks( + config, + parent, + capture_content, + to_tool_definitions(config.get("tools") or {}), + ) try: + # Inside the guard, because serialising the prompt raises on anything that is not + # JSON-serialisable and a raise out here would leave the root open: never ended, never + # exported, and the run gone from AI Config Monitoring with the feature_flag event on it. + if capture_content: + set_input_content_attributes( + span, + capture_content, + system_instructions=system_prompt, + messages=lang_chain_span_messages(initial_messages)[1], + ) base_model = llm if base_model is None: base_model = _make_default_chat_model(config) @@ -221,13 +235,25 @@ async def _call_impl( tools, **({"prompt": system_prompt} if system_prompt else {}), ) - result = await agent.ainvoke({"messages": initial_messages}) + result = await agent.ainvoke( + {"messages": initial_messages}, + config={"callbacks": span_callbacks.callbacks}, + ) msgs = ( result.get("messages", []) if isinstance(result, dict) else getattr(result, "messages", []) ) + run_usage = _run_usage_from_messages(msgs) + # The two sides do not see the same fields. A message carries usage_metadata and nothing + # else, while the callbacks read the LLMResult and fall back to llm_output.token_usage, + # which some providers use instead. When only the callbacks saw anything, they are the + # only record of what the run cost, and a successful root reporting zero while its own + # chat spans report real tokens is the one outcome neither figure can be right about. + if not run_usage.reported and span_callbacks.run_usage.reported: + run_usage = span_callbacks.run_usage + last_msg = msgs[-1] if msgs else None output = ( (last_msg.content if isinstance(last_msg.content, str) else "") @@ -235,50 +261,35 @@ async def _call_impl( else "" ) - total_input = sum( - (getattr(m, "usage_metadata", None) or {}).get("input_tokens", 0) - for m in msgs - ) - total_output = sum( - (getattr(m, "usage_metadata", None) or {}).get("output_tokens", 0) - for m in msgs + set_output_content_attributes( + span, + capture_content, + [ + SpanMessage( + role="assistant", + parts=[SpanMessagePart(type="text", content=output)], + ) + ], ) - - if span: - span.set_attribute( - "gen_ai.response.model", config.get("model", {}).get("name", "") - ) - span.set_attribute("gen_ai.usage.input_tokens", total_input) - span.set_attribute("gen_ai.usage.output_tokens", total_output) - span.set_attribute( - "gen_ai.usage.total_tokens", total_input + total_output - ) - span.add_event( - "gen_ai.content.completion", - { - "gen_ai.completion": output - if isinstance(output, str) - else json.dumps(output) - }, - ) - set_openllmetry_completion( - span, - output if isinstance(output, str) else json.dumps(output), - {"input_tokens": total_input, "output_tokens": total_output}, - ) - span.set_status(SpanStatusCode.OK) - span.end() + finish_root_span(span, config, run_usage.total) + succeed_span(span) return { "output": output, - "usage": {"input_tokens": total_input, "output_tokens": total_output}, + "usage": { + "input_tokens": run_usage.total.input, + "output_tokens": run_usage.total.output, + }, } except Exception as exc: - if span: - span.record_exception(exc) - span.set_status(SpanStatusCode.ERROR, str(exc)) - span.end() + span_callbacks.close_open_spans(exc) + # There is no `result` to sum, so the run total comes from the callbacks, which saw + # every turn that did complete. Those tokens were billed and the root is the only span + # a config-scoped cost query can find them on. + if span_callbacks.run_usage.reported: + finish_root_span(span, config, span_callbacks.run_usage.total) + fail_span(span, exc) raise def _stream_impl( @@ -289,7 +300,13 @@ def _stream_impl( history: list[dict[str, Any]] | None = None, ) -> AsyncGenerator[dict[str, Any], None]: return _stream_gen( - llm, config, user_input, tool_handlers or {}, variables or {}, history + llm, + config, + user_input, + tool_handlers or {}, + variables or {}, + history, + capture_content=capture_content, ) return create_handler(("*", "agent"), _call_impl, _stream_impl) # type: ignore[arg-type] @@ -302,54 +319,43 @@ async def _stream_gen( tool_handlers: dict[str, Any], variables: dict[str, Any], history: list[dict[str, Any]] | None = None, + *, + capture_content: bool = False, ) -> AsyncGenerator[dict[str, Any], None]: + """Streams the run, emitting the same span tree as the blocking path. + + A consumer that breaks out of ``async for``, or raises inside the loop body, makes this + generator run its ``finally`` without ever entering ``except``: ``GeneratorExit`` inherits from + ``BaseException``, so ``except Exception`` does not see it. Without the cleanup in ``finally`` + the root span, and any ``chat``/``execute_tool`` span LangChain's end callback never fired for, + is never ended, so it is never exported. + + ``ended`` stops the success, failure and abandonment paths from ending the same span twice. + """ import importlib - tracer_name = "@launchdarkly/ai-langchain-agents" - if _HAS_OTEL: - span = trace.get_tracer(tracer_name).start_span("langchain.agent.stream") - span.set_attribute("gen_ai.operation.name", "chat") - span.set_attribute( - "gen_ai.system", config.get("provider", {}).get("name", "langchain").lower() - ) - span.set_attribute( - "gen_ai.request.model", config.get("model", {}).get("name", "") - ) - set_ld_span_attributes(span, variables) - else: - span = None + span = start_root_span(config, variables) + parent = parent_context_of(span) system_prompt = _extract_system_prompt(config, variables, history) initial_messages = _build_initial_messages(config, user_input, variables) - if span: - prompt_text = "\n".join( - [ - *(["system: " + system_prompt] if system_prompt else []), - *[ - f"{getattr(m, 'type', type(m).__name__)}: {m.content}" - for m in initial_messages - ], - ] - ) - span.add_event("gen_ai.content.prompt", {"gen_ai.prompt": prompt_text}) - prompt_msgs: list[dict[str, str]] = [] - if system_prompt: - prompt_msgs.append({"role": "system", "content": system_prompt}) - prompt_msgs.extend( - [ - { - "role": getattr(m, "type", type(m).__name__), - "content": m.content - if isinstance(m.content, str) - else str(m.content), - } - for m in initial_messages - ] - ) - set_openllmetry_prompt(span, prompt_msgs) + span_callbacks = build_span_callbacks( + config, parent, capture_content, to_tool_definitions(config.get("tools") or {}) + ) + ended: set[int] = set() try: + # Inside the guard, because serialising the prompt raises on anything that is not + # JSON-serialisable. A raise out here would leave the root open with the `finally` never + # entered, so the run would vanish from AI Config Monitoring with its feature_flag event. + if capture_content: + set_input_content_attributes( + span, + capture_content, + system_instructions=system_prompt, + messages=lang_chain_span_messages(initial_messages)[1], + ) base_model = llm if base_model is None: base_model = _make_default_chat_model(config) @@ -363,11 +369,14 @@ async def _stream_gen( **({"prompt": system_prompt} if system_prompt else {}), ) - total_input = 0 - total_output = 0 + run_usage = create_run_usage() full_output = "" - async for step_state in agent.astream({"messages": initial_messages}): + # agent.astream() yields state updates per graph step: { [node_name]: { messages: [...] } } + async for step_state in agent.astream( + {"messages": initial_messages}, + config={"callbacks": span_callbacks.callbacks}, + ): for step_messages in ( step_state.values() if isinstance(step_state, dict) else [] ): @@ -377,52 +386,56 @@ async def _stream_gen( else [] ) for msg in msgs: - usage = getattr(msg, "usage_metadata", None) or {} - total_input += usage.get("input_tokens", 0) - total_output += usage.get("output_tokens", 0) + usage = getattr(msg, "usage_metadata", None) + if usage: + run_usage.add(lang_chain_span_usage(usage)) if getattr(msg, "type", None) == "ai": text = msg.content if isinstance(msg.content, str) else "" if text: yield {"type": "chunk", "text": text} full_output = text - if span: - span.set_attribute( - "gen_ai.response.model", config.get("model", {}).get("name", "") - ) - span.set_attribute("gen_ai.usage.input_tokens", total_input) - span.set_attribute("gen_ai.usage.output_tokens", total_output) - span.set_attribute("gen_ai.usage.total_tokens", total_input + total_output) - span.add_event( - "gen_ai.content.completion", - { - "gen_ai.completion": full_output - if isinstance(full_output, str) - else json.dumps(full_output) - }, - ) - set_openllmetry_completion( - span, - full_output - if isinstance(full_output, str) - else json.dumps(full_output), - {"input_tokens": total_input, "output_tokens": total_output}, - ) - span.set_status(SpanStatusCode.OK) - span.end() + set_output_content_attributes( + span, + capture_content, + [ + SpanMessage( + role="assistant", + parts=[SpanMessagePart(type="text", content=full_output)], + ) + ], + ) + finish_root_span(span, config, run_usage.total) + mark_ok(span) + end_span_once(span, ended) yield { "type": "done", "output": full_output, - "usage": {"input_tokens": total_input, "output_tokens": total_output}, + "usage": { + "input_tokens": run_usage.total.input, + "output_tokens": run_usage.total.output, + }, } except Exception as exc: - if span: - span.record_exception(exc) - span.set_status(SpanStatusCode.ERROR, str(exc)) - span.end() + span_callbacks.close_open_spans(exc) + if span_callbacks.run_usage.reported: + finish_root_span(span, config, span_callbacks.run_usage.total) + fail_span(span, exc, ended) raise + finally: + # A no-op on the success and failure paths, because both already ended their spans through + # `ended`. On abandonment it is the only chance to close the tree, including any chat or + # tool span whose LangChain end callback never fired, and to report what the completed + # turns already cost. An abandoned span is left UNSET rather than ERROR: stopping early is + # a normal thing for a consumer to do, and LaunchDarkly's own metrics record neither a + # success nor an error for it. + if span is not None and id(span) not in ended: + span_callbacks.abandon_open_spans(ended) + if span_callbacks.run_usage.reported: + finish_root_span(span, config, span_callbacks.run_usage.total) + end_span_once(span, ended, abandoned=True) def langchain_agents( @@ -432,7 +445,13 @@ def langchain_agents( **kwargs: Any, ) -> Any: """Convenience wrapper: creates a handler and calls config(...).invoke().""" + # Both are lifted out of kwargs: capture_content configures the handler, variables belong to + # the invocation. Leaving either in would pass it to config(), which takes neither, so a caller + # asking for content on spans got a TypeError instead of content. variables = kwargs.pop("variables", None) + capture_content = kwargs.pop("capture_content", False) return config( - key=config_key, handler=create_langchain_agents_handler(), **kwargs + key=config_key, + handler=create_langchain_agents_handler(capture_content=capture_content), + **kwargs, ).invoke(user_input, context, variables=variables) diff --git a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/spans.py b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/spans.py new file mode 100644 index 0000000..e302881 --- /dev/null +++ b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/spans.py @@ -0,0 +1,514 @@ +"""Span construction for the LangChain agents handler. + +Separate from ``handler.py`` so the span shape is readable on its own, and so the agent-invocation +code reads as agent invocation rather than as span bookkeeping around a LangGraph call. + +The shape is ``invoke_agent`` root, one ``chat {model}`` child per model turn, one +``execute_tool {name}`` child per tool call. Tool spans are siblings of the ``chat`` span, not +children of it: both take the same parent context, which is the root's. See TELEMETRY-CONTRACT.md +section 1. + +Unlike the Claude and OpenAI handlers, this one does not drive the provider call itself: LangGraph's +``create_react_agent`` does, and the only hook into its lifecycle is LangChain's callback protocol. +:func:`build_span_callbacks` is therefore the center of this module: it returns a +``BaseCallbackHandler`` that opens and closes ``chat`` and ``execute_tool`` spans as LangChain +dispatches ``on_chat_model_start`` / ``on_llm_end`` / ``on_tool_start`` / ``on_tool_end`` events, +keyed by the callback ``run_id`` so concurrent tool calls do not collide. +""" + +from __future__ import annotations + +from typing import Any + +from launchdarkly_ai_server import ( + AiConfigRep, + RunUsage, + SpanUsage, + ToolDefinitionInput, + create_run_usage, + end_span_once, + lang_chain_finish_reasons, + lang_chain_span_messages, + lang_chain_span_usage, + set_input_content_attributes, + set_ld_span_attributes, + set_model_identity_attributes, + set_output_content_attributes, + set_tool_call_content_attributes, + set_usage_span_attributes, +) + +try: + from opentelemetry import trace + from opentelemetry.trace import StatusCode as SpanStatusCode + + _HAS_OTEL = True +except ImportError: # pragma: no cover - exercised by the no-OTel install path + _HAS_OTEL = False + +try: + from langchain_core.callbacks import AsyncCallbackHandler + + _HAS_LANGCHAIN_CORE = True +except ImportError: # pragma: no cover + AsyncCallbackHandler = object # type: ignore[assignment,misc] + _HAS_LANGCHAIN_CORE = False + +TRACER_NAME = "@launchdarkly/ai-langchain-agents" + + +def model_name(config: AiConfigRep) -> str: + return str(config.get("model", {}).get("name", "")) + + +def serving_provider(config: AiConfigRep) -> str: + """The provider that actually serves the model. + + ``gen_ai.provider.name`` names who served the request, and its semconv enum has no + ``langchain`` member: LangChain is the framework, not the provider. This mirrors the choice + ``_make_default_chat_model`` makes, so the attribute agrees with the client that is really + used. It is a binary choice, not a passthrough of the configured name: the configured name + lower-cased if it equals ``anthropic``, otherwise ``openai``, no matter what else the config + names (Bedrock, Azure, Cohere, a typo, or nothing at all). + """ + provider = ((config.get("provider") or {}).get("name") or "").lower() + return "anthropic" if provider == "anthropic" else "openai" + + +# ─── Span starts ───────────────────────────────────────────────────────────── + + +def start_root_span(config: AiConfigRep, variables: dict[str, Any]) -> Any: + """Opens the ``invoke_agent`` root and returns it, or ``None`` when OTel is absent. + + The root is the only span carrying ``launchdarkly.*`` and the ``feature_flag`` event, so it is + the span a config-scoped query finds. Child spans must not carry them. + """ + if not _HAS_OTEL: + return None + span = trace.get_tracer(TRACER_NAME).start_span("invoke_agent") + span.set_attribute("gen_ai.operation.name", "invoke_agent") + set_model_identity_attributes( + span, serving_provider(config), model_name(config), legacy_system="langchain" + ) + set_ld_span_attributes(span, variables) + return span + + +def parent_context_of(span: Any) -> Any: + """The context a child span should be parented to. + + Explicit rather than a bare current context: the current context only carries this span while + a context manager has attached it, and these handlers open a plain span rather than an active + one, so a host app that installs its own tracer provider would otherwise get a flat trace. + """ + if not _HAS_OTEL or span is None: + return None + return trace.set_span_in_context(span) + + +def start_model_span(config: AiConfigRep, parent: Any) -> Any: + """Opens one ``chat {model}`` span for one model turn. + + The semantic conventions name an inference span ``{operation} {model}``, so the model belongs + in the name and not only in ``gen_ai.request.model``. A bare ``chat`` aggregates more neatly + but tells a reader nothing about which model ran, which matters most in exactly the case this + span exists for: a multi-turn run that switches models partway through. + """ + if not _HAS_OTEL: + return None + name = model_name(config) + span = trace.get_tracer(TRACER_NAME).start_span(f"chat {name}", context=parent) + span.set_attribute("gen_ai.operation.name", "chat") + set_model_identity_attributes( + span, serving_provider(config), name, legacy_system="langchain" + ) + return span + + +def start_tool_span(tool_name: str, tool_call_id: str, parent: Any) -> Any: + """Opens one ``execute_tool {name}`` span for one tool call.""" + if not _HAS_OTEL: + return None + span = trace.get_tracer(TRACER_NAME).start_span( + f"execute_tool {tool_name}", context=parent + ) + span.set_attribute("gen_ai.operation.name", "execute_tool") + span.set_attribute("gen_ai.tool.name", tool_name) + span.set_attribute("gen_ai.tool.call.id", tool_call_id) + return span + + +# ─── Span finishes ─────────────────────────────────────────────────────────── + + +def finish_root_span(span: Any, config: AiConfigRep, run_usage: SpanUsage) -> None: + """Writes the run-level identity and token totals onto the root. + + ``gen_ai.response.model`` is the requested name here, same as on a ``chat`` span: neither + LangChain handler resolves an alias to a different snapshot. See TELEMETRY-CONTRACT.md + section 2a. + """ + if span is None: + return + span.set_attribute("gen_ai.response.model", model_name(config)) + set_usage_span_attributes(span, run_usage) + + +def finish_model_span( + span: Any, + config: AiConfigRep, + raw_usage: dict[str, Any] | None, + finish_reasons: list[str] | None = None, +) -> None: + """Ends one ``chat`` span successfully. *finish_reasons* arrives already mapped.""" + if span is None: + return + span.set_attribute("gen_ai.response.model", model_name(config)) + if finish_reasons: + span.set_attribute("gen_ai.response.finish_reasons", finish_reasons) + # A span always carries the complete usage attribute set, zeros included, unlike the run + # accumulator: an absent attribute drops the span from every query that groups on usage. + set_usage_span_attributes(span, lang_chain_span_usage(raw_usage) or SpanUsage()) + span.set_status(SpanStatusCode.OK) + span.end() + + +def succeed_span(span: Any) -> None: + """Marks a span OK and ends it, for spans with nothing else to report.""" + if span is None: + return + span.set_status(SpanStatusCode.OK) + span.end() + + +def mark_ok(span: Any) -> None: + """Marks a span OK without ending it. + + The streaming path needs this: its ``finally`` owns every end, through ``end_span_once``, so a + success tail that ended the span itself would end it twice. + """ + if span is None: + return + span.set_status(SpanStatusCode.OK) + + +def fail_span(span: Any, error: BaseException, tracker: set[int] | None = None) -> None: + """Records the exception, sets ERROR, and ends the span. + + *tracker* is passed only from the streaming path, where a ``finally`` may race this to the same + span; elsewhere there is exactly one end and the tracker is unnecessary. + """ + if span is None: + return + span.record_exception(error) + span.set_status(SpanStatusCode.ERROR, str(error)) + if tracker is not None: + end_span_once(span, tracker) + else: + span.end() + + +# ─── LangChain result shapes as span shapes ────────────────────────────────── + + +def _get(obj: Any, key: str) -> Any: + if isinstance(obj, dict): + return obj.get(key) + return getattr(obj, key, None) + + +def generated_messages(output: Any) -> list[Any]: + """The generated messages of an ``LLMResult``, which is where a turn's real output lives.""" + generations = _get(output, "generations") + if not isinstance(generations, list): + return [] + flat = [item for group in generations for item in group] + return [msg for msg in (_get(gen, "message") for gen in flat) if msg is not None] + + +def extract_llm_usage(output: Any) -> dict[str, Any]: + """LangChain's ``LLMResult`` carries usage either on each generation's message + (``usage_metadata``) or, for some providers, in ``llm_output.token_usage``. Prefer the former, + fall back to the latter. + """ + generations = _get(output, "generations") + flat = ( + [item for group in generations for item in group] + if isinstance(generations, list) + else [] + ) + for gen in flat: + message = _get(gen, "message") + usage_metadata = _get(message, "usage_metadata") if message else None + if usage_metadata: + return dict(usage_metadata) + + llm_output = _get(output, "llm_output") or {} + token_usage = llm_output.get("token_usage") or llm_output.get("usage") + if token_usage: + # `or` would read a real 0 as missing, and with both counts at zero the bag came back all + # None, which lang_chain_span_usage reports as "the provider said nothing". A turn that + # completed and reported zero is not the same as a turn that reported nothing: only the + # second may leave the root without usage attributes. + return { + "input_tokens": _first_present( + token_usage, "prompt_tokens", "input_tokens" + ), + "output_tokens": _first_present( + token_usage, "completion_tokens", "output_tokens" + ), + } + return {} + + +def _first_present(bag: dict[str, Any], *keys: str) -> Any: + """The first key actually present, so a reported 0 is kept rather than skipped.""" + for key in keys: + if key in bag: + return bag[key] + return None + + +def to_tool_definitions(config_tools: dict[str, Any]) -> list[ToolDefinitionInput]: + """The catalog handed to the agent, so a ``chat`` span reports what the model could call.""" + return [ + ToolDefinitionInput( + name=tool.get("name", name), + description=tool.get("description"), + parameters=tool.get("parameters"), + ) + for name, tool in config_tools.items() + ] + + +# ─── The callback bridge ────────────────────────────────────────────────────── + + +class SpanCallbackHandler(AsyncCallbackHandler): + """Maps LangChain's async callback lifecycle onto the ``chat`` / ``execute_tool`` spans. + + Spans are keyed by the callback ``run_id`` (as ``str``) so concurrent tool calls, or a run that + somehow issues concurrent model calls, do not collide. + """ + + def __init__( + self, + config: AiConfigRep, + parent_context: Any, + capture_content: bool, + tool_definitions: list[ToolDefinitionInput], + run_usage: RunUsage, + ) -> None: + self._config = config + self._parent_context = parent_context + self._capture_content = capture_content + self._tool_definitions = tool_definitions + self.run_usage = run_usage + self.model_spans: dict[str, Any] = {} + self.tool_spans: dict[str, Any] = {} + + def _start_model(self, run_id: Any, messages: Any = None) -> None: + key = str(run_id) + if key in self.model_spans: + return + span = start_model_span(self._config, self._parent_context) + # Tracked before the content write, for the same reason as on_tool_start: serialising the + # conversation can raise, and a span created but never inserted is unreachable by every + # cleanup path there is. + self.model_spans[key] = span + if self._capture_content: + # `on_chat_model_start` hands over `list[list[BaseMessage]]`, one list per generation. + # The agent graph sends a single list; flattening keeps a multi-generation caller from + # losing turns. + flat: list[Any] = [] + if isinstance(messages, list): + for group in messages: + if isinstance(group, list): + flat.extend(group) + system_instructions, turn_messages = lang_chain_span_messages(flat) + set_input_content_attributes( + span, + self._capture_content, + system_instructions=system_instructions, + messages=turn_messages, + tool_definitions=self._tool_definitions, + ) + + async def on_chat_model_start( + self, + serialized: dict[str, Any], + messages: list[list[Any]], + *, + run_id: Any, + **kwargs: Any, + ) -> None: + self._start_model(run_id, messages) + + async def on_llm_start( + self, + serialized: dict[str, Any], + prompts: list[str], + *, + run_id: Any, + **kwargs: Any, + ) -> None: + self._start_model(run_id) + + async def on_llm_end(self, response: Any, *, run_id: Any, **kwargs: Any) -> None: + key = str(run_id) + span = self.model_spans.pop(key, None) + if span is None: + return + # Popped, so close_open_spans can no longer reach this span: whatever happens next, this + # method has to end it. Serialising conversation content raises on anything that is not + # JSON-serialisable, and a raise here would otherwise leak the span with nothing tracking it. + # Before the content write, because the provider has already billed these tokens: a content + # failure is our problem and must not report the run as having spent less than it did. + turn_usage_raw = extract_llm_usage(response) + self.run_usage.add(lang_chain_span_usage(turn_usage_raw)) + try: + if self._capture_content: + set_output_content_attributes( + span, + self._capture_content, + lang_chain_span_messages(generated_messages(response))[1], + ) + finish_model_span( + span, self._config, turn_usage_raw, lang_chain_finish_reasons(response) + ) + except Exception as exc: + fail_span(span, exc) + raise + + async def on_llm_error( + self, error: BaseException, *, run_id: Any, **kwargs: Any + ) -> None: + key = str(run_id) + span = self.model_spans.pop(key, None) + if span is None: + return + fail_span(span, error) + + async def on_tool_start( + self, + serialized: dict[str, Any], + input_str: str, + *, + run_id: Any, + name: str | None = None, + tool_call_id: str | None = None, + **kwargs: Any, + ) -> None: + tool_name = name or (serialized or {}).get("name") or "tool" + span = start_tool_span( + tool_name, tool_call_id or str(run_id), self._parent_context + ) + # Tracked before the content write, not after. Serialising the arguments can raise, and a + # span created but never inserted is unreachable by every cleanup path there is. + self.tool_spans[str(run_id)] = span + set_tool_call_content_attributes( + span, self._capture_content, arguments=kwargs.get("inputs") or input_str + ) + + async def on_tool_end(self, output: Any, *, run_id: Any, **kwargs: Any) -> None: + key = str(run_id) + span = self.tool_spans.pop(key, None) + if span is None: + return + # Same reason as on_llm_end: once popped, ending it is this method's job alone. + try: + set_tool_call_content_attributes( + span, self._capture_content, result=_tool_result_text(output) + ) + succeed_span(span) + except Exception as exc: + fail_span(span, exc) + raise + + async def on_tool_error( + self, error: BaseException, *, run_id: Any, **kwargs: Any + ) -> None: + key = str(run_id) + span = self.tool_spans.pop(key, None) + if span is None: + return + fail_span(span, error) + + def close_open_spans(self, error: BaseException) -> None: + """Fails every span this run opened but never closed, for the caller's failure path.""" + for span in self.model_spans.values(): + fail_span(span, error) + for span in self.tool_spans.values(): + fail_span(span, error) + self.model_spans.clear() + self.tool_spans.clear() + + def abandon_open_spans(self, ended: set[int]) -> None: + """Ends every span this run has open, for stream abandonment. + + Unlike :meth:`close_open_spans`, nothing failed: a consumer stopping early is normal. + ``end_span_once`` leaves each span at ``UNSET`` and marks ``launchdarkly.stream.abandoned``, + rather than recording a synthetic exception and setting ``ERROR``, which would make an early + consumer stop indistinguishable from a provider failure in a trace. + """ + for span in self.model_spans.values(): + end_span_once(span, ended, abandoned=True) + for span in self.tool_spans.values(): + end_span_once(span, ended, abandoned=True) + self.model_spans.clear() + self.tool_spans.clear() + + +def _tool_result_text(output: Any) -> Any: + """A ``ToolMessage`` carries the result in ``content``; anything else passes through.""" + content = _get(output, "content") + return content if content is not None else output + + +class SpanCallbacks: + """The bundle a handler call site needs: callbacks to hand LangChain, the run's accumulated + usage, and a way to close whatever spans a failure or an abandonment left open. + + A plain wrapper, not the callback handler itself, so the no-OTel / no-``langchain-core`` + fallback can satisfy the same shape with an empty callback list rather than a special case at + every call site. + """ + + def __init__( + self, callbacks: list[Any], run_usage: RunUsage, handler: Any = None + ) -> None: + self.callbacks = callbacks + self.run_usage = run_usage + self._handler = handler + + def close_open_spans(self, error: BaseException) -> None: + if self._handler is not None: + self._handler.close_open_spans(error) + + def abandon_open_spans(self, ended: set[int]) -> None: + if self._handler is not None: + self._handler.abandon_open_spans(ended) + + +def build_span_callbacks( + config: AiConfigRep, + parent_context: Any, + capture_content: bool = False, + tool_definitions: list[ToolDefinitionInput] | None = None, +) -> SpanCallbacks: + """Builds a LangChain callback handler that maps the agent's LLM and tool lifecycle onto OTel + spans: one ``chat`` child span per model turn, one ``execute_tool`` child span per tool call. + + Returns a :class:`SpanCallbacks` exposing ``.callbacks`` (to pass as + ``config={"callbacks": ...}``), ``.run_usage`` (the run's accumulated :class:`SpanUsage`), and + ``.close_open_spans(error)`` (for the caller's failure and abandonment paths, where a span this + callback opened may never see its matching end event). + """ + if not _HAS_OTEL or not _HAS_LANGCHAIN_CORE: + return SpanCallbacks([], create_run_usage()) + run_usage = create_run_usage() + handler = SpanCallbackHandler( + config, parent_context, capture_content, tool_definitions or [], run_usage + ) + return SpanCallbacks([handler], run_usage, handler) diff --git a/packages/langchain-agents/tests/test_handler.py b/packages/langchain-agents/tests/test_handler.py index fe35713..cdc34cb 100644 --- a/packages/langchain-agents/tests/test_handler.py +++ b/packages/langchain-agents/tests/test_handler.py @@ -7,12 +7,16 @@ from __future__ import annotations from collections.abc import AsyncIterator +from types import SimpleNamespace from typing import Any, ClassVar from unittest.mock import AsyncMock, MagicMock, patch +import pydantic import pytest +from langchain_core.language_models.chat_models import BaseChatModel import launchdarkly_ai_langchain_agents.handler as handler_mod +import launchdarkly_ai_langchain_agents.spans as spans_mod from launchdarkly_ai_langchain_agents.handler import ( _build_initial_messages, _extract_system_prompt, @@ -145,7 +149,7 @@ def _import_side_effect(name: str) -> Any: return __import__(name) with patch("importlib.import_module", side_effect=_import_side_effect): - with patch.object(handler_mod, "_HAS_OTEL", False): + with patch.object(spans_mod, "_HAS_OTEL", False): h = create_langchain_agents_handler(llm=MagicMock()) await h(_make_config(instructions="Be helpful."), "hi") @@ -450,328 +454,654 @@ async def test_no_tools_in_config_handler_never_invoked(self) -> None: # --------------------------------------------------------------------------- -class TestTelemetry: - @pytest.mark.asyncio - async def test_span_name(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span +class _FakeToolModel(BaseChatModel): + """A real ``BaseChatModel`` (so ``create_react_agent`` accepts it) that returns canned replies + in sequence and never touches the network. ``bind_tools`` is required by the agent graph and + the base class raises ``NotImplementedError`` for it. + """ - mocks = _make_langchain_mock() - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler( - llm=MagicMock(ainvoke=AsyncMock(return_value=mocks["_ai_msg"])) - ) - await h(_make_config(), "hi") + replies: list[Any] = pydantic.Field(default_factory=list) + fail_after: int | None = None + fail_with: Exception | None = None + calls: int = 0 - mock_trace.get_tracer.return_value.start_span.assert_called_with( - "langchain.agent" - ) + def bind_tools(self, tools: Any, **kwargs: Any) -> Any: + return self - @pytest.mark.asyncio - async def test_gen_ai_system(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span + def _generate( + self, messages: Any, stop: Any = None, run_manager: Any = None, **kwargs: Any + ) -> Any: + raise NotImplementedError - mocks = _make_langchain_mock() - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler( - llm=MagicMock(ainvoke=AsyncMock(return_value=mocks["_ai_msg"])) - ) - await h(_make_config(), "hi") + async def _agenerate( + self, messages: Any, stop: Any = None, run_manager: Any = None, **kwargs: Any + ) -> Any: + from langchain_core.outputs import ChatGeneration, ChatResult - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("gen_ai.system") == "langchain" + if self.fail_after is not None and self.calls >= self.fail_after: + raise self.fail_with or RuntimeError("model down") + reply = self.replies[min(self.calls, len(self.replies) - 1)] + self.calls += 1 + return ChatResult(generations=[ChatGeneration(message=reply)]) - @pytest.mark.asyncio - async def test_gen_ai_operation_name(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span + @property + def _llm_type(self) -> str: + return "fake-tool-model" - mocks = _make_langchain_mock() - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler( - llm=MagicMock(ainvoke=AsyncMock(return_value=mocks["_ai_msg"])) - ) - await h(_make_config(), "hi") - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("gen_ai.operation.name") == "chat" +def _ai_message( + content: str = "", + input_tokens: int = 10, + output_tokens: int = 5, + tool_calls: list[dict[str, Any]] | None = None, + response_metadata: dict[str, Any] | None = None, +) -> Any: + from langchain_core.messages import AIMessage + + return AIMessage( + content=content, + usage_metadata={ + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + }, + tool_calls=tool_calls or [], + response_metadata=response_metadata or {}, + ) - @pytest.mark.asyncio - async def test_span_status_ok(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - mocks = _make_langchain_mock() - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler( - llm=MagicMock(ainvoke=AsyncMock(return_value=mocks["_ai_msg"])) - ) - await h(_make_config(), "hi") +class RecordedSpan: + """A span that remembers what a handler did to it, so a test can assert on the whole thing.""" - mock_span.set_status.assert_called() + def __init__(self, name: str, context: Any = None) -> None: + self.name = name + self.context = context + self.attributes: dict[str, Any] = {} + self.events: list[tuple[str, dict[str, Any]]] = [] + self.statuses: list[Any] = [] + self.exceptions: list[BaseException] = [] + self.ended = 0 - @pytest.mark.asyncio - async def test_span_end_always_called(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span + def set_attribute(self, key: str, value: Any) -> None: + self.attributes[key] = value - mocks = _make_langchain_mock() - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler( - llm=MagicMock(ainvoke=AsyncMock(return_value=mocks["_ai_msg"])) - ) - await h(_make_config(), "hi") + def add_event(self, name: str, attributes: dict[str, Any] | None = None) -> None: + self.events.append((name, attributes or {})) - mock_span.end.assert_called() + def set_status(self, code: Any, description: str | None = None) -> None: + self.statuses.append(code) - @pytest.mark.asyncio - async def test_gen_ai_request_model(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span + def record_exception(self, exc: BaseException) -> None: + self.exceptions.append(exc) - mocks = _make_langchain_mock() - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler( - llm=MagicMock(ainvoke=AsyncMock(return_value=mocks["_ai_msg"])) - ) - await h(_make_config(model={"name": "gpt-4o"}), "hi") + def end(self) -> None: + self.ended += 1 - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("gen_ai.request.model") == "gpt-4o" - @pytest.mark.asyncio - async def test_gen_ai_content_prompt_event(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span +class SpanRecorder: + """Stands in for the ``trace`` module inside ``spans.py`` and records every span opened. - mocks = _make_langchain_mock() - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler( - llm=MagicMock(ainvoke=AsyncMock(return_value=mocks["_ai_msg"])) - ) - await h(_make_config(), "my question") + A single MagicMock cannot see a span tree at all: every span would be the same object, so a + parent and its children would be indistinguishable. + """ + + def __init__(self) -> None: + self.spans: list[RecordedSpan] = [] + + def get_tracer(self, name: str) -> SpanRecorder: + return self + + def start_span(self, name: str, context: Any = None) -> RecordedSpan: + span = RecordedSpan(name, context) + self.spans.append(span) + return span + + def set_span_in_context(self, span: RecordedSpan) -> Any: + return ("context-of", span) + + @property + def root(self) -> RecordedSpan: + return self.spans[0] + + def named(self, prefix: str) -> list[RecordedSpan]: + return [s for s in self.spans if s.name.startswith(prefix)] + + @property + def names(self) -> list[str]: + return [s.name for s in self.spans] + + +def _recording() -> Any: + """Patches the tracer that ``spans.py`` holds, and yields the recorder. + + ``AsyncCallbackHandler`` (the base class of ``SpanCallbackHandler``) stays real: LangChain's own + callback machinery decides when ``on_chat_model_start`` / ``on_llm_end`` / ``on_tool_start`` / + ``on_tool_end`` fire, and that dispatch is exactly what these tests need to prove, not something + to mock away. + """ + import launchdarkly_ai_langchain_agents.spans as spans_mod + + recorder = SpanRecorder() + return patch.object(spans_mod, "trace", recorder), recorder + + +BASE_CONFIG: dict[str, Any] = { + "model": {"name": "gpt-4o"}, + "provider": {"name": "OpenAI"}, + "instructions": "You are helpful.", +} + +TOOL_CONFIG: dict[str, Any] = { + **BASE_CONFIG, + "tools": { + "search": { + "description": "search the web", + "parameters": {"type": "object", "properties": {}}, + } + }, +} - event_calls = [ - c - for c in mock_span.add_event.call_args_list - if c[0][0] == "gen_ai.content.prompt" - ] - assert event_calls - # The gen_ai.prompt attribute must include the user input text - prompt_attr = event_calls[0][0][1].get("gen_ai.prompt", "") - assert "my question" in prompt_attr, ( - f"gen_ai.prompt must include user input 'my question', got: {prompt_attr!r}" - ) + +class TestSpanTree: + """TELEMETRY-CONTRACT.md section 1. A real langgraph agent and a real LangChain callback + dispatch, against a tracer this file controls.""" @pytest.mark.asyncio - async def test_token_attributes_set(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span + async def test_opens_a_root_span_named_invoke_agent(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("Hello!")]) + with ctx: + await create_langchain_agents_handler(llm)(BASE_CONFIG, "q") + assert rec.root.name == "invoke_agent" + assert rec.root.attributes["gen_ai.operation.name"] == "invoke_agent" - ai_msg = _make_ai_msg("answer", input_tokens=30, output_tokens=12) - mocks = _make_langchain_mock() + @pytest.mark.asyncio + async def test_emits_one_chat_child_per_model_turn(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("Hello!")]) + with ctx: + await create_langchain_agents_handler(llm)(BASE_CONFIG, "q") + chats = rec.named("chat ") + assert len(chats) == 1 + assert chats[0].name == "chat gpt-4o" + assert chats[0].attributes["gen_ai.operation.name"] == "chat" + assert chats[0].context == ("context-of", rec.root) - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler( - llm=MagicMock(ainvoke=AsyncMock(return_value=ai_msg)) - ) - await h(_make_config(), "hi") + @pytest.mark.asyncio + async def test_names_the_chat_span_after_the_model(self) -> None: + ctx, rec = _recording() + cfg = {**BASE_CONFIG, "model": {"name": "claude-sonnet-4-5"}} + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ctx: + await create_langchain_agents_handler(llm)(cfg, "q") + assert "chat claude-sonnet-4-5" in rec.names - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert ( - "gen_ai.usage.input_tokens" in calls - or "gen_ai.usage.output_tokens" in calls + @pytest.mark.asyncio + async def test_emits_a_chat_span_per_turn_of_a_tool_loop(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[ + _ai_message( + tool_calls=[{"name": "search", "args": {"q": "x"}, "id": "call_1"}] + ), + _ai_message("done"), + ] ) + with ctx: + await create_langchain_agents_handler(llm)( + TOOL_CONFIG, "q", {"search": AsyncMock(return_value="r")} + ) + assert len(rec.named("chat ")) == 2 @pytest.mark.asyncio - async def test_gen_ai_content_completion_event(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span + async def test_emits_an_execute_tool_span_per_tool_call(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[ + _ai_message( + tool_calls=[{"name": "search", "args": {"q": "x"}, "id": "call_1"}] + ), + _ai_message("done"), + ] + ) + with ctx: + await create_langchain_agents_handler(llm)( + TOOL_CONFIG, "q", {"search": AsyncMock(return_value="r")} + ) + tools = rec.named("execute_tool ") + assert len(tools) == 1 + assert tools[0].name == "execute_tool search" + assert tools[0].attributes["gen_ai.operation.name"] == "execute_tool" + assert tools[0].attributes["gen_ai.tool.name"] == "search" + assert tools[0].attributes["gen_ai.tool.call.id"] == "call_1" - mocks = _make_langchain_mock() - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler( - llm=MagicMock(ainvoke=AsyncMock(return_value=mocks["_ai_msg"])) - ) - await h(_make_config(), "hi") + @pytest.mark.asyncio + async def test_tool_spans_are_siblings_of_chat_not_children(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[ + _ai_message( + tool_calls=[{"name": "search", "args": {"q": "x"}, "id": "call_1"}] + ), + _ai_message("done"), + ] + ) + with ctx: + await create_langchain_agents_handler(llm)( + TOOL_CONFIG, "q", {"search": AsyncMock(return_value="r")} + ) + assert rec.named("execute_tool ")[0].context == ("context-of", rec.root) - event_calls = [ - c - for c in mock_span.add_event.call_args_list - if c[0][0] == "gen_ai.content.completion" - ] - assert event_calls + @pytest.mark.asyncio + async def test_every_span_is_ended(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[ + _ai_message( + tool_calls=[{"name": "search", "args": {"q": "x"}, "id": "call_1"}] + ), + _ai_message("done"), + ] + ) + with ctx: + await create_langchain_agents_handler(llm)( + TOOL_CONFIG, "q", {"search": AsyncMock(return_value="r")} + ) + assert [s.ended for s in rec.spans] == [1] * len(rec.spans) @pytest.mark.asyncio - async def test_gen_ai_response_model(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span + async def test_nests_the_chat_span_under_the_root_in_the_streaming_path( + self, + ) -> None: + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("Hello!")]) + with ctx: + async for _ in await create_langchain_agents_handler(llm).stream( + BASE_CONFIG, "q", {}, {} + ): + pass + chats = rec.named("chat ") + assert chats + assert chats[0].context == ("context-of", rec.root) + + +class TestRootSpanAttributes: + """TELEMETRY-CONTRACT.md sections 2, 2a and 9.""" - mocks = _make_langchain_mock() - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler( - llm=MagicMock(ainvoke=AsyncMock(return_value=mocks["_ai_msg"])) - ) - await h(_make_config(model={"name": "gpt-4o"}), "hi") + @pytest.mark.asyncio + async def test_gen_ai_provider_name_is_openai_for_a_non_anthropic_config( + self, + ) -> None: + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ctx: + await create_langchain_agents_handler(llm)(BASE_CONFIG, "q") + attrs = rec.root.attributes + # gen_ai.provider.name names who really served the model: ChatOpenAI here. gen_ai.system + # keeps the framework name so existing dashboards do not break. + assert attrs["gen_ai.provider.name"] == "openai" + assert attrs["gen_ai.system"] == "langchain" - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert "gen_ai.response.model" in calls - assert calls["gen_ai.response.model"] == "gpt-4o" + @pytest.mark.asyncio + async def test_gen_ai_provider_name_is_anthropic_when_config_names_it(self) -> None: + ctx, rec = _recording() + cfg = { + **BASE_CONFIG, + "provider": {"name": "Anthropic"}, + "model": {"name": "claude-sonnet-4-5"}, + } + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ctx: + await create_langchain_agents_handler(llm)(cfg, "q") + assert rec.root.attributes["gen_ai.provider.name"] == "anthropic" + assert rec.root.attributes["gen_ai.system"] == "langchain" @pytest.mark.asyncio - async def test_ld_span_attributes(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span + async def test_gen_ai_provider_name_falls_back_to_openai_for_anything_else( + self, + ) -> None: + # A binary choice, not a passthrough: Bedrock, Azure, Cohere, a typo, or nothing at all all + # report `openai`, because that mirrors which chat model class is really instantiated. + ctx, rec = _recording() + cfg = {**BASE_CONFIG, "provider": {"name": "Bedrock"}} + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ctx: + await create_langchain_agents_handler(llm)(cfg, "q") + assert rec.root.attributes["gen_ai.provider.name"] == "openai" - mocks = _make_langchain_mock() + @pytest.mark.asyncio + async def test_response_model_is_the_requested_name(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ctx: + await create_langchain_agents_handler(llm)(BASE_CONFIG, "q") + assert rec.root.attributes["gen_ai.response.model"] == "gpt-4o" + assert rec.root.attributes["gen_ai.request.model"] == "gpt-4o" + + @pytest.mark.asyncio + async def test_carries_the_launchdarkly_attributes_and_feature_flag_event( + self, + ) -> None: + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("hi")]) variables = { "__ld": { - "configKey": "my-config", - "variationKey": "v1", - "runId": "run-abc", + "configKey": "k", + "variationKey": "v", + "runId": "r", + "graphKey": "g", } } - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler( - llm=MagicMock(ainvoke=AsyncMock(return_value=mocks["_ai_msg"])) - ) - await h(_make_config(), "hi", variables=variables) + with ctx: + await create_langchain_agents_handler(llm)(BASE_CONFIG, "q", {}, variables) + attrs = rec.root.attributes + assert attrs["launchdarkly.operation.type"] == "gen_ai" + assert attrs["launchdarkly.config.key"] == "k" + assert attrs["launchdarkly.variation.key"] == "v" + assert attrs["launchdarkly.run.id"] == "r" + assert attrs["launchdarkly.graph.key"] == "g" + assert [n for n, _ in rec.root.events] == ["feature_flag"] - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("launchdarkly.operation.type") == "gen_ai" - assert calls.get("launchdarkly.config.key") == "my-config" - assert calls.get("launchdarkly.variation.key") == "v1" - assert calls.get("launchdarkly.run.id") == "run-abc" - assert "launchdarkly.graph.key" not in calls + @pytest.mark.asyncio + async def test_child_spans_carry_no_launchdarkly_identity(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[ + _ai_message( + tool_calls=[{"name": "search", "args": {"q": "x"}, "id": "call_1"}] + ), + _ai_message("done"), + ] + ) + variables = {"__ld": {"configKey": "k", "variationKey": "v", "runId": "r"}} + with ctx: + await create_langchain_agents_handler(llm)( + TOOL_CONFIG, "q", {"search": AsyncMock(return_value="r")}, variables + ) + for child in rec.spans[1:]: + assert not [k for k in child.attributes if k.startswith("launchdarkly.")] + assert "feature_flag" not in [n for n, _ in child.events] - async def test_ld_graph_key_set_when_present(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span + @pytest.mark.asyncio + async def test_carries_the_run_total_not_one_turn(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[ + _ai_message( + input_tokens=10, + output_tokens=1, + tool_calls=[{"name": "search", "args": {"q": "x"}, "id": "call_1"}], + ), + _ai_message("done", input_tokens=20, output_tokens=2), + ] + ) + with ctx: + await create_langchain_agents_handler(llm)( + TOOL_CONFIG, "q", {"search": AsyncMock(return_value="r")} + ) + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 30 + assert rec.root.attributes["gen_ai.usage.output_tokens"] == 3 + assert rec.root.attributes["gen_ai.usage.total_tokens"] == 33 - mocks = _make_langchain_mock() - variables = { - "__ld": { - "configKey": "my-config", - "variationKey": "v1", - "runId": "run-abc", - "graphKey": "my-graph", - } - } - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler( - llm=MagicMock(ainvoke=AsyncMock(return_value=mocks["_ai_msg"])) - ) - await h(_make_config(), "hi", variables=variables) - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("launchdarkly.graph.key") == "my-graph" +class TestChatSpanAttributes: + """TELEMETRY-CONTRACT.md sections 3, 5 and 8.""" + @pytest.mark.asyncio + async def test_writes_all_seven_usage_attributes_including_zeros(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[_ai_message("hi", input_tokens=10, output_tokens=5)] + ) + with ctx: + await create_langchain_agents_handler(llm)(BASE_CONFIG, "q") + attrs = rec.named("chat ")[0].attributes + assert attrs["gen_ai.usage.input_tokens"] == 10 + assert attrs["gen_ai.usage.output_tokens"] == 5 + assert attrs["gen_ai.usage.total_tokens"] == 15 + assert attrs["gen_ai.usage.cache_read.input_tokens"] == 0 + assert attrs["gen_ai.usage.cache_creation.input_tokens"] == 0 + assert attrs["gen_ai.usage.prompt_tokens"] == 10 + assert attrs["gen_ai.usage.completion_tokens"] == 5 -# --------------------------------------------------------------------------- -# §1.6 Error handling -# --------------------------------------------------------------------------- + @pytest.mark.asyncio + async def test_input_tokens_pass_through_untouched_cache_included(self) -> None: + # LangChain already includes cached tokens inside input_tokens, unlike Anthropic: the input + # figure must pass through as reported, and the cache figures are for parity only. This is + # the assertion that catches a fold applied in the wrong direction. + from langchain_core.messages import AIMessage + + reply = AIMessage( + content="hi", + usage_metadata={ + "input_tokens": 23554, + "output_tokens": 10, + "total_tokens": 23564, + "input_token_details": {"cache_read": 19971, "cache_creation": 3580}, + }, + ) + ctx, rec = _recording() + llm = _FakeToolModel(replies=[reply]) + with ctx: + await create_langchain_agents_handler(llm)(BASE_CONFIG, "q") + attrs = rec.named("chat ")[0].attributes + assert attrs["gen_ai.usage.input_tokens"] == 23554 + assert attrs["gen_ai.usage.cache_read.input_tokens"] == 19971 + assert attrs["gen_ai.usage.cache_creation.input_tokens"] == 3580 + assert attrs["gen_ai.usage.total_tokens"] == 23564 + 10 - 10 # input + output + + @pytest.mark.asyncio + async def test_reports_the_mapped_finish_reason(self) -> None: + # LangChain does not normalise the field; the handler reads response_metadata.stop_reason + # and maps Anthropic's `end_turn` onto semconv's `stop`. + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[_ai_message("hi", response_metadata={"stop_reason": "end_turn"})] + ) + with ctx: + await create_langchain_agents_handler(llm)(BASE_CONFIG, "q") + assert rec.named("chat ")[0].attributes["gen_ai.response.finish_reasons"] == [ + "stop" + ] + @pytest.mark.asyncio + async def test_maps_openai_finish_reason_too(self) -> None: + # This handler can serve either vendor; both spellings must map onto one vocabulary. + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[_ai_message("hi", response_metadata={"finish_reason": "stop"})] + ) + with ctx: + await create_langchain_agents_handler(llm)(BASE_CONFIG, "q") + assert rec.named("chat ")[0].attributes["gen_ai.response.finish_reasons"] == [ + "stop" + ] -class TestErrorHandling: @pytest.mark.asyncio - async def test_records_exception_on_span(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span + async def test_omits_the_finish_reason_when_the_provider_gives_none(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ctx: + await create_langchain_agents_handler(llm)(BASE_CONFIG, "q") + assert "gen_ai.response.finish_reasons" not in rec.named("chat ")[0].attributes - mocks = _make_langchain_mock() - mocks["_agent"].ainvoke = AsyncMock(side_effect=RuntimeError("lc error")) - mocks["langgraph.prebuilt"].create_react_agent = MagicMock( - return_value=mocks["_agent"] + @pytest.mark.asyncio + async def test_sets_status_ok_on_a_successful_turn(self) -> None: + from opentelemetry.trace import StatusCode + + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ctx: + await create_langchain_agents_handler(llm)(BASE_CONFIG, "q") + assert StatusCode.OK in rec.named("chat ")[0].statuses + + +class TestContentCapture: + """TELEMETRY-CONTRACT.md section 7.""" + + @pytest.mark.asyncio + async def test_emits_no_content_at_all_by_default(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ctx: + await create_langchain_agents_handler(llm)(BASE_CONFIG, "q") + for span in rec.spans: + content_keys = [ + k + for k in span.attributes + if k.startswith(("gen_ai.prompt", "gen_ai.completion")) + or k + in ( + "gen_ai.input.messages", + "gen_ai.output.messages", + "gen_ai.system_instructions", + "gen_ai.tool.definitions", + ) + ] + assert content_keys == [] + assert [n for n, _ in span.events if n.startswith("gen_ai.content")] == [] + + @pytest.mark.asyncio + async def test_puts_prompt_and_completion_on_the_root_when_enabled(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("Hello World")]) + with ctx: + await create_langchain_agents_handler(llm, capture_content=True)( + BASE_CONFIG, "q" + ) + root = rec.root + assert root.attributes["gen_ai.system_instructions"] + assert "gen_ai.input.messages" in root.attributes + assert root.attributes["gen_ai.completion.0.content"] == "Hello World" + assert "gen_ai.output.messages" in root.attributes + + @pytest.mark.asyncio + async def test_records_the_tool_catalog_on_the_chat_span_when_enabled(self) -> None: + import json + + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ctx: + await create_langchain_agents_handler(llm, capture_content=True)( + TOOL_CONFIG, "q", {"search": AsyncMock(return_value="r")} + ) + definitions = json.loads( + rec.named("chat ")[0].attributes["gen_ai.tool.definitions"] ) + assert definitions[0]["name"] == "search" + assert definitions[0]["type"] == "function" - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler( - llm=MagicMock(ainvoke=AsyncMock(return_value=None)) - ) - with pytest.raises(RuntimeError): - await h(_make_config(), "hi") + @pytest.mark.asyncio + async def test_records_tool_arguments_and_results_when_enabled(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[ + _ai_message( + tool_calls=[ + {"name": "search", "args": {"q": "weather"}, "id": "call_1"} + ] + ), + _ai_message("done"), + ] + ) + with ctx: + await create_langchain_agents_handler(llm, capture_content=True)( + TOOL_CONFIG, "q", {"search": AsyncMock(return_value="72F")} + ) + tool = rec.named("execute_tool ")[0] + assert "weather" in tool.attributes["gen_ai.tool.call.arguments"] + assert tool.attributes["gen_ai.tool.call.result"] == "72F" - mock_span.record_exception.assert_called() + +class TestErrorHandling: + """TELEMETRY-CONTRACT.md section 6.""" @pytest.mark.asyncio - async def test_sets_span_status_error(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span + async def test_fails_the_chat_span_when_the_provider_call_raises(self) -> None: + from opentelemetry.trace import StatusCode - mocks = _make_langchain_mock() - mocks["_agent"].ainvoke = AsyncMock(side_effect=RuntimeError("lc error")) + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[], fail_after=0, fail_with=RuntimeError("model down") + ) + with ctx, pytest.raises(RuntimeError, match="model down"): + await create_langchain_agents_handler(llm)(BASE_CONFIG, "q") + chat = rec.named("chat ")[0] + assert len(chat.exceptions) == 1 + assert StatusCode.ERROR in chat.statuses + assert chat.ended == 1 - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler( - llm=MagicMock(ainvoke=AsyncMock()) - ) - with pytest.raises(RuntimeError): - await h(_make_config(), "hi") + @pytest.mark.asyncio + async def test_fails_the_root_span_too(self) -> None: + from opentelemetry.trace import StatusCode - mock_span.set_status.assert_called() + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[], fail_after=0, fail_with=RuntimeError("model down") + ) + with ctx, pytest.raises(RuntimeError): + await create_langchain_agents_handler(llm)(BASE_CONFIG, "q") + assert len(rec.root.exceptions) == 1 + assert StatusCode.ERROR in rec.root.statuses + assert rec.root.ended == 1 @pytest.mark.asyncio - async def test_ends_span_on_error(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span + async def test_fails_the_execute_tool_span_when_a_tool_raises(self) -> None: + from opentelemetry.trace import StatusCode - mocks = _make_langchain_mock() - mocks["_agent"].ainvoke = AsyncMock(side_effect=RuntimeError("lc error")) + async def _boom(_args: Any) -> str: + raise RuntimeError("tool exploded") - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler( - llm=MagicMock(ainvoke=AsyncMock()) - ) - with pytest.raises(RuntimeError): - await h(_make_config(), "hi") + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[ + _ai_message( + tool_calls=[{"name": "search", "args": {"q": "x"}, "id": "call_1"}] + ) + ] + ) + with ctx, pytest.raises(Exception, match="tool exploded"): + await create_langchain_agents_handler(llm)( + TOOL_CONFIG, "q", {"search": _boom} + ) + tool = rec.named("execute_tool ")[0] + assert len(tool.exceptions) == 1 + assert StatusCode.ERROR in tool.statuses + assert tool.ended == 1 - mock_span.end.assert_called() + @pytest.mark.asyncio + async def test_reports_the_spend_of_completed_turns_on_a_failed_run(self) -> None: + # The first turn was billed. The root is the only span a config-scoped cost query finds it + # on. There is no `result` to sum on this path, so the total comes from the callbacks. + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[ + _ai_message( + input_tokens=40, + output_tokens=7, + tool_calls=[{"name": "search", "args": {"q": "x"}, "id": "call_1"}], + ) + ], + fail_after=1, + fail_with=RuntimeError("second turn died"), + ) + with ctx, pytest.raises(RuntimeError): + await create_langchain_agents_handler(llm)( + TOOL_CONFIG, "q", {"search": AsyncMock(return_value="r")} + ) + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 40 + assert rec.root.attributes["gen_ai.usage.output_tokens"] == 7 + + @pytest.mark.asyncio + async def test_writes_no_usage_when_no_turn_ever_reported_any(self) -> None: + # All-zero attributes would assert the run cost nothing; an absent attribute says "unknown". + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[], fail_after=0, fail_with=RuntimeError("died first") + ) + with ctx, pytest.raises(RuntimeError): + await create_langchain_agents_handler(llm)(BASE_CONFIG, "q") + assert "gen_ai.usage.input_tokens" not in rec.root.attributes @pytest.mark.asyncio async def test_rethrows_error(self) -> None: @@ -779,7 +1109,7 @@ async def test_rethrows_error(self) -> None: mocks["_agent"].ainvoke = AsyncMock(side_effect=RuntimeError("specific error")) with _patch_lc(mocks): - with patch.object(handler_mod, "_HAS_OTEL", False): + with patch.object(spans_mod, "_HAS_OTEL", False): h = create_langchain_agents_handler(llm=MagicMock()) with pytest.raises(RuntimeError, match="specific error"): await h(_make_config(), "hi") @@ -807,7 +1137,6 @@ def test_passes_config_key_user_input_and_context(self) -> None: assert "context" in sig.parameters def test_config_key_forwarded_as_key(self) -> None: - import launchdarkly_ai_langchain_agents.handler as handler_mod mock_config_instance = MagicMock() mock_config_fn = MagicMock(return_value=mock_config_instance) @@ -830,7 +1159,6 @@ def test_config_key_forwarded_as_key(self) -> None: ) def test_callable_without_extra_kwargs(self) -> None: - import launchdarkly_ai_langchain_agents.handler as handler_mod mock_config_instance = MagicMock() mock_config_fn = MagicMock(return_value=mock_config_instance) @@ -871,7 +1199,7 @@ async def _mock_astream(*a: Any, **kw: Any) -> AsyncIterator[Any]: mocks["_agent"].astream = _mock_astream with _patch_lc(mocks): - with patch.object(handler_mod, "_HAS_OTEL", False): + with patch.object(spans_mod, "_HAS_OTEL", False): h = create_langchain_agents_handler(llm=MagicMock()) gen = await h.stream(_make_config(), "hi") assert inspect.isasyncgen(gen) or hasattr(gen, "__aiter__") @@ -887,7 +1215,7 @@ async def _mock_astream(*a: Any, **kw: Any) -> AsyncIterator[Any]: mocks["_agent"].astream = _mock_astream with _patch_lc(mocks): - with patch.object(handler_mod, "_HAS_OTEL", False): + with patch.object(spans_mod, "_HAS_OTEL", False): h = create_langchain_agents_handler(llm=MagicMock()) events = [e async for e in await h.stream(_make_config(), "hi")] @@ -905,7 +1233,7 @@ async def _bad_astream(*a: Any, **kw: Any) -> AsyncIterator[Any]: mocks["_agent"].astream = _bad_astream with _patch_lc(mocks): - with patch.object(handler_mod, "_HAS_OTEL", False): + with patch.object(spans_mod, "_HAS_OTEL", False): h = create_langchain_agents_handler(llm=MagicMock()) with pytest.raises(RuntimeError, match="stream fail"): async for _ in await h.stream(_make_config(), "hi"): @@ -918,83 +1246,124 @@ async def _bad_astream(*a: Any, **kw: Any) -> AsyncIterator[Any]: class TestStreamingTelemetry: - @pytest.mark.asyncio - async def test_span_started_during_stream(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - mocks = _make_langchain_mock() - - async def _empty_astream(*a: Any, **kw: Any) -> AsyncIterator[Any]: - return - yield - - mocks["_agent"].astream = _empty_astream - - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler(llm=MagicMock()) - async for _ in await h.stream(_make_config(), "hi"): - pass - - mock_trace.get_tracer.return_value.start_span.assert_called_with( - "langchain.agent.stream" - ) + """TELEMETRY-CONTRACT.md sections 1 and 6. The streaming path emits the same tree.""" @pytest.mark.asyncio - async def test_ld_span_attributes_set_during_stream(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - mocks = _make_langchain_mock() - - async def _empty_astream(*a: Any, **kw: Any) -> AsyncIterator[Any]: - return - yield + async def test_opens_the_same_root_span_name_as_the_blocking_path(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ctx: + async for _ in await create_langchain_agents_handler(llm).stream( + BASE_CONFIG, "q", {}, {} + ): + pass + assert rec.root.name == "invoke_agent" + assert "chat gpt-4o" in rec.names - mocks["_agent"].astream = _empty_astream + @pytest.mark.asyncio + async def test_carries_the_launchdarkly_attributes_on_the_root(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("hi")]) variables = {"__ld": {"configKey": "k", "variationKey": "v", "runId": "r"}} + with ctx: + async for _ in await create_langchain_agents_handler(llm).stream( + BASE_CONFIG, "q", {}, variables + ): + pass + attrs = rec.root.attributes + assert attrs["launchdarkly.operation.type"] == "gen_ai" + assert attrs["launchdarkly.config.key"] == "k" + assert attrs["launchdarkly.variation.key"] == "v" + assert attrs["launchdarkly.run.id"] == "r" - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler(llm=MagicMock()) - async for _ in await h.stream( - _make_config(), "hi", None, variables - ): - pass - - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("launchdarkly.operation.type") == "gen_ai" - assert calls.get("launchdarkly.config.key") == "k" - assert calls.get("launchdarkly.variation.key") == "v" - assert calls.get("launchdarkly.run.id") == "r" + @pytest.mark.asyncio + async def test_ends_every_span_once_when_the_stream_completes(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ctx: + async for _ in await create_langchain_agents_handler(llm).stream( + BASE_CONFIG, "q", {}, {} + ): + pass + assert [s.ended for s in rec.spans] == [1] * len(rec.spans) + assert "launchdarkly.stream.abandoned" not in rec.root.attributes @pytest.mark.asyncio - async def test_span_ended_after_stream_completes(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span + async def test_writes_the_run_total_to_the_root(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[_ai_message("hi", input_tokens=11, output_tokens=4)] + ) + with ctx: + async for _ in await create_langchain_agents_handler(llm).stream( + BASE_CONFIG, "q", {}, {} + ): + pass + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 11 + assert rec.root.attributes["gen_ai.usage.total_tokens"] == 15 - mocks = _make_langchain_mock() + @pytest.mark.asyncio + async def test_an_abandoned_stream_still_ends_and_exports_every_span(self) -> None: + # A consumer that breaks out of `async for` makes this generator run `finally` without ever + # entering `except`: GeneratorExit is a BaseException. Without the cleanup there the root is + # never ended, so it is never exported. + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ctx: + gen = await create_langchain_agents_handler(llm).stream( + BASE_CONFIG, "q", {}, {} + ) + async for _ in gen: + break + await gen.aclose() + assert rec.root.ended == 1 + assert [s.ended for s in rec.spans] == [1] * len(rec.spans) - async def _empty_astream(*a: Any, **kw: Any) -> AsyncIterator[Any]: - return - yield + @pytest.mark.asyncio + async def test_an_abandoned_stream_is_marked_but_not_failed(self) -> None: + from opentelemetry.trace import StatusCode + + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ctx: + gen = await create_langchain_agents_handler(llm).stream( + BASE_CONFIG, "q", {}, {} + ) + async for _ in gen: + break + await gen.aclose() + assert rec.root.attributes["launchdarkly.stream.abandoned"] is True + assert StatusCode.ERROR not in rec.root.statuses + assert rec.root.exceptions == [] - mocks["_agent"].astream = _empty_astream + @pytest.mark.asyncio + async def test_fails_the_spans_when_the_stream_raises(self) -> None: + from opentelemetry.trace import StatusCode - with _patch_lc(mocks): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_langchain_agents_handler(llm=MagicMock()) - async for _ in await h.stream(_make_config(), "hi"): - pass + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[], fail_after=0, fail_with=RuntimeError("stream died") + ) + with ctx, pytest.raises(RuntimeError, match="stream died"): + async for _ in await create_langchain_agents_handler(llm).stream( + BASE_CONFIG, "q", {}, {} + ): + pass + assert StatusCode.ERROR in rec.root.statuses + assert rec.root.ended == 1 - mock_span.end.assert_called() + @pytest.mark.asyncio + async def test_emits_no_content_by_default_on_the_streaming_path(self) -> None: + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ctx: + async for _ in await create_langchain_agents_handler(llm).stream( + BASE_CONFIG, "q", {}, {} + ): + pass + for span in rec.spans: + assert [k for k in span.attributes if k.startswith("gen_ai.prompt")] == [] + assert [n for n, _ in span.events if n.startswith("gen_ai.content")] == [] # --------------------------------------------------------------------------- @@ -1007,7 +1376,7 @@ class TestOutputFormat: async def test_absent_output_format_no_change(self) -> None: mocks = _make_langchain_mock() with _patch_lc(mocks): - with patch.object(handler_mod, "_HAS_OTEL", False): + with patch.object(spans_mod, "_HAS_OTEL", False): h = create_langchain_agents_handler( llm=MagicMock(ainvoke=AsyncMock(return_value=mocks["_ai_msg"])) ) @@ -1028,7 +1397,7 @@ def _capture_agent(*args: Any, **kw: Any) -> Any: ) with _patch_lc(mocks): - with patch.object(handler_mod, "_HAS_OTEL", False): + with patch.object(spans_mod, "_HAS_OTEL", False): h = create_langchain_agents_handler( llm=MagicMock(ainvoke=AsyncMock(return_value=mocks["_ai_msg"])) ) @@ -1128,3 +1497,378 @@ def test_history_without_prior_system_prompt(self) -> None: assert system is not None assert "Conversation History:" in system assert "user: What is feature flagging?" in system + + +class TestAbandonOpenSpans: + """An early consumer stop must not look like a provider failure. + + The abandonment path used to reuse `close_open_spans`, which records a synthetic exception and + sets ERROR on every span still open. TELEMETRY-CONTRACT.md section 6 says an abandoned span stays + UNSET and carries `launchdarkly.stream.abandoned`, and `openai-agents` already did that. + + Tested directly on the callback handler rather than through the streaming path. Reaching the + state that matters, a chat or tool span still open at the break, needs a fake model that yields + mid-turn, and with the fixtures here LangGraph has already run every callback by the time the + first chunk reaches the consumer. A test driven through `stream` therefore passes whether or not + the fix is present, which is worse than no test. + """ + + def _handler_with_open_spans(self) -> tuple[Any, Any, Any]: + import launchdarkly_ai_langchain_agents.spans as spans_mod + from launchdarkly_ai_langchain_agents.spans import build_span_callbacks + + recorder = SpanRecorder() + with patch.object(spans_mod, "trace", recorder): + bundle = build_span_callbacks(BASE_CONFIG, None, capture_content=False) + handler = bundle._handler + chat = recorder.start_span("chat gpt-4o") + tool = recorder.start_span("execute_tool search") + handler.model_spans["run-1"] = chat + handler.tool_spans["run-2"] = tool + return bundle, chat, tool + + def test_marks_open_spans_abandoned_and_leaves_them_unset(self) -> None: + from opentelemetry.trace import StatusCode + + bundle, chat, tool = self._handler_with_open_spans() + bundle.abandon_open_spans(set()) + for span in (chat, tool): + assert span.ended == 1 + assert span.attributes["launchdarkly.stream.abandoned"] is True + assert StatusCode.ERROR not in span.statuses + assert span.exceptions == [] + + def test_close_open_spans_still_fails_them_for_a_real_error(self) -> None: + # The failure path keeps its behaviour; only abandonment changed. + from opentelemetry.trace import StatusCode + + bundle, chat, tool = self._handler_with_open_spans() + bundle.close_open_spans(RuntimeError("provider died")) + for span in (chat, tool): + assert span.ended == 1 + assert StatusCode.ERROR in span.statuses + assert len(span.exceptions) == 1 + assert "launchdarkly.stream.abandoned" not in span.attributes + + def test_abandoning_twice_ends_each_span_once(self) -> None: + bundle, chat, tool = self._handler_with_open_spans() + ended: set[int] = set() + bundle.abandon_open_spans(ended) + bundle.abandon_open_spans(ended) + assert chat.ended == 1 + assert tool.ended == 1 + + +class TestConvenienceWrapperForwardsCaptureContent: + """`capture_content` must reach the handler, not fall through into `config()`. + + `config()` takes no such argument, so leaving it in kwargs raised TypeError: a caller asking for + content on spans got an exception instead. Five of the six wrappers had this. + """ + + def _run(self, **kwargs: Any) -> dict[str, Any]: + + seen: dict[str, Any] = {} + + def _factory(*args: Any, capture_content: bool = False, **kw: Any) -> Any: + seen["capture_content"] = capture_content + return MagicMock() + + fake_config = MagicMock() + fake_config.return_value.invoke = MagicMock(return_value="ok") + with ( + patch.object(handler_mod, "create_langchain_agents_handler", _factory), + patch.object(handler_mod, "config", fake_config), + ): + handler_mod.langchain_agents("k", "q", {}, **kwargs) + seen["config_kwargs"] = fake_config.call_args.kwargs + return seen + + def test_capture_content_reaches_the_factory(self) -> None: + seen = self._run(capture_content=True) + assert seen["capture_content"] is True + # And it must not have been forwarded to config(), which does not accept it. + assert "capture_content" not in seen["config_kwargs"] + + def test_defaults_to_off(self) -> None: + assert self._run()["capture_content"] is False + + +class TestCallbackSpansNeverLeak: + """A raise inside a callback must not leave a span both untracked and unended. + + The end callbacks pop the span before doing work that can raise, so after the pop nothing else + can reach it and ending it is that callback's job alone. `on_tool_start` had the mirror problem: + it created the span before inserting it, so a raise in between left a span no cleanup path knew + about. + """ + + class _Unserialisable: + __slots__ = () + + def _recording_handler(self) -> Any: + """The patch must stay active while the callbacks run, not only while they are built.""" + import launchdarkly_ai_langchain_agents.spans as spans_mod + + recorder = SpanRecorder() + return patch.object(spans_mod, "trace", recorder), recorder + + @pytest.mark.asyncio + async def test_a_raise_in_on_tool_end_still_ends_the_span(self) -> None: + from opentelemetry.trace import StatusCode + + from launchdarkly_ai_langchain_agents.spans import build_span_callbacks + + ctx, recorder = self._recording_handler() + with ctx: + handler = build_span_callbacks( + BASE_CONFIG, None, capture_content=True + )._handler + await handler.on_tool_start({"name": "search"}, "{}", run_id="r1") + assert "r1" in handler.tool_spans + with pytest.raises(TypeError): + await handler.on_tool_end(self._Unserialisable(), run_id="r1") + + span = recorder.named("execute_tool ")[0] + assert span.ended == 1, "the tool span leaked" + assert StatusCode.ERROR in span.statuses + + @pytest.mark.asyncio + async def test_on_tool_start_tracks_the_span_before_it_can_raise(self) -> None: + from launchdarkly_ai_langchain_agents.spans import build_span_callbacks + + ctx, recorder = self._recording_handler() + with ctx: + bundle = build_span_callbacks(BASE_CONFIG, None, capture_content=True) + handler = bundle._handler + with pytest.raises(TypeError): + await handler.on_tool_start( + {"name": "search"}, "{}", run_id="r1", inputs=self._Unserialisable() + ) + # Tracked despite the raise, so the caller's cleanup can still close it. + assert "r1" in handler.tool_spans + bundle.abandon_open_spans(set()) + + assert recorder.named("execute_tool ")[0].ended == 1 + + +class TestModelSpanTrackedBeforeContent: + """`_start_model` must insert the span before writing content that can raise. + + A span created but never inserted is unreachable by close_open_spans, abandon_open_spans and the + end callbacks alike, so it never ends and never exports. `on_tool_start` already had this fix. + """ + + @pytest.mark.asyncio + async def test_a_raise_while_writing_input_still_leaves_the_span_reachable( + self, + ) -> None: + import launchdarkly_ai_langchain_agents.spans as spans_mod + from launchdarkly_ai_langchain_agents.spans import build_span_callbacks + + class _Exploding: + def _get_type(self) -> str: + raise TypeError("cannot narrow this message") + + recorder = SpanRecorder() + with patch.object(spans_mod, "trace", recorder): + bundle = build_span_callbacks(BASE_CONFIG, None, capture_content=True) + handler = bundle._handler + with pytest.raises(TypeError): + await handler.on_chat_model_start({}, [[_Exploding()]], run_id="r1") + assert "r1" in handler.model_spans, "the chat span was never tracked" + bundle.abandon_open_spans(set()) + + assert recorder.named("chat ")[0].ended == 1 + + +class TestZeroTokensAreStillReported: + """A reported 0 is not a missing count. + + `extract_llm_usage` read the llm_output fallback with `or`, so a genuine 0 was skipped. With both + counts at zero the bag came back all None, `lang_chain_span_usage` read that as "the provider + said nothing", and the run went unreported: a turn that completed and cost nothing became + indistinguishable from one that never reported. + """ + + def test_a_zero_prompt_count_survives(self) -> None: + from launchdarkly_ai_langchain_agents.spans import extract_llm_usage + + result = SimpleNamespace( + generations=[], + llm_output={"token_usage": {"prompt_tokens": 0, "completion_tokens": 7}}, + ) + assert extract_llm_usage(result) == {"input_tokens": 0, "output_tokens": 7} + + def test_both_zero_still_counts_as_reported(self) -> None: + from launchdarkly_ai_langchain_agents.spans import extract_llm_usage + from launchdarkly_ai_server import lang_chain_span_usage + + result = SimpleNamespace( + generations=[], + llm_output={"token_usage": {"prompt_tokens": 0, "completion_tokens": 0}}, + ) + raw = extract_llm_usage(result) + assert raw == {"input_tokens": 0, "output_tokens": 0} + # The distinction that matters downstream: this is a reported turn, not an absent one. + assert lang_chain_span_usage(raw) is not None + + def test_an_absent_count_is_still_none(self) -> None: + from launchdarkly_ai_langchain_agents.spans import extract_llm_usage + + result = SimpleNamespace(generations=[], llm_output={"token_usage": {}}) + assert extract_llm_usage(result) == {} + + +class TestSuccessAndFailureUsageAgree: + """A successful root must not report zero while its own chat spans report real tokens. + + The two sides read different fields. The success path sums `usage_metadata` off each message; the + callbacks read the whole `LLMResult` and fall back to `llm_output.token_usage`, which some + providers use instead. For those providers the chat spans carried the tokens and the successful + run's root, and the bag handed back to the caller, both stayed at zero. + """ + + @pytest.mark.asyncio + async def test_llm_output_only_usage_still_reaches_the_root(self) -> None: + # A provider that reports in llm_output and leaves usage_metadata empty. + class _LlmOutputOnlyModel(_FakeToolModel): + async def _agenerate( + self, + messages: Any, + stop: Any = None, + run_manager: Any = None, + **kwargs: Any, + ) -> Any: + from langchain_core.messages import AIMessage + from langchain_core.outputs import ChatGeneration, ChatResult + + self.calls += 1 + return ChatResult( + generations=[ChatGeneration(message=AIMessage(content="hi"))], + llm_output={ + "token_usage": {"prompt_tokens": 41, "completion_tokens": 9} + }, + ) + + ctx, rec = _recording() + result = None + with ctx: + result = await create_langchain_agents_handler(_LlmOutputOnlyModel())( + BASE_CONFIG, "q" + ) + chat = rec.named("chat ")[0].attributes + assert chat["gen_ai.usage.input_tokens"] == 41 + # The point of the test: the root and the returned bag must agree with the chat span. + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 41 + assert result["usage"] == {"input_tokens": 41, "output_tokens": 9} + + @pytest.mark.asyncio + async def test_usage_metadata_still_wins_when_both_are_present(self) -> None: + # The message-level sum stays authoritative where it has anything to say, so this change + # cannot double-count a provider that reports in both places. + ctx, rec = _recording() + llm = _FakeToolModel( + replies=[_ai_message("hi", input_tokens=10, output_tokens=5)] + ) + with ctx: + result = await create_langchain_agents_handler(llm)(BASE_CONFIG, "q") + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 10 + assert result["usage"] == {"input_tokens": 10, "output_tokens": 5} + + +class TestCallbackKeepsBilledTokens: + """The callback must report the tokens a turn spent even when writing its content fails. + + `on_llm_end` accumulated after the content write. A raise while serialising completion content + dropped a turn the provider had already billed, and that accumulator is what a failed run's root + reports and what a successful run falls back to when the messages carry no usage of their own. + """ + + @pytest.mark.asyncio + async def test_a_content_failure_does_not_lose_the_tokens_already_billed( + self, + ) -> None: + from langchain_core.outputs import ChatGeneration, LLMResult + + import launchdarkly_ai_langchain_agents.spans as spans_mod + + callbacks = spans_mod.build_span_callbacks(BASE_CONFIG, None, True, []) + handler = callbacks._handler + + result = LLMResult( + generations=[ + [ + ChatGeneration( + message=_ai_message("hi", input_tokens=61, output_tokens=14) + ) + ] + ], + ) + + def _explode(*_a: Any, **_k: Any) -> None: + raise TypeError("cannot serialise this content") + + await handler.on_llm_start({}, ["q"], run_id="r1") + with patch.object(spans_mod, "set_output_content_attributes", _explode): + with pytest.raises(TypeError): + await handler.on_llm_end(result, run_id="r1") + + # The turn was billed before the write that failed, so its tokens are still counted. + assert callbacks.run_usage.total.input == 61 + assert callbacks.run_usage.total.output == 14 + + +class TestInputWritesNeverLeakASpan: + """Serialising the prompt must not be able to strand the root span. + + The input content write ran before the guard that fails the root, so a raise there left it open: + never ended, never exported, so the run disappeared from AI Config Monitoring along with the + feature_flag event it carries. + """ + + @pytest.mark.asyncio + async def test_the_blocking_root_still_ends(self) -> None: + from opentelemetry.trace import StatusCode + + import launchdarkly_ai_langchain_agents.handler as handler_mod + + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ( + ctx, + patch.object( + handler_mod, + "set_input_content_attributes", + side_effect=TypeError("cannot serialise this prompt"), + ), + pytest.raises(TypeError), + ): + await create_langchain_agents_handler(llm, capture_content=True)( + BASE_CONFIG, "q" + ) + + assert rec.root.ended == 1, "the root span leaked" + assert StatusCode.ERROR in rec.root.statuses + + @pytest.mark.asyncio + async def test_the_streaming_root_still_ends(self) -> None: + import launchdarkly_ai_langchain_agents.handler as handler_mod + + ctx, rec = _recording() + llm = _FakeToolModel(replies=[_ai_message("hi")]) + with ( + ctx, + patch.object( + handler_mod, + "set_input_content_attributes", + side_effect=TypeError("cannot serialise this prompt"), + ), + pytest.raises(TypeError), + ): + async for _ in await create_langchain_agents_handler( + llm, capture_content=True + ).stream(BASE_CONFIG, "q"): + pass + + assert rec.root.ended == 1, "the root span leaked"