Kazuhiro Sera
459 lines
15 KiB
Python
459 lines
15 KiB
Python
from typing import Any
|
|
|
|
import httpx2
|
|
import litellm
|
|
import pytest
|
|
from httpx import Headers, Response
|
|
from litellm.exceptions import RateLimitError
|
|
from litellm.types.utils import Choices, Message, ModelResponse, Usage
|
|
from openai import APIConnectionError, omit
|
|
from openai.types.chat.chat_completion import ChatCompletion, Choice
|
|
from openai.types.chat.chat_completion_message import ChatCompletionMessage
|
|
from openai.types.completion_usage import CompletionUsage
|
|
|
|
from agents import Agent, function_tool, handoff
|
|
from agents.extensions.models.litellm_model import LitellmModel
|
|
from agents.model_settings import ModelSettings
|
|
from agents.models._retry_runtime import provider_managed_retries_disabled
|
|
from agents.models.interface import ModelTracing
|
|
from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel
|
|
from agents.retry import ModelRetryAdviceRequest, ModelRetrySettings
|
|
from agents.tool import Tool
|
|
|
|
|
|
@pytest.mark.allow_call_model_methods
|
|
@pytest.mark.asyncio
|
|
async def test_litellm_kwargs_forwarded(monkeypatch):
|
|
"""
|
|
Test that kwargs from ModelSettings are forwarded to litellm.acompletion.
|
|
"""
|
|
captured: dict[str, object] = {}
|
|
|
|
async def fake_acompletion(model, messages=None, **kwargs):
|
|
captured.update(kwargs)
|
|
msg = Message(role="assistant", content="test response")
|
|
choice = Choices(index=0, message=msg)
|
|
return ModelResponse(choices=[choice], usage=Usage(0, 0, 0))
|
|
|
|
monkeypatch.setattr(litellm, "acompletion", fake_acompletion)
|
|
|
|
settings = ModelSettings(
|
|
temperature=0.5,
|
|
extra_args={
|
|
"custom_param": "custom_value",
|
|
"seed": 42,
|
|
"stop": ["END"],
|
|
"logit_bias": {123: -100},
|
|
},
|
|
)
|
|
model = LitellmModel(model="test-model")
|
|
|
|
await model.get_response(
|
|
system_instructions=None,
|
|
input="test input",
|
|
model_settings=settings,
|
|
tools=[],
|
|
output_schema=None,
|
|
handoffs=[],
|
|
tracing=ModelTracing.DISABLED,
|
|
previous_response_id=None,
|
|
conversation_id=None,
|
|
)
|
|
|
|
# Verify that all kwargs were passed through
|
|
assert captured["custom_param"] == "custom_value"
|
|
assert captured["seed"] == 42
|
|
assert captured["stop"] == ["END"]
|
|
assert captured["logit_bias"] == {123: -100}
|
|
|
|
# Verify regular parameters are still passed
|
|
assert captured["temperature"] == 0.5
|
|
|
|
|
|
@pytest.mark.allow_call_model_methods
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("parallel_tool_calls", [True, False])
|
|
@pytest.mark.parametrize("tool_source", ["none", "function", "handoff"])
|
|
async def test_litellm_only_forwards_parallel_tool_calls_with_converted_tools(
|
|
monkeypatch, parallel_tool_calls: bool, tool_source: str
|
|
):
|
|
captured: dict[str, object] = {}
|
|
|
|
async def fake_acompletion(model, messages=None, **kwargs):
|
|
captured.update(kwargs)
|
|
message = Message(role="assistant", content="test response")
|
|
return ModelResponse(choices=[Choices(index=0, message=message)], usage=Usage(0, 0, 0))
|
|
|
|
monkeypatch.setattr(litellm, "acompletion", fake_acompletion)
|
|
|
|
tools: list[Tool] = (
|
|
[function_tool(lambda: "ok", name_override="test_tool")]
|
|
if tool_source == "function"
|
|
else []
|
|
)
|
|
handoffs = [handoff(Agent(name="handoff"))] if tool_source == "handoff" else []
|
|
|
|
await LitellmModel(model="test-model").get_response(
|
|
system_instructions=None,
|
|
input="test input",
|
|
model_settings=ModelSettings(parallel_tool_calls=parallel_tool_calls),
|
|
tools=tools,
|
|
output_schema=None,
|
|
handoffs=handoffs,
|
|
tracing=ModelTracing.DISABLED,
|
|
previous_response_id=None,
|
|
conversation_id=None,
|
|
)
|
|
|
|
expected_parallel_tool_calls = parallel_tool_calls if tool_source != "none" else None
|
|
assert captured["parallel_tool_calls"] is expected_parallel_tool_calls
|
|
assert (captured["tools"] is not None) is (tool_source != "none")
|
|
|
|
|
|
@pytest.mark.allow_call_model_methods
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("parallel_tool_calls", [True, False])
|
|
@pytest.mark.parametrize("tool_source", ["none", "function", "handoff"])
|
|
async def test_openai_only_forwards_parallel_tool_calls_with_converted_tools(
|
|
parallel_tool_calls: bool, tool_source: str
|
|
):
|
|
captured: dict[str, object] = {}
|
|
|
|
class MockChatCompletions:
|
|
async def create(self, **kwargs):
|
|
captured.update(kwargs)
|
|
msg = ChatCompletionMessage(role="assistant", content="test response")
|
|
return ChatCompletion(
|
|
id="test-id",
|
|
created=0,
|
|
model="gpt-4",
|
|
object="chat.completion",
|
|
choices=[Choice(index=0, message=msg, finish_reason="stop")],
|
|
usage=CompletionUsage(completion_tokens=5, prompt_tokens=10, total_tokens=15),
|
|
)
|
|
|
|
class MockChat:
|
|
def __init__(self):
|
|
self.completions = MockChatCompletions()
|
|
|
|
class MockClient:
|
|
def __init__(self):
|
|
self.chat = MockChat()
|
|
self.base_url = "https://api.openai.com/v1"
|
|
|
|
tools: list[Tool] = (
|
|
[function_tool(lambda: "ok", name_override="test_tool")]
|
|
if tool_source == "function"
|
|
else []
|
|
)
|
|
handoffs = [handoff(Agent(name="handoff"))] if tool_source == "handoff" else []
|
|
|
|
model = OpenAIChatCompletionsModel(model="gpt-4", openai_client=MockClient()) # type: ignore
|
|
await model.get_response(
|
|
system_instructions=None,
|
|
input="test input",
|
|
model_settings=ModelSettings(parallel_tool_calls=parallel_tool_calls),
|
|
tools=tools,
|
|
output_schema=None,
|
|
handoffs=handoffs,
|
|
tracing=ModelTracing.DISABLED,
|
|
previous_response_id=None,
|
|
conversation_id=None,
|
|
)
|
|
|
|
if tool_source == "none":
|
|
assert captured["parallel_tool_calls"] is omit
|
|
assert captured["tools"] is omit
|
|
else:
|
|
assert captured["parallel_tool_calls"] is parallel_tool_calls
|
|
assert captured["tools"] is not omit
|
|
|
|
|
|
@pytest.mark.allow_call_model_methods
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("use_dictionary", [False, True], ids=["model-settings", "dictionary"])
|
|
async def test_litellm_normalizes_dictionary_agent_model_settings(
|
|
monkeypatch, use_dictionary: bool
|
|
):
|
|
captured: dict[str, object] = {}
|
|
|
|
async def fake_acompletion(model, messages=None, **kwargs):
|
|
captured.update(kwargs)
|
|
message = Message(role="assistant", content="test response")
|
|
return ModelResponse(choices=[Choices(index=0, message=message)], usage=Usage(0, 0, 0))
|
|
|
|
monkeypatch.setattr(litellm, "acompletion", fake_acompletion)
|
|
settings: dict[str, Any] = {"temperature": 0.0, "reasoning": {"effort": "low"}}
|
|
model = LitellmModel(model="test-model")
|
|
agent = Agent(
|
|
name="test",
|
|
model=model,
|
|
model_settings=settings if use_dictionary else ModelSettings(**settings),
|
|
)
|
|
|
|
await model.get_response(
|
|
system_instructions=None,
|
|
input="test input",
|
|
model_settings=agent.model_settings,
|
|
tools=[],
|
|
output_schema=None,
|
|
handoffs=[],
|
|
tracing=ModelTracing.DISABLED,
|
|
previous_response_id=None,
|
|
conversation_id=None,
|
|
)
|
|
|
|
assert captured["temperature"] == 0.0
|
|
assert captured["reasoning_effort"] == "low"
|
|
|
|
|
|
@pytest.mark.allow_call_model_methods
|
|
@pytest.mark.asyncio
|
|
async def test_openai_chatcompletions_kwargs_forwarded(monkeypatch):
|
|
"""
|
|
Test that kwargs from ModelSettings are forwarded to OpenAI chat completions API.
|
|
"""
|
|
captured: dict[str, object] = {}
|
|
|
|
class MockChatCompletions:
|
|
async def create(self, **kwargs):
|
|
captured.update(kwargs)
|
|
msg = ChatCompletionMessage(role="assistant", content="test response")
|
|
choice = Choice(index=0, message=msg, finish_reason="stop")
|
|
return ChatCompletion(
|
|
id="test-id",
|
|
created=0,
|
|
model="gpt-4",
|
|
object="chat.completion",
|
|
choices=[choice],
|
|
usage=CompletionUsage(completion_tokens=5, prompt_tokens=10, total_tokens=15),
|
|
)
|
|
|
|
class MockChat:
|
|
def __init__(self):
|
|
self.completions = MockChatCompletions()
|
|
|
|
class MockClient:
|
|
def __init__(self):
|
|
self.chat = MockChat()
|
|
self.base_url = "https://api.openai.com/v1"
|
|
|
|
settings = ModelSettings(
|
|
temperature=0.7,
|
|
extra_args={
|
|
"seed": 123,
|
|
"logit_bias": {456: 10},
|
|
"stop": ["STOP", "END"],
|
|
"user": "test-user",
|
|
},
|
|
)
|
|
|
|
mock_client = MockClient()
|
|
model = OpenAIChatCompletionsModel(model="gpt-4", openai_client=mock_client) # type: ignore
|
|
|
|
await model.get_response(
|
|
system_instructions="Test system",
|
|
input="test input",
|
|
model_settings=settings,
|
|
tools=[],
|
|
output_schema=None,
|
|
handoffs=[],
|
|
tracing=ModelTracing.DISABLED,
|
|
previous_response_id=None,
|
|
)
|
|
|
|
# Verify that all kwargs were passed through
|
|
assert captured["seed"] == 123
|
|
assert captured["logit_bias"] == {456: 10}
|
|
assert captured["stop"] == ["STOP", "END"]
|
|
assert captured["user"] == "test-user"
|
|
|
|
# Verify regular parameters are still passed
|
|
assert captured["temperature"] == 0.7
|
|
|
|
|
|
@pytest.mark.allow_call_model_methods
|
|
@pytest.mark.asyncio
|
|
async def test_empty_kwargs_handling(monkeypatch):
|
|
"""
|
|
Test that empty or None kwargs are handled gracefully.
|
|
"""
|
|
captured: dict[str, object] = {}
|
|
|
|
async def fake_acompletion(model, messages=None, **kwargs):
|
|
captured.update(kwargs)
|
|
msg = Message(role="assistant", content="test response")
|
|
choice = Choices(index=0, message=msg)
|
|
return ModelResponse(choices=[choice], usage=Usage(0, 0, 0))
|
|
|
|
monkeypatch.setattr(litellm, "acompletion", fake_acompletion)
|
|
|
|
# Test with None kwargs
|
|
settings_none = ModelSettings(temperature=0.5, extra_args=None)
|
|
model = LitellmModel(model="test-model")
|
|
|
|
await model.get_response(
|
|
system_instructions=None,
|
|
input="test input",
|
|
model_settings=settings_none,
|
|
tools=[],
|
|
output_schema=None,
|
|
handoffs=[],
|
|
tracing=ModelTracing.DISABLED,
|
|
previous_response_id=None,
|
|
)
|
|
|
|
# Should work without error and include regular parameters
|
|
assert captured["temperature"] == 0.5
|
|
|
|
# Test with empty dict
|
|
captured.clear()
|
|
settings_empty = ModelSettings(temperature=0.3, extra_args={})
|
|
|
|
await model.get_response(
|
|
system_instructions=None,
|
|
input="test input",
|
|
model_settings=settings_empty,
|
|
tools=[],
|
|
output_schema=None,
|
|
handoffs=[],
|
|
tracing=ModelTracing.DISABLED,
|
|
previous_response_id=None,
|
|
)
|
|
|
|
# Should work without error and include regular parameters
|
|
assert captured["temperature"] == 0.3
|
|
|
|
|
|
@pytest.mark.allow_call_model_methods
|
|
@pytest.mark.asyncio
|
|
async def test_reasoning_effort_falls_back_to_extra_args(monkeypatch):
|
|
"""
|
|
Ensure reasoning_effort from extra_args is promoted when reasoning settings are missing.
|
|
"""
|
|
captured: dict[str, object] = {}
|
|
|
|
async def fake_acompletion(model, messages=None, **kwargs):
|
|
captured.update(kwargs)
|
|
msg = Message(role="assistant", content="test response")
|
|
choice = Choices(index=0, message=msg)
|
|
return ModelResponse(choices=[choice], usage=Usage(0, 0, 0))
|
|
|
|
monkeypatch.setattr(litellm, "acompletion", fake_acompletion)
|
|
|
|
# GitHub issue context: https://github.com/openai/openai-agents-python/issues/1764.
|
|
settings = ModelSettings(
|
|
extra_args={"reasoning_effort": "none", "custom_param": "custom_value"}
|
|
)
|
|
model = LitellmModel(model="test-model")
|
|
|
|
await model.get_response(
|
|
system_instructions=None,
|
|
input="test input",
|
|
model_settings=settings,
|
|
tools=[],
|
|
output_schema=None,
|
|
handoffs=[],
|
|
tracing=ModelTracing.DISABLED,
|
|
previous_response_id=None,
|
|
)
|
|
|
|
assert captured["reasoning_effort"] == "none"
|
|
assert captured["custom_param"] == "custom_value"
|
|
assert settings.extra_args == {"reasoning_effort": "none", "custom_param": "custom_value"}
|
|
|
|
|
|
@pytest.mark.allow_call_model_methods
|
|
@pytest.mark.asyncio
|
|
async def test_litellm_retry_settings_do_not_leak_and_disable_provider_retries_on_runner_retry(
|
|
monkeypatch,
|
|
):
|
|
"""Runner retries should disable LiteLLM's own retries without forwarding SDK retry config."""
|
|
|
|
captured: dict[str, object] = {}
|
|
|
|
async def fake_acompletion(model, messages=None, **kwargs):
|
|
captured.update(kwargs)
|
|
msg = Message(role="assistant", content="test response")
|
|
choice = Choices(index=0, message=msg)
|
|
return ModelResponse(choices=[choice], usage=Usage(0, 0, 0))
|
|
|
|
monkeypatch.setattr(litellm, "acompletion", fake_acompletion)
|
|
|
|
settings = ModelSettings(
|
|
retry=ModelRetrySettings(
|
|
max_retries=2,
|
|
backoff={"initial_delay": 0.25, "jitter": False},
|
|
),
|
|
extra_args={"max_retries": 7, "num_retries": 6, "custom_param": "custom_value"},
|
|
)
|
|
model = LitellmModel(model="test-model")
|
|
|
|
with provider_managed_retries_disabled(True):
|
|
await model.get_response(
|
|
system_instructions=None,
|
|
input="test input",
|
|
model_settings=settings,
|
|
tools=[],
|
|
output_schema=None,
|
|
handoffs=[],
|
|
tracing=ModelTracing.DISABLED,
|
|
previous_response_id=None,
|
|
conversation_id=None,
|
|
)
|
|
|
|
assert settings.retry is not None
|
|
assert settings.retry.backoff is not None
|
|
assert captured["custom_param"] == "custom_value"
|
|
assert captured["max_retries"] == 0
|
|
assert captured["num_retries"] == 0
|
|
assert "retry" not in captured
|
|
|
|
|
|
def test_litellm_get_retry_advice_uses_response_headers() -> None:
|
|
"""LiteLLM retry advice should expose OpenAI-compatible retry headers."""
|
|
|
|
model = LitellmModel(model="test-model")
|
|
error = RateLimitError(
|
|
message="rate limited",
|
|
llm_provider="openai",
|
|
model="gpt-4o-mini",
|
|
response=Response(
|
|
status_code=429,
|
|
headers=Headers({"x-should-retry": "true", "retry-after-ms": "250"}),
|
|
),
|
|
)
|
|
|
|
advice = model.get_retry_advice(
|
|
ModelRetryAdviceRequest(
|
|
error=error,
|
|
attempt=1,
|
|
stream=False,
|
|
)
|
|
)
|
|
|
|
assert advice is not None
|
|
assert advice.suggested is True
|
|
assert advice.retry_after == 0.25
|
|
|
|
|
|
def test_litellm_get_retry_advice_keeps_stateful_transport_failures_ambiguous() -> None:
|
|
model = LitellmModel(model="test-model")
|
|
error = APIConnectionError(
|
|
message="connection error",
|
|
request=httpx2.Request("POST", "https://api.openai.com/v1/responses"),
|
|
)
|
|
|
|
advice = model.get_retry_advice(
|
|
ModelRetryAdviceRequest(
|
|
error=error,
|
|
attempt=1,
|
|
stream=False,
|
|
previous_response_id="resp_prev",
|
|
)
|
|
)
|
|
|
|
assert advice is not None
|
|
assert advice.suggested is True
|
|
assert advice.replay_safety is None
|