diff --git a/src/agents/extensions/models/litellm_model.py b/src/agents/extensions/models/litellm_model.py index 8d39ad390..3743d82f2 100644 --- a/src/agents/extensions/models/litellm_model.py +++ b/src/agents/extensions/models/litellm_model.py @@ -39,7 +39,7 @@ from ...logger import logger from ...model_settings import ModelSettings from ...models.chatcmpl_converter import Converter -from ...models.chatcmpl_helpers import HEADERS +from ...models.chatcmpl_helpers import HEADERS, USER_AGENT_OVERRIDE from ...models.chatcmpl_stream_handler import ChatCmplStreamHandler from ...models.fake_id import FAKE_RESPONSES_ID from ...models.interface import Model, ModelTracing @@ -353,7 +353,7 @@ async def _fetch_response( stream_options=stream_options, reasoning_effort=reasoning_effort, top_logprobs=model_settings.top_logprobs, - extra_headers={**HEADERS, **(model_settings.extra_headers or {})}, + extra_headers=self._merge_headers(model_settings), api_key=self.api_key, base_url=self.base_url, **extra_kwargs, @@ -384,6 +384,13 @@ def _remove_not_given(self, value: Any) -> Any: return None return value + def _merge_headers(self, model_settings: ModelSettings): + merged = {**HEADERS, **(model_settings.extra_headers or {})} + ua_ctx = USER_AGENT_OVERRIDE.get() + if ua_ctx is not None: + merged["User-Agent"] = ua_ctx + return merged + class LitellmConverter: @classmethod diff --git a/src/agents/models/chatcmpl_helpers.py b/src/agents/models/chatcmpl_helpers.py index 0cee21ecc..51f2cc258 100644 --- a/src/agents/models/chatcmpl_helpers.py +++ b/src/agents/models/chatcmpl_helpers.py @@ -1,5 +1,7 @@ from __future__ import annotations +from contextvars import ContextVar + from openai import AsyncOpenAI from ..model_settings import ModelSettings @@ -8,6 +10,10 @@ _USER_AGENT = f"Agents/Python {__version__}" HEADERS = {"User-Agent": _USER_AGENT} +USER_AGENT_OVERRIDE: ContextVar[str | None] = ContextVar( + "openai_chatcompletions_user_agent_override", default=None +) + class ChatCmplHelpers: @classmethod diff --git a/src/agents/models/openai_chatcompletions.py b/src/agents/models/openai_chatcompletions.py index a50a1a8a5..ea355b325 100644 --- a/src/agents/models/openai_chatcompletions.py +++ b/src/agents/models/openai_chatcompletions.py @@ -25,7 +25,7 @@ from ..usage import Usage from ..util._json import _to_dump_compatible from .chatcmpl_converter import Converter -from .chatcmpl_helpers import HEADERS, ChatCmplHelpers +from .chatcmpl_helpers import HEADERS, USER_AGENT_OVERRIDE, ChatCmplHelpers from .chatcmpl_stream_handler import ChatCmplStreamHandler from .fake_id import FAKE_RESPONSES_ID from .interface import Model, ModelTracing @@ -306,7 +306,7 @@ async def _fetch_response( reasoning_effort=self._non_null_or_not_given(reasoning_effort), verbosity=self._non_null_or_not_given(model_settings.verbosity), top_logprobs=self._non_null_or_not_given(model_settings.top_logprobs), - extra_headers={**HEADERS, **(model_settings.extra_headers or {})}, + extra_headers=self._merge_headers(model_settings), extra_query=model_settings.extra_query, extra_body=model_settings.extra_body, metadata=self._non_null_or_not_given(model_settings.metadata), @@ -349,3 +349,10 @@ def _get_client(self) -> AsyncOpenAI: if self._client is None: self._client = AsyncOpenAI() return self._client + + def _merge_headers(self, model_settings: ModelSettings): + merged = {**HEADERS, **(model_settings.extra_headers or {})} + ua_ctx = USER_AGENT_OVERRIDE.get() + if ua_ctx is not None: + merged["User-Agent"] = ua_ctx + return merged diff --git a/src/agents/models/openai_responses.py b/src/agents/models/openai_responses.py index 9ca2d324f..5886b4833 100644 --- a/src/agents/models/openai_responses.py +++ b/src/agents/models/openai_responses.py @@ -2,6 +2,7 @@ import json from collections.abc import AsyncIterator +from contextvars import ContextVar from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal, cast, overload @@ -49,6 +50,11 @@ _USER_AGENT = f"Agents/Python {__version__}" _HEADERS = {"User-Agent": _USER_AGENT} +# Override for the User-Agent header used by the Responses API. +_USER_AGENT_OVERRIDE: ContextVar[str | None] = ContextVar( + "openai_responses_user_agent_override", default=None +) + class OpenAIResponsesModel(Model): """ @@ -312,7 +318,7 @@ async def _fetch_response( tool_choice=tool_choice, parallel_tool_calls=parallel_tool_calls, stream=stream, - extra_headers={**_HEADERS, **(model_settings.extra_headers or {})}, + extra_headers=self._merge_headers(model_settings), extra_query=model_settings.extra_query, extra_body=model_settings.extra_body, text=response_format, @@ -327,6 +333,13 @@ def _get_client(self) -> AsyncOpenAI: self._client = AsyncOpenAI() return self._client + def _merge_headers(self, model_settings: ModelSettings): + merged = {**_HEADERS, **(model_settings.extra_headers or {})} + ua_ctx = _USER_AGENT_OVERRIDE.get() + if ua_ctx is not None: + merged["User-Agent"] = ua_ctx + return merged + @dataclass class ConvertedTools: diff --git a/tests/models/test_litellm_user_agent.py b/tests/models/test_litellm_user_agent.py new file mode 100644 index 000000000..03f0f6b84 --- /dev/null +++ b/tests/models/test_litellm_user_agent.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +from agents import ModelSettings, ModelTracing, __version__ +from agents.models.chatcmpl_helpers import USER_AGENT_OVERRIDE + + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +@pytest.mark.parametrize("override_ua", [None, "test_user_agent"]) +async def test_user_agent_header_litellm(override_ua: str | None, monkeypatch): + called_kwargs: dict[str, Any] = {} + expected_ua = override_ua or f"Agents/Python {__version__}" + + import importlib + import sys + import types as pytypes + + litellm_fake: Any = pytypes.ModuleType("litellm") + + class DummyMessage: + role = "assistant" + content = "Hello" + tool_calls: list[Any] | None = None + + def get(self, _key, _default=None): + return None + + def model_dump(self): + return {"role": self.role, "content": self.content} + + class Choices: # noqa: N801 - mimic litellm naming + def __init__(self): + self.message = DummyMessage() + + class DummyModelResponse: + def __init__(self): + self.choices = [Choices()] + + async def acompletion(**kwargs): + nonlocal called_kwargs + called_kwargs = kwargs + return DummyModelResponse() + + utils_ns = pytypes.SimpleNamespace() + utils_ns.Choices = Choices + utils_ns.ModelResponse = DummyModelResponse + + litellm_types = pytypes.SimpleNamespace( + utils=utils_ns, + llms=pytypes.SimpleNamespace(openai=pytypes.SimpleNamespace(ChatCompletionAnnotation=dict)), + ) + litellm_fake.acompletion = acompletion + litellm_fake.types = litellm_types + + monkeypatch.setitem(sys.modules, "litellm", litellm_fake) + + litellm_mod = importlib.import_module("agents.extensions.models.litellm_model") + monkeypatch.setattr(litellm_mod, "litellm", litellm_fake, raising=True) + LitellmModel = litellm_mod.LitellmModel + + model = LitellmModel(model="gpt-4") + + if override_ua is not None: + token = USER_AGENT_OVERRIDE.set(override_ua) + else: + token = None + try: + await model.get_response( + system_instructions=None, + input="hi", + model_settings=ModelSettings(), + tools=[], + output_schema=None, + handoffs=[], + tracing=ModelTracing.DISABLED, + previous_response_id=None, + conversation_id=None, + prompt=None, + ) + finally: + if token is not None: + USER_AGENT_OVERRIDE.reset(token) + + assert "extra_headers" in called_kwargs + assert called_kwargs["extra_headers"]["User-Agent"] == expected_ua diff --git a/tests/test_openai_chatcompletions.py b/tests/test_openai_chatcompletions.py index d52d89b47..df44021a2 100644 --- a/tests/test_openai_chatcompletions.py +++ b/tests/test_openai_chatcompletions.py @@ -31,9 +31,10 @@ ModelTracing, OpenAIChatCompletionsModel, OpenAIProvider, + __version__, generation_span, ) -from agents.models.chatcmpl_helpers import ChatCmplHelpers +from agents.models.chatcmpl_helpers import USER_AGENT_OVERRIDE, ChatCmplHelpers from agents.models.fake_id import FAKE_RESPONSES_ID @@ -370,6 +371,60 @@ def test_store_param(): "Should respect explicitly set store=True" ) + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +@pytest.mark.parametrize("override_ua", [None, "test_user_agent"]) +async def test_user_agent_header_chat_completions(override_ua): + called_kwargs: dict[str, Any] = {} + expected_ua = override_ua or f"Agents/Python {__version__}" + + class DummyCompletions: + async def create(self, **kwargs): + nonlocal called_kwargs + called_kwargs = kwargs + msg = ChatCompletionMessage(role="assistant", content="Hello") + choice = Choice(index=0, finish_reason="stop", message=msg) + return ChatCompletion( + id="resp-id", + created=0, + model="fake", + object="chat.completion", + choices=[choice], + usage=None, + ) + + class DummyChatClient: + def __init__(self): + self.chat = type("_Chat", (), {"completions": DummyCompletions()})() + self.base_url = "https://api.openai.com" + + model = OpenAIChatCompletionsModel(model="gpt-4", openai_client=DummyChatClient()) # type: ignore + + if override_ua is not None: + token = USER_AGENT_OVERRIDE.set(override_ua) + else: + token = None + + try: + await model.get_response( + system_instructions=None, + input="hi", + model_settings=ModelSettings(), + tools=[], + output_schema=None, + handoffs=[], + tracing=ModelTracing.DISABLED, + previous_response_id=None, + conversation_id=None, + ) + finally: + if token is not None: + USER_AGENT_OVERRIDE.reset(token) + + assert "extra_headers" in called_kwargs + assert called_kwargs["extra_headers"]["User-Agent"] == expected_ua + client = AsyncOpenAI(base_url="http://www.notopenai.com") model_settings = ModelSettings() assert ChatCmplHelpers.get_store_param(client, model_settings) is None, ( diff --git a/tests/test_openai_responses.py b/tests/test_openai_responses.py new file mode 100644 index 000000000..81e16c03e --- /dev/null +++ b/tests/test_openai_responses.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from typing import Any + +import pytest +from openai.types.responses import ResponseCompletedEvent + +from agents import ModelSettings, ModelTracing, __version__ +from agents.models.openai_responses import _USER_AGENT_OVERRIDE as RESP_UA, OpenAIResponsesModel +from tests.fake_model import get_response_obj + + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +@pytest.mark.parametrize("override_ua", [None, "test_user_agent"]) +async def test_user_agent_header_responses(override_ua: str | None): + called_kwargs: dict[str, Any] = {} + expected_ua = override_ua or f"Agents/Python {__version__}" + + class DummyStream: + def __aiter__(self): + async def gen(): + yield ResponseCompletedEvent( + type="response.completed", + response=get_response_obj([]), + sequence_number=0, + ) + + return gen() + + class DummyResponses: + async def create(self, **kwargs): + nonlocal called_kwargs + called_kwargs = kwargs + return DummyStream() + + class DummyResponsesClient: + def __init__(self): + self.responses = DummyResponses() + + model = OpenAIResponsesModel(model="gpt-4", openai_client=DummyResponsesClient()) # type: ignore + + if override_ua is not None: + token = RESP_UA.set(override_ua) + else: + token = None + + try: + stream = model.stream_response( + system_instructions=None, + input="hi", + model_settings=ModelSettings(), + tools=[], + output_schema=None, + handoffs=[], + tracing=ModelTracing.DISABLED, + ) + async for _ in stream: + pass + finally: + if token is not None: + RESP_UA.reset(token) + + assert "extra_headers" in called_kwargs + assert called_kwargs["extra_headers"]["User-Agent"] == expected_ua