diff --git a/src/plyngent/agent/compact.py b/src/plyngent/agent/compact.py index 871ad91..6112155 100644 --- a/src/plyngent/agent/compact.py +++ b/src/plyngent/agent/compact.py @@ -159,6 +159,13 @@ async def summarize_messages( stream=False, ), ) + elif kind == "messages": + from .messages_dispatch import dispatch_messages + + response = cast( + "ChatCompletionResponse", + await dispatch_messages(cast("Any", client), param, stream=False), + ) else: msg = f"summarization for client.kind={kind!r} is not implemented" raise NotImplementedError(msg) diff --git a/src/plyngent/agent/loop.py b/src/plyngent/agent/loop.py index 25cf30b..3a25391 100644 --- a/src/plyngent/agent/loop.py +++ b/src/plyngent/agent/loop.py @@ -54,6 +54,7 @@ from .usage import resolve_round_usage, token_usage_from_api if TYPE_CHECKING: from collections.abc import AsyncIterator, Awaitable, Callable, Sequence + from plyngent.lmproto.anthropic.client import AnthropicClient from plyngent.lmproto.openai.client import OpenAIClient from plyngent.lmproto.openai_compatible.model import AnyChatMessage, AnyToolItem @@ -210,8 +211,13 @@ async def _dispatch_chat_completions( stream=stream, ) if kind == "messages": - msg = "Anthropic /messages dispatch is not wired in the agent loop yet" - raise NotImplementedError(msg) + from .messages_dispatch import dispatch_messages + + return await dispatch_messages( + cast("AnthropicClient", client), + param, + stream=stream, + ) msg = f"client kind {kind!r} is not supported by the agent loop" raise NotImplementedError(msg) diff --git a/src/plyngent/agent/messages_bridge.py b/src/plyngent/agent/messages_bridge.py new file mode 100644 index 0000000..5f68b1a --- /dev/null +++ b/src/plyngent/agent/messages_bridge.py @@ -0,0 +1,418 @@ +"""Convert agent chat history/tools to Anthropic Messages API shapes and back. + +Agent memory and events stay chat-completions-shaped; only the transport uses +``POST /messages``. OpenAI / DeepSeek paths never enter this module. +""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any, cast + +from msgspec import UNSET + +from plyngent.lmproto.anthropic.model import ( + AnthropicAssistantMessage, + AnthropicMessageResponse, + AnthropicMessagesParam, + AnthropicResponseText, + AnthropicTextContent, + AnthropicToolChoice, + AnthropicToolDefinition, + AnthropicToolResultContent, + AnthropicToolUseContent, + AnthropicUsage, + AnthropicUserMessage, +) +from plyngent.lmproto.openai_compatible.model import ( + AnyAssistantToolCall, + AssistantChatMessage, + AssistantFunctionTool, + AssistantFunctionToolCall, + ChatCompletionChoice, + ChatCompletionChunk, + ChatCompletionResponse, + ChatCompletionsParam, + ChunkChoice, + DeltaMessage, + DeveloperChatMessage, + StreamFunctionDelta, + StreamToolCallDelta, + SystemChatMessage, + ToolFunctionItem, + UserChatMessage, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + + from plyngent.lmproto.anthropic.model import AnthropicMessage + from plyngent.lmproto.openai_compatible.model import AnyChatMessage, AnyToolItem + from plyngent.typedef import Unset + + +def _parse_tool_arguments(arguments: str) -> dict[str, Any]: + """Decode tool-call argument JSON into an object; empty/invalid → ``{}``.""" + text = arguments.strip() if arguments else "" + if not text: + return {} + try: + parsed = json.loads(text) + except json.JSONDecodeError: + return {} + if isinstance(parsed, dict): + return cast("dict[str, Any]", parsed) + return {} + + +def _encode_tool_arguments(obj: dict[str, Any]) -> str: + return json.dumps(obj, ensure_ascii=False, separators=(",", ":")) + + +def _as_nonneg_int(value: object) -> int: + if isinstance(value, bool): + return 0 + if isinstance(value, int): + return max(0, value) + if isinstance(value, float): + return max(0, int(value)) + return 0 + + +def tool_items_to_anthropic_tools( + tools: Sequence[AnyToolItem] | None, +) -> list[AnthropicToolDefinition]: + """Map chat ``ToolFunctionItem`` list to Anthropic tool definitions.""" + if not tools: + return [] + result: list[AnthropicToolDefinition] = [] + for item in tools: + if not isinstance(item, ToolFunctionItem): + continue + fn = item.function + result.append( + AnthropicToolDefinition( + name=fn.name, + description=fn.description if fn.description is not UNSET else UNSET, + input_schema=fn.parameters if fn.parameters is not UNSET else UNSET, + ) + ) + return result + + +def _tool_choice_from_chat(param: ChatCompletionsParam) -> AnthropicToolChoice | Unset: + if param.tool_choice is UNSET: + return UNSET + choice = param.tool_choice + if not isinstance(choice, str): + return UNSET + if choice in {"none", "auto"}: + # Anthropic has no "none"; keep auto (callers can omit tools for hard none). + return AnthropicToolChoice(type="auto") + if choice == "required": + return AnthropicToolChoice(type="any") + return UNSET + + +def _assistant_to_anthropic(message: AssistantChatMessage) -> AnthropicAssistantMessage | None: + blocks: list[AnthropicTextContent | AnthropicToolUseContent] = [] + if isinstance(message.content, str) and message.content: + blocks.append(AnthropicTextContent(text=message.content)) + if message.tool_calls is not UNSET and message.tool_calls: + for call in message.tool_calls: + if not isinstance(call, AssistantFunctionToolCall): + continue + blocks.append( + AnthropicToolUseContent( + id=call.id, + name=call.function.name, + input=_parse_tool_arguments(call.function.arguments), + ) + ) + if not blocks: + return None + if len(blocks) == 1 and isinstance(blocks[0], AnthropicTextContent): + return AnthropicAssistantMessage(content=blocks[0].text) + return AnthropicAssistantMessage(content=blocks) + + +def chat_messages_to_anthropic( # noqa: C901 — multi-role conversion + messages: Sequence[AnyChatMessage], +) -> tuple[str | None, list[AnthropicMessage]]: + """Split leading system prompts; convert the rest to Anthropic messages. + + Mid-turn developer messages become user text with a ``[developer]`` prefix + (Anthropic has no developer role). Consecutive tool results are merged into + a single user message with multiple ``tool_result`` blocks, which is the + shape Anthropic expects after an assistant ``tool_use`` turn. + """ + system_parts: list[str] = [] + out: list[AnthropicMessage] = [] + pending_tool_results: list[AnthropicToolResultContent] = [] + + def flush_tool_results() -> None: + nonlocal pending_tool_results + if not pending_tool_results: + return + out.append(AnthropicUserMessage(content=list(pending_tool_results))) + pending_tool_results = [] + + for message in messages: + if isinstance(message, SystemChatMessage): + # Leading system only becomes top-level ``system``; mid-history + # system (unusual) is treated as developer-style control text. + if not out and not pending_tool_results: + if message.content.strip(): + system_parts.append(message.content) + else: + flush_tool_results() + if message.content.strip(): + out.append(AnthropicUserMessage(content=f"[system] {message.content}")) + continue + + if isinstance(message, DeveloperChatMessage): + flush_tool_results() + if message.content.strip(): + out.append(AnthropicUserMessage(content=f"[developer] {message.content}")) + continue + + if isinstance(message, UserChatMessage): + flush_tool_results() + out.append(AnthropicUserMessage(content=message.content)) + continue + + if isinstance(message, AssistantChatMessage): + flush_tool_results() + converted = _assistant_to_anthropic(message) + if converted is not None: + out.append(converted) + continue + + # ToolChatMessage + pending_tool_results.append( + AnthropicToolResultContent( + tool_use_id=message.tool_call_id, + content=message.content, + ) + ) + + flush_tool_results() + system = "\n\n".join(system_parts) if system_parts else None + return system, out + + +def anthropic_response_to_assistant(response: AnthropicMessageResponse) -> AssistantChatMessage: + """Map a completed Anthropic message to agent ``AssistantChatMessage``.""" + text_parts: list[str] = [] + tool_calls: list[AnyAssistantToolCall] = [] + for block in response.content: + if isinstance(block, AnthropicResponseText): + if block.text: + text_parts.append(block.text) + continue + # AnthropicResponseToolUse (remaining content arm) + tool_calls.append( + AssistantFunctionToolCall( + id=block.id, + function=AssistantFunctionTool( + name=block.name, + arguments=_encode_tool_arguments(block.input), + ), + ) + ) + text = "".join(text_parts) + return AssistantChatMessage( + content=text or None, + tool_calls=tool_calls or UNSET, + ) + + +def anthropic_stop_to_finish_reason( + stop_reason: str | None, + *, + has_tool_calls: bool, +) -> str: + """Map Anthropic ``stop_reason`` to a chat-style finish_reason.""" + reason = (stop_reason or "").lower() or None + if reason == "max_tokens": + return "length" + if reason in {"refusal", "content_filter"}: + return "content_filter" + if reason == "tool_use" or has_tool_calls: + return "tool_calls" + # end_turn, stop_sequence, None, … + return "stop" + + +def anthropic_usage_to_dict(usage: AnthropicUsage | dict[str, Any] | None) -> dict[str, Any] | None: + """Normalize Anthropic usage Struct/dict to chat-compatible usage fields.""" + if usage is None: + return None + if isinstance(usage, AnthropicUsage): + prompt = _as_nonneg_int(usage.input_tokens) + completion = _as_nonneg_int(usage.output_tokens) + else: + raw = cast("dict[str, object]", usage) + prompt = _as_nonneg_int(raw.get("input_tokens") or raw.get("prompt_tokens") or 0) + completion = _as_nonneg_int(raw.get("output_tokens") or raw.get("completion_tokens") or 0) + if prompt == 0 and completion == 0: + return None + return { + "prompt_tokens": prompt, + "completion_tokens": completion, + "total_tokens": prompt + completion, + "input_tokens": prompt, + "output_tokens": completion, + } + + +def anthropic_response_to_chat_completion( + response: AnthropicMessageResponse, +) -> ChatCompletionResponse: + """Wrap an Anthropic message as a synthetic chat completion for the agent loop.""" + assistant = anthropic_response_to_assistant(response) + has_tools = assistant.tool_calls is not UNSET and bool(assistant.tool_calls) + finish = anthropic_stop_to_finish_reason(response.stop_reason, has_tool_calls=has_tools) + usage = anthropic_usage_to_dict(response.usage) + return ChatCompletionResponse( + id=response.id, + object="chat.completion", + created=0, + model=response.model or "", + choices=[ + ChatCompletionChoice( + index=0, + message=assistant, + finish_reason=cast("Any", finish), + ) + ], + usage=cast("Any", usage) if usage is not None else UNSET, + ) + + +def finish_reason_chunk( + *, + model: str, + finish_reason: str, + created: int = 0, +) -> ChatCompletionChunk: + return ChatCompletionChunk( + id="anthropic-stream", + object="chat.completion.chunk", + created=created, + model=model, + choices=[ + ChunkChoice( + index=0, + delta=DeltaMessage(), + finish_reason=cast("Any", finish_reason), + ) + ], + ) + + +def text_delta_chunk(*, model: str, content: str, created: int = 0) -> ChatCompletionChunk: + return ChatCompletionChunk( + id="anthropic-stream", + object="chat.completion.chunk", + created=created, + model=model, + choices=[ + ChunkChoice( + index=0, + delta=DeltaMessage(content=content), + ) + ], + ) + + +def tool_call_delta_chunk( + *, + model: str, + index: int, + call_id: str | None = None, + name: str | None = None, + arguments: str | None = None, + created: int = 0, +) -> ChatCompletionChunk: + return ChatCompletionChunk( + id="anthropic-stream", + object="chat.completion.chunk", + created=created, + model=model, + choices=[ + ChunkChoice( + index=0, + delta=DeltaMessage( + tool_calls=[ + StreamToolCallDelta( + index=index, + id=call_id if call_id is not None else UNSET, + type="function", + function=StreamFunctionDelta( + name=name if name is not None else UNSET, + arguments=arguments if arguments is not None else UNSET, + ), + ) + ] + ), + ) + ], + ) + + +def usage_chunk(*, model: str, usage: dict[str, Any], created: int = 0) -> ChatCompletionChunk: + return ChatCompletionChunk( + id="anthropic-stream", + object="chat.completion.chunk", + created=created, + model=model, + choices=[], + usage=usage, + ) + + +def _max_tokens_from_param(param: ChatCompletionsParam) -> int: + if param.max_completion_tokens is not UNSET: + return int(param.max_completion_tokens) + if param.max_tokens is not UNSET: + return int(param.max_tokens) + return 8192 + + +def chat_param_to_anthropic_param(param: ChatCompletionsParam) -> AnthropicMessagesParam: + """Build :class:`AnthropicMessagesParam` from a chat completions param.""" + system, messages = chat_messages_to_anthropic(param.messages) + tools = tool_items_to_anthropic_tools(param.tools if param.tools is not UNSET else None) + + kwargs: dict[str, Any] = { + "model": param.model, + "max_tokens": _max_tokens_from_param(param), + "messages": messages, + } + if system: + kwargs["system"] = system + if tools: + kwargs["tools"] = tools + tool_choice = _tool_choice_from_chat(param) + if tool_choice is not UNSET: + kwargs["tool_choice"] = tool_choice + if param.parallel_tool_calls is not UNSET and param.parallel_tool_calls is False: + existing = kwargs.get("tool_choice") + if isinstance(existing, AnthropicToolChoice): + kwargs["tool_choice"] = AnthropicToolChoice( + type=existing.type, + name=existing.name, + disable_parallel_tool_use=True, + ) + else: + kwargs["tool_choice"] = AnthropicToolChoice( + type="auto", + disable_parallel_tool_use=True, + ) + if param.temperature is not UNSET: + kwargs["temperature"] = param.temperature + if param.top_p is not UNSET: + kwargs["top_p"] = param.top_p + + return AnthropicMessagesParam(**kwargs) diff --git a/src/plyngent/agent/messages_dispatch.py b/src/plyngent/agent/messages_dispatch.py new file mode 100644 index 0000000..ee99f33 --- /dev/null +++ b/src/plyngent/agent/messages_dispatch.py @@ -0,0 +1,253 @@ +"""Anthropic Messages dispatch: present ``/messages`` as chat completions to the loop. + +Agent memory and events stay chat-completions-shaped; only the transport uses +``POST /messages``. OpenAI / DeepSeek paths never enter this module. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Final + +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, + finish_reason_chunk, + text_delta_chunk, + tool_call_delta_chunk, + usage_chunk, +) +from plyngent.lmproto.anthropic.model import ( + AnthropicContentBlockDelta, + AnthropicContentBlockStart, + AnthropicErrorEvent, + AnthropicMessageDelta, + AnthropicMessageStart, + AnthropicMessageStop, + AnthropicRawContentBlock, +) + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from plyngent.lmproto.anthropic.client import AnthropicClient + from plyngent.lmproto.anthropic.model import AnthropicStreamEvent + from plyngent.lmproto.openai_compatible.model import ( + ChatCompletionChunk, + ChatCompletionResponse, + ChatCompletionsParam, + ) + + +class _ToolBlockState: + """Tracks an in-flight tool_use content block during SSE streaming.""" + + __slots__: Final = ("call_id", "index", "name", "tool_index") + index: int + call_id: str + name: str + tool_index: int + + def __init__(self, *, index: int, call_id: str, name: str, tool_index: int) -> None: + self.index = index + self.call_id = call_id + self.name = name + self.tool_index = tool_index + + +def _merge_usage(previous: dict[str, Any] | None, new: dict[str, Any]) -> dict[str, Any]: + """Merge message_start input tokens with message_delta output tokens.""" + if previous is None: + return new + prompt = int(previous.get("prompt_tokens", 0) or new.get("prompt_tokens", 0) or 0) + completion = int(new.get("completion_tokens", 0) or previous.get("completion_tokens", 0) or 0) + return { + "prompt_tokens": prompt, + "completion_tokens": completion, + "total_tokens": prompt + completion, + "input_tokens": prompt, + "output_tokens": completion, + } + + +def _error_detail(error: dict[str, Any]) -> str: + msg = error.get("message") + if isinstance(msg, str) and msg.strip(): + return msg + return str(error) + + +class _StreamState: + """Mutable accumulator while converting Anthropic SSE to chat chunks.""" + + model: str + tool_blocks: dict[int, _ToolBlockState] + next_tool_index: int + stop_reason: str | None + last_usage: dict[str, Any] | None + saw_message_stop: bool + has_tool_calls: bool + + def __init__(self, model: str) -> None: + self.model = model + self.tool_blocks = {} + self.next_tool_index = 0 + self.stop_reason = None + self.last_usage = None + self.saw_message_stop = False + self.has_tool_calls = False + + def start_tool(self, *, block_index: int, call_id: str, name: str) -> _ToolBlockState: + state = _ToolBlockState( + index=block_index, + call_id=call_id, + name=name, + tool_index=self.next_tool_index, + ) + self.tool_blocks[block_index] = state + self.next_tool_index += 1 + self.has_tool_calls = True + return state + + def ensure_tool(self, block_index: int) -> _ToolBlockState: + state = self.tool_blocks.get(block_index) + if state is not None: + return state + return self.start_tool(block_index=block_index, call_id=f"toolu_{block_index}", name="") + + +def _handle_tool_use_start( + event: AnthropicContentBlockStart, + state: _StreamState, +) -> list[ChatCompletionChunk]: + block = event.content_block + if block.type != "tool_use": + return [] + call_id = block.id or f"toolu_{event.index}" + name = block.name or "" + tool_state = state.start_tool(block_index=event.index, call_id=call_id, name=name) + return [ + tool_call_delta_chunk( + model=state.model, + index=tool_state.tool_index, + call_id=call_id, + name=name, + arguments="", + ) + ] + + +def _handle_message_delta(event: AnthropicMessageDelta, state: _StreamState) -> None: + raw_stop = event.delta.get("stop_reason") + if isinstance(raw_stop, str) and raw_stop: + state.stop_reason = raw_stop + usage = anthropic_usage_to_dict(event.usage) + if usage is not None: + state.last_usage = _merge_usage(state.last_usage, usage) + + +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): + detail = _error_detail(event.error) + msg = f"anthropic stream error: {detail}" + raise RuntimeError(msg) # noqa: TRY004 — stream transport failure + + if isinstance(event, AnthropicMessageStart): + usage = anthropic_usage_to_dict(event.message.usage) + if usage is not None: + state.last_usage = usage + return [] + + if isinstance(event, AnthropicContentBlockStart): + return _handle_tool_use_start(event, state) + + if isinstance(event, AnthropicContentBlockDelta): + return _chunks_for_content_delta(event, state) + + if isinstance(event, AnthropicMessageDelta): + _handle_message_delta(event, state) + return [] + + if isinstance(event, AnthropicMessageStop): + state.saw_message_stop = True + + return [] + + +def _chunks_for_content_delta( + event: AnthropicContentBlockDelta, + state: _StreamState, +) -> list[ChatCompletionChunk]: + delta: AnthropicRawContentBlock = event.delta + 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)] + + is_json_delta = delta.type in {"input_json_delta", "input_json"} or ( + isinstance(delta.partial_json, str) and bool(delta.partial_json) + ) + if not is_json_delta: + return [] + + tool_state = state.ensure_tool(event.index) + fragment = delta.partial_json if isinstance(delta.partial_json, str) else "" + if not fragment: + return [] + return [ + tool_call_delta_chunk( + model=state.model, + index=tool_state.tool_index, + arguments=fragment, + ) + ] + + +async def _stream_as_chat_chunks( + client: AnthropicClient, + param: ChatCompletionsParam, +) -> AsyncIterator[ChatCompletionChunk]: + """Yield chat-completions chunks from an Anthropic Messages SSE stream. + + Text deltas stream as they arrive. Tool calls stream as OpenAI-style + tool-call deltas (id/name on start, argument JSON fragments on delta). + A finish_reason (and optional usage) is emitted on ``message_delta`` / + ``message_stop``. Errors raise ``RuntimeError`` so the agent loop can + surface them as retryable failures. + """ + create = chat_param_to_anthropic_param(param) + stream = await client.messages(create, stream=True) + state = _StreamState(param.model) + + async for event in stream: + for chunk in _chunks_for_event(event, state): + yield chunk + if state.saw_message_stop: + break + + if not state.saw_message_stop and state.stop_reason is None and not state.has_tool_calls: + # Stream ended without a terminal — yield nothing more so the loop can + # treat empty + no finish_reason as a missing terminal glitch. + return + + finish = anthropic_stop_to_finish_reason(state.stop_reason, has_tool_calls=state.has_tool_calls) + yield finish_reason_chunk(model=state.model, finish_reason=finish) + if state.last_usage is not None: + yield usage_chunk(model=state.model, usage=state.last_usage) + + +async def dispatch_messages( + client: AnthropicClient, + param: ChatCompletionsParam, + *, + stream: bool = False, +) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]: + """Run one Anthropic Messages turn and return a chat-completions-shaped result.""" + create = chat_param_to_anthropic_param(param) + if stream: + return _stream_as_chat_chunks(client, param) + response = await client.messages(create, stream=False) + return anthropic_response_to_chat_completion(response) + + +__all__ = ["dispatch_messages"] diff --git a/src/plyngent/config/models.py b/src/plyngent/config/models.py index 7d6c331..8a181f8 100644 --- a/src/plyngent/config/models.py +++ b/src/plyngent/config/models.py @@ -147,7 +147,7 @@ class ModelConfig(Struct, omit_defaults=True): ``preset`` / ``url`` override the parent provider for this model only. ``preset`` keeps the same meanings as provider presets: ``openai`` → /responses, ``openai-compatible`` → /chat/completions, - ``anthropic`` → /messages (native support planned), ``deepseek`` → DeepSeek. + ``anthropic`` → /messages, ``deepseek`` → DeepSeek. """ text: bool = True diff --git a/tests/test_agent/test_messages_bridge.py b/tests/test_agent/test_messages_bridge.py new file mode 100644 index 0000000..0b76c8f --- /dev/null +++ b/tests/test_agent/test_messages_bridge.py @@ -0,0 +1,200 @@ +"""Tests for Anthropic Messages chat ↔ native conversion.""" + +from __future__ import annotations + +from msgspec import UNSET + +from plyngent.agent.messages_bridge import ( + anthropic_response_to_assistant, + anthropic_response_to_chat_completion, + anthropic_stop_to_finish_reason, + anthropic_usage_to_dict, + chat_messages_to_anthropic, + chat_param_to_anthropic_param, + tool_items_to_anthropic_tools, +) +from plyngent.lmproto.anthropic.model import ( + AnthropicAssistantMessage, + AnthropicMessageResponse, + AnthropicResponseText, + AnthropicResponseToolUse, + AnthropicToolResultContent, + AnthropicToolUseContent, + AnthropicUsage, + AnthropicUserMessage, +) +from plyngent.lmproto.openai_compatible.model import ( + AssistantChatMessage, + AssistantFunctionTool, + AssistantFunctionToolCall, + ChatCompletionsParam, + DeveloperChatMessage, + SystemChatMessage, + ToolChatMessage, + ToolFunction, + ToolFunctionItem, + UserChatMessage, +) + + +def test_tool_items_to_anthropic_tools() -> None: + items = [ + ToolFunctionItem( + function=ToolFunction( + name="read_file", + description="Read a file", + parameters={"type": "object", "properties": {"path": {"type": "string"}}}, + ) + ) + ] + tools = tool_items_to_anthropic_tools(items) + assert len(tools) == 1 + assert tools[0].name == "read_file" + assert tools[0].input_schema is not UNSET + + +def test_chat_messages_system_and_tool_round() -> None: + messages = [ + SystemChatMessage(content="You are helpful."), + UserChatMessage(content="hi"), + AssistantChatMessage( + content=None, + tool_calls=[ + AssistantFunctionToolCall( + id="call_1", + function=AssistantFunctionTool(name="read_file", arguments='{"path":"a"}'), + ) + ], + ), + ToolChatMessage(content="file body", tool_call_id="call_1"), + ] + system, items = chat_messages_to_anthropic(messages) + assert system == "You are helpful." + assert len(items) == 3 + assert isinstance(items[0], AnthropicUserMessage) + assert items[0].content == "hi" + assert isinstance(items[1], AnthropicAssistantMessage) + assert isinstance(items[1].content, list) + assert any(isinstance(b, AnthropicToolUseContent) and b.id == "call_1" for b in items[1].content) + assert isinstance(items[2], AnthropicUserMessage) + assert isinstance(items[2].content, list) + assert isinstance(items[2].content[0], AnthropicToolResultContent) + assert items[2].content[0].tool_use_id == "call_1" + assert items[2].content[0].content == "file body" + + +def test_developer_becomes_user_prefix() -> None: + system, items = chat_messages_to_anthropic( + [ + SystemChatMessage(content="persona"), + UserChatMessage(content="go"), + DeveloperChatMessage(content="checkpoint"), + ] + ) + assert system == "persona" + assert len(items) == 2 + assert isinstance(items[1], AnthropicUserMessage) + assert items[1].content == "[developer] checkpoint" + + +def test_merge_consecutive_tool_results() -> None: + messages = [ + UserChatMessage(content="go"), + AssistantChatMessage( + content=None, + tool_calls=[ + AssistantFunctionToolCall( + id="c1", + function=AssistantFunctionTool(name="a", arguments="{}"), + ), + AssistantFunctionToolCall( + id="c2", + function=AssistantFunctionTool(name="b", arguments="{}"), + ), + ], + ), + ToolChatMessage(content="r1", tool_call_id="c1"), + ToolChatMessage(content="r2", tool_call_id="c2"), + ] + _, items = chat_messages_to_anthropic(messages) + assert len(items) == 3 # user, assistant, one merged tool-result user + assert isinstance(items[2], AnthropicUserMessage) + assert isinstance(items[2].content, list) + assert len(items[2].content) == 2 + + +def test_anthropic_response_to_assistant_with_tools() -> None: + response = AnthropicMessageResponse( + id="msg_1", + model="claude-test", + stop_reason="tool_use", + content=[ + AnthropicResponseText(text="calling"), + AnthropicResponseToolUse(id="tu_1", name="add", input={"a": 1, "b": 2}), + ], + usage=AnthropicUsage(input_tokens=12, output_tokens=4), + ) + assistant = anthropic_response_to_assistant(response) + assert assistant.content == "calling" + assert assistant.tool_calls is not UNSET + assert len(assistant.tool_calls) == 1 + call = assistant.tool_calls[0] + assert isinstance(call, AssistantFunctionToolCall) + assert call.id == "tu_1" + assert call.function.name == "add" + assert '"a":1' in call.function.arguments or '"a": 1' in call.function.arguments + + +def test_anthropic_response_to_chat_completion_usage() -> None: + response = AnthropicMessageResponse( + id="msg_1", + model="claude-test", + stop_reason="end_turn", + content=[AnthropicResponseText(text="done")], + usage=AnthropicUsage(input_tokens=9, output_tokens=3), + ) + chat = anthropic_response_to_chat_completion(response) + assert chat.choices[0].message.content == "done" + assert chat.choices[0].finish_reason == "stop" + assert chat.usage is not UNSET + assert chat.usage["prompt_tokens"] == 9 + assert chat.usage["completion_tokens"] == 3 + + +def test_stop_reason_mapping() -> None: + assert anthropic_stop_to_finish_reason("end_turn", has_tool_calls=False) == "stop" + assert anthropic_stop_to_finish_reason("tool_use", has_tool_calls=True) == "tool_calls" + assert anthropic_stop_to_finish_reason("max_tokens", has_tool_calls=False) == "length" + assert anthropic_stop_to_finish_reason("refusal", has_tool_calls=False) == "content_filter" + assert anthropic_stop_to_finish_reason(None, has_tool_calls=True) == "tool_calls" + + +def test_usage_to_dict() -> None: + assert anthropic_usage_to_dict(AnthropicUsage(input_tokens=1, output_tokens=2)) == { + "prompt_tokens": 1, + "completion_tokens": 2, + "total_tokens": 3, + "input_tokens": 1, + "output_tokens": 2, + } + assert anthropic_usage_to_dict(AnthropicUsage()) is None + + +def test_chat_param_to_anthropic_param() -> None: + param = ChatCompletionsParam( + model="claude-test", + messages=[ + SystemChatMessage(content="Be brief."), + UserChatMessage(content="hi"), + ], + tools=[ToolFunctionItem(function=ToolFunction(name="f", description="d", parameters={"type": "object"}))], + temperature=0.2, + max_completion_tokens=1024, + ) + body = chat_param_to_anthropic_param(param) + assert body.model == "claude-test" + assert body.max_tokens == 1024 + assert body.system == "Be brief." + assert body.tools is not UNSET + assert len(body.tools) == 1 + assert body.temperature == 0.2 diff --git a/tests/test_agent/test_messages_dispatch.py b/tests/test_agent/test_messages_dispatch.py new file mode 100644 index 0000000..c121f09 --- /dev/null +++ b/tests/test_agent/test_messages_dispatch.py @@ -0,0 +1,300 @@ +"""Tests for kind-based Anthropic Messages dispatch in the agent loop.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal, cast, overload + +from msgspec import UNSET + +from plyngent.agent import ( + AssistantMessageEvent, + TextDeltaEvent, + ToolCallEvent, + ToolRegistry, + ToolResultEvent, + UsageEvent, + run_chat_loop, + tool, +) +from plyngent.agent.messages_dispatch import dispatch_messages +from plyngent.lmproto.anthropic.model import ( + AnthropicContentBlockDelta, + AnthropicContentBlockStart, + AnthropicMessageDelta, + AnthropicMessageResponse, + AnthropicMessagesParam, + AnthropicMessageStart, + AnthropicMessageStop, + AnthropicRawContentBlock, + AnthropicResponseText, + AnthropicResponseToolUse, + AnthropicStreamEvent, + AnthropicToolResultContent, + AnthropicUsage, + AnthropicUserMessage, +) +from plyngent.lmproto.openai_compatible.model import ( + AssistantChatMessage, + ChatCompletionsParam, + UserChatMessage, +) + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Sequence + + +def _message_response( + *, + text: str = "hello", + tool_uses: list[tuple[str, str, dict[str, Any]]] | None = None, + stop_reason: str = "end_turn", + usage: AnthropicUsage | None = None, +) -> AnthropicMessageResponse: + content: list[Any] = [] + if text: + content.append(AnthropicResponseText(text=text)) + if tool_uses: + for call_id, name, inp in tool_uses: + content.append(AnthropicResponseToolUse(id=call_id, name=name, input=inp)) + stop_reason = "tool_use" + return AnthropicMessageResponse( + id="msg_test", + model="claude-test", + content=content, + stop_reason=stop_reason, + usage=usage or AnthropicUsage(input_tokens=10, output_tokens=5), + ) + + +class ScriptedMessagesClient: + """Fake AnthropicClient that scripts ``messages()`` results.""" + + kind: str = "messages" + calls: list[AnthropicMessagesParam] + _non_stream: list[AnthropicMessageResponse] + _stream_events: list[list[AnthropicStreamEvent]] + + def __init__( + self, + *, + non_stream: Sequence[AnthropicMessageResponse] | None = None, + stream_events: Sequence[Sequence[AnthropicStreamEvent]] | None = None, + ) -> None: + self.calls = [] + self._non_stream = list(non_stream or ()) + self._stream_events = [list(events) for events in (stream_events or ())] + + @overload + async def messages( + self, param: AnthropicMessagesParam, *, stream: Literal[False] = False + ) -> AnthropicMessageResponse: ... + + @overload + async def messages( + self, param: AnthropicMessagesParam, *, stream: Literal[True] + ) -> AsyncIterator[AnthropicStreamEvent]: ... + + async def messages( + self, param: AnthropicMessagesParam, *, stream: bool = False + ) -> AnthropicMessageResponse | AsyncIterator[AnthropicStreamEvent]: + self.calls.append(param) + if stream: + 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: + yield event + + return _gen() + if not self._non_stream: + msg = "no more scripted messages" + raise RuntimeError(msg) + return self._non_stream.pop(0) + + +async def test_dispatch_messages_non_stream_text() -> None: + client = ScriptedMessagesClient(non_stream=[_message_response(text="hi there")]) + param = ChatCompletionsParam( + model="claude-test", + messages=[UserChatMessage(content="hello")], + ) + result = await dispatch_messages(cast("Any", client), param, stream=False) + assert result.choices[0].message.content == "hi there" + assert result.choices[0].finish_reason == "stop" + assert len(client.calls) == 1 + assert client.calls[0].max_tokens == 8192 + + +async def test_dispatch_messages_stream_text_and_usage() -> None: + events: list[AnthropicStreamEvent] = [ + AnthropicMessageStart( + message=AnthropicMessageResponse( + id="msg_s", + model="claude-test", + content=[], + usage=AnthropicUsage(input_tokens=8, output_tokens=0), + ) + ), + AnthropicContentBlockStart( + index=0, + content_block=AnthropicRawContentBlock(type="text", text=""), + ), + AnthropicContentBlockDelta( + index=0, + delta=AnthropicRawContentBlock(type="text_delta", text="stream"), + ), + AnthropicContentBlockDelta( + index=0, + delta=AnthropicRawContentBlock(type="text_delta", text="ed"), + ), + AnthropicMessageDelta( + delta={"stop_reason": "end_turn"}, + usage=AnthropicUsage(input_tokens=0, output_tokens=4), + ), + AnthropicMessageStop(), + ] + client = ScriptedMessagesClient(stream_events=[events]) + param = ChatCompletionsParam( + model="claude-test", + messages=[UserChatMessage(content="hi")], + ) + stream = await dispatch_messages(cast("Any", client), 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 == ["stream", "ed"] + assert any(c.choices and c.choices[0].finish_reason == "stop" for c in chunks) + assert any(getattr(c, "usage", UNSET) not in (UNSET, None) for c in chunks) + + +async def test_run_chat_loop_messages_kind_text() -> None: + client = ScriptedMessagesClient(non_stream=[_message_response(text="done")]) + messages: list[Any] = [UserChatMessage(content="go")] + events = [ + event + async for event in run_chat_loop( + cast("Any", client), + messages, + model="claude-test", + stream=False, + ) + ] + assert any(isinstance(e, TextDeltaEvent) and e.content == "done" for e in events) + assert any(isinstance(e, AssistantMessageEvent) for e in events) + assert any(isinstance(e, UsageEvent) and e.usage.prompt_tokens == 10 for e in events) + assert isinstance(messages[-1], AssistantChatMessage) + assert messages[-1].content == "done" + + +async def test_run_chat_loop_messages_kind_tool_round() -> None: + first = _message_response( + text="", + tool_uses=[("call_1", "add", {"a": 1, "b": 2})], + ) + second = _message_response(text="sum is 3") + client = ScriptedMessagesClient(non_stream=[first, second]) + + @tool(register=False) + def add(a: int, b: int) -> str: + return str(a + b) + + registry = ToolRegistry([add]) + messages: list[Any] = [UserChatMessage(content="add 1 and 2")] + events = [ + event + async for event in run_chat_loop( + cast("Any", client), + messages, + model="claude-test", + tools=registry, + stream=False, + ) + ] + assert any(isinstance(e, ToolCallEvent) for e in events) + assert any(isinstance(e, ToolResultEvent) and "3" in e.message.content for e in events) + assert any(isinstance(e, AssistantMessageEvent) and e.message.content == "sum is 3" for e in events) + assert len(client.calls) == 2 + # Second turn must include tool_result in a user message. + second_msgs = client.calls[1].messages + assert any( + isinstance(m, AnthropicUserMessage) + and isinstance(m.content, list) + and any(isinstance(b, AnthropicToolResultContent) and b.tool_use_id == "call_1" for b in m.content) + for m in second_msgs + ) + + +async def test_stream_tool_call_deltas() -> None: + events: list[AnthropicStreamEvent] = [ + AnthropicMessageStart( + message=AnthropicMessageResponse( + id="msg_s", + model="claude-test", + content=[], + usage=AnthropicUsage(input_tokens=5, output_tokens=0), + ) + ), + AnthropicContentBlockStart( + index=0, + content_block=AnthropicRawContentBlock(type="tool_use", id="tu_9", name="add"), + ), + AnthropicContentBlockDelta( + index=0, + delta=AnthropicRawContentBlock(type="input_json_delta", partial_json='{"a":'), + ), + AnthropicContentBlockDelta( + index=0, + delta=AnthropicRawContentBlock(type="input_json_delta", partial_json="1}"), + ), + AnthropicMessageDelta( + delta={"stop_reason": "tool_use"}, + usage=AnthropicUsage(input_tokens=0, output_tokens=6), + ), + AnthropicMessageStop(), + ] + client = ScriptedMessagesClient(stream_events=[events]) + + @tool(register=False) + def add(a: int) -> str: + return str(a) + + registry = ToolRegistry([add]) + # Second streamed round after tool execution. + client._stream_events.append( + [ + AnthropicMessageStart( + message=AnthropicMessageResponse( + id="msg_2", + model="claude-test", + content=[], + usage=AnthropicUsage(input_tokens=12, output_tokens=0), + ) + ), + AnthropicContentBlockDelta( + index=0, + delta=AnthropicRawContentBlock(type="text_delta", text="ok"), + ), + AnthropicMessageDelta( + delta={"stop_reason": "end_turn"}, + usage=AnthropicUsage(input_tokens=0, output_tokens=1), + ), + AnthropicMessageStop(), + ] + ) + + messages: list[Any] = [UserChatMessage(content="go")] + events_out = [ + event + async for event in run_chat_loop( + cast("Any", client), + messages, + model="claude-test", + tools=registry, + stream=True, + ) + ] + assert any(isinstance(e, ToolCallEvent) for e in events_out) + assert any(isinstance(e, ToolResultEvent) for e in events_out) + assert any(isinstance(e, AssistantMessageEvent) and e.message.content == "ok" for e in events_out)