diff --git a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py index 32fce47..7e451ae 100644 --- a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py +++ b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py @@ -3,23 +3,46 @@ import asyncio import json from collections.abc import AsyncGenerator +from types import SimpleNamespace from typing import Any from launchdarkly_ai_server import ( AiConfigRep, LDContext, ProviderHandler, + SpanMessage, + SpanMessagePart, + SpanUsage, config, create_handler, + create_run_usage, + end_span_once, + lang_chain_finish_reasons, + lang_chain_span_messages, + lang_chain_span_usage, + number_or_zero, parse_template, - set_ld_span_attributes, - set_openllmetry_completion, - set_openllmetry_prompt, + set_input_content_attributes, + set_output_content_attributes, + set_tool_call_content_attributes, +) + +from .spans import ( + fail_span, + finish_model_span, + finish_root_span, + mark_ok, + parent_context_of, + start_model_span, + start_root_span, + start_tool_span, + succeed_span, + to_tool_definitions, ) try: - from opentelemetry import trace - from opentelemetry.trace import StatusCode as SpanStatusCode + from opentelemetry import trace # noqa: F401 + from opentelemetry.trace import StatusCode as SpanStatusCode # noqa: F401 _HAS_OTEL = True except ImportError: @@ -27,6 +50,10 @@ def _build_tools(config_tools: dict[str, Any]) -> list[dict[str, Any]]: + # Not filtered to the tools that have a registered handler, unlike the TypeScript SDK's + # `buildTools`. That difference predates this span work and changes what the model is offered, + # not what the span reports, so it stays as it is: the catalog recorded below (via + # `to_tool_definitions`) is the catalog actually sent. return [ { "type": "function", @@ -95,6 +122,54 @@ def _build_messages( return messages +def _assistant_output_messages( + content: Any, tool_calls: list[Any] | None +) -> list[SpanMessage]: + """One assistant reply as a canonical output message, tool calls included. + + Built as a duck-typed stand-in rather than a real ``AIMessage``, matching the TypeScript + handler's ``assistantOutput``: :func:`lang_chain_span_messages` only reads ``_get_type()``, + ``content`` and ``tool_calls``, so constructing a real message here would be extra ceremony for + fields nothing downstream reads. + """ + msg = SimpleNamespace(content=content, tool_calls=tool_calls or []) + msg._get_type = lambda: "ai" + _, messages = lang_chain_span_messages([msg]) + return messages + + +def _with_get_type(msg: Any) -> Any: + """Adapts one LangChain message to the interface the client's ``lang_chain_span_messages`` + narrows on. + + Works around a version-skew bug in the shared client helper rather than fixing it there: + ``lang_chain_span_messages`` reads a message's role off ``_get_type()``, which older LangChain + releases exposed as the canonical accessor. The ``langchain-core`` release this package + actually depends on replaced it with a plain ``type`` field and dropped the method entirely, so + every real ``SystemMessage``/``HumanMessage``/``AIMessage`` reaching the client helper + unmodified is misclassified as role ``user`` with no error raised: ``getattr(raw, + '_get_type', None)`` returns ``None`` for a missing attribute rather than raising, and the + caller has no way to tell "the method is absent" from "this message really has no type". + Reported in this package's TELEMETRY-CONTRACT.md report rather than patched in + ``packages/client``, which is out of scope for this change. + """ + if callable(getattr(msg, "_get_type", None)): + return msg + msg_type = getattr(msg, "type", None) + shim = SimpleNamespace( + content=getattr(msg, "content", None), + tool_calls=getattr(msg, "tool_calls", None), + tool_call_id=getattr(msg, "tool_call_id", None), + ) + shim._get_type = lambda: msg_type + return shim + + +def _span_messages(messages: list[Any]) -> tuple[str | None, list[SpanMessage]]: + """``lang_chain_span_messages``, after adapting each message. See :func:`_with_get_type`.""" + return lang_chain_span_messages([_with_get_type(m) for m in messages]) + + def _is_coroutine(fn: Any) -> bool: return asyncio.iscoroutinefunction(fn) @@ -119,14 +194,89 @@ def _make_default_chat_model(config: AiConfigRep, importlib: Any) -> Any: return lc_openai.ChatOpenAI(model=model_name or "gpt-4o") -def create_langchain_messages_handler(llm: Any = None) -> ProviderHandler: +async def _run_structured_turn( + base_model: Any, + config: AiConfigRep, + messages: list[Any], + output_format: dict[str, Any], + parent: Any, + *, + capture_content: bool, + run_usage: Any, +) -> Any: + """Runs one structured-output turn under its own ``chat`` child span. Returns the parsed value. + + Used both for the outputFormat-only path and for the final turn of a tool loop that also has an + outputFormat, mirroring the TypeScript handler's ``runStructuredTurn``. + """ + model_span = start_model_span(config, parent) + if capture_content: + system_instructions, span_messages = _span_messages(messages) + set_input_content_attributes( + model_span, + capture_content, + system_instructions=system_instructions, + messages=span_messages, + ) + try: + structured_model = base_model.with_structured_output( + output_format, include_raw=True + ) + result = await structured_model.ainvoke(messages) + except Exception as exc: + fail_span(model_span, exc) + raise + + raw: Any = ( + result.get("raw") if isinstance(result, dict) else getattr(result, "raw", None) + ) + raw_usage = getattr(raw, "usage_metadata", None) or {} + parsed: Any = ( + result.get("parsed") + if isinstance(result, dict) + else getattr(result, "parsed", None) + ) + if capture_content: + set_output_content_attributes( + model_span, + capture_content, + [ + SpanMessage( + role="assistant", + parts=[ + SpanMessagePart( + type="text", + content=parsed + if isinstance(parsed, str) + else json.dumps(parsed), + ) + ], + ) + ], + ) + finish_model_span( + model_span, + config, + lang_chain_span_usage(raw_usage) or SpanUsage(), + lang_chain_finish_reasons(raw), + ) + run_usage.add(lang_chain_span_usage(raw_usage)) + return parsed + + +def create_langchain_messages_handler( + llm: Any = None, *, capture_content: bool = False +) -> ProviderHandler: """ Creates a ``ProviderHandler`` for LangChain (chat models). Requires ``langchain-openai`` or another LangChain integration to be installed. Pass *llm* to use a specific chat model; omit to default to ``ChatOpenAI(model=)`` resolved at call time. + + 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. """ - tracer_name = "@launchdarkly/ai-langchain-messages" async def _call_impl( config: AiConfigRep, @@ -140,93 +290,80 @@ async def _call_impl( th = tool_handlers or {} vs = variables or {} - base_model = llm - if base_model is None: - base_model = _make_default_chat_model(config, importlib) - - if _HAS_OTEL: - span = trace.get_tracer(tracer_name).start_span("langchain.invoke") - 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) initial_messages = _build_messages(config, user_input, vs, history) - - if span: - prompt_text = "\n".join( - f"{getattr(m, 'type', type(m).__name__)}: {m.content if isinstance(m.content, str) else json.dumps(m.content)}" - for m in initial_messages - ) - span.add_event("gen_ai.content.prompt", {"gen_ai.prompt": prompt_text}) - set_openllmetry_prompt( + if capture_content: + system_instructions, span_messages = _span_messages(initial_messages) + set_input_content_attributes( span, - [ - { - "role": getattr(m, "type", type(m).__name__), - "content": m.content - if isinstance(m.content, str) - else json.dumps(m.content), - } - for m in initial_messages - ], + capture_content, + system_instructions=system_instructions, + messages=span_messages, ) + # Outside the try, so the failure path can still report the spend of the turns that + # completed before it. + run_usage = create_run_usage() try: + base_model = ( + llm if llm is not None else _make_default_chat_model(config, importlib) + ) + tool_defs = _build_tools(config.get("tools") or {}) output_format = config.get("outputFormat") - provider_name = config.get("provider", {}).get("name", "openai").lower() - is_openai = provider_name == "openai" - # Structured output path — only when no tools are present. + # CASE 1: outputFormat only, no tools -> withStructuredOutput (all providers). # LangChain cannot apply with_structured_output and bind_tools to the same model. if output_format and not tool_defs: - structured_model = base_model.with_structured_output( - output_format, include_raw=True + parsed = await _run_structured_turn( + base_model, + config, + initial_messages, + output_format, + parent, + capture_content=capture_content, + run_usage=run_usage, ) - result = await structured_model.ainvoke(initial_messages) - raw_usage = getattr(result.get("raw"), "usage_metadata", None) or {} - input_tokens = raw_usage.get("input_tokens", 0) - output_tokens = raw_usage.get("output_tokens", 0) - if span: - span.set_attribute( - "gen_ai.response.model", config.get("model", {}).get("name", "") - ) - span.set_attribute("gen_ai.usage.input_tokens", input_tokens) - span.set_attribute("gen_ai.usage.output_tokens", output_tokens) - span.set_attribute( - "gen_ai.usage.total_tokens", input_tokens + output_tokens - ) - span.add_event( - "gen_ai.content.completion", - {"gen_ai.completion": json.dumps(result.get("parsed"))}, - ) - set_openllmetry_completion( - span, - json.dumps(result.get("parsed")), - {"input_tokens": input_tokens, "output_tokens": output_tokens}, - ) - span.set_status(SpanStatusCode.OK) - span.end() + output_str = parsed if isinstance(parsed, str) else json.dumps(parsed) + set_output_content_attributes( + span, + capture_content, + [ + SpanMessage( + role="assistant", + parts=[SpanMessagePart(type="text", content=output_str)], + ) + ], + ) + finish_root_span(span, config, run_usage.total) + succeed_span(span) return { - "output": result.get("parsed"), + "output": parsed, "usage": { - "input_tokens": input_tokens, - "output_tokens": output_tokens, + "input_tokens": run_usage.total.input, + "output_tokens": run_usage.total.output, }, } + # CASE 2: tools present -> agentic loop, then withStructuredOutput for the final + # response when outputFormat is set. + provider_name = str( + (config.get("provider") or {}).get("name") or "openai" + ).lower() + is_openai = provider_name == "openai" + # For OpenAI models with both outputFormat and tools: bind response_format so the # final text response is structured JSON. The client layer parses the returned string. + # + # When this binding applies, the tool loop's own final reply is already structured, so + # the structured follow-up turn below must not run: it would discard that reply, bill a + # second turn, and leave the two strategies fighting over the same output. + format_is_bound = False bound_model = base_model if output_format and tool_defs and is_openai: + format_is_bound = True bound_model = base_model.bind( response_format={ "type": "json_schema", @@ -238,27 +375,73 @@ async def _call_impl( } ) - active_model = ( - bound_model.bind_tools(tool_defs) if tool_defs else bound_model - ) + tool_model = bound_model.bind_tools(tool_defs) if tool_defs else bound_model + tool_definitions = to_tool_definitions(tool_defs) conversation_messages = list(initial_messages) - total_input = 0 - total_output = 0 - output = "" + output: Any = "" steps = 0 + msgs_mod = importlib.import_module("langchain_core.messages") + ToolMessage = msgs_mod.ToolMessage + while True: - response = await active_model.ainvoke(conversation_messages) - usage = getattr(response, "usage_metadata", None) or {} - total_input += usage.get("input_tokens", 0) - total_output += usage.get("output_tokens", 0) - conversation_messages.append(response) + model_span = start_model_span(config, parent) + if capture_content: + system_instructions, span_messages = _span_messages( + conversation_messages + ) + set_input_content_attributes( + model_span, + capture_content, + system_instructions=system_instructions, + messages=span_messages, + tool_definitions=tool_definitions, + ) + # The output write and the finish live inside the guard too. Serialising completion + # content raises on anything that is not JSON-serialisable, and a raise out here + # would leave this chat span open: the blocking path has no `finally` to recover it. + try: + response = await tool_model.ainvoke(conversation_messages) + + usage = getattr(response, "usage_metadata", None) or {} + tool_calls = getattr(response, "tool_calls", None) or [] + if capture_content: + set_output_content_attributes( + model_span, + capture_content, + _assistant_output_messages(response.content, tool_calls), + ) + finish_model_span( + model_span, + config, + lang_chain_span_usage(usage) or SpanUsage(), + lang_chain_finish_reasons(response), + ) + except Exception as exc: + fail_span(model_span, exc) + raise + run_usage.add(lang_chain_span_usage(usage)) - tool_calls = getattr(response, "tool_calls", []) or [] if not tool_calls: - output = ( - response.content if isinstance(response.content, str) else "" - ) + # Only when response_format was not already bound above. Anthropic and any other + # non-OpenAI provider reach the model through this second turn, because binding + # response_format is an OpenAI-only mechanism. + if output_format and not format_is_bound: + output = await _run_structured_turn( + base_model, + config, + conversation_messages, + output_format, + parent, + capture_content=capture_content, + run_usage=run_usage, + ) + else: + output = ( + response.content + if isinstance(response.content, str) + else "" + ) break if steps >= _MAX_STEPS: @@ -267,23 +450,37 @@ async def _call_impl( ) steps += 1 - import importlib - - msgs_mod = importlib.import_module("langchain_core.messages") - ToolMessage = msgs_mod.ToolMessage + conversation_messages.append(response) tool_results: list[Any] = [] for tc in tool_calls: - handler_fn = th.get(tc["name"]) - if not handler_fn or not callable(handler_fn): - raise ValueError( - f'No handler registered for tool "{tc["name"]}"' - ) - result_val = ( - await handler_fn(tc["args"]) - if _is_coroutine(handler_fn) - else handler_fn(tc["args"]) + tool_span = start_tool_span( + tc["name"], tc.get("id") or tc["name"], parent + ) + set_tool_call_content_attributes( + tool_span, capture_content, arguments=tc.get("args") ) + try: + handler_fn = th.get(tc["name"]) + if not handler_fn or not callable(handler_fn): + raise ValueError( + f'No handler registered for tool "{tc["name"]}"' + ) + result_val = ( + await handler_fn(tc["args"]) + if _is_coroutine(handler_fn) + else handler_fn(tc["args"]) + ) + # Inside the try on purpose. Serialising a tool result can raise, most easily + # when capture_content is on and the result is not JSON-serialisable, and a + # raise out here would leave this span open: nothing else knows it exists. + set_tool_call_content_attributes( + tool_span, capture_content, result=result_val + ) + succeed_span(tool_span) + except Exception as exc: + fail_span(tool_span, exc) + raise tool_results.append( ToolMessage( tool_call_id=tc.get("id") or tc["name"], @@ -292,41 +489,32 @@ async def _call_impl( ) conversation_messages.extend(tool_results) - 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() - + output_str = output if isinstance(output, str) else json.dumps(output) + set_output_content_attributes( + span, + capture_content, + [ + SpanMessage( + role="assistant", + parts=[SpanMessagePart(type="text", content=output_str)], + ) + ], + ) + 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() + # The turns that completed were billed, and the root is the only span a config-scoped + # cost query can find them on. + if run_usage.reported: + finish_root_span(span, config, run_usage.total) + fail_span(span, exc) raise def _stream_impl( @@ -337,7 +525,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(("*", "messages"), _call_impl, _stream_impl) # type: ignore[arg-type] @@ -350,87 +544,150 @@ 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 is never ended, so it is never exported, and the whole run disappears from AI + Config Monitoring along with the ``feature_flag`` event it carries. + + ``ended`` stops the success, failure and abandonment paths from ending the same span twice. + + ``open_chunk_stream`` holds LangChain's own ``astream`` generator, closed by hand in the same + ``finally``. LangChain's ``BaseChatModel.astream`` is itself an async generator that awaits + ``run_manager.on_llm_error(...)`` inside a ``except BaseException`` block, the same shape that + makes ``claude_agent_sdk``'s generator unsafe to abandon bare: a consumer breaking out of this + generator leaves that inner one suspended mid-await with nothing to resume it until Python's + garbage collector finalizes it, which can raise rather than clean up quietly. Calling + ``aclose()`` on it here, on the same object the loop was iterating, drives that cleanup + immediately instead of leaving it to GC. + """ import importlib - base_model = llm - if base_model is None: - base_model = _make_default_chat_model(config, importlib) + base_model = llm if llm is not None else _make_default_chat_model(config, importlib) - tracer_name = "@launchdarkly/ai-langchain-messages" - if _HAS_OTEL: - span = trace.get_tracer(tracer_name).start_span("langchain.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) initial_messages = _build_messages(config, user_input, variables, history) - if span: - prompt_text = "\n".join( - f"{getattr(m, 'type', type(m).__name__)}: {m.content if isinstance(m.content, str) else json.dumps(m.content)}" - for m in initial_messages - ) - span.add_event("gen_ai.content.prompt", {"gen_ai.prompt": prompt_text}) - set_openllmetry_prompt( + if capture_content: + system_instructions, span_messages = _span_messages(initial_messages) + set_input_content_attributes( span, - [ - { - "role": getattr(m, "type", type(m).__name__), - "content": m.content - if isinstance(m.content, str) - else json.dumps(m.content), - } - for m in initial_messages - ], + capture_content, + system_instructions=system_instructions, + messages=span_messages, ) tool_defs = _build_tools(config.get("tools") or {}) - active_model = base_model.bind_tools(tool_defs) if tool_defs else base_model + tool_model = base_model.bind_tools(tool_defs) if tool_defs else base_model + tool_definitions = to_tool_definitions(tool_defs) conversation_messages = list(initial_messages) - total_input = 0 - total_output = 0 - full_output = "" + full_output: Any = "" steps = 0 - try: - import importlib + ended: set[int] = set() + open_model_span: Any = None + open_chunk_stream: Any = None + # Outside the try, so the failure and abandonment paths can still report the spend. + run_usage = create_run_usage() - msgs_mod = importlib.import_module("langchain_core.messages") - AIMessage = msgs_mod.AIMessage - ToolMessage = msgs_mod.ToolMessage + msgs_mod = importlib.import_module("langchain_core.messages") + AIMessage = msgs_mod.AIMessage + ToolMessage = msgs_mod.ToolMessage + try: while True: - chunk_stream = active_model.astream(conversation_messages) + model_span = start_model_span(config, parent) + open_model_span = model_span + if capture_content: + system_instructions, span_messages = _span_messages( + conversation_messages + ) + set_input_content_attributes( + model_span, + capture_content, + system_instructions=system_instructions, + messages=span_messages, + tool_definitions=tool_definitions, + ) + accumulated_content = "" accumulated_tool_calls: list[Any] = [] - turn_input = 0 - turn_output = 0 - - async for chunk in chunk_stream: - text = chunk.content if isinstance(chunk.content, str) else "" - if text: - yield {"type": "chunk", "text": text} - accumulated_content += text - usage = getattr(chunk, "usage_metadata", None) - if usage: - turn_input += usage.get("input_tokens", 0) - turn_output += usage.get("output_tokens", 0) - chunk_tools = getattr(chunk, "tool_calls", []) or [] - if chunk_tools: - accumulated_tool_calls = chunk_tools - - total_input += turn_input - total_output += turn_output + # The cache breakdown has to be accumulated alongside the scalars. LangChain reports it + # per chunk in `usage_metadata.input_token_details`, and synthesizing a usage bag + # without it would make the streaming span emit `cache_read = 0` where the blocking path + # emits the real figure: a zero that reads as "no cached tokens" rather than "not + # reported". + turn_usage = SpanUsage() + # Carried forward chunk by chunk for the same reason as the cache breakdown above: only + # the terminal chunk reports it, and dropping it would make the streaming span omit a + # finish reason where the blocking path emits the real one. + finish_reasons: list[str] | None = None + usage_reported = False + + try: + chunk_stream = tool_model.astream(conversation_messages) + open_chunk_stream = chunk_stream + async for chunk in chunk_stream: + text = chunk.content if isinstance(chunk.content, str) else "" + if text: + yield {"type": "chunk", "text": text} + accumulated_content += text + usage = getattr(chunk, "usage_metadata", None) + if usage: + usage_reported = True + details = usage.get("input_token_details") or {} + turn_usage.input += number_or_zero(usage.get("input_tokens")) + turn_usage.output += number_or_zero(usage.get("output_tokens")) + turn_usage.cache_read += number_or_zero( + details.get("cache_read") + ) + turn_usage.cache_creation += number_or_zero( + details.get("cache_creation") + ) + chunk_tools = getattr(chunk, "tool_calls", None) or [] + if chunk_tools: + accumulated_tool_calls = chunk_tools + finish_reasons = lang_chain_finish_reasons(chunk) or finish_reasons + open_chunk_stream = None + except Exception as exc: + # The tracker matters here: the outer `except` also fails `open_model_span`, which + # still points at this span because the line that clears it is unreachable on this + # path. + fail_span(model_span, exc, ended) + open_model_span = None + raise + + if capture_content: + set_output_content_attributes( + model_span, + capture_content, + _assistant_output_messages( + accumulated_content, accumulated_tool_calls + ), + ) + # finish_model_span ends the span. Clearing open_model_span is what stops the + # `finally` from ending it a second time. + finish_model_span(model_span, config, turn_usage, finish_reasons) + open_model_span = None + # The already-mapped figures, not a bag rebuilt from them: this path summed the chunks + # into a SpanUsage to begin with, so re-parsing its own output would be a round trip + # whose only effect is another chance to disagree with itself. + # + # Only when a chunk actually carried usage. Adding unconditionally marks the run as + # having reported, so a later failure or abandonment writes all-zero totals on the root + # and claims the run cost nothing, which is the one thing `reported` exists to prevent. + # The blocking path gets this for free, because lang_chain_span_usage returns None for a + # bag the provider never filled. + run_usage.add(turn_usage if usage_reported else None) if not accumulated_tool_calls: - full_output += accumulated_content + full_output = (full_output or "") + accumulated_content break if steps >= _MAX_STEPS: @@ -439,7 +696,7 @@ async def _stream_gen( ) steps += 1 - full_output += accumulated_content + full_output = (full_output or "") + accumulated_content assistant_msg = AIMessage( content=accumulated_content, tool_calls=accumulated_tool_calls ) @@ -447,14 +704,33 @@ async def _stream_gen( tool_results: list[Any] = [] for tc in accumulated_tool_calls: - handler_fn = tool_handlers.get(tc["name"]) - if not handler_fn or not callable(handler_fn): - raise ValueError(f'No handler registered for tool "{tc["name"]}"') - result_val = ( - await handler_fn(tc["args"]) - if _is_coroutine(handler_fn) - else handler_fn(tc["args"]) + tool_span = start_tool_span( + tc["name"], tc.get("id") or tc["name"], parent ) + set_tool_call_content_attributes( + tool_span, capture_content, arguments=tc.get("args") + ) + try: + handler_fn = tool_handlers.get(tc["name"]) + if not handler_fn or not callable(handler_fn): + raise ValueError( + f'No handler registered for tool "{tc["name"]}"' + ) + result_val = ( + await handler_fn(tc["args"]) + if _is_coroutine(handler_fn) + else handler_fn(tc["args"]) + ) + # Inside the try on purpose. Serialising a tool result can raise, most easily + # when capture_content is on and the result is not JSON-serialisable, and a + # raise out here would leave this span open: nothing else knows it exists. + set_tool_call_content_attributes( + tool_span, capture_content, result=result_val + ) + succeed_span(tool_span) + except Exception as exc: + fail_span(tool_span, exc, ended) + raise tool_results.append( ToolMessage( tool_call_id=tc.get("id") or tc["name"], content=str(result_val) @@ -462,43 +738,58 @@ async def _stream_gen( ) conversation_messages.extend(tool_results) - 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() + full_output_str = ( + full_output if isinstance(full_output, str) else json.dumps(full_output) + ) + set_output_content_attributes( + span, + capture_content, + [ + SpanMessage( + role="assistant", + parts=[SpanMessagePart(type="text", content=full_output_str)], + ) + ], + ) + 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}, + "output": full_output_str, + "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() + if run_usage.reported: + finish_root_span(span, config, 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` and already drove `chunk_stream` to completion. On abandonment this is the only + # chance to close the tree, to close LangChain's own generator, 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, so ERROR would put two dashboards in disagreement. + # Spans first, and the vendor generator after. aclose() can raise, and doing it first would + # take the whole teardown with it: the root would never end, never export, and the run would + # vanish from AI Config Monitoring along with the feature_flag event this block exists to + # protect. Its own failure is not worth losing the trace over, so it is contained. + if open_model_span is not None: + end_span_once(open_model_span, ended, abandoned=True) + if span is not None and id(span) not in ended and run_usage.reported: + finish_root_span(span, config, run_usage.total) + end_span_once(span, ended, abandoned=True) + if open_chunk_stream is not None: + try: + await open_chunk_stream.aclose() + except Exception: # pragma: no cover - best-effort vendor teardown + pass def langchain_messages( @@ -509,7 +800,15 @@ def langchain_messages( **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_messages_handler(llm=llm), **kwargs + key=config_key, + handler=create_langchain_messages_handler( + llm=llm, capture_content=capture_content + ), + **kwargs, ).invoke(user_input, context, variables=variables) diff --git a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/spans.py b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/spans.py new file mode 100644 index 0000000..2710409 --- /dev/null +++ b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/spans.py @@ -0,0 +1,227 @@ +"""Span construction for the LangChain messages handler. + +Separate from ``handler.py`` so the span shape is readable on its own, and so the tool loop reads as +the tool loop rather than as span bookkeeping with a provider call in the middle. + +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. + +Two provider keys, two different values, on purpose. ``gen_ai.system`` is the literal string +``langchain`` on every span this package opens, because that is what the handler shipped before +the span hierarchy landed. ``gen_ai.provider.name`` names *who served the model*, and semconv's +enum has no ``langchain`` member, so it follows :func:`serving_provider` instead: whichever chat +model class the handler actually instantiates. See TELEMETRY-CONTRACT.md section 9. +""" + +from __future__ import annotations + +from typing import Any + +from launchdarkly_ai_server import ( + AiConfigRep, + SpanUsage, + ToolDefinitionInput, + set_ld_span_attributes, + set_model_identity_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 + +TRACER_NAME = "@launchdarkly/ai-langchain-messages" + +#: The literal value ``gen_ai.system`` carries on every span this package opens. Not derived from +#: the configured provider: TypeScript's LangChain handlers keep this constant regardless of which +#: model actually served the call, because LangChain (the framework) is what shipped the span, not +#: Anthropic or OpenAI (the provider). See TELEMETRY-CONTRACT.md section 9. +LEGACY_SYSTEM = "langchain" + + +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, because LangChain is the framework, not the provider. This mirrors the + choice the handler's model-resolution logic makes: ``ChatAnthropic`` for a configured provider + of ``"anthropic"``, ``ChatOpenAI`` for everything else, including Bedrock, Azure, Cohere, a + typo, or an unset value. Not a passthrough of the configured name. See TELEMETRY-CONTRACT.md + section 9. + """ + provider = str((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 + ) + 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) + 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, unlike ``openai-messages``: LangChain does + not hand this handler a resolved model id to report instead. See TELEMETRY-CONTRACT.md section + 2a. + + The per-turn ``chat`` children carry the same attributes for their own turn, but the root is the + only span a config-scoped query finds, so leaving the totals off it means such a query returns + nothing at all: summing the children requires having already found them. + """ + 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, + usage: SpanUsage, + finish_reasons: list[str] | None = None, +) -> None: + """Ends one ``chat`` span successfully. + + *usage* is the caller's responsibility to default to zeros when LangChain reported nothing: a + ``chat`` span always writes the complete attribute set, unlike the root, which withholds it when + no turn ever reported. See TELEMETRY-CONTRACT.md section 8. + + *finish_reasons* arrives already mapped through :func:`lang_chain_finish_reasons`. + """ + if span is None: + return + span.set_attribute("gen_ai.response.model", model_name(config)) + # A list because one response may hold several choices; the providers used here return one. + if finish_reasons: + span.set_attribute("gen_ai.response.finish_reasons", finish_reasons) + set_usage_span_attributes(span, usage) + 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. Ending twice is silently ignored by + the OTel SDK but recorded as a diagnostic error, and would also hide a genuine leak. + """ + 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: + from launchdarkly_ai_server import end_span_once + + end_span_once(span, tracker) + else: + span.end() + + +# ─── Provider shapes as span shapes ────────────────────────────────────────── + + +def to_tool_definitions(tools: list[dict[str, Any]]) -> list[ToolDefinitionInput]: + """The catalog as bound to the model, so the span reports what it could actually call. + + *tools* are this package's own OpenAI-function-style tool dicts (``{"type": "function", + "function": {...}}``), not the AI Config's tool type: see ``_build_tools`` in ``handler.py``. + """ + return [ + ToolDefinitionInput( + name=str(t.get("function", {}).get("name", "")), + description=t.get("function", {}).get("description"), + parameters=t.get("function", {}).get("parameters"), + ) + for t in tools + ] diff --git a/packages/langchain-messages/tests/test_handler.py b/packages/langchain-messages/tests/test_handler.py index 1c57c9f..bf8141c 100644 --- a/packages/langchain-messages/tests/test_handler.py +++ b/packages/langchain-messages/tests/test_handler.py @@ -1,7 +1,6 @@ """ Tests for launchdarkly-ai-langchain-messages handler. -Covers §1.1–1.9 and §1.x (LangChain-specific extras). -Reference: TESTING.md §1, §1.x +Covers §1.1-1.10 (generic handler tests) plus TELEMETRY-CONTRACT.md sections 1-9. """ from __future__ import annotations @@ -12,47 +11,144 @@ import pytest -CONFIG = { - "model": {"name": "gpt-4o"}, - "provider": {"name": "LangChain"}, - "instructions": "Be helpful.", -} +# --------------------------------------------------------------------------- +# Fake LangChain message helpers +# +# Deliberately plain objects rather than MagicMock: MagicMock answers every attribute access with +# a fresh Mock rather than raising AttributeError, so `lang_chain_finish_reasons`'s `_get(obj, key)` +# (a `getattr(obj, key, None)`) never falls through to its default, and a finish reason silently +# stops being derivable from the mock message. +# --------------------------------------------------------------------------- -def _make_ai_message( - content: str = "Hello", - tool_calls: list[dict] | None = None, - input_tokens: int = 10, - output_tokens: int = 5, -) -> MagicMock: - msg = MagicMock() - msg.content = content - msg.tool_calls = tool_calls or [] - msg.usage_metadata = {"input_tokens": input_tokens, "output_tokens": output_tokens} - msg._getType = lambda: "ai" - return msg +class FakeAIMessage: + def __init__( + self, + content: str = "Hello", + tool_calls: list[dict[str, Any]] | None = None, + input_tokens: int = 10, + output_tokens: int = 5, + cache_read: int | None = None, + cache_creation: int | None = None, + finish_reason: str | None = None, + ) -> None: + self.content = content + self.tool_calls = tool_calls or [] + usage: dict[str, Any] = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + } + if cache_read is not None or cache_creation is not None: + details: dict[str, Any] = {} + if cache_read is not None: + details["cache_read"] = cache_read + if cache_creation is not None: + details["cache_creation"] = cache_creation + usage["input_token_details"] = details + self.usage_metadata = usage + self.response_metadata = ( + {"finish_reason": finish_reason} if finish_reason else {} + ) + + def _get_type(self) -> str: + return "ai" def _make_llm(response_content: str = "Hello") -> MagicMock: """Creates a mock LangChain LLM.""" llm = MagicMock() - ai_msg = _make_ai_message(response_content) + ai_msg = FakeAIMessage(response_content) llm.ainvoke = AsyncMock(return_value=ai_msg) llm.bind_tools = MagicMock(return_value=llm) llm.with_structured_output = MagicMock(return_value=llm) + llm.bind = MagicMock(return_value=llm) - async def _astream(msgs: Any) -> AsyncGenerator: - chunk = MagicMock() - chunk.content = response_content - chunk.usage_metadata = {"input_tokens": 5, "output_tokens": 3} - chunk.tool_calls = [] + async def _astream(msgs: Any) -> AsyncGenerator[Any, None]: + chunk = FakeAIMessage(response_content, input_tokens=5, output_tokens=3) yield chunk llm.astream = _astream return llm -def _make_tracer_patch(mock_span: MagicMock) -> tuple[MagicMock, MagicMock]: +CONFIG = { + "model": {"name": "gpt-4o"}, + "provider": {"name": "OpenAI"}, + "instructions": "Be helpful.", +} + + +# --------------------------------------------------------------------------- +# Span recording, mirroring launchdarkly_ai_claude_messages' test approach. +# --------------------------------------------------------------------------- + + +class RecordedSpan: + """A span that remembers what a handler did to it, so a test can assert on the whole thing.""" + + 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 + + def set_attribute(self, key: str, value: Any) -> None: + self.attributes[key] = value + + def add_event(self, name: str, attributes: dict[str, Any] | None = None) -> None: + self.events.append((name, attributes or {})) + + def set_status(self, code: Any, description: str | None = None) -> None: + self.statuses.append(code) + + def record_exception(self, exc: BaseException) -> None: + self.exceptions.append(exc) + + def end(self) -> None: + self.ended += 1 + + +class SpanRecorder: + """Stands in for the ``trace`` module inside ``spans.py`` and records every span opened.""" + + 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.""" + import launchdarkly_ai_langchain_messages.spans as spans_mod + + recorder = SpanRecorder() + return patch.object(spans_mod, "trace", recorder), recorder + + +def _make_tracer_patch(mock_span: MagicMock) -> Any: mock_tracer = MagicMock() mock_tracer.start_span = MagicMock(return_value=mock_span) mock_trace_mod = MagicMock() @@ -69,8 +165,7 @@ class TestFactory: def test_returns_callable(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - llm = _make_llm() - h = create_langchain_messages_handler(llm=llm) + h = create_langchain_messages_handler(llm=_make_llm()) assert callable(h) def test_attaches_provides_for(self) -> None: @@ -79,7 +174,7 @@ def test_attaches_provides_for(self) -> None: h = create_langchain_messages_handler(llm=_make_llm()) assert h.provides_for is not None - def test_provides_for_values_are_correct(self) -> None: + def test_provides_for_is_the_wildcard_provider(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler h = create_langchain_messages_handler(llm=_make_llm()) @@ -106,11 +201,6 @@ async def test_path_a_instructions(self) -> None: h = create_langchain_messages_handler(llm=llm) await h(CONFIG, "hi", {}, {}) call_args = llm.ainvoke.call_args[0][0] - _types = [ - m._getType() if hasattr(m, "_getType") else type(m).__name__ - for m in call_args - ] - # First message should be a system message assert any( "system" in str(m.__class__.__name__).lower() or "System" in str(type(m)) for m in call_args @@ -143,10 +233,8 @@ async def test_path_b_messages_system_extracted(self) -> None: config = { "model": {"name": "gpt-4o"}, - "provider": {"name": "LangChain"}, - "messages": [ - {"role": "system", "content": "Be a poet"}, - ], + "provider": {"name": "OpenAI"}, + "messages": [{"role": "system", "content": "Be a poet"}], } llm = _make_llm() h = create_langchain_messages_handler(llm=llm) @@ -160,16 +248,13 @@ async def test_path_b_user_input_appended_as_final_turn(self) -> None: config = { "model": {"name": "gpt-4o"}, - "provider": {"name": "LangChain"}, + "provider": {"name": "OpenAI"}, "instructions": "be helpful", } llm = _make_llm() h = create_langchain_messages_handler(llm=llm) await h(config, "final-input", {}, {}) - # call_args captures a mutable list; verify "final-input" appears in the messages - call_args_list = llm.ainvoke.call_args_list - assert len(call_args_list) > 0 - messages_sent = call_args_list[0][0][0] + messages_sent = llm.ainvoke.call_args_list[0][0][0] all_content = " ".join(str(getattr(m, "content", "")) for m in messages_sent) assert "final-input" in all_content @@ -180,19 +265,12 @@ async def test_path_c_empty_user_input_no_throw(self) -> None: h = create_langchain_messages_handler(llm=llm) await h(CONFIG, "", {}, {}) - async def test_path_c_undefined_user_input_no_throw(self) -> None: - from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - - llm = _make_llm() - h = create_langchain_messages_handler(llm=llm) - await h(CONFIG, None, {}, {}) # type: ignore[arg-type] - async def test_path_b_variable_substitution_in_system_message(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler config = { "model": {"name": "gpt-4o"}, - "provider": {"name": "LangChain"}, + "provider": {"name": "OpenAI"}, "messages": [{"role": "system", "content": "Hello {{name}}"}], } llm = _make_llm() @@ -206,7 +284,7 @@ async def test_path_c_both_instructions_and_messages_messages_wins(self) -> None from launchdarkly_ai_langchain_messages import create_langchain_messages_handler config = { - **CONFIG, # has instructions = "Be helpful." + **CONFIG, "messages": [{"role": "system", "content": "from-messages"}], } llm = _make_llm() @@ -280,10 +358,10 @@ class TestToolExecutionLoop: async def test_single_tool_call_then_done(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - tool_ai_msg = _make_ai_message( + tool_ai_msg = FakeAIMessage( tool_calls=[{"name": "search", "id": "tc1", "args": {"q": "test"}}] ) - final_ai_msg = _make_ai_message("final answer") + final_ai_msg = FakeAIMessage("final answer") llm = _make_llm() llm.ainvoke = AsyncMock(side_effect=[tool_ai_msg, final_ai_msg]) config = { @@ -301,7 +379,7 @@ async def test_single_tool_call_then_done(self) -> None: async def test_tool_not_found_throws(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - tool_ai_msg = _make_ai_message( + tool_ai_msg = FakeAIMessage( tool_calls=[{"name": "unknown_tool", "id": "tc1", "args": {}}] ) llm = _make_llm() @@ -317,7 +395,7 @@ async def test_tool_not_found_throws(self) -> None: async def test_tool_handler_throws_propagates(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - tool_ai_msg = _make_ai_message( + tool_ai_msg = FakeAIMessage( tool_calls=[{"name": "t1", "id": "tc1", "args": {}}] ) llm = _make_llm() @@ -344,13 +422,9 @@ async def test_no_tools_in_config_handler_never_invoked(self) -> None: async def test_multiple_consecutive_tool_calls(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - msg1 = _make_ai_message( - tool_calls=[{"name": "t1", "id": "tc1", "args": {"x": 1}}] - ) - msg2 = _make_ai_message( - tool_calls=[{"name": "t2", "id": "tc2", "args": {"y": 2}}] - ) - msg3 = _make_ai_message("final") + msg1 = FakeAIMessage(tool_calls=[{"name": "t1", "id": "tc1", "args": {"x": 1}}]) + msg2 = FakeAIMessage(tool_calls=[{"name": "t2", "id": "tc2", "args": {"y": 2}}]) + msg3 = FakeAIMessage("final") llm = _make_llm() llm.ainvoke = AsyncMock(side_effect=[msg1, msg2, msg3]) cfg = { @@ -370,259 +444,551 @@ async def test_multiple_consecutive_tool_calls(self) -> None: # --------------------------------------------------------------------------- -# §1.5 Telemetry +# TELEMETRY-CONTRACT.md section 1: span tree # --------------------------------------------------------------------------- -class TestTelemetry: - async def test_span_name_blocking(self) -> None: - mock_span = MagicMock() - mock_trace_mod, mock_tracer = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod +class TestSpanTree: + async def test_opens_a_root_span_named_invoke_agent(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_langchain_messages import ( - create_langchain_messages_handler, + with ctx: + await create_langchain_messages_handler(llm=_make_llm())( + CONFIG, "q", {}, {} ) + assert rec.root.name == "invoke_agent" + assert rec.root.attributes["gen_ai.operation.name"] == "invoke_agent" - h = create_langchain_messages_handler(llm=_make_llm()) - await h(CONFIG, "q", {}, {}) - mock_tracer.start_span.assert_called_with("langchain.invoke") + async def test_emits_one_chat_child_per_model_turn(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - async def test_gen_ai_system(self) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod + with ctx: + await create_langchain_messages_handler(llm=_make_llm())( + 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) + + async def test_emits_a_chat_span_per_turn_of_a_tool_loop(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_langchain_messages import ( - create_langchain_messages_handler, + llm = _make_llm() + llm.ainvoke = AsyncMock( + side_effect=[ + FakeAIMessage(tool_calls=[{"name": "myTool", "id": "tc1", "args": {}}]), + FakeAIMessage("done"), + ] + ) + cfg = {**CONFIG, "tools": {"myTool": {"type": "function", "parameters": {}}}} + with ctx: + await create_langchain_messages_handler(llm=llm)( + cfg, "q", {"myTool": lambda _: "result"}, {} ) + assert len(rec.named("chat ")) == 2 - h = create_langchain_messages_handler(llm=_make_llm()) - await h(CONFIG, "q", {}, {}) - attrs = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert attrs.get("gen_ai.system") == "langchain" + async def test_emits_an_execute_tool_span_per_tool_call(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - async def test_span_end_always_called(self) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod + llm = _make_llm() + llm.ainvoke = AsyncMock( + side_effect=[ + FakeAIMessage(tool_calls=[{"name": "myTool", "id": "tu1", "args": {}}]), + FakeAIMessage("done"), + ] + ) + cfg = {**CONFIG, "tools": {"myTool": {"type": "function", "parameters": {}}}} + with ctx: + await create_langchain_messages_handler(llm=llm)( + cfg, "q", {"myTool": lambda _: "result"}, {} + ) + tools = rec.named("execute_tool ") + assert len(tools) == 1 + assert tools[0].name == "execute_tool myTool" + assert tools[0].attributes["gen_ai.operation.name"] == "execute_tool" + assert tools[0].attributes["gen_ai.tool.name"] == "myTool" + assert tools[0].attributes["gen_ai.tool.call.id"] == "tu1" + + async def test_tool_spans_are_siblings_of_chat_not_children(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_langchain_messages import ( - create_langchain_messages_handler, + llm = _make_llm() + llm.ainvoke = AsyncMock( + side_effect=[ + FakeAIMessage(tool_calls=[{"name": "myTool", "id": "tu1", "args": {}}]), + FakeAIMessage("done"), + ] + ) + cfg = {**CONFIG, "tools": {"myTool": {"type": "function", "parameters": {}}}} + with ctx: + await create_langchain_messages_handler(llm=llm)( + cfg, "q", {"myTool": lambda _: "r"}, {} ) + assert rec.named("execute_tool ")[0].context == ("context-of", rec.root) - h = create_langchain_messages_handler(llm=_make_llm()) - await h(CONFIG, "q", {}, {}) - mock_span.end.assert_called_once() + async def test_every_span_is_ended(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - async def test_gen_ai_operation_name(self) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod + llm = _make_llm() + llm.ainvoke = AsyncMock( + side_effect=[ + FakeAIMessage(tool_calls=[{"name": "myTool", "id": "tu1", "args": {}}]), + FakeAIMessage("done"), + ] + ) + cfg = {**CONFIG, "tools": {"myTool": {"type": "function", "parameters": {}}}} + with ctx: + await create_langchain_messages_handler(llm=llm)( + cfg, "q", {"myTool": lambda _: "r"}, {} + ) + assert [s.ended for s in rec.spans] == [1] * len(rec.spans) - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_langchain_messages import ( - create_langchain_messages_handler, + +# --------------------------------------------------------------------------- +# TELEMETRY-CONTRACT.md sections 2, 2a and 9: root span / model identity +# --------------------------------------------------------------------------- + + +class TestRootSpanAttributes: + async def test_gen_ai_system_is_the_literal_langchain(self) -> None: + # TELEMETRY-CONTRACT.md section 9: Python used to set this to the configured provider, + # lower-cased. TypeScript's LangChain handlers keep it the constant `langchain` regardless. + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + cfg = {**CONFIG, "provider": {"name": "Anthropic"}} + with ctx: + await create_langchain_messages_handler(llm=_make_llm())(cfg, "q", {}, {}) + assert rec.root.attributes["gen_ai.system"] == "langchain" + + async def test_gen_ai_provider_name_is_anthropic_only_for_anthropic(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + cfg = {**CONFIG, "provider": {"name": "Anthropic"}} + with ctx: + await create_langchain_messages_handler(llm=_make_llm())(cfg, "q", {}, {}) + assert rec.root.attributes["gen_ai.provider.name"] == "anthropic" + + @pytest.mark.parametrize( + "provider_name", ["OpenAI", "Bedrock", "Azure", "Cohere", "Typo", ""] + ) + async def test_gen_ai_provider_name_is_openai_for_everything_else( + self, provider_name: str + ) -> None: + # Not a passthrough. Bedrock, Azure, Cohere, a typo and an unset value all report `openai`, + # mirroring the chat model class the handler actually instantiates. Section 9. + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + cfg = {**CONFIG, "provider": {"name": provider_name}} + with ctx: + await create_langchain_messages_handler(llm=_make_llm())(cfg, "q", {}, {}) + assert rec.root.attributes["gen_ai.provider.name"] == "openai" + + async def test_writes_the_requested_model(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + with ctx: + await create_langchain_messages_handler(llm=_make_llm())( + CONFIG, "q", {}, {} ) + assert rec.root.attributes["gen_ai.request.model"] == "gpt-4o" - h = create_langchain_messages_handler(llm=_make_llm()) - await h(CONFIG, "q", {}, {}) - attrs = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert attrs.get("gen_ai.operation.name") == "chat" + async def test_response_model_is_the_requested_name(self) -> None: + # LangChain does not resolve an alias to a different snapshot in this handler. Section 2a. + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - async def test_gen_ai_request_model(self) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod + with ctx: + await create_langchain_messages_handler(llm=_make_llm())( + CONFIG, "q", {}, {} + ) + assert rec.root.attributes["gen_ai.response.model"] == "gpt-4o" - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_langchain_messages import ( - create_langchain_messages_handler, + async def test_carries_the_launchdarkly_attributes_and_feature_flag_event( + self, + ) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + variables = { + "__ld": { + "configKey": "k", + "variationKey": "v", + "runId": "r", + "graphKey": "g", + } + } + with ctx: + await create_langchain_messages_handler(llm=_make_llm())( + 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"] + + async def test_child_spans_carry_no_launchdarkly_identity(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - h = create_langchain_messages_handler(llm=_make_llm()) - await h(CONFIG, "q", {}, {}) - attrs = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert attrs.get("gen_ai.request.model") == CONFIG["model"]["name"] + llm = _make_llm() + llm.ainvoke = AsyncMock( + side_effect=[ + FakeAIMessage(tool_calls=[{"name": "myTool", "id": "tu1", "args": {}}]), + FakeAIMessage("done"), + ] + ) + cfg = {**CONFIG, "tools": {"myTool": {"type": "function", "parameters": {}}}} + variables = {"__ld": {"configKey": "k", "variationKey": "v", "runId": "r"}} + with ctx: + await create_langchain_messages_handler(llm=llm)( + cfg, "q", {"myTool": lambda _: "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_gen_ai_content_prompt_event(self) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod + async def test_carries_the_run_total_not_one_turn(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_langchain_messages import ( - create_langchain_messages_handler, + llm = _make_llm() + llm.ainvoke = AsyncMock( + side_effect=[ + FakeAIMessage( + tool_calls=[{"name": "myTool", "id": "tu1", "args": {}}], + input_tokens=10, + output_tokens=1, + ), + FakeAIMessage("done", input_tokens=20, output_tokens=2), + ] + ) + cfg = {**CONFIG, "tools": {"myTool": {"type": "function", "parameters": {}}}} + with ctx: + await create_langchain_messages_handler(llm=llm)( + cfg, "q", {"myTool": lambda _: "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 - h = create_langchain_messages_handler(llm=_make_llm()) - await h(CONFIG, "q", {}, {}) - event_names = [c[0][0] for c in mock_span.add_event.call_args_list] - assert "gen_ai.content.prompt" in event_names - async def test_gen_ai_content_completion_event(self) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod +# --------------------------------------------------------------------------- +# TELEMETRY-CONTRACT.md sections 3, 5 and 8: chat span attributes +# --------------------------------------------------------------------------- + + +class TestChatSpanAttributes: + async def test_writes_all_seven_usage_attributes_including_zeros(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_langchain_messages import ( - create_langchain_messages_handler, + with ctx: + await create_langchain_messages_handler(llm=_make_llm())( + 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 + + async def test_passes_the_input_figure_through_without_folding_cache_into_it( + self, + ) -> None: + # TELEMETRY-CONTRACT.md section 8: LangChain already counts cached tokens inside + # `input_tokens`, unlike Anthropic. Adding the cache buckets on top here would double-count. + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - h = create_langchain_messages_handler(llm=_make_llm()) - await h(CONFIG, "q", {}, {}) - event_names = [c[0][0] for c in mock_span.add_event.call_args_list] - assert "gen_ai.content.completion" in event_names + llm = _make_llm() + llm.ainvoke = AsyncMock( + return_value=FakeAIMessage( + input_tokens=23554, + output_tokens=100, + cache_read=19971, + cache_creation=3580, + ) + ) + with ctx: + await create_langchain_messages_handler(llm=llm)(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 + + async def test_reports_the_mapped_finish_reason(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - async def test_token_attributes_set(self) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod + llm = _make_llm() + llm.ainvoke = AsyncMock(return_value=FakeAIMessage(finish_reason="stop")) + with ctx: + await create_langchain_messages_handler(llm=llm)(CONFIG, "q", {}, {}) + assert rec.named("chat ")[0].attributes["gen_ai.response.finish_reasons"] == [ + "stop" + ] + + async def test_maps_the_anthropic_word_through_the_table(self) -> None: + # LangChain can serve an Anthropic model, and this handler does use the mapping table, + # unlike the two OpenAI handlers. Section 5. + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + llm = _make_llm() + llm.ainvoke = AsyncMock(return_value=FakeAIMessage(finish_reason="end_turn")) + with ctx: + await create_langchain_messages_handler(llm=llm)(CONFIG, "q", {}, {}) + assert rec.named("chat ")[0].attributes["gen_ai.response.finish_reasons"] == [ + "stop" + ] + + async def test_omits_the_finish_reason_when_the_provider_gives_none(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_langchain_messages import ( - create_langchain_messages_handler, + with ctx: + await create_langchain_messages_handler(llm=_make_llm())( + CONFIG, "q", {}, {} ) + assert "gen_ai.response.finish_reasons" not in rec.named("chat ")[0].attributes - h = create_langchain_messages_handler(llm=_make_llm()) - await h(CONFIG, "q", {}, {}) - attrs = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert "gen_ai.usage.input_tokens" in attrs - assert "gen_ai.usage.output_tokens" in attrs - assert "gen_ai.usage.total_tokens" in attrs - - async def test_gen_ai_response_model(self) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod + async def test_sets_status_ok_on_a_successful_turn(self) -> None: + from opentelemetry.trace import StatusCode - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_langchain_messages import ( - create_langchain_messages_handler, + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + with ctx: + await create_langchain_messages_handler(llm=_make_llm())( + CONFIG, "q", {}, {} ) + assert rec.named("chat ")[0].statuses == [StatusCode.OK] - h = create_langchain_messages_handler(llm=_make_llm()) - await h(CONFIG, "q", {}, {}) - attrs = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert "gen_ai.response.model" in attrs - assert attrs["gen_ai.response.model"] == CONFIG["model"]["name"] - async def test_ld_span_attributes(self) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod +# --------------------------------------------------------------------------- +# TELEMETRY-CONTRACT.md section 7: content capture +# --------------------------------------------------------------------------- - variables = { - "__ld": { - "configKey": "my-config", - "variationKey": "v1", - "runId": "run-abc", - } - } - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_langchain_messages import ( - create_langchain_messages_handler, + +class TestContentCapture: + async def test_emits_no_content_at_all_by_default(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + with ctx: + await create_langchain_messages_handler(llm=_make_llm())( + 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")] == [] - h = create_langchain_messages_handler(llm=_make_llm()) - await h(CONFIG, "q", {}, variables) - attrs = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert attrs.get("launchdarkly.operation.type") == "gen_ai" - assert attrs.get("launchdarkly.config.key") == "my-config" - assert attrs.get("launchdarkly.variation.key") == "v1" - assert attrs.get("launchdarkly.run.id") == "run-abc" - assert "launchdarkly.graph.key" not in attrs - - async def test_ld_graph_key_set_when_present(self) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod + async def test_puts_prompt_and_completion_on_spans_when_enabled(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - variables = { - "__ld": { - "configKey": "my-config", - "variationKey": "v1", - "runId": "run-abc", - "graphKey": "my-graph", - } + with ctx: + await create_langchain_messages_handler( + llm=_make_llm(), capture_content=True + )(CONFIG, "q", {}, {}) + chat = rec.named("chat ")[0] + assert chat.attributes["gen_ai.prompt.0.role"] == "system" + assert chat.attributes["gen_ai.prompt.0.content"] == "Be helpful." + assert "gen_ai.input.messages" in chat.attributes + assert chat.attributes["gen_ai.completion.0.content"] == "Hello" + assert "gen_ai.output.messages" in chat.attributes + + async def test_records_the_tool_catalog_on_the_chat_span_when_enabled(self) -> None: + import json + + cfg = { + **CONFIG, + "tools": {"myTool": {"description": "d", "parameters": {"type": "object"}}}, } - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_langchain_messages import ( - create_langchain_messages_handler, + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + llm = _make_llm() + llm.ainvoke = AsyncMock(return_value=FakeAIMessage("done")) + with ctx: + await create_langchain_messages_handler(llm=llm, capture_content=True)( + cfg, "q", {"myTool": lambda _: "r"}, {} + ) + definitions = json.loads( + rec.named("chat ")[0].attributes["gen_ai.tool.definitions"] + ) + assert definitions[0]["name"] == "myTool" + assert definitions[0]["type"] == "function" + + async def test_records_tool_arguments_and_results_when_enabled(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + llm = _make_llm() + llm.ainvoke = AsyncMock( + side_effect=[ + FakeAIMessage( + tool_calls=[ + {"name": "myTool", "id": "tu1", "args": {"city": "NYC"}} + ] + ), + FakeAIMessage("done"), + ] + ) + cfg = {**CONFIG, "tools": {"myTool": {"type": "function", "parameters": {}}}} + with ctx: + await create_langchain_messages_handler(llm=llm, capture_content=True)( + cfg, "q", {"myTool": lambda _: "72F"}, {} ) + tool = rec.named("execute_tool ")[0] + assert tool.attributes["gen_ai.tool.call.arguments"] == '{"city": "NYC"}' + assert tool.attributes["gen_ai.tool.call.result"] == "72F" + + async def test_still_writes_the_legacy_content_events_when_enabled(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - h = create_langchain_messages_handler(llm=_make_llm()) - await h(CONFIG, "q", {}, variables) - attrs = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert attrs.get("launchdarkly.graph.key") == "my-graph" + with ctx: + await create_langchain_messages_handler( + llm=_make_llm(), capture_content=True + )(CONFIG, "q", {}, {}) + names = [n for n, _ in rec.named("chat ")[0].events] + assert "gen_ai.content.prompt" in names + assert "gen_ai.content.completion" in names # --------------------------------------------------------------------------- -# §1.6 Error handling +# TELEMETRY-CONTRACT.md section 6: errors # --------------------------------------------------------------------------- class TestErrorHandling: - async def test_rethrows_error(self) -> None: + async def test_fails_the_chat_span_when_the_provider_call_raises(self) -> None: + from opentelemetry.trace import StatusCode + + ctx, rec = _recording() from launchdarkly_ai_langchain_messages import create_langchain_messages_handler llm = _make_llm() - llm.ainvoke = AsyncMock(side_effect=RuntimeError("rethrown")) - h = create_langchain_messages_handler(llm=llm) - with pytest.raises(RuntimeError, match="rethrown"): - await h(CONFIG, "q", {}, {}) + llm.ainvoke = AsyncMock(side_effect=RuntimeError("api error")) + with ctx, pytest.raises(RuntimeError): + await create_langchain_messages_handler(llm=llm)(CONFIG, "q", {}, {}) + chat = rec.named("chat ")[0] + assert len(chat.exceptions) == 1 + assert StatusCode.ERROR in chat.statuses + assert chat.ended == 1 + + async def test_fails_the_root_span_too(self) -> None: + from opentelemetry.trace import StatusCode - async def test_records_exception_on_span(self) -> None: + ctx, rec = _recording() from launchdarkly_ai_langchain_messages import create_langchain_messages_handler llm = _make_llm() - llm.ainvoke = AsyncMock(side_effect=RuntimeError("fail")) - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod - - with patch.object(handler_mod, "trace", mock_trace_mod): - h = create_langchain_messages_handler(llm=llm) - with pytest.raises(RuntimeError): - await h(CONFIG, "q", {}, {}) - mock_span.record_exception.assert_called_once() + llm.ainvoke = AsyncMock(side_effect=RuntimeError("api error")) + with ctx, pytest.raises(RuntimeError): + await create_langchain_messages_handler(llm=llm)(CONFIG, "q", {}, {}) + assert len(rec.root.exceptions) == 1 + assert StatusCode.ERROR in rec.root.statuses + assert rec.root.ended == 1 + + async def test_fails_the_execute_tool_span_when_a_tool_raises(self) -> None: + from opentelemetry.trace import StatusCode - async def test_sets_span_status_error(self) -> None: + ctx, rec = _recording() from launchdarkly_ai_langchain_messages import create_langchain_messages_handler llm = _make_llm() - llm.ainvoke = AsyncMock(side_effect=RuntimeError("fail")) - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod + llm.ainvoke = AsyncMock( + return_value=FakeAIMessage( + tool_calls=[{"name": "myTool", "id": "tu1", "args": {}}] + ) + ) + cfg = {**CONFIG, "tools": {"myTool": {"type": "function", "parameters": {}}}} - with patch.object(handler_mod, "trace", mock_trace_mod): - h = create_langchain_messages_handler(llm=llm) - with pytest.raises(RuntimeError): - await h(CONFIG, "q", {}, {}) - from opentelemetry.trace import StatusCode + def _boom(_: Any) -> Any: + raise RuntimeError("tool exploded") - status_codes = [c[0][0] for c in mock_span.set_status.call_args_list] - assert StatusCode.ERROR in status_codes + with ctx, pytest.raises(RuntimeError, match="tool exploded"): + await create_langchain_messages_handler(llm=llm)( + cfg, "q", {"myTool": _boom}, {} + ) + tool = rec.named("execute_tool ")[0] + assert len(tool.exceptions) == 1 + assert StatusCode.ERROR in tool.statuses + assert tool.ended == 1 - async def test_ends_span_on_error(self) -> None: + async def test_reports_the_spend_of_completed_turns_on_a_failed_run(self) -> None: + ctx, rec = _recording() from launchdarkly_ai_langchain_messages import create_langchain_messages_handler llm = _make_llm() - llm.ainvoke = AsyncMock(side_effect=RuntimeError("fail")) - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod + llm.ainvoke = AsyncMock( + side_effect=[ + FakeAIMessage( + tool_calls=[{"name": "myTool", "id": "tu1", "args": {}}], + input_tokens=40, + output_tokens=7, + ), + RuntimeError("second turn died"), + ] + ) + cfg = {**CONFIG, "tools": {"myTool": {"type": "function", "parameters": {}}}} + with ctx, pytest.raises(RuntimeError): + await create_langchain_messages_handler(llm=llm)( + cfg, "q", {"myTool": lambda _: "r"}, {} + ) + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 40 + assert rec.root.attributes["gen_ai.usage.output_tokens"] == 7 - with patch.object(handler_mod, "trace", mock_trace_mod): - h = create_langchain_messages_handler(llm=llm) - with pytest.raises(RuntimeError): - await h(CONFIG, "q", {}, {}) - mock_span.end.assert_called_once() + async def test_writes_no_usage_when_no_turn_ever_reported_any(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + llm = _make_llm() + llm.ainvoke = AsyncMock(side_effect=RuntimeError("died on the first call")) + with ctx, pytest.raises(RuntimeError): + await create_langchain_messages_handler(llm=llm)(CONFIG, "q", {}, {}) + assert "gen_ai.usage.input_tokens" not in rec.root.attributes + + async def test_rethrows_error(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + llm = _make_llm() + llm.ainvoke = AsyncMock(side_effect=RuntimeError("rethrown")) + h = create_langchain_messages_handler(llm=llm) + with pytest.raises(RuntimeError, match="rethrown"): + await h(CONFIG, "q", {}, {}) # --------------------------------------------------------------------------- @@ -646,7 +1012,7 @@ async def test_with_structured_output_called_when_set(self) -> None: llm = _make_llm() structured_llm = MagicMock() structured_llm.ainvoke = AsyncMock( - return_value={"parsed": {"ok": True}, "raw": MagicMock(usage_metadata={})} + return_value={"parsed": {"ok": True}, "raw": FakeAIMessage("")} ) llm.with_structured_output = MagicMock(return_value=structured_llm) h = create_langchain_messages_handler(llm=llm) @@ -660,7 +1026,7 @@ async def test_returns_parsed_object_when_set(self) -> None: llm = _make_llm() structured_llm = MagicMock() structured_llm.ainvoke = AsyncMock( - return_value={"parsed": {"result": 42}, "raw": MagicMock(usage_metadata={})} + return_value={"parsed": {"result": 42}, "raw": FakeAIMessage("")} ) llm.with_structured_output = MagicMock(return_value=structured_llm) h = create_langchain_messages_handler(llm=llm) @@ -684,6 +1050,7 @@ async def test_does_not_throw_when_both_output_format_and_tools_set(self) -> Non "tools": {"t1": {"name": "t1", "type": "function", "parameters": {}}}, } llm = _make_llm() + llm.ainvoke = AsyncMock(return_value=FakeAIMessage("")) h = create_langchain_messages_handler(llm=llm) result = await h(config, "q", {}, {}) assert "output" in result @@ -693,8 +1060,7 @@ async def test_token_usage_from_usage_metadata_when_structured_output(self) -> N config = {**CONFIG, "outputFormat": {"type": "object"}} llm = _make_llm() - raw_msg = MagicMock() - raw_msg.usage_metadata = {"input_tokens": 12, "output_tokens": 8} + raw_msg = FakeAIMessage("", input_tokens=12, output_tokens=8) structured_llm = MagicMock() structured_llm.ainvoke = AsyncMock( return_value={"parsed": {"x": 1}, "raw": raw_msg} @@ -707,7 +1073,7 @@ async def test_token_usage_from_usage_metadata_when_structured_output(self) -> N # --------------------------------------------------------------------------- -# §1.7 Convenience export — §1.x.6 +# §1.7 Convenience export # --------------------------------------------------------------------------- @@ -723,7 +1089,7 @@ def test_calls_through_to_model_call(self) -> None: from launchdarkly_ai_langchain_messages.handler import langchain_messages ctx = {"kind": "user", "key": "u1"} - langchain_messages("my-flag", "hello", ctx, llm=_make_llm()) + langchain_messages("my-flag", "hello", ctx) mock_config_fn.assert_called_once() call_kwargs = mock_config_fn.call_args.kwargs @@ -748,14 +1114,13 @@ def test_callable_without_extra_kwargs(self) -> None: ctx = {"kind": "user", "key": "u1"} langchain_messages("my-flag", "hello", ctx) - mock_config_fn.assert_called_once() mock_config_instance.invoke.assert_called_once_with( "hello", ctx, variables=None ) # --------------------------------------------------------------------------- -# §1.8 Streaming — §1.x.7 and §1.x.8 +# §1.8 Streaming # --------------------------------------------------------------------------- @@ -766,97 +1131,64 @@ def _make_streaming_llm( llm = MagicMock() llm.bind_tools = MagicMock(return_value=llm) - async def _astream(msgs: Any) -> AsyncGenerator: + async def _astream(msgs: Any) -> AsyncGenerator[Any, None]: for c in chunks: - chunk = MagicMock() - chunk.content = c - chunk.usage_metadata = { - "input_tokens": input_tok, - "output_tokens": output_tok, - } - chunk.tool_calls = [] - yield chunk + yield FakeAIMessage(c, input_tokens=input_tok, output_tokens=output_tok) llm.astream = _astream - # ainvoke needed for tool loop fallback (unused here) - llm.ainvoke = AsyncMock(return_value=_make_ai_message("")) + llm.ainvoke = AsyncMock(return_value=FakeAIMessage("")) return llm async def test_stream_defined_and_async_generator(self) -> None: - import launchdarkly_ai_langchain_messages.handler as handler_mod from launchdarkly_ai_langchain_messages import create_langchain_messages_handler llm = self._make_streaming_llm(["hi"]) - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_langchain_messages_handler(llm=llm) + h = create_langchain_messages_handler(llm=llm) assert h.has_stream gen = await h.stream(CONFIG, "q") assert hasattr(gen, "__aiter__") async def test_yields_chunk_events(self) -> None: - import launchdarkly_ai_langchain_messages.handler as handler_mod from launchdarkly_ai_langchain_messages import create_langchain_messages_handler llm = self._make_streaming_llm(["hello ", "world"]) - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_langchain_messages_handler(llm=llm) - events = [e async for e in await h.stream(CONFIG, "q")] + h = create_langchain_messages_handler(llm=llm) + events = [e async for e in await h.stream(CONFIG, "q")] chunks = [e for e in events if e.get("type") == "chunk"] assert len(chunks) == 2 assert chunks[0]["text"] == "hello " assert chunks[1]["text"] == "world" async def test_yields_exactly_one_done_event(self) -> None: - import launchdarkly_ai_langchain_messages.handler as handler_mod from launchdarkly_ai_langchain_messages import create_langchain_messages_handler llm = self._make_streaming_llm(["x"]) - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_langchain_messages_handler(llm=llm) - events = [e async for e in await h.stream(CONFIG, "q")] - done_events = [e for e in events if e.get("type") == "done"] - assert len(done_events) == 1 + h = create_langchain_messages_handler(llm=llm) + events = [e async for e in await h.stream(CONFIG, "q")] + assert len([e for e in events if e.get("type") == "done"]) == 1 async def test_done_event_carries_accumulated_output(self) -> None: - import launchdarkly_ai_langchain_messages.handler as handler_mod from launchdarkly_ai_langchain_messages import create_langchain_messages_handler llm = self._make_streaming_llm(["hello ", "world"]) - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_langchain_messages_handler(llm=llm) - events = [e async for e in await h.stream(CONFIG, "q")] + h = create_langchain_messages_handler(llm=llm) + events = [e async for e in await h.stream(CONFIG, "q")] done = next(e for e in events if e.get("type") == "done") assert done["output"] == "hello world" async def test_done_usage(self) -> None: - import launchdarkly_ai_langchain_messages.handler as handler_mod from launchdarkly_ai_langchain_messages import create_langchain_messages_handler llm = self._make_streaming_llm(["text"], input_tok=7, output_tok=3) - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_langchain_messages_handler(llm=llm) - events = [e async for e in await h.stream(CONFIG, "q")] + h = create_langchain_messages_handler(llm=llm) + events = [e async for e in await h.stream(CONFIG, "q")] done = next(e for e in events if e.get("type") == "done") - assert done["usage"]["input_tokens"] > 0 or done["usage"]["output_tokens"] > 0 - - async def test_streaming_span_name(self) -> None: - from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - - mock_span = MagicMock() - mock_trace_mod, mock_tracer = _make_tracer_patch(mock_span) - mock_tracer.start_span = MagicMock(return_value=mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod - - llm = self._make_streaming_llm(["hi"]) - with patch.object(handler_mod, "trace", mock_trace_mod): - h = create_langchain_messages_handler(llm=llm) - _events = [e async for e in await h.stream(CONFIG, "q")] - span_names = [c[0][0] for c in mock_tracer.start_span.call_args_list] - assert "langchain.stream" in span_names + assert done["usage"]["input_tokens"] == 7 + assert done["usage"]["output_tokens"] == 3 # --------------------------------------------------------------------------- -# §1.5 Streaming telemetry (Appendix A.5 — do not patch _HAS_OTEL=False) +# TELEMETRY-CONTRACT.md sections 1 and 6: the streaming path emits the same tree. # --------------------------------------------------------------------------- @@ -865,48 +1197,141 @@ def _make_streaming_llm(self, chunks: list[str]) -> MagicMock: llm = MagicMock() llm.bind_tools = MagicMock(return_value=llm) - async def _astream(msgs: Any) -> AsyncGenerator: + async def _astream(msgs: Any) -> AsyncGenerator[Any, None]: for c in chunks: - chunk = MagicMock() - chunk.content = c - chunk.usage_metadata = {"input_tokens": 5, "output_tokens": 3} - chunk.tool_calls = [] - yield chunk + yield FakeAIMessage(c, input_tokens=5, output_tokens=3) llm.astream = _astream - llm.ainvoke = AsyncMock(return_value=_make_ai_message("")) + llm.ainvoke = AsyncMock(return_value=FakeAIMessage("")) return llm - async def test_ld_span_attributes_set_during_stream(self) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod + async def test_opens_the_same_root_span_name_as_the_blocking_path(self) -> None: + # A consumer must not be able to tell from the trace which path ran. + ctx, rec = _recording() from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + llm = self._make_streaming_llm(["hi"]) + with ctx: + async for _ in await create_langchain_messages_handler(llm=llm).stream( + CONFIG, "q" + ): + pass + assert rec.root.name == "invoke_agent" + assert "chat gpt-4o" in rec.names + + async def test_carries_the_launchdarkly_attributes_on_the_root(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + llm = self._make_streaming_llm(["hi"]) variables = {"__ld": {"configKey": "k", "variationKey": "v", "runId": "r"}} + with ctx: + async for _ in await create_langchain_messages_handler(llm=llm).stream( + CONFIG, "q", None, 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" + + async def test_ends_every_span_once_when_the_stream_completes(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + llm = self._make_streaming_llm(["hi"]) - with patch.object(handler_mod, "trace", mock_trace_mod): - h = create_langchain_messages_handler(llm=llm) - async for _ in await h.stream(CONFIG, "q", None, variables): + with ctx: + async for _ in await create_langchain_messages_handler(llm=llm).stream( + CONFIG, "q" + ): pass - attrs = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert attrs.get("launchdarkly.operation.type") == "gen_ai" - assert attrs.get("launchdarkly.config.key") == "k" - assert attrs.get("launchdarkly.variation.key") == "v" - assert attrs.get("launchdarkly.run.id") == "r" - - async def test_span_ended_after_stream_completes(self) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_langchain_messages.handler as handler_mod + assert [s.ended for s in rec.spans] == [1] * len(rec.spans) + assert "launchdarkly.stream.abandoned" not in rec.root.attributes + + async def test_writes_the_run_total_to_the_root(self) -> None: + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + llm = self._make_streaming_llm(["hi"]) + with ctx: + async for _ in await create_langchain_messages_handler(llm=llm).stream( + CONFIG, "q" + ): + pass + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 5 + assert rec.root.attributes["gen_ai.usage.total_tokens"] == 8 + + async def test_an_abandoned_stream_still_ends_and_exports_every_span(self) -> None: + # A consumer that breaks out mid-stream makes the 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, and the whole run vanishes from AI Config + # Monitoring along with the feature_flag event it carries. + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + llm = self._make_streaming_llm(["one", "two", "three"]) + with ctx: + gen = await create_langchain_messages_handler(llm=llm).stream(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 test_an_abandoned_stream_is_marked_but_not_failed(self) -> None: + # Stopping early is normal, and LaunchDarkly's own metrics record neither a success nor an + # error for it, so ERROR here would put two dashboards in disagreement about one run. + from opentelemetry.trace import StatusCode + + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + llm = self._make_streaming_llm(["one", "two", "three"]) + with ctx: + gen = await create_langchain_messages_handler(llm=llm).stream(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 == [] + + async def test_fails_the_spans_when_the_stream_raises(self) -> None: + from opentelemetry.trace import StatusCode + + ctx, rec = _recording() + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + + async def _astream(msgs: Any) -> AsyncGenerator[Any, None]: + raise RuntimeError("stream died") + yield # pragma: no cover - unreachable, keeps this an async generator + + llm.astream = _astream + with ctx, pytest.raises(RuntimeError, match="stream died"): + async for _ in await create_langchain_messages_handler(llm=llm).stream( + CONFIG, "q" + ): + pass + assert StatusCode.ERROR in rec.root.statuses + assert rec.root.ended == 1 + + async def test_emits_no_content_by_default_on_the_streaming_path(self) -> None: + ctx, rec = _recording() from launchdarkly_ai_langchain_messages import create_langchain_messages_handler llm = self._make_streaming_llm(["hi"]) - with patch.object(handler_mod, "trace", mock_trace_mod): - h = create_langchain_messages_handler(llm=llm) - async for _ in await h.stream(CONFIG, "q"): + with ctx: + async for _ in await create_langchain_messages_handler(llm=llm).stream( + CONFIG, "q" + ): pass - mock_span.end.assert_called() + 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")] == [] # --------------------------------------------------------------------------- @@ -915,11 +1340,8 @@ async def test_span_ended_after_stream_completes(self) -> None: class TestMaxStepsCap: - """TESTING.md §1.10: The tool loop must break with an error after MAX_STEPS (5) iterations.""" - def _make_tool_call_llm(self) -> MagicMock: - """Returns an LLM that always responds with a tool call.""" - tool_msg = _make_ai_message( + tool_msg = FakeAIMessage( content="", tool_calls=[{"id": "tc_1", "name": "myTool", "args": {}}], input_tokens=1, @@ -931,72 +1353,53 @@ def _make_tool_call_llm(self) -> MagicMock: return llm async def test_invoke_throws_after_max_steps(self) -> None: - import launchdarkly_ai_langchain_messages.handler as handler_mod from launchdarkly_ai_langchain_messages import create_langchain_messages_handler llm = self._make_tool_call_llm() cfg = {**CONFIG, "tools": {"myTool": {"type": "function", "parameters": {}}}} - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_langchain_messages_handler(llm=llm) - with pytest.raises(RuntimeError, match="maximum number of steps"): - await h(cfg, "q", {"myTool": lambda _: "result"}) + h = create_langchain_messages_handler(llm=llm) + with pytest.raises(RuntimeError, match="maximum number of steps"): + await h(cfg, "q", {"myTool": lambda _: "result"}) async def test_invoke_succeeds_at_exactly_max_steps(self) -> None: - import launchdarkly_ai_langchain_messages.handler as handler_mod from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - tool_msg = _make_ai_message( + tool_msg = FakeAIMessage( content="", tool_calls=[{"id": "tc_1", "name": "myTool", "args": {}}], input_tokens=1, output_tokens=1, ) - final_msg = _make_ai_message("Done", input_tokens=1, output_tokens=1) + final_msg = FakeAIMessage("Done", input_tokens=1, output_tokens=1) llm = MagicMock() - llm.ainvoke = AsyncMock( - side_effect=[ - tool_msg, - tool_msg, - tool_msg, - tool_msg, - tool_msg, - tool_msg, - tool_msg, - tool_msg, - tool_msg, - tool_msg, - final_msg, - ] - ) + llm.ainvoke = AsyncMock(side_effect=[tool_msg] * 10 + [final_msg]) llm.bind_tools = MagicMock(return_value=llm) cfg = {**CONFIG, "tools": {"myTool": {"type": "function", "parameters": {}}}} - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_langchain_messages_handler(llm=llm) - result = await h(cfg, "q", {"myTool": lambda _: "result"}) + h = create_langchain_messages_handler(llm=llm) + result = await h(cfg, "q", {"myTool": lambda _: "result"}) assert result["output"] == "Done" async def test_stream_throws_after_max_steps(self) -> None: - import launchdarkly_ai_langchain_messages.handler as handler_mod from launchdarkly_ai_langchain_messages import create_langchain_messages_handler - async def _tool_chunk_stream(_msgs: Any) -> AsyncGenerator: - chunk = MagicMock() - chunk.content = "" - chunk.usage_metadata = {"input_tokens": 1, "output_tokens": 1} - chunk.tool_calls = [{"id": "tc_1", "name": "myTool", "args": {}}] - yield chunk + async def _tool_chunk_stream(_msgs: Any) -> AsyncGenerator[Any, None]: + yield FakeAIMessage( + "", + tool_calls=[{"id": "tc_1", "name": "myTool", "args": {}}], + input_tokens=1, + output_tokens=1, + ) llm = MagicMock() llm.astream = _tool_chunk_stream llm.bind_tools = MagicMock(return_value=llm) cfg = {**CONFIG, "tools": {"myTool": {"type": "function", "parameters": {}}}} - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_langchain_messages_handler(llm=llm) - with pytest.raises(RuntimeError, match="maximum number of steps"): - async for _ in await h.stream(cfg, "q", {"myTool": lambda _: "result"}): - pass + h = create_langchain_messages_handler(llm=llm) + with pytest.raises(RuntimeError, match="maximum number of steps"): + async for _ in await h.stream(cfg, "q", {"myTool": lambda _: "result"}): + pass # --------------------------------------------------------------------------- @@ -1017,76 +1420,232 @@ async def test_history_inserted_between_config_messages_and_user_input( config = { "model": {"name": "gpt-4o"}, - "provider": {"name": "LangChain"}, - "messages": [ - {"role": "user", "content": "First"}, - {"role": "assistant", "content": "Second"}, - ], + "provider": {"name": "OpenAI"}, + "messages": [{"role": "system", "content": "base"}], } llm = _make_llm() h = create_langchain_messages_handler(llm=llm) - await h(config, "Third", {}, {}, self.SAMPLE_HISTORY) + await h(config, "final question", {}, {}, self.SAMPLE_HISTORY) call_args = llm.ainvoke.call_args[0][0] - contents = [getattr(m, "content", "") for m in call_args] - assert contents[0] == "First" - assert contents[1] == "Second" - assert contents[2] == "What is feature flagging?" - assert contents[3] == "Feature flagging is a technique..." - assert contents[4] == "Third" + contents = [str(getattr(m, "content", "")) for m in call_args] + assert contents == [ + "base", + "What is feature flagging?", + "Feature flagging is a technique...", + "final question", + ] async def test_history_with_instructions_path(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler llm = _make_llm() h = create_langchain_messages_handler(llm=llm) - await h(CONFIG, "my question", {}, {}, self.SAMPLE_HISTORY) + await h(CONFIG, "final question", {}, {}, self.SAMPLE_HISTORY) call_args = llm.ainvoke.call_args[0][0] - non_system = [ - m - for m in call_args - if not ( - "system" in str(type(m).__name__).lower() or "System" in str(type(m)) - ) - ] - contents = [getattr(m, "content", "") for m in non_system] - assert contents[0] == "What is feature flagging?" - assert contents[1] == "Feature flagging is a technique..." - assert contents[2] == "my question" + contents = [str(getattr(m, "content", "")) for m in call_args] + assert "What is feature flagging?" in contents + assert "final question" in contents async def test_empty_history_treated_like_no_history(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler llm = _make_llm() h = create_langchain_messages_handler(llm=llm) - await h(CONFIG, "hi", {}, {}, []) - msgs_with_empty = llm.ainvoke.call_args[0][0] - contents_with_empty = [getattr(m, "content", "") for m in msgs_with_empty] - - llm2 = _make_llm() - h2 = create_langchain_messages_handler(llm=llm2) - await h2(CONFIG, "hi", {}, {}) - msgs_without = llm2.ainvoke.call_args[0][0] - contents_without = [getattr(m, "content", "") for m in msgs_without] - - assert contents_with_empty == contents_without + await h(CONFIG, "q", {}, {}, []) + call_args = llm.ainvoke.call_args[0][0] + assert len(call_args) == 2 # system + user async def test_system_role_in_history_filtered_out(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler history_with_system = [ - {"role": "user", "content": "Hello"}, - {"role": "system", "content": "You are evil"}, - {"role": "assistant", "content": "Hi there"}, + *self.SAMPLE_HISTORY, + {"role": "system", "content": "ignored"}, ] llm = _make_llm() h = create_langchain_messages_handler(llm=llm) await h(CONFIG, "q", {}, {}, history_with_system) call_args = llm.ainvoke.call_args[0][0] - history_contents = [ - getattr(m, "content", "") - for m in call_args - if getattr(m, "content", "") in ("Hello", "You are evil", "Hi there") - ] - assert "You are evil" not in history_contents - assert "Hello" in history_contents - assert "Hi there" in history_contents + contents = [str(getattr(m, "content", "")) for m in call_args] + assert "ignored" not in contents + + +# --------------------------------------------------------------------------- +# TELEMETRY-CONTRACT.md section 6: reported is not the same as reported zero +# --------------------------------------------------------------------------- + + +class TestStreamingUsageReported: + """A stream whose chunks carried no usage must not mark the run as having reported. + + Adding an all-zero turn to the accumulator makes a later failure or abandonment write all-zero + totals on the root, which asserts the run cost nothing. That is a different claim from "unknown", + and it is the one thing the reported flag exists to prevent. The blocking path gets this for + free, because `lang_chain_span_usage` returns None for a bag the provider never filled. + """ + + def _two_turn_llm(self, *, first_turn_usage: bool) -> MagicMock: + """One turn that completes with a tool call, then a second turn that dies. + + The first turn is what puts something in the accumulator. Whether it carried usage is the + variable under test, and a turn that dies mid-iteration never reaches the accumulator at + all, which is why a single failing turn cannot exercise this. + """ + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + state = {"turn": 0} + + async def _astream(msgs: Any) -> AsyncGenerator[Any, None]: + state["turn"] += 1 + if state["turn"] == 1: + msg = FakeAIMessage( + "calling", + tool_calls=[{"name": "search", "id": "c1", "args": {}}], + input_tokens=12, + output_tokens=3, + ) + if not first_turn_usage: + msg.usage_metadata = None + yield msg + return + raise RuntimeError("second turn died") + + llm.astream = _astream + llm.ainvoke = AsyncMock(return_value=FakeAIMessage("")) + return llm + + async def test_a_completed_turn_with_no_usage_does_not_mark_reported(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + ctx, rec = _recording() + llm = self._two_turn_llm(first_turn_usage=False) + with ctx, pytest.raises(RuntimeError): + async for _ in await create_langchain_messages_handler(llm=llm).stream( + CONFIG, "q", {"search": lambda _: "ok"} + ): + pass + # Zeros here would assert the run cost nothing. It is unknown, so nothing is written. + assert "gen_ai.usage.input_tokens" not in rec.root.attributes + + async def test_a_completed_turn_with_usage_still_reaches_the_root(self) -> None: + # The guard must not cost the run its real numbers. + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + ctx, rec = _recording() + llm = self._two_turn_llm(first_turn_usage=True) + with ctx, pytest.raises(RuntimeError): + async for _ in await create_langchain_messages_handler(llm=llm).stream( + CONFIG, "q", {"search": lambda _: "ok"} + ): + pass + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 12 + + +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]: + import launchdarkly_ai_langchain_messages.handler as handler_mod + + 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_messages_handler", _factory), + patch.object(handler_mod, "config", fake_config), + ): + handler_mod.langchain_messages("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 TestChatSpanAndTeardownNeverLeak: + """Two ways the run could vanish from the trace, both reachable through content serialisation.""" + + @pytest.mark.asyncio + async def test_an_unserialisable_completion_still_ends_the_chat_span(self) -> None: + # The blocking path has no `finally`, so a raise outside the guard leaves the span open with + # nothing able to recover it. + from opentelemetry.trace import StatusCode + + class _Exploding: + """Shaped like an AIMessage, but reading its content raises.""" + + usage_metadata: ClassVar[dict[str, Any]] = { + "input_tokens": 5, + "output_tokens": 1, + } + tool_calls: ClassVar[list[Any]] = [] + response_metadata: ClassVar[dict[str, Any]] = {} + + @property + def content(self) -> Any: + raise TypeError("cannot serialise this content") + + from launchdarkly_ai_langchain_messages import ( + create_langchain_messages_handler, + ) + + ctx, rec = _recording() + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + llm.ainvoke = AsyncMock(return_value=_Exploding()) + with ctx, pytest.raises(TypeError): + await create_langchain_messages_handler(llm=llm, capture_content=True)( + CONFIG, "q", {}, {} + ) + chat = rec.named("chat ") + assert len(chat) == 1 + assert chat[0].ended == 1, "the chat span leaked" + assert StatusCode.ERROR in chat[0].statuses + + @pytest.mark.asyncio + async def test_an_aclose_failure_does_not_cost_the_run_its_root_span(self) -> None: + # aclose() runs after span teardown, so its own failure cannot take the trace with it. + from launchdarkly_ai_langchain_messages import ( + create_langchain_messages_handler, + ) + + ctx, rec = _recording() + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + + class _BadStream: + def __aiter__(self) -> Any: + return self + + async def __anext__(self) -> Any: + return FakeAIMessage("chunk") + + async def aclose(self) -> None: + raise RuntimeError("vendor teardown exploded") + + llm.astream = MagicMock(return_value=_BadStream()) + llm.ainvoke = AsyncMock(return_value=FakeAIMessage("")) + with ctx: + gen = await create_langchain_messages_handler(llm=llm).stream(CONFIG, "q") + async for _ in gen: + break + await gen.aclose() + assert rec.root.ended == 1, ( + "the root span was lost to a vendor teardown failure" + ) + assert rec.root.attributes["launchdarkly.stream.abandoned"] is True