diff --git a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py index 0c6d851..abd7ee8 100644 --- a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py +++ b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py @@ -1,6 +1,19 @@ """ OpenAI Agents handler — uses the openai-agents SDK (``agents`` package). Mirrors the TypeScript @launchdarkly/ai-openai-agents handler. + +Span construction lives in ``spans.py``. This module drives the ``agents`` SDK's own ``Runner`` +and reports each per-turn and per-tool-call boundary the SDK gives us, through ``RunHooks``. + +The TypeScript handler intercepts model calls by wrapping the Agents SDK's ``Model`` / +``ModelProvider`` interfaces (``SpanningModel`` / ``SpanningModelProvider``) and tool calls by +listening for the ``Agent``'s ``agent_tool_start`` / ``agent_tool_end`` events. The Python +``agents`` SDK exposes the same two boundaries more directly, as ``RunHooks`` callbacks +(``on_llm_start`` / ``on_llm_end`` for a model turn, ``on_tool_start`` / ``on_tool_end`` for a tool +call), and both fire identically on the blocking and the streaming path. Using them is simpler than +building a parallel ``Model`` wrapper and produces the same span tree, because ``Model.get_response`` +in this SDK discards the same information (see ``spans.derive_finish_reason``) no matter which +interface intercepts it. """ from __future__ import annotations @@ -13,27 +26,54 @@ AiConfigRep, LDContext, ProviderHandler, + RunUsage, + SpanUsage, 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, + text_message, ) -try: - from opentelemetry import trace - from opentelemetry.trace import StatusCode as SpanStatusCode +from .spans import ( + derive_finish_reason, + fail_span, + finish_model_span, + finish_root_span, + mark_ok, + parent_context_of, + start_model_span, + start_root_span, + start_tool_span, + succeed_span, + to_request_span_messages, + to_response_span_messages, + to_span_usage, + to_tool_definitions, +) - _HAS_OTEL = True -except ImportError: - _HAS_OTEL = False +try: + # The `agents` SDK validates that anything passed as `hooks=` is a `RunHooksBase` instance, so + # `_SpanningHooks` below has to actually subclass it rather than merely duck-type it. Imported + # once at module load, unlike the rest of this handler's dynamic `importlib.import_module` + # calls, because a class statement needs its base at class-definition time. + from agents.lifecycle import RunHooksBase as _RunHooksBase +except ImportError: # pragma: no cover - `agents` is a hard dependency of this package + _RunHooksBase = object # type: ignore[assignment,misc] def _build_agent_tools( config_tools: dict[str, Any], tool_handlers: dict[str, Any], ) -> list[Any]: + # Not filtered to the tools that have a registered handler, unlike the TypeScript SDK, which + # excludes an unregistered tool from the catalog entirely. This difference predates the span + # work and changes what the model is offered, not what a span reports, so it is left as it is; + # `to_tool_definitions` below records the catalog actually sent, whatever it is. import importlib agents_mod = importlib.import_module("agents") @@ -113,6 +153,13 @@ def _build_agent_and_prompt( tools = _build_agent_tools(config.get("tools") or {}, tool_handlers) + # `outputFormat` is not wired into `Agent(output_type=...)` here, matching this handler's + # pre-existing behaviour (and unlike the TypeScript handler, which does wire it): the Python + # Agents SDK requires a concrete Python type for `output_type`, not a raw JSON Schema dict (see + # `utils.build_output_type`'s docstring). Fixing that gap is a "what is sent to the model" + # change, not a telemetry one, so it is left alone. `_call_impl` still returns the parsed + # `final_output` object as-is when `outputFormat` is configured, matching the pre-existing + # return-shape contract. agent = Agent( name="assistant", model=config.get("model", {}).get("name", "gpt-4o"), @@ -122,28 +169,184 @@ def _build_agent_and_prompt( return agent, prompt, instructions -def _sum_usage(raw_responses: list[Any]) -> tuple[int, int, int]: - input_tokens = sum( - getattr(r.usage, "input_tokens", 0) - for r in raw_responses - if hasattr(r, "usage") - ) - output_tokens = sum( - getattr(r.usage, "output_tokens", 0) - for r in raw_responses - if hasattr(r, "usage") - ) - total_tokens = sum( - getattr(r.usage, "total_tokens", 0) - for r in raw_responses - if hasattr(r, "usage") - ) - return input_tokens, output_tokens, total_tokens +class _SpanningHooks(_RunHooksBase): + """Opens and closes one ``chat`` span per model turn and one ``execute_tool`` span per tool + call, all parented to the root's context. + Subclasses ``agents.RunHooksBase``: the ``Runner`` validates that ``hooks=`` is an instance of + it before running, so a merely duck-typed object is rejected outright. -def create_openai_agent_handler() -> ProviderHandler: - """Creates a ``ProviderHandler`` for OpenAI via the openai-agents SDK.""" - tracer_name = "@launchdarkly/ai-openai-agents" + Owns the run's usage accumulator too, so a run that raises mid-flight still has the completed + turns' spend somewhere the caller can read it from, via :attr:`run_usage`. + """ + + def __init__( + self, + config: AiConfigRep, + parent: Any, + capture_content: bool, + run_usage: Any, + ) -> None: + self.config = config + self.parent = parent + self.capture_content = capture_content + self.run_usage = run_usage + self.open_model_span: Any = None + self.open_tool_spans: dict[str, Any] = {} + + async def on_llm_start( + self, context: Any, agent: Any, system_prompt: str | None, input_items: Any + ) -> None: + span = start_model_span(self.config, self.parent) + self.open_model_span = span + if self.capture_content: + set_input_content_attributes( + span, + self.capture_content, + system_instructions=system_prompt, + messages=to_request_span_messages(input_items), + tool_definitions=to_tool_definitions(getattr(agent, "tools", []) or []), + ) + + async def on_llm_end(self, context: Any, agent: Any, response: Any) -> None: + span = self.open_model_span + self.open_model_span = None + output = getattr(response, "output", None) or [] + finish_reason = derive_finish_reason(output) + usage = to_span_usage(getattr(response, "usage", None)) + # Accounting before the span, and not conditional on it. run_usage is what the caller gets + # back as its usage bag and what LaunchDarkly's own metrics bill from, and neither depends on + # a span existing. Returning early here dropped the tokens of every turn on an install + # without the otel extra, which is a billing figure rather than a telemetry one. + self.run_usage.add(usage) + if span is None: + return + if self.capture_content: + set_output_content_attributes( + span, + self.capture_content, + to_response_span_messages(output, finish_reason), + ) + finish_model_span(span, self.config, usage, finish_reason) + + async def on_tool_start(self, context: Any, agent: Any, tool: Any) -> None: + call_id = getattr(context, "tool_call_id", None) or getattr( + tool, "name", "tool" + ) + name = getattr(context, "tool_name", None) or getattr(tool, "name", "tool") + span = start_tool_span(str(name), str(call_id), self.parent) + if self.capture_content: + # `tool_arguments` is the raw JSON args string the model produced. Passed through as a + # string rather than parsed and re-serialised, per TELEMETRY-CONTRACT.md section 7: a + # string passes through unchanged. + args = getattr(context, "tool_arguments", None) + if args is not None: + set_tool_call_content_attributes( + span, self.capture_content, arguments=args + ) + self.open_tool_spans[str(call_id)] = span + + async def on_tool_end( + self, context: Any, agent: Any, tool: Any, result: Any + ) -> None: + call_id = str(getattr(context, "tool_call_id", None)) + span = self.open_tool_spans.pop(call_id, None) + if span is None: + return + if self.capture_content: + set_tool_call_content_attributes(span, self.capture_content, result=result) + succeed_span(span) + + def close_open_spans(self, error: BaseException) -> None: + """Fails every span this run has open, for the crash path. + + A tool handler's own exception, or the run raising for any other reason, propagates before + ``on_tool_end``/``on_llm_end`` ever fire, so those spans are still open when the caller's + ``except`` runs. Mirrors the TypeScript handler's ``attachToolSpanHooks(...).closeOpenSpans``. + """ + for span in self.open_tool_spans.values(): + fail_span(span, error) + self.open_tool_spans.clear() + if self.open_model_span is not None: + fail_span(self.open_model_span, error) + self.open_model_span = None + + def abandon_open_spans(self, ended: set[int]) -> None: + """Ends every span this run has open, for stream abandonment. + + Unlike :meth:`close_open_spans`, nothing failed: a consumer stopping early is normal. + ``end_span_once`` leaves each span at ``UNSET`` and marks ``launchdarkly.stream.abandoned``, + rather than recording an exception and setting ``ERROR``. + """ + for span in self.open_tool_spans.values(): + end_span_once(span, ended, abandoned=True) + self.open_tool_spans.clear() + if self.open_model_span is not None: + end_span_once(self.open_model_span, ended, abandoned=True) + self.open_model_span = None + + +def _write_failed_run_usage( + span: Any, + config: AiConfigRep, + error: BaseException, + run_usage: RunUsage, +) -> None: + """Writes what a failed run spent onto the root, without counting it twice. + + Two sources describe the same spend, and they overlap. The run hooks add each turn as it + finishes, so by the time the run raises they already hold every completed turn. The exception + also carries the SDK's own aggregate over those same turns. + + Adding the aggregate to the accumulator therefore roughly doubles the reported cost of any run + that failed after paid turns, which MaxTurnsExceeded does by definition. The aggregate is the + authoritative figure, so it replaces the accumulator rather than adding to it. + + When the error carries no aggregate, which is what a tool handler's own error looks like, the + accumulator is all there is and is used instead. Nothing is written when neither has anything: + all-zero attributes would assert the run cost nothing, which a run that died on its first call + cannot claim. + """ + spent = _usage_from_error(error) + # An aggregate that reports no tokens is not an aggregate. RunContextWrapper.usage defaults to an + # empty Usage, so the attribute is present from the moment the run starts, and an exception before + # the first paid call carries a full set of zeros. Writing those would assert the run cost + # nothing, which is the one claim this function's docstring says it must not make. + if spent is not None and (spent.input or spent.output): + finish_root_span(span, config, spent) + elif run_usage.reported: + finish_root_span(span, config, run_usage.total) + + +def _usage_from_error(error: BaseException) -> SpanUsage | None: + """The run's spend at the point it raised, when the SDK attached one. + + ``AgentsException`` carries ``run_data.context_wrapper.usage``, the same aggregate the success + path reads off ``result.state.usage``. The tokens a failed run already spent were really + billed, and the root is the only span a config-scoped cost query can find them on, so dropping + this would silently zero out a run that failed after several paid turns. + + Read structurally rather than with a provider import: a tool handler's own error propagates + unwrapped and carries no ``run_data``, and that case has to yield ``None`` so nothing is + written, rather than asserting the run spent nothing. + """ + run_data = getattr(error, "run_data", None) + if run_data is None: + return None + context_wrapper = getattr(run_data, "context_wrapper", None) + usage = getattr(context_wrapper, "usage", None) if context_wrapper else None + if usage is None: + return None + return to_span_usage(usage) + + +def create_openai_agent_handler(*, capture_content: bool = False) -> ProviderHandler: + """Creates a ``ProviderHandler`` for OpenAI via the openai-agents SDK. + + Set *capture_content* to put prompts, model output, tool arguments and tool results on the + emitted spans. It defaults to off. Conversation content is PII, so a run emits only metadata, + meaning models, token counts, timings and tool names, until a caller asks for more. + """ async def _call_impl( config: AiConfigRep, @@ -160,75 +363,49 @@ async def _call_impl( th = tool_handlers or {} vs = variables or {} - if _HAS_OTEL: - span = trace.get_tracer(tracer_name).start_span("openai.agent.run") - 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) agent, prompt, instructions = _build_agent_and_prompt( config, user_input, th, vs, history ) - - if span: - prompt_text = ( - f"system: {instructions}\n\n" if instructions else "" - ) + prompt - span.add_event("gen_ai.content.prompt", {"gen_ai.prompt": prompt_text}) - prompt_msgs: list[dict[str, str]] = [] - if instructions: - prompt_msgs.append({"role": "system", "content": instructions}) - prompt_msgs.append({"role": "user", "content": prompt}) - set_openllmetry_prompt(span, prompt_msgs) - + run_usage = create_run_usage() + hooks = _SpanningHooks(config, parent, capture_content, run_usage) try: - result = await Runner.run(agent, prompt) + # Inside the guard, because serialising the prompt raises on anything that is not + # JSON-serialisable and a raise out here would leave the root open: never ended, never + # exported, and the run gone from AI Config Monitoring with the feature_flag event on it. + set_input_content_attributes( + span, + capture_content, + system_instructions=instructions, + messages=[text_message("user", prompt)], + ) + result = await Runner.run(agent, prompt, hooks=hooks) final_output = result.final_output - input_tokens, output_tokens, total_tokens = _sum_usage(result.raw_responses) - - if span: - span.set_attribute( - "gen_ai.response.model", config.get("model", {}).get("name", "") - ) - span.set_attribute("gen_ai.usage.input_tokens", input_tokens) - span.set_attribute("gen_ai.usage.output_tokens", output_tokens) - span.set_attribute("gen_ai.usage.total_tokens", total_tokens) - span.add_event( - "gen_ai.content.completion", - { - "gen_ai.completion": final_output - if isinstance(final_output, str) - else json.dumps(final_output) - }, - ) - set_openllmetry_completion( - span, - final_output - if isinstance(final_output, str) - else json.dumps(final_output), - {"input_tokens": input_tokens, "output_tokens": output_tokens}, - ) - span.set_status(SpanStatusCode.OK) - span.end() - + set_output_content_attributes( + span, + capture_content, + [text_message("assistant", _stringify_output(final_output))], + ) + finish_root_span(span, config, run_usage.total) + succeed_span(span) output = ( - final_output if config.get("outputFormat") else str(final_output or "") + final_output + if config.get("outputFormat") + else _stringify_output(final_output) ) return { "output": output, - "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens}, + "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() + hooks.close_open_spans(exc) + _write_failed_run_usage(span, config, exc, run_usage) + fail_span(span, exc) raise def _stream_impl( @@ -239,51 +416,78 @@ def _stream_impl( history: list[dict[str, Any]] | None = None, ) -> AsyncGenerator[dict[str, Any], None]: return _stream_gen( - config, user_input, tool_handlers or {}, variables or {}, history + config, + user_input, + tool_handlers or {}, + variables or {}, + history, + capture_content=capture_content, ) return create_handler(("OpenAI", "agent"), _call_impl, _stream_impl) # type: ignore[arg-type] +def _stringify_output(value: Any) -> str: + if isinstance(value, str): + return value + if value is None: + return "" + try: + return json.dumps(value) + except (TypeError, ValueError): + return str(value) + + async def _stream_gen( config: AiConfigRep, user_input: str, tool_handlers: dict[str, Any], variables: dict[str, Any], history: list[dict[str, Any]] | None = None, + *, + capture_content: bool = False, ) -> AsyncGenerator[dict[str, Any], None]: + """Streams the run, emitting the same span tree as the blocking path. + + A consumer that breaks out of ``async for``, or raises inside the loop body, makes this + generator run its ``finally`` without ever entering ``except``: ``GeneratorExit`` inherits from + ``BaseException``, so ``except Exception`` does not see it. Without the cleanup in ``finally`` + the root span, and any still-open ``chat``/``execute_tool`` span, would never be ended. + + Ending the spans is not the whole story here: the ``Runner.run_streamed`` background task keeps + the agent run going, and spending tokens, after a consumer stops reading. ``finally`` also + cancels the streamed run itself, mirroring the TypeScript handler's ``AbortController``. + """ import importlib agents_mod = importlib.import_module("agents") Runner = agents_mod.Runner - tracer_name = "@launchdarkly/ai-openai-agents" - if _HAS_OTEL: - span = trace.get_tracer(tracer_name).start_span("openai.agent.run.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 + span = start_root_span(config, variables) + parent = parent_context_of(span) agent, prompt, instructions = _build_agent_and_prompt( config, user_input, tool_handlers, variables, history ) - - if span: - prompt_text = (f"system: {instructions}\n\n" if instructions else "") + prompt - span.add_event("gen_ai.content.prompt", {"gen_ai.prompt": prompt_text}) - prompt_msgs: list[dict[str, str]] = [] - if instructions: - prompt_msgs.append({"role": "system", "content": instructions}) - prompt_msgs.append({"role": "user", "content": prompt}) - set_openllmetry_prompt(span, prompt_msgs) + ended: set[int] = set() + run_usage = create_run_usage() + hooks = _SpanningHooks(config, parent, capture_content, run_usage) + streamed: Any = None + # Whether the vendor's stream ran to the end. Tracked separately from the span, because the + # teardown below has to cancel the Runner on an install with no spans at all. + stream_completed = False try: - streamed = Runner.run_streamed(agent, prompt) + # Inside the guard, because serialising the prompt raises on anything that is not + # JSON-serialisable. A raise out here would leave the root open with the `finally` never + # entered, so the run would vanish from AI Config Monitoring with its feature_flag event. + set_input_content_attributes( + span, + capture_content, + system_instructions=instructions, + messages=[text_message("user", prompt)], + ) + streamed = Runner.run_streamed(agent, prompt, hooks=hooks) full_output = "" async for event in streamed.stream_events(): @@ -298,46 +502,53 @@ async def _stream_gen( yield {"type": "chunk", "text": delta} full_output += delta - input_tokens, output_tokens, total_tokens = _sum_usage(streamed.raw_responses) - final_output = streamed.final_output or full_output + stream_completed = True + final_output = streamed.final_output + output = ( + _stringify_output(final_output) if final_output is not None else full_output + ) - if span: - span.set_attribute( - "gen_ai.response.model", config.get("model", {}).get("name", "") - ) - span.set_attribute("gen_ai.usage.input_tokens", input_tokens) - span.set_attribute("gen_ai.usage.output_tokens", output_tokens) - span.set_attribute("gen_ai.usage.total_tokens", total_tokens) - span.add_event( - "gen_ai.content.completion", - { - "gen_ai.completion": str(final_output) - if isinstance(final_output, str) - else json.dumps(final_output) - }, - ) - set_openllmetry_completion( - span, - str(final_output) - if isinstance(final_output, str) - else json.dumps(final_output), - {"input_tokens": input_tokens, "output_tokens": output_tokens}, - ) - span.set_status(SpanStatusCode.OK) - span.end() + set_output_content_attributes( + span, capture_content, [text_message("assistant", output)] + ) + finish_root_span(span, config, run_usage.total) + mark_ok(span) + end_span_once(span, ended) yield { "type": "done", - "output": str(final_output), - "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens}, + "output": 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() + hooks.close_open_spans(exc) + _write_failed_run_usage(span, config, exc, run_usage) + fail_span(span, exc, ended) raise + finally: + # A no-op on the success and failure paths: both already ended their spans through + # `ended`. On abandonment it is the only chance to close the tree, and the only chance to + # stop the vendor's run — breaking out of `async for` above only stops us reading; the + # Runner's own background task keeps calling the model and spending tokens until told to + # stop. + # Cancelling the Runner is not telemetry, so it does not sit behind a span. Without the otel + # extra there is no root span at all, and the guard below would have skipped the cancel + # entirely: an early break would leave the vendor's background task calling the model and + # spending money, which is the failure this teardown exists to prevent. + if streamed is not None and not stream_completed: + try: + streamed.cancel() + except Exception: # pragma: no cover - best-effort teardown + pass + if span is not None and id(span) not in ended: + hooks.abandon_open_spans(ended) + if run_usage.reported: + finish_root_span(span, config, run_usage.total) + end_span_once(span, ended, abandoned=True) def openai_agents( @@ -347,7 +558,13 @@ def openai_agents( **kwargs: Any, ) -> Any: """Convenience wrapper: creates a handler and calls config(...).invoke().""" + # Both are lifted out of kwargs: capture_content configures the handler, variables belong to + # the invocation. Leaving either in would pass it to config(), which takes neither, so a caller + # asking for content on spans got a TypeError instead of content. variables = kwargs.pop("variables", None) + capture_content = kwargs.pop("capture_content", False) return config( - key=config_key, handler=create_openai_agent_handler(), **kwargs + key=config_key, + handler=create_openai_agent_handler(capture_content=capture_content), + **kwargs, ).invoke(user_input, context, variables=variables) diff --git a/packages/openai-agents/src/launchdarkly_ai_openai_agents/spans.py b/packages/openai-agents/src/launchdarkly_ai_openai_agents/spans.py new file mode 100644 index 0000000..1fc92d7 --- /dev/null +++ b/packages/openai-agents/src/launchdarkly_ai_openai_agents/spans.py @@ -0,0 +1,373 @@ +"""Span construction for the OpenAI Agents handler. + +Separate from ``handler.py`` so the span shape is readable on its own. 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. + +``gen_ai.response.model`` is the *requested* name everywhere on this handler, on both the root and +every ``chat`` span. The openai-agents SDK never hands back a resolved model name the way the +Responses API does for ``openai-messages``, and TELEMETRY-CONTRACT.md section 2a is explicit that +this handler must not invent that behaviour. + +Finish reasons are derived, not mapped, because the Responses API this handler sits on has no +per-message ``finish_reason`` field. See :func:`derive_finish_reason` for the important limitation +this port has relative to TELEMETRY-CONTRACT.md section 5a: only two of its four steps are +reachable through the ``openai-agents`` Python SDK's public ``ModelResponse``. +""" + +from __future__ import annotations + +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_usage_span_attributes, + text_message, +) + +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-agents" + +PROVIDER = "openai" + + +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: these handlers open a plain span rather than an + active one, so there is no ambient span for a bare ``context.active()`` to inherit. + """ + 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.""" + 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_call_id: str, parent: Any) -> Any: + """Opens one ``execute_tool {name}`` span for one tool call.""" + if not _HAS_OTEL: + return None + span = trace.get_tracer(TRACER_NAME).start_span( + f"execute_tool {tool_name}", context=parent + ) + span.set_attribute("gen_ai.operation.name", "execute_tool") + span.set_attribute("gen_ai.tool.name", tool_name) + span.set_attribute("gen_ai.tool.call.id", tool_call_id) + return span + + +# ─── Span finishes ─────────────────────────────────────────────────────────── + + +def finish_root_span(span: Any, config: AiConfigRep, usage: SpanUsage) -> None: + """Writes the run-level identity and token totals onto the root. + + ``gen_ai.response.model`` is the requested name here: this handler never resolves an answering + model. 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, usage) + + +def finish_model_span( + span: Any, + config: AiConfigRep, + usage: SpanUsage, + finish_reason: str | None = None, +) -> None: + """Ends one ``chat`` span successfully.""" + if span is None: + return + span.set_attribute("gen_ai.response.model", model_name(config)) + # A list because one response may hold several choices; the Responses API returns one. + 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's ``finally`` owns every end.""" + 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: + """One turn's ``agents.Usage`` (or the run-level equivalent) as a ``SpanUsage``. + + OpenAI reports cached tokens *inside* the input total, unlike Anthropic's separate buckets, so + nothing is added on top here: TELEMETRY-CONTRACT.md section 8 puts ``openai-agents`` in the + "pass through" row. OpenAI has no cache-creation concept, so that figure is always 0. + + Accepts ``None`` so a turn (or an error path) that never reported usage still yields all-zero + numbers rather than raising. + """ + if usage is None: + return SpanUsage() + details = _attr(usage, "input_tokens_details") + cached = _attr(details, "cached_tokens") if details is not None else None + return SpanUsage( + input=number_or_zero(_attr(usage, "input_tokens")), + output=number_or_zero(_attr(usage, "output_tokens")), + cache_read=number_or_zero(cached), + cache_creation=0, + ) + + +# ─── Finish reasons ─────────────────────────────────────────────────────────── + + +def derive_finish_reason(output: list[Any] | None) -> str | None: + """Derives a semconv finish reason from one turn's output items. + + TELEMETRY-CONTRACT.md section 5a specifies a four-step precedence for the Responses API: + + 1. Any output item is a function call → ``tool_calls``. + 2. Otherwise, the response status is ``incomplete`` → ``length`` or ``content_filter``, + depending on ``incomplete_details.reason``. + 3. Otherwise, the response status is ``completed`` → ``stop``. + 4. Otherwise, no attribute. + + Only step 1 and a version of step 3 are implemented. The openai-agents Python SDK's public + ``Model.get_response`` return type, ``agents.items.ModelResponse``, carries only ``output``, + ``usage``, ``response_id`` and ``request_id``. Unlike the TypeScript ``@openai/agents`` + package, whose ``ModelResponse.providerData`` holds the entire raw OpenAI response (status and + ``incomplete_details`` included), the Python SDK's ``OpenAIResponsesModel.get_response`` + discards the raw response's ``status`` before constructing ``ModelResponse``, on both the + blocking and the streaming path. No wrapper built on top of the public ``Model`` or + ``ModelProvider`` interfaces, nor a ``RunHooks`` callback, can recover it, because the loss + happens inside ``OpenAIResponsesModel`` itself, one layer below anything this handler can + intercept without reimplementing the model call. + + So step 2 cannot be implemented without either forking the vendor SDK or re-issuing the + Responses API call ourselves, and the latter would violate "do not change what the model is + sent." Treated as an open question in the report rather than guessed at: a turn that is cut off + by length or moderation reports ``stop`` here, identically to one that finished normally, + which is wrong but is the closest available approximation given the data this SDK exposes. + """ + items = output or [] + if any(_attr(item, "type") == "function_call" for item in items): + return "tool_calls" + if items: + return "stop" + return None + + +# ─── Provider shapes as span shapes ────────────────────────────────────────── + + +def _attr(obj: Any, name: str) -> Any: + """Reads a field off a provider object or a plain dict, whichever the caller holds. + + Request-side items are typically plain dicts (``TResponseInputItem`` is a ``TypedDict``); + response-side items are typically pydantic models. Both reach these converters. + """ + if obj is None: + return None + if isinstance(obj, dict): + return obj.get(name) + return getattr(obj, name, None) + + +def _text_of_content_blocks(content: Any) -> str: + """Flattens a Responses API ``content`` list (or a bare string) to plain text.""" + if isinstance(content, str): + return content + if not isinstance(content, list): + return "" + parts: list[str] = [] + for block in content: + block_type = _attr(block, "type") + if block_type in ("input_text", "output_text", "text"): + parts.append(str(_attr(block, "text") or "")) + return "".join(parts) + + +def _reasoning_text(item: Any) -> str: + """A reasoning item's summary, joined. Falls back to ``content`` when ``summary`` is absent.""" + summary = _attr(item, "summary") + if isinstance(summary, list) and summary: + return "\n".join(str(_attr(entry, "text") or "") for entry in summary) + return _text_of_content_blocks(_attr(item, "content")) + + +def item_parts(item: Any) -> list[SpanMessagePart]: + """Converts one Responses API item (request or response side) into canonical span parts. + + Item kinds a span has no part for are dropped rather than emitted malformed. + """ + 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 == "function_call_output": + call_id = _attr(item, "call_id") + return [ + SpanMessagePart( + type="tool_call_response", + id=call_id if isinstance(call_id, str) else None, + result=_attr(item, "output"), + ) + ] + if item_type == "reasoning": + text = _reasoning_text(item) + return [SpanMessagePart(type="reasoning", content=text)] if text else [] + + content = _attr(item, "content") + if isinstance(content, str): + return [SpanMessagePart(type="text", content=content)] if content else [] + text = _text_of_content_blocks(content) + return [SpanMessagePart(type="text", content=text)] if text else [] + + +def to_request_span_messages(input_items: Any) -> list[SpanMessage]: + """A model turn's input, whether the SDK passed a bare string or a list of items.""" + if isinstance(input_items, str): + return [text_message("user", input_items)] if input_items else [] + if not isinstance(input_items, list): + return [] + messages: list[SpanMessage] = [] + for item in input_items: + role = _attr(item, "role") + if not isinstance(role, str): + role = ( + "tool" if _attr(item, "type") == "function_call_output" else "assistant" + ) + messages.append(SpanMessage(role=role, parts=item_parts(item))) + return messages + + +def to_response_span_messages( + output_items: list[Any] | None, finish_reason: str | None +) -> list[SpanMessage]: + """One turn's output items as canonical span messages. + + ``finish_reason`` is attached only to the last message, mirroring the shape of a single + Responses API turn: one logical assistant turn, whatever it took to produce it. + """ + items = output_items or [] + messages: list[SpanMessage] = [] + for index, item in enumerate(items): + role = _attr(item, "role") + if not isinstance(role, str): + role = "assistant" + is_last = index == len(items) - 1 + messages.append( + SpanMessage( + role=role, + parts=item_parts(item), + finish_reason=finish_reason if is_last else None, + ) + ) + return messages + + +def to_tool_definitions(tools: list[Any]) -> list[ToolDefinitionInput]: + """The catalog as the ``Agent`` actually holds it, so the span reports what the model could + actually call. + + Filtered to ``FunctionTool``-shaped entries (``name`` / ``description`` / + ``params_json_schema``): the Agents SDK also allows hosted and computer tools this handler + never builds, and those have no JSON Schema to report. + """ + definitions: list[ToolDefinitionInput] = [] + for tool in tools: + params = _attr(tool, "params_json_schema") + name = _attr(tool, "name") + if not isinstance(name, str): + continue + definitions.append( + ToolDefinitionInput( + name=name, + description=_attr(tool, "description"), + parameters=params, + ) + ) + return definitions diff --git a/packages/openai-agents/tests/test_handler.py b/packages/openai-agents/tests/test_handler.py index a629483..da344e9 100644 --- a/packages/openai-agents/tests/test_handler.py +++ b/packages/openai-agents/tests/test_handler.py @@ -1,38 +1,186 @@ """ Tests for launchdarkly-ai-openai-agents handler. -Covers §1.1–1.9 (generic) and OpenAI-agents-specific extras. -Reference: TESTING.md §1, §2.x (OpenAI) + +TELEMETRY-CONTRACT.md sections 1-9 for the span tree, and TESTING.md §1 for the generic handler +behaviours. Rewritten from the pre-span-work version: this handler drives the ``agents`` SDK's own +``Runner``, which owns the per-turn loop, so a test drives spans by calling the ``RunHooks`` +callbacks (``on_llm_start`` / ``on_llm_end`` / ``on_tool_start`` / ``on_tool_end``) the way the real +``Runner`` would, rather than by mocking a client response directly. """ from __future__ import annotations +import json from collections.abc import AsyncIterator +from types import SimpleNamespace from typing import Any, ClassVar from unittest.mock import AsyncMock, MagicMock, patch import pytest import launchdarkly_ai_openai_agents.handler as handler_mod +import launchdarkly_ai_openai_agents.spans as spans_mod from launchdarkly_ai_openai_agents.handler import ( _build_agent_and_prompt, + _build_agent_tools, create_openai_agent_handler, + openai_agents, ) from launchdarkly_ai_openai_agents.utils import build_output_type # --------------------------------------------------------------------------- -# Helpers +# Fake `agents` SDK # --------------------------------------------------------------------------- -def _make_config(**kwargs: Any) -> dict[str, Any]: - base = {"model": {"name": "gpt-4o"}, "provider": {"name": "OpenAI"}} - base.update(kwargs) - return base +class FakeFunctionTool: + def __init__(self, **kw: Any) -> None: + self.name = kw["name"] + self.description = kw.get("description") + self.params_json_schema = kw.get("params_json_schema") + self.on_invoke_tool = kw.get("on_invoke_tool") + + +class FakeAgent: + def __init__(self, **kw: Any) -> None: + self.kwargs = kw + self.name = kw.get("name") + self.model = kw.get("model") + self.instructions = kw.get("instructions") + self.tools = kw.get("tools", []) + self.output_type = kw.get("output_type") + + +class FakeRunResult: + def __init__(self, final_output: Any, run_data: Any = None) -> None: + self.final_output = final_output + self.run_data = run_data + + +async def _drive_turns( + hooks: Any, agent: Any, prompt: str, turns: list[dict[str, Any]] +) -> None: + """Calls the ``RunHooks`` callbacks the way the real ``Runner.run`` would, one turn at a time.""" + for turn in turns: + await hooks.on_llm_start(MagicMock(), agent, agent.instructions, prompt) + for call in turn.get("tool_calls", []): + ctx = SimpleNamespace( + tool_name=call["name"], + tool_call_id=call["id"], + tool_arguments=json.dumps(call.get("args", {})), + ) + tool_obj = SimpleNamespace(name=call["name"]) + await hooks.on_tool_start(ctx, agent, tool_obj) + if "error" in call: + raise call["error"] + await hooks.on_tool_end(ctx, agent, tool_obj, call.get("result", "ok")) + if turn.get("error") is not None: + raise turn["error"] + response = SimpleNamespace( + output=turn.get("output", []), usage=turn.get("usage", {}) + ) + await hooks.on_llm_end(MagicMock(), agent, response) + + +def _make_run(turns: list[dict[str, Any]], final_output: str = "done") -> Any: + async def run(agent: Any, prompt: str, hooks: Any = None, **kw: Any) -> Any: + await _drive_turns(hooks, agent, prompt, turns) + return FakeRunResult(final_output) + + return run + + +class FakeStreamedResult: + """Interleaves hook calls with delta events, the way ``RunResultStreaming`` really does. + + Breaking out of ``stream_events()`` early therefore leaves whichever turn was in flight + without its ``on_llm_end`` (or ``on_tool_end``), exactly as an abandoned real run would. + """ + + def __init__( + self, + agent: Any, + prompt: str, + hooks: Any, + turns: list[dict[str, Any]], + final_output: str, + ) -> None: + self.agent = agent + self.prompt = prompt + self.hooks = hooks + self.turns = turns + self.final_output = final_output + self.cancelled = False + + async def stream_events(self) -> AsyncIterator[Any]: + for turn in self.turns: + await self.hooks.on_llm_start( + MagicMock(), self.agent, self.agent.instructions, self.prompt + ) + for delta in turn.get("deltas", []): + yield SimpleNamespace( + type="raw_response_event", + data=SimpleNamespace( + type="response.output_text.delta", delta=delta + ), + ) + for call in turn.get("tool_calls", []): + ctx = SimpleNamespace( + tool_name=call["name"], + tool_call_id=call["id"], + tool_arguments=json.dumps(call.get("args", {})), + ) + tool_obj = SimpleNamespace(name=call["name"]) + await self.hooks.on_tool_start(ctx, self.agent, tool_obj) + await self.hooks.on_tool_end( + ctx, self.agent, tool_obj, call.get("result", "ok") + ) + if turn.get("error") is not None: + raise turn["error"] + response = SimpleNamespace( + output=turn.get("output", []), usage=turn.get("usage", {}) + ) + await self.hooks.on_llm_end(MagicMock(), self.agent, response) + + def cancel(self, mode: str = "immediate") -> None: + self.cancelled = True + + +def _make_run_streamed(turns: list[dict[str, Any]], final_output: str = "done") -> Any: + def run_streamed(agent: Any, prompt: str, hooks: Any = None, **kw: Any) -> Any: + return FakeStreamedResult(agent, prompt, hooks, turns, final_output) + + return run_streamed + + +def _fake_agents_module(run: Any = None, run_streamed: Any = None) -> Any: + mod = SimpleNamespace() + mod.FunctionTool = FakeFunctionTool + mod.Agent = FakeAgent + + class Runner: + pass + + Runner.run = staticmethod(run or _make_run([{"output": [], "usage": {}}])) + Runner.run_streamed = staticmethod( + run_streamed or _make_run_streamed([{"output": [], "usage": {}}]) + ) + mod.Runner = Runner + return mod + + +def _patched_agents(agents_mod: Any) -> Any: + return patch( + "importlib.import_module", + side_effect=lambda n: agents_mod if n == "agents" else __import__(n), + ) def _make_run_result( output: str = "hello", input_tokens: int = 10, output_tokens: int = 5 ) -> Any: + """The pre-span-era flat mock result, kept for the restored non-telemetry tests that predate + the ``RunHooks``-driven span work and never touch spans or usage attribution.""" usage = MagicMock() usage.input_tokens = input_tokens usage.output_tokens = output_tokens @@ -47,7 +195,15 @@ def _make_run_result( return result +async def _empty_async_gen() -> AsyncIterator[Any]: + return + yield + + def _mock_agents_module(run_result: Any) -> Any: + """The pre-span-era fully-flat ``agents`` mock: ``Runner.run``/``run_streamed`` never invoke + ``hooks``. Kept for the restored tests that only care about prompt/tool wiring, not spans. + """ mock = MagicMock() mock.Agent = MagicMock(return_value=MagicMock()) mock.Runner.run = AsyncMock(return_value=run_result) @@ -64,9 +220,112 @@ def _mock_agents_module(run_result: Any) -> Any: return mock -async def _empty_async_gen() -> AsyncIterator[Any]: - return - yield +def _make_config(**kwargs: Any) -> dict[str, Any]: + base: dict[str, Any] = {"model": {"name": "gpt-4o"}, "provider": {"name": "OpenAI"}} + base.update(kwargs) + return base + + +CONFIG = _make_config(instructions="Be helpful.") + + +def _text_output(text: str) -> list[dict[str, Any]]: + return [{"type": "message", "role": "assistant", "content": text}] + + +def _tool_call_output(name: str, call_id: str, arguments: str = "{}") -> dict[str, Any]: + return { + "type": "function_call", + "call_id": call_id, + "name": name, + "arguments": arguments, + } + + +def _usage( + input_tokens: int = 10, output_tokens: int = 5, cached: int | None = None +) -> dict[str, Any]: + usage: dict[str, Any] = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + } + if cached is not None: + usage["input_tokens_details"] = {"cached_tokens": cached} + return usage + + +# --------------------------------------------------------------------------- +# 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. + + A single ``MagicMock`` cannot see a span tree at all, because every span is the same object and + a parent and its children are indistinguishable. This keeps one object per span. + """ + + 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_openai_agents.spans as spans_mod + + recorder = SpanRecorder() + return patch.object(spans_mod, "trace", recorder), recorder # --------------------------------------------------------------------------- @@ -76,23 +335,20 @@ async def _empty_async_gen() -> AsyncIterator[Any]: class TestFactory: def test_returns_callable(self) -> None: - h = create_openai_agent_handler() - assert callable(h) + assert callable(create_openai_agent_handler()) def test_attaches_provides_for(self) -> None: h = create_openai_agent_handler() - assert hasattr(h, "provides_for") + assert h.provides_for == ("OpenAI", "agent") + + def test_multiple_calls_return_independent_instances(self) -> None: + assert create_openai_agent_handler() is not create_openai_agent_handler() def test_provides_for_values_are_correct(self) -> None: h = create_openai_agent_handler() pf = h.provides_for assert "OpenAI" in pf or "openai" in str(pf).lower() - def test_multiple_calls_return_independent_instances(self) -> None: - h1 = create_openai_agent_handler() - h2 = create_openai_agent_handler() - assert h1 is not h2 - # --------------------------------------------------------------------------- # §1.2 Prompt construction @@ -203,69 +459,73 @@ def test_path_c_instructions_takes_priority_over_messages(self) -> None: _, _, instructions = _build_agent_and_prompt(config, "q", {}, {}) assert instructions == "Use instructions." + def test_instructions_path(self) -> None: + config = _make_config(instructions="Be concise.") + _, prompt, instructions = _build_agent_and_prompt(config, "hi", {}, {}) + assert instructions == "Be concise." + assert prompt == "hi" -# --------------------------------------------------------------------------- -# §1.3 Tool conversion -# --------------------------------------------------------------------------- - - -class TestToolConversion: - @pytest.mark.asyncio - async def test_all_fields_forwarded(self) -> None: - run_result = _make_run_result("out") - agents_mock = _mock_agents_module(run_result) - - captured_tools: list[Any] = [] - - def _capture_tool(**kw: Any) -> Any: - captured_tools.append(kw) - return MagicMock() - - agents_mock.FunctionTool = MagicMock(side_effect=_capture_tool) + def test_variable_substitution(self) -> None: + config = _make_config(instructions="Hello {{name}}!") + _, _, instructions = _build_agent_and_prompt(config, "hi", {}, {"name": "Bob"}) + assert instructions == "Hello Bob!" - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - h = create_openai_agent_handler() - config = _make_config( - tools={"my-tool": {"description": "does stuff", "parameters": {}}} - ) - await h(config, "hi", {"my-tool": AsyncMock(return_value="ok")}) + def test_messages_path_extracts_system(self) -> None: + config = _make_config( + messages=[ + {"role": "system", "content": "System msg."}, + {"role": "user", "content": "prior turn"}, + ] + ) + _, prompt, instructions = _build_agent_and_prompt(config, "question", {}, {}) + assert instructions == "System msg." + assert "prior turn" in prompt + assert "question" in prompt - names = [t.get("name") for t in captured_tools] - assert "my-tool" in names + def test_instructions_take_priority_over_messages(self) -> None: + config = _make_config( + instructions="From instructions.", + messages=[{"role": "system", "content": "From messages."}], + ) + _, _, instructions = _build_agent_and_prompt(config, "q", {}, {}) + assert instructions == "From instructions." - @pytest.mark.asyncio - async def test_multiple_tools_all_included(self) -> None: - run_result = _make_run_result("out") - agents_mock = _mock_agents_module(run_result) + def test_none_user_input_becomes_empty_string(self) -> None: + config = _make_config(instructions="Be helpful.") + _, prompt, _ = _build_agent_and_prompt(config, None, {}, {}) + assert prompt == "" - captured_tools: list[str] = [] - def _capture_tool(**kw: Any) -> Any: - captured_tools.append(kw.get("name")) - return MagicMock() +# --------------------------------------------------------------------------- +# §1.3 / §1.4 Tool conversion and execution +# --------------------------------------------------------------------------- - agents_mock.FunctionTool = MagicMock(side_effect=_capture_tool) - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - h = create_openai_agent_handler() - config = _make_config( - tools={ - "tool-a": {"description": "a", "parameters": {}}, - "tool-b": {"description": "b", "parameters": {}}, +class TestToolConversion: + def test_all_fields_forwarded(self) -> None: + tools = _build_agent_tools( + { + "my-tool": { + "description": "does stuff", + "parameters": {"type": "object"}, } - ) - await h(config, "hi", {"tool-a": AsyncMock(), "tool-b": AsyncMock()}) + }, + {"my-tool": AsyncMock()}, + ) + assert len(tools) == 1 + assert tools[0].name == "my-tool" + assert tools[0].description == "does stuff" + + def test_multiple_tools_all_included(self) -> None: + tools = _build_agent_tools( + {"tool-a": {"description": "a"}, "tool-b": {"description": "b"}}, + {"tool-a": AsyncMock(), "tool-b": AsyncMock()}, + ) + assert {t.name for t in tools} == {"tool-a", "tool-b"} - assert "tool-a" in captured_tools - assert "tool-b" in captured_tools + def test_empty_tools_no_tools_built(self) -> None: + assert _build_agent_tools({}, {}) == [] - @pytest.mark.asyncio async def test_empty_tools_no_tools_sent(self) -> None: run_result = _make_run_result("out") agents_mock = _mock_agents_module(run_result) @@ -284,17 +544,27 @@ async def test_empty_tools_no_tools_sent(self) -> None: if captured: assert not captured[0].get("tools") + async def test_tool_not_found_throws(self) -> None: + tools = _build_agent_tools({"my-tool": {"description": "d"}}, {}) + with pytest.raises(ValueError, match="No handler"): + await tools[0].on_invoke_tool(MagicMock(), "{}") + + async def test_tool_handler_throws_propagates(self) -> None: + async def _bad(args: Any) -> str: + raise RuntimeError("handler error") + + tools = _build_agent_tools({"my-tool": {"description": "d"}}, {"my-tool": _bad}) + with pytest.raises(RuntimeError, match="handler error"): + await tools[0].on_invoke_tool(MagicMock(), "{}") + # --------------------------------------------------------------------------- -# §1.4 Tool execution loop +# §1.4 Tool execution loop (pre-span-era; restored, not telemetry) # --------------------------------------------------------------------------- class TestToolExecutionLoop: - @pytest.mark.asyncio async def test_tool_not_found_execute_callback_throws(self) -> None: - from launchdarkly_ai_openai_agents.handler import _build_agent_tools - agents_mock = MagicMock() captured_fns: list[Any] = [] agents_mock.tool = MagicMock( @@ -311,32 +581,6 @@ async def test_tool_not_found_execute_callback_throws(self) -> None: with pytest.raises(ValueError, match="No handler"): await captured_fns[0]({}) - @pytest.mark.asyncio - async def test_tool_handler_throws_propagates(self) -> None: - from launchdarkly_ai_openai_agents.handler import _build_agent_tools - - agents_mock = MagicMock() - captured_fns: list[Any] = [] - agents_mock.tool = MagicMock( - side_effect=lambda **kw: lambda fn: (captured_fns.append(fn), fn)[1] - ) - - async def _bad_handler(args: Any) -> str: - raise RuntimeError("handler error") - - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - _build_agent_tools( - {"my-tool": {"description": "d"}}, {"my-tool": _bad_handler} - ) - - if captured_fns: - with pytest.raises(RuntimeError, match="handler error"): - await captured_fns[0]({}) - - @pytest.mark.asyncio async def test_no_tools_in_config_tool_builder_never_called(self) -> None: run_result = _make_run_result("out") agents_mock = _mock_agents_module(run_result) @@ -356,405 +600,450 @@ async def test_no_tools_in_config_tool_builder_never_called(self) -> None: # --------------------------------------------------------------------------- -# §1.5 Telemetry +# TELEMETRY-CONTRACT.md section 1: span tree # --------------------------------------------------------------------------- -class TestTelemetry: - @pytest.mark.asyncio - async def test_span_name_blocking(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - run_result = _make_run_result("out") - agents_mock = _mock_agents_module(run_result) +class TestSpanTree: + async def test_opens_a_root_span_named_invoke_agent(self) -> None: + ctx, rec = _recording() + agents_mod = _fake_agents_module( + run=_make_run([{"output": _text_output("hi")}]) + ) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler()(CONFIG, "q", {}, {}) + assert rec.root.name == "invoke_agent" + assert rec.root.attributes["gen_ai.operation.name"] == "invoke_agent" + + async def test_emits_one_chat_child_per_model_turn(self) -> None: + ctx, rec = _recording() + agents_mod = _fake_agents_module( + run=_make_run([{"output": _text_output("hi")}]) + ) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_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" + assert chats[0].context == ("context-of", rec.root) + + async def test_emits_a_chat_span_per_turn_of_a_tool_loop(self) -> None: + turns = [ + { + "output": [_tool_call_output("search", "call_1")], + "tool_calls": [{"name": "search", "id": "call_1"}], + }, + {"output": _text_output("done")}, + ] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_make_run(turns)) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler()( + CONFIG, "q", {"search": AsyncMock()}, {} + ) + assert len(rec.named("chat ")) == 2 + + async def test_emits_an_execute_tool_span_per_tool_call(self) -> None: + turns = [ + { + "output": [_tool_call_output("search", "call_1")], + "tool_calls": [{"name": "search", "id": "call_1"}], + }, + {"output": _text_output("done")}, + ] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_make_run(turns)) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler()( + CONFIG, "q", {"search": AsyncMock()}, {} + ) + tools = rec.named("execute_tool ") + assert len(tools) == 1 + assert tools[0].name == "execute_tool search" + assert tools[0].attributes["gen_ai.operation.name"] == "execute_tool" + assert tools[0].attributes["gen_ai.tool.name"] == "search" + assert tools[0].attributes["gen_ai.tool.call.id"] == "call_1" + + async def test_tool_spans_are_siblings_of_chat_not_children(self) -> None: + turns = [ + { + "output": [_tool_call_output("search", "call_1")], + "tool_calls": [{"name": "search", "id": "call_1"}], + }, + {"output": _text_output("done")}, + ] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_make_run(turns)) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler()( + CONFIG, "q", {"search": AsyncMock()}, {} + ) + assert rec.named("execute_tool ")[0].context == ("context-of", rec.root) + + async def test_every_span_is_ended(self) -> None: + turns = [ + { + "output": [_tool_call_output("search", "call_1")], + "tool_calls": [{"name": "search", "id": "call_1"}], + }, + {"output": _text_output("done")}, + ] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_make_run(turns)) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler()( + CONFIG, "q", {"search": AsyncMock()}, {} + ) + assert [s.ended for s in rec.spans] == [1] * len(rec.spans) + + async def test_children_carry_no_launchdarkly_attributes(self) -> None: + turns = [ + { + "output": [_tool_call_output("search", "call_1")], + "tool_calls": [{"name": "search", "id": "call_1"}], + }, + {"output": _text_output("done")}, + ] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_make_run(turns)) + vs = {"__ld": {"configKey": "c", "variationKey": "v", "runId": "r"}} + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler()( + CONFIG, "q", {"search": AsyncMock()}, vs + ) + for span in rec.spans[1:]: + assert [k for k in span.attributes if k.startswith("launchdarkly.")] == [] + assert [n for n, _ in rec.root.events] == ["feature_flag"] + assert all( + "feature_flag" not in [n for n, _ in s.events] for s in rec.spans[1:] + ) - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - await h(_make_config(), "hi") - mock_trace.get_tracer.return_value.start_span.assert_called_with( - "openai.agent.run" - ) +# --------------------------------------------------------------------------- +# TELEMETRY-CONTRACT.md sections 2 / 2a: root span +# --------------------------------------------------------------------------- - @pytest.mark.asyncio - async def test_gen_ai_system(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - run_result = _make_run_result("out") - agents_mock = _mock_agents_module(run_result) +class TestRootSpanAttributes: + async def test_writes_both_provider_keys_and_the_requested_model(self) -> None: + ctx, rec = _recording() + agents_mod = _fake_agents_module( + run=_make_run([{"output": _text_output("hi")}]) + ) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_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_requested_name_not_a_resolved_one( + self, + ) -> None: + # Section 2a: openai-agents never resolves an answering model. + ctx, rec = _recording() + agents_mod = _fake_agents_module( + run=_make_run([{"output": _text_output("hi")}]) + ) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler()(CONFIG, "q", {}, {}) + assert rec.root.attributes["gen_ai.response.model"] == "gpt-4o" - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - await h(_make_config(), "hi") + async def test_run_totals_on_the_root(self) -> None: + turns = [ + {"output": _text_output("a"), "usage": _usage(10, 5)}, + ] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_make_run(turns)) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler()(CONFIG, "q", {}, {}) + attrs = rec.root.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 + + async def test_feature_flag_event_on_root_only(self) -> None: + ctx, rec = _recording() + agents_mod = _fake_agents_module( + run=_make_run([{"output": _text_output("hi")}]) + ) + vs = { + "__ld": { + "configKey": "my-config", + "variationKey": "v1", + "runId": "run-abc", + "environmentId": "env-1", + } + } + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler()(CONFIG, "q", {}, vs) + assert rec.root.events == [ + ( + "feature_flag", + { + "feature_flag.key": "my-config", + "feature_flag.provider.name": "LaunchDarkly", + "feature_flag.set.id": "env-1", + }, + ) + ] + assert rec.root.attributes["launchdarkly.config.key"] == "my-config" + assert rec.root.attributes["launchdarkly.variation.key"] == "v1" + assert rec.root.attributes["launchdarkly.run.id"] == "run-abc" - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("gen_ai.system") == "openai" - @pytest.mark.asyncio - async def test_gen_ai_operation_name(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span +# --------------------------------------------------------------------------- +# TELEMETRY-CONTRACT.md sections 3, 5a, 8: chat span attributes +# --------------------------------------------------------------------------- - run_result = _make_run_result("out") - agents_mock = _mock_agents_module(run_result) - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), +class TestChatSpanAttributes: + async def test_passes_input_through_without_folding_cache_into_it(self) -> None: + # Section 8: OpenAI already counts cached tokens inside input_tokens. + turns = [{"output": _text_output("hi"), "usage": _usage(50, 5, cached=30)}] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_make_run(turns)) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler()(CONFIG, "q", {}, {}) + chat = rec.named("chat ")[0] + assert chat.attributes["gen_ai.usage.input_tokens"] == 50 + assert chat.attributes["gen_ai.usage.output_tokens"] == 5 + assert chat.attributes["gen_ai.usage.total_tokens"] == 55 + assert chat.attributes["gen_ai.usage.cache_read.input_tokens"] == 30 + assert chat.attributes["gen_ai.usage.cache_creation.input_tokens"] == 0 + + async def test_derives_total_from_input_plus_output_not_the_provider_total( + self, + ) -> None: + turns = [{"output": _text_output("hi"), "usage": _usage(10, 5)}] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_make_run(turns)) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler()(CONFIG, "q", {}, {}) + chat = rec.named("chat ")[0] + assert chat.attributes["gen_ai.usage.total_tokens"] == 15 + + async def test_finish_reason_tool_calls_when_output_holds_a_function_call( + self, + ) -> None: + turns = [ + { + "output": [_tool_call_output("search", "call_1")], + "tool_calls": [{"name": "search", "id": "call_1"}], + }, + {"output": _text_output("done")}, + ] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_make_run(turns)) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler()( + CONFIG, "q", {"search": AsyncMock()}, {} + ) + chats = rec.named("chat ") + assert chats[0].attributes["gen_ai.response.finish_reasons"] == ["tool_calls"] + assert chats[1].attributes["gen_ai.response.finish_reasons"] == ["stop"] + + async def test_omits_the_finish_reason_when_there_is_no_output(self) -> None: + turns = [{"output": []}] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_make_run(turns)) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler()(CONFIG, "q", {}, {}) + chat = rec.named("chat ")[0] + assert "gen_ai.response.finish_reasons" not in chat.attributes + + async def test_fails_the_chat_span_when_the_model_turn_throws(self) -> None: + err = RuntimeError("model down") + turns = [{"output": [], "error": err}] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_make_run(turns)) + with ( + ctx, + _patched_agents(agents_mod), + pytest.raises(RuntimeError, match="model down"), ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - await h(_make_config(), "hi") + await create_openai_agent_handler()(CONFIG, "q", {}, {}) + chat = rec.named("chat ")[0] + assert chat.exceptions == [err] + from opentelemetry.trace import StatusCode + + assert StatusCode.ERROR in chat.statuses + assert chat.ended == 1 + + async def test_root_still_reports_partial_usage_on_failure(self) -> None: + err = RuntimeError("second turn failed") + turns = [ + {"output": _text_output("partial"), "usage": _usage(10, 5)}, + ] - 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 run(agent: Any, prompt: str, hooks: Any = None, **kw: Any) -> Any: + await _drive_turns(hooks, agent, prompt, turns) + await hooks.on_llm_start(MagicMock(), agent, agent.instructions, prompt) + raise err + + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=run) + with ( + ctx, + _patched_agents(agents_mod), + pytest.raises(RuntimeError, match="second turn failed"), + ): + await create_openai_agent_handler()(CONFIG, "q", {}, {}) + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 10 + assert rec.root.attributes["gen_ai.usage.output_tokens"] == 5 + from opentelemetry.trace import StatusCode - @pytest.mark.asyncio - async def test_gen_ai_request_model(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span + assert StatusCode.ERROR in rec.root.statuses - run_result = _make_run_result("out") - agents_mock = _mock_agents_module(run_result) - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), +class TestToolSpanAttributes: + async def test_fails_an_open_tool_span_when_the_run_crashes_mid_flight( + self, + ) -> None: + err = RuntimeError("run boom") + turns = [ + { + "output": [_tool_call_output("search", "call_1")], + "tool_calls": [{"name": "search", "id": "call_1", "error": err}], + } + ] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_make_run(turns)) + with ( + ctx, + _patched_agents(agents_mod), + pytest.raises(RuntimeError, match="run boom"), ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - await h(_make_config(model={"name": "gpt-4o"}), "hi") + await create_openai_agent_handler()( + CONFIG, "q", {"search": AsyncMock()}, {} + ) + tool_span = rec.named("execute_tool ")[0] + from opentelemetry.trace import StatusCode - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("gen_ai.request.model") == "gpt-4o" + assert StatusCode.ERROR in tool_span.statuses + assert tool_span.ended == 1 - @pytest.mark.asyncio - async def test_token_attributes_set(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - run_result = _make_run_result("out", input_tokens=20, output_tokens=8) - agents_mock = _mock_agents_module(run_result) +# --------------------------------------------------------------------------- +# TELEMETRY-CONTRACT.md section 7: content capture +# --------------------------------------------------------------------------- - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - await h(_make_config(), "hi") - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("gen_ai.usage.input_tokens") == 20 - assert calls.get("gen_ai.usage.output_tokens") == 8 +class TestContentCapture: + async def test_emits_no_content_at_all_by_default(self) -> None: + ctx, rec = _recording() + agents_mod = _fake_agents_module( + run=_make_run([{"output": _text_output("answer")}]) + ) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler()(CONFIG, "my question", {}, {}) + for span in rec.spans: + assert [k for k in span.attributes if k.startswith("gen_ai.prompt")] == [] + assert [ + k for k in span.attributes if k.startswith("gen_ai.completion") + ] == [] + assert "gen_ai.input.messages" not in span.attributes + assert "gen_ai.output.messages" not in span.attributes + 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) -> None: + ctx, rec = _recording() + agents_mod = _fake_agents_module( + run=_make_run([{"output": _text_output("answer")}]) + ) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler(capture_content=True)( + CONFIG, "my question", {}, {} + ) + assert "gen_ai.input.messages" in rec.root.attributes + assert "gen_ai.output.messages" in rec.root.attributes + prompt_written = json.dumps(rec.root.attributes) + assert "my question" in prompt_written + + async def test_tool_call_arguments_and_result_gated(self) -> None: + turns = [ + { + "output": [_tool_call_output("search", "call_1", '{"q": "x"}')], + "tool_calls": [ + { + "name": "search", + "id": "call_1", + "args": {"q": "x"}, + "result": "found", + } + ], + }, + {"output": _text_output("done")}, + ] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_make_run(turns)) + with ctx, _patched_agents(agents_mod): + await create_openai_agent_handler(capture_content=True)( + CONFIG, "q", {"search": AsyncMock()}, {} + ) + tool_span = rec.named("execute_tool ")[0] + assert tool_span.attributes["gen_ai.tool.call.arguments"] == '{"q": "x"}' + assert tool_span.attributes["gen_ai.tool.call.result"] == "found" - @pytest.mark.asyncio - async def test_span_status_ok(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - run_result = _make_run_result("out") - agents_mock = _mock_agents_module(run_result) +# --------------------------------------------------------------------------- +# §1.6 Error handling (top-level, no partial usage) +# --------------------------------------------------------------------------- - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), + +class TestErrorHandling: + async def test_records_exception_sets_error_ends_span_and_rethrows(self) -> None: + err = RuntimeError("provider error") + + async def run(agent: Any, prompt: str, hooks: Any = None, **kw: Any) -> Any: + raise err + + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=run) + with ( + ctx, + _patched_agents(agents_mod), + pytest.raises(RuntimeError, match="provider error"), ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - await h(_make_config(), "hi") + await create_openai_agent_handler()(CONFIG, "hi", {}, {}) + assert rec.root.exceptions == [err] + from opentelemetry.trace import StatusCode + + assert StatusCode.ERROR in rec.root.statuses + assert rec.root.ended == 1 + + async def test_usage_from_agents_exception_reaches_the_root(self) -> None: + err = RuntimeError("max turns exceeded") + err.run_data = SimpleNamespace( + context_wrapper=SimpleNamespace( + usage=SimpleNamespace(input_tokens=100, output_tokens=20) + ) + ) - mock_span.set_status.assert_called() + async def run(agent: Any, prompt: str, hooks: Any = None, **kw: Any) -> Any: + raise err - @pytest.mark.asyncio - async def test_span_end_always_called(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=run) + with ctx, _patched_agents(agents_mod), pytest.raises(RuntimeError): + await create_openai_agent_handler()(CONFIG, "hi", {}, {}) + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 100 + assert rec.root.attributes["gen_ai.usage.output_tokens"] == 20 - run_result = _make_run_result("out") - agents_mock = _mock_agents_module(run_result) - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - await h(_make_config(), "hi") +# --------------------------------------------------------------------------- +# §1.7 Convenience export +# --------------------------------------------------------------------------- - mock_span.end.assert_called() - @pytest.mark.asyncio - async def test_gen_ai_content_prompt_event(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - run_result = _make_run_result("out") - agents_mock = _mock_agents_module(run_result) - - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - await h(_make_config(), "my question") - - event_calls = [ - c - for c in mock_span.add_event.call_args_list - if c[0][0] == "gen_ai.content.prompt" - ] - assert event_calls - # The gen_ai.prompt attribute must include the user input text - prompt_attr = event_calls[0][0][1].get("gen_ai.prompt", "") - assert "my question" in prompt_attr, ( - f"gen_ai.prompt must include user input 'my question', got: {prompt_attr!r}" - ) - - @pytest.mark.asyncio - async def test_gen_ai_content_completion_event(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - run_result = _make_run_result("final answer") - agents_mock = _mock_agents_module(run_result) - - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - await h(_make_config(), "hi") - - event_calls = [ - c - for c in mock_span.add_event.call_args_list - if c[0][0] == "gen_ai.content.completion" - ] - assert event_calls - - @pytest.mark.asyncio - async def test_gen_ai_response_model(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - run_result = _make_run_result("out") - agents_mock = _mock_agents_module(run_result) - - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - await h(_make_config(model={"name": "gpt-4o"}), "hi") - - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert "gen_ai.response.model" in calls - assert calls["gen_ai.response.model"] == "gpt-4o" - - @pytest.mark.asyncio - async def test_ld_span_attributes(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - run_result = _make_run_result("out") - agents_mock = _mock_agents_module(run_result) - variables = { - "__ld": { - "configKey": "my-config", - "variationKey": "v1", - "runId": "run-abc", - } - } - - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - await h(_make_config(), "hi", variables=variables) - - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("launchdarkly.operation.type") == "gen_ai" - assert calls.get("launchdarkly.config.key") == "my-config" - assert calls.get("launchdarkly.variation.key") == "v1" - assert calls.get("launchdarkly.run.id") == "run-abc" - assert "launchdarkly.graph.key" not in calls - - async def test_ld_graph_key_set_when_present(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - run_result = _make_run_result("out") - agents_mock = _mock_agents_module(run_result) - variables = { - "__ld": { - "configKey": "my-config", - "variationKey": "v1", - "runId": "run-abc", - "graphKey": "my-graph", - } - } - - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - await h(_make_config(), "hi", variables=variables) - - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("launchdarkly.graph.key") == "my-graph" - - -# --------------------------------------------------------------------------- -# §1.6 Error handling -# --------------------------------------------------------------------------- - - -class TestErrorHandling: - @pytest.mark.asyncio - async def test_records_exception_on_span(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - agents_mock = MagicMock() - agents_mock.Agent = MagicMock(return_value=MagicMock()) - agents_mock.Runner.run = AsyncMock(side_effect=RuntimeError("provider error")) - agents_mock.tool = MagicMock(side_effect=lambda **kw: lambda fn: fn) - agents_mock.handoff = MagicMock() - - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - with pytest.raises(RuntimeError): - await h(_make_config(), "hi") - - mock_span.record_exception.assert_called() - - @pytest.mark.asyncio - async def test_sets_span_status_error(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - agents_mock = MagicMock() - agents_mock.Agent = MagicMock(return_value=MagicMock()) - agents_mock.Runner.run = AsyncMock(side_effect=RuntimeError("fail")) - agents_mock.tool = MagicMock(side_effect=lambda **kw: lambda fn: fn) - - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - with pytest.raises(RuntimeError): - await h(_make_config(), "hi") - - mock_span.set_status.assert_called() - - @pytest.mark.asyncio - async def test_ends_span_on_error(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - agents_mock = MagicMock() - agents_mock.Agent = MagicMock(return_value=MagicMock()) - agents_mock.Runner.run = AsyncMock(side_effect=RuntimeError("fail")) - agents_mock.tool = MagicMock(side_effect=lambda **kw: lambda fn: fn) - - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - with pytest.raises(RuntimeError): - await h(_make_config(), "hi") - - mock_span.end.assert_called() - - @pytest.mark.asyncio - async def test_rethrows_error(self) -> None: - agents_mock = MagicMock() - agents_mock.Agent = MagicMock(return_value=MagicMock()) - agents_mock.Runner.run = AsyncMock(side_effect=RuntimeError("specific")) - agents_mock.tool = MagicMock(side_effect=lambda **kw: lambda fn: fn) - - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_openai_agent_handler() - with pytest.raises(RuntimeError, match="specific"): - await h(_make_config(), "hi") - - -# --------------------------------------------------------------------------- -# §1.7 Convenience export -# --------------------------------------------------------------------------- - - -class TestConvenienceExport: - def test_calls_through_to_model_call(self) -> None: - from launchdarkly_ai_openai_agents.handler import openai_agents - - assert callable(openai_agents) +class TestConvenienceExport: + def test_calls_through_to_model_call(self) -> None: + assert callable(openai_agents) def test_passes_config_key_user_input_and_context(self) -> None: import inspect - from launchdarkly_ai_openai_agents.handler import openai_agents - sig = inspect.signature(openai_agents) assert "config_key" in sig.parameters assert "user_input" in sig.parameters @@ -768,31 +1057,23 @@ def test_config_key_forwarded_as_key(self) -> None: mock_config_instance.invoke = MagicMock(return_value="result") with patch.object(handler_mod, "config", mock_config_fn): - from launchdarkly_ai_openai_agents.handler import openai_agents - ctx = {"kind": "user", "key": "u1"} openai_agents("my-flag", "hello", ctx) mock_config_fn.assert_called_once() call_kwargs = mock_config_fn.call_args.kwargs assert call_kwargs.get("key") == "my-flag" - handler = call_kwargs.get("handler") - assert handler is not None - assert handler.provides_for == ("OpenAI", "agent") + assert call_kwargs["handler"].provides_for == ("OpenAI", "agent") mock_config_instance.invoke.assert_called_once_with( "hello", ctx, variables=None ) def test_callable_without_extra_kwargs(self) -> None: - import launchdarkly_ai_openai_agents.handler as handler_mod - mock_config_instance = MagicMock() mock_config_fn = MagicMock(return_value=mock_config_instance) mock_config_instance.invoke = MagicMock(return_value="result") with patch.object(handler_mod, "config", mock_config_fn): - from launchdarkly_ai_openai_agents.handler import openai_agents - ctx = {"kind": "user", "key": "u1"} openai_agents("my-flag", "hello", ctx) @@ -803,16 +1084,14 @@ def test_callable_without_extra_kwargs(self) -> None: # --------------------------------------------------------------------------- -# §1.8 Streaming +# Streaming # --------------------------------------------------------------------------- class TestStreaming: - def test_stream_is_defined(self) -> None: - h = create_openai_agent_handler() - assert hasattr(h, "stream") + async def test_stream_is_defined(self) -> None: + assert hasattr(create_openai_agent_handler(), "stream") - @pytest.mark.asyncio async def test_stream_returns_async_generator(self) -> None: import inspect @@ -834,12 +1113,11 @@ async def _empty_stream() -> AsyncIterator[Any]: "importlib.import_module", side_effect=lambda n: agents_mock if n == "agents" else __import__(n), ): - with patch.object(handler_mod, "_HAS_OTEL", False): + with patch.object(spans_mod, "_HAS_OTEL", False): h = create_openai_agent_handler() gen = await h.stream(_make_config(), "hi") assert inspect.isasyncgen(gen) or hasattr(gen, "__aiter__") - @pytest.mark.asyncio async def test_yields_exactly_one_done_event(self) -> None: agents_mock = MagicMock() agents_mock.Agent = MagicMock(return_value=MagicMock()) @@ -859,114 +1137,115 @@ async def _empty_stream() -> AsyncIterator[Any]: "importlib.import_module", side_effect=lambda n: agents_mock if n == "agents" else __import__(n), ): - with patch.object(handler_mod, "_HAS_OTEL", False): + with patch.object(spans_mod, "_HAS_OTEL", False): h = create_openai_agent_handler() events = [e async for e in await h.stream(_make_config(), "hi")] done_events = [e for e in events if e.get("type") == "done"] assert len(done_events) == 1 - -# --------------------------------------------------------------------------- -# §1.5 Streaming telemetry (Appendix A.5 — do not patch _HAS_OTEL=False) -# --------------------------------------------------------------------------- - - -class TestStreamingTelemetry: - def _make_streamed_mock(self) -> Any: - async def _empty_stream() -> AsyncIterator[Any]: - return - yield - - streamed_result = MagicMock() - streamed_result.stream_events = _empty_stream - streamed_result.raw_responses = [] - streamed_result.final_output = "done" - return streamed_result - - @pytest.mark.asyncio - async def test_span_started_during_stream(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - agents_mock = MagicMock() - agents_mock.Agent = MagicMock(return_value=MagicMock()) - agents_mock.tool = MagicMock(side_effect=lambda **kw: lambda fn: fn) - agents_mock.Runner.run_streamed = MagicMock( - return_value=self._make_streamed_mock() - ) - - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - async for _ in await h.stream(_make_config(), "hi"): - pass - - mock_trace.get_tracer.return_value.start_span.assert_called_with( - "openai.agent.run.stream" - ) - - @pytest.mark.asyncio - async def test_ld_span_attributes_set_during_stream(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - agents_mock = MagicMock() - agents_mock.Agent = MagicMock(return_value=MagicMock()) - agents_mock.tool = MagicMock(side_effect=lambda **kw: lambda fn: fn) - agents_mock.Runner.run_streamed = MagicMock( - return_value=self._make_streamed_mock() + async def test_yields_chunks_then_exactly_one_done_event(self) -> None: + turns = [{"deltas": ["Hello", " world"], "output": _text_output("Hello world")}] + agents_mod = _fake_agents_module( + run_streamed=_make_run_streamed(turns, "Hello world") ) - variables = {"__ld": {"configKey": "k", "variationKey": "v", "runId": "r"}} - - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - async for _ in await h.stream( - _make_config(), "hi", None, variables - ): - pass - - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("launchdarkly.operation.type") == "gen_ai" - assert calls.get("launchdarkly.config.key") == "k" - assert calls.get("launchdarkly.variation.key") == "v" - assert calls.get("launchdarkly.run.id") == "r" - - @pytest.mark.asyncio - async def test_span_ended_after_stream_completes(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - agents_mock = MagicMock() - agents_mock.Agent = MagicMock(return_value=MagicMock()) - agents_mock.tool = MagicMock(side_effect=lambda **kw: lambda fn: fn) - agents_mock.Runner.run_streamed = MagicMock( - return_value=self._make_streamed_mock() - ) - - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), + with _patched_agents(agents_mod): + events = [ + e + async for e in await create_openai_agent_handler().stream( + CONFIG, "q", {}, {} + ) + ] + chunks = [e["text"] for e in events if e["type"] == "chunk"] + assert chunks == ["Hello", " world"] + done = [e for e in events if e["type"] == "done"] + assert len(done) == 1 + assert done[0]["output"] == "Hello world" + + async def test_opens_the_same_root_span_name_as_the_blocking_path(self) -> None: + ctx, rec = _recording() + turns = [{"deltas": ["hi"], "output": _text_output("hi")}] + agents_mod = _fake_agents_module(run_streamed=_make_run_streamed(turns, "hi")) + with ctx, _patched_agents(agents_mod): + async for _ in await create_openai_agent_handler().stream( + CONFIG, "q", {}, {} + ): + pass + assert rec.root.name == "invoke_agent" + assert len(rec.named("chat ")) == 1 + + async def test_ends_every_span_once_when_the_stream_completes(self) -> None: + ctx, rec = _recording() + turns = [{"deltas": ["hi"], "output": _text_output("hi")}] + agents_mod = _fake_agents_module(run_streamed=_make_run_streamed(turns, "hi")) + with ctx, _patched_agents(agents_mod): + async for _ in await create_openai_agent_handler().stream( + CONFIG, "q", {}, {} + ): + pass + assert [s.ended for s in rec.spans] == [1] * len(rec.spans) + + async def test_an_abandoned_stream_ends_every_span_but_stays_unset(self) -> None: + from opentelemetry.trace import StatusCode + + ctx, rec = _recording() + turns = [{"deltas": ["one", "two", "three"], "output": _text_output("done")}] + streamed_holder: dict[str, Any] = {} + + def run_streamed(agent: Any, prompt: str, hooks: Any = None, **kw: Any) -> Any: + streamed = FakeStreamedResult(agent, prompt, hooks, turns, "done") + streamed_holder["streamed"] = streamed + return streamed + + agents_mod = _fake_agents_module(run_streamed=run_streamed) + with ctx, _patched_agents(agents_mod): + gen = await create_openai_agent_handler().stream(CONFIG, "q", {}, {}) + async for _ in gen: + break + await gen.aclose() + + assert [s.ended for s in rec.spans] == [1] * len(rec.spans) + assert rec.root.attributes["launchdarkly.stream.abandoned"] is True + assert StatusCode.ERROR not in rec.root.statuses + assert rec.root.exceptions == [] + # The chat span left open by the abandoned turn is also ended cleanly, not failed. + chat = rec.named("chat ")[0] + assert StatusCode.ERROR not in chat.statuses + # Cancelling the vendor's own run is the other half of teardown: ending our span does not + # stop the Runner's background task from spending more tokens. + assert streamed_holder["streamed"].cancelled is True + + async def test_fails_the_spans_when_the_stream_raises(self) -> None: + from opentelemetry.trace import StatusCode + + err = RuntimeError("stream died") + turns = [{"deltas": [], "output": [], "error": err}] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run_streamed=_make_run_streamed(turns, "")) + with ( + ctx, + _patched_agents(agents_mod), + pytest.raises(RuntimeError, match="stream died"), ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_openai_agent_handler() - async for _ in await h.stream(_make_config(), "hi"): - pass - - mock_span.end.assert_called() + async for _ in await create_openai_agent_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) -> None: + turns = [{"deltas": ["hi"], "output": _text_output("hi")}] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run_streamed=_make_run_streamed(turns, "hi")) + with ctx, _patched_agents(agents_mod): + async for _ in await create_openai_agent_handler().stream( + CONFIG, "q", {}, {} + ): + pass + for span in rec.spans: + assert [k for k in span.attributes if k.startswith("gen_ai.prompt")] == [] + assert [n for n, _ in span.events if n.startswith("gen_ai.content")] == [] # --------------------------------------------------------------------------- @@ -976,8 +1255,7 @@ async def test_span_ended_after_stream_completes(self) -> None: class TestOutputFormat: def test_absent_output_format_no_change(self) -> None: - result = build_output_type(None) - assert result is None + assert build_output_type(None) is None def test_output_format_sets_output_type_on_agent(self) -> None: schema = {"type": "object", "properties": {"name": {"type": "string"}}} @@ -996,6 +1274,24 @@ def test_output_format_returns_parsed_object(self) -> None: assert result is not None assert result["schema"]["properties"]["count"]["type"] == "integer" + async def test_output_format_returns_parsed_object_from_result(self) -> None: + agents_mod = _fake_agents_module( + run=_make_run([{"output": _text_output("ignored")}]) + ) + + async def run(agent: Any, prompt: str, hooks: Any = None, **kw: Any) -> Any: + await _drive_turns( + hooks, agent, prompt, [{"output": _text_output("ignored")}] + ) + return FakeRunResult({"score": 9}) + + agents_mod.Runner.run = staticmethod(run) + with _patched_agents(agents_mod): + result = await create_openai_agent_handler()( + _make_config(instructions="x", outputFormat={"type": "object"}), "hello" + ) + assert result["output"] == {"score": 9} + # --------------------------------------------------------------------------- # §1.2 Path C — None user_input must not produce None prompt @@ -1006,7 +1302,6 @@ class TestNoneUserInput: """TESTING.md §1.2 Path C: When user_input is None, the prompt passed to Runner.run must be '' (empty string), not None.""" - @pytest.mark.asyncio async def test_none_user_input_instructions_path_prompt_is_empty_string( self, ) -> None: @@ -1017,7 +1312,10 @@ async def test_none_user_input_instructions_path_prompt_is_empty_string( run_result = _make_run_result("ok") agents_mock = _mock_agents_module(run_result) - async def _spy_run(agent: Any, prompt: Any) -> Any: + # `hooks` is accepted (and ignored) here because `_call_impl` now always passes + # `hooks=hooks` to `Runner.run` — a genuine signature change from the pre-span-work + # handler this test predates, per the assignment's adaptation rule. + async def _spy_run(agent: Any, prompt: Any, hooks: Any = None) -> Any: captured_prompts.append(prompt) return run_result @@ -1027,7 +1325,7 @@ async def _spy_run(agent: Any, prompt: Any) -> Any: "importlib.import_module", side_effect=lambda n: agents_mock if n == "agents" else __import__(n), ): - with patch.object(handler_mod, "_HAS_OTEL", False): + with patch.object(spans_mod, "_HAS_OTEL", False): h = create_openai_agent_handler() await h(_make_config(instructions="Be helpful."), None) @@ -1051,62 +1349,383 @@ class TestHistory: {"role": "assistant", "content": "Feature flagging is a technique..."}, ] - def _mock_agents(self) -> Any: - mock = MagicMock() - mock.Agent = MagicMock(return_value=MagicMock()) - mock.tool = MagicMock(side_effect=lambda **kw: lambda fn: fn) - return mock - def test_history_appended_to_instructions(self) -> None: config = _make_config(instructions="Be concise.") - agents_mock = self._mock_agents() - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - _, _, instructions = _build_agent_and_prompt( - config, "hi", {}, {}, self.SAMPLE_HISTORY - ) + _, _, instructions = _build_agent_and_prompt( + config, "hi", {}, {}, self.SAMPLE_HISTORY + ) assert instructions is not None assert "Conversation History:" in instructions assert "Be concise." in instructions def test_history_format_is_correct(self) -> None: config = _make_config(instructions="Be helpful.") - agents_mock = self._mock_agents() - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - _, _, instructions = _build_agent_and_prompt( - config, "hi", {}, {}, self.SAMPLE_HISTORY - ) + _, _, instructions = _build_agent_and_prompt( + config, "hi", {}, {}, self.SAMPLE_HISTORY + ) assert instructions is not None assert "user: What is feature flagging?" in instructions assert "assistant: Feature flagging is a technique..." in instructions def test_empty_history_treated_like_no_history(self) -> None: config = _make_config(instructions="Be concise.") - agents_mock = self._mock_agents() - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - _, _, instr_with_empty = _build_agent_and_prompt(config, "hi", {}, {}, []) - _, _, instr_without = _build_agent_and_prompt(config, "hi", {}, {}) + _, _, instr_with_empty = _build_agent_and_prompt(config, "hi", {}, {}, []) + _, _, instr_without = _build_agent_and_prompt(config, "hi", {}, {}) assert instr_with_empty == instr_without assert "Conversation History:" not in (instr_with_empty or "") def test_history_without_prior_instructions(self) -> None: config = _make_config() - agents_mock = self._mock_agents() - with patch( - "importlib.import_module", - side_effect=lambda n: agents_mock if n == "agents" else __import__(n), - ): - _, _, instructions = _build_agent_and_prompt( - config, "hi", {}, {}, self.SAMPLE_HISTORY - ) + _, _, instructions = _build_agent_and_prompt( + config, "hi", {}, {}, self.SAMPLE_HISTORY + ) assert instructions is not None assert "Conversation History:" in instructions assert "user: What is feature flagging?" in instructions + + +# --------------------------------------------------------------------------- +# TELEMETRY-CONTRACT.md section 6: what a failed run reports it spent +# --------------------------------------------------------------------------- + + +class _AgentsError(Exception): + """An AgentsException-shaped error: it carries the SDK's own run aggregate. + + MaxTurnsExceeded is the common case, and by definition it happens after paid turns. + """ + + def __init__(self, input_tokens_: int, output_tokens_: int) -> None: + super().__init__("max turns exceeded") + + class _Usage: + # Matches agents.Usage, which is snake_case. + input_tokens = input_tokens_ + output_tokens = output_tokens_ + input_tokens_details = None + + class _Ctx: + usage = _Usage() + + class _RunData: + context_wrapper = _Ctx() + + self.run_data = _RunData() + + +class TestFailedRunUsage: + async def test_reports_the_sdk_aggregate_once_not_twice(self) -> None: + # The hooks already added every completed turn by the time the run raises, and the exception + # carries the SDK's aggregate over those same turns. Adding one to the other roughly doubles + # the reported cost of any run that failed after paid turns. + turns = [ + { + "output": _text_output("one"), + "usage": {"input_tokens": 30, "output_tokens": 5}, + }, + { + "output": _text_output("two"), + "usage": {"input_tokens": 40, "output_tokens": 7}, + }, + ] + + async def run(agent: Any, prompt: str, hooks: Any = None, **kw: Any) -> Any: + await _drive_turns(hooks, agent, prompt, turns) + raise _AgentsError(70, 12) + + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=run) + with ctx, _patched_agents(agents_mod), pytest.raises(_AgentsError): + await create_openai_agent_handler()(CONFIG, "q", {}, {}) + + # 70 and 12 are the aggregate. 140 and 24 would be the aggregate counted twice. + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 70 + assert rec.root.attributes["gen_ai.usage.output_tokens"] == 12 + + async def test_falls_back_to_the_turns_it_saw_when_the_error_carries_nothing( + self, + ) -> None: + # A tool handler's own error propagates unwrapped and has no run_data, so the accumulated + # turns are the only record of what the run spent. + turns = [ + { + "output": _text_output("one"), + "usage": {"input_tokens": 30, "output_tokens": 5}, + } + ] + + async def run(agent: Any, prompt: str, hooks: Any = None, **kw: Any) -> Any: + await _drive_turns(hooks, agent, prompt, turns) + raise RuntimeError("tool exploded") + + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=run) + with ctx, _patched_agents(agents_mod), pytest.raises(RuntimeError): + await create_openai_agent_handler()(CONFIG, "q", {}, {}) + + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 30 + + async def test_writes_nothing_when_the_run_died_before_any_turn(self) -> None: + # All-zero attributes would assert the run cost nothing, which is a different claim. + async def run(agent: Any, prompt: str, hooks: Any = None, **kw: Any) -> Any: + raise RuntimeError("died immediately") + + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=run) + with ctx, _patched_agents(agents_mod), pytest.raises(RuntimeError): + await create_openai_agent_handler()(CONFIG, "q", {}, {}) + + assert "gen_ai.usage.input_tokens" not in rec.root.attributes + + +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. + """ + + async def test_capture_content_reaches_the_factory(self) -> None: + import launchdarkly_ai_openai_agents.handler as handler_mod + + seen: dict[str, Any] = {} + + def _factory(*, capture_content: bool = False) -> 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_agent_handler", _factory), + patch.object(handler_mod, "config", fake_config), + ): + handler_mod.openai_agents("k", "q", {}, 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 fake_config.call_args.kwargs + + async def test_defaults_to_off(self) -> None: + import launchdarkly_ai_openai_agents.handler as handler_mod + + seen: dict[str, Any] = {} + + def _factory(*, capture_content: bool = False) -> 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_agent_handler", _factory), + patch.object(handler_mod, "config", fake_config), + ): + handler_mod.openai_agents("k", "q", {}) + + assert seen["capture_content"] is False + + +class TestUsageIsNotCoupledToSpans: + """Token accounting must not depend on a span existing. + + `on_llm_end` returned early when there was no open chat span, before accumulating the turn's + usage. On an install without the `otel` extra no span is ever created, so every turn's tokens were + dropped and the handler returned zeros. That is a billing figure, not a telemetry one: it is what + the caller reads back and what LaunchDarkly's own metrics record for the AI Config. + """ + + async def test_the_usage_bag_is_right_with_telemetry_switched_off(self) -> None: + import launchdarkly_ai_openai_agents.spans as spans_mod + + turns = [ + {"output": _text_output("a"), "usage": _usage(10, 5)}, + {"output": _text_output("b"), "usage": _usage(7, 2)}, + ] + agents_mod = _fake_agents_module(run=_make_run(turns)) + with patch.object(spans_mod, "_HAS_OTEL", False), _patched_agents(agents_mod): + result = await create_openai_agent_handler()(CONFIG, "q", {}, {}) + assert result["usage"] == {"input_tokens": 17, "output_tokens": 7} + + async def test_the_usage_bag_agrees_with_the_root_span_when_telemetry_is_on( + self, + ) -> None: + # The same two turns, with spans. The bag and the span must not be able to disagree. + turns = [ + {"output": _text_output("a"), "usage": _usage(10, 5)}, + {"output": _text_output("b"), "usage": _usage(7, 2)}, + ] + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_make_run(turns)) + with ctx, _patched_agents(agents_mod): + result = await create_openai_agent_handler()(CONFIG, "q", {}, {}) + assert result["usage"]["input_tokens"] == 17 + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 17 + + +class TestAFailedRunDoesNotClaimItCostNothing: + """A run that died before its first paid call must leave the root's usage attributes absent. + + RunContextWrapper.usage defaults to an empty Usage, so the aggregate the exception carries is + present from the start and reads as a full set of zeros. Writing those says the run cost nothing, + which is different from not knowing what it cost, and a config-scoped cost query cannot tell the + two apart once the zeros are on the span. + """ + + async def test_an_error_before_any_turn_writes_no_usage_attributes(self) -> None: + class _EmptyUsage: + input_tokens = 0 + output_tokens = 0 + input_tokens_details = None + + class _Boom(Exception): + def __init__(self) -> None: + super().__init__("died on the first call") + self.run_data = SimpleNamespace( + context_wrapper=SimpleNamespace(usage=_EmptyUsage()) + ) + + async def _raising_run(*_a: Any, **_k: Any) -> Any: + raise _Boom() + + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_raising_run) + with ctx, _patched_agents(agents_mod), pytest.raises(_Boom): + await create_openai_agent_handler()(CONFIG, "q", {}, {}) + + attrs = rec.root.attributes + assert "gen_ai.usage.input_tokens" not in attrs + assert "gen_ai.usage.output_tokens" not in attrs + assert "gen_ai.usage.total_tokens" not in attrs + + async def test_an_error_after_a_paid_turn_still_reports_that_spend(self) -> None: + # The other side of the same branch: real tokens must survive the failure. + class _RealUsage: + input_tokens = 12 + output_tokens = 4 + input_tokens_details = None + + class _Boom(Exception): + def __init__(self) -> None: + super().__init__("died after a paid turn") + self.run_data = SimpleNamespace( + context_wrapper=SimpleNamespace(usage=_RealUsage()) + ) + + async def _raising_run(*_a: Any, **_k: Any) -> Any: + raise _Boom() + + ctx, rec = _recording() + agents_mod = _fake_agents_module(run=_raising_run) + with ctx, _patched_agents(agents_mod), pytest.raises(_Boom): + await create_openai_agent_handler()(CONFIG, "q", {}, {}) + assert rec.root.attributes["gen_ai.usage.input_tokens"] == 12 + + +class TestCancellationDoesNotDependOnTelemetry: + """Stopping the vendor's run is not telemetry, so no span may gate it. + + The cancel sat inside a `span is not None` guard. Without the `otel` extra there is no root span + at all, so an early consumer break never cancelled the Runner and its background task kept calling + the model and spending money: the exact failure this teardown exists to prevent, reintroduced by + an install choice that has nothing to do with tracing. + """ + + async def test_an_abandoned_stream_cancels_the_runner_with_telemetry_off( + self, + ) -> None: + import launchdarkly_ai_openai_agents.spans as spans_mod + + streamed_holder: dict[str, Any] = {} + # Deltas, so there is something to break out of: a stream with no chunks runs to the end and + # is not abandoned at all. + turns = [{"deltas": ["one", "two", "three"], "output": _text_output("done")}] + + def run_streamed(agent: Any, prompt: Any, hooks: Any = None, **_k: Any) -> Any: + streamed = FakeStreamedResult(agent, prompt, hooks, turns, "done") + streamed_holder["streamed"] = streamed + return streamed + + agents_mod = _fake_agents_module(run_streamed=run_streamed) + with patch.object(spans_mod, "_HAS_OTEL", False), _patched_agents(agents_mod): + gen = await create_openai_agent_handler().stream(CONFIG, "q", {}, {}) + async for _ in gen: + break + await gen.aclose() + + assert streamed_holder["streamed"].cancelled is True + + async def test_a_stream_read_to_the_end_is_not_cancelled(self) -> None: + # The other side of the branch: a completed run must not be cancelled after the fact. + streamed_holder: dict[str, Any] = {} + turns = [{"deltas": ["one", "two"], "output": _text_output("done")}] + + def run_streamed(agent: Any, prompt: Any, hooks: Any = None, **_k: Any) -> Any: + streamed = FakeStreamedResult(agent, prompt, hooks, turns, "done") + streamed_holder["streamed"] = streamed + return streamed + + agents_mod = _fake_agents_module(run_streamed=run_streamed) + with _patched_agents(agents_mod): + async for _ in await create_openai_agent_handler().stream( + CONFIG, "q", {}, {} + ): + pass + + assert streamed_holder["streamed"].cancelled is False + + +class TestInputWritesNeverLeakASpan: + """Serialising the prompt must not be able to strand the root span. + + The input content write ran before the guard that fails the root, so a raise there left it open: + never ended, never exported, so the run disappeared from AI Config Monitoring along with the + feature_flag event it carries. + """ + + async def test_the_blocking_root_still_ends(self) -> None: + from opentelemetry.trace import StatusCode + + import launchdarkly_ai_openai_agents.handler as handler_mod + + ctx, rec = _recording() + agents_mod = _fake_agents_module( + run=_make_run([{"output": _text_output("a"), "usage": _usage(1, 1)}]) + ) + with ( + ctx, + _patched_agents(agents_mod), + patch.object( + handler_mod, + "set_input_content_attributes", + side_effect=TypeError("cannot serialise this prompt"), + ), + pytest.raises(TypeError), + ): + await create_openai_agent_handler(capture_content=True)(CONFIG, "q", {}, {}) + + assert rec.root.ended == 1, "the root span leaked" + assert StatusCode.ERROR in rec.root.statuses + + async def test_the_streaming_root_still_ends(self) -> None: + import launchdarkly_ai_openai_agents.handler as handler_mod + + ctx, rec = _recording() + turns = [{"deltas": ["one"], "output": _text_output("done")}] + agents_mod = _fake_agents_module(run_streamed=_make_run_streamed(turns, "done")) + with ( + ctx, + _patched_agents(agents_mod), + patch.object( + handler_mod, + "set_input_content_attributes", + side_effect=TypeError("cannot serialise this prompt"), + ), + pytest.raises(TypeError), + ): + async for _ in await create_openai_agent_handler( + capture_content=True + ).stream(CONFIG, "q", {}, {}): + pass + + assert rec.root.ended == 1, "the root span leaked"