mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 15:37:55 +08:00
lmproto: carry thinking options on the wire
This commit is contained in:
@@ -34,6 +34,8 @@ class AnthropicClient:
|
|||||||
decoder: msgspec.json.Decoder[AnthropicMessageResponse]
|
decoder: msgspec.json.Decoder[AnthropicMessageResponse]
|
||||||
stream_decoder: msgspec.json.Decoder[AnthropicStreamEvent]
|
stream_decoder: msgspec.json.Decoder[AnthropicStreamEvent]
|
||||||
models_decoder: msgspec.json.Decoder[AnthropicModelsResponse]
|
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_key: str
|
||||||
_api_version: str
|
_api_version: str
|
||||||
|
|
||||||
@@ -43,6 +45,7 @@ class AnthropicClient:
|
|||||||
self.decoder = msgspec.json.Decoder(AnthropicMessageResponse)
|
self.decoder = msgspec.json.Decoder(AnthropicMessageResponse)
|
||||||
self.stream_decoder = msgspec.json.Decoder(AnthropicStreamEvent)
|
self.stream_decoder = msgspec.json.Decoder(AnthropicStreamEvent)
|
||||||
self.models_decoder = msgspec.json.Decoder(AnthropicModelsResponse)
|
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_key = config.api_key
|
||||||
self._api_version = config.anthropic_version
|
self._api_version = config.anthropic_version
|
||||||
|
|
||||||
|
|||||||
@@ -20,3 +20,6 @@ class AnthropicConfig:
|
|||||||
base_url: str = DEFAULT_ANTHROPIC_BASE_URL
|
base_url: str = DEFAULT_ANTHROPIC_BASE_URL
|
||||||
anthropic_version: str = DEFAULT_ANTHROPIC_VERSION
|
anthropic_version: str = DEFAULT_ANTHROPIC_VERSION
|
||||||
timeout: HttpTimeout = (DEFAULT_HTTP_CONNECT_TIMEOUT, DEFAULT_HTTP_READ_TIMEOUT)
|
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
|
||||||
|
|||||||
@@ -62,6 +62,18 @@ class AnthropicMetadata(Struct, omit_defaults=True):
|
|||||||
user_id: str | Unset = UNSET
|
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):
|
class AnthropicMessagesParam(Struct, omit_defaults=True):
|
||||||
model: str
|
model: str
|
||||||
max_tokens: int = 8192
|
max_tokens: int = 8192
|
||||||
@@ -75,6 +87,7 @@ class AnthropicMessagesParam(Struct, omit_defaults=True):
|
|||||||
temperature: float | Unset = UNSET
|
temperature: float | Unset = UNSET
|
||||||
top_p: float | Unset = UNSET
|
top_p: float | Unset = UNSET
|
||||||
top_k: int | Unset = UNSET
|
top_k: int | Unset = UNSET
|
||||||
|
thinking: AnthropicThinkingConfig | Unset = UNSET
|
||||||
|
|
||||||
|
|
||||||
class AnthropicUsage(Struct, omit_defaults=True):
|
class AnthropicUsage(Struct, omit_defaults=True):
|
||||||
|
|||||||
@@ -16,16 +16,17 @@ if TYPE_CHECKING:
|
|||||||
from .model import ChatCompletionsParam
|
from .model import ChatCompletionsParam
|
||||||
|
|
||||||
|
|
||||||
def _inject_thinking(data: bytes) -> bytes:
|
def _inject_thinking(data: bytes, param: ChatCompletionsParam) -> bytes:
|
||||||
"""Inject ``thinking: {type: "enabled"}`` into the encoded request body.
|
"""Inject the ``thinking`` flag into the encoded request body.
|
||||||
|
|
||||||
DeepSeek's reasoning model requires the ``thinking`` parameter to be
|
DeepSeek's reasoning model requires the ``thinking`` parameter to be
|
||||||
explicitly set (default is ``enabled``). The agent loop constructs the
|
explicitly set (default is ``enabled``); ``reasoning_effort = "none"`` is the
|
||||||
base ``ChatCompletionsParam`` which lacks this field, so we add it
|
way to ask for it off. The agent loop constructs the base
|
||||||
after msgspec encoding.
|
``ChatCompletionsParam`` which lacks this field, so we add it after msgspec
|
||||||
|
encoding.
|
||||||
"""
|
"""
|
||||||
body = json.loads(data)
|
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")
|
return json.dumps(body, separators=(",", ":")).encode("utf-8")
|
||||||
|
|
||||||
|
|
||||||
@@ -49,7 +50,7 @@ class DeepseekOpenAIClient(BaseOpenAIClient):
|
|||||||
self, param: ChatCompletionsParam, *, stream: bool = False
|
self, param: ChatCompletionsParam, *, stream: bool = False
|
||||||
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
|
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
|
||||||
param = coerce_chat_completions_param_any(msgspec.structs.replace(param, stream=stream))
|
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:
|
if stream:
|
||||||
resp = await self.session.post(
|
resp = await self.session.post(
|
||||||
"/chat/completions",
|
"/chat/completions",
|
||||||
|
|||||||
@@ -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"
|
||||||
Reference in New Issue
Block a user