diff --git a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py index 9144e26..2daa65f 100644 --- a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py +++ b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py @@ -9,17 +9,40 @@ AiConfigRep, LDContext, ProviderHandler, + RunUsage, + SpanMessage, + SpanMessagePart, config, create_handler, + create_run_usage, + 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, +) + +from .spans import ( + fail_span, + finish_model_span, + finish_reason_of, + finish_root_span, + mark_ok, + model_name, + parent_context_of, + set_response_output_content, + split_input_messages, + start_model_span, + start_root_span, + start_tool_span, + succeed_span, + to_span_usage, + 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,9 @@ 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. 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 on the span is the catalog actually sent. return [ { "type": "function", @@ -71,6 +97,10 @@ def _build_input_messages( return result +def _json_schema_format(schema: dict[str, Any]) -> dict[str, Any]: + return {"type": "json_schema", "name": "output", "schema": schema, "strict": False} + + def _is_coroutine(fn: Any) -> bool: return asyncio.iscoroutinefunction(fn) @@ -78,18 +108,67 @@ def _is_coroutine(fn: Any) -> bool: _MAX_STEPS = 10 -def create_openai_messages_handler() -> ProviderHandler: +async def _run_model_turn( + client: Any, + config: AiConfigRep, + params: dict[str, Any], + tool_definitions: list[Any], + *, + capture_content: bool, + parent: Any, + run_usage: RunUsage, +) -> Any: + """Runs one provider turn under its own ``chat`` child span. + + Written before the call, so an in-flight or failed turn still shows what it was asked. Returns + the raw provider response so the caller can inspect its output items. + """ + model_span = start_model_span(config, parent) + # Everything that touches this span sits inside the try, including the content writes on both + # sides of the call. Serialising conversation content raises on anything that is not + # JSON-serialisable, and a raise outside the guard would leave this span open forever: only the + # root gets failed, and nothing else knows the chat span exists. + try: + if capture_content: + system_instructions, messages = split_input_messages(params["input"]) + set_input_content_attributes( + model_span, + capture_content, + system_instructions=system_instructions, + messages=messages, + tool_definitions=tool_definitions, + ) + + response = await client.responses.create(**params) + + set_response_output_content(model_span, capture_content, response) + finish_reason = finish_reason_of(response) + response_model = getattr(response, "model", None) or model_name(config) + usage = to_span_usage(getattr(response, "usage", None)) + finish_model_span(model_span, response_model, usage, finish_reason) + except Exception as exc: + fail_span(model_span, exc) + raise + # `to_span_usage` of an absent bag is still a real object, so a turn that completed without + # reported usage counts as reported: the call happened, whatever the provider said. + run_usage.add(usage) + return response + + +def create_openai_messages_handler(*, capture_content: bool = False) -> ProviderHandler: """ Creates a ``ProviderHandler`` for OpenAI (responses API). Requires ``openai`` 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. """ import importlib openai_mod = importlib.import_module("openai") client = openai_mod.AsyncOpenAI() - tracer_name = "@launchdarkly/ai-openai-messages" - async def _call_impl( config: AiConfigRep, user_input: str = "", @@ -100,55 +179,50 @@ async def _call_impl( th = tool_handlers or {} vs = variables or {} - if _HAS_OTEL: - span = trace.get_tracer(tracer_name).start_span("openai.response") - span.set_attribute("gen_ai.operation.name", "chat") - span.set_attribute("gen_ai.system", "openai") - 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) - tools = _build_tools(config.get("tools") or {}) - input_messages = _build_input_messages(config, user_input, vs, history) + # Declared out here, not inside the `try`, so the failure path can still report the tokens + # the run had already spent. + run_usage = create_run_usage() - if span: - span.add_event( - "gen_ai.content.prompt", {"gen_ai.prompt": json.dumps(input_messages)} - ) - set_openllmetry_prompt( + try: + tools = _build_tools(config.get("tools") or {}) + input_messages = _build_input_messages(config, user_input, vs, history) + tool_definitions = to_tool_definitions(tools) + + root_system, root_messages = split_input_messages(input_messages) + set_input_content_attributes( span, - [{"role": m["role"], "content": m["content"]} for m in input_messages], + capture_content, + system_instructions=root_system, + messages=root_messages, ) - try: - kwargs: dict[str, Any] = { + params: dict[str, Any] = { "model": config["model"]["name"], "input": input_messages, } if tools: - kwargs["tools"] = tools + params["tools"] = tools if config.get("outputFormat"): - kwargs["text"] = { - "format": { - "type": "json_schema", - "name": "output", - "schema": config["outputFormat"], - "strict": False, - } - } + params["text"] = {"format": _json_schema_format(config["outputFormat"])} + + response = await _run_model_turn( + client, + config, + params, + tool_definitions, + capture_content=capture_content, + parent=parent, + run_usage=run_usage, + ) - response = await client.responses.create(**kwargs) - total_input = getattr(response.usage, "input_tokens", 0) or 0 - total_output = getattr(response.usage, "output_tokens", 0) or 0 steps = 0 - while True: tool_calls = [ item - for item in (response.output or []) + for item in (getattr(response, "output", None) or []) if getattr(item, "type", None) == "function_call" ] if not tool_calls: @@ -162,15 +236,32 @@ async def _call_impl( tool_outputs = [] for tc in tool_calls: - args = json.loads(tc.arguments) - handler_fn = th.get(tc.name) - if not handler_fn: - raise ValueError(f'No handler registered for tool "{tc.name}"') - result = ( - await handler_fn(args) - if _is_coroutine(handler_fn) - else handler_fn(args) + tool_span = start_tool_span(tc.name, tc.call_id, parent) + set_tool_call_content_attributes( + tool_span, capture_content, arguments=tc.arguments ) + try: + args = json.loads(tc.arguments) + 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 = ( + await handler_fn(args) + if _is_coroutine(handler_fn) + else handler_fn(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 + ) + succeed_span(tool_span) + except Exception as exc: + fail_span(tool_span, exc) + raise tool_outputs.append( { "type": "function_call_output", @@ -179,51 +270,43 @@ async def _call_impl( } ) - response = await client.responses.create( - model=config["model"]["name"], - previous_response_id=response.id, - input=tool_outputs, - ) - total_input += getattr(response.usage, "input_tokens", 0) or 0 - total_output += getattr(response.usage, "output_tokens", 0) or 0 - - output = getattr(response, "output_text", None) or "" - - 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", + response = await _run_model_turn( + client, + config, { - "gen_ai.completion": output - if isinstance(output, str) - else json.dumps(output) + "model": config["model"]["name"], + "previous_response_id": response.id, + "input": tool_outputs, }, + tool_definitions, + capture_content=capture_content, + parent=parent, + run_usage=run_usage, ) - 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 = getattr(response, "output_text", None) or "" + set_output_content_attributes( + span, capture_content, _final_output_messages(output) + ) + response_model = getattr(response, "model", None) or model_name(config) + finish_root_span(span, response_model, run_usage.total) + succeed_span(span) + # Cache keys are deliberately omitted: OpenAI's input already includes them, and + # `parse_usage` would otherwise fold them in a second time. 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() + # Report what the turns that did complete already cost. Falls back to the requested + # model name rather than tracking the last answering model, matching the TypeScript + # SDK's blocking failure path. + if run_usage.reported: + finish_root_span(span, model_name(config), run_usage.total) + fail_span(span, exc) raise def _stream_impl( @@ -234,12 +317,26 @@ 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(("OpenAI", "messages"), _call_impl, _stream_impl) # type: ignore[arg-type] +def _final_output_messages(output: str) -> list[SpanMessage]: + return [ + SpanMessage( + role="assistant", parts=[SpanMessagePart(type="text", content=output)] + ) + ] + + async def _stream_gen( client: Any, config: AiConfigRep, @@ -247,68 +344,98 @@ 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-openai-messages" - if _HAS_OTEL: - span = trace.get_tracer(tracer_name).start_span("openai.response.stream") - span.set_attribute("gen_ai.operation.name", "chat") - span.set_attribute("gen_ai.system", "openai") - 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. - tools = _build_tools(config.get("tools") or {}) - input_messages = _build_input_messages(config, user_input, variables, history) + 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. + """ + span = start_root_span(config, variables) + parent = parent_context_of(span) - if span: - span.add_event( - "gen_ai.content.prompt", {"gen_ai.prompt": json.dumps(input_messages)} - ) - set_openllmetry_prompt( - span, [{"role": m["role"], "content": m["content"]} for m in input_messages] + ended: set[int] = set() + open_model_span: Any = None + # Outside the try, so the failure and abandonment paths can still report the spend and the + # model that answered. + run_usage = create_run_usage() + last_response_model = model_name(config) + + try: + tools = _build_tools(config.get("tools") or {}) + input_messages = _build_input_messages(config, user_input, variables, history) + tool_definitions = to_tool_definitions(tools) + + root_system, root_messages = split_input_messages(input_messages) + set_input_content_attributes( + span, + capture_content, + system_instructions=root_system, + messages=root_messages, ) - total_input = 0 - total_output = 0 - full_output = "" - previous_response_id: str | None = None - current_input: Any = input_messages - steps = 0 + full_output = "" + previous_response_id: str | None = None + current_input: Any = input_messages + steps = 0 - try: while True: + model_span = start_model_span(config, parent) + open_model_span = model_span + if capture_content: + turn_system, turn_messages = split_input_messages(current_input) + set_input_content_attributes( + model_span, + capture_content, + system_instructions=turn_system, + messages=turn_messages, + tool_definitions=tool_definitions, + ) + stream_params: dict[str, Any] = { "model": config["model"]["name"], "input": current_input, } if previous_response_id: stream_params["previous_response_id"] = previous_response_id + # Tools are forwarded on every streaming turn, not only the first, unlike the blocking + # path and 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. if tools: stream_params["tools"] = tools - stream = client.responses.stream(**stream_params) - async with stream as s: - async for event in s: - if getattr(event, "type", None) == "response.output_text.delta": - text = getattr(event, "delta", "") - full_output += text - yield {"type": "chunk", "text": text} - - final_resp = await s.get_final_response() - - total_input += ( - getattr(getattr(final_resp, "usage", None), "input_tokens", 0) or 0 - ) - total_output += ( - getattr(getattr(final_resp, "usage", None), "output_tokens", 0) or 0 + try: + stream = client.responses.stream(**stream_params) + async with stream as s: + async for event in s: + if getattr(event, "type", None) == "response.output_text.delta": + text = getattr(event, "delta", "") + full_output += text + yield {"type": "chunk", "text": text} + + final_resp = await s.get_final_response() + except Exception as exc: + fail_span(model_span, exc, ended) + open_model_span = None + raise + + last_response_model = getattr(final_resp, "model", None) or model_name( + config ) + set_response_output_content(model_span, capture_content, final_resp) + finish_reason = finish_reason_of(final_resp) + usage = to_span_usage(getattr(final_resp, "usage", None)) + finish_model_span(model_span, last_response_model, usage, finish_reason) + open_model_span = None + run_usage.add(usage) tool_calls = [ item - for item in (getattr(final_resp, "output", []) or []) + for item in (getattr(final_resp, "output", None) or []) if getattr(item, "type", None) == "function_call" ] if not tool_calls: @@ -323,15 +450,30 @@ async def _stream_gen( previous_response_id = getattr(final_resp, "id", None) tool_outputs = [] for tc in tool_calls: - args = json.loads(tc.arguments) - handler_fn = tool_handlers.get(tc.name) - if not handler_fn: - raise ValueError(f'No handler registered for tool "{tc.name}"') - result = ( - await handler_fn(args) - if _is_coroutine(handler_fn) - else handler_fn(args) + tool_span = start_tool_span(tc.name, tc.call_id, parent) + set_tool_call_content_attributes( + tool_span, capture_content, arguments=tc.arguments ) + try: + args = json.loads(tc.arguments) + 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 = ( + await handler_fn(args) + if _is_coroutine(handler_fn) + else handler_fn(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 + ) + succeed_span(tool_span) + except Exception as exc: + fail_span(tool_span, exc, ended) + raise tool_outputs.append( { "type": "function_call_output", @@ -341,40 +483,40 @@ async def _stream_gen( ) current_input = tool_outputs - if span: - 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, _final_output_messages(full_output) + ) + finish_root_span(span, last_response_model, run_usage.total) + mark_ok(span) + end_span_once(span, ended) yield { "type": "done", "output": full_output, - "usage": {"input_tokens": total_input, "output_tokens": total_output}, + "usage": { + "input_tokens": run_usage.total.input, + "output_tokens": run_usage.total.output, + }, } except Exception as exc: - if span: - span.record_exception(exc) - span.set_status(SpanStatusCode.ERROR, str(exc)) - span.end() + if open_model_span is not None: + fail_span(open_model_span, exc, ended) + if run_usage.reported: + finish_root_span(span, last_response_model, 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, last_response_model, run_usage.total) + end_span_once(span, ended, abandoned=True) def openai_messages( @@ -384,7 +526,13 @@ def openai_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_openai_messages_handler(), **kwargs + key=config_key, + handler=create_openai_messages_handler(capture_content=capture_content), + **kwargs, ).invoke(user_input, context, variables=variables) diff --git a/packages/openai-messages/src/launchdarkly_ai_openai_messages/spans.py b/packages/openai-messages/src/launchdarkly_ai_openai_messages/spans.py new file mode 100644 index 0000000..95b73d6 --- /dev/null +++ b/packages/openai-messages/src/launchdarkly_ai_openai_messages/spans.py @@ -0,0 +1,357 @@ +"""Span construction for the OpenAI 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 + +import json +from typing import Any + +from launchdarkly_ai_server import ( + AiConfigRep, + SpanMessage, + SpanMessagePart, + SpanUsage, + ToolDefinitionInput, + number_or_zero, + set_ld_span_attributes, + set_model_identity_attributes, + set_output_content_attributes, + set_usage_span_attributes, +) + +try: + from opentelemetry import trace + from opentelemetry.trace import StatusCode as SpanStatusCode + + _HAS_OTEL = True +except ImportError: # pragma: no cover - exercised by the no-OTel install path + _HAS_OTEL = False + +TRACER_NAME = "@launchdarkly/ai-openai-messages" + +#: OpenAI serves every model behind this handler, so the provider name is a constant. +PROVIDER = "openai" + + +def model_name(config: AiConfigRep) -> str: + return str(config.get("model", {}).get("name", "")) + + +def _attr(obj: Any, name: str) -> Any: + """Reads a field off a provider object or a plain dict, whichever the caller holds. + + The Responses API hands back objects; input items built by the handler and by the tool loop are + plain dicts, and both shapes reach these converters. + """ + if isinstance(obj, dict): + return obj.get(name) + return getattr(obj, name, None) + + +# ─── 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 + 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. + + Named after the *requested* model, like every other handler's ``chat`` span. The value written + to ``gen_ai.response.model`` when the turn finishes is the model that actually answered; see + :func:`finish_model_span` and TELEMETRY-CONTRACT.md section 2a. + """ + 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, 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", call_id) + return span + + +# ─── Span finishes ─────────────────────────────────────────────────────────── + + +def finish_root_span(span: Any, response_model: str, run_usage: SpanUsage) -> None: + """Writes the run-level identity and token totals onto the root. + + ``response_model`` is the model that answered, not the one requested: OpenAI resolves an alias + such as ``gpt-4o`` to a dated snapshot, and the ``chat`` children already report the real value. + A root copying ``config.model.name`` would contradict its own children. The caller supplies the + fallback to the requested name; see TELEMETRY-CONTRACT.md section 2a. + """ + if span is None: + return + span.set_attribute("gen_ai.response.model", response_model) + set_usage_span_attributes(span, run_usage) + + +def finish_model_span( + span: Any, + response_model: str, + usage: SpanUsage, + finish_reason: str | None = None, +) -> None: + """Ends one ``chat`` span successfully. *finish_reason* arrives already derived.""" + if span is None: + return + span.set_attribute("gen_ai.response.model", response_model) + if finish_reason: + span.set_attribute("gen_ai.response.finish_reasons", [finish_reason]) + 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. + """ + 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() + + +# ─── Usage ─────────────────────────────────────────────────────────────────── + + +def to_span_usage(usage: Any) -> SpanUsage: + """This turn's usage as a ``SpanUsage``, with OpenAI's cache rule applied. + + OpenAI reports cached tokens *within* the input total (a subset), so unlike Anthropic they are + not added on top: ``cached_tokens`` is surfaced as ``cache_read`` for cross-handler parity only. + OpenAI has no cache-creation concept, so that count is always 0. + """ + details = _attr(usage, "input_tokens_details") + return SpanUsage( + input=number_or_zero(_attr(usage, "input_tokens")), + output=number_or_zero(_attr(usage, "output_tokens")), + cache_read=number_or_zero(_attr(details, "cached_tokens")), + cache_creation=0, + ) + + +# ─── Finish reasons ────────────────────────────────────────────────────────── + + +def finish_reason_of(response: Any) -> str | None: + """Maps a Responses result onto semconv's ``finish_reasons`` vocabulary. + + The Responses API has no per-message finish reason of its own: it reports a run ``status`` plus, + on an incomplete run, a machine-readable cause. The two OpenAI handlers do not use the shared + mapping table at all; see TELEMETRY-CONTRACT.md section 5a. + + The function-call check comes first on purpose. A live seven-turn capture put status + ``completed`` on every turn, including the six that stopped to call a tool, so status alone made + the attribute worthless. + """ + output = _attr(response, "output") or [] + if any(_attr(item, "type") == "function_call" for item in output): + return "tool_calls" + status = _attr(response, "status") + if status == "incomplete": + details = _attr(response, "incomplete_details") + reason = _attr(details, "reason") if details is not None else None + return "length" if reason == "max_output_tokens" else "content_filter" + if status == "completed": + return "stop" + return None + + +# ─── Provider shapes as span shapes ────────────────────────────────────────── + + +def split_input_messages(items: list[Any]) -> tuple[str | None, list[SpanMessage]]: + """Splits the Responses input list into system instructions and conversation turns. + + The system message is lifted out so it lands on ``gen_ai.system_instructions`` rather than being + buried mid-conversation; ``set_input_content_attributes`` puts it back as message 0 of the flat + carrier, which has no separate slot for it. + """ + system: list[str] = [] + messages: list[SpanMessage] = [] + + for raw in items: + role = _attr(raw, "role") + if role in ("system", "developer"): + system.append(str(_attr(raw, "content") or "")) + continue + + item_type = _attr(raw, "type") + if item_type == "function_call_output": + call_id = _attr(raw, "call_id") + messages.append( + SpanMessage( + role="tool", + parts=[ + SpanMessagePart( + type="tool_call_response", + id=call_id if isinstance(call_id, str) else None, + result=_attr(raw, "output"), + ) + ], + ) + ) + continue + if item_type == "function_call": + call_id = _attr(raw, "call_id") + messages.append( + SpanMessage( + role="assistant", + parts=[ + SpanMessagePart( + type="tool_call", + id=call_id if isinstance(call_id, str) else None, + name=str(_attr(raw, "name") or ""), + arguments=_attr(raw, "arguments"), + ) + ], + ) + ) + continue + + content = _attr(raw, "content") + text = content if isinstance(content, str) else json.dumps(content) + messages.append( + SpanMessage( + role=role if isinstance(role, str) else "user", + parts=[SpanMessagePart(type="text", content=text)], + ) + ) + + return ("\n".join(system) if system else None, messages) + + +def output_item_parts(item: Any) -> list[SpanMessagePart]: + """Converts one Responses output item into canonical span message parts.""" + item_type = _attr(item, "type") + if item_type == "function_call": + call_id = _attr(item, "call_id") + return [ + SpanMessagePart( + type="tool_call", + id=call_id if isinstance(call_id, str) else None, + name=str(_attr(item, "name") or ""), + arguments=_attr(item, "arguments"), + ) + ] + if item_type == "reasoning": + summary = _attr(item, "summary") + text = ( + "\n".join(str(_attr(entry, "text") or "") for entry in summary) + if isinstance(summary, list) + else "" + ) + return [SpanMessagePart(type="reasoning", content=text)] if text else [] + + content = _attr(item, "content") + if not isinstance(content, list): + return [] + return [ + SpanMessagePart(type="text", content=str(_attr(block, "text") or "")) + for block in content + if _attr(block, "type") == "output_text" + ] + + +def set_response_output_content(span: Any, capture: bool, response: Any) -> None: + """Records what the model produced on this turn, gated on *capture*.""" + if not capture: + return + finish_reason = finish_reason_of(response) + output = _attr(response, "output") or [] + messages = [ + SpanMessage( + role=str(_attr(item, "role") or "assistant"), + parts=output_item_parts(item), + finish_reason=finish_reason, + ) + for item in output + ] + set_output_content_attributes(span, capture, 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("parameters"), + ) + for t in tools + ] diff --git a/packages/openai-messages/tests/test_handler.py b/packages/openai-messages/tests/test_handler.py index 3b76f48..65b39f7 100644 --- a/packages/openai-messages/tests/test_handler.py +++ b/packages/openai-messages/tests/test_handler.py @@ -1,7 +1,7 @@ """ Tests for launchdarkly-ai-openai-messages handler. -Covers §1.1–1.9. -Reference: TESTING.md §1 +Covers §1.1-1.9. +Reference: TESTING.md §1, TELEMETRY-CONTRACT.md """ from __future__ import annotations @@ -21,32 +21,68 @@ } +def _message_item(text: str) -> MagicMock: + item = MagicMock() + item.type = "message" + item.role = "assistant" + block = MagicMock() + block.type = "output_text" + block.text = text + item.content = [block] + return item + + +def _function_call_item( + name: str, call_id: str = "call_1", args: dict | None = None +) -> MagicMock: + item = MagicMock() + item.type = "function_call" + item.name = name + item.call_id = call_id + item.arguments = json.dumps(args or {}) + return item + + def _make_response( output_text: str = "Hello", tool_calls: list[Any] | None = None, input_tokens: int = 10, output_tokens: int = 5, resp_id: str = "resp-1", + model: str = "gpt-4o", + status: str = "completed", + include_output_message: bool = True, + cache_read: int | None = None, ) -> MagicMock: r = MagicMock() r.id = resp_id - r.model = "gpt-4o" + r.model = model r.output_text = output_text + r.status = status + r.incomplete_details = None r.usage = MagicMock() r.usage.input_tokens = input_tokens r.usage.output_tokens = output_tokens + if cache_read is not None: + r.usage.input_tokens_details = MagicMock(cached_tokens=cache_read) + else: + r.usage.input_tokens_details = MagicMock(cached_tokens=0) items: list[MagicMock] = [] for tc in tool_calls or []: - item = MagicMock() - item.type = "function_call" - item.name = tc["name"] - item.call_id = tc["call_id"] - item.arguments = json.dumps(tc.get("args", {})) - items.append(item) + items.append(_function_call_item(tc["name"], tc["call_id"], tc.get("args", {}))) + if not tool_calls and include_output_message and output_text: + items.append(_message_item(output_text)) r.output = items return r +CONFIG = { + "model": {"name": "gpt-4o"}, + "provider": {"name": "OpenAI"}, + "instructions": "Be helpful.", +} + + @pytest.fixture def mock_openai(mocker): mock_client = MagicMock() @@ -56,14 +92,6 @@ def mock_openai(mocker): return mock_client -def _make_tracer_patch(mock_span: MagicMock) -> tuple[MagicMock, MagicMock]: - mock_tracer = MagicMock() - mock_tracer.start_span = MagicMock(return_value=mock_span) - mock_trace_mod = MagicMock() - mock_trace_mod.get_tracer = MagicMock(return_value=mock_tracer) - return mock_trace_mod, mock_tracer - - # --------------------------------------------------------------------------- # §1.1 Factory function and metadata # --------------------------------------------------------------------------- @@ -379,200 +407,537 @@ async def test_multiple_consecutive_tool_calls( # --------------------------------------------------------------------------- -# §1.5 Telemetry +# §1.5 Telemetry — span recording # --------------------------------------------------------------------------- -class TestTelemetry: - async def test_span_name(self, mock_openai: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, mock_tracer = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod +class RecordedSpan: + """A span that remembers what a handler did to it, so a test can assert on the whole thing.""" - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler + 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 - h = create_openai_messages_handler() - await h(CONFIG, "q", {}, {}) - mock_tracer.start_span.assert_called_with("openai.response") + def set_attribute(self, key: str, value: Any) -> None: + self.attributes[key] = value - async def test_gen_ai_system(self, mock_openai: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod + def add_event(self, name: str, attributes: dict[str, Any] | None = None) -> None: + self.events.append((name, attributes or {})) - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler + def set_status(self, code: Any, description: str | None = None) -> None: + self.statuses.append(code) - h = create_openai_messages_handler() - 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") == "openai" + def record_exception(self, exc: BaseException) -> None: + self.exceptions.append(exc) - async def test_gen_ai_request_model(self, mock_openai: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod + def end(self) -> None: + self.ended += 1 - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler - h = create_openai_messages_handler() - 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") == "gpt-4o" +class SpanRecorder: + """Stands in for the ``trace`` module inside ``spans.py`` and records every span opened.""" - async def test_token_attributes_set(self, mock_openai: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod + def __init__(self) -> None: + self.spans: list[RecordedSpan] = [] - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler + def get_tracer(self, name: str) -> SpanRecorder: + return self - h = create_openai_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 + def start_span(self, name: str, context: Any = None) -> RecordedSpan: + span = RecordedSpan(name, context) + self.spans.append(span) + return span - async def test_span_status_ok(self, mock_openai: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod + def set_span_in_context(self, span: RecordedSpan) -> Any: + return ("context-of", span) - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler + @property + def root(self) -> RecordedSpan: + return self.spans[0] - h = create_openai_messages_handler() - await h(CONFIG, "q", {}, {}) - from opentelemetry.trace import StatusCode + def named(self, prefix: str) -> list[RecordedSpan]: + return [s for s in self.spans if s.name.startswith(prefix)] - mock_span.set_status.assert_called_with(StatusCode.OK) + @property + def names(self) -> list[str]: + return [s.name for s in self.spans] - async def test_span_end_always_called(self, mock_openai: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler +def _recording() -> Any: + """Patches the tracer that ``spans.py`` holds, and yields the recorder.""" + import launchdarkly_ai_openai_messages.spans as spans_mod - h = create_openai_messages_handler() - await h(CONFIG, "q", {}, {}) - mock_span.end.assert_called_once() + recorder = SpanRecorder() + return patch.object(spans_mod, "trace", recorder), recorder - async def test_gen_ai_operation_name(self, mock_openai: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler +def _make_tracer_patch(mock_span: MagicMock) -> Any: + """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() + mock_trace_mod.get_tracer = MagicMock(return_value=mock_tracer) + return mock_trace_mod, mock_tracer - h = create_openai_messages_handler() - 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_gen_ai_content_prompt_event(self, mock_openai: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod +class TestSpanTree: + """TELEMETRY-CONTRACT.md section 1.""" - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler + async def test_opens_a_root_span_named_invoke_agent( + self, mock_openai: MagicMock + ) -> None: + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler - h = create_openai_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.prompt" in event_names + with ctx: + await create_openai_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_content_completion_event( + async def test_emits_one_chat_child_per_model_turn( self, mock_openai: MagicMock ) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler + with ctx: + await create_openai_messages_handler()(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" + # Parented to the root, not to nothing. + assert chats[0].context == ("context-of", rec.root) - h = create_openai_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 + async def test_names_the_chat_span_after_the_requested_model( + self, mock_openai: MagicMock + ) -> None: + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler - async def test_total_tokens_attribute(self, mock_openai: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod + cfg = {**CONFIG, "model": {"name": "gpt-4o-mini"}} + with ctx: + await create_openai_messages_handler()(cfg, "q", {}, {}) + assert "chat gpt-4o-mini" in rec.names - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler + async def test_emits_a_chat_span_per_turn_of_a_tool_loop( + self, mock_openai: MagicMock + ) -> None: + mock_openai.responses.create = AsyncMock( + side_effect=[ + _make_response(tool_calls=[{"name": "myTool", "call_id": "tu1"}]), + _make_response(output_text="done"), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler - h = create_openai_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 + with ctx: + await create_openai_messages_handler()( + CONFIG, "q", {"myTool": lambda _: "result"}, {} + ) + assert len(rec.named("chat ")) == 2 - async def test_gen_ai_response_model(self, mock_openai: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod + async def test_emits_an_execute_tool_span_per_tool_call( + self, mock_openai: MagicMock + ) -> None: + mock_openai.responses.create = AsyncMock( + side_effect=[ + _make_response(tool_calls=[{"name": "myTool", "call_id": "tu1"}]), + _make_response(output_text="done"), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler + with ctx: + await create_openai_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_openai_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_tool_spans_are_siblings_of_chat_not_children( + self, mock_openai: MagicMock + ) -> None: + mock_openai.responses.create = AsyncMock( + side_effect=[ + _make_response(tool_calls=[{"name": "myTool", "call_id": "tu1"}]), + _make_response(output_text="done"), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler - async def test_ld_span_attributes(self, mock_openai: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod + with ctx: + await create_openai_messages_handler()( + CONFIG, "q", {"myTool": lambda _: "r"}, {} + ) + assert rec.named("execute_tool ")[0].context == ("context-of", rec.root) + + async def test_every_span_is_ended(self, mock_openai: MagicMock) -> None: + mock_openai.responses.create = AsyncMock( + side_effect=[ + _make_response(tool_calls=[{"name": "myTool", "call_id": "tu1"}]), + _make_response(output_text="done"), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_messages_handler()( + CONFIG, "q", {"myTool": lambda _: "r"}, {} + ) + assert [s.ended for s in rec.spans] == [1] * len(rec.spans) + + +class TestRootSpanAttributes: + """TELEMETRY-CONTRACT.md sections 2 and 2a.""" + + async def test_writes_both_provider_keys_and_the_requested_model( + self, mock_openai: MagicMock + ) -> None: + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_messages_handler()(CONFIG, "q", {}, {}) + attrs = rec.root.attributes + assert attrs["gen_ai.system"] == "openai" + assert attrs["gen_ai.provider.name"] == "openai" + assert attrs["gen_ai.request.model"] == "gpt-4o" + + async def test_response_model_is_the_model_that_answered( + self, mock_openai: MagicMock + ) -> None: + # OpenAI resolves an alias like `gpt-4o` to a dated snapshot. openai-messages is the only + # handler whose root reports the answering model rather than the requested one. Section 2a. + mock_openai.responses.create = AsyncMock( + return_value=_make_response(model="gpt-4o-2024-08-06") + ) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_messages_handler()(CONFIG, "q", {}, {}) + assert rec.root.attributes["gen_ai.response.model"] == "gpt-4o-2024-08-06" + + async def test_carries_the_launchdarkly_attributes_and_feature_flag_event( + self, mock_openai: MagicMock + ) -> None: + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler variables = { "__ld": { - "configKey": "my-config", - "variationKey": "v1", - "runId": "run-abc", + "configKey": "k", + "variationKey": "v", + "runId": "r", + "graphKey": "g", } } - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler + with ctx: + await create_openai_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_openai: MagicMock + ) -> None: + mock_openai.responses.create = AsyncMock( + side_effect=[ + _make_response(tool_calls=[{"name": "myTool", "call_id": "tu1"}]), + _make_response(output_text="done"), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler - h = create_openai_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 - - async def test_ld_graph_key_set_when_present(self, mock_openai: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod + variables = {"__ld": {"configKey": "k", "variationKey": "v", "runId": "r"}} + with ctx: + await create_openai_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] - variables = { - "__ld": { - "configKey": "my-config", - "variationKey": "v1", - "runId": "run-abc", - "graphKey": "my-graph", - } + async def test_carries_the_run_total_not_one_turn( + self, mock_openai: MagicMock + ) -> None: + mock_openai.responses.create = AsyncMock( + side_effect=[ + _make_response( + tool_calls=[{"name": "myTool", "call_id": "tu1"}], + input_tokens=10, + output_tokens=1, + ), + _make_response(output_text="done", input_tokens=20, output_tokens=2), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_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, 5a and 8.""" + + async def test_writes_all_seven_usage_attributes_including_zeros( + self, mock_openai: MagicMock + ) -> None: + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_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_reports_cached_tokens_as_cache_read_without_folding_into_input( + self, mock_openai: MagicMock + ) -> None: + # OpenAI already counts cached tokens inside input_tokens: this is the assertion that + # catches a fold in the Anthropic direction, which would double-count. + mock_openai.responses.create = AsyncMock( + return_value=_make_response(input_tokens=50, output_tokens=5, cache_read=30) + ) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_messages_handler()(CONFIG, "q", {}, {}) + attrs = rec.named("chat ")[0].attributes + assert attrs["gen_ai.usage.input_tokens"] == 50 + assert attrs["gen_ai.usage.total_tokens"] == 55 + assert attrs["gen_ai.usage.cache_read.input_tokens"] == 30 + # OpenAI has no cache-creation concept; still emitted, as 0, so the set is always complete. + assert attrs["gen_ai.usage.cache_creation.input_tokens"] == 0 + + async def test_derives_tool_calls_before_checking_status( + self, mock_openai: MagicMock + ) -> None: + # A live capture put status `completed` on every turn including the ones that stopped to + # call a tool, so the function-call check must run first. Section 5a. + mock_openai.responses.create = AsyncMock( + side_effect=[ + _make_response( + tool_calls=[{"name": "myTool", "call_id": "tu1"}], + status="completed", + ), + _make_response(output_text="done", status="completed"), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_messages_handler()( + CONFIG, "q", {"myTool": lambda _: "r"}, {} + ) + first = rec.named("chat ")[0] + assert first.attributes["gen_ai.response.finish_reasons"] == ["tool_calls"] + + async def test_derives_stop_from_completed_status( + self, mock_openai: MagicMock + ) -> None: + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_messages_handler()(CONFIG, "q", {}, {}) + assert rec.named("chat ")[0].attributes["gen_ai.response.finish_reasons"] == [ + "stop" + ] + + async def test_derives_length_from_incomplete_max_output_tokens( + self, mock_openai: MagicMock + ) -> None: + resp = _make_response(status="incomplete", include_output_message=False) + resp.incomplete_details = MagicMock(reason="max_output_tokens") + mock_openai.responses.create = AsyncMock(return_value=resp) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_messages_handler()(CONFIG, "q", {}, {}) + assert rec.named("chat ")[0].attributes["gen_ai.response.finish_reasons"] == [ + "length" + ] + + async def test_derives_content_filter_from_incomplete_other_reason( + self, mock_openai: MagicMock + ) -> None: + resp = _make_response(status="incomplete", include_output_message=False) + resp.incomplete_details = MagicMock(reason="content_filter") + mock_openai.responses.create = AsyncMock(return_value=resp) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_messages_handler()(CONFIG, "q", {}, {}) + assert rec.named("chat ")[0].attributes["gen_ai.response.finish_reasons"] == [ + "content_filter" + ] + + async def test_writes_no_finish_reason_for_an_unrecognised_status( + self, mock_openai: MagicMock + ) -> None: + # No passthrough for the two OpenAI handlers: an unrecognised status drops the attribute. + resp = _make_response(status="cancelled", include_output_message=False) + mock_openai.responses.create = AsyncMock(return_value=resp) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_messages_handler()(CONFIG, "q", {}, {}) + assert "gen_ai.response.finish_reasons" not in rec.named("chat ")[0].attributes + + async def test_sets_status_ok_on_a_successful_turn( + self, mock_openai: MagicMock + ) -> None: + from opentelemetry.trace import StatusCode + + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_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_openai: MagicMock + ) -> None: + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_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_openai: MagicMock + ) -> None: + mock_openai.responses.create = AsyncMock( + return_value=_make_response(output_text="Hello World") + ) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_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_openai: MagicMock + ) -> None: + cfg = { + **CONFIG, + "tools": {"myTool": {"description": "d", "parameters": {"type": "object"}}}, } - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler - h = create_openai_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_openai_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_openai: MagicMock + ) -> None: + mock_openai.responses.create = AsyncMock( + side_effect=[ + _make_response( + tool_calls=[ + {"name": "myTool", "call_id": "tu1", "args": {"city": "NYC"}} + ] + ), + _make_response(output_text="done"), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_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_openai: MagicMock + ) -> None: + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + await create_openai_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 # --------------------------------------------------------------------------- @@ -581,50 +946,96 @@ async def test_ld_graph_key_set_when_present(self, mock_openai: MagicMock) -> No class TestErrorHandling: - async def test_records_exception_on_span(self, mock_openai: MagicMock) -> None: - mock_openai.responses.create = AsyncMock(side_effect=RuntimeError("api err")) - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod + """TELEMETRY-CONTRACT.md section 6.""" - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler + async def test_fails_the_chat_span_when_the_provider_call_raises( + self, mock_openai: MagicMock + ) -> None: + from opentelemetry.trace import StatusCode - h = create_openai_messages_handler() - with pytest.raises(RuntimeError): - await h(CONFIG, "q", {}, {}) - mock_span.record_exception.assert_called_once() - - async def test_ends_span_on_error(self, mock_openai: MagicMock) -> None: - mock_openai.responses.create = AsyncMock(side_effect=RuntimeError("api err")) - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod + mock_openai.responses.create = AsyncMock(side_effect=RuntimeError("api error")) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler + with ctx, pytest.raises(RuntimeError): + await create_openai_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_openai_messages_handler() - with pytest.raises(RuntimeError): - await h(CONFIG, "q", {}, {}) - mock_span.end.assert_called_once() - - async def test_sets_span_status_error(self, mock_openai: MagicMock) -> None: - mock_openai.responses.create = AsyncMock(side_effect=RuntimeError("api err")) - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod + async def test_fails_the_root_span_too(self, mock_openai: MagicMock) -> None: + from opentelemetry.trace import StatusCode - with patch.object(handler_mod, "trace", mock_trace_mod): - from launchdarkly_ai_openai_messages import create_openai_messages_handler + mock_openai.responses.create = AsyncMock(side_effect=RuntimeError("api error")) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler - h = create_openai_messages_handler() - with pytest.raises(RuntimeError): - await h(CONFIG, "q", {}, {}) + with ctx, pytest.raises(RuntimeError): + await create_openai_messages_handler()(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, mock_openai: MagicMock + ) -> None: from opentelemetry.trace import StatusCode - status_codes = [c[0][0] for c in mock_span.set_status.call_args_list] - assert StatusCode.ERROR in status_codes + mock_openai.responses.create = AsyncMock( + return_value=_make_response( + tool_calls=[{"name": "myTool", "call_id": "tu1"}] + ) + ) + + def _boom(_: Any) -> Any: + raise RuntimeError("tool exploded") + + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx, pytest.raises(RuntimeError, match="tool exploded"): + await create_openai_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_reports_the_spend_of_completed_turns_on_a_failed_run( + self, mock_openai: MagicMock + ) -> None: + mock_openai.responses.create = AsyncMock( + side_effect=[ + _make_response( + tool_calls=[{"name": "myTool", "call_id": "tu1"}], + input_tokens=40, + output_tokens=7, + ), + RuntimeError("second turn died"), + ] + ) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx, pytest.raises(RuntimeError): + await create_openai_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 + + async def test_writes_no_usage_when_no_turn_ever_reported_any( + self, mock_openai: MagicMock + ) -> None: + mock_openai.responses.create = AsyncMock( + side_effect=RuntimeError("died on the first call") + ) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx, pytest.raises(RuntimeError): + await create_openai_messages_handler()(CONFIG, "q", {}, {}) + assert "gen_ai.usage.input_tokens" not in rec.root.attributes async def test_rethrows_error(self, mock_openai: MagicMock) -> None: mock_openai.responses.create = AsyncMock(side_effect=RuntimeError("rethrown")) @@ -636,7 +1047,7 @@ async def test_rethrows_error(self, mock_openai: MagicMock) -> None: # --------------------------------------------------------------------------- -# §1.9 Structured output (outputFormat) — first-class json_schema +# §1.9 Structured output (outputFormat) # --------------------------------------------------------------------------- @@ -661,16 +1072,6 @@ async def test_output_format_uses_text_format_json_schema( assert "text" in kwargs assert kwargs["text"]["format"]["type"] == "json_schema" - async def test_absent_output_format_text_format_not_sent( - self, mock_openai: MagicMock - ) -> None: - from launchdarkly_ai_openai_messages import create_openai_messages_handler - - h = create_openai_messages_handler() - await h(CONFIG, "q", {}, {}) - kwargs = mock_openai.responses.create.call_args.kwargs - assert "text" not in kwargs - # --------------------------------------------------------------------------- # §1.7 Convenience export @@ -726,7 +1127,10 @@ def test_callable_without_extra_kwargs(self, mock_openai: MagicMock) -> None: def _make_openai_stream_context( - chunks: list[str], input_tok: int = 5, output_tok: int = 3 + chunks: list[str], + input_tok: int = 5, + output_tok: int = 3, + output: list[Any] | None = None, ) -> Any: """Returns a mock OpenAI stream context manager.""" events = [] @@ -737,9 +1141,13 @@ def _make_openai_stream_context( events.append(e) final_resp = MagicMock() - final_resp.output = [] + final_resp.output = output if output is not None else [] + final_resp.status = "completed" + final_resp.incomplete_details = None final_resp.usage = MagicMock(input_tokens=input_tok, output_tokens=output_tok) + final_resp.usage.input_tokens_details = MagicMock(cached_tokens=0) final_resp.id = "resp-stream" + final_resp.model = "gpt-4o" class _FakeStream: def __aiter__(self) -> AsyncIterator[Any]: @@ -854,27 +1262,25 @@ async def _bad_ctx() -> AsyncGenerator[Any, None]: async def test_tools_forwarded_on_second_streaming_turn( self, mock_openai: MagicMock ) -> None: - """§1.8 — tools must appear in stream_params on every streaming turn. + """§1.8 - tools must appear in stream_params on every streaming turn. - When the first streaming turn returns a tool call and a second streaming - turn is required to send the tool result, the ``tools`` parameter must - be present in the second ``responses.stream()`` call too — not just the - first. Without this, the model loses tool access after the first turn. + This is a pre-existing Python-only behaviour that diverges from the TypeScript SDK (which + does not resend tools after the first turn). It changes what the model is offered, not what + the span reports, so this test only pins that the behaviour is unchanged by the span work. """ import launchdarkly_ai_openai_messages.handler as handler_mod from launchdarkly_ai_openai_messages import create_openai_messages_handler - # -- First streaming turn: one text chunk then a tool call ----------- - tool_call_item = MagicMock() - tool_call_item.type = "function_call" - tool_call_item.name = "my-tool" - tool_call_item.call_id = "call-1" - tool_call_item.arguments = '{"q": "x"}' + tool_call_item = _function_call_item("my-tool", "call-1", {"q": "x"}) first_final = MagicMock() first_final.output = [tool_call_item] + first_final.status = "completed" + first_final.incomplete_details = None first_final.usage = MagicMock(input_tokens=3, output_tokens=1) + first_final.usage.input_tokens_details = MagicMock(cached_tokens=0) first_final.id = "resp-first" + first_final.model = "gpt-4o" class _FirstStream: def __aiter__(self) -> AsyncIterator[Any]: @@ -889,12 +1295,15 @@ async def _iter(self) -> AsyncIterator[Any]: async def get_final_response(self) -> Any: return first_final - # -- Second streaming turn: final text response ----------------------- second_final = MagicMock() second_final.output = [] second_final.output_text = "done" + second_final.status = "completed" + second_final.incomplete_details = None second_final.usage = MagicMock(input_tokens=4, output_tokens=2) + second_final.usage.input_tokens_details = MagicMock(cached_tokens=0) second_final.id = "resp-second" + second_final.model = "gpt-4o" class _SecondStream: def __aiter__(self) -> AsyncIterator[Any]: @@ -957,8 +1366,6 @@ class TestNoneUserInput: async def test_none_user_input_instructions_path_no_none_content( self, mock_openai: MagicMock ) -> None: - """When instructions path is taken and user_input=None, no message in - the API call may have content=None.""" from launchdarkly_ai_openai_messages import create_openai_messages_handler captured: list[Any] = [] @@ -1071,11 +1478,13 @@ async def _iter() -> AsyncGenerator: # --------------------------------------------------------------------------- -# §1.5 Streaming telemetry (Appendix A.5 — do not patch _HAS_OTEL=False) +# §1.5 Streaming telemetry (do not patch _HAS_OTEL=False) # --------------------------------------------------------------------------- class TestStreamingTelemetry: + """TELEMETRY-CONTRACT.md sections 1 and 6. The streaming path emits the same tree.""" + def _patch_stream( self, mock_openai: MagicMock, @@ -1087,53 +1496,127 @@ def _patch_stream( return_value=_make_openai_stream_context(chunks, input_tok, output_tok) ) - async def test_span_started_during_stream(self, mock_openai: MagicMock) -> None: - mock_span = MagicMock() - mock_trace_mod, mock_tracer = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod + async def test_opens_the_same_root_span_name_as_the_blocking_path( + self, mock_openai: MagicMock + ) -> None: + self._patch_stream(mock_openai, ["hi"]) + ctx, rec = _recording() from launchdarkly_ai_openai_messages import create_openai_messages_handler - self._patch_stream(mock_openai, ["hi"]) - with patch.object(handler_mod, "trace", mock_trace_mod): - h = create_openai_messages_handler() - async for _ in await h.stream(CONFIG, "q"): + with ctx: + async for _ in await create_openai_messages_handler().stream(CONFIG, "q"): pass - mock_tracer.start_span.assert_called_with("openai.response.stream") + assert rec.root.name == "invoke_agent" + assert "chat gpt-4o" in rec.names - async def test_ld_span_attributes_set_during_stream( + async def test_carries_the_launchdarkly_attributes_on_the_root( self, mock_openai: MagicMock ) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod + self._patch_stream(mock_openai, ["hi"]) + ctx, rec = _recording() from launchdarkly_ai_openai_messages import create_openai_messages_handler variables = {"__ld": {"configKey": "k", "variationKey": "v", "runId": "r"}} + with ctx: + async for _ in await create_openai_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_openai: MagicMock + ) -> None: self._patch_stream(mock_openai, ["hi"]) - with patch.object(handler_mod, "trace", mock_trace_mod): - h = create_openai_messages_handler() - async for _ in await h.stream(CONFIG, "q", None, variables): + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + async for _ in await create_openai_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_openai: MagicMock ) -> None: - mock_span = MagicMock() - mock_trace_mod, _ = _make_tracer_patch(mock_span) - import launchdarkly_ai_openai_messages.handler as handler_mod + self._patch_stream(mock_openai, ["hi"], input_tok=11, output_tok=4) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + async for _ in await create_openai_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_openai: MagicMock + ) -> None: + self._patch_stream(mock_openai, ["one", "two", "three"]) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + gen = await create_openai_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_openai: MagicMock + ) -> None: + from opentelemetry.trace import StatusCode + + self._patch_stream(mock_openai, ["one", "two", "three"]) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + gen = await create_openai_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_openai: MagicMock + ) -> None: + from opentelemetry.trace import StatusCode + + mock_openai.responses.stream = MagicMock( + side_effect=RuntimeError("stream died") + ) + ctx, rec = _recording() from launchdarkly_ai_openai_messages import create_openai_messages_handler + with ctx, pytest.raises(RuntimeError, match="stream died"): + async for _ in await create_openai_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_openai: MagicMock + ) -> None: self._patch_stream(mock_openai, ["hi"]) - with patch.object(handler_mod, "trace", mock_trace_mod): - h = create_openai_messages_handler() - async for _ in await h.stream(CONFIG, "q"): + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx: + async for _ in await create_openai_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")] == [] # --------------------------------------------------------------------------- @@ -1216,3 +1699,74 @@ async def test_system_role_in_history_filtered_out( if m.get("content") in ("Hello", "You are evil", "Hi there") ] assert "system" not in history_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_openai_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_openai_messages_handler", _factory), + patch.object(handler_mod, "config", fake_config), + ): + handler_mod.openai_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 TestChatSpanNeverLeaks: + """A raise while recording conversation content must not leave the chat span open. + + The content writes on both sides of the provider call used to sit outside the try that fails the + span. A raise there failed only the root, and the chat span was never ended, so the exporter + never saw the turn. + """ + + async def test_an_unserialisable_output_still_ends_the_chat_span( + self, mock_openai: MagicMock + ) -> None: + from opentelemetry.trace import StatusCode + + class _Exploding: + model = "gpt-4o" + usage = None + + @property + def output(self) -> Any: + raise TypeError("cannot serialise this response") + + mock_openai.responses.create = AsyncMock(return_value=_Exploding()) + ctx, rec = _recording() + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + with ctx, pytest.raises(TypeError): + await create_openai_messages_handler(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