mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 15:37:55 +08:00
core/agent+anthropic: retry a refused thinking budget as adaptive thinking
This commit is contained in:
@@ -17,6 +17,7 @@ from plyngent.lmproto.anthropic.model import (
|
||||
AnthropicAssistantMessage,
|
||||
AnthropicMessageResponse,
|
||||
AnthropicMessagesParam,
|
||||
AnthropicOutputConfig,
|
||||
AnthropicResponseText,
|
||||
AnthropicServerToolUseContent,
|
||||
AnthropicTextContent,
|
||||
@@ -384,7 +385,7 @@ def _thinking_from_param(
|
||||
*,
|
||||
thinking_budget_tokens: int,
|
||||
) -> AnthropicThinkingConfig | Unset:
|
||||
"""Build the ``thinking`` block for *param* (UNSET = send nothing).
|
||||
"""Build the extended-thinking block for *param* (UNSET = send nothing).
|
||||
|
||||
``reasoning_effort = "none"`` disables thinking; any other level enables it
|
||||
with the explicit *thinking_budget_tokens*, else with the budget that level
|
||||
@@ -423,11 +424,76 @@ def _thinking_and_max_tokens(
|
||||
return max_tokens, thinking
|
||||
|
||||
|
||||
# Effort level → Anthropic ``output_config.effort`` on an adaptive-thinking
|
||||
# model. ``minimal`` has no API level of its own and lands on the lowest one;
|
||||
# unknown levels fall back to "medium" (the budgets above do the same).
|
||||
_ANTHROPIC_ADAPTIVE_EFFORTS: dict[str, str] = {
|
||||
"minimal": "low",
|
||||
"low": "low",
|
||||
"medium": "medium",
|
||||
"high": "high",
|
||||
"xhigh": "xhigh",
|
||||
"max": "max",
|
||||
}
|
||||
_ANTHROPIC_FALLBACK_ADAPTIVE_EFFORT = "medium"
|
||||
|
||||
|
||||
def _adaptive_thinking_from_param(
|
||||
param: ChatCompletionsParam,
|
||||
) -> tuple[AnthropicThinkingConfig | Unset, AnthropicOutputConfig | Unset]:
|
||||
"""Build the adaptive-thinking blocks for *param* (UNSET = send nothing).
|
||||
|
||||
The shape for a model that rejects ``thinking.type.enabled``: adaptive
|
||||
thinking, steered by ``output_config.effort`` rather than a token budget.
|
||||
``reasoning_effort = "none"`` still disables thinking; a configured
|
||||
``thinking_budget_tokens`` has no adaptive equivalent and is left out.
|
||||
"""
|
||||
effort = param.reasoning_effort
|
||||
if effort is UNSET:
|
||||
return UNSET, UNSET
|
||||
if effort == "none":
|
||||
return AnthropicThinkingConfig(type="disabled"), UNSET
|
||||
level = _ANTHROPIC_ADAPTIVE_EFFORTS.get(str(effort), _ANTHROPIC_FALLBACK_ADAPTIVE_EFFORT)
|
||||
return AnthropicThinkingConfig(type="adaptive"), AnthropicOutputConfig(effort=level)
|
||||
|
||||
|
||||
def extended_thinking_rejected(message: str) -> bool:
|
||||
"""Whether a ``/messages`` failure is the API refusing extended thinking.
|
||||
|
||||
The documented wording is ``"thinking.type.enabled" is not supported for this
|
||||
model. Use "thinking.type.adaptive" and "output_config.effort" …``; gateways
|
||||
paraphrase it (``claude-opus-5-5 requires adaptive thinking; omit thinking or
|
||||
use thinking.type=adaptive …``), so match the one word they all carry —
|
||||
callers only get here after a request that sent extended thinking.
|
||||
"""
|
||||
return "adaptive" in message.lower()
|
||||
|
||||
|
||||
def _thinking_for_param(
|
||||
param: ChatCompletionsParam,
|
||||
*,
|
||||
thinking_budget_tokens: int,
|
||||
adaptive: bool,
|
||||
) -> tuple[int, AnthropicThinkingConfig | Unset, AnthropicOutputConfig | Unset]:
|
||||
"""Resolve *param*'s ``max_tokens`` and thinking blocks for one shape.
|
||||
|
||||
Extended thinking needs ``max_tokens`` raised above its budget (see
|
||||
:func:`_thinking_and_max_tokens`); adaptive thinking carries no budget, so
|
||||
the caller's ``max_tokens`` stands.
|
||||
"""
|
||||
if adaptive:
|
||||
thinking, output_config = _adaptive_thinking_from_param(param)
|
||||
return _max_tokens_from_param(param), thinking, output_config
|
||||
max_tokens, thinking = _thinking_and_max_tokens(param, thinking_budget_tokens=thinking_budget_tokens)
|
||||
return max_tokens, thinking, UNSET
|
||||
|
||||
|
||||
def chat_param_to_anthropic_param(
|
||||
param: ChatCompletionsParam,
|
||||
*,
|
||||
thinking_budget_tokens: int = 0,
|
||||
provider_tools: Sequence[dict[str, Any]] | None = None,
|
||||
adaptive_thinking: bool = False,
|
||||
) -> AnthropicMessagesParam:
|
||||
"""Build :class:`AnthropicMessagesParam` from a chat completions param.
|
||||
|
||||
@@ -435,13 +501,22 @@ def chat_param_to_anthropic_param(
|
||||
(0 = derive it from the param's ``reasoning_effort``); *provider_tools* are
|
||||
hosted tools (web_search, …) as opaque dicts, merged after local function
|
||||
tools and never executed by the registry.
|
||||
|
||||
*adaptive_thinking* asks for the thinking shape of a model that rejects
|
||||
extended thinking: ``thinking.type = "adaptive"`` with the effort in
|
||||
``output_config`` instead of a ``budget_tokens`` budget (see
|
||||
:func:`_adaptive_thinking_from_param`).
|
||||
"""
|
||||
system, messages = chat_messages_to_anthropic(param.messages)
|
||||
tools: list[AnthropicToolDefinition | dict[str, Any]] = list(
|
||||
tool_items_to_anthropic_tools(param.tools if param.tools is not UNSET else None)
|
||||
)
|
||||
tools.extend(provider_tools_to_anthropic_tools(provider_tools))
|
||||
max_tokens, thinking = _thinking_and_max_tokens(param, thinking_budget_tokens=thinking_budget_tokens)
|
||||
max_tokens, thinking, output_config = _thinking_for_param(
|
||||
param,
|
||||
thinking_budget_tokens=thinking_budget_tokens,
|
||||
adaptive=adaptive_thinking,
|
||||
)
|
||||
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": param.model,
|
||||
@@ -450,6 +525,8 @@ def chat_param_to_anthropic_param(
|
||||
}
|
||||
if thinking is not UNSET:
|
||||
kwargs["thinking"] = thinking
|
||||
if output_config is not UNSET:
|
||||
kwargs["output_config"] = output_config
|
||||
if system:
|
||||
kwargs["system"] = system
|
||||
if tools:
|
||||
|
||||
@@ -11,11 +11,14 @@ from __future__ import annotations
|
||||
import json
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
from msgspec import UNSET
|
||||
|
||||
from plyngent.agent.messages_bridge import (
|
||||
anthropic_response_to_chat_completion,
|
||||
anthropic_stop_to_finish_reason,
|
||||
anthropic_usage_to_dict,
|
||||
chat_param_to_anthropic_param,
|
||||
extended_thinking_rejected,
|
||||
)
|
||||
from plyngent.agent.stream_chunks import (
|
||||
finish_reason_chunk,
|
||||
@@ -40,7 +43,7 @@ if TYPE_CHECKING:
|
||||
from collections.abc import AsyncIterator, Sequence
|
||||
|
||||
from plyngent.lmproto.anthropic.client import AnthropicClient
|
||||
from plyngent.lmproto.anthropic.model import AnthropicStreamEvent
|
||||
from plyngent.lmproto.anthropic.model import AnthropicMessagesParam, AnthropicStreamEvent
|
||||
from plyngent.lmproto.openai_compatible.model import (
|
||||
ChatCompletionChunk,
|
||||
ChatCompletionResponse,
|
||||
@@ -53,6 +56,48 @@ def _thinking_budget(client: AnthropicClient) -> int:
|
||||
return max(0, int(getattr(client, "thinking_budget_tokens", 0)))
|
||||
|
||||
|
||||
def _create_param(
|
||||
client: AnthropicClient,
|
||||
param: ChatCompletionsParam,
|
||||
provider_tools: Sequence[dict[str, Any]] | None,
|
||||
*,
|
||||
adaptive: bool | None = None,
|
||||
) -> AnthropicMessagesParam:
|
||||
"""Build the ``/messages`` request for *param*.
|
||||
|
||||
*adaptive* picks the thinking shape; ``None`` asks the client whether this
|
||||
session already found *param*'s model rejecting extended thinking.
|
||||
"""
|
||||
use_adaptive = client.needs_adaptive_thinking(param.model) if adaptive is None else adaptive
|
||||
return chat_param_to_anthropic_param(
|
||||
param,
|
||||
thinking_budget_tokens=_thinking_budget(client),
|
||||
provider_tools=provider_tools,
|
||||
adaptive_thinking=use_adaptive,
|
||||
)
|
||||
|
||||
|
||||
def _retry_with_adaptive_thinking(
|
||||
client: AnthropicClient,
|
||||
param: ChatCompletionsParam,
|
||||
provider_tools: Sequence[dict[str, Any]] | None,
|
||||
error: Exception,
|
||||
) -> AnthropicMessagesParam | None:
|
||||
"""The request to retry with when the model refused extended thinking.
|
||||
|
||||
``None`` when the failure is not that refusal, or when the request never sent
|
||||
extended thinking. The model is remembered on the client, so only its first
|
||||
request pays the extra round trip.
|
||||
"""
|
||||
if not extended_thinking_rejected(str(error)):
|
||||
return None
|
||||
create = _create_param(client, param, provider_tools)
|
||||
if create.thinking is UNSET or create.thinking.type != "enabled":
|
||||
return None
|
||||
client.note_adaptive_thinking(param.model)
|
||||
return _create_param(client, param, provider_tools, adaptive=True)
|
||||
|
||||
|
||||
class _ToolBlockState:
|
||||
"""Tracks an in-flight tool_use content block during SSE streaming."""
|
||||
|
||||
@@ -365,6 +410,33 @@ def _chunks_for_content_delta(
|
||||
]
|
||||
|
||||
|
||||
async def _stream_events(
|
||||
client: AnthropicClient,
|
||||
param: ChatCompletionsParam,
|
||||
*,
|
||||
provider_tools: Sequence[dict[str, Any]] | None = None,
|
||||
) -> AsyncIterator[AnthropicStreamEvent]:
|
||||
"""Open the Messages stream, retrying once with adaptive thinking if refused.
|
||||
|
||||
A refusal arrives before the stream yields anything, so the retry cannot
|
||||
duplicate events the caller has already seen.
|
||||
"""
|
||||
create = _create_param(client, param, provider_tools)
|
||||
saw_event = False
|
||||
try:
|
||||
async for event in await client.messages(create, stream=True):
|
||||
saw_event = True
|
||||
yield event
|
||||
except RuntimeError as exc:
|
||||
if saw_event:
|
||||
raise
|
||||
retry = _retry_with_adaptive_thinking(client, param, provider_tools, exc)
|
||||
if retry is None:
|
||||
raise
|
||||
async for event in await client.messages(retry, stream=True):
|
||||
yield event
|
||||
|
||||
|
||||
async def _stream_as_chat_chunks(
|
||||
client: AnthropicClient,
|
||||
param: ChatCompletionsParam,
|
||||
@@ -379,15 +451,9 @@ async def _stream_as_chat_chunks(
|
||||
``message_stop``. Errors raise ``RuntimeError`` so the agent loop can
|
||||
surface them as retryable failures.
|
||||
"""
|
||||
create = chat_param_to_anthropic_param(
|
||||
param,
|
||||
thinking_budget_tokens=_thinking_budget(client),
|
||||
provider_tools=provider_tools,
|
||||
)
|
||||
stream = await client.messages(create, stream=True)
|
||||
state = _StreamState(param.model)
|
||||
|
||||
async for event in stream:
|
||||
async for event in _stream_events(client, param, provider_tools=provider_tools):
|
||||
for chunk in _chunks_for_event(event, state):
|
||||
yield chunk
|
||||
if state.saw_message_stop:
|
||||
@@ -419,14 +485,16 @@ async def dispatch_messages(
|
||||
surface's own server-tool definition is built from them and the results
|
||||
arrive inside the same response, so nothing is executed locally.
|
||||
"""
|
||||
create = chat_param_to_anthropic_param(
|
||||
param,
|
||||
thinking_budget_tokens=_thinking_budget(client),
|
||||
provider_tools=provider_tools,
|
||||
)
|
||||
if stream:
|
||||
return _stream_as_chat_chunks(client, param, provider_tools=provider_tools)
|
||||
response = await client.messages(create, stream=False)
|
||||
create = _create_param(client, param, provider_tools)
|
||||
try:
|
||||
response = await client.messages(create, stream=False)
|
||||
except RuntimeError as exc:
|
||||
retry = _retry_with_adaptive_thinking(client, param, provider_tools, exc)
|
||||
if retry is None:
|
||||
raise
|
||||
response = await client.messages(retry, stream=False)
|
||||
return anthropic_response_to_chat_completion(response)
|
||||
|
||||
|
||||
|
||||
@@ -36,6 +36,9 @@ class AnthropicClient:
|
||||
models_decoder: msgspec.json.Decoder[AnthropicModelsResponse]
|
||||
# Explicit ``thinking`` budget for this client (0 = derive from the effort).
|
||||
thinking_budget_tokens: int = 0
|
||||
# Models this session found to reject extended thinking; the dispatch retries
|
||||
# their requests with adaptive thinking instead (see ``messages_dispatch``).
|
||||
adaptive_thinking_models: set[str]
|
||||
_api_key: str
|
||||
_api_version: str
|
||||
|
||||
@@ -46,9 +49,18 @@ class AnthropicClient:
|
||||
self.stream_decoder = msgspec.json.Decoder(AnthropicStreamEvent)
|
||||
self.models_decoder = msgspec.json.Decoder(AnthropicModelsResponse)
|
||||
self.thinking_budget_tokens = max(0, config.thinking_budget_tokens)
|
||||
self.adaptive_thinking_models = set()
|
||||
self._api_key = config.api_key
|
||||
self._api_version = config.anthropic_version
|
||||
|
||||
def needs_adaptive_thinking(self, model: str) -> bool:
|
||||
"""Whether *model* already rejected extended thinking in this session."""
|
||||
return model in self.adaptive_thinking_models
|
||||
|
||||
def note_adaptive_thinking(self, model: str) -> None:
|
||||
"""Remember that *model* requires adaptive thinking, saving the retry."""
|
||||
self.adaptive_thinking_models.add(model)
|
||||
|
||||
def _headers(self) -> dict[str, str]:
|
||||
return {
|
||||
"Content-Type": "application/json",
|
||||
|
||||
@@ -66,21 +66,34 @@ class AnthropicMetadata(Struct, omit_defaults=True):
|
||||
|
||||
|
||||
class AnthropicThinkingConfig(Struct, omit_defaults=True):
|
||||
"""``thinking`` request block (extended thinking).
|
||||
"""``thinking`` request block (extended or adaptive thinking).
|
||||
|
||||
``budget_tokens`` is the reasoning budget the model may spend before it
|
||||
answers; the API requires it to be at least 1024 and smaller than
|
||||
``max_tokens`` (see ``agent.messages_bridge``, which keeps both valid).
|
||||
``enabled`` is extended thinking: ``budget_tokens`` is the reasoning budget
|
||||
the model may spend before it answers; the API requires it to be at least
|
||||
1024 and smaller than ``max_tokens`` (see ``agent.messages_bridge``, which
|
||||
keeps both valid). Models from Claude 4.7 on reject ``enabled`` and take
|
||||
``adaptive`` instead, where the depth comes from ``output_config.effort``;
|
||||
``disabled`` turns thinking off where the model allows it.
|
||||
|
||||
``type`` carries no default: ``omit_defaults`` would drop a defaulted one,
|
||||
and the server reads the block as malformed without it — DeepSeek answers
|
||||
``thinking: missing field ``type``.
|
||||
"""
|
||||
|
||||
type: Literal["enabled", "disabled"]
|
||||
type: Literal["enabled", "disabled", "adaptive"]
|
||||
budget_tokens: int | Unset = UNSET
|
||||
|
||||
|
||||
class AnthropicOutputConfig(Struct, omit_defaults=True):
|
||||
"""``output_config`` request block: how much the model spends on a turn.
|
||||
|
||||
``effort`` is the adaptive-thinking knob (``low`` … ``max``); it replaces a
|
||||
``thinking.budget_tokens`` budget on the models that no longer take one.
|
||||
"""
|
||||
|
||||
effort: str | Unset = UNSET
|
||||
|
||||
|
||||
class AnthropicMessagesParam(Struct, omit_defaults=True):
|
||||
model: str
|
||||
max_tokens: int = 8192
|
||||
@@ -97,6 +110,7 @@ class AnthropicMessagesParam(Struct, omit_defaults=True):
|
||||
top_p: float | Unset = UNSET
|
||||
top_k: int | Unset = UNSET
|
||||
thinking: AnthropicThinkingConfig | Unset = UNSET
|
||||
output_config: AnthropicOutputConfig | Unset = UNSET
|
||||
|
||||
|
||||
class AnthropicUsage(Struct, omit_defaults=True):
|
||||
|
||||
@@ -14,6 +14,7 @@ from plyngent.agent.messages_bridge import (
|
||||
anthropic_usage_to_dict,
|
||||
chat_messages_to_anthropic,
|
||||
chat_param_to_anthropic_param,
|
||||
extended_thinking_rejected,
|
||||
provider_tools_to_anthropic_tools,
|
||||
tool_items_to_anthropic_tools,
|
||||
)
|
||||
@@ -387,3 +388,59 @@ def test_chat_param_to_anthropic_param_thinking_disabled() -> None:
|
||||
assert body.thinking is not UNSET
|
||||
assert body.thinking.type == "disabled"
|
||||
assert body.thinking.budget_tokens is UNSET
|
||||
|
||||
|
||||
def test_chat_param_to_anthropic_param_adaptive_thinking_from_effort() -> None:
|
||||
body = chat_param_to_anthropic_param(_effort_param("high"), adaptive_thinking=True)
|
||||
assert body.thinking is not UNSET
|
||||
assert body.thinking.type == "adaptive"
|
||||
assert body.thinking.budget_tokens is UNSET
|
||||
assert body.output_config is not UNSET
|
||||
assert body.output_config.effort == "high"
|
||||
# Adaptive thinking spends no fixed budget, so max_tokens stays as asked.
|
||||
assert body.max_tokens == 8192
|
||||
|
||||
|
||||
def test_chat_param_to_anthropic_param_adaptive_effort_levels() -> None:
|
||||
def effort_for(level: ReasoningEffort) -> str:
|
||||
body = chat_param_to_anthropic_param(_effort_param(level), adaptive_thinking=True)
|
||||
assert body.output_config is not UNSET
|
||||
return str(body.output_config.effort)
|
||||
|
||||
assert effort_for("minimal") == "low" # no such API level; the lowest one
|
||||
assert effort_for("xhigh") == "xhigh"
|
||||
|
||||
|
||||
def test_chat_param_to_anthropic_param_adaptive_ignores_a_configured_budget() -> None:
|
||||
body = chat_param_to_anthropic_param(_effort_param("medium"), thinking_budget_tokens=20000, adaptive_thinking=True)
|
||||
assert body.thinking is not UNSET
|
||||
assert body.thinking.type == "adaptive"
|
||||
assert body.thinking.budget_tokens is UNSET
|
||||
assert body.max_tokens == 8192
|
||||
|
||||
|
||||
def test_chat_param_to_anthropic_param_adaptive_disabled_carries_no_effort() -> None:
|
||||
body = chat_param_to_anthropic_param(_effort_param("none"), adaptive_thinking=True)
|
||||
assert body.thinking is not UNSET
|
||||
assert body.thinking.type == "disabled"
|
||||
assert body.output_config is UNSET
|
||||
|
||||
|
||||
def test_chat_param_to_anthropic_param_adaptive_without_effort_sends_nothing() -> None:
|
||||
body = chat_param_to_anthropic_param(_effort_param(), adaptive_thinking=True)
|
||||
assert body.thinking is UNSET
|
||||
assert body.output_config is UNSET
|
||||
|
||||
|
||||
def test_extended_thinking_rejected_matches_the_api_and_gateway_wording() -> None:
|
||||
documented = (
|
||||
'messages HTTP 400: {"error":{"message":"\\"thinking.type.enabled\\" is not supported for this '
|
||||
'model. Use \\"thinking.type.adaptive\\" and \\"output_config.effort\\" to control thinking behavior."}}'
|
||||
)
|
||||
gateway = (
|
||||
'messages HTTP 400: {"error":{"message":"claude-opus-5-5 requires adaptive thinking; omit thinking '
|
||||
'or use thinking.type=adaptive and output_config.effort"}}'
|
||||
)
|
||||
assert extended_thinking_rejected(documented)
|
||||
assert extended_thinking_rejected(gateway)
|
||||
assert not extended_thinking_rejected('messages HTTP 400: {"error":{"message":"max_tokens: required"}}')
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Literal, cast, overload
|
||||
|
||||
import pytest
|
||||
from msgspec import UNSET
|
||||
|
||||
from plyngent.agent import (
|
||||
@@ -45,6 +46,14 @@ if TYPE_CHECKING:
|
||||
from collections.abc import AsyncIterator, Sequence
|
||||
|
||||
|
||||
# What the API answers a request that sends extended thinking to a model that
|
||||
# takes adaptive thinking only (a gateway paraphrase of the documented wording).
|
||||
_ADAPTIVE_REQUIRED_400 = (
|
||||
'messages HTTP 400: {"error":{"type":"error","message":"claude-opus-5-5 requires adaptive '
|
||||
'thinking; omit thinking or use thinking.type=adaptive and output_config.effort"}}'
|
||||
)
|
||||
|
||||
|
||||
def _message_response(
|
||||
*,
|
||||
text: str = "hello",
|
||||
@@ -73,21 +82,43 @@ class ScriptedMessagesClient:
|
||||
|
||||
kind: str = "messages"
|
||||
calls: list[AnthropicMessagesParam]
|
||||
# Mirrors AnthropicClient: explicit thinking budget (0 = derive from effort).
|
||||
# Mirrors AnthropicClient: explicit thinking budget (0 = derive from effort),
|
||||
# and the models the dispatch learned to send adaptive thinking to.
|
||||
thinking_budget_tokens: int
|
||||
adaptive_thinking_models: set[str]
|
||||
_non_stream: list[AnthropicMessageResponse]
|
||||
_stream_events: list[list[AnthropicStreamEvent]]
|
||||
_rejects_extended_thinking: bool
|
||||
_error: str | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
non_stream: Sequence[AnthropicMessageResponse] | None = None,
|
||||
stream_events: Sequence[Sequence[AnthropicStreamEvent]] | None = None,
|
||||
rejects_extended_thinking: bool = False,
|
||||
error: str | None = None,
|
||||
) -> None:
|
||||
self.calls = []
|
||||
self.thinking_budget_tokens = 0
|
||||
self.adaptive_thinking_models = set()
|
||||
self._non_stream = list(non_stream or ())
|
||||
self._stream_events = [list(events) for events in (stream_events or ())]
|
||||
self._rejects_extended_thinking = rejects_extended_thinking
|
||||
self._error = error
|
||||
|
||||
def needs_adaptive_thinking(self, model: str) -> bool:
|
||||
return model in self.adaptive_thinking_models
|
||||
|
||||
def note_adaptive_thinking(self, model: str) -> None:
|
||||
self.adaptive_thinking_models.add(model)
|
||||
|
||||
def _reject(self, param: AnthropicMessagesParam) -> None:
|
||||
"""Raise what the API raises for a request it refuses to run."""
|
||||
if self._error is not None:
|
||||
raise RuntimeError(self._error)
|
||||
if self._rejects_extended_thinking and param.thinking is not UNSET and param.thinking.type == "enabled":
|
||||
raise RuntimeError(_ADAPTIVE_REQUIRED_400)
|
||||
|
||||
@overload
|
||||
async def messages(
|
||||
@@ -107,13 +138,16 @@ class ScriptedMessagesClient:
|
||||
if not self._stream_events:
|
||||
msg = "no more scripted stream events"
|
||||
raise RuntimeError(msg)
|
||||
events = self._stream_events.pop(0)
|
||||
|
||||
async def _gen() -> AsyncIterator[AnthropicStreamEvent]:
|
||||
for event in events:
|
||||
# The real client reads the error body on the first iteration,
|
||||
# so a refused request consumes no scripted events.
|
||||
self._reject(param)
|
||||
for event in self._stream_events.pop(0):
|
||||
yield event
|
||||
|
||||
return _gen()
|
||||
self._reject(param)
|
||||
if not self._non_stream:
|
||||
msg = "no more scripted messages"
|
||||
raise RuntimeError(msg)
|
||||
@@ -602,3 +636,68 @@ async def test_search_turn_reports_the_prompt_the_delta_counts() -> None:
|
||||
assert isinstance(usage, dict)
|
||||
assert usage["prompt_tokens"] == 12468
|
||||
assert usage["completion_tokens"] == 198
|
||||
|
||||
|
||||
def _thinking_param(model: str = "claude-opus-5-5", effort: str = "high") -> ChatCompletionsParam:
|
||||
return ChatCompletionsParam(
|
||||
model=model,
|
||||
messages=[UserChatMessage(content="hi")],
|
||||
reasoning_effort=cast("Any", effort),
|
||||
)
|
||||
|
||||
|
||||
async def test_dispatch_messages_retries_with_adaptive_thinking() -> None:
|
||||
client = ScriptedMessagesClient(non_stream=[_message_response(text="done")], rejects_extended_thinking=True)
|
||||
result = await dispatch_messages(cast("Any", client), _thinking_param(), stream=False)
|
||||
assert result.choices[0].message.content == "done"
|
||||
assert len(client.calls) == 2
|
||||
refused, retried = client.calls
|
||||
assert refused.thinking is not UNSET
|
||||
assert refused.thinking.type == "enabled"
|
||||
assert retried.thinking is not UNSET
|
||||
assert retried.thinking.type == "adaptive"
|
||||
assert retried.output_config is not UNSET
|
||||
assert retried.output_config.effort == "high"
|
||||
|
||||
|
||||
async def test_dispatch_messages_remembers_the_refusing_model() -> None:
|
||||
client = ScriptedMessagesClient(
|
||||
non_stream=[_message_response(text="one"), _message_response(text="two")],
|
||||
rejects_extended_thinking=True,
|
||||
)
|
||||
param = _thinking_param()
|
||||
_ = await dispatch_messages(cast("Any", client), param, stream=False)
|
||||
_ = await dispatch_messages(cast("Any", client), param, stream=False)
|
||||
# Only the model's first request pays the refused attempt.
|
||||
assert [call.thinking.type for call in client.calls] == ["enabled", "adaptive", "adaptive"]
|
||||
|
||||
|
||||
async def test_dispatch_messages_stream_retries_with_adaptive_thinking() -> None:
|
||||
events: list[AnthropicStreamEvent] = [
|
||||
AnthropicContentBlockStart(index=0, content_block=AnthropicRawContentBlock(type="text", text="")),
|
||||
AnthropicContentBlockDelta(index=0, delta=AnthropicRawContentBlock(type="text_delta", text="streamed")),
|
||||
AnthropicMessageDelta(delta={"stop_reason": "end_turn"}, usage=AnthropicUsage(output_tokens=3)),
|
||||
AnthropicMessageStop(),
|
||||
]
|
||||
client = ScriptedMessagesClient(stream_events=[events], rejects_extended_thinking=True)
|
||||
stream = await dispatch_messages(cast("Any", client), _thinking_param(), stream=True)
|
||||
chunks = [chunk async for chunk in cast("Any", stream)]
|
||||
texts = [c.choices[0].delta.content for c in chunks if c.choices and isinstance(c.choices[0].delta.content, str)]
|
||||
assert texts == ["streamed"]
|
||||
assert len(client.calls) == 2
|
||||
assert client.calls[1].thinking.type == "adaptive"
|
||||
|
||||
|
||||
async def test_dispatch_messages_does_not_retry_other_failures() -> None:
|
||||
client = ScriptedMessagesClient(error='messages HTTP 400: {"error":{"message":"max_tokens: required"}}')
|
||||
with pytest.raises(RuntimeError, match="max_tokens"):
|
||||
_ = await dispatch_messages(cast("Any", client), _thinking_param(), stream=False)
|
||||
assert len(client.calls) == 1
|
||||
|
||||
|
||||
async def test_dispatch_messages_does_not_retry_a_request_without_extended_thinking() -> None:
|
||||
client = ScriptedMessagesClient(error=_ADAPTIVE_REQUIRED_400)
|
||||
param = ChatCompletionsParam(model="claude-opus-5-5", messages=[UserChatMessage(content="hi")])
|
||||
with pytest.raises(RuntimeError, match="adaptive thinking"):
|
||||
_ = await dispatch_messages(cast("Any", client), param, stream=False)
|
||||
assert len(client.calls) == 1
|
||||
|
||||
Reference in New Issue
Block a user