wehub-resource-sync
404 lines
15 KiB
Python
404 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, cast
|
|
|
|
import httpx
|
|
import litellm
|
|
import pytest
|
|
from litellm.types.utils import Choices, Message, ModelResponse, Usage
|
|
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.extensions.models.litellm_model import LitellmModel
|
|
from agents.items import TResponseInputItem
|
|
from agents.model_settings import ModelSettings
|
|
from agents.models.chatcmpl_converter import Converter
|
|
from agents.models.interface import ModelTracing
|
|
from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel
|
|
from agents.models.reasoning_content_replay import ReasoningContentReplayContext
|
|
|
|
REASONING_CONTENT_MODEL_A = "reasoning-content-model-a"
|
|
REASONING_CONTENT_MODEL_B = "reasoning-content-model-b"
|
|
# The converter currently keys Anthropic thinking-block reconstruction off the model name,
|
|
# so this test model keeps the "anthropic" substring while staying otherwise generic.
|
|
REASONING_CONTENT_MODEL_C = "reasoning-content-model-c-anthropic"
|
|
|
|
|
|
def _second_turn_input_items(model_name: str) -> list[TResponseInputItem]:
|
|
return cast(
|
|
list[TResponseInputItem],
|
|
[
|
|
{"role": "user", "content": "What's the weather in Tokyo?"},
|
|
{
|
|
"id": "__fake_id__",
|
|
"summary": [
|
|
{"text": "I should call the weather tool first.", "type": "summary_text"}
|
|
],
|
|
"type": "reasoning",
|
|
"content": None,
|
|
"encrypted_content": None,
|
|
"status": None,
|
|
"provider_data": {"model": model_name, "response_id": "chatcmpl-test"},
|
|
},
|
|
{
|
|
"arguments": '{"city": "Tokyo"}',
|
|
"call_id": "call_weather_123",
|
|
"name": "get_weather",
|
|
"type": "function_call",
|
|
"id": "__fake_id__",
|
|
"status": None,
|
|
"provider_data": {"model": model_name},
|
|
},
|
|
{
|
|
"type": "function_call_output",
|
|
"call_id": "call_weather_123",
|
|
"output": "The weather in Tokyo is sunny and 22°C.",
|
|
},
|
|
],
|
|
)
|
|
|
|
|
|
def _second_turn_input_items_with_message(model_name: str) -> list[TResponseInputItem]:
|
|
return cast(
|
|
list[TResponseInputItem],
|
|
[
|
|
{"role": "user", "content": "What's the weather in Tokyo?"},
|
|
{
|
|
"id": "__fake_id__",
|
|
"summary": [
|
|
{"text": "I should call the weather tool first.", "type": "summary_text"}
|
|
],
|
|
"type": "reasoning",
|
|
"content": None,
|
|
"encrypted_content": None,
|
|
"status": None,
|
|
"provider_data": {"model": model_name, "response_id": "chatcmpl-test"},
|
|
},
|
|
{
|
|
"id": "__fake_id__",
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"status": "completed",
|
|
"content": [
|
|
{
|
|
"type": "output_text",
|
|
"text": "I'll call the weather tool now.",
|
|
"annotations": [],
|
|
"logprobs": [],
|
|
}
|
|
],
|
|
"provider_data": {"model": model_name, "response_id": "chatcmpl-test"},
|
|
},
|
|
{
|
|
"arguments": '{"city": "Tokyo"}',
|
|
"call_id": "call_weather_123",
|
|
"name": "get_weather",
|
|
"type": "function_call",
|
|
"id": "__fake_id__",
|
|
"status": None,
|
|
"provider_data": {"model": model_name},
|
|
},
|
|
{
|
|
"type": "function_call_output",
|
|
"call_id": "call_weather_123",
|
|
"output": "The weather in Tokyo is sunny and 22°C.",
|
|
},
|
|
],
|
|
)
|
|
|
|
|
|
def _second_turn_input_items_with_file_search(model_name: str) -> list[TResponseInputItem]:
|
|
return cast(
|
|
list[TResponseInputItem],
|
|
[
|
|
{"role": "user", "content": "Find notes about Tokyo weather."},
|
|
{
|
|
"id": "__fake_id__",
|
|
"summary": [
|
|
{"text": "I should search the knowledge base first.", "type": "summary_text"}
|
|
],
|
|
"type": "reasoning",
|
|
"content": None,
|
|
"encrypted_content": None,
|
|
"status": None,
|
|
"provider_data": {"model": model_name, "response_id": "chatcmpl-test"},
|
|
},
|
|
{
|
|
"id": "__fake_file_search_id__",
|
|
"queries": ["Tokyo weather"],
|
|
"status": "completed",
|
|
"type": "file_search_call",
|
|
},
|
|
],
|
|
)
|
|
|
|
|
|
def _second_turn_input_items_with_message_then_reasoning(
|
|
model_name: str,
|
|
) -> list[TResponseInputItem]:
|
|
return cast(
|
|
list[TResponseInputItem],
|
|
[
|
|
{"role": "user", "content": "What's the weather in Tokyo?"},
|
|
{
|
|
"id": "__fake_id__",
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"status": "completed",
|
|
"content": [
|
|
{
|
|
"type": "output_text",
|
|
"text": "I'll call the weather tool now.",
|
|
"annotations": [],
|
|
"logprobs": [],
|
|
}
|
|
],
|
|
"provider_data": {"model": model_name, "response_id": "chatcmpl-test"},
|
|
},
|
|
{
|
|
"id": "__fake_id__",
|
|
"summary": [
|
|
{"text": "I should call the weather tool first.", "type": "summary_text"}
|
|
],
|
|
"type": "reasoning",
|
|
"content": None,
|
|
"encrypted_content": None,
|
|
"status": None,
|
|
"provider_data": {"model": model_name, "response_id": "chatcmpl-test"},
|
|
},
|
|
{
|
|
"arguments": '{"city": "Tokyo"}',
|
|
"call_id": "call_weather_123",
|
|
"name": "get_weather",
|
|
"type": "function_call",
|
|
"id": "__fake_id__",
|
|
"status": None,
|
|
"provider_data": {"model": model_name},
|
|
},
|
|
{
|
|
"type": "function_call_output",
|
|
"call_id": "call_weather_123",
|
|
"output": "The weather in Tokyo is sunny and 22°C.",
|
|
},
|
|
],
|
|
)
|
|
|
|
|
|
def _second_turn_input_items_with_thinking_blocks(model_name: str) -> list[TResponseInputItem]:
|
|
return cast(
|
|
list[TResponseInputItem],
|
|
[
|
|
{"role": "user", "content": "What's the weather in Tokyo?"},
|
|
{
|
|
"id": "__fake_id__",
|
|
"summary": [
|
|
{"text": "I should call the weather tool first.", "type": "summary_text"}
|
|
],
|
|
"type": "reasoning",
|
|
"content": [
|
|
{
|
|
"type": "reasoning_text",
|
|
"text": "First, I need to inspect the request.",
|
|
}
|
|
],
|
|
"encrypted_content": "test-signature",
|
|
"status": None,
|
|
"provider_data": {"model": model_name, "response_id": "chatcmpl-test"},
|
|
},
|
|
{
|
|
"arguments": '{"city": "Tokyo"}',
|
|
"call_id": "call_weather_123",
|
|
"name": "get_weather",
|
|
"type": "function_call",
|
|
"id": "__fake_id__",
|
|
"status": None,
|
|
"provider_data": {"model": model_name},
|
|
},
|
|
{
|
|
"type": "function_call_output",
|
|
"call_id": "call_weather_123",
|
|
"output": "The weather in Tokyo is sunny and 22°C.",
|
|
},
|
|
],
|
|
)
|
|
|
|
|
|
def _assistant_with_tool_calls(messages: list[Any]) -> dict[str, Any]:
|
|
for msg in messages:
|
|
if isinstance(msg, dict) and msg.get("role") == "assistant" and msg.get("tool_calls"):
|
|
return msg
|
|
raise AssertionError("Expected an assistant message with tool_calls.")
|
|
|
|
|
|
def test_converter_keeps_default_reasoning_replay_behavior_for_non_default_model() -> None:
|
|
messages = Converter.items_to_messages(
|
|
_second_turn_input_items(REASONING_CONTENT_MODEL_A),
|
|
model=REASONING_CONTENT_MODEL_A,
|
|
)
|
|
|
|
assistant = _assistant_with_tool_calls(messages)
|
|
assert "reasoning_content" not in assistant
|
|
|
|
|
|
def test_converter_preserves_reasoning_content_across_output_message_with_hook() -> None:
|
|
def should_replay_reasoning_content(_context: ReasoningContentReplayContext) -> bool:
|
|
return True
|
|
|
|
messages = Converter.items_to_messages(
|
|
_second_turn_input_items_with_message(REASONING_CONTENT_MODEL_A),
|
|
model=REASONING_CONTENT_MODEL_A,
|
|
should_replay_reasoning_content=should_replay_reasoning_content,
|
|
)
|
|
|
|
assistant = _assistant_with_tool_calls(messages)
|
|
assert assistant["content"] == "I'll call the weather tool now."
|
|
assert assistant["reasoning_content"] == "I should call the weather tool first."
|
|
|
|
|
|
def test_converter_replays_reasoning_content_when_reasoning_follows_message_with_hook() -> None:
|
|
def should_replay_reasoning_content(_context: ReasoningContentReplayContext) -> bool:
|
|
return True
|
|
|
|
messages = Converter.items_to_messages(
|
|
_second_turn_input_items_with_message_then_reasoning(REASONING_CONTENT_MODEL_A),
|
|
model=REASONING_CONTENT_MODEL_A,
|
|
should_replay_reasoning_content=should_replay_reasoning_content,
|
|
)
|
|
|
|
assistant = _assistant_with_tool_calls(messages)
|
|
assert assistant["content"] == "I'll call the weather tool now."
|
|
assert assistant["reasoning_content"] == "I should call the weather tool first."
|
|
|
|
|
|
def test_converter_replays_reasoning_content_for_file_search_call_with_hook() -> None:
|
|
def should_replay_reasoning_content(_context: ReasoningContentReplayContext) -> bool:
|
|
return True
|
|
|
|
messages = Converter.items_to_messages(
|
|
_second_turn_input_items_with_file_search(REASONING_CONTENT_MODEL_A),
|
|
model=REASONING_CONTENT_MODEL_A,
|
|
should_replay_reasoning_content=should_replay_reasoning_content,
|
|
)
|
|
|
|
assistant = _assistant_with_tool_calls(messages)
|
|
assert assistant["reasoning_content"] == "I should search the knowledge base first."
|
|
assert assistant["tool_calls"][0]["function"]["name"] == "file_search_call"
|
|
|
|
|
|
def test_converter_replays_reasoning_content_with_thinking_blocks_and_hook() -> None:
|
|
def should_replay_reasoning_content(_context: ReasoningContentReplayContext) -> bool:
|
|
return True
|
|
|
|
messages = Converter.items_to_messages(
|
|
_second_turn_input_items_with_thinking_blocks(REASONING_CONTENT_MODEL_C),
|
|
model=REASONING_CONTENT_MODEL_C,
|
|
preserve_thinking_blocks=True,
|
|
should_replay_reasoning_content=should_replay_reasoning_content,
|
|
)
|
|
|
|
assistant = _assistant_with_tool_calls(messages)
|
|
assert assistant["reasoning_content"] == "I should call the weather tool first."
|
|
assert assistant["content"][0]["type"] == "thinking"
|
|
assert assistant["content"][0]["thinking"] == "First, I need to inspect the request."
|
|
|
|
|
|
@pytest.mark.allow_call_model_methods
|
|
@pytest.mark.asyncio
|
|
async def test_openai_chatcompletions_hook_can_enable_reasoning_content_replay() -> None:
|
|
captured: dict[str, Any] = {}
|
|
contexts: list[ReasoningContentReplayContext] = []
|
|
|
|
def should_replay_reasoning_content(context: ReasoningContentReplayContext) -> bool:
|
|
contexts.append(context)
|
|
return context.model == REASONING_CONTENT_MODEL_B
|
|
|
|
class MockChatCompletions:
|
|
async def create(self, **kwargs):
|
|
captured.update(kwargs)
|
|
msg = ChatCompletionMessage(role="assistant", content="done")
|
|
choice = Choice(index=0, message=msg, finish_reason="stop")
|
|
return ChatCompletion(
|
|
id="test-id",
|
|
created=0,
|
|
model=REASONING_CONTENT_MODEL_B,
|
|
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 = httpx.URL("https://example.com/v1/")
|
|
|
|
model = OpenAIChatCompletionsModel(
|
|
model=REASONING_CONTENT_MODEL_B,
|
|
openai_client=cast(Any, MockClient()),
|
|
should_replay_reasoning_content=should_replay_reasoning_content,
|
|
)
|
|
|
|
await model.get_response(
|
|
system_instructions=None,
|
|
input=_second_turn_input_items(REASONING_CONTENT_MODEL_B),
|
|
model_settings=ModelSettings(),
|
|
tools=[],
|
|
output_schema=None,
|
|
handoffs=[],
|
|
tracing=ModelTracing.DISABLED,
|
|
previous_response_id=None,
|
|
)
|
|
|
|
assistant = _assistant_with_tool_calls(cast(list[dict[str, Any]], captured["messages"]))
|
|
assert assistant["reasoning_content"] == "I should call the weather tool first."
|
|
assert len(contexts) == 1
|
|
assert contexts[0].model == REASONING_CONTENT_MODEL_B
|
|
assert contexts[0].base_url == "https://example.com/v1"
|
|
assert contexts[0].reasoning.origin_model == REASONING_CONTENT_MODEL_B
|
|
|
|
|
|
@pytest.mark.allow_call_model_methods
|
|
@pytest.mark.asyncio
|
|
async def test_litellm_hook_can_enable_reasoning_content_replay(monkeypatch) -> None:
|
|
captured: dict[str, Any] = {}
|
|
contexts: list[ReasoningContentReplayContext] = []
|
|
|
|
def should_replay_reasoning_content(context: ReasoningContentReplayContext) -> bool:
|
|
contexts.append(context)
|
|
return context.model == REASONING_CONTENT_MODEL_B
|
|
|
|
async def fake_acompletion(model, messages=None, **kwargs):
|
|
captured["messages"] = messages
|
|
msg = Message(role="assistant", content="done")
|
|
choice = Choices(index=0, message=msg)
|
|
return ModelResponse(choices=[choice], usage=Usage(0, 0, 0))
|
|
|
|
monkeypatch.setattr(litellm, "acompletion", fake_acompletion)
|
|
|
|
model = LitellmModel(
|
|
model=REASONING_CONTENT_MODEL_B,
|
|
should_replay_reasoning_content=should_replay_reasoning_content,
|
|
)
|
|
|
|
await model.get_response(
|
|
system_instructions=None,
|
|
input=_second_turn_input_items(REASONING_CONTENT_MODEL_B),
|
|
model_settings=ModelSettings(),
|
|
tools=[],
|
|
output_schema=None,
|
|
handoffs=[],
|
|
tracing=ModelTracing.DISABLED,
|
|
previous_response_id=None,
|
|
)
|
|
|
|
assistant = _assistant_with_tool_calls(cast(list[dict[str, Any]], captured["messages"]))
|
|
assert assistant["reasoning_content"] == "I should call the weather tool first."
|
|
assert len(contexts) == 1
|
|
assert contexts[0].model == REASONING_CONTENT_MODEL_B
|
|
assert contexts[0].base_url is None
|
|
assert contexts[0].reasoning.origin_model == REASONING_CONTENT_MODEL_B
|