openai--openai-agents-python/tests/models/test_litellm_extra_body.py

import logging

import litellm
import pytest
from litellm.types.utils import Choices, Message, ModelResponse, Usage

from agents.extensions.models.litellm_model import LitellmModel
from agents.model_settings import ModelSettings
from agents.models.interface import ModelTracing


@pytest.mark.allow_call_model_methods
@pytest.mark.asyncio
async def test_extra_body_is_forwarded(monkeypatch):
    """
    Forward `extra_body` via LiteLLM's dedicated kwarg.

    This ensures that provider-specific request fields stay nested under `extra_body`
    so LiteLLM can merge them into the upstream request body itself.
    """
    captured: dict[str, object] = {}

    async def fake_acompletion(model, messages=None, **kwargs):
        captured.update(kwargs)
        msg = Message(role="assistant", content="ok")
        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.1, extra_body={"cached_content": "some_cache", "foo": 123}
    )
    model = LitellmModel(model="test-model")

    await model.get_response(
        system_instructions=None,
        input=[],
        model_settings=settings,
        tools=[],
        output_schema=None,
        handoffs=[],
        tracing=ModelTracing.DISABLED,
        previous_response_id=None,
    )

    assert captured["extra_body"] == {"cached_content": "some_cache", "foo": 123}
    assert "cached_content" not in captured
    assert "foo" not in captured


@pytest.mark.allow_call_model_methods
@pytest.mark.asyncio
async def test_extra_body_reasoning_effort_is_promoted(monkeypatch):
    """
    Ensure reasoning_effort from extra_body is promoted to the top-level parameter.
    """
    captured: dict[str, object] = {}

    async def fake_acompletion(model, messages=None, **kwargs):
        captured.update(kwargs)
        msg = Message(role="assistant", content="ok")
        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_body={"reasoning_effort": "none", "cached_content": "some_cache"}
    )
    model = LitellmModel(model="test-model")

    await model.get_response(
        system_instructions=None,
        input=[],
        model_settings=settings,
        tools=[],
        output_schema=None,
        handoffs=[],
        tracing=ModelTracing.DISABLED,
        previous_response_id=None,
    )

    assert captured["reasoning_effort"] == "none"
    assert captured["extra_body"] == {"cached_content": "some_cache"}
    assert settings.extra_body == {"reasoning_effort": "none", "cached_content": "some_cache"}


@pytest.mark.allow_call_model_methods
@pytest.mark.asyncio
async def test_reasoning_effort_prefers_model_settings(monkeypatch):
    """
    Verify explicit ModelSettings.reasoning takes precedence over extra_body entries.
    """
    from openai.types.shared import Reasoning

    captured: dict[str, object] = {}

    async def fake_acompletion(model, messages=None, **kwargs):
        captured.update(kwargs)
        msg = Message(role="assistant", content="ok")
        choice = Choices(index=0, message=msg)
        return ModelResponse(choices=[choice], usage=Usage(0, 0, 0))

    monkeypatch.setattr(litellm, "acompletion", fake_acompletion)
    settings = ModelSettings(
        reasoning=Reasoning(effort="low"),
        extra_body={"reasoning_effort": "high"},
    )
    model = LitellmModel(model="test-model")

    await model.get_response(
        system_instructions=None,
        input=[],
        model_settings=settings,
        tools=[],
        output_schema=None,
        handoffs=[],
        tracing=ModelTracing.DISABLED,
        previous_response_id=None,
    )

    # reasoning_effort is string when no summary is provided (backward compatible)
    assert captured["reasoning_effort"] == "low"
    assert "extra_body" not in captured
    assert settings.extra_body == {"reasoning_effort": "high"}


@pytest.mark.allow_call_model_methods
@pytest.mark.asyncio
async def test_extra_body_reasoning_effort_overrides_extra_args(monkeypatch):
    """
    Ensure extra_body reasoning_effort wins over extra_args when both are provided.
    """
    captured: dict[str, object] = {}

    async def fake_acompletion(model, messages=None, **kwargs):
        captured.update(kwargs)
        msg = Message(role="assistant", content="ok")
        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_body={"reasoning_effort": "none"},
        extra_args={"reasoning_effort": "low", "custom_param": "custom"},
    )
    model = LitellmModel(model="test-model")

    await model.get_response(
        system_instructions=None,
        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"
    assert "extra_body" not in captured
    assert settings.extra_args == {"reasoning_effort": "low", "custom_param": "custom"}


@pytest.mark.allow_call_model_methods
@pytest.mark.asyncio
async def test_extra_body_metadata_stays_nested(monkeypatch):
    """
    Keep extra_body metadata nested even when top-level metadata is also set.

    LiteLLM resolves top-level metadata and extra_body separately. Flattening the nested
    metadata dict loses the caller's intended request shape for OpenAI-compatible proxies.
    """
    captured: dict[str, object] = {}

    async def fake_acompletion(model, messages=None, **kwargs):
        captured.update(kwargs)
        msg = Message(role="assistant", content="ok")
        choice = Choices(index=0, message=msg)
        return ModelResponse(choices=[choice], usage=Usage(0, 0, 0))

    monkeypatch.setattr(litellm, "acompletion", fake_acompletion)
    settings = ModelSettings(
        metadata={"sdk": "agents"},
        extra_body={
            "metadata": {"trace_user_id": "user-123", "generation_id": "gen-456"},
            "cached_content": "some_cache",
        },
    )
    model = LitellmModel(model="test-model")

    await model.get_response(
        system_instructions=None,
        input=[],
        model_settings=settings,
        tools=[],
        output_schema=None,
        handoffs=[],
        tracing=ModelTracing.DISABLED,
        previous_response_id=None,
    )

    assert captured["metadata"] == {"sdk": "agents"}
    assert captured["extra_body"] == {
        "metadata": {"trace_user_id": "user-123", "generation_id": "gen-456"},
        "cached_content": "some_cache",
    }


@pytest.mark.allow_call_model_methods
@pytest.mark.asyncio
@pytest.mark.parametrize(
    "model_name",
    [
        "openai/gpt-5-mini",
        "anthropic/claude-sonnet-4-5",
        "gemini/gemini-2.5-pro",
    ],
)
async def test_reasoning_summary_uses_scalar_effort_and_warns(
    monkeypatch, caplog: pytest.LogCaptureFixture, model_name: str
):
    """
    Ensure reasoning.summary does not change the LiteLLM chat-completions argument shape.

    LitellmModel should continue to pass a scalar reasoning_effort value and warn that summary
    is ignored on this path, regardless of the provider encoded in the model string.
    """
    from openai.types.shared import Reasoning

    captured: dict[str, object] = {}

    async def fake_acompletion(model, messages=None, **kwargs):
        captured.update(kwargs)
        msg = Message(role="assistant", content="ok")
        choice = Choices(index=0, message=msg)
        return ModelResponse(choices=[choice], usage=Usage(0, 0, 0))

    monkeypatch.setattr(litellm, "acompletion", fake_acompletion)
    settings = ModelSettings(
        reasoning=Reasoning(effort="medium", summary="auto"),
    )
    model = LitellmModel(model=model_name)

    with caplog.at_level(logging.WARNING, logger="openai.agents"):
        await model.get_response(
            system_instructions=None,
            input=[],
            model_settings=settings,
            tools=[],
            output_schema=None,
            handoffs=[],
            tracing=ModelTracing.DISABLED,
            previous_response_id=None,
        )

    assert captured["reasoning_effort"] == "medium"
    warning_messages = [
        record.message
        for record in caplog.records
        if "does not forward Reasoning.summary" in record.message
    ]
    assert len(warning_messages) == 1