diff --git a/tensorrt_llm/serve/anthropic_adapter.py b/tensorrt_llm/serve/anthropic_adapter.py index 489f4d1e0a32..cc07ae2f1b51 100644 --- a/tensorrt_llm/serve/anthropic_adapter.py +++ b/tensorrt_llm/serve/anthropic_adapter.py @@ -24,6 +24,8 @@ chat templates, or the engine. """ +import base64 +import hashlib import json import re import traceback @@ -50,6 +52,7 @@ AnthropicMessagesResponse, AnthropicMessageStartEvent, AnthropicMessageStopEvent, + AnthropicSignatureDelta, AnthropicStopReason, AnthropicTextBlock, AnthropicTextDelta, @@ -612,6 +615,21 @@ def convert_usage(usage: Optional[UsageInfo]) -> AnthropicUsage: return anthropic_usage +def thinking_signature(thinking: str) -> str: + """The signature of a thinking block this server generated. + + Anthropic clients replay a thinking block in later turns only together + with its signature, and some turn an unsigned one into a text block -- the + pi coding agent does -- so the model then reads its own reasoning back as + an answer it gave. Anthropic signs with an encrypted proof of origin; this + signature is only the base64 SHA-256 digest of the text. It proves nothing + and the server does not check it on replay. It exists so that a thinking + block round-trips as thinking. + """ + digest = hashlib.sha256(thinking.encode("utf-8")).digest() + return base64.b64encode(digest).decode("ascii") + + def convert_chat_response(chat_response: ChatCompletionResponse) -> AnthropicMessagesResponse: """Translate a non-streaming chat completion into an Anthropic message.""" content: List[Any] = [] @@ -623,7 +641,9 @@ def convert_chat_response(chat_response: ChatCompletionResponse) -> AnthropicMes message = choice.message reasoning = message.reasoning_content or message.reasoning if reasoning: - content.append(AnthropicThinkingBlock(thinking=reasoning)) + content.append( + AnthropicThinkingBlock(thinking=reasoning, signature=thinking_signature(reasoning)) + ) if message.content and message.content.strip(): content.append(AnthropicTextBlock(text=message.content)) for tool_call in message.tool_calls: @@ -723,6 +743,8 @@ def __init__(self, model: str): self.message_started = False self.block_index = -1 self.open_block_type: Optional[str] = None + # Text of the open thinking block, signed when the block closes. + self.thinking_parts: List[str] = [] # Upstream index of the tool call whose block is currently open. self.open_tool_key: Optional[Any] = None # Per-upstream-index tool call state. Insertion order is preserved and @@ -750,10 +772,23 @@ def __init__(self, model: str): def _close_block(self) -> List[str]: if self.open_block_type is None: return [] - event = AnthropicContentBlockStopEvent(index=self.block_index) + frames = [] + if self.open_block_type == "thinking": + # As Anthropic streams it: the signature is the thinking block's + # last delta, once its whole text is known. + signature = thinking_signature("".join(self.thinking_parts)) + self.thinking_parts = [] + frames.append( + anthropic_sse( + AnthropicContentBlockDeltaEvent( + index=self.block_index, delta=AnthropicSignatureDelta(signature=signature) + ) + ) + ) + frames.append(anthropic_sse(AnthropicContentBlockStopEvent(index=self.block_index))) self.open_block_type = None self.open_tool_key = None - return [anthropic_sse(event)] + return frames def _open_block(self, block: Any, block_type: str) -> List[str]: if block_type != "text": @@ -780,7 +815,11 @@ def _ensure_block(self, block_type: str) -> List[str]: if block_type == "text": frames.extend(self._open_block(AnthropicTextBlock(text=""), "text")) elif block_type == "thinking": - frames.extend(self._open_block(AnthropicThinkingBlock(thinking=""), "thinking")) + # The signature arrives as the block's last delta; the opening block + # carries it empty, as a thinking block always has the field. + frames.extend( + self._open_block(AnthropicThinkingBlock(thinking="", signature=""), "thinking") + ) else: raise ValueError(f"unexpected block type {block_type}") return frames @@ -939,6 +978,7 @@ def process_chunk(self, chunk: ChatCompletionStreamResponse) -> List[str]: reasoning_delta = delta.reasoning_content or delta.reasoning if reasoning_delta: frames.extend(self._ensure_block("thinking")) + self.thinking_parts.append(reasoning_delta) frames.append( anthropic_sse( AnthropicContentBlockDeltaEvent( diff --git a/tests/unittest/llmapi/apps/test_anthropic_adapter.py b/tests/unittest/llmapi/apps/test_anthropic_adapter.py index 82502a55c184..c9cb41d5cd1a 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_adapter.py +++ b/tests/unittest/llmapi/apps/test_anthropic_adapter.py @@ -19,6 +19,8 @@ """ import asyncio +import base64 +import hashlib import json import pytest @@ -53,6 +55,11 @@ MODEL = "test-model" +def thinking_signature(thinking: str) -> str: + """The signature the server is expected to put on a thinking block.""" + return base64.b64encode(hashlib.sha256(thinking.encode("utf-8")).digest()).decode("ascii") + + def make_request(**overrides) -> AnthropicMessagesRequest: payload = { "model": MODEL, @@ -749,6 +756,18 @@ def test_reasoning_becomes_thinking_block(): assert resp.content[1].type == "text" +def test_thinking_block_is_signed(): + """Clients replay a thinking block as thinking only when it carries a signature. + + Without one, the pi coding agent sends the block back as text, and the + model then reads its reasoning as an answer it gave. + """ + message = ChatMessage(role="assistant", content="answer", reasoning_content="step by step") + resp = convert_chat_response(make_chat_response(message)) + assert resp.content[0].signature == thinking_signature("step by step") + assert resp.model_dump(exclude_none=True)["content"][0]["signature"] + + def test_empty_content_gets_placeholder_text_block(): resp = convert_chat_response(make_chat_response(ChatMessage(role="assistant", content=None))) assert len(resp.content) == 1 @@ -1191,6 +1210,100 @@ def test_stream_thinking_then_text(): assert starts == ["thinking", "text"] +def test_stream_thinking_block_ends_with_its_signature(): + events = run_reframer( + [ + chunk({"reasoning_content": "first "}), + chunk({"reasoning_content": "second"}), + chunk({"content": "answer"}), + chunk({}, finish_reason="stop"), + ] + ) + assert_event_invariants(events) + thinking = [ + (name, payload) + for name, payload in events + if name in ("content_block_delta", "content_block_stop") and payload["index"] == 0 + ] + assert [name for name, _ in thinking] == [ + "content_block_delta", + "content_block_delta", + "content_block_delta", + "content_block_stop", + ] + assert thinking[2][1]["delta"] == { + "type": "signature_delta", + "signature": thinking_signature("first second"), + } + signed = [ + p + for n, p in events + if n == "content_block_delta" and p["delta"]["type"] == "signature_delta" + ] + assert len(signed) == 1, "only the thinking block is signed" + + +def test_stream_thinking_block_opens_with_an_empty_signature(): + """A thinking block always has a signature field, even before its last delta fills it.""" + events = run_reframer( + [ + chunk({"reasoning_content": "first"}), + chunk({"content": "answer"}), + chunk({}, finish_reason="stop"), + ] + ) + assert_event_invariants(events) + starts = [p["content_block"] for n, p in events if n == "content_block_start"] + assert starts[0] == {"type": "thinking", "thinking": "", "signature": ""} + assert "signature" not in starts[1] + + +def test_stream_thinking_left_open_is_signed_when_the_stream_ends(): + events = run_reframer( + [ + chunk({"reasoning_content": "unfinished"}), + chunk({}, finish_reason="length"), + ] + ) + assert_event_invariants(events) + deltas = [p["delta"] for n, p in events if n == "content_block_delta"] + assert deltas[-1] == {"type": "signature_delta", "signature": thinking_signature("unfinished")} + + +def test_stream_error_signs_the_open_thinking_block_before_the_error(): + reframer = AnthropicStreamReframer(model=MODEL) + frames = reframer.process_chunk(chunk({"reasoning_content": "partial"})) + frames += reframer.error("engine failed") + events = parse_frames(frames) + (delta_name, delta), (stop_name, stop), (error_name, error) = events[-3:] + assert [delta_name, stop_name, error_name] == [ + "content_block_delta", + "content_block_stop", + "error", + ] + assert delta["delta"] == {"type": "signature_delta", "signature": thinking_signature("partial")} + assert stop["index"] == delta["index"] + assert error["error"]["message"] == "engine failed" + + +def test_stream_each_thinking_block_is_signed_with_its_own_text(): + events = run_reframer( + [ + chunk({"reasoning_content": "one"}), + chunk({"content": "text"}), + chunk({"reasoning_content": "two"}), + chunk({}, finish_reason="stop"), + ] + ) + assert_event_invariants(events) + signatures = [ + p["delta"]["signature"] + for n, p in events + if n == "content_block_delta" and p["delta"]["type"] == "signature_delta" + ] + assert signatures == [thinking_signature("one"), thinking_signature("two")] + + def test_stream_empty_generation_still_valid(): events = run_reframer([chunk({}, finish_reason="stop")]) assert_event_invariants(events)