feat: #2669 add opt-in reasoning content replay for chat completion models (#2670)

This commit is contained in:
Kazuhiro Sera
2026-03-20 00:20:36 -07:00
committed by GitHub
parent f0df572428
commit 34ff8481bb
6 changed files with 581 additions and 19 deletions
@@ -49,6 +49,7 @@ from ...models.chatcmpl_stream_handler import ChatCmplStreamHandler
from ...models.fake_id import FAKE_RESPONSES_ID
from ...models.interface import Model, ModelTracing
from ...models.openai_responses import Converter as OpenAIResponsesConverter
from ...models.reasoning_content_replay import ShouldReplayReasoningContent
from ...retry import ModelRetryAdvice, ModelRetryAdviceRequest
from ...tool import Tool
from ...tracing import generation_span
@@ -146,10 +147,12 @@ class LitellmModel(Model):
model: str,
base_url: str | None = None,
api_key: str | None = None,
should_replay_reasoning_content: ShouldReplayReasoningContent | None = None,
):
self.model = model
self.base_url = base_url
self.api_key = api_key
self.should_replay_reasoning_content = should_replay_reasoning_content
def get_retry_advice(self, request: ModelRetryAdviceRequest) -> ModelRetryAdvice | None:
# LiteLLM exceptions mirror OpenAI-style status/header fields.
@@ -383,9 +386,11 @@ class LitellmModel(Model):
converted_messages = Converter.items_to_messages(
input,
base_url=self.base_url,
preserve_thinking_blocks=preserve_thinking_blocks,
preserve_tool_output_all_content=True,
model=self.model,
should_replay_reasoning_content=self.should_replay_reasoning_content,
)
# Fix message ordering: reorder to ensure tool_use comes before tool_result.
+48 -18
View File
@@ -55,6 +55,12 @@ from ..tool import (
ensure_tool_choice_supports_backend,
)
from .fake_id import FAKE_RESPONSES_ID
from .reasoning_content_replay import (
ReasoningContentReplayContext,
ReasoningContentSource,
ShouldReplayReasoningContent,
default_should_replay_reasoning_content,
)
ResponseInputContentWithAudioParam = Union[
ResponseInputContentParam,
@@ -422,6 +428,8 @@ class Converter:
model: str | None = None,
preserve_thinking_blocks: bool = False,
preserve_tool_output_all_content: bool = False,
base_url: str | None = None,
should_replay_reasoning_content: ShouldReplayReasoningContent | None = None,
) -> list[ChatCompletionMessageParam]:
"""
Convert a sequence of 'Item' objects into a list of ChatCompletionMessageParam.
@@ -441,6 +449,12 @@ class Converter:
When True, all content types including images are preserved. This is useful
for model providers (e.g. Anthropic via LiteLLM) that support processing
non-text content in tool results.
base_url: The request base URL, if the caller knows the concrete endpoint.
This is used by reasoning-content replay hooks to distinguish direct
provider calls from proxy or gateway requests.
should_replay_reasoning_content: Optional hook that decides whether a
reasoning item should be replayed into the next assistant message as
`reasoning_content`.
Rules:
- EasyInputMessage or InputMessage (role=user) => ChatCompletionUserMessageParam
@@ -464,8 +478,9 @@ class Converter:
current_assistant_msg: ChatCompletionAssistantMessageParam | None = None
pending_thinking_blocks: list[dict[str, str]] | None = None
pending_reasoning_content: str | None = None # For DeepSeek reasoning_content
normalized_base_url = base_url.rstrip("/") if base_url is not None else None
def flush_assistant_message() -> None:
def flush_assistant_message(*, clear_pending_reasoning_content: bool = True) -> None:
nonlocal current_assistant_msg, pending_reasoning_content
if current_assistant_msg is not None:
# The API doesn't support empty arrays for tool_calls
@@ -475,7 +490,15 @@ class Converter:
pending_reasoning_content = None
result.append(current_assistant_msg)
current_assistant_msg = None
else:
elif clear_pending_reasoning_content:
pending_reasoning_content = None
def apply_pending_reasoning_content(
assistant_msg: ChatCompletionAssistantMessageParam,
) -> None:
nonlocal pending_reasoning_content
if pending_reasoning_content:
assistant_msg["reasoning_content"] = pending_reasoning_content # type: ignore[typeddict-unknown-key]
pending_reasoning_content = None
def ensure_assistant_message() -> ChatCompletionAssistantMessageParam:
@@ -485,6 +508,8 @@ class Converter:
current_assistant_msg["content"] = None
current_assistant_msg["tool_calls"] = []
apply_pending_reasoning_content(current_assistant_msg)
return current_assistant_msg
for item in items:
@@ -553,7 +578,9 @@ class Converter:
# 3) response output message => assistant
elif resp_msg := cls.maybe_response_output_message(item):
flush_assistant_message()
# A reasoning item can be followed by an assistant message and then tool calls
# in the same turn, so preserve pending reasoning_content across this flush.
flush_assistant_message(clear_pending_reasoning_content=False)
new_asst = ChatCompletionAssistantMessageParam(role="assistant")
contents = resp_msg["content"]
@@ -594,6 +621,7 @@ class Converter:
pending_thinking_blocks = None # Clear after using
new_asst["tool_calls"] = []
apply_pending_reasoning_content(new_asst)
current_assistant_msg = new_asst
# 4) function/file-search calls => attach to assistant
@@ -619,11 +647,6 @@ class Converter:
elif func_call := cls.maybe_function_tool_call(item):
asst = ensure_assistant_message()
# If we have pending reasoning content for DeepSeek, add it to the assistant message
if pending_reasoning_content:
asst["reasoning_content"] = pending_reasoning_content # type: ignore[typeddict-unknown-key]
pending_reasoning_content = None # Clear after using
# If we have pending thinking blocks, use them as the content
# This is required for Anthropic API tool calls with interleaved thinking
if pending_thinking_blocks:
@@ -708,6 +731,7 @@ class Converter:
item_provider_data: dict[str, Any] = reasoning_item.get("provider_data", {}) # type: ignore[assignment]
item_model = item_provider_data.get("model", "")
should_replay = False
if (
model
@@ -740,17 +764,23 @@ class Converter:
# This preserves the original behavior
pending_thinking_blocks = reconstructed_thinking_blocks
# DeepSeek requires reasoning_content field in assistant messages with tool calls
# Items may not all originate from DeepSeek, so need to check for model match.
# For backward compatibility, if provider_data is missing, ignore the check.
elif (
model
and "deepseek" in model.lower()
and (
(item_model and "deepseek" in item_model.lower())
or item_provider_data == {}
if model is not None:
replay_context = ReasoningContentReplayContext(
model=model,
base_url=normalized_base_url,
reasoning=ReasoningContentSource(
item=reasoning_item,
origin_model=item_model or None,
provider_data=item_provider_data,
),
)
):
should_replay = (
should_replay_reasoning_content(replay_context)
if should_replay_reasoning_content is not None
else default_should_replay_reasoning_content(replay_context)
)
if should_replay:
summary_items = reasoning_item.get("summary", [])
if summary_items:
reasoning_texts = []
+9 -1
View File
@@ -39,6 +39,7 @@ from .chatcmpl_stream_handler import ChatCmplStreamHandler
from .fake_id import FAKE_RESPONSES_ID
from .interface import Model, ModelTracing
from .openai_responses import Converter as OpenAIResponsesConverter
from .reasoning_content_replay import ShouldReplayReasoningContent
if TYPE_CHECKING:
from ..model_settings import ModelSettings
@@ -53,9 +54,11 @@ class OpenAIChatCompletionsModel(Model):
self,
model: str | ChatModel,
openai_client: AsyncOpenAI,
should_replay_reasoning_content: ShouldReplayReasoningContent | None = None,
) -> None:
self.model = model
self._client = openai_client
self.should_replay_reasoning_content = should_replay_reasoning_content
def _non_null_or_omit(self, value: Any) -> Any:
return value if value is not None else omit
@@ -314,7 +317,12 @@ class OpenAIChatCompletionsModel(Model):
prompt: ResponsePromptParam | None = None,
) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]:
self._validate_official_openai_input_content_types(input)
converted_messages = Converter.items_to_messages(input, model=self.model)
converted_messages = Converter.items_to_messages(
input,
model=self.model,
base_url=str(self._client.base_url),
should_replay_reasoning_content=self.should_replay_reasoning_content,
)
if system_instructions:
converted_messages.insert(
@@ -0,0 +1,59 @@
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any, Callable
@dataclass
class ReasoningContentSource:
"""The reasoning item being considered for replay into the next request."""
item: Any
"""The raw reasoning item."""
origin_model: str | None
"""The model that originally produced the reasoning item, if known."""
provider_data: Mapping[str, Any]
"""Provider-specific metadata captured on the reasoning item."""
@dataclass
class ReasoningContentReplayContext:
"""Context passed to reasoning-content replay hooks."""
model: str
"""The model that will receive the next Chat Completions request."""
base_url: str | None
"""The request base URL, if the SDK knows the concrete endpoint."""
reasoning: ReasoningContentSource
"""The reasoning item candidate being evaluated for replay."""
ShouldReplayReasoningContent = Callable[[ReasoningContentReplayContext], bool]
def default_should_replay_reasoning_content(context: ReasoningContentReplayContext) -> bool:
"""Return whether the SDK should replay reasoning content by default."""
if "deepseek" not in context.model.lower():
return False
origin_model = context.reasoning.origin_model
# Replay only when the current request targets DeepSeek and the reasoning item either
# came from a DeepSeek model or predates provider tracking. This avoids mixing reasoning
# content from a different model family into the DeepSeek assistant message.
return (
origin_model is not None and "deepseek" in origin_model.lower()
) or context.reasoning.provider_data == {}
__all__ = [
"ReasoningContentReplayContext",
"ReasoningContentSource",
"ShouldReplayReasoningContent",
"default_should_replay_reasoning_content",
]
@@ -0,0 +1,403 @@
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
+57
View File
@@ -248,6 +248,63 @@ def test_anthropic_thinking_blocks_with_tool_calls():
assert cast(list[Any], tool_calls)[0]["function"]["name"] == "get_weather"
def test_items_to_messages_preserves_positional_bool_arguments():
"""
Preserve positional compatibility for the released items_to_messages signature.
"""
message = InternalChatCompletionMessage(
role="assistant",
content="I'll check the weather for you.",
reasoning_content="The user wants weather information, I need to call the weather function",
thinking_blocks=[
{
"type": "thinking",
"thinking": (
"The user is asking about weather. "
"Let me use the weather tool to get this information."
),
"signature": "TestSignature123",
}
],
tool_calls=[
ChatCompletionMessageToolCall(
id="call_123",
type="function",
function=Function(name="get_weather", arguments='{"city": "Tokyo"}'),
)
],
)
output_items = Converter.message_to_output_items(message)
items_as_dicts: list[dict[str, Any]] = []
for item in output_items:
if hasattr(item, "model_dump"):
items_as_dicts.append(item.model_dump())
else:
items_as_dicts.append(cast(dict[str, Any], item))
messages = Converter.items_to_messages(
items_as_dicts, # type: ignore[arg-type]
"anthropic/claude-4-opus",
True,
True,
)
assistant_messages = [
msg for msg in messages if msg.get("role") == "assistant" and msg.get("tool_calls")
]
assert len(assistant_messages) == 1, "Should have exactly one assistant message with tool calls"
assistant_msg = assistant_messages[0]
content = assistant_msg.get("content")
assert isinstance(content, list) and len(content) > 0, (
"Positional bool arguments should still preserve thinking blocks"
)
assert content[0].get("type") == "thinking", (
"The third positional argument must continue to map to preserve_thinking_blocks"
)
def test_anthropic_thinking_blocks_without_tool_calls():
"""
Test for models with extended thinking WITHOUT tool calls.