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

agents/__init__.py

@@ -6,7 +6,7 @@ from openai import AsyncOpenAI
 
 from . import _config
 from .agent import Agent, ToolsToFinalOutputFunction, ToolsToFinalOutputResult
-from .agent_output import AgentOutputSchema
+from .agent_output import AgentOutputSchema, AgentOutputSchemaBase
 from .computer import AsyncComputer, Button, Computer, Environment
 from .exceptions import (
     AgentsException,
@@ -158,6 +158,7 @@ __all__ = [
     "OpenAIProvider",
     "OpenAIResponsesModel",
     "AgentOutputSchema",
+    "AgentOutputSchemaBase",
     "Computer",
     "AsyncComputer",
     "Environment",

agents/_run_impl.py

@@ -29,7 +29,7 @@ from openai.types.responses.response_input_param import ComputerCallOutput
 from openai.types.responses.response_reasoning_item import ResponseReasoningItem
 
 from .agent import Agent, ToolsToFinalOutputResult
-from .agent_output import AgentOutputSchema
+from .agent_output import AgentOutputSchemaBase
 from .computer import AsyncComputer, Computer
 from .exceptions import AgentsException, ModelBehaviorError, UserError
 from .guardrail import InputGuardrail, InputGuardrailResult, OutputGuardrail, OutputGuardrailResult
@@ -195,7 +195,7 @@ class RunImpl:
         pre_step_items: list[RunItem],
         new_response: ModelResponse,
         processed_response: ProcessedResponse,
-        output_schema: AgentOutputSchema | None,
+        output_schema: AgentOutputSchemaBase | None,
         hooks: RunHooks[TContext],
         context_wrapper: RunContextWrapper[TContext],
         run_config: RunConfig,
@@ -335,7 +335,7 @@ class RunImpl:
         agent: Agent[Any],
         all_tools: list[Tool],
         response: ModelResponse,
-        output_schema: AgentOutputSchema | None,
+        output_schema: AgentOutputSchemaBase | None,
         handoffs: list[Handoff],
     ) -> ProcessedResponse:
         items: list[RunItem] = []

agents/agent.py

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, cast
 
 from typing_extensions import NotRequired, TypeAlias, TypedDict
 
+from .agent_output import AgentOutputSchemaBase
 from .guardrail import InputGuardrail, OutputGuardrail
 from .handoffs import Handoff
 from .items import ItemHelpers
@@ -141,8 +142,14 @@ class Agent(Generic[TContext]):
     Runs only if the agent produces a final output.
     """
 
-    output_type: type[Any] | None = None
-    """The type of the output object. If not provided, the output will be `str`."""
+    output_type: type[Any] | AgentOutputSchemaBase | None = None
+    """The type of the output object. If not provided, the output will be `str`. In most cases,
+    you should pass a regular Python type (e.g. a dataclass, Pydantic model, TypedDict, etc).
+    You can customize this in two ways:
+    1. If you want non-strict schemas, pass `AgentOutputSchema(MyClass, strict_json_schema=False)`.
+    2. If you want to use a custom JSON schema (i.e. without using the SDK's automatic schema)
+       creation, subclass and pass an `AgentOutputSchemaBase` subclass.
+    """
 
     hooks: AgentHooks[TContext] | None = None
     """A class that receives callbacks on various lifecycle events for this agent.

agents/agent_output.py

@@ -1,3 +1,4 @@
+import abc
 from dataclasses import dataclass
 from typing import Any
 
@@ -12,8 +13,46 @@ from .util import _error_tracing, _json
 _WRAPPER_DICT_KEY = "response"
 
 
+class AgentOutputSchemaBase(abc.ABC):
+    """An object that captures the JSON schema of the output, as well as validating/parsing JSON
+    produced by the LLM into the output type.
+    """
+
+    @abc.abstractmethod
+    def is_plain_text(self) -> bool:
+        """Whether the output type is plain text (versus a JSON object)."""
+        pass
+
+    @abc.abstractmethod
+    def name(self) -> str:
+        """The name of the output type."""
+        pass
+
+    @abc.abstractmethod
+    def json_schema(self) -> dict[str, Any]:
+        """Returns the JSON schema of the output. Will only be called if the output type is not
+        plain text.
+        """
+        pass
+
+    @abc.abstractmethod
+    def is_strict_json_schema(self) -> bool:
+        """Whether the JSON schema is in strict mode. Strict mode constrains the JSON schema
+        features, but guarantees valis JSON. See here for details:
+        https://platform.openai.com/docs/guides/structured-outputs#supported-schemas
+        """
+        pass
+
+    @abc.abstractmethod
+    def validate_json(self, json_str: str) -> Any:
+        """Validate a JSON string against the output type. You must return the validated object,
+        or raise a `ModelBehaviorError` if the JSON is invalid.
+        """
+        pass
+
+
 @dataclass(init=False)
-class AgentOutputSchema:
+class AgentOutputSchema(AgentOutputSchemaBase):
     """An object that captures the JSON schema of the output, as well as validating/parsing JSON
     produced by the LLM into the output type.
     """
@@ -32,7 +71,7 @@ class AgentOutputSchema:
     _output_schema: dict[str, Any]
     """The JSON schema of the output."""
 
-    strict_json_schema: bool
+    _strict_json_schema: bool
     """Whether the JSON schema is in strict mode. We **strongly** recommend setting this to True,
     as it increases the likelihood of correct JSON input.
     """
@@ -45,7 +84,7 @@ class AgentOutputSchema:
                 setting this to True, as it increases the likelihood of correct JSON input.
         """
         self.output_type = output_type
-        self.strict_json_schema = strict_json_schema
+        self._strict_json_schema = strict_json_schema
 
         if output_type is None or output_type is str:
             self._is_wrapped = False
@@ -70,24 +109,35 @@ class AgentOutputSchema:
             self._type_adapter = TypeAdapter(output_type)
             self._output_schema = self._type_adapter.json_schema()
 
-        if self.strict_json_schema:
-            self._output_schema = ensure_strict_json_schema(self._output_schema)
+        if self._strict_json_schema:
+            try:
+                self._output_schema = ensure_strict_json_schema(self._output_schema)
+            except UserError as e:
+                raise UserError(
+                    "Strict JSON schema is enabled, but the output type is not valid. "
+                    "Either make the output type strict, or pass output_schema_strict=False to "
+                    "your Agent()"
+                ) from e
 
     def is_plain_text(self) -> bool:
         """Whether the output type is plain text (versus a JSON object)."""
         return self.output_type is None or self.output_type is str
 
+    def is_strict_json_schema(self) -> bool:
+        """Whether the JSON schema is in strict mode."""
+        return self._strict_json_schema
+
     def json_schema(self) -> dict[str, Any]:
         """The JSON schema of the output type."""
         if self.is_plain_text():
             raise UserError("Output type is plain text, so no JSON schema is available")
         return self._output_schema
 
-    def validate_json(self, json_str: str, partial: bool = False) -> Any:
+    def validate_json(self, json_str: str) -> Any:
         """Validate a JSON string against the output type. Returns the validated object, or raises
         a `ModelBehaviorError` if the JSON is invalid.
         """
-        validated = _json.validate_json(json_str, self._type_adapter, partial)
+        validated = _json.validate_json(json_str, self._type_adapter, partial=False)
         if self._is_wrapped:
             if not isinstance(validated, dict):
                 _error_tracing.attach_error_to_current_span(
@@ -113,7 +163,7 @@ class AgentOutputSchema:
             return validated[_WRAPPER_DICT_KEY]
         return validated
 
-    def output_type_name(self) -> str:
+    def name(self) -> str:
         """The name of the output type."""
         return _type_to_str(self.output_type)
 

agents/extensions/models/litellm_model.py

@@ -0,0 +1,382 @@
+from __future__ import annotations
+
+import dataclasses
+import json
+import time
+from collections.abc import AsyncIterator
+from typing import Any, Literal, cast, overload
+
+import litellm.types
+
+from agents.exceptions import ModelBehaviorError
+
+try:
+    import litellm
+except ImportError as _e:
+    raise ImportError(
+        "`litellm` is required to use the LitellmModel. You can install it via the optional "
+        "dependency group: `pip install 'openai-agents[litellm]'`."
+    ) from _e
+
+from openai import NOT_GIVEN, AsyncStream, NotGiven
+from openai.types.chat import ChatCompletionChunk, ChatCompletionMessageToolCall
+from openai.types.chat.chat_completion_message import (
+    Annotation,
+    AnnotationURLCitation,
+    ChatCompletionMessage,
+)
+from openai.types.chat.chat_completion_message_tool_call import Function
+from openai.types.responses import Response
+
+from ... import _debug
+from ...agent_output import AgentOutputSchemaBase
+from ...handoffs import Handoff
+from ...items import ModelResponse, TResponseInputItem, TResponseStreamEvent
+from ...logger import logger
+from ...model_settings import ModelSettings
+from ...models.chatcmpl_converter import Converter
+from ...models.chatcmpl_helpers import HEADERS
+from ...models.chatcmpl_stream_handler import ChatCmplStreamHandler
+from ...models.fake_id import FAKE_RESPONSES_ID
+from ...models.interface import Model, ModelTracing
+from ...tool import Tool
+from ...tracing import generation_span
+from ...tracing.span_data import GenerationSpanData
+from ...tracing.spans import Span
+from ...usage import Usage
+
+
+class LitellmModel(Model):
+    """This class enables using any model via LiteLLM. LiteLLM allows you to acess OpenAPI,
+    Anthropic, Gemini, Mistral, and many other models.
+    See supported models here: [litellm models](https://docs.litellm.ai/docs/providers).
+    """
+
+    def __init__(
+        self,
+        model: str,
+        base_url: str | None = None,
+        api_key: str | None = None,
+    ):
+        self.model = model
+        self.base_url = base_url
+        self.api_key = api_key
+
+    async def get_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,
+        previous_response_id: str | None,
+    ) -> ModelResponse:
+        with generation_span(
+            model=str(self.model),
+            model_config=dataclasses.asdict(model_settings)
+            | {"base_url": str(self.base_url or ""), "model_impl": "litellm"},
+            disabled=tracing.is_disabled(),
+        ) as span_generation:
+            response = await self._fetch_response(
+                system_instructions,
+                input,
+                model_settings,
+                tools,
+                output_schema,
+                handoffs,
+                span_generation,
+                tracing,
+                stream=False,
+            )
+
+            assert isinstance(response.choices[0], litellm.types.utils.Choices)
+
+            if _debug.DONT_LOG_MODEL_DATA:
+                logger.debug("Received model response")
+            else:
+                logger.debug(
+                    f"LLM resp:\n{json.dumps(response.choices[0].message.model_dump(), indent=2)}\n"
+                )
+
+            if hasattr(response, "usage"):
+                response_usage = response.usage
+                usage = (
+                    Usage(
+                        requests=1,
+                        input_tokens=response_usage.prompt_tokens,
+                        output_tokens=response_usage.completion_tokens,
+                        total_tokens=response_usage.total_tokens,
+                    )
+                    if response.usage
+                    else Usage()
+                )
+            else:
+                usage = Usage()
+                logger.warning("No usage information returned from Litellm")
+
+            if tracing.include_data():
+                span_generation.span_data.output = [response.choices[0].message.model_dump()]
+            span_generation.span_data.usage = {
+                "input_tokens": usage.input_tokens,
+                "output_tokens": usage.output_tokens,
+            }
+
+            items = Converter.message_to_output_items(
+                LitellmConverter.convert_message_to_openai(response.choices[0].message)
+            )
+
+            return ModelResponse(
+                output=items,
+                usage=usage,
+                response_id=None,
+            )
+
+    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,
+        *,
+        previous_response_id: str | None,
+    ) -> AsyncIterator[TResponseStreamEvent]:
+        with generation_span(
+            model=str(self.model),
+            model_config=dataclasses.asdict(model_settings)
+            | {"base_url": str(self.base_url or ""), "model_impl": "litellm"},
+            disabled=tracing.is_disabled(),
+        ) as span_generation:
+            response, stream = await self._fetch_response(
+                system_instructions,
+                input,
+                model_settings,
+                tools,
+                output_schema,
+                handoffs,
+                span_generation,
+                tracing,
+                stream=True,
+            )
+
+            final_response: Response | None = None
+            async for chunk in ChatCmplStreamHandler.handle_stream(response, stream):
+                yield chunk
+
+                if chunk.type == "response.completed":
+                    final_response = chunk.response
+
+            if tracing.include_data() and final_response:
+                span_generation.span_data.output = [final_response.model_dump()]
+
+            if final_response and final_response.usage:
+                span_generation.span_data.usage = {
+                    "input_tokens": final_response.usage.input_tokens,
+                    "output_tokens": final_response.usage.output_tokens,
+                }
+
+    @overload
+    async def _fetch_response(
+        self,
+        system_instructions: str | None,
+        input: str | list[TResponseInputItem],
+        model_settings: ModelSettings,
+        tools: list[Tool],
+        output_schema: AgentOutputSchemaBase | None,
+        handoffs: list[Handoff],
+        span: Span[GenerationSpanData],
+        tracing: ModelTracing,
+        stream: Literal[True],
+    ) -> tuple[Response, AsyncStream[ChatCompletionChunk]]: ...
+
+    @overload
+    async def _fetch_response(
+        self,
+        system_instructions: str | None,
+        input: str | list[TResponseInputItem],
+        model_settings: ModelSettings,
+        tools: list[Tool],
+        output_schema: AgentOutputSchemaBase | None,
+        handoffs: list[Handoff],
+        span: Span[GenerationSpanData],
+        tracing: ModelTracing,
+        stream: Literal[False],
+    ) -> litellm.types.utils.ModelResponse: ...
+
+    async def _fetch_response(
+        self,
+        system_instructions: str | None,
+        input: str | list[TResponseInputItem],
+        model_settings: ModelSettings,
+        tools: list[Tool],
+        output_schema: AgentOutputSchemaBase | None,
+        handoffs: list[Handoff],
+        span: Span[GenerationSpanData],
+        tracing: ModelTracing,
+        stream: bool = False,
+    ) -> litellm.types.utils.ModelResponse | tuple[Response, AsyncStream[ChatCompletionChunk]]:
+        converted_messages = Converter.items_to_messages(input)
+
+        if system_instructions:
+            converted_messages.insert(
+                0,
+                {
+                    "content": system_instructions,
+                    "role": "system",
+                },
+            )
+        if tracing.include_data():
+            span.span_data.input = converted_messages
+
+        parallel_tool_calls = (
+            True
+            if model_settings.parallel_tool_calls and tools and len(tools) > 0
+            else False
+            if model_settings.parallel_tool_calls is False
+            else None
+        )
+        tool_choice = Converter.convert_tool_choice(model_settings.tool_choice)
+        response_format = Converter.convert_response_format(output_schema)
+
+        converted_tools = [Converter.tool_to_openai(tool) for tool in tools] if tools else []
+
+        for handoff in handoffs:
+            converted_tools.append(Converter.convert_handoff_tool(handoff))
+
+        if _debug.DONT_LOG_MODEL_DATA:
+            logger.debug("Calling LLM")
+        else:
+            logger.debug(
+                f"Calling Litellm model: {self.model}\n"
+                f"{json.dumps(converted_messages, indent=2)}\n"
+                f"Tools:\n{json.dumps(converted_tools, indent=2)}\n"
+                f"Stream: {stream}\n"
+                f"Tool choice: {tool_choice}\n"
+                f"Response format: {response_format}\n"
+            )
+
+        reasoning_effort = model_settings.reasoning.effort if model_settings.reasoning else None
+
+        stream_options = None
+        if stream and model_settings.include_usage is not None:
+            stream_options = {"include_usage": model_settings.include_usage}
+
+        extra_kwargs = {}
+        if model_settings.extra_query:
+            extra_kwargs["extra_query"] = model_settings.extra_query
+        if model_settings.metadata:
+            extra_kwargs["metadata"] = model_settings.metadata
+
+        ret = await litellm.acompletion(
+            model=self.model,
+            messages=converted_messages,
+            tools=converted_tools or None,
+            temperature=model_settings.temperature,
+            top_p=model_settings.top_p,
+            frequency_penalty=model_settings.frequency_penalty,
+            presence_penalty=model_settings.presence_penalty,
+            max_tokens=model_settings.max_tokens,
+            tool_choice=self._remove_not_given(tool_choice),
+            response_format=self._remove_not_given(response_format),
+            parallel_tool_calls=parallel_tool_calls,
+            stream=stream,
+            stream_options=stream_options,
+            reasoning_effort=reasoning_effort,
+            extra_headers=HEADERS,
+            api_key=self.api_key,
+            base_url=self.base_url,
+            **extra_kwargs,
+        )
+
+        if isinstance(ret, litellm.types.utils.ModelResponse):
+            return ret
+
+        response = Response(
+            id=FAKE_RESPONSES_ID,
+            created_at=time.time(),
+            model=self.model,
+            object="response",
+            output=[],
+            tool_choice=cast(Literal["auto", "required", "none"], tool_choice)
+            if tool_choice != NOT_GIVEN
+            else "auto",
+            top_p=model_settings.top_p,
+            temperature=model_settings.temperature,
+            tools=[],
+            parallel_tool_calls=parallel_tool_calls or False,
+            reasoning=model_settings.reasoning,
+        )
+        return response, ret
+
+    def _remove_not_given(self, value: Any) -> Any:
+        if isinstance(value, NotGiven):
+            return None
+        return value
+
+
+class LitellmConverter:
+    @classmethod
+    def convert_message_to_openai(
+        cls, message: litellm.types.utils.Message
+    ) -> ChatCompletionMessage:
+        if message.role != "assistant":
+            raise ModelBehaviorError(f"Unsupported role: {message.role}")
+
+        tool_calls = (
+            [LitellmConverter.convert_tool_call_to_openai(tool) for tool in message.tool_calls]
+            if message.tool_calls
+            else None
+        )
+
+        provider_specific_fields = message.get("provider_specific_fields", None)
+        refusal = (
+            provider_specific_fields.get("refusal", None) if provider_specific_fields else None
+        )
+
+        return ChatCompletionMessage(
+            content=message.content,
+            refusal=refusal,
+            role="assistant",
+            annotations=cls.convert_annotations_to_openai(message),
+            audio=message.get("audio", None),  # litellm deletes audio if not present
+            tool_calls=tool_calls,
+        )
+
+    @classmethod
+    def convert_annotations_to_openai(
+        cls, message: litellm.types.utils.Message
+    ) -> list[Annotation] | None:
+        annotations: list[litellm.types.llms.openai.ChatCompletionAnnotation] | None = message.get(
+            "annotations", None
+        )
+        if not annotations:
+            return None
+
+        return [
+            Annotation(
+                type="url_citation",
+                url_citation=AnnotationURLCitation(
+                    start_index=annotation["url_citation"]["start_index"],
+                    end_index=annotation["url_citation"]["end_index"],
+                    url=annotation["url_citation"]["url"],
+                    title=annotation["url_citation"]["title"],
+                ),
+            )
+            for annotation in annotations
+        ]
+
+    @classmethod
+    def convert_tool_call_to_openai(
+        cls, tool_call: litellm.types.utils.ChatCompletionMessageToolCall
+    ) -> ChatCompletionMessageToolCall:
+        return ChatCompletionMessageToolCall(
+            id=tool_call.id,
+            type="function",
+            function=Function(
+                name=tool_call.function.name or "", arguments=tool_call.function.arguments
+            ),
+        )

agents/extensions/models/litellm_provider.py

@@ -0,0 +1,21 @@
+from ...models.interface import Model, ModelProvider
+from .litellm_model import LitellmModel
+
+DEFAULT_MODEL: str = "gpt-4.1"
+
+
+class LitellmProvider(ModelProvider):
+    """A ModelProvider that uses LiteLLM to route to any model provider. You can use it via:
+    ```python
+    Runner.run(agent, input, run_config=RunConfig(model_provider=LitellmProvider()))
+    ```
+    See supported models here: [litellm models](https://docs.litellm.ai/docs/providers).
+
+    NOTE: API keys must be set via environment variables. If you're using models that require
+    additional configuration (e.g. Azure API base or version), those must also be set via the
+    environment variables that LiteLLM expects. If you have more advanced needs, we recommend
+    copy-pasting this class and making any modifications you need.
+    """
+
+    def get_model(self, model_name: str | None) -> Model:
+        return LitellmModel(model_name or DEFAULT_MODEL)

agents/extensions/visualization.py

@@ -132,6 +132,6 @@ def draw_graph(agent: Agent, filename: Optional[str] = None) -> graphviz.Source:
     graph = graphviz.Source(dot_code)
 
     if filename:
-        graph.render(filename, format="png")
+        graph.render(filename, format="png", cleanup=True)
 
     return graph

agents/mcp/server.py

@@ -137,9 +137,10 @@ class _MCPServerWithClientSession(MCPServer, abc.ABC):
         async with self._cleanup_lock:
             try:
                 await self.exit_stack.aclose()
-                self.session = None
             except Exception as e:
                 logger.error(f"Error cleaning up server: {e}")
+            finally:
+                self.session = None
 
 
 class MCPServerStdioParams(TypedDict):