diff --git a/packages/client/src/launchdarkly_ai_server/responses.py b/packages/client/src/launchdarkly_ai_server/responses.py new file mode 100644 index 00000000..ef2665e1 --- /dev/null +++ b/packages/client/src/launchdarkly_ai_server/responses.py @@ -0,0 +1,57 @@ +"""Map OpenAI Responses stream events onto the AG-UI text and tool subset. + +Both OpenAI handler packages call :func:`map_responses_event`. They only +offer function tools, so hosted and MCP item types are ignored until those +tools are wired. +""" + +from __future__ import annotations + +from typing import Any + +from .agui import TurnEmitter + + +def _call_id(event: Any, item_to_call: dict[str, str]) -> str: + direct = getattr(event, "call_id", None) + if isinstance(direct, str) and direct: + return direct + item_id = getattr(event, "item_id", None) + if isinstance(item_id, str) and item_id in item_to_call: + return item_to_call[item_id] + return item_id if isinstance(item_id, str) else "" + + +def map_responses_event( + event: Any, emitter: TurnEmitter, item_to_call: dict[str, str] +) -> list[dict[str, Any]]: + """Map one Responses event. Unknown events, including hosted tools, are ignored.""" + kind = getattr(event, "type", None) + if kind == "response.created": + return emitter.end_text() + if kind == "response.output_text.delta": + delta = getattr(event, "delta", "") + if not isinstance(delta, str) or not delta: + return [] + return emitter.text(delta) + if kind == "response.output_item.added": + item = getattr(event, "item", None) + if getattr(item, "type", None) != "function_call": + return [] + call_id = getattr(item, "call_id", None) or getattr(item, "id", None) or "" + name = getattr(item, "name", None) or "" + item_id = getattr(item, "id", None) + if isinstance(item_id, str) and isinstance(call_id, str) and call_id: + item_to_call[item_id] = call_id + return emitter.tool_start(str(call_id), str(name)) + if kind == "response.function_call_arguments.delta": + delta = getattr(event, "delta", "") + if not isinstance(delta, str): + return [] + return emitter.tool_args(_call_id(event, item_to_call), delta) + if kind == "response.function_call_arguments.done": + return emitter.tool_end( + _call_id(event, item_to_call), + arguments=getattr(event, "arguments", None), + ) + return [] diff --git a/packages/client/tests/test_agui.py b/packages/client/tests/test_agui.py index 915e6944..68613a08 100644 --- a/packages/client/tests/test_agui.py +++ b/packages/client/tests/test_agui.py @@ -149,6 +149,16 @@ def test_fail_open_masks_a_call_that_never_finished(self) -> None: assert events[-1]["content"] == "" assert events[-1]["metadata"]["launchdarkly"]["toolFailed"] is True + def test_tool_end_extends_a_streamed_argument_prefix(self) -> None: + emitter = TurnEmitter(capture_content=True) + events = emitter.tool_start("a", "search") + events.extend(emitter.tool_args("a", '{"q":')) + events.extend(emitter.tool_end("a", arguments={"q": "x"})) + deltas = [ + event["delta"] for event in events if event["type"] == "TOOL_CALL_ARGS" + ] + assert "".join(deltas) == '{"q":"x"}' + def test_a_second_start_for_the_same_id_is_dropped(self) -> None: emitter = TurnEmitter(capture_content=True) assert emitter.tool_start("a", "search") 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 d840e019..4a74a9eb 100644 --- a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py +++ b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py @@ -45,7 +45,9 @@ set_tool_call_content_attributes, text_message, ) +from launchdarkly_ai_server.agui import TurnEmitter from launchdarkly_ai_server.lifecycle import effective_capture_content +from launchdarkly_ai_server.responses import map_responses_event from ._version import PACKAGE_NAME, __version__ from .spans import ( @@ -250,10 +252,16 @@ def __init__( self.run_usage = run_usage self.open_model_span: Any = None self.open_tool_spans: dict[str, Any] = {} + self.emitter = TurnEmitter(capture_content=capture_content) + self.pending: list[dict[str, Any]] = [] + self.item_to_call: dict[str, str] = {} async def on_llm_start( self, context: Any, agent: Any, system_prompt: str | None, input_items: Any ) -> None: + # Text is closed from response.created, in the order the consumer + # reads raw events. Closing here runs ahead of unread deltas and + # splits one turn across two messages. span = start_model_span(self.config, self.parent) self.open_model_span = span if self.capture_content: @@ -320,6 +328,16 @@ async def on_tool_start(self, context: Any, agent: Any, tool: Any) -> None: # abandon_open_spans both walk this dict, so a span missing from it never ends and never # exports. self.open_tool_spans[str(call_id)] = span + self.pending.extend(self.emitter.tool_start(str(call_id), str(name))) + # The Runner can invoke this hook from its background task while + # argument deltas are still unread. Closing now would drop the rest + # and skip the full argument string, because a prefix is already + # recorded. arguments.done closes a call that is still streaming. + if not self.emitter.has_argument_deltas(str(call_id)): + tool_args = getattr(context, "tool_arguments", None) + self.pending.extend( + self.emitter.tool_end(str(call_id), arguments=tool_args) + ) if self.capture_content: # `tool_arguments` is the raw JSON args string the model produced. Parsed here rather # than passed through, per TELEMETRY-CONTRACT.md section 12: arguments hold the object @@ -336,7 +354,19 @@ async def on_tool_start(self, context: Any, agent: Any, tool: Any) -> None: async def on_tool_end( self, context: Any, agent: Any, tool: Any, result: Any ) -> None: - span = self.open_tool_spans.pop(self._tool_span_key(context, tool), None) + call_id = self._tool_span_key(context, tool) + # The hook can run before unread argument deltas. Close with the full + # argument string first, so tool_result does not seal a prefix. + # Text stays open: this runs before unread text deltas, and closing + # here would start a second message for the rest of the turn. + self.pending.extend( + self.emitter.tool_end( + str(call_id), + arguments=getattr(context, "tool_arguments", None), + ) + ) + self.pending.extend(self.emitter.tool_result(str(call_id), str(result))) + span = self.open_tool_spans.pop(call_id, None) if span is None: return # Once popped, ending it is this method's job alone: nothing else knows the span exists. A @@ -383,6 +413,26 @@ def abandon_open_spans(self, ended: set[int], cancelled: bool = False) -> None: self.open_model_span = None +def _take_finished_tools(hooks: _SpanningHooks) -> list[dict[str, Any]]: + """Return tool events recorded by the hook, text closed before each result. + + The hook runs before unread deltas from the same model call. The caller + applies those deltas first. Tool results then follow ``TEXT_MESSAGE_END`` + so the client does not attach them to the next message. + """ + if not hooks.pending: + return [] + earlier: list[dict[str, Any]] = [] + results: list[dict[str, Any]] = [] + for event in hooks.pending: + if event.get("type") == "TOOL_CALL_RESULT": + results.append(event) + else: + earlier.append(event) + hooks.pending.clear() + return [*earlier, *hooks.emitter.end_text(), *results] + + def _write_failed_run_usage( span: Any, config: AiConfigRep, @@ -654,16 +704,28 @@ async def _stream_gen( full_output = "" async for event in streamed.stream_events(): - if event.type == "raw_response_event": - raw = getattr(event, "data", None) - if ( - raw is not None - and getattr(raw, "type", None) == "response.output_text.delta" - ): - delta = getattr(raw, "delta", "") - if isinstance(delta, str) and delta: - yield {"type": "chunk", "text": delta} - full_output += delta + raw = ( + getattr(event, "data", None) + if event.type == "raw_response_event" + else None + ) + # The next model call is queued only after this turn's tools have + # run. Unread deltas from this turn are still ahead of that event, + # so the tool result waits until they have been applied. + if getattr(raw, "type", None) == "response.created": + for rich in _take_finished_tools(hooks): + yield rich + if raw is not None: + for rich in map_responses_event(raw, hooks.emitter, hooks.item_to_call): + if rich.get("type") == "TEXT_MESSAGE_CONTENT": + full_output += rich.get("delta") or "" + yield rich + for rich in _take_finished_tools(hooks): + yield rich + for rich in hooks.emitter.end_text(): + yield rich + for rich in hooks.emitter.fail_open(): + yield rich stream_completed = True final_output = streamed.final_output @@ -697,6 +759,12 @@ async def _stream_gen( hooks.close_open_spans(exc) _write_failed_run_usage(span, config, exc, run_usage) fail_span(span, exc, ended) + for rich in _take_finished_tools(hooks): + yield rich + for rich in hooks.emitter.end_text(): + yield rich + for rich in hooks.emitter.fail_open(): + yield rich raise finally: # A no-op on the success and failure paths: both already ended their spans through diff --git a/packages/openai-agents/tests/test_handler.py b/packages/openai-agents/tests/test_handler.py index e9eb7f40..f6d83bf2 100644 --- a/packages/openai-agents/tests/test_handler.py +++ b/packages/openai-agents/tests/test_handler.py @@ -117,6 +117,10 @@ async def stream_events(self) -> AsyncIterator[Any]: await self.hooks.on_llm_start( MagicMock(), self.agent, self.agent.instructions, self.prompt ) + yield SimpleNamespace( + type="raw_response_event", + data=SimpleNamespace(type="response.created"), + ) for delta in turn.get("deltas", []): yield SimpleNamespace( type="raw_response_event", @@ -135,6 +139,13 @@ async def stream_events(self) -> AsyncIterator[Any]: await self.hooks.on_tool_end( ctx, self.agent, tool_obj, call.get("result", "ok") ) + for delta in turn.get("late_deltas", []): + yield SimpleNamespace( + type="raw_response_event", + data=SimpleNamespace( + type="response.output_text.delta", delta=delta + ), + ) if turn.get("error") is not None: raise turn["error"] response = SimpleNamespace( @@ -1387,12 +1398,181 @@ async def test_yields_chunks_then_exactly_one_done_event(self) -> None: CONFIG, "q", {}, {} ) ] - chunks = [e["text"] for e in events if e["type"] == "chunk"] - assert chunks == ["Hello", " world"] + deltas = [e["delta"] for e in events if e["type"] == "TEXT_MESSAGE_CONTENT"] + assert deltas == ["Hello", " world"] + assert events[0]["type"] == "TEXT_MESSAGE_START" + assert any(e["type"] == "TEXT_MESSAGE_END" for e in events) done = [e for e in events if e["type"] == "done"] assert len(done) == 1 assert done[0]["output"] == "Hello world" + async def test_tool_lifecycle_hides_payloads_until_capture_is_on(self) -> None: + turns = [ + { + "deltas": ["Looking"], + "tool_calls": [ + { + "name": "search", + "id": "c1", + "args": {"q": "x"}, + "result": "secret", + } + ], + "output": [], + } + ] + + async def _events(capture: bool) -> list[dict[str, Any]]: + agents_mod = _fake_agents_module( + run_streamed=_make_run_streamed(turns, "done") + ) + with _patched_agents(agents_mod): + return [ + event + async for event in await create_openai_agent_handler( + capture_content=capture + ).stream(CONFIG, "q", {}, {}) + ] + + hidden = await _events(False) + start = next(event for event in hidden if event["type"] == "TOOL_CALL_START") + assert start["toolCallId"] == "c1" + assert start["toolCallName"] == "search" + assert not any(event["type"] == "TOOL_CALL_ARGS" for event in hidden) + result = next(event for event in hidden if event["type"] == "TOOL_CALL_RESULT") + assert result["content"] == "" + + shown = await _events(True) + args = next(event for event in shown if event["type"] == "TOOL_CALL_ARGS") + assert args["delta"] == json.dumps({"q": "x"}) + shown_result = next( + event for event in shown if event["type"] == "TOOL_CALL_RESULT" + ) + assert shown_result["content"] == "secret" + types = [event["type"] for event in shown] + assert types.index("TEXT_MESSAGE_CONTENT") < types.index("TEXT_MESSAGE_END") + assert types.index("TEXT_MESSAGE_END") < types.index("TOOL_CALL_RESULT") + + async def test_unread_text_stays_on_the_message_the_tool_result_closes( + self, + ) -> None: + turns = [ + { + "deltas": ["Hel"], + "tool_calls": [ + {"name": "search", "id": "c1", "args": {}, "result": "found"} + ], + "late_deltas": ["lo"], + "output": [], + } + ] + agents_mod = _fake_agents_module(run_streamed=_make_run_streamed(turns, "done")) + with _patched_agents(agents_mod): + events = [ + event + async for event in await create_openai_agent_handler( + capture_content=True + ).stream(CONFIG, "q", {}, {}) + ] + contents = [ + event for event in events if event["type"] == "TEXT_MESSAGE_CONTENT" + ] + assert [event["delta"] for event in contents] == ["Hel", "lo"] + assert contents[0]["messageId"] == contents[1]["messageId"] + types = [event["type"] for event in events] + assert types.index("TEXT_MESSAGE_END") < types.index("TOOL_CALL_RESULT") + assert types.index("TOOL_CALL_START") < types.index("TEXT_MESSAGE_END") + + async def test_a_streaming_prefix_is_not_closed_by_the_tool_hook(self) -> None: + from launchdarkly_ai_openai_agents.handler import _SpanningHooks + from launchdarkly_ai_server import create_run_usage + + hooks = _SpanningHooks(CONFIG, None, True, create_run_usage()) + hooks.emitter.tool_start("c1", "search") + hooks.emitter.tool_args("c1", '{"q":') + ctx = SimpleNamespace( + tool_name="search", + tool_call_id="c1", + tool_arguments='{"q":"x"}', + ) + await hooks.on_tool_start( + ctx, SimpleNamespace(), SimpleNamespace(name="search") + ) + assert not any(event["type"] == "TOOL_CALL_END" for event in hooks.pending) + rest = hooks.emitter.tool_args("c1", '"}') + assert rest[0]["delta"] == '"}' + assert any( + event["type"] == "TOOL_CALL_END" for event in hooks.emitter.tool_end("c1") + ) + + async def test_tool_end_publishes_arguments_the_consumer_has_not_read(self) -> None: + from launchdarkly_ai_openai_agents.handler import _SpanningHooks + from launchdarkly_ai_server import create_run_usage + + hooks = _SpanningHooks(CONFIG, None, True, create_run_usage()) + hooks.emitter.tool_start("c1", "search") + hooks.emitter.tool_args("c1", '{"q":') + ctx = SimpleNamespace( + tool_name="search", + tool_call_id="c1", + tool_arguments='{"q":"x"}', + ) + await hooks.on_tool_end( + ctx, SimpleNamespace(), SimpleNamespace(name="search"), "found" + ) + deltas = [ + event["delta"] + for event in hooks.pending + if event["type"] == "TOOL_CALL_ARGS" + ] + assert "".join(['{"q":', *deltas]) == '{"q":"x"}' + assert any(event["type"] == "TOOL_CALL_RESULT" for event in hooks.pending) + assert hooks.emitter.tool_args("c1", "more") == [] + + async def test_a_tool_hook_leaves_text_open_until_the_deltas_are_read(self) -> None: + from launchdarkly_ai_openai_agents.handler import ( + _SpanningHooks, + _take_finished_tools, + ) + from launchdarkly_ai_server import create_run_usage + + hooks = _SpanningHooks(CONFIG, None, True, create_run_usage()) + opened = hooks.emitter.text("Hel") + ctx = SimpleNamespace( + tool_name="search", + tool_call_id="c1", + tool_arguments="{}", + ) + await hooks.on_tool_start( + ctx, SimpleNamespace(), SimpleNamespace(name="search") + ) + await hooks.on_tool_end( + ctx, SimpleNamespace(), SimpleNamespace(name="search"), "found" + ) + assert not any(event["type"] == "TEXT_MESSAGE_END" for event in hooks.pending) + more = hooks.emitter.text("lo") + assert more[0]["type"] == "TEXT_MESSAGE_CONTENT" + assert more[0]["messageId"] == opened[0]["messageId"] + released = _take_finished_tools(hooks) + types = [event["type"] for event in released] + assert types.index("TOOL_CALL_START") < types.index("TEXT_MESSAGE_END") + assert types.index("TEXT_MESSAGE_END") < types.index("TOOL_CALL_RESULT") + + def test_a_new_response_closes_text_in_read_order(self) -> None: + from launchdarkly_ai_server.agui import TurnEmitter + from launchdarkly_ai_server.responses import map_responses_event + + emitter = TurnEmitter(capture_content=True) + opened = emitter.text("Let me ") + closed = map_responses_event( + SimpleNamespace(type="response.created"), emitter, {} + ) + assert closed[0]["type"] == "TEXT_MESSAGE_END" + assert closed[0]["messageId"] == opened[0]["messageId"] + nxt = emitter.text("Answer") + assert nxt[0]["type"] == "TEXT_MESSAGE_START" + assert nxt[0]["messageId"] != opened[0]["messageId"] + 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")}] @@ -1465,6 +1645,25 @@ async def test_fails_the_spans_when_the_stream_raises(self) -> None: assert StatusCode.ERROR in rec.root.statuses assert rec.root.ended == 1 + async def test_a_failed_turn_closes_the_open_text_message(self) -> None: + err = RuntimeError("next turn failed") + turns = [{"deltas": ["Let me "], "error": err}] + agents_mod = _fake_agents_module(run_streamed=_make_run_streamed(turns, "")) + seen: list[dict[str, Any]] = [] + with ( + _patched_agents(agents_mod), + pytest.raises(RuntimeError, match="next turn failed"), + ): + async for event in await create_openai_agent_handler().stream( + CONFIG, "q", {}, {} + ): + seen.append(event) + starts = [event for event in seen if event["type"] == "TEXT_MESSAGE_START"] + ends = [event for event in seen if event["type"] == "TEXT_MESSAGE_END"] + assert len(starts) == 1 + assert len(ends) == 1 + assert ends[0]["messageId"] == starts[0]["messageId"] + async def test_emits_no_content_by_default_on_the_streaming_path(self) -> None: turns = [{"deltas": ["hi"], "output": _text_output("hi")}] ctx, rec = _recording() diff --git a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py index 9d59a2d6..7dcaf91d 100644 --- a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py +++ b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py @@ -27,7 +27,9 @@ set_output_content_attributes, set_tool_call_content_attributes, ) +from launchdarkly_ai_server.agui import TurnEmitter from launchdarkly_ai_server.lifecycle import effective_capture_content +from launchdarkly_ai_server.responses import map_responses_event from ._version import PACKAGE_NAME, __version__ from .spans import ( @@ -476,6 +478,7 @@ async def _stream_gen( # stream; a CancelledError means something cancelled the run, usually a timeout, and the # consumer chose nothing. The blocking path already tells these apart. cancelled = False + emitter = TurnEmitter(capture_content=capture_content) try: tools = _build_tools(config.get("tools") or {}) input_messages = _build_input_messages(config, user_input, variables, history) @@ -493,6 +496,7 @@ async def _stream_gen( previous_response_id: str | None = None current_input: Any = input_messages steps = 0 + item_to_call: dict[str, str] = {} while True: model_span = start_model_span(config, parent) @@ -523,10 +527,10 @@ async def _stream_gen( stream = client.responses.stream(**stream_params) async with stream as s: async for event in s: - if getattr(event, "type", None) == "response.output_text.delta": - text = getattr(event, "delta", "") - full_output += text - yield {"type": "chunk", "text": text} + for rich in map_responses_event(event, emitter, item_to_call): + if rich.get("type") == "TEXT_MESSAGE_CONTENT": + full_output += rich.get("delta") or "" + yield rich final_resp = await s.get_final_response() @@ -556,6 +560,10 @@ async def _stream_gen( if getattr(item, "type", None) == "function_call" ] if not tool_calls: + for rich in emitter.end_text(): + yield rich + for rich in emitter.fail_open(): + yield rich break if steps >= _MAX_STEPS: @@ -565,6 +573,13 @@ async def _stream_gen( steps += 1 previous_response_id = getattr(final_resp, "id", None) + for tc in tool_calls: + for rich in emitter.tool_start(str(tc.call_id), str(tc.name)): + yield rich + for rich in emitter.tool_end(str(tc.call_id), arguments=tc.arguments): + yield rich + for rich in emitter.end_text(): + yield rich tool_outputs = [] for tc in tool_calls: tool_span = start_tool_span(tc.name, tc.call_id, parent) @@ -597,6 +612,9 @@ async def _stream_gen( # Cleared on both paths that end the span, and deliberately not in a `finally`: # a `finally` would also clear it for a BaseException, which is the one case where # the span is still open and the outer `finally` is the only thing left to close it. + # The model is sent str(result). The stream reports that same text. + for rich in emitter.tool_result(str(tc.call_id), str(result)): + yield rich tool_outputs.append( { "type": "function_call_output", @@ -631,6 +649,10 @@ async def _stream_gen( if run_usage.reported: finish_root_span(span, last_response_model, run_usage.total) fail_span(span, exc, ended) + for rich in emitter.end_text(): + yield rich + for rich in emitter.fail_open(): + yield rich raise finally: # A no-op on the success and failure paths, because both already ended their spans through diff --git a/packages/openai-messages/tests/test_handler.py b/packages/openai-messages/tests/test_handler.py index 564765a7..6ae392de 100644 --- a/packages/openai-messages/tests/test_handler.py +++ b/packages/openai-messages/tests/test_handler.py @@ -1281,10 +1281,279 @@ async def test_yields_chunk_events(self, mock_openai: MagicMock) -> None: with patch.object(spans_mod, "_HAS_OTEL", False): h = create_openai_messages_handler() events = [e async for e in await h.stream(CONFIG, "q")] - chunks = [e for e in events if e.get("type") == "chunk"] - assert len(chunks) == 2 - assert chunks[0]["text"] == "hello " - assert chunks[1]["text"] == "world" + deltas = [e["delta"] for e in events if e.get("type") == "TEXT_MESSAGE_CONTENT"] + assert deltas == ["hello ", "world"] + assert events[0]["type"] == "TEXT_MESSAGE_START" + assert events[-2]["type"] == "TEXT_MESSAGE_END" + assert events[-1]["type"] == "done" + + async def test_tool_lifecycle_hides_payloads_until_capture_is_on( + self, mock_openai: MagicMock + ) -> None: + import launchdarkly_ai_openai_messages.spans as spans_mod + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + tool_call = _function_call_item("search", "c1", {"q": "x"}) + + async def search(args: dict[str, Any]) -> str: + return "secret" + + async def _events(capture: bool) -> list[dict[str, Any]]: + mock_openai.responses.stream = MagicMock( + side_effect=[ + _make_openai_stream_context(["Looking"], output=[tool_call]), + _make_openai_stream_context(["Done"]), + ] + ) + with patch.object(spans_mod, "_HAS_OTEL", False): + handler = create_openai_messages_handler(capture_content=capture) + return [ + event + async for event in await handler.stream( + CONFIG, "q", {"search": search} + ) + ] + + hidden = await _events(False) + start = next(event for event in hidden if event["type"] == "TOOL_CALL_START") + assert start["toolCallId"] == "c1" + assert start["toolCallName"] == "search" + assert start["parentMessageId"] + assert not any(event["type"] == "TOOL_CALL_ARGS" for event in hidden) + result = next(event for event in hidden if event["type"] == "TOOL_CALL_RESULT") + assert result["content"] == "" + + shown = await _events(True) + args = next(event for event in shown if event["type"] == "TOOL_CALL_ARGS") + assert args["delta"] == json.dumps({"q": "x"}) + shown_result = next( + event for event in shown if event["type"] == "TOOL_CALL_RESULT" + ) + assert shown_result["content"] == "secret" + + async def test_a_failed_call_closes_its_siblings( + self, mock_openai: MagicMock + ) -> None: + import launchdarkly_ai_openai_messages.spans as spans_mod + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + calls = [ + _function_call_item("boom", "c1", {}), + _function_call_item("later", "c2", {}), + ] + mock_openai.responses.stream = MagicMock( + return_value=_make_openai_stream_context(["x"], output=calls) + ) + + def boom(_args: dict[str, Any]) -> str: + raise RuntimeError("nope") + + with patch.object(spans_mod, "_HAS_OTEL", False): + handler = create_openai_messages_handler() + seen: list[dict[str, Any]] = [] + with pytest.raises(RuntimeError, match="nope"): + async for event in await handler.stream( + CONFIG, "q", {"boom": boom, "later": lambda _args: "ok"} + ): + seen.append(event) + results = { + event["toolCallId"] for event in seen if event["type"] == "TOOL_CALL_RESULT" + } + assert results == {"c1", "c2"} + + def test_parallel_argument_deltas_stay_on_their_call(self) -> None: + from types import SimpleNamespace + + from launchdarkly_ai_server.agui import TurnEmitter + from launchdarkly_ai_server.responses import map_responses_event + + emitter = TurnEmitter(capture_content=True) + ids: dict[str, str] = {} + events: list[dict[str, Any]] = [] + + def feed(event: Any) -> None: + events.extend(map_responses_event(event, emitter, ids)) + + feed( + SimpleNamespace( + type="response.output_item.added", + item=SimpleNamespace( + type="function_call", id="fc_1", call_id="call_1", name="search" + ), + ) + ) + feed( + SimpleNamespace( + type="response.output_item.added", + item=SimpleNamespace( + type="function_call", id="fc_2", call_id="call_2", name="other" + ), + ) + ) + feed( + SimpleNamespace( + type="response.function_call_arguments.delta", + item_id="fc_1", + delta='{"q":', + ) + ) + feed( + SimpleNamespace( + type="response.function_call_arguments.delta", + item_id="fc_2", + delta='{"n":', + ) + ) + feed( + SimpleNamespace( + type="response.function_call_arguments.delta", + item_id="fc_1", + delta='"x"}', + ) + ) + feed( + SimpleNamespace( + type="response.function_call_arguments.delta", + item_id="fc_2", + delta="1}", + ) + ) + by_id: dict[str, str] = {} + for event in events: + if event["type"] != "TOOL_CALL_ARGS": + continue + by_id[event["toolCallId"]] = ( + by_id.get(event["toolCallId"], "") + event["delta"] + ) + assert by_id == {"call_1": '{"q":"x"}', "call_2": '{"n":1}'} + + async def test_an_opened_call_without_a_function_result_is_closed( + self, mock_openai: MagicMock + ) -> None: + from types import SimpleNamespace + + import launchdarkly_ai_openai_messages.spans as spans_mod + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + added = SimpleNamespace( + type="response.output_item.added", + item=SimpleNamespace( + type="function_call", call_id="c1", id="item-1", name="search" + ), + ) + + class _Stream: + def __aiter__(self) -> AsyncIterator[Any]: + return self._iter() + + async def _iter(self) -> AsyncIterator[Any]: + yield added + + async def get_final_response(self) -> Any: + final = MagicMock() + final.output = [] + final.status = "completed" + final.incomplete_details = None + final.usage = MagicMock(input_tokens=1, output_tokens=1) + final.usage.input_tokens_details = MagicMock(cached_tokens=0) + final.id = "resp-stream" + final.model = "gpt-4o" + return final + + @asynccontextmanager + async def _mgr() -> AsyncGenerator[Any, None]: + yield _Stream() + + mock_openai.responses.stream = MagicMock(return_value=_mgr()) + with patch.object(spans_mod, "_HAS_OTEL", False): + handler = create_openai_messages_handler() + events = [event async for event in await handler.stream(CONFIG, "q")] + assert any( + event["type"] == "TOOL_CALL_RESULT" and event["toolCallId"] == "c1" + for event in events + ) + + async def test_a_provider_error_closes_an_opened_call( + self, mock_openai: MagicMock + ) -> None: + from types import SimpleNamespace + + import launchdarkly_ai_openai_messages.spans as spans_mod + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + added = SimpleNamespace( + type="response.output_item.added", + item=SimpleNamespace( + type="function_call", call_id="c1", id="item-1", name="search" + ), + ) + + class _Stream: + def __aiter__(self) -> AsyncIterator[Any]: + return self._iter() + + async def _iter(self) -> AsyncIterator[Any]: + yield added + raise RuntimeError("stream broke") + + async def get_final_response(self) -> Any: + raise AssertionError("final response is not read") + + @asynccontextmanager + async def _mgr() -> AsyncGenerator[Any, None]: + yield _Stream() + + mock_openai.responses.stream = MagicMock(return_value=_mgr()) + seen: list[dict[str, Any]] = [] + with ( + patch.object(spans_mod, "_HAS_OTEL", False), + pytest.raises(RuntimeError, match="stream broke"), + ): + handler = create_openai_messages_handler() + async for event in await handler.stream(CONFIG, "q"): + seen.append(event) + assert any( + event["type"] == "TOOL_CALL_RESULT" and event["toolCallId"] == "c1" + for event in seen + ) + + async def test_a_provider_error_closes_the_open_text( + self, mock_openai: MagicMock + ) -> None: + from types import SimpleNamespace + + import launchdarkly_ai_openai_messages.spans as spans_mod + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + delta = SimpleNamespace(type="response.output_text.delta", delta="Hello") + + class _Stream: + def __aiter__(self) -> AsyncIterator[Any]: + return self._iter() + + async def _iter(self) -> AsyncIterator[Any]: + yield delta + raise RuntimeError("stream broke") + + async def get_final_response(self) -> Any: + raise AssertionError("final response is not read") + + @asynccontextmanager + async def _mgr() -> AsyncGenerator[Any, None]: + yield _Stream() + + mock_openai.responses.stream = MagicMock(return_value=_mgr()) + seen: list[dict[str, Any]] = [] + with ( + patch.object(spans_mod, "_HAS_OTEL", False), + pytest.raises(RuntimeError, match="stream broke"), + ): + handler = create_openai_messages_handler() + async for event in await handler.stream(CONFIG, "q"): + seen.append(event) + start = next(event for event in seen if event["type"] == "TEXT_MESSAGE_START") + end = next(event for event in seen if event["type"] == "TEXT_MESSAGE_END") + assert end["messageId"] == start["messageId"] async def test_yields_exactly_one_done_event(self, mock_openai: MagicMock) -> None: import launchdarkly_ai_openai_messages.spans as spans_mod