diff --git a/src/plyngent/agent/loop.py b/src/plyngent/agent/loop.py index b3cda4f..9a8105e 100644 --- a/src/plyngent/agent/loop.py +++ b/src/plyngent/agent/loop.py @@ -136,11 +136,16 @@ async def _execute_tool_calls( def _assistant_has_payload(assistant: AssistantChatMessage) -> bool: - """True if the model produced text, reasoning, or tool calls.""" + """True if the model produced text, reasoning, provider blocks, or tool calls.""" if assistant.tool_calls is not UNSET and assistant.tool_calls: return True if isinstance(assistant.content, str) and assistant.content.strip(): return True + # A round that only searched (a hosted ``web_search``) kept the search + # blocks as provider blocks: the model did produce something, even with no + # text of its own yet. + if assistant.provider_blocks is not UNSET and assistant.provider_blocks: + return True reasoning = assistant.reasoning_content return bool(isinstance(reasoning, str) and reasoning.strip()) @@ -283,6 +288,23 @@ def _take_reasoning_delta( return [content], [ReasoningDeltaEvent(content=fresh)] if fresh else [] +def _accumulate_deltas( + delta: DeltaMessage, + *, + tool_deltas: list[StreamToolCallDelta], + provider_blocks: list[dict[str, Any]], +) -> None: + """Collect the carriers a round keeps from one chunk's delta. + + Tool calls merge after the stream; provider-opaque blocks (a hosted search's + blocks, a thinking block) were rebuilt by the bridge and ride along as-is. + """ + if delta.tool_calls is not UNSET and delta.tool_calls: + tool_deltas.extend(delta.tool_calls) + if delta.provider_blocks is not UNSET and delta.provider_blocks: + provider_blocks.extend(delta.provider_blocks) + + async def _stream_round( client: AnyLLMClient, param: ChatCompletionsParam, @@ -306,6 +328,7 @@ async def _stream_round( ) content_parts: list[str] = [] reasoning_parts: list[str] = [] + provider_blocks: list[dict[str, Any]] = [] tool_deltas: list[StreamToolCallDelta] = [] last_api_usage: object = UNSET finish_reason: str | None = None @@ -332,8 +355,7 @@ async def _stream_round( if isinstance(delta.content, str) and delta.content: content_parts.append(delta.content) yield TextDeltaEvent(content=delta.content) - if delta.tool_calls is not UNSET and delta.tool_calls: - tool_deltas.extend(delta.tool_calls) + _accumulate_deltas(delta, tool_deltas=tool_deltas, provider_blocks=provider_blocks) full_content = "".join(content_parts) full_reasoning = "".join(reasoning_parts) @@ -347,6 +369,7 @@ async def _stream_round( content=full_content or None, tool_calls=tool_calls, reasoning_content=full_reasoning or UNSET, + provider_blocks=provider_blocks or UNSET, ) # Usage-only final chunks leave choices empty; a non-empty usage after # content still is not a finish_reason. Prefer explicit finish_reason. diff --git a/src/plyngent/agent/messages_dispatch.py b/src/plyngent/agent/messages_dispatch.py index 7ff6dfb..a0f8c8d 100644 --- a/src/plyngent/agent/messages_dispatch.py +++ b/src/plyngent/agent/messages_dispatch.py @@ -8,7 +8,8 @@ enter it. from __future__ import annotations -from typing import TYPE_CHECKING, Any, Final +import json +from typing import TYPE_CHECKING, Any, Final, cast from plyngent.agent.messages_bridge import ( anthropic_response_to_chat_completion, @@ -18,6 +19,8 @@ from plyngent.agent.messages_bridge import ( ) from plyngent.agent.stream_chunks import ( finish_reason_chunk, + provider_blocks_chunk, + reasoning_delta_chunk, text_delta_chunk, tool_call_delta_chunk, usage_chunk, @@ -25,6 +28,7 @@ from plyngent.agent.stream_chunks import ( from plyngent.lmproto.anthropic.model import ( AnthropicContentBlockDelta, AnthropicContentBlockStart, + AnthropicContentBlockStop, AnthropicErrorEvent, AnthropicMessageDelta, AnthropicMessageStart, @@ -65,6 +69,92 @@ class _ToolBlockState: self.tool_index = tool_index +# Block types the chat-shaped round carries itself: text becomes the message's +# content and ``tool_use`` a local tool call. Every other type is the surface's +# own (``thinking``, ``server_tool_use``, ``web_search_tool_result``) and is +# rebuilt so the round can hand it back. +_LOCAL_BLOCK_TYPES: Final = frozenset({"text", "tool_use"}) + +# The two fields a *delta* owns; a block's own JSON skips them. +_DELTA_ONLY_FIELDS: Final = frozenset({"text", "partial_json"}) + + +def _raw_block_json(block: AnthropicRawContentBlock) -> dict[str, Any]: + """The block's own JSON: every field the stream set, under its own name. + + Field-driven, so a field added to the block model comes along without a + second list to keep in sync. + """ + out: dict[str, Any] = {"type": block.type} + for name in block.__struct_fields__: + if name in _DELTA_ONLY_FIELDS: + continue + value: object = getattr(block, name) + if value is not None: + out[name] = value + return out + + +def _parse_input_fragments(fragments: list[str]) -> dict[str, Any]: + """The ``input`` a ``server_tool_use`` streamed as ``partial_json`` chunks.""" + try: + parsed = json.loads("".join(fragments)) + except json.JSONDecodeError: + return {} + if isinstance(parsed, dict): + return cast("dict[str, Any]", parsed) + return {} + + +class _ProviderBlockState: + """A surface-owned block rebuilt from its start payload and fragments. + + Nothing local asked for it, but it goes back on the wire with the round (see + ``AssistantChatMessage.provider_blocks``): a ``thinking`` block's text and + its opaque ``signature``, and the search query a ``server_tool_use`` streams + as ``partial_json``. + """ + + __slots__: Final = ("json_fragments", "parts", "raw", "signature") + raw: dict[str, Any] + parts: list[str] + json_fragments: list[str] + signature: str + + def __init__(self, raw: dict[str, Any]) -> None: + self.raw = raw + self.parts = [] + self.json_fragments = [] + self.signature = "" + + def add(self, delta: AnthropicRawContentBlock) -> str | None: + """Fold one fragment in; return the thinking text it carried, if any.""" + if isinstance(delta.partial_json, str) and delta.partial_json: + self.json_fragments.append(delta.partial_json) + return None + if isinstance(delta.signature, str) and delta.signature: + self.signature = delta.signature + return None + # Anthropic streams a thinking fragment as ``thinking``; DeepSeek's + # compat also spells the same fragment ``text``. + fragment = delta.thinking if isinstance(delta.thinking, str) else delta.text + if not isinstance(fragment, str) or not fragment: + return None + self.parts.append(fragment) + return fragment + + def finish(self) -> dict[str, Any]: + """The block to hand back, fragments folded into its start payload.""" + block = dict(self.raw) + if self.parts: + block["thinking"] = f"{block.get('thinking', '')}{''.join(self.parts)}" + if self.json_fragments: + block["input"] = _parse_input_fragments(self.json_fragments) + if self.signature: + block["signature"] = self.signature + return block + + def _merge_usage(previous: dict[str, Any] | None, new: dict[str, Any]) -> dict[str, Any]: """Merge ``message_start`` usage with the later ``message_delta`` usage. @@ -99,6 +189,7 @@ class _StreamState: model: str tool_blocks: dict[int, _ToolBlockState] block_types: dict[int, str] + provider_blocks: dict[int, _ProviderBlockState] next_tool_index: int stop_reason: str | None last_usage: dict[str, Any] | None @@ -109,6 +200,7 @@ class _StreamState: self.model = model self.tool_blocks = {} self.block_types = {} + self.provider_blocks = {} self.next_tool_index = 0 self.stop_reason = None self.last_usage = None @@ -140,13 +232,17 @@ def _handle_block_start( ) -> list[ChatCompletionChunk]: """Record the block's type; only ``tool_use`` becomes a local tool call. - ``server_tool_use`` (a search the model's backend runs), its - ``web_search_tool_result``, and ``thinking`` stream here too — recording - their type keeps the argument fragments that follow from being mistaken for - a local call's arguments (see ``_chunks_for_content_delta``). + A surface-owned block (``server_tool_use``, its ``web_search_tool_result``, + ``thinking``) is rebuilt from here on so the round can hand it back (see + ``AssistantChatMessage.provider_blocks``); recording its type also keeps the + fragments that follow from being mistaken for a local call's arguments (see + ``_chunks_for_content_delta``). """ block = event.content_block state.block_types[event.index] = block.type + if block.type not in _LOCAL_BLOCK_TYPES: + state.provider_blocks[event.index] = _ProviderBlockState(_raw_block_json(block)) + return [] if block.type != "tool_use": return [] call_id = block.id or f"toolu_{event.index}" @@ -172,6 +268,17 @@ def _handle_message_delta(event: AnthropicMessageDelta, state: _StreamState) -> state.last_usage = _merge_usage(state.last_usage, usage) +def _chunks_for_block_stop( + event: AnthropicContentBlockStop, + state: _StreamState, +) -> list[ChatCompletionChunk]: + """Emit a surface-owned block once its fragments are folded in.""" + provider = state.provider_blocks.pop(event.index, None) + if provider is None: + return [] + return [provider_blocks_chunk(model=state.model, blocks=[provider.finish()])] + + def _chunks_for_event(event: AnthropicStreamEvent, state: _StreamState) -> list[ChatCompletionChunk]: """Convert one Anthropic stream event into zero or more chat chunks.""" if isinstance(event, AnthropicErrorEvent): @@ -188,6 +295,9 @@ def _chunks_for_event(event: AnthropicStreamEvent, state: _StreamState) -> list[ if isinstance(event, AnthropicContentBlockStart): return _handle_block_start(event, state) + if isinstance(event, AnthropicContentBlockStop): + return _chunks_for_block_stop(event, state) + if isinstance(event, AnthropicContentBlockDelta): return _chunks_for_content_delta(event, state) @@ -201,11 +311,31 @@ def _chunks_for_event(event: AnthropicStreamEvent, state: _StreamState) -> list[ return [] +def _chunks_for_provider_delta( + delta: AnthropicRawContentBlock, + provider: _ProviderBlockState, + *, + model: str, +) -> list[ChatCompletionChunk]: + """One fragment of a surface-owned block: its reasoning text, or nothing.""" + fragment = provider.add(delta) + if fragment is None: + return [] + return [reasoning_delta_chunk(model=model, content=fragment)] + + def _chunks_for_content_delta( event: AnthropicContentBlockDelta, state: _StreamState, ) -> list[ChatCompletionChunk]: delta: AnthropicRawContentBlock = event.delta + provider = state.provider_blocks.get(event.index) + if provider is not None: + # A surface-owned block's fragments never become local output here: the + # thinking text streams as reasoning, while the search query and the + # signature only go back with the round. + return _chunks_for_provider_delta(delta, provider, model=state.model) + if delta.type in {"text_delta", "text"} and isinstance(delta.text, str) and delta.text: return [text_delta_chunk(model=state.model, content=delta.text)] diff --git a/tests/test_agent/test_messages_dispatch.py b/tests/test_agent/test_messages_dispatch.py index 98db65a..c43d97d 100644 --- a/tests/test_agent/test_messages_dispatch.py +++ b/tests/test_agent/test_messages_dispatch.py @@ -8,6 +8,7 @@ from msgspec import UNSET from plyngent.agent import ( AssistantMessageEvent, + ReasoningDeltaEvent, TextDeltaEvent, ToolCallEvent, ToolRegistry, @@ -396,9 +397,9 @@ async def test_dispatch_stream_tool_calls_without_stop_yields_no_finish() -> Non assert not any(c.choices and c.choices[0].finish_reason not in (None, UNSET) for c in chunks) -def _search_stream_events() -> list[AnthropicStreamEvent]: +def _search_stream_events(*, text: bool = True) -> list[AnthropicStreamEvent]: """The stream a DeepSeek search turn sends: thinking, search, results, text.""" - return [ + events: list[AnthropicStreamEvent] = [ AnthropicMessageStart( message=AnthropicMessageResponse( id="msg_s", @@ -407,11 +408,16 @@ def _search_stream_events() -> list[AnthropicStreamEvent]: usage=AnthropicUsage(input_tokens=164, output_tokens=0), ) ), - AnthropicContentBlockStart(index=0, content_block=AnthropicRawContentBlock(type="thinking", text="")), + AnthropicContentBlockStart( + index=0, + content_block=AnthropicRawContentBlock(type="thinking", thinking=""), + ), + AnthropicContentBlockDelta(index=0, delta=AnthropicRawContentBlock(type="thinking_delta", thinking="Search ")), AnthropicContentBlockDelta( index=0, - delta=AnthropicRawContentBlock(type="thinking_delta", text="Search."), + delta=AnthropicRawContentBlock(type="thinking_delta", thinking="the docs."), ), + AnthropicContentBlockDelta(index=0, delta=AnthropicRawContentBlock(type="signature_delta", signature="sig-1")), AnthropicContentBlockStop(index=0), AnthropicContentBlockStart( index=1, @@ -430,23 +436,39 @@ def _search_stream_events() -> list[AnthropicStreamEvent]: index=2, content_block=AnthropicRawContentBlock( type="web_search_tool_result", - id="call_00_1", - input={"content": [{"type": "web_search_result", "url": "https://api-docs.deepseek.com"}]}, + tool_use_id="call_00_1", + content=[ + { + "type": "web_search_result", + "url": "https://api-docs.deepseek.com", + "encrypted_content": "sd918tZXSvVp", + } + ], ), ), AnthropicContentBlockStop(index=2), - AnthropicContentBlockStart(index=3, content_block=AnthropicRawContentBlock(type="text", text="")), - AnthropicContentBlockDelta( - index=3, - delta=AnthropicRawContentBlock(type="text_delta", text="DeepSeek-V4 is out."), - ), - AnthropicContentBlockStop(index=3), - AnthropicMessageDelta( - delta={"stop_reason": "end_turn"}, - usage=AnthropicUsage(input_tokens=0, output_tokens=12), - ), - AnthropicMessageStop(), ] + if text: + events.extend( + [ + AnthropicContentBlockStart(index=3, content_block=AnthropicRawContentBlock(type="text", text="")), + AnthropicContentBlockDelta( + index=3, + delta=AnthropicRawContentBlock(type="text_delta", text="DeepSeek-V4 is out."), + ), + AnthropicContentBlockStop(index=3), + ] + ) + events.extend( + [ + AnthropicMessageDelta( + delta={"stop_reason": "end_turn"}, + usage=AnthropicUsage(input_tokens=0, output_tokens=12), + ), + AnthropicMessageStop(), + ] + ) + return events async def test_search_turn_streams_its_text_without_a_local_tool_call() -> None: @@ -476,10 +498,58 @@ async def test_search_turn_streams_its_text_without_a_local_tool_call() -> None: assert not any(isinstance(event, ToolCallEvent) for event in events) assert not any(isinstance(event, ToolResultEvent) for event in events) assert any(isinstance(event, TextDeltaEvent) and event.content == "DeepSeek-V4 is out." for event in events) + # The chain of thought streams as reasoning, fragment by fragment. + assert [event.content for event in events if isinstance(event, ReasoningDeltaEvent)] == ["Search ", "the docs."] + # ...and the round keeps the search blocks, rebuilt, for the next request. + assistant = next(event.message for event in events if isinstance(event, AssistantMessageEvent)) + assert assistant.reasoning_content == "Search the docs." + assert assistant.provider_blocks == [ + {"type": "thinking", "thinking": "Search the docs.", "signature": "sig-1"}, + { + "type": "server_tool_use", + "id": "call_00_1", + "name": "web_search", + "input": {"query": "DeepSeek-V4"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "call_00_1", + "content": [ + { + "type": "web_search_result", + "url": "https://api-docs.deepseek.com", + "encrypted_content": "sd918tZXSvVp", + } + ], + }, + ] # One round: the search answered inside the response, so nothing was sent back. assert len(client.calls) == 1 +async def test_search_only_round_is_not_an_empty_completion() -> None: + """A search turn with no text yet still produced payload. + + A budget-capped or ``max_uses_exceeded`` search can end the turn on its + blocks alone; calling that an empty completion failed a round the model had + actually used. + """ + client = ScriptedMessagesClient(stream_events=[_search_stream_events(text=False)]) + events = [ + event + async for event in run_chat_loop( + cast("Any", client), + [UserChatMessage(content="what is DeepSeek-V4?")], + model="deepseek-flash", + stream=True, + ) + ] + assistant = next(event.message for event in events if isinstance(event, AssistantMessageEvent)) + assert assistant.content is None + assert assistant.provider_blocks is not UNSET + assert len(assistant.provider_blocks) == 3 + + async def test_dispatch_messages_sends_provider_tools_as_server_tools() -> None: """Hosted tools reach the Anthropic request in that surface's own shape.""" client = ScriptedMessagesClient(non_stream=[_message_response(text="ok")])