diff --git a/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py b/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py index 0d2f3ca..c2bbabb 100644 --- a/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py +++ b/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py @@ -9,17 +9,38 @@ AiConfigRep, LDContext, ProviderHandler, + SpanMessage, + SpanMessagePart, config, create_handler, + end_span_once, 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, + to_semconv_finish_reason, +) + +from .spans import ( + RawRunUsage, + fail_span, + finish_model_span, + finish_root_span, + mark_ok, + parent_context_of, + raw_usage_of, + start_model_span, + start_root_span, + start_tool_span, + succeed_span, + to_span_messages, + to_span_parts, + 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: @@ -95,19 +116,46 @@ async def _run_tool_loop( messages: list[dict[str, Any]], system: str | None, tool_handlers: dict[str, Any], -) -> tuple[str, int, int]: - """Runs the Anthropic messages loop, handling tool calls. Returns (output, input_tokens, output_tokens).""" + *, + capture_content: bool = False, + parent: Any = None, + run_usage: RawRunUsage, +) -> tuple[str, dict[str, Any]]: + """Runs the model, then its tools, until the model stops asking for tools. + + Emits one ``chat`` span per model turn and one ``execute_tool`` span per tool call. Both are + parented to *parent*, which is the root's context, so tool spans are siblings of the ``chat`` + span rather than children of it. + + *run_usage* is owned by the caller rather than created here, so a turn that raises does not take + the run's spend with it. + """ + # Not filtered to the tools that have a registered handler, unlike the TypeScript SDK. 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 is the catalog actually sent. tools = _build_tools(config.get("tools") or {}) max_tokens = (config.get("model", {}).get("parameters") or {}).get( "max_tokens", 1024 ) conversation = list(messages) - total_input = 0 - total_output = 0 output = "" steps = 0 + tool_definitions = to_tool_definitions(tools) + while True: + model_span = start_model_span(config, parent) + # Written before the call, so an in-flight or failed turn still shows what it was asked. + # `conversation` grows with each turn, which is what makes a `chat` span self-contained. + if capture_content: + set_input_content_attributes( + model_span, + capture_content, + system_instructions=system, + messages=to_span_messages(conversation), + tool_definitions=tool_definitions, + ) + kwargs: dict[str, Any] = { "model": config["model"]["name"], "max_tokens": max_tokens, @@ -118,9 +166,31 @@ async def _run_tool_loop( if tools: kwargs["tools"] = tools - resp = await client.messages.create(**kwargs) - total_input += resp.usage.input_tokens - total_output += resp.usage.output_tokens + try: + resp = await client.messages.create(**kwargs) + except Exception as exc: + fail_span(model_span, exc) + raise + + raw_usage = raw_usage_of(getattr(resp, "usage", None)) + # Mapped once, into a local, so the span attribute and the output message cannot disagree: + # Anthropic's `end_turn` is semconv's `stop`, and this handler is not the only one whose + # spans a consumer groups by that value. + finish_reason = to_semconv_finish_reason(getattr(resp, "stop_reason", None)) + if capture_content: + set_output_content_attributes( + model_span, + capture_content, + [ + SpanMessage( + role="assistant", + parts=to_span_parts(resp.content), + finish_reason=finish_reason, + ) + ], + ) + finish_model_span(model_span, config, raw_usage, finish_reason) + run_usage.add_turn(raw_usage) if resp.stop_reason != "tool_use": output = "".join( @@ -139,20 +209,35 @@ async def _run_tool_loop( for block in resp.content: if block.type != "tool_use": continue - handler_fn = tool_handlers.get(block.name) - if not handler_fn or not callable(handler_fn): - raise ValueError(f'No handler registered for tool "{block.name}"') - result = ( - await handler_fn(block.input) - if _is_coroutine(handler_fn) - else handler_fn(block.input) + tool_span = start_tool_span(block.name, block.id, parent) + set_tool_call_content_attributes( + tool_span, capture_content, arguments=block.input ) + try: + handler_fn = tool_handlers.get(block.name) + if not handler_fn or not callable(handler_fn): + raise ValueError(f'No handler registered for tool "{block.name}"') + result = ( + await handler_fn(block.input) + if _is_coroutine(handler_fn) + else handler_fn(block.input) + ) + # 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 forever: nothing else knows it exists. + set_tool_call_content_attributes( + tool_span, capture_content, result=result + ) + succeed_span(tool_span) + except Exception as exc: + fail_span(tool_span, exc) + raise tool_results.append( {"type": "tool_result", "tool_use_id": block.id, "content": str(result)} ) conversation.append({"role": "user", "content": tool_results}) - return output, total_input, total_output + return output, run_usage.total def _is_coroutine(fn: Any) -> bool: @@ -162,18 +247,21 @@ def _is_coroutine(fn: Any) -> bool: _MAX_STEPS = 10 -def create_claude_messages_handler() -> ProviderHandler: - """ - Creates a ``ProviderHandler`` for Anthropic Claude (messages API). +def create_claude_messages_handler(*, capture_content: bool = False) -> ProviderHandler: + """Creates a ``ProviderHandler`` for Anthropic Claude (messages API). + Requires ``anthropic`` to be installed as a peer dependency. + + 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. Turning this + on sends the text of every request and response to whatever collector the SDK is pointed at. """ import importlib anthropic_mod = importlib.import_module("anthropic") client = anthropic_mod.AsyncAnthropic() - tracer_name = "@launchdarkly/ai-claude-messages" - async def _call_impl( config: AiConfigRep, user_input: str = "", @@ -184,64 +272,52 @@ async def _call_impl( th = tool_handlers or {} vs = variables or {} - if _HAS_OTEL: - tracer = trace.get_tracer(tracer_name) - span = tracer.start_span("claude.messages") - span.set_attribute("gen_ai.operation.name", "chat") - span.set_attribute("gen_ai.system", "anthropic") - 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) messages, system = _build_messages(config, user_input, vs, history=history) - if span: - prompt_text = (f"system: {system}\n" if system else "") + "\n".join( - f"{m['role']}: {m['content']}" for m in messages - ) - span.add_event("gen_ai.content.prompt", {"gen_ai.prompt": prompt_text}) - prompt_msgs = ( - [{"role": "system", "content": system}] if system else [] - ) + [{"role": m["role"], "content": m["content"]} for m in messages] - set_openllmetry_prompt(span, prompt_msgs) + set_input_content_attributes( + span, + capture_content, + system_instructions=system, + messages=to_span_messages(messages), + ) + # Outside the try, so the failure path can still report the spend of the turns that + # completed before it. + run_usage = RawRunUsage() try: - output, inp, out = await _run_tool_loop( - client, config, messages, system, th + output, usage = await _run_tool_loop( + client, + config, + messages, + system, + th, + capture_content=capture_content, + parent=parent, + run_usage=run_usage, ) - if span: - span.set_attribute( - "gen_ai.response.model", config.get("model", {}).get("name", "") - ) - span.set_attribute("gen_ai.usage.input_tokens", inp) - span.set_attribute("gen_ai.usage.output_tokens", out) - span.set_attribute("gen_ai.usage.total_tokens", inp + out) - 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": inp, "output_tokens": out}, - ) - span.set_status(SpanStatusCode.OK) - span.end() - return { - "output": output, - "usage": {"input_tokens": inp, "output_tokens": out}, - } + set_output_content_attributes( + span, + capture_content, + [ + SpanMessage( + role="assistant", + parts=[SpanMessagePart(type="text", content=output)], + ) + ], + ) + finish_root_span(span, config, usage) + succeed_span(span) + # Raw usage, cache fields intact. `parse_usage` folds them exactly once; handing back a + # pre-folded figure alongside the fields would count the cache twice. + return {"output": output, "usage": usage} 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( @@ -252,7 +328,13 @@ def _stream_impl( history: list[dict[str, Any]] | None = None, ) -> AsyncGenerator[dict[str, Any], None]: return _stream_gen( - client, config, user_input, tool_handlers or {}, variables or {}, history + client, + config, + user_input, + tool_handlers or {}, + variables or {}, + history, + capture_content=capture_content, ) return create_handler(("Anthropic", "messages"), _call_impl, _stream_impl) # type: ignore[arg-type] @@ -265,44 +347,59 @@ 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]: - tracer_name = "@launchdarkly/ai-claude-messages" - if _HAS_OTEL: - span = trace.get_tracer(tracer_name).start_span("claude.messages.stream") - span.set_attribute("gen_ai.operation.name", "chat") - span.set_attribute("gen_ai.system", "anthropic") - span.set_attribute( - "gen_ai.request.model", config.get("model", {}).get("name", "") - ) - set_ld_span_attributes(span, variables) - else: - span = 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. + """ + span = start_root_span(config, variables) + parent = parent_context_of(span) messages, system = _build_messages( config, user_input, variables, include_output_format=False, history=history ) - if span: - prompt_text = (f"system: {system}\n" if system else "") + "\n".join( - f"{m['role']}: {m['content']}" for m in messages - ) - span.add_event("gen_ai.content.prompt", {"gen_ai.prompt": prompt_text}) - prompt_msgs = ([{"role": "system", "content": system}] if system else []) + [ - {"role": m["role"], "content": m["content"]} for m in messages - ] - set_openllmetry_prompt(span, prompt_msgs) + set_input_content_attributes( + span, + capture_content, + system_instructions=system, + messages=to_span_messages(messages), + ) tools = _build_tools(config.get("tools") or {}) + tool_definitions = to_tool_definitions(tools) max_tokens = (config.get("model", {}).get("parameters") or {}).get( "max_tokens", 1024 ) conversation = list(messages) - total_input = 0 - total_output = 0 full_output = "" steps = 0 + ended: set[int] = set() + open_model_span: Any = None + # Outside the try, so the failure and abandonment paths can still report the spend. + run_usage = RawRunUsage() + try: while True: + model_span = start_model_span(config, parent) + open_model_span = model_span + if capture_content: + set_input_content_attributes( + model_span, + capture_content, + system_instructions=system, + messages=to_span_messages(conversation), + tool_definitions=tool_definitions, + ) + kwargs: dict[str, Any] = { "model": config["model"]["name"], "max_tokens": max_tokens, @@ -313,23 +410,47 @@ async def _stream_gen( if tools: kwargs["tools"] = tools - stream = client.messages.stream(**kwargs) - async with stream as s: - async for event in s: - if ( - hasattr(event, "type") - and event.type == "content_block_delta" - and hasattr(event, "delta") - and getattr(event.delta, "type", None) == "text_delta" - ): - text = event.delta.text - full_output += text - yield {"type": "chunk", "text": text} - - final_msg = await s.get_final_message() - - total_input += final_msg.usage.input_tokens - total_output += final_msg.usage.output_tokens + try: + stream = client.messages.stream(**kwargs) + async with stream as s: + async for event in s: + if ( + hasattr(event, "type") + and event.type == "content_block_delta" + and hasattr(event, "delta") + and getattr(event.delta, "type", None) == "text_delta" + ): + text = event.delta.text + full_output += text + yield {"type": "chunk", "text": text} + + final_msg = await s.get_final_message() + except Exception as exc: + fail_span(model_span, exc, ended) + open_model_span = None + raise + + raw_usage = raw_usage_of(getattr(final_msg, "usage", None)) + finish_reason = to_semconv_finish_reason( + getattr(final_msg, "stop_reason", None) + ) + if capture_content: + set_output_content_attributes( + model_span, + capture_content, + [ + SpanMessage( + role="assistant", + parts=to_span_parts(final_msg.content), + finish_reason=finish_reason, + ) + ], + ) + # 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, raw_usage, finish_reason) + open_model_span = None + run_usage.add_turn(raw_usage) if final_msg.stop_reason != "tool_use": break @@ -345,14 +466,29 @@ async def _stream_gen( for block in final_msg.content: if block.type != "tool_use": continue - handler_fn = tool_handlers.get(block.name) - if not handler_fn or not callable(handler_fn): - raise ValueError(f'No handler registered for tool "{block.name}"') - result = ( - await handler_fn(block.input) - if _is_coroutine(handler_fn) - else handler_fn(block.input) + tool_span = start_tool_span(block.name, block.id, parent) + set_tool_call_content_attributes( + tool_span, capture_content, arguments=block.input ) + try: + handler_fn = tool_handlers.get(block.name) + if not handler_fn or not callable(handler_fn): + raise ValueError( + f'No handler registered for tool "{block.name}"' + ) + result = ( + await handler_fn(block.input) + if _is_coroutine(handler_fn) + else handler_fn(block.input) + ) + # Inside the try, for the same reason as the blocking path above. + set_tool_call_content_attributes( + tool_span, capture_content, result=result + ) + succeed_span(tool_span) + except Exception as exc: + fail_span(tool_span, exc, ended) + raise tool_results.append( { "type": "tool_result", @@ -362,43 +498,42 @@ async def _stream_gen( ) conversation.append({"role": "user", "content": 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() + set_output_content_attributes( + span, + capture_content, + [ + SpanMessage( + role="assistant", + parts=[SpanMessagePart(type="text", content=full_output)], + ) + ], + ) + finish_root_span(span, config, run_usage.total) + mark_ok(span) + end_span_once(span, ended) yield { "type": "done", "output": full_output, - "usage": {"input_tokens": total_input, "output_tokens": total_output}, + "usage": run_usage.total, } 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`. On abandonment it is the only chance to close the tree, 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. + 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) def claude_messages( @@ -408,7 +543,13 @@ def claude_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_claude_messages_handler(), **kwargs + key=config_key, + handler=create_claude_messages_handler(capture_content=capture_content), + **kwargs, ).invoke(user_input, context, variables=variables) diff --git a/packages/claude-messages/src/launchdarkly_ai_claude_messages/spans.py b/packages/claude-messages/src/launchdarkly_ai_claude_messages/spans.py new file mode 100644 index 0000000..0d6e550 --- /dev/null +++ b/packages/claude-messages/src/launchdarkly_ai_claude_messages/spans.py @@ -0,0 +1,332 @@ +"""Span construction for the Claude 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. +""" + +from __future__ import annotations + +from typing import Any + +from launchdarkly_ai_server import ( + AiConfigRep, + SpanMessage, + SpanMessagePart, + SpanUsage, + ToolDefinitionInput, + add_cached_tokens_to_input, + number_or_zero, + 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-claude-messages" + +#: Anthropic serves every model behind this handler, so the provider name is a constant. +PROVIDER = "anthropic" + + +def model_name(config: AiConfigRep) -> str: + return str(config.get("model", {}).get("name", "")) + + +# ─── 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, PROVIDER, model_name(config)) + 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 + # `set_span_in_context` with no context argument reads the current one, which is what the + # TypeScript SDK's `trace.setSpan(context.active(), span)` does. + 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, PROVIDER, name) + return span + + +def start_tool_span(tool_name: str, tool_use_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_use_id) + return span + + +# ─── Span finishes ─────────────────────────────────────────────────────────── + + +def finish_root_span(span: Any, config: AiConfigRep, raw_usage: dict[str, Any]) -> None: + """Writes the run-level identity and token totals onto the root. + + 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. + + ``gen_ai.response.model`` is the requested name here. Anthropic does not resolve an alias to a + different snapshot, so there is no other value to report. See TELEMETRY-CONTRACT.md section 2a. + """ + if span is None: + return + span.set_attribute("gen_ai.response.model", model_name(config)) + set_usage_span_attributes(span, add_cached_tokens_to_input(raw_usage)) + + +def finish_model_span( + span: Any, + config: AiConfigRep, + raw_usage: dict[str, Any], + finish_reason: str | None = None, +) -> None: + """Ends one ``chat`` span successfully. *finish_reason* arrives already mapped.""" + if span is None: + return + span.set_attribute("gen_ai.response.model", model_name(config)) + # A list because one response may hold several choices; Anthropic returns one. + if finish_reason: + span.set_attribute("gen_ai.response.finish_reasons", [finish_reason]) + set_usage_span_attributes(span, add_cached_tokens_to_input(raw_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_span_parts(content: Any) -> list[SpanMessagePart]: + """Converts Anthropic content blocks into canonical span parts. + + A plain string is one text part. Unknown block types are dropped rather than guessed at. + """ + if isinstance(content, str): + return [SpanMessagePart(type="text", content=content)] + if not isinstance(content, list): + return [] + + parts: list[SpanMessagePart] = [] + for block in content: + block_type = _attr(block, "type") + if block_type == "text": + parts.append( + SpanMessagePart(type="text", content=str(_attr(block, "text") or "")) + ) + elif block_type == "thinking": + parts.append( + SpanMessagePart( + type="reasoning", content=str(_attr(block, "thinking") or "") + ) + ) + elif block_type == "tool_use": + block_id = _attr(block, "id") + parts.append( + SpanMessagePart( + type="tool_call", + id=block_id if isinstance(block_id, str) else None, + name=str(_attr(block, "name") or ""), + arguments=_attr(block, "input"), + ) + ) + elif block_type == "tool_result": + use_id = _attr(block, "tool_use_id") + parts.append( + SpanMessagePart( + type="tool_call_response", + id=use_id if isinstance(use_id, str) else None, + result=_attr(block, "content"), + ) + ) + return parts + + +def to_span_messages(messages: list[dict[str, Any]]) -> list[SpanMessage]: + return [ + SpanMessage(role=str(m.get("role", "")), parts=to_span_parts(m.get("content"))) + for m in messages + ] + + +def to_tool_definitions(tools: list[dict[str, Any]]) -> list[ToolDefinitionInput]: + """The catalog as sent, so the span reports what the model could actually call.""" + return [ + ToolDefinitionInput( + name=str(t.get("name", "")), + description=t.get("description"), + parameters=t.get("input_schema"), + ) + for t in tools + ] + + +def _attr(obj: Any, name: str) -> Any: + """Reads a field off a provider object or a plain dict, whichever the caller holds. + + The provider SDK hands back objects; the tool loop appends dicts to the same conversation, and + both reach these converters. + """ + if isinstance(obj, dict): + return obj.get(name) + return getattr(obj, name, None) + + +# ─── The run accumulator, in Anthropic's own field names ───────────────────── + + +class RawRunUsage: + """The run's accumulated usage, in Anthropic's own field names. + + Kept unfolded, with the cache figures beside the input total rather than added into it, so + ``parse_usage`` can fold once and derive the breakdown. Pre-folding here would hide the cache + detail from callers, and :func:`add_cached_tokens_to_input` would then count it twice. + + That is also why this is not the client package's ``RunUsage``: that one accumulates + ``SpanUsage``, whose ``input`` is already cache-inclusive, and handing one back as this handler's + return value would double-count the cache downstream. Named differently on purpose, because the + two differ in exactly the way that matters. + + Created by the caller that owns the root span rather than by the tool loop, so a loop that + raises still leaves the run's spend somewhere the root can read it. + """ + + def __init__(self) -> None: + self.total: dict[str, Any] = { + "input_tokens": 0, + "output_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + } + self._turns = 0 + + @property + def reported(self) -> bool: + """Whether any turn reported usage. + + Separates "no call completed" from "a call completed and reported zero". Only the second may + reach a span: all-zero attributes assert the run cost nothing, which a run whose first call + died mid-flight cannot claim. + """ + return self._turns > 0 + + def add_turn(self, raw_usage: dict[str, Any]) -> None: + self._turns += 1 + for key in self.total: + self.total[key] += number_or_zero(raw_usage.get(key)) + + +def raw_usage_of(usage: Any) -> dict[str, Any]: + """Anthropic's usage object as a plain dict, tolerating an absent or partial one. + + Read through :func:`number_or_zero` at the point of use rather than coerced here, so an unknown + cache spelling the provider adds later still reaches ``parse_usage``. + """ + if usage is None: + return {} + if isinstance(usage, dict): + return dict(usage) + fields = ( + "input_tokens", + "output_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", + ) + out: dict[str, Any] = {} + for field in fields: + value = getattr(usage, field, None) + if value is not None: + out[field] = value + return out + + +def span_usage_of(raw_usage: dict[str, Any]) -> SpanUsage: + """This turn's usage as a ``SpanUsage``, with Anthropic's cache rule applied.""" + return add_cached_tokens_to_input(raw_usage) diff --git a/packages/claude-messages/tests/test_handler.py b/packages/claude-messages/tests/test_handler.py index ccd4d90..10d8a6b 100644 --- a/packages/claude-messages/tests/test_handler.py +++ b/packages/claude-messages/tests/test_handler.py @@ -34,18 +34,40 @@ def _tool_use_block(name: str, id: str = "tu1", input: dict | None = None) -> Ma return b +class _Usage: + """Anthropic's usage object, with only the fields Anthropic actually sets. + + A bare MagicMock would answer every cache attribute with a mock, which is not what a real + response looks like and would let a handler read a cache field that was never reported. + """ + + def __init__( + self, + input_tokens: int, + output_tokens: int, + cache_read: int | None = None, + cache_creation: int | None = None, + ) -> None: + self.input_tokens = input_tokens + self.output_tokens = output_tokens + if cache_read is not None: + self.cache_read_input_tokens = cache_read + if cache_creation is not None: + self.cache_creation_input_tokens = cache_creation + + def _anthropic_response( content: list[Any], stop_reason: str = "end_turn", input_tokens: int = 10, output_tokens: int = 5, + cache_read: int | None = None, + cache_creation: int | None = None, ) -> MagicMock: r = MagicMock() r.content = content r.stop_reason = stop_reason - r.usage = MagicMock() - r.usage.input_tokens = input_tokens - r.usage.output_tokens = output_tokens + r.usage = _Usage(input_tokens, output_tokens, cache_read, cache_creation) return r @@ -438,10 +460,81 @@ async def test_multiple_consecutive_tool_calls( # --------------------------------------------------------------------------- # §1.5 Telemetry # --------------------------------------------------------------------------- +# Span recording +# --------------------------------------------------------------------------- + + +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. + + Replaces the old single-MagicMock approach, which could not see a span tree at all: every span + was the same object, so a parent and its children were indistinguishable. + """ + + def __init__(self) -> None: + self.spans: list[RecordedSpan] = [] + + def get_tracer(self, name: str) -> SpanRecorder: + return self + + def start_span(self, name: str, context: Any = None) -> RecordedSpan: + span = RecordedSpan(name, context) + self.spans.append(span) + return span + + def set_span_in_context(self, span: RecordedSpan) -> Any: + return ("context-of", span) + + @property + def root(self) -> RecordedSpan: + return self.spans[0] + + def named(self, prefix: str) -> list[RecordedSpan]: + return [s for s in self.spans if s.name.startswith(prefix)] + + @property + def names(self) -> list[str]: + return [s.name for s in self.spans] + + +def _recording() -> Any: + """Patches the tracer that ``spans.py`` holds, and yields the recorder.""" + import launchdarkly_ai_claude_messages.spans as spans_mod + + recorder = SpanRecorder() + return patch.object(spans_mod, "trace", recorder), recorder def _make_tracer_patch(mock_span: MagicMock) -> Any: - """Creates a patched trace module targeting the handler's imported `trace`.""" + """Kept for the tests that only need to know a span was opened.""" mock_tracer = MagicMock() mock_tracer.start_span = MagicMock(return_value=mock_span) mock_trace_mod = MagicMock() @@ -449,202 +542,467 @@ def _make_tracer_patch(mock_span: MagicMock) -> Any: return mock_trace_mod, mock_tracer -class TestTelemetry: - async def test_span_name(self, mock_anthropic: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, mock_tracer = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod +class TestSpanTree: + """TELEMETRY-CONTRACT.md section 1.""" - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + async def test_opens_a_root_span_named_invoke_agent( + self, mock_anthropic: MagicMock + ) -> None: + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - h = create_claude_messages_handler() - await h(CONFIG, "q", {}, {}) - mock_tracer.start_span.assert_called_with("claude.messages") + with ctx: + await create_claude_messages_handler()(CONFIG, "q", {}, {}) + assert rec.root.name == "invoke_agent" + assert rec.root.attributes["gen_ai.operation.name"] == "invoke_agent" - async def test_gen_ai_system(self, mock_anthropic: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + async def test_emits_one_chat_child_per_model_turn( + self, mock_anthropic: MagicMock + ) -> None: + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + with ctx: + await create_claude_messages_handler()(CONFIG, "q", {}, {}) + chats = rec.named("chat ") + assert len(chats) == 1 + assert chats[0].name == "chat claude-3-sonnet-20240229" + assert chats[0].attributes["gen_ai.operation.name"] == "chat" + # Parented to the root, not to nothing. + assert chats[0].context == ("context-of", rec.root) - h = create_claude_messages_handler() - await h(CONFIG, "q", {}, {}) - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("gen_ai.system") == "anthropic" + async def test_names_the_chat_span_after_the_model( + self, mock_anthropic: MagicMock + ) -> None: + # The semantic conventions name an inference span `{operation} {model}`. A bare `chat` tells + # a reader nothing about which model ran. + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - async def test_gen_ai_request_model(self, mock_anthropic: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + cfg = {**CONFIG, "model": {"name": "claude-opus-4"}} + with ctx: + await create_claude_messages_handler()(cfg, "q", {}, {}) + assert "chat claude-opus-4" in rec.names - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + async def test_emits_a_chat_span_per_turn_of_a_tool_loop( + self, mock_anthropic: MagicMock + ) -> None: + mock_anthropic.messages.create = AsyncMock( + side_effect=[ + _anthropic_response( + [_tool_use_block("myTool", "tu1", {"a": 1})], stop_reason="tool_use" + ), + _anthropic_response([_text_block("done")]), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - h = create_claude_messages_handler() - await h(CONFIG, "q", {}, {}) - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("gen_ai.request.model") == CONFIG["model"]["name"] + with ctx: + await create_claude_messages_handler()( + CONFIG, "q", {"myTool": lambda _: "result"}, {} + ) + assert len(rec.named("chat ")) == 2 - async def test_token_attributes_set(self, mock_anthropic: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + async def test_emits_an_execute_tool_span_per_tool_call( + self, mock_anthropic: MagicMock + ) -> None: + mock_anthropic.messages.create = AsyncMock( + side_effect=[ + _anthropic_response( + [_tool_use_block("myTool", "tu1", {"a": 1})], stop_reason="tool_use" + ), + _anthropic_response([_text_block("done")]), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + with ctx: + await create_claude_messages_handler()( + CONFIG, "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" - h = create_claude_messages_handler() - 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 + async def test_tool_spans_are_siblings_of_chat_not_children( + self, mock_anthropic: MagicMock + ) -> None: + # Both take the root's context. See TELEMETRY-CONTRACT.md section 1. + mock_anthropic.messages.create = AsyncMock( + side_effect=[ + _anthropic_response( + [_tool_use_block("myTool", "tu1")], stop_reason="tool_use" + ), + _anthropic_response([_text_block("done")]), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - async def test_span_status_ok(self, mock_anthropic: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + with ctx: + await create_claude_messages_handler()( + CONFIG, "q", {"myTool": lambda _: "r"}, {} + ) + assert rec.named("execute_tool ")[0].context == ("context-of", rec.root) - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + async def test_every_span_is_ended(self, mock_anthropic: MagicMock) -> None: + mock_anthropic.messages.create = AsyncMock( + side_effect=[ + _anthropic_response( + [_tool_use_block("myTool", "tu1")], stop_reason="tool_use" + ), + _anthropic_response([_text_block("done")]), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - h = create_claude_messages_handler() - await h(CONFIG, "q", {}, {}) - from opentelemetry.trace import StatusCode + with ctx: + await create_claude_messages_handler()( + CONFIG, "q", {"myTool": lambda _: "r"}, {} + ) + assert [s.ended for s in rec.spans] == [1] * len(rec.spans) - mock_span.set_status.assert_called_with(StatusCode.OK) - async def test_span_end_always_called(self, mock_anthropic: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod +class TestRootSpanAttributes: + """TELEMETRY-CONTRACT.md sections 2 and 2a.""" - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + async def test_writes_both_provider_keys_and_the_requested_model( + self, mock_anthropic: MagicMock + ) -> None: + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - h = create_claude_messages_handler() - await h(CONFIG, "q", {}, {}) - mock_span.end.assert_called_once() + with ctx: + await create_claude_messages_handler()(CONFIG, "q", {}, {}) + attrs = rec.root.attributes + assert attrs["gen_ai.system"] == "anthropic" + assert attrs["gen_ai.provider.name"] == "anthropic" + assert attrs["gen_ai.request.model"] == "claude-3-sonnet-20240229" - async def test_gen_ai_operation_name(self, mock_anthropic: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + async def test_response_model_is_the_requested_name( + self, mock_anthropic: MagicMock + ) -> None: + # Anthropic does not resolve an alias to a different snapshot, so there is no other value + # to report. Section 2a. + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + with ctx: + await create_claude_messages_handler()(CONFIG, "q", {}, {}) + assert ( + rec.root.attributes["gen_ai.response.model"] == "claude-3-sonnet-20240229" + ) - h = create_claude_messages_handler() - await h(CONFIG, "q", {}, {}) - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("gen_ai.operation.name") == "chat" + async def test_carries_the_launchdarkly_attributes_and_feature_flag_event( + self, mock_anthropic: MagicMock + ) -> None: + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - async def test_gen_ai_content_prompt_event(self, mock_anthropic: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + variables = { + "__ld": { + "configKey": "k", + "variationKey": "v", + "runId": "r", + "graphKey": "g", + } + } + with ctx: + await create_claude_messages_handler()(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, mock_anthropic: MagicMock + ) -> None: + # The root is the only span a config-scoped query finds; children must not duplicate it. + mock_anthropic.messages.create = AsyncMock( + side_effect=[ + _anthropic_response( + [_tool_use_block("myTool", "tu1")], stop_reason="tool_use" + ), + _anthropic_response([_text_block("done")]), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + variables = {"__ld": {"configKey": "k", "variationKey": "v", "runId": "r"}} + with ctx: + await create_claude_messages_handler()( + CONFIG, "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] - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + async def test_carries_the_run_total_not_one_turn( + self, mock_anthropic: MagicMock + ) -> None: + mock_anthropic.messages.create = AsyncMock( + side_effect=[ + _anthropic_response( + [_tool_use_block("myTool", "tu1")], + stop_reason="tool_use", + input_tokens=10, + output_tokens=1, + ), + _anthropic_response( + [_text_block("done")], input_tokens=20, output_tokens=2 + ), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - h = create_claude_messages_handler() - await h(CONFIG, "user input", {}, {}) - event_names = [c[0][0] for c in mock_span.add_event.call_args_list] - assert "gen_ai.content.prompt" in event_names + with ctx: + await create_claude_messages_handler()( + CONFIG, "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 + + +class TestChatSpanAttributes: + """TELEMETRY-CONTRACT.md sections 3, 5 and 8.""" - async def test_gen_ai_content_completion_event( + async def test_writes_all_seven_usage_attributes_including_zeros( self, mock_anthropic: MagicMock ) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + with ctx: + await create_claude_messages_handler()(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_folds_cache_tokens_into_the_input_total( + self, mock_anthropic: MagicMock + ) -> None: + # Anthropic reports cache beside input, so the real input is the sum of all three. This is + # the assertion that catches a fold in the wrong direction. + mock_anthropic.messages.create = AsyncMock( + return_value=_anthropic_response( + [_text_block("hi")], + input_tokens=3, + output_tokens=10, + cache_read=19971, + cache_creation=3580, + ) + ) + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - h = create_claude_messages_handler() - 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 + with ctx: + await create_claude_messages_handler()(CONFIG, "q", {}, {}) + attrs = rec.named("chat ")[0].attributes + assert attrs["gen_ai.usage.input_tokens"] == 23554 + assert attrs["gen_ai.usage.cache_read.input_tokens"] == 19971 + assert attrs["gen_ai.usage.cache_creation.input_tokens"] == 3580 + assert attrs["gen_ai.usage.total_tokens"] == 23564 - async def test_total_tokens_attribute(self, mock_anthropic: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + async def test_yields_raw_usage_so_parse_usage_folds_exactly_once( + self, mock_anthropic: MagicMock + ) -> None: + # The returned bag keeps the cache fields unfolded. Returning a pre-folded input alongside + # them would count the cache twice downstream. + mock_anthropic.messages.create = AsyncMock( + return_value=_anthropic_response( + [_text_block("hi")], + input_tokens=3, + output_tokens=1, + cache_read=100, + cache_creation=50, + ) + ) + from launchdarkly_ai_claude_messages import create_claude_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + result = await create_claude_messages_handler()(CONFIG, "q", {}, {}) + assert result["usage"]["input_tokens"] == 3 + assert result["usage"]["cache_read_input_tokens"] == 100 + assert result["usage"]["cache_creation_input_tokens"] == 50 - h = create_claude_messages_handler() - 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.total_tokens" in attrs - assert attrs["gen_ai.usage.total_tokens"] == attrs.get( - "gen_ai.usage.input_tokens", 0 - ) + attrs.get("gen_ai.usage.output_tokens", 0) - - async def test_gen_ai_response_model(self, mock_anthropic: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + async def test_reports_the_mapped_finish_reason( + self, mock_anthropic: MagicMock + ) -> None: + # Anthropic's `end_turn` is semconv's `stop`. + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + with ctx: + await create_claude_messages_handler()(CONFIG, "q", {}, {}) + assert rec.named("chat ")[0].attributes["gen_ai.response.finish_reasons"] == [ + "stop" + ] - h = create_claude_messages_handler() - 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_maps_tool_use_to_tool_calls(self, mock_anthropic: MagicMock) -> None: + mock_anthropic.messages.create = AsyncMock( + side_effect=[ + _anthropic_response( + [_tool_use_block("myTool", "tu1")], stop_reason="tool_use" + ), + _anthropic_response([_text_block("done")]), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - async def test_ld_span_attributes(self, mock_anthropic: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + with ctx: + await create_claude_messages_handler()( + CONFIG, "q", {"myTool": lambda _: "r"}, {} + ) + first = rec.named("chat ")[0] + assert first.attributes["gen_ai.response.finish_reasons"] == ["tool_calls"] - variables = { - "__ld": { - "configKey": "my-config", - "variationKey": "v1", - "runId": "run-abc", - } - } - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + async def test_omits_the_finish_reason_when_the_provider_gives_none( + self, mock_anthropic: MagicMock + ) -> None: + mock_anthropic.messages.create = AsyncMock( + return_value=_anthropic_response([_text_block("hi")], stop_reason=None) + ) + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - h = create_claude_messages_handler() - 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 + with ctx: + await create_claude_messages_handler()(CONFIG, "q", {}, {}) + assert "gen_ai.response.finish_reasons" not in rec.named("chat ")[0].attributes - async def test_ld_graph_key_set_when_present( + async def test_sets_status_ok_on_a_successful_turn( self, mock_anthropic: MagicMock ) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + from opentelemetry.trace import StatusCode - variables = { - "__ld": { - "configKey": "my-config", - "variationKey": "v1", - "runId": "run-abc", - "graphKey": "my-graph", - } + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + with ctx: + await create_claude_messages_handler()(CONFIG, "q", {}, {}) + assert StatusCode.OK in rec.named("chat ")[0].statuses + + +class TestContentCapture: + """TELEMETRY-CONTRACT.md section 7.""" + + async def test_emits_no_content_at_all_by_default( + self, mock_anthropic: MagicMock + ) -> None: + # Conversation content is PII. This is the assertion worth pinning hardest. + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + with ctx: + await create_claude_messages_handler()(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")] == [] + + async def test_puts_prompt_and_completion_on_spans_when_enabled( + self, mock_anthropic: MagicMock + ) -> None: + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + with ctx: + await create_claude_messages_handler(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 World" + assert "gen_ai.output.messages" in chat.attributes + + async def test_records_the_tool_catalog_on_the_chat_span_when_enabled( + self, mock_anthropic: MagicMock + ) -> None: + import json + + cfg = { + **CONFIG, + "tools": {"myTool": {"description": "d", "parameters": {"type": "object"}}}, } - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - h = create_claude_messages_handler() - 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_claude_messages_handler(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, mock_anthropic: MagicMock + ) -> None: + mock_anthropic.messages.create = AsyncMock( + side_effect=[ + _anthropic_response( + [_tool_use_block("myTool", "tu1", {"city": "NYC"})], + stop_reason="tool_use", + ), + _anthropic_response([_text_block("done")]), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + with ctx: + await create_claude_messages_handler(capture_content=True)( + CONFIG, "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, mock_anthropic: MagicMock + ) -> None: + # Redundant and deprecated, but every published version emitted them. + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + with ctx: + await create_claude_messages_handler(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 # --------------------------------------------------------------------------- @@ -653,56 +1011,104 @@ async def test_ld_graph_key_set_when_present( class TestErrorHandling: - async def test_records_exception_on_span(self, mock_anthropic: MagicMock) -> None: + """TELEMETRY-CONTRACT.md section 6.""" + + async def test_fails_the_chat_span_when_the_provider_call_raises( + self, mock_anthropic: MagicMock + ) -> None: + from opentelemetry.trace import StatusCode + mock_anthropic.messages.create = AsyncMock( side_effect=RuntimeError("api error") ) - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + with ctx, pytest.raises(RuntimeError): + await create_claude_messages_handler()(CONFIG, "q", {}, {}) + chat = rec.named("chat ")[0] + assert len(chat.exceptions) == 1 + assert StatusCode.ERROR in chat.statuses + assert chat.ended == 1 - h = create_claude_messages_handler() - with pytest.raises(RuntimeError): - await h(CONFIG, "q", {}, {}) - mock_span.record_exception.assert_called_once() + async def test_fails_the_root_span_too(self, mock_anthropic: MagicMock) -> None: + from opentelemetry.trace import StatusCode - async def test_sets_span_status_error(self, mock_anthropic: MagicMock) -> None: mock_anthropic.messages.create = AsyncMock( side_effect=RuntimeError("api error") ) - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + with ctx, pytest.raises(RuntimeError): + await create_claude_messages_handler()(CONFIG, "q", {}, {}) + assert len(rec.root.exceptions) == 1 + assert StatusCode.ERROR in rec.root.statuses + assert rec.root.ended == 1 - h = create_claude_messages_handler() - with pytest.raises(RuntimeError): - await h(CONFIG, "q", {}, {}) + async def test_fails_the_execute_tool_span_when_a_tool_raises( + self, mock_anthropic: MagicMock + ) -> None: from opentelemetry.trace import StatusCode - status_calls = [c[0][0] for c in mock_span.set_status.call_args_list] - assert StatusCode.ERROR in status_calls + mock_anthropic.messages.create = AsyncMock( + return_value=_anthropic_response( + [_tool_use_block("myTool", "tu1")], stop_reason="tool_use" + ) + ) + + def _boom(_: Any) -> Any: + raise RuntimeError("tool exploded") + + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + with ctx, pytest.raises(RuntimeError, match="tool exploded"): + await create_claude_messages_handler()(CONFIG, "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, mock_anthropic: MagicMock) -> None: + async def test_reports_the_spend_of_completed_turns_on_a_failed_run( + self, mock_anthropic: MagicMock + ) -> None: + # The first turn was billed. The root is the only span a config-scoped cost query finds it on. mock_anthropic.messages.create = AsyncMock( - side_effect=RuntimeError("api error") + side_effect=[ + _anthropic_response( + [_tool_use_block("myTool", "tu1")], + stop_reason="tool_use", + input_tokens=40, + output_tokens=7, + ), + RuntimeError("second turn died"), + ] ) - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_claude_messages import create_claude_messages_handler + with ctx, pytest.raises(RuntimeError): + await create_claude_messages_handler()( + CONFIG, "q", {"myTool": lambda _: "r"}, {} + ) + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 40 + assert rec.root.attributes["gen_ai.usage.output_tokens"] == 7 - h = create_claude_messages_handler() - 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, mock_anthropic: MagicMock + ) -> None: + # All-zero attributes would assert the run cost nothing, which a run whose first call died + # mid-flight cannot claim. An absent attribute correctly says "unknown". + mock_anthropic.messages.create = AsyncMock( + side_effect=RuntimeError("died on the first call") + ) + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + with ctx, pytest.raises(RuntimeError): + await create_claude_messages_handler()(CONFIG, "q", {}, {}) + assert "gen_ai.usage.input_tokens" not in rec.root.attributes async def test_rethrows_error(self, mock_anthropic: MagicMock) -> None: mock_anthropic.messages.create = AsyncMock(side_effect=RuntimeError("rethrown")) @@ -836,7 +1242,7 @@ def _make_stream_context( events = [_make_stream_event(c) for c in chunks] final_msg = MagicMock() final_msg.stop_reason = "end_turn" - final_msg.usage = MagicMock(input_tokens=input_tok, output_tokens=output_tok) + final_msg.usage = _Usage(input_tok, output_tok) final_msg.content = [] class _FakeStream: @@ -1101,6 +1507,8 @@ async def _iter() -> AsyncGenerator: class TestStreamingTelemetry: + """TELEMETRY-CONTRACT.md sections 1 and 6. The streaming path emits the same tree.""" + def _patch_stream( self, mock_anthropic: MagicMock, @@ -1111,53 +1519,134 @@ def _patch_stream( ctx, _ = _make_stream_context(chunks, input_tok, output_tok) mock_anthropic.messages.stream = MagicMock(return_value=ctx) - async def test_span_started_during_stream(self, mock_anthropic: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, mock_tracer = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + async def test_opens_the_same_root_span_name_as_the_blocking_path( + self, mock_anthropic: MagicMock + ) -> None: + # A consumer must not be able to tell from the trace which path ran. + self._patch_stream(mock_anthropic, ["hi"]) + ctx, rec = _recording() from launchdarkly_ai_claude_messages import create_claude_messages_handler - self._patch_stream(mock_anthropic, ["hi"]) - with patch.object(handler_mod, "trace", mock_trace_mod): - h = create_claude_messages_handler() - async for _ in await h.stream(CONFIG, "q"): + with ctx: + async for _ in await create_claude_messages_handler().stream(CONFIG, "q"): pass - mock_tracer.start_span.assert_called_with("claude.messages.stream") + assert rec.root.name == "invoke_agent" + assert "chat claude-3-sonnet-20240229" in rec.names - async def test_ld_span_attributes_set_during_stream( + async def test_carries_the_launchdarkly_attributes_on_the_root( self, mock_anthropic: MagicMock ) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + self._patch_stream(mock_anthropic, ["hi"]) + ctx, rec = _recording() from launchdarkly_ai_claude_messages import create_claude_messages_handler variables = {"__ld": {"configKey": "k", "variationKey": "v", "runId": "r"}} + with ctx: + async for _ in await create_claude_messages_handler().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, mock_anthropic: MagicMock + ) -> None: self._patch_stream(mock_anthropic, ["hi"]) - with patch.object(handler_mod, "trace", mock_trace_mod): - h = create_claude_messages_handler() - async for _ in await h.stream(CONFIG, "q", None, variables): + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + with ctx: + async for _ in await create_claude_messages_handler().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" + assert [s.ended for s in rec.spans] == [1] * len(rec.spans) + assert "launchdarkly.stream.abandoned" not in rec.root.attributes - async def test_span_ended_after_stream_completes( + async def test_writes_the_run_total_to_the_root( self, mock_anthropic: MagicMock ) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_claude_messages.handler as handler_mod + self._patch_stream(mock_anthropic, ["hi"], input_tok=11, output_tok=4) + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + with ctx: + async for _ in await create_claude_messages_handler().stream(CONFIG, "q"): + pass + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 11 + assert rec.root.attributes["gen_ai.usage.total_tokens"] == 15 + + async def test_an_abandoned_stream_still_ends_and_exports_every_span( + self, mock_anthropic: MagicMock + ) -> 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. + self._patch_stream(mock_anthropic, ["one", "two", "three"]) + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + with ctx: + gen = await create_claude_messages_handler().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, mock_anthropic: MagicMock + ) -> 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 + + self._patch_stream(mock_anthropic, ["one", "two", "three"]) + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + with ctx: + gen = await create_claude_messages_handler().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, mock_anthropic: MagicMock + ) -> None: + from opentelemetry.trace import StatusCode + + mock_anthropic.messages.stream = MagicMock( + side_effect=RuntimeError("stream died") + ) + ctx, rec = _recording() from launchdarkly_ai_claude_messages import create_claude_messages_handler + with ctx, pytest.raises(RuntimeError, match="stream died"): + async for _ in await create_claude_messages_handler().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, mock_anthropic: MagicMock + ) -> None: self._patch_stream(mock_anthropic, ["hi"]) - with patch.object(handler_mod, "trace", mock_trace_mod): - h = create_claude_messages_handler() - async for _ in await h.stream(CONFIG, "q"): + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + with ctx: + async for _ in await create_claude_messages_handler().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")] == [] # --------------------------------------------------------------------------- @@ -1238,3 +1727,79 @@ async def test_system_role_in_history_filtered_out( msgs = mock_anthropic.messages.create.call_args.kwargs["messages"] roles = [m["role"] for m in msgs] assert "system" not in roles + + +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_claude_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_claude_messages_handler", _factory), + patch.object(handler_mod, "config", fake_config), + ): + handler_mod.claude_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 TestToolSpanNeverLeaks: + """A raise while recording a tool result must not leave its span open. + + Serialising a tool result can raise, most easily when capture_content is on and the result is + not JSON-serialisable. The success-side content write used to sit outside the try, so that raise + skipped both the finish and the failure path: only the root was marked ERROR, and the tool span + was never ended, so the exporter never saw it. + """ + + async def test_an_unserialisable_tool_result_still_ends_the_tool_span( + self, mock_anthropic: MagicMock + ) -> None: + from opentelemetry.trace import StatusCode + + mock_anthropic.messages.create = AsyncMock( + side_effect=[ + _anthropic_response( + [_tool_use_block("myTool", "tu1")], stop_reason="tool_use" + ), + _anthropic_response([_text_block("done")]), + ] + ) + + class _Unserialisable: + __slots__ = () + + ctx, rec = _recording() + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + with ctx, pytest.raises(TypeError): + await create_claude_messages_handler(capture_content=True)( + CONFIG, "q", {"myTool": lambda _: _Unserialisable()}, {} + ) + + tool_spans = rec.named("execute_tool ") + assert len(tool_spans) == 1 + assert tool_spans[0].ended == 1, "the tool span leaked" + assert StatusCode.ERROR in tool_spans[0].statuses