openai-agents 0.0.11__py3-none-any.whl → 0.0.12__py3-none-any.whl

agents/models/chatcmpl_stream_handler.py

@@ -0,0 +1,290 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from dataclasses import dataclass, field
+
+from openai import AsyncStream
+from openai.types.chat import ChatCompletionChunk
+from openai.types.completion_usage import CompletionUsage
+from openai.types.responses import (
+    Response,
+    ResponseCompletedEvent,
+    ResponseContentPartAddedEvent,
+    ResponseContentPartDoneEvent,
+    ResponseCreatedEvent,
+    ResponseFunctionCallArgumentsDeltaEvent,
+    ResponseFunctionToolCall,
+    ResponseOutputItem,
+    ResponseOutputItemAddedEvent,
+    ResponseOutputItemDoneEvent,
+    ResponseOutputMessage,
+    ResponseOutputRefusal,
+    ResponseOutputText,
+    ResponseRefusalDeltaEvent,
+    ResponseTextDeltaEvent,
+    ResponseUsage,
+)
+from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails
+
+from ..items import TResponseStreamEvent
+from .fake_id import FAKE_RESPONSES_ID
+
+
+@dataclass
+class StreamingState:
+    started: bool = False
+    text_content_index_and_output: tuple[int, ResponseOutputText] | None = None
+    refusal_content_index_and_output: tuple[int, ResponseOutputRefusal] | None = None
+    function_calls: dict[int, ResponseFunctionToolCall] = field(default_factory=dict)
+
+
+class ChatCmplStreamHandler:
+    @classmethod
+    async def handle_stream(
+        cls,
+        response: Response,
+        stream: AsyncStream[ChatCompletionChunk],
+    ) -> AsyncIterator[TResponseStreamEvent]:
+        usage: CompletionUsage | None = None
+        state = StreamingState()
+
+        async for chunk in stream:
+            if not state.started:
+                state.started = True
+                yield ResponseCreatedEvent(
+                    response=response,
+                    type="response.created",
+                )
+
+            usage = chunk.usage
+
+            if not chunk.choices or not chunk.choices[0].delta:
+                continue
+
+            delta = chunk.choices[0].delta
+
+            # Handle text
+            if delta.content:
+                if not state.text_content_index_and_output:
+                    # Initialize a content tracker for streaming text
+                    state.text_content_index_and_output = (
+                        0 if not state.refusal_content_index_and_output else 1,
+                        ResponseOutputText(
+                            text="",
+                            type="output_text",
+                            annotations=[],
+                        ),
+                    )
+                    # Start a new assistant message stream
+                    assistant_item = ResponseOutputMessage(
+                        id=FAKE_RESPONSES_ID,
+                        content=[],
+                        role="assistant",
+                        type="message",
+                        status="in_progress",
+                    )
+                    # Notify consumers of the start of a new output message + first content part
+                    yield ResponseOutputItemAddedEvent(
+                        item=assistant_item,
+                        output_index=0,
+                        type="response.output_item.added",
+                    )
+                    yield ResponseContentPartAddedEvent(
+                        content_index=state.text_content_index_and_output[0],
+                        item_id=FAKE_RESPONSES_ID,
+                        output_index=0,
+                        part=ResponseOutputText(
+                            text="",
+                            type="output_text",
+                            annotations=[],
+                        ),
+                        type="response.content_part.added",
+                    )
+                # Emit the delta for this segment of content
+                yield ResponseTextDeltaEvent(
+                    content_index=state.text_content_index_and_output[0],
+                    delta=delta.content,
+                    item_id=FAKE_RESPONSES_ID,
+                    output_index=0,
+                    type="response.output_text.delta",
+                )
+                # Accumulate the text into the response part
+                state.text_content_index_and_output[1].text += delta.content
+
+            # Handle refusals (model declines to answer)
+            if delta.refusal:
+                if not state.refusal_content_index_and_output:
+                    # Initialize a content tracker for streaming refusal text
+                    state.refusal_content_index_and_output = (
+                        0 if not state.text_content_index_and_output else 1,
+                        ResponseOutputRefusal(refusal="", type="refusal"),
+                    )
+                    # Start a new assistant message if one doesn't exist yet (in-progress)
+                    assistant_item = ResponseOutputMessage(
+                        id=FAKE_RESPONSES_ID,
+                        content=[],
+                        role="assistant",
+                        type="message",
+                        status="in_progress",
+                    )
+                    # Notify downstream that assistant message + first content part are starting
+                    yield ResponseOutputItemAddedEvent(
+                        item=assistant_item,
+                        output_index=0,
+                        type="response.output_item.added",
+                    )
+                    yield ResponseContentPartAddedEvent(
+                        content_index=state.refusal_content_index_and_output[0],
+                        item_id=FAKE_RESPONSES_ID,
+                        output_index=0,
+                        part=ResponseOutputText(
+                            text="",
+                            type="output_text",
+                            annotations=[],
+                        ),
+                        type="response.content_part.added",
+                    )
+                # Emit the delta for this segment of refusal
+                yield ResponseRefusalDeltaEvent(
+                    content_index=state.refusal_content_index_and_output[0],
+                    delta=delta.refusal,
+                    item_id=FAKE_RESPONSES_ID,
+                    output_index=0,
+                    type="response.refusal.delta",
+                )
+                # Accumulate the refusal string in the output part
+                state.refusal_content_index_and_output[1].refusal += delta.refusal
+
+            # Handle tool calls
+            # Because we don't know the name of the function until the end of the stream, we'll
+            # save everything and yield events at the end
+            if delta.tool_calls:
+                for tc_delta in delta.tool_calls:
+                    if tc_delta.index not in state.function_calls:
+                        state.function_calls[tc_delta.index] = ResponseFunctionToolCall(
+                            id=FAKE_RESPONSES_ID,
+                            arguments="",
+                            name="",
+                            type="function_call",
+                            call_id="",
+                        )
+                    tc_function = tc_delta.function
+
+                    state.function_calls[tc_delta.index].arguments += (
+                        tc_function.arguments if tc_function else ""
+                    ) or ""
+                    state.function_calls[tc_delta.index].name += (
+                        tc_function.name if tc_function else ""
+                    ) or ""
+                    state.function_calls[tc_delta.index].call_id += tc_delta.id or ""
+
+        function_call_starting_index = 0
+        if state.text_content_index_and_output:
+            function_call_starting_index += 1
+            # Send end event for this content part
+            yield ResponseContentPartDoneEvent(
+                content_index=state.text_content_index_and_output[0],
+                item_id=FAKE_RESPONSES_ID,
+                output_index=0,
+                part=state.text_content_index_and_output[1],
+                type="response.content_part.done",
+            )
+
+        if state.refusal_content_index_and_output:
+            function_call_starting_index += 1
+            # Send end event for this content part
+            yield ResponseContentPartDoneEvent(
+                content_index=state.refusal_content_index_and_output[0],
+                item_id=FAKE_RESPONSES_ID,
+                output_index=0,
+                part=state.refusal_content_index_and_output[1],
+                type="response.content_part.done",
+            )
+
+        # Actually send events for the function calls
+        for function_call in state.function_calls.values():
+            # First, a ResponseOutputItemAdded for the function call
+            yield ResponseOutputItemAddedEvent(
+                item=ResponseFunctionToolCall(
+                    id=FAKE_RESPONSES_ID,
+                    call_id=function_call.call_id,
+                    arguments=function_call.arguments,
+                    name=function_call.name,
+                    type="function_call",
+                ),
+                output_index=function_call_starting_index,
+                type="response.output_item.added",
+            )
+            # Then, yield the args
+            yield ResponseFunctionCallArgumentsDeltaEvent(
+                delta=function_call.arguments,
+                item_id=FAKE_RESPONSES_ID,
+                output_index=function_call_starting_index,
+                type="response.function_call_arguments.delta",
+            )
+            # Finally, the ResponseOutputItemDone
+            yield ResponseOutputItemDoneEvent(
+                item=ResponseFunctionToolCall(
+                    id=FAKE_RESPONSES_ID,
+                    call_id=function_call.call_id,
+                    arguments=function_call.arguments,
+                    name=function_call.name,
+                    type="function_call",
+                ),
+                output_index=function_call_starting_index,
+                type="response.output_item.done",
+            )
+
+        # Finally, send the Response completed event
+        outputs: list[ResponseOutputItem] = []
+        if state.text_content_index_and_output or state.refusal_content_index_and_output:
+            assistant_msg = ResponseOutputMessage(
+                id=FAKE_RESPONSES_ID,
+                content=[],
+                role="assistant",
+                type="message",
+                status="completed",
+            )
+            if state.text_content_index_and_output:
+                assistant_msg.content.append(state.text_content_index_and_output[1])
+            if state.refusal_content_index_and_output:
+                assistant_msg.content.append(state.refusal_content_index_and_output[1])
+            outputs.append(assistant_msg)
+
+            # send a ResponseOutputItemDone for the assistant message
+            yield ResponseOutputItemDoneEvent(
+                item=assistant_msg,
+                output_index=0,
+                type="response.output_item.done",
+            )
+
+        for function_call in state.function_calls.values():
+            outputs.append(function_call)
+
+        final_response = response.model_copy()
+        final_response.output = outputs
+        final_response.usage = (
+            ResponseUsage(
+                input_tokens=usage.prompt_tokens,
+                output_tokens=usage.completion_tokens,
+                total_tokens=usage.total_tokens,
+                output_tokens_details=OutputTokensDetails(
+                    reasoning_tokens=usage.completion_tokens_details.reasoning_tokens
+                    if usage.completion_tokens_details
+                    and usage.completion_tokens_details.reasoning_tokens
+                    else 0
+                ),
+                input_tokens_details=InputTokensDetails(
+                    cached_tokens=usage.prompt_tokens_details.cached_tokens
+                    if usage.prompt_tokens_details and usage.prompt_tokens_details.cached_tokens
+                    else 0
+                ),
+            )
+            if usage
+            else None
+        )
+
+        yield ResponseCompletedEvent(
+            response=final_response,
+            type="response.completed",
+        )

agents/models/interface.py

@@ -5,7 +5,7 @@ import enum
 from collections.abc import AsyncIterator
 from typing import TYPE_CHECKING
 
-from ..agent_output import AgentOutputSchema
+from ..agent_output import AgentOutputSchemaBase
 from ..handoffs import Handoff
 from ..items import ModelResponse, TResponseInputItem, TResponseStreamEvent
 from ..tool import Tool
@@ -41,7 +41,7 @@ class Model(abc.ABC):
         input: str | list[TResponseInputItem],
         model_settings: ModelSettings,
         tools: list[Tool],
-        output_schema: AgentOutputSchema | None,
+        output_schema: AgentOutputSchemaBase | None,
         handoffs: list[Handoff],
         tracing: ModelTracing,
         *,
@@ -72,7 +72,7 @@ class Model(abc.ABC):
         input: str | list[TResponseInputItem],
         model_settings: ModelSettings,
         tools: list[Tool],
-        output_schema: AgentOutputSchema | None,
+        output_schema: AgentOutputSchemaBase | None,
         handoffs: list[Handoff],
         tracing: ModelTracing,
         *,

agents/models/multi_provider.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+from openai import AsyncOpenAI
+
+from ..exceptions import UserError
+from .interface import Model, ModelProvider
+from .openai_provider import OpenAIProvider
+
+
+class MultiProviderMap:
+    """A map of model name prefixes to ModelProviders."""
+
+    def __init__(self):
+        self._mapping: dict[str, ModelProvider] = {}
+
+    def has_prefix(self, prefix: str) -> bool:
+        """Returns True if the given prefix is in the mapping."""
+        return prefix in self._mapping
+
+    def get_mapping(self) -> dict[str, ModelProvider]:
+        """Returns a copy of the current prefix -> ModelProvider mapping."""
+        return self._mapping.copy()
+
+    def set_mapping(self, mapping: dict[str, ModelProvider]):
+        """Overwrites the current mapping with a new one."""
+        self._mapping = mapping
+
+    def get_provider(self, prefix: str) -> ModelProvider | None:
+        """Returns the ModelProvider for the given prefix.
+
+        Args:
+            prefix: The prefix of the model name e.g. "openai" or "my_prefix".
+        """
+        return self._mapping.get(prefix)
+
+    def add_provider(self, prefix: str, provider: ModelProvider):
+        """Adds a new prefix -> ModelProvider mapping.
+
+        Args:
+            prefix: The prefix of the model name e.g. "openai" or "my_prefix".
+            provider: The ModelProvider to use for the given prefix.
+        """
+        self._mapping[prefix] = provider
+
+    def remove_provider(self, prefix: str):
+        """Removes the mapping for the given prefix.
+
+        Args:
+            prefix: The prefix of the model name e.g. "openai" or "my_prefix".
+        """
+        del self._mapping[prefix]
+
+
+class MultiProvider(ModelProvider):
+    """This ModelProvider maps to a Model based on the prefix of the model name. By default, the
+    mapping is:
+    - "openai/" prefix or no prefix -> OpenAIProvider. e.g. "openai/gpt-4.1", "gpt-4.1"
+    - "litellm/" prefix -> LitellmProvider. e.g. "litellm/openai/gpt-4.1"
+
+    You can override or customize this mapping.
+    """
+
+    def __init__(
+        self,
+        *,
+        provider_map: MultiProviderMap | None = None,
+        openai_api_key: str | None = None,
+        openai_base_url: str | None = None,
+        openai_client: AsyncOpenAI | None = None,
+        openai_organization: str | None = None,
+        openai_project: str | None = None,
+        openai_use_responses: bool | None = None,
+    ) -> None:
+        """Create a new OpenAI provider.
+
+        Args:
+            provider_map: A MultiProviderMap that maps prefixes to ModelProviders. If not provided,
+                we will use a default mapping. See the documentation for this class to see the
+                default mapping.
+            openai_api_key: The API key to use for the OpenAI provider. If not provided, we will use
+                the default API key.
+            openai_base_url: The base URL to use for the OpenAI provider. If not provided, we will
+                use the default base URL.
+            openai_client: An optional OpenAI client to use. If not provided, we will create a new
+                OpenAI client using the api_key and base_url.
+            openai_organization: The organization to use for the OpenAI provider.
+            openai_project: The project to use for the OpenAI provider.
+            openai_use_responses: Whether to use the OpenAI responses API.
+        """
+        self.provider_map = provider_map
+        self.openai_provider = OpenAIProvider(
+            api_key=openai_api_key,
+            base_url=openai_base_url,
+            openai_client=openai_client,
+            organization=openai_organization,
+            project=openai_project,
+            use_responses=openai_use_responses,
+        )
+
+        self._fallback_providers: dict[str, ModelProvider] = {}
+
+    def _get_prefix_and_model_name(self, model_name: str | None) -> tuple[str | None, str | None]:
+        if model_name is None:
+            return None, None
+        elif "/" in model_name:
+            prefix, model_name = model_name.split("/", 1)
+            return prefix, model_name
+        else:
+            return None, model_name
+
+    def _create_fallback_provider(self, prefix: str) -> ModelProvider:
+        if prefix == "litellm":
+            from ..extensions.models.litellm_provider import LitellmProvider
+
+            return LitellmProvider()
+        else:
+            raise UserError(f"Unknown prefix: {prefix}")
+
+    def _get_fallback_provider(self, prefix: str | None) -> ModelProvider:
+        if prefix is None or prefix == "openai":
+            return self.openai_provider
+        elif prefix in self._fallback_providers:
+            return self._fallback_providers[prefix]
+        else:
+            self._fallback_providers[prefix] = self._create_fallback_provider(prefix)
+            return self._fallback_providers[prefix]
+
+    def get_model(self, model_name: str | None) -> Model:
+        """Returns a Model based on the model name. The model name can have a prefix, ending with
+        a "/", which will be used to look up the ModelProvider. If there is no prefix, we will use
+        the OpenAI provider.
+
+        Args:
+            model_name: The name of the model to get.
+
+        Returns:
+            A Model.
+        """
+        prefix, model_name = self._get_prefix_and_model_name(model_name)
+
+        if prefix and self.provider_map and (provider := self.provider_map.get_provider(prefix)):
+            return provider.get_model(model_name)
+        else:
+            return self._get_fallback_provider(prefix).get_model(model_name)