From 445aba7e5f8c92be9eefeec187c6669a8d8281ca Mon Sep 17 00:00:00 2001 From: worldmozara Date: Tue, 6 Oct 2026 21:07:04 +0800 Subject: [PATCH] lmproto: carry thinking options on the wire --- src/plyngent/lmproto/anthropic/client.py | 3 + src/plyngent/lmproto/anthropic/config.py | 3 + src/plyngent/lmproto/anthropic/model.py | 13 +++ .../lmproto/deepseek/openai_compat/client.py | 15 ++-- tests/test_lmproto/test_deepseek_chat.py | 79 +++++++++++++++++++ 5 files changed, 106 insertions(+), 7 deletions(-) create mode 100644 tests/test_lmproto/test_deepseek_chat.py diff --git a/src/plyngent/lmproto/anthropic/client.py b/src/plyngent/lmproto/anthropic/client.py index 345f029..5128506 100644 --- a/src/plyngent/lmproto/anthropic/client.py +++ b/src/plyngent/lmproto/anthropic/client.py @@ -34,6 +34,8 @@ class AnthropicClient: decoder: msgspec.json.Decoder[AnthropicMessageResponse] stream_decoder: msgspec.json.Decoder[AnthropicStreamEvent] models_decoder: msgspec.json.Decoder[AnthropicModelsResponse] + # Explicit ``thinking`` budget for this client (0 = derive from the effort). + thinking_budget_tokens: int = 0 _api_key: str _api_version: str @@ -43,6 +45,7 @@ class AnthropicClient: self.decoder = msgspec.json.Decoder(AnthropicMessageResponse) 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._api_key = config.api_key self._api_version = config.anthropic_version diff --git a/src/plyngent/lmproto/anthropic/config.py b/src/plyngent/lmproto/anthropic/config.py index ef57991..3db3482 100644 --- a/src/plyngent/lmproto/anthropic/config.py +++ b/src/plyngent/lmproto/anthropic/config.py @@ -20,3 +20,6 @@ class AnthropicConfig: base_url: str = DEFAULT_ANTHROPIC_BASE_URL anthropic_version: str = DEFAULT_ANTHROPIC_VERSION timeout: HttpTimeout = (DEFAULT_HTTP_CONNECT_TIMEOUT, DEFAULT_HTTP_READ_TIMEOUT) + # Explicit ``thinking.budget_tokens`` for reasoning models (0 = derive it + # from the configured reasoning_effort; see config/reasoning.py). + thinking_budget_tokens: int = 0 diff --git a/src/plyngent/lmproto/anthropic/model.py b/src/plyngent/lmproto/anthropic/model.py index 95f0950..89dcb54 100644 --- a/src/plyngent/lmproto/anthropic/model.py +++ b/src/plyngent/lmproto/anthropic/model.py @@ -62,6 +62,18 @@ class AnthropicMetadata(Struct, omit_defaults=True): user_id: str | Unset = UNSET +class AnthropicThinkingConfig(Struct, omit_defaults=True): + """``thinking`` request block (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). + """ + + type: Literal["enabled", "disabled"] = "enabled" + budget_tokens: int | Unset = UNSET + + class AnthropicMessagesParam(Struct, omit_defaults=True): model: str max_tokens: int = 8192 @@ -75,6 +87,7 @@ class AnthropicMessagesParam(Struct, omit_defaults=True): temperature: float | Unset = UNSET top_p: float | Unset = UNSET top_k: int | Unset = UNSET + thinking: AnthropicThinkingConfig | Unset = UNSET class AnthropicUsage(Struct, omit_defaults=True): diff --git a/src/plyngent/lmproto/deepseek/openai_compat/client.py b/src/plyngent/lmproto/deepseek/openai_compat/client.py index dce1692..01ba74d 100644 --- a/src/plyngent/lmproto/deepseek/openai_compat/client.py +++ b/src/plyngent/lmproto/deepseek/openai_compat/client.py @@ -16,16 +16,17 @@ if TYPE_CHECKING: from .model import ChatCompletionsParam -def _inject_thinking(data: bytes) -> bytes: - """Inject ``thinking: {type: "enabled"}`` into the encoded request body. +def _inject_thinking(data: bytes, param: ChatCompletionsParam) -> bytes: + """Inject the ``thinking`` flag into the encoded request body. DeepSeek's reasoning model requires the ``thinking`` parameter to be - explicitly set (default is ``enabled``). The agent loop constructs the - base ``ChatCompletionsParam`` which lacks this field, so we add it - after msgspec encoding. + explicitly set (default is ``enabled``); ``reasoning_effort = "none"`` is the + way to ask for it off. The agent loop constructs the base + ``ChatCompletionsParam`` which lacks this field, so we add it after msgspec + encoding. """ body = json.loads(data) - body["thinking"] = {"type": "enabled"} + body["thinking"] = {"type": "disabled" if param.reasoning_effort == "none" else "enabled"} return json.dumps(body, separators=(",", ":")).encode("utf-8") @@ -49,7 +50,7 @@ class DeepseekOpenAIClient(BaseOpenAIClient): self, param: ChatCompletionsParam, *, stream: bool = False ) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]: param = coerce_chat_completions_param_any(msgspec.structs.replace(param, stream=stream)) - data = _inject_thinking(self.encoder.encode(param)) + data = _inject_thinking(self.encoder.encode(param), param) if stream: resp = await self.session.post( "/chat/completions", diff --git a/tests/test_lmproto/test_deepseek_chat.py b/tests/test_lmproto/test_deepseek_chat.py new file mode 100644 index 0000000..146b176 --- /dev/null +++ b/tests/test_lmproto/test_deepseek_chat.py @@ -0,0 +1,79 @@ +"""Tests for the DeepSeek chat-completions client's ``thinking`` flag.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any, cast + +import pytest +from msgspec import UNSET + +from plyngent.lmproto.deepseek import DeepseekOpenAIClient +from plyngent.lmproto.openai_compatible.config import OpenAIConfig +from plyngent.lmproto.openai_compatible.model import ChatCompletionsParam, UserChatMessage + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from plyngent.lmproto.openai_compatible.model import ReasoningEffort + from plyngent.typedef import Unset + + +async def _posted_body( + monkeypatch: pytest.MonkeyPatch, + effort: ReasoningEffort | Unset = UNSET, +) -> dict[str, Any]: + """Run one streamed request and return the JSON body the client posted.""" + client = DeepseekOpenAIClient(OpenAIConfig(access_key_or_token="sk-test", base_url="https://api.deepseek.com")) + captured: dict[str, Any] = {} + + class _Resp: + status_code = 200 + + async def iter_lines(self) -> AsyncIterator[bytes]: + yield b"data: [DONE]" + + def close(self) -> None: + return None + + async def fake_post(path: str, **kwargs: object) -> _Resp: + assert path == "/chat/completions" + captured["data"] = kwargs.get("data") + return _Resp() + + monkeypatch.setattr(client.session, "post", fake_post) + param = ChatCompletionsParam( + model="deepseek-v4-pro", + messages=[UserChatMessage(content="hi")], + reasoning_effort=effort, + ) + # Production dispatches the base param through the OpenAI-compatible surface + # (not the DeepSeek-specific annotation), so mirror that here. + cc: Any = client + stream = cast("AsyncIterator[Any]", await cc.chat_completions(param, stream=True)) + async for _chunk in stream: + pass + data = captured["data"] + assert isinstance(data, bytes) + return json.loads(data) + + +@pytest.mark.asyncio +async def test_thinking_enabled_without_configured_effort(monkeypatch: pytest.MonkeyPatch) -> None: + body = await _posted_body(monkeypatch) + assert body["thinking"] == {"type": "enabled"} + assert "reasoning_effort" not in body + + +@pytest.mark.asyncio +async def test_thinking_enabled_with_effort(monkeypatch: pytest.MonkeyPatch) -> None: + body = await _posted_body(monkeypatch, "high") + assert body["thinking"] == {"type": "enabled"} + assert body["reasoning_effort"] == "high" + + +@pytest.mark.asyncio +async def test_thinking_disabled_for_none_effort(monkeypatch: pytest.MonkeyPatch) -> None: + body = await _posted_body(monkeypatch, "none") + assert body["thinking"] == {"type": "disabled"} + assert body["reasoning_effort"] == "none"