diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 403c731de..39052cfc8 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -23,4 +23,5 @@ /tests/google_adk_agents/ @temporalio/sdk @temporalio/ai-sdk /tests/langgraph_plugin/ @temporalio/sdk @temporalio/ai-sdk /tests/langsmith_tracing/ @temporalio/sdk @temporalio/ai-sdk +/tests/openai_agents/ @temporalio/sdk @temporalio/ai-sdk /tests/strands_plugin/ @temporalio/sdk @temporalio/ai-sdk diff --git a/openai_agents/README.md b/openai_agents/README.md index 9404278fd..405747b4f 100644 --- a/openai_agents/README.md +++ b/openai_agents/README.md @@ -36,3 +36,4 @@ Each directory contains a complete example with its own README for detailed inst - **[Customer Service](./customer_service/README.md)** - Interactive customer service agent with escalation capabilities, demonstrating conversational workflows. - **[Reasoning Content](./reasoning_content/README.md)** - Example of how to retrieve the thought process of reasoning models. - **[Financial Research Agent](./financial_research_agent/README.md)** - Multi-agent financial research system with planner, search, analyst, writer, and verifier agents collaborating. +- **[Streaming](./streaming/README.md)** - `Runner.run_streamed` with buffered token streaming to external subscribers via `temporalio.contrib.workflow_streams`. **Experimental.** diff --git a/openai_agents/streaming/README.md b/openai_agents/streaming/README.md new file mode 100644 index 000000000..79e379489 --- /dev/null +++ b/openai_agents/streaming/README.md @@ -0,0 +1,122 @@ +# Streaming OpenAI Agents + +> **Experimental.** These samples use the streaming support in +> `temporalio.contrib.openai_agents` together with +> `temporalio.contrib.workflow_streams`. Both are experimental and their APIs +> may change in future versions. + +*Adapted from the [OpenAI Agents SDK basic examples](https://github.com/openai/openai-agents-python/tree/main/examples/basic)* + +Before running these examples, be sure to review the [prerequisites and background on the integration](../README.md). + +The OpenAI Agents SDK streams model output via `Runner.run_streamed`, which +yields events as the model produces them. Inside a Temporal workflow the model +call runs in an activity, so the workflow cannot iterate the live HTTP stream +directly. Instead the plugin runs `model.stream_response()` in a streaming +activity, and that activity publishes each event to the workflow's +[`WorkflowStream`](../../workflow_streams/README.md) so external subscribers +see events as they are produced. + +Publishing is batched: the activity coalesces events over +`ModelActivityParameters.streaming_batch_interval` (default 100ms) before +signalling the workflow. Call this **buffered token streaming** — deltas reach +subscribers within a batch window of being produced, not on every byte. At +typical model speeds one batch carries several tokens, so output arrives in +small bursts rather than glyph-by-glyph. Lower the interval for smoother +output at the cost of more signals. + +Two things to know before reading the samples: + +* `streaming_topic` is **required** for `Runner.run_streamed`. If it is unset, + `run_streamed` raises before scheduling any activity. +* The workflow must host a `WorkflowStream`, constructed in `@workflow.init` so + the publish-signal handler is registered before the activity publishes. + Without one, the publishes are unhandled and silently dropped. + +## Running the Examples + +First, start the worker (supports both examples): + +```bash +uv run openai_agents/streaming/run_worker.py +``` + +Then run either example in another terminal. + +### `stream_text` — buffered text deltas + +Adapted from [`examples/basic/stream_text.py`][upstream-text]. The workflow +just calls `Runner.run_streamed`; the subscriber renders the +`ResponseTextDeltaEvent`s the streaming activity publishes on the `events` +topic. + +Subscribers receive **native OpenAI events** (`TResponseStreamEvent`), because +the activity publishes them straight from `Model.stream_response`. That differs +from `stream_events()` inside the workflow, which yields the agents-SDK +`StreamEvent` union — raw model events arrive there wrapped as +`RawResponsesStreamEvent.data`. + +[upstream-text]: https://github.com/openai/openai-agents-python/blob/main/examples/basic/stream_text.py + +```bash +uv run openai_agents/streaming/run_stream_text_workflow.py +``` + +### `stream_items` — agent-level events with a tool call + +Adapted from [`examples/basic/stream_items.py`][upstream-items]. Renders agent +updates, tool calls, tool outputs, and message outputs as a play-by-play. + +The agents SDK builds those higher-level events from the model output, so they +exist only inside the workflow — the streaming activity never sees them. This +workflow therefore does its own publishing: it iterates +`result.stream_events()` and forwards each event of interest to an `items` +topic as a small serializable `ItemEvent`. (The agents-SDK event types carry +the originating `Agent`, which holds tool callables and so cannot be +serialized.) `stream_events()` resolves a turn at a time — each model call is +one activity — so a multi-turn run like this one reaches the subscriber +progressively rather than in one lump. + +[upstream-items]: https://github.com/openai/openai-agents-python/blob/main/examples/basic/stream_items.py + +```bash +uv run openai_agents/streaming/run_stream_items_workflow.py +``` + +## How it works + +1. The workflow constructs a `WorkflowStream` in `@workflow.init`. +2. `OpenAIAgentsPlugin` is configured with `streaming_topic="events"`, which + routes `Runner.run_streamed` to `invoke_model_activity_streaming`. +3. Inside that activity each event from the live HTTP stream is both collected + (returned to the workflow when the activity completes) and published to the + stream via `WorkflowStreamClient.from_within_activity()`. +4. Just before returning, the workflow publishes a terminator on a separate + `done` topic, then sleeps briefly so the subscriber's next poll can drain + the tail of the stream — the log lives in workflow memory and disappears + when the run completes. +5. External code subscribes with + `WorkflowStreamClient.create(...).subscribe([...], result_type=RawValue)` + and breaks on the terminator. `RawValue` keeps the payloads undecoded so + each topic can be decoded against its own type. If the workflow reaches a + terminal state without publishing a terminator (a failure, say), the + iterator exhausts on its own and the following `handle.result()` raises. + +In the workflow, `stream_events()` resolves only after the model activity +returns, so the workflow itself does not see deltas as they arrive — the +streaming benefit is for external observers. + +## Notes + +* Streaming is incompatible with `use_local_activity=True`: local activities + support neither heartbeats nor the workflow stream signal channel. +* The streaming activity heartbeats on a background task, so set + `heartbeat_timeout` well below `start_to_close_timeout` to detect a stuck + model call early. +* Delivery is at-least-once per activity attempt. An attempt that fails + mid-response leaves its events on the stream and the retry publishes a second + sequence; `stream_events()` in the workflow only sees the final successful + attempt. The [workflow_streams module + documentation](https://github.com/temporalio/sdk-python/blob/main/temporalio/contrib/workflow_streams/README.md) + covers the trade and the conventional `RETRY` event pattern for surfacing it + to consumers. diff --git a/openai_agents/streaming/__init__.py b/openai_agents/streaming/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openai_agents/streaming/activities/__init__.py b/openai_agents/streaming/activities/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openai_agents/streaming/activities/joke_activities.py b/openai_agents/streaming/activities/joke_activities.py new file mode 100644 index 000000000..7fe1c4992 --- /dev/null +++ b/openai_agents/streaming/activities/joke_activities.py @@ -0,0 +1,11 @@ +from __future__ import annotations + +import random + +from temporalio import activity + + +@activity.defn +async def how_many_jokes() -> int: + """Return a random integer of jokes to tell between 1 and 10 (inclusive).""" + return random.randint(1, 10) diff --git a/openai_agents/streaming/run_stream_items_workflow.py b/openai_agents/streaming/run_stream_items_workflow.py new file mode 100644 index 000000000..93698d358 --- /dev/null +++ b/openai_agents/streaming/run_stream_items_workflow.py @@ -0,0 +1,65 @@ +"""Start StreamItemsWorkflow and render its run as a play-by-play.""" + +from __future__ import annotations + +import asyncio +import uuid + +from temporalio.client import Client +from temporalio.common import RawValue +from temporalio.contrib.openai_agents import OpenAIAgentsPlugin +from temporalio.contrib.workflow_streams import WorkflowStreamClient + +from openai_agents.streaming.shared import ( + TASK_QUEUE, + TOPIC_DONE, + TOPIC_ITEMS, + ItemEvent, +) +from openai_agents.streaming.workflows.stream_items_workflow import ( + StreamItemsInput, + StreamItemsWorkflow, +) + + +async def main() -> None: + client = await Client.connect( + "localhost:7233", + plugins=[OpenAIAgentsPlugin()], + ) + + workflow_id = f"stream-items-{uuid.uuid4().hex[:8]}" + handle = await client.start_workflow( + StreamItemsWorkflow.run, + StreamItemsInput(), + id=workflow_id, + task_queue=TASK_QUEUE, + ) + + stream = WorkflowStreamClient.create(client, workflow_id) + converter = client.data_converter.payload_converter + + print("=== Run starting ===") + # result_type=RawValue so the two topics can be decoded per item.topic. + # The raw model events the streaming activity publishes on TOPIC_EVENTS are + # on the stream too; this subscriber just isn't interested in them. + async for item in stream.subscribe([TOPIC_ITEMS, TOPIC_DONE], result_type=RawValue): + if item.topic == TOPIC_DONE: + break + event = converter.from_payload(item.data.payload, ItemEvent) + if event.kind == "agent_updated": + print(f"Agent updated: {event.detail}") + elif event.kind == "tool_call": + print(f"-- Tool was called: {event.detail}") + elif event.kind == "tool_output": + print(f"-- Tool output: {event.detail}") + elif event.kind == "message_output": + print(f"-- Message output:\n {event.detail}") + + result = await handle.result() + print("=== Run complete ===") + print(result) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/openai_agents/streaming/run_stream_text_workflow.py b/openai_agents/streaming/run_stream_text_workflow.py new file mode 100644 index 000000000..6b05a2e44 --- /dev/null +++ b/openai_agents/streaming/run_stream_text_workflow.py @@ -0,0 +1,70 @@ +"""Start StreamTextWorkflow and render its model output as it streams.""" + +from __future__ import annotations + +import asyncio +import uuid +from typing import Any, cast + +from agents.items import TResponseStreamEvent +from openai.types.responses import ResponseTextDeltaEvent +from temporalio.client import Client +from temporalio.common import RawValue +from temporalio.contrib.openai_agents import OpenAIAgentsPlugin +from temporalio.contrib.workflow_streams import WorkflowStreamClient + +from openai_agents.streaming.shared import TASK_QUEUE, TOPIC_DONE, TOPIC_EVENTS +from openai_agents.streaming.workflows.stream_text_workflow import ( + StreamTextInput, + StreamTextWorkflow, +) + +# TResponseStreamEvent is a typing.Annotated union rather than a class, so it +# needs a cast to satisfy from_payload's type[T] signature. The plugin's +# pydantic converter resolves the union's discriminator at runtime. +EVENT_TYPE = cast(type, TResponseStreamEvent) + + +async def main() -> None: + # The plugin's data converter is what decodes the OpenAI event payloads + # published on TOPIC_EVENTS. + client = await Client.connect( + "localhost:7233", + plugins=[OpenAIAgentsPlugin()], + ) + + workflow_id = f"stream-text-{uuid.uuid4().hex[:8]}" + handle = await client.start_workflow( + StreamTextWorkflow.run, + StreamTextInput(prompt="Please tell me 5 jokes."), + id=workflow_id, + task_queue=TASK_QUEUE, + ) + + stream = WorkflowStreamClient.create(client, workflow_id) + converter = client.data_converter.payload_converter + + # A single iterator over both topics — one subscriber, no cancellation race + # between concurrent ones. result_type=RawValue delivers the underlying + # Payload so heterogeneous topics can be decoded per item.topic. The loop + # ends on the in-band terminator, or by the iterator exhausting if the + # workflow reaches a terminal state without publishing one (e.g. on + # failure); either way handle.result() below surfaces the outcome. + async for item in stream.subscribe( + [TOPIC_EVENTS, TOPIC_DONE], result_type=RawValue + ): + if item.topic == TOPIC_DONE: + break + # Subscribers receive native OpenAI events, not the agents-SDK + # StreamEvent wrappers that stream_events() yields in the workflow. + event: Any = converter.from_payload(item.data.payload, EVENT_TYPE) + if isinstance(event, ResponseTextDeltaEvent): + print(event.delta, end="", flush=True) + + result = await handle.result() + print("\n--- final result ---") + print(result) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/openai_agents/streaming/run_worker.py b/openai_agents/streaming/run_worker.py new file mode 100644 index 000000000..99a754702 --- /dev/null +++ b/openai_agents/streaming/run_worker.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import asyncio +import logging +from datetime import timedelta + +from temporalio.client import Client +from temporalio.contrib.openai_agents import ( + ModelActivityParameters, + OpenAIAgentsPlugin, +) +from temporalio.worker import Worker + +from openai_agents.streaming.activities.joke_activities import how_many_jokes +from openai_agents.streaming.shared import TASK_QUEUE, TOPIC_EVENTS +from openai_agents.streaming.workflows.stream_items_workflow import ( + StreamItemsWorkflow, +) +from openai_agents.streaming.workflows.stream_text_workflow import ( + StreamTextWorkflow, +) + + +async def main() -> None: + logging.basicConfig(level=logging.INFO) + client = await Client.connect( + "localhost:7233", + plugins=[ + OpenAIAgentsPlugin( + model_params=ModelActivityParameters( + # The streaming activity heartbeats on a background task, + # so a heartbeat_timeout well under start_to_close_timeout + # detects a stuck model call early. + heartbeat_timeout=timedelta(seconds=10), + start_to_close_timeout=timedelta(minutes=5), + # Required for Runner.run_streamed: the topic the streaming + # activity publishes raw model events to. + streaming_topic=TOPIC_EVENTS, + ), + ), + ], + ) + + worker = Worker( + client, + task_queue=TASK_QUEUE, + workflows=[StreamTextWorkflow, StreamItemsWorkflow], + activities=[how_many_jokes], + ) + await worker.run() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/openai_agents/streaming/shared.py b/openai_agents/streaming/shared.py new file mode 100644 index 000000000..467ec3ce1 --- /dev/null +++ b/openai_agents/streaming/shared.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from dataclasses import dataclass +from datetime import timedelta + +TASK_QUEUE = "openai-agents-streaming-task-queue" + +# Topic the streaming activity publishes raw model stream events to. Must match +# OpenAIAgentsPlugin(model_params=ModelActivityParameters(streaming_topic=...)). +# Events on this topic are native OpenAI `TResponseStreamEvent`s, not the +# agents-SDK `StreamEvent` wrappers that `stream_events()` yields. +TOPIC_EVENTS = "events" + +# Topic the stream_items workflow publishes its own higher-level events to. The +# agents SDK builds those from the model output inside the workflow, so the +# workflow — not the activity — is what publishes them. +TOPIC_ITEMS = "items" + +# Topic the workflow publishes a terminator to once Runner.run_streamed has +# finished. Subscribers watch both topics and break on the terminator, rather +# than racing handle.result() against their next poll. +TOPIC_DONE = "done" + +# How long a workflow holds its run open after publishing the terminator, so a +# subscriber's next poll can drain the tail of the stream. The log lives in +# workflow memory, so it disappears when the run completes. +DRAIN_INTERVAL = timedelta(milliseconds=500) + + +@dataclass +class ItemEvent: + """One step of a run, as published on TOPIC_ITEMS. + + The agents-SDK event types (`RunItemStreamEvent` and friends) carry the + originating `Agent`, which holds tool callables and so is not + serializable. Samples publish their own flattened event instead. + """ + + kind: str + """One of "agent_updated", "tool_call", "tool_output", "message_output".""" + + detail: str diff --git a/openai_agents/streaming/workflows/__init__.py b/openai_agents/streaming/workflows/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openai_agents/streaming/workflows/stream_items_workflow.py b/openai_agents/streaming/workflows/stream_items_workflow.py new file mode 100644 index 000000000..8a1b873d1 --- /dev/null +++ b/openai_agents/streaming/workflows/stream_items_workflow.py @@ -0,0 +1,94 @@ +"""Streaming counterpart to the OpenAI Agents SDK ``stream_items.py`` example. + +Adapted from https://github.com/openai/openai-agents-python/blob/main/examples/basic/stream_items.py + +The upstream example renders higher-level events as they arrive: agent +updates, tool calls, tool outputs, and message outputs. Those are built by the +agents SDK from the model's output, so unlike the raw model events in +``stream_text_workflow`` they exist only inside the workflow — the streaming +activity never sees them. + +So this workflow does the publishing itself: it iterates +``result.stream_events()`` and forwards each interesting event to its own +topic. ``stream_events()`` resolves a turn at a time (each model call is one +activity), so a multi-turn run like this one — model call, tool call, model +call — reaches the subscriber as a play-by-play rather than in one lump. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import timedelta + +from agents import Agent, ItemHelpers, Runner +from temporalio import workflow +from temporalio.contrib import openai_agents as temporal_agents +from temporalio.contrib.workflow_streams import WorkflowStream, WorkflowStreamState + +from openai_agents.streaming.activities.joke_activities import how_many_jokes +from openai_agents.streaming.shared import ( + DRAIN_INTERVAL, + TOPIC_DONE, + TOPIC_ITEMS, + ItemEvent, +) + + +@dataclass +class StreamItemsInput: + prompt: str = "Hello" + # Carries stream state across continue-as-new. None on a fresh start. + stream_state: WorkflowStreamState | None = None + + +@workflow.defn +class StreamItemsWorkflow: + @workflow.init + def __init__(self, input: StreamItemsInput) -> None: + # Construct the stream from @workflow.init so the publish-Signal + # handler is registered before the streaming activity publishes. + self.stream = WorkflowStream(prior_state=input.stream_state) + self.items = self.stream.topic(TOPIC_ITEMS, type=ItemEvent) + self.done = self.stream.topic(TOPIC_DONE, type=bool) + + @workflow.run + async def run(self, input: StreamItemsInput) -> str: + agent = Agent( + name="Joker", + instructions=( + "First call the `how_many_jokes` tool, then tell that many jokes." + ), + tools=[ + temporal_agents.workflow.activity_as_tool( + how_many_jokes, start_to_close_timeout=timedelta(seconds=10) + ) + ], + ) + result = Runner.run_streamed(agent, input=input.prompt) + + messages: list[str] = [] + async for event in result.stream_events(): + if event.type == "agent_updated_stream_event": + self.items.publish( + ItemEvent(kind="agent_updated", detail=event.new_agent.name) + ) + elif event.type == "run_item_stream_event": + item = event.item + if item.type == "tool_call_item": + name = getattr(item.raw_item, "name", "Unknown Tool") + self.items.publish(ItemEvent(kind="tool_call", detail=name)) + elif item.type == "tool_call_output_item": + self.items.publish( + ItemEvent(kind="tool_output", detail=str(item.output)) + ) + elif item.type == "message_output_item": + text = ItemHelpers.text_message_output(item) + messages.append(text) + self.items.publish(ItemEvent(kind="message_output", detail=text)) + + self.done.publish(True) + # Brief pause so the subscriber's next poll can drain the tail of the + # stream — the log lives in workflow memory and is gone once this run + # completes. + await workflow.sleep(DRAIN_INTERVAL) + return "\n\n".join(messages) if messages else result.final_output diff --git a/openai_agents/streaming/workflows/stream_text_workflow.py b/openai_agents/streaming/workflows/stream_text_workflow.py new file mode 100644 index 000000000..dbd797f99 --- /dev/null +++ b/openai_agents/streaming/workflows/stream_text_workflow.py @@ -0,0 +1,79 @@ +"""Streaming counterpart to the OpenAI Agents SDK ``stream_text.py`` example. + +Adapted from https://github.com/openai/openai-agents-python/blob/main/examples/basic/stream_text.py + +The upstream example calls ``Runner.run_streamed`` and iterates raw +``ResponseTextDeltaEvent``s as they arrive over HTTP. Inside a Temporal +workflow the model call runs in an activity, so the workflow cannot iterate +the live HTTP stream directly. The plugin's streaming support runs +``model.stream_response()`` inside the activity and publishes each event to +the workflow's stream, where external subscribers see them as they are +produced. + +The workflow itself only needs to: + +1. host a ``WorkflowStream`` so the streaming activity has somewhere to + publish to; +2. call ``Runner.run_streamed`` (rather than ``Runner.run``) so the agents + framework drives the streaming activity. + +``stream_events()`` inside the workflow resolves only once the activity +returns, so in-workflow consumption is over the final list — not +deltas-as-they-arrive. Streaming is for external observers. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from agents import Agent, Runner +from openai.types.responses import ResponseTextDeltaEvent +from temporalio import workflow +from temporalio.contrib.workflow_streams import WorkflowStream, WorkflowStreamState + +from openai_agents.streaming.shared import DRAIN_INTERVAL, TOPIC_DONE + + +@dataclass +class StreamTextInput: + prompt: str + # Carries stream state across continue-as-new. None on a fresh start. + stream_state: WorkflowStreamState | None = None + + +@workflow.defn +class StreamTextWorkflow: + @workflow.init + def __init__(self, input: StreamTextInput) -> None: + # Construct the stream from @workflow.init so the publish-Signal + # handler is registered before the streaming activity publishes. + # Without a stream, those signals are unhandled and dropped. + self.stream = WorkflowStream(prior_state=input.stream_state) + self.done = self.stream.topic(TOPIC_DONE, type=bool) + + @workflow.run + async def run(self, input: StreamTextInput) -> str: + agent = Agent( + name="Joker", + instructions="You are a helpful assistant.", + ) + result = Runner.run_streamed(agent, input=input.prompt) + + # The workflow only sees these events once the activity returns, so + # the loop just counts them. External subscribers receive them as the + # activity publishes them. + deltas = 0 + async for event in result.stream_events(): + if event.type == "raw_response_event" and isinstance( + event.data, ResponseTextDeltaEvent + ): + deltas += 1 + workflow.logger.info("collected %d delta events", deltas) + + # In-band terminator so the subscriber can stop without racing the + # workflow's completion, then a brief pause to let its next poll + # deliver the tail of the stream — the log lives in workflow memory + # and is gone once this run completes. + self.done.publish(True) + await workflow.sleep(DRAIN_INTERVAL) + return result.final_output diff --git a/tests/openai_agents/__init__.py b/tests/openai_agents/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/openai_agents/_mock_model.py b/tests/openai_agents/_mock_model.py new file mode 100644 index 000000000..58a235c32 --- /dev/null +++ b/tests/openai_agents/_mock_model.py @@ -0,0 +1,150 @@ +"""Scripted streaming model for openai_agents sample tests. + +Each entry in the script drives one ``stream_response`` call — that is, one +model activity: a ``str`` streams that text as deltas followed by a terminal +``ResponseCompletedEvent``, a :class:`ToolCall` emits a function call so the +agent runs the tool and comes back for another turn. + +The plugin takes a ``model_provider`` directly, so nothing needs patching. +""" + +from __future__ import annotations + +import itertools +from collections.abc import AsyncIterator +from dataclasses import dataclass +from typing import Any + +from agents import ( + AgentOutputSchemaBase, + Handoff, + Model, + ModelProvider, + ModelResponse, + ModelSettings, + ModelTracing, + Tool, + TResponseInputItem, +) +from agents.items import TResponseStreamEvent +from openai.types.responses import ( + Response, + ResponseCompletedEvent, + ResponseFunctionToolCall, + ResponseOutputMessage, + ResponseOutputText, + ResponseTextDeltaEvent, +) + +# Chunk size for text deltas. Small enough that any test text streams as +# several events, so a subscriber that reassembles them is doing real work. +_DELTA_CHARS = 12 + + +@dataclass +class ToolCall: + """Script entry for a turn that calls a tool instead of answering.""" + + name: str + arguments: str = "{}" + + +def _message(text: str) -> ResponseOutputMessage: + return ResponseOutputMessage( + id="msg_mock", + content=[ResponseOutputText(text=text, annotations=[], type="output_text")], + role="assistant", + status="completed", + type="message", + ) + + +def _response(output: list[Any]) -> Response: + return Response( + id="resp_mock", + created_at=0.0, + model="mock-model", + object="response", + output=output, + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + status="completed", + ) + + +class ScriptedStreamingModel(Model): + """Model that replays a fixed script of turns, one per streamed call.""" + + def __init__(self, script: list[str | ToolCall]) -> None: + self._script = list(script) + self._calls = itertools.count() + + def _next_turn(self) -> str | ToolCall: + if not self._script: + raise AssertionError("ScriptedStreamingModel script exhausted") + return self._script.pop(0) + + async def get_response(self, *args: Any, **kwargs: Any) -> ModelResponse: + """Unimplemented: this mock exists for Runner.run_streamed.""" + raise NotImplementedError + + async def stream_response( + self, + system_instructions: str | None, + input: str | list[TResponseInputItem], + model_settings: ModelSettings, + tools: list[Tool], + output_schema: AgentOutputSchemaBase | None, + handoffs: list[Handoff], + tracing: ModelTracing, + **kwargs: Any, + ) -> AsyncIterator[TResponseStreamEvent]: + turn = self._next_turn() + seq = itertools.count() + + if isinstance(turn, ToolCall): + call = self._tool_call(turn, next(self._calls)) + yield ResponseCompletedEvent( + response=_response([call]), + sequence_number=next(seq), + type="response.completed", + ) + return + + for start in range(0, len(turn), _DELTA_CHARS): + yield ResponseTextDeltaEvent( + content_index=0, + delta=turn[start : start + _DELTA_CHARS], + item_id="msg_mock", + logprobs=[], + output_index=0, + sequence_number=next(seq), + type="response.output_text.delta", + ) + yield ResponseCompletedEvent( + response=_response([_message(turn)]), + sequence_number=next(seq), + type="response.completed", + ) + + @staticmethod + def _tool_call(turn: ToolCall, index: int) -> ResponseFunctionToolCall: + return ResponseFunctionToolCall( + arguments=turn.arguments, + call_id=f"call_mock_{index}", + name=turn.name, + type="function_call", + id=f"fc_mock_{index}", + status="completed", + ) + + +class ScriptedModelProvider(ModelProvider): + """Hands out one shared model so the script advances across turns.""" + + def __init__(self, script: list[str | ToolCall]) -> None: + self._model = ScriptedStreamingModel(script) + + def get_model(self, model_name: str | None) -> Model: + return self._model diff --git a/tests/openai_agents/streaming_test.py b/tests/openai_agents/streaming_test.py new file mode 100644 index 000000000..6e8edf54f --- /dev/null +++ b/tests/openai_agents/streaming_test.py @@ -0,0 +1,153 @@ +import uuid +from datetime import timedelta +from typing import Any, cast + +from agents.items import TResponseStreamEvent +from openai.types.responses import ResponseTextDeltaEvent +from temporalio.client import Client +from temporalio.common import RawValue +from temporalio.contrib.openai_agents import ModelActivityParameters, OpenAIAgentsPlugin +from temporalio.contrib.workflow_streams import WorkflowStreamClient +from temporalio.worker import Worker + +from openai_agents.streaming.activities.joke_activities import how_many_jokes +from openai_agents.streaming.shared import ( + TOPIC_DONE, + TOPIC_EVENTS, + TOPIC_ITEMS, + ItemEvent, +) +from openai_agents.streaming.workflows.stream_items_workflow import ( + StreamItemsInput, + StreamItemsWorkflow, +) +from openai_agents.streaming.workflows.stream_text_workflow import ( + StreamTextInput, + StreamTextWorkflow, +) +from tests.openai_agents._mock_model import ScriptedModelProvider, ToolCall + +JOKES = ( + "Why did the developer go broke? He used up all his cache. " + "Why do programmers prefer dark mode? Light attracts bugs." +) + +# TResponseStreamEvent is a typing.Annotated union rather than a class, so it +# needs a cast to satisfy from_payload's type[T] signature. +EVENT_TYPE = cast(type, TResponseStreamEvent) + +# Fast polling so a test does not spend most of its time waiting on cooldowns. +POLL_COOLDOWN = timedelta(milliseconds=50) + + +def _client_with_plugin(client: Client, script: list[str | ToolCall]) -> Client: + config = client.config() + config["plugins"] = [ + *config["plugins"], + OpenAIAgentsPlugin( + model_params=ModelActivityParameters( + start_to_close_timeout=timedelta(seconds=30), + heartbeat_timeout=timedelta(seconds=10), + streaming_topic=TOPIC_EVENTS, + ), + model_provider=ScriptedModelProvider(script), + ), + ] + return Client(**config) + + +async def test_stream_text(client: Client) -> None: + client = _client_with_plugin(client, [JOKES]) + task_queue = f"openai-agents-stream-text-{uuid.uuid4()}" + workflow_id = f"stream-text-{uuid.uuid4()}" + + async with Worker( + client, + task_queue=task_queue, + workflows=[StreamTextWorkflow], + max_cached_workflows=0, + ): + handle = await client.start_workflow( + StreamTextWorkflow.run, + StreamTextInput(prompt="Please tell me 2 jokes."), + id=workflow_id, + task_queue=task_queue, + ) + + # Same shape as run_stream_text_workflow.py: one iterator over both + # topics, RawValue payloads decoded per item.topic, break on the + # terminator the workflow publishes. + stream = WorkflowStreamClient.create(client, workflow_id) + converter = client.data_converter.payload_converter + deltas: list[str] = [] + saw_terminator = False + async for item in stream.subscribe( + [TOPIC_EVENTS, TOPIC_DONE], + result_type=RawValue, + poll_cooldown=POLL_COOLDOWN, + ): + if item.topic == TOPIC_DONE: + saw_terminator = True + break + event: Any = converter.from_payload(item.data.payload, EVENT_TYPE) + if isinstance(event, ResponseTextDeltaEvent): + deltas.append(event.delta) + + result = await handle.result() + + assert saw_terminator, "subscriber exited without seeing the terminator" + # Subscribers see native OpenAI events, so the deltas arrive unwrapped and + # reassemble into exactly what the workflow returns. + assert len(deltas) > 1, "expected the text to arrive as several deltas" + assert "".join(deltas) == JOKES + assert result == JOKES + + +async def test_stream_items(client: Client) -> None: + # Turn one calls the tool, turn two answers with the jokes. + client = _client_with_plugin(client, [ToolCall("how_many_jokes"), JOKES]) + task_queue = f"openai-agents-stream-items-{uuid.uuid4()}" + workflow_id = f"stream-items-{uuid.uuid4()}" + + async with Worker( + client, + task_queue=task_queue, + workflows=[StreamItemsWorkflow], + activities=[how_many_jokes], + max_cached_workflows=0, + ): + handle = await client.start_workflow( + StreamItemsWorkflow.run, + StreamItemsInput(), + id=workflow_id, + task_queue=task_queue, + ) + + stream = WorkflowStreamClient.create(client, workflow_id) + converter = client.data_converter.payload_converter + events: list[ItemEvent] = [] + saw_terminator = False + async for item in stream.subscribe( + [TOPIC_ITEMS, TOPIC_DONE], + result_type=RawValue, + poll_cooldown=POLL_COOLDOWN, + ): + if item.topic == TOPIC_DONE: + saw_terminator = True + break + events.append(converter.from_payload(item.data.payload, ItemEvent)) + + result = await handle.result() + + assert saw_terminator, "subscriber exited without seeing the terminator" + assert [e.kind for e in events] == [ + "agent_updated", + "tool_call", + "tool_output", + "message_output", + ] + assert events[0].detail == "Joker" + assert events[1].detail == "how_many_jokes" + assert 1 <= int(events[2].detail) <= 10 + assert events[3].detail == JOKES + assert result == JOKES