diff --git a/src/agentscope/model/_openai_response/_models/gpt-5.4.yaml b/src/agentscope/model/_openai_response/_models/gpt-5.4.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5a110a459fa46ea7ef730e00d276817c66dae367 --- /dev/null +++ b/src/agentscope/model/_openai_response/_models/gpt-5.4.yaml @@ -0,0 +1,22 @@ +name: gpt-5.4 +label: GPT-5.4 (Responses API) +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + +context_size: 1050000 +output_size: 128000 + +parameter_overrides: + max_tokens: + maximum: 128000 + reasoning_effort: + hidden: true diff --git a/src/agentscope/model/_openai_response/_models/gpt-5.5.yaml b/src/agentscope/model/_openai_response/_models/gpt-5.5.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fb4a9cf996191a50ad75e8f08d67d4649b50c832 --- /dev/null +++ b/src/agentscope/model/_openai_response/_models/gpt-5.5.yaml @@ -0,0 +1,26 @@ +name: gpt-5.5 +label: GPT-5.5 +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + - audio/mp3 + - audio/wav + +output_types: + - text/plain + +context_size: 1050000 +output_size: 128000 + +parameter_overrides: + max_tokens: + maximum: 128000 + thinking_enable: + hidden: true + reasoning_effort: + hidden: true diff --git a/src/agentscope/model/_openai_response/_models/o3.yaml b/src/agentscope/model/_openai_response/_models/o3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2f772a465e67b6c2b7ad1f62d72998c57b290e75 --- /dev/null +++ b/src/agentscope/model/_openai_response/_models/o3.yaml @@ -0,0 +1,24 @@ +name: o3 +label: o3 (Responses API) +status: active + +input_types: + - text/plain + - application/x-thinking + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + - application/x-thinking + +context_size: 200000 +output_size: 100000 + +parameter_overrides: + max_tokens: + maximum: 100000 + temperature: + hidden: true diff --git a/src/agentscope/model/_openai_response/_models/o4-mini.yaml b/src/agentscope/model/_openai_response/_models/o4-mini.yaml new file mode 100644 index 0000000000000000000000000000000000000000..52cda3f4e4ab5cecf4c8416c4262bd3d60bfe7a5 --- /dev/null +++ b/src/agentscope/model/_openai_response/_models/o4-mini.yaml @@ -0,0 +1,24 @@ +name: o4-mini +label: o4-mini (Responses API) +status: active + +input_types: + - text/plain + - application/x-thinking + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + - application/x-thinking + +context_size: 200000 +output_size: 100000 + +parameter_overrides: + max_tokens: + maximum: 100000 + temperature: + hidden: true diff --git a/src/agentscope/model/_xai/__init__.py b/src/agentscope/model/_xai/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a13f85a3a23f08b4a494a0281d74838a5e1387b4 --- /dev/null +++ b/src/agentscope/model/_xai/__init__.py @@ -0,0 +1,9 @@ +# -*- coding: utf-8 -*- +"""The xAI LLM API modules.""" + +from ._model import XAICredential, XAIChatModel + +__all__ = [ + "XAICredential", + "XAIChatModel", +] diff --git a/src/agentscope/model/_xai/_model.py b/src/agentscope/model/_xai/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..9ad0f6b7cd9d409fbc9a66fda6e71768561e16de --- /dev/null +++ b/src/agentscope/model/_xai/_model.py @@ -0,0 +1,455 @@ +# -*- coding: utf-8 -*- +"""The xAI chat model implementation using the official xai_sdk.""" +from datetime import datetime +from typing import Any, AsyncGenerator, List, Literal, TYPE_CHECKING, Type + +from pydantic import BaseModel, Field + +from .._base import ChatModelBase, _TOOL_CHOICE_LITERAL_MODES +from .._model_response import ChatResponse +from .._model_usage import ChatUsage +from ...credential import XAICredential +from ...formatter import XAIChatFormatter +from ...message import ( + Msg, + TextBlock, + ThinkingBlock, + ToolCallBlock, +) +from ...tool import ToolChoice + +if TYPE_CHECKING: + from xai_sdk import AsyncClient + from xai_sdk.chat import Response +else: + AsyncClient = Any + Response = Any + + +class XAIChatModel(ChatModelBase): + """The xAI chat model using the official ``xai_sdk`` gRPC client. + + This model provides native access to xAI-specific features such as + server-side agentic tools (web search, X search, code execution) and + reasoning effort control, which are not available through the + OpenAI-compatible REST endpoint. + """ + + class Parameters(BaseModel): + """The parameters for the xAI chat model.""" + + max_tokens: int | None = Field( + default=None, + title="Max Tokens", + description="The maximum number of tokens for the LLM output.", + gt=0, + ) + """The maximum number of tokens to generate.""" + + thinking_enable: bool = Field( + default=False, + title="Thinking", + description=( + "Whether to enable reasoning for models that support " + "extended thinking (e.g. ``grok-3-mini``). Use " + "reasoning_effort to control the depth of reasoning." + ), + ) + + reasoning_effort: Literal["low", "medium", "high"] | None = Field( + default=None, + title="Reasoning Effort", + description=( + "Controls the depth of reasoning for models that support " + "extended thinking (e.g. ``grok-3-mini``). Set to " + "``'low'``, ``'medium'``, or ``'high'`` to enable reasoning " + "with the corresponding effort level. ``None`` disables " + "reasoning." + ), + ) + + temperature: float | None = Field( + default=None, + title="Temperature", + description="The temperature for the LLM output.", + ge=0, + le=2, + ) + """The sampling temperature.""" + + top_p: float | None = Field( + default=None, + title="Top P", + description="The top-p nucleus sampling value.", + gt=0, + le=1, + ) + """The top-p sampling parameter.""" + + type: Literal["xai_chat"] = "xai_chat" + """The type of the chat model.""" + + def __init__( + self, + credential: XAICredential, + model: str, + parameters: "XAIChatModel.Parameters | None" = None, + stream: bool = True, + max_retries: int = 3, + retry_delay: float = 1.0, + context_size: int = 131072, + formatter: XAIChatFormatter | None = None, + client_kwargs: dict[str, Any] | None = None, + ) -> None: + """Initialize the xAI chat model. + + Args: + credential (`XAICredential`): + The xAI credential used to authenticate API calls. + model (`str`): + The xAI model name, e.g. ``grok-3`` or ``grok-3-mini``. + parameters (`XAIChatModel.Parameters | None`, defaults to \ + `None`): + The xAI API parameters. When ``None``, the default + parameters will be used. + stream (`bool`, defaults to `True`): + Whether to enable streaming output. + max_retries (`int`, defaults to `3`): + The maximum number of retries for the xAI API. + retry_delay (`float`, defaults to `1.0`): + Seconds to sleep between retry attempts. + context_size (`int`, defaults to `131072`): + The model context size used for context compression. + formatter (`XAIChatFormatter | None`, defaults to `None`): + The formatter that converts ``Msg`` objects to xai_sdk + proto messages. When ``None``, an ``XAIChatFormatter`` + instance will be used. + client_kwargs (`dict[str, Any] | None`, defaults to `None`): + Extra keyword arguments forwarded to ``xai_sdk.AsyncClient`` + (e.g. ``timeout``, ``metadata``, ``channel_options``). Keys + that overlap with credential-derived arguments (such as + ``api_key`` or ``api_host``) take precedence over the + credential values. + """ + super().__init__( + credential=credential, + model=model, + parameters=parameters or self.Parameters(), + stream=stream, + max_retries=max_retries, + retry_delay=retry_delay, + context_size=context_size, + ) + self.formatter = formatter or XAIChatFormatter() + self.client_kwargs = client_kwargs or {} + + @classmethod + def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]: + import grpc + + # xai_sdk uses grpc.aio under the hood; transport/API failures surface + # as grpc.aio.AioRpcError (subclass of grpc.RpcError). We retry the + # whole class because the retry mechanism here only filters by type, + # not by status code — so 4xx-equivalents (UNAUTHENTICATED, INVALID + # _ARGUMENT) will also be retried a few times before failing, which + # we accept. Note xai_sdk additionally retries UNAVAILABLE 5 times at + # the gRPC layer with exponential backoff before raising to us. + return (grpc.RpcError,) + + async def _call_api( + self, + model_name: str, + messages: list[Msg], + tools: list[dict] | None = None, + tool_choice: ToolChoice | None = None, + **generate_kwargs: Any, + ) -> ChatResponse | AsyncGenerator[ChatResponse, None]: + """Call the xAI API using the official ``xai_sdk`` gRPC client. + + Args: + model_name (`str`): + The model name to use for this call. + messages (`list`): + A list of ``Msg`` objects representing the conversation. + tools (`list[dict]`, default `None`): + The tools JSON schemas. + tool_choice (`ToolChoice | None`, optional): + Controls which (if any) tool is called by the model. + **generate_kwargs (`Any`): + Extra keyword arguments forwarded to the API. + + Returns: + `ChatResponse | AsyncGenerator[ChatResponse, None]`: + A ``ChatResponse`` when streaming is disabled, or an async + generator of ``ChatResponse`` objects when streaming is + enabled. + """ + from xai_sdk import AsyncClient + + client = AsyncClient( + **{ + "api_key": self.credential.api_key.get_secret_value(), + "api_host": self.credential.api_host, + **self.client_kwargs, + }, + ) + + xai_messages = await self.formatter.format(messages) + + xai_tools, xai_tool_choice = self._format_tools(tools, tool_choice) + + create_kwargs: dict[str, Any] = {"model": model_name} + if self.parameters.max_tokens is not None: + create_kwargs["max_tokens"] = self.parameters.max_tokens + if self.parameters.temperature is not None: + create_kwargs["temperature"] = self.parameters.temperature + if self.parameters.top_p is not None: + create_kwargs["top_p"] = self.parameters.top_p + if ( + self.parameters.thinking_enable + and self.parameters.reasoning_effort + ): + create_kwargs[ + "reasoning_effort" + ] = self.parameters.reasoning_effort + if xai_tools: + create_kwargs["tools"] = xai_tools + if xai_tool_choice is not None: + create_kwargs["tool_choice"] = xai_tool_choice + + create_kwargs.update(generate_kwargs) + + chat = client.chat.create(**create_kwargs) + for xai_msg in xai_messages: + chat.append(xai_msg) + + start_datetime = datetime.now() + + if self.stream: + return self._parse_stream_response(start_datetime, chat, client) + + try: + response = await chat.sample() + finally: + await client.close() + + return self._parse_completion_response(start_datetime, response) + + def _format_tools( + self, + tools: list[dict] | None, + tool_choice: ToolChoice | None, + ) -> tuple[list | None, Any]: + """Validate, filter, and format tools and tool_choice for the xAI API. + + When ``tool_choice.tools`` is specified the schemas list is filtered + to only those tools. When ``tool_choice.mode`` is a specific tool name + (str) the model is forced to call exactly that tool without needing to + filter the list, preserving prompt-cache efficiency. + + Args: + tools (`list[dict] | None`, optional): + The raw tool schemas. + tool_choice (`ToolChoice | None`, optional): + The tool choice configuration. + + Returns: + `tuple[list | None, Any]`: + A tuple of (xai_tools, xai_tool_choice) ready for the + ``xai_sdk`` client. + """ + from xai_sdk.chat import required_tool, tool + + if tool_choice and tools: + self._validate_tool_choice(tool_choice, tools) + if tool_choice.tools: + allowed = set(tool_choice.tools) + tools = [t for t in tools if t["function"]["name"] in allowed] + + xai_tools = None + if tools: + xai_tools = [] + for t in tools: + if t.get("type") == "function" and "function" in t: + fn = t["function"] + xai_tools.append( + tool( + name=fn["name"], + description=fn.get("description", ""), + parameters=fn.get("parameters", {}), + ), + ) + + if not tool_choice: + return xai_tools, None + + mode = tool_choice.mode + + if mode in _TOOL_CHOICE_LITERAL_MODES: + return xai_tools, mode + + return xai_tools, required_tool(mode) + + async def _parse_stream_response( + self, + start_datetime: datetime, + chat: Any, + client: AsyncClient, + ) -> AsyncGenerator[ChatResponse, None]: + """Parse the xAI streaming response from ``xai_sdk``. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + chat (`Any`): + The ``xai_sdk`` chat session object. + client (`Any`): + The ``xai_sdk.AsyncClient`` instance; closed when the + generator is exhausted or abandoned. + + Yields: + `ChatResponse`: + Incremental ``ChatResponse`` objects with ``is_last=False`` + followed by a final one with ``is_last=True``. + """ + acc_text = TextBlock(text="") + acc_thinking = ThinkingBlock(thinking="") + last_response = None + response_id: str | None = None + + try: + async for response, chunk in chat.stream(): + if response_id is None: + response_id = getattr(response, "id", None) or None + + delta_text: str = chunk.content or "" + delta_thinking: str = chunk.reasoning_content or "" + + delta_contents: List[TextBlock | ThinkingBlock] = [] + + if delta_thinking: + acc_thinking.thinking += delta_thinking + delta_contents.append( + ThinkingBlock( + id=acc_thinking.id, + thinking=delta_thinking, + ), + ) + if delta_text: + acc_text.text += delta_text + delta_contents.append( + TextBlock(id=acc_text.id, text=delta_text), + ) + + if delta_contents: + _kwargs: dict[str, Any] = { + "content": delta_contents, + "is_last": False, + } + if response_id: + _kwargs["id"] = response_id + yield ChatResponse(**_kwargs) + + last_response = response + + finally: + await client.close() + + final_contents: List[TextBlock | ToolCallBlock | ThinkingBlock] = [] + if acc_thinking.thinking: + final_contents.append(acc_thinking) + if acc_text.text: + final_contents.append(acc_text) + + if last_response is not None: + for tc in last_response.tool_calls or []: + final_contents.append( + ToolCallBlock( + id=tc.id, + name=tc.function.name, + input=tc.function.arguments, + ), + ) + + usage = None + if last_response is not None and last_response.usage is not None: + u = last_response.usage + usage = ChatUsage( + input_tokens=u.prompt_tokens, + output_tokens=u.completion_tokens, + time=(datetime.now() - start_datetime).total_seconds(), + cache_input_tokens=getattr( + u, + "cached_prompt_text_tokens", + 0, + ), + ) + + final_kwargs: dict[str, Any] = { + "content": final_contents, + "usage": usage, + "is_last": True, + } + if response_id: + final_kwargs["id"] = response_id + yield ChatResponse(**final_kwargs) + + def _parse_completion_response( + self, + start_datetime: datetime, + response: Response, + ) -> ChatResponse: + """Parse the xAI non-streaming response from ``xai_sdk``. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`Any`): + The ``xai_sdk`` ``Response`` object. + + Returns: + `ChatResponse`: + A single ``ChatResponse`` with ``is_last=True``. + """ + content_blocks: List[TextBlock | ToolCallBlock | ThinkingBlock] = [] + + if response.reasoning_content: + content_blocks.append( + ThinkingBlock(thinking=response.reasoning_content), + ) + if response.content: + content_blocks.append(TextBlock(text=response.content)) + + for tc in response.tool_calls or []: + content_blocks.append( + ToolCallBlock( + id=tc.id, + name=tc.function.name, + input=tc.function.arguments, + ), + ) + + usage = None + if response.usage is not None: + u = response.usage + usage = ChatUsage( + input_tokens=u.prompt_tokens, + output_tokens=u.completion_tokens, + time=(datetime.now() - start_datetime).total_seconds(), + cache_input_tokens=getattr( + u, + "cached_prompt_text_tokens", + 0, + ), + ) + + resp_kwargs: dict[str, Any] = { + "content": content_blocks, + "is_last": True, + "usage": usage, + } + response_id = getattr(response, "id", None) + if response_id: + resp_kwargs["id"] = response_id + + return ChatResponse(**resp_kwargs) diff --git a/src/agentscope/model/_xai/_models/grok-3-fast.yaml b/src/agentscope/model/_xai/_models/grok-3-fast.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6447a0899eab30a72e371a56094eb62ae405fd6a --- /dev/null +++ b/src/agentscope/model/_xai/_models/grok-3-fast.yaml @@ -0,0 +1,19 @@ +name: grok-3-fast +label: Grok 3 Fast +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + +output_types: + - text/plain + - application/x-thinking + +context_size: 131072 +output_size: 16000 + +parameter_overrides: + max_tokens: + maximum: 16000 diff --git a/src/agentscope/model/_xai/_models/grok-3-mini.yaml b/src/agentscope/model/_xai/_models/grok-3-mini.yaml new file mode 100644 index 0000000000000000000000000000000000000000..eaddca82de0f831c05670eef56a08ef2b915538f --- /dev/null +++ b/src/agentscope/model/_xai/_models/grok-3-mini.yaml @@ -0,0 +1,17 @@ +name: grok-3-mini +label: Grok 3 Mini +status: active + +input_types: + - text/plain + +output_types: + - text/plain + - application/x-thinking + +context_size: 131072 +output_size: 16000 + +parameter_overrides: + max_tokens: + maximum: 16000 diff --git a/src/agentscope/model/_xai/_models/grok-3.yaml b/src/agentscope/model/_xai/_models/grok-3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0dd32acd8b418f1da1fd4438c68b05400fa9d29a --- /dev/null +++ b/src/agentscope/model/_xai/_models/grok-3.yaml @@ -0,0 +1,19 @@ +name: grok-3 +label: Grok 3 +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + +output_types: + - text/plain + - application/x-thinking + +context_size: 131072 +output_size: 16000 + +parameter_overrides: + max_tokens: + maximum: 16000 diff --git a/src/agentscope/model/_xai/_models/grok-4.3.yaml b/src/agentscope/model/_xai/_models/grok-4.3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..49c3b5c7fb799a186bbe53f80b704f98d66587a0 --- /dev/null +++ b/src/agentscope/model/_xai/_models/grok-4.3.yaml @@ -0,0 +1,19 @@ +name: grok-4.3 +label: Grok 4.3 +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + +output_types: + - text/plain + - application/x-thinking + +context_size: 1000000 +output_size: 30000 + +parameter_overrides: + max_tokens: + maximum: 30000 diff --git a/src/agentscope/permission/__init__.py b/src/agentscope/permission/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d42dd912969cbb26768e7669cc88a576de466cd9 --- /dev/null +++ b/src/agentscope/permission/__init__.py @@ -0,0 +1,18 @@ +# -*- coding: utf-8 -*- +"""The tool permission related types and functions.""" + +from ._context import PermissionContext, AdditionalWorkingDirectory +from ._decision import PermissionDecision +from ._engine import PermissionEngine +from ._rule import PermissionRule +from ._types import PermissionMode, PermissionBehavior + +__all__ = [ + "PermissionContext", + "AdditionalWorkingDirectory", + "PermissionDecision", + "PermissionEngine", + "PermissionRule", + "PermissionMode", + "PermissionBehavior", +] diff --git a/src/agentscope/permission/_context.py b/src/agentscope/permission/_context.py new file mode 100644 index 0000000000000000000000000000000000000000..f5c07513b4f9a462279d32d681cd27754b27528e --- /dev/null +++ b/src/agentscope/permission/_context.py @@ -0,0 +1,46 @@ +# -*- coding: utf-8 -*- +"""The permission context module.""" +from pydantic import BaseModel, Field + +from ._rule import PermissionRule +from ._types import PermissionMode + + +class AdditionalWorkingDirectory(BaseModel): + """An additional directory included in permission scope. + + Working directories are used to determine which file paths should be + automatically allowed in ACCEPT_EDITS mode. + """ + + path: str + """Absolute path to the directory.""" + + source: str + """Where this directory permission originated from + (e.g., 'userSettings', 'session').""" + + +class PermissionContext(BaseModel): + """Context for permission checking. + + Contains the permission mode, working directories, and all configured + permission rules organized by behavior type (allow, deny, ask). + """ + + mode: PermissionMode = PermissionMode.DEFAULT + """The current permission mode.""" + + working_directories: dict[str, AdditionalWorkingDirectory] = Field( + default_factory=dict, + ) + """Additional directories allowed for file operations, keyed by path.""" + + allow_rules: dict[str, list[PermissionRule]] = Field(default_factory=dict) + """Rules that allow tool execution, keyed by tool name.""" + + deny_rules: dict[str, list[PermissionRule]] = Field(default_factory=dict) + """Rules that deny tool execution, keyed by tool name.""" + + ask_rules: dict[str, list[PermissionRule]] = Field(default_factory=dict) + """Rules that require user confirmation, keyed by tool name.""" diff --git a/src/agentscope/permission/_decision.py b/src/agentscope/permission/_decision.py new file mode 100644 index 0000000000000000000000000000000000000000..3f5a66bdb24ba56a8b11c9eb6dc7e9459648eba1 --- /dev/null +++ b/src/agentscope/permission/_decision.py @@ -0,0 +1,68 @@ +# -*- coding: utf-8 -*- +"""The permission decision result.""" +from dataclasses import dataclass +from typing import Any + +from ._rule import PermissionRule +from ._types import PermissionBehavior + + +@dataclass +class PermissionDecision: + """Decision result from permission checking. + + Represents the outcome of a permission check, including whether + the action should be allowed, denied, or require user confirmation. + """ + + behavior: PermissionBehavior + """The permission behavior decision.""" + + message: str + """Human-readable message describing the decision.""" + + decision_reason: str | None = None + """Optional explanation for why this decision was made.""" + + updated_input: dict[str, Any] | None = None + """Optional modified input data (e.g., sanitized paths).""" + + suggested_rules: list[PermissionRule] | None = None + """Optional list of suggested permission rules for user to apply.""" + + bypass_immune: bool = False + """Whether this decision is immune to being silenced by allow rules + ("bypass-immune"). + + Only meaningful when :attr:`behavior` is :attr:`PermissionBehavior.ASK`. + A tool sets this to ``True`` to signal that the operation is + dangerous enough that **no allow rule** may convert the ASK into an + ALLOW — the user must explicitly confirm in-the-moment. In + :attr:`PermissionMode.DONT_ASK` where no user is available, a + bypass-immune ASK is converted to DENY rather than silently allowed. + + Per-mode handling of a ``bypass_immune=True`` ASK: + + - ``DEFAULT`` / ``ACCEPT_EDITS``: honored — allow rules cannot + override. + - ``EXPLORE``: not applicable (the engine resolves EXPLORE via + :meth:`ToolBase.check_read_only` and does not invoke + :meth:`ToolBase.check_permissions`). + - ``BYPASS``: **intentionally ignored** — BYPASS's contract is + "the user has opted out of safety prompts; only deny / ask + rules remain as guardrails." Use deny rules in BYPASS to + enforce specific protections. + - ``DONT_ASK``: converted to DENY (no user available). + + Default is ``False``: a regular ASK that may be overridden by an + allow rule in DEFAULT / ACCEPT_EDITS, and is silently allowed by + BYPASS's fallback. Tools should set this only for genuine safety + checks (e.g. writes to dangerous paths, ``rm -rf /``, command + injection patterns) — not for "I'd prefer user input" cases. + + Note: this field is internal metadata for the permission engine. + Callers handling the decision (agent loop, HITL backend, UI) treat + a bypass-immune ASK the same as a regular ASK — both prompt the + user. The distinction only governs whether engine-level rules / + modes may override it before reaching the caller. + """ diff --git a/src/agentscope/permission/_engine.py b/src/agentscope/permission/_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..15aeb6f39dc29cfc4af89e16bfac5364954e6f26 --- /dev/null +++ b/src/agentscope/permission/_engine.py @@ -0,0 +1,729 @@ +# -*- coding: utf-8 -*- +"""The permission engine for checking and enforcing permission rules.""" +from typing import Any, List, TYPE_CHECKING + +from ._context import PermissionContext +from ._rule import PermissionRule +from ._decision import PermissionDecision, PermissionBehavior +from ._types import PermissionMode +from .._utils._common import _execute_async_or_sync_func + +if TYPE_CHECKING: + from ..tool import ToolBase +else: + ToolBase = "ToolBase" + + +class PermissionEngine: + """Engine for checking and enforcing permission rules. + + Evaluates tool execution requests against configured permission rules. + Matching strategy is delegated to each tool's :meth:`ToolBase.match_rule`: + + - Bash tools: substring / prefix wildcard matching against the command + - Write/Read/Edit tools: glob matching against file paths + - Other tools: generic pattern matching (or tool-name-level only) + + Each :class:`PermissionMode` has its own ``_check_`` method so + that mode policies are self-contained and readable in isolation. See + :meth:`check_permission` for the dispatcher and the individual methods + for per-mode evaluation order. + """ + + def __init__( + self, + context: PermissionContext, + ) -> None: + """Initialize the permission engine. + + Args: + context (`PermissionContext`): + The permission context containing rules and mode + + Example: + >>> context = PermissionContext(mode=PermissionMode.ACCEPT_EDITS) + >>> engine = PermissionEngine(context) + """ + self.context = context + + def add_rule(self, rule: PermissionRule) -> None: + """Add a permission rule to the context. + + Args: + rule (`PermissionRule`): + The permission rule to add + + Example: + >>> engine.add_rule(PermissionRule( + ... tool_name="Bash", + ... rule_content="git:*", + ... behavior=PermissionBehavior.ALLOW, + ... )) + """ + + if rule.behavior == PermissionBehavior.ALLOW: + if rule.tool_name not in self.context.allow_rules: + self.context.allow_rules[rule.tool_name] = [] + self.context.allow_rules[rule.tool_name].append(rule) + elif rule.behavior == PermissionBehavior.DENY: + if rule.tool_name not in self.context.deny_rules: + self.context.deny_rules[rule.tool_name] = [] + self.context.deny_rules[rule.tool_name].append(rule) + elif rule.behavior == PermissionBehavior.ASK: + if rule.tool_name not in self.context.ask_rules: + self.context.ask_rules[rule.tool_name] = [] + self.context.ask_rules[rule.tool_name].append(rule) + + async def check_permission( + self, + tool: ToolBase, + tool_input: dict[str, Any], + ) -> PermissionDecision: + """Check permission for a tool execution request. + + Dispatches to a per-mode private method so each mode's policy + is self-contained and readable in isolation: + + - DEFAULT → :meth:`_check_default` + - EXPLORE → :meth:`_check_explore` + - ACCEPT_EDITS → :meth:`_check_accept_edits` + - BYPASS → :meth:`_check_bypass` + - DONT_ASK → :meth:`_check_dont_ask` + + Args: + tool (`ToolBase`): + The tool instance being called. + tool_input (`dict[str, Any]`): + The tool input data, used for rule matching and + tool-specific checks. + + Returns: + `PermissionDecision`: + Decision indicating whether to allow, deny, or ask. + """ + mode = self.context.mode + if mode == PermissionMode.DEFAULT: + return await self._check_default(tool, tool_input) + if mode == PermissionMode.EXPLORE: + return await self._check_explore(tool, tool_input) + if mode == PermissionMode.ACCEPT_EDITS: + return await self._check_accept_edits(tool, tool_input) + if mode == PermissionMode.BYPASS: + return await self._check_bypass(tool, tool_input) + if mode == PermissionMode.DONT_ASK: + return await self._check_dont_ask(tool, tool_input) + raise ValueError(f"Unknown permission mode: {mode}") + + async def _check_default( + self, + tool: ToolBase, + tool_input: dict[str, Any], + ) -> PermissionDecision: + """Permission check for :attr:`PermissionMode.DEFAULT`. + + Every operation requires explicit permission unless either an + allow rule matches or the tool's own ``check_permissions`` + explicitly returns ALLOW (e.g. ``Bash`` auto-allows recognized + read-only commands like ``ls``/``git status``). Evaluation order: + + 1. Deny rules → DENY + 2. Ask rules → ASK (with suggestions) + 3. ``tool.check_permissions``: + - ALLOW / DENY → returned as-is + - Safety ASK (bypass-immune) → returned with suggestions; cannot + be overridden by allow rules + - Non-safety ASK / PASSTHROUGH → continue + 4. Allow rules → ALLOW + 5. Default → ASK (with suggestions) + + Args: + tool (`ToolBase`): + The tool instance being called. + tool_input (`dict[str, Any]`): + The tool input data. + + Returns: + `PermissionDecision`: + The final decision. + """ + # step 1: deny rules — highest priority + deny = await self._check_deny_rules(tool, tool_input) + if deny: + return deny + + # step 2: ask rules + ask = await self._check_ask_rules(tool, tool_input) + if ask: + ask.suggested_rules = await self._generate_suggestions( + tool, + tool_input, + ) + return ask + + # step 3: tool's own check_permissions + tool_decision = await tool.check_permissions(tool_input, self.context) + # step 3a: tool ALLOW / DENY returned as-is + if tool_decision.behavior in ( + PermissionBehavior.ALLOW, + PermissionBehavior.DENY, + ): + return tool_decision + # step 3b: safety ASK is bypass-immune — allow rules can't override + if self._is_safety_ask(tool_decision): + tool_decision.suggested_rules = await self._generate_suggestions( + tool, + tool_input, + ) + return tool_decision + + # step 4: allow rules + allow = await self._check_allow_rules(tool, tool_input) + if allow: + return allow + + # step 5: default — ASK the user + default = PermissionDecision( + behavior=PermissionBehavior.ASK, + message=f"Permission required for {tool.name}", + decision_reason=f"Mode: {self.context.mode.value}", + ) + default.suggested_rules = await self._generate_suggestions( + tool, + tool_input, + ) + return default + + async def _check_explore( + self, + tool: ToolBase, + tool_input: dict[str, Any], + ) -> PermissionDecision: + """Permission check for :attr:`PermissionMode.EXPLORE`. + + Read-only mode — modifications are categorically denied. Evaluation + order: + + 1. Deny rules → DENY + 2. Ask rules → ASK (with suggestions) + 3. :meth:`ToolBase.check_read_only` (input-aware): + - True → ALLOW + - False → DENY + + ``tool.check_permissions`` is not invoked: EXPLORE is fully + resolved by the read-only verdict, so safety ASK paths (e.g. + ``rm -rf /``) are subsumed into the broader DENY. Allow rules are + intentionally not consulted — EXPLORE's read-only guarantee + cannot be granted away by a user-configured rule. + + Args: + tool (`ToolBase`): + The tool instance being called. + tool_input (`dict[str, Any]`): + The tool input data. + + Returns: + `PermissionDecision`: + ALLOW for read-only invocations, DENY otherwise. + """ + # step 1: deny rules + deny = await self._check_deny_rules(tool, tool_input) + if deny: + return deny + + # step 2: ask rules + ask = await self._check_ask_rules(tool, tool_input) + if ask: + ask.suggested_rules = await self._generate_suggestions( + tool, + tool_input, + ) + return ask + + # step 3: read-only verdict decides everything (ALLOW or DENY) + if await tool.check_read_only(tool_input): + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message=( + f"Permission granted for {tool.name} " + f"(explore mode - read-only invocation)" + ), + decision_reason="Explore mode allows read-only operations", + ) + return PermissionDecision( + behavior=PermissionBehavior.DENY, + message=( + f"Permission denied for {tool.name} " + f"(explore mode is read-only)" + ), + decision_reason="Explore mode does not allow modifications", + ) + + async def _check_accept_edits( + self, + tool: ToolBase, + tool_input: dict[str, Any], + ) -> PermissionDecision: + """Permission check for :attr:`PermissionMode.ACCEPT_EDITS`. + + Edits within working directories are auto-allowed by each tool's + own ``check_permissions``; other operations follow the normal + flow. Evaluation order: + + 1. Deny rules → DENY + 2. Ask rules → ASK (with suggestions) + 3. :meth:`ToolBase.check_read_only` → True → ALLOW (fast path) + 4. ``tool.check_permissions``: + - ALLOW (e.g. ``Write`` to a file in the working directory) / + DENY → returned as-is + - Safety ASK (bypass-immune) → returned with suggestions + - Non-safety ASK / PASSTHROUGH → continue + 5. Allow rules → ALLOW + 6. Default → ASK (with suggestions) + + Args: + tool (`ToolBase`): + The tool instance being called. + tool_input (`dict[str, Any]`): + The tool input data. + + Returns: + `PermissionDecision`: + The final decision. + """ + # step 1: deny rules + deny = await self._check_deny_rules(tool, tool_input) + if deny: + return deny + + # step 2: ask rules + ask = await self._check_ask_rules(tool, tool_input) + if ask: + ask.suggested_rules = await self._generate_suggestions( + tool, + tool_input, + ) + return ask + + # step 3: read-only fast path — ALLOW without invoking the tool + if await tool.check_read_only(tool_input): + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message=( + f"Permission granted for {tool.name} " + f"(accept edits mode - read-only invocation)" + ), + decision_reason="Accept edits mode allows read-only " + "operations", + ) + + # step 4: tool's own check_permissions (working-directory check + # for Write/Edit, path-checked auto-allow for Bash, ...) + tool_decision = await tool.check_permissions(tool_input, self.context) + # step 4a: tool ALLOW / DENY returned as-is + if tool_decision.behavior in ( + PermissionBehavior.ALLOW, + PermissionBehavior.DENY, + ): + return tool_decision + # step 4b: safety ASK is bypass-immune + if self._is_safety_ask(tool_decision): + tool_decision.suggested_rules = await self._generate_suggestions( + tool, + tool_input, + ) + return tool_decision + + # step 5: allow rules + allow = await self._check_allow_rules(tool, tool_input) + if allow: + return allow + + # step 6: default — ASK the user + default = PermissionDecision( + behavior=PermissionBehavior.ASK, + message=f"Permission required for {tool.name}", + decision_reason=f"Mode: {self.context.mode.value}", + ) + default.suggested_rules = await self._generate_suggestions( + tool, + tool_input, + ) + return default + + async def _check_bypass( + self, + tool: ToolBase, + tool_input: dict[str, Any], + ) -> PermissionDecision: + """Permission check for :attr:`PermissionMode.BYPASS`. + + BYPASS is the "fully trusted" mode: the user has explicitly + opted out of safety prompts. All tool-emitted safety ASKs + (``rm -rf /``, write to ``~/.bashrc``, command-injection + patterns, dangerous sed, etc.) are **skipped** — only + user-configured deny / ask rules and tool-emitted DENY remain + as guardrails. The :attr:`PermissionDecision.bypass_immune` + field has no effect in BYPASS by design. + + Use BYPASS only for sandboxed / containerized environments or + when you fully trust the agent's behavior. For unattended + execution where safety still matters, use + :attr:`PermissionMode.DONT_ASK` instead — it converts safety + ASKs to DENY rather than skipping them. + + Evaluation order: + + 1. Deny rules → DENY + 2. Ask rules → ASK (with suggestions; honors explicit user intent) + 3. ``tool.check_permissions``: + - ALLOW / DENY → returned as-is + - ASK (including bypass-immune safety ASKs) → falls through + - PASSTHROUGH → falls through + 4. Allow rules → ALLOW + 5. Fallback → ALLOW (BYPASS) + + Args: + tool (`ToolBase`): + The tool instance being called. + tool_input (`dict[str, Any]`): + The tool input data. + + Returns: + `PermissionDecision`: + The final decision. + """ + # step 1: deny rules + deny = await self._check_deny_rules(tool, tool_input) + if deny: + return deny + + # step 2: ask rules (honor explicit user intent to be prompted) + ask = await self._check_ask_rules(tool, tool_input) + if ask: + ask.suggested_rules = await self._generate_suggestions( + tool, + tool_input, + ) + return ask + + # step 3: tool's own check_permissions — ALLOW / DENY returned; + # any ASK (including bypass-immune safety ASK) is intentionally + # NOT honored here, per BYPASS's "skip safety prompts" contract. + tool_decision = await tool.check_permissions(tool_input, self.context) + if tool_decision.behavior in ( + PermissionBehavior.ALLOW, + PermissionBehavior.DENY, + ): + return tool_decision + + # step 4: allow rules + allow = await self._check_allow_rules(tool, tool_input) + if allow: + return allow + + # step 5: bypass fallback — ALLOW everything else + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message=f"Permission granted for {tool.name} (bypass mode)", + decision_reason="Bypass mode allows all operations", + ) + + async def _check_dont_ask( + self, + tool: ToolBase, + tool_input: dict[str, Any], + ) -> PermissionDecision: + """Permission check for :attr:`PermissionMode.DONT_ASK`. + + Used when no user is available to answer prompts (scheduled + tasks, background runs). Invariant: this method must never + return :attr:`PermissionBehavior.ASK` — every code path that + would otherwise ASK is converted to DENY via + :meth:`_convert_ask_to_deny`. Evaluation order: + + 1. Deny rules → DENY + 2. Ask rules → DENY (converted, with suggestions preserved) + 3. ``tool.check_permissions``: + - ALLOW / DENY → returned as-is + - Safety ASK → DENY (converted, with suggestions preserved) + - Non-safety ASK / PASSTHROUGH → continue + 4. Allow rules → ALLOW + 5. Default → DENY (user not available to answer) + + Args: + tool (`ToolBase`): + The tool instance being called. + tool_input (`dict[str, Any]`): + The tool input data. + + Returns: + `PermissionDecision`: + The final decision (never ASK). + """ + # step 1: deny rules + deny = await self._check_deny_rules(tool, tool_input) + if deny: + return deny + + # step 2: ask rules — converted to DENY (no user available) + ask = await self._check_ask_rules(tool, tool_input) + if ask: + ask.suggested_rules = await self._generate_suggestions( + tool, + tool_input, + ) + return self._convert_ask_to_deny(tool, ask) + + # step 3: tool's own check_permissions + tool_decision = await tool.check_permissions(tool_input, self.context) + # step 3a: tool ALLOW / DENY returned as-is + if tool_decision.behavior in ( + PermissionBehavior.ALLOW, + PermissionBehavior.DENY, + ): + return tool_decision + # step 3b: safety ASK converted to DENY (no user available) + if self._is_safety_ask(tool_decision): + tool_decision.suggested_rules = await self._generate_suggestions( + tool, + tool_input, + ) + return self._convert_ask_to_deny(tool, tool_decision) + + # step 4: allow rules + allow = await self._check_allow_rules(tool, tool_input) + if allow: + return allow + + # step 5: default — DENY (no user available to confirm) + return PermissionDecision( + behavior=PermissionBehavior.DENY, + message=( + f"Permission denied for {tool.name} " + f"(dont_ask mode - user not available)" + ), + decision_reason="User is not available to answer permission " + "prompts", + ) + + @staticmethod + def _convert_ask_to_deny( + tool: ToolBase, + ask_decision: PermissionDecision, + ) -> PermissionDecision: + """Convert an ASK decision into a DENY for DONT_ASK mode. + + DONT_ASK's invariant is "never return ASK" — the user is not + available to answer prompts. This helper turns whatever produced + the ASK (an ASK rule, a safety check) into a DENY while + preserving traceability by carrying the original reason and + ``suggested_rules`` forward; callers (e.g. a UI surfacing the + scheduled-task failure) can still show the user what rule they + could add to unblock the operation in the future. + + Args: + tool (`ToolBase`): + The tool whose invocation is being denied. + ask_decision (`PermissionDecision`): + The original ASK decision to convert. + + Returns: + `PermissionDecision`: + A DENY decision with the original ASK's reason and + suggestions attached. + """ + return PermissionDecision( + behavior=PermissionBehavior.DENY, + message=( + f"Permission denied for {tool.name} " + f"(dont_ask mode - ASK converted to DENY, " + f"user not available)" + ), + decision_reason=( + f"DONT_ASK mode converted ASK to DENY. " + f"Original reason: {ask_decision.decision_reason}" + ), + suggested_rules=ask_decision.suggested_rules, + ) + + @staticmethod + def _is_safety_ask(decision: PermissionDecision) -> bool: + """Whether a decision is a bypass-immune safety ASK. + + A safety ASK is an ASK that a tool has explicitly marked with + :attr:`PermissionDecision.bypass_immune` ``= True``. Tools emit + these for dangerous operations (e.g. write to ``~/.bashrc``, + ``rm -rf /``, command-injection patterns) that must be surfaced + to the user regardless of allow rules in + ``DEFAULT``/``ACCEPT_EDITS``. ``BYPASS`` mode intentionally + skips this check (see :meth:`_check_bypass`); ``DONT_ASK`` + converts the ASK to DENY (see :meth:`_check_dont_ask`). + + Args: + decision (`PermissionDecision`): + The decision returned by a tool's ``check_permissions``. + + Returns: + `bool`: + True iff ``behavior == ASK`` and ``bypass_immune`` is set. + """ + return ( + decision.behavior == PermissionBehavior.ASK + and decision.bypass_immune + ) + + async def _check_deny_rules( + self, + tool: ToolBase, + input_data: dict[str, Any], + ) -> PermissionDecision | None: + """Check if any deny rules match the request. + + Args: + tool (`ToolBase`): + The tool instance being called + input_data (`dict[str, Any]`): + The tool input data + + Returns: + `PermissionDecision | None`: + DENY decision if a rule matches, None otherwise + """ + rules = self.context.deny_rules.get(tool.name, []) + for rule in rules: + if await self._rule_matches(tool, rule, input_data): + return PermissionDecision( + behavior=PermissionBehavior.DENY, + message=f"Permission to use {tool.name} has been denied", + decision_reason=f"Rule: {rule.rule_content}", + ) + return None + + async def _check_ask_rules( + self, + tool: ToolBase, + input_data: dict[str, Any], + ) -> PermissionDecision | None: + """Check if any ask rules match the request. + + Args: + tool (`ToolBase`): + The tool instance being called (used for tool-specific checks) + input_data (`dict[str, Any]`): + The tool input data + + Returns: + `PermissionDecision | None`: + ASK decision if a rule matches, None otherwise + """ + rules = self.context.ask_rules.get(tool.name, []) + for rule in rules: + if await self._rule_matches(tool, rule, input_data): + return PermissionDecision( + behavior=PermissionBehavior.ASK, + message=f"Permission required for {tool.name}", + decision_reason=f"Rule: {rule.rule_content}", + ) + return None + + async def _check_allow_rules( + self, + tool: ToolBase, + input_data: dict[str, Any], + ) -> PermissionDecision | None: + """Check if any allow rules match the request. + + Args: + tool (`ToolBase`): + The tool instance being called (used for tool-specific checks) + input_data (`dict[str, Any]`): + The tool input data + + Returns: + `PermissionDecision | None`: + ALLOW decision if a rule matches, None otherwise + """ + rules = self.context.allow_rules.get(tool.name, []) + for rule in rules: + if await self._rule_matches(tool, rule, input_data): + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message=f"Permission granted for {tool.name}", + updated_input=input_data, + ) + return None + + async def _rule_matches( + self, + tool: ToolBase, + rule: PermissionRule, + input_data: dict[str, Any], + ) -> bool: + """Check if a rule matches the input data. + + The matching strategy depends on the tool type: + - Bash: Substring matching against the command + - Write/Read: Glob pattern matching against file paths + - Other: Generic pattern matching + + Args: + rule (`PermissionRule`): + The permission rule to check + input_data (`dict[str, Any]`): + The tool input data + + Returns: + `bool`: + True if the rule matches, False otherwise + """ + # Empty rule_content matches everything + if not rule.rule_content: + return True + + # Try to use tool's match_rule method if available. + # ``_execute_async_or_sync_func`` keeps backward compatibility + # with third-party tools that still override match_rule with a + # sync ``def`` (the framework's signature is now ``async def``). + return await _execute_async_or_sync_func( + tool.match_rule, + rule.rule_content, + input_data, + ) + + async def _generate_suggestions( + self, + tool: ToolBase, + tool_input: dict[str, Any], + ) -> List[PermissionRule]: + """Generate suggested permission rules from a tool call. + + This method analyzes the tool call and generates broader permission + suggestions that the user can apply to avoid future confirmations. + + Strategy: + - For Bash: Extract command prefix (e.g., "npm run" -> "npm run:*") + - For File operations: Extract directory (e.g., + "src/file.py" -> "src/**") + - For other tools: Generate exact match rule + + Args: + tool (`ToolBase`): + The tool instance being called (used for tool-specific + suggestions) + tool_input (`dict[str, Any]`): + The tool input data (used for generating suggestions) + + Returns: + `List[PermissionRule]`: + List of suggested permission rules (usually 1, max 5 for + compound commands) + """ + + # Try to use tool's generate_suggestions method if available. + # ``_execute_async_or_sync_func`` keeps backward compatibility + # with third-party tools that still override this method with + # a sync ``def`` (the framework's signature is now ``async def``). + return await _execute_async_or_sync_func( + tool.generate_suggestions, + tool_input, + ) diff --git a/src/agentscope/permission/_rule.py b/src/agentscope/permission/_rule.py new file mode 100644 index 0000000000000000000000000000000000000000..8def2957872b690651cee7e3d3a84800fd5bcd9c --- /dev/null +++ b/src/agentscope/permission/_rule.py @@ -0,0 +1,36 @@ +# -*- coding: utf-8 -*- +"""Permission rule model for tool usage.""" +from pydantic import BaseModel + +from ._types import PermissionBehavior + + +class PermissionRule(BaseModel): + """Permission rule for tool usage. + + A permission rule defines whether a specific tool or tool operation + should be allowed, denied, or require user confirmation. The + rule_content field has different semantics depending on the tool_name: + + - For "Bash": rule_content is a substring pattern matched against the + command Example: rule_content="npm install" matches "npm install express" + + - For "Write"/"Read": rule_content is a glob pattern matched against file + paths Example: rule_content="src/**" matches "src/main.py" + + - For other tools: rule_content is a tool-specific filter pattern + """ + + tool_name: str + """The name of the tool this rule applies to (e.g., "Bash", + "Write", "Read").""" + + rule_content: str | None + """Optional filter pattern - semantics depend on tool_name.""" + + behavior: PermissionBehavior + """The permission behavior ("allow", "deny", or "ask").""" + + source: str + """Where this rule originated from (e.g., "userSettings", + "projectSettings").""" diff --git a/src/agentscope/permission/_types.py b/src/agentscope/permission/_types.py new file mode 100644 index 0000000000000000000000000000000000000000..7eb4cddd5201edd3ecb1159022d962d3d5735074 --- /dev/null +++ b/src/agentscope/permission/_types.py @@ -0,0 +1,102 @@ +# -*- coding: utf-8 -*- +# pylint: disable=line-too-long +"""Permission system types and engine for tool usage control. + +This module implements a permission system that controls tool execution based +on configurable rules. The permission system supports different matching +strategies depending on the tool type: + +- For Bash tools: rule_content is a substring pattern matched against commands +- For Write/Read tools: rule_content is a glob pattern matched against file + paths +- For other tools: rule_content uses generic matching logic +""" + +from enum import Enum + + +class PermissionMode(Enum): + """The mode of permission. + + Permission modes control how the system handles tool execution requests. + Different modes are suitable for different scenarios: + + +---------------+--------------------------------------------------+--------------------------------+ + | Mode | Behavior | Use Case | + +===============+==================================================+================================+ + | DEFAULT | Every operation asks for permission unless: | Default mode, most secure | + | | - an allow rule matches, OR | | + | | - the tool's ``check_permissions`` explicitly | | + | | returns ALLOW for the invocation (currently | | + | | only ``Bash`` auto-allows recognized | | + | | read-only commands such as ``ls``/``git | | + | | status``). Read/Glob/Grep return PASSTHROUGH | | + | | and fall through to the default ASK unless | | + | | an allow rule matches. | | + +---------------+--------------------------------------------------+--------------------------------+ + | ACCEPT_EDITS | - Auto-allow file writes in working directories | User present, rapid iteration | + | | - Auto-allow file reads in working directories | development | + | | - Auto-allow filesystem commands (mkdir, rm, | | + | | mv, cp, ...) **only when all target paths** | | + | | **resolve inside a working directory** | | + | | - Other operations follow normal rules | | + +---------------+--------------------------------------------------+--------------------------------+ + | EXPLORE | Read-only mode: | Exploring codebase, planning | + | | - Allow: read-only tools (``Read``/``Grep``/ | implementation | + | | ``Glob``) and read-only bash commands | | + | | (e.g. ``ls``, ``git status``) | | + | | - Deny: any modification tool / command | | + | | - User-configured DENY or ASK rules take | | + | | precedence over the read-only auto-allow | | + +---------------+--------------------------------------------------+--------------------------------+ + | BYPASS | Skip all permission checks except explicit | Sandboxed environments | + | | user-configured deny / ask rules and tool | (container, VM), unattended | + | | DENY. **Safety ASKs from tools are NOT** | runs where you fully trust | + | | **enforced** — including ``rm -rf /``, writes | the agent | + | | to ``~/.bashrc``, command-injection patterns, | | + | | etc. Use deny rules to protect specific paths. | | + | | For unattended runs that still need safety, | | + | | prefer DONT_ASK. | | + +---------------+--------------------------------------------------+--------------------------------+ + | DONT_ASK | Convert every ASK (including safety ASKs and | Scheduled tasks, background | + | | ASK-rule hits) to DENY. Safe-by-default for | execution when user is away | + | | unattended execution. | | + +---------------+--------------------------------------------------+--------------------------------+ + + Attributes: + DEFAULT: Default mode - explicit permission per action. The + only auto-allow path is the tool's own ``check_permissions`` + returning ALLOW (currently just ``Bash`` for recognized + read-only commands like ``ls``/``git status``). + ACCEPT_EDITS: Accept edits mode - automatically allows file + edits within working directories (including filesystem + bash commands whose every target is in a working dir). + EXPLORE: Explore mode - read-only; modifications are denied. + BYPASS: Bypass mode - skips safety checks; relies on user + deny / ask rules as the only guardrail. + DONT_ASK: Don't ask mode - converts all ASK decisions to DENY + (for unattended execution). + """ # noqa: E501 + + DEFAULT = "default" + ACCEPT_EDITS = "accept_edits" + EXPLORE = "explore" + BYPASS = "bypass" + DONT_ASK = "dont_ask" + + +class PermissionBehavior(Enum): + """The behavior of permission. + + Attributes: + ALLOW: Allow the operation + DENY: Deny the operation + ASK: Ask the user for permission + PASSTHROUGH: Let the permission engine continue with rule matching + (used by tools to defer decision to the engine) + """ + + ALLOW = "allow" + DENY = "deny" + ASK = "ask" + PASSTHROUGH = "passthrough" diff --git a/src/agentscope/py.typed b/src/agentscope/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/agentscope/rag/__init__.py b/src/agentscope/rag/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f62ccb30fb2b95e2ab9c8dbaa95aec95e5528546 --- /dev/null +++ b/src/agentscope/rag/__init__.py @@ -0,0 +1,35 @@ +# -*- coding: utf-8 -*- +"""The retrieval-augmented generation (RAG) module in AgentScope.""" + +from ._chunker import ApproxTokenChunker, ChunkerBase +from ._document import ( + Section, + Chunk, +) +from ._parser import ImageParser, ParserBase, PDFParser, PPTParser, TextParser +from ._vdb import ( + DocumentSummary, + VectorStoreBase, + VectorRecord, + VectorSearchResult, + QdrantStore, +) +from ._knowledge import KnowledgeBase + +__all__ = [ + "ApproxTokenChunker", + "ChunkerBase", + "Chunk", + "DocumentSummary", + "ImageParser", + "ParserBase", + "PDFParser", + "PPTParser", + "TextParser", + "Section", + "VectorStoreBase", + "VectorRecord", + "VectorSearchResult", + "QdrantStore", + "KnowledgeBase", +] diff --git a/src/agentscope/rag/_chunker/__init__.py b/src/agentscope/rag/_chunker/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d002f6e31a5ac8493aa227f42a771b8ffc69155a --- /dev/null +++ b/src/agentscope/rag/_chunker/__init__.py @@ -0,0 +1,10 @@ +# -*- coding: utf-8 -*- +"""Chunker implementations for the RAG indexing pipeline.""" + +from ._approx_token_chunker import ApproxTokenChunker +from ._base import ChunkerBase + +__all__ = [ + "ApproxTokenChunker", + "ChunkerBase", +] diff --git a/src/agentscope/rag/_chunker/_approx_token_chunker.py b/src/agentscope/rag/_chunker/_approx_token_chunker.py new file mode 100644 index 0000000000000000000000000000000000000000..b8138a2124c811b2ae468110f4d9d4870c00ae11 --- /dev/null +++ b/src/agentscope/rag/_chunker/_approx_token_chunker.py @@ -0,0 +1,172 @@ +# -*- coding: utf-8 -*- +"""A chunker that splits text by an approximate token count. + +The token count is approximated as ``len(text.encode("utf-8")) // 4``, +which avoids a hard dependency on any tokenizer library while staying +within the right order of magnitude for most LLM tokenizers. +""" +from bisect import bisect_right +from itertools import accumulate + +from ._base import ChunkerBase +from .._document import Chunk, Section +from ...message import TextBlock, DataBlock + + +class ApproxTokenChunker(ChunkerBase): + """A chunker based on an approximate token counting strategy. + + Text sections are sliced into pieces of at most ``chunk_size`` + approximate tokens, with ``overlap`` approximate tokens shared + between two consecutive pieces. The token count of a string is + approximated as ``len(text.encode("utf-8")) // 4``, so no + tokenizer dependency is required. + + Sections carrying a :class:`~agentscope.message.DataBlock` + (images, video, etc.) are passed through unchanged as a single + chunk. + + .. note:: Chunks never span across two input Sections, as + required by :class:`ChunkerBase`. + """ + + def __init__(self, chunk_size: int = 512, overlap: int = 50) -> None: + """Initialize the approx token chunker. + + Args: + chunk_size (`int`, defaults to `512`): + The maximum number of approximate tokens per chunk. + Must be a positive integer. + overlap (`int`, defaults to `50`): + The number of approximate tokens shared between two + consecutive chunks. Must be non-negative and smaller + than ``chunk_size``. + + Raises: + `ValueError`: + If ``chunk_size`` is not positive, or ``overlap`` is + negative or not smaller than ``chunk_size``. + """ + if chunk_size <= 0: + raise ValueError( + f"chunk_size must be positive, got {chunk_size}.", + ) + if overlap < 0 or overlap >= chunk_size: + raise ValueError( + "overlap must satisfy 0 <= overlap < chunk_size, " + f"got overlap={overlap}, chunk_size={chunk_size}.", + ) + + self.chunk_size = chunk_size + self.overlap = overlap + + async def chunk(self, sections: list[Section]) -> list[Chunk]: + """Chunk the input sections into smaller chunks based on an approx + token counting strategy. + + Args: + sections (`list[Section]`): + A list of sections to chunk. + + Returns: + `list[Chunk]`: + A list of chunks, with ``chunk_index`` numbered + ``0..N-1`` and ``total_chunks`` set to ``N`` on every + chunk. + """ + chunks: list[Chunk] = [] + for section in sections: + contents: list[TextBlock | DataBlock] + if isinstance(section.content, TextBlock): + contents = [ + TextBlock(text=piece) + for piece in self._split_text(section.content.text) + ] + else: + # DataBlock pass-through: never slice multimodal data + contents = [section.content] + + chunks.extend( + Chunk( + content=content, + source=section.source, + chunk_index=0, # renumbered below + total_chunks=0, # renumbered below + metadata=dict(section.metadata), + ) + for content in contents + ) + + for index, chunk in enumerate(chunks): + chunk.chunk_index = index + chunk.total_chunks = len(chunks) + + return chunks + + def _split_text(self, text: str) -> list[str]: + """Split text into pieces of at most ``chunk_size`` approx tokens. + + Consecutive pieces share approximately ``overlap`` tokens. + + Args: + text (`str`): + The text to split. + + Returns: + `list[str]`: + The text pieces, in document order. + """ + if self._approx_count_tokens(text) <= self.chunk_size: + return [text] + + # Cumulative UTF-8 byte length after each character, so that + # the byte length of text[i:j] == byte_offsets[j] - byte_offsets[i] + byte_offsets = [0, *accumulate(len(c.encode("utf-8")) for c in text)] + + chunk_bytes = self.chunk_size * 4 + overlap_bytes = self.overlap * 4 + + pieces: list[str] = [] + start = 0 + while start < len(text): + # The largest end such that the slice fits the byte budget + end = ( + bisect_right( + byte_offsets, + byte_offsets[start] + chunk_bytes, + ) + - 1 + ) + # Always make progress, even for characters whose UTF-8 + # encoding exceeds the budget on their own + end = max(end, start + 1) + pieces.append(text[start:end]) + + if end >= len(text): + break + + # Step back by the overlap budget, ensuring forward progress + next_start = ( + bisect_right( + byte_offsets, + byte_offsets[end] - overlap_bytes, + ) + - 1 + ) + start = max(next_start, start + 1) + + return pieces + + @staticmethod + def _approx_count_tokens(text: str) -> int: + """The approx count of tokens. + + Args: + text (`str`): + The text to be counted. + + Returns: + `int`: + The approx count of tokens. + """ + return len(text.encode("utf-8")) // 4 diff --git a/src/agentscope/rag/_chunker/_base.py b/src/agentscope/rag/_chunker/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..d981f2ec4f5b58a45cca4440bc494ac9bef908cb --- /dev/null +++ b/src/agentscope/rag/_chunker/_base.py @@ -0,0 +1,62 @@ +# -*- coding: utf-8 -*- +"""Abstract base class for chunkers. + +A :class:`ChunkerBase` subclass takes the :class:`Section` list +produced by a :class:`~agentscope.rag.ParserBase` and splits the +content into final :class:`Chunk` objects suitable for embedding and +storage in a vector database. + +Chunkers are **format-agnostic** — they operate on the unified +``TextBlock | DataBlock`` content carried in each Section. Long +:class:`TextBlock` content is sliced according to a chunking strategy +(by character count, by tokens, by semantic boundaries, etc.); short +text and :class:`DataBlock` content are passed through unchanged. + +Chunkers **never combine content across Section boundaries**. This +guarantee preserves the structural metadata attached by the Parser +(page numbers, slide indices, embedded-image isolation, etc.). +""" +from abc import ABC, abstractmethod + +from .._document import Chunk, Section + + +class ChunkerBase(ABC): + """Abstract base class for chunkers. + + Subclasses implement a specific chunking strategy (by character + count, by token count, by semantic boundary, etc.). The + chunker is configured once at construction time and reused + across many ``chunk()`` calls within the same knowledge base. + + Subclasses must guarantee: + + - **No cross-Section merging**: every output :class:`Chunk` is + derived from exactly one input :class:`Section`. + - **DataBlock pass-through**: a Section whose content is a + :class:`~agentscope.message.DataBlock` becomes a single Chunk + with the same content; multimodal data is never sliced. + - **Continuous indexing**: ``chunk_index`` runs from ``0`` to + ``total_chunks - 1`` across the entire output list, even + when the input contains many Sections. + - **Consistent total_chunks**: every output Chunk carries the + same ``total_chunks`` value (the length of the output list). + - **Metadata inheritance**: each output Chunk's ``source`` and + ``metadata`` are copied from its parent Section. + """ + + @abstractmethod + async def chunk(self, sections: list[Section]) -> list[Chunk]: + """Split a list of Sections into Chunks. + + Args: + sections (`list[Section]`): + The Sections produced by a :class:`ParserBase`, in + document order. + + Returns: + `list[Chunk]`: + The final chunks, in document order, with + ``chunk_index`` numbered ``0..N-1`` and + ``total_chunks`` set to ``N`` on every chunk. + """ diff --git a/src/agentscope/rag/_document.py b/src/agentscope/rag/_document.py new file mode 100644 index 0000000000000000000000000000000000000000..ad805ea920069a5671a9fc8164e7de7b4006dfaf --- /dev/null +++ b/src/agentscope/rag/_document.py @@ -0,0 +1,102 @@ +# -*- coding: utf-8 -*- +"""Data structures used in the RAG indexing pipeline. + +The indexing pipeline has two stages, each producing its own +structured output: + +1. :class:`Section` — produced by a :class:`ParserBase` from a raw + file. Each ``Section`` represents one "natural boundary" of the + source (a PDF page, a PPTX slide, an embedded image, a Markdown + heading section, etc.). A ``Chunker`` never combines content + across two ``Section`` instances, so ``Section`` is also a hard + boundary that prevents leakage of format-specific structure into + downstream chunks. + +2. :class:`Chunk` — produced by a :class:`ChunkerBase` from one or + more ``Section`` instances. Each ``Chunk`` is the final unit + that gets embedded and inserted into the vector store. + +Neither structure is persisted on its own — they are transient +in-memory carriers between pipeline stages. Persistence happens at +the :class:`~agentscope.rag.VectorRecord` and +``KnowledgeDocumentRecord`` layers. +""" +from typing import Any + +from pydantic import BaseModel, Field + +from ..message import TextBlock, DataBlock + + +class Section(BaseModel): + """A single natural section produced by a :class:`ParserBase`. + + A ``Section`` represents one logical region of the source file. + The :class:`ChunkerBase` guarantees that no resulting + :class:`Chunk` ever spans content from two different sections. + + The granularity of a ``Section`` is format-specific: + + - **PDF**: one section per page (plus separate sections for + embedded images). + - **PPTX**: one section per slide. + - **Markdown**: one section per top-level heading, or the entire + file if unstructured. + - **TXT / image / video**: one section for the whole file. + """ + + content: TextBlock | DataBlock + """The section content. Text sections use :class:`TextBlock`; + multimodal sections (images, video, etc.) use :class:`DataBlock`.""" + + source: str + """The source filename (e.g. ``"report.pdf"``). Carried through + to every downstream :class:`Chunk` and into the vector store + metadata for citation / display.""" + + metadata: dict[str, Any] = Field(default_factory=dict) + """Format-specific metadata written by the parser. Examples: + + - PDFParser: ``{"page": 3}`` + - PPTXParser: ``{"slide": 2}`` + - ExcelParser: ``{"sheet": "Q3 Sales"}`` + + These keys are not part of any retrieval / pipeline contract — + they are passed through verbatim to the vector store metadata for + later citation. Each Chunk inherits this dict from its parent + Section. + """ + + +class Chunk(BaseModel): + """A final indexable chunk produced by a :class:`ChunkerBase`. + + Each ``Chunk`` corresponds to one record in the vector store. + The required structural fields (``source``, ``chunk_index``, + ``total_chunks``) enable downstream features such as "expand + context around a hit" during retrieval. + """ + + content: TextBlock | DataBlock + """The chunk content (sliced from a text :class:`Section`, or a + multimodal :class:`DataBlock` passed through unchanged).""" + + source: str + """The source filename — inherited from the parent + :class:`Section`. Used for display / citation.""" + + chunk_index: int + """The 0-based index of this chunk **within the document**. + Sequential across all sections of the same source file. Used to + locate neighbouring chunks for "context expansion" at query time. + """ + + total_chunks: int + """The total number of chunks produced from the same source file. + Together with :attr:`chunk_index` lets callers know whether a hit + is near the start / end of the document, and bounds the + expansion range.""" + + metadata: dict[str, Any] = Field(default_factory=dict) + """Format-specific metadata inherited from the parent + :class:`Section`. See :attr:`Section.metadata`.""" diff --git a/src/agentscope/rag/_knowledge.py b/src/agentscope/rag/_knowledge.py new file mode 100644 index 0000000000000000000000000000000000000000..7649c328d896693bf1e0387d00eea9d9b8c55a16 --- /dev/null +++ b/src/agentscope/rag/_knowledge.py @@ -0,0 +1,384 @@ +# -*- coding: utf-8 -*- +"""Runtime handle for a single knowledge base. + +A :class:`KnowledgeBase` instance is the **single algorithmic source of +truth** for talking to one knowledge base: it pairs an embedding model +with a vector-store collection (optionally scoped by a payload +``metadata_filter``) and exposes the four operations a caller ever +needs — :meth:`search`, :meth:`insert_document`, +:meth:`delete_document`, :meth:`list_documents`. + +The handle is *narrow on purpose* — it carries the resolved runtime +state (embedding model + vector store + scope) and delegates every +operation to the bound :class:`VectorStoreBase`. Document parsing, +chunking, credential resolution, dimension policy validation, and +persistence of knowledge-base records all belong one layer up +(service-side :class:`KnowledgeBaseManagerBase` for hosted +deployments; the caller directly otherwise). + +The backing collection is created on first use — each operation +transparently calls :meth:`ensure_collection`, which is itself +idempotent and memoised after the first success, so the only cost is +one extra round-trip on the very first call against a fresh +deployment. + +``metadata_filter`` is the defense-in-depth scoping mechanism for +co-locating multiple logical knowledge bases inside the same physical +collection — typically multi-tenant deployments where every record +carries a ``{"tenant_id": "..."}`` payload. It is set once at +construction time and **always** applied: search/list never escape +it, and insert forces it onto every chunk's metadata so a malicious or +buggy parser cannot rebind a record into another scope. +""" + +import asyncio + +from ._document import Chunk +from ._vdb import VectorRecord, VectorSearchResult, VectorStoreBase +from .._utils._common import _generate_id +from ..embedding import EmbeddingModelBase +from ..message import DataBlock, TextBlock +from ._vdb import DocumentSummary + + +class KnowledgeBase: + """Runtime handle for one knowledge base. + + Binds an embedding model and a vector-store collection together so + callers can retrieve / insert / delete / list documents without + repeating the wiring. Cheap to construct (no I/O); the collection + itself is created lazily on the first operation, so a fresh + deployment "just works" without an explicit setup step. + + .. code-block:: python + + kb = KnowledgeBase( + name="company-handbook", + description="Internal HR and onboarding documents.", + embedding_model=embedding_model, + vector_store=vector_store, + collection="handbook", + ) + await kb.insert_document(chunks) + results = await kb.search(["What is the PTO policy?"]) + """ + + name: str + """Agent-oriented knowledge base name — used by tool descriptions + and frontend rendering.""" + + description: str + """Agent-oriented knowledge base description — what this knowledge + base contains and when to retrieve from it.""" + + def __init__( + self, + name: str, + description: str, + embedding_model: EmbeddingModelBase, + vector_store: VectorStoreBase, + collection: str, + metadata_filter: dict | None = None, + ) -> None: + """Initialize the runtime handle. + + Args: + name (`str`): + Agent-oriented knowledge base name. Surfaced to the + LLM (via tool descriptions) and to the front-end. + description (`str`): + Agent-oriented description. Should answer "what is in + this knowledge base and when should I search it?" — the + LLM uses it to decide whether to call the search tool + in agentic mode. + embedding_model (`EmbeddingModelBase`): + The embedding model used to embed both queries and + inserted chunks. Must be the same model used at + indexing time and at retrieval time, otherwise vectors + will not be comparable. + vector_store (`VectorStoreBase`): + The shared vector-store connection. The store must + already be entered (its own ``__aenter__`` already + called) before any operation on this handle runs. + collection (`str`): + The physical collection backing this knowledge base. + Created lazily on the first operation; see + :meth:`ensure_collection`. + metadata_filter (`dict | None`, optional): + Defense-in-depth payload filter. When set: + + - :meth:`search` and :meth:`list_documents` restrict + results to records whose payload matches every + ``key == value`` pair; + - :meth:`insert_document` forces these keys onto every + inserted chunk's metadata, overriding caller-supplied + values, so records cannot leak into another scope. + + ``None`` disables filtering — the default for + deployments where every knowledge base owns its + collection outright. + """ + self.name = name + self.description = description + self._embedding_model = embedding_model + self._vector_store = vector_store + self._collection = collection + self._metadata_filter = metadata_filter + # Memoise the "collection exists" check after the first + # successful ensure_collection so subsequent operations avoid + # the extra round-trip. + self._collection_ready = False + + # ------------------------------------------------------------------ + # Read-only accessors + # ------------------------------------------------------------------ + + @property + def embedding_model(self) -> EmbeddingModelBase: + """The bound embedding model.""" + return self._embedding_model + + @property + def vector_store(self) -> VectorStoreBase: + """The bound vector store.""" + return self._vector_store + + @property + def collection(self) -> str: + """The physical collection backing this knowledge base.""" + return self._collection + + @property + def metadata_filter(self) -> dict | None: + """The defense-in-depth payload filter, or ``None``.""" + return self._metadata_filter + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def ensure_collection(self) -> None: + """Idempotently create the backing collection if missing. + + Called transparently at the top of every public operation — + callers should not need to invoke it themselves. Memoised on + the instance after the first success, so subsequent calls are + a single ``if`` check. + + Looks up the collection via + :meth:`VectorStoreBase.has_collection` and creates it with the + embedding model's :attr:`~EmbeddingModelBase.dimensions` when + absent. + + Raises whatever the backend raises if the collection exists at + an incompatible dimension (the backend is the authority on + that; we do not double-check here). + """ + if self._collection_ready: + return + if not await self._vector_store.has_collection(self._collection): + await self._vector_store.create_collection( + self._collection, + dimensions=self._embedding_model.dimensions, + ) + self._collection_ready = True + + # ------------------------------------------------------------------ + # Search + # ------------------------------------------------------------------ + + async def search( + self, + queries: list[str | TextBlock | DataBlock], + top_k: int = 5, + score_threshold: float | None = None, + ) -> list[VectorSearchResult]: + """Search the knowledge base with one or more queries. + + All queries are embedded in a single batch, then searched + concurrently against the bound collection (with + :attr:`metadata_filter` applied). Hits are deduplicated by + ``(document_id, chunk_index)`` keeping the best score, + optionally filtered by ``score_threshold``, sorted by + descending score, and truncated to ``top_k``. + + Args: + queries (`list[str | TextBlock | DataBlock]`): + Query inputs. Text may be either bare ``str`` or + :class:`TextBlock`; :class:`DataBlock` items are + **silently dropped** when the bound embedding model + does not declare ``supports_multimodal`` — text-only + models would otherwise reject them. Callers can + therefore pass a mixed list without per-KB filtering. + top_k (`int`, defaults to ``5``): + Maximum number of results returned across all queries + (after dedup). + score_threshold (`float | None`, optional): + Minimum similarity score for a hit to be retained. + Only meaningful for similarity metrics where higher is + better (cosine / dot-product). ``None`` disables + filtering. + + Returns: + `list[VectorSearchResult]`: + At most ``top_k`` deduplicated hits ordered by + descending similarity score. Empty when there are no + queries the bound embedding model can consume. + """ + if not queries: + return [] + + if not self._embedding_model.supports_multimodal: + queries = [q for q in queries if not isinstance(q, DataBlock)] + if not queries: + return [] + + await self.ensure_collection() + response = await self._embedding_model(queries) + + results_per_query = await asyncio.gather( + *( + self._vector_store.search( + collection=self._collection, + query_vector=vector, + top_k=top_k, + metadata_filter=self._metadata_filter, + ) + for vector in response.embeddings + ), + ) + + best: dict[tuple[str, int], VectorSearchResult] = {} + for results in results_per_query: + for result in results: + if ( + score_threshold is not None + and result.score < score_threshold + ): + continue + # ``(document_id, chunk_index)`` is the stable identity + # of a chunk: it survives reindex (block UUIDs do not) + # and uniquely names "this slice of that document" + # regardless of which query surfaced it. + key = (result.document_id, result.chunk.chunk_index) + if key not in best or result.score > best[key].score: + best[key] = result + + merged = sorted( + best.values(), + key=lambda result: result.score, + reverse=True, + ) + return merged[:top_k] + + # ------------------------------------------------------------------ + # Document management + # ------------------------------------------------------------------ + + async def insert_document( + self, + chunks: list[Chunk], + document_id: str | None = None, + document_metadata: dict | None = None, + ) -> str: + """Embed and insert a list of chunks as a single source document. + + All chunks share the resolved ``document_id``; + :meth:`delete_document` later removes them as a unit. Each + chunk's metadata is merged in this precedence (highest wins): + + 1. :attr:`metadata_filter` keys — defense-in-depth scoping, so + a chunk can never be inserted with a payload that escapes + the filter (any escape would silently disappear at retrieve + time anyway, but failing closed at insert is clearer). + 2. The chunk's pre-existing ``metadata`` — parser-supplied. + 3. ``document_metadata`` — document-level fields propagated + down (filename, media type, upload time, ...). + + Args: + chunks (`list[Chunk]`): + The pre-chunked document content (already produced by + a parser + chunker pipeline). An empty list is a + no-op. + document_id (`str | None`, optional): + The document identifier. When ``None`` a fresh UUID + hex is generated and returned so the caller can record + it for future :meth:`delete_document` calls. + document_metadata (`dict | None`, optional): + Document-level metadata (filename, media type, size, + upload time, ...). Merged into each chunk's + ``metadata``. + + Returns: + `str`: + The (possibly generated) document id. + + Raises: + `RuntimeError`: + If the embedding model returns a number of vectors + that does not match the number of chunks. + """ + if not chunks: + return document_id or _generate_id() + document_id = document_id or _generate_id() + + await self.ensure_collection() + + # Precedence: metadata_filter wins (security boundary), then + # chunk metadata, then document_metadata. See docstring. + for chunk in chunks: + chunk.metadata = { + **(document_metadata or {}), + **chunk.metadata, + **(self._metadata_filter or {}), + } + + response = await self._embedding_model( + [chunk.content for chunk in chunks], + ) + + if len(response.embeddings) != len(chunks): + raise RuntimeError( + f"Embedding model returned {len(response.embeddings)} " + f"vectors for {len(chunks)} chunks.", + ) + + records = [ + VectorRecord( + vector=vector, + document_id=document_id, + chunk=chunk, + ) + for vector, chunk in zip(response.embeddings, chunks) + ] + await self._vector_store.insert(self._collection, records) + return document_id + + async def delete_document(self, document_id: str) -> None: + """Remove every record for one source document. + + Args: + document_id (`str`): + The source document id whose records should be removed. + """ + await self.ensure_collection() + await self._vector_store.delete( + self._collection, + document_id, + ) + + async def list_documents(self) -> list["DocumentSummary"]: + """List all distinct source documents in this knowledge base. + + Filtered by :attr:`metadata_filter` when set, so callers only + ever see documents within their own scope. + + Returns: + `list[DocumentSummary]`: + One summary per indexed document, in unspecified order. + """ + await self.ensure_collection() + return await self._vector_store.list_documents( + self._collection, + metadata_filter=self._metadata_filter, + ) diff --git a/src/agentscope/rag/_parser/__init__.py b/src/agentscope/rag/_parser/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fa8e496ea474f3de2b189764a3e9a3778e8a705d --- /dev/null +++ b/src/agentscope/rag/_parser/__init__.py @@ -0,0 +1,16 @@ +# -*- coding: utf-8 -*- +"""File parser implementations for the RAG indexing pipeline.""" + +from ._base import ParserBase +from ._image import ImageParser +from ._pdf import PDFParser +from ._ppt import PPTParser +from ._text import TextParser + +__all__ = [ + "ParserBase", + "PDFParser", + "PPTParser", + "ImageParser", + "TextParser", +] diff --git a/src/agentscope/rag/_parser/_base.py b/src/agentscope/rag/_parser/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..b05f55d3bb3ca71e0038dbf65bccbe8961abb9f0 --- /dev/null +++ b/src/agentscope/rag/_parser/_base.py @@ -0,0 +1,116 @@ +# -*- coding: utf-8 -*- +"""Abstract base class for file parsers. + +A :class:`ParserBase` subclass handles **one file format**. Its job +is to read a file's raw bytes and produce a list of +:class:`~agentscope.rag.Section` objects, each representing a natural +boundary of the source (e.g. one PDF page, one PPTX slide, one +embedded image). + +Parsers **do not chunk text**. Long text is left intact inside the +Section; splitting happens later in a +:class:`~agentscope.rag.ChunkerBase`. Parsers also do not need to +worry about output size — only about preserving the structural +boundaries that downstream consumers must not cross. +""" +import mimetypes +from abc import ABC, abstractmethod + +from .._document import Section + + +class ParserBase(ABC): + """Abstract base class for file-format parsers. + + Each subclass handles a single file format (or a related family, + e.g. all plain-text MIME types). Subclasses are typically + instantiated once and reused across many ``parse()`` calls. + + Subclasses should be stateless or thread-safe — a single + instance may be invoked concurrently from multiple agent runs. + + Subclasses must declare :attr:`supported_media_types` so that the + KnowledgeBaseManager can route uploaded files to the right parser + based on standard IANA media types (RFC 6838). + """ + + supported_media_types: list[str] + """Standard IANA media types (RFC 6838) this parser handles, + e.g. ``["application/pdf"]`` or + ``["text/plain", "text/markdown"]``. Used by the + KnowledgeBaseManager to select a parser for an uploaded file.""" + + @classmethod + def supported_extensions(cls) -> list[str]: + """Filename extensions (including the leading ``.``) this parser + can produce uploads for. + + The base implementation derives extensions from + :attr:`supported_media_types` via + :func:`mimetypes.guess_all_extensions` — good enough for clean + IANA types like ``application/pdf``. Subclasses **should + override** this when the default reverse-lookup is noisy + (``text/plain`` resolves to ``.bat`` / ``.c`` / ``.pl`` and a + dozen other developer extensions no KB user wants in the file + picker) or when a media type has no registered extension at all + (``application/x-yaml`` returns the empty list). + + The result is consumed by the front-end's ```` and + by the client-side filename guard; it is **not** consulted for + media-type routing — that always goes through + :attr:`supported_media_types`. + + Returns: + `list[str]`: + Deduplicated, sorted extensions (each starting with + ``.``). May be empty when no media type resolves. + """ + out: set[str] = set() + for media_type in cls.supported_media_types: + out.update(mimetypes.guess_all_extensions(media_type)) + return sorted(out) + + @abstractmethod + async def parse( + self, + file: bytes | str, + filename: str, + ) -> list[Section]: + """Parse a file into a list of :class:`Section` objects. + + The ``file`` argument is a union covering the three call sites + a parser sees in practice: + + - ``bytes`` — the raw payload, as handed in by HTTP uploads + and blob-store reads. + - ``str`` for binary parsers (PDF, PPT, image, …) — a + **filesystem path** to the file to read. The parser opens + the path itself; callers do not need to read the bytes first. + - ``str`` for :class:`TextParser` — disambiguated at runtime: + if the string names an existing file on disk it is treated + as a path and the file is decoded with the configured + encoding; otherwise it is treated as pre-decoded text. + + Args: + file (`bytes | str`): + The file content or a path to it (see above). + filename (`str`): + The original filename (e.g. ``"report.pdf"``). Used + for error messages and copied into each Section's + :attr:`Section.source` field for downstream display + / citation. + + Returns: + `list[Section]`: + One Section per natural boundary in the source file. + For unstructured formats (plain text, image, video), + a single Section may cover the whole file. Sections + are returned in document order. + + Raises: + `TypeError`: If the subclass does not accept the supplied + ``file`` form. + `FileNotFoundError`: If a binary parser is handed a + ``str`` that does not name an existing file. + `ValueError`: If the file cannot be parsed. + """ diff --git a/src/agentscope/rag/_parser/_image.py b/src/agentscope/rag/_parser/_image.py new file mode 100644 index 0000000000000000000000000000000000000000..8e8063b2fa77d3aed520cbff51406a5656a2ed61 --- /dev/null +++ b/src/agentscope/rag/_parser/_image.py @@ -0,0 +1,92 @@ +# -*- coding: utf-8 -*- +"""Image file parser. + +A single :class:`Section` carrying the raw image bytes as a +base64-encoded :class:`DataBlock`. No OCR, no captioning — the +section is the image, ready to flow through to a multimodal +embedding model unchanged. +""" +import base64 + +from ...message import Base64Source, DataBlock +from .._document import Section +from ._base import ParserBase +from ._utils import _guess_image_media_type + + +class ImageParser(ParserBase): + """Parser for image files. + + Wraps the entire file as a single :class:`Section` holding a + :class:`DataBlock` with the image's base64-encoded bytes. This is + the input shape a multimodal embedding model expects; the + surrounding pipeline (chunker, vector store) treats the section + opaquely. + + The IANA media type is sniffed from the bytes' magic number so + callers do not need to pass it explicitly. + """ + + supported_media_types: list[str] = [ + "image/png", + "image/jpeg", + "image/gif", + "image/bmp", + "image/webp", + ] + + @classmethod + def supported_extensions(cls) -> list[str]: + """Return the canonical image extensions.""" + return [ + ".bmp", + ".gif", + ".jpeg", + ".jpg", + ".png", + ".webp", + ] + + async def parse( + self, + file: bytes | str, + filename: str, + ) -> list[Section]: + """Wrap the image bytes in a single :class:`Section`. + + Args: + file (`bytes | str`): + Either the raw image bytes, or a filesystem path to + the image file. + filename (`str`): + The source filename, copied into + :attr:`Section.source`. + + Returns: + `list[Section]`: + A one-element list whose section's ``content`` is a + :class:`DataBlock` with the base64-encoded image data. + + Raises: + `FileNotFoundError`: If ``file`` is a ``str`` pointing to + a path that does not exist. + """ + if isinstance(file, str): + with open(file, "rb") as fp: + file = fp.read() + + media_type = _guess_image_media_type(file) + data = base64.b64encode(file).decode("utf-8") + return [ + Section( + content=DataBlock( + source=Base64Source( + media_type=media_type, + data=data, + ), + name=filename, + ), + source=filename, + metadata={"media_type": media_type}, + ), + ] diff --git a/src/agentscope/rag/_parser/_pdf.py b/src/agentscope/rag/_parser/_pdf.py new file mode 100644 index 0000000000000000000000000000000000000000..92389f418d166a1ac56612fa82f44072377707f9 --- /dev/null +++ b/src/agentscope/rag/_parser/_pdf.py @@ -0,0 +1,92 @@ +# -*- coding: utf-8 -*- +"""PDF file parser. + +One :class:`Section` per page so a downstream +:class:`~agentscope.rag.ChunkerBase` never combines text across page +boundaries. Each section's :attr:`Section.metadata` carries the +page number (starting at 1) for later citation. +""" +import io + +from ...message import TextBlock +from .._document import Section +from ._base import ParserBase + + +class PDFParser(ParserBase): + """Parser for PDF files. + + Returns one :class:`Section` per page. Empty / image-only pages + still produce a Section (with empty text) so the page → section + correspondence stays exact — downstream chunkers naturally drop + empty content. + + Requires :mod:`pypdf`; install with ``pip install pypdf`` (or via + the ``agentscope[rag]`` extra). + """ + + supported_media_types: list[str] = ["application/pdf"] + + @classmethod + def supported_extensions(cls) -> list[str]: + """Return the canonical ``.pdf`` extension.""" + return [".pdf"] + + async def parse( + self, + file: bytes | str, + filename: str, + ) -> list[Section]: + """Read the PDF bytes and return one Section per page. + + Args: + file (`bytes | str`): + Either the raw PDF bytes, or a filesystem path to + the PDF file. + filename (`str`): + The source filename, copied verbatim into each + Section's :attr:`Section.source` field. + + Returns: + `list[Section]`: + One Section per page, in document order. Each + section's metadata holds ``{"page": }``. + + Raises: + `FileNotFoundError`: If ``file`` is a ``str`` pointing to + a path that does not exist. + `ImportError`: If :mod:`pypdf` is not installed. + `ValueError`: If the bytes cannot be parsed as PDF. + """ + if isinstance(file, str): + with open(file, "rb") as fp: + file = fp.read() + + try: + from pypdf import PdfReader + from pypdf.errors import PdfReadError + except ImportError as e: + raise ImportError( + "Please install pypdf to use the PDF parser. " + "You can install it by `pip install pypdf` (or " + "`pip install agentscope[rag]`).", + ) from e + + try: + reader = PdfReader(io.BytesIO(file)) + except PdfReadError as e: + raise ValueError( + f"Failed to parse {filename!r} as PDF: {e}", + ) from e + + sections: list[Section] = [] + for page_idx, page in enumerate(reader.pages, start=1): + text = page.extract_text() or "" + sections.append( + Section( + content=TextBlock(text=text), + source=filename, + metadata={"page": page_idx}, + ), + ) + return sections diff --git a/src/agentscope/rag/_parser/_ppt.py b/src/agentscope/rag/_parser/_ppt.py new file mode 100644 index 0000000000000000000000000000000000000000..6b0cc4a8aa903d19c96ea5b8e1de91676366de5d --- /dev/null +++ b/src/agentscope/rag/_parser/_ppt.py @@ -0,0 +1,341 @@ +# -*- coding: utf-8 -*- +"""PowerPoint (.pptx) file parser. + +Walks the deck slide-by-slide and emits one :class:`Section` per +contiguous content block — adjacent text runs (and, by default, +tables) are merged into a single text section; embedded images are +emitted as their own :class:`DataBlock` sections. Each section's +metadata carries the slide index (starting at 1) for later citation. + +Mirrors the knob set of the v1 ``PowerPointReader``: +``include_image``, ``separate_table``, ``table_format``, +``slide_prefix``, ``slide_suffix``. Chunking is **not** done here — +long text stays intact inside a section and is split downstream by a +:class:`~agentscope.rag.ChunkerBase`. +""" +import base64 +import io +from typing import Any, Literal + +from ..._logging import logger +from ...message import Base64Source, DataBlock, TextBlock +from .._document import Section +from ._base import ParserBase +from ._utils import ( + _guess_image_media_type, + _table_to_json, + _table_to_markdown, +) + + +def _extract_table_rows(table: Any) -> list[list[str]]: + """Read a python-pptx table into a 2-D ``list[list[str]]``. + + Args: + table (`Any`): + The python-pptx ``Table`` object. + + Returns: + `list[list[str]]`: + One inner list per row; per-cell line breaks are + normalised to ``\\n``. + """ + rows: list[list[str]] = [] + for row in table.rows: + cells: list[str] = [] + for cell in row.cells: + text = cell.text.strip() + text = text.replace("\r\n", "\n").replace("\r", "\n") + cells.append(text) + rows.append(cells) + return rows + + +def _extract_image_bytes(shape: Any) -> bytes | None: + """Return the embedded image bytes for a picture shape, or ``None``. + + Args: + shape (`Any`): + A python-pptx shape. Non-picture shapes return ``None``. + + Returns: + `bytes | None`: + The raw image bytes, or ``None`` when ``shape`` is not a + picture / the bytes are unreadable. + """ + try: + from pptx.enum.shapes import MSO_SHAPE_TYPE + + picture_type = MSO_SHAPE_TYPE.PICTURE + except ImportError: + # MSO_SHAPE_TYPE.PICTURE numeric value used as the fallback + # so the parser still works against pptx builds where the + # enum import path has moved. + picture_type = 13 + + if shape.shape_type != picture_type: + return None + try: + return shape.image.blob + except Exception as e: # pylint: disable=broad-except + logger.warning("Failed to extract image from PPT shape: %s", e) + return None + + +class PPTParser(ParserBase): + """Parser for PowerPoint ``.pptx`` files. + + Slide order is preserved. Within a slide, shapes are visited in + document order and grouped into a minimum number of sections: + + - **Text** shapes and **table** shapes contribute to one running + text section. When ``separate_table=True`` a table closes the + running section and starts a new one of its own. + - **Picture** shapes emit a standalone :class:`Section` whose + ``content`` is a :class:`DataBlock` holding the base64-encoded + image bytes. + + ``slide_prefix`` / ``slide_suffix`` are wrapped around each + slide's text content (the prefix is prepended to the slide's + first text section, the suffix is appended to its last text + section). Use ``None`` on either to disable wrapping. + """ + + supported_media_types: list[str] = [ + "application/vnd.openxmlformats-officedocument.presentationml" + ".presentation", + ] + + @classmethod + def supported_extensions(cls) -> list[str]: + """Return ``[".pptx"]`` — the only format ``python-pptx`` + reads.""" + return [".pptx"] + + def __init__( + self, + include_image: bool = True, + separate_table: bool = False, + table_format: Literal["markdown", "json"] = "markdown", + slide_prefix: str | None = "", + slide_suffix: str | None = "", + ) -> None: + """Initialize the PowerPoint parser. + + Args: + include_image (`bool`, defaults to ``True``): + When ``True``, picture shapes are emitted as + :class:`DataBlock` sections. Set to ``False`` to keep + a text-only index. + separate_table (`bool`, defaults to ``False``): + When ``True``, each table becomes its own text + section, never merged with surrounding text. + table_format (`Literal["markdown", "json"]`, defaults to + ``"markdown"``): + How to render tables. ``"markdown"`` uses pipe-table + syntax; ``"json"`` emits a JSON array prefixed with a + ```` marker — choose JSON when cells + contain newlines that would corrupt Markdown layout. + slide_prefix (`str | None`, defaults to + ``""``): + Prepended to the first text section of each slide. + Supports the ``{index}`` placeholder (starting at 1). Use + ``None`` to disable. + slide_suffix (`str | None`, defaults to ``""``): + Appended to the last text section of each slide. Use + ``None`` to disable. + + Raises: + `ValueError`: If ``table_format`` is not one of + ``"markdown"`` / ``"json"``. + """ + if table_format not in ("markdown", "json"): + raise ValueError( + "The table_format must be one of 'markdown' or 'json', " + f"got {table_format!r}.", + ) + self.include_image = include_image + self.separate_table = separate_table + self.table_format = table_format + self.slide_prefix = slide_prefix + self.slide_suffix = slide_suffix + + async def parse( + self, + file: bytes | str, + filename: str, + ) -> list[Section]: + """Parse a PPTX file into a list of :class:`Section` objects. + + Args: + file (`bytes | str`): + Either the raw PPTX bytes, or a filesystem path to + the PPTX file. + filename (`str`): + The source filename, copied into each Section's + :attr:`Section.source`. + + Returns: + `list[Section]`: + Sections in deck order. Text sections carry + ``{"slide": }``; image sections add + ``{"media_type": "image/..."}``. + + Raises: + `FileNotFoundError`: If ``file`` is a ``str`` pointing to + a path that does not exist. + `ImportError`: If :mod:`python-pptx` is not installed. + `ValueError`: If the bytes cannot be parsed. + """ + if isinstance(file, str): + with open(file, "rb") as fp: + file = fp.read() + + try: + from pptx import Presentation + except ImportError as e: + raise ImportError( + "Please install python-pptx to use the PowerPoint " + "parser. You can install it by " + "`pip install python-pptx` (or " + "`pip install agentscope[rag]`).", + ) from e + + try: + prs = Presentation(io.BytesIO(file)) + except Exception as e: # pylint: disable=broad-except + raise ValueError( + f"Failed to parse {filename!r} as PPTX: {e}", + ) from e + + sections: list[Section] = [] + for slide_idx, slide in enumerate(prs.slides): + sections.extend( + self._parse_slide(slide, slide_idx, filename), + ) + return sections + + # ------------------------------------------------------------------ + # Slide-level parsing + # ------------------------------------------------------------------ + + def _parse_slide( + self, + slide: Any, + slide_idx: int, + filename: str, + ) -> list[Section]: + """Walk one slide and return its ordered sections.""" + slide_no = slide_idx + 1 + prefix = ( + self.slide_prefix.format(index=slide_no) + if self.slide_prefix is not None + else "" + ) + + slide_sections: list[Section] = [] + # ``text_buffer`` accumulates the running text section; it is + # flushed whenever an image arrives, the slide ends, or + # ``separate_table`` forces a break around a table shape. + text_buffer: list[str] = [] + + def flush_text() -> None: + if not text_buffer: + return + slide_sections.append( + Section( + content=TextBlock(text="\n".join(text_buffer)), + source=filename, + metadata={"slide": slide_no}, + ), + ) + text_buffer.clear() + + # Slide prefix lives at the very top of the first text section. + if prefix: + text_buffer.append(prefix) + + for shape in slide.shapes: + # 1. Pictures — flush running text, emit a DataBlock section. + if self.include_image: + image_bytes = _extract_image_bytes(shape) + if image_bytes is not None: + flush_text() + media_type = _guess_image_media_type(image_bytes) + data = base64.b64encode(image_bytes).decode("utf-8") + slide_sections.append( + Section( + content=DataBlock( + source=Base64Source( + media_type=media_type, + data=data, + ), + name=filename, + ), + source=filename, + metadata={ + "slide": slide_no, + "media_type": media_type, + }, + ), + ) + continue + + # 2. Tables — render to text; optionally flush around them. + if getattr(shape, "has_table", False): + try: + rows = _extract_table_rows(shape.table) + except Exception as e: # pylint: disable=broad-except + logger.warning( + "Failed to extract table from slide %d: %s", + slide_no, + e, + ) + continue + rendered = ( + _table_to_markdown(rows) + if self.table_format == "markdown" + else _table_to_json(rows) + ) + if not rendered: + continue + if self.separate_table: + flush_text() + slide_sections.append( + Section( + content=TextBlock(text=rendered), + source=filename, + metadata={"slide": slide_no}, + ), + ) + else: + text_buffer.append(rendered) + continue + + # 3. Text frames. + if getattr(shape, "has_text_frame", False): + try: + parts = [ + para.text.strip() + for para in shape.text_frame.paragraphs + if para.text.strip() + ] + except Exception as e: # pylint: disable=broad-except + logger.warning( + "Failed to extract text from shape in slide %d: %s", + slide_no, + e, + ) + continue + if parts: + text_buffer.append("\n".join(parts)) + + # Suffix goes onto the last text section of this slide. If the + # slide ends on an image (text_buffer empty, no prior text + # section in this slide), create a small text-only carrier so + # the suffix is preserved. + if self.slide_suffix is not None: + text_buffer.append(self.slide_suffix) + + flush_text() + return slide_sections diff --git a/src/agentscope/rag/_parser/_text.py b/src/agentscope/rag/_parser/_text.py new file mode 100644 index 0000000000000000000000000000000000000000..26e229525e0ad3d0b3edcea62abd816d8d3d3893 --- /dev/null +++ b/src/agentscope/rag/_parser/_text.py @@ -0,0 +1,125 @@ +# -*- coding: utf-8 -*- +"""Plain-text file parser.""" +import os + +from ...message import TextBlock +from .._document import Section +from ._base import ParserBase + + +class TextParser(ParserBase): + """Parser for plain-text file formats. + + Reads the entire file as UTF-8 text and returns a single + :class:`Section`. No internal boundaries are inferred — the file + is treated as one unstructured blob, leaving all splitting to a + downstream :class:`~agentscope.rag.ChunkerBase`. + + Supports a fixed set of standard text-based IANA media types + (``text/plain``, ``text/markdown``, ``text/csv``, …). Use + ``TextParser.supported_media_types`` to enumerate them. + """ + + supported_media_types: list[str] = [ + "text/plain", + "text/markdown", + "text/csv", + "text/html", + "text/x-rst", + "application/json", + "application/xml", + "application/x-yaml", + ] + """Standard IANA media types this parser handles.""" + + @classmethod + def supported_extensions(cls) -> list[str]: + """Return the human-friendly text extensions. + + Override the base reverse-lookup because + :func:`mimetypes.guess_all_extensions` returns a long tail of + developer-tool extensions for ``text/plain`` (``.bat`` / + ``.c`` / ``.pl`` / ``.ksh`` / …) that have no place in a KB + file picker, and returns the empty list for + ``application/x-yaml``. + """ + return [ + ".csv", + ".htm", + ".html", + ".json", + ".markdown", + ".md", + ".rst", + ".txt", + ".xml", + ".yaml", + ".yml", + ] + + def __init__(self, encoding: str = "utf-8") -> None: + """Initialize the text parser. + + Args: + encoding (`str`, defaults to ``"utf-8"``): + The text encoding used to decode the file bytes. + """ + self.encoding = encoding + + async def parse( + self, + file: bytes | str, + filename: str, + ) -> list[Section]: + """Read the file as text and return a single :class:`Section`. + + Args: + file (`bytes | str`): + The file content. ``bytes`` is decoded with the + configured encoding. ``str`` is disambiguated at + runtime: if it names an existing file on disk the + file is read and decoded; otherwise it is used + verbatim as pre-decoded text — letting local-mode + callers skip the encode → decode round trip. + filename (`str`): + The source filename, copied verbatim into + :attr:`Section.source`. + + Returns: + `list[Section]`: + Always a one-element list containing the entire file + contents. + + Raises: + `ValueError`: If the bytes cannot be decoded with the + configured encoding. + """ + if isinstance(file, str): + if os.path.isfile(file): + with open(file, "rb") as fp: + raw = fp.read() + try: + text = raw.decode(self.encoding) + except UnicodeDecodeError as e: + raise ValueError( + f"Failed to decode {filename!r} as " + f"{self.encoding!r}: {e}", + ) from e + else: + text = file + else: + try: + text = file.decode(self.encoding) + except UnicodeDecodeError as e: + raise ValueError( + f"Failed to decode {filename!r} as " + f"{self.encoding!r}: {e}", + ) from e + + return [ + Section( + content=TextBlock(text=text), + source=filename, + metadata={}, + ), + ] diff --git a/src/agentscope/rag/_parser/_utils.py b/src/agentscope/rag/_parser/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..64545b0cdea97f0e2ccecfd091daa581aad9b244 --- /dev/null +++ b/src/agentscope/rag/_parser/_utils.py @@ -0,0 +1,99 @@ +# -*- coding: utf-8 -*- +"""Shared helpers for binary parsers. + +Three small utilities used by :class:`PDFParser`, :class:`ImageParser`, +and :class:`PPTParser`: + +- :func:`_guess_image_media_type` — sniff the IANA media type from raw + image bytes by looking at the magic number. Used to populate the + ``media_type`` field of :class:`Base64Source` for embedded images. +- :func:`_table_to_markdown` — render a 2-D ``list[list[str]]`` as a + Markdown pipe-table; the default rendering for table content. +- :func:`_table_to_json` — render the same shape as a JSON array + prefixed with a one-line ```` marker; used when the + caller picks ``table_format="json"`` to avoid Markdown's + multi-line-cell ambiguity. +""" +import json + + +def _guess_image_media_type(data: bytes) -> str: + """Sniff the IANA media type of an image from its magic number. + + Args: + data (`bytes`): + The raw image bytes. + + Returns: + `str`: + The IANA media type (e.g. ``"image/png"``). Falls back to + ``"image/jpeg"`` when no signature matches — JPEG is the + most permissive default and matches what every consumer can + decode. + """ + signatures = { + b"\x89PNG\r\n\x1a\n": "image/png", + b"\xff\xd8": "image/jpeg", + b"GIF87a": "image/gif", + b"GIF89a": "image/gif", + b"BM": "image/bmp", + } + for signature, media_type in signatures.items(): + if data.startswith(signature): + return media_type + # WebP: ``RIFF`` at offset 0 + ``WEBP`` at offset 8. + if len(data) > 12 and data[:4] == b"RIFF" and data[8:12] == b"WEBP": + return "image/webp" + return "image/jpeg" + + +def _table_to_markdown(table_data: list[list[str]]) -> str: + """Render a 2-D table as a Markdown pipe-table. + + Args: + table_data (`list[list[str]]`): + The table data; ``table_data[0]`` is the header row. + + Returns: + `str`: + The Markdown rendering, or the empty string when + ``table_data`` is empty / column-less. + """ + if not table_data: + return "" + + num_cols = len(table_data[0]) + if num_cols == 0: + return "" + + lines = [ + "| " + " | ".join(table_data[0]) + " |", + "| " + " | ".join(["---"] * num_cols) + " |", + ] + for row in table_data[1:]: + # Pad short rows so column counts match the header. + padded = list(row) + [""] * max(0, num_cols - len(row)) + lines.append("| " + " | ".join(padded[:num_cols]) + " |") + return "\n".join(lines) + "\n" + + +def _table_to_json(table_data: list[list[str]]) -> str: + """Render a 2-D table as JSON prefixed by a one-line marker. + + The ```` marker lets the LLM (and any + structure-aware downstream renderer) tell at a glance that this + block is tabular data, not free text — which is otherwise + indistinguishable from a raw JSON dump in the middle of a chunk. + + Args: + table_data (`list[list[str]]`): + The table data. + + Returns: + `str`: + ``"...\\n"``. + """ + return ( + "A table loaded as a JSON array:\n" + + json.dumps(table_data, ensure_ascii=False) + ) diff --git a/src/agentscope/rag/_vdb/__init__.py b/src/agentscope/rag/_vdb/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8c2f5d2aeb2952012b9bc9a77111a92a8eea118d --- /dev/null +++ b/src/agentscope/rag/_vdb/__init__.py @@ -0,0 +1,18 @@ +# -*- coding: utf-8 -*- +"""The vector store classes in AgentScope.""" + +from ._vector_store import ( + DocumentSummary, + VectorRecord, + VectorSearchResult, + VectorStoreBase, +) +from ._qdrant import QdrantStore + +__all__ = [ + "DocumentSummary", + "VectorStoreBase", + "VectorRecord", + "VectorSearchResult", + "QdrantStore", +] diff --git a/src/agentscope/rag/_vdb/_qdrant.py b/src/agentscope/rag/_vdb/_qdrant.py new file mode 100644 index 0000000000000000000000000000000000000000..731fb19ac1d20ec8e2856e63219b68dc5446f577 --- /dev/null +++ b/src/agentscope/rag/_vdb/_qdrant.py @@ -0,0 +1,392 @@ +# -*- coding: utf-8 -*- +"""Qdrant implementation of the vector store backend. + +Built on the official ``qdrant-client`` SDK using its fully +asynchronous client (:class:`~qdrant_client.AsyncQdrantClient`), so all +operations are non-blocking and safe to call from the application's +event loop. + +The same class supports all Qdrant deployment modes through the +constructor arguments: + +- ``location=":memory:"`` — in-process, ephemeral (ideal for tests) +- ``path="/path/to/db"`` — in-process, persisted to local disk +- ``url="http://localhost:6333"`` — remote Qdrant server / cloud +""" +import uuid +from typing import TYPE_CHECKING, Any, Literal + +from ._vector_store import ( + DocumentSummary, + VectorRecord, + VectorSearchResult, + VectorStoreBase, +) +from .._document import Chunk + +if TYPE_CHECKING: + from qdrant_client import AsyncQdrantClient + + +class QdrantStore(VectorStoreBase): + """Vector store backend backed by `Qdrant `_. + + Each knowledge base maps to one Qdrant collection. Every point + payload stores the owning ``document_id`` plus the serialized + :class:`~agentscope.rag.Chunk`, which is reconstructed on + retrieval. + + .. note:: The ``qdrant-client`` package is required. Install it + with ``pip install qdrant-client``. + + .. code-block:: python + + # In-memory (tests / prototyping) + store = QdrantStore(location=":memory:") + + # Remote server + store = QdrantStore( + url="http://localhost:6333", + api_key="...", + ) + + async with store: + await store.create_collection("kb-1", dimensions=768) + + """ + + def __init__( + self, + location: str | None = None, + url: str | None = None, + path: str | None = None, + api_key: str | None = None, + distance: Literal["Cosine", "Dot", "Euclid", "Manhattan"] = "Cosine", + client_kwargs: dict[str, Any] | None = None, + ) -> None: + """Initialize the Qdrant vector store. + + Args: + location (`str | None`, optional): + Pass ``":memory:"`` for an ephemeral in-process + instance. Mutually exclusive with ``url`` and ``path``. + url (`str | None`, optional): + The URL of a remote Qdrant server, e.g. + ``"http://localhost:6333"``. + path (`str | None`, optional): + A local directory for an in-process, on-disk instance. + api_key (`str | None`, optional): + The API key for Qdrant Cloud or a secured server. + distance (`Literal["Cosine", "Dot", "Euclid", "Manhattan"]`, \ + defaults to ``"Cosine"``): + The distance metric used when creating collections. + client_kwargs (`dict[str, Any] | None`, optional): + Extra keyword arguments forwarded to + :class:`~qdrant_client.AsyncQdrantClient`. + """ + self._location = location + self._url = url + self._path = path + self._api_key = api_key + self._distance = distance + self._client_kwargs = client_kwargs or {} + self._client: "AsyncQdrantClient | None" = None + + def get_client(self) -> "AsyncQdrantClient": + """Lazily create and cache the async Qdrant client. + + Returns: + `AsyncQdrantClient`: + The shared async client instance. + """ + if self._client is None: + from qdrant_client import AsyncQdrantClient + + self._client = AsyncQdrantClient( + location=self._location, + url=self._url, + path=self._path, + api_key=self._api_key, + **self._client_kwargs, + ) + return self._client + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: Any, + ) -> None: + """Exit the async context — close the underlying client.""" + if self._client is not None: + await self._client.close() + self._client = None + + # ------------------------------------------------------------------ + # Collection management + # ------------------------------------------------------------------ + + async def create_collection( + self, + name: str, + dimensions: int, + ) -> None: + """Create a new Qdrant collection. + + No-op if the collection already exists. + + Args: + name (`str`): + The collection name. Typically, the knowledge base ID. + dimensions (`int`): + The fixed vector dimensionality for this collection. + """ + from qdrant_client import models + + client = self.get_client() + if await client.collection_exists(name): + return + await client.create_collection( + collection_name=name, + vectors_config=models.VectorParams( + size=dimensions, + distance=models.Distance(self._distance), + ), + ) + + async def delete_collection(self, name: str) -> None: + """Delete a collection and all its data. + + Args: + name (`str`): + The collection name to delete. + """ + await self.get_client().delete_collection(name) + + async def has_collection(self, name: str) -> bool: + """Check whether a collection exists. + + Args: + name (`str`): + The collection name to check. + + Returns: + `bool`: ``True`` if the collection exists. + """ + return await self.get_client().collection_exists(name) + + # ------------------------------------------------------------------ + # Data operations + # ------------------------------------------------------------------ + + async def insert( + self, + collection: str, + records: list[VectorRecord], + ) -> None: + """Insert records into a collection. + + Each point payload stores the :attr:`VectorRecord.document_id` + under the ``document_id`` key and the serialized + :class:`Chunk` under the ``chunk`` key, so that :meth:`delete` + can remove all records of one document. + + Args: + collection (`str`): + The target collection name. + records (`list[VectorRecord]`): + The records to insert (each carrying a + :class:`Chunk` and its embedding vector). + """ + + from qdrant_client import models + + if not records: + return + await self.get_client().upsert( + collection_name=collection, + points=[ + models.PointStruct( + id=str(uuid.uuid4()), + vector=record.vector, + payload={ + "document_id": record.document_id, + "chunk": record.chunk.model_dump(mode="json"), + }, + ) + for record in records + ], + ) + + async def delete( + self, + collection: str, + document_id: str, + ) -> None: + """Delete all records belonging to one source document. + + Matches the ``document_id`` payload key written by + :meth:`insert` from :attr:`VectorRecord.document_id`. + + Args: + collection (`str`): + The target collection name. + document_id (`str`): + The source document ID whose records should be + removed. + """ + from qdrant_client import models + + await self.get_client().delete( + collection_name=collection, + points_selector=models.FilterSelector( + filter=models.Filter( + must=[ + models.FieldCondition( + key="document_id", + match=models.MatchValue(value=document_id), + ), + ], + ), + ), + ) + + # ------------------------------------------------------------------ + # Search + # ------------------------------------------------------------------ + + async def search( + self, + collection: str, + query_vector: list[float], + top_k: int = 5, + metadata_filter: dict[str, Any] | None = None, + ) -> list[VectorSearchResult]: + """Find the most similar records to a query vector. + + Args: + collection (`str`): + The collection to search. + query_vector (`list[float]`): + The query embedding vector. + top_k (`int`, defaults to ``5``): + Maximum number of results to return. + metadata_filter (`dict[str, Any] | None`, optional): + If provided, restrict the search to records whose + ``chunk.metadata`` matches every ``key == value`` pair + in this dict (translated into a Qdrant ``must`` payload + filter against ``chunk.metadata.``). + + Returns: + `list[VectorSearchResult]`: + Results ordered by descending similarity score. + """ + response = await self.get_client().query_points( + collection_name=collection, + query=query_vector, + limit=top_k, + with_payload=True, + query_filter=self._build_metadata_filter(metadata_filter), + ) + return [ + VectorSearchResult( + score=point.score, + document_id=point.payload["document_id"], + chunk=Chunk.model_validate(point.payload["chunk"]), + ) + for point in response.points + ] + + # ------------------------------------------------------------------ + # Document listing + # ------------------------------------------------------------------ + + async def list_documents( + self, + collection: str, + metadata_filter: dict[str, Any] | None = None, + ) -> list[DocumentSummary]: + """List all distinct source documents indexed in a collection. + + Scrolls the collection in payload-only mode (vectors disabled) + and aggregates by ``document_id``. The first chunk encountered + for each document supplies the ``source`` filename and the + document-level ``metadata``. + + Args: + collection (`str`): + The target collection name. + metadata_filter (`dict[str, Any] | None`, optional): + If provided, restrict aggregation to records whose + ``chunk.metadata`` matches every ``key == value`` pair. + + Returns: + `list[DocumentSummary]`: + One summary per distinct ``document_id``. + """ + client = self.get_client() + query_filter = self._build_metadata_filter(metadata_filter) + summaries: dict[str, DocumentSummary] = {} + offset: Any = None + + while True: + points, next_offset = await client.scroll( + collection_name=collection, + scroll_filter=query_filter, + limit=256, + offset=offset, + with_payload=True, + with_vectors=False, + ) + for point in points: + doc_id = point.payload["document_id"] + summary = summaries.get(doc_id) + if summary is None: + chunk_payload = point.payload["chunk"] + summaries[doc_id] = DocumentSummary( + document_id=doc_id, + source=chunk_payload.get("source", ""), + chunk_count=1, + metadata=dict(chunk_payload.get("metadata", {})), + ) + else: + summary.chunk_count += 1 + if next_offset is None: + break + offset = next_offset + + return list(summaries.values()) + + @staticmethod + def _build_metadata_filter( + metadata_filter: dict[str, Any] | None, + ) -> Any: + """Translate a flat ``{key: value}`` filter into a Qdrant filter. + + Each ``key`` is matched against the corresponding nested path + ``chunk.metadata.`` written by :meth:`insert`. Returns + ``None`` when ``metadata_filter`` is empty so that callers + skip the filter argument entirely. + + Args: + metadata_filter (`dict[str, Any] | None`): + The flat filter, or ``None`` for no filter. + + Returns: + `qdrant_client.models.Filter | None`: + A Qdrant ``Filter`` object, or ``None``. + """ + if not metadata_filter: + return None + + from qdrant_client import models + + return models.Filter( + must=[ + models.FieldCondition( + key=f"chunk.metadata.{key}", + match=models.MatchValue(value=value), + ) + for key, value in metadata_filter.items() + ], + ) diff --git a/src/agentscope/rag/_vdb/_vector_store.py b/src/agentscope/rag/_vdb/_vector_store.py new file mode 100644 index 0000000000000000000000000000000000000000..347f4ff77a09cb0b852be52154c3f496bf786ada --- /dev/null +++ b/src/agentscope/rag/_vdb/_vector_store.py @@ -0,0 +1,291 @@ +# -*- coding: utf-8 -*- +"""Abstract base class for vector store backends. + +A :class:`VectorStoreBase` instance is the single connection point to +one vector database deployment. It is created once at application +startup, passed into ``create_app(vector_store=...)``, and shared +across all requests for the lifetime of the process — similar to +:class:`~agentscope.app.storage.StorageBase` and +:class:`~agentscope.app.message_bus.MessageBus`. + +Each **knowledge base** maps to one **collection** inside the vector +store. Collections are isolated: different knowledge bases never +share a collection, so retrieval is always scoped to a single +collection without cross-collection filtering. + +Lifecycle is managed via the async context manager protocol +(``__aenter__`` / ``__aexit__``), which the app lifespan calls +automatically. +""" +from abc import ABC, abstractmethod +from typing import Any, Self + +from pydantic import BaseModel, Field + +from .._document import Chunk + + +class VectorRecord(BaseModel): + """A single record to insert into a vector store collection. + + Pairs a :class:`Chunk` (the business payload — content, source, + structural metadata) with its dense embedding vector. ``Chunk`` + is intentionally not extended with an ``embedding`` field so its + semantics stay stable across the indexing pipeline; instead the + vector lives in this wrapper whose only purpose is "I am about + to be inserted into a vector database." + """ + + vector: list[float] + """The dense embedding vector for :attr:`chunk`.""" + + document_id: str + """The ID of the source document this record belongs to. + Assigned by the knowledge base layer when the document is + registered. Backends must persist it at insertion time so that + :meth:`VectorStoreBase.delete` can remove all records of one + document as a unit.""" + + chunk: Chunk + """The business payload — content, source, structural metadata.""" + + +class VectorSearchResult(BaseModel): + """A single result returned by a similarity search. + + Pairs the matched :class:`Chunk` with its similarity score. + ``Chunk`` is intentionally not extended with a ``score`` field so + its semantics stay stable; instead the score lives in this + wrapper whose only purpose is "I am a query hit." + """ + + score: float + """Similarity score. Higher = more similar for cosine / + dot-product; lower = more similar for L2 distance.""" + + document_id: str + """The ID of the source document the matched chunk belongs to — + the same value carried by :attr:`VectorRecord.document_id` at + insertion time. Lets callers cite, group, or delete the source + document of a hit.""" + + chunk: Chunk + """The matched business payload.""" + + +class DocumentSummary(BaseModel): + """A lightweight description of one source document inside a collection. + + Aggregated by :meth:`VectorStoreBase.list_documents` from the + records of each ``document_id`` — the vector store is the single + source of truth for "what documents exist in a knowledge base". + """ + + document_id: str + """The source document identifier — the same value carried by + :attr:`VectorRecord.document_id` at insertion time.""" + + source: str + """The original filename, taken from the first chunk encountered. + All chunks of the same document share the same filename so any + chunk yields the same value.""" + + chunk_count: int + """The total number of chunks indexed for this document.""" + + metadata: dict[str, Any] = Field(default_factory=dict) + """Document-level metadata propagated from the parser / uploader + (media type, size, upload time, ...). Taken from the first chunk + encountered.""" + + +class VectorStoreBase(ABC): + """Abstract base class for vector store backends. + + Subclasses implement the concrete connection and query logic for a + specific vector database (Chroma, Milvus, Qdrant, FAISS, etc.). + + A single instance is shared across the entire application. The + underlying client SDK is expected to handle connection pooling and + thread safety internally. + """ + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def __aenter__(self) -> Self: + """Enter the async context — open connections if needed. + + The default implementation is a no-op. Subclasses that need + explicit connection setup should override this. + + Returns: + `VectorStoreBase`: ``self``. + """ + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: Any, + ) -> None: + """Exit the async context — close connections if needed. + + The default implementation is a no-op. + """ + + # ------------------------------------------------------------------ + # Collection management + # ------------------------------------------------------------------ + + @abstractmethod + async def create_collection( + self, + name: str, + dimensions: int, + ) -> None: + """Create a new collection (vector index). + + If the collection already exists, implementations should raise + or silently no-op depending on the backend's semantics. + + Args: + name (`str`): + The collection name. Typically, the knowledge base ID. + dimensions (`int`): + The fixed vector dimensionality for this collection. + All vectors inserted later must have this many elements. + """ + + @abstractmethod + async def delete_collection(self, name: str) -> None: + """Delete a collection and all its data. + + Args: + name (`str`): + The collection name to delete. + """ + + @abstractmethod + async def has_collection(self, name: str) -> bool: + """Check whether a collection exists. + + Args: + name (`str`): + The collection name to check. + + Returns: + `bool`: ``True`` if the collection exists. + """ + + # ------------------------------------------------------------------ + # Data operations + # ------------------------------------------------------------------ + + @abstractmethod + async def insert( + self, + collection: str, + records: list[VectorRecord], + ) -> None: + """Insert records into a collection. + + Args: + collection (`str`): + The target collection name. + records (`list[VectorRecord]`): + The records to insert (each carrying a + :class:`Chunk` and its embedding vector). + """ + + @abstractmethod + async def delete( + self, + collection: str, + document_id: str, + ) -> None: + """Delete all records belonging to one source document. + + Identifies records by the :attr:`VectorRecord.document_id` + field that backends persist at insertion time. This matches + the typical RAG workflow where a user uploads or removes a + file as a unit. + + Args: + collection (`str`): + The target collection name. + document_id (`str`): + The source document ID whose records should be + removed. + """ + + # ------------------------------------------------------------------ + # Search + # ------------------------------------------------------------------ + + @abstractmethod + async def search( + self, + collection: str, + query_vector: list[float], + top_k: int = 5, + metadata_filter: dict[str, Any] | None = None, + ) -> list[VectorSearchResult]: + """Find the most similar records to a query vector. + + Args: + collection (`str`): + The collection to search. + query_vector (`list[float]`): + The query embedding vector. + top_k (`int`, defaults to ``5``): + Maximum number of results to return. + metadata_filter (`dict[str, Any] | None`, optional): + If provided, restrict the search to records whose + ``chunk.metadata`` matches every ``key == value`` pair + in this dict. Backends translate this into a native + payload filter. Used for defense-in-depth + cross-tenant scoping when an isolation strategy + co-locates multiple knowledge bases inside the same + collection. + + Returns: + `list[VectorSearchResult]`: + Results ordered by descending similarity score. + """ + + # ------------------------------------------------------------------ + # Document listing + # ------------------------------------------------------------------ + + @abstractmethod + async def list_documents( + self, + collection: str, + metadata_filter: dict[str, Any] | None = None, + ) -> list[DocumentSummary]: + """List all distinct source documents indexed in a collection. + + Aggregates records by :attr:`VectorRecord.document_id` and + returns one :class:`DocumentSummary` per document. Backends + are free to use whatever scrolling / aggregation primitive + they expose; this method is expected to be O(documents) not + O(chunks) on backends that support payload-only scans. + + Args: + collection (`str`): + The target collection name. + metadata_filter (`dict[str, Any] | None`, optional): + If provided, restrict aggregation to records whose + ``chunk.metadata`` matches every ``key == value`` pair + in this dict. Used together with the search-time + filter when an isolation strategy co-locates multiple + knowledge bases inside the same collection. + + Returns: + `list[DocumentSummary]`: + One summary per distinct ``document_id``, in + unspecified order. + """ diff --git a/src/agentscope/skill/__init__.py b/src/agentscope/skill/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5c438a326d658c827d957ee25106463fc1f30625 --- /dev/null +++ b/src/agentscope/skill/__init__.py @@ -0,0 +1,11 @@ +# -*- coding: utf-8 -*- +"""The skill related classes and functions.""" + +from ._base import SkillLoaderBase, Skill +from ._local_loader import LocalSkillLoader + +__all__ = [ + "Skill", + "SkillLoaderBase", + "LocalSkillLoader", +] diff --git a/src/agentscope/skill/_base.py b/src/agentscope/skill/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..ac8cd8f5d593c2363eebf9d94e0dd9cd0dcb1895 --- /dev/null +++ b/src/agentscope/skill/_base.py @@ -0,0 +1,29 @@ +# -*- coding: utf-8 -*- +"""The skill loader base class.""" +from abc import abstractmethod, ABC +from dataclasses import dataclass + + +@dataclass +class Skill: + """The agent skill class""" + + name: str + """The name of the skill.""" + description: str + """The description of the skill.""" + dir: str + """The directory of the agent skill.""" + markdown: str + """The markdown content of the agent skill.""" + updated_at: float + """The last updated time of the skill.""" + + +class SkillLoaderBase(ABC): + """The base class for skill loader.""" + + @abstractmethod + async def list_skills(self) -> list[Skill]: + """List all the skills that can be loaded by this loader.""" + raise NotImplementedError diff --git a/src/agentscope/skill/_local_loader.py b/src/agentscope/skill/_local_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..6e739b0424ab4e9eb9639de703abcb8459348004 --- /dev/null +++ b/src/agentscope/skill/_local_loader.py @@ -0,0 +1,171 @@ +# -*- coding: utf-8 -*- +"""The local skill loader class.""" +import asyncio +import os + +import aiofiles +import aiofiles.ospath +import frontmatter + +from ._base import SkillLoaderBase +from .._logging import logger +from ..skill import Skill + + +class LocalSkillLoader(SkillLoaderBase): + """The skill loader that loads skills from a local directory.""" + + def __init__(self, directory: str, scan_subdir: bool = False) -> None: + """Initialize the loader with the directory. + + Args: + directory (`str`): + The directory to load skills from. + scan_subdir (`bool`, defaults to False): + Whether to scan subdirectories. Defaults to False (only + scan current directory). + """ + self.directory = os.path.abspath(directory) + self.scan_subdir = scan_subdir + self._cache: dict[str, Skill] = {} + + async def _load_single_skill(self, skill_root: str) -> Skill | None: + """Load a single skill from a skill root directory. + + Args: + skill_root (`str`): The skill root directory containing SKILL.md. + + Returns: + `Skill | None`: A Skill object or None if loading failed. + """ + skill_md_path = os.path.join(skill_root, "SKILL.md") + + try: + # Check if SKILL.md exists + if not await aiofiles.ospath.isfile(skill_md_path): + return None + + # Get file modification time + updated_at = await aiofiles.ospath.getmtime(skill_md_path) + + # Check cache: if cached skill exists and updated_at matches, + # return cached + if skill_root in self._cache: + cached_skill = self._cache[skill_root] + if cached_skill.updated_at == updated_at: + return cached_skill + + # Read and parse SKILL.md + async with aiofiles.open( + skill_md_path, + "r", + encoding="utf-8", + ) as f: + content_str = await f.read() + content = frontmatter.loads(content_str) + + name = content.get("name") + description = content.get("description") + + if not name or not description: + logger.warning( + "SKILL.md in %s is missing required fields " + "(name or description). Skipping.", + skill_root, + ) + return None + + skill = Skill( + name=str(name), + description=str(description), + dir=skill_root, + markdown=content.content, + updated_at=updated_at, + ) + + # Update cache + self._cache[skill_root] = skill + + return skill + + except Exception as e: + logger.warning( + "Failed to load skill from %s: %s", + skill_root, + str(e), + ) + return None + + async def list_skills(self) -> list[Skill]: + """List all the available skills from the directory. + + This method will: + 1. Search for SKILL.md in the current directory + 2. If scan_subdir is True, search for SKILL.md in all subdirectories + 3. Load all SKILL.md files concurrently + + Returns: + `list[Skill]`: A list of Skill objects. + """ + try: + # Check if directory exists + if not await aiofiles.ospath.isdir(self.directory): + logger.warning( + "Skill directory %s does not exist.", + self.directory, + ) + return [] + + # Find all directories containing SKILL.md + def _find_skill_dirs() -> list[str]: + """Find all directories containing SKILL.md file.""" + dirs = [] + + if os.path.isfile(os.path.join(self.directory, "SKILL.md")): + dirs.append(self.directory) + + if self.scan_subdir: + for root, _, filenames in os.walk(self.directory): + if root == self.directory: + continue + if "SKILL.md" in filenames: + dirs.append(root) + + return dirs + + skill_dirs = await asyncio.to_thread(_find_skill_dirs) + + if not skill_dirs: + logger.info( + "No SKILL.md files found in %s", + self.directory, + ) + return [] + + # Load all skills concurrently + tasks = [ + self._load_single_skill(skill_dir) for skill_dir in skill_dirs + ] + results = await asyncio.gather(*tasks, return_exceptions=True) + + # Filter out None results and exceptions + skills: list = [] + for i, result in enumerate(results): + if isinstance(result, Exception): + logger.warning( + "Failed to load skill from %s: %s", + skill_dirs[i], + str(result), + ) + elif result is not None: + skills.append(result) + + return skills + + except Exception as e: + logger.warning( + "Failed to list skills from directory %s: %s", + self.directory, + str(e), + ) + return [] diff --git a/src/agentscope/state/__init__.py b/src/agentscope/state/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..23e1ee09cd33a10f907f97dd2b3b1c0e66382320 --- /dev/null +++ b/src/agentscope/state/__init__.py @@ -0,0 +1,11 @@ +# -*- coding: utf-8 -*- +"""The agent state module in agentscope.""" + +from ._state import AgentState, TaskContext +from ._task import Task + +__all__ = [ + "Task", + "TaskContext", + "AgentState", +] diff --git a/src/agentscope/state/_state.py b/src/agentscope/state/_state.py new file mode 100644 index 0000000000000000000000000000000000000000..545de62d9b19ee10555008bba4b034ad3357affb --- /dev/null +++ b/src/agentscope/state/_state.py @@ -0,0 +1,221 @@ +# -*- coding: utf-8 -*- +"""The agent state class.""" +from typing import Any + +from pydantic import BaseModel, Field + +import aiofiles.os + +from .._utils._common import _generate_id +from ._task import Task +from ..message import ( + TextBlock, + DataBlock, + Msg, + ToolCallBlock, + ToolResultBlock, + HintBlock, +) +from ..permission import PermissionContext + + +class ReadCacheEntry(BaseModel): + """The read file cache.""" + + lines: list[str] + updated_at: float + bytes: float + file_path: str + + +class ToolContext(BaseModel): + """The tool context, e.g. tool cache""" + + max_cache_files: int = Field(default=100, gt=1) + """The maximum number of cached files.""" + max_cache_bytes: float = Field(default=25000, gt=10000) + """The maximum size of the accumulated read file cache.""" + read_file_cache: list[ReadCacheEntry] = Field(default_factory=list) + """The cache for Read/Write/Edit file tools.""" + + activated_groups: list[str] = Field(default_factory=list) + """The names of the activated tool groups, each group contains a set of + tools.""" + + async def get_cache(self, file_path: str) -> ReadCacheEntry | None: + """Get cached file content if still valid. + + Args: + file_path: The absolute path of the file. + + Returns: + The cached entry if valid, otherwise None. + """ + + # Find the cache entry + for entry in self.read_file_cache: + if entry.file_path == file_path: + # Check if cache is still valid + try: + updated_at = await aiofiles.os.path.getmtime(file_path) + if updated_at == entry.updated_at: + return entry + else: + # Cache is outdated, remove it + self.read_file_cache.remove(entry) + return None + except Exception: + # File might not exist anymore + self.read_file_cache.remove(entry) + return None + return None + + async def cache_file(self, file_path: str, lines: list[str]) -> None: + """Cache file content with LRU eviction. + + Args: + file_path: The absolute path of the file. + lines: The lines of the file content. + """ + try: + updated_at = await aiofiles.os.path.getmtime(file_path) + except Exception: + # Cannot get mtime, skip caching + return + + # Calculate size in KB + new_entry_bytes = ( + sum(len(line.encode("utf-8")) for line in lines) / 1024 + ) + + # Remove existing cache for this file if present + self.read_file_cache = [ + entry + for entry in self.read_file_cache + if entry.file_path != file_path + ] + + # Evict the oldest entries if exceeding max_cache_files + while len(self.read_file_cache) >= self.max_cache_files: + self.read_file_cache.pop(0) + + # Evict the oldest entries if exceeding max_cache_bytes + current_size = sum(entry.bytes for entry in self.read_file_cache) + while ( + self.read_file_cache + and current_size + new_entry_bytes > self.max_cache_bytes + ): + removed = self.read_file_cache.pop(0) + current_size -= removed.bytes + + # Add new entry to the end (most recent) + self.read_file_cache.append( + ReadCacheEntry( + lines=lines, + updated_at=updated_at, + bytes=new_entry_bytes, + file_path=file_path, + ), + ) + + async def clean_file_cache( + self, + reserved_file_paths: set[str] | None = None, + ) -> None: + """Drop read caches whose paths are not in ``reserved_file_paths``. + + Args: + reserved_file_paths: File paths from Read calls that remain in the + context. Caches for these files are kept; all others are + evicted. + """ + reserved_file_paths = reserved_file_paths or set() + + self.read_file_cache = [ + entry + for entry in self.read_file_cache + if entry.file_path in reserved_file_paths + ] + + +class TaskContext(BaseModel): + """The task context.""" + + tasks: list[Task] = Field(default_factory=lambda: []) + """The task context.""" + + +class AgentState(BaseModel): + """The agent state that should be saved and loaded from storage.""" + + session_id: str = Field(default_factory=_generate_id) + """The session id of the agent. Normally, each session will maintain one + independent agent state for each agent.""" + + summary: str | list[TextBlock | DataBlock] = "" + """The compressed summary of the context, which will be prepended to the + context when feed into the LLM.""" + context: list[Msg] = Field(default_factory=list) + """The uncompressed conversation context, that will be feed into the LLM""" + reply_id: str = Field(default_factory=_generate_id) + """The id of the current reply, which is also used as the id of the + final message of the reply.""" + cur_iter: int = 0 + """The current iteration of the agent's reasoning-acting loop.""" + + # ================================================================= + # The permission context + # ================================================================= + permission_context: PermissionContext = Field( + default_factory=PermissionContext, + ) + """The permission context that will be passed to the toolkit to determine + the tool permissions.""" + + # ================================================================= + # The tool context + # ================================================================= + tool_context: ToolContext = Field(default_factory=ToolContext) + + # ================================================================= + # The tasks context + # ================================================================= + tasks_context: TaskContext = Field(default_factory=TaskContext) + """The task context that records the agent tasks.""" + + # ================================================================= + # The middleware context + # ================================================================= + middle_context: dict[str, Any] = Field(default_factory=dict) + """The context that allow the middlewares to store/get data across + different replies.""" + + def append_context( + self, + name: str, + blocks: list[ + TextBlock | DataBlock | HintBlock | ToolCallBlock | ToolResultBlock + ], + ) -> None: + """Append the given blocks to the agent's own message with the current + `reply_id`. If such message doesn't exist, a new assistant message + with agent's name and current reply ID will be created. + """ + # If append to the latest message + if ( + self.context + and self.context[-1].role == "assistant" + and self.context[-1].name == name + and self.context[-1].id == self.reply_id + ): + self.context[-1].content.extend(blocks) + else: + # Create a new assistant message with the current reply ID + self.context.append( + Msg( + id=self.reply_id, + role="assistant", + name=name, + content=blocks, + ), + ) diff --git a/src/agentscope/state/_task.py b/src/agentscope/state/_task.py new file mode 100644 index 0000000000000000000000000000000000000000..aef442868897be6b6398b96dcf3bf11c780a55ef --- /dev/null +++ b/src/agentscope/state/_task.py @@ -0,0 +1,39 @@ +# -*- coding: utf-8 -*- +"""The task class.""" +from datetime import datetime +from typing import Any, Literal + +from pydantic import BaseModel, Field + +from .._utils._common import _generate_id + + +class Task(BaseModel): + """The agent task.""" + + subject: str + """The subject of the task.""" + + description: str + """The task description.""" + + metadata: dict[str, Any] + """The additional metadata of the task.""" + + created_at: str = Field(default_factory=lambda: datetime.now().isoformat()) + """The created timestamp.""" + + state: Literal["pending", "in_progress", "completed"] = "pending" + """The task state.""" + + id: str = Field(default_factory=_generate_id) + """The task identifier.""" + + owner: str | None = None + """The owner of the task.""" + + blocks: list[str] = Field(default_factory=lambda: []) + """The task ids blocked by this task.""" + + blocked_by: list[str] = Field(default_factory=lambda: []) + """The task ids blocking this task.""" diff --git a/src/agentscope/tool/__init__.py b/src/agentscope/tool/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b39bcb0f613bc4b9282f80015a3d3b73a766f8b3 --- /dev/null +++ b/src/agentscope/tool/__init__.py @@ -0,0 +1,58 @@ +# -*- coding: utf-8 -*- +"""The tool module in agentscope.""" + +from ._types import ToolChoice, Function, RegisteredTool +from ._response import ToolResponse, ToolChunk +from ._toolkit import Toolkit +from ._base import ToolBase, ParamsBase, ToolMiddlewareBase +from ._adapters import MCPTool, FunctionTool +from ._builtin import ( + ResetTools, + Bash, + Edit, + Glob, + Grep, + Read, + Write, + BackendBase, + ExecResult, + LocalBackend, +) +from ._task import ( + TaskUpdate, + TaskGet, + TaskList, + TaskCreate, +) +from ._tool_group import ToolGroup + +__all__ = [ + # Basic tool related types and functions + "ToolChoice", + "Function", + "ToolBase", + "ParamsBase", + "ToolMiddlewareBase", + "MCPTool", + "FunctionTool", + "ToolGroup", + "Toolkit", + "ToolChunk", + "ToolResponse", + "RegisteredTool", + # Builtin tools + "BackendBase", + "LocalBackend", + "ExecResult", + "ResetTools", + "Bash", + "Edit", + "Glob", + "Grep", + "Read", + "Write", + "TaskUpdate", + "TaskGet", + "TaskList", + "TaskCreate", +] diff --git a/src/agentscope/tool/_adapters.py b/src/agentscope/tool/_adapters.py new file mode 100644 index 0000000000000000000000000000000000000000..93ca74e553872c92f9b486fe0d9f8e07e12e3d79 --- /dev/null +++ b/src/agentscope/tool/_adapters.py @@ -0,0 +1,394 @@ +# -*- coding: utf-8 -*- +"""Adapters to convert functions and MCP tools to ToolProtocol.""" +import inspect +import json +import re +from contextlib import _AsyncGeneratorContextManager +from datetime import timedelta +from typing import Callable, Any, AsyncGenerator, Generator + +from mcp import ClientSession +import mcp + +from ._types import Function +from ._base import ToolBase, ToolMiddlewareBase +from ..permission import ( + PermissionBehavior, + PermissionDecision, +) +from ._response import ToolChunk +from ._utils import _extract_func_description, _extract_input_schema +from .._logging import logger +from ..message import ( + TextBlock, + DataBlock, + Base64Source, + URLSource, + ToolResultState, +) + + +class FunctionTool(ToolBase): + """Adapter to convert a Python function to ToolProtocol. + + This class wraps a regular Python function and makes it compatible with + the ToolProtocol interface. It automatically extracts metadata from the + function's signature and docstring, and normalizes the return value to + ToolChunk or AsyncGenerator[ToolChunk, None]. + """ + + is_external_tool: bool = False + """If this tool is an external tool, which doesn't need to implement the + __call__ method and the agent will yield the external tool call event.""" + is_mcp: bool = False + """If this tool is an MCP tool, which will be used in the permission""" + mcp_name: str | None = None + """The name of the MCP server this tool belongs to, which is required if + this tool is an MCP tool.""" + + def __init__( + self, + func: Function, + name: str | None = None, + description: str | None = None, + is_concurrency_safe: bool = True, + is_read_only: bool = False, + is_state_injected: bool = False, + middlewares: list[ToolMiddlewareBase] | None = None, + ) -> None: + """Initialize the FunctionTool. + + Args: + func (`Callable`): + The Python function to wrap. + name (`str | None`, optional): + Custom tool name. If None, uses the function name. + description (`str | None`, optional): + Custom tool description. If None, extracts from docstring. + is_concurrency_safe (`bool`, optional): + Whether this tool is safe to call concurrently. + is_read_only (`bool`, optional): + Whether this tool only reads data without side effects. + is_state_injected (`bool`, optional): + Whether this tool requires agent state injection. + middlewares (`list[ToolMiddlewareBase] | None`, optional): + Tool middlewares wrapping the tool execution. + """ + super().__init__(middlewares=middlewares) + self.name = name or func.__name__ + self.description = description or _extract_func_description( + func.__doc__ or "", + ) + self.input_schema = _extract_input_schema(func) + self.is_concurrency_safe = is_concurrency_safe + self.is_read_only = is_read_only + self.is_state_injected = is_state_injected + self.is_external_tool = False + self.is_mcp = False + self._func = func + + async def check_permissions( + self, + *_args: Any, + **_kwargs: Any, + ) -> PermissionDecision: + """Check permissions for the tool usage. + + Default implementation allows all operations. + + Returns: + `PermissionDecision`: + Permission decision (default: allowed). + """ + return PermissionDecision( + behavior=PermissionBehavior.ASK, + message="Custom function tools must be explicitly allowed " + "by the user.", + ) + + async def call( + self, + **kwargs: Any, + ) -> ToolChunk | AsyncGenerator[ToolChunk, None]: + """Invoke the wrapped function in an async style. + + Returns: + `ToolChunk` or `AsyncGenerator[ToolChunk, None]`: + The normalized result of the function execution. + """ + if inspect.iscoroutinefunction(self._func): + result = await self._func(**kwargs) + else: + result = self._func(**kwargs) + + if isinstance(result, AsyncGenerator): + + async def _stream() -> AsyncGenerator[ToolChunk, None]: + async for chunk in result: + if isinstance(chunk, ToolChunk): + yield chunk + else: + yield self._convert_func_result_to_chunk(chunk) + + return _stream() + + if isinstance(result, Generator): + + async def _stream() -> AsyncGenerator[ToolChunk, None]: + for chunk in result: + if isinstance(chunk, ToolChunk): + yield chunk + else: + yield self._convert_func_result_to_chunk(chunk) + + return _stream() + + return self._convert_func_result_to_chunk(result) + + @staticmethod + def _convert_func_result_to_chunk( + result: Any, + ) -> ToolChunk: + if isinstance(result, ToolChunk): + return result + if isinstance(result, str): + text = result + else: + try: + text = json.dumps(result, ensure_ascii=False) + except (TypeError, ValueError): + text = str(result) + return ToolChunk( + content=[TextBlock(text=text)], + state=ToolResultState.RUNNING, + ) + + +class MCPTool(ToolBase): + """Adapter to convert an MCP tool to ToolProtocol. + + This class wraps an MCP tool and makes it compatible with the ToolProtocol + interface. It handles the conversion between MCP's result format and + AgentScope's ToolChunk format. + """ + + is_mcp: bool = True + """Whether this tool is an MCP tool.""" + is_state_injected: bool = False + """The mcp tools is prohibited state injection for safety reason.""" + + def __init__( + self, + mcp_name: str, + tool: mcp.types.Tool, + client_gen: Callable[..., _AsyncGeneratorContextManager[Any]] + | None = None, + session: Any | None = None, + timeout: float | None = None, + middlewares: list[ToolMiddlewareBase] | None = None, + ) -> None: + """Initialize the MCPTool. + + Args: + mcp_name (`str`): + The name of the MCP server instance. + tool (`mcp.types.Tool`): + The MCP tool definition. + client_gen (`Callable[..., _AsyncGeneratorContextManager[Any]] \ + | None`, optional): + The MCP client generator function for stateless clients. + Either this or ``session`` must be provided. + session (`mcp.ClientSession | None`, optional): + The MCP client session for stateful clients. + Either this or ``client_gen`` must be provided. + timeout (`float | None`, optional): + The timeout in seconds for tool execution. + middlewares (`list[ToolMiddlewareBase] | None`, optional): + Tool middlewares wrapping the tool execution. + """ + super().__init__(middlewares=middlewares) + self.mcp_name = mcp_name + + # LLM providers enforce ^[a-zA-Z0-9_-]+$ on tool names. + # mcp_name is validated in MCPClient.model_post_init; + # tool.name comes from the MCP server and may contain dots, + # colons, etc. — replace illegal chars with "x" (not "_") + # to avoid collisions with the "__" separator. + # self._tool.name retains the original for server-side calls. + sanitized_tool = re.sub(r"[^a-zA-Z0-9_-]", "x", tool.name) + self.name = f"mcp__{mcp_name}__{sanitized_tool}" + if sanitized_tool != tool.name: + logger.debug( + "MCP tool name sanitized: '%s' -> '%s'.", + tool.name, + self.name, + ) + + self.description = tool.description or "" + + # Preserve the full inputSchema (including $defs, anyOf, oneOf, etc.) + # rather than only copying "properties" and "required", which would + # silently drop any nested type definitions that the LLM needs to + # resolve $ref pointers. + _schema = dict(tool.inputSchema) if tool.inputSchema else {} + _schema.setdefault("type", "object") + _schema.setdefault("properties", {}) + _schema.setdefault("required", []) + self.input_schema = _schema + + # By default + self.is_concurrency_safe = False + self.is_external_tool = False + + # Extract is_read_only from MCP tool annotations + self.is_read_only = False + if tool.annotations and hasattr(tool.annotations, "readOnlyHint"): + self.is_read_only = tool.annotations.readOnlyHint or False + + # Store MCP tool and connection info + self._tool = tool + self._client_gen = client_gen + self._session = session + + if timeout: + self._timeout = timedelta(seconds=timeout) + else: + self._timeout = None + + # Validate that either client_gen or session is provided + if (client_gen is None and session is None) or ( + client_gen is not None and session is not None + ): + raise ValueError( + "Either client_gen or session must be provided, but not both.", + ) + + async def check_permissions( + self, + *_args: Any, + **_kwargs: Any, + ) -> PermissionDecision: + """Check permissions for the MCP tool usage. + + Default implementation allows all operations. + + Returns: + `PermissionDecision`: + Permission decision (default: ask for confirmation). + """ + if self.is_read_only: + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message="This is a read-only MCP tool. Allowing execution.", + ) + return PermissionDecision( + behavior=PermissionBehavior.ASK, + message="MCP tools must be explicitly allowed by the user.", + ) + + async def call( + self, + **kwargs: Any, + ) -> ToolChunk: + """Invoke the MCP tool and convert the result to ToolChunk. + + Args: + **kwargs: Arguments to pass to the MCP tool. + + Returns: + `ToolChunk`: The converted tool execution result. + """ + + # Call the MCP tool + if self._client_gen: + # Stateless client: create temporary session + async with self._client_gen() as cli: + read_stream, write_stream = cli[0], cli[1] + async with ClientSession(read_stream, write_stream) as session: + await session.initialize() + result = await session.call_tool( + self._tool.name, + arguments=kwargs, + read_timeout_seconds=self._timeout, + ) + else: + # Stateful client: use existing session + result = await self._session.call_tool( + self._tool.name, + arguments=kwargs, + read_timeout_seconds=self._timeout, + ) + + # Convert MCP result to AgentScope blocks + return ToolChunk( + content=self._convert_mcp_content_to_blocks(result.content), + state=ToolResultState.ERROR + if result.isError + else ToolResultState.RUNNING, + ) + + @staticmethod + def _convert_mcp_content_to_blocks( + mcp_content_blocks: list, + ) -> list[TextBlock | DataBlock]: + """Convert MCP content to AgentScope blocks. + + Args: + mcp_content_blocks (`list`): + The MCP content blocks to convert. + + Returns: + `list[TextBlock | DataBlock]`: Converted AgentScope blocks. + """ + + as_content = [] + for content in mcp_content_blocks: + if isinstance(content, mcp.types.TextContent): + as_content.append(TextBlock(text=content.text)) + elif isinstance( + content, + (mcp.types.ImageContent, mcp.types.AudioContent), + ): + as_content.append( + DataBlock( + source=Base64Source( + type="base64", + media_type=content.mimeType, + data=content.data, + ), + ), + ) + + elif isinstance(content, mcp.types.EmbeddedResource): + if isinstance( + content.resource, + mcp.types.TextResourceContents, + ): + as_content.append( + TextBlock( + text=content.resource.model_dump_json(indent=2), + ), + ) + else: + logger.error( + "Unsupported EmbeddedResource content type: %s. " + "Skipping this content.", + type(content.resource), + ) + + elif isinstance(content, mcp.types.ResourceContents): + as_content.append( + DataBlock( + source=URLSource( + media_type=content.mimeType, + url=content.uri, + ), + ), + ) + + else: + logger.warning( + "Unsupported content type: %s. Skipping this content.", + type(content), + ) + return as_content diff --git a/src/agentscope/tool/_base.py b/src/agentscope/tool/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..bed2bb245b873d7ab0c4a4873f964a8bbb596bb0 --- /dev/null +++ b/src/agentscope/tool/_base.py @@ -0,0 +1,451 @@ +# -*- coding: utf-8 -*- +# pylint: disable=unused-argument +"""The tool protocol in agentscope.""" +import inspect +import os +from abc import abstractmethod, ABC +from pathlib import Path +from typing import AsyncGenerator, Any, Callable, List + +from pydantic import BaseModel + +from ._constants import DEFAULT_DANGEROUS_FILES, DEFAULT_DANGEROUS_DIRECTORIES +from ..permission import ( + PermissionContext, + PermissionDecision, + PermissionRule, + PermissionBehavior, +) +from ._response import ToolChunk +from ._utils import _remove_title_field + + +class ParamsBase(BaseModel): + """A base class for tool parameters that remove the title field from the + exported JSON schema. + """ + + @classmethod + def model_json_schema(cls, *args: Any, **kwargs: Any) -> dict: + """An override implementation to remove the title field from the + exported schema. + """ + return _remove_title_field(super().model_json_schema(*args, **kwargs)) + + +class ToolMiddlewareBase(ABC): + """Base class for tool middlewares. + + A tool middleware wraps the execution of a tool in an onion fashion: the + first registered middleware is the outermost layer and runs its pre-logic + before any inner layer, then its post-logic after all inner layers have + completed. Subclass this and implement :meth:`on_tool_call` — the signature + is already spelled out, so second-party developers only need to fill in the + body without reasoning about the wrapping protocol. + + Streaming and non-streaming tools are unified: ``next_handler`` always + returns an async generator, so a middleware never needs to know whether the + underlying tool yields a stream of chunks or returns a single chunk. + + Example: + ```python + class LoggingMiddleware(ToolMiddlewareBase): + async def on_tool_call(self, tool, input_kwargs, next_handler): + print(f"Calling {tool.name} with {input_kwargs}") + async for chunk in next_handler(**input_kwargs): + yield chunk + print(f"Finished {tool.name}") + + tool = MyTool(middlewares=[LoggingMiddleware()]) + ``` + """ + + @abstractmethod + async def on_tool_call( + self, + tool: "ToolBase", + input_kwargs: dict[str, Any], + next_handler: Callable[..., AsyncGenerator[ToolChunk, None]], + ) -> AsyncGenerator[ToolChunk, None]: + """Intercept a single tool invocation. + + Add pre-/post-logic around ``next_handler``, rewrite the tool inputs by + passing modified keyword arguments to ``next_handler``, or transform + the yielded chunks. + + Args: + tool (`ToolBase`): + The tool instance being invoked. + input_kwargs (`dict[str, Any]`): + The tool's input arguments for this invocation. Pass them on + via ``next_handler(**input_kwargs)``; mutate or replace them to + change what the inner layers and the tool itself receive. + next_handler (`Callable[..., AsyncGenerator[ToolChunk, None]]`): + Call it as ``next_handler(**input_kwargs)`` to run the next + layer. It always returns an async generator, regardless of + whether the underlying tool is streaming or not. + + Yields: + `ToolChunk`: + The chunks produced by this tool invocation. + """ + + +class ToolBase(ABC): + """The tool protocol.""" + + name: str + """The name presented to the agent.""" + description: str + """The agent-oriented tool description.""" + input_schema: dict[str, Any] + """The input schema of the tool, following JSON schema format.""" + is_concurrency_safe: bool + """If this tool is concurrency safe.""" + is_read_only: bool + """If this tool is read-only, which will be used in the permission + checking.""" + is_external_tool: bool = False + """If this tool is an external tool, which doesn't need to implement the + __call__ method and the agent will yield the external tool call event.""" + is_state_injected: bool = False + """If this tool requires agent state to be injected when called. If `True`, + the state will be injected by an argument named `_agent_state`. Note your + tool should be able to accept such argument. + """ + is_mcp: bool = False + """If this tool is an MCP tool, which will be used in the permission""" + mcp_name: str | None = None + """The name of the MCP server this tool belongs to, which is required if + this tool is an MCP tool.""" + + # Class attributes for dangerous path checking + dangerous_files: list[str] = DEFAULT_DANGEROUS_FILES + """List of dangerous files that should be protected from auto-editing.""" + dangerous_directories: list[str] = DEFAULT_DANGEROUS_DIRECTORIES + """List of dangerous directories that should be protected from + auto-editing.""" + + def __init__( + self, + middlewares: List["ToolMiddlewareBase"] | None = None, + ) -> None: + """Initialize the tool with optional middlewares. + + Args: + middlewares (`List[ToolMiddlewareBase] | None`, optional): + A list of :class:`ToolMiddlewareBase` instances wrapping the + tool execution in an onion fashion. Defaults to an empty list. + """ + self._middlewares: List["ToolMiddlewareBase"] = ( + middlewares if middlewares is not None else [] + ) + + async def call( + self, + *args: Any, + **kwargs: Any, + ) -> ToolChunk | AsyncGenerator[ToolChunk, None]: + """Execute the tool logic. + + This is the new override point for tool implementations. + Subclasses should override this method instead of + :meth:`__call__`. The base implementation raises + :exc:`NotImplementedError` for non-external tools and + :exc:`RuntimeError` for external tools. + + Args: + **kwargs: Tool input arguments. + + Returns: + `ToolChunk | AsyncGenerator[ToolChunk, None]`: + A single :class:`~agentscope.tool.ToolChunk` or an + async generator that yields them. + """ + if not self.is_external_tool: + raise NotImplementedError( + f"{self.__class__.__name__} does not implement call", + ) + + raise RuntimeError( + f"{self.__class__.__name__} is an external tool and should not " + f"be called directly", + ) + + async def __call__( + self, + *args: Any, + **kwargs: Any, + ) -> ToolChunk | AsyncGenerator[ToolChunk, None]: + """Invoke the tool, layering any registered middlewares around + :meth:`call`. + + Tools are always invoked with keyword arguments only. ``*args`` is + accepted in the signature solely to stay Liskov-compatible with + subclasses that override ``__call__`` with their own positional + parameters; any positional argument actually passed here is rejected + (raising :exc:`TypeError`) so it fails loudly instead of being silently + dropped. + + Middlewares are applied in an onion fashion: the first registered + middleware is the outermost layer and runs its pre-logic before + any inner layers, then its post-logic after all inner layers + have completed. + """ + if args: + raise TypeError( + f"{type(self).__name__} must be called with keyword arguments " + f"only, but got {len(args)} positional argument(s).", + ) + # ``getattr`` with a default so the no-middleware path keeps working + # even if a subclass overrides ``__init__`` without calling + # ``super().__init__()``. + middlewares = getattr(self, "_middlewares", []) + if not middlewares: + if inspect.isasyncgenfunction(self.call): + return self.call(**kwargs) + return await self.call(**kwargs) + + async def execute_chain( + index: int = 0, + **chain_kwargs: Any, + ) -> AsyncGenerator[ToolChunk, None]: + """Execute the tool middleware chain.""" + if index >= len(middlewares): + # Innermost layer: run the tool's own ``call``. ``call`` is + # always async but comes in two shapes — an async generator + # function (e.g. ``Bash``) or a coroutine returning a single + # ``ToolChunk`` / an async generator (e.g. ``FunctionTool``). + # Normalize both into a single stream so middlewares never have + # to distinguish them. + if inspect.isasyncgenfunction(self.call): + async for chunk in self.call(**chain_kwargs): + yield chunk + else: + result = await self.call(**chain_kwargs) + if isinstance(result, AsyncGenerator): + async for chunk in result: + yield chunk + else: + yield result + else: + mw = middlewares[index] + input_kwargs = dict(chain_kwargs) + + async def next_handler( + **kw: Any, + ) -> AsyncGenerator[ToolChunk, None]: + async for chunk in execute_chain(index + 1, **kw): + yield chunk + + async for chunk in mw.on_tool_call( + tool=self, + input_kwargs=input_kwargs, + next_handler=next_handler, + ): + yield chunk + + return execute_chain(**kwargs) + + @abstractmethod + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permissions for the tool usage.""" + + async def check_read_only( + self, + tool_input: dict[str, Any], + ) -> bool: + """Decide whether this specific invocation is read-only. + + Returns the static :attr:`is_read_only` attribute by default. + Subclasses with input-dependent semantics (e.g. ``Bash``) should + override this to inspect ``tool_input`` — for example, ``Bash`` is + statically marked as not read-only but ``ls -a`` is in fact read-only. + + Should be cheap — the permission engine may call this before the + full :meth:`check_permissions` flow. + + Args: + tool_input (`dict[str, Any]`): + The tool input data for this invocation. + + Returns: + `bool`: + ``True`` if this invocation is read-only, ``False`` otherwise. + """ + return self.is_read_only + + async def match_rule( + self, + rule_content: str | None, + tool_input: dict[str, Any], + ) -> bool: + """Check if a permission rule matches the tool input. + + .. note:: This is an optional method. A rule with no content (``None``) + is a tool-name-level rule that matches every invocation; a rule + with content requires the tool to override this method with its + own matching logic, otherwise it returns ``False``. + + This means: + - ``_FunctionTool`` and ``MCPTool`` (which do not override this) + can still be controlled at the tool-name level via rules like + ``{"tool_name": "my_tool", "rule_content": None}``. + - Specific tools (Bash, Read, Write, Edit, Glob, Grep) override + this method to support fine-grained pattern matching. + + Args: + rule_content (`str | None`): + The rule pattern to match. ``None`` means "match all + invocations of this tool" (tool-name-level rule). + tool_input (`dict[str, Any]`): + The tool input data + + Returns: + `bool`: + True if the rule matches, False otherwise + """ + # None rule_content = tool-name-level rule, matches everything + return rule_content is None + + async def generate_suggestions( + self, + tool_input: dict[str, Any], + ) -> List[PermissionRule]: + """Generate suggested permission rules for the tool input. + + .. note:: Suggest a single tool-name-level rule (``rule_content=None``) + that allows all invocations of this tool. Tools can override this to + provide finer-grained suggestions. + + For example: + - File tools (Read/Write/Edit): suggest a glob pattern covering the + parent directory (e.g., "src/main.py" -> "src/**") + - Bash: suggest command prefix patterns (e.g., "git commit -m 'xxx'" + -> "git commit:*") + - Grep/Glob: suggest patterns based on search paths + + Args: + tool_input (`dict[str, Any]`): + The tool input data + + Returns: + `List[PermissionRule]`: + List of suggested permission rules (usually 1, max 5 for + compound operations) + """ + return [ + PermissionRule( + tool_name=self.name, + rule_content=None, + behavior=PermissionBehavior.ALLOW, + source="suggested", + ), + ] + + def _path_in_allowed_working_path( + self, + file_path: str, + context: PermissionContext, + ) -> bool: + """Check if a file path is within any allowed working directory. + + A "working directory" is the process's current directory plus any + entries in :attr:`PermissionContext.working_directories`. Paths + are compared via :func:`os.path.realpath` so that aliases like + macOS's ``/tmp`` → ``/private/tmp`` and symlinked working + directories compare equal on both sides. + + Used by tools that conditionally auto-allow file operations in + :attr:`PermissionMode.ACCEPT_EDITS` (e.g. Write, Edit, and the + filesystem-command branch of Bash). + + Args: + file_path (`str`): + The file path to check. + context (`PermissionContext`): + The permission context containing the working directories. + + Returns: + `bool`: + True if ``file_path`` is within any allowed working + directory. + """ + current_dir = os.getcwd() + additional_dirs = list(context.working_directories.keys()) + all_working_dirs = [current_dir] + additional_dirs + + abs_file_path = os.path.realpath(os.path.expanduser(file_path)) + + for working_dir in all_working_dirs: + abs_working_dir = os.path.realpath( + os.path.expanduser(working_dir), + ) + try: + os.path.relpath(abs_file_path, abs_working_dir) + if ( + abs_file_path.startswith(abs_working_dir + os.sep) + or abs_file_path == abs_working_dir + ): + return True + except ValueError: + # On Windows, relpath raises ValueError if paths are on + # different drives. + continue + + return False + + def _is_dangerous_path(self, file_path: str) -> bool: + """Check if a file path is dangerous (sensitive file or directory). + + A path is considered dangerous if: + 1. The filename matches a dangerous file (e.g., .bashrc, .gitconfig) + 2. Any path segment matches a dangerous directory (e.g., .git, .ssh) + + Case-insensitive matching is used to prevent bypasses on + case-insensitive filesystems (macOS, Windows). + + Args: + file_path (`str`): + The file path to check + + Returns: + `bool`: + True if the path is dangerous and should require explicit + permission + + Example: + >>> self._is_dangerous_path("/home/user/.bashrc") + True + >>> self._is_dangerous_path("/home/user/.git/config") + True + >>> self._is_dangerous_path("/home/user/project/main.py") + False + """ + + # Normalize path + abs_path = os.path.abspath(os.path.expanduser(file_path)) + + # Split path into segments + path_parts = Path(abs_path).parts + path_parts_lower = [p.lower() for p in path_parts] + + # Check if filename matches dangerous files (case-insensitive) + filename = os.path.basename(abs_path) + filename_lower = filename.lower() + for dangerous_file in self.dangerous_files: + if filename_lower == dangerous_file.lower(): + return True + + # Check if any path segment matches dangerous directories + # (case-insensitive) + for dangerous_dir in self.dangerous_directories: + dangerous_dir_lower = dangerous_dir.lower() + if dangerous_dir_lower in path_parts_lower: + return True + + return False diff --git a/src/agentscope/tool/_builtin/__init__.py b/src/agentscope/tool/_builtin/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e6daefa4de6ebf6a618db4f38669ae99dd217168 --- /dev/null +++ b/src/agentscope/tool/_builtin/__init__.py @@ -0,0 +1,26 @@ +# -*- coding: utf-8 -*- +"""The builtin tools in agentscope.""" + +from ._backend import BackendBase, ExecResult, LocalBackend +from ._bash import Bash +from ._edit import Edit +from ._glob import Glob +from ._grep import Grep +from ._meta import ResetTools +from ._read import Read +from ._skill import SkillViewer +from ._write import Write + +__all__ = [ + "ResetTools", + "SkillViewer", + "Bash", + "Edit", + "Glob", + "Grep", + "Read", + "Write", + "BackendBase", + "LocalBackend", + "ExecResult", +] diff --git a/src/agentscope/tool/_builtin/_backend.py b/src/agentscope/tool/_builtin/_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..cbbc3f19f5e41fd35cbd8546f71fd0defc23824b --- /dev/null +++ b/src/agentscope/tool/_builtin/_backend.py @@ -0,0 +1,733 @@ +# -*- coding: utf-8 -*- +"""Backend abstraction for builtin tools. + +Provides a :class:`BackendBase` abstract base class that captures the +core I/O primitives shared across all six builtin tools (Bash, Read, +Write, Edit, Grep, Glob). + +Every backend implements exactly **three** abstract primitives whose +mechanism genuinely differs per environment: + +* :meth:`BackendBase.exec_shell` — run a program from an argv list + (no shell; callers needing shell features wrap with ``sh -c``). +* :meth:`BackendBase.read_file` — read raw bytes. +* :meth:`BackendBase.write_file` — write raw bytes. + +All remaining filesystem operations (``file_exists``, ``is_dir``, +``list_dir``, ``stat_mtime``, ``delete_path``) are derived on the base +class from ``exec_shell`` and work out-of-the-box for any remote +backend. A backend that has a cheaper native path (e.g. +:class:`LocalBackend` using ``os.*``) simply overrides them. + +Concrete implementations: + +* :class:`LocalBackend` — default; uses ``asyncio`` subprocesses, + ``aiofiles``, and ``os.*`` for host-local I/O. Injected automatically + when no explicit backend is given. +* ``DockerBackend`` — uses ``aiodocker`` exec / archive APIs. +* ``E2BBackend`` — uses the E2B SDK ``commands`` / ``files`` APIs. + +By accepting a ``BackendBase`` parameter, each builtin tool can +operate identically in local, Docker, and E2B workspaces without any +workspace-specific branching inside the tool code itself. +""" + +from __future__ import annotations + +import asyncio +import os +import posixpath +import shlex +import shutil +from abc import ABC, abstractmethod +from dataclasses import dataclass +from types import ModuleType +from typing import Any + +import aiofiles + +# ── data class ───────────────────────────────────────────────────────── + + +@dataclass(frozen=True, slots=True) +class ExecResult: + """Result of running a shell command via a backend. + + Attributes: + exit_code: Process exit code. ``-1`` conventionally indicates + an internal failure (timeout, connection error, …). + stdout: Raw bytes captured from standard output. + stderr: Raw bytes captured from standard error. + """ + + exit_code: int + stdout: bytes + stderr: bytes + + def ok(self) -> bool: + """Whether the command exited successfully. + + Returns: + `bool`: + ``True`` iff the command exited with code ``0``. + """ + return self.exit_code == 0 + + +# ── helpers ──────────────────────────────────────────────────────────── + + +def _normalize_newlines(text: str) -> str: + """Normalize Windows/old-Mac line endings to ``\\n``. + + Converts ``\\r\\n`` (Windows) and lone ``\\r`` (classic Mac) to a + single ``\\n``. Builtin tools read files as raw bytes (so binary + payloads survive intact); when the bytes are decoded as text for + line-based caching, editing, or matching, the line endings must be + normalized so that content written on Windows behaves identically to + content written on POSIX. + + Args: + text (`str`): + Decoded file contents. + + Returns: + `str`: + The text with all line endings collapsed to ``\\n``. + """ + return text.replace("\r\n", "\n").replace("\r", "\n") + + +# ── base class ───────────────────────────────────────────────────────── + + +class BackendBase(ABC): + """Filesystem + subprocess interface consumed by builtin tools. + + Subclasses must implement three abstract primitives — ``exec_shell``, + ``read_file``, and ``write_file`` — which are the only operations + whose mechanism genuinely differs per environment. The remaining + filesystem helpers are implemented here on top of ``exec_shell`` and + work for any backend whose shell is POSIX-like; a backend with a + cheaper native path may override them (see :class:`LocalBackend`). + """ + + #: Path-manipulation module whose semantics match the backend's + #: environment. Used by :meth:`join_path`, :meth:`dirname`, + #: :meth:`isabs`, :meth:`normpath`, and :meth:`abspath` to ensure + #: correct behavior when the host OS and the backend OS differ + #: (e.g. a Windows host driving a Linux Docker container). + #: + #: Defaults to :mod:`posixpath`, which is correct for any backend + #: whose environment is Linux/macOS (Docker, E2B, …). Subclasses + #: targeting a different environment override this attribute, e.g. + #: :class:`LocalBackend` sets it to :mod:`os.path`, and a future + #: Windows-container backend would set it to :mod:`ntpath`. + #: + #: .. important:: + #: + #: Only **pure string operations** on this module are safe to + #: call (``join``, ``split``, ``dirname``, ``basename``, + #: ``normpath``, ``isabs``, ``splitext``, ``splitdrive``, …). + #: + #: Do **not** call functions that touch the filesystem or + #: environment variables — ``exists``, ``isfile``, ``isdir``, + #: ``getmtime``, ``realpath``, ``expanduser``, ``expandvars``, + #: and parameterless ``abspath`` — because those read the + #: **host** process's filesystem / ``$HOME`` / ``cwd``, which + #: is meaningless (and a silent bug) for remote backends. Use + #: the async I/O methods on the backend instead + #: (:meth:`file_exists`, :meth:`is_dir`, :meth:`stat_mtime`, + #: …) or the :meth:`abspath` wrapper below, which requires an + #: explicit ``cwd`` argument. + _path_module: ModuleType = posixpath + + # ── path manipulation helpers (pure string ops) ──────────────── + + def join_path(self, path: str, *paths: str) -> str: + """Join one or more path components using the backend's separator. + + Args: + path (`str`): + The first path component. + *paths (`str`): + Additional components to join onto ``path``. + + Returns: + `str`: + The joined path, using the backend environment's path + separator. + """ + return self._path_module.join(path, *paths) + + def dirname(self, path: str) -> str: + """Return the directory component of ``path``. + + Args: + path (`str`): + A path inside the backend's environment. + + Returns: + `str`: + Everything up to (but not including) the last path + separator. Empty string if ``path`` has no separator. + """ + return self._path_module.dirname(path) + + def isabs(self, path: str) -> bool: + """Return ``True`` if ``path`` is absolute in the backend. + + Args: + path (`str`): + A path inside the backend's environment. + + Returns: + `bool`: + ``True`` iff ``path`` is absolute under the backend + environment's path semantics. + """ + return self._path_module.isabs(path) + + def normpath(self, path: str) -> str: + """Normalize ``path`` (collapse ``..``, ``.``, duplicate seps). + + Pure string operation — does not touch the filesystem. + + Args: + path (`str`): + A path inside the backend's environment. + + Returns: + `str`: + The normalized path. + """ + return self._path_module.normpath(path) + + def abspath(self, path: str, *, cwd: str) -> str: + """Return an absolute, normalized version of ``path``. + + Unlike :func:`os.path.abspath`, this helper **never** reads + the host process's working directory: when ``path`` is + relative it is joined with the explicitly supplied ``cwd``, + which must itself be a path that is meaningful inside the + backend's environment. This avoids the silent bug where the + host's ``os.getcwd()`` leaks into paths that will actually be + used on a remote backend. + + Args: + path (`str`): + A path inside the backend's environment. + cwd (`str`): + Directory to resolve a relative ``path`` against. + Ignored when ``path`` is already absolute. + + Returns: + `str`: + An absolute, normalized path. + """ + if self._path_module.isabs(path): + return self._path_module.normpath(path) + return self._path_module.normpath( + self._path_module.join(cwd, path), + ) + + # ── abstract primitives ──────────────────────────────────────── + + @abstractmethod + async def exec_shell( + self, + command: list[str], + *, + cwd: str | None = None, + timeout: float | None = None, + ) -> ExecResult: + """Run a program directly from an argument vector. + + *command* is an executable followed by its arguments — it is + **not** passed through a shell, so callers never have to quote + or escape arguments and there is no platform-specific quoting + bug. This makes the primitive portable to Windows, where POSIX + single-quote escaping (``shlex.quote``) is not understood by + ``cmd.exe``. + + Callers that genuinely need shell features (pipes, redirects, + ``&&``) must wrap their command line explicitly, e.g. + ``["/bin/sh", "-c", command_line]``. + + Args: + command (`list[str]`): + Executable path/name followed by its arguments. + cwd (`str | None`, optional): + Working directory to run the command in. When ``None`` + the backend's default working directory is used. + timeout (`float | None`, optional): + Maximum number of seconds to wait. When ``None`` the + call waits indefinitely. On timeout the result carries + an ``exit_code`` of ``-1``. + + Returns: + `ExecResult`: + The captured exit code, stdout, and stderr. + """ + + @abstractmethod + async def read_file(self, path: str) -> bytes: + """Read the full contents of ``path`` as raw bytes. + + Args: + path (`str`): + Path to the file inside the backend's environment. + + Returns: + `bytes`: + The raw file contents. + """ + + @abstractmethod + async def write_file(self, path: str, data: bytes) -> None: + """Write ``data`` to ``path``, creating parent directories. + + Args: + path (`str`): + Destination path inside the backend's environment. + data (`bytes`): + The raw bytes to write. + """ + + # ── derived filesystem ops (shell-based defaults) ────────────── + + async def getcwd(self) -> str: + """Return the backend environment's current working directory. + + This is the directory that bare ``exec_shell`` invocations + (those with ``cwd=None``) execute in. Tools should call this + — instead of :func:`os.getcwd` — whenever they need a default + path that is meaningful inside the backend, because the host + process's cwd is meaningless for remote backends + (Docker / E2B). + + The default implementation runs ``pwd`` via :meth:`exec_shell`, + which works for any POSIX-like backend. Backends with cheaper + native access (e.g. :class:`LocalBackend`, or remote backends + that already track their workdir) should override it. + + Returns: + `str`: + The backend's current working directory. + """ + result = await self.exec_shell(["pwd"]) + return result.stdout.decode("utf-8", errors="replace").strip() + + async def expanduser(self, path: str) -> str: + """Expand a leading ``~`` / ``~/`` to the backend's home directory. + + Tools should call this — instead of :func:`os.path.expanduser` + — whenever they need to expand ``~`` in a path that lives + inside the backend's environment, because the host process's + ``$HOME`` is meaningless for remote backends. + + The default implementation queries ``$HOME`` via + :meth:`exec_shell` (POSIX-only). Only the leading ``~`` / + ``~/foo`` form is expanded; ``~user/...`` is not supported by + the default and is returned unchanged. Backends with cheaper + native access should override (e.g. :class:`LocalBackend`). + + Args: + path (`str`): + A path inside the backend's environment, possibly + starting with ``~``. + + Returns: + `str`: + ``path`` with a leading ``~`` / ``~/`` expanded. If + ``path`` does not start with ``~``, or starts with + ``~user`` (unsupported), it is returned unchanged. + """ + if not path or path[0] != "~": + return path + # ``~user/...`` form — not supported by the default impl. + if len(path) > 1 and path[1] not in ("/", self._path_module.sep): + return path + result = await self.exec_shell(["printenv", "HOME"]) + home = result.stdout.decode("utf-8", errors="replace").strip() + if not home: + return path + return home + path[1:] + + async def file_exists(self, path: str) -> bool: + """Return ``True`` if ``path`` exists (file or directory). + + Args: + path (`str`): + Path to test inside the backend's environment. + + Returns: + `bool`: + ``True`` if the path exists, ``False`` otherwise. + """ + result = await self.exec_shell(["test", "-e", path]) + return result.ok() + + async def is_dir(self, path: str) -> bool: + """Return ``True`` if ``path`` is an existing directory. + + Args: + path (`str`): + Path to test inside the backend's environment. + + Returns: + `bool`: + ``True`` if the path is an existing directory. + """ + result = await self.exec_shell(["test", "-d", path]) + return result.ok() + + async def list_dir( + self, + path: str, + *, + recursive: bool = False, + ) -> list[str]: + """List entries under ``path``. + + Output is NUL-delimited (``find -print0`` / ``-printf '%f\\0'``) + and split on ``\\0`` so that file names containing spaces or + newlines are handled correctly. ``find -printf`` is a GNU + extension; backends running on non-GNU userlands should override + this method. + + Args: + path (`str`): + Directory to list inside the backend's environment. + recursive (`bool`, optional): + When ``True``, return all files underneath ``path`` as + paths (like ``find path -type f``). When ``False`` + (default), return the immediate children's base names + (like ``os.listdir``). + + Returns: + `list[str]`: + The matched entries, or an empty list if ``path`` does + not exist or cannot be listed. + """ + if recursive: + command = ["find", path, "-type", "f", "-print0"] + else: + command = [ + "find", + path, + "-mindepth", + "1", + "-maxdepth", + "1", + "-printf", + "%f\\0", + ] + result = await self.exec_shell(command) + if not result.ok(): + return [] + return [ + part.decode("utf-8", errors="surrogateescape") + for part in result.stdout.split(b"\0") + if part + ] + + async def stat_mtime(self, path: str) -> float | None: + """Return the modification time of ``path``, or ``None``. + + Tries GNU ``stat -c %Y`` first and falls back to BSD + ``stat -f %m`` so the same call works across coreutils and + BSD/macOS userlands. The two attempts are combined with ``||``, + so this default wraps a ``sh -c`` script; backends without a + POSIX shell (e.g. :class:`LocalBackend`) override it. + + Args: + path (`str`): + Path to stat inside the backend's environment. + + Returns: + `float | None`: + The modification time as a POSIX timestamp, or ``None`` + if the path does not exist or cannot be stat'd. + """ + quoted = shlex.quote(path) + script = ( + f"stat -c %Y {quoted} 2>/dev/null || " + f"stat -f %m {quoted} 2>/dev/null" + ) + result = await self.exec_shell(["sh", "-c", script]) + if not result.ok(): + return None + try: + return float( + result.stdout.decode("utf-8", errors="replace").strip(), + ) + except ValueError: + return None + + async def delete_path(self, path: str) -> None: + """Delete ``path`` (file or directory tree). + + If ``path`` does not exist the call is a silent no-op (like + ``rm -rf``). Handles both files and directories (recursively). + + Args: + path (`str`): + Path to delete inside the backend's environment. + """ + await self.exec_shell(["rm", "-rf", path]) + + +# ── local backend ────────────────────────────────────────────────────── + + +def _subprocess_creation_kwargs() -> dict[str, Any]: + """Return platform-specific subprocess creation options. + + Returns: + `dict[str, Any]`: + Extra keyword arguments for ``create_subprocess_shell``. + Empty on POSIX; on Windows it sets ``creationflags`` to + suppress a console window. + """ + if os.name != "nt": + return {} + + import subprocess + + return { + "creationflags": getattr( + subprocess, + "CREATE_NO_WINDOW", + 0x08000000, + ), + } + + +class LocalBackend(BackendBase): + """Host-local :class:`BackendBase` implementation. + + Uses ``asyncio.create_subprocess_exec``, ``aiofiles``, and the + ``os`` module. This is the default backend injected when no + explicit one is given to a builtin tool. Commands are spawned + directly from their argument vector (no shell), which avoids the + POSIX-vs-``cmd.exe`` quoting mismatch and makes the backend work on + Windows. The derived filesystem helpers are overridden with native + ``os.*`` calls — faster and more robust than shelling out, and + portable to Windows where ``test`` / ``find`` / ``stat`` are + unavailable. + """ + + # Use the host OS's path semantics (Windows or POSIX) instead of + # the base class default (``posixpath``), so path helpers behave + # correctly when running on a Windows host. + _path_module = os.path + + async def exec_shell( + self, + command: list[str], + *, + cwd: str | None = None, + timeout: float | None = None, + ) -> ExecResult: + """Run a program via ``asyncio.create_subprocess_exec``. + + The program is spawned directly from *command* without an + intervening shell, so no argument quoting is required and the + same code path works on POSIX and Windows. + + Args: + command (`list[str]`): + Executable path/name followed by its arguments. + cwd (`str | None`, optional): + Working directory for the subprocess. When ``None`` the + current process working directory is used. + timeout (`float | None`, optional): + Maximum number of seconds to wait before the process is + killed and an ``exit_code`` of ``-1`` is returned. + + Returns: + `ExecResult`: + The captured exit code, stdout, and stderr. If the + executable cannot be found or spawned, ``exit_code`` is + ``127`` (matching a shell's "command not found"), with + the OS error message on stderr. + """ + kwargs = _subprocess_creation_kwargs() + if cwd is not None: + kwargs["cwd"] = cwd + + try: + process = await asyncio.create_subprocess_exec( + *command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + **kwargs, + ) + except (FileNotFoundError, NotADirectoryError, OSError) as exc: + # The executable could not be found or spawned. A shell would + # have returned 127 ("command not found"); mirror that so + # callers see a normal non-zero ExecResult instead of an + # exception. + return ExecResult( + exit_code=127, + stdout=b"", + stderr=str(exc).encode("utf-8"), + ) + + try: + stdout, stderr = await asyncio.wait_for( + process.communicate(), + timeout=timeout, + ) + except asyncio.TimeoutError: + process.kill() + await process.communicate() + return ExecResult(exit_code=-1, stdout=b"", stderr=b"timed out") + + return ExecResult( + exit_code=process.returncode or 0, + stdout=stdout, + stderr=stderr, + ) + + async def read_file(self, path: str) -> bytes: + """Read a local file as raw bytes. + + Args: + path (`str`): + Path to the local file. + + Returns: + `bytes`: + The raw file contents. + """ + async with aiofiles.open(path, mode="rb") as f: + return await f.read() + + async def write_file(self, path: str, data: bytes) -> None: + """Write *data* to a local file, creating parent dirs. + + Args: + path (`str`): + Destination path on the local filesystem. + data (`bytes`): + The raw bytes to write. + """ + parent = os.path.dirname(path) + if parent: + os.makedirs(parent, exist_ok=True) + async with aiofiles.open(path, mode="wb") as f: + await f.write(data) + + async def getcwd(self) -> str: + """Return the host process's current working directory. + + Returns: + `str`: + ``os.getcwd()`` — avoids spawning a ``pwd`` subprocess. + """ + return os.getcwd() + + async def expanduser(self, path: str) -> str: + """Expand ``~`` using the host process's ``$HOME``. + + Args: + path (`str`): + A local path, possibly starting with ``~``. + + Returns: + `str`: + ``os.path.expanduser(path)`` — avoids spawning a + subprocess. + """ + return os.path.expanduser(path) + + async def file_exists(self, path: str) -> bool: + """Check if a local path exists. + + Args: + path (`str`): + Path to test. + + Returns: + `bool`: + ``True`` if the path exists. + """ + return os.path.exists(path) + + async def is_dir(self, path: str) -> bool: + """Check if a local path is a directory. + + Args: + path (`str`): + Path to test. + + Returns: + `bool`: + ``True`` if the path is an existing directory. + """ + return os.path.isdir(path) + + async def list_dir( + self, + path: str, + *, + recursive: bool = False, + ) -> list[str]: + """List local directory entries. + + Mirrors the base contract using native ``os`` calls. + + Args: + path (`str`): + Directory to list. + recursive (`bool`, optional): + When ``True``, return file paths underneath ``path`` + (``os.walk``). When ``False`` (default), return the + immediate children's base names (``os.listdir``). + + Returns: + `list[str]`: + The matched entries. + """ + if recursive: + results: list[str] = [] + for root, _dirs, files in os.walk(path): + for f in files: + results.append(os.path.join(root, f)) + return results + return os.listdir(path) + + async def stat_mtime(self, path: str) -> float | None: + """Return the modification time of a local file. + + Args: + path (`str`): + Path to stat. + + Returns: + `float | None`: + The modification time as a POSIX timestamp, or ``None`` + if the path does not exist or cannot be stat'd. + """ + try: + return os.stat(path).st_mtime + except (OSError, FileNotFoundError): + return None + + async def delete_path(self, path: str) -> None: + """Delete a local file or directory tree. + + No-op if *path* does not exist. + + Args: + path (`str`): + Path to delete. + """ + if not os.path.exists(path): + return + if os.path.isdir(path): + shutil.rmtree(path) + else: + os.remove(path) diff --git a/src/agentscope/tool/_builtin/_bash.py b/src/agentscope/tool/_builtin/_bash.py new file mode 100644 index 0000000000000000000000000000000000000000..3cb303cc75433773e5beefcf57f671c0f0286f58 --- /dev/null +++ b/src/agentscope/tool/_builtin/_bash.py @@ -0,0 +1,778 @@ +# -*- coding: utf-8 -*- +"""The bash tool in agentscope.""" +import os +from typing import AsyncGenerator, Any, List +import re + +from ._bash_parser import BashCommandParser +from .._base import ToolBase, ToolMiddlewareBase +from .._constants import ( + DEFAULT_DANGEROUS_FILES, + DEFAULT_DANGEROUS_DIRECTORIES, +) +from ...permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, + PermissionMode, + PermissionRule, +) +from ...message import TextBlock, ToolResultState +from .._response import ToolChunk +from ._backend import BackendBase + + +class Bash(ToolBase): + """The bash tool.""" + + name: str = "Bash" + """The tool name presented to the agent.""" + + description: str = """Executes a bash command and returns its output. + +The working directory persists between commands, but shell state does +not. The shell environment is initialized from the user's profile +(bash or zsh). + +IMPORTANT: Avoid using this tool to run `find`, `grep`, `cat`, `head`, +`tail`, `sed`, `awk`, or `echo` commands, unless explicitly instructed +or after you have verified that a dedicated tool cannot accomplish your +task. Instead, use the appropriate dedicated tool as this will provide +a much better experience for the user: + + - File search: Use Glob (NOT find or ls) + - Content search: Use Grep (NOT grep or rg) + - Read files: Use Read (NOT cat/head/tail) + - Edit files: Use Edit (NOT sed/awk) + - Write files: Use Write (NOT echo >/cat < None: + """Initialize the bash tool. + + Args: + dangerous_files (`list[str]`, optional): + Sensitive files that require explicit user confirmation, + even in BYPASS mode. Matched by basename + (case-insensitive). Defaults to `DEFAULT_DANGEROUS_FILES`. + Pass a custom list to fully replace the defaults, or `[]` + to disable the filename check. + dangerous_directories (`list[str]`, optional): + Sensitive directories that require explicit user + confirmation. Matched when any path segment equals an + entry (case-insensitive). Defaults to + `DEFAULT_DANGEROUS_DIRECTORIES`. Pass a custom list to + fully replace the defaults, or `[]` to disable the + directory check. + cwd (`str | os.PathLike[str] | None`, optional): + The working directory used when executing bash commands. + middlewares (`List[ToolMiddlewareBase] | None`, optional): + Tool middlewares wrapping the tool execution. + backend (`BackendBase | None`, optional): + The sandbox backend to use for shell execution. When + ``None``, a :class:`LocalBackend` is created. + """ + from ._backend import LocalBackend + + super().__init__(middlewares=middlewares) + self._bash_parser = BashCommandParser() + + self.dangerous_files = list(dangerous_files) + self.dangerous_directories = list(dangerous_directories) + self._cwd = os.fspath(cwd) if cwd is not None else None + self._backend = backend or LocalBackend() + + async def check_read_only( + self, + tool_input: dict[str, Any], + ) -> bool: + """Decide whether this specific bash invocation is read-only. + + Inspects the command and returns ``True`` for known-safe read-only + commands (e.g. ``ls``, ``cat``, ``grep``, ``git status``). The + static :attr:`is_read_only` class attribute is ``False`` because + Bash can execute arbitrary commands; this method overrides that + with a per-invocation answer. + """ + command = tool_input.get("command", "") + if not command: + return self.is_read_only + return self._bash_parser.is_read_only_command(command) + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permissions for bash command execution. + + This method implements Bash-specific permission checks: + + 0. Injection risk check (bypass-immune safety ASK if command + contains dynamic expansion like ``$(...)`` or ``<(...)``) + 1. Read-only command check — auto-ALLOW in **every mode** + (including DEFAULT) for known-safe read-only commands + (``ls``, ``pwd``, ``git status``, ``cat``, etc.). This is + the static counterpart to :meth:`check_read_only`. + 2. Dangerous command pattern check (bypass-immune safety ASK) + 3. Sed in-place constraint check (bypass-immune safety ASK) + 4. Dangerous path check for config files (bypass-immune safety + ASK) + 5. Dangerous removal path check for system dirs (bypass-immune + safety ASK) + 6. ACCEPT_EDITS auto-allow for ``mkdir``/``touch``/``rm``/ + ``rmdir``/``mv``/``cp``/``sed`` — only when **every** + target path resolves inside a working directory + 7. PASSTHROUGH (engine continues with rule matching) + + "Bypass-immune" decisions set + :attr:`PermissionDecision.bypass_immune` so they cannot be + silenced by allow rules in DEFAULT mode. In BYPASS mode all + bypass-immune ASKs are intentionally skipped — see + :attr:`PermissionMode.BYPASS`. + + Args: + tool_input (`dict[str, Any]`): + The tool input containing "command" key + context (`PermissionContext`): + The permission context with mode and rules + + Returns: + `PermissionDecision`: + ALLOW for safe operations, ASK for dangerous operations, + PASSTHROUGH to let Engine continue with rule matching + """ + + command = tool_input.get("command", "") + if not command: + return PermissionDecision( + behavior=PermissionBehavior.PASSTHROUGH, + message="Empty command", + ) + + # 0. Injection check: detect dynamic shell structures that cannot be + # statically analyzed (command substitution, process substitution, + # control flow, etc.). Must run before read-only check so that + # `$(rm -rf /)` inside an otherwise-safe command is caught. + injection_reason = self._bash_parser.check_injection_risk(command) + if injection_reason: + return PermissionDecision( + behavior=PermissionBehavior.ASK, + message=f"Permission required: {injection_reason}", + decision_reason="Safety check: command contains dynamic " + "expansion that cannot be statically analyzed", + bypass_immune=True, + ) + + # 1. Check if command is read-only (auto-allow) + if self._bash_parser.is_read_only_command(command): + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message="Permission granted for read-only command", + decision_reason="Read-only command is allowed", + ) + + # 2. Check for dangerous commands (safety check, bypass-immune) + dangerous_pattern = self._bash_parser.check_dangerous_command(command) + if dangerous_pattern: + return PermissionDecision( + behavior=PermissionBehavior.ASK, + message=f"Permission required: Command contains dangerous " + f"pattern: {dangerous_pattern}", + decision_reason="Safety check: dangerous command pattern " + "detected", + bypass_immune=True, + ) + + # 3. Check for sed constraints (safety check, bypass-immune) + sed_error = self._bash_parser.check_sed_constraints( + command, + self.dangerous_files, + ) + if sed_error: + return PermissionDecision( + behavior=PermissionBehavior.ASK, + message=f"Permission required: {sed_error}", + decision_reason="Safety check: sed in-place modification " + "of dangerous file", + bypass_immune=True, + ) + + # 4. Check for dangerous paths in sensitive config files/dirs + # (safety check, bypass-immune) + dangerous_paths = self._extract_dangerous_paths_from_bash(command) + if dangerous_paths: + paths_str = ", ".join(dangerous_paths) + return PermissionDecision( + behavior=PermissionBehavior.ASK, + message=f"Permission required: Bash command operates on " + f"sensitive paths: {paths_str}", + decision_reason="Safety check: dangerous file or " + "directory in bash command", + bypass_immune=True, + ) + + # 5. Check for dangerous removal paths: rm/rmdir targeting system + # critical directories like /, /usr, /etc, ~ (bypass-immune). + # Checked separately from step 4 because these paths are not in the + # dangerous_files/directories lists — they are system-level paths + # that should never be removed regardless of user configuration. + removal_path = await self._check_dangerous_removal_path(command) + if removal_path: + return PermissionDecision( + behavior=PermissionBehavior.ASK, + message=f"Dangerous removal operation detected: " + f"'{removal_path}'\n\nThis command would remove a critical " + f"system directory. This requires explicit approval and " + f"cannot be auto-allowed by permission rules.", + decision_reason="Safety check: dangerous removal of " + "critical system path", + bypass_immune=True, + ) + + # 6. ACCEPT_EDITS auto-allow for filesystem commands whose targets + # all live inside a working directory. Mirrors Write/Edit's strict + # working-directory check — we never auto-allow a bash command that + # would touch a path outside the configured working set (e.g. + # ``cp /etc/hosts /tmp/x`` must not pass even though ``cp`` is in + # the auto-allow list). + if context.mode == PermissionMode.ACCEPT_EDITS: + filesystem_commands = { + "mkdir", + "touch", + "rm", + "rmdir", + "mv", + "cp", + "sed", + } + base_command = ( + command.strip().split()[0] if command.strip() else "" + ) + + if base_command in filesystem_commands: + # Collect every target path: file arguments AND output + # redirections. ``extract_file_paths`` includes both. + target_paths = [ + path + for _cmd, path in self._bash_parser.extract_file_paths( + command, + ) + ] + # Conservative: only auto-allow when we extracted at least + # one target AND every target resolves inside a working + # directory. An empty list means the parser found nothing + # actionable (or the command has no args) — in that case + # we fall through to PASSTHROUGH rather than blindly + # allowing. + if target_paths and all( + self._path_in_allowed_working_path(path, context) + for path in target_paths + ): + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message=f"Permission granted for '{base_command}' " + f"command (accept edits mode - filesystem command, " + f"all targets in working directory)", + decision_reason=( + f"Filesystem command '{base_command}' is " + f"auto-allowed in accept edits mode because " + f"all target paths are within a working " + f"directory" + ), + ) + + # 7. Passthrough to let Engine continue with rule matching + return PermissionDecision( + behavior=PermissionBehavior.PASSTHROUGH, + message=f"Execute bash command: {command}", + ) + + async def match_rule( + self, + rule_content: str | None, + tool_input: dict[str, Any], + ) -> bool: + r"""Match Bash command using regex-based wildcard matching. + + Implements wildcard matching with escape sequences: + - Supports \* for literal asterisk and \\ for literal backslash + - Special optimization: "git *" matches both "git" and "git add" + - Prefix pattern (e.g., "git:*"): matches commands starting with "git " + - Wildcard pattern: converts to regex with proper escape handling + - Substring pattern: exact substring matching + - If rule_content is None, matches all invocations + (tool-name-level rule) + + Args: + rule_content: The command pattern to match, or None to match all + tool_input: Must contain a "command" key with the command string + + Returns: + True if pattern matches the command + """ + # None = tool-name-level rule, matches everything + if rule_content is None: + return True + + command = tool_input.get("command", "") + + # Check if pattern is a prefix pattern (ends with :*) + if rule_content.endswith(":*"): + prefix = rule_content[:-2].strip() + return command.startswith(prefix + " ") or command == prefix + + # Check if pattern has unescaped wildcards + def has_wildcards(pattern: str) -> bool: + """Check if pattern contains unescaped * wildcards.""" + i = 0 + while i < len(pattern): + if pattern[i] == "\\": + i += 2 # Skip escaped character + elif pattern[i] == "*": + return True + else: + i += 1 + return False + + if not has_wildcards(rule_content): + # No wildcards, but may have escape sequences + # Convert escape sequences for matching + pattern = rule_content + pattern = pattern.replace("\\\\", "\x00BACKSLASH\x00") + pattern = pattern.replace("\\*", "*") + pattern = pattern.replace("\x00BACKSLASH\x00", "\\") + # Use substring matching with converted pattern + return pattern in command + + # Convert wildcard pattern to regex with escape handling + # Use placeholders for escaped sequences + ESCAPED_STAR = "\x00ESCAPED_STAR\x00" + ESCAPED_BACKSLASH = "\x00ESCAPED_BACKSLASH\x00" + + pattern = rule_content + # Replace \\ with placeholder + pattern = pattern.replace("\\\\", ESCAPED_BACKSLASH) + # Replace \* with placeholder + pattern = pattern.replace("\\*", ESCAPED_STAR) + + # Manually escape regex special characters (except *) + # Don't use re.escape() as it escapes spaces too + special_chars = r".^$+?{}[]|()" + for char in special_chars: + pattern = pattern.replace(char, "\\" + char) + + # Convert * to regex .* (match any characters) + pattern = pattern.replace("*", ".*") + + # Restore escaped sequences + pattern = pattern.replace(ESCAPED_STAR, r"\*") + pattern = pattern.replace(ESCAPED_BACKSLASH, r"\\") + + # Special optimization: "git *" should match both "git" and "git add" + # Pattern: if ends with .*, make it optional + if pattern.endswith(".*"): + base_pattern = pattern[:-2] # Remove .* + # Try exact match first (handles trailing space) + base_pattern = base_pattern.rstrip() + if re.fullmatch(base_pattern, command): + return True + + # Full regex match + try: + return bool(re.fullmatch(pattern, command)) + except re.error: + # Invalid regex, fall back to substring matching + return rule_content.replace("*", "") in command + + async def generate_suggestions( + self, + tool_input: dict[str, Any], + ) -> List["PermissionRule"]: + """Generate suggested permission rules for Bash commands. + + Generates prefix rules based on command + subcommand (two words). + For example, "git commit -m 'xxx'" generates "git commit:*". + + Args: + tool_input (`dict[str, Any]`): + The tool input data containing "command" key + + Returns: + `List[PermissionRule]`: + List of suggested permission rules based on command prefixes + """ + + command = tool_input.get("command", "") + if not command: + return [] + + # Use bash parser to extract command prefixes + prefixes = self._bash_parser.extract_command_prefixes( + command, + max_prefixes=5, + ) + + if not prefixes: + # Cannot extract any prefix, return empty + return [] + + # Generate rules for each prefix + rules = [] + for prefix in prefixes: + rules.append( + PermissionRule( + tool_name="Bash", + rule_content=f"{prefix}:*", + behavior=PermissionBehavior.ALLOW, + source="suggested", + ), + ) + + return rules + + def _extract_dangerous_paths_from_bash( + self, + command: str, + ) -> list[str]: + """Extract dangerous paths from a bash command using tree-sitter. + + Checks for dangerous paths in: + - File-manipulating commands (rm, mv, cp, chmod, chown, sed, touch) + - Output redirections (>, >>) + + Args: + command (`str`): + The bash command string + + Returns: + `list[str]`: + List of dangerous paths found in the command + """ + dangerous_paths = [] + + # Use tree-sitter to extract file paths + file_paths = self._bash_parser.extract_file_paths(command) + + for _cmd_name, path in file_paths: + if self._is_dangerous_path(path): + dangerous_paths.append(path) + + return dangerous_paths + + async def _check_dangerous_removal_path(self, command: str) -> str | None: + """Check if a rm/rmdir command targets a critical system path. + + Detects commands like `rm -rf /`, `rm -rf /usr`, `rmdir ~` that + would destroy critical system directories. Unlike _is_dangerous_path + (which checks against a configurable list of sensitive config files), + this checks against a fixed set of system-level paths that must + never be removed regardless of user configuration. + + Dangerous paths are: + - Root directory (/) + - Home directory (~) + - Wildcard alone (*) or as dir/* (removes everything) + - Direct children of root (/usr, /etc, /tmp, /var, etc.) + + Args: + command (`str`): + The bash command string + + Returns: + `str | None`: + The dangerous path if found, None otherwise + """ + tokens = command.strip().split() + if not tokens: + return None + + # Find rm or rmdir subcommands (handle compound commands) + try: + tree = self._bash_parser.parser.parse(bytes(command, "utf8")) + subcommands = self._bash_parser.split_compound_command( + tree.root_node, + command, + ) + except Exception: + subcommands = [command] + + # Check each subcommand for rm/rmdir + for subcmd in subcommands: + subcmd_tokens = subcmd.strip().split() + if not subcmd_tokens: + continue + base = subcmd_tokens[0] + if base not in ("rm", "rmdir"): + continue + + # Collect non-flag arguments as potential paths + i = 1 + while i < len(subcmd_tokens): + tok = subcmd_tokens[i] + # Skip flags + if tok.startswith("-"): + i += 1 + continue + path = tok.strip("'\"") + if await self._is_dangerous_removal_path(path): + return path + i += 1 + + return None + + async def _is_dangerous_removal_path(self, path: str) -> bool: + """Check if a path is a critical system directory that must not be + removed. + + All path resolution is performed via the backend so that the + check operates on the **backend environment's** ``$HOME`` / + ``cwd`` / path semantics, not the host process's. + + Args: + path (`str`): + The path to check (may be relative, absolute, or contain ~) + + Returns: + `bool`: + True if removing this path would be catastrophic + """ + + # Bare wildcard + if path in ("*", "./*", "/"): + return True + # Ends with /* — removes everything in a directory + if path.endswith("/*") or path.endswith("\\*"): + return True + + # Expand tilde and resolve to an absolute path inside the + # backend environment. Don't resolve symlinks — ``/tmp`` is a + # symlink on macOS but is still a root-child and should be + # flagged. + expanded = await self._backend.expanduser(path) + backend_cwd = await self._backend.getcwd() + abs_path = self._backend.abspath(expanded, cwd=backend_cwd) + + # Home directory + home = await self._backend.expanduser("~") + if abs_path == home: + return True + + # Root itself: ``dirname(root) == root`` on both POSIX + # (``"/"``) and Windows (``"C:\\"``), so this check is + # path-flavor agnostic. + parent = self._backend.dirname(abs_path) + if abs_path == parent: + return True + + # Direct children of root (e.g. ``/usr``, ``/etc``, ``/tmp``): + # the *parent* of these is the root, where + # ``dirname(parent) == parent``. + if self._backend.dirname(parent) == parent: + return True + + return False + + async def call( # type: ignore[override] # pylint: disable=unused-argument + self, + command: str, + description: str = "", + timeout: int = 120000, + ) -> AsyncGenerator[ToolChunk, None]: + """Execute the bash and return the output. + + Args: + command: The bash command to execute. + description: Optional description of what the command does. + timeout: Timeout in milliseconds (default: 120000, max: 600000). + + Yields: + ToolChunk: The tool execution result with stdout/stderr content. + """ + + # Clamp timeout to max 600000ms and convert to seconds + timeout_ms = min(timeout, 600000) + timeout_sec = timeout_ms / 1000.0 + + try: + # ``command`` is a full shell command line (it may contain + # pipes, redirects, ``&&``, …), so wrap it in a shell — the + # backend primitive runs the argv directly without one. Pick + # the platform's native shell so the Windows experience that + # ``main`` had (commands interpreted by ``cmd.exe``) is + # preserved; POSIX hosts use ``/bin/sh``. + if os.name == "nt": + shell_command = ["cmd", "/c", command] + else: + shell_command = ["/bin/sh", "-c", command] + result = await self._backend.exec_shell( + shell_command, + cwd=self._cwd, + timeout=timeout_sec, + ) + + # Decode and normalize line endings + stdout = result.stdout.decode( + "utf-8", + errors="replace", + ).replace("\r\n", "\n") + stderr = result.stderr.decode( + "utf-8", + errors="replace", + ).replace("\r\n", "\n") + + # Check for timeout (backend returns exit_code=-1, + # stderr=b"timed out") + if result.exit_code == -1 and result.stderr == b"timed out": + error_msg = ( + f"Command timed out after {timeout_ms}ms: {command}" + ) + yield ToolChunk( + content=[TextBlock(text=error_msg)], + state=ToolResultState.ERROR, + is_last=True, + ) + return + + # Combine output + output = stdout + if stderr: + if output: + output += "\n" + output += stderr + + # Truncate if exceeds 30000 characters + if len(output) > 30000: + output = output[:30000] + "\n... (output truncated)" + + # Check exit code + if not result.ok(): + # Command failed + error_result = f"Command failed: {command}\n" + if stdout: + error_result += f"\nStdout:\n{stdout}" + if stderr: + error_result += f"\nStderr:\n{stderr}" + + # Truncate error message if needed + if len(error_result) > 30000: + error_result = ( + error_result[:30000] + "\n... (output truncated)" + ) + + yield ToolChunk( + content=[TextBlock(text=error_result)], + state=ToolResultState.ERROR, + is_last=True, + ) + else: + # Command succeeded - note: ToolChunk uses "running" state + # which will be converted to "finished" in ToolResponse + yield ToolChunk( + content=[TextBlock(text=output)], + state=ToolResultState.RUNNING, + is_last=True, + ) + + except Exception as e: + # Other errors + error_msg = f"Command failed: {command}\nError: {str(e)}" + yield ToolChunk( + content=[TextBlock(text=error_msg)], + state=ToolResultState.ERROR, + is_last=True, + ) diff --git a/src/agentscope/tool/_builtin/_bash_parser.py b/src/agentscope/tool/_builtin/_bash_parser.py new file mode 100644 index 0000000000000000000000000000000000000000..bb55fd7c43f0e18eb80d5503ccb771b8a38c1acb --- /dev/null +++ b/src/agentscope/tool/_builtin/_bash_parser.py @@ -0,0 +1,884 @@ +# -*- coding: utf-8 -*- +"""Bash command parser using tree-sitter for precise syntax analysis. + +This module provides utilities to parse Bash commands and extract meaningful +information for permission rule generation, including: +- Splitting compound commands (&&, ||, ;, |) +- Extracting command prefixes (e.g., "npm run" from "npm run build") +- Extracting file paths from commands for dangerous path detection +- Extracting output redirections +- Checking if commands are read-only +""" + +from typing import List, Optional, Set, Tuple + +import re +import shlex + +import tree_sitter_bash as tsbash +from tree_sitter import Language, Parser, Node + +from .._constants import DANGEROUS_NODE_TYPES, DANGEROUS_COMMANDS + + +# Commands that are considered safe and don't require permission rules +SAFE_COMMANDS: Set[str] = { + "echo", + "cat", + "ls", + "pwd", + "cd", + "true", + "false", + "printf", + "grep", + "tee", +} + +# Safe environment variables that can be skipped when extracting command prefix +SAFE_ENV_VARS = { + "NODE_ENV", + "PYTHONUNBUFFERED", + "RUST_LOG", + "LANG", + "TERM", + "NO_COLOR", + "FORCE_COLOR", + "DEBUG", + "VERBOSE", + "CI", + "PATH", + "HOME", + "USER", + "SHELL", + "EDITOR", + "PAGER", + "TZ", + "LC_ALL", + "LC_CTYPE", + "COLUMNS", + "LINES", + "CLICOLOR", + "CLICOLOR_FORCE", +} + +# Read-only git commands +GIT_READ_ONLY_COMMANDS = { + "git status", + "git log", + "git diff", + "git show", + "git branch", + "git tag", + "git remote", + "git ls-files", + "git ls-tree", + "git cat-file", + "git rev-parse", + "git rev-list", + "git describe", + "git shortlog", + "git blame", + "git grep", + "git reflog", + "git config --get", + "git config --list", +} + +# Read-only commands for various tools +READ_ONLY_COMMANDS = { + # Basic file operations + "ls", + "cat", + "head", + "tail", + "less", + "more", + "file", + "stat", + "wc", + "grep", + "rg", + "ag", + "ack", + "find", + "tree", + "pwd", + "which", + "whereis", + "type", + # Git commands + *GIT_READ_ONLY_COMMANDS, + # Docker read-only + "docker ps", + "docker images", + "docker inspect", + "docker logs", + "docker version", + "docker info", + # GitHub CLI read-only + "gh repo view", + "gh issue list", + "gh pr list", + "gh status", + # Python/Node tools + "python --version", + "python -V", + "node --version", + "node -v", + "npm list", + "npm ls", + "pip list", + "pip show", +} + + +class BashCommandParser: + """Parse Bash commands using tree-sitter for accurate syntax analysis.""" + + def __init__(self) -> None: + """Initialize the parser with tree-sitter-bash language.""" + self.parser = Parser(Language(tsbash.language())) + + def is_read_only_command(self, command: str) -> bool: + """Check if a command is read-only (safe to auto-allow). + + For compound commands (&&, ||, ;, |), ALL subcommands must be + read-only for the entire command to be considered read-only. + + Commands with output redirections (>, >>) are NOT considered read-only. + + Args: + command (`str`): + The bash command string + + Returns: + `bool`: + True if the command (and all subcommands) are read-only, + False otherwise + """ + # Normalize command (strip leading/trailing whitespace) + cmd = command.strip() + + # Check for output redirections - these are NOT read-only + if ">" in cmd: + return False + + # Check if it's a compound command + if any(op in cmd for op in ["&&", "||", ";", "|"]): + # Split into subcommands and check each one + try: + tree = self.parser.parse(bytes(cmd, "utf8")) + root = tree.root_node + subcommands = self.split_compound_command(root, cmd) + + # All subcommands must be read-only + for subcmd in subcommands: + if not self._is_single_command_read_only(subcmd.strip()): + return False + return True + except Exception: + # If parsing fails, be conservative + return False + + # Single command - check directly + return self._is_single_command_read_only(cmd) + + def _is_single_command_read_only(self, cmd: str) -> bool: + """Check if a single (non-compound) command is read-only. + + Args: + cmd (`str`): + A single command string (no &&, ||, ;, |) + + Returns: + `bool`: + True if the command is read-only, False otherwise + """ + # Check exact match in read-only commands + if cmd in READ_ONLY_COMMANDS: + return True + + # Check if it starts with a read-only prefix + for readonly_cmd in READ_ONLY_COMMANDS: + if cmd == readonly_cmd or cmd.startswith(readonly_cmd + " "): + return True + + # Check base command for simple read-only operations + tokens = cmd.split() + if tokens: + base_cmd = tokens[0] + # Skip environment variables + i = 0 + while i < len(tokens) and "=" in tokens[i]: + i += 1 + if i < len(tokens): + base_cmd = tokens[i] + + # Check if base command is in safe commands + if base_cmd in SAFE_COMMANDS: + return True + + return False + + def extract_file_paths( + self, + command: str, + ) -> List[Tuple[str, str]]: + """Extract file paths from a bash command using tree-sitter. + + Returns paths that are arguments to file-manipulating commands + (rm, mv, cp, chmod, chown, etc.) and output redirection targets. + + Args: + command (`str`): + The bash command string + + Returns: + `List[Tuple[str, str]]`: + List of tuples (command_name, file_path) + """ + paths = [] + + try: + # Parse command to AST + tree = self.parser.parse(bytes(command, "utf8")) + root = tree.root_node + + # Extract paths from commands + self._extract_paths_from_node(root, command, paths) + + except Exception: + # Fallback to simple token-based extraction + paths = self._extract_paths_fallback(command) + + return paths + + def _extract_paths_from_node( + self, + node: Node, + command: str, + paths: List[Tuple[str, str]], + ) -> None: + """Recursively extract file paths from AST nodes. + + Args: + node (`Node`): + The AST node to process + command (`str`): + The original command string + paths (`List[Tuple[str, str]]`): + List to append (command_name, path) tuples to + """ + # Check for redirections + if node.type == "file_redirect": + # Extract the target file + for child in node.children: + if child.type == "word": + path = command[child.start_byte : child.end_byte] + paths.append(("redirect", path.strip("'\""))) + # Check for commands + if node.type == "command": + # Extract command name and arguments + cmd_name = None + args = [] + + for child in node.children: + if child.type == "command_name": + cmd_name = command[child.start_byte : child.end_byte] + elif child.type == "word" and cmd_name: + arg = command[child.start_byte : child.end_byte] + args.append(arg.strip("'\"")) + + # Check if this is a file-manipulating command + if cmd_name in [ + "rm", + "mv", + "cp", + "chmod", + "chown", + "chgrp", + "touch", + "ln", + "sed", + "mkdir", + "rmdir", + ]: + # Extract file arguments (skip flags) + for arg in args: + if not arg.startswith("-"): + paths.append((cmd_name, arg)) + + # Recursively process children + for child in node.children: + self._extract_paths_from_node(child, command, paths) + + def _extract_paths_fallback( + self, + command: str, + ) -> List[Tuple[str, str]]: + """Fallback path extraction using simple token parsing. + + Args: + command (`str`): + The bash command string + + Returns: + `List[Tuple[str, str]]`: + List of tuples (command_name, file_path) + """ + paths = [] + tokens = command.split() + i = 0 + + while i < len(tokens): + token = tokens[i] + + # Check for output redirections + if token in [">", ">>", "2>", "&>"]: + if i + 1 < len(tokens): + path = tokens[i + 1].strip("'\"") + paths.append(("redirect", path)) + i += 2 + continue + + # Check for file-manipulating commands + if token in [ + "rm", + "mv", + "cp", + "chmod", + "chown", + "sed", + "touch", + "mkdir", + "rmdir", + ]: + cmd_name = token + # Look for file arguments after this command + j = i + 1 + while j < len(tokens): + arg = tokens[j].strip("'\"") + # Skip flags + if arg.startswith("-"): + j += 1 + continue + # This is a file argument + paths.append((cmd_name, arg)) + j += 1 + break + + i += 1 + + return paths + + def extract_redirections(self, command: str) -> List[str]: + """Extract output redirection targets from a bash command. + + Args: + command (`str`): + The bash command string + + Returns: + `List[str]`: + List of file paths that are redirection targets + """ + redirections = [] + + try: + # Parse command to AST + tree = self.parser.parse(bytes(command, "utf8")) + root = tree.root_node + + # Extract redirections + self._extract_redirections_from_node(root, command, redirections) + + except Exception: + # Fallback to simple extraction + tokens = command.split() + for i, token in enumerate(tokens): + if token in [">", ">>", "2>", "&>"] and i + 1 < len(tokens): + path = tokens[i + 1].strip("'\"") + redirections.append(path) + + return redirections + + def _extract_redirections_from_node( + self, + node: Node, + command: str, + redirections: List[str], + ) -> None: + """Recursively extract redirections from AST nodes. + + Args: + node (`Node`): + The AST node to process + command (`str`): + The original command string + redirections (`List[str]`): + List to append redirection targets to + """ + if node.type == "file_redirect": + # Extract the target file + for child in node.children: + if child.type == "word": + path = command[child.start_byte : child.end_byte] + redirections.append(path.strip("'\"")) + + # Recursively process children + for child in node.children: + self._extract_redirections_from_node(child, command, redirections) + + def extract_command_prefixes( + self, + command: str, + max_prefixes: int = 5, + ) -> List[str]: + """Extract command prefixes from a bash command. + + Automatically handles compound commands (&&, ||, ;, |) and extracts + prefixes from each subcommand. Returns deduplicated list of prefixes. + + Args: + command (`str`): + The bash command string (may be compound) + max_prefixes (`int`): + Maximum number of prefixes to return (default: 5) + + Returns: + `List[str]`: + List of command prefixes (deduplicated), e.g., ["npm run", + "git commit"] + + Examples: + >>> parser.extract_command_prefixes("git add . && git commit") + ['git add', 'git commit'] + >>> parser.extract_command_prefixes("npm run build") + ['npm run'] + >>> parser.extract_command_prefixes("ls -la") + [] + """ + if not command or not command.strip(): + return [] + + # Parse command to AST + tree = self.parser.parse(bytes(command, "utf8")) + root = tree.root_node + + # Split compound commands + subcommands = self.split_compound_command(root, command) + + # Extract prefixes from each subcommand + prefixes = [] + seen = set() + + for subcmd in subcommands[:max_prefixes]: + prefix = self._extract_command_prefix(subcmd) + if prefix and prefix not in seen: + prefixes.append(prefix) + seen.add(prefix) + + if len(prefixes) >= max_prefixes: + break + + return prefixes + + def split_compound_command(self, root: Node, command: str) -> List[str]: + """Split compound commands using tree-sitter for precise parsing. + + Recognizes: &&, ||, ;, | + + Args: + root (`Node`): + The root AST node + command (`str`): + The original command string + + Returns: + `List[str]`: + List of individual subcommands + """ + subcommands = [] + + def extract_commands(node: Node) -> None: + """Recursively extract commands from AST.""" + if node.type == "command": + # Extract command text + cmd_text = command[node.start_byte : node.end_byte] + subcommands.append(cmd_text) + elif node.type in ["list", "pipeline", "command_list"]: + # Recursively process compound structures + for child in node.children: + if child.type not in ["&&", "||", ";", "|", "|&"]: + extract_commands(child) + else: + # Continue traversing + for child in node.children: + extract_commands(child) + + extract_commands(root) + return subcommands if subcommands else [command] + + def _extract_command_prefix( + self, + subcmd: str, + ) -> Optional[str]: + """Extract command prefix (first two words) from a subcommand. + + Logic: + 1. Skip safe environment variable assignments + 2. Extract command name and first subcommand + 3. Verify the second word looks like a subcommand (not a flag) + + Args: + subcmd (`str`): + The subcommand string to extract prefix from + + Returns: + `Optional[str]`: + Command prefix (e.g., "npm run") or None if cannot extract + """ + # Parse the subcommand + tree = self.parser.parse(bytes(subcmd, "utf8")) + root = tree.root_node + + # Find the first simple_command node + simple_cmd = self._find_first_simple_command(root) + if not simple_cmd: + return None + + # Extract command parts + parts = [] + env_vars = [] + + for child in simple_cmd.children: + if child.type == "variable_assignment": + # Environment variable assignment + var_name = subcmd[child.start_byte : child.end_byte].split( + "=", + )[0] + env_vars.append(var_name) + elif child.type == "command_name": + # Command name + parts.append(subcmd[child.start_byte : child.end_byte]) + elif child.type == "word" and len(parts) >= 1: + # Argument (might be a flag or subcommand) + word = subcmd[child.start_byte : child.end_byte] + parts.append(word) + # Stop after we have command + first argument + if len(parts) >= 2: + break + + # Check if environment variables are safe + if env_vars and not all(v in SAFE_ENV_VARS for v in env_vars): + return None + + # Check if the command is a safe command that doesn't need permission + if parts and parts[0].lower() in SAFE_COMMANDS: + return None + + # Return first two words + if len(parts) >= 2: + return " ".join(parts[:2]) + + return None + + def _find_first_simple_command(self, node: Node) -> Optional[Node]: + """Recursively find the first command node in AST. + + Args: + node (`Node`): + The AST node to search from + + Returns: + `Optional[Node]`: + The first command node found, or None + """ + if node.type == "command": + return node + + for child in node.children: + result = self._find_first_simple_command(child) + if result: + return result + + return None + + def check_dangerous_command(self, command: str) -> Optional[str]: + """Check if command contains dangerous patterns. + + Uses word-boundary aware matching to avoid false positives like + 'git add' matching 'dd' pattern. + + Args: + command (`str`): + The bash command to check + + Returns: + `Optional[str]`: + The matched dangerous pattern if found, None otherwise + """ + + # Normalize command for matching + normalized = " ".join(command.split()) + + # Check each dangerous pattern + for pattern in DANGEROUS_COMMANDS: + # For single-word patterns like "dd", use word boundary matching + # to avoid false positives (e.g., "git add" shouldn't match "dd") + if " " not in pattern and len(pattern) <= 4: + # Single word pattern - use word boundaries + regex = r"\b" + re.escape(pattern) + r"\b" + if re.search(regex, normalized): + return pattern + else: + # Multi-word pattern or longer pattern - use substring match + if pattern in normalized: + return pattern + + return None + + # pylint: disable=too-many-return-statements, too-many-branches + def check_sed_constraints( + self, + command: str, + dangerous_files: List[str], + ) -> str | None: + """Check if sed command violates safety constraints. + + Implements allowlist/denylist system: + - Allowlist: Line printing (sed -n 'Np') and substitution (sed 's///') + - Denylist: Dangerous operations (w/W/e/E), file writes, command + execution + + Args: + command: The bash command to check + dangerous_files: List of dangerous file patterns + + Returns: + Error message if dangerous sed operation found, None otherwise + """ + + if "sed" not in command: + return None + + # Parse command using shlex + try: + tokens = shlex.split(command) + except ValueError: + return "sed command has invalid shell syntax" + + # Find sed command position + sed_idx = None + for i, token in enumerate(tokens): + if token == "sed" or token.endswith("/sed"): + sed_idx = i + break + + if sed_idx is None: + return None + + # Extract flags and expressions + args = tokens[sed_idx + 1 :] + flags = [] + expressions = [] + file_args = [] + i = 0 + found_first_expr = False + + while i < len(args): + arg = args[i] + + # Handle flags + if arg.startswith("-") and not arg.startswith("--"): + # Combined flags like -nE + flag_chars = arg[1:] + for char in flag_chars: + flags.append(char) + # -i flag may have optional backup extension argument + # But only skip if next arg doesn't look like an expression + if "i" in flag_chars and i + 1 < len(args): + next_arg = args[i + 1] + # Skip backup extension only if it's not an expression + # or file + if ( + not next_arg.startswith("-") + and not next_arg.startswith("s") + and "." not in next_arg + ): + i += 1 # Skip backup extension + elif arg == "--in-place": + flags.append("i") + if i + 1 < len(args): + next_arg = args[i + 1] + if ( + not next_arg.startswith("-") + and not next_arg.startswith("s") + and "." not in next_arg + ): + i += 1 + elif arg in ["-e", "--expression"]: + if i + 1 < len(args): + expressions.append(args[i + 1]) + i += 1 + elif not arg.startswith("-"): + # First non-flag, non-option arg is expression (if no -e used) + if not found_first_expr: + expressions.append(arg) + found_first_expr = True + else: + file_args.append(arg) + + i += 1 + + # If no expressions found, command is invalid + if not expressions: + return "sed command missing expression" + + # Validate flags - only allow specific flags + allowed_flags = {"n", "E", "e", "i"} + for flag in flags: + if flag not in allowed_flags: + return f"sed flag -{flag} not allowed" + + # Check allowlist patterns + has_n_flag = "n" in flags + has_i_flag = "i" in flags + + for expr in expressions: + # Denylist checks first - dangerous operations + # Check for write operations (w, W) - must be at end or followed + # by space/filename + if ( + re.search(r"/[wW]\s+\S+", expr) + or expr.endswith("/w") + or expr.endswith("/W") + ): + return "sed write operation (w/W) not allowed" + + # Check for execute operations (e, E) - must be at end or + # followed by space + if ( + re.search(r"/[eE](?:\s|$)", expr) + or expr.endswith("/e") + or expr.endswith("/E") + ): + return "sed execute operation (e/E) not allowed" + + # Check for dangerous patterns + if "{" in expr or "}" in expr: + return "sed curly braces not allowed" + if expr.startswith("!"): + return "sed negation (!) not allowed" + if "#" in expr and not expr.startswith("s#"): + return "sed comments not allowed" + + # Pattern 1: Line printing with -n flag (sed -n 'Np' or 'N,Mp') + if has_n_flag: + # Match: number followed by 'p', or range 'N,Mp' + if re.match( + r"^\d+p$", + expr, + ) or re.match( + r"^\d+,\d+p$", + expr, + ): + continue + + # Pattern 2: Substitution command + # (sed 's/pattern/replacement/flags') + if ( + expr.startswith("s/") + or expr.startswith("s|") + or expr.startswith("s#") + ): + delimiter = expr[1] + parts = expr[2:].split(delimiter) + if len(parts) >= 2: + # Valid substitution + # Check substitution flags (g, p, number, etc.) + if len(parts) > 2: + sub_flags = parts[2] + # Allow common substitution flags + if all(c in "gp0123456789" for c in sub_flags): + continue + else: + continue + + # If we reach here, expression doesn't match allowlist + return f"sed expression '{expr}' not in allowlist" + + # Check -i flag with dangerous files + if has_i_flag and file_args: + for file_path in file_args: + for dangerous_file in dangerous_files: + if dangerous_file in file_path or file_path.endswith( + dangerous_file, + ): + return f"sed -i modifying dangerous file: {file_path}" + + return None + + def check_injection_risk(self, command: str) -> Optional[str]: + """Check if command contains structures that cannot be statically + analyzed. + + This detects command substitution, process substitution, complex + expansions, control flow, and other dynamic shell features that + make it impossible to determine the command's behavior without + execution. + + Args: + command (`str`): + The bash command to check + + Returns: + `Optional[str]`: + Reason string if command is too complex, None if it can be + statically analyzed + + Examples: + >>> parser.check_injection_risk("ls -la") + None + >>> parser.check_injection_risk("rm $(find . -name '*.tmp')") + "Command contains command_substitution which cannot be statically + analyzed" + >>> parser.check_injection_risk("for f in *.txt; do cat $f; done") + "Command contains for_statement which cannot be statically + analyzed" + """ + + try: + tree = self.parser.parse(bytes(command, "utf8")) + return self._walk_for_dangerous_nodes(tree.root_node) + except Exception: + # If parsing fails, be conservative and require review + return "Command parsing failed, cannot verify safety" + + def _walk_for_dangerous_nodes(self, node: Node) -> Optional[str]: + """Recursively walk AST to find dangerous node types. + + Args: + node (`Node`): + The AST node to check + + Returns: + `Optional[str]`: + Reason string if dangerous node found, None otherwise + """ + + # Check if this node is a dangerous type + if node.type in DANGEROUS_NODE_TYPES: + return ( + f"Command contains {node.type} which cannot be " + f"statically analyzed" + ) + + # Recursively check children + for child in node.children: + result = self._walk_for_dangerous_nodes(child) + if result: + return result + + return None diff --git a/src/agentscope/tool/_builtin/_edit.py b/src/agentscope/tool/_builtin/_edit.py new file mode 100644 index 0000000000000000000000000000000000000000..436278f01a976c3e2514d7d70ac1eb36f9831b34 --- /dev/null +++ b/src/agentscope/tool/_builtin/_edit.py @@ -0,0 +1,416 @@ +# -*- coding: utf-8 -*- +"""The edit tool in agentscope.""" +import difflib +import fnmatch +from typing import Any, List + +from .._base import ToolBase, ToolMiddlewareBase +from .._constants import ( + DEFAULT_DANGEROUS_FILES, + DEFAULT_DANGEROUS_DIRECTORIES, +) +from ...permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, + PermissionMode, + PermissionRule, +) +from .._response import ToolChunk +from ...message import TextBlock, ToolResultState +from ...state import AgentState +from ._backend import BackendBase, _normalize_newlines + + +class Edit(ToolBase): + """The edit tool for performing exact string replacements in files.""" + + name: str = "Edit" + """The tool name presented to the agent.""" + + description: str = """Performs exact string replacements in files. + +Usage: +- You must use your `Read` tool at least once in the conversation + before editing. This tool will error if you attempt an edit without + reading the file. +- When editing text from Read tool output, ensure you preserve the + exact indentation (tabs/spaces) as it appears AFTER the line number + prefix. The line number prefix format is: line number + tab. + Everything after that is the actual file content to match. Never + include any part of the line number prefix in the old_string or + new_string. +- ALWAYS prefer editing existing files in the codebase. NEVER write + new files unless explicitly required. +- Only use emojis if the user explicitly requests it. Avoid adding + emojis to files unless asked. +- The edit will FAIL if `old_string` is not unique in the file.""" # noqa: E501 + """The description presented to the agent.""" + + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "file_path": { + "type": "string", + "description": "The absolute path to the file to edit.", + }, + "old_string": { + "type": "string", + "description": ( + "The exact string to replace. Must match exactly " + "including whitespace and indentation." + ), + }, + "new_string": { + "type": "string", + "description": "The string to replace old_string with.", + }, + "replace_all": { + "type": "boolean", + "description": ( + "If true, replace all occurrences. If false " + "(default), only replace if there is exactly one " + "occurrence." + ), + "default": False, + }, + }, + "required": ["file_path", "old_string", "new_string"], + } + + is_mcp: bool = False + is_read_only: bool = False + is_concurrency_safe: bool = False + is_external_tool: bool = False + is_state_injected: bool = True + + def __init__( # pylint: disable=dangerous-default-value + self, + dangerous_files: list[str] = DEFAULT_DANGEROUS_FILES, + dangerous_directories: list[str] = DEFAULT_DANGEROUS_DIRECTORIES, + middlewares: List[ToolMiddlewareBase] | None = None, + backend: BackendBase | None = None, + ) -> None: + """Initialize the edit tool. + + Args: + dangerous_files (`list[str]`, optional): + Sensitive files that require explicit user confirmation, + even in BYPASS mode. Matched by basename + (case-insensitive). Defaults to `DEFAULT_DANGEROUS_FILES`. + Pass a custom list to fully replace the defaults, or `[]` + to disable the filename check. + dangerous_directories (`list[str]`, optional): + Sensitive directories that require explicit user + confirmation. Matched when any path segment equals an + entry (case-insensitive). Defaults to + `DEFAULT_DANGEROUS_DIRECTORIES`. Pass a custom list to + fully replace the defaults, or `[]` to disable the + directory check. + middlewares (`List[ToolMiddlewareBase] | None`, optional): + Tool middlewares wrapping the tool execution. + backend (`BackendBase | None`, optional): + The sandbox backend to use for file I/O. When ``None``, + a :class:`LocalBackend` is created. + """ + from ._backend import LocalBackend + + super().__init__(middlewares=middlewares) + self.dangerous_files = list(dangerous_files) + self.dangerous_directories = list(dangerous_directories) + self._backend = backend or LocalBackend() + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permissions for file editing. + + This method implements Edit-specific permission checks: + 1. Dangerous path check (safety check, bypass-immune) + 2. ACCEPT_EDITS mode check for files in working directories + + Args: + tool_input (`dict[str, Any]`): + The tool input containing "file_path" key + context (`PermissionContext`): + The permission context with mode and rules + + Returns: + `PermissionDecision`: + ASK for dangerous paths, ALLOW for safe operations in + ACCEPT_EDITS mode, PASSTHROUGH otherwise + """ + + file_path = tool_input.get("file_path") + if not file_path: + return PermissionDecision( + behavior=PermissionBehavior.PASSTHROUGH, + message="No file path provided", + ) + + # 1. Check for dangerous paths (safety check, bypass-immune) + if self._is_dangerous_path(file_path): + return PermissionDecision( + behavior=PermissionBehavior.ASK, + message=f"Permission required: Edit operation on " + f"sensitive file {file_path}", + decision_reason="Safety check: dangerous file or directory", + bypass_immune=True, + ) + + # 2. Check ACCEPT_EDITS mode for files in working directories + if context.mode == PermissionMode.ACCEPT_EDITS: + if self._path_in_allowed_working_path(file_path, context): + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message=f"Permission granted for editing {file_path} " + f"(accept edits mode - in working directory)", + decision_reason="File is in working directory and not " + "a dangerous path", + ) + + # 3. Return PASSTHROUGH to let PermissionEngine check allow rules + # This ensures allow rules can grant Edit permissions + return PermissionDecision( + behavior=PermissionBehavior.PASSTHROUGH, + message="", + ) + + async def match_rule( + self, + rule_content: str | None, + tool_input: dict[str, Any], + ) -> bool: + """Check if a permission rule matches the file path. + + Matches rule_content as a glob pattern against the "file_path" + parameter using fnmatch. If rule_content is None, matches all + invocations (tool-name-level rule). + + Args: + rule_content (`str | None`): + Glob pattern to match against the file path (e.g., "src/**"), + or None to match all invocations + tool_input (`dict[str, Any]`): + The tool input data containing "file_path" key + + Returns: + `bool`: + True if the glob pattern matches the file path, False otherwise + """ + if rule_content is None: + return True + + file_path = tool_input.get("file_path", "") + if not file_path: + return False + return fnmatch.fnmatch(file_path, rule_content) + + async def generate_suggestions( + self, + tool_input: dict[str, Any], + ) -> List[PermissionRule]: + """Generate suggested permission rules for the file path. + + Suggests a glob pattern covering the parent directory of the file, + allowing the user to grant permission for the entire directory at once. + + Args: + tool_input (`dict[str, Any]`): + The tool input data containing "file_path" key + + Returns: + `List[PermissionRule]`: + A single suggested rule covering the parent directory + (e.g., file "/src/main.py" -> rule "src/**") + """ + file_path = tool_input.get("file_path", "") + if not file_path: + return [] + + parent = self._backend.dirname(file_path) + # Glob patterns are POSIX-style strings (matched by fnmatch), + # not real filesystem paths — do NOT use backend.join_path here. + pattern = (parent.rstrip("/\\") + "/**") if parent else "**" + + return [ + PermissionRule( + tool_name=self.name, + rule_content=pattern, + behavior=PermissionBehavior.ALLOW, + source="suggested", + ), + ] + + async def call( # type: ignore[override] + self, + file_path: str, + old_string: str, + new_string: str, + replace_all: bool = False, + _agent_state: AgentState | None = None, + ) -> ToolChunk: + """Execute the edit and return the result.""" + # Validate file_path is absolute + if not self._backend.isabs(file_path): + return ToolChunk( + content=[ + TextBlock( + text=( + f"Error: file_path must be an absolute " + f"path, got: {file_path}" + ), + ), + ], + state=ToolResultState.ERROR, + is_last=True, + ) + + # Check file exists + if not await self._backend.file_exists(file_path): + return ToolChunk( + content=[ + TextBlock(text=f"Error: File not found: {file_path}"), + ], + state=ToolResultState.ERROR, + is_last=True, + ) + + # Check old_string != new_string + if old_string == new_string: + return ToolChunk( + content=[ + TextBlock( + text=( + "Error: old_string and new_string are " + "identical. No changes to make." + ), + ), + ], + state=ToolResultState.ERROR, + is_last=True, + ) + + content = None + if _agent_state is not None: + cache = await _agent_state.tool_context.get_cache(file_path) + if cache is None: + # Haven't read this file before + return ToolChunk( + content=[ + TextBlock( + text="Error: To edit a file, you must first read " + "it using the Read tool.", + ), + ], + state=ToolResultState.ERROR, + is_last=True, + ) + content = "".join(cache.lines) + else: + # No state provided, read from backend + try: + raw = await self._backend.read_file(file_path) + # Normalize CRLF/CR to match the cached-content path and + # the LF-based old_string the caller supplies. + content = _normalize_newlines( + raw.decode("utf-8", errors="replace"), + ) + except Exception as e: + return ToolChunk( + content=[TextBlock(text=f"Error reading file: {str(e)}")], + state=ToolResultState.ERROR, + is_last=True, + ) + + # Count occurrences + occurrences = content.count(old_string) + + # If occurrences == 0, raise error + if occurrences == 0: + return ToolChunk( + content=[ + TextBlock( + text=f"Error: old_string not found in {file_path}", + ), + ], + state=ToolResultState.ERROR, + is_last=True, + ) + + # If occurrences > 1 and not replace_all, raise error + if occurrences > 1 and not replace_all: + return ToolChunk( + content=[ + TextBlock( + text=( + f"Error: old_string appears {occurrences} " + f"times in {file_path}. Set replace_all=true " + f"to replace all occurrences, or make " + f"old_string more specific." + ), + ), + ], + state=ToolResultState.ERROR, + is_last=True, + ) + + # Perform replacement + if replace_all: + updated_content = content.replace(old_string, new_string) + else: + updated_content = content.replace( + old_string, + new_string, + 1, + ) + + # Write updated content back to file via backend + try: + await self._backend.write_file( + file_path, + updated_content.encode("utf-8"), + ) + except Exception as e: + return ToolChunk( + content=[TextBlock(text=f"Error writing file: {str(e)}")], + state=ToolResultState.ERROR, + is_last=True, + ) + + # Return success message + replacement_msg = ( + f"all {occurrences} occurrences" if replace_all else "1 occurrence" + ) + + # Build a unified diff of the change with absolute line numbers so the + # web UI can render it with real line numbers and proper inter-hunk + # gaps. The diff is kept in ``metadata`` only (not in the textual + # output) so it does not bloat the LLM context. + diff_text = "".join( + difflib.unified_diff( + content.splitlines(keepends=True), + updated_content.splitlines(keepends=True), + fromfile=f"a/{file_path}", + tofile=f"b/{file_path}", + n=3, + ), + ) + + return ToolChunk( + content=[ + TextBlock( + text=f"Successfully replaced {replacement_msg} " + f"in {file_path}", + ), + ], + state=ToolResultState.RUNNING, + is_last=True, + metadata={ + "diff": diff_text, + "file_path": file_path, + "occurrences": occurrences if replace_all else 1, + }, + ) diff --git a/src/agentscope/tool/_builtin/_glob.py b/src/agentscope/tool/_builtin/_glob.py new file mode 100644 index 0000000000000000000000000000000000000000..79ee6801a15357fbac66e4de7b01cd51a1fdd3f4 --- /dev/null +++ b/src/agentscope/tool/_builtin/_glob.py @@ -0,0 +1,305 @@ +# -*- coding: utf-8 -*- +"""The glob tool in agentscope.""" + +from __future__ import annotations + +import fnmatch +import json +import sys +from typing import TYPE_CHECKING, Any, List + +from ...message import TextBlock, ToolResultState +from ...permission import ( + PermissionBehavior, + PermissionContext, + PermissionDecision, + PermissionRule, +) +from .._base import ToolBase, ToolMiddlewareBase +from .._response import ToolChunk + +if TYPE_CHECKING: + from ._backend import BackendBase + + +def _default_glob_helper_path() -> str: + """Resolve the on-disk path of the bundled ``_glob_helper.py`` script. + + Used by :class:`Glob` when no explicit ``glob_helper_path`` is + provided (i.e. the local-workspace case). The path is obtained via + :mod:`importlib.resources` so it works for both editable and + installed packages. + """ + import importlib.resources as _res + + ref = _res.files("agentscope.tool._builtin._scripts").joinpath( + "_glob_helper.py", + ) + # as_posix() on a MultiplexedPath / PosixPath gives a str path + return str(ref) + + +class Glob(ToolBase): + """The glob tool for fast file pattern matching.""" + + name: str = "Glob" + """The tool name presented to the agent.""" + + description: str = """Fast file pattern matching tool that works with +any codebase size. + +Supports glob patterns like "**/*.js" or "src/**/*.ts" and returns +matching file paths sorted by modification time (newest first). + +Use this tool when you need to find files by pattern across the +codebase.""" # ignore: E501 + """The description presented to the agent.""" + + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "pattern": { + "type": "string", + "description": "The glob pattern to match against " + "(e.g., '**/*.py', 'src/**/*.ts')", + }, + "path": { + "type": "string", + "description": "The base directory to search from " + "(defaults to current working directory)", + }, + }, + "required": ["pattern"], + } + + is_mcp: bool = False + is_read_only: bool = True + is_concurrency_safe: bool = True + is_external_tool: bool = False + is_state_injected: bool = False + + def __init__( + self, + backend: BackendBase | None = None, + glob_helper_path: str | None = None, + middlewares: List[ToolMiddlewareBase] | None = None, + ) -> None: + """Initialize the glob tool. + + Args: + middlewares (`List[ToolMiddlewareBase] | None`, optional): + Tool middlewares wrapping the tool execution. + backend (`BackendBase | None`, optional): + The sandbox backend to use. When ``None``, a + :class:`LocalBackend` is created automatically. + glob_helper_path (`str | None`, optional): + Filesystem path (inside the backend's environment) to + the ``_glob_helper.py`` script. When ``None``, the + path is resolved from the installed package resources + (suitable for :class:`LocalBackend`). Remote backends + (Docker, E2B) should pass the path where the script + was deployed during workspace initialization. + """ + from ._backend import LocalBackend + + super().__init__(middlewares=middlewares) + self._backend = backend or LocalBackend() + # When running against the host, invoke the helper with the + # current interpreter (``sys.executable``) rather than assuming + # ``python3`` is on PATH. + self._is_local = isinstance(self._backend, LocalBackend) + self._glob_helper_path = ( + glob_helper_path + if glob_helper_path is not None + else _default_glob_helper_path() + ) + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permissions for glob pattern matching. + + Glob is a read-only tool. Return PASSTHROUGH to let the engine + handle EXPLORE mode and rule matching. + """ + return PermissionDecision( + behavior=PermissionBehavior.PASSTHROUGH, + message="Glob pattern matching is read-only.", + ) + + async def match_rule( + self, + rule_content: str | None, + tool_input: dict[str, Any], + ) -> bool: + """Check if a permission rule matches the glob pattern or path. + + Matches rule_content as a glob pattern against the "pattern" or "path" + parameters. This allows rules to match either the search pattern itself + or the directory being searched. If rule_content is None, matches all + invocations (tool-name-level rule). + + Args: + rule_content (`str | None`): + Glob pattern to match (e.g., "src/**" to match searches in + src), or None to match all invocations + tool_input (`dict[str, Any]`): + The tool input data containing "pattern" and optional "path" + + Returns: + `bool`: + True if the rule matches the pattern or path, False otherwise + """ + # None = tool-name-level rule, matches everything + if rule_content is None: + return True + + # Try matching against the search path first + path = tool_input.get("path", "") + if path and fnmatch.fnmatch(path, rule_content): + return True + + # Fall back to matching against the pattern itself + pattern = tool_input.get("pattern", "") + if pattern and fnmatch.fnmatch(pattern, rule_content): + return True + + return False + + async def generate_suggestions( + self, + tool_input: dict[str, Any], + ) -> List[PermissionRule]: + """Generate suggested permission rules for the glob search. + + Suggests a rule based on the search path. If no path is provided, + suggests a rule for the current directory. + + Args: + tool_input (`dict[str, Any]`): + The tool input data containing optional "path" key + + Returns: + `List[PermissionRule]`: + A single suggested rule covering the search directory + """ + backend_cwd = await self._backend.getcwd() + path = tool_input.get("path") or backend_cwd + + # Normalize path and build a glob pattern. Glob patterns are + # POSIX-style strings (matched by fnmatch), not real filesystem + # paths — do NOT use backend.join_path here. + abs_path = self._backend.abspath(path, cwd=backend_cwd) + pattern = abs_path.rstrip("/\\") + "/**" + + return [ + PermissionRule( + tool_name=self.name, + rule_content=pattern, + behavior=PermissionBehavior.ALLOW, + source="suggested", + ), + ] + + async def call( # type: ignore[override] + self, + pattern: str, + path: str | None = None, + ) -> ToolChunk: + """Execute the glob pattern matching and return the results. + + Invokes the standalone ``_glob_helper.py`` script via + ``exec_shell``. The script performs high-performance + ``os.walk`` + ``os.scandir`` matching and returns results + sorted by modification time (newest first) as JSON. + + This unified path works identically across Local, Docker, + and E2B backends. + + Args: + pattern (`str`): + The glob pattern to match against (e.g. ``**/*.py``). + path (`str | None`, optional): + Base directory to search from. Defaults to the current + working directory when ``None``. + + Returns: + `ToolChunk`: + On success, the matched file paths joined by newlines + (or a "no files found" message). If the base directory + is missing or the helper fails, an error chunk with + ``ToolResultState.ERROR``. + """ + base_dir = path if path else await self._backend.getcwd() + + # The base must be an existing directory; a regular file would + # otherwise be accepted here and fail later with a confusing + # error from the helper. + if not await self._backend.is_dir(base_dir): + return ToolChunk( + content=[ + TextBlock(text=f"Directory not found: {base_dir}"), + ], + state=ToolResultState.ERROR, + is_last=True, + ) + + # Invoke the glob helper script via exec_shell as an argv list + # (run directly, without a shell, so no platform-specific + # quoting is needed). Use the current interpreter locally + # (``python3`` may be absent, e.g. on Windows or venvs exposing + # only ``python``); remote backends run inside Linux images + # where ``python3`` is the safe choice. + python = sys.executable if self._is_local else "python3" + command = [ + python, + self._glob_helper_path, + "--pattern", + pattern, + "--base-dir", + base_dir, + ] + result = await self._backend.exec_shell(command, timeout=30.0) + + # A non-zero exit means the helper itself failed (missing + # interpreter/script, permission error, …) — surface it rather + # than masking it as an empty match. + if not result.ok(): + stderr = result.stderr.decode("utf-8", errors="replace").strip() + return ToolChunk( + content=[ + TextBlock( + text=f"Glob helper failed: {stderr}" + if stderr + else "Glob helper failed with no error output.", + ), + ], + state=ToolResultState.ERROR, + is_last=True, + ) + + try: + matches = json.loads( + result.stdout.decode("utf-8", errors="replace"), + ) + except (json.JSONDecodeError, ValueError): + matches = [] + + if len(matches) == 0: + return ToolChunk( + content=[ + TextBlock( + text=f"No files found matching pattern: {pattern}", + ), + ], + state=ToolResultState.RUNNING, + is_last=True, + ) + + return ToolChunk( + content=[TextBlock(text="\n".join(matches))], + state=ToolResultState.RUNNING, + is_last=True, + ) diff --git a/src/agentscope/tool/_builtin/_grep.py b/src/agentscope/tool/_builtin/_grep.py new file mode 100644 index 0000000000000000000000000000000000000000..20c3cb5e1fac15251304e17ea512a1ed0580f85d --- /dev/null +++ b/src/agentscope/tool/_builtin/_grep.py @@ -0,0 +1,488 @@ +# -*- coding: utf-8 -*- +"""The grep tool in agentscope.""" +import fnmatch +from typing import Any, List, Literal + +from .._base import ToolBase, ToolMiddlewareBase +from ...permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, + PermissionRule, +) +from .._response import ToolChunk +from ...message import TextBlock, ToolResultState +from ._backend import BackendBase + +# Version control system directories to exclude from searches +VCS_DIRECTORIES_TO_EXCLUDE = [ + ".git", + ".svn", + ".hg", + ".bzr", + ".jj", + ".sl", +] + +# Default cap on grep results when head_limit is unspecified +DEFAULT_HEAD_LIMIT = 250 + + +class RipgrepTimeoutError(Exception): + """Custom error class for ripgrep timeouts.""" + + def __init__(self, message: str, partial_results: list[str]): + super().__init__(message) + self.partial_results = partial_results + + +class Grep(ToolBase): + """The grep tool for searching file contents using ripgrep.""" + + name: str = "Grep" + """The tool name presented to the agent.""" + + description: str = """A powerful search tool built on ripgrep + + Usage: +- ALWAYS use Grep for search tasks. NEVER invoke `grep` or `rg` as a Bash command. The Grep tool has been optimized for correct permissions and access. +- Supports full regex syntax (e.g., "log.*Error", "function\\s+\\w+") +- Filter files with glob parameter (e.g., "*.js", "**/*.tsx") or type parameter (e.g., "js", "py", "rust") +- Output modes: "content" shows matching lines, "files_with_matches" shows only file paths (default), "count" shows match counts per file +- Context lines: use context parameter or -A/-B/-C for lines after/before/around matches +- Case-insensitive search: set i to true +- Multiline regex: set multiline to true for patterns spanning multiple lines +- Limit results: use head_limit to cap the number of results returned""" # noqa: E501 + """The description presented to the agent.""" + + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "pattern": { + "type": "string", + "description": "The regular expression pattern to search " + "for in file contents.", + }, + "path": { + "type": "string", + "description": "File or directory to search in. Defaults " + "to current working directory.", + }, + "output_mode": { + "type": "string", + "enum": ["content", "files_with_matches", "count"], + "description": "Output mode: 'content' shows matching lines " + "(supports -A/-B/-C context, -n line numbers, " + "head_limit), 'files_with_matches' shows file " + "paths (supports head_limit), 'count' shows " + "match counts (supports head_limit). " + "Defaults to 'files_with_matches'.", + "default": "files_with_matches", + }, + "glob": { + "type": "string", + "description": "Glob pattern to filter files (e.g., '*.js', " + "'*.{ts,tsx}').", + }, + "type": { + "type": "string", + "description": "File type to search (rg --type). " + "Common types: js, py, rust, go, java, etc.", + }, + "-A": { + "type": "integer", + "description": "Number of lines to show after each match. " + "Requires output_mode: 'content'.", + }, + "-B": { + "type": "integer", + "description": "Number of lines to show before each match. " + "Requires output_mode: 'content'.", + }, + "-C": { + "type": "integer", + "description": "Alias for context.", + }, + "context": { + "type": "integer", + "description": "Number of context lines to show before and " + "after matches. Requires output_mode: " + "'content'.", + }, + "n": { + "type": "boolean", + "description": "Show line numbers in output. Requires " + "output_mode: 'content'. Defaults to true.", + "default": True, + }, + "i": { + "type": "boolean", + "description": "Case insensitive search.", + "default": False, + }, + "case_insensitive": { + "type": "boolean", + "description": "Case insensitive search (alias for i).", + "default": False, + }, + "multiline": { + "type": "boolean", + "description": "Enable multiline mode where . matches " + "newlines and patterns can span lines. " + "Default: false.", + "default": False, + }, + "head_limit": { + "type": "integer", + "description": "Limit output to first N lines/entries. " + "Defaults to 250 when unspecified. " + "Pass 0 for unlimited.", + "minimum": 0, + }, + "offset": { + "type": "integer", + "description": "Skip first N lines/entries before applying " + "head_limit. Defaults to 0.", + "default": 0, + "minimum": 0, + }, + }, + "required": ["pattern"], + } + + is_mcp: bool = False + is_read_only: bool = True + is_concurrency_safe: bool = True + is_external_tool: bool = False + is_state_injected: bool = False + + def __init__( + self, + middlewares: List[ToolMiddlewareBase] | None = None, + backend: BackendBase | None = None, + ) -> None: + """Initialize the grep tool. + + Args: + middlewares (`List[ToolMiddlewareBase] | None`, optional): + Tool middlewares wrapping the tool execution. + backend (`BackendBase | None`, optional): + The sandbox backend to use for shell execution. When + ``None``, a :class:`LocalBackend` is created. + Ripgrep is always invoked via ``exec_shell`` so that + the same code path works for local, Docker, and E2B + backends. + """ + from ._backend import LocalBackend + + super().__init__(middlewares=middlewares) + self._backend = backend or LocalBackend() + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permissions for grep search.""" + return PermissionDecision( + behavior=PermissionBehavior.PASSTHROUGH, + message="Grep search is read-only.", + ) + + async def match_rule( + self, + rule_content: str | None, + tool_input: dict[str, Any], + ) -> bool: + """Check if a permission rule matches the grep search path. + + Matches rule_content as a glob pattern against the "path" parameter. + If no path is given, falls back to the current working directory. + If rule_content is None, matches all invocations (tool-name-level + rule). + + Args: + rule_content (`str | None`): + Glob pattern to match against the search path (e.g., "src/**"), + or None to match all invocations + tool_input (`dict[str, Any]`): + The tool input data containing optional "path" key + + Returns: + `bool`: + True if the glob pattern matches the search path, False + otherwise + """ + # None = tool-name-level rule, matches everything + if rule_content is None: + return True + + path = tool_input.get("path", "") + if not path: + path = await self._backend.getcwd() + return fnmatch.fnmatch(path, rule_content) + + async def generate_suggestions( + self, + tool_input: dict[str, Any], + ) -> List[PermissionRule]: + """Generate suggested permission rules for the grep search path. + + Suggests a rule based on the search path. If no path is provided, + suggests a rule for the current directory. + + Args: + tool_input (`dict[str, Any]`): + The tool input data containing optional "path" key + + Returns: + `List[PermissionRule]`: + A single suggested rule covering the search directory + """ + backend_cwd = await self._backend.getcwd() + path = tool_input.get("path") or backend_cwd + + abs_path = self._backend.abspath(path, cwd=backend_cwd) + # Glob patterns are POSIX-style strings (matched by fnmatch), + # not real filesystem paths — do NOT use backend.join_path here. + pattern = abs_path.rstrip("/\\") + "/**" + + return [ + PermissionRule( + tool_name=self.name, + rule_content=pattern, + behavior=PermissionBehavior.ALLOW, + source="suggested", + ), + ] + + def _apply_head_limit( + self, + items: list[str], + limit: int | None, + offset: int = 0, + ) -> tuple[list[str], int | None]: + """Apply head_limit and offset to a list of items. + + Returns (sliced_items, applied_limit_if_truncated). + """ + if limit == 0: + return items[offset:], None + effective_limit = limit if limit is not None else DEFAULT_HEAD_LIMIT + sliced = items[offset : offset + effective_limit] + was_truncated = len(items) - offset > effective_limit + return sliced, (effective_limit if was_truncated else None) + + async def _run_ripgrep( + self, + args: list[str], + search_path: str, + timeout: int = 30, + ) -> list[str]: + """Run ripgrep and return output lines. + + Builds an argument vector and dispatches it through + ``backend.exec_shell`` (which runs the program directly, without + a shell), so the same code path works for local, Docker, and E2B + backends and needs no platform-specific argument quoting. + """ + command = ["rg", *args, search_path] + + result = await self._backend.exec_shell( + command, + timeout=float(timeout), + ) + + if result.exit_code == -1 and result.stderr == b"timed out": + raise RipgrepTimeoutError( + f"Ripgrep search timed out after {timeout} seconds. " + "Try searching a more specific path or pattern.", + [], + ) + + # returncode 0 = matches found, 1 = no matches + if result.exit_code not in (0, 1): + error_msg = result.stderr.decode( + "utf-8", + errors="ignore", + ).strip() + raise RuntimeError( + f"ripgrep error (code {result.exit_code}): {error_msg}", + ) + + raw = result.stdout.decode("utf-8", errors="ignore") + + lines = [ + line.rstrip("\r") for line in raw.split("\n") if line.rstrip("\r") + ] + return lines + + async def call( # type: ignore[override] + self, + pattern: str, + path: str | None = None, + output_mode: Literal[ + "content", + "files_with_matches", + "count", + ] = "files_with_matches", + glob: str | None = None, + type: str | None = None, # pylint: disable=redefined-builtin + i: bool = False, + case_insensitive: bool = False, + context: int | None = None, + multiline: bool = False, + head_limit: int | None = None, + offset: int = 0, + n: bool = True, + **kwargs: Any, + ) -> ToolChunk: + """Execute the grep search using ripgrep. + + Args: + pattern: The regex pattern to search for + path: The directory or file path to search in + output_mode: Output mode ('content', 'files_with_matches', 'count') + glob: Glob pattern to filter files + type: File type to filter by (rg --type) + i: Case-insensitive search (rg -i) + case_insensitive: Alias for i (backward compatibility) + context: Number of context lines around matches + multiline: Enable multiline regex matching + head_limit: Maximum number of results to return + (default 250, 0=unlimited) + offset: Skip first N results + n: Show line numbers (content mode only, default True) + **kwargs: Additional parameters (-A, -B, -C) + """ + search_path = path or await self._backend.getcwd() + + if head_limit is not None and head_limit < 0: + return ToolChunk( + content=[ + TextBlock(text="Error: head_limit must be non-negative."), + ], + state=ToolResultState.ERROR, + is_last=True, + ) + + if offset < 0: + return ToolChunk( + content=[ + TextBlock(text="Error: offset must be non-negative."), + ], + state=ToolResultState.ERROR, + is_last=True, + ) + + args: list[str] = ["--hidden"] + + # Exclude VCS directories + for vcs_dir in VCS_DIRECTORIES_TO_EXCLUDE: + args.extend(["--glob", f"!{vcs_dir}"]) + + # Limit line length to prevent base64/minified content + args.extend(["--max-columns", "500"]) + + # Multiline mode + if multiline: + args.extend(["-U", "--multiline-dotall"]) + + # Case-insensitive (support both i and case_insensitive + # for compatibility) + if i or case_insensitive: + args.append("-i") + + # Output mode flags + if output_mode == "files_with_matches": + args.append("-l") + elif output_mode == "count": + args.append("-c") + + # Line numbers (content mode only) + if n and output_mode == "content": + args.append("-n") + + # Context flags (content mode only) + if output_mode == "content": + A = kwargs.get("-A") + B = kwargs.get("-B") + C = kwargs.get("-C") + + if context is not None: + args.extend(["-C", str(context)]) + elif C is not None: + args.extend(["-C", str(C)]) + else: + if B is not None: + args.extend(["-B", str(B)]) + if A is not None: + args.extend(["-A", str(A)]) + + # Pattern — use -e if it starts with a dash + if pattern.startswith("-"): + args.extend(["-e", pattern]) + else: + args.append(pattern) + + # File type filter + if type is not None: + args.extend(["--type", type]) + + # Glob filter + if glob is not None: + raw_patterns = glob.split() + glob_patterns: list[str] = [] + for raw in raw_patterns: + if "{" in raw and "}" in raw: + glob_patterns.append(raw) + else: + glob_patterns.extend(p for p in raw.split(",") if p) + for gp in glob_patterns: + args.extend(["--glob", gp]) + + try: + results = await self._run_ripgrep(args, search_path) + except RipgrepTimeoutError as e: + return ToolChunk( + content=[TextBlock(text=str(e))], + state=ToolResultState.ERROR, + is_last=True, + ) + except RuntimeError as e: + return ToolChunk( + content=[TextBlock(text=str(e))], + state=ToolResultState.ERROR, + is_last=True, + ) + + if not results: + return ToolChunk( + content=[ + TextBlock(text=f"No matches found for pattern: {pattern}"), + ], + state=ToolResultState.SUCCESS, + is_last=True, + ) + + limited, applied_limit = self._apply_head_limit( + results, + head_limit, + offset, + ) + + suffix = "" + if applied_limit is not None: + suffix = ( + f"\n\n[Showing results with pagination = " + f"limit: {applied_limit}" + ) + if offset: + suffix += f", offset: {offset}" + suffix += "]" + + return ToolChunk( + content=[TextBlock(text="\n".join(limited) + suffix)], + state=ToolResultState.SUCCESS, + is_last=True, + ) diff --git a/src/agentscope/tool/_builtin/_meta.py b/src/agentscope/tool/_builtin/_meta.py new file mode 100644 index 0000000000000000000000000000000000000000..9c9291faf3c7b6a294b0d7da4391901286a40638 --- /dev/null +++ b/src/agentscope/tool/_builtin/_meta.py @@ -0,0 +1,130 @@ +# -*- coding: utf-8 -*- +"""The meta tool class.""" +from typing import Any, List + +from pydantic import Field, create_model +from jinja2 import Template + +from .._tool_group import ToolGroup +from ...permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, +) +from .._response import ToolChunk +from .._base import ToolBase, ToolMiddlewareBase +from ...exception import DeveloperOrientedException +from ...message import TextBlock +from ...state import AgentState + + +class ResetTools(ToolBase): + """A meta tool allows agent to self-manage its equipped tools by + activating or deactivating tool groups dynamically.""" + + name: str = "reset_tools" + description: str = ( + "This tool allows you to reset your equipped tools based on your " + "current task requirements. These tools are organized into different " + "groups, and you can activate/deactivate them by specifying the " + "boolean values for each group in the input.\n\n" + "**Important: The input booleans are the final state of their " + "corresponding tool groups, not incremental changes.** Any group not " + "explicitly set to True will be deactivated, regardless of its " + "previous state.\n\n" + "**Best practice**: Actively manage your tool groups——activate only " + "what you need for the current task, and promptly deactivate groups " + "as soon as they are no longer needed to conserve context space.\n\n" + "This tool will return the usage instructions for the activated tool " + "groups, which you **MUST pay attention to and follow**. You can " + "also reuse this tool to re-check the instructions." + ) + is_mcp: bool = False + is_read_only: bool = False + is_concurrency_safe: bool = True + is_external_tool: bool = False + is_state_injected: bool = True + + def __init__( + self, + groups: list[ToolGroup], + response_template: str, + middlewares: List[ToolMiddlewareBase] | None = None, + ) -> None: + """Initialize the meta tool with the current tool groups.""" + super().__init__(middlewares=middlewares) + self.groups = groups + self.response_template = response_template + + @property + def input_schema(self) -> dict[str, Any]: # type: ignore[override] + """Dynamically generate the input schema based on the current + available tool groups.""" + fields = {} + for group in self.groups: + if group.name == "basic": + continue + fields[group.name] = ( + bool, + Field( + default=False, + description=group.description, + ), + ) + + model = create_model("_DynamicModel", **fields) + schema = model.model_json_schema() + return schema + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """The meta tool is always allowed to be called.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message="The meta tool is always allowed to be called.", + ) + + async def call( + self, + _agent_state: AgentState, + **kwargs: Any, + ) -> ToolChunk: + """Activate or deactivate tool groups based on the input arguments, + and return their usage instructions.""" + if _agent_state is None: + raise DeveloperOrientedException( + "Error: ResetTools requires state to be provided.", + ) + + # Deactivate all tool groups first + _agent_state.tool_context.activated_groups.clear() + + to_activate = [] + for key, value in kwargs.items(): + if not isinstance(value, bool): + return ToolChunk( + content=[ + TextBlock( + text=f"Invalid arguments: the argument {key} " + f"should be a bool value, but got {type(value)}.", + ), + ], + ) + + if value: + to_activate.append(key) + + _agent_state.tool_context.activated_groups.extend(to_activate) + + template = Template(self.response_template) + activated_groups = [_ for _ in self.groups if _.name in to_activate] + return ToolChunk( + content=[ + TextBlock( + text=template.render(groups=activated_groups), + ), + ], + ) diff --git a/src/agentscope/tool/_builtin/_read.py b/src/agentscope/tool/_builtin/_read.py new file mode 100644 index 0000000000000000000000000000000000000000..27b1ccff75b4825f892332080821b4fb430de7c0 --- /dev/null +++ b/src/agentscope/tool/_builtin/_read.py @@ -0,0 +1,285 @@ +# -*- coding: utf-8 -*- +"""The read tool in agentscope.""" +import fnmatch +from typing import Any, List + +from .._base import ToolBase, ToolMiddlewareBase +from ...permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, + PermissionRule, +) +from .._response import ToolChunk +from ...message import TextBlock, ToolResultState +from ...state import AgentState +from ._backend import BackendBase, _normalize_newlines + + +class Read(ToolBase): + """The read tool.""" + + name: str = "Read" + """The tool name presented to the agent.""" + + # pylint: disable=line-too-long + description: str = """Reads a file from the local filesystem. You can access any file directly by using this tool. +Assume this tool is able to read all files on the machine. If the User provides a path to a file assume that path is valid. It is okay to read a file that does not exist; an error will be returned. + +Usage: +- The file_path parameter must be an absolute path, not a relative path +- By default, it reads up to 2000 lines starting from the beginning of the file +- You can optionally specify a line offset and limit (especially handy for long files), but it's recommended to read the whole file by not providing these parameters +- Results are returned using cat -n format, with line numbers starting at 1 +- This tool allows you to read images (eg PNG, JPG, etc). When reading an image file the contents are presented visually as you're a multimodal LLM. +- This tool can read PDF files (.pdf). For large PDFs (more than 10 pages), you MUST provide the pages parameter to read specific pages.""" # noqa: E501 + """The description presented to the agent.""" + + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "file_path": { + "type": "string", + "description": "The absolute path to the file to read.", + }, + "offset": { + "type": "integer", + "description": "Optional 1-based line number to start reading " + "from (default: 1)", + "default": 1, + "minimum": 1, + }, + "limit": { + "type": "integer", + "description": "Optional maximum number of lines to read " + "(default: 2000, max: 2000)", + "default": 2000, + "maximum": 2000, + "minimum": 1, + }, + }, + "required": ["file_path"], + } + + is_mcp: bool = False + is_read_only: bool = True + is_concurrency_safe: bool = True + is_external_tool: bool = False + is_state_injected: bool = True + + def __init__( + self, + max_line_characters: int = 2000, + middlewares: List[ToolMiddlewareBase] | None = None, + backend: BackendBase | None = None, + ) -> None: + """Initialize the read tool. + + Args: + max_line_characters (`int`, defaults to 2000): + The maximum number of characters to include for each line when + reading files. Lines longer than this will be truncated with + a "[truncated]" suffix. This prevents overwhelming the agent + with excessively long lines while still providing useful + content. + middlewares (`List[ToolMiddlewareBase] | None`, optional): + Tool middlewares wrapping the tool execution. + backend (`BackendBase | None`, optional): + The sandbox backend to use for file I/O. When ``None``, + a :class:`LocalBackend` is created. + """ + from ._backend import LocalBackend + + super().__init__(middlewares=middlewares) + self._max_line_characters = max_line_characters + self._backend = backend or LocalBackend() + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permissions for file reading. + + Read is a read-only tool. In EXPLORE mode the engine already handles + the ALLOW via _check_explore_mode, so here we just return PASSTHROUGH + to let the engine continue with rule matching. + """ + return PermissionDecision( + behavior=PermissionBehavior.PASSTHROUGH, + message="File reading is read-only.", + ) + + async def match_rule( + self, + rule_content: str | None, + tool_input: dict[str, Any], + ) -> bool: + """Check if a permission rule matches the file path. + + Matches rule_content as a glob pattern against the "file_path" + parameter using fnmatch. If rule_content is None, matches all + invocations (tool-name-level rule). + + Args: + rule_content (`str | None`): + Glob pattern to match against the file path (e.g., "src/**"), + or None to match all invocations + tool_input (`dict[str, Any]`): + The tool input data containing "file_path" key + + Returns: + `bool`: + True if the glob pattern matches the file path, False otherwise + """ + # None = tool-name-level rule, matches everything + if rule_content is None: + return True + + file_path = tool_input.get("file_path", "") + if not file_path: + return False + return fnmatch.fnmatch(file_path, rule_content) + + async def generate_suggestions( + self, + tool_input: dict[str, Any], + ) -> List[PermissionRule]: + """Generate suggested permission rules for the file path. + + Suggests a glob pattern covering the parent directory of the file, + allowing the user to grant permission for the entire directory at once. + + Args: + tool_input (`dict[str, Any]`): + The tool input data containing "file_path" key + + Returns: + `List[PermissionRule]`: + A single suggested rule covering the parent directory + (e.g., file "/src/main.py" -> rule "src/**") + """ + file_path = tool_input.get("file_path", "") + if not file_path: + return [] + + parent = self._backend.dirname(file_path) + # Glob patterns are POSIX-style strings (matched by fnmatch), + # not real filesystem paths — do NOT use backend.join_path here. + pattern = (parent.rstrip("/\\") + "/**") if parent else "**" + + return [ + PermissionRule( + tool_name=self.name, + rule_content=pattern, + behavior=PermissionBehavior.ALLOW, + source="suggested", + ), + ] + + async def call( # type: ignore[override] + self, + file_path: str, + offset: int = 1, + limit: int = 2000, + _agent_state: AgentState | None = None, + ) -> ToolChunk: + """Read the file and return the content with line numbers.""" + + # Validate file_path is absolute + if not self._backend.isabs(file_path): + return ToolChunk( + content=[ + TextBlock( + text=f"Error: file_path must be an absolute path, " + f"got: {file_path}", + ), + ], + state=ToolResultState.ERROR, + is_last=True, + ) + + # Check file exists + if not await self._backend.file_exists(file_path): + return ToolChunk( + content=[ + TextBlock(text=f"Error: File does not exist: {file_path}"), + ], + state=ToolResultState.ERROR, + is_last=True, + ) + + # Check it's not a directory + if await self._backend.is_dir(file_path): + return ToolChunk( + content=[ + TextBlock( + text=f"Error: Path is a directory, not a file: " + f"{file_path}", + ), + ], + state=ToolResultState.ERROR, + is_last=True, + ) + + try: + # Read file content via backend + lines = None + if _agent_state is not None: + cache = await _agent_state.tool_context.get_cache(file_path) + if cache is not None: + lines = cache.lines + + if lines is None: + raw = await self._backend.read_file(file_path) + content_str = raw.decode("utf-8", errors="replace") + # Normalize CRLF/CR so cached lines end in "\n" regardless + # of the platform the file was written on (Windows text + # files use "\r\n"). + content_str = _normalize_newlines(content_str) + lines = content_str.splitlines(keepends=True) + + # Cache file if state is provided + if _agent_state is not None: + await _agent_state.tool_context.cache_file( + file_path=file_path, + lines=lines, + ) + + # Apply offset and limit (offset is 1-based) + start_idx = offset - 1 + end_idx = start_idx + limit + selected_lines = lines[start_idx:end_idx] + + # Format with line numbers (6-char padded + tab + content) + formatted_lines = [] + for i, line in enumerate(selected_lines, start=offset): + # Remove trailing newline if present + line_content = line.rstrip("\n\r") + + # Truncate lines longer than 2000 chars + if len(line_content) > self._max_line_characters: + line_content = ( + line_content[: self._max_line_characters] + + "[truncated]" + ) + + # Format: 6-char padded line number + tab + content + formatted_line = f"{i:6d}\t{line_content}" + formatted_lines.append(formatted_line) + + # Join all lines + result = "\n".join(formatted_lines) + + return ToolChunk( + content=[TextBlock(text=result)], + state=ToolResultState.RUNNING, + is_last=True, + ) + + except Exception as e: + return ToolChunk( + content=[TextBlock(text=f"Error reading file: {str(e)}")], + state=ToolResultState.ERROR, + is_last=True, + ) diff --git a/src/agentscope/tool/_builtin/_scripts/__init__.py b/src/agentscope/tool/_builtin/_scripts/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d04f77726cc3c97e8b47830c8ba2c1ceea14a3b3 --- /dev/null +++ b/src/agentscope/tool/_builtin/_scripts/__init__.py @@ -0,0 +1,9 @@ +# -*- coding: utf-8 -*- +"""Standalone helper scripts shipped as package resources. + +Scripts in this package are deployed into remote workspaces (Docker / +E2B) at initialization time and invoked via ``exec_shell``. They must +remain importable *without* ``agentscope`` installed — the host reads +them as raw bytes via :mod:`importlib.resources` and ships them into +the workspace environment. +""" diff --git a/src/agentscope/tool/_builtin/_scripts/_glob_helper.py b/src/agentscope/tool/_builtin/_scripts/_glob_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..72eefd95c399cca4f4b3989ed17ea72ff51d364b --- /dev/null +++ b/src/agentscope/tool/_builtin/_scripts/_glob_helper.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""Standalone glob helper script for agentscope builtin tools. + +This script is designed to run **without** agentscope installed. It is +deployed into remote workspaces (Docker / E2B) at initialization time +and invoked via ``exec_shell`` by the :class:`Glob` tool. + +Usage:: + + python3 _glob_helper.py --pattern '**/*.py' --base-dir /workspace + +Output: a JSON array of matching file paths, sorted by modification +time (newest first). Exits with code 0 on success (even when no +matches are found — the array is simply empty). +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import sys + +# ── glob matching (mirrors the logic from Glob tool) ────────────── + + +def _glob_part_to_regex(part: str) -> re.Pattern[str]: + """Convert a single glob pattern segment to a compiled regex. + + Translates glob wildcards (``*`` → ``.*``, ``?`` → ``.``) and + escapes regex meta-characters so that a segment like ``*.py`` + becomes the anchored pattern ``^.*\\.py$``. + + Args: + part: One path segment of a glob pattern (e.g. ``*.py`` + or ``test_*``). + + Returns: + A compiled :class:`re.Pattern` anchored with ``^…$``. + """ + regex_str = "" + for c in part: + if c == "*": + regex_str += ".*" + elif c == "?": + regex_str += "." + elif c in ".^$+{}[]|()\\": + regex_str += "\\" + c + else: + regex_str += c + return re.compile(f"^{regex_str}$") + + +def _collect_all(current_dir: str, results: list[str]) -> None: + """Recursively collect all file paths under *current_dir*. + + Uses :func:`os.walk` to traverse the directory tree. + ``PermissionError`` and ``OSError`` are silently ignored so that + inaccessible subtrees do not abort the entire glob operation. + + Args: + current_dir: Root directory to walk. + results: Accumulator list; matched file paths are appended + in-place. + """ + try: + for root, _dirs, files in os.walk(current_dir): + for fname in files: + results.append(os.path.join(root, fname)) + except (PermissionError, OSError): + pass + + +def _match_parts( + parts: list[str], + part_index: int, + current_dir: str, + results: list[str], +) -> None: + """Recursively match glob pattern *parts* against directory entries. + + Walks the filesystem starting from *current_dir*, consuming one + pattern segment per directory level. The ``**`` segment is handled + specially: it matches zero or more intermediate directories by + recursing into every subdirectory while keeping the same + *part_index*, and also advancing to the next segment in the + current directory. + + Args: + parts: The glob pattern split into path segments + (e.g. ``["src", "**", "*.py"]``). + part_index: Index into *parts* indicating which segment is + being matched at this recursion level. + current_dir: The directory currently being scanned. + results: Accumulator list; matched file paths are appended + in-place. + """ + if part_index >= len(parts): + return + + part = parts[part_index] + is_last = part_index == len(parts) - 1 + + if part == "**": + if is_last: + _collect_all(current_dir, results) + else: + _match_parts(parts, part_index + 1, current_dir, results) + try: + with os.scandir(current_dir) as entries: + for entry in entries: + if entry.is_dir(follow_symlinks=False): + _match_parts( + parts, + part_index, + entry.path, + results, + ) + except (PermissionError, OSError): + pass + else: + regex = _glob_part_to_regex(part) + try: + with os.scandir(current_dir) as entries: + for entry in entries: + if regex.match(entry.name): + full_path = entry.path + if is_last: + if entry.is_file(follow_symlinks=False): + results.append(full_path) + elif entry.is_dir(follow_symlinks=False): + _match_parts( + parts, + part_index + 1, + full_path, + results, + ) + except (PermissionError, OSError): + pass + + +def glob_match(pattern: str, base_dir: str) -> list[str]: + """Match files against a glob pattern starting from *base_dir*. + + Splits *pattern* on path separators (``/`` or ``\\``) and + delegates to :func:`_match_parts` for recursive directory + traversal. Supports ``*`` (any characters within a segment), + ``?`` (single character), and ``**`` (zero or more directories). + + Args: + pattern: Glob pattern such as ``"**/*.py"`` or + ``"src/utils/*.txt"``. + base_dir: Absolute path of the directory to search from. + + Returns: + A list of absolute file paths that match *pattern*. The + list is unsorted; callers should sort as needed (e.g. by + modification time). + """ + results: list[str] = [] + parts = [p for p in re.split(r"[\\/]+", pattern) if p] + _match_parts(parts, 0, base_dir, results) + return results + + +# ── entry point ─────────────────────────────────────────────────── + + +def main() -> None: + """CLI entry point: parse ``--pattern`` and ``--base-dir``, run + the glob, and print results as a JSON array to stdout. + + The results are sorted by file modification time (newest first). + If *base_dir* does not exist, an empty JSON array ``[]`` is + printed and the process exits with code 0. + """ + parser = argparse.ArgumentParser( + description="Glob file matching with mtime sorting.", + ) + parser.add_argument( + "--pattern", + required=True, + help="Glob pattern (e.g. '**/*.py')", + ) + parser.add_argument( + "--base-dir", + required=True, + help="Base directory to search from", + ) + args = parser.parse_args() + + if not os.path.isdir(args.base_dir): + # Empty result for non-existent directory (caller handles + # the "directory not found" error message). + json.dump([], sys.stdout) + return + + matches = glob_match(args.pattern, args.base_dir) + + # Sort by modification time, newest first. + def _mtime(path: str) -> float: + try: + return os.stat(path).st_mtime + except (OSError, FileNotFoundError): + return 0.0 + + matches.sort(key=_mtime, reverse=True) + + json.dump(matches, sys.stdout) + + +if __name__ == "__main__": + main() diff --git a/src/agentscope/tool/_builtin/_skill.py b/src/agentscope/tool/_builtin/_skill.py new file mode 100644 index 0000000000000000000000000000000000000000..54eb1b72a343bd97c080fb67d5f28e1c4d951e1b --- /dev/null +++ b/src/agentscope/tool/_builtin/_skill.py @@ -0,0 +1,127 @@ +# -*- coding: utf-8 -*- +"""The builtin skill viewer tool.""" +from typing import Any, Callable, Awaitable, List + +from ...exception import DeveloperOrientedException +from ...permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, +) +from .._response import ToolChunk +from .._base import ToolBase, ToolMiddlewareBase +from ...skill import Skill +from ...message import TextBlock, ToolResultState +from ...state import AgentState + + +class SkillViewer(ToolBase): + """The builtin skill viewer tool.""" + + name: str = "Skill" + """The name of the skill viewer tool to the agent.""" + + description = ( + "Retrieve a skill within the conversation. " + "When users asks you to perform tasks, check if any of the available " + "skills match. " + "Skills provide specialized capabilities and domain knowledge." + ) + """The tool description of the skill viewer tool to the agent.""" + + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "skill": { + "type": "string", + "description": "The exact name of the skill to view. ", + }, + }, + "required": ["skill"], + } + """The input schema of the skill viewer tool.""" + + is_concurrency_safe: bool = True + """The skill viewer is concurrency safe.""" + + is_external_tool: bool = False + """The skill viewer is not an external tool.""" + + is_state_injected: bool = True + """The skill viewer require state injection to access the activated tool + group.""" + + is_read_only: bool = True + """The skill viewer is read-only.""" + + is_mcp: bool = False + """The skill viewer is not an MCP tool.""" + + mcp_name: str | None = None + """The skill viewer does not belong to any MCP server.""" + + def __init__( + self, + get_skills_method: Callable[..., Awaitable[dict[str, Skill]]], + middlewares: List[ToolMiddlewareBase] | None = None, + ) -> None: + """Initialize the skill viewer with the list of skills. + + Args: + get_skills_method (`Callable[..., dict[str, Skill]]`): + An async method that returns the current skills of the agent. + middlewares (`List[ToolMiddlewareBase] | None`, optional): + Tool middlewares wrapping the tool execution. + """ + super().__init__(middlewares=middlewares) + self._get_skills_method = get_skills_method + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """The skill viewer is always allowed to be called.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message="The skill viewer is always allowed to be called.", + ) + + async def call( + self, + skill: str, + _agent_state: AgentState, + ) -> ToolChunk: + """View the details of the skill with the given name. + + Args: + skill (`str`): + The name of the skill to be viewed. + + Returns: + `ToolChunk`: + The details of the skill. + """ + if not isinstance(_agent_state, AgentState): + raise DeveloperOrientedException( + f"Expected AgentState but got {type(_agent_state)} " + "instead for the Skill viewer tool.", + ) + + # View the activated skills + skills = await self._get_skills_method( + _agent_state.tool_context.activated_groups, + ) + target_skill = skills.get(skill) + if not target_skill: + return ToolChunk( + content=[ + TextBlock( + text=f"SkillNotFoundError: Skill '{skill}' " + f"not found.", + ), + ], + state=ToolResultState.ERROR, + ) + + return ToolChunk(content=[TextBlock(text=target_skill.markdown)]) diff --git a/src/agentscope/tool/_builtin/_write.py b/src/agentscope/tool/_builtin/_write.py new file mode 100644 index 0000000000000000000000000000000000000000..5ce94dd0e0f00c72f59eb0826fb7cbc1c83dac07 --- /dev/null +++ b/src/agentscope/tool/_builtin/_write.py @@ -0,0 +1,328 @@ +# -*- coding: utf-8 -*- +"""The write tool in agentscope.""" +import difflib +import fnmatch +from pathlib import Path +from typing import Any, List + +from .._base import ToolBase, ToolMiddlewareBase +from .._constants import ( + DEFAULT_DANGEROUS_FILES, + DEFAULT_DANGEROUS_DIRECTORIES, +) +from ...permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, + PermissionMode, + PermissionRule, +) +from .._response import ToolChunk +from ...message import TextBlock, ToolResultState +from ...state import AgentState +from ._backend import BackendBase + + +class Write(ToolBase): + """The write tool.""" + + name: str = "Write" + """The tool name presented to the agent.""" + + # pylint: disable=line-too-long + description: str = """Writes a file to the local filesystem. + +Usage: +- This tool will overwrite the existing file if there is one at the provided path. +- If this is an existing file, you MUST use the Read tool first to read the file's contents. This tool will fail if you did not read the file first. +- ALWAYS prefer editing existing files in the codebase. NEVER write new files unless explicitly required. +- NEVER proactively create documentation files (*.md) or README files. Only create documentation files if explicitly requested by the User. +- Only use emojis if the user explicitly requests it. Avoid writing emojis to files unless asked.""" # noqa: E501 + """The description presented to the agent.""" + + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "file_path": { + "type": "string", + "description": "The absolute path to the file to write " + "(must be absolute, not relative)", + }, + "content": { + "type": "string", + "description": "The content to write to the file", + }, + }, + "required": ["file_path", "content"], + } + + is_mcp: bool = False + is_read_only: bool = False + is_concurrency_safe: bool = False + is_external_tool: bool = False + is_state_injected: bool = True + + def __init__( # pylint: disable=dangerous-default-value + self, + dangerous_files: list[str] = DEFAULT_DANGEROUS_FILES, + dangerous_directories: list[str] = DEFAULT_DANGEROUS_DIRECTORIES, + middlewares: List[ToolMiddlewareBase] | None = None, + backend: BackendBase | None = None, + ) -> None: + """Initialize the write tool. + + Args: + dangerous_files (`list[str]`, optional): + Sensitive files that require explicit user confirmation, + even in BYPASS mode. Matched by basename + (case-insensitive). Defaults to `DEFAULT_DANGEROUS_FILES`. + Pass a custom list to fully replace the defaults, or `[]` + to disable the filename check. + dangerous_directories (`list[str]`, optional): + Sensitive directories that require explicit user + confirmation. Matched when any path segment equals an + entry (case-insensitive). Defaults to + `DEFAULT_DANGEROUS_DIRECTORIES`. Pass a custom list to + fully replace the defaults, or `[]` to disable the + directory check. + middlewares (`List[ToolMiddlewareBase] | None`, optional): + Tool middlewares wrapping the tool execution. + backend (`BackendBase | None`, optional): + The sandbox backend to use for file I/O. When ``None``, + a :class:`LocalBackend` is created. + """ + from ._backend import LocalBackend + + super().__init__(middlewares=middlewares) + self.dangerous_files = list(dangerous_files) + self.dangerous_directories = list(dangerous_directories) + + self._backend = backend or LocalBackend() + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permissions for file writing. + + This method implements Write-specific permission checks: + 1. Dangerous path check (safety check, bypass-immune) + 2. ACCEPT_EDITS mode check for files in working directories + + Args: + tool_input (`dict[str, Any]`): + The tool input containing "file_path" key + context (`PermissionContext`): + The permission context with mode and rules + + Returns: + `PermissionDecision`: + ASK for dangerous paths, ALLOW for safe operations in + ACCEPT_EDITS mode, PASSTHROUGH otherwise + """ + + file_path = tool_input.get("file_path") + if not file_path: + return PermissionDecision( + behavior=PermissionBehavior.PASSTHROUGH, + message="No file path provided", + ) + + # 1. Check for dangerous paths (safety check, bypass-immune) + if self._is_dangerous_path(file_path): + return PermissionDecision( + behavior=PermissionBehavior.ASK, + message=f"Permission required: Write operation on " + f"sensitive file {file_path}", + decision_reason="Safety check: dangerous file or directory", + bypass_immune=True, + ) + + # 2. Check ACCEPT_EDITS mode for files in working directories + if context.mode == PermissionMode.ACCEPT_EDITS: + if self._path_in_allowed_working_path(file_path, context): + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message=f"Permission granted for writing {file_path} " + f"(accept edits mode - in working directory)", + decision_reason="File is in working directory and not " + "a dangerous path", + ) + + # 3. Return PASSTHROUGH to let PermissionEngine check allow rules + # This ensures allow rules can grant Write permissions + return PermissionDecision( + behavior=PermissionBehavior.PASSTHROUGH, + message="", + ) + + async def match_rule( + self, + rule_content: str | None, + tool_input: dict[str, Any], + ) -> bool: + """Check if a permission rule matches the file path. + + Matches rule_content as a glob pattern against the "file_path" + parameter using fnmatch. If rule_content is None, matches all + invocations (tool-name-level rule). + + Args: + rule_content (`str | None`): + Glob pattern to match against the file path (e.g., "src/**"), + or None to match all invocations + tool_input (`dict[str, Any]`): + The tool input data containing "file_path" key + + Returns: + `bool`: + True if the glob pattern matches the file path, False otherwise + """ + if rule_content is None: + return True + + file_path = tool_input.get("file_path", "") + if not file_path: + return False + return fnmatch.fnmatch(file_path, rule_content) + + async def generate_suggestions( + self, + tool_input: dict[str, Any], + ) -> List[PermissionRule]: + """Generate suggested permission rules for the file path. + + Suggests a glob pattern covering the parent directory of the file, + allowing the user to grant permission for the entire directory at once. + + Args: + tool_input (`dict[str, Any]`): + The tool input data containing "file_path" key + + Returns: + `List[PermissionRule]`: + A single suggested rule covering the parent directory + (e.g., file "/src/main.py" -> rule "src/**") + """ + file_path = tool_input.get("file_path", "") + if not file_path: + return [] + + parent = self._backend.dirname(file_path) + # Glob patterns are POSIX-style strings (matched by fnmatch), + # not real filesystem paths — do NOT use backend.join_path here. + pattern = (parent.rstrip("/\\") + "/**") if parent else "**" + + return [ + PermissionRule( + tool_name=self.name, + rule_content=pattern, + behavior=PermissionBehavior.ALLOW, + source="suggested", + ), + ] + + async def call( # type: ignore[override] + self, + file_path: str, + content: str, + _agent_state: AgentState | None = None, + ) -> ToolChunk: + """Write content to a file and return the result.""" + # Validate that file_path is absolute + if not self._backend.isabs(file_path): + return ToolChunk( + content=[ + TextBlock( + text=f"Error: file_path must be an absolute path, " + f"got: {file_path}", + ), + ], + state=ToolResultState.ERROR, + is_last=True, + ) + + # Check if file exists, it must be read first if it exists + if ( + await self._backend.file_exists(file_path) + and _agent_state is not None + ): + cache = await _agent_state.tool_context.get_cache(file_path) + if cache is None: + return ToolChunk( + content=[ + TextBlock( + text=f"Error: File {file_path} exists but has not " + f"been read yet. You must read the file first " + f"before writing to it.", + ), + ], + state=ToolResultState.ERROR, + is_last=True, + ) + + # Capture the pre-write content (if any) so we can compute a unified + # diff for the web UI. For brand-new files this stays as an empty + # string, which produces a clean "new file" diff (``--- /dev/null``). + # Track ``file_existed`` separately from ``previous_content`` because + # an *existing* empty file overwrite is not the same as creating a + # new file — the diff header must reflect that. + file_existed = await self._backend.file_exists(file_path) + previous_content = "" + if file_existed: + try: + previous_content = ( + await self._backend.read_file(file_path) + ).decode("utf-8") + except Exception: # pylint: disable=broad-except + # Binary or unreadable file — fall back to empty so we still + # render a best-effort "add" diff in the UI. + previous_content = "" + + # Create parent directories if they don't exist + parent_dir = Path(file_path).parent + await self._backend.exec_shell( + ["mkdir", "-p", str(parent_dir)], + ) + + # Write content to file (backend handles parent dir creation) + await self._backend.write_file( + file_path, + content.encode("utf-8"), + ) + + # Count lines in content + line_count = len(content.split("\n")) + + # Build the unified diff between previous and new content. When the + # file is brand new, ``unified_diff`` over an empty old side naturally + # produces a single "all add" hunk starting at line 1. + diff_text = "".join( + difflib.unified_diff( + previous_content.splitlines(keepends=True), + content.splitlines(keepends=True), + fromfile=( + "/dev/null" if not file_existed else f"a/{file_path}" + ), + tofile=f"b/{file_path}", + n=3, + ), + ) + + # Return success message + return ToolChunk( + content=[ + TextBlock( + text=f"The file {file_path} has been written successfully " + f"({line_count} lines).", + ), + ], + state=ToolResultState.RUNNING, + is_last=True, + metadata={ + "diff": diff_text, + "file_path": file_path, + "occurrences": 1, + }, + ) diff --git a/src/agentscope/tool/_constants.py b/src/agentscope/tool/_constants.py new file mode 100644 index 0000000000000000000000000000000000000000..1ee2c937383a5810948001e366e1ff5ff8db1268 --- /dev/null +++ b/src/agentscope/tool/_constants.py @@ -0,0 +1,110 @@ +# -*- coding: utf-8 -*- +"""Constants for tool permission system.""" + +DEFAULT_DANGEROUS_FILES = [ + ".gitconfig", + ".gitmodules", + ".bashrc", + ".bash_profile", + ".zshrc", + ".zprofile", + ".profile", + ".ssh/config", + ".ssh/authorized_keys", + ".netrc", + ".npmrc", + ".pypirc", + ".env", + ".envrc", + ".env.local", + ".env.development", + ".env.development.local", + ".env.test", + ".env.test.local", + ".env.staging", + ".env.production", + ".env.production.local", +] +# Built-in list of dangerous files that should be protected from auto-editing. +# +# These files can be used for code execution, credential storage, or data +# exfiltration: +# - Shell configuration files: .bashrc, .zshrc, .profile, etc. +# - Git configuration: .gitconfig, .gitmodules +# - SSH configuration: .ssh/config, .ssh/authorized_keys +# - Credential files: .netrc, .npmrc, .pypirc +# - Environment / secret files: .env and common variants (.envrc for direnv, +# .env.local / .env.production / etc. for framework-specific overrides) + + +DEFAULT_DANGEROUS_DIRECTORIES = [ + ".git", + ".vscode", + ".idea", + ".ssh", +] +# Built-in list of dangerous directories that should be protected from +# auto-editing. +# +# These directories contain sensitive configuration or executable files: +# - .git: Git repository metadata +# - .vscode: VS Code configuration +# - .idea: JetBrains IDE configuration +# - .ssh: SSH keys and configuration + + +DANGEROUS_COMMANDS = [ + "rm -rf", + "sudo rm", + "dd", + "mkfs", + "fdisk", + "format", + "chmod 777", + "chmod -R 777", + "chown -R", + "kill -9", + "> /dev/", +] +# Built-in list of dangerous command patterns that require explicit +# user approval. +# +# These commands can cause data loss, system damage, or security issues: +# - rm -rf: Recursive force deletion +# - sudo rm: Deletion with elevated privileges +# - dd: Direct disk operations +# - mkfs: Format filesystem +# - fdisk: Disk partitioning +# - format: Format disk +# - chmod 777: Overly permissive file permissions +# - chown -R: Recursive ownership changes +# - kill -9: Force kill processes +# - > /dev/: Writing to device files + + +DANGEROUS_NODE_TYPES = { + "command_substitution", # $(...) or `...` + "process_substitution", # <(...) + "expansion", # ${VAR:-default} complex expansion + "subshell", # (...) + "for_statement", # for loops + "while_statement", # while loops + "until_statement", # until loops + "if_statement", # if conditionals + "case_statement", # case statements + "function_definition", # function definitions + "test_command", # [[ ... ]] test commands +} +# Node types that indicate commands cannot be statically analyzed. +# These either execute arbitrary code or expand to values we can't +# determine statically. +# +# When these are detected, the command requires user review because: +# - Command substitution $(cmd) executes arbitrary commands +# - Process substitution <(cmd) creates dynamic file descriptors +# - Complex expansions ${VAR:-default} have runtime-dependent values +# - Control flow (if/while/for) has conditional execution paths +# - Subshells (...) create new execution contexts +# +# Note: simple_expansion ($VAR) is handled separately with allowlist +# for known-safe environment variables ($HOME, $PWD, etc.) diff --git a/src/agentscope/tool/_response.py b/src/agentscope/tool/_response.py new file mode 100644 index 0000000000000000000000000000000000000000..5abed2e5f5a0510928b8e4cad08717da21dc8cb7 --- /dev/null +++ b/src/agentscope/tool/_response.py @@ -0,0 +1,164 @@ +# -*- coding: utf-8 -*- +"""The tool response class.""" +import base64 +import binascii +from typing import List, Literal, Self + +from pydantic import BaseModel, Field + +from .._utils._common import _generate_id +from ..message import DataBlock, TextBlock, Base64Source, ToolResultState + + +def _merge_base64_chunks(existing: str, incoming: str) -> str: + """Merge independently encoded base64 chunks without corrupting padding.""" + try: + merged = base64.b64decode( + existing, + validate=True, + ) + base64.b64decode(incoming, validate=True) + except (binascii.Error, ValueError): + # Keep compatibility with callers/tests that used placeholder strings + # instead of valid base64 payloads. + return existing + incoming + + return base64.b64encode(merged).decode("ascii") + + +class ToolChunk(BaseModel): + """The tool result chunk from a tool execution.""" + + content: List[TextBlock | DataBlock] + """The chunk data blocks, note for one multimodal data, the DataBlock + instance should have the same block id, so that the agent can group them + together.""" + + state: ToolResultState = ToolResultState.RUNNING + """The execution state of the tool chunk.""" + + is_last: bool = True + """Whether this is the last response in a stream tool execution.""" + + metadata: dict = Field(default_factory=dict) + """The metadata to be accessed within the agent, so that we don't need to + parse the tool result block.""" + + id: str = Field(default_factory=_generate_id) + """The identity of the tool response.""" + + +class ToolResponse(BaseModel): + """The tool response from a tool execution, which contains the completed + tool result (compared to ToolChunk).""" + + content: List[TextBlock | DataBlock] = Field(default_factory=list) + """The completed tool result data blocks.""" + + state: Literal[ + ToolResultState.ERROR, + ToolResultState.DENIED, + ToolResultState.INTERRUPTED, + ToolResultState.SUCCESS, + ] = ToolResultState.SUCCESS + """The execution state of the tool response.""" + + metadata: dict = Field(default_factory=dict) + """The metadata to be accessed within the agent, so that we don't need to + parse the tool result block.""" + + id: str = Field(default_factory=_generate_id) + """The identity of the tool response.""" + + def append_chunk(self, chunk: ToolChunk) -> Self: + """Append a tool chunk to the current tool response, accumulate the + data blocks and update the state and metadata.""" + + # Update content blocks + current_ids_to_index = { + _.id: index for index, _ in enumerate(self.content) + } + for chunk_block in chunk.content: + if chunk_block.id in current_ids_to_index: + # Append to the existing block + target_block = self.content[ + current_ids_to_index[chunk_block.id] + ] + + if isinstance(target_block, TextBlock) and isinstance( + chunk_block, + TextBlock, + ): + target_block.text += chunk_block.text + elif isinstance(target_block, DataBlock) and isinstance( + chunk_block, + DataBlock, + ): + if isinstance( + target_block.source, + Base64Source, + ) and isinstance(chunk_block.source, Base64Source): + # Accumulate independently encoded base64 chunks. + target_block.source.data = _merge_base64_chunks( + target_block.source.data, + chunk_block.source.data, + ) + # Update the newest media type and name if provided + target_block.name = ( + chunk_block.name or target_block.name + ) + target_block.source.media_type = ( + chunk_block.source.media_type + or target_block.source.media_type + ) + else: + raise ValueError( + "Cannot append DataBlock with URL source or " + f"different source types: {target_block.source} " + f"vs {chunk_block.source}", + ) + else: + # For different block types with the same ID, we just + # append the new block with a new ID to avoid the conflict + new_chunk_block = chunk_block.model_copy(deep=True) + new_chunk_block.id = _generate_id() + self.content.append(new_chunk_block) + + else: + # Append a copy to avoid modifying the original chunk + self.content.append(chunk_block.model_copy(deep=True)) + + # Update the index mapping for the new block + current_ids_to_index[chunk_block.id] = len(self.content) - 1 + + # Update id, state and metadata + # Only reserve the failure state and keep the previous state if not + # worse. + if chunk.state == ToolResultState.ERROR: + self.state = ToolResultState.ERROR + elif chunk.state == "interrupted": + self.state = ToolResultState.INTERRUPTED + elif chunk.state == ToolResultState.DENIED: + self.state = ToolResultState.DENIED + + self.metadata.update(chunk.metadata) + + # Post-processing: merge consecutive TextBlocks + # DataBlocks are kept separate and only merged by explicit id matching + merged_content: List[TextBlock | DataBlock] = [] + for block in self.content: + if isinstance(block, TextBlock) and merged_content: + # Check if the last block is also a TextBlock + last_block = merged_content[-1] + if isinstance(last_block, TextBlock): + # Merge consecutive TextBlocks + last_block.text += block.text + else: + # Last block is DataBlock, append current TextBlock + merged_content.append(block) + else: + # First block or current block is DataBlock + merged_content.append(block) + + self.content = merged_content + + return self diff --git a/src/agentscope/tool/_task/__init__.py b/src/agentscope/tool/_task/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8c1cb1cbd59ffcfe51e1a10c0ab203066ee72f1f --- /dev/null +++ b/src/agentscope/tool/_task/__init__.py @@ -0,0 +1,13 @@ +# -*- coding: utf-8 -*- +"""Task planning tools for agents.""" +from ._create_task import TaskCreate +from ._get_task import TaskGet +from ._list_task import TaskList +from ._update_task import TaskUpdate + +__all__ = [ + "TaskCreate", + "TaskGet", + "TaskList", + "TaskUpdate", +] diff --git a/src/agentscope/tool/_task/_create_task.py b/src/agentscope/tool/_task/_create_task.py new file mode 100644 index 0000000000000000000000000000000000000000..6ba46202e3ceb8fd5d5f088d5d277ac17ed114ad --- /dev/null +++ b/src/agentscope/tool/_task/_create_task.py @@ -0,0 +1,132 @@ +# -*- coding: utf-8 -*- +"""The creating task tool class.""" +from typing import Any + +from pydantic import BaseModel, Field + +from ._task_tool_base import _TaskToolBase +from .._response import ToolChunk +from ...state import AgentState, Task +from ...exception import DeveloperOrientedException +from ...message import TextBlock, ToolResultState + + +class _TaskCreateParams(BaseModel): + """The params of the creating task tool.""" + + subject: str = Field(description="A brief title for the task") + description: str = Field(description="What needs to be done") + metadata: dict[str, Any] | None = Field( + default=None, + description="Arbitrary metadata to attach to the task", + ) + + +class TaskCreate(_TaskToolBase): + """Create a task for the agent to perform.""" + + name: str = "TaskCreate" + + description: str = """Use this tool to create a structured task list for \ +your current session. This helps you track progress, organize complex tasks, \ +and demonstrate thoroughness to the user. +It also helps the user understand the progress of the task and overall \ +progress of their requests. + +## When to Use This Tool +Use this tool proactively in these scenarios: + +- Complex multi-step tasks - When a task requires 3 or more distinct steps \ +or actions +- Non-trivial and complex tasks - Tasks that require careful planning or \ +multiple operations +- Plan mode - When using plan mode, create a task list to track the work +- User explicitly requests todo list - When the user directly asks you to \ +use the todo list +- User provides multiple tasks - When users provide a list of things to be \ +done (numbered or comma-separated) +- After receiving new instructions - Immediately capture user requirements \ +as tasks +- When you start working on a task - Mark it as in_progress BEFORE \ +beginning work +- After completing a task - Mark it as completed and add any new follow-up \ +tasks discovered during implementation + +## When NOT to Use This Tool + +Skip using this tool when: +- There is only a single, straightforward task +- The task is trivial and tracking it provides no organizational benefit +- The task can be completed in less than 3 trivial steps +- The task is purely conversational or informational + +NOTE that you should not use this tool if there is only one trivial task to \ +do. In this case you are better off just doing the task directly. + +## Task Fields + +- **subject**: A brief, actionable title in imperative form (e.g., \ +"Fix authentication bug in login flow") +- **description**: What needs to be done + +All tasks are created with status `pending`. + +## Tips + +- Create tasks with clear, specific subjects that describe the outcome +- After creating tasks, use TaskUpdate to set up dependencies \ +(blocks/blockedBy) if needed +- Check TaskList first to avoid creating duplicate tasks""" + + input_schema: dict = _TaskCreateParams.model_json_schema() + + async def call( + self, + _agent_state: AgentState, + subject: str, + description: str, + metadata: dict[str, Any] | None = None, + ) -> ToolChunk: + """Create the subtask and add it into the agent state.""" + if not isinstance(_agent_state, AgentState): + # Expose error to the developer + raise DeveloperOrientedException( + f"Error: TaskCreate requires AgentState to be provided, got " + f"{_agent_state} instead.", + ) + + try: + # Derive the next sequential id from existing tasks. + # Existing ids that look numeric are considered; any + # non-numeric ids (e.g. legacy UUIDs) are ignored. + max_numeric = 0 + for t in _agent_state.tasks_context.tasks: + try: + max_numeric = max(max_numeric, int(t.id)) + except (ValueError, TypeError): + pass + next_id = str(max_numeric + 1) + + task = Task( + id=next_id, + subject=subject, + description=description, + metadata=metadata or {}, + ) + _agent_state.tasks_context.tasks.append(task) + + return ToolChunk( + content=[ + TextBlock( + text=f"Task (id={next_id}) created successfully: " + f"{task.subject}", + ), + ], + ) + except Exception as e: + return ToolChunk( + content=[ + TextBlock(text=f"CreateTaskError: {e}"), + ], + state=ToolResultState.ERROR, + ) diff --git a/src/agentscope/tool/_task/_get_task.py b/src/agentscope/tool/_task/_get_task.py new file mode 100644 index 0000000000000000000000000000000000000000..d7826fda7a44be3792e689c0edd9118a203ef1a9 --- /dev/null +++ b/src/agentscope/tool/_task/_get_task.py @@ -0,0 +1,100 @@ +# -*- coding: utf-8 -*- +"""The get task tool class.""" +from pydantic import BaseModel, Field + +from ._task_tool_base import _TaskToolBase +from .._response import ToolChunk +from ...state import AgentState +from ...exception import DeveloperOrientedException +from ...message import TextBlock, ToolResultState + + +class _TaskGetParams(BaseModel): + """The params of the get task.""" + + task_id: str = Field(description="The ID of the task to retrieve") + + +class TaskGet(_TaskToolBase): + """Retrieve a task by its ID from the task list.""" + + name: str = "TaskGet" + + description: str = """Use this tool to retrieve a task by its ID from the task list. + +## When to Use This Tool + +- When you need the full description and context before starting work on a task +- To understand task dependencies (what it blocks, what blocks it) +- After being assigned a task, to get complete requirements + +## Output + +Returns full task details: +- **subject**: Task title +- **description**: Detailed requirements and context +- **status**: 'pending', 'in_progress', or 'completed' +- **blocks**: Tasks waiting on this one to complete +- **blockedBy**: Tasks that must complete before this one can start + +## Tips + +- After fetching a task, verify its blockedBy list is empty before beginning work. +- Use TaskList to see all tasks in summary form.""" # noqa: E501 + + input_schema: dict = _TaskGetParams.model_json_schema() + + async def call( + self, + task_id: str, + _agent_state: AgentState, + ) -> ToolChunk: + """Retrieve a task by its ID.""" + if not isinstance(_agent_state, AgentState): + # Expose error to the developer + raise DeveloperOrientedException( + f"Error: TaskGet requires AgentState to be provided, got " + f"{_agent_state} instead.", + ) + + # Find the task by ID + task = None + for t in _agent_state.tasks_context.tasks: + if t.id == task_id: + task = t + break + + if task is None: + return ToolChunk( + content=[ + TextBlock(text="Task not found"), + ], + state=ToolResultState.ERROR, + ) + + # Build the response + lines = [ + f"Task (id={task.id}): {task.subject}", + f"Status: {task.state}", + f"Description: {task.description}", + ] + + if task.owner: + lines.append(f"Owner: {task.owner}") + + if task.blocked_by: + blocked_by_str = ", ".join([f"#{bid}" for bid in task.blocked_by]) + lines.append(f"Blocked by: {blocked_by_str}") + + if task.blocks: + blocks_str = ", ".join([f"#{bid}" for bid in task.blocks]) + lines.append(f"Blocks: {blocks_str}") + + if task.metadata: + lines.append(f"Metadata: {task.metadata}") + + return ToolChunk( + content=[ + TextBlock(text="\n".join(lines)), + ], + ) diff --git a/src/agentscope/tool/_task/_list_task.py b/src/agentscope/tool/_task/_list_task.py new file mode 100644 index 0000000000000000000000000000000000000000..1d334555daae31616feec8afdd239fc7f668cfad --- /dev/null +++ b/src/agentscope/tool/_task/_list_task.py @@ -0,0 +1,74 @@ +# -*- coding: utf-8 -*- +"""The task list tool class.""" + +from ._task_tool_base import _TaskToolBase +from .._response import ToolChunk +from .._base import ParamsBase +from ...state import AgentState +from ...exception import DeveloperOrientedException +from ...message import TextBlock + + +class _TaskListParams(ParamsBase): + """The params of the list task params.""" + + +class TaskList(_TaskToolBase): + """List tasks for the agent to perform.""" + + name: str = "TaskList" + + # pylint: disable=line-too-long + description: str = """Use this tool to list all tasks in the task list. + +## When to Use This Tool +- To see what tasks are available to work on (status: 'pending', no owner, not blocked) +- To check overall progress on the project +- To find tasks that are blocked and need dependencies resolved +- After completing a task, to check for newly unblocked work or claim the next available task +- **Prefer working on tasks in ID order** (lowest ID first) when multiple tasks are available, as earlier tasks often set up context for later ones + +## Output + +Returns a summary of each task: +- **id**: Task identifier (use with TaskGet, TaskUpdate) +- **subject**: Brief description of the task +- **status**: 'pending', 'in_progress', or 'completed' +- **owner**: Agent ID if assigned, empty if available +- **blockedBy**: List of open task IDs that must be resolved first (tasks with blockedBy cannot be claimed until dependencies resolve) + +Use TaskGet with a specific task ID to view full details including description and comments.""" # noqa: E501 + + input_schema: dict = _TaskListParams.model_json_schema() + + async def call(self, _agent_state: AgentState) -> ToolChunk: + """List tasks for the agent to perform.""" + if not isinstance(_agent_state, AgentState): + # Expose error to the developer + raise DeveloperOrientedException( + f"Error: TaskList requires AgentState to be provided, got " + f"{_agent_state} instead.", + ) + + if len(_agent_state.tasks_context.tasks) == 0: + return ToolChunk( + content=[TextBlock(text="No tasks available.")], + ) + + tasks = [] + for task in _agent_state.tasks_context.tasks: + owner = f"({task.owner})" if task.owner else "" + blocked = ( + f'[blocked by {", ".join(task.blocked_by)}]' + if task.blocked_by + else "" + ) + tasks.append( + f"{task.id} [{task.state}] {task.subject}{owner}{blocked}", + ) + + return ToolChunk( + content=[ + TextBlock(text="\n".join(tasks)), + ], + ) diff --git a/src/agentscope/tool/_task/_task_tool_base.py b/src/agentscope/tool/_task/_task_tool_base.py new file mode 100644 index 0000000000000000000000000000000000000000..4ce4271e8daaa0a1df06428bb34a7c4bfdd58206 --- /dev/null +++ b/src/agentscope/tool/_task/_task_tool_base.py @@ -0,0 +1,42 @@ +# -*- coding: utf-8 -*- +"""The task tool base class, providing unified interface and permission +check for builtin task related tools.""" +from typing import Any + +from .._base import ToolBase +from ...permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, +) + + +class _TaskToolBase(ToolBase): + name: str + + description: str + + input_schema: dict + + is_concurrency_safe: bool = True + + is_read_only: bool = False + + is_state_injected: bool = True + + is_external_tool: bool = False + + is_mcp: bool = False + + mcp_name: str | None = None + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permission for the tool usage.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message=f"{self.name} is always allowed to be called.", + ) diff --git a/src/agentscope/tool/_task/_update_task.py b/src/agentscope/tool/_task/_update_task.py new file mode 100644 index 0000000000000000000000000000000000000000..c295d890b2fe69e8306543215d65a464a4d96c79 --- /dev/null +++ b/src/agentscope/tool/_task/_update_task.py @@ -0,0 +1,293 @@ +# -*- coding: utf-8 -*- +"""The task updated tool class.""" +from typing import Literal + +from pydantic import BaseModel, Field + +from ._task_tool_base import _TaskToolBase +from .._response import ToolChunk +from ...state import AgentState +from ...exception import DeveloperOrientedException +from ...message import TextBlock, ToolResultState + + +class _TaskUpdateParams(BaseModel): + """The params of the update task.""" + + task_id: str = Field(description="The task id.") + subject: str | None = Field( + default=None, + description="New subject for the task", + ) + description: str | None = Field( + default=None, + description="New description for the task", + ) + add_blocks: list[str] | None = Field( + default=None, + description="Task IDs that this task blocks", + ) + status: Literal[ + "pending", + "in_progress", + "completed", + "deleted", + ] | None = Field( + default=None, + description="New status for the task", + ) + add_blocked_by: list[str] | None = Field( + default=None, + description="Task IDs that block this task", + ) + owner: str | None = Field( + default=None, + description="New owner for the task", + ) + metadata: dict | None = Field( + default=None, + description="Metadata keys to merge into the task. " + "Set a key to null to delete it.", + ) + + +class TaskUpdate(_TaskToolBase): + """The tool to update the agent task.""" + + name: str = "TaskUpdate" + + description: str = """Use this tool to update a task in the task list. + +## When to Use This Tool + +**Mark tasks as resolved:** +- When you have completed the work described in a task +- When a task is no longer needed or has been superseded +- IMPORTANT: Always mark your assigned tasks as resolved when you finish them +- After resolving, call TaskList to find your next task + +- ONLY mark a task as completed when you have FULLY accomplished it +- If you encounter errors, blockers, or cannot finish, keep the task as in_progress +- When blocked, create a new task describing what needs to be resolved +- Never mark a task as completed if: + - Tests are failing + - Implementation is partial + - You encountered unresolved errors + - You couldn't find necessary files or dependencies + +**Delete tasks:** +- When a task is no longer relevant or was created in error +- Setting status to `deleted` permanently removes the task + +**Update task details:** +- When requirements change or become clearer +- When establishing dependencies between tasks + +## Fields You Can Update + +- **status**: The task status (see Status Workflow below) +- **subject**: Change the task title (imperative form, e.g., "Run tests") +- **description**: Change the task description +- **owner**: Change the task owner (agent name) +- **metadata**: Merge metadata keys into the task (set a key to null to delete it) +- **add_blocks**: Mark tasks that cannot start until this one completes +- **add_blocked_by**: Mark tasks that must complete before this one can start + +## Status Workflow + +Status progresses: `pending` → `in_progress` → `completed` + +Use `deleted` to permanently remove a task. + +## Staleness + +Make sure to read a task's latest state using `TaskGet` before updating it. + +## Examples + +Mark task as in progress when starting work: +```json +{"task_id": "1", "status": "in_progress"} +``` + +Mark task as completed after finishing work: +```json +{"task_id": "1", "status": "completed"} +``` + +Delete a task: +```json +{"task_id": "1", "status": "deleted"} +``` + +Claim a task by setting owner: +```json +{"task_id": "1", "owner": "my-name"} +``` + +Set up task dependencies: +```json +{"task_id": "2", "add_blocked_by": ["1"]} +```""" # noqa: E501 + + input_schema: dict = _TaskUpdateParams.model_json_schema() + + async def call( + self, + _agent_state: AgentState, + task_id: str, + subject: str | None = None, + description: str | None = None, + add_blocks: list[str] | None = None, + status: Literal["pending", "completed", "in_progress", "deleted"] + | None = None, + add_blocked_by: list[str] | None = None, + owner: str | None = None, + metadata: dict | None = None, + ) -> ToolChunk: + """Update the agent task.""" + if not isinstance(_agent_state, AgentState): + # Expose error to the developer + raise DeveloperOrientedException( + f"Error: {self.name} requires AgentState to be provided, got " + f"{_agent_state} instead.", + ) + + index = None + for i, task in enumerate(_agent_state.tasks_context.tasks): + if task.id == task_id: + index = i + + if index is None: + return ToolChunk( + content=[ + TextBlock( + text=f"TaskNotFoundError: " + f"The task (id={task_id}) does not exist.", + ), + ], + state=ToolResultState.ERROR, + ) + + updated_fields = [] + + if subject: + updated_fields.append("subject") + _agent_state.tasks_context.tasks[index].subject = subject + + if description is not None: + updated_fields.append("description") + _agent_state.tasks_context.tasks[index].description = description + + existed_ids = [_.id for _ in _agent_state.tasks_context.tasks] + if add_blocks: + current_blocks = _agent_state.tasks_context.tasks[index].blocks + new_blocks = [ + _ + for _ in add_blocks + if _ not in current_blocks and _ in existed_ids + ] + if new_blocks: + updated_fields.append("add_blocks") + for block_id in new_blocks: + self._update_block_relation( + task_id, + block_id, + _agent_state, + ) + + if add_blocked_by is not None: + current_blocked_by = _agent_state.tasks_context.tasks[ + index + ].blocked_by + new_blocked_by = [ + _ + for _ in add_blocked_by + if _ not in current_blocked_by and _ in existed_ids + ] + if new_blocked_by: + updated_fields.append("add_blocked_by") + for blocked_by_id in new_blocked_by: + self._update_block_relation( + blocked_by_id, + task_id, + _agent_state, + ) + + if status: + if status == "deleted": + # Permanently remove the task + _agent_state.tasks_context.tasks.pop(index) + # Remove task id from all the blocks and blocked_by + for task in _agent_state.tasks_context.tasks: + if task_id in task.blocks: + task.blocks.remove(task_id) + + if task_id in task.blocked_by: + task.blocked_by.remove(task_id) + return ToolChunk( + content=[ + TextBlock( + text=f"Task (id={task_id}) has been deleted.", + ), + ], + ) + # Update the status + updated_fields.append("status") + _agent_state.tasks_context.tasks[index].state = status + + if owner is not None: + updated_fields.append("owner") + _agent_state.tasks_context.tasks[index].owner = owner + + if metadata: + updated_fields.append("metadata") + for k, v in metadata.items(): + if v is None: + _agent_state.tasks_context.tasks[index].metadata.pop( + k, + None, + ) + else: + _agent_state.tasks_context.tasks[index].metadata[k] = v + + if updated_fields: + res = f'Update task (id={task_id}) {", ".join(updated_fields)}.' + else: + res = ( + f"No updates were made to the task (id={task_id}). " + f"Make sure you provided at least one field to update and " + f"the values are correct." + ) + + if _agent_state.tasks_context.tasks[index].state == "completed": + res += ( + "\n\nTask completed. Call TaskList now to find your next " + "available task or see if your work unblocked others." + ) + + return ToolChunk(content=[TextBlock(text=res)]) + + @staticmethod + def _update_block_relation( + block_id: str, + blocked_by_id: str, + _agent_state: AgentState, + ) -> None: + """Update the block relationship between the tasks. + + Args: + block_id (`str`): + The id of the task that blocks the other tasks. + blocked_by_id (`str`): + The id of the task blocked by the task of `block_id`. + _agent_state (`AgentState`): + The agent state to update. + """ + # Update the blocks + for task in _agent_state.tasks_context.tasks: + if task.id == block_id and blocked_by_id not in task.blocks: + task.blocks.append(blocked_by_id) + + if task.id == blocked_by_id and block_id not in task.blocked_by: + task.blocked_by.append(block_id) diff --git a/src/agentscope/tool/_tool_group.py b/src/agentscope/tool/_tool_group.py new file mode 100644 index 0000000000000000000000000000000000000000..f6d879c5cb07f8bff65d5eb42494982cdbbccad4 --- /dev/null +++ b/src/agentscope/tool/_tool_group.py @@ -0,0 +1,108 @@ +# -*- coding: utf-8 -*- +"""The tool group class.""" +from typing import Literal, Sequence + +from ..mcp import MCPClient +from ._base import ToolBase +from ..skill import SkillLoaderBase, Skill, LocalSkillLoader + + +class ToolGroup: + """A group of related tools, mcps, and skills that an agent can activate, + deactivate and use together. The tool group is activated by the meta tool + `ResetTools`. + + In high-code scenarios, the tools argument accepts any child classes for + ToolBase class, and the tool groups supports serialization. + """ + + name: Literal["basic"] | str + """Note the "basic" group is special and represents the default tool group + that will be always be activated for the agent.""" + + description: str + """A description of the tool group from an agent-oriented perspective, + outlining its capabilities and the conditions under which it should be + activated.""" + + instructions: str | None + """Instructions included in the meta tool's result upon activation of + this tool group, guiding the agent on how to properly use the meta tool.""" + + tools: list[ToolBase] + """The tools in this group.""" + + skills_or_loaders: list[Skill | SkillLoaderBase] + """The skills in this group.""" + + mcps: list[MCPClient] + """The mcps in this group.""" + + def __init__( + self, + name: Literal["basic"] | str, + description: str | None = None, + instructions: str | None = None, + tools: list[ToolBase] | None = None, + skills_or_loaders: Sequence[str | Skill | SkillLoaderBase] + | None = None, + mcps: list[MCPClient] | None = None, + ) -> None: + """Initialize the tool group. + + Args: + name (`Literal["basic"] | str): + The name of the tool group. + description (`str | None`): + The description of the tool group. + instructions (`str | None`, optional): + Instructions included in the meta tool's result upon + activation of this tool group, guiding the agent on how to + properly use the meta tool. + tools (`list[ToolBase] | None`, optional): + The tools in this group. + skills_or_loaders (`list[str | Skill | SkillLoaderBase] | None`, \ + optional): + The skill paths, data, and loaders to access skills in this + group. + mcps (`list[MCPClient] | None`, optional): + The mcps in this group. + """ + if name != "basic" and description is None: + raise ValueError( + f"The tool group description is required for tool group " + f"'{name}' (Only the 'basic' tool group can have an optional " + f"description).", + ) + + self.name = name + self.description = description or "" + self.instructions = instructions + self.tools = tools or [] + self.mcps = mcps or [] + + # Skill + self.skills_or_loaders = [] + for _ in skills_or_loaders or []: + if isinstance(_, str): + self.skills_or_loaders.append(LocalSkillLoader(directory=_)) + + elif isinstance(_, (Skill, SkillLoaderBase)): + self.skills_or_loaders.append(_) + + else: + raise TypeError( + f"Invalid skill or loader: {_}. Must be a skill, " + f"skill loader, or directory path.", + ) + + async def list_skills(self) -> list[Skill]: + """List all the skills in this tool group.""" + skills = [] + for skill_or_loader in self.skills_or_loaders: + if isinstance(skill_or_loader, Skill): + skills.append(skill_or_loader) + elif isinstance(skill_or_loader, SkillLoaderBase): + skills.extend(await skill_or_loader.list_skills()) + + return skills diff --git a/src/agentscope/tool/_toolkit.py b/src/agentscope/tool/_toolkit.py new file mode 100644 index 0000000000000000000000000000000000000000..3957be63f974ff178991c11a244afa5214f61dc0 --- /dev/null +++ b/src/agentscope/tool/_toolkit.py @@ -0,0 +1,614 @@ +# -*- coding: utf-8 -*- +"""The toolkit class for tool calls in AgentScope.""" +import asyncio +import inspect +from collections import OrderedDict +from typing import ( + AsyncGenerator, + Type, + Generator, + Sequence, +) + +import mcp +from jinja2 import Template +from pydantic import ( + BaseModel, + Field, + create_model, +) + +from ._builtin import ResetTools, SkillViewer +from ._base import ToolBase +from ._response import ToolResponse, ToolChunk +from ..skill import SkillLoaderBase, Skill +from ._types import RegisteredTool +from .._utils._common import _json_loads_with_repair +from ..exception import ( + DeveloperOrientedException, + ToolNotFoundError, + ToolGroupInactiveError, +) +from ..mcp import MCPClient +from ..message import ( + ToolCallBlock, + TextBlock, + ToolResultState, +) +from ._tool_group import ToolGroup +from .._logging import logger +from ..state import AgentState + + +# pylint: disable=line-too-long +DEFAULT_META_TOOL_RESPONSE_TEMPLATE = """{% if groups | length == 0 %}All tool groups are currently deactivated.{% else %}The currently activated tool group(s): {{ groups | map(attribute='name') | join(', ') }}.{% if groups | selectattr('instructions', 'ne', None) | list | length > 0 %} + +The tool instructions are a collection of suggestions, rules and notifications about how to use the tools in the activated groups. +{% for group in groups %}{% if group.instructions %}{{ group.instructions }}{% endif %}{% endfor %} +{% endif %}{% endif %}""" # noqa: E501 + + +DEFAULT_SKILL_INSTRUCTION = """ +Skills are a collection of instructions, scripts, and resources to extend your capabilities. + +**IMPORTANT**: Skills are NOT tools, and you cannot call a skill directly. To use a skill, you MUST use the `{{ skill_viewer }}` tool to read the skill's full instructions, and then follow those instructions to use the tools and resources provided by the skill. + +# Available Skills:{% for skill in skills %} + +{{ skill.name }} +{{ skill.description }} +{{ skill.dir }} +{% endfor %} + +""" # noqa: E501 + + +class Toolkit: + """Toolkit is the core module to register, manage and delete tool + functions, MCP clients, Agent skills in AgentScope. + + About tool functions: + + - Register and parse JSON schemas from their docstrings automatically. + - Group-wise tools management, and agentic tools activation/deactivation. + - Extend the tool function JSON schema dynamically with Pydantic BaseModel. + - Tool function execution with unified streaming interface. + + About MCP clients: + + - Register tool functions from MCP clients directly. + - Client-level tool functions removal. + + About Agent skills: + + - Register agent skills from the given directory. + - Provide prompt for the registered skills to the agent. + """ + + def __init__( + self, + tools: list[ToolBase] | None = None, + skills_or_loaders: Sequence[str | Skill | SkillLoaderBase] + | None = None, + mcps: list[MCPClient] | None = None, + tool_groups: list[ToolGroup] | None = None, + meta_tool_response_template: str = DEFAULT_META_TOOL_RESPONSE_TEMPLATE, + skill_instruction_template: str = DEFAULT_SKILL_INSTRUCTION, + ) -> None: + """Initialize the toolkit. + + Args: + tools (`list[ToolBase] | None`, optional): + The tool objects that belong to the "basic" tool group. + skills_or_loaders (`list[str | Skill | SkillLoaderBase] | None`, \ + optional): + The agent skill directories to be registered in the "basic" + tool group. + mcps (`list[MCPClient] | None`, optional): + The mcp clients to be registered in the "basic" tool group. + tool_groups (`list[ToolGroup] | None`, optional): + The tool groups to be registered. + meta_tool_response_template (`str`, optional): + The template for meta tool responses. + skill_instruction_template (`str`): + A Jinja2 template for generating the agent skill instruction. + """ + + if tool_groups is not None and any( + _.name == "basic" for _ in tool_groups + ): + raise ValueError( + "The 'basic' tool group is reserved for the default tool " + "group. Don't include 'basic' in the tool_groups argument " + "when you also provide tools, skills or mcps in the " + "constructor.", + ) + + self.tool_groups = [ + ToolGroup( + name="basic", + tools=tools or [], + skills_or_loaders=skills_or_loaders or [], + mcps=mcps or [], + ), + ] + (tool_groups or []) + + # Check name conflict for tool groups + if len(set(_.name for _ in self.tool_groups)) != len( + self.tool_groups, + ): + raise ValueError( + "Tool groups must not contain duplicate tool groups.", + ) + + # The stateful MCP clients should be initialized already + for group in self.tool_groups: + for client in group.mcps: + if client.is_stateful and not client.is_connected: + raise ValueError( + f"The MCP client '{client.name}' is stateful, but " + f"not connected.", + ) + + self.meta_tool_response_template = meta_tool_response_template + self.skill_instruction_template = skill_instruction_template + + self.builtin_meta_tool = RegisteredTool( + tool=ResetTools( + # An inference value for groups so that it can generate the + # corresponding input schema. + groups=self.tool_groups, + response_template=meta_tool_response_template, + ), + ) + + self.builtin_skill_viewer = RegisteredTool( + tool=SkillViewer( + get_skills_method=self._get_available_skills, + ), + ) + + async def get_tool_schemas( + self, + groups: list[str] | None = None, + ) -> list[dict]: + """Get the JSON schemas of the currently available tool functions + based on the given activated tool groups. + + .. note:: The preset keyword arguments is removed from the JSON + schema, and the extended model is applied if it is set. + + Args: + groups (`list[str] | None`, optional): + A list of group names to filter the tool function. The "basic" + group will always be included regardless of the filter. If not + provided, only the "basic" group will be included. + + Example: + .. code-block:: JSON + :caption: Example of tool function JSON schemas + + [ + { + "type": "function", + "function": { + "name": "google_search", + "description": "Search on Google.", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query." + } + }, + "required": ["query"] + } + } + }, + ... + ] + + Returns: + `list[dict]`: + A list of function JSON schemas. + """ + function_schemas = [] + + # Get all available tools + tools_dict = await self._get_available_tools(groups) + for tool in tools_dict.values(): + function_schemas.append(tool.get_tool_schema()) + + return function_schemas + + async def call_tool( + self, + tool_call: ToolCallBlock, + state: AgentState, + ) -> AsyncGenerator[ToolChunk | ToolResponse, None]: + """Call the tool function, return the incremental tool result in + a ToolChunk stream, and finally return the complete tool result in a + ToolResponse object. **Note the accumulation process occurs within this + function, so the tool functions only need to return/yield the + ToolChunk objects in an incremental manner.** + + Args: + tool_call (`ToolCallBlock`): + A tool call block. + state (`AgentState`): + The current agent state, used to state injection. + + Yields: + `ToolChunk | ToolResponse`: + The incremental tool result in a ToolChunk stream, and finally + the complete tool result in a ToolResponse object. + """ + tool_response = ToolResponse(id=tool_call.id) + + # Check + available_tools = await self._get_available_tools( + state.tool_context.activated_groups, + ) + + if tool_call.name not in available_tools: + all_tools = await self._get_available_tools( + groups=[_.name for _ in self.tool_groups], + ) + # Not activate + if tool_call.name in all_tools: + group_name = all_tools[tool_call.name].group + chunk = ToolChunk( + content=[ + TextBlock( + text=( + "ToolGroupInactiveError: The tool " + f"'{tool_call.name}' in group '{group_name}' " + "is currently inactive. You should first " + "activate the group by calling the " + f"'{self.builtin_meta_tool.tool.name}' tool." + ), + ), + ], + state=ToolResultState.ERROR, + ) + yield chunk + yield tool_response.append_chunk(chunk) + return + + # Not exist + chunk = ToolChunk( + content=[ + TextBlock( + text=f"ToolNotFoundError: The tool named " + f"'{tool_call.name}' doesn't exist.", + ), + ], + state=ToolResultState.ERROR, + ) + yield chunk + yield tool_response.append_chunk(chunk) + return + + # Obtain the tool function + tool_func = available_tools[tool_call.name].tool + + # Async function + try: + # Prepare keyword arguments + kwargs = _json_loads_with_repair(tool_call.input) + + # State injection + if ( + tool_func.is_state_injected + and not tool_func.is_mcp + and not tool_func.is_external_tool + ): + kwargs["_agent_state"] = state + + if inspect.iscoroutinefunction(tool_func.__call__): + res = await tool_func(**kwargs) + else: + # When `tool_func.original_func` is Async generator function or + # Sync function + res = tool_func(**kwargs) + + if isinstance(res, ToolChunk): + yield res + tool_response.append_chunk(res) + + # If return an async generator + elif isinstance(res, AsyncGenerator): + async for chunk in res: + yield chunk + tool_response.append_chunk(chunk) + + # If return a sync generator + elif isinstance(res, Generator): + for chunk in res: + yield chunk + tool_response.append_chunk(chunk) + + else: + raise DeveloperOrientedException( + "The tool function must return a ToolChunk object, or an " + "AsyncGenerator/Generator of ToolChunk objects, " + f"but got {type(res)}.", + ) + + except mcp.shared.exceptions.McpError as e: + chunk = ToolChunk( + content=[ + TextBlock( + type="text", + text=f"Error occurred when calling MCP tool: {e}", + ), + ], + state=ToolResultState.ERROR, + ) + yield chunk + tool_response.append_chunk(chunk) + + except Exception as e: + # Raise the developer-oriented exception + if isinstance(e, DeveloperOrientedException): + raise e from None + + # The exceptions should be handled by the agent + chunk = ToolChunk( + content=[ + TextBlock( + type="text", + text=str(e), + ), + ], + state=ToolResultState.ERROR, + ) + yield chunk + tool_response.append_chunk(chunk) + + except asyncio.CancelledError: + chunk = ToolChunk( + content=[ + TextBlock( + type="text", + text="" + "The tool call has been interrupted " + "by the user." + "", + ), + ], + state=ToolResultState.INTERRUPTED, + ) + yield chunk + tool_response.append_chunk(chunk) + + finally: + # Finally, yield the complete tool response + yield tool_response + + async def _get_available_skills( + self, + groups: list[str] | None = None, + ) -> dict[str, Skill]: + """A unified method to collect all skills from the registered skill + loaders. + + Args: + groups (`list[str] | None`, optional): + A list of group names to filter the skill loaders. The "basic" + group will always be included regardless of the filter. If not + provided, only the "basic" group will be included. + + Returns: + `dict[str, Skill]` + A dictionary of skill name and their corresponding Skill + objects. + """ + groups_filter = ["basic"] + (groups or []) + + skills = OrderedDict() + for group in self.tool_groups: + if group.name not in groups_filter: + continue + + for skill in await group.list_skills(): + if skill.name in skills: + logger.warning( + "Duplicate skill name '%s' found in group '%s', " + "overwriting it.", + skill.name, + group.name, + ) + skills[skill.name] = skill + + return skills + + async def get_skill_instructions( + self, + activated_groups: list[str] | None = None, + ) -> str | None: + """Get the prompt for all registered agent skills, which can be + attached to the system prompt for the agent. + + The prompt is consisted of an overall instruction and the detailed + descriptions of each skill, including its name, description, and + directory. + + .. note:: If no skill is registered, None will be returned. + + Args: + activated_groups (`list[str] | None`, optional): + The currently activated tool groups. If omitted, all groups + will be included for backwards compatibility. + + Returns: + `str | None`: + The combined prompt for registered agent skills, or None + if no skill is registered. + """ + if activated_groups is None: + activated_groups = [_.name for _ in self.tool_groups] + + skills = await self._get_available_skills( + activated_groups, + ) + + # If no skills were collected, return None + if len(skills) == 0: + return None + + # Generate the skill instruction prompt with the template + template = Template(self.skill_instruction_template) + + return template.render( + skills=skills.values(), + skill_viewer=self.builtin_skill_viewer.tool.name, + ) + + async def _get_available_tools( + self, + groups: list[str] | None, + ) -> dict[str, RegisteredTool]: + """Return the currently available tools based on the given + activated tool groups. Tools in the ``"basic"`` group are always + included. When at least one tool group is registered, the built-in + meta tool is also included. + + Args: + groups (`list[str]`): + The list of currently activated tool group names. + + Returns: + `dict[str, RegisteredTool]`: + The dictionary of available tool name and their corresponding + RegisteredTool objects. + """ + available_tools = {} + + # Built-in skill viewers + skills = await self._get_available_skills(groups) + if len(skills): + available_tools[ + self.builtin_skill_viewer.tool.name + ] = self.builtin_skill_viewer + + # Builtin meta tool is only included when there is at least one tool + # group + if ( + len(self.tool_groups) == 1 + and self.tool_groups[0].name != "basic" + or len(self.tool_groups) > 1 + ): + available_tools[ + self.builtin_meta_tool.tool.name + ] = self.builtin_meta_tool + + # The tools in the activated groups and the "basic" group are included + groups_filter = ["basic"] + (groups or []) + for group in self.tool_groups: + if group.name not in groups_filter: + continue + + cache_tools = [] + # Python tools + for tool in group.tools: + cache_tools.append(tool) + + # MCP tools + for client in group.mcps: + tools = await client.list_tools() + cache_tools.extend(tools) + + # Append cached tools into the available tools and solve the name + # conflict + for tool in cache_tools: + if tool.name in available_tools: + logger.warning( + "Duplicate tool name '%s' found in group '%s', " + "overwriting it.", + tool.name, + group.name, + ) + available_tools[tool.name] = RegisteredTool( + tool=tool, + group=group.name, + ) + + return available_tools + + async def check_tool_available( + self, + tool_name: str, + activated_groups: list[str], + ) -> ToolBase: + """Check if the tool is available now. If not, raise the + agent-oriented exception. + + Args: + tool_name (`str`): + The name of the tool to be checked. + activated_groups (`list[str]`): + The currently activated tool groups. + + Returns: + `ToolBase`: + If the tool is available, return the corresponding ToolBase + object. Otherwise, raise the agent-oriented exception with the + error message. + """ + tools = await self._get_available_tools(activated_groups) + if tool_name not in tools: + raise ToolNotFoundError( + f"ToolNotFoundError: The tool named '{tool_name}' doesn't " + f"exist.", + ) + + group_name = tools[tool_name].group + if group_name != "basic" and group_name not in activated_groups: + raise ToolGroupInactiveError( + f"ToolGroupInactiveError: The tool '{tool_name}' in group " + f"'{group_name}' is currently inactive. " + f"You should first activate the group by calling the " + f"'{self.builtin_meta_tool.tool.name}' tool.", + ) + + return tools[tool_name].tool + + async def get_tool(self, name: str) -> ToolBase | None: + """Get tool instance by its name. + + Args: + name (`str`): + The name of the tool to be checked. + + Returns: + `ToolBase | None`: + The tool instance, or `None` if no tool is found. + """ + tools = await self._get_available_tools( + [_.name for _ in self.tool_groups], + ) + registered_tool = tools.get(name, None) + if registered_tool is None: + return None + return registered_tool.tool + + def _get_meta_tool_schema(self) -> Type[BaseModel]: + """Get the meta tool schema based on the current tool groups.""" + fields = {} + for group in self.tool_groups: + if group.name == "basic": + continue + fields[group.name] = ( + bool, + Field( + default=False, + description=group.description, + ), + ) + return create_model("_DynamicModel", **fields) + + def clear(self) -> None: + """Clear the registered tools, skills and MCPs.""" + self.tool_groups.clear() diff --git a/src/agentscope/tool/_types.py b/src/agentscope/tool/_types.py new file mode 100644 index 0000000000000000000000000000000000000000..bf2f1514b37f3c8648128b160d351d8348028c4e --- /dev/null +++ b/src/agentscope/tool/_types.py @@ -0,0 +1,203 @@ +# -*- coding: utf-8 -*- +"""The types for the tool module in AgentScope.""" +from copy import deepcopy +from dataclasses import dataclass, field +from typing import ( + Literal, + Type, + Any, + TypeAlias, + Coroutine, + AsyncGenerator, + Generator, + Awaitable, + Callable, +) + +from pydantic import BaseModel + +from ._response import ToolChunk +from ._base import ToolBase +from ._utils import _remove_title_field + + +@dataclass +class RegisteredTool: + """The registered tool function class, used to store the tool function and + its registration information.""" + + tool: ToolBase + """The original tool function.""" + + # Execution related fields + extended_model: Type[BaseModel] | None = field(init=False, default=None) + """The base model used to extend the JSON schema of the original tool + function, so that we can dynamically adjust the tool function.""" + + # Tools management fields + group: str | Literal["basic"] = "basic" + """The belonging group of the tool function""" + original_name: str | None = field(default=None) + """The original name of the tool function when it has been renamed.""" + + def __post_init__(self) -> None: + """Validate the registered tool function after initialization.""" + # validate schema + if self.tool.input_schema is not None: + if not ( + isinstance(self.tool.input_schema, dict) + and self.tool.input_schema.get("type") == "object" + and isinstance(self.tool.input_schema.get("properties"), dict) + ): + raise ValueError( + f"Invalid input_schema: {self.tool.input_schema}. ", + ) + + def get_tool_schema( + self, + extended_model: Type[BaseModel] | None = None, + ) -> dict: + """Get the JSON schema of the tool function via the following steps: + + 1. Remove preset_kwargs from the JSON schema, since they are not + exposed to the agent. + 2. If extended_model is provided, merge its schema with the + current function schema. + + Args: + extended_model (`Type[BaseModel] | None`, optional): + The dynamic BaseModel used to extend the original function. If + provided, the given BaseModel will be merged into the original + function schema instead of the extended_model field. + + Returns: + `dict`: The JSON schema of the tool function. + """ + input_schema = deepcopy(self.tool.input_schema) + _remove_title_field(input_schema) + function_schema: dict = { + "type": "function", + "function": { + "name": self.tool.name, + "description": self.tool.description, + "parameters": input_schema, + }, + } + + extended_model = extended_model or self.extended_model + + if extended_model is None: + return function_schema + + # Merge the extended model with the original JSON schema + extended_schema = extended_model.model_json_schema() + + _remove_title_field(extended_schema) + + # Merge properties from extended schema + for key, value in extended_schema["properties"].items(): + if key in function_schema["function"]["parameters"]["properties"]: + raise ValueError( + f"The field `{key}` already exists in the original " + f"function schema of `{self.tool.name}`. Try to use a " + "different name.", + ) + + function_schema["function"]["parameters"]["properties"][ + key + ] = value + + if key in extended_schema.get("required", []): + if "required" not in function_schema["function"]["parameters"]: + function_schema["function"]["parameters"]["required"] = [] + function_schema["function"]["parameters"]["required"].append( + key, + ) + + # Merge $defs from extended schema to support nested models + if "$defs" in extended_schema: + merged_params = function_schema["function"]["parameters"] + if "$defs" not in merged_params: + merged_params["$defs"] = {} + + # Check for conflicts and merge $defs + for def_key, def_value in extended_schema["$defs"].items(): + def_value_copy = deepcopy(def_value) + _remove_title_field( + def_value_copy, + ) # pylint: disable=protected-access + + if def_key in merged_params["$defs"]: + # Check if the two definitions are from the same BaseModel + # by comparing their content + # Create copies and remove title fields for comparison + + existing_def_copy = deepcopy( + merged_params["$defs"][def_key], + ) + _remove_title_field(existing_def_copy) + + if existing_def_copy != def_value_copy: + # The definitions are different, raise an error + raise ValueError( + f"The $defs key `{def_key}` conflicts with " + f"existing definition in function schema of " + f"`{self.tool.name}`.", + ) + # The definitions are the same (from the same BaseModel), + # skip merging this key + continue + + merged_params["$defs"][def_key] = def_value_copy + + return function_schema + + +# The function types that can be registered as tools in AgentScope. +Function: TypeAlias = ( + # Sync function + Callable[..., ToolChunk] + | + # Async function + Callable[..., Awaitable[ToolChunk]] + | + # Sync generator function + Callable[..., Generator[ToolChunk, None, None]] + | + # Async generator function + Callable[..., AsyncGenerator[ToolChunk, None]] + | + # Async function that returns async generator + Callable[..., Coroutine[Any, Any, AsyncGenerator[ToolChunk, None]]] + | + # Async function that returns sync generator + Callable[..., Coroutine[Any, Any, Generator[ToolChunk, None, None]]] +) + + +class ToolChoice(BaseModel): + """The tool choice configuration. + + Attributes: + mode: The tool choice mode. Supports: + + * ``"auto"`` – the model decides whether to call a tool. + * ``"none"`` – the model must not call any tool. + * ``"required"`` – the model must call at least one tool. + * ``str`` (a tool name) – the model **must** call exactly that + tool (forced single-tool call). The name is validated against + ``tools`` (if provided) or against the full tools list passed to + the model. + + tools: An optional list of tool names. When specified, the tool + schemas forwarded to the model are filtered to only those tools. + This also acts as a validation whitelist for ``mode`` when it + is a specific tool name (str): the name must appear in this + list. Prefer using ``mode=`` (str) over + ``tools=[""]`` when the goal is a forced single-tool + call without changing the available tool set, as the former + avoids schema-list changes that would invalidate prompt caches. + """ + + mode: Literal["auto", "none", "required"] | str + tools: list[str] | None = None diff --git a/src/agentscope/tool/_utils.py b/src/agentscope/tool/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5e4046c3a73200446c70a15ac971e0c17148cde4 --- /dev/null +++ b/src/agentscope/tool/_utils.py @@ -0,0 +1,160 @@ +# -*- coding: utf-8 -*- +"""The tool module utils.""" +import inspect +from typing import Any, Dict, Callable + +from docstring_parser import parse +from pydantic import Field, create_model, ConfigDict + + +def _remove_title_field(schema: dict) -> dict: + """Remove the title field from the JSON schema to avoid + misleading the LLM.""" + # The top level title field + if "title" in schema: + schema.pop("title") + + # properties + if "properties" in schema: + for prop in schema["properties"].values(): + if isinstance(prop, dict): + _remove_title_field(prop) + + # items + if "items" in schema and isinstance(schema["items"], dict): + _remove_title_field(schema["items"]) + + # additionalProperties + if "additionalProperties" in schema and isinstance( + schema["additionalProperties"], + dict, + ): + _remove_title_field(schema["additionalProperties"]) + + # $defs — referenced sub-schemas, e.g. Pydantic models used as parameter + # types generate "$defs": {"SubModel": {"title": "SubModel", ...}}. + # These titles are auto-generated noise just like property titles, and + # should be removed for the same reason. + if "$defs" in schema and isinstance(schema["$defs"], dict): + for def_schema in schema["$defs"].values(): + if isinstance(def_schema, dict): + _remove_title_field(def_schema) + + return schema + + +def _extract_func_description(docstring: str) -> str: + """Extract the function description from the docstring. + + Args: + docstring (`str`): + The docstring to extract the function description from. + + Returns: + `str`: + The extracted function description. + """ + parsed_docstring = parse(docstring or "") + descriptions = [] + if parsed_docstring.short_description is not None: + descriptions.append(parsed_docstring.short_description) + + if parsed_docstring.long_description is not None: + descriptions.append(parsed_docstring.long_description) + + return "\n".join(descriptions) + + +def _extract_input_schema( + tool_func: Callable, + include_var_positional: bool = False, + include_var_keyword: bool = False, +) -> dict: + """Extract input schema from the tool function's docstring + + Args: + tool_func (`ToolFunction`): + The tool function to extract the JSON schema from. + include_var_positional (`bool`): + Whether to include variable positional arguments in the JSON + schema. + include_var_keyword (`bool`): + Whether to include variable keyword arguments in the JSON schema. + + Returns: + `dict`: + The extracted input JSON schema. + """ + docstring = parse(tool_func.__doc__ or "") + params_docstring = {_.arg_name: _.description for _ in docstring.params} + + # Create a dynamic model with the function signature + fields = {} + for name, param in inspect.signature(tool_func).parameters.items(): + # Skip the `self` and `cls` parameters + if name in ["self", "cls"]: + continue + + # Handle `**kwargs` + if param.kind == inspect.Parameter.VAR_KEYWORD: + if not include_var_keyword: + continue + + fields[name] = ( + Dict[str, Any] + if param.annotation == inspect.Parameter.empty + else Dict[str, param.annotation], # type: ignore + Field( + description=params_docstring.get( + f"**{name}", + params_docstring.get(name, None), + ), + default={} + if param.default is param.empty + else param.default, + ), + ) + + elif param.kind == inspect.Parameter.VAR_POSITIONAL: + if not include_var_positional: + continue + + fields[name] = ( + list[Any] + if param.annotation == inspect.Parameter.empty + else list[param.annotation], # type: ignore + Field( + description=params_docstring.get( + f"*{name}", + params_docstring.get(name, None), + ), + default=[] + if param.default is param.empty + else param.default, + ), + ) + + else: + fields[name] = ( + Any + if param.annotation == inspect.Parameter.empty + else param.annotation, + Field( + description=params_docstring.get(name, None), + default=... + if param.default is param.empty + else param.default, + ), + ) + + base_model = create_model( + "_StructuredOutputDynamicClass", + __config__=ConfigDict(arbitrary_types_allowed=True), + **fields, + ) + params_json_schema = base_model.model_json_schema() + + # Remove the title from the json schema + _remove_title_field(params_json_schema) + + return params_json_schema diff --git a/src/agentscope/tts/__init__.py b/src/agentscope/tts/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4e6e4d82315807d775746d7178197d0ae2dcf757 --- /dev/null +++ b/src/agentscope/tts/__init__.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- +"""The TTS (Text-to-Speech) module in AgentScope.""" + +from ._tts_base import TTSModelBase +from ._tts_model_card import TTSModelCard +from ._tts_response import TTSResponse, TTSUsage +from ._dashscope import ( + DashScopeTTSModel, + DashScopeRealtimeTTSModel, + DashScopeCosyVoiceRealtimeTTSModel, +) + +__all__ = [ + "TTSModelBase", + "TTSModelCard", + "TTSResponse", + "TTSUsage", + "DashScopeTTSModel", + "DashScopeRealtimeTTSModel", + "DashScopeCosyVoiceRealtimeTTSModel", +] diff --git a/src/agentscope/tts/_dashscope/__init__.py b/src/agentscope/tts/_dashscope/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..93215d9756d94503858d963ac31870594bbb3e65 --- /dev/null +++ b/src/agentscope/tts/_dashscope/__init__.py @@ -0,0 +1,12 @@ +# -*- coding: utf-8 -*- +"""The DashScope TTS module.""" + +from ._model import DashScopeTTSModel +from ._realtime_model import DashScopeRealtimeTTSModel +from ._cosyvoice_realtime_model import DashScopeCosyVoiceRealtimeTTSModel + +__all__ = [ + "DashScopeTTSModel", + "DashScopeRealtimeTTSModel", + "DashScopeCosyVoiceRealtimeTTSModel", +] diff --git a/src/agentscope/tts/_dashscope/_cosyvoice_models/cosyvoice-v3-flash.yaml b/src/agentscope/tts/_dashscope/_cosyvoice_models/cosyvoice-v3-flash.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d258a84541d6cbf02bba0b0aeaf6e9a9457ba1ac --- /dev/null +++ b/src/agentscope/tts/_dashscope/_cosyvoice_models/cosyvoice-v3-flash.yaml @@ -0,0 +1,25 @@ +name: cosyvoice-v3-flash +label: CosyVoice-v3-Flash +status: active +realtime: true + +input_types: + - text/plain + +output_types: + - audio/wav + +voices: + # Source: https://help.aliyun.com/en/model-studio/cosyvoice-voice-list + - longanyang + - longanhuan + - longhuhu_v3 + - longyingmu_v3 + - longxiaochun + - longxiaoxia + - longlaotie + - longshuo + - longjing + - longshu + +parameter_overrides: {} diff --git a/src/agentscope/tts/_dashscope/_cosyvoice_models/cosyvoice-v3-plus.yaml b/src/agentscope/tts/_dashscope/_cosyvoice_models/cosyvoice-v3-plus.yaml new file mode 100644 index 0000000000000000000000000000000000000000..977f27e5a26190691242147e143d572e2dd6dbb0 --- /dev/null +++ b/src/agentscope/tts/_dashscope/_cosyvoice_models/cosyvoice-v3-plus.yaml @@ -0,0 +1,25 @@ +name: cosyvoice-v3-plus +label: CosyVoice-v3-Plus +status: active +realtime: true + +input_types: + - text/plain + +output_types: + - audio/wav + +voices: + # Source: https://help.aliyun.com/en/model-studio/cosyvoice-voice-list + - longanyang + - longanhuan + - longhuhu_v3 + - longyingmu_v3 + - longxiaochun + - longxiaoxia + - longlaotie + - longshuo + - longjing + - longshu + +parameter_overrides: {} diff --git a/src/agentscope/tts/_dashscope/_cosyvoice_realtime_model.py b/src/agentscope/tts/_dashscope/_cosyvoice_realtime_model.py new file mode 100644 index 0000000000000000000000000000000000000000..d1884104cd4371ee83181a8d595d1f3068c989fb --- /dev/null +++ b/src/agentscope/tts/_dashscope/_cosyvoice_realtime_model.py @@ -0,0 +1,507 @@ +# -*- coding: utf-8 -*- +"""DashScope CosyVoice Realtime TTS model implementation. + +Uses the older ``dashscope.audio.tts_v2.SpeechSynthesizer`` SDK for models +such as ``cosyvoice-v3-plus``, ``cosyvoice-v3-flash``, ``sambert``, etc. +""" +import asyncio +import base64 +import os +import threading +from typing import Any, AsyncGenerator, Literal, TYPE_CHECKING + +from pydantic import BaseModel, Field + +from .._tts_base import TTSModelBase +from .._tts_response import TTSResponse +from ..._logging import logger +from ..._utils._audio import _build_streaming_wav_header +from ...credential import DashScopeCredential +from ...message import DataBlock, Base64Source + +if TYPE_CHECKING: + from dashscope.audio.tts_v2 import ResultCallback + + +_MEDIA_TYPE = "audio/wav" +_SAMPLE_RATE = 24000 +_CHANNELS = 1 +_BITS_PER_SAMPLE = 16 + + +def _make_cosyvoice_callback_class() -> type["ResultCallback"]: + """Create the DashScope CosyVoice TTS callback class lazily.""" + from dashscope.audio.tts_v2 import ResultCallback + + class _CosyVoiceCallback(ResultCallback): + """Internal callback that accumulates PCM audio from the WebSocket + and exposes incremental deltas.""" + + def __init__(self) -> None: + """Initialize callback with audio buffer and synchronization + events.""" + super().__init__() + self.chunk_event = threading.Event() + self.finish_event = threading.Event() + self._pcm_bytes: bytearray = bytearray() + self._consumed: int = 0 + + def on_open(self) -> None: + """Handle WebSocket open — reset audio state.""" + self._pcm_bytes = bytearray() + self._consumed = 0 + self.finish_event.clear() + self.chunk_event.clear() + + def on_data(self, data: bytes) -> None: + """Handle incoming PCM audio data.""" + if data: + self._pcm_bytes += data + if not self.chunk_event.is_set(): + self.chunk_event.set() + + def on_complete(self) -> None: + """Handle synthesis completion.""" + self.finish_event.set() + self.chunk_event.set() + + def on_close(self) -> None: + """Handle WebSocket close.""" + self.finish_event.set() + self.chunk_event.set() + + def on_error(self, message: Any) -> None: + """Handle synthesis error.""" + logger.error("CosyVoice TTS error: %s", message) + self.finish_event.set() + self.chunk_event.set() + + def _take_delta(self, header: bool = False) -> bytes | None: + """Return new PCM bytes since last call, or None if empty.""" + new_data = self._pcm_bytes[self._consumed :] + if not new_data: + return None + self._consumed = len(self._pcm_bytes) + if header: + return _build_streaming_wav_header( + sample_rate=_SAMPLE_RATE, + channels=_CHANNELS, + bits_per_sample=_BITS_PER_SAMPLE, + ) + bytes(new_data) + return bytes(new_data) + + def get_audio_response(self, block: bool) -> TTSResponse: + """Return incremental audio delta.""" + if block: + self.finish_event.wait() + delta = self._take_delta(header=self._consumed == 0) + if delta: + return TTSResponse( + content=DataBlock( + source=Base64Source( + data=base64.b64encode(delta).decode("ascii"), + media_type=_MEDIA_TYPE, + ), + ), + ) + return TTSResponse(content=None) + + async def get_audio_chunks(self) -> AsyncGenerator[TTSResponse, None]: + """Yield incremental WAV audio chunks as they arrive.""" + header_sent = self._consumed > 0 + while True: + if self.finish_event.is_set(): + delta = self._take_delta(header=not header_sent) + if delta: + yield TTSResponse( + content=DataBlock( + source=Base64Source( + data=base64.b64encode(delta).decode( + "ascii", + ), + media_type=_MEDIA_TYPE, + ), + ), + is_last=True, + ) + else: + yield TTSResponse(content=None, is_last=True) + self.reset() + break + + if self.chunk_event.is_set(): + self.chunk_event.clear() + else: + await asyncio.to_thread(self.chunk_event.wait, 30) + + if self.finish_event.is_set(): + continue + + delta = self._take_delta(header=not header_sent) + if delta: + header_sent = True + yield TTSResponse( + content=DataBlock( + source=Base64Source( + data=base64.b64encode(delta).decode("ascii"), + media_type=_MEDIA_TYPE, + ), + ), + is_last=False, + ) + + def reset(self) -> None: + """Reset internal state for the next utterance.""" + self.finish_event.clear() + self.chunk_event.clear() + self._pcm_bytes = bytearray() + self._consumed = 0 + + def has_audio_data(self) -> bool: + """Return whether any audio data has been received.""" + return bool(self._pcm_bytes) + + return _CosyVoiceCallback + + +class DashScopeCosyVoiceRealtimeTTSModel(TTSModelBase): + """DashScope CosyVoice Realtime TTS model using the + ``SpeechSynthesizer`` streaming API. + + Supports streaming input: text can be pushed incrementally via + :meth:`push`, and :meth:`synthesize` finalizes the current utterance. + + Supported models include ``cosyvoice-v3-plus``, ``cosyvoice-v3-flash``, + ``sambert``, etc. For more details see the `official document + `_. + + .. note:: Only one streaming input request can be active at a time. + + .. note:: Unlike ``DashScopeRealtimeTTSModel`` (Qwen3) which produces + audio at token-level granularity, CosyVoice server automatically + segments incoming text into sentences and synthesizes per-sentence. + Audio is returned via callback only after a complete sentence boundary + is detected. This means :meth:`push` may return empty responses until + enough text accumulates to form a sentence. Calling + :meth:`synthesize` forces synthesis of all remaining text (including + incomplete sentences). See `official docs + `_. + """ + + class Parameters(BaseModel): + """Frontend-exposed parameters for CosyVoice Realtime TTS models.""" + + voice: str = Field( + default="longanyang", + title="Voice", + description="The voice to use for synthesis.", + ) + + type: Literal[ + "dashscope_cosyvoice_realtime_tts" + ] = "dashscope_cosyvoice_realtime_tts" + + realtime: bool = True + + _MODELS_DIR = os.path.join(os.path.dirname(__file__), "_cosyvoice_models") + + @classmethod + def list_models( + cls, + custom_yaml_dir: str | None = None, + ) -> list: + """List CosyVoice model cards from the dedicated YAML directory.""" + return super().list_models( + custom_yaml_dir=custom_yaml_dir or cls._MODELS_DIR, + ) + + def __init__( + self, + credential: DashScopeCredential, + model: str = "cosyvoice-v3-plus", + parameters: "DashScopeCosyVoiceRealtimeTTSModel.Parameters | None" = ( + None + ), + stream: bool = True, + cold_start_length: int | None = None, + cold_start_words: int | None = None, + max_retries: int = 3, + retry_delay: float = 5.0, + ) -> None: + """Initialize the DashScope CosyVoice Realtime TTS model. + + Args: + credential (`DashScopeCredential`): + The DashScope credential. + model (`str`, defaults to ``"cosyvoice-v3-plus"``): + The CosyVoice model name, e.g. ``"cosyvoice-v3-plus"``, + ``"cosyvoice-v3-flash"``, ``"sambert"``. + parameters (`Parameters | None`, defaults to `None`): + The TTS parameters (voice, etc.). + stream (`bool`, defaults to `True`): + Whether :meth:`synthesize` returns a streaming async generator. + cold_start_length (`int | None`, defaults to `None`): + Minimum character count before the first text chunk is sent + to the synthesizer. + cold_start_words (`int | None`, defaults to `None`): + Minimum word count (split by spaces) before the first text + chunk is sent. + max_retries (`int`, defaults to `3`): + Max retry attempts on synthesis failure. + retry_delay (`float`, defaults to `5.0`): + Initial retry delay in seconds (exponential backoff). + """ + super().__init__( + credential=credential, + model=model, + parameters=parameters, + stream=stream, + ) + self.cold_start_length = cold_start_length + self.cold_start_words = cold_start_words + self.max_retries = max_retries + self.retry_delay = retry_delay + + self._synthesizer: Any = None + self._callback: Any = None + self._connected = False + self._cold_start_buffer: str = "" + self._cold_start_done: bool = False + self._accumulated_text: str = "" + + def _create_synthesizer(self) -> None: + """Create a fresh SpeechSynthesizer and callback.""" + import dashscope + from dashscope.audio.tts_v2 import SpeechSynthesizer, AudioFormat + + dashscope.api_key = self.credential.api_key.get_secret_value() + + callback_cls = _make_cosyvoice_callback_class() + self._callback = callback_cls() + self._synthesizer = SpeechSynthesizer( + model=self.model, + voice=self.parameters.voice, + format=AudioFormat.PCM_24000HZ_MONO_16BIT, + callback=self._callback, + ) + + async def connect(self) -> None: + """Initialize the SpeechSynthesizer.""" + if self._connected: + return + self._create_synthesizer() + self._connected = True + + async def close(self) -> None: + """Close the SpeechSynthesizer.""" + if not self._connected: + return + self._connected = False + try: + if self._synthesizer is not None: + self._synthesizer.close() + except Exception: + pass + + async def _reconnect(self) -> None: + """Reconnect by recreating the synthesizer.""" + try: + if self._synthesizer is not None: + self._synthesizer.close() + except Exception: + pass + self._connected = False + self._cold_start_buffer = "" + self._cold_start_done = False + await self.connect() + + async def push( + self, + text: str, + **kwargs: Any, + ) -> TTSResponse: + """Push an incremental text delta for realtime synthesis. + + .. note:: The CosyVoice server automatically segments text into + sentences before synthesizing. Audio is only produced after a + complete sentence is detected, so this method often returns an + empty response (``content=None``) for partial sentences. Remaining + audio is force-synthesized when :meth:`synthesize` is called. + See `CosyVoice Python SDK docs + `_. + + Args: + text (`str`): + An incremental text chunk (delta) to append. + **kwargs (`Any`): + Additional keyword arguments (unused). + + Returns: + `TTSResponse`: + Audio accumulated so far, or empty if not yet available. + """ + if not self._connected: + raise RuntimeError( + "TTS model is not connected. Call `connect()` first.", + ) + + if not text: + return TTSResponse(content=None) + + self._accumulated_text += text + + if self._cold_start_done: + text_to_send = text + else: + self._cold_start_buffer += text + if ( + self.cold_start_length + and len(self._cold_start_buffer) < self.cold_start_length + ) or ( + self.cold_start_words + and len(self._cold_start_buffer.split()) + < self.cold_start_words + ): + return self._callback.get_audio_response(block=False) + text_to_send = self._cold_start_buffer + self._cold_start_buffer = "" + self._cold_start_done = True + + try: + self._synthesizer.streaming_call(text_to_send) + except Exception: + return TTSResponse(content=None) + + return self._callback.get_audio_response(block=False) + + async def synthesize( + self, + text: str | None = None, + **kwargs: Any, + ) -> TTSResponse | AsyncGenerator[TTSResponse, None]: + """Finalize synthesis for the current utterance. + + If text was previously pushed via :meth:`push`, this flushes any + remaining buffered text, calls ``streaming_complete()``, and waits + for audio. If ``text`` is provided, it is appended before finalizing. + + Args: + text (`str | None`, defaults to `None`): + Optional additional text to append before finalizing. + **kwargs (`Any`): + Additional keyword arguments (unused). + + Returns: + `TTSResponse | AsyncGenerator[TTSResponse, None]`: + A single response when ``stream=False``, or an async generator + of incremental chunks when ``stream=True``. + """ + if not self._connected: + raise RuntimeError( + "TTS model is not connected. Call `connect()` first.", + ) + + if text is not None: + self._accumulated_text += text + + unsent = self._cold_start_buffer + if text is not None: + unsent += text + self._cold_start_buffer = "" + + full_text = self._accumulated_text + delay = self.retry_delay + + try: + if not full_text and not unsent: + if self.stream: + + async def _empty_gen() -> AsyncGenerator[ + TTSResponse, + None, + ]: + yield TTSResponse(content=None) + + return _empty_gen() + return TTSResponse(content=None) + + for attempt in range(self.max_retries): + try: + if unsent: + self._synthesizer.streaming_call(unsent) + + self._synthesizer.streaming_complete() + + finished = await asyncio.to_thread( + self._callback.finish_event.wait, + 30, + ) + + if not finished: + logger.warning( + "CosyVoice TTS: timed out waiting for synthesis " + "completion (30s)", + ) + if attempt < self.max_retries - 1: + await asyncio.sleep(delay) + await self._reconnect() + unsent = full_text + delay *= 2 + continue + raise RuntimeError( + "CosyVoice TTS synthesis timed out after 30s", + ) + + if full_text and not self._callback.has_audio_data(): + if attempt < self.max_retries - 1: + logger.warning( + "CosyVoice TTS: no audio received, retrying " + "(%d/%d) in %.1fs...", + attempt + 1, + self.max_retries, + delay, + ) + await asyncio.sleep(delay) + await self._reconnect() + unsent = full_text + delay *= 2 + continue + raise RuntimeError( + f"CosyVoice TTS synthesis failed: no audio after " + f"{self.max_retries} attempts", + ) + break + + except RuntimeError: + raise + except Exception as e: + if attempt < self.max_retries - 1: + logger.warning( + "CosyVoice TTS error, retrying (%d/%d) in " + "%.1fs: %s", + attempt + 1, + self.max_retries, + delay, + e, + ) + await asyncio.sleep(delay) + await self._reconnect() + unsent = full_text + delay *= 2 + else: + raise + + if self.stream: + return self._callback.get_audio_chunks() + + response = self._callback.get_audio_response(block=True) + self._callback.reset() + return response + finally: + self._reset_state() + + def _reset_state(self) -> None: + """Reset per-utterance tracking state.""" + self._cold_start_buffer = "" + self._cold_start_done = False + self._accumulated_text = "" diff --git a/src/agentscope/tts/_dashscope/_model.py b/src/agentscope/tts/_dashscope/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..62887e4d6fb3769887c0f57bb60d135fcfc6441a --- /dev/null +++ b/src/agentscope/tts/_dashscope/_model.py @@ -0,0 +1,252 @@ +# -*- coding: utf-8 -*- +"""DashScope TTS model implementation using MultiModalConversation API.""" +import base64 +import io +import wave +from datetime import datetime +from typing import ( + Any, + AsyncGenerator, + Generator, + Literal, + TYPE_CHECKING, +) + +from pydantic import BaseModel, Field + +from .._tts_base import TTSModelBase +from .._tts_response import TTSResponse, TTSUsage +from ..._utils._audio import _build_streaming_wav_header +from ...credential import DashScopeCredential +from ...message import DataBlock, Base64Source + +if TYPE_CHECKING: + from dashscope.api_entities.dashscope_response import ( + MultiModalConversationResponse, + ) + + +# DashScope TTS returns raw PCM (24kHz, mono, 16-bit). We wrap it as WAV +# on the way out so the frontend can play it: streaming deltas get a +# streaming WAV header on the first chunk; non-streaming returns a +# self-contained fixed-size WAV. +_TTS_SAMPLE_RATE = 24000 +_TTS_CHANNELS = 1 +_TTS_BITS_PER_SAMPLE = 16 +_DEFAULT_MEDIA_TYPE = "audio/wav" +_SENTINEL = object() + + +def _parse_usage(usage: Any, elapsed: float) -> TTSUsage | None: + """Extract a TTSUsage from the DashScope usage object, or None.""" + if usage is None: + return None + return TTSUsage( + input_tokens=getattr(usage, "input_tokens", 0) or 0, + output_tokens=getattr(usage, "output_tokens", 0) or 0, + time=elapsed, + ) + + +class DashScopeTTSModel(TTSModelBase): + """DashScope TTS model implementation using the MultiModalConversation + API. For more details please see the `official document + `_. + """ + + class Parameters(BaseModel): + """Frontend-exposed parameters for DashScope TTS models.""" + + voice: str = Field( + default="Cherry", + title="Voice", + description="The voice to use for synthesis.", + ) + + type: Literal["dashscope_tts"] = "dashscope_tts" + """The type of the TTS model.""" + + realtime: bool = False + + def __init__( + self, + credential: DashScopeCredential, + model: str = "qwen3-tts-flash", + parameters: "DashScopeTTSModel.Parameters | None" = None, + stream: bool = True, + ) -> None: + """Initialize the DashScope TTS model. + + .. note:: More details about the parameters, such as ``model`` + and ``voice``, can be found in the `official document + `_. + + Args: + credential (`DashScopeCredential`): + The DashScope credential used to authenticate the API call. + model (`str`, defaults to ``"qwen3-tts-flash"``): + The TTS model name. Supported models include + ``qwen3-tts-flash``, ``qwen-tts``, etc. + parameters (`DashScopeTTSModel.Parameters | None`, defaults to \ + `None`): + The TTS parameters (voice, language, etc.). When ``None``, + the default parameters will be used. + stream (`bool`, defaults to `True`): + Whether to use streaming output. When `True`, + :meth:`synthesize` returns an async generator yielding + ``TTSResponse`` chunks; when `False`, it returns a single + aggregated ``TTSResponse``. + """ + super().__init__( + credential=credential, + model=model, + parameters=parameters, + stream=stream, + ) + + async def synthesize( + self, + text: str | None = None, + **kwargs: Any, + ) -> TTSResponse | AsyncGenerator[TTSResponse, None]: + """Call the DashScope TTS API to synthesize speech from text. + + Args: + text (`str | None`, optional): + The text to be synthesized. + **kwargs (`Any`): + Additional keyword arguments to pass to the TTS API call. + + Returns: + `TTSResponse | AsyncGenerator[TTSResponse, None]`: + A single ``TTSResponse`` when ``stream=False``, or an async + generator yielding ``TTSResponse`` chunks when ``stream=True``. + """ + if not text: + return TTSResponse(content=None) + + import dashscope + + response = dashscope.MultiModalConversation.call( + model=self.model, + api_key=self.credential.api_key.get_secret_value(), + text=text, + voice=self.parameters.voice, + stream=True, + **kwargs, + ) + + if self.stream: + return self._parse_into_async_generator(response) + + return self._aggregate_sync(response) + + @staticmethod + def _aggregate_sync( + response: Generator["MultiModalConversationResponse", None, None], + ) -> TTSResponse: + """Aggregate all streaming chunks into a single self-contained WAV.""" + start_datetime = datetime.now() + audio_bytes = bytearray() + usage = None + for chunk in response: + if chunk.usage is not None: + usage = chunk.usage + if chunk.output is not None: + audio = chunk.output.audio + if audio and audio.data: + audio_bytes += base64.b64decode(audio.data) + elapsed = (datetime.now() - start_datetime).total_seconds() + + buf = io.BytesIO() + with wave.open(buf, "wb") as wav: + wav.setnchannels(_TTS_CHANNELS) + wav.setsampwidth(_TTS_BITS_PER_SAMPLE // 8) + wav.setframerate(_TTS_SAMPLE_RATE) + wav.writeframes(bytes(audio_bytes)) + + return TTSResponse( + content=DataBlock( + source=Base64Source( + data=base64.b64encode(buf.getvalue()).decode("ascii"), + media_type=_DEFAULT_MEDIA_TYPE, + ), + ), + usage=_parse_usage(usage, elapsed), + ) + + @staticmethod + async def _parse_into_async_generator( + response: Generator["MultiModalConversationResponse", None, None], + ) -> AsyncGenerator[TTSResponse, None]: + """Parse the streaming TTS response into an async generator. + + Each yielded ``TTSResponse`` carries an **incremental** WAV chunk: + the first chunk is prefixed with a streaming WAV/RIFF header so the + frontend can start playback immediately (without waiting for + end-of-stream); subsequent chunks are raw PCM bytes appended to that + open stream. The final response has ``is_last=True``. + + Args: + response (`Generator[MultiModalConversationResponse, None, None]`): + The streaming response from the DashScope TTS API. + + Yields: + `TTSResponse`: + A ``TTSResponse`` for each incremental audio chunk; the final + response has ``is_last=True``. + """ + pending: TTSResponse | None = None + header_sent = False + usage = None + start_datetime = datetime.now() + it = iter(response) + while True: + chunk = next(it, _SENTINEL) + if chunk is _SENTINEL: + break + if chunk.usage is not None: + usage = chunk.usage + if chunk.output is None: + continue + audio = chunk.output.audio + if not audio or not audio.data: + continue + delta_bytes = base64.b64decode(audio.data) + if not delta_bytes: + continue + if not header_sent: + payload = ( + _build_streaming_wav_header( + sample_rate=_TTS_SAMPLE_RATE, + channels=_TTS_CHANNELS, + bits_per_sample=_TTS_BITS_PER_SAMPLE, + ) + + delta_bytes + ) + header_sent = True + else: + payload = delta_bytes + if pending is not None: + yield pending + pending = TTSResponse( + content=DataBlock( + source=Base64Source( + data=base64.b64encode(payload).decode("ascii"), + media_type=_DEFAULT_MEDIA_TYPE, + ), + ), + is_last=False, + ) + elapsed = (datetime.now() - start_datetime).total_seconds() + + if pending is not None: + pending.is_last = True + pending.usage = _parse_usage(usage, elapsed) + yield pending + else: + yield TTSResponse( + content=None, + is_last=True, + usage=_parse_usage(usage, elapsed), + ) diff --git a/src/agentscope/tts/_dashscope/_models/qwen3-tts-flash-realtime.yaml b/src/agentscope/tts/_dashscope/_models/qwen3-tts-flash-realtime.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d8cc9c75db07afed7a2988ba5bed760f411b502a --- /dev/null +++ b/src/agentscope/tts/_dashscope/_models/qwen3-tts-flash-realtime.yaml @@ -0,0 +1,19 @@ +name: qwen3-tts-flash-realtime +label: Qwen3-TTS-Flash-Realtime +status: active +realtime: true + +input_types: + - text/plain + +output_types: + - audio/wav + +voices: + # Source: https://help.aliyun.com/zh/model-studio/qwen-tts-voice-list + - Cherry + - Serena + - Ethan + - Chelsie + +parameter_overrides: {} diff --git a/src/agentscope/tts/_dashscope/_models/qwen3-tts-flash.yaml b/src/agentscope/tts/_dashscope/_models/qwen3-tts-flash.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4b26ab3af64030e07b779a7f32711dacf74402cd --- /dev/null +++ b/src/agentscope/tts/_dashscope/_models/qwen3-tts-flash.yaml @@ -0,0 +1,18 @@ +name: qwen3-tts-flash +label: Qwen3-TTS-Flash +status: active + +input_types: + - text/plain + +output_types: + - audio/wav + +voices: + # Source: https://help.aliyun.com/zh/model-studio/qwen-tts-voice-list + - Cherry + - Serena + - Ethan + - Chelsie + +parameter_overrides: {} diff --git a/src/agentscope/tts/_dashscope/_realtime_model.py b/src/agentscope/tts/_dashscope/_realtime_model.py new file mode 100644 index 0000000000000000000000000000000000000000..f763af5064a852511ecf95a9aa77f0bda2df1de3 --- /dev/null +++ b/src/agentscope/tts/_dashscope/_realtime_model.py @@ -0,0 +1,468 @@ +# -*- coding: utf-8 -*- +"""DashScope Realtime TTS model implementation.""" +import asyncio +import base64 +import threading +from typing import Any, AsyncGenerator, Literal, TYPE_CHECKING + +from pydantic import BaseModel, Field + +from .._tts_base import TTSModelBase +from .._tts_response import TTSResponse +from ..._logging import logger +from ..._utils._audio import _build_streaming_wav_header +from ...credential import DashScopeCredential +from ...message import DataBlock, Base64Source + +if TYPE_CHECKING: + from dashscope.audio.qwen_tts_realtime import ( + QwenTtsRealtime, + QwenTtsRealtimeCallback, + ) + + +_MEDIA_TYPE = "audio/wav" +_SAMPLE_RATE = 24000 +_CHANNELS = 1 +_BITS_PER_SAMPLE = 16 + + +def _make_callback_class() -> type["QwenTtsRealtimeCallback"]: + """Create the DashScope realtime TTS callback class lazily to avoid + importing dashscope at module level.""" + from dashscope.audio.qwen_tts_realtime import QwenTtsRealtimeCallback + + class _Callback(QwenTtsRealtimeCallback): + """Internal callback that accumulates PCM audio from the WebSocket. + + Audio data is stored as raw bytes (decoded from base64) to allow + incremental delta extraction without base64 boundary issues. + """ + + def __init__(self) -> None: + """Initialize callback with audio buffer and synchronization + events.""" + super().__init__() + self.chunk_event = threading.Event() + self.finish_event = threading.Event() + self._pcm_bytes: bytearray = bytearray() + self._consumed: int = 0 + + def on_event(self, response: dict[str, Any]) -> None: + """Handle incoming WebSocket events from the TTS service.""" + try: + event_type = response.get("type") + + if event_type == "session.created": + self._pcm_bytes = bytearray() + self._consumed = 0 + self.finish_event.clear() + self.chunk_event.clear() + + elif event_type == "response.audio.delta": + audio_data = response.get("delta") + if audio_data: + if isinstance(audio_data, bytes): + self._pcm_bytes += audio_data + else: + self._pcm_bytes += base64.b64decode(audio_data) + if not self.chunk_event.is_set(): + self.chunk_event.set() + + elif event_type == "session.finished": + self.chunk_event.set() + self.finish_event.set() + + except Exception: + logger.exception("Error in TTS WebSocket callback") + self.finish_event.set() + + def on_close(self, close_status_code: int, close_msg: str) -> None: + """Handle WebSocket connection closure.""" + self.finish_event.set() + self.chunk_event.set() + if close_status_code: + logger.warning( + "TTS WebSocket closed with code %s: %s", + close_status_code, + close_msg, + ) + + def _take_delta(self, header: bool = False) -> bytes | None: + """Return new PCM bytes since last call, or None if empty. + + Args: + header: If True, prepend a streaming WAV header to the + first returned chunk. + """ + new_data = self._pcm_bytes[self._consumed :] + if not new_data: + return None + self._consumed = len(self._pcm_bytes) + if header: + return _build_streaming_wav_header( + sample_rate=_SAMPLE_RATE, + channels=_CHANNELS, + bits_per_sample=_BITS_PER_SAMPLE, + ) + bytes(new_data) + return bytes(new_data) + + def get_audio_response(self, block: bool) -> TTSResponse: + """Return incremental audio delta (non-blocking or blocking).""" + if block: + self.finish_event.wait() + delta = self._take_delta(header=self._consumed == 0) + if delta: + return TTSResponse( + content=DataBlock( + source=Base64Source( + data=base64.b64encode(delta).decode("ascii"), + media_type=_MEDIA_TYPE, + ), + ), + ) + return TTSResponse(content=None) + + async def get_audio_chunks(self) -> AsyncGenerator[TTSResponse, None]: + """Yield incremental WAV audio chunks as they arrive.""" + header_sent = self._consumed > 0 + while True: + if self.finish_event.is_set(): + delta = self._take_delta(header=not header_sent) + if delta: + yield TTSResponse( + content=DataBlock( + source=Base64Source( + data=base64.b64encode(delta).decode( + "ascii", + ), + media_type=_MEDIA_TYPE, + ), + ), + is_last=True, + ) + else: + yield TTSResponse(content=None, is_last=True) + self.reset() + break + + if self.chunk_event.is_set(): + self.chunk_event.clear() + else: + await asyncio.to_thread(self.chunk_event.wait) + + if self.finish_event.is_set(): + continue + + delta = self._take_delta(header=not header_sent) + if delta: + header_sent = True + yield TTSResponse( + content=DataBlock( + source=Base64Source( + data=base64.b64encode(delta).decode("ascii"), + media_type=_MEDIA_TYPE, + ), + ), + is_last=False, + ) + + def reset(self) -> None: + """Reset internal state for the next utterance.""" + self.finish_event.clear() + self.chunk_event.clear() + self._pcm_bytes = bytearray() + self._consumed = 0 + + def has_audio_data(self) -> bool: + """Return whether any audio data has been received.""" + return bool(self._pcm_bytes) + + return _Callback + + +class DashScopeRealtimeTTSModel(TTSModelBase): + """DashScope Realtime TTS model using the QwenTtsRealtime WebSocket API. + + This model supports streaming input: text can be pushed incrementally + via :meth:`push`, and :meth:`synthesize` finalizes the current utterance. + + For more details see the `official document + `_. + + .. note:: Only one streaming input request can be active at a time. + """ + + class Parameters(BaseModel): + """Frontend-exposed parameters for DashScope Realtime TTS models.""" + + voice: str = Field( + default="Cherry", + title="Voice", + description="The voice to use for synthesis.", + ) + + type: Literal["dashscope_realtime_tts"] = "dashscope_realtime_tts" + + realtime: bool = True + + def __init__( + self, + credential: DashScopeCredential, + model: str = "qwen3-tts-flash-realtime", + parameters: "DashScopeRealtimeTTSModel.Parameters | None" = None, + stream: bool = True, + cold_start_length: int | None = None, + cold_start_words: int | None = None, + max_retries: int = 3, + retry_delay: float = 5.0, + ) -> None: + """Initialize the DashScope Realtime TTS model. + + Args: + credential (`DashScopeCredential`): + The DashScope credential. + model (`str`, defaults to ``"qwen3-tts-flash-realtime"``): + The realtime TTS model name. + parameters (`Parameters | None`, defaults to `None`): + The TTS parameters (voice, etc.). + stream (`bool`, defaults to `True`): + Whether :meth:`synthesize` returns a streaming async generator. + cold_start_length (`int | None`, defaults to `None`): + Minimum character count before the first text chunk is sent. + cold_start_words (`int | None`, defaults to `None`): + Minimum word count before the first text chunk is sent. + max_retries (`int`, defaults to `3`): + Max retry attempts on WebSocket failure. + retry_delay (`float`, defaults to `5.0`): + Initial retry delay in seconds (exponential backoff). + """ + super().__init__( + credential=credential, + model=model, + parameters=parameters, + stream=stream, + ) + self.cold_start_length = cold_start_length + self.cold_start_words = cold_start_words + self.max_retries = max_retries + self.retry_delay = retry_delay + + self._tts_client: QwenTtsRealtime | None = None + self._callback: Any = None + self._connected = False + self._cold_start_buffer: str = "" + self._cold_start_done: bool = False + self._accumulated_text: str = "" + + def _create_client(self) -> None: + """Create a fresh TTS client and callback.""" + import dashscope + from dashscope.audio.qwen_tts_realtime import QwenTtsRealtime + + dashscope.api_key = self.credential.api_key.get_secret_value() + + callback_cls = _make_callback_class() + self._callback = callback_cls() + self._tts_client = QwenTtsRealtime( + model=self.model, + callback=self._callback, + ) + + async def connect(self) -> None: + """Establish the WebSocket connection.""" + if self._connected: + return + + self._create_client() + self._tts_client.connect() + self._tts_client.update_session( + voice=self.parameters.voice, + mode="server_commit", + ) + self._connected = True + + async def close(self) -> None: + """Close the WebSocket connection.""" + if not self._connected: + return + self._connected = False + try: + self._tts_client.close() + except Exception: + pass + + async def _reconnect(self) -> None: + """Reconnect by recreating the client.""" + try: + self._tts_client.close() + except Exception: + pass + self._connected = False + self._cold_start_buffer = "" + self._cold_start_done = False + await self.connect() + + async def push( + self, + text: str, + **kwargs: Any, + ) -> TTSResponse: + """Push an incremental text delta for realtime synthesis. + + Args: + text (`str`): + An incremental text chunk (delta) to append. + **kwargs (`Any`): + Additional keyword arguments (unused). + + Returns: + `TTSResponse`: + Audio accumulated so far, or empty if not yet available. + """ + from websocket import WebSocketConnectionClosedException + + if not self._connected: + raise RuntimeError( + "TTS model is not connected. Call `connect()` first.", + ) + + if not text: + return TTSResponse(content=None) + + self._accumulated_text += text + + if not self._cold_start_done: + self._cold_start_buffer += text + ready = True + if ( + self.cold_start_length + and len(self._cold_start_buffer) < self.cold_start_length + ): + ready = False + if ( + ready + and self.cold_start_words + and len(self._cold_start_buffer.split()) + < self.cold_start_words + ): + ready = False + if ready: + try: + self._tts_client.append_text(self._cold_start_buffer) + except WebSocketConnectionClosedException: + return TTSResponse(content=None) + self._cold_start_buffer = "" + self._cold_start_done = True + else: + try: + self._tts_client.append_text(text) + except WebSocketConnectionClosedException: + return TTSResponse(content=None) + + return self._callback.get_audio_response(block=False) + + async def synthesize( + self, + text: str | None = None, + **kwargs: Any, + ) -> TTSResponse | AsyncGenerator[TTSResponse, None]: + """Finalize synthesis for the current utterance. + + If text was previously pushed via :meth:`push`, this flushes any + remaining buffered text, commits, and waits for audio. If ``text`` + is provided, it is appended before committing. + + Args: + text (`str | None`, defaults to `None`): + Optional additional text to append before finalizing. + If ``None``, finalizes previously pushed text. + **kwargs (`Any`): + Additional keyword arguments (unused). + + Returns: + `TTSResponse | AsyncGenerator[TTSResponse, None]`: + A single response when ``stream=False``, or an async generator + of incremental chunks when ``stream=True``. + """ + from websocket import WebSocketConnectionClosedException + + if not self._connected: + raise RuntimeError( + "TTS model is not connected. Call `connect()` first.", + ) + + if text is not None: + self._accumulated_text += text + + unsent = self._cold_start_buffer + if text is not None: + unsent += text + self._cold_start_buffer = "" + + full_text = self._accumulated_text + delay = self.retry_delay + + for attempt in range(self.max_retries): + try: + if unsent: + self._tts_client.append_text(unsent) + + self._tts_client.commit() + self._tts_client.finish() + + await asyncio.to_thread( + self._callback.finish_event.wait, + ) + + if full_text and not self._callback.has_audio_data(): + if attempt < self.max_retries - 1: + logger.warning( + "TTS: no audio received, retrying (%d/%d) in " + "%.1fs...", + attempt + 1, + self.max_retries, + delay, + ) + await asyncio.sleep(delay) + await self._reconnect() + unsent = full_text + delay *= 2 + continue + self._reset_state() + raise RuntimeError( + f"TTS synthesis failed: no audio after " + f"{self.max_retries} attempts", + ) + break + + except WebSocketConnectionClosedException: + if attempt < self.max_retries - 1: + logger.warning( + "TTS WebSocket closed, retrying (%d/%d) in %.1fs...", + attempt + 1, + self.max_retries, + delay, + ) + await asyncio.sleep(delay) + await self._reconnect() + unsent = full_text + delay *= 2 + else: + self._reset_state() + raise + + self._reset_state() + + if self.stream: + return self._callback.get_audio_chunks() + + response = self._callback.get_audio_response(block=True) + self._callback.reset() + return response + + def _reset_state(self) -> None: + """Reset per-utterance tracking state.""" + self._cold_start_buffer = "" + self._cold_start_done = False + self._accumulated_text = "" diff --git a/src/agentscope/tts/_tts_base.py b/src/agentscope/tts/_tts_base.py new file mode 100644 index 0000000000000000000000000000000000000000..6e7b0c908997d5569a71aa2f3548a1b2f125ff05 --- /dev/null +++ b/src/agentscope/tts/_tts_base.py @@ -0,0 +1,212 @@ +# -*- coding: utf-8 -*- +"""The TTS model base class.""" +import inspect +from abc import abstractmethod +from pathlib import Path +from typing import TYPE_CHECKING, Any, AsyncGenerator + +from pydantic import BaseModel + +from ._tts_response import TTSResponse +from .._logging import logger +from ..credential import CredentialBase + +if TYPE_CHECKING: + from ._tts_model_card import TTSModelCard + + +class TTSModelBase: + """Base class for TTS models in AgentScope. + + This base class provides a unified abstraction for both non-realtime and + realtime (streaming-input) TTS models, governed by the + ``realtime`` flag. + + For non-realtime TTS models, only :meth:`synthesize` needs to be + implemented. For realtime TTS models, the lifecycle is managed via the + async context manager (``async with model: ...``) or by calling + :meth:`connect` / :meth:`close` manually; :meth:`push` appends text + chunks and returns whatever audio is currently available, while + :meth:`synthesize` blocks until the full speech has been synthesized. + """ + + class Parameters(BaseModel): + """Base parameter schema for TTS models. Subclasses should override + this with provider-specific parameters.""" + + credential: CredentialBase + """The credential used to authenticate against the TTS provider.""" + + model: str + """The name of the TTS model.""" + + parameters: BaseModel + """The TTS model parameters.""" + + stream: bool + """Whether to use streaming output if supported by the model.""" + + realtime: bool = False + """Whether the TTS model supports realtime (streaming-input) mode.""" + + def __init__( + self, + credential: CredentialBase, + model: str, + parameters: BaseModel | None = None, + stream: bool = True, + ) -> None: + """Initialize the TTS model base class. + + Args: + credential (`CredentialBase`): + The credential used to authenticate against the TTS provider. + model (`str`): + The name of the TTS model. + parameters (`BaseModel | None`, defaults to `None`): + The TTS model parameters. + stream (`bool`, defaults to `True`): + Whether to use streaming output if supported by the model. + """ + self.credential = credential + self.model = model + self.parameters = parameters or self.Parameters() + self.stream = stream + + async def __aenter__(self) -> "TTSModelBase": + """Enter the async context manager and initialize resources if + needed.""" + if self.realtime: + await self.connect() + return self + + async def __aexit__( + self, + exc_type: Any, + exc_value: Any, + traceback: Any, + ) -> None: + """Exit the async context manager and clean up resources if needed.""" + if self.realtime: + await self.close() + + async def connect(self) -> None: + """Connect to the TTS model and initialize resources. + + .. note:: Only relevant for realtime TTS models — realtime subclasses + must override this. The default is a no-op so that non-realtime + models can ignore the lifecycle hooks; :meth:`__aenter__` only + calls it when ``realtime`` is True. + """ + return + + async def close(self) -> None: + """Close the connection to the TTS model and clean up resources. + + .. note:: Only relevant for realtime TTS models — realtime subclasses + must override this. See :meth:`connect` for the rationale behind + the no-op default. + """ + return + + async def push( # pylint: disable=unused-argument + self, + text: str, + **kwargs: Any, + ) -> TTSResponse: + """Append text to be synthesized and return the received TTS response. + This method is non-blocking and may return an empty response if no + audio is available yet. + + To receive all the synthesized speech, call :meth:`synthesize` after + pushing all the text chunks. + + .. note:: Only relevant for realtime TTS models — realtime subclasses + must override this. Non-realtime models should call + :meth:`synthesize` directly and never reach this method. + + Args: + text (`str`): + The text chunk to be synthesized. + **kwargs (`Any`): + Additional keyword arguments to pass to the TTS API call. + + Returns: + `TTSResponse`: + The TTSResponse containing the audio block. + """ + return TTSResponse(content=None) + + @classmethod + def list_models( + cls, + custom_yaml_dir: str | None = None, + ) -> list["TTSModelCard"]: + """List candidate TTS models by scanning YAML model cards. + + Args: + custom_yaml_dir (`str | None`): + The custom YAML directory. If ``None``, uses the ``_models`` + directory next to the concrete subclass's source file. + + Returns: + `list[TTSModelCard]`: + A list of TTS model cards. + """ + from ._tts_model_card import TTSModelCard + + if custom_yaml_dir is None: + subclass_file = Path(inspect.getfile(cls)) + yaml_dir = subclass_file.parent / "_models" + else: + yaml_dir = Path(custom_yaml_dir) + + yaml_files = list(yaml_dir.glob("*.yaml")) + + model_cards = [] + for yaml_file in yaml_files: + try: + card = TTSModelCard.from_yaml( + yaml_path=str(yaml_file), + parameter_class=cls.Parameters, + ) + if card.realtime != cls.realtime: + continue + model_cards.append(card) + except Exception as e: + logger.warning( + "Warning: Failed to load %s: %s", + yaml_file, + str(e), + ) + continue + + return model_cards + + @abstractmethod + async def synthesize( + self, + text: str | None = None, + **kwargs: Any, + ) -> TTSResponse | AsyncGenerator[TTSResponse, None]: + """Synthesize speech from the appended text. Different from + :meth:`push`, this method blocks until the full speech has been + synthesized. + + Args: + text (`str | None`, defaults to `None`): + The text to be synthesized. If `None`, this method will + wait for all previously pushed text to be synthesized and + return the last synthesized TTSResponse. + **kwargs (`Any`): + Additional keyword arguments to pass to the TTS API call. + + Returns: + `TTSResponse | AsyncGenerator[TTSResponse, None]`: + A single TTSResponse containing the full audio when + ``stream=False``. When ``stream=True``, an async generator + yielding TTSResponse chunks where each chunk carries an + **incremental** audio delta (not a cumulative buffer); the + full audio is the concatenation of every chunk's decoded + bytes. The final yielded chunk has ``is_last=True``. + """ diff --git a/src/agentscope/tts/_tts_model_card.py b/src/agentscope/tts/_tts_model_card.py new file mode 100644 index 0000000000000000000000000000000000000000..ed59e8358384efc0f858778d8362c2795a42a59b --- /dev/null +++ b/src/agentscope/tts/_tts_model_card.py @@ -0,0 +1,134 @@ +# -*- coding: utf-8 -*- +"""The TTS model card class.""" +import copy +from datetime import datetime +from typing import Literal, Self, Type + +import yaml +from pydantic import BaseModel, Field + + +class TTSModelCard(BaseModel): + """The model card class for TTS models.""" + + type: Literal["tts_model"] = "tts_model" + """The model card type discriminator.""" + + name: str = Field(description="The name of the TTS model") + """The model name.""" + + label: str = Field(description="The model label.") + """The model label used for frontend rendering.""" + + status: Literal["active", "deprecated", "sunset"] = Field( + default="active", + title="Status", + description="The model status", + ) + """The model status.""" + + deprecated_at: datetime | None = Field( + default=None, + description="The model deprecation date and time.", + title="Deprecation date", + ) + """The model deprecated at.""" + + input_types: list[str] = Field( + description="The supported model input types.", + title="Input types", + default=["text/plain"], + ) + """The model supported input types.""" + + output_types: list[str] = Field( + description="The supported model output types.", + title="Output types", + default=["audio/wav"], + ) + """The model supported output types.""" + + realtime: bool = Field( + default=False, + description="Whether the model supports streaming input (realtime).", + title="Realtime", + ) + """Whether the model supports streaming input.""" + + parameter_schema: dict + """The parameters schema.""" + + parameters_overrides: dict[str, dict] + """The parameter overrides.""" + + @classmethod + def from_yaml( + cls, + yaml_path: str, + parameter_class: Type[BaseModel], + ) -> Self: + """Read a TTS model card from a YAML file. + + Args: + yaml_path (`str`): + Path to the YAML file + parameter_class (`Type[BaseModel]`): + The parameter class (e.g., DashScopeTTSModel.Parameters) + + Returns: + `TTSModelCard`: + TTSModelCard instance with merged parameter schema + """ + with open(yaml_path, "r", encoding="utf-8") as file: + config = yaml.safe_load(file) + + base_schema = parameter_class.model_json_schema() + properties = copy.deepcopy(base_schema.get("properties", {})) + + # Auto-inject: populate voice enum from YAML voices list + voices = config.get("voices", []) + if voices and "voice" in properties: + properties["voice"] = { + **properties["voice"], + "default": voices[0], + "enum": voices, + } + + # Apply parameter_overrides + overrides = config.get("parameter_overrides", {}) + for param_name, override in overrides.items(): + if override is None: + properties.pop(param_name, None) + continue + + if isinstance(override, dict): + if override.get("hidden"): + properties.pop(param_name, None) + continue + + if param_name in properties: + properties[param_name] = { + **properties[param_name], + **override, + } + + required = [ + r for r in base_schema.get("required", []) if r in properties + ] + final_schema = { + "type": "object", + "properties": properties, + "required": required, + } + + return cls( + name=config["name"], + label=config["label"], + status=config.get("status", "active"), + deprecated_at=config.get("deprecated_at"), + input_types=config.get("input_types", ["text/plain"]), + output_types=config.get("output_types", ["audio/wav"]), + realtime=config.get("realtime", False), + parameter_schema=final_schema, + parameters_overrides=config.get("parameter_overrides", {}), + ) diff --git a/src/agentscope/tts/_tts_response.py b/src/agentscope/tts/_tts_response.py new file mode 100644 index 0000000000000000000000000000000000000000..677d7c29be391e26bfaf21e29e6f4c4589f8a7b9 --- /dev/null +++ b/src/agentscope/tts/_tts_response.py @@ -0,0 +1,56 @@ +# -*- coding: utf-8 -*- +"""The TTS response module.""" +from dataclasses import dataclass, field +from typing import Literal + +from .._utils._common import _get_timestamp +from .._utils._mixin import DictMixin +from ..message import DataBlock +from ..types import JSONSerializableObject + + +@dataclass +class TTSUsage(DictMixin): + """The usage of a TTS model API invocation.""" + + input_tokens: int + """The number of input tokens.""" + + output_tokens: int + """The number of output tokens.""" + + time: float + """The time used in seconds.""" + + type: Literal["tts"] = field(default_factory=lambda: "tts") + """The type of the usage, must be `tts`.""" + + +@dataclass +class TTSResponse(DictMixin): + """The response of TTS models.""" + + content: DataBlock | None + """The audio chunk produced by the TTS model. The audio format is + indicated by ``content.source.media_type`` (e.g. + ``"audio/pcm;rate=24000"``, ``"audio/mpeg"``).""" + + id: str = field(default_factory=lambda: _get_timestamp(True)) + """The unique identifier of the response.""" + + created_at: str = field(default_factory=_get_timestamp) + """When the response was created.""" + + type: Literal["tts"] = field(default_factory=lambda: "tts") + """The type of the response, which is always 'tts'.""" + + usage: TTSUsage | None = field(default_factory=lambda: None) + """The usage information of the TTS response, if available.""" + + metadata: dict[str, JSONSerializableObject] | None = field( + default_factory=lambda: None, + ) + """The metadata of the TTS response.""" + + is_last: bool = field(default_factory=lambda: True) + """Whether this is the last response in a stream of TTS responses.""" diff --git a/src/agentscope/types/__init__.py b/src/agentscope/types/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ac6d3885996823b0cdacbd0f47ed4f6c457024b8 --- /dev/null +++ b/src/agentscope/types/__init__.py @@ -0,0 +1,20 @@ +# -*- coding: utf-8 -*- +"""The types in agentscope""" + +from ._hook import ( + AgentHookTypes, + ReActAgentHookTypes, +) +from ._object import Embedding +from ._json import ( + JSONPrimitive, + JSONSerializableObject, +) + +__all__ = [ + "AgentHookTypes", + "ReActAgentHookTypes", + "Embedding", + "JSONPrimitive", + "JSONSerializableObject", +] diff --git a/src/agentscope/types/_hook.py b/src/agentscope/types/_hook.py new file mode 100644 index 0000000000000000000000000000000000000000..924fe658f0f323bd24cd42ca03d65a6490690136 --- /dev/null +++ b/src/agentscope/types/_hook.py @@ -0,0 +1,25 @@ +# -*- coding: utf-8 -*- +"""The agent hooks types.""" +from typing import Literal + +AgentHookTypes = ( + str + | Literal[ + "pre_reply", + "post_reply", + "pre_print", + "post_print", + "pre_observe", + "post_observe", + ] +) + +ReActAgentHookTypes = ( + AgentHookTypes + | Literal[ + "pre_reasoning", + "post_reasoning", + "pre_acting", + "post_acting", + ] +) diff --git a/src/agentscope/types/_json.py b/src/agentscope/types/_json.py new file mode 100644 index 0000000000000000000000000000000000000000..e23755e0c38e86f28eaf7552586d0a9ab2dc606a --- /dev/null +++ b/src/agentscope/types/_json.py @@ -0,0 +1,14 @@ +# -*- coding: utf-8 -*- +"""The JSON related types""" +from typing import TypeAlias + +JSONPrimitive: TypeAlias = str | int | float | bool | None + +JSONSerializableObject: TypeAlias = ( + JSONPrimitive + | list["JSONSerializableObject"] + | dict[ + str, + "JSONSerializableObject", + ] +) diff --git a/src/agentscope/types/_object.py b/src/agentscope/types/_object.py new file mode 100644 index 0000000000000000000000000000000000000000..85a9b347cb9e82c682217a0ef7160e725069f47e --- /dev/null +++ b/src/agentscope/types/_object.py @@ -0,0 +1,5 @@ +# -*- coding: utf-8 -*- +"""The object types in agentscope.""" +from typing import List + +Embedding = List[float] diff --git a/src/agentscope/workspace/__init__.py b/src/agentscope/workspace/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4656194a9c8eec50d718c2e81281a934af3f671a --- /dev/null +++ b/src/agentscope/workspace/__init__.py @@ -0,0 +1,20 @@ +# -*- coding: utf-8 -*- +"""The workspace module in agentscope.""" + + +from ._base import WorkspaceBase +from ._local_workspace import LocalWorkspace +from ._offload_protocol import Offloader +from ._docker import DockerBackend, DockerWorkspace +from ._e2b import E2BWorkspace, E2BBackend + + +__all__ = [ + "WorkspaceBase", + "LocalWorkspace", + "DockerBackend", + "DockerWorkspace", + "E2BBackend", + "E2BWorkspace", + "Offloader", +] diff --git a/src/agentscope/workspace/_base.py b/src/agentscope/workspace/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..dc43dd8a3291386f842a2019259ab4a598afc06c --- /dev/null +++ b/src/agentscope/workspace/_base.py @@ -0,0 +1,230 @@ +# -*- coding: utf-8 -*- +"""WorkspaceBase — abstract interface for agent workspaces. + +A workspace provides: + +- **Resources** — skills available to the agent. +- **Tools** — MCPs and built-in tools for operating on resources. +- **Offload** — persistence of compressed context and tool results + for agentic retrieval. + +Three concrete implementations: + +- `LocalWorkspace` — local filesystem. +- `DockerWorkspace` — Docker container. +- `E2BWorkspace` — E2B cloud sandbox. + +Consumers: + +- **Agent** — calls ``list_mcps``, ``list_skills``, ``list_tools``, + ``offload_context``, ``offload_tool_result``. +- **User** — dynamically adds/removes MCPs and skills via + ``add_mcp``, ``remove_mcp``, ``add_skill``, ``remove_skill``. +- **Developer** — manages lifecycle via ``initialize`` / ``close``. +- **Backend consumers** access the active backend via ``get_backend``. +""" + +from abc import abstractmethod +from typing import Self + +from .._utils._common import _generate_id +from ..mcp import MCPClient +from ..message import Msg, ToolResultBlock +from ..skill import Skill +from ..tool import BackendBase, ToolBase + + +class WorkspaceBase: + """Abstract base class for all workspace implementations. + + Subclasses provide concrete behaviour for one execution backend + (local filesystem, Docker container, E2B sandbox). The base class + only fixes the lifecycle contract (``initialize`` / ``close`` / + ``reset``), the ``async with`` protocol, and the discovery / + offload / add-remove method signatures consumed by ``Agent`` and + by the workspace manager layer. + + State held on the base class is intentionally minimal: + ``workspace_id`` (stable identifier, generated if not given) and + ``is_alive`` (lifecycle flag). All backend-specific state lives on + the subclass. + """ + + workspace_id: str + """Unique identifier for this workspace instance.""" + + workdir: str + """Agent-visible root directory for workspace file operations.""" + + is_alive: bool + """If the workspace is still operational.""" + + _backend: BackendBase | None + """Current execution backend, available through :meth:`get_backend`.""" + + def __init__(self, workspace_id: str | None) -> None: + """Initialize the workspace base instance.""" + self.workspace_id = workspace_id or _generate_id() + self.is_alive = False + + # ── lifecycle (developer) ────────────────────────────────────── + + @abstractmethod + async def initialize(self) -> None: + """Provision resources, connect MCP servers, copy skills.""" + + @abstractmethod + async def close(self) -> None: + """Release all resources and connections.""" + + async def reset(self) -> None: + """Reset the workspace to a clean state. + + Closes and removes all registered MCPs, deletes all skills, + and wipes per-session state (offloaded context / tool results + and any data files). Constructor-time ``default_mcps`` and + ``skill_paths`` are **not** re-seeded — reset returns the + workspace to an empty state, not its initial state. + + The default implementation is a no-op. Subclasses with user + state must override this. + """ + + def get_backend(self) -> BackendBase: + """Return the workspace's active filesystem/execution backend. + + Docker and E2B workspaces may replace their backend when reconnecting, + so callers should resolve it from the workspace when beginning an + operation rather than retaining a stale private ``_backend`` value. + + Raises: + RuntimeError: + If the workspace has not been initialized or has no active + backend. + """ + backend = getattr(self, "_backend", None) + if not isinstance(backend, BackendBase): + raise RuntimeError( + f"{type(self).__name__} has no active backend. " + "Initialize the workspace before requesting its backend.", + ) + return backend + + async def __aenter__(self) -> Self: + """Context manager support for ``async with``. Calls ``initialize()`` + and returns the workspace instance. + """ + await self.initialize() + self.is_alive = True + return self + + async def __aexit__(self, *exc: object) -> None: + """Context manager support for ``async with``. Calls ``close()`` + and returns the workspace instance. + """ + await self.close() + self.is_alive = False + + # ── instructions ─────────────────────────────────────────────── + + @abstractmethod + async def get_instructions(self) -> str: + """Workspace-specific system prompt fragment.""" + + # ── for Agent: tool & resource discovery ─────────────────────── + + @abstractmethod + async def list_tools(self) -> list[ToolBase]: + """Built-in tools scoped to this workspace.""" + + @abstractmethod + async def list_mcps(self) -> list[MCPClient]: + """Active MCP clients (each provides its own tools).""" + + @abstractmethod + async def list_skills(self) -> list[Skill]: + """Skills available in this workspace.""" + + # ── for Agent: offload ───────────────────────────────────────── + + @abstractmethod + async def offload_context( + self, + session_id: str, + msgs: list[Msg], + ) -> str: + """Persist compressed context for agentic retrieval. + + Args: + session_id: Unique session identifier used to + partition offloaded data. + msgs: Conversation messages to offload. + + Returns: + Path or identifier for the offloaded data. + """ + + @abstractmethod + async def offload_tool_result( + self, + session_id: str, + tool_result: ToolResultBlock, + ) -> str: + """Persist a tool result for agentic retrieval. + + Args: + session_id: Unique session identifier used to + partition offloaded data. + tool_result: The tool result block to offload. + + Returns: + Path or identifier for the offloaded data. + """ + + # ── for User: dynamic MCP management ─────────────────────────── + + @abstractmethod + async def add_mcp(self, mcp_client: MCPClient) -> None: + """Dynamically register a new MCP server. + + Args: + mcp_client: An :class:`MCPClient` instance describing + the MCP server to add. + + Raises: + ValueError: If an MCP with the same name already exists. + """ + + @abstractmethod + async def remove_mcp(self, name: str) -> None: + """Dynamically remove an MCP server by name. + + Args: + name: Name of the MCP server to remove. + """ + + # ── for User: dynamic skill management ───────────────────────── + + @abstractmethod + async def add_skill(self, skill_path: str) -> None: + """Add a skill from a local directory path. + + The directory must contain a ``SKILL.md`` with ``name`` + and ``description`` in its YAML front matter. + + Args: + skill_path: Absolute or relative path to the skill + directory on the local filesystem. + """ + + @abstractmethod + async def remove_skill(self, name: str) -> None: + """Remove a skill by its agent-facing name. + + Args: + name: The ``name`` field from the skill's + ``SKILL.md`` front matter. + + Raises: + KeyError: If the skill is not found in the workspace. + """ diff --git a/src/agentscope/workspace/_docker/Dockerfile.install_pypi.template b/src/agentscope/workspace/_docker/Dockerfile.install_pypi.template new file mode 100644 index 0000000000000000000000000000000000000000..ed2d895ca2938b86bf38a1c9f9406b7300a93b3b --- /dev/null +++ b/src/agentscope/workspace/_docker/Dockerfile.install_pypi.template @@ -0,0 +1,5 @@ +# ``--no-deps``: see Dockerfile.install_src.template for the rationale. +# Gateway only needs ``agentscope.mcp.MCPClient`` and its lightweight +# transitive imports (mcp, pydantic, httpx, fastapi), all of which are +# pulled in by requirements.txt above. +RUN uv pip install --no-deps "agentscope=={agentscope_version}" diff --git a/src/agentscope/workspace/_docker/Dockerfile.install_src.template b/src/agentscope/workspace/_docker/Dockerfile.install_src.template new file mode 100644 index 0000000000000000000000000000000000000000..bc1d65b8eacd2136f8d7f848b634b15b420a45d1 --- /dev/null +++ b/src/agentscope/workspace/_docker/Dockerfile.install_src.template @@ -0,0 +1,12 @@ +# TODO(release): when agentscope is on PyPI, drop this whole block and use +# Dockerfile.install_pypi.template instead. Copying the source tree is a +# transitional dev-only path so the in-container gateway can `import agentscope`. +COPY agentscope_src /tmp/agentscope_src +# ``--no-deps``: the gateway only imports ``agentscope.mcp.MCPClient``, +# whose transitive needs (mcp, pydantic, httpx, fastapi) are already +# installed via requirements.txt above. Skipping agentscope's full +# dependency tree avoids dragging in heavy/Rust-built packages +# (ripgrep, tree_sitter, opentelemetry, anthropic/openai SDKs, …) +# which would otherwise require a C/Rust toolchain in the slim image. +RUN uv pip install --no-deps /tmp/agentscope_src \ + && rm -rf /tmp/agentscope_src diff --git a/src/agentscope/workspace/_docker/Dockerfile.node_copy.template b/src/agentscope/workspace/_docker/Dockerfile.node_copy.template new file mode 100644 index 0000000000000000000000000000000000000000..2731985c6639db15f47007a193287f672dc01c70 --- /dev/null +++ b/src/agentscope/workspace/_docker/Dockerfile.node_copy.template @@ -0,0 +1,4 @@ +COPY --from=node_stage /usr/local/bin/node /usr/local/bin/node +COPY --from=node_stage /usr/local/lib/node_modules /usr/local/lib/node_modules +RUN ln -sf /usr/local/lib/node_modules/npm/bin/npm-cli.js /usr/local/bin/npm \ + && ln -sf /usr/local/lib/node_modules/npm/bin/npx-cli.js /usr/local/bin/npx diff --git a/src/agentscope/workspace/_docker/Dockerfile.node_from.template b/src/agentscope/workspace/_docker/Dockerfile.node_from.template new file mode 100644 index 0000000000000000000000000000000000000000..396843f3c7ad7eee5565db0771fd0af0abc10ff9 --- /dev/null +++ b/src/agentscope/workspace/_docker/Dockerfile.node_from.template @@ -0,0 +1 @@ +FROM node:{node_version}-slim AS node_stage \ No newline at end of file diff --git a/src/agentscope/workspace/_docker/Dockerfile.template b/src/agentscope/workspace/_docker/Dockerfile.template new file mode 100644 index 0000000000000000000000000000000000000000..7531ca95d509e5921051ab26ec470a5504566dd5 --- /dev/null +++ b/src/agentscope/workspace/_docker/Dockerfile.template @@ -0,0 +1,47 @@ +# syntax=docker/dockerfile:1.6 +{node_from_block}FROM {base_image} + +# Install uv via Astral's official shell installer. +# We use this rather than ``COPY --from=ghcr.io/astral-sh/uv`` to avoid +# depending on GitHub Container Registry at build time (the GHCR blob CDN +# has been observed to time out from some networks). ``curl`` and +# ``ca-certificates`` are not present in ``python:*-slim`` so we apt them +# in first, then drop the apt cache to keep the layer small. +# ``UV_INSTALL_DIR`` redirects the installer's default ``~/.local/bin`` +# target to ``/usr/local/bin`` so ``uv`` lands directly on PATH; +# ``INSTALLER_NO_MODIFY_PATH=1`` suppresses the installer's attempt to +# patch shell rc files (we don't need it — PATH already covers it). +RUN apt-get update \ + && apt-get install -y --no-install-recommends curl ca-certificates ripgrep \ + && rm -rf /var/lib/apt/lists/* \ + && curl -LsSf https://astral.sh/uv/install.sh \ + | env UV_INSTALL_DIR=/usr/local/bin INSTALLER_NO_MODIFY_PATH=1 sh + +{node_copy_block}ENV UV_PROJECT_ENVIRONMENT={gateway_home}/.venv \ + UV_LINK_MODE=copy \ + PATH={gateway_home}/.venv/bin:$PATH +WORKDIR {gateway_home} +RUN uv venv {gateway_home}/.venv + +# TODO(release): once agentscope ships an extra such as `agentscope[gateway]` +# bundling fastapi/uvicorn/mcp, drop requirements.txt and fold these pins into +# the install step below — keeping a single `uv pip install agentscope[gateway]`. +COPY requirements.txt {gateway_home}/requirements.txt +# Note: aiodocker drives the legacy build endpoint, which does not +# support BuildKit cache mounts (``--mount=type=cache,...``). uv's +# in-package cache still works for single-build deduping; we just lose +# inter-build layer cache reuse — acceptable since the image tag hash +# already gates re-builds at the layer level. +RUN uv pip install -r {gateway_home}/requirements.txt + +{install_agentscope_block} +# Standalone gateway script. We invoke this directly (not via +# ``python -m``) so Python does not auto-import +# ``agentscope.workspace.__init__`` and pull in its heavy module +# graph (skill, tool, …). The script imports only +# ``agentscope.mcp.MCPClient``. +COPY _mcp_gateway_app.py {gateway_home}/_mcp_gateway_app.py +# Glob helper script used by the builtin Glob tool. +COPY _glob_helper.py {gateway_home}/_glob_helper.py + +WORKDIR {container_workdir} diff --git a/src/agentscope/workspace/_docker/__init__.py b/src/agentscope/workspace/_docker/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ba1ef2c438a1d198b7fde88c2e428f39efc81565 --- /dev/null +++ b/src/agentscope/workspace/_docker/__init__.py @@ -0,0 +1,13 @@ +# -*- coding: utf-8 -*- +"""Docker-backed workspace. + +The container image is built on demand from a content-hashed Dockerfile +(see :mod:`._make_dockerfile`); a tag cache hit skips the build. MCP +servers run *inside* the container behind a FastAPI gateway and are +reached over HTTP — see :mod:`agentscope.workspace._gateway_client`. +""" + +from ._docker_workspace import DockerWorkspace +from ._docker_backend import DockerBackend + +__all__ = ["DockerWorkspace", "DockerBackend"] diff --git a/src/agentscope/workspace/_docker/_docker_backend.py b/src/agentscope/workspace/_docker/_docker_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..b15a8a8826cab6e5a2a4d35ba3346f574b436fe2 --- /dev/null +++ b/src/agentscope/workspace/_docker/_docker_backend.py @@ -0,0 +1,193 @@ +# -*- coding: utf-8 -*- +"""Docker container :class:`BackendBase` implementation. + +Wraps the ``aiodocker`` container APIs (``exec``, ``get_archive``, +``put_archive``) into the three backend primitives (``exec_shell``, +``read_file``, ``write_file``) so that builtin tools (Bash, Read, +Write, Edit, Grep, Glob) can operate inside a Docker container +transparently. All derived filesystem helpers (``file_exists``, +``is_dir``, ``list_dir``, ``stat_mtime``, ``delete_path``) are +inherited from :class:`BackendBase`, which implements them via +``exec_shell``. +""" + +from __future__ import annotations + +import asyncio +import io +import posixpath +import tarfile +from typing import Any + +from ...tool import BackendBase, ExecResult + + +class DockerBackend(BackendBase): + """Backend that delegates to a running Docker container. + + Only the three abstract primitives (``exec_shell``, ``read_file``, + ``write_file``) are implemented here; the derived filesystem helpers + are inherited from :class:`BackendBase`. + + Args: + container (`Any`): + An ``aiodocker`` container object (must already be started). + workdir (`str`): + Default working directory for ``exec_shell`` calls inside + the container. + """ + + def __init__(self, container: Any, workdir: str) -> None: + """Initialize the Docker backend. + + Args: + container (`Any`): + A started ``aiodocker`` container object. + workdir (`str`): + Default working directory for ``exec_shell`` calls + inside the container. + """ + self._container = container + self._workdir = workdir + + # ── exec ─────────────────────────────────────────────────────── + + async def getcwd(self) -> str: + """Return the container's default working directory. + + Overrides the base class default (which would shell out to + ``pwd``) with the cached ``workdir`` supplied at construction, + avoiding a per-call container ``exec`` round-trip. + + Returns: + `str`: + The container's default working directory. + """ + return self._workdir + + async def exec_shell( + self, + command: list[str], + *, + cwd: str | None = None, + timeout: float | None = None, + ) -> ExecResult: + """Run a program directly inside the container. + + *command* is an argv list executed via the container ``exec`` + API without an intervening shell. Callers needing shell features + wrap their command line as ``["sh", "-c", line]``. + + Args: + command (`list[str]`): + Executable path/name followed by its arguments. + cwd (`str | None`, optional): + Working directory inside the container. When ``None`` + the backend's default ``workdir`` is used. + timeout (`float | None`, optional): + Maximum number of seconds to wait before returning an + ``exit_code`` of ``-1``. When ``None`` the call waits + indefinitely. + + Returns: + `ExecResult`: + The captured exit code, stdout, and stderr. + """ + + async def _run() -> ExecResult: + exec_obj = await self._container.exec( + cmd=command, + workdir=cwd or self._workdir, + ) + stdout_parts: list[bytes] = [] + stderr_parts: list[bytes] = [] + async with exec_obj.start() as stream: + while True: + msg = await stream.read_out() + if msg is None: + break + if msg.stream == 1: + stdout_parts.append(msg.data) + else: + stderr_parts.append(msg.data) + inspect = await exec_obj.inspect() + code = inspect.get("ExitCode", -1) + if code is None: + code = -1 + return ExecResult( + exit_code=int(code), + stdout=b"".join(stdout_parts), + stderr=b"".join(stderr_parts), + ) + + if timeout is None: + return await _run() + try: + return await asyncio.wait_for(_run(), timeout=timeout) + except asyncio.TimeoutError: + return ExecResult( + exit_code=-1, + stdout=b"", + stderr=b"timed out", + ) + + # ── file I/O ─────────────────────────────────────────────────── + + async def read_file(self, path: str) -> bytes: + """Fetch a file from the container via ``get_archive``. + + Args: + path (`str`): + Path to the file inside the container. + + Returns: + `bytes`: + The raw file contents. + + Raises: + `FileNotFoundError`: + If the path does not exist inside the container. + """ + from aiodocker import exceptions as aiodocker_exceptions + + try: + tar = await self._container.get_archive(path) + except aiodocker_exceptions.DockerError as exc: + if exc.status == 404: + raise FileNotFoundError( + f"not found in container: {path}", + ) from exc + raise + + try: + for member in tar.getmembers(): + if member.isfile(): + f = tar.extractfile(member) + if f: + return f.read() + finally: + tar.close() + raise FileNotFoundError(f"not found in container: {path}") + + async def write_file(self, path: str, data: bytes) -> None: + """Write raw bytes to a file inside the container. + + Creates the parent directory first since ``put_archive`` + requires it to exist. + + Args: + path (`str`): + Destination path inside the container. + data (`bytes`): + The raw bytes to write. + """ + parent = posixpath.dirname(path) or "/" + name = posixpath.basename(path) + + await self.exec_shell(["mkdir", "-p", parent]) + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w") as tf: + info = tarfile.TarInfo(name=name) + info.size = len(data) + tf.addfile(info, io.BytesIO(data)) + await self._container.put_archive(parent, buf.getvalue()) diff --git a/src/agentscope/workspace/_docker/_docker_workspace.py b/src/agentscope/workspace/_docker/_docker_workspace.py new file mode 100644 index 0000000000000000000000000000000000000000..28d5b16ec3fcddd4ba277b45905f072e1333db8f --- /dev/null +++ b/src/agentscope/workspace/_docker/_docker_workspace.py @@ -0,0 +1,1126 @@ +# -*- coding: utf-8 -*- +"""DockerWorkspace — sandboxed workspace backed by a Docker container. + +Architecture +------------ + +* Container lifecycle (build + run + stop) via **aiodocker**. +* MCP servers run *inside* the container behind a FastAPI gateway + (see :mod:`agentscope.workspace._mcp_gateway`); the host talks to it + over HTTP via :class:`GatewayClient` / :class:`GatewayMCPClient`. +* Optional bind-mounted host ``workdir`` makes the workspace + persistent — ``.mcp`` (registered MCPs), ``skills/``, ``sessions/`` + and ``data/`` survive restarts. Without ``workdir`` the container + is ephemeral. +* Image is content-hashed by Dockerfile + COPY payloads + (see :mod:`._make_dockerfile`); a cache hit skips the build. + +Persistence model mirrors :class:`agentscope.workspace.LocalWorkspace`: +on each :meth:`initialize`, MCPs are restored from ``/.mcp`` +if it exists (otherwise ``default_mcps`` are used and persisted). +Every :meth:`add_mcp` / :meth:`remove_mcp` rewrites the file. + +The gateway bearer token is freshly generated on each ``initialize`` +and shipped into the container via the gateway config file — it is +*not* persisted. +""" + +import asyncio +import base64 +import hashlib +import io +import json +import mimetypes +import os +import posixpath +import shlex +import shutil +import sys +import tarfile +import uuid +from copy import deepcopy +from typing import Any + +from pydantic import AnyUrl + +from ..._logging import logger +from ...mcp import MCPClient +from ...message import ( + Base64Source, + DataBlock, + Msg, + TextBlock, + ToolResultBlock, + URLSource, +) +from ...skill import Skill +from ...tool import ToolBase +from .._base import WorkspaceBase +from .._gateway_client import ( + GatewayClient, + GatewayMCPClient, +) +from ._docker_backend import DockerBackend +from ._make_dockerfile import ( + CONTAINER_DATA_DIR, + CONTAINER_SESSIONS_DIR, + CONTAINER_SKILLS_DIR, + CONTAINER_WORKDIR, + DEFAULT_BASE_IMAGE, + DEFAULT_GATEWAY_PORT, + GATEWAY_CONFIG, + GATEWAY_HOME, + GATEWAY_LOG, + GATEWAY_SCRIPT, + GATEWAY_VENV, + GLOB_HELPER_SCRIPT, + prepare_build_context, +) + +_DEFAULT_INSTRUCTIONS = """ +You have a Docker-based workspace. All tool calls execute **inside the +container** at ``{workdir}``. + +Layout: + +``` +{workdir} +├── data/ # offloaded multimodal files +├── skills/ # reusable skills +└── sessions/ # session context and tool results +``` + +Use the MCP-provided tools to interact with the container's filesystem +and processes. +""" + + +# ── the workspace ────────────────────────────────────────────────── + + +class DockerWorkspace(WorkspaceBase): + """Workspace backed by a Docker container. + + ``default_mcps`` and ``skill_paths`` are seed-time inputs and are + not retained as instance state past :meth:`initialize`. + """ + + def __init__( + self, + *, + workspace_id: str | None = None, + base_image: str = DEFAULT_BASE_IMAGE, + host_workdir: str | None = None, + node_version: str | None = None, + extra_pip: list[str] | None = None, + gateway_port: int = DEFAULT_GATEWAY_PORT, + env: dict[str, str] | None = None, + instructions: str = _DEFAULT_INSTRUCTIONS, + default_mcps: list[MCPClient] | None = None, + skill_paths: list[str] | None = None, + **kwargs: Any, + ) -> None: + """Construct a :class:`DockerWorkspace`. + + The workspace is *not* started here; call :meth:`initialize` + (or use the workspace as an ``async`` context manager). + + Args: + workspace_id (`str | None`, optional): + Existing workspace identifier to adopt. ``None`` + generates a fresh UUID. When the same + ``workspace_id`` is paired with a persistent + ``workdir``, restarts are stable across processes. + base_image (`str`, defaults to `DEFAULT_BASE_IMAGE`): + Base Docker image. Must provide ``python3`` in + ``$PATH`` (e.g. ``"python:3.11-slim"``). The image is + rebuilt on top of this base via the dynamic + Dockerfile (uv venv + agentscope install + optional + node + ``extra_pip``). + host_workdir (`str | None`, optional): + Host directory bind-mounted to ``/workspace`` inside + the container. ``None`` makes the workspace + ephemeral — files written under ``/workspace`` live + only in the container's writable layer and are lost + on :meth:`close`. When set, the directory is created + on demand and the ``.mcp`` / ``skills/`` / + ``sessions/`` / ``data/`` layout is mirrored + host-side. + node_version (`str | None`, optional): + Major Node.js version to bake into the image (e.g. + ``"20"``). ``None`` skips Node entirely. + extra_pip (`list[str] | None`, optional): + Extra Python packages to install into the gateway + venv at image-build time. + gateway_port (`int`, defaults to `DEFAULT_GATEWAY_PORT`): + TCP port the gateway listens on inside the + container; always exposed to a randomly assigned + host port. + env (`dict[str, str] | None`, optional): + Environment variables to set inside the container. + instructions (`str`, defaults to `_DEFAULT_INSTRUCTIONS`): + System-prompt fragment template returned by + :meth:`get_instructions`. Supports the ``{workdir}`` + placeholder, replaced with the container-side path. + default_mcps (`list[MCPClient] | None`, optional): + Initial MCPs registered on first :meth:`initialize`. + On subsequent restarts with a persistent ``workdir`` + these are ignored in favour of the persisted + ``/.mcp`` file. + skill_paths (`list[str] | None`, optional): + Local skill directories seeded into + ``/skills`` on first :meth:`initialize` + (only when ``workdir`` is set; subsequent starts + treat the host directory as the source of truth). + """ + super().__init__(workspace_id=workspace_id) + + # Backwards compatibility: accept legacy ``workdir`` kwarg. + if "workdir" in kwargs: + logger.warning( + "DockerWorkspace parameter 'workdir' is deprecated, " + "use 'host_workdir' instead.", + ) + if host_workdir is None: + host_workdir = kwargs.pop("workdir") + + # ── serializable config ───────────────────────────────── + self.workdir = CONTAINER_WORKDIR + self.base_image = base_image + self.host_workdir = host_workdir + self.node_version = node_version + self.extra_pip: list[str] = list(extra_pip or []) + self.gateway_port = gateway_port + self.env: dict[str, str] = dict(env or {}) + self.instructions = instructions + + # ── seed-only ─────────────────────────────────────────── + self.default_mcps: list[MCPClient] = list(default_mcps or []) + self.skill_paths: list[str] = list(skill_paths or []) + + # ── runtime state ─────────────────────────────────────── + self._client: Any = None # aiodocker.Docker + self._container: Any = None + self._backend: DockerBackend | None = None + self._port_mapping: dict[int, int] = {} + self._image_tag: str = "" + self._gateway: GatewayClient | None = None + self._gateway_token: str = "" + self._mcps: list[MCPClient] = [] + self._gateway_clients: dict[str, GatewayMCPClient] = {} + self._mcp_lock = asyncio.Lock() + self._skill_lock = asyncio.Lock() + + # ── lifecycle ─────────────────────────────────────────────── + + async def initialize(self) -> None: + """Build / reuse the image, start the container, launch the gateway. + + Steps: + + 1. Build the workspace image (or reuse a tag-cache hit). + 2. Restore MCPs from ``/.mcp`` if present, else seed + from ``default_mcps``. + 3. Mint a fresh gateway bearer token (not persisted). + 4. Start the container with the gateway port mapped to a host + port and ``workdir`` (if any) bind-mounted. + 5. Drop ``gateway.config.json`` into the container, launch the + gateway via ``python -m agentscope.workspace._mcp_gateway``, + and wait for ``/health`` to return 200. + 6. Pull the gateway-side MCP view back as + :class:`GatewayMCPClient` instances. + 7. Persist ``.mcp`` and seed skills (only when ``workdir`` is + set). + + Idempotent — calling on an already-alive workspace is a no-op. + + Raises: + RuntimeError: If the image build fails, the gateway port + fails to bind, or the gateway does not become healthy + within 30 seconds. + """ + if self.is_alive: + return + + import aiodocker + + self._client = aiodocker.Docker() + + await self._build_or_reuse_image() + + self._mcps = await self._restore_or_seed_mcps() + + self._gateway_token = uuid.uuid4().hex + + await self._create_and_start_container() + + await self._write_gateway_config() + await self._start_gateway_process() + + host_port = self._port_mapping[self.gateway_port] + self._gateway = GatewayClient( + base_url=f"http://127.0.0.1:{host_port}", + token=self._gateway_token, + timeout=30.0, + ) + await self._wait_for_gateway() + + # Pull back the gateway-side MCP view as GatewayMCPClient instances. + # The gateway loaded these from the config we just wrote, so the set + # matches self._mcps name-for-name. + self._gateway_clients = { + c.name: c for c in await self._gateway.list_mcps() + } + + if self.host_workdir is not None: + await self._save_mcp_file() + await self._seed_skills() + + self.is_alive = True + + async def reset(self) -> None: + """Return the workspace to an empty state. + + Deregisters every MCP from the gateway (``DELETE /mcps/{name}`` + for each), clears the local handles, and wipes ``.mcp``, + ``skills/``, ``sessions/``, and ``data/`` inside the container. + The gateway process keeps running with no upstream MCPs. + ``default_mcps`` / ``skill_paths`` are not re-seeded. + """ + if self._backend is None: + raise RuntimeError( + "DockerWorkspace is not initialized: its container " + "backend is unavailable. Use 'async with workspace:' " + "or call 'await workspace.initialize()' before " + "'reset()'.", + ) + + async with self._mcp_lock, self._skill_lock: + for gw_client in list(self._gateway_clients.values()): + try: + await gw_client.close() + except Exception as e: + logger.warning( + "MCP %r close failed during reset: %s", + gw_client.name, + e, + ) + self._gateway_clients.clear() + self._mcps = [] + + for path in ( + CONTAINER_SESSIONS_DIR, + CONTAINER_DATA_DIR, + CONTAINER_SKILLS_DIR, + ): + await self._backend.delete_path(path) + + # Rewrite ``.mcp`` to an empty list so a future restart does + # not fall back to ``default_mcps`` (which would only happen + # if the file were missing). + if self.host_workdir is not None: + await self._save_mcp_file() + + async def close(self) -> None: + """Stop and remove the container; release the aiodocker client. + + The cached image and the host ``workdir`` are intentionally + left behind — a subsequent :meth:`initialize` reuses both. + Errors during gateway/container teardown are swallowed so + that ``close`` is always safe to call (e.g. from + ``__aexit__``). + """ + if self._gateway is not None: + try: + await self._gateway.aclose() + except Exception: + pass + self._gateway = None + self._gateway_clients.clear() + + if self._container is not None: + # Linux native docker preserves container-side ownership on + # bind-mounted host paths verbatim — files written by the + # in-container root process land on host as root, so a + # non-root host user (CI runner, IDE user) cannot remove + # them. macOS Docker Desktop / Windows Docker remap uids + # transparently, so this only matters on Linux. Best-effort: + # exec failures are swallowed, and we still tear the + # container down. + if ( + self.host_workdir is not None + and sys.platform == "linux" + and self._backend is not None + ): + try: + await self._backend.exec_shell( + [ + "chown", + "-R", + f"{os.getuid()}:{os.getgid()}", + CONTAINER_WORKDIR, + ], + timeout=10.0, + ) + except Exception: + pass + try: + await self._container.kill() + except Exception: + pass + try: + await self._container.delete(force=True) + except Exception: + pass + self._container = None + self._backend = None + + if self._client is not None: + try: + await self._client.close() + except Exception: + pass + self._client = None + + self.is_alive = False + + # ── instructions ──────────────────────────────────────────── + + async def get_instructions(self) -> str: + """Return the system-prompt fragment for this workspace. + + The configured ``instructions`` template is formatted with the + container-side ``{workdir}`` (i.e. ``/workspace``), since the + agent always sees container-internal paths. + """ + return self.instructions.format(workdir=CONTAINER_WORKDIR) + + # ── tool / MCP / skill discovery ──────────────────────────── + + async def list_tools(self) -> list[ToolBase]: + """Built-in tools exposed by the workspace itself. + + Returns the six builtin tools (Bash, Read, Write, Edit, Grep, + Glob), each backed by the workspace's :class:`DockerBackend` + that executes inside the container. + + Raises: + `RuntimeError`: + If the workspace has not been initialized yet (the + container-backed backend is unavailable). Without this + guard the builtin tools would silently fall back to a + :class:`LocalBackend` and run on the host instead of + inside the container. + """ + if self._backend is None: + raise RuntimeError( + "DockerWorkspace is not initialized: its container " + "backend is unavailable. Use 'async with workspace:' " + "or call 'await workspace.initialize()' before " + "'list_tools()'.", + ) + + from ...tool._builtin import Bash, Edit, Glob, Grep, Read, Write + + return [ + Bash(cwd=CONTAINER_WORKDIR, backend=self._backend), + Edit(backend=self._backend), + Glob(backend=self._backend, glob_helper_path=GLOB_HELPER_SCRIPT), + Grep(backend=self._backend), + Read(backend=self._backend), + Write(backend=self._backend), + ] + + async def list_mcps(self) -> list[MCPClient]: + """Return one :class:`GatewayMCPClient` per registered MCP. + + Each entry's ``name`` matches the upstream MCP server name and + all of its protocol calls (connect / close / list_tools / + get_tool / tool ``__call__``) are routed over HTTP to the + in-container gateway. + """ + return list(self._gateway_clients.values()) + + async def list_skills(self) -> list[Skill]: + """Enumerate skills by scanning ``skills/`` inside the container. + + For each ``SKILL.md`` found, parses the YAML front-matter and + yields a :class:`Skill`. Files missing a ``name`` or + ``description`` field are skipped. + + Returns: + Skills available to the agent. Empty when the directory + is missing or contains no parseable ``SKILL.md`` files. + """ + import frontmatter as fm + + result = await self._backend.exec_shell( + [ + "sh", + "-c", + f"find {CONTAINER_SKILLS_DIR} -name SKILL.md " + f"2>/dev/null || true", + ], + ) + if not result.ok(): + return [] + listing = result.stdout.decode(errors="replace").strip() + if not listing: + return [] + + skills: list[Skill] = [] + for md_path in (line.strip() for line in listing.split("\n")): + if not md_path: + continue + try: + raw = await self._backend.read_file(md_path) + doc = fm.loads(raw.decode("utf-8")) + name = doc.get("name") + desc = doc.get("description") + if not name or not desc: + continue + skills.append( + Skill( + name=str(name), + description=str(desc), + dir=posixpath.dirname(md_path), + markdown=doc.content or "", + updated_at=0.0, + ), + ) + except Exception as e: + logger.warning("Failed to load skill %s: %s", md_path, e) + return skills + + # ── dynamic MCP management ────────────────────────────────── + + async def add_mcp(self, mcp_client: MCPClient) -> None: + """Register a new MCP server on the in-container gateway. + + Serialises the supplied client, registers it on the gateway + (which spawns the upstream MCP session inside the container), + and adds the corresponding :class:`GatewayMCPClient` handle to + :meth:`list_mcps`. The change is persisted to ``.mcp`` when + ``workdir`` is set. + + Args: + mcp_client: An :class:`MCPClient` describing the upstream + server (stdio / HTTP / SSE config). Its + ``model_dump()`` is what the gateway consumes, so any + ``MCPClient`` subclass is accepted. + + Raises: + ValueError: If an MCP with the same name is already + registered in this workspace. + RuntimeError: If the gateway rejects the registration + (e.g. upstream command not found inside the + container). + """ + async with self._mcp_lock: + if mcp_client.name in self._gateway_clients: + raise ValueError( + f"MCP {mcp_client.name!r} already exists in workspace.", + ) + spec = mcp_client.model_dump(mode="json") + gw_client = self._gateway.make_client(spec) + await gw_client.connect() + self._mcps.append(mcp_client) + self._gateway_clients[gw_client.name] = gw_client + if self.host_workdir is not None: + await self._save_mcp_file() + + async def remove_mcp(self, name: str) -> None: + """Unregister an MCP server by name. + + Tells the gateway to close the upstream session and drops the + :class:`GatewayMCPClient` handle from :meth:`list_mcps`. The + change is persisted to ``.mcp`` when ``workdir`` is set. + + Args: + name: The ``name`` field of the registered MCP. If no + MCP by that name exists, a warning is logged and the + call is a no-op (matching the silent behaviour of the + local workspace). + """ + async with self._mcp_lock: + gw_client = self._gateway_clients.pop(name, None) + if gw_client is None: + logger.warning("MCP %r not found in workspace", name) + return + try: + await gw_client.close() + except Exception as e: + logger.warning("MCP %r close failed: %s", name, e) + self._mcps = [m for m in self._mcps if m.name != name] + if self.host_workdir is not None: + await self._save_mcp_file() + + # ── dynamic skill management ──────────────────────────────── + + async def add_skill(self, skill_path: str) -> None: + """Copy a local skill directory into ``skills/`` inside the container. + + The directory must contain a ``SKILL.md`` with ``name`` and + ``description`` fields in its YAML front matter (validated + host-side before any container I/O). The directory is + tarred and uploaded via ``put_archive``; a directory of the + same basename already present in the container is rejected + rather than overwritten. + + Args: + skill_path: Absolute or relative path to a skill + directory on the host filesystem. + + Raises: + ValueError: If ``SKILL.md`` is missing, or a directory + with the same basename already exists in the + container's ``skills/``. + """ + skill_md = os.path.join(skill_path, "SKILL.md") + if not os.path.isfile(skill_md): + raise ValueError( + f"Invalid skill at {skill_path!r}: SKILL.md not found", + ) + + async with self._skill_lock: + await self._backend.exec_shell( + ["mkdir", "-p", CONTAINER_SKILLS_DIR], + ) + dir_name = os.path.basename(os.path.abspath(skill_path)) + + # Refuse to overwrite an existing directory of the same name — + # mirrors the conflict-rejection behaviour of ``mkdir`` here + # rather than LocalWorkspace's full hash-dedup index. + check = await self._backend.exec_shell( + ["test", "-e", CONTAINER_SKILLS_DIR + "/" + dir_name], + ) + if check.ok(): + raise ValueError( + f"Skill directory {dir_name!r} already exists in " + f"{CONTAINER_SKILLS_DIR}", + ) + + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w") as tf: + tf.add(skill_path, arcname=dir_name) + await self._container.put_archive( + CONTAINER_SKILLS_DIR, + buf.getvalue(), + ) + logger.info( + "DockerWorkspace: added skill %r at %s/%s", + dir_name, + CONTAINER_SKILLS_DIR, + dir_name, + ) + + async def remove_skill(self, name: str) -> None: + """Delete a skill directory by its agent-facing name. + + Looks up the skill by the ``name`` field of its + ``SKILL.md``, then ``rm -rf`` its directory inside the + container. + + Args: + name: The agent-facing skill name (the ``name`` value in + the SKILL.md front matter, *not* the directory name). + + Raises: + KeyError: If no skill with that ``name`` is found. + RuntimeError: If the in-container ``rm -rf`` returns a + non-zero exit code. + """ + skills = await self.list_skills() + target_dir: str | None = None + for s in skills: + if s.name == name: + target_dir = s.dir + break + if target_dir is None: + available = [s.name for s in skills] + raise KeyError( + f"Skill {name!r} not found. Available: {available}", + ) + await self._backend.delete_path(target_dir) + + # ── offload ───────────────────────────────────────────────── + + async def offload_context( + self, + session_id: str, + msgs: list[Msg], + ) -> str: + """Persist a batch of messages as JSONL inside the container. + + Each message is appended to + ``sessions//context.jsonl``. Inline base64 + :class:`DataBlock` payloads are extracted into the shared + ``data/`` directory and rewritten as ``file://`` URL blocks + before serialisation, keeping the JSONL line size bounded. + + Args: + session_id (`str`): + Session-scope key used to partition offloaded files + (one subdirectory per session). + msgs (`list[Msg]`): + Messages to append. Not mutated — a deep copy is + used internally so the caller's blocks remain + base64-inline. + + Returns: + `str`: + The container-side path of the JSONL file that + received the new lines. + """ + base = f"{CONTAINER_SESSIONS_DIR}/{session_id}" + path = f"{base}/context.jsonl" + + copied = deepcopy(msgs) + lines: list[str] = [] + for msg in copied: + if not isinstance(msg.content, str): + content = [] + for block in msg.content: + if isinstance(block, DataBlock) and isinstance( + block.source, + Base64Source, + ): + block = await self._offload_data_block(block) + content.append(block) + msg.content = content + lines.append(msg.model_dump_json()) + + await self._backend.exec_shell(["mkdir", "-p", base]) + existing = b"" + try: + existing = await self._backend.read_file(path) + except (FileNotFoundError, OSError): + pass + await self._backend.write_file( + path, + existing + ("\n".join(lines) + "\n").encode("utf-8"), + ) + return path + + async def offload_tool_result( + self, + session_id: str, + tool_result: ToolResultBlock, + ) -> str: + """Persist a single tool result as a flat text file. + + Writes ``sessions//tool_result-.txt`` inside + the container. Text blocks are concatenated verbatim; + :class:`DataBlock` items are emitted as + ```` placeholders, + with inline base64 payloads first offloaded to ``data/``. + + Args: + session_id (`str`): + Session-scope key used to partition offloaded files. + tool_result (`ToolResultBlock`): + The tool result block to persist. + + Returns: + `str`: + The container-side path of the offloaded text file. + """ + base = f"{CONTAINER_SESSIONS_DIR}/{session_id}" + path = f"{base}/tool_result-{tool_result.id}.txt" + + parts: list[str] = [] + if isinstance(tool_result.output, str): + parts.append(tool_result.output) + else: + for block in tool_result.output: + if isinstance(block, TextBlock): + parts.append(block.text) + elif isinstance(block, DataBlock): + if isinstance(block.source, Base64Source): + d = await self._offload_data_block(block) + url = str(d.source.url) + else: + url = str(block.source.url) + parts.append( + f"", + ) + + await self._backend.exec_shell(["mkdir", "-p", base]) + await self._backend.write_file( + path, + "".join(parts).encode("utf-8"), + ) + return path + + # ── internals: image build ────────────────────────────────── + + async def _build_or_reuse_image(self) -> None: + """Build the workspace image, or reuse a tag-cache hit. + + The tag is a content hash of the rendered Dockerfile plus + every file copied into the build context. ``self._image_tag`` + is populated unconditionally so that the container-creation + step has a stable reference, even on cache hits. + + Raises: + RuntimeError: If a build error message comes through the + docker stream. + """ + ctx_dir, tag, _ = prepare_build_context( + base_image=self.base_image, + gateway_home=GATEWAY_HOME, + container_workdir=CONTAINER_WORKDIR, + node_version=self.node_version, + extra_pip=self.extra_pip, + ) + self._image_tag = tag + + try: + try: + await self._client.images.inspect(tag) + logger.info("DockerWorkspace: image cache hit %r", tag) + return + except Exception: + pass + + logger.info("DockerWorkspace: building image %r", tag) + # The Docker daemon's POST /build endpoint requires the + # build context as a tar archive in the request body. + # docker-py hides this behind a ``path=`` shortcut that + # tars the directory for you; aiodocker does *not* — we + # have to tar ``ctx_dir`` ourselves and hand it over via + # ``fileobj``. ``arcname="."`` puts every entry at the + # tar root so the daemon finds ``./Dockerfile`` (and the + # ``COPY`` source files) without an extra prefix. + tar_buf = io.BytesIO() + with tarfile.open(fileobj=tar_buf, mode="w") as tf: + tf.add(str(ctx_dir), arcname=".") + tar_buf.seek(0) + # ``encoding="identity"`` tells aiodocker the body is a + # plain (uncompressed) tar — without it, aiodocker would + # gzip our already-tarred bytes and the daemon would + # reject the malformed stream. + stream = self._client.images.build( + fileobj=tar_buf, + encoding="identity", + tag=tag, + stream=True, + rm=True, + ) + # Buffer recent stream lines so that a failing RUN step's + # stderr is included in the RuntimeError below — the + # daemon's ``error`` chunk only carries a one-line summary + # ("command returned non-zero code: 1") and the actual + # diagnostic is in the preceding ``stream`` chunks. + tail: list[str] = [] + tail_max = 200 + async for chunk in stream: + if isinstance(chunk, dict): + if "stream" in chunk: + msg = str(chunk["stream"]).rstrip() + if msg: + logger.debug("[docker build] %s", msg) + tail.append(msg) + if len(tail) > tail_max: + del tail[: len(tail) - tail_max] + if "error" in chunk: + log = "\n".join(tail) + raise RuntimeError( + f"docker build failed: {chunk['error']}\n" + f"--- last {len(tail)} build log lines ---\n" + f"{log}", + ) + finally: + shutil.rmtree(ctx_dir, ignore_errors=True) + + # ── internals: container lifecycle ────────────────────────── + + async def _create_and_start_container(self) -> None: + """Create + start the workspace container. + + Wires up: + + * ``Cmd: ["sleep", "infinity"]`` so the container stays up + even when the gateway is restarted. + * ``ExposedPorts`` / ``PortBindings`` for ``gateway_port`` + (random host port, ``127.0.0.1`` only). + * Optional bind mount ``host workdir → /workspace``. + * ``agentscope.workspace.id`` label for later discovery. + + Resolves the assigned host port into ``self._port_mapping`` + and pre-creates the in-container persistence directories + (``data/`` / ``skills/`` / ``sessions/``). + + Raises: + RuntimeError: If the gateway port did not bind to any + host port (typically a docker daemon issue). + """ + config: dict[str, Any] = { + "Image": self._image_tag, + "Cmd": ["sleep", "infinity"], + "WorkingDir": CONTAINER_WORKDIR, + "Labels": { + "agentscope.workspace": "true", + "agentscope.workspace.id": self.workspace_id, + }, + "ExposedPorts": {f"{self.gateway_port}/tcp": {}}, + } + if self.env: + config["Env"] = [f"{k}={v}" for k, v in self.env.items()] + + host_config: dict[str, Any] = { + "PortBindings": { + f"{self.gateway_port}/tcp": [ + {"HostIp": "127.0.0.1", "HostPort": ""}, + ], + }, + } + if self.host_workdir is not None: + os.makedirs(self.host_workdir, exist_ok=True) + host_config["Binds"] = [ + f"{os.path.abspath(self.host_workdir)}:{CONTAINER_WORKDIR}:rw", + ] + config["HostConfig"] = host_config + + self._container = await self._client.containers.create_or_replace( + name=f"as_ws_{self.workspace_id}", + config=config, + ) + await self._container.start() + + info = await self._container.show() + ports_info = info.get("NetworkSettings", {}).get("Ports") or {} + bindings = ports_info.get(f"{self.gateway_port}/tcp", []) or [] + if not bindings: + raise RuntimeError( + f"gateway port {self.gateway_port} did not bind to a " + "host port", + ) + self._port_mapping[self.gateway_port] = int(bindings[0]["HostPort"]) + + # Create the backend now that the container is running. All + # subsequent container I/O in this workspace goes through it. + self._backend = DockerBackend(self._container, CONTAINER_WORKDIR) + + # Ensure the in-container persistence dirs exist (also makes a + # newly-bind-mounted host workdir agentscope-shaped on first use). + await self._backend.exec_shell( + [ + "mkdir", + "-p", + CONTAINER_DATA_DIR, + CONTAINER_SKILLS_DIR, + CONTAINER_SESSIONS_DIR, + ], + ) + + async def _restore_or_seed_mcps(self) -> list[MCPClient]: + """Decide the MCP set to ship to the gateway on startup. + + * No ``workdir`` → return ``default_mcps`` (purely ephemeral). + * ``workdir`` set, ``/.mcp`` missing → return + ``default_mcps`` and let the next ``_save_mcp_file`` write + it. + * ``/.mcp`` present → :meth:`MCPClient.model_validate` + each entry and return them. A read / parse error is + logged and the call falls back to ``default_mcps`` rather + than crashing the whole workspace. + + Returns: + The MCPClient instances to register on the gateway. + """ + if self.host_workdir is None: + return list(self.default_mcps) + host_mcp = os.path.join(self.host_workdir, ".mcp") + if not os.path.isfile(host_mcp): + return list(self.default_mcps) + try: + with open(host_mcp, encoding="utf-8") as f: + data = json.load(f) + return [MCPClient.model_validate(m) for m in data] + except Exception as e: + logger.warning( + "DockerWorkspace: failed to read %s, falling back to " + "default_mcps: %s", + host_mcp, + e, + ) + return list(self.default_mcps) + + async def _save_mcp_file(self) -> None: + """Persist ``self._mcps`` to ``/.mcp`` (host-side JSON). + + No-op when ``workdir`` is ``None``. Failures are logged but + not raised — losing the persistence file should not + propagate as an MCP-add/remove error to the caller. + """ + if self.host_workdir is None: + return + host_mcp = os.path.join(self.host_workdir, ".mcp") + try: + os.makedirs(self.host_workdir, exist_ok=True) + with open(host_mcp, "w", encoding="utf-8") as f: + json.dump( + [m.model_dump() for m in self._mcps], + f, + indent=2, + ensure_ascii=False, + ) + except Exception as e: + logger.warning( + "DockerWorkspace: failed to save %s: %s", + host_mcp, + e, + ) + + async def _write_gateway_config(self) -> None: + """Drop the gateway's ``--config`` JSON into the container. + + The file at :data:`GATEWAY_CONFIG` carries the freshly minted + bearer token plus the MCP server specs the gateway should + bring up at start. This is the *only* path the bearer token + crosses — it never lands on host disk. + """ + cfg = { + "token": self._gateway_token, + "servers": [m.model_dump(mode="json") for m in self._mcps], + } + await self._backend.exec_shell( + ["mkdir", "-p", GATEWAY_HOME], + ) + await self._backend.write_file( + GATEWAY_CONFIG, + json.dumps(cfg, indent=2, ensure_ascii=False).encode("utf-8"), + ) + + async def _start_gateway_process(self) -> None: + """Launch the gateway inside the container as a detached process. + + Runs ``nohup python --config + --port `` from the baked-in venv. The script is + invoked by path (not via ``python -m``) so Python does not + auto-import ``agentscope.workspace.__init__`` and the heavy + module graph it pulls in. stdout/stderr are redirected to + :data:`GATEWAY_LOG` so :meth:`_wait_for_gateway` can dump the + tail when startup fails. + + We do not block on this exec call; readiness is detected via + the ``/health`` poll instead. + """ + cmd = ( + f"nohup {shlex.quote(GATEWAY_VENV + '/bin/python')} -u " + f"{shlex.quote(GATEWAY_SCRIPT)} " + f"--config {shlex.quote(GATEWAY_CONFIG)} " + f"--port {self.gateway_port} " + f"> {shlex.quote(GATEWAY_LOG)} 2>&1 &" + ) + # Detach: we don't await stream completion, just kick it off. + await self._backend.exec_shell(["sh", "-c", cmd]) + + async def _wait_for_gateway(self, timeout: float = 30.0) -> None: + """Block until the gateway answers ``/health`` with 200. + + Uses an exponentially-backed-off poll capped at 1 s. When + the deadline expires, attempts to read the gateway log and + surfaces the tail in the raised error so callers can see the + actual startup failure. + + Args: + timeout: Maximum seconds to wait for readiness. + + Raises: + RuntimeError: If the gateway does not become healthy + before the deadline. + """ + assert self._gateway is not None + deadline = asyncio.get_event_loop().time() + timeout + delay = 0.1 + while asyncio.get_event_loop().time() < deadline: + if await self._gateway.health(): + return + await asyncio.sleep(delay) + delay = min(delay * 1.5, 1.0) + # Last-ditch: dump the gateway log to help debug startup failures. + try: + log = await self._backend.read_file(GATEWAY_LOG) + tail = log[-2000:].decode(errors="replace") + except Exception: + tail = "" + raise RuntimeError( + f"gateway did not become healthy within {timeout}s. " + f"Tail of {GATEWAY_LOG}:\n{tail}", + ) + + async def _seed_skills(self) -> None: + """Copy ``self.skill_paths`` into ``skills/`` once, on first init. + + Skips seeding when (a) ``workdir`` is unset, (b) ``skill_paths`` + is empty, or (c) the host-side ``skills/`` directory already + contains entries — meaning the user (or a prior init) is the + source of truth and we should not append duplicates. + + Failures on individual paths are logged and skipped rather + than raised, so that one bad skill cannot block startup. + """ + if not self.skill_paths or self.host_workdir is None: + return + skills_host = os.path.join(self.host_workdir, "skills") + # The bind mount lazily materialises ``/skills`` on + # host the first time the container writes into it, so on a + # fresh workdir the host path may not exist yet — create it + # so the next check has something to inspect. + os.makedirs(skills_host, exist_ok=True) + if os.listdir(skills_host): + # already seeded (or user pre-populated) — leave as-is. + return + for path in self.skill_paths: + try: + await self.add_skill(path) + except Exception as e: + logger.warning( + "DockerWorkspace: skip skill %r: %s", + path, + e, + ) + + # ── internals: data offload ──────────────────────────────── + + async def _offload_data_block(self, block: DataBlock) -> DataBlock: + """Persist a base64 :class:`DataBlock` under ``data/``. + + The decoded payload is stored at + ``data/.``, where ```` is + guessed from the block's media type. Hashing the *base64* + text rather than the decoded bytes lets a second offload of + the same block short-circuit (same key, same file). + + Args: + block: An inline-base64 data block. + + Returns: + A new :class:`DataBlock` whose source is a ``file://`` + URL pointing at the persisted file inside the container. + Blocks already backed by a :class:`URLSource` are returned + unchanged (nothing to persist). + """ + if not isinstance(block.source, Base64Source): + return block + h = hashlib.sha256(block.source.data.encode()).hexdigest() + ext = mimetypes.guess_extension(block.source.media_type) or ".bin" + path = f"{CONTAINER_DATA_DIR}/{h}{ext}" + await self._backend.exec_shell( + ["mkdir", "-p", CONTAINER_DATA_DIR], + ) + await self._backend.write_file( + path, + base64.b64decode(block.source.data), + ) + return DataBlock( + id=block.id, + name=block.name, + source=URLSource( + url=AnyUrl(f"file://{path}"), + media_type=block.source.media_type, + ), + ) diff --git a/src/agentscope/workspace/_docker/_make_dockerfile.py b/src/agentscope/workspace/_docker/_make_dockerfile.py new file mode 100644 index 0000000000000000000000000000000000000000..e568c4c0af2bfc4d5613d966040024b8543e374d --- /dev/null +++ b/src/agentscope/workspace/_docker/_make_dockerfile.py @@ -0,0 +1,280 @@ +# -*- coding: utf-8 -*- +"""Dockerfile generation + build-context preparation for DockerWorkspace. + +The image is keyed by content hash of the Dockerfile text plus all +files COPYed into it. If the tag already exists locally the build is +skipped; otherwise the caller builds with the prepared context. + +Two install modes are supported, picked automatically from the host's +``agentscope`` install: + +* **released** — when ``agentscope`` is in site-packages, the container + installs the same version from PyPI. +* **dev** — when running from a source checkout, the project tree is + copied into the build context and installed via ``uv pip install + /tmp/agentscope_src``. This is a transitional path used until + agentscope is published; it is marked with TODOs in the templates so + it is easy to delete later. + +Public functions: + +* :func:`render_dockerfile` — substitute placeholders in + ``Dockerfile.template`` and return the rendered text. +* :func:`compute_image_tag` — sha256 of Dockerfile + COPY files; + returns ``agentscope-workspace:<12hex>``. +* :func:`prepare_build_context` — assemble a temp directory holding + the Dockerfile, ``requirements.txt``, and (in dev mode) the + agentscope source tree. Returns ``(ctx_dir, tag, copy_files)``. +""" + +import hashlib +import importlib.resources as _res +import shutil +import tempfile +from pathlib import Path + +from .._utils import ( + _GATEWAY_BASE_REQUIREMENTS, + _agentscope_source_root, + _agentscope_version, + _is_released_install, + _is_source_ignored, + _read_gateway_script_bytes, + _read_glob_helper_bytes, +) + +# ── shared constants (also imported by _docker_workspace) ────────── + +DEFAULT_BASE_IMAGE = "python:3.11-slim" +DEFAULT_GATEWAY_PORT = 5600 + +CONTAINER_WORKDIR = "/workspace" +CONTAINER_DATA_DIR = f"{CONTAINER_WORKDIR}/data" +CONTAINER_SKILLS_DIR = f"{CONTAINER_WORKDIR}/skills" +CONTAINER_SESSIONS_DIR = f"{CONTAINER_WORKDIR}/sessions" +CONTAINER_MCP_FILE = f"{CONTAINER_WORKDIR}/.mcp" + +GATEWAY_HOME = "/root/.agentscope" +GATEWAY_VENV = f"{GATEWAY_HOME}/.venv" +GATEWAY_CONFIG = f"{GATEWAY_HOME}/gateway.config.json" +GATEWAY_LOG = f"{GATEWAY_HOME}/gateway.log" +# Standalone gateway script copied into the image. We invoke this +# directly rather than via ``python -m agentscope.workspace._mcp_gateway`` +# so Python does not auto-import ``agentscope.workspace.__init__`` (which +# pulls in skill/tool/local_workspace/docker_workspace and their heavy +# transitive dependencies). +GATEWAY_SCRIPT = f"{GATEWAY_HOME}/_mcp_gateway_app.py" +# Standalone glob helper script used by the Glob builtin tool. +GLOB_HELPER_SCRIPT = f"{GATEWAY_HOME}/_glob_helper.py" + +IMAGE_REPO = "agentscope-workspace" + +# ── template loading ─────────────────────────────────────────────── + +_TEMPLATE_PKG = "agentscope.workspace._docker" +_DOCKERFILE_TEMPLATE = "Dockerfile.template" +_DOCKERFILE_NODE_FROM_TEMPLATE = "Dockerfile.node_from.template" +_DOCKERFILE_NODE_COPY_TEMPLATE = "Dockerfile.node_copy.template" +_DOCKERFILE_INSTALL_PYPI_TEMPLATE = "Dockerfile.install_pypi.template" +_DOCKERFILE_INSTALL_SRC_TEMPLATE = "Dockerfile.install_src.template" + + +def _read_template(name: str) -> str: + """Read a packaged template file as text.""" + return _res.files(_TEMPLATE_PKG).joinpath(name).read_text(encoding="utf-8") + + +# ── source-tree packaging (dev mode) ─────────────────────────────── + + +def _source_ignore(_dir: str, names: list[str]) -> list[str]: + """``shutil.copytree`` ignore filter — defers + to :func:`_is_source_ignored`.""" + return [n for n in names if _is_source_ignored(n)] + + +def _hash_directory(root: Path) -> bytes: + """Hash a directory tree's contents in a stable order.""" + h = hashlib.sha256() + for p in sorted(root.rglob("*")): + if not p.is_file(): + continue + rel = p.relative_to(root).as_posix().encode("utf-8") + h.update(b"\x00") + h.update(rel) + h.update(b"\x00") + h.update(p.read_bytes()) + return h.digest() + + +# ── public API ───────────────────────────────────────────────────── + + +def render_dockerfile( + *, + base_image: str = DEFAULT_BASE_IMAGE, + gateway_home: str = GATEWAY_HOME, + container_workdir: str = CONTAINER_WORKDIR, + node_version: str | None = None, + install_agentscope_block: str = "", +) -> str: + """Render the Dockerfile by substituting into the template files. + + Args: + base_image: Base image (must already provide ``python3``). + gateway_home: In-container directory for the gateway venv, + script and config. + container_workdir: Container-side workdir; bind-mounted from + the host when the workspace's ``workdir`` is set, else + an empty in-image directory. + node_version: When given (e.g. ``"20"``) a ``node`` and ``npm`` + of that version are copied from the official Node slim + image. ``None`` skips Node installation. + install_agentscope_block: Pre-rendered block (no surrounding + blank lines) that installs ``agentscope`` into the gateway + venv. Built by :func:`prepare_build_context` from one of + ``Dockerfile.install_{pypi,src}.template``. + + Returns: + The full Dockerfile text. + """ + if node_version: + # Normalise trailing whitespace so the main template's surrounding + # newlines fully control inter-section spacing — the template files + # themselves are not relied on for exact terminal newlines. + nf_raw = _read_template(_DOCKERFILE_NODE_FROM_TEMPLATE).format( + node_version=node_version, + ) + nc_raw = _read_template(_DOCKERFILE_NODE_COPY_TEMPLATE) + node_from_block = nf_raw.rstrip() + "\n" + node_copy_block = nc_raw.rstrip() + "\n\n" + else: + node_from_block = "" + node_copy_block = "" + + return _read_template(_DOCKERFILE_TEMPLATE).format( + base_image=base_image, + gateway_home=gateway_home, + container_workdir=container_workdir, + node_from_block=node_from_block, + node_copy_block=node_copy_block, + install_agentscope_block=install_agentscope_block.rstrip() + "\n", + ) + + +def _render_requirements(extra_pip: list[str]) -> str: + """Render ``requirements.txt`` content for the gateway venv.""" + pinned = list(_GATEWAY_BASE_REQUIREMENTS) + list(extra_pip or []) + return "\n".join(pinned) + "\n" + + +def compute_image_tag( + dockerfile_text: str, + copy_files: dict[str, bytes], +) -> str: + """Hash the Dockerfile and COPY payloads into a deterministic tag. + + Args: + dockerfile_text: Full Dockerfile text. + copy_files: Mapping of context-relative filename → bytes for + every file referenced by a ``COPY`` instruction. Directory + payloads (e.g. the agentscope source tree in dev mode) are + represented as a single synthetic key whose value is the + tree's content hash. + + Returns: + Tag of the form ``agentscope-workspace:<12 hex chars>``. + """ + h = hashlib.sha256() + h.update(b"DOCKERFILE\x00") + h.update(dockerfile_text.encode("utf-8")) + for name in sorted(copy_files): + h.update(b"\x00FILE\x00") + h.update(name.encode("utf-8")) + h.update(b"\x00") + h.update(copy_files[name]) + return f"{IMAGE_REPO}:{h.hexdigest()[:12]}" + + +def prepare_build_context( + *, + base_image: str = DEFAULT_BASE_IMAGE, + gateway_home: str = GATEWAY_HOME, + container_workdir: str = CONTAINER_WORKDIR, + node_version: str | None = None, + extra_pip: list[str] | None = None, +) -> tuple[Path, str, dict[str, bytes]]: + """Assemble a temporary build context directory. + + Writes Dockerfile, ``requirements.txt`` and the agentscope payload + (source tree in dev mode, nothing extra in released mode) into a + fresh temp dir. The caller is responsible for removing the + directory after the build completes. + + Returns: + ``(ctx_dir, tag, copy_files)`` — ``ctx_dir`` holds the + materialised files; ``tag`` is the deterministic image tag; + ``copy_files`` is the same mapping that was hashed into the + tag (handy for callers that want to recompute / verify). + """ + extra_pip_list = list(extra_pip or []) + + released = _is_released_install() + if released: + version = _agentscope_version() + install_block = _read_template( + _DOCKERFILE_INSTALL_PYPI_TEMPLATE, + ).format(agentscope_version=version) + source_root: Path | None = None + else: + # TODO(release): drop this branch once agentscope is on PyPI; the + # released path above subsumes it. The dev branch copies the project + # tree into the build context so the in-container venv can install + # the same code the host is running. + install_block = _read_template(_DOCKERFILE_INSTALL_SRC_TEMPLATE) + source_root = _agentscope_source_root() + + dockerfile_text = render_dockerfile( + base_image=base_image, + gateway_home=gateway_home, + container_workdir=container_workdir, + node_version=node_version, + install_agentscope_block=install_block, + ) + requirements_text = _render_requirements(extra_pip_list) + + # Read helper scripts once — we both hash them into the image tag + # (so edits invalidate the image cache) and write them into the + # build context so the Dockerfile can ``COPY`` them. + gateway_script_bytes = _read_gateway_script_bytes() + glob_helper_bytes = _read_glob_helper_bytes() + + copy_files: dict[str, bytes] = { + "requirements.txt": requirements_text.encode("utf-8"), + "_mcp_gateway_app.py": gateway_script_bytes, + "_glob_helper.py": glob_helper_bytes, + } + if source_root is not None: + # Synthetic entry: the directory tree is too large to inline, so we + # hash it once and stash the digest under a stable key. Image-tag + # determinism depends only on this digest, not on the temp path. + copy_files["agentscope_src/"] = _hash_directory(source_root) + + tag = compute_image_tag(dockerfile_text, copy_files) + + ctx_dir = Path(tempfile.mkdtemp(prefix="as-ws-build-")) + (ctx_dir / "Dockerfile").write_text(dockerfile_text, encoding="utf-8") + (ctx_dir / "requirements.txt").write_bytes( + copy_files["requirements.txt"], + ) + (ctx_dir / "_mcp_gateway_app.py").write_bytes(gateway_script_bytes) + (ctx_dir / "_glob_helper.py").write_bytes(glob_helper_bytes) + if source_root is not None: + shutil.copytree( + source_root, + ctx_dir / "agentscope_src", + ignore=_source_ignore, + symlinks=False, + ) + + return ctx_dir, tag, copy_files diff --git a/src/agentscope/workspace/_e2b/__init__.py b/src/agentscope/workspace/_e2b/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..393b979e7f11c0c0afdbef13f71556466c25af11 --- /dev/null +++ b/src/agentscope/workspace/_e2b/__init__.py @@ -0,0 +1,12 @@ +# -*- coding: utf-8 -*- +"""E2B-backed workspace package. + +Re-exports :class:`E2BWorkspace` so callers can write +``from agentscope.workspace._e2b import E2BWorkspace`` without having to +poke at the underlying module layout. +""" + +from ._e2b_workspace import E2BWorkspace +from ._e2b_backend import E2BBackend + +__all__ = ["E2BWorkspace", "E2BBackend"] diff --git a/src/agentscope/workspace/_e2b/_bootstrap.py b/src/agentscope/workspace/_e2b/_bootstrap.py new file mode 100644 index 0000000000000000000000000000000000000000..23aa5c0d66598dcf68afc49c7bfe424aba6d40cd --- /dev/null +++ b/src/agentscope/workspace/_e2b/_bootstrap.py @@ -0,0 +1,204 @@ +# -*- coding: utf-8 -*- +"""Bootstrap helpers for :class:`E2BWorkspace` first-time provisioning. + +Unlike :class:`DockerWorkspace`, E2B has no image-build phase: it +attaches to a pre-built template (``base`` by default). The first +time we create a sandbox for a given ``workspace_id``, we run a +sequence of shell commands to install ``uv``, the gateway venv, the +gateway base requirements, and the ``agentscope`` package. After this +the gateway script is uploaded into the sandbox. + +These artefacts persist across ``sandbox.pause()`` / ``resume`` so +the bootstrap cost is paid exactly once per sandbox lifetime. + +Two install modes mirror Docker: + +* **released** — ``agentscope`` is in site-packages on the host; + install the same version from PyPI inside the sandbox. +* **dev** — running from a source checkout; tar the project tree on + the host, upload it as a single blob, untar inside the sandbox, and + ``uv pip install --no-deps`` it. + +``--no-deps`` is mandatory: the gateway only imports +``agentscope.mcp.MCPClient`` whose transitive needs are +``mcp / pydantic / httpx`` (already installed via the gateway base +requirements). Pulling agentscope's full dep tree drags in heavy / +Rust-built packages (ripgrep, tree_sitter, opentelemetry, openai, +anthropic, dashscope) for which the sandbox base image typically has +no compiler, exactly the same trap we hit on Docker. +""" + +import io +import tarfile + +from ..._logging import logger +from .._utils import ( + _GATEWAY_BASE_REQUIREMENTS, + _agentscope_source_root, + _is_source_ignored, +) + +# ── shared constants ─────────────────────────────────────────────── + +#: Default E2B template. Matches the SDK's ``base`` template — has +#: Ubuntu + python3 + curl out of the box, which is everything the +#: bootstrap needs. +DEFAULT_TEMPLATE = "base" + +#: Default keep-alive timeout in seconds for newly-created sandboxes. +DEFAULT_TIMEOUT = 300 + +#: Default port the in-sandbox gateway listens on. +DEFAULT_GATEWAY_PORT = 5600 + +#: Sandbox-side runtime user (E2B base image runs as ``user``, not +#: ``root``). All the per-workspace paths sit under its ``$HOME``. +SANDBOX_USER_HOME = "/home/user" + +# Workspace-side persistent layout — mirrors the DockerWorkspace one. +SANDBOX_WORKDIR = f"{SANDBOX_USER_HOME}/workspace" +SANDBOX_DATA_DIR = f"{SANDBOX_WORKDIR}/data" +SANDBOX_SKILLS_DIR = f"{SANDBOX_WORKDIR}/skills" +SANDBOX_SESSIONS_DIR = f"{SANDBOX_WORKDIR}/sessions" +SANDBOX_MCP_FILE = f"{SANDBOX_WORKDIR}/.mcp" + +# Gateway home — venv, script, config, logs. +GATEWAY_HOME = f"{SANDBOX_USER_HOME}/.agentscope" +GATEWAY_VENV = f"{GATEWAY_HOME}/.venv" +GATEWAY_VENV_PY = f"{GATEWAY_VENV}/bin/python" +GATEWAY_SCRIPT = f"{GATEWAY_HOME}/_mcp_gateway_app.py" +# Standalone glob helper script used by the builtin Glob tool. +GLOB_HELPER_SCRIPT = f"{GATEWAY_HOME}/_glob_helper.py" +GATEWAY_CONFIG = f"{GATEWAY_HOME}/gateway.config.json" +GATEWAY_LOG = f"{GATEWAY_HOME}/gateway.log" + +# uv binary lives under the user's local bin since we cannot write to +# /usr/local/bin without sudo on the default E2B image. +UV_BIN = f"{SANDBOX_USER_HOME}/.local/bin/uv" + +#: Sandbox metadata key used to map workspace_id → sandbox_id. The +#: manager filters ``list_sandboxes`` by this key on cache miss to +#: locate (and resume) an existing sandbox. +METADATA_WORKSPACE_ID_KEY = "agentscope.workspace.id" + +#: Tarball drop point for dev-mode ``agentscope`` source uploads. +#: Lives under ``GATEWAY_HOME`` (user-owned) instead of ``/tmp`` +#: because the E2B ``files.write`` API runs as ``user`` and ``/tmp`` +#: can have restrictive permissions after a sandbox pause/resume cycle. +DEV_SRC_TAR = f"{GATEWAY_HOME}/agentscope_src.tar" +DEV_SRC_DIR = f"{GATEWAY_HOME}/agentscope_src" + +# ── source tarball (dev mode only) ───────────────────────────────── + + +def _tar_filter(info: tarfile.TarInfo) -> tarfile.TarInfo | None: + """``tarfile.add`` filter — skip caches, hidden files and heavy dirs. + + Returning ``None`` for an entry tells :meth:`tarfile.add` to also + *stop recursing* into it, so the root archive entry (``arcname="."``) + must be passed through unchanged — otherwise the entire tree is + excluded and the resulting tarball is effectively empty. + """ + if info.name == ".": + return info + name = info.name.split("/")[-1] + if _is_source_ignored(name): + return None + return info + + +def build_source_tarball() -> bytes: + """Tar up the agentscope source tree for dev-mode upload. + + Returns the tar bytes (uncompressed); the sandbox will untar with + ``tar -xf`` (no ``-z``). + """ + root = _agentscope_source_root() + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w") as tf: + tf.add(str(root), arcname=".", filter=_tar_filter) + return buf.getvalue() + + +# ── bootstrap command sequence ───────────────────────────────────── + + +def bootstrap_commands( + *, + extra_pip: list[str] | None = None, + install_agentscope_cmd: str, +) -> list[str]: + """Return the shell commands to provision a fresh E2B sandbox. + + Args: + extra_pip: Extra Python packages to install into the gateway + venv alongside the base requirements. + install_agentscope_cmd: Shell command that installs + ``agentscope`` into the gateway venv. Built per install + mode by :func:`render_install_agentscope_cmd` (released) + or :func:`render_install_agentscope_cmd_dev` (dev). + + Returns: + A list of shell command strings, to be executed in order. + Each must exit 0; a non-zero exit aborts the bootstrap. + """ + pip_pkgs = list(_GATEWAY_BASE_REQUIREMENTS) + list(extra_pip or []) + pip_args = " ".join(pip_pkgs) + + return [ + # 1. Persistent layout. mkdir -p is cheap on resume too — keeps + # the command idempotent in case bootstrap is re-run. + f"mkdir -p {SANDBOX_DATA_DIR} {SANDBOX_SKILLS_DIR} " + f"{SANDBOX_SESSIONS_DIR} {GATEWAY_HOME}", + # 2. Install ripgrep for the Grep builtin tool. The base E2B + # image runs as non-root, so we use sudo. ``apt-get update`` + # + install is idempotent — safe to re-run on resume. + "sudo apt-get update -qq " + "&& sudo apt-get install -y --no-install-recommends ripgrep " + "&& sudo rm -rf /var/lib/apt/lists/*", + # 3. Astral uv — same shell installer as Docker. The base E2B + # image already ships ``curl``; we land uv at + # ``$HOME/.local/bin`` since the sandbox user has no sudo by + # default. ``INSTALLER_NO_MODIFY_PATH=1`` suppresses shell + # rc edits — we always invoke uv by full path here anyway. + f"curl -LsSf https://astral.sh/uv/install.sh " + f"| env UV_INSTALL_DIR={SANDBOX_USER_HOME}/.local/bin " + f"INSTALLER_NO_MODIFY_PATH=1 sh", + # 4. Gateway venv + base requirements. + f"{UV_BIN} venv {GATEWAY_VENV}", + f"{UV_BIN} pip install --python {GATEWAY_VENV_PY} {pip_args}", + # 5. agentscope itself (mode-dependent). + install_agentscope_cmd, + ] + + +def render_install_agentscope_cmd_released(version: str) -> str: + """Released-mode install: pin the same version as the host.""" + return ( + f"{UV_BIN} pip install --python {GATEWAY_VENV_PY} " + f"--no-deps 'agentscope=={version}'" + ) + + +def render_install_agentscope_cmd_dev() -> str: + """Dev-mode install: untar the uploaded tarball and ``pip install``. + + Caller is responsible for uploading the source tarball to + :data:`DEV_SRC_TAR` before running this command. + """ + return ( + f"mkdir -p {DEV_SRC_DIR} && " + f"tar -xf {DEV_SRC_TAR} -C {DEV_SRC_DIR} && " + f"{UV_BIN} pip install --python {GATEWAY_VENV_PY} " + f"--no-deps {DEV_SRC_DIR} && " + f"rm -rf {DEV_SRC_TAR} {DEV_SRC_DIR}" + ) + + +def log_bootstrap_attempt(workspace_id: str, mode: str) -> None: + """Single info-level log so first-time bootstraps are easy to spot.""" + logger.info( + "E2BWorkspace: bootstrapping sandbox for workspace_id=%r (mode=%s)", + workspace_id, + mode, + ) diff --git a/src/agentscope/workspace/_e2b/_e2b_backend.py b/src/agentscope/workspace/_e2b/_e2b_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..e2e8d0e328eada92477f33229353bf35dc86760b --- /dev/null +++ b/src/agentscope/workspace/_e2b/_e2b_backend.py @@ -0,0 +1,163 @@ +# -*- coding: utf-8 -*- +"""E2B sandbox :class:`BackendBase` implementation. + +Wraps the E2B SDK's ``commands.run`` and ``files.*`` APIs into the +three backend primitives (``exec_shell``, ``read_file``, +``write_file``) so that builtin tools (Bash, Read, Write, Edit, Grep, +Glob) can operate inside an E2B cloud sandbox transparently. All +derived filesystem helpers (``file_exists``, ``is_dir``, ``list_dir``, +``stat_mtime``, ``delete_path``) are inherited from +:class:`BackendBase`, which implements them via ``exec_shell``. +""" + +from __future__ import annotations + +import posixpath +import shlex +from typing import Any + +from ...tool import BackendBase, ExecResult + + +class E2BBackend(BackendBase): + """Backend that delegates to a running E2B sandbox. + + Only the three abstract primitives (``exec_shell``, ``read_file``, + ``write_file``) are implemented here; the derived filesystem helpers + are inherited from :class:`BackendBase`. + + Args: + sandbox (`Any`): + An ``e2b.AsyncSandbox`` object (must already be started / + connected). + workdir (`str`): + Default working directory for ``exec_shell`` calls inside + the sandbox. + """ + + def __init__(self, sandbox: Any, workdir: str) -> None: + """Initialize the E2B backend. + + Args: + sandbox (`Any`): + A started / connected ``e2b.AsyncSandbox`` object. + workdir (`str`): + Default working directory for ``exec_shell`` calls + inside the sandbox. + """ + self._sandbox = sandbox + self._workdir = workdir + + # ── exec ─────────────────────────────────────────────────────── + + async def getcwd(self) -> str: + """Return the sandbox's default working directory. + + Overrides the base class default (which would shell out to + ``pwd``) with the cached ``workdir`` supplied at construction, + avoiding a per-call sandbox round-trip. + + Returns: + `str`: + The sandbox's default working directory. + """ + return self._workdir + + async def exec_shell( + self, + command: list[str], + *, + cwd: str | None = None, + timeout: float | None = None, + ) -> ExecResult: + """Run a program inside the sandbox via ``commands.run``. + + *command* is an argv list. The E2B ``commands.run`` API takes a + single shell command line, so the argv is POSIX-quoted back into + a string before dispatch (the sandbox is always Linux). Callers + needing shell features pass ``["sh", "-c", line]``. + + Args: + command (`list[str]`): + Executable path/name followed by its arguments. + cwd (`str | None`, optional): + Working directory inside the sandbox. When ``None`` the + backend's default ``workdir`` is used. + timeout (`float | None`, optional): + Maximum number of seconds to wait. When ``None`` the + SDK default applies. + + Returns: + `ExecResult`: + The captured exit code, stdout, and stderr. A non-zero + command exit is reported as a normal result; transport + errors yield an ``exit_code`` of ``-1``. + """ + from e2b import CommandExitException + + command_line = " ".join(shlex.quote(arg) for arg in command) + kwargs: dict[str, Any] = {"cwd": cwd or self._workdir} + if timeout is not None: + kwargs["timeout"] = timeout + try: + res = await self._sandbox.commands.run(command_line, **kwargs) + return ExecResult( + exit_code=int(res.exit_code or 0), + stdout=(res.stdout or "").encode("utf-8"), + stderr=(res.stderr or "").encode("utf-8"), + ) + except CommandExitException as e: + return ExecResult( + exit_code=int(e.exit_code or 1), + stdout=(e.stdout or "").encode("utf-8"), + stderr=(e.stderr or "").encode("utf-8"), + ) + except Exception as e: # noqa: BLE001 + return ExecResult( + exit_code=-1, + stdout=b"", + stderr=str(e).encode("utf-8"), + ) + + # ── file I/O ─────────────────────────────────────────────────── + + async def read_file(self, path: str) -> bytes: + """Read a file from the sandbox via ``files.read``. + + Args: + path (`str`): + Path to the file inside the sandbox. + + Returns: + `bytes`: + The raw file contents. + + Raises: + `FileNotFoundError`: + If the path does not exist inside the sandbox. + """ + from e2b import FileNotFoundException + + try: + data = await self._sandbox.files.read(path, format="bytes") + except FileNotFoundException as exc: + raise FileNotFoundError( + f"not found in sandbox: {path}", + ) from exc + return bytes(data) + + async def write_file(self, path: str, data: bytes) -> None: + """Write *data* to a file inside the sandbox. + + Creates parent directories via ``exec_shell`` first. + + Args: + path (`str`): + Destination path inside the sandbox. + data (`bytes`): + The raw bytes to write. + """ + parent = posixpath.dirname(path) + if parent: + await self.exec_shell(["mkdir", "-p", parent]) + await self._sandbox.files.write(path, data) diff --git a/src/agentscope/workspace/_e2b/_e2b_workspace.py b/src/agentscope/workspace/_e2b/_e2b_workspace.py new file mode 100644 index 0000000000000000000000000000000000000000..769c9ef5484402e7eddc44a79729adc7ec8c6c91 --- /dev/null +++ b/src/agentscope/workspace/_e2b/_e2b_workspace.py @@ -0,0 +1,1012 @@ +# -*- coding: utf-8 -*- +"""E2BWorkspace — sandboxed workspace backed by an E2B cloud sandbox. + +Architecture +------------ + +Mirrors :class:`agentscope.workspace.DockerWorkspace` but swaps the +Docker engine for the E2B SDK (``e2b.AsyncSandbox``): + +* **Lifecycle.** ``initialize()`` looks up an existing sandbox by + metadata and either resumes it (``connect(sandbox_id=...)`` + auto-resumes paused sandboxes) or creates a fresh one and runs the + bootstrap shell sequence. ``close()`` calls ``sandbox.pause()`` so + the sandbox filesystem (skills, ``.mcp``, sessions, data) survives + for the next ``initialize()``. There is no ``kill()`` path in this + iteration. +* **Persistence.** Sandbox filesystem state is the persistence layer — + there is no host-side ``workdir`` parameter. Pausing keeps the disk; + resuming brings it back wholesale. +* **Bootstrap.** First-time provisioning installs uv + a gateway venv + + agentscope (``--no-deps``) and uploads the gateway script. We + detect whether bootstrap has already happened via a single + ``files.exists(GATEWAY_SCRIPT)`` probe so the cost is paid exactly + once per sandbox lifetime. +* **MCP gateway.** Identical to Docker: a FastAPI process inside the + sandbox, host-side talks to it over HTTPS via E2B's proxy + (``sandbox.get_host(port)`` + ``X-Access-Token`` header). +* **Service-layer index.** The host stores only ``workspace_id``; + the sandbox carries ``METADATA_WORKSPACE_ID_KEY = workspace_id`` in + its E2B metadata. Manager code calls ``AsyncSandbox.list(query=...)`` + with that filter to find the sandbox on cache miss. + +Configuration is per-instance: every workspace owns one sandbox. The +manager handles cache, TTL eviction and metadata-based reattachment. +""" + +import asyncio +import base64 +import hashlib +import json +import mimetypes +import os +import posixpath +import shlex +import uuid +from copy import deepcopy +from typing import Any + +from pydantic import AnyUrl + +from ..._logging import logger +from ...mcp import MCPClient +from ...message import ( + Base64Source, + DataBlock, + Msg, + TextBlock, + ToolResultBlock, + URLSource, +) +from ...skill import Skill +from ...tool import ToolBase +from .._base import WorkspaceBase +from .._gateway_client import ( + GatewayClient, + GatewayMCPClient, +) +from .._utils import ( + _agentscope_version, + _is_released_install, + _read_gateway_script_bytes, + _read_glob_helper_bytes, +) +from ._bootstrap import ( + DEFAULT_GATEWAY_PORT, + DEFAULT_TEMPLATE, + DEFAULT_TIMEOUT, + DEV_SRC_TAR, + GATEWAY_CONFIG, + GATEWAY_HOME, + GATEWAY_LOG, + GATEWAY_SCRIPT, + GATEWAY_VENV_PY, + GLOB_HELPER_SCRIPT, + METADATA_WORKSPACE_ID_KEY, + SANDBOX_DATA_DIR, + SANDBOX_MCP_FILE, + SANDBOX_SESSIONS_DIR, + SANDBOX_SKILLS_DIR, + SANDBOX_WORKDIR, + bootstrap_commands, + build_source_tarball, + log_bootstrap_attempt, + render_install_agentscope_cmd_dev, + render_install_agentscope_cmd_released, +) +from ._e2b_backend import E2BBackend + +_DEFAULT_INSTRUCTIONS = """ +You have an E2B-based cloud workspace. All tool calls execute **inside +the sandbox** at ``{workdir}``. + +Layout: + +``` +{workdir} +├── data/ # offloaded multimodal files +├── skills/ # reusable skills +└── sessions/ # session context and tool results +``` + +Use the MCP-provided tools to interact with the sandbox's filesystem +and processes. +""" + + +# ── the workspace ────────────────────────────────────────────────── + + +class E2BWorkspace(WorkspaceBase): + """Workspace backed by an E2B cloud sandbox. + + ``default_mcps`` and ``skill_paths`` are seed-time inputs and are + not retained as instance state past :meth:`initialize`. + """ + + def __init__( + self, + *, + workspace_id: str | None = None, + template: str = DEFAULT_TEMPLATE, + api_key: str = "", + domain: str = "", + timeout_seconds: int = DEFAULT_TIMEOUT, + gateway_port: int = DEFAULT_GATEWAY_PORT, + env: dict[str, str] | None = None, + sandbox_metadata: dict[str, str] | None = None, + extra_pip: list[str] | None = None, + instructions: str = _DEFAULT_INSTRUCTIONS, + default_mcps: list[MCPClient] | None = None, + skill_paths: list[str] | None = None, + ) -> None: + """Construct an :class:`E2BWorkspace`. + + The sandbox is *not* started here; call :meth:`initialize` + (or use the workspace as an ``async`` context manager). + + Args: + workspace_id (`str | None`, optional): + Stable identifier; doubles as the value stored in the + sandbox's ``agentscope.workspace.id`` metadata for + later reattachment. ``None`` generates a fresh UUID. + template (`str`, defaults to `DEFAULT_TEMPLATE`): + E2B template id. Defaults to ``"base"`` — the stock + Ubuntu image with python3 + curl, which is enough for + the bootstrap to install uv on top. + api_key (`str`, defaults to `""`): + E2B API key. ``""`` falls back to the ``E2B_API_KEY`` + env var. + domain (`str`, defaults to `""`): + Optional custom E2B domain (self-hosted etc.). + timeout_seconds (`int`, defaults to `DEFAULT_TIMEOUT`): + Sandbox keep-alive timeout passed to ``create`` / + ``connect``. + gateway_port (`int`, defaults to `DEFAULT_GATEWAY_PORT`): + TCP port the in-sandbox gateway listens on. + env (`dict[str, str] | None`, optional): + Environment variables baked into the sandbox at + create time (``envs`` parameter on the SDK side). + sandbox_metadata (`dict[str, str] | None`, optional): + Extra metadata merged with + ``{METADATA_WORKSPACE_ID_KEY: workspace_id}``. Useful + for attaching ``user_id`` / ``agent_id`` for E2B + dashboard filtering. + extra_pip (`list[str] | None`, optional): + Extra Python packages to install into the gateway + venv during bootstrap. + instructions (`str`, defaults to `_DEFAULT_INSTRUCTIONS`): + System-prompt fragment template returned by + :meth:`get_instructions`. + default_mcps (`list[MCPClient] | None`, optional): + Initial MCPs registered on first :meth:`initialize`. + Subsequent restarts read ``$workdir/.mcp`` instead. + skill_paths (`list[str] | None`, optional): + Local skill directories seeded into + ``$workdir/skills`` on first :meth:`initialize`. + """ + super().__init__(workspace_id=workspace_id) + + # ── serializable config ───────────────────────────────── + self.workdir = SANDBOX_WORKDIR + self.template = template + self.api_key = api_key + self.domain = domain + self.timeout_seconds = timeout_seconds + self.gateway_port = gateway_port + self.env: dict[str, str] = dict(env or {}) + self.sandbox_metadata: dict[str, str] = dict(sandbox_metadata or {}) + self.extra_pip: list[str] = list(extra_pip or []) + self.instructions = instructions + + # ── seed-only ─────────────────────────────────────────── + self.default_mcps: list[MCPClient] = list(default_mcps or []) + self.skill_paths: list[str] = list(skill_paths or []) + + # ── runtime state ─────────────────────────────────────── + self._sandbox: Any = None # e2b.AsyncSandbox + self._backend: E2BBackend | None = None + self._gateway: GatewayClient | None = None + self._gateway_token: str = "" + self._mcps: list[MCPClient] = [] + self._gateway_clients: dict[str, GatewayMCPClient] = {} + self._mcp_lock = asyncio.Lock() + self._skill_lock = asyncio.Lock() + + # ── lifecycle ─────────────────────────────────────────────── + + @property + def sandbox_id(self) -> str | None: + """E2B sandbox id, or ``None`` if not started.""" + return self._sandbox.sandbox_id if self._sandbox else None + + async def initialize(self) -> None: + """Reattach or create the sandbox, then start the gateway. + + Steps: + + 1. Look up an existing sandbox via + ``AsyncSandbox.list(query=SandboxQuery(metadata=...))``. + If found, ``AsyncSandbox.connect(sandbox_id=...)`` reattaches + — it auto-resumes paused sandboxes. + 2. If not found, ``AsyncSandbox.create(...)`` provisions a + fresh sandbox tagged with our metadata and runs bootstrap + (uv → gateway venv → agentscope ``--no-deps`` → upload + gateway script). + 3. If bootstrap output is missing on a reattached sandbox + (e.g. a previous bootstrap was interrupted), run bootstrap + again — detected by ``files.exists(GATEWAY_SCRIPT)``. + 4. Restore MCPs from ``$workdir/.mcp`` if present, else seed + from ``default_mcps``. + 5. Mint a fresh gateway bearer token (not persisted). + 6. Kill any leftover gateway process, drop a fresh + ``gateway.config.json`` into the sandbox, launch the + gateway, wait for ``/health``. + 7. Pull the gateway-side MCP view back as + :class:`GatewayMCPClient` instances. + + Idempotent — a no-op when already alive. + """ + if self.is_alive: + return + + await self._attach_or_create_sandbox() + self._backend = E2BBackend(self._sandbox, workdir=SANDBOX_WORKDIR) + + # If the gateway script is missing, the sandbox is fresh (or + # a prior bootstrap was interrupted). Re-running bootstrap is + # safe because every step is idempotent (mkdir -p, uv venv, + # uv pip install). + if not await self._sandbox.files.exists(GATEWAY_SCRIPT): + # The backend pins ``cwd=SANDBOX_WORKDIR`` so the very + # first bootstrap command (which itself is ``mkdir -p``) + # would fail before it ran when the dir does not yet + # exist. Use ``cwd="/"`` to break the chicken-and-egg — + # ``mkdir -p`` itself never fails on an already-existing + # directory. + await self._backend.exec_shell( + ["mkdir", "-p", SANDBOX_WORKDIR], + cwd="/", + ) + await self._run_bootstrap() + + self._mcps = await self._restore_or_seed_mcps() + + self._gateway_token = uuid.uuid4().hex + + # Stop any stale gateway from a previous resume cycle. Each + # init mints a new bearer token, so an old gateway listening + # on the port would happily accept old-token requests but + # reject new ones — kill it before starting the new one. + await self._backend.exec_shell( + ["sh", "-c", "pkill -f _mcp_gateway_app.py || true"], + ) + + await self._write_gateway_config() + await self._start_gateway_process() + + host = self._sandbox.get_host(self.gateway_port) + self._gateway = GatewayClient( + base_url=f"https://{host}", + token=self._gateway_token, + timeout=30.0, + extra_headers=self._sandbox_proxy_headers(), + ) + await self._wait_for_gateway() + + self._gateway_clients = { + c.name: c for c in await self._gateway.list_mcps() + } + + # Persist the MCP set unconditionally so a freshly seeded + # ``self._mcps`` (default_mcps path) is rewritten as the + # canonical ``.mcp`` for the next restart, and a restored set + # is round-tripped harmlessly. ``_seed_skills`` itself is + # idempotent — it short-circuits when the sandbox-side + # ``skills/`` already has entries. + await self._save_mcp_file() + await self._seed_skills() + + self.is_alive = True + + async def reset(self) -> None: + """Return the workspace to an empty state. + + Mirrors :meth:`DockerWorkspace.reset`: deregisters every MCP + from the gateway, clears the local handles, and wipes + ``.mcp``, ``skills/``, ``sessions/``, and ``data/`` inside the + sandbox. The gateway process keeps running with no upstream + MCPs. ``default_mcps`` / ``skill_paths`` are not re-seeded. + """ + if self._backend is None: + raise RuntimeError( + "E2BWorkspace is not initialized: its sandbox backend " + "is unavailable. Use 'async with workspace:' or call " + "'await workspace.initialize()' before 'reset()'.", + ) + + async with self._mcp_lock, self._skill_lock: + for gw_client in list(self._gateway_clients.values()): + try: + await gw_client.close() + except Exception as e: + logger.warning( + "MCP %r close failed during reset: %s", + gw_client.name, + e, + ) + self._gateway_clients.clear() + self._mcps = [] + + for path in ( + SANDBOX_SESSIONS_DIR, + SANDBOX_DATA_DIR, + SANDBOX_SKILLS_DIR, + ): + await self._backend.delete_path(path) + + # Rewrite ``.mcp`` to an empty list so a future restart does + # not fall back to ``default_mcps``. + await self._save_mcp_file() + + async def close(self) -> None: + """Pause the sandbox and release host-side resources. + + ``sandbox.pause()`` (not ``kill()``) keeps the sandbox's + filesystem so the next :meth:`initialize` can reattach to it + via metadata lookup. The host-side gateway client is closed + first so its connection pool is released cleanly. + + Errors during teardown are swallowed so ``close`` is always + safe to call (e.g. from ``__aexit__``). + """ + if self._gateway is not None: + try: + await self._gateway.aclose() + except Exception: + pass + self._gateway = None + self._gateway_clients.clear() + + if self._sandbox is not None: + try: + await self._sandbox.pause() + except Exception as e: + logger.warning("E2BWorkspace: pause failed: %s", e) + self._sandbox = None + self._backend = None + + self.is_alive = False + + # ── instructions ──────────────────────────────────────────── + + async def get_instructions(self) -> str: + """Return the system-prompt fragment for this workspace. + + Substitutes ``{workdir}`` in the configured template with + the sandbox-side path (``/home/user/workspace``). The agent + always sees sandbox-internal paths. + """ + return self.instructions.format(workdir=SANDBOX_WORKDIR) + + # ── tool / MCP / skill discovery ──────────────────────────── + + async def list_tools(self) -> list[ToolBase]: + """Built-in tools backed by the E2B sandbox. + + Returns the six builtin tools (Bash, Read, Write, Edit, Grep, + Glob), each backed by the workspace's :class:`E2BBackend` + that executes inside the sandbox. + + Raises: + `RuntimeError`: + If the workspace has not been initialized yet (the + sandbox-backed backend is unavailable). Without this + guard the builtin tools would silently fall back to a + :class:`LocalBackend` and run on the host instead of + inside the sandbox. + """ + if self._backend is None: + raise RuntimeError( + "E2BWorkspace is not initialized: its sandbox backend " + "is unavailable. Use 'async with workspace:' or call " + "'await workspace.initialize()' before 'list_tools()'.", + ) + + from ...tool._builtin import Bash, Edit, Glob, Grep, Read, Write + + return [ + Bash(cwd=SANDBOX_WORKDIR, backend=self._backend), + Edit(backend=self._backend), + Glob(backend=self._backend, glob_helper_path=GLOB_HELPER_SCRIPT), + Grep(backend=self._backend), + Read(backend=self._backend), + Write(backend=self._backend), + ] + + async def list_mcps(self) -> list[MCPClient]: + """Return one :class:`GatewayMCPClient` per registered MCP. + + Each entry's ``name`` matches the upstream MCP server name and + all of its protocol calls are routed over HTTPS to the + in-sandbox gateway. + """ + return list(self._gateway_clients.values()) + + async def list_skills(self) -> list[Skill]: + """Enumerate skills by scanning ``skills/`` inside the sandbox. + + Reads each ``SKILL.md`` via the SDK's ``files.read`` and parses + the YAML front-matter. Files missing ``name`` or ``description`` + are skipped. + """ + import frontmatter as fm + + result = await self._backend.exec_shell( + [ + "sh", + "-c", + f"find {SANDBOX_SKILLS_DIR} -name SKILL.md " + f"2>/dev/null || true", + ], + ) + if not result.ok(): + return [] + listing = result.stdout.decode(errors="replace").strip() + if not listing: + return [] + + skills: list[Skill] = [] + for md_path in (line.strip() for line in listing.split("\n")): + if not md_path: + continue + try: + raw = await self._backend.read_file(md_path) + doc = fm.loads(raw.decode("utf-8")) + name = doc.get("name") + desc = doc.get("description") + if not name or not desc: + continue + skills.append( + Skill( + name=str(name), + description=str(desc), + dir=posixpath.dirname(md_path), + markdown=doc.content or "", + updated_at=0.0, + ), + ) + except Exception as e: + logger.warning("Failed to load skill %s: %s", md_path, e) + return skills + + # ── dynamic MCP management ────────────────────────────────── + + async def add_mcp(self, mcp_client: MCPClient) -> None: + """Register a new MCP server on the in-sandbox gateway. + + Mirrors :meth:`DockerWorkspace.add_mcp` but persists ``.mcp`` + unconditionally — the sandbox filesystem is always + persistent for E2B. + """ + async with self._mcp_lock: + if mcp_client.name in self._gateway_clients: + raise ValueError( + f"MCP {mcp_client.name!r} already exists in workspace.", + ) + spec = mcp_client.model_dump(mode="json") + assert self._gateway is not None + gw_client = self._gateway.make_client(spec) + await gw_client.connect() + self._mcps.append(mcp_client) + self._gateway_clients[gw_client.name] = gw_client + await self._save_mcp_file() + + async def remove_mcp(self, name: str) -> None: + """Unregister an MCP server by name. + + Mirrors :meth:`DockerWorkspace.remove_mcp`. + """ + async with self._mcp_lock: + gw_client = self._gateway_clients.pop(name, None) + if gw_client is None: + logger.warning("MCP %r not found in workspace", name) + return + try: + await gw_client.close() + except Exception as e: + logger.warning("MCP %r close failed: %s", name, e) + self._mcps = [m for m in self._mcps if m.name != name] + await self._save_mcp_file() + + # ── dynamic skill management ──────────────────────────────── + + async def add_skill(self, skill_path: str) -> None: + """Upload a local skill directory into ``skills/`` inside the sandbox. + + The directory must contain a ``SKILL.md`` with ``name`` and + ``description`` in its YAML front matter. A directory of the + same basename already in the sandbox is rejected rather than + overwritten. + """ + skill_md = os.path.join(skill_path, "SKILL.md") + if not os.path.isfile(skill_md): + raise ValueError( + f"Invalid skill at {skill_path!r}: SKILL.md not found", + ) + + async with self._skill_lock: + await self._backend.exec_shell( + ["mkdir", "-p", SANDBOX_SKILLS_DIR], + ) + dir_name = os.path.basename(os.path.abspath(skill_path)) + + check = await self._backend.exec_shell( + ["test", "-e", SANDBOX_SKILLS_DIR + "/" + dir_name], + ) + if check.ok(): + raise ValueError( + f"Skill directory {dir_name!r} already exists in " + f"{SANDBOX_SKILLS_DIR}", + ) + + for root, _dirs, files in os.walk(skill_path): + for fname in files: + local = os.path.join(root, fname) + rel = os.path.relpath(local, skill_path) + remote = f"{SANDBOX_SKILLS_DIR}/{dir_name}/{rel}" + with open(local, "rb") as f: + data = f.read() + await self._backend.write_file(remote, data) + + logger.info( + "E2BWorkspace: added skill %r at %s/%s", + dir_name, + SANDBOX_SKILLS_DIR, + dir_name, + ) + + async def remove_skill(self, name: str) -> None: + """Delete a skill directory by its agent-facing name.""" + skills = await self.list_skills() + target_dir: str | None = None + for s in skills: + if s.name == name: + target_dir = s.dir + break + if target_dir is None: + available = [s.name for s in skills] + raise KeyError( + f"Skill {name!r} not found. Available: {available}", + ) + await self._backend.delete_path(target_dir) + + # ── offload ───────────────────────────────────────────────── + + async def offload_context( + self, + session_id: str, + msgs: list[Msg], + ) -> str: + """Persist a batch of messages as JSONL inside the sandbox. + + Same shape as :meth:`DockerWorkspace.offload_context`: each + :class:`Msg` becomes a line; inline base64 :class:`DataBlock` + payloads are extracted into ``data/`` and replaced with + ``file://`` URL blocks. + """ + base = f"{SANDBOX_SESSIONS_DIR}/{session_id}" + path = f"{base}/context.jsonl" + + copied = deepcopy(msgs) + lines: list[str] = [] + for msg in copied: + if not isinstance(msg.content, str): + content = [] + for block in msg.content: + if isinstance(block, DataBlock) and isinstance( + block.source, + Base64Source, + ): + block = await self._offload_data_block(block) + content.append(block) + msg.content = content + lines.append(msg.model_dump_json()) + + await self._backend.exec_shell(["mkdir", "-p", base]) + existing = b"" + try: + existing = await self._backend.read_file(path) + except FileNotFoundError: + pass + await self._backend.write_file( + path, + existing + ("\n".join(lines) + "\n").encode("utf-8"), + ) + return path + + async def offload_tool_result( + self, + session_id: str, + tool_result: ToolResultBlock, + ) -> str: + """Persist a single tool result as a flat text file.""" + base = f"{SANDBOX_SESSIONS_DIR}/{session_id}" + path = f"{base}/tool_result-{tool_result.id}.txt" + + parts: list[str] = [] + if isinstance(tool_result.output, str): + parts.append(tool_result.output) + else: + for block in tool_result.output: + if isinstance(block, TextBlock): + parts.append(block.text) + elif isinstance(block, DataBlock): + if isinstance(block.source, Base64Source): + d = await self._offload_data_block(block) + url = str(d.source.url) + else: + url = str(block.source.url) + parts.append( + f"", + ) + + await self._backend.exec_shell(["mkdir", "-p", base]) + await self._backend.write_file( + path, + "".join(parts).encode("utf-8"), + ) + return path + + # ── internals: sandbox attach / create ───────────────────── + + async def _attach_or_create_sandbox(self) -> None: + """Reattach to an existing sandbox by metadata, or create one. + + Resolution rule: a single sandbox is expected per + ``workspace_id``. If multiple are returned (e.g. a leaked + running + paused pair after an unclean shutdown) we attach to + the newest by ``started_at`` and log a warning — manual + cleanup is left to the operator. + + Always blocks until the sandbox's envd answers + :meth:`AsyncSandbox.is_running` so the caller can issue + ``commands`` / ``files`` calls without hitting transient + "not yet routable" errors — typical on a paused sandbox that + has just been auto-resumed via ``connect``. + """ + from e2b import AsyncSandbox + + existing = await self._find_existing_sandbox() + + api_opts = self._api_opts() + if existing is not None: + self._sandbox = await AsyncSandbox.connect( + sandbox_id=existing.sandbox_id, + timeout=self.timeout_seconds, + **api_opts, + ) + else: + merged_metadata = { + METADATA_WORKSPACE_ID_KEY: self.workspace_id, + **self.sandbox_metadata, + } + create_kwargs: dict[str, Any] = { + "template": self.template, + "timeout": self.timeout_seconds, + "metadata": merged_metadata, + **api_opts, + } + if self.env: + create_kwargs["envs"] = self.env + + self._sandbox = await AsyncSandbox.create(**create_kwargs) + + await self._wait_until_running() + + async def _wait_until_running(self, timeout: float = 30.0) -> None: + """Poll ``self._sandbox.is_running()`` until it answers ``True``. + + ``AsyncSandbox.create`` / ``AsyncSandbox.connect`` can return + before the in-sandbox envd is routable. Subsequent + ``commands.run`` / ``files.exists`` calls against an + unrouted envd surface as transient SDK errors. We poll envd's + own ``/health`` (which is what :meth:`AsyncSandbox.is_running` + wraps — 502 → ``False``, 200 → ``True``) until it goes green. + + Args: + timeout (`float`, defaults to `30.0`): + Hard ceiling in seconds. Raises :class:`RuntimeError` + if envd is still not routable after this long. + """ + deadline = asyncio.get_event_loop().time() + timeout + delay = 0.1 + while asyncio.get_event_loop().time() < deadline: + try: + if await self._sandbox.is_running(): + return + except Exception as e: # noqa: BLE001 + # SDK can raise on transient network / proxy errors + # while the sandbox is still provisioning. Treat as + # "not yet" and keep polling. + logger.debug( + "E2BWorkspace: is_running probe error (will retry): %s", + e, + ) + await asyncio.sleep(delay) + delay = min(delay * 1.5, 1.0) + raise RuntimeError( + f"E2B sandbox did not become ready within {timeout}s " + f"(workspace_id={self.workspace_id!r})", + ) + + async def _find_existing_sandbox(self) -> Any: + """List sandboxes filtered by ``workspace_id`` metadata. + + Returns the most recent :class:`SandboxInfo` (paused or + running) or ``None`` if no match exists. + """ + from e2b import AsyncSandbox + from e2b.api.client.models.sandbox_state import SandboxState + from e2b.sandbox.sandbox_api import SandboxQuery + + query = SandboxQuery( + metadata={METADATA_WORKSPACE_ID_KEY: self.workspace_id}, + state=[SandboxState.PAUSED, SandboxState.RUNNING], + ) + + candidates: list[Any] = [] + paginator = AsyncSandbox.list(query=query, **self._api_opts()) + while paginator.has_next: + try: + page = await paginator.next_items() + except Exception as e: + logger.warning( + "E2BWorkspace: list sandboxes failed: %s", + e, + ) + break + candidates.extend(page) + + if not candidates: + return None + if len(candidates) > 1: + logger.warning( + "E2BWorkspace: %d sandboxes match workspace_id=%r; " + "attaching to most recent", + len(candidates), + self.workspace_id, + ) + candidates.sort(key=lambda s: s.started_at, reverse=True) + return candidates[0] + + def _api_opts(self) -> dict[str, Any]: + """Common ``api_key`` / ``domain`` opts forwarded to E2B SDK calls.""" + opts: dict[str, Any] = {} + if self.api_key: + opts["api_key"] = self.api_key + if self.domain: + opts["domain"] = self.domain + return opts + + def _sandbox_proxy_headers(self) -> dict[str, str]: + """Headers required by the E2B proxy to reach the gateway port. + + E2B's edge proxy gates non-default ports behind the + ``X-Access-Token`` header tied to the sandbox. The token is + exposed on the sandbox object as ``traffic_access_token``. + """ + if self._sandbox is None: + return {} + token = getattr(self._sandbox, "traffic_access_token", None) + if not token: + return {} + return {"X-Access-Token": token} + + # ── internals: bootstrap ──────────────────────────────────── + + async def _run_bootstrap(self) -> None: + """Provision a fresh sandbox: uv → venv → agentscope → script. + + Each command runs through :meth:`_exec`; a non-zero exit + raises :class:`RuntimeError` with the captured stderr so + startup failures are visible in logs (mirroring the docker + build-tail strategy). + """ + if _is_released_install(): + log_bootstrap_attempt(self.workspace_id, "released") + install_cmd = render_install_agentscope_cmd_released( + _agentscope_version(), + ) + else: + log_bootstrap_attempt(self.workspace_id, "dev") + tar_bytes = build_source_tarball() + await self._backend.write_file(DEV_SRC_TAR, tar_bytes) + install_cmd = render_install_agentscope_cmd_dev() + + commands = bootstrap_commands( + extra_pip=self.extra_pip, + install_agentscope_cmd=install_cmd, + ) + for cmd in commands: + r = await self._backend.exec_shell( + ["sh", "-c", cmd], + timeout=600.0, + ) + if not r.ok(): + raise RuntimeError( + f"E2BWorkspace bootstrap failed (exit {r.exit_code}) " + f"for: {cmd!r}\n" + f"stderr: {r.stderr.decode(errors='replace')}\n" + f"stdout: {r.stdout.decode(errors='replace')}", + ) + + # Upload helper scripts used by builtin tools. + await self._backend.write_file( + GLOB_HELPER_SCRIPT, + _read_glob_helper_bytes(), + ) + + # Upload the gateway script last so its presence is the + # idempotency marker we probe in :meth:`initialize`. + await self._backend.write_file( + GATEWAY_SCRIPT, + _read_gateway_script_bytes(), + ) + + # ── internals: gateway lifecycle ──────────────────────────── + + async def _restore_or_seed_mcps(self) -> list[MCPClient]: + """Decide the MCP set to ship to the gateway on startup. + + * ``$workdir/.mcp`` missing → return ``default_mcps``. + * ``.mcp`` present → :meth:`MCPClient.model_validate` each + entry. Read / parse error → log and fall back to + ``default_mcps``. + """ + try: + raw = await self._backend.read_file(SANDBOX_MCP_FILE) + except FileNotFoundError: + return list(self.default_mcps) + try: + data = json.loads(raw.decode("utf-8")) + return [MCPClient.model_validate(m) for m in data] + except Exception as e: + logger.warning( + "E2BWorkspace: failed to parse %s, falling back to " + "default_mcps: %s", + SANDBOX_MCP_FILE, + e, + ) + return list(self.default_mcps) + + async def _save_mcp_file(self) -> None: + """Persist ``self._mcps`` to ``$workdir/.mcp`` inside the sandbox. + + Failures are logged but not raised. + """ + payload = json.dumps( + [m.model_dump(mode="json") for m in self._mcps], + indent=2, + ensure_ascii=False, + ) + try: + await self._backend.exec_shell( + ["mkdir", "-p", SANDBOX_WORKDIR], + ) + await self._backend.write_file( + SANDBOX_MCP_FILE, + payload.encode("utf-8"), + ) + except Exception as e: + logger.warning( + "E2BWorkspace: failed to save %s: %s", + SANDBOX_MCP_FILE, + e, + ) + + async def _write_gateway_config(self) -> None: + """Drop the gateway's ``--config`` JSON into the sandbox.""" + cfg = { + "token": self._gateway_token, + "servers": [m.model_dump(mode="json") for m in self._mcps], + } + await self._backend.exec_shell( + ["mkdir", "-p", GATEWAY_HOME], + ) + await self._backend.write_file( + GATEWAY_CONFIG, + json.dumps(cfg, indent=2, ensure_ascii=False).encode("utf-8"), + ) + + async def _start_gateway_process(self) -> None: + """Launch the gateway inside the sandbox as a detached process.""" + cmd = ( + f"nohup {shlex.quote(GATEWAY_VENV_PY)} -u " + f"{shlex.quote(GATEWAY_SCRIPT)} " + f"--config {shlex.quote(GATEWAY_CONFIG)} " + f"--port {self.gateway_port} " + f"> {shlex.quote(GATEWAY_LOG)} 2>&1 &" + ) + await self._backend.exec_shell(["sh", "-c", cmd]) + + async def _wait_for_gateway(self, timeout: float = 30.0) -> None: + """Block until the gateway answers ``/health`` with 200.""" + assert self._gateway is not None + deadline = asyncio.get_event_loop().time() + timeout + delay = 0.1 + while asyncio.get_event_loop().time() < deadline: + if await self._gateway.health(): + return + await asyncio.sleep(delay) + delay = min(delay * 1.5, 1.0) + try: + log = await self._backend.read_file(GATEWAY_LOG) + tail = log[-2000:].decode(errors="replace") + except Exception: + tail = "" + raise RuntimeError( + f"gateway did not become healthy within {timeout}s. " + f"Tail of {GATEWAY_LOG}:\n{tail}", + ) + + async def _seed_skills(self) -> None: + """Copy ``self.skill_paths`` into ``skills/`` once, on first init. + + Skips seeding when ``skill_paths`` is empty or the sandbox-side + ``skills/`` already contains entries — meaning the user (or a + prior init) is the source of truth. + """ + if not self.skill_paths: + return + listing = await self._backend.exec_shell( + [ + "sh", + "-c", + f"ls -A {shlex.quote(SANDBOX_SKILLS_DIR)} " + f"2>/dev/null || true", + ], + ) + if listing.ok() and listing.stdout.strip(): + return + for path in self.skill_paths: + try: + await self.add_skill(path) + except Exception as e: + logger.warning( + "E2BWorkspace: skip skill %r: %s", + path, + e, + ) + + # ── internals: data offload ──────────────────────────────── + + async def _offload_data_block(self, block: DataBlock) -> DataBlock: + """Persist a base64 :class:`DataBlock` under ``data/``. + + Mirrors :meth:`DockerWorkspace._offload_data_block` exactly, + only the I/O primitive differs. Hashing the *base64* text + rather than decoded bytes keeps the key short-circuit: a + repeat offload of the same block writes the same file. + """ + if not isinstance(block.source, Base64Source): + return block + h = hashlib.sha256(block.source.data.encode()).hexdigest() + ext = mimetypes.guess_extension(block.source.media_type) or ".bin" + path = f"{SANDBOX_DATA_DIR}/{h}{ext}" + await self._backend.exec_shell( + ["mkdir", "-p", SANDBOX_DATA_DIR], + ) + await self._backend.write_file( + path, + base64.b64decode(block.source.data), + ) + return DataBlock( + id=block.id, + name=block.name, + source=URLSource( + url=AnyUrl(f"file://{path}"), + media_type=block.source.media_type, + ), + ) diff --git a/src/agentscope/workspace/_gateway_client.py b/src/agentscope/workspace/_gateway_client.py new file mode 100644 index 0000000000000000000000000000000000000000..9679e8e3da9d0147f0d50ad0a41612a95f574f65 --- /dev/null +++ b/src/agentscope/workspace/_gateway_client.py @@ -0,0 +1,634 @@ +# -*- coding: utf-8 -*- +"""Host-side client for the in-workspace MCP gateway. + +Three classes live here: + +* :class:`GatewayClient` — workspace-side facade over the gateway's + ``/health`` and ``/mcps`` endpoints. Used by ``DockerWorkspace`` (and + later ``E2BWorkspace``) for top-level operations. + +* :class:`GatewayMCPClient` — an :class:`MCPClient` subclass whose + protocol behaviour is replaced by HTTP calls to the gateway. The + field surface is identical to ``MCPClient`` (instances are built from + ``MCPClient.model_dump()`` data the gateway returns), so callers + that ``model_dump()`` it round-trip cleanly. Local stdio/HTTP + session machinery is bypassed: ``model_post_init`` is a no-op, + ``connect`` POSTs to ``/mcps``, ``close`` DELETEs ``/mcps/{name}``, + and ``list_tools`` / ``get_tool`` fetch / wrap upstream tools. + +* :class:`GatewayMCPTool` — :class:`ToolBase` subclass whose + ``__call__`` posts to ``/mcps/{name}/tools/{tool}`` and reconstructs + the returned ``ToolChunk``. +""" + +import contextlib +from typing import Any, AsyncIterator + +import httpx +import mcp.types +from pydantic import PrivateAttr + +from ..mcp import MCPClient +from ..message import ToolResultState +from ..permission import ( + PermissionBehavior, + PermissionDecision, +) +from ..tool import ToolBase, ToolChunk + + +# ── tool ─────────────────────────────────────────────────────────── + + +class GatewayMCPTool(ToolBase): + """An MCP tool whose ``__call__`` is a single HTTP POST to the gateway. + + Mirrors :class:`agentscope.tool.MCPTool` field-by-field so the toolkit + treats it identically (same ``name`` format, same permission policy) + — only the call path changes. + """ + + is_mcp: bool = True + is_state_injected: bool = False + + def __init__( + self, + mcp_name: str, + tool: mcp.types.Tool, + gateway_url: str, + token: str, + http: httpx.AsyncClient | None = None, + timeout: float | None = None, + ) -> None: + """Build a gateway-backed MCP tool. + + The instance mirrors the field surface of + :class:`agentscope.tool.MCPTool` (``name``, ``description``, + ``input_schema``, ``is_read_only``, …) so the host-side toolkit + cannot tell the difference between a local MCP tool and one that + forwards through the in-container gateway. + + Args: + mcp_name: Name of the upstream MCP server this tool belongs + to. Used both for the visible ``mcp__{mcp}__{tool}`` + name and for the gateway URL path. + tool: Raw upstream tool descriptor as returned by the + gateway. Its ``name`` is the upstream-side identifier + (no ``mcp__`` prefix), ``inputSchema`` is forwarded + verbatim, and ``annotations.readOnlyHint`` drives the + permission policy. + gateway_url: Host-visible base URL of the gateway, e.g. + ``http://127.0.0.1:``. Trailing slash is + stripped. + token: Bearer token the host injected into the gateway's + config; sent as ``Authorization: Bearer …`` on every + call. + http: Shared :class:`httpx.AsyncClient` to reuse connection + pooling across many tool calls. When ``None`` each + ``__call__`` creates and disposes a one-shot client. + timeout: Per-call HTTP timeout in seconds. Only consulted + when ``http`` is ``None`` (the shared client carries + its own timeout). + """ + self.mcp_name = mcp_name + self.name = f"mcp__{mcp_name}__{tool.name}" + self.description = tool.description or "" + + schema = dict(tool.inputSchema) if tool.inputSchema else {} + schema.setdefault("type", "object") + schema.setdefault("properties", {}) + schema.setdefault("required", []) + self.input_schema = schema + + self.is_concurrency_safe = False + self.is_external_tool = False + + self.is_read_only = False + if tool.annotations and hasattr(tool.annotations, "readOnlyHint"): + self.is_read_only = tool.annotations.readOnlyHint or False + + self._tool = tool + self._gateway_url = gateway_url.rstrip("/") + self._token = token + self._http = http + self._timeout = timeout + + async def check_permissions( + self, + *_args: Any, + **_kwargs: Any, + ) -> PermissionDecision: + """Default policy: read-only tools auto-allow, everything else + defers to the user via ``ASK``. Mirrors + :class:`agentscope.tool.MCPTool.check_permissions` so toolkit + callers see identical behaviour through the gateway. + """ + if self.is_read_only: + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message="This is a read-only MCP tool. Allowing execution.", + ) + return PermissionDecision( + behavior=PermissionBehavior.ASK, + message="MCP tools must be explicitly allowed by the user.", + ) + + async def __call__(self, **kwargs: Any) -> ToolChunk: + """Invoke the upstream tool by POSTing to + ``/mcps/{mcp}/tools/{tool}`` on the gateway. + + Args: + **kwargs: Tool arguments forwarded as the JSON body's + ``arguments`` field; the gateway re-dispatches them to + the upstream MCP session. + + Returns: + `ToolChunk`: + The reconstructed chunk returned by the upstream tool. + 4xx / 5xx responses are surfaced as a + ``ToolChunk(state=ERROR)`` so the agent loop can reason + about the failure instead of crashing. + + Raises: + RuntimeError: If the gateway returns 2xx but no ``chunk`` + payload (protocol violation on the gateway side). + """ + url = ( + f"{self._gateway_url}/mcps/{self.mcp_name}" + f"/tools/{self._tool.name}" + ) + headers = _bearer_headers(self._token) + async with _http_session(self._http, self._timeout) as http: + resp = await http.post( + url, + json={"arguments": kwargs}, + headers=headers, + ) + if resp.status_code >= 400: + # Surface gateway-side error as a failed ToolChunk so + # the agent loop can reason about it instead of crashing. + detail = _safe_detail(resp) + return ToolChunk( + content=[{"type": "text", "text": detail}], + state=ToolResultState.ERROR, + ) + payload = resp.json() + chunk_dict = payload.get("chunk") + if chunk_dict is None: + raise RuntimeError( + f"gateway returned no chunk for {self.name!r}", + ) + return ToolChunk.model_validate(chunk_dict) + + +# ── pseudo MCP client ────────────────────────────────────────────── + + +class GatewayMCPClient(MCPClient): + """An :class:`MCPClient` whose protocol logic is replaced by HTTP. + + Constructed from the dict returned by ``GET /mcps`` (or freshly from + user input via :meth:`GatewayClient.make_client`). The local MCP + machinery is short-circuited entirely: + + * ``model_post_init`` does nothing (parent's ``_initialize_client`` + is never called — no stdio context manager is built). + * ``connect`` POSTs to ``/mcps`` to register-and-start the upstream + server inside the gateway. + * ``close`` DELETEs ``/mcps/{name}``. + * ``list_tools`` / ``get_tool`` fetch and wrap upstream tools. + """ + + _gateway_url: str = PrivateAttr(default="") + _gateway_token: str = PrivateAttr(default="") + _http_timeout: float | None = PrivateAttr(default=None) + _http: httpx.AsyncClient | None = PrivateAttr(default=None) + + def model_post_init(self, __context: Any) -> None: + """Skip the parent's stdio/HTTP client preparation. + + For a real :class:`MCPClient`, ``model_post_init`` builds the + local stdio context manager (or wires up an HTTP client) so + the in-process session can be opened. For + :class:`GatewayMCPClient` all MCP-side work happens inside the + gateway container; the host-side proxy needs no local session + machinery, so this override is a no-op. + """ + return + + # ── lifecycle ───────────────────────────────────────────────── + + def attach( + self, + *, + gateway_url: str, + token: str, + http: httpx.AsyncClient | None, + timeout: float | None, + connected: bool = False, + ) -> None: + """Wire this client to a gateway transport. + + :class:`GatewayMCPClient` is normally produced by + ``model_validate(spec)`` over a dict returned by the gateway's + ``GET /mcps`` endpoint — that step recovers the public field + surface but leaves all transport-related private attributes + empty. ``attach`` injects them in a single call so subsequent + :meth:`connect`, :meth:`close`, :meth:`list_raw_tools`, and + :meth:`get_tool` can talk to the gateway. It is the only + supported way to populate the transport state from outside the + class — encapsulating the writes here keeps callers free of + ``protected-access`` warnings. + + Args: + gateway_url: Host-visible base URL of the gateway (e.g. + ``http://127.0.0.1:``). Trailing slash is + stripped before storage. + token: Bearer token the host generated for this gateway. + Sent as ``Authorization: Bearer …`` on every request. + http: Shared :class:`httpx.AsyncClient` provided by the + owning :class:`GatewayClient` so connection pooling is + shared across all derived clients and tools. When + ``None`` each call creates a one-shot client. + timeout: Default per-request timeout in seconds, used only + when ``http`` is ``None``. + connected: When ``True``, mark this client as already + connected (i.e. the gateway is already maintaining the + upstream session). Used by + :meth:`GatewayClient.list_mcps` for clients that came + back from the gateway as registered. Leave ``False`` + when the caller will call :meth:`connect` themselves. + """ + self._gateway_url = gateway_url.rstrip("/") + self._gateway_token = token + self._http = http + self._http_timeout = timeout + if connected: + self._is_connected = True + + async def connect(self) -> None: + """Register this MCP on the gateway via ``POST /mcps``. + + Stateless MCPs are a no-op (the gateway invokes them on + demand). Stateful MCPs are registered and started inside the + gateway container; this method blocks until the gateway has + confirmed the upstream connection. + + Raises: + RuntimeError: If the client is already connected, or if + the gateway returns a 4xx/5xx response. + """ + if not self.is_stateful: + return + if self._is_connected: + raise RuntimeError( + f"MCP {self.name!r} is already connected. " + "Call close() before reconnecting.", + ) + body = self.model_dump(mode="json") + async with _http_session(self._http, self._http_timeout) as http: + resp = await http.post( + f"{self._gateway_url}/mcps", + json=body, + headers=_bearer_headers(self._gateway_token), + ) + if resp.status_code >= 400: + raise RuntimeError( + f"gateway failed to add MCP {self.name!r}: " + f"{_safe_detail(resp)}", + ) + self._is_connected = True + + async def close(self, ignore_errors: bool = True) -> None: + """Deregister this MCP from the gateway via + ``DELETE /mcps/{name}``. + + Stateless MCPs are a no-op. For stateful MCPs the gateway + closes the upstream session before responding. + + Args: + ignore_errors: When ``True`` (the default), suppress both + "not connected" precondition failures and + gateway-side 4xx/5xx responses; when ``False`` such + conditions raise :class:`RuntimeError`. Mirrors + :meth:`MCPClient.close` so callers can use the same + shutdown idiom regardless of transport. + """ + if not self.is_stateful: + return + if not self._is_connected: + if ignore_errors: + return + raise RuntimeError( + f"MCP {self.name!r} is not connected. Call connect() first.", + ) + try: + async with _http_session(self._http, self._http_timeout) as http: + resp = await http.delete( + f"{self._gateway_url}/mcps/{self.name}", + headers=_bearer_headers(self._gateway_token), + ) + if resp.status_code >= 400 and not ignore_errors: + raise RuntimeError( + f"gateway failed to remove MCP {self.name!r}: " + f"{_safe_detail(resp)}", + ) + except Exception: + if not ignore_errors: + raise + self._is_connected = False + + # ── tool discovery ──────────────────────────────────────────── + + async def list_raw_tools(self) -> list[mcp.types.Tool]: + """Fetch the upstream tool list via ``GET /mcps/{name}/tools``. + + Returns the raw :class:`mcp.types.Tool` descriptors the gateway + forwarded — i.e. with their **upstream** names (no ``mcp__`` + prefix) so the inherited :meth:`list_tools` / :meth:`get_tool` + path can re-wrap them through :meth:`_wrap_tool` exactly as a + local :class:`MCPClient` would. The full unfiltered list is + cached on ``_cached_tools`` first; the returned list then has + ``enable_tools`` / ``disable_tools`` filtering applied + identically to :meth:`MCPClient.list_raw_tools`. + + Returns: + `list[mcp.types.Tool]`: + The upstream-named, post-filter tool descriptors. + + Raises: + httpx.HTTPStatusError: If the gateway returns a non-2xx + response. + """ + async with _http_session(self._http, self._http_timeout) as http: + resp = await http.get( + f"{self._gateway_url}/mcps/{self.name}/tools", + headers=_bearer_headers(self._gateway_token), + ) + resp.raise_for_status() + data = resp.json() + + raw_tools = [mcp.types.Tool.model_validate(d) for d in data] + self._cached_tools = raw_tools + + # Honour the same enable/disable filtering MCPClient does locally — + # gateway returns the unfiltered upstream view. + if self.enable_tools is not None: + raw_tools = [t for t in raw_tools if t.name in self.enable_tools] + if self.disable_tools is not None: + raw_tools = [ + t for t in raw_tools if t.name not in self.disable_tools + ] + return raw_tools + + async def get_tool( # type: ignore[override] + self, + name: str, + ) -> GatewayMCPTool: + """Look up a single tool by upstream name and wrap it. + + Falls back to :meth:`list_raw_tools` on cache miss, then + searches ``_cached_tools`` (which holds the **unfiltered** + upstream view) so tools that ``enable_tools`` / + ``disable_tools`` would have hidden are still resolvable — + matching :meth:`MCPClient.get_tool`'s behaviour. + + The wrapped tool inherits the gateway-wide HTTP timeout + (``_http_timeout``, set via :meth:`attach`). Per-call execution + timeout is the upstream MCP server's responsibility — it is + carried in :attr:`MCPClient.execution_timeout`, serialised into + the spec, and reconstructed on the gateway side; the host has + no need to override it. + + Args: + name: Upstream tool name (no ``mcp__`` prefix). The + returned :class:`GatewayMCPTool` exposes the prefixed + form via its own ``name`` attribute. + + Returns: + `GatewayMCPTool`: + A fresh wrapper around the upstream descriptor, ready + to be ``await``-ed or registered with a toolkit. + + Raises: + ValueError: If no tool with that upstream name exists on + the gateway side. + """ + if self._cached_tools is None: + await self.list_raw_tools() + for raw in self._cached_tools or []: + if raw.name == name: + return self._wrap_tool(raw) + raise ValueError( + f"Tool {name!r} not found in MCP {self.name!r}.", + ) + + # ── helpers ─────────────────────────────────────────────────── + + def _wrap_tool(self, tool: mcp.types.Tool) -> GatewayMCPTool: + """Build a :class:`GatewayMCPTool` bound to this client's + gateway transport. Always uses the client-wide + ``_http_timeout`` — there is no per-call override path. + + Args: + tool: Raw upstream tool descriptor (typically pulled out of + ``_cached_tools``). + """ + return GatewayMCPTool( + mcp_name=self.name, + tool=tool, + gateway_url=self._gateway_url, + token=self._gateway_token, + http=self._http, + timeout=self._http_timeout, + ) + + +# ── workspace-side facade ────────────────────────────────────────── + + +class GatewayClient: + """Workspace-side facade over the in-container MCP gateway. + + Owns a shared :class:`httpx.AsyncClient` so all derived + :class:`GatewayMCPClient` and :class:`GatewayMCPTool` instances + share connection pooling. + + The gateway ``base_url`` is the host-visible URL (e.g. + ``http://127.0.0.1:`` after Docker port mapping); the + ``token`` is the bearer the host generated and shipped into the + container's gateway config. + """ + + def __init__( + self, + base_url: str, + token: str, + timeout: float | None = None, + extra_headers: dict[str, str] | None = None, + ) -> None: + """Build a workspace-side gateway facade. + + Args: + base_url: Host-visible base URL of the gateway, e.g. + ``http://127.0.0.1:`` after Docker port + mapping or ``https://.e2b.dev`` for E2B. + Trailing slash is stripped. + token: Bearer token shared with the gateway via its config + file. Sent as ``Authorization: Bearer …`` on every + request and propagated to every derived + :class:`GatewayMCPClient` / :class:`GatewayMCPTool` via + :meth:`make_client`. + timeout: Default HTTP timeout in seconds, applied to the + shared :class:`httpx.AsyncClient` (see + :meth:`_client`) and propagated to every derived + :class:`GatewayMCPClient` via :meth:`make_client`. + There is no per-call override path; per-tool execution + timeouts live on :attr:`MCPClient.execution_timeout` + and are honoured upstream inside the gateway. + extra_headers: Default headers applied to every request + through the shared httpx client (in addition to the + per-call bearer). :class:`E2BWorkspace` uses this to + inject E2B's ``X-Access-Token`` proxy header. + """ + self.base_url = base_url.rstrip("/") + self.token = token + self.timeout = timeout + self.extra_headers: dict[str, str] = dict(extra_headers or {}) + self._http: httpx.AsyncClient | None = None + + def _client(self) -> httpx.AsyncClient: + """Return (lazily creating on first use) the shared + :class:`httpx.AsyncClient` reused by every gateway request.""" + if self._http is None: + self._http = httpx.AsyncClient( + timeout=self.timeout, + headers=self.extra_headers or None, + ) + return self._http + + def _headers(self) -> dict[str, str]: + """Build the bearer-auth header dict for direct + :class:`GatewayClient` calls (``/health``, ``/mcps``).""" + return _bearer_headers(self.token) + + async def health(self) -> bool: + """Probe ``/health`` — used by the workspace to wait for readiness.""" + try: + resp = await self._client().get(f"{self.base_url}/health") + except Exception: + return False + return resp.status_code == 200 + + async def list_mcps(self) -> list[GatewayMCPClient]: + """Fetch every MCP currently registered on the gateway. + + The returned clients are marked as already connected (via + :meth:`GatewayMCPClient.attach`'s ``connected=True``) because + the gateway is already maintaining their upstream sessions — + the host should not invoke :meth:`GatewayMCPClient.connect` + again. + + Returns: + `list[GatewayMCPClient]`: + One transport-wired client per registered MCP. The + workspace's :meth:`list_mcps` implementation surfaces + this list straight to its consumer. + + Raises: + httpx.HTTPStatusError: If the gateway returns a non-2xx + response. + """ + resp = await self._client().get( + f"{self.base_url}/mcps", + headers=self._headers(), + ) + resp.raise_for_status() + return [self.make_client(spec, connected=True) for spec in resp.json()] + + def make_client( + self, + spec: dict[str, Any], + *, + connected: bool = False, + ) -> GatewayMCPClient: + """Build a :class:`GatewayMCPClient` wired to this gateway. + + Reconstructs the public field surface from ``spec`` via + :meth:`MCPClient.model_validate`, then hands the + transport-related private state to the new client through + :meth:`GatewayMCPClient.attach`. Doing the wiring through + ``attach`` keeps the writes inside the target class and avoids + ``protected-access`` warnings on every assignment. + + Args: + spec: A dict produced by ``MCPClient.model_dump(mode="json")`` + — typically the body returned by the gateway's + ``GET /mcps`` endpoint, or built from user input by + ``DockerWorkspace.add_mcp``. + connected: When ``True``, mark the new client as already + connected so :meth:`GatewayMCPClient.connect` need not + run again. Set by :meth:`list_mcps` for clients that + came back from the gateway already registered. Leave + ``False`` for fresh clients the caller will explicitly + ``await client.connect()`` on (the ``add_mcp`` path). + + Returns: + `GatewayMCPClient`: + A pydantic-valid client whose transport state is fully + populated. Stateful clients still require an explicit + ``await client.connect()`` unless ``connected=True``. + """ + client = GatewayMCPClient.model_validate(spec) + client.attach( + gateway_url=self.base_url, + token=self.token, + http=self._client(), + timeout=self.timeout, + connected=connected, + ) + return client + + async def aclose(self) -> None: + """Close the shared HTTP client.""" + if self._http is not None: + await self._http.aclose() + self._http = None + + +# ── module-private utilities ─────────────────────────────────────── + + +def _bearer_headers(token: str) -> dict[str, str]: + if token: + return {"Authorization": f"Bearer {token}"} + return {} + + +@contextlib.asynccontextmanager +async def _http_session( + shared: httpx.AsyncClient | None, + timeout: float | None, +) -> AsyncIterator[httpx.AsyncClient]: + """Yield a shared httpx client when injected, else a one-shot client. + + The shared variant lets the workspace pool connections across many + tool calls without each call paying the TLS/handshake cost. + """ + if shared is not None: + yield shared + else: + async with httpx.AsyncClient(timeout=timeout) as http: + yield http + + +def _safe_detail(resp: httpx.Response) -> str: + """Best-effort extraction of an HTTPException-style detail from a + response.""" + try: + body = resp.json() + except Exception: + return f"HTTP {resp.status_code}: {resp.text[:200]}" + if isinstance(body, dict) and "detail" in body: + return f"HTTP {resp.status_code}: {body['detail']}" + return f"HTTP {resp.status_code}: {str(body)[:200]}" diff --git a/src/agentscope/workspace/_local_workspace.py b/src/agentscope/workspace/_local_workspace.py new file mode 100644 index 0000000000000000000000000000000000000000..5ee2d5bb06978a54c03b60efbaf24e15d026ea2a --- /dev/null +++ b/src/agentscope/workspace/_local_workspace.py @@ -0,0 +1,1079 @@ +# -*- coding: utf-8 -*- +"""The local workspace class.""" + +import asyncio +import base64 +import hashlib +import json +import mimetypes +import os +import re +import shutil +from copy import deepcopy +from pathlib import Path +from typing import TypedDict + +import frontmatter +from pydantic import AnyUrl + +from .._logging import logger +from ..mcp import MCPClient +from ..message import ( + Base64Source, + DataBlock, + Msg, + TextBlock, + ToolResultBlock, + URLSource, +) +from ..skill import Skill +from ..tool import ( + Bash, + Edit, + Glob, + Grep, + Read, + ToolBase, + Write, +) +from ..tool._builtin._backend import LocalBackend +from ._base import WorkspaceBase + + +class _SkillEntry(TypedDict): + """A single entry in the .skills index file.""" + + hash: str + """SHA-256 hash of the skill's SKILL.md content.""" + skill_name: str + """The name exposed to the agent (may differ from the directory name).""" + + +class _SkillsFile(TypedDict): + """Schema of the .skills index file stored inside skills_dir.""" + + skills_dir_mtime: float + """mtime of skills_dir at the time the index was last written.""" + skills: dict[str, _SkillEntry] + """Mapping from directory name (relative to skills_dir) to skill entry.""" + + +def _sanitize_dir_name(name: str) -> str: + """Sanitize a skill name into a safe directory name. + + Allowed characters: ASCII letters, digits, CJK unified ideographs, + hyphens, and underscores. Everything else is replaced with ``_``. + + Args: + name (`str`): + The raw skill name from SKILL.md frontmatter. + + Returns: + `str`: + A sanitized string safe to use as a directory name on Windows, + macOS, and Linux. + """ + return re.sub(r"[^\w一-鿿-]", "_", name) + + +_DEFAULT_WORKSPACE_INSTRUCTIONS = """ +You have access to a local workspace at {workdir} with the following structure: + +``` +{workdir} +├── data/ # offloaded multimodal files (images, etc.) +├── skills/ # reusable skills, each in its own subdirectory +└── sessions/ # session context and tool results +``` + +This workspace is your personal working environment for completing various tasks. +You are responsible for keeping it clean, structured, and easy to navigate over time. + +### Project Directory +- Create a dedicated subdirectory for each task or project under the workspace root. +- Name the directory concisely and descriptively, e.g. `20240315_web-scraper`, so it remains identifiable long after creation. +- Always create a `README.md` at the project root documenting: + - What the project is about + - When it was created + - Key decisions or context that would help you resume work later + - The changes you have made (and when) + +### Version Control +- It is recommended to initialize a `git` repository in each project directory + to track changes and allow rollbacks. +- Always create a `.gitignore` before the first commit to exclude unwanted files + (e.g. virtual environments, cache, secrets). + +### Python Environment +- If a project requires Python, use `uv` to create an isolated virtual environment + inside the project directory: + ```shell + uv venv && uv pip install ... + ``` +- Never install packages into a shared or global environment — each project must + manage its own dependencies to avoid conflicts. +""" # noqa: E501 + + +class LocalWorkspace(WorkspaceBase): + # pylint: disable=line-too-long + """Local-directory workspace. + + Layout:: + + {workdir}/ + ├── .mcp # persisted MCP client configs (JSON array) + ├── data/ # offloaded multimodal files + ├── skills/ # skill subdirectories + └── sessions/ # per-session context and tool-result files + """ # noqa: E501 + + def __init__( + self, + *, + workdir: str, + workspace_id: str | None = None, + default_mcps: list[MCPClient] | None = None, + skill_paths: list[str] | None = None, + instructions: str = _DEFAULT_WORKSPACE_INSTRUCTIONS, + ) -> None: + """Construct a :class:`LocalWorkspace`. + + Args: + workdir (`str`): + Filesystem path to the workspace root. Created on + demand. Always resolved to an absolute path. + workspace_id (`str | None`, optional): + Existing workspace identifier to adopt. ``None`` + generates a fresh UUID. + default_mcps (`list[MCPClient] | None`, optional): + MCP clients seeded into a brand-new workspace. + Ignored on subsequent restarts that already have a + persisted ``/.mcp`` file. + skill_paths (`list[str] | None`, optional): + Local skill directories seeded into + ``/skills`` on first :meth:`initialize`. + instructions (`str`, defaults to \ + `_DEFAULT_WORKSPACE_INSTRUCTIONS`): + System-prompt fragment template returned by + :meth:`get_instructions`. Supports the ``{workdir}`` + placeholder. + """ + super().__init__(workspace_id=workspace_id) + + # ── serializable config ───────────────────────────────── + self.workdir = os.path.abspath(workdir) + self.instructions = instructions.format(workdir=self.workdir) + + # ── seed-only ─────────────────────────────────────────── + self.default_mcps: list[MCPClient] = list(default_mcps or []) + self.skill_paths: list[str] = list(skill_paths or []) + + # ── runtime state ─────────────────────────────────────── + self._backend = LocalBackend() + self._mcps: list[MCPClient] = [] + + self._skill_lock = asyncio.Lock() + self._mcp_lock = asyncio.Lock() + + async def initialize(self) -> None: + """Initialise the workspace. + + MCP state is restored from ``.mcp`` if it exists; otherwise + ``default_mcps`` are used and persisted so the next start picks + them up from disk. ``skill_paths`` are seeded on first use. + + Idempotent: a no-op when the workspace is already alive. + """ + if self.is_alive: + return + + os.makedirs(self.workdir, exist_ok=True) + + # Restore or seed MCPs + mcp_file = os.path.join(self.workdir, ".mcp") + if await self._backend.file_exists(mcp_file): + raw = await self._backend.read_file(mcp_file) + raw_list = json.loads(raw.decode("utf-8")) + for m in raw_list: + try: + self._mcps.append(MCPClient.model_validate(m)) + except Exception as e: + logger.warning( + "Skipping invalid MCP entry '%s': %s", + m.get("name", "?"), + e, + ) + else: + self._mcps = list(self.default_mcps) + await self._save_mcp_file() + + failed: list[MCPClient] = [] + for mcp in self._mcps: + if mcp.is_stateful and not mcp.is_connected: + try: + await mcp.connect() + except Exception as e: + logger.warning( + "Failed to connect stateful MCP '%s': %s, removing.", + mcp.name, + e, + ) + failed.append(mcp) + for mcp in failed: + self._mcps.remove(mcp) + + # Seed skills + skills_dir = os.path.join(self.workdir, "skills") + os.makedirs(skills_dir, exist_ok=True) + + skills_file = await self._load_skills_file(skills_dir) + existing: dict[str, _SkillEntry] = skills_file["skills"] + + # Build fast-lookup sets from the current index + existing_hashes: set[str] = {e["hash"] for e in existing.values()} + existing_agent_names: set[str] = { + e["skill_name"] for e in existing.values() + } + existing_dir_names: set[str] = set(existing.keys()) + + updated = False + for skill_path in self.skill_paths: + result = await self._validate_and_hash_skill(skill_path) + if result is None: + continue + + _, raw_name, skill_hash = result + + # Skip if already present (by content hash) + if skill_hash in existing_hashes: + logger.info( + "Skill '%s' (hash: %s...) already exists, skipping", + raw_name, + skill_hash[:8], + ) + continue + + # Resolve agent-facing name conflict + agent_name = raw_name + counter = 1 + while agent_name in existing_agent_names: + agent_name = f"{raw_name} ({counter})" + counter += 1 + + # Resolve directory name conflict + base_dir = _sanitize_dir_name(raw_name) + dir_name = base_dir + counter = 1 + while dir_name in existing_dir_names: + dir_name = f"{base_dir}_{counter}" + counter += 1 + + dest_path = os.path.join(skills_dir, dir_name) + + # Defensive path-traversal check + if not os.path.realpath(dest_path).startswith( + os.path.realpath(skills_dir) + os.sep, + ): + logger.warning( + "Skill '%s' resolves outside skills_dir, skipping", + raw_name, + ) + continue + + try: + await asyncio.to_thread( + shutil.copytree, + skill_path, + dest_path, + dirs_exist_ok=False, + ) + except Exception as e: + logger.warning( + "Failed to copy skill '%s' from %s: %s", + raw_name, + skill_path, + str(e), + ) + continue + + logger.info( + "Copied skill '%s' (agent name: '%s') from %s to %s", + raw_name, + agent_name, + skill_path, + dest_path, + ) + + entry: _SkillEntry = {"hash": skill_hash, "skill_name": agent_name} + existing[dir_name] = entry + existing_hashes.add(skill_hash) + existing_agent_names.add(agent_name) + existing_dir_names.add(dir_name) + updated = True + + if updated: + skills_file["skills"] = existing + mtime = await self._backend.stat_mtime(skills_dir) + skills_file["skills_dir_mtime"] = ( + mtime if mtime is not None else 0.0 + ) + await self._save_skills_file(skills_dir, skills_file) + + self.is_alive = True + + async def get_instructions(self) -> str: + """Get the workspace instructions.""" + return self.instructions + + async def _load_skills_file(self, skills_dir: str) -> _SkillsFile: + """Load the .skills index file, returning an empty structure if absent. + + Args: + skills_dir (`str`): The skills directory path. + + Returns: + `_SkillsFile`: The parsed index, or a fresh empty structure. + """ + path = os.path.join(skills_dir, ".skills") + if not await self._backend.file_exists(path): + return {"skills_dir_mtime": 0.0, "skills": {}} + + try: + raw = await self._backend.read_file(path) + data = json.loads(raw.decode("utf-8")) + return _SkillsFile( + skills_dir_mtime=float(data.get("skills_dir_mtime", 0.0)), + skills=data.get("skills", {}), + ) + except Exception as e: + logger.warning("Failed to load .skills from %s: %s", path, str(e)) + return {"skills_dir_mtime": 0.0, "skills": {}} + + async def _save_skills_file( + self, + skills_dir: str, + data: _SkillsFile, + ) -> None: + """Persist the .skills index file. + + Args: + skills_dir (`str`): The skills directory path. + data (`_SkillsFile`): The index to write. + """ + path = os.path.join(skills_dir, ".skills") + try: + await self._backend.write_file( + path, + json.dumps(data, indent=2, ensure_ascii=False).encode("utf-8"), + ) + except Exception as e: + logger.warning("Failed to save .skills to %s: %s", path, str(e)) + + async def _validate_skill( + self, + skill_path: str, + ) -> tuple[str, str, str] | None: + """Validate if a skill path contains a valid SKILL.md file. + + Args: + skill_path (`str`): + The path to the skill directory. + + Returns: + `tuple[str, str, str] | None`: + A tuple of (name, description, skill_md_content) if valid, + None otherwise. + """ + skill_md_path = os.path.join(skill_path, "SKILL.md") + + try: + # Check if SKILL.md exists + if not await self._backend.file_exists(skill_md_path): + logger.warning( + "Invalid skill at %s: SKILL.md not found", + skill_path, + ) + return None + + # Read and parse SKILL.md + raw = await self._backend.read_file(skill_md_path) + content_str = raw.decode("utf-8") + + # Parse frontmatter + content = frontmatter.loads(content_str) + name = content.get("name") + description = content.get("description") + + if not name or not description: + logger.warning( + "Invalid skill at %s: SKILL.md missing required " + "fields (name or description)", + skill_path, + ) + return None + + return str(name), str(description), content_str + + except Exception as e: + logger.warning( + "Failed to validate skill at %s: %s", + skill_path, + str(e), + ) + return None + + async def _validate_and_hash_skill( + self, + skill_path: str, + ) -> tuple[str, str, str] | None: + """Validate a skill and compute its hash. + + Args: + skill_path (`str`): + The path to the skill directory. + + Returns: + `tuple[str, str, str] | None`: + A tuple of (skill_path, skill_name, skill_hash) if valid, + None otherwise. + """ + validation_result = await self._validate_skill(skill_path) + if validation_result is None: + return None + + skill_name, _, skill_md_content = validation_result + + # Compute hash + skill_hash = hashlib.sha256( + skill_md_content.encode("utf-8"), + ).hexdigest() + + return skill_path, skill_name, skill_hash + + async def _offload_data_block(self, data_block: DataBlock) -> DataBlock: + """Offload the data block by persisting it as local files. + + Uses the backend to write the decoded binary data, avoiding + embedding large base64-encoded data directly in the offload files, + keeping them lightweight and readable. + + Args: + data_block (`DataBlock`): + The data block with base64 source. + + Returns: + `DataBlock`: + A new data block with the same metadata but with the source + replaced by the local file path where the data is stored. + """ + if isinstance(data_block.source, URLSource): + return data_block + + # Use the full SHA-256 hex digest (256-bit) as the filename stem. + # A full hash collision is computationally infeasible, so an existing + # file with the same name is guaranteed to have identical content — + # no need to read and compare bytes. + hash_str = hashlib.sha256(data_block.source.data.encode()).hexdigest() + ext = mimetypes.guess_extension(data_block.source.media_type) or ".bin" + data_dir = os.path.join(self.workdir, "data") + path = os.path.join(data_dir, f"{hash_str}{ext}") + + # Reuse the existing file directly — same hash ⟹ same content. + if not await self._backend.file_exists(path): + await self._backend.write_file( + path, + base64.b64decode(data_block.source.data), + ) + + return DataBlock( + id=data_block.id, + name=data_block.name, + source=URLSource( + url=AnyUrl(Path(path).as_uri()), + media_type=data_block.source.media_type, + ), + ) + + async def offload_context( + self, + session_id: str, + msgs: list[Msg], + ) -> str: + """Offload the compressed messages into the local directory for + further processing. + + Args: + session_id (`str`): + The session id. + msgs (`list[Msg]`): + The messages to offload. + + Returns: + `str`: + The file path to the offloaded message. + """ + base = os.path.join(self.workdir, "sessions", session_id) + path = os.path.join(base, "context.jsonl") + + copied_msgs = deepcopy(msgs) + lines: list[str] = [] + for msg in copied_msgs: + if not isinstance(msg.content, str): + content = [] + for block in msg.content: + if isinstance(block, DataBlock) and isinstance( + block.source, + Base64Source, + ): + content.append(await self._offload_data_block(block)) + else: + content.append(block) + msg.content = content + lines.append(msg.model_dump_json()) + + payload = "\n".join(lines) + "\n" + + # Read existing content if any, then append. ``write_file`` + # creates parent directories, so no explicit mkdir is needed. + existing = b"" + try: + existing = await self._backend.read_file(path) + except (FileNotFoundError, OSError): + pass + await self._backend.write_file( + path, + existing + payload.encode("utf-8"), + ) + return path + + async def offload_tool_result( + self, + session_id: str, + tool_result: ToolResultBlock, + ) -> str: + """Offload the tool results into the local directory for agentic + retrieval. + + Args: + session_id (`str`): + The session id. + tool_result (`ToolResultBlock`): + The tool result. + + Returns: + `str`: + The file path to the offloaded tool results. + """ + base = os.path.join(self.workdir, "sessions", session_id) + path = os.path.join(base, f"tool_result-{tool_result.id}.txt") + + # Avoid filename conflict + index = 1 + while await self._backend.file_exists(path): + path = os.path.join( + base, + f"tool_result-{tool_result.id}({index}).txt", + ) + index += 1 + + parts: list[str] = [] + if isinstance(tool_result.output, str): + parts.append(tool_result.output) + else: + for block in tool_result.output: + if isinstance(block, TextBlock): + parts.append(block.text) + elif isinstance(block, DataBlock): + if isinstance(block.source, Base64Source): + data_block = await self._offload_data_block(block) + url = data_block.source.url + else: + url = block.source.url + parts.append( + f"", + ) + + await self._backend.write_file( + path, + "".join(parts).encode("utf-8"), + ) + return path + + async def close(self) -> None: + """Close every stateful MCP attached to this workspace. + + ``LocalWorkspace`` itself owns no resources (the workdir is + the persistence layer and is left untouched), but stdio / + stateful HTTP MCPs hold long-lived sessions that have to be + closed explicitly. Stateless HTTP MCPs are skipped — they + spin up an ad-hoc session per call and have nothing to close. + """ + async with self._mcp_lock: + for mcp in self._mcps: + if mcp.is_stateful and mcp.is_connected: + try: + await mcp.close() + except Exception as e: + logger.warning( + ( + "Failed to close MCP %r " + "when closing local workspace: %s" + ), + mcp.name, + e, + ) + self.is_alive = False + + async def reset(self) -> None: + """Return the workspace to an empty state. + + Closes and drops all MCPs (including the persisted ``.mcp``) + and deletes ``skills/``, ``sessions/``, and ``data/``. + ``default_mcps`` and ``skill_paths`` are not re-seeded. + """ + async with self._mcp_lock: + for mcp in self._mcps: + if mcp.is_stateful and mcp.is_connected: + try: + await mcp.close() + except Exception as e: + logger.warning( + "MCP %r close failed during reset: %s", + mcp.name, + e, + ) + self._mcps = [] + + mcp_file = os.path.join(self.workdir, ".mcp") + await self._backend.delete_path(mcp_file) + + async with self._skill_lock: + skills_path = os.path.join(self.workdir, "skills") + await self._backend.delete_path(skills_path) + + for sub in ("sessions", "data"): + path = os.path.join(self.workdir, sub) + await self._backend.delete_path(path) + + async def list_tools(self) -> list[ToolBase]: + """List all tools available in the workspace. + + Returns the six builtin tools (Bash, Read, Write, Edit, Grep, + Glob), each backed by the workspace's :class:`LocalBackend`. + """ + return [ + Bash(cwd=self.workdir, backend=self._backend), + Edit(backend=self._backend), + Glob(backend=self._backend), + Grep(backend=self._backend), + Read(backend=self._backend), + Write(backend=self._backend), + ] + + async def list_skills(self) -> list[Skill]: + """List all skills available in the workspace. + + The method uses the .skills index for agent-facing names, compares the + skills directory mtime to detect manual additions/removals since the + last write, and reconciles the index when a change is found. + + Returns: + `list[Skill]`: + A list of Skill objects found in the workspace. + """ + skills_dir = os.path.join(self.workdir, "skills") + async with self._skill_lock: + if not await self._backend.is_dir(skills_dir): + return [] + + skills_file = await self._load_skills_file(skills_dir) + current_mtime = await self._backend.stat_mtime(skills_dir) + if current_mtime is None: + current_mtime = 0.0 + + # Detect if the skills directory has changed since last indexing + if current_mtime != skills_file["skills_dir_mtime"]: + skills_file = await self._reconcile_skills_dir( + skills_dir, + skills_file, + current_mtime, + ) + + # Load skills from disk using the index for the agent-facing name + tasks = [ + self._load_single_skill( + os.path.join(skills_dir, dir_name), + entry["skill_name"], + ) + for dir_name, entry in skills_file["skills"].items() + ] + results = await asyncio.gather(*tasks, return_exceptions=True) + + skills: list = [] + for dir_name, result in zip(skills_file["skills"], results): + if isinstance(result, Exception): + logger.warning( + "Failed to load skill from %s: %s", + dir_name, + str(result), + ) + elif result is not None: + skills.append(result) + + return skills + + async def _reconcile_skills_dir( + self, + skills_dir: str, + skills_file: _SkillsFile, + current_mtime: float, + ) -> _SkillsFile: + """Reconcile the .skills index after the skills directory has changed. + + Handles: + - Manually deleted subdirectories: removed from the index. + - Manually added subdirectories: validated and added with conflict + resolution for both directory name and agent-facing skill name. + + Args: + skills_dir (`str`): Path to the skills directory. + skills_file (`_SkillsFile`): The current (stale) index. + current_mtime (`float`): The freshly-read mtime of skills_dir. + + Returns: + `_SkillsFile`: The updated index (also persisted to disk). + """ + existing: dict[str, _SkillEntry] = skills_file["skills"] + original_mtime = skills_file["skills_dir_mtime"] + + # Collect actual subdirectories on disk + entries = await self._backend.list_dir(skills_dir) + actual_dirs: set[str] = set() + for d in entries: + dir_path = os.path.join(skills_dir, d) + if await self._backend.is_dir(dir_path): + actual_dirs.add(d) + + indexed_dirs = set(existing.keys()) + + updated = False + + # Remove entries for directories that no longer exist + for removed in indexed_dirs - actual_dirs: + logger.info( + "Skill directory '%s' removed, updating index", + removed, + ) + del existing[removed] + updated = True + + # Add entries for directories not yet in the index + existing_agent_names: set[str] = { + e["skill_name"] for e in existing.values() + } + existing_hashes: set[str] = {e["hash"] for e in existing.values()} + + for new_dir in actual_dirs - indexed_dirs: + skill_path = os.path.join(skills_dir, new_dir) + result = await self._validate_and_hash_skill(skill_path) + if result is None: + continue + + _, raw_name, skill_hash = result + + if skill_hash in existing_hashes: + logger.info( + "Manually added skill '%s' already tracked by hash, " + "skipping", + new_dir, + ) + continue + + agent_name = raw_name + counter = 1 + while agent_name in existing_agent_names: + agent_name = f"{raw_name} ({counter})" + counter += 1 + + entry: _SkillEntry = {"hash": skill_hash, "skill_name": agent_name} + existing[new_dir] = entry + existing_agent_names.add(agent_name) + existing_hashes.add(skill_hash) + updated = True + logger.info( + "Manually added skill '%s' indexed as agent name '%s'", + new_dir, + agent_name, + ) + + skills_file["skills"] = existing + skills_file["skills_dir_mtime"] = current_mtime + + # Save if index changed OR if mtime needs updating + # (mtime change without index change means non-skill files were + # added/removed, we still need to record the new mtime to avoid + # re-reconciling on every list_skills call) + if updated or current_mtime != original_mtime: + await self._save_skills_file(skills_dir, skills_file) + + return skills_file + + async def _load_single_skill( + self, + skill_dir: str, + skill_name: str, + ) -> Skill | None: + """Load a single skill from disk using the agent-facing name from + the index. + + Args: + skill_dir (`str`): + The skill directory path containing SKILL.md. + skill_name (`str`): + The agent-facing name stored in the .skills index. + + Returns: + `Skill | None`: + A Skill object or None if the SKILL.md is missing/invalid. + """ + skill_md_path = os.path.join(skill_dir, "SKILL.md") + + try: + if not await self._backend.file_exists(skill_md_path): + return None + + updated_at = await self._backend.stat_mtime(skill_md_path) + if updated_at is None: + updated_at = 0.0 + + raw = await self._backend.read_file(skill_md_path) + content_str = raw.decode("utf-8") + content = frontmatter.loads(content_str) + + description = content.get("description") + if not description: + logger.warning( + "SKILL.md in %s is missing 'description'. Skipping.", + skill_dir, + ) + return None + + return Skill( + name=skill_name, + description=str(description), + dir=skill_dir, + markdown=content.content, + updated_at=updated_at, + ) + + except Exception as e: + logger.warning( + "Failed to load skill from %s: %s", + skill_dir, + str(e), + ) + return None + + async def list_mcps(self) -> list[MCPClient]: + """Return all MCP clients attached to this workspace.""" + return self._mcps + + async def _save_mcp_file(self) -> None: + """Persist the current MCP client list to ``.mcp`` in workdir.""" + mcp_file = os.path.join(self.workdir, ".mcp") + try: + # callers have lock. + await self._backend.write_file( + mcp_file, + json.dumps( + [m.model_dump() for m in self._mcps], + indent=2, + ensure_ascii=False, + ).encode("utf-8"), + ) + except Exception as e: + logger.warning("Failed to save .mcp to %s: %s", mcp_file, str(e)) + + async def add_mcp(self, mcp_client: MCPClient) -> None: + """Add an MCP client, connect it if stateful, and persist. + + Args: + mcp_client: The MCP client to add. + """ + async with self._mcp_lock: + if mcp_client.is_stateful and not mcp_client.is_connected: + await mcp_client.connect() + self._mcps.append(mcp_client) + await self._save_mcp_file() + + async def remove_mcp(self, name: str) -> None: + """Remove an MCP client by name, disconnecting it if stateful. + + Args: + name: The ``name`` field of the client to remove. + """ + async with self._mcp_lock: + for i, mcp in enumerate(self._mcps): + if mcp.name == name: + if mcp.is_stateful and mcp.is_connected: + await mcp.close() + self._mcps.pop(i) + await self._save_mcp_file() + return + logger.warning("MCP client %r not found in workspace", name) + + async def add_skill(self, skill_path: str) -> None: + """Add a skill to the workspace by copying from the given path. + + The skill directory must contain a valid ``SKILL.md`` file with + ``name`` and ``description`` frontmatter fields. Duplicate skills + (identified by the SHA-256 hash of ``SKILL.md``) are silently skipped. + Name and directory conflicts are resolved by appending a numeric + suffix. + + Args: + skill_path (`str`): + Absolute or relative path to the skill directory to copy. + + Raises: + ValueError: If the skill at ``skill_path`` is invalid (missing or + malformed ``SKILL.md``). + """ + skills_dir = os.path.join(self.workdir, "skills") + async with self._skill_lock: + os.makedirs(skills_dir, exist_ok=True) + + result = await self._validate_and_hash_skill(skill_path) + if result is None: + raise ValueError( + f"Invalid skill at {skill_path!r}: missing or malformed " + "SKILL.md (requires 'name' and 'description' fields).", + ) + + _, raw_name, skill_hash = result + + skills_file = await self._load_skills_file(skills_dir) + existing: dict[str, _SkillEntry] = skills_file["skills"] + + existing_hashes: set[str] = {e["hash"] for e in existing.values()} + if skill_hash in existing_hashes: + logger.info( + "Skill '%s' (hash: %s...) already exists, skipping", + raw_name, + skill_hash[:8], + ) + return + + existing_agent_names: set[str] = { + e["skill_name"] for e in existing.values() + } + existing_dir_names: set[str] = set(existing.keys()) + + # Resolve agent-facing name conflict + agent_name = raw_name + counter = 1 + while agent_name in existing_agent_names: + agent_name = f"{raw_name} ({counter})" + counter += 1 + + # Resolve directory name conflict + base_dir = _sanitize_dir_name(raw_name) + dir_name = base_dir + counter = 1 + while dir_name in existing_dir_names: + dir_name = f"{base_dir}_{counter}" + counter += 1 + + dest_path = os.path.join(skills_dir, dir_name) + + if not os.path.realpath(dest_path).startswith( + os.path.realpath(skills_dir) + os.sep, + ): + raise ValueError( + f"Skill path {skill_path!r} resolves outside skills_dir.", + ) + + await asyncio.to_thread( + shutil.copytree, + skill_path, + dest_path, + dirs_exist_ok=False, + ) + + logger.info( + "Copied skill '%s' (agent name: '%s') from %s to %s", + raw_name, + agent_name, + skill_path, + dest_path, + ) + + existing[dir_name] = {"hash": skill_hash, "skill_name": agent_name} + skills_file["skills"] = existing + mtime = await self._backend.stat_mtime(skills_dir) + skills_file["skills_dir_mtime"] = ( + mtime if mtime is not None else 0.0 + ) + await self._save_skills_file(skills_dir, skills_file) + + async def remove_skill(self, name: str) -> None: + """Remove a skill from the workspace by its agent-facing name. + + The skill directory is deleted from disk and the ``.skills`` index is + updated. If no skill with the given name is found, a warning is + logged and the method returns without error. + + Args: + name (`str`): + The agent-facing name of the skill to remove (as stored in the + ``.skills`` index, i.e. the ``name`` field from ``SKILL.md`` + possibly with a numeric suffix for de-duplication). + """ + skills_dir = os.path.join(self.workdir, "skills") + async with self._skill_lock: + if not await self._backend.is_dir(skills_dir): + logger.warning( + "Skills directory does not exist; cannot remove skill %r", + name, + ) + return + + skills_file = await self._load_skills_file(skills_dir) + existing: dict[str, _SkillEntry] = skills_file["skills"] + + target_dir: str | None = None + for dir_name, entry in existing.items(): + if entry["skill_name"] == name: + target_dir = dir_name + break + + if target_dir is None: + logger.warning("Skill %r not found in workspace", name) + return + + skill_dir_path = os.path.join(skills_dir, target_dir) + if await self._backend.is_dir(skill_dir_path): + await self._backend.delete_path(skill_dir_path) + logger.info( + "Removed skill '%s' from %s", + name, + skill_dir_path, + ) + else: + logger.warning( + ( + "Skill directory %r not found on disk; " + "removing index entry" + ), + skill_dir_path, + ) + + del existing[target_dir] + skills_file["skills"] = existing + mtime = await self._backend.stat_mtime(skills_dir) + skills_file["skills_dir_mtime"] = ( + mtime if mtime is not None else 0.0 + ) + await self._save_skills_file(skills_dir, skills_file) diff --git a/src/agentscope/workspace/_mcp_gateway/__init__.py b/src/agentscope/workspace/_mcp_gateway/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..808f156a5fb78c3273eab3bc6e73782347125d11 --- /dev/null +++ b/src/agentscope/workspace/_mcp_gateway/__init__.py @@ -0,0 +1,11 @@ +# -*- coding: utf-8 -*- +"""In-workspace MCP gateway package. + +The gateway is a single self-contained script that runs *inside* the +workspace environment (Docker / E2B). It is copied into the container +at image build time and executed by the workspace at startup. + +The script must remain importable without ``agentscope`` installed — +the host reads it as raw text via :mod:`importlib.resources` and ships +it into the container. +""" diff --git a/src/agentscope/workspace/_mcp_gateway/__main__.py b/src/agentscope/workspace/_mcp_gateway/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d0cc3a38c3738fdaabc7a37338493f621b04ee1 --- /dev/null +++ b/src/agentscope/workspace/_mcp_gateway/__main__.py @@ -0,0 +1,14 @@ +# -*- coding: utf-8 -*- +"""``python -m agentscope.workspace._mcp_gateway`` entry point. + +Inside the workspace container the gateway is launched as:: + + python -m agentscope.workspace._mcp_gateway --config --port + +so this module simply forwards to :func:`_mcp_gateway_app.main`. +""" + +from ._mcp_gateway_app import main + +if __name__ == "__main__": + main() diff --git a/src/agentscope/workspace/_mcp_gateway/_mcp_gateway_app.py b/src/agentscope/workspace/_mcp_gateway/_mcp_gateway_app.py new file mode 100644 index 0000000000000000000000000000000000000000..453bec12c304f6609420f7d46917e1343a42aba3 --- /dev/null +++ b/src/agentscope/workspace/_mcp_gateway/_mcp_gateway_app.py @@ -0,0 +1,242 @@ +# -*- coding: utf-8 -*- +"""In-workspace MCP gateway — FastAPI router over agentscope MCPClients. + +Runs *inside* the workspace environment as a standalone script +(``python /path/to/_mcp_gateway_app.py``). Reads ``--config`` JSON, +instantiates one :class:`agentscope.mcp.MCPClient` per configured server, +and exposes per-server HTTP endpoints. Each call is forwarded to the +underlying ``MCPClient`` (which owns the upstream session). + +The script uses an absolute import for ``agentscope.mcp`` (rather than +a package-relative import) so it can be invoked directly without +loading ``agentscope.workspace.__init__`` — the latter eagerly imports +heavy modules (skill, tool, …) that are unnecessary for the gateway +and would force their dependencies into the in-container venv. + +Endpoints +--------- + + GET /health # liveness, no auth + GET /mcps # [{name, tools}, ...] + POST /mcps # body: MCPClient.model_dump() + DELETE /mcps/{name} + GET /mcps/{name}/tools # upstream tool schemas + POST /mcps/{name}/tools/{tool} # body: {arguments: {...}} + +Auth: every endpoint except ``/health`` requires +``Authorization: Bearer `` when a token is configured. + +Config schema:: + + { + "token": "bearer-token", + "servers": [, ...] + } +""" + +import argparse +import asyncio +import json +from typing import Any + +from fastapi import Depends, FastAPI, HTTPException, Request +from fastapi.responses import PlainTextResponse + +from agentscope.mcp import MCPClient + + +# ── gateway state ────────────────────────────────────────────────── + + +class _State: + """Mutable runtime state shared by FastAPI routes.""" + + def __init__(self) -> None: + self.clients: dict[str, MCPClient] = {} + self.token: str = "" + self.lock = asyncio.Lock() + + +def _make_auth_dep(state: _State) -> Any: + """Build a Bearer-token auth dependency closed over the state. + + No-op when ``state.token`` is empty. + """ + + async def _auth(request: Request) -> None: + if not state.token: + return + header = request.headers.get("authorization", "") + if header != f"Bearer {state.token}": + raise HTTPException(status_code=401, detail="unauthorized") + + return _auth + + +# ── client construction ─────────────────────────────────────────── + + +async def _build_client(spec: dict[str, Any]) -> MCPClient: + """Validate a config / request body into an :class:`MCPClient`, + then connect if stateful so subsequent ``list_raw_tools`` / + ``get_tool`` work without re-spawning the upstream session. + """ + client = MCPClient.model_validate(spec) + if client.is_stateful: + await client.connect() + # Prime the tool cache so /mcps/{name}/tools is cheap and stable. + await client.list_raw_tools() + return client + + +# ── FastAPI app ──────────────────────────────────────────────────── + + +def _build_app(state: _State) -> FastAPI: + """Build the FastAPI app with all routes wired against ``state``.""" + app = FastAPI(title="agentscope-workspace-mcp-gateway") + auth = Depends(_make_auth_dep(state)) + + @app.get("/health") + async def _health() -> PlainTextResponse: + return PlainTextResponse("ok") + + @app.get("/mcps", dependencies=[auth]) + async def _list_mcps() -> list[dict[str, Any]]: + # Dump the full MCPClient field set so the host can rebuild + # `GatewayMCPClient.model_validate(spec)` losslessly. + return [c.model_dump(mode="json") for c in state.clients.values()] + + @app.post("/mcps", dependencies=[auth]) + async def _add_mcp(request: Request) -> dict[str, Any]: + body = await request.json() + name = body.get("name", "") + if not name: + raise HTTPException(400, "name required") + async with state.lock: + if name in state.clients: + raise HTTPException(409, f"{name!r} already exists") + try: + client = await _build_client(body) + except HTTPException: + raise + except Exception as e: # noqa: BLE001 + raise HTTPException( + 500, + f"connect failed: {e}", + ) from e + state.clients[name] = client + return {"ok": True} + + @app.delete("/mcps/{name}", dependencies=[auth]) + async def _remove_mcp(name: str) -> dict[str, Any]: + async with state.lock: + client = state.clients.pop(name, None) + if client is None: + raise HTTPException(404, f"{name!r} not found") + if client.is_stateful and client.is_connected: + await client.close() + return {"ok": True} + + @app.get("/mcps/{name}/tools", dependencies=[auth]) + async def _list_tools(name: str) -> list[dict[str, Any]]: + client = state.clients.get(name) + if client is None: + raise HTTPException(404, f"{name!r} not found") + # Send raw mcp.types.Tool over the wire so the host-side + # GatewayMCPClient can re-wrap them via the standard MCPClient + # path (preserves inputSchema, annotations.readOnlyHint, ...). + raw = await client.list_raw_tools() + return [t.model_dump(mode="json") for t in raw] + + @app.post("/mcps/{name}/tools/{tool}", dependencies=[auth]) + async def _call_tool( + name: str, + tool: str, + request: Request, + ) -> dict[str, Any]: + client = state.clients.get(name) + if client is None: + raise HTTPException(404, f"{name!r} not found") + body = await request.json() + arguments = body.get("arguments") or {} + try: + tool_obj = await client.get_tool(tool) + chunk = await tool_obj(**arguments) + except ValueError as e: + raise HTTPException(404, str(e)) from e + except Exception as e: # noqa: BLE001 + raise HTTPException(500, str(e)) from e + # ToolChunk is a pydantic model — let host reconstruct it. + return {"chunk": chunk.model_dump(mode="json")} + + return app + + +# ── lifecycle ────────────────────────────────────────────────────── + + +async def _connect_initial( + state: _State, + server_cfgs: list[dict[str, Any]], +) -> None: + """Connect every server listed in the static config file.""" + for cfg in server_cfgs: + client = await _build_client(cfg) + if client.name in state.clients: + if client.is_stateful and client.is_connected: + await client.close() + raise ValueError( + f"Duplicated server name in config: {client.name!r}", + ) + state.clients[client.name] = client + print(f"[gateway] connected {client.name!r}", flush=True) + + +async def _run(config_path: str, port: int) -> None: + """Read config, connect upstreams, start uvicorn, clean up on exit.""" + with open(config_path, encoding="utf-8") as f: + config = json.load(f) + + state = _State() + state.token = config.get("token", "") or "" + await _connect_initial(state, config.get("servers", []) or []) + + app = _build_app(state) + print( + f"[gateway] serving {len(state.clients)} MCPs on :{port}", + flush=True, + ) + + import uvicorn + + uvi_cfg = uvicorn.Config( + app, + host="0.0.0.0", # noqa: S104 — gateway listens inside container + port=port, + log_level="warning", + ) + server = uvicorn.Server(uvi_cfg) + try: + await server.serve() + finally: + for client in list(state.clients.values()): + if client.is_stateful and client.is_connected: + await client.close() + + +def main() -> None: + """CLI entry point — invoked via + ``python -m agentscope.workspace._mcp_gateway``. + """ + parser = argparse.ArgumentParser( + description="In-workspace MCP gateway (FastAPI)", + ) + parser.add_argument("--config", required=True) + parser.add_argument("--port", type=int, default=5600) + args = parser.parse_args() + asyncio.run(_run(args.config, args.port)) + + +if __name__ == "__main__": + main() diff --git a/src/agentscope/workspace/_offload_protocol.py b/src/agentscope/workspace/_offload_protocol.py new file mode 100644 index 0000000000000000000000000000000000000000..dbd819934b153d418cb9ccc2169d931ab67c1862 --- /dev/null +++ b/src/agentscope/workspace/_offload_protocol.py @@ -0,0 +1,45 @@ +# -*- coding: utf-8 -*- +"""The offload protocol.""" +from typing import Protocol + +from ..message import Msg, ToolResultBlock + + +class Offloader(Protocol): + """The offloader protocol.""" + + async def offload_context( + self, + session_id: str, + msgs: list[Msg], + ) -> str: + """Offload compressed context to workspace-accessible storage. + + Args: + session_id (`str`): + The session id. + msgs (`list[Msg]`): + The messages to offload. + + Returns: + `str`: + The offloaded context reference. + """ + + async def offload_tool_result( + self, + session_id: str, + tool_result: ToolResultBlock, + ) -> str: + """Offload a tool result to workspace-accessible storage. + + Args: + session_id (`str`): + The session id. + tool_result (`ToolResultBlock`): + The tool result. + + Returns: + `str`: + The offloaded context reference. + """ diff --git a/src/agentscope/workspace/_utils.py b/src/agentscope/workspace/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..53151467cd0347872bf09b14cedcc41f93202d24 --- /dev/null +++ b/src/agentscope/workspace/_utils.py @@ -0,0 +1,184 @@ +# -*- coding: utf-8 -*- +"""Host-side helpers shared by Docker + E2B backends. + +Pure functions for detecting the local ``agentscope`` install (released +vs dev), iterating the source tree with a stable ignore set, and +reading the gateway script bundled with the package. No Docker / E2B +SDK dependency lives here — both backends import the same helpers so +their install logic stays byte-for-byte in sync. + +All names in this module are package-private (leading underscore). +External code should not import from here directly; the two backends +that consume these helpers live next to it under +``agentscope.workspace``. +""" + +import importlib.resources as _res +from pathlib import Path + +# ── shared constants ─────────────────────────────────────────────── + +#: Minimum Python packages the gateway script needs at runtime. +#: Both Docker (image build) and E2B (sandbox bootstrap) install this +#: same tuple into the gateway venv before adding ``agentscope`` itself. +_GATEWAY_BASE_REQUIREMENTS: tuple[str, ...] = ( + "mcp", + "uvicorn", + "fastapi", +) + +#: Basename set excluded when packaging the agentscope source tree +#: into a build / bootstrap context (dev-mode only — released installs +#: pull from PyPI and never copy the tree). +_SOURCE_IGNORE_NAMES: frozenset[str] = frozenset( + { + "__pycache__", + "node_modules", + "build", + "dist", + "venv", + "workdir", + "examples", + "tests", + "docs", + "assets", + "scripts", + "dump.rdb", + "uv.lock", + }, +) + + +def _is_source_ignored(name: str) -> bool: + """Whether a single basename should be excluded from the source payload. + + Args: + name (`str`): + Basename to check (e.g. ``"__pycache__"``, ``".git"``, + ``"foo.pyc"``). + + Returns: + `bool`: + ``True`` if the name matches the ignore set / hidden-file + rule / cache-extension rule; ``False`` otherwise. + """ + if name.startswith("."): + return True + if name in _SOURCE_IGNORE_NAMES: + return True + return name.endswith(".pyc") or name.endswith(".egg-info") + + +# ── agentscope install detection ─────────────────────────────────── + + +def _agentscope_module_path() -> Path: + """Return the filesystem path of the imported ``agentscope`` package. + + Returns: + `Path`: + The directory containing ``agentscope/__init__.py``. + """ + import agentscope # local import — keeps module import cheap + + file = getattr(agentscope, "__file__", None) + if not file: + raise RuntimeError( + "agentscope has no __file__ attribute; cannot locate package", + ) + return Path(file).resolve().parent + + +def _is_released_install() -> bool: + """Return ``True`` if the imported ``agentscope`` lives in site-packages. + + Used to pick between PyPI install (released) and source-tree + upload (dev) when provisioning the gateway venv inside a + container or sandbox. + """ + parts = _agentscope_module_path().parts + return "site-packages" in parts or "dist-packages" in parts + + +def _agentscope_version() -> str: + """Return the installed ``agentscope`` version string. + + Falls back to :func:`importlib.metadata.version` when the package + has no ``__version__`` attribute. + """ + import agentscope + + version = getattr(agentscope, "__version__", None) + if not version: + try: + from importlib.metadata import version as _v + + version = _v("agentscope") + except Exception as e: # noqa: BLE001 + raise RuntimeError( + "cannot determine agentscope version", + ) from e + return version + + +def _agentscope_source_root() -> Path: + """Locate the project root containing ``pyproject.toml`` + ``src/``. + + Only valid in dev mode. Walks up from the package directory until + a ``pyproject.toml`` is found alongside a ``src/`` (or + ``agentscope/``) directory. + + Returns: + `Path`: + The project root path. + """ + pkg = _agentscope_module_path() + for parent in [pkg, *pkg.parents]: + if (parent / "pyproject.toml").is_file() and ( + (parent / "src").is_dir() or (parent / "agentscope").is_dir() + ): + return parent + raise RuntimeError( + f"cannot locate agentscope project root from {pkg}", + ) + + +# ── gateway script ───────────────────────────────────────────────── + + +def _read_gateway_script_bytes() -> bytes: + """Read the standalone gateway script as bytes via ``importlib.resources``. + + The script ships at + ``agentscope/workspace/_mcp_gateway/_mcp_gateway_app.py``. Both + backends copy it to a fixed in-container / in-sandbox path so the + launch command can invoke it directly, avoiding ``python -m`` and + the heavy ``agentscope.workspace.__init__`` import graph. + """ + return ( + _res.files("agentscope.workspace._mcp_gateway") + .joinpath("_mcp_gateway_app.py") + .read_bytes() + ) + + +# ── builtin tool helper scripts ─────────────────────────────────── + + +def _read_glob_helper_bytes() -> bytes: + """Read the standalone glob helper script as bytes. + + The script ships at + ``agentscope/tool/_builtin/_scripts/_glob_helper.py``. Both Docker + and E2B backends copy it into the workspace so the :class:`Glob` + tool can invoke it uniformly via ``exec_shell``. + + Returns: + `bytes`: + The raw contents of the ``_glob_helper.py`` script. + """ + return ( + _res.files("agentscope.tool._builtin._scripts") + .joinpath("_glob_helper.py") + .read_bytes() + ) diff --git a/start.sh b/start.sh new file mode 100644 index 0000000000000000000000000000000000000000..f5f6b54f9297e92691cb6e55aca08ecf44026558 --- /dev/null +++ b/start.sh @@ -0,0 +1,3 @@ +#!/bin/bash +redis-server --daemonize yes +python hf_app.py diff --git a/tests/agent_basic_test.py b/tests/agent_basic_test.py new file mode 100644 index 0000000000000000000000000000000000000000..90c893b166ec4676fbf2793047442220bff6d140 --- /dev/null +++ b/tests/agent_basic_test.py @@ -0,0 +1,1369 @@ +# -*- coding: utf-8 -*- +"""The basic test of the agent class.""" +from typing import Any +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import AnyString, MockModel + +from agentscope.agent import Agent +from agentscope.model import ChatResponse +from agentscope.tool import ( + ToolBase, + Toolkit, + ToolChunk, +) +from agentscope.permission import ( + PermissionDecision, + PermissionBehavior, + PermissionContext, +) +from agentscope.message import TextBlock, ToolCallBlock, UserMsg + + +class MockSequentialTool(ToolBase): + """A mock tool that is not concurrency safe (sequential execution).""" + + name: str = "mock_sequential_tool" + description: str = "A mock sequential tool for testing" + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "input": {"type": "string", "description": "Input string"}, + }, + "required": ["input"], + } + is_concurrency_safe: bool = False + is_read_only: bool = True + is_external_tool: bool = False + is_mcp: bool = False + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permissions for the tool usage.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + decision_reason="Mock tool always allows", + message="Mock tool always allows", + ) + + # pylint: disable=redefined-builtin + async def __call__(self, input: str, **kwargs: Any) -> ToolChunk: + """Execute the tool.""" + return ToolChunk( + content=[TextBlock(text=f"Sequential result: {input}")], + ) + + +class MockConcurrentTool(ToolBase): + """A mock tool that is concurrency safe (concurrent execution).""" + + name: str = "mock_concurrent_tool" + description: str = "A mock concurrent tool for testing" + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "input": {"type": "string", "description": "Input string"}, + }, + "required": ["input"], + } + is_concurrency_safe: bool = True + is_read_only: bool = True + is_external_tool: bool = False + is_mcp: bool = False + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permissions for the tool usage.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + decision_reason="Mock tool always allows", + message="Mock tool always allows", + ) + + # pylint: disable=redefined-builtin + async def __call__(self, input: str, **kwargs: Any) -> ToolChunk: + """Execute the tool.""" + return ToolChunk( + content=[TextBlock(text=f"Concurrent result: {input}")], + ) + + +class AgentBasicTest(IsolatedAsyncioTestCase): + """The basic test of the agent class.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.model = MockModel() + self.agent = Agent( + name="Friday", + system_prompt="You are a helpful assistant.", + model=self.model, + toolkit=Toolkit(), + ) + + def _get_event_base(self, reply_id: str) -> dict: + """Get the dict with the basic fields for event assertion.""" + return { + "id": AnyString(), + "created_at": AnyString(), + "metadata": {}, + "reply_id": reply_id, + } + + def _get_msg_base(self) -> dict: + """Get the dict with the basic fields for message assertion.""" + return { + "id": AnyString(), + "created_at": AnyString(), + "finished_at": None, + "metadata": {}, + "name": "Friday", + "role": "assistant", + "usage": None, + } + + async def test_default_configs_are_not_shared_between_agents( + self, + ) -> None: + """Agents created with defaults own independent config objects.""" + agent_1 = Agent( + name="agent-1", + system_prompt="You are agent 1.", + model=MockModel(), + ) + agent_2 = Agent( + name="agent-2", + system_prompt="You are agent 2.", + model=MockModel(), + ) + + self.assertIsNot(agent_1.model_config, agent_2.model_config) + self.assertIsNot(agent_1.context_config, agent_2.context_config) + self.assertIsNot(agent_1.react_config, agent_2.react_config) + + agent_1.model_config.max_retries = 3 + agent_1.context_config.tool_result_limit = 123 + agent_1.react_config.max_iters = 2 + + self.assertNotEqual( + agent_1.model_config.max_retries, + agent_2.model_config.max_retries, + ) + self.assertNotEqual( + agent_1.context_config.tool_result_limit, + agent_2.context_config.tool_result_limit, + ) + self.assertNotEqual( + agent_1.react_config.max_iters, + agent_2.react_config.max_iters, + ) + + async def test_streaming_reasoning(self) -> None: + """Test the streaming model inference without tool calls generated, + only text in model response. + + Test both the reply and replyStream interfaces. + """ + # Set up mock responses for streaming + self.model.set_responses( + [ + [ + ChatResponse( + content=[TextBlock(text="Hello")], + is_last=False, + ), + ChatResponse( + content=[TextBlock(text=" world")], + is_last=False, + ), + ChatResponse(content=[TextBlock(text="!")], is_last=False), + ChatResponse( + content=[TextBlock(text="Hello world!")], + is_last=True, + ), + ], + ], + ) + + # Test replyStream interface + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content="Hi"), + ): + events.append(event.model_dump()) + + # Verify events + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + + expected_events = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + { + "type": "MODEL_CALL_START", + "model_name": "mock-model", + }, + { + "type": "TEXT_BLOCK_START", + "block_id": AnyString(), + }, + { + "type": "TEXT_BLOCK_DELTA", + "block_id": AnyString(), + "delta": "Hello", + }, + { + "type": "TEXT_BLOCK_DELTA", + "block_id": AnyString(), + "delta": " world", + }, + { + "type": "TEXT_BLOCK_DELTA", + "block_id": AnyString(), + "delta": "!", + }, + { + "type": "TEXT_BLOCK_END", + "block_id": AnyString(), + }, + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + { + "type": "REPLY_END", + "session_id": session_id, + }, + ] + + basic_dict = self._get_event_base(reply_id) + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + # Assert context after reply_stream + msg_base = self._get_msg_base() + expected_context = [ + { + **msg_base, + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hi", + }, + ], + "finished_at": AnyString(), + }, + { + **msg_base, + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hello world!", + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + self.assertListEqual(context_dicts, expected_context) + + # Test reply interface + self.model.cnt = 0 # Reset mock model response index + msg = await self.agent.reply(UserMsg(name="user", content="Hi again")) + self.assertDictEqual( + msg.model_dump(), + { + "name": "Friday", + "role": "assistant", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hello world!", + }, + ], + "created_at": AnyString(), + "finished_at": None, + "id": AnyString(), + "metadata": {}, + "usage": None, + }, + ) + + # Assert context after reply + expected_context_after_reply = [ + { + **msg_base, + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hi", + }, + ], + "metadata": {}, + "finished_at": AnyString(), + }, + { + **msg_base, + "name": "Friday", + "role": "assistant", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hello world!", + }, + ], + "metadata": {}, + }, + { + **msg_base, + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hi again", + }, + ], + "metadata": {}, + "finished_at": AnyString(), + }, + { + **msg_base, + "name": "Friday", + "role": "assistant", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hello world!", + }, + ], + "metadata": {}, + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + self.assertListEqual(context_dicts, expected_context_after_reply) + + async def test_non_streaming_reasoning(self) -> None: + """Test the non-streaming model inference without tool calls generated, + only text in model response. + + Test both the reply and replyStream interfaces. + """ + # Set up mock response for non-streaming + self.model.set_responses( + [ + ChatResponse( + content=[TextBlock(text="Hello world!")], + is_last=True, + usage=None, + ), + ], + ) + + # Test replyStream interface + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content="Hi"), + ): + events.append(event.model_dump()) + + # Verify events + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + + expected_events = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + { + "type": "MODEL_CALL_START", + "model_name": "mock-model", + }, + { + "type": "TEXT_BLOCK_START", + "block_id": AnyString(), + }, + { + "type": "TEXT_BLOCK_DELTA", + "block_id": AnyString(), + "delta": "Hello world!", + }, + { + "type": "TEXT_BLOCK_END", + "block_id": AnyString(), + }, + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + { + "type": "REPLY_END", + "session_id": session_id, + }, + ] + + basic_dict = self._get_event_base(reply_id) + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + # Assert context after reply_stream + msg_base = self._get_msg_base() + expected_context = [ + { + **msg_base, + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hi", + }, + ], + "metadata": {}, + "finished_at": AnyString(), + }, + { + **msg_base, + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hello world!", + }, + ], + "metadata": {}, + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + self.assertListEqual(context_dicts, expected_context) + + # Test reply interface + self.model.set_responses( + [ + ChatResponse( + content=[TextBlock(text="Hello world!")], + is_last=True, + usage=None, + ), + ], + ) + + msg = await self.agent.reply(UserMsg(name="user", content="Hi again")) + self.assertDictEqual( + msg.model_dump(), + { + "name": "Friday", + "role": "assistant", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hello world!", + }, + ], + "created_at": AnyString(), + "finished_at": None, + "id": AnyString(), + "metadata": {}, + "usage": None, + }, + ) + + # Assert context after reply + expected_context_after_reply = [ + { + **msg_base, + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hi", + }, + ], + "metadata": {}, + "finished_at": AnyString(), + }, + { + **msg_base, + "name": "Friday", + "role": "assistant", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hello world!", + }, + ], + "metadata": {}, + }, + { + **msg_base, + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hi again", + }, + ], + "metadata": {}, + "finished_at": AnyString(), + }, + { + **msg_base, + "name": "Friday", + "role": "assistant", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hello world!", + }, + ], + "metadata": {}, + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + self.assertListEqual(context_dicts, expected_context_after_reply) + + async def test_streaming_sequential_tool_calls(self) -> None: + """Test the streaming model inference with tool calls generated. + + Test only the replyStream interface. The `is_concurrent_safe` of the + registered tools should be False to make sure the tool calls are + executed sequentially. + + To assert: + 1. The events (by assert the dict generated by model_dump) + 2. The final reply message + 3. The agent state (Before and after replyStream) + """ + # Register sequential tools + seq_tool = MockSequentialTool() + self.agent.toolkit = Toolkit(tools=[seq_tool]) + + # Create tool call IDs + tool_call_id_1 = "tool_call_1" + tool_call_id_2 = "tool_call_2" + + # Set up mock responses with tool calls + text_block = TextBlock(text="I'll call the tool") + tool_call_1 = ToolCallBlock( + id=tool_call_id_1, + name="mock_sequential_tool", + input='{"input": "test1"}', + ) + tool_call_1_part1 = ToolCallBlock( + id=tool_call_id_1, + name="mock_sequential_tool", + input='{"input": ', + ) + tool_call_1_part2 = ToolCallBlock( + id=tool_call_id_1, + name="mock_sequential_tool", + input='"test1"}', + ) + tool_call_2 = ToolCallBlock( + id=tool_call_id_2, + name="mock_sequential_tool", + input='{"input": "test2"}', + ) + self.model.set_responses( + [ + [ + ChatResponse( + content=[text_block, tool_call_1_part1], + is_last=False, + ), + ChatResponse( + content=[tool_call_1_part2], + is_last=False, + ), + ChatResponse( + content=[tool_call_2], + is_last=False, + ), + ChatResponse( + content=[text_block, tool_call_1, tool_call_2], + is_last=True, + ), + ], + [ + ChatResponse( + content=[TextBlock(text="ended")], + is_last=False, + ), + ChatResponse( + content=[TextBlock(text="ended")], + is_last=True, + ), + ], + ], + ) + + # Collect all events + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content="Test"), + ): + events.append(event.model_dump(mode="json")) + + # Verify events + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + + # Expected events for sequential tool calls + expected_events = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + {"type": "TEXT_BLOCK_START", "block_id": AnyString()}, + { + "type": "TEXT_BLOCK_DELTA", + "block_id": AnyString(), + "delta": "I'll call the tool", + }, + { + "type": "TOOL_CALL_START", + "tool_call_id": tool_call_id_1, + "tool_call_name": "mock_sequential_tool", + }, + { + "type": "TOOL_CALL_DELTA", + "tool_call_id": tool_call_id_1, + "delta": '{"input": ', + }, + {"type": "TEXT_BLOCK_END", "block_id": AnyString()}, + { + "type": "TOOL_CALL_DELTA", + "tool_call_id": tool_call_id_1, + "delta": '"test1"}', + }, + { + "type": "TOOL_CALL_START", + "tool_call_id": tool_call_id_2, + "tool_call_name": "mock_sequential_tool", + }, + { + "type": "TOOL_CALL_DELTA", + "tool_call_id": tool_call_id_2, + "delta": '{"input": "test2"}', + }, + {"type": "TOOL_CALL_END", "tool_call_id": tool_call_id_1}, + {"type": "TOOL_CALL_END", "tool_call_id": tool_call_id_2}, + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + { + "type": "TOOL_RESULT_START", + "tool_call_id": tool_call_id_1, + "tool_call_name": "mock_sequential_tool", + }, + { + "type": "TOOL_RESULT_TEXT_DELTA", + "tool_call_id": tool_call_id_1, + "delta": "Sequential result: test1", + }, + { + "type": "TOOL_RESULT_END", + "tool_call_id": tool_call_id_1, + "state": "success", + }, + { + "type": "TOOL_RESULT_START", + "tool_call_id": tool_call_id_2, + "tool_call_name": "mock_sequential_tool", + }, + { + "type": "TOOL_RESULT_TEXT_DELTA", + "tool_call_id": tool_call_id_2, + "delta": "Sequential result: test2", + }, + { + "type": "TOOL_RESULT_END", + "tool_call_id": tool_call_id_2, + "state": "success", + }, + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + {"type": "TEXT_BLOCK_START", "block_id": AnyString()}, + { + "type": "TEXT_BLOCK_DELTA", + "block_id": AnyString(), + "delta": "ended", + }, + {"type": "TEXT_BLOCK_END", "block_id": AnyString()}, + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + {"type": "REPLY_END", "session_id": session_id}, + ] + + basic_dict = self._get_event_base(reply_id) + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + # Assert context after reply_stream + msg_base = self._get_msg_base() + expected_context = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Test", + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "I'll call the tool", + }, + { + "type": "tool_call", + "id": tool_call_id_1, + "name": "mock_sequential_tool", + "input": '{"input": "test1"}', + "state": "finished", + "suggested_rules": [], + }, + { + "type": "tool_call", + "id": tool_call_id_2, + "name": "mock_sequential_tool", + "input": '{"input": "test2"}', + "state": "finished", + "suggested_rules": [], + }, + { + "type": "tool_result", + "id": AnyString(), + "output": [ + { + "type": "text", + "id": AnyString(), + "text": "Sequential result: test1", + }, + ], + "name": "mock_sequential_tool", + "state": "success", + "metadata": {}, + }, + { + "type": "tool_result", + "id": AnyString(), + "output": [ + { + "type": "text", + "id": AnyString(), + "text": "Sequential result: test2", + }, + ], + "name": "mock_sequential_tool", + "state": "success", + "metadata": {}, + }, + { + "type": "text", + "id": AnyString(), + "text": "ended", + }, + ], + }, + ] + + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + async def test_streaming_concurrent_tool_calls(self) -> None: + """Test the streaming model inference with tool calls generated. + + Test only the replyStream interface. The `is_concurrent_safe` of the + registered tools should be True to make sure the tool calls are + executed concurrently. + + To assert: + 1. The events (by assert the dict generated by model_dump) + 2. The final reply message + 3. The agent state (Before and after replyStream) + """ + # Register concurrent tools + conc_tool = MockConcurrentTool() + + self.agent.toolkit = Toolkit( + tools=[conc_tool], + ) + + # Create tool call IDs + tool_call_id_1 = "tool_call_1" + tool_call_id_2 = "tool_call_2" + + # Set up mock responses with tool calls + self.model.set_responses( + [ + [ + ChatResponse( + content=[ + ToolCallBlock( + id=tool_call_id_1, + name="mock_concurrent_tool", + input='{"input": "test1"}', + ), + ToolCallBlock( + id=tool_call_id_2, + name="mock_concurrent_tool", + input='{"input": "test2"}', + ), + ], + is_last=True, + usage=None, + ), + ], + [ + ChatResponse( + content=[TextBlock(text="All done")], + is_last=False, + ), + ChatResponse( + content=[TextBlock(text="All done")], + is_last=True, + ), + ], + ], + ) + + # Collect all events + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content="Test"), + ): + events.append(event.model_dump()) + + # Verify events + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + + # For concurrent execution, the order of tool results may vary + # Split events into: prefix (before concurrent), concurrent part, + # suffix (after concurrent) + + # Expected prefix events (before tool execution) + expected_prefix = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + ] + + # Expected concurrent events (order may vary) + expected_concurrent = [ + { + "type": "TOOL_RESULT_START", + "tool_call_id": tool_call_id_1, + "tool_call_name": "mock_concurrent_tool", + }, + { + "type": "TOOL_RESULT_TEXT_DELTA", + "tool_call_id": tool_call_id_1, + "delta": "Concurrent result: test1", + }, + { + "type": "TOOL_RESULT_END", + "tool_call_id": tool_call_id_1, + "state": "success", + }, + { + "type": "TOOL_RESULT_START", + "tool_call_id": tool_call_id_2, + "tool_call_name": "mock_concurrent_tool", + }, + { + "type": "TOOL_RESULT_TEXT_DELTA", + "tool_call_id": tool_call_id_2, + "delta": "Concurrent result: test2", + }, + { + "type": "TOOL_RESULT_END", + "tool_call_id": tool_call_id_2, + "state": "success", + }, + ] + + # Expected suffix events (final model call with pure text) + expected_suffix = [ + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + {"type": "TEXT_BLOCK_START", "block_id": AnyString()}, + { + "type": "TEXT_BLOCK_DELTA", + "block_id": AnyString(), + "delta": "All done", + }, + {"type": "TEXT_BLOCK_END", "block_id": AnyString()}, + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + {"type": "REPLY_END", "session_id": session_id}, + ] + + # Assert prefix events (fixed order) + basic_dict = self._get_event_base(reply_id) + prefix_len = len(expected_prefix) + self.assertListEqual( + events[:prefix_len], + [{**basic_dict, **_} for _ in expected_prefix], + ) + + # Assert concurrent events (order may vary) + concurrent_len = len(expected_concurrent) + concurrent_events = events[prefix_len : prefix_len + concurrent_len] + + # Check length matches + self.assertEqual(len(concurrent_events), len(expected_concurrent)) + + # Check each expected event is in the actual events + for expected_event in expected_concurrent: + expected_with_base = {**basic_dict, **expected_event} + self.assertIn(expected_with_base, concurrent_events) + + # Assert suffix events (fixed order) + suffix_events = events[prefix_len + concurrent_len :] + self.assertListEqual( + suffix_events, + [{**basic_dict, **_} for _ in expected_suffix], + ) + + # Assert context after reply_stream + msg_base = self._get_msg_base() + expected_context = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Test", + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": tool_call_id_1, + "name": "mock_concurrent_tool", + "input": '{"input": "test1"}', + "state": "finished", + "suggested_rules": [], + }, + { + "type": "tool_call", + "id": tool_call_id_2, + "name": "mock_concurrent_tool", + "input": '{"input": "test2"}', + "state": "finished", + "suggested_rules": [], + }, + { + "type": "tool_result", + "id": AnyString(), + "output": [ + { + "type": "text", + "id": AnyString(), + "text": "Concurrent result: test1", + }, + ], + "name": "mock_concurrent_tool", + "state": "success", + "metadata": {}, + }, + { + "type": "tool_result", + "id": AnyString(), + "output": [ + { + "type": "text", + "id": AnyString(), + "text": "Concurrent result: test2", + }, + ], + "name": "mock_concurrent_tool", + "state": "success", + "metadata": {}, + }, + { + "type": "text", + "id": AnyString(), + "text": "All done", + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + async def test_streaming_mixed_tool_calls(self) -> None: + """Test the streaming model inference with both sequential and + concurrent tool calls generated. + + Test only the replyStream interface. + + To assert: + 1. The events (by assert the dict generated by model_dump) + 2. The final reply message + 3. The agent state (Before and after replyStream) + """ + # Register both sequential and concurrent tools + seq_tool = MockSequentialTool() + conc_tool = MockConcurrentTool() + + self.agent.toolkit = Toolkit( + tools=[seq_tool, conc_tool], + ) + + # Create tool call IDs + tool_call_id_1 = "tool_call_1" + tool_call_id_2 = "tool_call_2" + tool_call_id_3 = "tool_call_3" + + # Set up mock responses with mixed tool calls + self.model.set_responses( + [ + [ + ChatResponse( + content=[ + ToolCallBlock( + id=tool_call_id_1, + name="mock_sequential_tool", + input='{"input": "seq1"}', + ), + ToolCallBlock( + id=tool_call_id_2, + name="mock_concurrent_tool", + input='{"input": "conc1"}', + ), + ToolCallBlock( + id=tool_call_id_3, + name="mock_concurrent_tool", + input='{"input": "conc2"}', + ), + ], + is_last=True, + usage=None, + ), + ], + [ + ChatResponse( + content=[TextBlock(text="All done")], + is_last=False, + ), + ChatResponse( + content=[TextBlock(text="All done")], + is_last=True, + ), + ], + ], + ) + + # Collect all events + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content="Test"), + ): + events.append(event.model_dump()) + + # Verify events + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + + # For mixed execution: sequential tool first, then concurrent tools + # Split events into: prefix (before tool execution), sequential part, + # concurrent part, suffix (final model call) + + # Expected prefix events (before tool execution) + expected_prefix = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + ] + + # Expected sequential tool execution events + expected_sequential = [ + { + "type": "TOOL_RESULT_START", + "tool_call_id": tool_call_id_1, + "tool_call_name": "mock_sequential_tool", + }, + { + "type": "TOOL_RESULT_TEXT_DELTA", + "tool_call_id": tool_call_id_1, + "delta": "Sequential result: seq1", + }, + { + "type": "TOOL_RESULT_END", + "tool_call_id": tool_call_id_1, + "state": "success", + "metadata": {}, + }, + ] + + # Expected concurrent events (order may vary) + expected_concurrent: list[dict[str, Any]] = [ + { + "type": "TOOL_RESULT_START", + "tool_call_id": tool_call_id_2, + "tool_call_name": "mock_concurrent_tool", + }, + { + "type": "TOOL_RESULT_TEXT_DELTA", + "tool_call_id": tool_call_id_2, + "delta": "Concurrent result: conc1", + }, + { + "type": "TOOL_RESULT_END", + "tool_call_id": tool_call_id_2, + "state": "success", + }, + { + "type": "TOOL_RESULT_START", + "tool_call_id": tool_call_id_3, + "tool_call_name": "mock_concurrent_tool", + }, + { + "type": "TOOL_RESULT_TEXT_DELTA", + "tool_call_id": tool_call_id_3, + "delta": "Concurrent result: conc2", + }, + { + "type": "TOOL_RESULT_END", + "tool_call_id": tool_call_id_3, + "state": "success", + "metadata": {}, + }, + ] + + # Expected suffix events (final model call with pure text) + expected_suffix = [ + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + {"type": "TEXT_BLOCK_START", "block_id": AnyString()}, + { + "type": "TEXT_BLOCK_DELTA", + "block_id": AnyString(), + "delta": "All done", + }, + {"type": "TEXT_BLOCK_END", "block_id": AnyString()}, + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + {"type": "REPLY_END", "session_id": session_id}, + ] + + # Assert prefix events (fixed order) + basic_dict = self._get_event_base(reply_id) + prefix_len = len(expected_prefix) + self.assertListEqual( + events[:prefix_len], + [{**basic_dict, **_} for _ in expected_prefix], + ) + + # Assert sequential events (fixed order) + sequential_len = len(expected_sequential) + self.assertListEqual( + events[prefix_len : prefix_len + sequential_len], + [{**basic_dict, **_} for _ in expected_sequential], + ) + + # Assert concurrent events (order may vary) + concurrent_len = len(expected_concurrent) + concurrent_events = events[ + prefix_len + + sequential_len : prefix_len + + sequential_len + + concurrent_len + ] + + # Check length matches + self.assertEqual(len(concurrent_events), len(expected_concurrent)) + + # Check each expected event is in the actual events + for expected_event in expected_concurrent: + expected_with_base = {**basic_dict, **expected_event} + self.assertIn(expected_with_base, concurrent_events) + + # Assert suffix events (fixed order) + suffix_events = events[prefix_len + sequential_len + concurrent_len :] + self.assertListEqual( + suffix_events, + [{**basic_dict, **_} for _ in expected_suffix], + ) + + # Assert context after reply_stream + msg_base = self._get_msg_base() + expected_context = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Test", + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": tool_call_id_1, + "name": "mock_sequential_tool", + "input": '{"input": "seq1"}', + "state": "finished", + "suggested_rules": [], + }, + { + "type": "tool_call", + "id": tool_call_id_2, + "name": "mock_concurrent_tool", + "input": '{"input": "conc1"}', + "state": "finished", + "suggested_rules": [], + }, + { + "type": "tool_call", + "id": tool_call_id_3, + "name": "mock_concurrent_tool", + "input": '{"input": "conc2"}', + "state": "finished", + "suggested_rules": [], + }, + { + "type": "tool_result", + "id": AnyString(), + "name": "mock_sequential_tool", + "state": "success", + "output": [ + { + "type": "text", + "id": AnyString(), + "text": "Sequential result: seq1", + }, + ], + "metadata": {}, + }, + { + "type": "tool_result", + "id": AnyString(), + "name": "mock_concurrent_tool", + "state": "success", + "output": [ + { + "type": "text", + "id": AnyString(), + "text": "Concurrent result: conc1", + }, + ], + "metadata": {}, + }, + { + "type": "tool_result", + "id": AnyString(), + "name": "mock_concurrent_tool", + "state": "success", + "output": [ + { + "type": "text", + "id": AnyString(), + "text": "Concurrent result: conc2", + }, + ], + "metadata": {}, + }, + { + "type": "text", + "id": AnyString(), + "text": "All done", + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + async def asyncTearDown(self) -> None: + """The async teardown method.""" diff --git a/tests/agui_protocol_test.py b/tests/agui_protocol_test.py new file mode 100644 index 0000000000000000000000000000000000000000..5529147d23905af44f1598f549b1ea2a000a141b --- /dev/null +++ b/tests/agui_protocol_test.py @@ -0,0 +1,707 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Test cases for AGUI protocol middleware.""" + +import json +from typing import AsyncGenerator +from unittest.async_case import IsolatedAsyncioTestCase +from unittest.mock import MagicMock + +from fastapi import FastAPI +from fastapi.responses import JSONResponse +from fastapi.responses import StreamingResponse +from fastapi.testclient import TestClient + +from agentscope.app.middleware import AGUIProtocolMiddleware +from agentscope.event import ( + ConfirmResult, + DataBlockDeltaEvent, + DataBlockEndEvent, + DataBlockStartEvent, + ExceedMaxItersEvent, + ExternalExecutionResultEvent, + ModelCallEndEvent, + ModelCallStartEvent, + ReplyEndEvent, + ReplyStartEvent, + RequireExternalExecutionEvent, + RequireUserConfirmEvent, + TextBlockDeltaEvent, + TextBlockEndEvent, + TextBlockStartEvent, + ThinkingBlockDeltaEvent, + ThinkingBlockEndEvent, + ThinkingBlockStartEvent, + ToolCallDeltaEvent, + ToolCallEndEvent, + ToolCallStartEvent, + ToolResultDataDeltaEvent, + ToolResultEndEvent, + ToolResultStartEvent, + ToolResultTextDeltaEvent, + UserConfirmResultEvent, +) +from agentscope.message import ToolCallBlock, ToolResultBlock, ToolResultState + + +async def _collect_stream( + mw: AGUIProtocolMiddleware, + chunks: list[str], +) -> str: + """Collect converted stream chunks as text.""" + + async def _stream() -> AsyncGenerator[str, None]: + """Yield the provided chunks.""" + for chunk in chunks: + yield chunk + + out: list[str] = [] + async for item in mw._convert_stream(_stream()): + out.append(item.decode("utf-8")) + return "".join(out) + + +class AGUIProtocolStreamTest(IsolatedAsyncioTestCase): + """Test stream-level conversion behavior.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.mw = AGUIProtocolMiddleware(app=MagicMock()) + + async def test_raw_json_stream_is_converted(self) -> None: + """Test raw AgentEvent JSON stream conversion.""" + event = ReplyStartEvent( + session_id="sess_1", + reply_id="reply_1", + name="agent", + ) + + body = await _collect_stream(self.mw, [event.model_dump_json()]) + data = json.loads(body) + + self.assertEqual(data["type"], "RUN_STARTED") + self.assertEqual(data["threadId"], "sess_1") + self.assertEqual(data["runId"], "reply_1") + + async def test_sse_data_frame_is_converted(self) -> None: + """Test AgentEvent JSON inside an SSE data frame is converted.""" + event = ReplyStartEvent( + session_id="sess_1", + reply_id="reply_1", + name="agent", + ) + + body = await _collect_stream( + self.mw, + [f"data: {event.model_dump_json()}\n\n"], + ) + self.assertTrue(body.startswith("data: ")) + + data = json.loads(body.removeprefix("data: ").strip()) + self.assertEqual(data["type"], "RUN_STARTED") + self.assertEqual(data["threadId"], "sess_1") + self.assertEqual(data["runId"], "reply_1") + self.assertNotIn("session_id", data) + + async def test_sse_heartbeat_is_passed_through(self) -> None: + """Test SSE heartbeat frames are not modified.""" + self.assertEqual( + await _collect_stream(self.mw, [":\n\n"]), + ":\n\n", + ) + + async def test_sse_data_frame_with_crlf_is_converted(self) -> None: + """Test AgentEvent JSON inside a CRLF SSE frame is converted.""" + event = ReplyStartEvent( + session_id="sess_1", + reply_id="reply_1", + name="agent", + ) + + body = await _collect_stream( + self.mw, + [f"data: {event.model_dump_json()}\r\n\r\n"], + ) + self.assertTrue(body.startswith("data: ")) + self.assertTrue(body.endswith("\r\n\r\n")) + + data = json.loads(body.removeprefix("data: ").strip()) + self.assertEqual(data["type"], "RUN_STARTED") + self.assertEqual(data["threadId"], "sess_1") + self.assertEqual(data["runId"], "reply_1") + + async def test_fastapi_sse_response_is_converted(self) -> None: + """Test middleware converts a real FastAPI SSE response.""" + app = FastAPI() + app.add_middleware(AGUIProtocolMiddleware) + + @app.get("/sessions/sess_1/stream") + async def stream() -> StreamingResponse: + event = ReplyStartEvent( + session_id="sess_1", + reply_id="reply_1", + name="agent", + ) + + async def gen() -> AsyncGenerator[str, None]: + yield f"data: {event.model_dump_json()}\n\n" + + return StreamingResponse(gen(), media_type="text/event-stream") + + client = TestClient(app) + with client.stream("GET", "/sessions/sess_1/stream") as response: + body = "".join(response.iter_text()) + + self.assertTrue(body.startswith("data: ")) + data = json.loads(body.removeprefix("data: ").strip()) + self.assertEqual(data["type"], "RUN_STARTED") + self.assertEqual(data["threadId"], "sess_1") + self.assertEqual(data["runId"], "reply_1") + self.assertNotIn("session_id", data) + + async def test_fastapi_json_response_is_not_converted(self) -> None: + """Test non-SSE responses are outside protocol conversion scope.""" + app = FastAPI() + app.add_middleware(AGUIProtocolMiddleware) + + @app.get("/event") + def event() -> JSONResponse: + agent_event = ReplyStartEvent( + session_id="sess_1", + reply_id="reply_1", + name="agent", + ) + return JSONResponse(agent_event.model_dump(mode="json")) + + client = TestClient(app) + response = client.get("/event") + + self.assertEqual(response.status_code, 200) + data = response.json() + self.assertEqual(data["type"], "REPLY_START") + self.assertEqual(data["session_id"], "sess_1") + self.assertNotIn("RUN_STARTED", response.text) + + async def asyncTearDown(self) -> None: + """The async teardown method.""" + + +class AGUIProtocolLifecycleTest(IsolatedAsyncioTestCase): + """Test lifecycle event conversions.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.mw = AGUIProtocolMiddleware(app=MagicMock()) + + async def test_reply_start_to_run_started(self) -> None: + """Test ReplyStartEvent -> RUN_STARTED.""" + event = ReplyStartEvent( + session_id="sess_1", + reply_id="reply_1", + name="agent", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "RUN_STARTED") + self.assertEqual(result["threadId"], "sess_1") + self.assertEqual(result["runId"], "reply_1") + self.assertNotIn("name", result) + + async def test_reply_end_to_run_finished(self) -> None: + """Test ReplyEndEvent -> RUN_FINISHED.""" + event = ReplyEndEvent( + session_id="sess_1", + reply_id="reply_1", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "RUN_FINISHED") + self.assertEqual(result["threadId"], "sess_1") + self.assertEqual(result["runId"], "reply_1") + + async def test_exceed_max_iters_to_run_error(self) -> None: + """Test ExceedMaxItersEvent -> RUN_ERROR.""" + event = ExceedMaxItersEvent( + reply_id="reply_1", + name="my_agent", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "RUN_ERROR") + self.assertIn("my_agent", result["message"]) + self.assertEqual(result["code"], "exceed_max_iters") + + async def asyncTearDown(self) -> None: + """The async teardown method.""" + + +class AGUIProtocolStepTest(IsolatedAsyncioTestCase): + """Test model call -> step event conversions.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.mw = AGUIProtocolMiddleware(app=MagicMock()) + + async def test_model_call_start_to_step_started(self) -> None: + """Test ModelCallStartEvent -> STEP_STARTED.""" + event = ModelCallStartEvent( + reply_id="reply_1", + model_name="gpt-4", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "STEP_STARTED") + self.assertEqual(result["stepName"], "gpt-4") + + async def test_model_call_end_to_step_finished(self) -> None: + """Test ModelCallEndEvent -> STEP_FINISHED with matching step_name.""" + start_event = ModelCallStartEvent( + reply_id="reply_1", + model_name="gpt-4", + ) + self.mw._convert_to_protocol(start_event) + + end_event = ModelCallEndEvent( + reply_id="reply_1", + input_tokens=100, + output_tokens=50, + ) + result = self.mw._convert_to_protocol(end_event) + + self.assertEqual(result["type"], "STEP_FINISHED") + self.assertEqual(result["stepName"], "gpt-4") + + async def asyncTearDown(self) -> None: + """The async teardown method.""" + + +class AGUIProtocolTextMessageTest(IsolatedAsyncioTestCase): + """Test text block -> text message event conversions.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.mw = AGUIProtocolMiddleware(app=MagicMock()) + + async def test_text_block_start(self) -> None: + """Test TextBlockStartEvent -> TEXT_MESSAGE_START.""" + event = TextBlockStartEvent( + reply_id="reply_1", + block_id="block_1", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "TEXT_MESSAGE_START") + self.assertEqual(result["messageId"], "block_1") + + async def test_text_block_delta(self) -> None: + """Test TextBlockDeltaEvent -> TEXT_MESSAGE_CONTENT.""" + event = TextBlockDeltaEvent( + reply_id="reply_1", + block_id="block_1", + delta="Hello, ", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "TEXT_MESSAGE_CONTENT") + self.assertEqual(result["messageId"], "block_1") + self.assertEqual(result["delta"], "Hello, ") + + async def test_text_block_end(self) -> None: + """Test TextBlockEndEvent -> TEXT_MESSAGE_END.""" + event = TextBlockEndEvent( + reply_id="reply_1", + block_id="block_1", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "TEXT_MESSAGE_END") + self.assertEqual(result["messageId"], "block_1") + + async def asyncTearDown(self) -> None: + """The async teardown method.""" + + +class AGUIProtocolReasoningTest(IsolatedAsyncioTestCase): + """Test thinking block -> reasoning message event conversions.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.mw = AGUIProtocolMiddleware(app=MagicMock()) + + async def test_thinking_block_start(self) -> None: + """Test ThinkingBlockStartEvent -> REASONING_MESSAGE_START.""" + event = ThinkingBlockStartEvent( + reply_id="reply_1", + block_id="think_1", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "REASONING_MESSAGE_START") + self.assertEqual(result["messageId"], "think_1") + self.assertEqual(result["role"], "reasoning") + + async def test_thinking_block_delta(self) -> None: + """Test ThinkingBlockDeltaEvent -> REASONING_MESSAGE_CONTENT.""" + event = ThinkingBlockDeltaEvent( + reply_id="reply_1", + block_id="think_1", + delta="Let me think...", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "REASONING_MESSAGE_CONTENT") + self.assertEqual(result["messageId"], "think_1") + self.assertEqual(result["delta"], "Let me think...") + + async def test_thinking_block_end(self) -> None: + """Test ThinkingBlockEndEvent -> REASONING_MESSAGE_END.""" + event = ThinkingBlockEndEvent( + reply_id="reply_1", + block_id="think_1", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "REASONING_MESSAGE_END") + self.assertEqual(result["messageId"], "think_1") + + async def asyncTearDown(self) -> None: + """The async teardown method.""" + + +class AGUIProtocolToolCallTest(IsolatedAsyncioTestCase): + """Test tool call event conversions.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.mw = AGUIProtocolMiddleware(app=MagicMock()) + + async def test_tool_call_start(self) -> None: + """Test ToolCallStartEvent -> TOOL_CALL_START.""" + event = ToolCallStartEvent( + reply_id="reply_1", + tool_call_id="tc_1", + tool_call_name="search", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "TOOL_CALL_START") + self.assertEqual(result["toolCallId"], "tc_1") + self.assertEqual(result["toolCallName"], "search") + self.assertEqual(result["parentMessageId"], "reply_1") + + async def test_tool_call_delta(self) -> None: + """Test ToolCallDeltaEvent -> TOOL_CALL_ARGS.""" + event = ToolCallDeltaEvent( + reply_id="reply_1", + tool_call_id="tc_1", + delta='{"query": "hello"}', + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "TOOL_CALL_ARGS") + self.assertEqual(result["toolCallId"], "tc_1") + self.assertEqual(result["delta"], '{"query": "hello"}') + + async def test_tool_call_end(self) -> None: + """Test ToolCallEndEvent -> TOOL_CALL_END.""" + event = ToolCallEndEvent( + reply_id="reply_1", + tool_call_id="tc_1", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "TOOL_CALL_END") + self.assertEqual(result["toolCallId"], "tc_1") + + async def asyncTearDown(self) -> None: + """The async teardown method.""" + + +class AGUIProtocolToolResultTest(IsolatedAsyncioTestCase): + """Test tool result event conversions.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.mw = AGUIProtocolMiddleware(app=MagicMock()) + + async def test_tool_result_end_with_buffered_content(self) -> None: + """Test that ToolResultEndEvent carries accumulated text content.""" + self.mw._convert_to_protocol( + ToolResultTextDeltaEvent( + reply_id="reply_1", + tool_call_id="tc_1", + delta="partial ", + ), + ) + self.mw._convert_to_protocol( + ToolResultTextDeltaEvent( + reply_id="reply_1", + tool_call_id="tc_1", + delta="result", + ), + ) + + result = self.mw._convert_to_protocol( + ToolResultEndEvent( + reply_id="reply_1", + tool_call_id="tc_1", + state=ToolResultState.SUCCESS, + ), + ) + + self.assertEqual(result["type"], "TOOL_CALL_RESULT") + self.assertEqual(result["toolCallId"], "tc_1") + self.assertEqual(result["messageId"], "reply_1") + self.assertEqual(result["content"], "partial result") + + async def test_tool_result_end_fallback_to_state(self) -> None: + """Test that ToolResultEndEvent falls back to state when no buffer.""" + result = self.mw._convert_to_protocol( + ToolResultEndEvent( + reply_id="reply_1", + tool_call_id="tc_1", + state=ToolResultState.ERROR, + ), + ) + + self.assertEqual(result["type"], "TOOL_CALL_RESULT") + self.assertEqual(result["content"], "error") + + async def test_tool_result_start_to_custom(self) -> None: + """Test ToolResultStartEvent -> CUSTOM.""" + event = ToolResultStartEvent( + reply_id="reply_1", + tool_call_id="tc_1", + tool_call_name="search", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "CUSTOM") + self.assertEqual(result["name"], "tool_result_start") + self.assertIsInstance(result["value"], dict) + + async def test_tool_result_text_delta_to_custom(self) -> None: + """Test ToolResultTextDeltaEvent -> CUSTOM.""" + event = ToolResultTextDeltaEvent( + reply_id="reply_1", + tool_call_id="tc_1", + delta="partial result", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "CUSTOM") + self.assertEqual(result["name"], "tool_result_text_delta") + self.assertEqual(result["value"]["delta"], "partial result") + + async def test_tool_result_data_delta_to_custom(self) -> None: + """Test ToolResultDataDeltaEvent -> CUSTOM.""" + event = ToolResultDataDeltaEvent( + reply_id="reply_1", + tool_call_id="tc_1", + media_type="image/png", + data="base64data", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "CUSTOM") + self.assertEqual(result["name"], "tool_result_data_delta") + self.assertEqual(result["value"]["media_type"], "image/png") + + async def asyncTearDown(self) -> None: + """The async teardown method.""" + + +class AGUIProtocolDataBlockTest(IsolatedAsyncioTestCase): + """Test data block -> custom event conversions.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.mw = AGUIProtocolMiddleware(app=MagicMock()) + + async def test_data_block_start(self) -> None: + """Test DataBlockStartEvent -> CUSTOM.""" + event = DataBlockStartEvent( + reply_id="reply_1", + block_id="db_1", + media_type="image/png", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "CUSTOM") + self.assertEqual(result["name"], "data_block_start") + self.assertEqual(result["value"]["media_type"], "image/png") + + async def test_data_block_delta(self) -> None: + """Test DataBlockDeltaEvent -> CUSTOM.""" + event = DataBlockDeltaEvent( + reply_id="reply_1", + block_id="db_1", + data="base64chunk", + media_type="image/png", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "CUSTOM") + self.assertEqual(result["name"], "data_block_delta") + self.assertEqual(result["value"]["data"], "base64chunk") + + async def test_data_block_end(self) -> None: + """Test DataBlockEndEvent -> CUSTOM.""" + event = DataBlockEndEvent( + reply_id="reply_1", + block_id="db_1", + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "CUSTOM") + self.assertEqual(result["name"], "data_block_end") + + async def asyncTearDown(self) -> None: + """The async teardown method.""" + + +class AGUIProtocolPermissionTest(IsolatedAsyncioTestCase): + """Test permission-related event conversions.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.mw = AGUIProtocolMiddleware(app=MagicMock()) + + async def test_require_user_confirm(self) -> None: + """Test RequireUserConfirmEvent -> CUSTOM.""" + event = RequireUserConfirmEvent( + reply_id="reply_1", + tool_calls=[ + ToolCallBlock( + id="tc_1", + name="bash", + input='{"command": "ls"}', + ), + ], + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "CUSTOM") + self.assertEqual(result["name"], "require_user_confirm") + self.assertIn("tool_calls", result["value"]) + + async def test_require_external_execution(self) -> None: + """Test RequireExternalExecutionEvent -> CUSTOM.""" + event = RequireExternalExecutionEvent( + reply_id="reply_1", + tool_calls=[ + ToolCallBlock( + id="tc_2", + name="external_api", + input="{}", + ), + ], + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "CUSTOM") + self.assertEqual(result["name"], "require_external_execution") + + async def test_user_confirm_result(self) -> None: + """Test UserConfirmResultEvent -> CUSTOM.""" + event = UserConfirmResultEvent( + reply_id="reply_1", + confirm_results=[ + ConfirmResult( + confirmed=True, + tool_call=ToolCallBlock( + id="tc_1", + name="bash", + input='{"command": "ls"}', + ), + ), + ], + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "CUSTOM") + self.assertEqual(result["name"], "user_confirm_result") + + async def test_external_execution_result(self) -> None: + """Test ExternalExecutionResultEvent -> CUSTOM.""" + event = ExternalExecutionResultEvent( + reply_id="reply_1", + execution_results=[ + ToolResultBlock( + id="tc_2", + name="external_api", + output="result data", + ), + ], + ) + result = self.mw._convert_to_protocol(event) + + self.assertEqual(result["type"], "CUSTOM") + self.assertEqual(result["name"], "external_execution_result") + + async def asyncTearDown(self) -> None: + """The async teardown method.""" + + +class AGUIProtocolCamelCaseTest(IsolatedAsyncioTestCase): + """Verify that all output dicts use camelCase keys as AGUI requires.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.mw = AGUIProtocolMiddleware(app=MagicMock()) + + def _assert_no_snake_case_keys(self, d: dict, context: str) -> None: + """Assert none of the top-level keys contain underscores.""" + for key in d: + if key == "value": + continue + self.assertNotIn( + "_", + key, + f"Key '{key}' in {context} should be camelCase", + ) + + async def test_all_standard_events_produce_camel_case(self) -> None: + """Test that all directly-mapped events produce camelCase keys.""" + events = [ + ReplyStartEvent( + session_id="s", + reply_id="r", + name="a", + ), + ReplyEndEvent(session_id="s", reply_id="r"), + ModelCallStartEvent(reply_id="r", model_name="m"), + ModelCallEndEvent( + reply_id="r", + input_tokens=1, + output_tokens=1, + ), + TextBlockStartEvent(reply_id="r", block_id="b"), + TextBlockDeltaEvent(reply_id="r", block_id="b", delta="x"), + TextBlockEndEvent(reply_id="r", block_id="b"), + ThinkingBlockStartEvent(reply_id="r", block_id="b"), + ThinkingBlockDeltaEvent(reply_id="r", block_id="b", delta="x"), + ThinkingBlockEndEvent(reply_id="r", block_id="b"), + ToolCallStartEvent( + reply_id="r", + tool_call_id="t", + tool_call_name="n", + ), + ToolCallDeltaEvent( + reply_id="r", + tool_call_id="t", + delta="x", + ), + ToolCallEndEvent(reply_id="r", tool_call_id="t"), + ExceedMaxItersEvent(reply_id="r", name="a"), + ] + + for event in events: + result = self.mw._convert_to_protocol(event) + self._assert_no_snake_case_keys( + result, + type(event).__name__, + ) + + async def asyncTearDown(self) -> None: + """The async teardown method.""" diff --git a/tests/app_lifespan_dedicated_test.py b/tests/app_lifespan_dedicated_test.py new file mode 100644 index 0000000000000000000000000000000000000000..30fab478744f1189bd884a874e37065e1f1c9fac --- /dev/null +++ b/tests/app_lifespan_dedicated_test.py @@ -0,0 +1,332 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""End-to-end wiring test for dedicated-deployment knowledge-base upload. + +Boots the FastAPI app with ``enable_index_worker=False`` so the API +process does NOT host an :class:`IndexWorker`. Dispatch happens +through the message bus: a ``MessageBusDispatcher`` writes the task +to the shared queue and publishes the signal. In a real deployment +a separate worker process would pick it up; here we run an +:class:`IndexTaskConsumer` against a stub worker inside the test +to confirm the producer side reaches the bus correctly. + +What this test guards: + +- the API's lifespan accepts ``enable_index_worker=False`` without + raising; +- an upload returns ``201`` and persists a ``pending`` record; +- the message bus carries the dispatch out of the API process so + a worker subscribed on the same bus can consume it; +- the consumer's ``worker.process`` is invoked exactly once per + dispatch, with the same ids the API recorded in storage. + +We deliberately do NOT drive the document to ``ready`` here — that +path is already covered by the embedded-mode upload test. The +purpose of this test is to lock down the bus hop introduced by +:class:`MessageBusDispatcher`. +""" +import asyncio +import tempfile +from typing import Any +from unittest.async_case import IsolatedAsyncioTestCase + +import fakeredis.aioredis +from fastapi.testclient import TestClient + +from agentscope.app import create_app +from agentscope.app._service import IndexTaskConsumer +from agentscope.app.rag.blob_store import LocalBlobStore +from agentscope.app.rag.knowledge_base_manager import ( + KnowledgeBaseManagerBase, + KnowledgeBaseNotFoundError, +) +from agentscope.app.rag.knowledge_base_manager._dimension_policy import ( + DimensionPolicy, + DimensionPolicyKind, +) +from agentscope.app.message_bus import RedisMessageBus +from agentscope.app.storage import ( + EmbeddingModelConfig, + KnowledgeBaseRecord, + RedisStorage, +) +from agentscope.app.workspace_manager._base import WorkspaceManagerBase +from agentscope.rag import VectorStoreBase +from agentscope.rag._vdb._vector_store import ( + DocumentSummary, + VectorRecord, + VectorSearchResult, +) + + +# ---------------------------------------------------------------------- +# Test doubles — borrowed in spirit from service_knowledge_base_upload_test +# but trimmed: the dedicated-mode test does not drive embedding. +# ---------------------------------------------------------------------- + + +class _FakeVectorStore(VectorStoreBase): + """Bare minimum to satisfy create_app's vector-store wiring.""" + + def __init__(self) -> None: + self._collections: dict[str, list[VectorRecord]] = {} + + async def create_collection(self, name: str, dimensions: int) -> None: + self._collections.setdefault(name, []) + + async def delete_collection(self, name: str) -> None: + self._collections.pop(name, None) + + async def has_collection(self, name: str) -> bool: + return name in self._collections + + async def insert( + self, + collection: str, + records: list[VectorRecord], + ) -> None: + self._collections.setdefault(collection, []).extend(records) + + async def delete(self, collection: str, document_id: str) -> None: + self._collections[collection] = [ + r + for r in self._collections.get(collection, []) + if r.document_id != document_id + ] + + async def search( + self, + collection: str, + query_vector: list[float], + top_k: int = 5, + metadata_filter: dict[str, Any] | None = None, + ) -> list[VectorSearchResult]: + return [] + + async def list_documents( + self, + collection: str, + metadata_filter: dict[str, Any] | None = None, + ) -> list[DocumentSummary]: + return [] + + +class _FakeKbManager(KnowledgeBaseManagerBase): + """KB manager that resolves knowledge bases via storage only. + + Returns a noop knowledge for ``get_knowledge`` because dedicated + mode does not exercise embedding in this test — the worker stub + intercepts ``process`` before the knowledge call. + """ + + async def get_dimension_policy(self) -> DimensionPolicy: + return DimensionPolicy(kind=DimensionPolicyKind.ANY, dimension=None) + + async def create_knowledge_base( + self, + user_id: str, + name: str, + description: str, + embedding_model_config: EmbeddingModelConfig, + ) -> KnowledgeBaseRecord: + raise NotImplementedError + + async def delete_knowledge_base( + self, + user_id: str, + knowledge_base_id: str, + ) -> bool: + return False + + async def get_knowledge( + self, + user_id: str, + knowledge_base_id: str, + ) -> Any: + record = await self._storage.get_knowledge_base( + user_id, + knowledge_base_id, + ) + if record is None: + raise KnowledgeBaseNotFoundError( + f"Knowledge base {knowledge_base_id!r} not found.", + ) + raise NotImplementedError # unused in this test + + +class _NoopWorkspaceManager(WorkspaceManagerBase): + """Workspace manager that does nothing.""" + + async def get_workspace(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError + + async def create_workspace(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError + + async def close(self, workspace_id: str) -> None: + return None + + async def close_all(self) -> None: + return None + + +def _make_storage(fr: fakeredis.aioredis.FakeRedis) -> RedisStorage: + class _FakeStorage(RedisStorage): + async def __aenter__(self) -> "_FakeStorage": # type: ignore[override] + self._client = fr + return self + + async def aclose(self) -> None: + self._client = None + + return _FakeStorage() + + +def _make_bus(fr: fakeredis.aioredis.FakeRedis) -> RedisMessageBus: + class _FakeBus(RedisMessageBus): + async def __aenter__(self) -> "_FakeBus": # type: ignore[override] + self._client = fr + return self + + async def aclose(self) -> None: + self._client = None + + return _FakeBus() + + +class _RecordingWorker: + """Stub worker that records each ``process`` invocation. + + Stands in for :class:`IndexWorker` in the test so we can verify + that the dispatch hop landed without spinning up the full + parse → chunk → embed pipeline. + """ + + def __init__(self) -> None: + self.calls: list[dict] = [] + self.notify = asyncio.Event() + + async def process( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + ) -> None: + """Record the dispatched task and signal the test. + + Stands in for :class:`IndexWorker.process_one` so the lifespan + tests can assert that the API process forwarded the right + ``user_id`` / ``knowledge_base_id`` / ``document_id`` triple. + + Args: + user_id (`str`): + The owning user id. + knowledge_base_id (`str`): + The parent knowledge base id. + document_id (`str`): + The document id to "process". + """ + self.calls.append( + { + "user_id": user_id, + "knowledge_base_id": knowledge_base_id, + "document_id": document_id, + }, + ) + self.notify.set() + + +class DedicatedModeUploadFlowTest(IsolatedAsyncioTestCase): + """The producer side reaches the bus; a separate consumer sees it.""" + + async def asyncSetUp(self) -> None: + self._tmp = tempfile.TemporaryDirectory() + self._fr = fakeredis.aioredis.FakeRedis(decode_responses=True) + self._vector_store = _FakeVectorStore() + storage = _make_storage(self._fr) + self._api_message_bus = _make_bus(self._fr) + + self._app = create_app( + storage=storage, + message_bus=self._api_message_bus, + workspace_manager=_NoopWorkspaceManager(), + knowledge_base_manager=_FakeKbManager( + storage=storage, + vector_store=self._vector_store, + ), + blob_store=LocalBlobStore(root_dir=self._tmp.name), + enable_index_worker=False, + ) + + # Seed a knowledge base directly so we don't have to mock the + # manager's create flow over HTTP. + kb_record = KnowledgeBaseRecord( + user_id="user-1", + name="kb", + description="", + embedding_model_config=EmbeddingModelConfig( + type="openai_credential", + credential_id="cred-1", + model="text-embedding-3-small", + dimensions=1, + ), + collection_name="", + ) + kb_record.collection_name = f"kb_{kb_record.id}" + await self._vector_store.create_collection( + kb_record.collection_name, + 1, + ) + storage._client = self._fr + await storage.upsert_knowledge_base("user-1", kb_record) + storage._client = None + self._kb_id = kb_record.id + + async def asyncTearDown(self) -> None: + await self._fr.aclose() + self._tmp.cleanup() + + async def test_upload_dispatches_through_message_bus(self) -> None: + """An upload in dedicated mode reaches a separate consumer.""" + # The consumer's bus is a SEPARATE RedisMessageBus instance + # bound to the same fakeredis store. Production wiring would + # be two TCP-connected clients; here they share the in-memory + # backend, which exercises the bus contract correctly. + consumer_bus = _make_bus(self._fr) + worker = _RecordingWorker() + + async with consumer_bus, IndexTaskConsumer( + message_bus=consumer_bus, + worker=worker, + ): + headers = {"X-User-ID": "user-1"} + with TestClient(self._app) as client: + files = { + "file": ( + "hello.txt", + b"hello world\n" * 16, + "text/plain", + ), + } + resp = client.post( + f"/knowledge_bases/{self._kb_id}/documents", + files=files, + headers=headers, + ) + self.assertEqual(resp.status_code, 201, resp.text) + body = resp.json() + document_id = body["document_id"] + + # The consumer's worker should see the dispatch. + await asyncio.wait_for(worker.notify.wait(), timeout=5.0) + + self.assertEqual( + worker.calls, + [ + { + "user_id": "user-1", + "knowledge_base_id": self._kb_id, + "document_id": document_id, + }, + ], + ) diff --git a/tests/backend_docker_test.py b/tests/backend_docker_test.py new file mode 100644 index 0000000000000000000000000000000000000000..932a08be6ca13a76398eefc8c2e85e567ec893b3 --- /dev/null +++ b/tests/backend_docker_test.py @@ -0,0 +1,218 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Test cases for :class:`DockerBackend`. + +Validates that the three backend primitives (``exec_shell``, +``read_file``, ``write_file``) and the inherited shell-based filesystem +helpers behave correctly inside a real Docker container. + +The suite is skipped unless: + +* the Docker daemon is reachable (probed via ``docker info``), and +* the host is Linux — the project only validates the Docker backend on + Ubuntu/Linux CI runners; spinning containers up on other hosts is + out of scope for these tests. + +A live container is obtained by initializing a :class:`DockerWorkspace` +and reusing its already-wired :class:`DockerBackend` (``ws._backend``), +which avoids duplicating the container bring-up logic here. +""" + +import shutil +import subprocess +import sys +import tempfile +import unittest +import uuid +from unittest.async_case import IsolatedAsyncioTestCase + +from agentscope.tool import ExecResult +from agentscope.workspace import DockerWorkspace, DockerBackend +from agentscope.workspace._docker._make_dockerfile import CONTAINER_WORKDIR + + +# ── availability checks ──────────────────────────────────────────── + + +def _docker_available() -> bool: + """Return ``True`` iff the Docker daemon is reachable. + + Probes via the ``docker`` CLI (cheap, synchronous) so the result can + gate the module at import time. + """ + if shutil.which("docker") is None: + return False + try: + result = subprocess.run( + ["docker", "info"], + capture_output=True, + timeout=5, + check=False, + ) + return result.returncode == 0 + except (subprocess.TimeoutExpired, FileNotFoundError, OSError): + return False + + +_IS_LINUX = sys.platform == "linux" +_DOCKER_OK = _IS_LINUX and _docker_available() +_SKIP_REASON = "Docker backend is only tested on Linux with a live daemon" + + +@unittest.skipUnless(_DOCKER_OK, _SKIP_REASON) +class TestDockerBackend(IsolatedAsyncioTestCase): + """Test cases for ``DockerBackend`` against a live container.""" + + async def asyncSetUp(self) -> None: + """Start a workspace and reuse its wired backend. + + The ``workspace_id`` is randomised so concurrent / repeated runs + do not collide on the deterministic container name. + """ + # pylint: disable=consider-using-with + self.temp_dir = tempfile.TemporaryDirectory() + self.workspace = DockerWorkspace( + workspace_id=f"test-{uuid.uuid4().hex[:8]}", + host_workdir=self.temp_dir.name, + ) + await self.workspace.initialize() + self.backend = self.workspace._backend + self.assertIsInstance(self.backend, DockerBackend) + + async def asyncTearDown(self) -> None: + """Stop the container and drop the temp host dir.""" + try: + await self.workspace.close() + finally: + self.temp_dir.cleanup() + + # ── exec ─────────────────────────────────────────────────────── + + async def test_exec_returns_stdout(self) -> None: + """A program's stdout/exit code are captured into ``ExecResult``.""" + result = await self.backend.exec_shell(["echo", "hello world"]) + self.assertIsInstance(result, ExecResult) + self.assertTrue(result.ok()) + self.assertEqual(result.stdout.decode().strip(), "hello world") + + async def test_exec_nonzero_exit(self) -> None: + """A non-zero exit is reported as a normal (non-raising) result.""" + result = await self.backend.exec_shell( + ["sh", "-c", "echo oops >&2; exit 4"], + ) + self.assertEqual(result.exit_code, 4) + self.assertIn("oops", result.stderr.decode()) + + async def test_exec_argv_not_shell_split(self) -> None: + """A single argv element with metacharacters reaches the program + intact (no shell interposed by the primitive).""" + tricky = "a b $(echo x) | ;" + result = await self.backend.exec_shell(["echo", tricky]) + self.assertTrue(result.ok()) + self.assertEqual(result.stdout.decode().rstrip("\n"), tricky) + + async def test_exec_cwd_default_is_workdir(self) -> None: + """With no explicit ``cwd`` the container workdir is used.""" + result = await self.backend.exec_shell(["pwd"]) + self.assertTrue(result.ok()) + self.assertEqual(result.stdout.decode().strip(), CONTAINER_WORKDIR) + + async def test_exec_timeout_returns_minus_one(self) -> None: + """A command exceeding ``timeout`` reports the -1 sentinel.""" + result = await self.backend.exec_shell( + ["sleep", "10"], + timeout=0.5, + ) + self.assertEqual(result.exit_code, -1) + self.assertEqual(result.stderr, b"timed out") + + # ── file I/O ─────────────────────────────────────────────────── + + async def test_write_then_read_roundtrip(self) -> None: + """Bytes written into the container are read back verbatim.""" + path = f"{CONTAINER_WORKDIR}/roundtrip.txt" + payload = b"hello\nworld\n" + await self.backend.write_file(path, payload) + self.assertEqual(await self.backend.read_file(path), payload) + + async def test_write_creates_parent_dirs(self) -> None: + """``write_file`` creates missing parent directories.""" + path = f"{CONTAINER_WORKDIR}/a/b/c/file.txt" + await self.backend.write_file(path, b"x") + self.assertEqual(await self.backend.read_file(path), b"x") + + async def test_write_preserves_binary(self) -> None: + """Raw bytes (incl. NULs and high bytes) survive the round-trip.""" + path = f"{CONTAINER_WORKDIR}/bin.dat" + payload = b"a\r\nb\x00\xffc" + await self.backend.write_file(path, payload) + self.assertEqual(await self.backend.read_file(path), payload) + + async def test_read_missing_file_raises(self) -> None: + """Reading a non-existent file raises ``FileNotFoundError``.""" + with self.assertRaises(FileNotFoundError): + await self.backend.read_file(f"{CONTAINER_WORKDIR}/nope.txt") + + # ── derived filesystem helpers (shell-based) ─────────────────── + + async def test_file_exists_and_is_dir(self) -> None: + """``file_exists`` / ``is_dir`` reflect the container filesystem.""" + path = f"{CONTAINER_WORKDIR}/f.txt" + await self.backend.write_file(path, b"x") + self.assertTrue(await self.backend.file_exists(path)) + self.assertTrue(await self.backend.file_exists(CONTAINER_WORKDIR)) + self.assertTrue(await self.backend.is_dir(CONTAINER_WORKDIR)) + self.assertFalse(await self.backend.is_dir(path)) + self.assertFalse( + await self.backend.file_exists( + f"{CONTAINER_WORKDIR}/missing", + ), + ) + + async def test_list_dir(self) -> None: + """Non-recursive ``list_dir`` returns immediate child base names.""" + base = f"{CONTAINER_WORKDIR}/listing" + await self.backend.write_file(f"{base}/a.txt", b"x") + await self.backend.write_file(f"{base}/b.txt", b"x") + entries = await self.backend.list_dir(base) + self.assertEqual(sorted(entries), ["a.txt", "b.txt"]) + + async def test_list_dir_recursive(self) -> None: + """Recursive ``list_dir`` returns file paths underneath the root.""" + base = f"{CONTAINER_WORKDIR}/rec" + await self.backend.write_file(f"{base}/top.txt", b"x") + await self.backend.write_file(f"{base}/sub/nested.txt", b"x") + entries = await self.backend.list_dir(base, recursive=True) + basenames = sorted(e.rsplit("/", 1)[-1] for e in entries) + self.assertEqual(basenames, ["nested.txt", "top.txt"]) + + async def test_stat_mtime(self) -> None: + """``stat_mtime`` returns a float for an existing path, None else.""" + path = f"{CONTAINER_WORKDIR}/stat.txt" + await self.backend.write_file(path, b"x") + mtime = await self.backend.stat_mtime(path) + self.assertIsInstance(mtime, float) + self.assertIsNone( + await self.backend.stat_mtime( + f"{CONTAINER_WORKDIR}/missing", + ), + ) + + async def test_delete_path(self) -> None: + """``delete_path`` removes files and trees; missing is a no-op.""" + path = f"{CONTAINER_WORKDIR}/to_delete.txt" + await self.backend.write_file(path, b"x") + await self.backend.delete_path(path) + self.assertFalse(await self.backend.file_exists(path)) + + tree = f"{CONTAINER_WORKDIR}/tree" + await self.backend.write_file(f"{tree}/deep/f.txt", b"x") + await self.backend.delete_path(tree) + self.assertFalse(await self.backend.file_exists(tree)) + + # Deleting a non-existent path must not raise. + await self.backend.delete_path(f"{CONTAINER_WORKDIR}/missing") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/backend_e2b_test.py b/tests/backend_e2b_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e0f197baba3aadbb557b40e024b8a5e0c4a4a58a --- /dev/null +++ b/tests/backend_e2b_test.py @@ -0,0 +1,153 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Test cases for :class:`E2BBackend`. + +Validates that the three backend primitives (``exec_shell``, +``read_file``, ``write_file``) and the inherited shell-based filesystem +helpers behave correctly inside a real E2B cloud sandbox. + +The whole module is skipped unless the ``E2B_API_KEY`` environment +variable is set, because every test requires a live E2B sandbox. CI +runs without E2B credentials are therefore unaffected; when a key *is* +present the tests exercise the real ``commands.run`` / ``files.*`` APIs. + +A live sandbox is obtained by initializing an :class:`E2BWorkspace` and +reusing its already-wired :class:`E2BBackend` (``ws._backend``), which +avoids duplicating the sandbox bring-up logic here. +""" + +import os +import unittest +from unittest.async_case import IsolatedAsyncioTestCase + +from agentscope.tool import ExecResult +from agentscope.workspace import E2BWorkspace +from agentscope.workspace import E2BBackend +from agentscope.workspace._e2b._bootstrap import SANDBOX_WORKDIR + + +# ── E2B availability check ───────────────────────────────────────── + +_E2B_API_KEY = os.getenv("E2B_API_KEY", "") +_SKIP_REASON = "E2B_API_KEY environment variable is not set" + + +@unittest.skipUnless(_E2B_API_KEY, _SKIP_REASON) +class TestE2BBackend(IsolatedAsyncioTestCase): + """Test cases for ``E2BBackend`` against a live sandbox. + + Each test creates a real E2B cloud sandbox via ``E2BWorkspace`` and + tears it down (``close`` → sandbox pause) afterwards. + """ + + async def asyncSetUp(self) -> None: + """Start a workspace and reuse its wired backend.""" + self.workspace = E2BWorkspace(api_key=_E2B_API_KEY) + await self.workspace.initialize() + self.backend = self.workspace._backend + self.assertIsInstance(self.backend, E2BBackend) + + async def asyncTearDown(self) -> None: + """Pause / close the sandbox.""" + await self.workspace.close() + + # ── exec ─────────────────────────────────────────────────────── + + async def test_exec_returns_stdout(self) -> None: + """A program's stdout/exit code are captured into ``ExecResult``.""" + result = await self.backend.exec_shell(["echo", "hello world"]) + self.assertIsInstance(result, ExecResult) + self.assertTrue(result.ok()) + self.assertEqual(result.stdout.decode().strip(), "hello world") + + async def test_exec_nonzero_exit(self) -> None: + """A non-zero command exit is reported as a normal result.""" + result = await self.backend.exec_shell( + ["sh", "-c", "echo oops >&2; exit 4"], + ) + self.assertEqual(result.exit_code, 4) + self.assertIn("oops", result.stderr.decode()) + + async def test_exec_argv_quoting_preserved(self) -> None: + """An argv element with spaces / metacharacters survives the + POSIX-quote round-trip the backend does for ``commands.run``.""" + tricky = "a b c | ;" + result = await self.backend.exec_shell(["echo", tricky]) + self.assertTrue(result.ok()) + self.assertEqual(result.stdout.decode().rstrip("\n"), tricky) + + async def test_exec_cwd_default_is_workdir(self) -> None: + """With no explicit ``cwd`` the sandbox workdir is used.""" + result = await self.backend.exec_shell(["pwd"]) + self.assertTrue(result.ok()) + self.assertEqual(result.stdout.decode().strip(), SANDBOX_WORKDIR) + + # ── file I/O ─────────────────────────────────────────────────── + + async def test_write_then_read_roundtrip(self) -> None: + """Bytes written into the sandbox are read back verbatim.""" + path = f"{SANDBOX_WORKDIR}/roundtrip.txt" + payload = b"hello\nworld\n" + await self.backend.write_file(path, payload) + self.assertEqual(await self.backend.read_file(path), payload) + + async def test_write_creates_parent_dirs(self) -> None: + """``write_file`` creates missing parent directories.""" + path = f"{SANDBOX_WORKDIR}/a/b/c/file.txt" + await self.backend.write_file(path, b"x") + self.assertEqual(await self.backend.read_file(path), b"x") + + async def test_read_missing_file_raises(self) -> None: + """Reading a non-existent file raises ``FileNotFoundError``.""" + with self.assertRaises(FileNotFoundError): + await self.backend.read_file(f"{SANDBOX_WORKDIR}/nope.txt") + + # ── derived filesystem helpers (shell-based) ─────────────────── + + async def test_file_exists_and_is_dir(self) -> None: + """``file_exists`` / ``is_dir`` reflect the sandbox filesystem.""" + path = f"{SANDBOX_WORKDIR}/f.txt" + await self.backend.write_file(path, b"x") + self.assertTrue(await self.backend.file_exists(path)) + self.assertTrue(await self.backend.is_dir(SANDBOX_WORKDIR)) + self.assertFalse(await self.backend.is_dir(path)) + self.assertFalse( + await self.backend.file_exists(f"{SANDBOX_WORKDIR}/missing"), + ) + + async def test_list_dir(self) -> None: + """Non-recursive ``list_dir`` returns immediate child base names.""" + base = f"{SANDBOX_WORKDIR}/listing" + await self.backend.write_file(f"{base}/a.txt", b"x") + await self.backend.write_file(f"{base}/b.txt", b"x") + entries = await self.backend.list_dir(base) + self.assertEqual(sorted(entries), ["a.txt", "b.txt"]) + + async def test_stat_mtime(self) -> None: + """``stat_mtime`` returns a float for an existing path, None else.""" + path = f"{SANDBOX_WORKDIR}/stat.txt" + await self.backend.write_file(path, b"x") + mtime = await self.backend.stat_mtime(path) + self.assertIsInstance(mtime, float) + self.assertIsNone( + await self.backend.stat_mtime(f"{SANDBOX_WORKDIR}/missing"), + ) + + async def test_delete_path(self) -> None: + """``delete_path`` removes files and trees; missing is a no-op.""" + path = f"{SANDBOX_WORKDIR}/to_delete.txt" + await self.backend.write_file(path, b"x") + await self.backend.delete_path(path) + self.assertFalse(await self.backend.file_exists(path)) + + tree = f"{SANDBOX_WORKDIR}/tree" + await self.backend.write_file(f"{tree}/deep/f.txt", b"x") + await self.backend.delete_path(tree) + self.assertFalse(await self.backend.file_exists(tree)) + + # Deleting a non-existent path must not raise. + await self.backend.delete_path(f"{SANDBOX_WORKDIR}/missing") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/backend_local_test.py b/tests/backend_local_test.py new file mode 100644 index 0000000000000000000000000000000000000000..b79666db343455bf1b0d6ca94e8c5f64de3705c2 --- /dev/null +++ b/tests/backend_local_test.py @@ -0,0 +1,329 @@ +# -*- coding: utf-8 -*- +"""Test cases for :class:`LocalBackend` and the backend helpers. + +Exercises the three abstract primitives (``exec_shell``, ``read_file``, +``write_file``) plus the derived filesystem helpers (``file_exists``, +``is_dir``, ``list_dir``, ``stat_mtime``, ``delete_path``) of the +host-local backend, and the module-level ``normalize_newlines`` helper. + +``LocalBackend`` is designed to run on every platform (it spawns +programs from an argv list without a shell and implements the +filesystem helpers with native ``os.*`` calls), so the bulk of this +module runs on Windows too. Only the handful of cases that genuinely +rely on a POSIX shell / POSIX-only utilities are skipped on Windows. +""" + +import os +import sys +import tempfile +import unittest +from unittest.async_case import IsolatedAsyncioTestCase + +from agentscope.tool import ExecResult, LocalBackend +from agentscope.tool._builtin._backend import _normalize_newlines + +_IS_WINDOWS = sys.platform == "win32" + + +class TestNormalizeNewlines(unittest.TestCase): + """Unit tests for the ``normalize_newlines`` helper (pure, no I/O).""" + + def test_crlf_collapsed_to_lf(self) -> None: + """Windows ``\\r\\n`` is collapsed to a single ``\\n``.""" + self.assertEqual( + _normalize_newlines("a\r\nb\r\nc"), + "a\nb\nc", + ) + + def test_lone_cr_collapsed_to_lf(self) -> None: + """Classic-Mac lone ``\\r`` is collapsed to ``\\n``.""" + self.assertEqual(_normalize_newlines("a\rb\rc"), "a\nb\nc") + + def test_mixed_endings(self) -> None: + """A mix of ``\\r\\n``, ``\\r`` and ``\\n`` normalizes uniformly.""" + self.assertEqual( + _normalize_newlines("a\r\nb\rc\nd"), + "a\nb\nc\nd", + ) + + def test_plain_lf_unchanged(self) -> None: + """Text already using ``\\n`` is returned unchanged.""" + self.assertEqual(_normalize_newlines("a\nb\nc"), "a\nb\nc") + + def test_no_double_collapse(self) -> None: + """``\\r\\n`` becomes exactly one ``\\n`` (not two).""" + self.assertEqual(_normalize_newlines("a\r\n\r\nb"), "a\n\nb") + + +class TestLocalBackendExec(IsolatedAsyncioTestCase): + """Test cases for ``LocalBackend.exec_shell``.""" + + async def asyncSetUp(self) -> None: + """Build a fresh backend per test.""" + self.backend = LocalBackend() + + async def test_exec_returns_stdout(self) -> None: + """A program's stdout/exit code are captured into ``ExecResult``. + + Uses the current interpreter so the test is portable across + platforms (no reliance on ``echo`` / shell builtins). + """ + result = await self.backend.exec_shell( + [sys.executable, "-c", "print('hello world')"], + ) + self.assertIsInstance(result, ExecResult) + self.assertTrue(result.ok()) + self.assertEqual(result.exit_code, 0) + self.assertEqual(result.stdout.decode().strip(), "hello world") + self.assertEqual(result.stderr, b"") + + async def test_exec_captures_stderr_and_exit_code(self) -> None: + """A non-zero exit and stderr bytes are reported faithfully.""" + result = await self.backend.exec_shell( + [ + sys.executable, + "-c", + "import sys; sys.stderr.write('boom'); sys.exit(3)", + ], + ) + self.assertFalse(result.ok()) + self.assertEqual(result.exit_code, 3) + self.assertEqual(result.stderr.decode().strip(), "boom") + + async def test_exec_argv_not_split_by_shell(self) -> None: + """Arguments are passed verbatim (no shell word-splitting/globbing). + + A single argument containing spaces and shell metacharacters must + reach the program intact, proving no shell is interposed. + """ + tricky = "a b $(echo x) | & ; '\"" + result = await self.backend.exec_shell( + [sys.executable, "-c", "import sys; print(sys.argv[1])", tricky], + ) + self.assertTrue(result.ok()) + self.assertEqual(result.stdout.decode().rstrip("\r\n"), tricky) + + async def test_exec_cwd_is_respected(self) -> None: + """``cwd`` sets the subprocess working directory.""" + with tempfile.TemporaryDirectory() as tmp: + result = await self.backend.exec_shell( + [sys.executable, "-c", "import os; print(os.getcwd())"], + cwd=tmp, + ) + self.assertTrue(result.ok()) + self.assertEqual( + os.path.realpath(result.stdout.decode().strip()), + os.path.realpath(tmp), + ) + + async def test_missing_executable_returns_127(self) -> None: + """An unspawnable executable yields exit code 127 (not an exception). + + Mirrors a shell's "command not found" so callers see a normal + non-zero ``ExecResult``. + """ + result = await self.backend.exec_shell( + ["this_executable_does_not_exist_xyz", "--nope"], + ) + self.assertEqual(result.exit_code, 127) + self.assertEqual(result.stdout, b"") + self.assertNotEqual(result.stderr, b"") + + async def test_timeout_returns_minus_one(self) -> None: + """A command exceeding ``timeout`` is killed and reports -1. + + The sentinel ``ExecResult(exit_code=-1, stderr=b"timed out")`` is + what Grep/Glob check for, so it is asserted exactly. + """ + result = await self.backend.exec_shell( + [sys.executable, "-c", "import time; time.sleep(10)"], + timeout=0.2, + ) + self.assertEqual(result.exit_code, -1) + self.assertEqual(result.stderr, b"timed out") + + +class TestLocalBackendFileIO(IsolatedAsyncioTestCase): + """Test cases for ``read_file`` / ``write_file`` round-trips.""" + + async def asyncSetUp(self) -> None: + """Build a backend and a temp dir per test.""" + # pylint: disable=consider-using-with + self.backend = LocalBackend() + self.temp_dir = tempfile.TemporaryDirectory() + + async def asyncTearDown(self) -> None: + """Drop the temp dir.""" + self.temp_dir.cleanup() + + async def test_write_then_read_roundtrip(self) -> None: + """Bytes written are read back verbatim.""" + path = os.path.join(self.temp_dir.name, "file.txt") + payload = b"hello\nworld\n" + await self.backend.write_file(path, payload) + self.assertEqual(await self.backend.read_file(path), payload) + + async def test_write_creates_parent_dirs(self) -> None: + """``write_file`` creates missing parent directories.""" + path = os.path.join(self.temp_dir.name, "a", "b", "c", "file.txt") + await self.backend.write_file(path, b"x") + self.assertTrue(os.path.exists(path)) + self.assertEqual(await self.backend.read_file(path), b"x") + + async def test_write_preserves_binary_and_crlf(self) -> None: + """Raw bytes (incl. ``\\r\\n`` and NULs) survive the round-trip. + + The backend deals in raw bytes; newline normalization happens + later in the text layer, never here. + """ + path = os.path.join(self.temp_dir.name, "bin.dat") + payload = b"a\r\nb\x00\xffc" + await self.backend.write_file(path, payload) + self.assertEqual(await self.backend.read_file(path), payload) + + async def test_read_missing_file_raises(self) -> None: + """Reading a non-existent file raises ``FileNotFoundError``.""" + path = os.path.join(self.temp_dir.name, "nope.txt") + with self.assertRaises(FileNotFoundError): + await self.backend.read_file(path) + + +class TestLocalBackendFilesystemHelpers(IsolatedAsyncioTestCase): + """Test cases for the derived filesystem helpers (native ``os.*``).""" + + async def asyncSetUp(self) -> None: + """Build a backend and a temp dir per test.""" + # pylint: disable=consider-using-with + self.backend = LocalBackend() + self.temp_dir = tempfile.TemporaryDirectory() + + async def asyncTearDown(self) -> None: + """Drop the temp dir.""" + self.temp_dir.cleanup() + + async def test_file_exists(self) -> None: + """``file_exists`` is True for files and dirs, False otherwise.""" + path = os.path.join(self.temp_dir.name, "f.txt") + await self.backend.write_file(path, b"x") + self.assertTrue(await self.backend.file_exists(path)) + self.assertTrue(await self.backend.file_exists(self.temp_dir.name)) + self.assertFalse( + await self.backend.file_exists( + os.path.join(self.temp_dir.name, "missing"), + ), + ) + + async def test_is_dir(self) -> None: + """``is_dir`` distinguishes directories from files.""" + path = os.path.join(self.temp_dir.name, "f.txt") + await self.backend.write_file(path, b"x") + self.assertTrue(await self.backend.is_dir(self.temp_dir.name)) + self.assertFalse(await self.backend.is_dir(path)) + self.assertFalse( + await self.backend.is_dir( + os.path.join(self.temp_dir.name, "missing"), + ), + ) + + async def test_list_dir_shallow(self) -> None: + """Non-recursive ``list_dir`` returns immediate child base names.""" + for name in ("a.txt", "b.txt"): + await self.backend.write_file( + os.path.join(self.temp_dir.name, name), + b"x", + ) + os.makedirs(os.path.join(self.temp_dir.name, "sub")) + entries = await self.backend.list_dir(self.temp_dir.name) + self.assertEqual(sorted(entries), ["a.txt", "b.txt", "sub"]) + + async def test_list_dir_recursive(self) -> None: + """Recursive ``list_dir`` returns file paths underneath the root.""" + await self.backend.write_file( + os.path.join(self.temp_dir.name, "top.txt"), + b"x", + ) + await self.backend.write_file( + os.path.join(self.temp_dir.name, "sub", "nested.txt"), + b"x", + ) + entries = await self.backend.list_dir( + self.temp_dir.name, + recursive=True, + ) + basenames = sorted(os.path.basename(e) for e in entries) + self.assertEqual(basenames, ["nested.txt", "top.txt"]) + + async def test_stat_mtime(self) -> None: + """``stat_mtime`` returns a float for an existing path, None else.""" + path = os.path.join(self.temp_dir.name, "f.txt") + await self.backend.write_file(path, b"x") + mtime = await self.backend.stat_mtime(path) + self.assertIsInstance(mtime, float) + self.assertIsNone( + await self.backend.stat_mtime( + os.path.join(self.temp_dir.name, "missing"), + ), + ) + + async def test_delete_path_file(self) -> None: + """``delete_path`` removes a single file.""" + path = os.path.join(self.temp_dir.name, "f.txt") + await self.backend.write_file(path, b"x") + await self.backend.delete_path(path) + self.assertFalse(os.path.exists(path)) + + async def test_delete_path_tree(self) -> None: + """``delete_path`` removes a directory tree recursively.""" + nested = os.path.join(self.temp_dir.name, "d", "e") + os.makedirs(nested) + await self.backend.write_file( + os.path.join(nested, "f.txt"), + b"x", + ) + target = os.path.join(self.temp_dir.name, "d") + await self.backend.delete_path(target) + self.assertFalse(os.path.exists(target)) + + async def test_delete_path_missing_is_noop(self) -> None: + """Deleting a non-existent path is a silent no-op (like rm -rf).""" + # Must not raise. + await self.backend.delete_path( + os.path.join(self.temp_dir.name, "missing"), + ) + + +@unittest.skipIf( + _IS_WINDOWS, + "POSIX shell (/bin/sh) is not available on Windows", +) +class TestLocalBackendShellWrapping(IsolatedAsyncioTestCase): + """Cases that explicitly use ``sh -c`` (POSIX-only, skipped on Windows). + + The backend primitive never invokes a shell itself; these tests cover + the documented escape hatch where a *caller* wraps a command line as + ``["/bin/sh", "-c", line]`` to use pipes / redirects / ``&&``. + """ + + async def asyncSetUp(self) -> None: + """Build a backend per test.""" + self.backend = LocalBackend() + + async def test_sh_c_pipeline(self) -> None: + """A piped command line runs when wrapped in ``sh -c``.""" + result = await self.backend.exec_shell( + ["/bin/sh", "-c", "printf 'a\\nb\\na\\n' | sort | uniq"], + ) + self.assertTrue(result.ok()) + self.assertEqual(result.stdout.decode().split(), ["a", "b"]) + + async def test_sh_c_and_chain(self) -> None: + """``&&`` chaining works through ``sh -c``.""" + result = await self.backend.exec_shell( + ["/bin/sh", "-c", "true && echo chained"], + ) + self.assertTrue(result.ok()) + self.assertEqual(result.stdout.decode().strip(), "chained") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/blob_store_s3_test.py b/tests/blob_store_s3_test.py new file mode 100644 index 0000000000000000000000000000000000000000..31121782d4b4afd142822fc12e538766c00651af --- /dev/null +++ b/tests/blob_store_s3_test.py @@ -0,0 +1,159 @@ +# -*- coding: utf-8 -*- +"""Tests for :class:`S3BlobStore`. + +Uses ``moto.server.ThreadedMotoServer``: a real HTTP S3 endpoint +served in-process on a thread. The aioboto3 client talks to it over +TCP the way it would talk to AWS / MinIO / R2 — so the test +exercises the real wire path (and the ``endpoint_url`` plumbing the +production deployment will rely on for non-AWS services). + +We deliberately avoid ``moto.mock_aws()`` because that decorator +patches ``botocore``'s endpoint layer, which does not play nicely +with ``aiobotocore``: the patched response object is sync, but +aiobotocore awaits ``http_response.content`` and gets a TypeError. +The threaded server bypasses every monkey-patch path and gives us +end-to-end coverage. + +Properties under test: + +- A blob written through ``write_stream`` reads back identical + bytes via ``open``. +- The URI emitted is ``s3://{bucket}/{key}`` and round-trips + through ``exists`` / ``delete``. +- ``exists`` distinguishes present from missing without raising. +- ``delete`` is idempotent. +- Bad URIs raise ``ValueError`` before any network call. +- A store used outside ``async with`` raises ``RuntimeError`` — + callers MUST manage lifecycle via :class:`AsyncExitStack`. +""" +import io +import os +import socket +from typing import TYPE_CHECKING +from unittest import IsolatedAsyncioTestCase + +import boto3 +from moto.server import ThreadedMotoServer + +if TYPE_CHECKING: + from agentscope.app.rag.blob_store import S3BlobStore + + +_BUCKET = "test-blobs" +_REGION = "us-east-1" + + +def _pick_port() -> int: + """Find an unused TCP port to bind moto on.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +class S3BlobStoreTest(IsolatedAsyncioTestCase): + """Round-trip + URI parsing + streaming behaviour.""" + + def setUp(self) -> None: + # ThreadedMotoServer needs explicit credentials — without + # them aiobotocore looks for the AWS config chain and may + # fail on CI runners without credentials baked in. + self._env_overrides = { + "AWS_ACCESS_KEY_ID": "testing", + "AWS_SECRET_ACCESS_KEY": "testing", + "AWS_SECURITY_TOKEN": "testing", + "AWS_SESSION_TOKEN": "testing", + "AWS_DEFAULT_REGION": _REGION, + } + self._saved_env = {k: os.environ.get(k) for k in self._env_overrides} + for k, v in self._env_overrides.items(): + os.environ[k] = v + + port = _pick_port() + self._endpoint_url = f"http://127.0.0.1:{port}" + self._server = ThreadedMotoServer(ip_address="127.0.0.1", port=port) + self._server.start() + + client = boto3.client( + "s3", + region_name=_REGION, + endpoint_url=self._endpoint_url, + ) + client.create_bucket(Bucket=_BUCKET) + + def tearDown(self) -> None: + self._server.stop() + for k, prev in self._saved_env.items(): + if prev is None: + os.environ.pop(k, None) + else: + os.environ[k] = prev + + def _store(self) -> "S3BlobStore": + """Build a fresh :class:`S3BlobStore` pointing at the moto endpoint. + + Returns: + `S3BlobStore`: + The store under test, configured for the moto-backed + ``_BUCKET`` in ``_REGION``. + """ + # Import inside the method so the skip decorator can fire + # without the import path running. + from agentscope.app.rag.blob_store import S3BlobStore + + return S3BlobStore( + bucket=_BUCKET, + region_name=_REGION, + endpoint_url=self._endpoint_url, + use_ssl=False, + ) + + async def test_write_open_round_trip(self) -> None: + """A blob written through ``write_stream`` reads back the + same bytes via ``open`` and the URI is the documented + ``s3://{bucket}/{key}`` shape.""" + async with self._store() as store: + payload = b"hello world\n" * 100 # ~1.2 KiB + uri = await store.write_stream( + "kb/abc/doc-1", + io.BytesIO(payload), + ) + self.assertEqual(uri, "s3://test-blobs/kb/abc/doc-1") + + received: list[bytes] = [] + async with store.open(uri) as fp: + while True: + chunk = await fp.read(64) + if not chunk: + break + received.append(chunk) + self.assertEqual(b"".join(received), payload) + + async def test_exists_and_delete(self) -> None: + """``exists`` flips from ``True`` to ``False`` after + ``delete``; ``delete`` is idempotent on a missing key.""" + async with self._store() as store: + uri = await store.write_stream( + "kb/abc/doc-2", + io.BytesIO(b"data"), + ) + self.assertTrue(await store.exists(uri)) + await store.delete(uri) + self.assertFalse(await store.exists(uri)) + # Second delete is a no-op. + await store.delete(uri) + + async def test_bad_uri_scheme_rejected(self) -> None: + """Non-``s3://`` URIs are refused before any network call.""" + async with self._store() as store: + with self.assertRaises(ValueError): + await store.exists("local://kb/abc/doc-3") + with self.assertRaises(ValueError): + await store.delete("s3://wrong-bucket/kb/abc/doc-4") + + async def test_methods_require_aenter(self) -> None: + """Calling a blob method without entering the context fails + loudly — production code MUST manage the lifecycle through + ``async with`` / ``AsyncExitStack``.""" + store = self._store() # not entered + with self.assertRaises(RuntimeError): + await store.write_stream("kb/abc/doc-5", io.BytesIO(b"")) diff --git a/tests/builtin_bash_test.py b/tests/builtin_bash_test.py new file mode 100644 index 0000000000000000000000000000000000000000..96729c0641dcc59854a379bc3b5167b71c3eb659 --- /dev/null +++ b/tests/builtin_bash_test.py @@ -0,0 +1,528 @@ +# -*- coding: utf-8 -*- +"""Bash tool test case.""" + +import os +import sys +import unittest +from unittest.async_case import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, MagicMock, patch + +from agentscope.message import TextBlock +from agentscope.permission import ( + PermissionBehavior, + PermissionContext, + PermissionRule, +) +from agentscope.tool import Bash, ToolChunk +from agentscope.tool._builtin._backend import ( + _subprocess_creation_kwargs, +) + + +class BashSubprocessKwargsTest(unittest.TestCase): + """Test platform-specific subprocess kwargs.""" + + def test_non_windows_subprocess_kwargs_are_empty(self) -> None: + """On non-Windows the helper returns no extra subprocess kwargs.""" + with patch( + "agentscope.tool._builtin._backend.os.name", + "posix", + ): + self.assertEqual(_subprocess_creation_kwargs(), {}) + + def test_windows_subprocess_kwargs_hide_console(self) -> None: + """On Windows the helper sets ``creationflags`` to hide the console.""" + with patch("agentscope.tool._builtin._backend.os.name", "nt"): + self.assertEqual( + _subprocess_creation_kwargs(), + {"creationflags": 0x08000000}, + ) + + +class BashCwdTest(IsolatedAsyncioTestCase): + """Test Bash working-directory wiring.""" + + async def test_cwd_is_passed_to_subprocess(self) -> None: + """The constructor-level cwd should be used for each command.""" + process = MagicMock() + process.returncode = 0 + process.communicate = AsyncMock(return_value=(b"ok\n", b"")) + + create_process = AsyncMock(return_value=process) + with patch( + "agentscope.tool._builtin._backend." + "asyncio.create_subprocess_exec", + create_process, + ): + chunks = [] + async for chunk in await Bash(cwd="workspace")(command="pwd"): + chunks.append(chunk) + + # cwd is forwarded, and the command line is wrapped in the + # platform's native shell (the backend primitive runs an argv + # without a shell): ``cmd /c`` on Windows, ``/bin/sh -c`` else. + self.assertEqual(create_process.call_args.kwargs["cwd"], "workspace") + expected_argv = ( + ("cmd", "/c", "pwd") + if os.name == "nt" + else ("/bin/sh", "-c", "pwd") + ) + self.assertEqual( + create_process.call_args.args, + expected_argv, + ) + self.assertEqual(chunks[0].state, "running") + + +@unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", +) +class BashToolTest(IsolatedAsyncioTestCase): + """The bash tool test case.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.bash_tool = Bash() + + async def test_tool_properties(self) -> None: + """Test bash tool properties.""" + self.assertEqual(self.bash_tool.name, "Bash") + self.assertIsInstance(self.bash_tool.description, str) + self.assertIsInstance(self.bash_tool.input_schema, dict) + self.assertFalse(self.bash_tool.is_mcp) + self.assertFalse(self.bash_tool.is_read_only) + self.assertFalse(self.bash_tool.is_concurrency_safe) + + async def test_check_permissions(self) -> None: + """Test bash tool permission checking.""" + + context = PermissionContext() + tool_input = {"command": "echo hello"} + decision = await self.bash_tool.check_permissions(tool_input, context) + + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + async def test_simple_command(self) -> None: + """Test executing a simple bash command.""" + chunks = [] + async for chunk in await self.bash_tool(command="echo 'Hello World'"): + chunks.append(chunk) + + self.assertEqual(len(chunks), 1) + self.assertIsInstance(chunks[0], ToolChunk) + self.assertEqual(chunks[0].state, "running") + self.assertTrue(chunks[0].is_last) + self.assertEqual(len(chunks[0].content), 1) + self.assertIsInstance(chunks[0].content[0], TextBlock) + self.assertIn("Hello World", chunks[0].content[0].text) + + async def test_command_with_error(self) -> None: + """Test executing a command that fails.""" + chunks = [] + async for chunk in await self.bash_tool(command="exit 1"): + chunks.append(chunk) + + self.assertEqual(len(chunks), 1) + self.assertEqual(chunks[0].state, "error") + self.assertTrue(chunks[0].is_last) + + @unittest.skipIf( + sys.platform == "win32", + "sleep command not available on Windows", + ) + async def test_command_timeout(self) -> None: + """Test command timeout.""" + chunks = [] + async for chunk in await self.bash_tool( + command="sleep 10", + timeout=100, # 100ms timeout + ): + chunks.append(chunk) + + self.assertEqual(len(chunks), 1) + self.assertEqual(chunks[0].state, "error") + self.assertIn("timed out", chunks[0].content[0].text.lower()) + + +@unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", +) +class BashToolInjectionCheckTest(IsolatedAsyncioTestCase): + """Test injection detection in Bash tool permission checks.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.bash_tool = Bash() + self.context = PermissionContext() + + async def test_command_substitution_blocked(self) -> None: + """Test that command substitution is blocked.""" + test_cases = [ + "ls $(pwd)", + "rm $(find . -name '*.tmp')", + "cat `which python`", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + decision = await self.bash_tool.check_permissions( + {"command": cmd}, + self.context, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + self.assertIn("command_substitution", decision.message) + + async def test_control_flow_blocked(self) -> None: + """Test that control flow structures are blocked.""" + + test_cases = [ + "for f in *.txt; do cat $f; done", + "while read line; do echo $line; done < file.txt", + "if [ -f file.txt ]; then cat file.txt; fi", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + decision = await self.bash_tool.check_permissions( + {"command": cmd}, + self.context, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + self.assertIn( + "cannot be statically analyzed", + decision.message, + ) + + async def test_subshell_blocked(self) -> None: + """Test that subshells are blocked.""" + + cmd = "(cd /tmp && ls)" + decision = await self.bash_tool.check_permissions( + {"command": cmd}, + self.context, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + self.assertIn("subshell", decision.message) + + async def test_injection_check_before_readonly(self) -> None: + """Test that injection check runs before read-only check.""" + + # ls is read-only, but $(rm -rf /) is dangerous + cmd = "ls $(rm -rf /)" + decision = await self.bash_tool.check_permissions( + {"command": cmd}, + self.context, + ) + # Should be blocked by injection check, not allowed as read-only + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + self.assertIn("command_substitution", decision.message) + + async def test_safe_commands_pass(self) -> None: + """Test that safe commands pass injection check.""" + + safe_commands = [ + "ls -la", + "cat file.txt", + "git status", + "echo 'hello world'", + ] + for cmd in safe_commands: + with self.subTest(cmd=cmd): + decision = await self.bash_tool.check_permissions( + {"command": cmd}, + self.context, + ) + # Should pass injection check (either ALLOW or PASSTHROUGH) + self.assertNotEqual(decision.behavior, PermissionBehavior.ASK) + if decision.behavior == PermissionBehavior.ASK: + self.assertNotIn( + "cannot be statically analyzed", + decision.message, + ) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.bash_tool = None + self.context = None + + +@unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", +) +class BashToolMatchRuleTest(IsolatedAsyncioTestCase): + """Test cases for Bash tool match_rule and generate_suggestions.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.bash_tool = Bash() + + async def test_match_rule_prefix_pattern(self) -> None: + """Test match_rule with prefix patterns (e.g., git:*).""" + # Test exact command match + self.assertTrue( + await self.bash_tool.match_rule( + "git:*", + {"command": "git"}, + ), + ) + + # Test command with arguments + self.assertTrue( + await self.bash_tool.match_rule( + "git:*", + {"command": "git status"}, + ), + ) + + # Test non-matching command + self.assertFalse( + await self.bash_tool.match_rule( + "git:*", + {"command": "npm install"}, + ), + ) + + async def test_match_rule_wildcard_pattern(self) -> None: + """Test match_rule with wildcard patterns.""" + # Test wildcard matching + self.assertTrue( + await self.bash_tool.match_rule( + "git * -m *", + {"command": "git commit -m 'test'"}, + ), + ) + + # Test non-matching wildcard + self.assertFalse( + await self.bash_tool.match_rule( + "git * -m *", + {"command": "git status"}, + ), + ) + + async def test_match_rule_substring_pattern(self) -> None: + """Test match_rule with substring patterns.""" + # Test substring matching + self.assertTrue( + await self.bash_tool.match_rule( + "install", + {"command": "npm install package"}, + ), + ) + + # Test non-matching substring + self.assertFalse( + await self.bash_tool.match_rule( + "install", + {"command": "npm run build"}, + ), + ) + + async def test_match_rule_escaped_characters(self) -> None: + """Test match_rule with escaped characters.""" + # Test escaped asterisk + self.assertTrue( + await self.bash_tool.match_rule( + r"echo \*", + {"command": "echo *"}, + ), + ) + + # Test escaped backslash + self.assertTrue( + await self.bash_tool.match_rule( + r"echo \\", + {"command": "echo \\"}, + ), + ) + + async def test_generate_suggestions(self) -> None: + """Test generate_suggestions for bash commands.""" + + # Test two-word command + suggestions = await self.bash_tool.generate_suggestions( + {"command": "git commit -m 'test'"}, + ) + + self.assertIsInstance(suggestions, list) + self.assertGreater(len(suggestions), 0) + self.assertIsInstance(suggestions[0], PermissionRule) + + # Should suggest "git commit:*" + suggestion_contents = [s.rule_content for s in suggestions] + self.assertIn("git commit:*", suggestion_contents) + + async def test_generate_suggestions_single_word(self) -> None: + """Test generate_suggestions for single-word commands.""" + suggestions = await self.bash_tool.generate_suggestions( + {"command": "npm install"}, + ) + + self.assertGreater(len(suggestions), 0) + + # Should suggest "npm install:*" + suggestion_contents = [s.rule_content for s in suggestions] + self.assertIn("npm install:*", suggestion_contents) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.bash_tool = None + + +@unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", +) +class BashToolDangerousRemovalTest(IsolatedAsyncioTestCase): + """Test dangerous removal path detection in Bash tool.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.bash_tool = Bash() + self.context = PermissionContext() + + async def test_rm_root_blocked(self) -> None: + """Test that rm -rf / is blocked.""" + + cmd = "rm -rf /" + decision = await self.bash_tool.check_permissions( + {"command": cmd}, + self.context, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + # Can be blocked by either dangerous command pattern or dangerous + # removal path check + self.assertTrue( + "Dangerous removal operation" in decision.message + or "dangerous pattern" in decision.message, + ) + + async def test_rm_root_children_blocked(self) -> None: + """Test that rm -rf /usr, /etc, etc. are blocked.""" + + test_cases = [ + "rm -rf /usr", + "rm -rf /etc", + "rm -rf /tmp", + "rm -rf /var", + "rm -rf /bin", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + decision = await self.bash_tool.check_permissions( + {"command": cmd}, + self.context, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + # Can be blocked by either dangerous command pattern or + # dangerous removal path check + self.assertTrue( + "Dangerous removal operation" in decision.message + or "dangerous pattern" in decision.message, + ) + + async def test_rm_home_blocked(self) -> None: + """Test that rm -rf ~ is blocked.""" + + cmd = "rm -rf ~" + decision = await self.bash_tool.check_permissions( + {"command": cmd}, + self.context, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + # Can be blocked by either dangerous command pattern or dangerous + # removal path check + self.assertTrue( + "Dangerous removal operation" in decision.message + or "dangerous pattern" in decision.message, + ) + + async def test_rm_wildcard_blocked(self) -> None: + """Test that rm -rf * and rm -rf /* are blocked.""" + + test_cases = [ + "rm -rf *", + "rm -rf /*", + "rm -rf /tmp/*", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + decision = await self.bash_tool.check_permissions( + {"command": cmd}, + self.context, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + # Can be blocked by either dangerous command pattern or + # dangerous removal path check + self.assertTrue( + "Dangerous removal operation" in decision.message + or "dangerous pattern" in decision.message, + ) + + async def test_rmdir_dangerous_paths_blocked(self) -> None: + """Test that rmdir on dangerous paths is blocked.""" + + test_cases = [ + "rmdir /", + "rmdir /usr", + "rmdir ~", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + decision = await self.bash_tool.check_permissions( + {"command": cmd}, + self.context, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + self.assertIn("Dangerous removal operation", decision.message) + + async def test_safe_rm_commands_pass(self) -> None: + """Test that safe rm commands pass dangerous removal check.""" + + safe_commands = [ + "rm file.txt", + "rm -f temp.log", + "rm -rf /tmp/my_project/build", + "rm -rf ./node_modules", + ] + for cmd in safe_commands: + with self.subTest(cmd=cmd): + decision = await self.bash_tool.check_permissions( + {"command": cmd}, + self.context, + ) + # Should not be blocked by dangerous removal check + # (may still be blocked by other checks) + if decision.behavior == PermissionBehavior.ASK: + self.assertNotIn( + "Dangerous removal operation", + decision.message, + ) + + async def test_compound_commands_with_dangerous_removal(self) -> None: + """Test compound commands containing dangerous removal.""" + + test_cases = [ + "ls && rm -rf /", + "cd /tmp && rm -rf /usr", + "echo start; rm -rf ~; echo end", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + decision = await self.bash_tool.check_permissions( + {"command": cmd}, + self.context, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + # Can be blocked by either dangerous command pattern or + # dangerous removal path check + self.assertTrue( + "Dangerous removal operation" in decision.message + or "dangerous pattern" in decision.message, + ) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.bash_tool = None + self.context = None diff --git a/tests/builtin_edit_test.py b/tests/builtin_edit_test.py new file mode 100644 index 0000000000000000000000000000000000000000..3cf2f5b32375aaf1e74b8ee7d6fcc78af542dc05 --- /dev/null +++ b/tests/builtin_edit_test.py @@ -0,0 +1,182 @@ +# -*- coding: utf-8 -*- +"""Edit tool test case.""" +import os +import tempfile +from unittest.async_case import IsolatedAsyncioTestCase + +from agentscope.tool import Edit +from agentscope.permission import ( + PermissionContext, + PermissionBehavior, + PermissionRule, +) + + +class EditToolTest(IsolatedAsyncioTestCase): + """The edit tool test case.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.edit_tool = Edit() + # Create a temporary file for testing + self.temp_file = tempfile.NamedTemporaryFile( + mode="w", + delete=False, + suffix=".txt", + ) + self.temp_file.write("Hello World\nThis is a test\n") + self.temp_file.close() + + async def asyncTearDown(self) -> None: + """Clean up temporary files.""" + if os.path.exists(self.temp_file.name): + os.unlink(self.temp_file.name) + + async def test_tool_properties(self) -> None: + """Test edit tool properties.""" + self.assertEqual(self.edit_tool.name, "Edit") + self.assertIsInstance(self.edit_tool.description, str) + self.assertIsInstance(self.edit_tool.input_schema, dict) + self.assertFalse(self.edit_tool.is_mcp) + self.assertFalse(self.edit_tool.is_read_only) + self.assertFalse(self.edit_tool.is_concurrency_safe) + + async def test_check_permissions(self) -> None: + """Test edit tool permission checking. + + Edit tool should return PASSTHROUGH for non-dangerous paths, + allowing PermissionEngine to check allow rules. + """ + context = PermissionContext() + tool_input = {"file_path": "/tmp/test.txt"} + decision = await self.edit_tool.check_permissions(tool_input, context) + + self.assertEqual(decision.behavior, PermissionBehavior.PASSTHROUGH) + + async def test_simple_edit(self) -> None: + """Test simple file editing.""" + chunk = await self.edit_tool( + file_path=self.temp_file.name, + old_string="Hello World", + new_string="Hello Python", + ) + + self.assertEqual(chunk.state, "running") + self.assertTrue(chunk.is_last) + + # Verify file content + with open(self.temp_file.name, "r", encoding="utf-8") as f: + content = f.read() + self.assertIn("Hello Python", content) + self.assertNotIn("Hello World", content) + + async def test_edit_not_found(self) -> None: + """Test editing with string not found.""" + chunk = await self.edit_tool( + file_path=self.temp_file.name, + old_string="NonExistent", + new_string="Something", + ) + + self.assertEqual(chunk.state, "error") + self.assertIn("not found", chunk.content[0].text) + + async def test_edit_multiple_occurrences(self) -> None: + """Test editing with multiple occurrences.""" + # Write file with duplicate content + with open(self.temp_file.name, "w", encoding="utf-8") as f: + f.write("test\ntest\ntest\n") + + chunk = await self.edit_tool( + file_path=self.temp_file.name, + old_string="test", + new_string="replaced", + ) + + # Should fail without replace_all + self.assertEqual(chunk.state, "error") + + async def test_edit_replace_all(self) -> None: + """Test editing with replace_all flag.""" + # Write file with duplicate content + with open(self.temp_file.name, "w", encoding="utf-8") as f: + f.write("test\ntest\ntest\n") + + chunk = await self.edit_tool( + file_path=self.temp_file.name, + old_string="test", + new_string="replaced", + replace_all=True, + ) + + self.assertEqual(chunk.state, "running") + + # Verify all occurrences replaced + with open(self.temp_file.name, "r", encoding="utf-8") as f: + content = f.read() + self.assertEqual(content.count("replaced"), 3) + self.assertEqual(content.count("test"), 0) + + async def test_match_rule_glob_pattern(self) -> None: + """Test match_rule with glob patterns.""" + # Test exact match + self.assertTrue( + await self.edit_tool.match_rule( + "test.py", + {"file_path": "test.py"}, + ), + ) + + # Test wildcard pattern + self.assertTrue( + await self.edit_tool.match_rule( + "*.py", + {"file_path": "test.py"}, + ), + ) + + # Test directory pattern + self.assertTrue( + await self.edit_tool.match_rule( + "/tmp/**", + {"file_path": "/tmp/test.py"}, + ), + ) + + # Test non-matching pattern + self.assertFalse( + await self.edit_tool.match_rule( + "*.txt", + {"file_path": "test.py"}, + ), + ) + + # Test empty file_path + self.assertFalse( + await self.edit_tool.match_rule( + "*.py", + {"file_path": ""}, + ), + ) + + async def test_generate_suggestions(self) -> None: + """Test generate_suggestions for file operations.""" + + # Test suggestion for file in subdirectory + suggestions = await self.edit_tool.generate_suggestions( + {"file_path": "/tmp/project/src/main.py"}, + ) + + self.assertIsInstance(suggestions, list) + self.assertGreater(len(suggestions), 0) + self.assertIsInstance(suggestions[0], PermissionRule) + + # Should suggest parent directory pattern + suggestion_contents = [s.rule_content for s in suggestions] + self.assertIn("/tmp/project/src/**", suggestion_contents) + + # Test suggestion for file in root + suggestions = await self.edit_tool.generate_suggestions( + {"file_path": "/test.py"}, + ) + self.assertGreater(len(suggestions), 0) diff --git a/tests/builtin_file_cache_test.py b/tests/builtin_file_cache_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e448e5a1d167fdfa1f3f0be5913b268ffb1674e8 --- /dev/null +++ b/tests/builtin_file_cache_test.py @@ -0,0 +1,432 @@ +# -*- coding: utf-8 -*- +"""File cache test case for Read/Write/Edit tools.""" +import os +import tempfile +from unittest.async_case import IsolatedAsyncioTestCase + +from agentscope.state import AgentState +from agentscope.tool import Read, Write, Edit + + +class FileCacheTest(IsolatedAsyncioTestCase): + """Test file cache functionality for Read/Write/Edit tools.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.read_tool = Read() + self.write_tool = Write() + self.edit_tool = Edit() + self.state = AgentState() + + # Create a temporary directory + self.temp_dir = tempfile.mkdtemp() + self.test_file = os.path.join(self.temp_dir, "test.txt") + + async def asyncTearDown(self) -> None: + """Clean up temporary files.""" + import shutil + + if os.path.exists(self.temp_dir): + shutil.rmtree(self.temp_dir) + + async def test_edit_without_read(self) -> None: + """Test Edit fails when file not read first.""" + # Create a file + with open(self.test_file, "w", encoding="utf-8") as f: + f.write("Hello World\n") + + # Try to edit without reading first + chunk = await self.edit_tool( + file_path=self.test_file, + old_string="Hello", + new_string="Hi", + _agent_state=self.state, + ) + + # Should fail with error + self.assertEqual(chunk.state, "error") + self.assertIn("must first read", chunk.content[0].text) + + async def test_write_without_read(self) -> None: + """Test Write fails when existing file not read first.""" + # Create a file + with open(self.test_file, "w", encoding="utf-8") as f: + f.write("Existing content\n") + + # Try to write without reading first + chunk = await self.write_tool( + file_path=self.test_file, + content="New content\n", + _agent_state=self.state, + ) + + # Should fail with error + self.assertEqual(chunk.state, "error") + self.assertIn("has not been read yet", chunk.content[0].text) + + async def test_write_new_file_without_read(self) -> None: + """Test Write succeeds for new file without reading.""" + new_file = os.path.join(self.temp_dir, "new_file.txt") + + # Write to a new file (doesn't exist yet) + chunk = await self.write_tool( + file_path=new_file, + content="New file content\n", + _agent_state=self.state, + ) + + # Should succeed + self.assertEqual(chunk.state, "running") + self.assertTrue(os.path.exists(new_file)) + + with open(new_file, "r", encoding="utf-8") as f: + content = f.read() + self.assertEqual(content, "New file content\n") + + async def test_edit_after_read(self) -> None: + """Test Edit succeeds after reading file.""" + # Create a file + with open(self.test_file, "w", encoding="utf-8") as f: + f.write("Hello World\n") + + # Read the file first + read_chunk = await self.read_tool( + file_path=self.test_file, + _agent_state=self.state, + ) + self.assertEqual(read_chunk.state, "running") + + # Now edit should succeed + edit_chunk = await self.edit_tool( + file_path=self.test_file, + old_string="Hello", + new_string="Hi", + _agent_state=self.state, + ) + + self.assertEqual(edit_chunk.state, "running") + + # Verify file was edited + with open(self.test_file, "r", encoding="utf-8") as f: + content = f.read() + self.assertEqual(content, "Hi World\n") + + async def test_write_after_read(self) -> None: + """Test Write succeeds after reading file.""" + # Create a file + with open(self.test_file, "w", encoding="utf-8") as f: + f.write("Old content\n") + + # Read the file first + read_chunk = await self.read_tool( + file_path=self.test_file, + _agent_state=self.state, + ) + self.assertEqual(read_chunk.state, "running") + + # Now write should succeed + write_chunk = await self.write_tool( + file_path=self.test_file, + content="New content\n", + _agent_state=self.state, + ) + + self.assertEqual(write_chunk.state, "running") + + # Verify file was overwritten + with open(self.test_file, "r", encoding="utf-8") as f: + content = f.read() + self.assertEqual(content, "New content\n") + + async def test_cache_invalidation_after_file_deletion(self) -> None: + """Test cache handles file deletion gracefully.""" + # Create a file + with open(self.test_file, "w", encoding="utf-8") as f: + f.write("Test content\n") + + # Read the file to cache it + read_chunk = await self.read_tool( + file_path=self.test_file, + _agent_state=self.state, + ) + self.assertEqual(read_chunk.state, "running") + + # Verify cache exists + self.assertEqual(len(self.state.tool_context.read_file_cache), 1) + self.assertEqual( + self.state.tool_context.read_file_cache[0].file_path, + self.test_file, + ) + + # Delete the file + os.unlink(self.test_file) + + # Try to edit - should fail with "File not found" error + edit_chunk = await self.edit_tool( + file_path=self.test_file, + old_string="Test", + new_string="New", + _agent_state=self.state, + ) + + # Should fail with "File not found" error + self.assertEqual(edit_chunk.state, "error") + self.assertIn("not found", edit_chunk.content[0].text.lower()) + + # Try to read again - should also fail + read_chunk2 = await self.read_tool( + file_path=self.test_file, + _agent_state=self.state, + ) + self.assertEqual(read_chunk2.state, "error") + self.assertIn("does not exist", read_chunk2.content[0].text) + + async def test_cache_invalidation_after_file_modification(self) -> None: + """Test cache detects file modification.""" + # Create a file + with open(self.test_file, "w", encoding="utf-8") as f: + f.write("Original content\n") + + # Read the file to cache it + read_chunk = await self.read_tool( + file_path=self.test_file, + _agent_state=self.state, + ) + self.assertEqual(read_chunk.state, "running") + + # Modify the file externally + import time + + time.sleep(0.1) # Ensure mtime changes + with open(self.test_file, "w", encoding="utf-8") as f: + f.write("Modified content\n") + + # Try to edit - should fail because cache is stale + edit_chunk = await self.edit_tool( + file_path=self.test_file, + old_string="Original", + new_string="New", + _agent_state=self.state, + ) + + # Should fail with error + self.assertEqual(edit_chunk.state, "error") + self.assertIn("must first read", edit_chunk.content[0].text) + + async def test_cache_lru_eviction(self) -> None: + """Test LRU cache eviction when max_cache_files is exceeded.""" + # Set a small cache limit + self.state.tool_context.max_cache_files = 3 + + # Create and read 4 files + files = [] + for i in range(4): + file_path = os.path.join(self.temp_dir, f"file{i}.txt") + with open(file_path, "w", encoding="utf-8") as f: + f.write(f"Content {i}\n") + files.append(file_path) + + # Read each file + await self.read_tool( + file_path=file_path, + _agent_state=self.state, + ) + + # Cache should only have 3 files (oldest evicted) + self.assertEqual(len(self.state.tool_context.read_file_cache), 3) + + # The first file should have been evicted + cached_paths = [ + entry.file_path + for entry in self.state.tool_context.read_file_cache + ] + self.assertNotIn(files[0], cached_paths) + self.assertIn(files[1], cached_paths) + self.assertIn(files[2], cached_paths) + self.assertIn(files[3], cached_paths) + + async def test_cache_without_state(self) -> None: + """Test tools work without state (fallback mode).""" + # Create a file + with open(self.test_file, "w", encoding="utf-8") as f: + f.write("Hello World\n") + + # Edit without state should work (fallback to reading from disk) + edit_chunk = await self.edit_tool( + file_path=self.test_file, + old_string="Hello", + new_string="Hi", + _agent_state=None, + ) + + self.assertEqual(edit_chunk.state, "running") + + # Verify file was edited + with open(self.test_file, "r", encoding="utf-8") as f: + content = f.read() + self.assertEqual(content, "Hi World\n") + + async def test_multiple_reads_update_cache(self) -> None: + """Test reading same file multiple times updates cache.""" + # Create a file + with open(self.test_file, "w", encoding="utf-8") as f: + f.write("Content v1\n") + + # Read the file + await self.read_tool( + file_path=self.test_file, + _agent_state=self.state, + ) + self.assertEqual(len(self.state.tool_context.read_file_cache), 1) + + # Read again - should update cache, not add duplicate + await self.read_tool( + file_path=self.test_file, + _agent_state=self.state, + ) + self.assertEqual(len(self.state.tool_context.read_file_cache), 1) + + async def test_read_cache_lines_contain_newlines(self) -> None: + """Test that cached lines from Read retain trailing newlines. + + readlines() includes the newline character in each line. The cache + must store them as-is so that "".join(lines) reconstructs the exact + original file content. + """ + with open(self.test_file, "w", encoding="utf-8") as f: + f.write("line1\nline2\nline3\n") + + await self.read_tool( + file_path=self.test_file, + _agent_state=self.state, + ) + + cache = await self.state.tool_context.get_cache(self.test_file) + self.assertIsNotNone(cache) + # Each line from readlines() ends with \n + self.assertEqual(cache.lines, ["line1\n", "line2\n", "line3\n"]) + # "".join reconstructs the exact original content + self.assertEqual("".join(cache.lines), "line1\nline2\nline3\n") + + async def test_edit_multiline_match_from_cache(self) -> None: + """Test Edit correctly matches multi-line old_string from + cached content. + + Regression test for the bug where "\n".join(cache.lines) doubled the + newlines (e.g. "line1\n\nline2\n" instead of "line1\nline2\n"), + making multi-line old_string matching fail. + """ + with open(self.test_file, "w", encoding="utf-8") as f: + f.write("line1\nline2\nline3\n") + + await self.read_tool( + file_path=self.test_file, + _agent_state=self.state, + ) + + # This multi-line old_string must match exactly in the reconstructed + # content; with the bug it would not be found. + chunk = await self.edit_tool( + file_path=self.test_file, + old_string="line1\nline2", + new_string="replaced", + _agent_state=self.state, + ) + + self.assertEqual(chunk.state, "running") + + with open(self.test_file, "r", encoding="utf-8") as f: + content = f.read() + self.assertEqual(content, "replaced\nline3\n") + + async def test_edit_single_line_match_from_cache(self) -> None: + """Test Edit correctly matches single-line old_string from cache.""" + with open(self.test_file, "w", encoding="utf-8") as f: + f.write("Hello World\nThis is a test\n") + + await self.read_tool( + file_path=self.test_file, + _agent_state=self.state, + ) + + chunk = await self.edit_tool( + file_path=self.test_file, + old_string="Hello World", + new_string="Hello Python", + _agent_state=self.state, + ) + + self.assertEqual(chunk.state, "running") + + with open(self.test_file, "r", encoding="utf-8") as f: + content = f.read() + self.assertEqual(content, "Hello Python\nThis is a test\n") + + async def test_write_invalidates_cache(self) -> None: + """Test that Write updates the cache so Edit can use it afterwards.""" + with open(self.test_file, "w", encoding="utf-8") as f: + f.write("original content\n") + + # Read to populate cache + await self.read_tool( + file_path=self.test_file, + _agent_state=self.state, + ) + + # Overwrite with Write (mtime changes, old cache becomes stale) + write_chunk = await self.write_tool( + file_path=self.test_file, + content="new content\n", + _agent_state=self.state, + ) + self.assertEqual(write_chunk.state, "running") + + # The old cache entry is now stale; Edit should require a new Read + edit_chunk = await self.edit_tool( + file_path=self.test_file, + old_string="new content", + new_string="updated content", + _agent_state=self.state, + ) + self.assertEqual(edit_chunk.state, "error") + self.assertIn("must first read", edit_chunk.content[0].text) + + async def test_write_cache_stale_then_reread(self) -> None: + """Test workflow: Read -> Write -> Read -> Edit works correctly.""" + with open(self.test_file, "w", encoding="utf-8") as f: + f.write("original\n") + + # First read + await self.read_tool( + file_path=self.test_file, + _agent_state=self.state, + ) + + # Overwrite + import time + + time.sleep(0.01) # ensure mtime changes + await self.write_tool( + file_path=self.test_file, + content="rewritten\n", + _agent_state=self.state, + ) + + # Re-read to refresh cache + await self.read_tool( + file_path=self.test_file, + _agent_state=self.state, + ) + + # Now Edit should succeed against the new content + edit_chunk = await self.edit_tool( + file_path=self.test_file, + old_string="rewritten", + new_string="final", + _agent_state=self.state, + ) + self.assertEqual(edit_chunk.state, "running") + + with open(self.test_file, "r", encoding="utf-8") as f: + content = f.read() + self.assertEqual(content, "final\n") diff --git a/tests/builtin_glob_test.py b/tests/builtin_glob_test.py new file mode 100644 index 0000000000000000000000000000000000000000..7845f07d3736a3bf507450d37a1356be8f467733 --- /dev/null +++ b/tests/builtin_glob_test.py @@ -0,0 +1,236 @@ +# -*- coding: utf-8 -*- +"""Glob tool test case.""" +import os +import tempfile +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import AnyString +from agentscope.tool import Glob +from agentscope.permission import ( + PermissionContext, + PermissionBehavior, + PermissionRule, +) + + +class GlobToolTest(IsolatedAsyncioTestCase): + """The glob tool test case.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.glob_tool = Glob() + # Create a temporary directory with test files + self.temp_dir = tempfile.mkdtemp() + + # Create test files + with open( + os.path.join(self.temp_dir, "test1.py"), + "w", + encoding="utf-8", + ): + pass + with open( + os.path.join(self.temp_dir, "test2.py"), + "w", + encoding="utf-8", + ): + pass + with open( + os.path.join(self.temp_dir, "test.txt"), + "w", + encoding="utf-8", + ): + pass + + # Create subdirectory + sub_dir = os.path.join(self.temp_dir, "subdir") + os.makedirs(sub_dir) + with open(os.path.join(sub_dir, "test3.py"), "w", encoding="utf-8"): + pass + + async def asyncTearDown(self) -> None: + """Clean up temporary files.""" + import shutil + + if os.path.exists(self.temp_dir): + shutil.rmtree(self.temp_dir) + + async def test_tool_properties(self) -> None: + """Test glob tool properties.""" + self.assertEqual(self.glob_tool.name, "Glob") + self.assertIsInstance(self.glob_tool.description, str) + self.assertIsInstance(self.glob_tool.input_schema, dict) + self.assertFalse(self.glob_tool.is_mcp) + self.assertTrue(self.glob_tool.is_read_only) + self.assertTrue(self.glob_tool.is_concurrency_safe) + + async def test_check_permissions(self) -> None: + """Test glob tool permission checking.""" + context = PermissionContext() + tool_input = {"pattern": "*.py"} + decision = await self.glob_tool.check_permissions(tool_input, context) + + # Read/Glob/Grep are read-only, return PASSTHROUGH + self.assertEqual(decision.behavior, PermissionBehavior.PASSTHROUGH) + + async def test_simple_pattern(self) -> None: + """Test simple glob pattern.""" + chunk = await self.glob_tool( + pattern="*.py", + path=self.temp_dir, + ) + + self.assertEqual(chunk.state, "running") + + # Should find test1.py and test2.py + content = chunk.content[0].text + self.assertIn("test1.py", content) + self.assertIn("test2.py", content) + self.assertNotIn("test.txt", content) + + async def test_recursive_pattern(self) -> None: + """Test recursive glob pattern.""" + chunk = await self.glob_tool( + pattern="**/*.py", + path=self.temp_dir, + ) + + content = chunk.content[0].text + + # Should find all .py files including in subdirectory + self.assertIn("test1.py", content) + self.assertIn("test2.py", content) + self.assertIn("test3.py", content) + + async def test_windows_style_separator_pattern(self) -> None: + """Test glob patterns that use backslashes as path separators.""" + chunk = await self.glob_tool( + pattern=r"subdir\*.py", + path=self.temp_dir, + ) + + self.assertDictEqual( + chunk.model_dump(), + { + "content": [ + { + "type": "text", + "text": os.path.join( + self.temp_dir, + "subdir", + "test3.py", + ), + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + async def test_no_matches(self) -> None: + """Test pattern with no matches.""" + chunk = await self.glob_tool( + pattern="*.nonexistent", + path=self.temp_dir, + ) + + self.assertEqual(chunk.state, "running") + self.assertIn("No files found", chunk.content[0].text) + + async def test_match_rule_path(self) -> None: + """Test match_rule with path patterns.""" + # Test matching explicit path + self.assertTrue( + await self.glob_tool.match_rule( + self.temp_dir, + {"path": self.temp_dir, "pattern": "*.py"}, + ), + ) + + # Test wildcard pattern matching path + parent_dir = os.path.dirname(self.temp_dir) + self.assertTrue( + await self.glob_tool.match_rule( + parent_dir + "/**", + {"path": self.temp_dir, "pattern": "*.py"}, + ), + ) + + # Test non-matching path + self.assertFalse( + await self.glob_tool.match_rule( + "/some/other/path/**", + {"path": self.temp_dir, "pattern": "*.py"}, + ), + ) + + async def test_match_rule_pattern(self) -> None: + """Test match_rule with pattern matching.""" + # Test matching against the pattern itself + self.assertTrue( + await self.glob_tool.match_rule( + "*.py", + {"pattern": "*.py"}, + ), + ) + + # Test wildcard pattern matching + self.assertTrue( + await self.glob_tool.match_rule( + "**/*.py", + {"pattern": "src/**/*.py"}, + ), + ) + + # Test non-matching pattern + self.assertFalse( + await self.glob_tool.match_rule( + "*.txt", + {"pattern": "*.py"}, + ), + ) + + async def test_match_rule_path_priority(self) -> None: + """Test that path matching takes priority over pattern matching.""" + # If path matches, should return True even if pattern doesn't + self.assertTrue( + await self.glob_tool.match_rule( + self.temp_dir, + {"path": self.temp_dir, "pattern": "*.txt"}, + ), + ) + + async def test_generate_suggestions_with_path(self) -> None: + """Test generate_suggestions for glob with explicit path.""" + + suggestions = await self.glob_tool.generate_suggestions( + {"path": self.temp_dir, "pattern": "*.py"}, + ) + + self.assertIsInstance(suggestions, list) + self.assertGreater(len(suggestions), 0) + self.assertIsInstance(suggestions[0], PermissionRule) + + # Should suggest directory pattern + abs_path = os.path.abspath(self.temp_dir) + expected_pattern = abs_path.rstrip("/") + "/**" + suggestion_contents = [s.rule_content for s in suggestions] + self.assertIn(expected_pattern, suggestion_contents) + + async def test_generate_suggestions_defaults_to_cwd(self) -> None: + """Test generate_suggestions defaults to cwd when no path provided.""" + + suggestions = await self.glob_tool.generate_suggestions( + {"pattern": "*.py"}, + ) + + self.assertIsInstance(suggestions, list) + self.assertGreater(len(suggestions), 0) + + cwd = os.getcwd() + expected_pattern = os.path.abspath(cwd).rstrip("/") + "/**" + suggestion_contents = [s.rule_content for s in suggestions] + self.assertIn(expected_pattern, suggestion_contents) diff --git a/tests/builtin_grep_test.py b/tests/builtin_grep_test.py new file mode 100644 index 0000000000000000000000000000000000000000..7f1266cbba9c7ecff91080bc5614b146e85d5123 --- /dev/null +++ b/tests/builtin_grep_test.py @@ -0,0 +1,254 @@ +# -*- coding: utf-8 -*- +"""Grep tool test case.""" +import os +import tempfile +from unittest.async_case import IsolatedAsyncioTestCase + +from agentscope.message import ToolResultState +from agentscope.tool import Grep +from agentscope.permission import ( + PermissionContext, + PermissionBehavior, + PermissionRule, +) + + +class GrepToolTest(IsolatedAsyncioTestCase): + """The grep tool test case.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.grep_tool = Grep() + # Create a temporary directory with test files + self.temp_dir = tempfile.mkdtemp() + + # Create test files + with open( + os.path.join(self.temp_dir, "test1.py"), + "w", + encoding="utf-8", + ) as f: + f.write("def hello():\n print('Hello World')\n") + + with open( + os.path.join(self.temp_dir, "test2.py"), + "w", + encoding="utf-8", + ) as f: + f.write("def goodbye():\n print('Goodbye')\n") + + with open( + os.path.join(self.temp_dir, "test.txt"), + "w", + encoding="utf-8", + ) as f: + f.write("This is a text file\nHello from text\n") + + # Create subdirectory with files for glob pattern testing + subdir = os.path.join(self.temp_dir, "subdir") + os.makedirs(subdir) + with open( + os.path.join(subdir, "nested.py"), + "w", + encoding="utf-8", + ) as f: + f.write("def nested():\n print('Nested')\n") + + async def asyncTearDown(self) -> None: + """Clean up temporary files.""" + import shutil + + if os.path.exists(self.temp_dir): + shutil.rmtree(self.temp_dir) + + async def test_tool_properties(self) -> None: + """Test grep tool properties.""" + self.assertEqual(self.grep_tool.name, "Grep") + self.assertIsInstance(self.grep_tool.description, str) + self.assertIsInstance(self.grep_tool.input_schema, dict) + self.assertFalse(self.grep_tool.is_mcp) + self.assertTrue(self.grep_tool.is_read_only) + self.assertTrue(self.grep_tool.is_concurrency_safe) + + async def test_check_permissions(self) -> None: + """Test grep tool permission checking.""" + context = PermissionContext() + tool_input = {"pattern": "hello"} + decision = await self.grep_tool.check_permissions(tool_input, context) + + # Read/Glob/Grep are read-only, return PASSTHROUGH + self.assertEqual(decision.behavior, PermissionBehavior.PASSTHROUGH) + + async def test_simple_search(self) -> None: + """Test simple grep search.""" + chunk = await self.grep_tool( + pattern="Hello", + path=self.temp_dir, + output_mode="files_with_matches", + ) + + self.assertEqual(chunk.state, ToolResultState.SUCCESS) + + content = chunk.content[0].text + # Should find files containing "Hello" + self.assertIn("test1.py", content) + self.assertIn("test.txt", content) + + async def test_content_mode(self) -> None: + """Test grep with content output mode.""" + chunk = await self.grep_tool( + pattern="def", + path=self.temp_dir, + output_mode="content", + type="py", + ) + + content = chunk.content[0].text + + # Should show matching lines + self.assertIn("def hello", content) + self.assertIn("def goodbye", content) + + async def test_case_insensitive(self) -> None: + """Test case-insensitive search.""" + chunk = await self.grep_tool( + pattern="HELLO", + path=self.temp_dir, + case_insensitive=True, + output_mode="files_with_matches", + ) + + content = chunk.content[0].text + self.assertIn("test1.py", content) + + async def test_no_matches(self) -> None: + """Test search with no matches.""" + chunk = await self.grep_tool( + pattern="NonExistentPattern", + path=self.temp_dir, + ) + + self.assertIn("No matches found", chunk.content[0].text) + + async def test_type_filter(self) -> None: + """Test filtering by file type.""" + chunk = await self.grep_tool( + pattern="Hello", + path=self.temp_dir, + type="py", + output_mode="files_with_matches", + ) + + content = chunk.content[0].text + + # Should only find .py files + self.assertIn("test1.py", content) + self.assertNotIn("test.txt", content) + + async def test_invalid_regex(self) -> None: + """Test grep with invalid regex pattern.""" + chunk = await self.grep_tool( + pattern="[invalid(regex", + path=self.temp_dir, + ) + + self.assertEqual(chunk.state, "error") + # ripgrep returns its own error message for regex parse errors + self.assertIn("regex parse error", chunk.content[0].text) + + async def test_glob_pattern_with_subdirs(self) -> None: + """Test glob pattern matching with subdirectories like **/*.py.""" + chunk = await self.grep_tool( + pattern="def", + path=self.temp_dir, + glob="**/*.py", + output_mode="files_with_matches", + ) + + content = chunk.content[0].text + + # Should find all .py files including in subdirectories + self.assertIn("test1.py", content) + self.assertIn("test2.py", content) + self.assertIn("nested.py", content) + # Should not find .txt files + self.assertNotIn("test.txt", content) + + async def test_match_rule_path(self) -> None: + """Test match_rule with search path patterns.""" + # Test matching explicit path + self.assertTrue( + await self.grep_tool.match_rule( + self.temp_dir, + {"path": self.temp_dir}, + ), + ) + + # Test wildcard pattern matching path + parent_dir = os.path.dirname(self.temp_dir) + self.assertTrue( + await self.grep_tool.match_rule( + parent_dir + "/**", + {"path": self.temp_dir}, + ), + ) + + # Test non-matching path + self.assertFalse( + await self.grep_tool.match_rule( + "/some/other/path/**", + {"path": self.temp_dir}, + ), + ) + + async def test_match_rule_defaults_to_cwd(self) -> None: + """Test match_rule defaults to cwd when no path is provided.""" + cwd = os.getcwd() + + # When no path provided, should match against cwd + self.assertTrue( + await self.grep_tool.match_rule( + cwd, + {"pattern": "hello"}, + ), + ) + + # Should not match a different path + self.assertFalse( + await self.grep_tool.match_rule( + "/some/other/path", + {"pattern": "hello"}, + ), + ) + + async def test_generate_suggestions_with_path(self) -> None: + """Test generate_suggestions for grep with explicit path.""" + + suggestions = await self.grep_tool.generate_suggestions( + {"path": self.temp_dir, "pattern": "hello"}, + ) + + self.assertIsInstance(suggestions, list) + self.assertGreater(len(suggestions), 0) + self.assertIsInstance(suggestions[0], PermissionRule) + + # Should suggest directory pattern + abs_path = os.path.abspath(self.temp_dir) + expected_pattern = abs_path.rstrip("/") + "/**" + suggestion_contents = [s.rule_content for s in suggestions] + self.assertIn(expected_pattern, suggestion_contents) + + async def test_generate_suggestions_defaults_to_cwd(self) -> None: + """Test generate_suggestions defaults to cwd when no path provided.""" + + suggestions = await self.grep_tool.generate_suggestions( + {"pattern": "hello"}, + ) + + self.assertIsInstance(suggestions, list) + self.assertGreater(len(suggestions), 0) + + cwd = os.getcwd() + expected_pattern = os.path.abspath(cwd).rstrip("/") + "/**" + suggestion_contents = [s.rule_content for s in suggestions] + self.assertIn(expected_pattern, suggestion_contents) diff --git a/tests/builtin_read_test.py b/tests/builtin_read_test.py new file mode 100644 index 0000000000000000000000000000000000000000..2cf5ccc1412d2939e35f85682bcf909a36231ec6 --- /dev/null +++ b/tests/builtin_read_test.py @@ -0,0 +1,188 @@ +# -*- coding: utf-8 -*- +"""Read tool test case.""" +import os +import tempfile +from unittest.async_case import IsolatedAsyncioTestCase + +from agentscope.tool import ToolChunk, Read +from agentscope.permission import ( + PermissionContext, + PermissionBehavior, + PermissionRule, +) +from agentscope.message import TextBlock + + +class ReadToolTest(IsolatedAsyncioTestCase): + """The read tool test case.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.read_tool = Read() + # Create a temporary file for testing + self.temp_file = tempfile.NamedTemporaryFile( + mode="w", + delete=False, + suffix=".txt", + ) + # Write multiple lines + for i in range(1, 11): + self.temp_file.write(f"Line {i}\n") + self.temp_file.close() + + async def asyncTearDown(self) -> None: + """Clean up temporary files.""" + if os.path.exists(self.temp_file.name): + os.unlink(self.temp_file.name) + + async def test_tool_properties(self) -> None: + """Test read tool properties.""" + self.assertEqual(self.read_tool.name, "Read") + self.assertIsInstance(self.read_tool.description, str) + self.assertIsInstance(self.read_tool.input_schema, dict) + self.assertFalse(self.read_tool.is_mcp) + self.assertTrue(self.read_tool.is_read_only) + self.assertTrue(self.read_tool.is_concurrency_safe) + + async def test_check_permissions(self) -> None: + """Test read tool permission checking.""" + context = PermissionContext() + tool_input = {"file_path": "/tmp/test.txt"} + decision = await self.read_tool.check_permissions(tool_input, context) + + # Read/Glob/Grep are read-only, return PASSTHROUGH + self.assertEqual(decision.behavior, PermissionBehavior.PASSTHROUGH) + + async def test_simple_read(self) -> None: + """Test simple file reading.""" + chunk = await self.read_tool(file_path=self.temp_file.name) + + self.assertIsInstance(chunk, ToolChunk) + self.assertEqual(chunk.state, "running") + self.assertEqual(len(chunk.content), 1) + self.assertIsInstance(chunk.content[0], TextBlock) + + content = chunk.content[0].text + # Should contain all lines with line numbers + self.assertIn("Line 1", content) + self.assertIn("Line 10", content) + + async def test_read_with_offset(self) -> None: + """Test reading with offset.""" + chunk = await self.read_tool( + file_path=self.temp_file.name, + offset=5, + ) + + self.assertEqual(chunk.state, "running") + content = chunk.content[0].text + + # Should start from line 5 + self.assertIn("Line 5", content) + # Line 1 should not appear (but Line 10 contains "1", + # so check more specifically) + lines = content.split("\n") + line_numbers = [ + int(line.split("\t")[0].strip()) for line in lines if line.strip() + ] + self.assertNotIn(1, line_numbers) + self.assertIn(5, line_numbers) + + async def test_read_with_limit(self) -> None: + """Test reading with limit.""" + chunk = await self.read_tool( + file_path=self.temp_file.name, + offset=1, + limit=3, + ) + + self.assertEqual(chunk.state, "running") + content = chunk.content[0].text + + # Should only read 3 lines + self.assertIn("Line 1", content) + self.assertIn("Line 2", content) + self.assertIn("Line 3", content) + self.assertNotIn("Line 4", content) + + async def test_read_nonexistent_file(self) -> None: + """Test reading a non-existent file.""" + chunk = await self.read_tool(file_path="/nonexistent/file.txt") + + self.assertEqual(chunk.state, "error") + self.assertIn("does not exist", chunk.content[0].text) + + async def test_read_directory(self) -> None: + """Test reading a directory (should fail).""" + temp_dir = tempfile.mkdtemp() + try: + chunk = await self.read_tool(file_path=temp_dir) + + self.assertEqual(chunk.state, "error") + self.assertIn("directory", chunk.content[0].text.lower()) + finally: + os.rmdir(temp_dir) + + async def test_match_rule_glob_pattern(self) -> None: + """Test match_rule with glob patterns.""" + # Test exact match + self.assertTrue( + await self.read_tool.match_rule( + "test.py", + {"file_path": "test.py"}, + ), + ) + + # Test wildcard pattern + self.assertTrue( + await self.read_tool.match_rule( + "*.py", + {"file_path": "test.py"}, + ), + ) + + # Test directory pattern + self.assertTrue( + await self.read_tool.match_rule( + "/tmp/**", + {"file_path": "/tmp/test.py"}, + ), + ) + + # Test non-matching pattern + self.assertFalse( + await self.read_tool.match_rule( + "*.txt", + {"file_path": "test.py"}, + ), + ) + + # Test empty file_path + self.assertFalse( + await self.read_tool.match_rule( + "*.py", + {"file_path": ""}, + ), + ) + + async def test_generate_suggestions(self) -> None: + """Test generate_suggestions for file operations.""" + + # Test suggestion for file in subdirectory + suggestions = await self.read_tool.generate_suggestions( + {"file_path": "/tmp/project/src/main.py"}, + ) + + self.assertIsInstance(suggestions, list) + self.assertGreater(len(suggestions), 0) + self.assertIsInstance(suggestions[0], PermissionRule) + + # Should suggest parent directory pattern + suggestion_contents = [s.rule_content for s in suggestions] + self.assertIn("/tmp/project/src/**", suggestion_contents) + + # Test suggestion for file in root + suggestions = await self.read_tool.generate_suggestions( + {"file_path": "/test.py"}, + ) + self.assertGreater(len(suggestions), 0) diff --git a/tests/builtin_write_test.py b/tests/builtin_write_test.py new file mode 100644 index 0000000000000000000000000000000000000000..7cff3b3b6bb92a317d017ee5c21ad90bbf3a2555 --- /dev/null +++ b/tests/builtin_write_test.py @@ -0,0 +1,212 @@ +# -*- coding: utf-8 -*- +"""Write tool test case.""" +import os +import tempfile +from unittest.async_case import IsolatedAsyncioTestCase + +from agentscope.tool import Write +from agentscope.permission import ( + PermissionContext, + PermissionBehavior, + PermissionRule, +) +from agentscope.state import AgentState +from agentscope.message import ToolResultState + + +class WriteToolTest(IsolatedAsyncioTestCase): + """The write tool test case.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.write_tool = Write() + self.temp_dir = tempfile.mkdtemp() + + async def asyncTearDown(self) -> None: + """Clean up temporary files.""" + import shutil + + if os.path.exists(self.temp_dir): + shutil.rmtree(self.temp_dir) + + async def test_tool_properties(self) -> None: + """Test write tool properties.""" + self.assertEqual(self.write_tool.name, "Write") + self.assertIsInstance(self.write_tool.description, str) + self.assertIsInstance(self.write_tool.input_schema, dict) + self.assertFalse(self.write_tool.is_mcp) + self.assertFalse(self.write_tool.is_read_only) + self.assertFalse(self.write_tool.is_concurrency_safe) + + async def test_check_permissions(self) -> None: + """Test write tool permission checking. + + Write tool should return PASSTHROUGH for non-dangerous paths, + allowing PermissionEngine to check allow rules. + """ + context = PermissionContext() + tool_input = {"file_path": "/tmp/test.txt"} + decision = await self.write_tool.check_permissions(tool_input, context) + + self.assertEqual(decision.behavior, PermissionBehavior.PASSTHROUGH) + + async def test_simple_write(self) -> None: + """Test simple file writing.""" + file_path = os.path.join(self.temp_dir, "test.txt") + content = "Hello World\nThis is a test\n" + + chunk = await self.write_tool( + file_path=file_path, + content=content, + ) + + self.assertEqual(chunk.state, "running") + self.assertTrue(chunk.is_last) + + # Verify file was created and content is correct + self.assertTrue(os.path.exists(file_path)) + with open(file_path, "r", encoding="utf-8") as f: + written_content = f.read() + self.assertEqual(written_content, content) + + async def test_write_creates_directory(self) -> None: + """Test that write creates parent directories.""" + file_path = os.path.join(self.temp_dir, "subdir", "test.txt") + content = "Test content" + + chunk = await self.write_tool( + file_path=file_path, + content=content, + ) + + self.assertEqual(chunk.state, "running") + + # Verify directory and file were created + self.assertTrue(os.path.exists(file_path)) + with open(file_path, "r", encoding="utf-8") as f: + written_content = f.read() + self.assertEqual(written_content, content) + + async def test_write_overwrites_existing(self) -> None: + """Test that write overwrites existing files.""" + file_path = os.path.join(self.temp_dir, "test.txt") + + # Write initial content + with open(file_path, "w", encoding="utf-8") as f: + f.write("Initial content") + + # Overwrite with new content + new_content = "New content" + chunk = await self.write_tool( + file_path=file_path, + content=new_content, + ) + + self.assertEqual(len([chunk]), 1) + + # Verify content was overwritten + with open(file_path, "r", encoding="utf-8") as f: + written_content = f.read() + self.assertEqual(written_content, new_content) + self.assertNotIn("Initial", written_content) + + async def test_write_empty_content(self) -> None: + """Test writing empty content.""" + file_path = os.path.join(self.temp_dir, "empty.txt") + + chunk = await self.write_tool( + file_path=file_path, + content="", + ) + + self.assertEqual(len([chunk]), 1) + self.assertTrue(os.path.exists(file_path)) + + with open(file_path, "r", encoding="utf-8") as f: + content = f.read() + self.assertEqual(content, "") + + async def test_overwrite_existing_without_prior_read_errors(self) -> None: + """Overwriting an existing file via state-injected call requires + the file to have been read first (cached in tool_context). + """ + file_path = os.path.join(self.temp_dir, "existing.txt") + with open(file_path, "w", encoding="utf-8") as f: + f.write("original") + + state = AgentState() + chunk = await self.write_tool( + file_path=file_path, + content="new", + _agent_state=state, + ) + + self.assertEqual(chunk.state, ToolResultState.ERROR) + self.assertIn("has not been read", chunk.content[0].text) + # File must not have been mutated + with open(file_path, "r", encoding="utf-8") as f: + self.assertEqual(f.read(), "original") + + async def test_match_rule_glob_pattern(self) -> None: + """Test match_rule with glob patterns.""" + # Test exact match + self.assertTrue( + await self.write_tool.match_rule( + "test.py", + {"file_path": "test.py"}, + ), + ) + + # Test wildcard pattern + self.assertTrue( + await self.write_tool.match_rule( + "*.py", + {"file_path": "test.py"}, + ), + ) + + # Test directory pattern + self.assertTrue( + await self.write_tool.match_rule( + "/tmp/**", + {"file_path": "/tmp/test.py"}, + ), + ) + + # Test non-matching pattern + self.assertFalse( + await self.write_tool.match_rule( + "*.txt", + {"file_path": "test.py"}, + ), + ) + + # Test empty file_path + self.assertFalse( + await self.write_tool.match_rule( + "*.py", + {"file_path": ""}, + ), + ) + + async def test_generate_suggestions(self) -> None: + """Test generate_suggestions for file operations.""" + + # Test suggestion for file in subdirectory + suggestions = await self.write_tool.generate_suggestions( + {"file_path": "/tmp/project/src/main.py"}, + ) + + self.assertIsInstance(suggestions, list) + self.assertGreater(len(suggestions), 0) + self.assertIsInstance(suggestions[0], PermissionRule) + + # Should suggest parent directory pattern + suggestion_contents = [s.rule_content for s in suggestions] + self.assertIn("/tmp/project/src/**", suggestion_contents) + + # Test suggestion for file in root + suggestions = await self.write_tool.generate_suggestions( + {"file_path": "/test.py"}, + ) + self.assertGreater(len(suggestions), 0) diff --git a/tests/compress_context_test.py b/tests/compress_context_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ed25866b6affb2bf88a2f0d675cbaa6b9cc7b588 --- /dev/null +++ b/tests/compress_context_test.py @@ -0,0 +1,1083 @@ +# -*- coding: utf-8 -*- +"""A template test case.""" +# pylint: disable=protected-access +import json +import os +import tempfile +from typing import Any + +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import MockModel, AnyString + +from agentscope.model import StructuredResponse +from agentscope.agent import Agent, ContextConfig +from agentscope.state import AgentState +from agentscope.message import ( + UserMsg, + AssistantMsg, + TextBlock, + ToolCallBlock, + HintBlock, + Msg, +) +from agentscope.tool import Toolkit + + +class RecordingStructuredMockModel(MockModel): + """A mock model that records structured-output compression calls.""" + + def __init__( + self, + *args: Any, + fail_structured_output_times: int = 0, + force_compression_overflow: bool = False, + **kwargs: Any, + ) -> None: + """Initialize the recording mock model.""" + super().__init__(*args, **kwargs) + self.recorded_structured_messages: list[list[Msg]] = [] + self._fail_structured_output_times = fail_structured_output_times + self._force_compression_overflow = force_compression_overflow + self._compression_count_calls = 0 + + async def count_tokens( + self, + messages: list[Msg], + tools: list[dict] | None, + ) -> int: + """Force the overflow branch when counting compression messages.""" + is_compression_count = bool( + tools + and tools[0].get("function", {}).get("name") + == "generate_structured_output", + ) + if self._force_compression_overflow and is_compression_count: + self._compression_count_calls += 1 + if self._compression_count_calls == 1: + return self.context_size + 1 + return 1 + return await super().count_tokens(messages, tools) + + async def _call_api_with_structured_output( + self, + model_name: str, + messages: list[Msg], + structured_model: Any, + **kwargs: Any, + ) -> StructuredResponse: + """Record the structured-output call and optionally fail first.""" + self.recorded_structured_messages.append(list(messages)) + if self._fail_structured_output_times > 0: + self._fail_structured_output_times -= 1 + raise RuntimeError("simulated compression overflow") + return await super()._call_api_with_structured_output( + model_name, + messages, + structured_model, + **kwargs, + ) + + +def _has_instruction_hint( + messages: list[Msg], + instructions: HintBlock, +) -> bool: + """Return True if messages contain instructions as an assistant hint.""" + for msg in messages: + if msg.role != "assistant": + continue + for hint_block in msg.get_content_blocks("hint"): + if hint_block.id == instructions.id: + return hint_block.hint == instructions.hint + return False + + +class ContextCompressionTest(IsolatedAsyncioTestCase): + """The template test case.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + + async def test_split_function(self) -> None: + """The template test.""" + agent = Agent( + name="Friday", + system_prompt="".join(["0" for _ in range(60 * 4)]), + model=MockModel(), + context_config=ContextConfig( + trigger_ratio=0.8, + reserve_ratio=0.1, + ), + state=AgentState( + session_id="123", + context=[ + UserMsg( + "User", + "".join(["1" for _ in range(30 * 4)]), + id="1", + ), + AssistantMsg( + "Friday", + "".join(["2" for _ in range(10 * 4)]), + id="2", + ), + UserMsg( + "User", + "".join(["3" for _ in range(10 * 4)]), + id="3", + ), + ], + ), + toolkit=Toolkit(), + ) + + # When the length of last two messages is exactly appropriate + to_compress, to_reserve = await agent._split_context_for_compression( + to_reserved_tokens=80, + tools=[], + ) + + self.assertListEqual( + [_.id for _ in to_compress], + ["1"], + ) + self.assertListEqual( + [_.id for _ in to_reserve], + ["2", "3"], + ) + + # When one message is in the dividing line + agent.state.context = [ + UserMsg("User", "".join(["2" for _ in range(30 * 4)]), id="1"), + AssistantMsg( + "Friday", + "".join(["3" for _ in range(15 * 4)]), + id="2", + ), + UserMsg("User", "".join(["3" for _ in range(10 * 4)]), id="3"), + ] + + to_compress, to_reserve = await agent._split_context_for_compression( + to_reserved_tokens=80, + tools=[], + ) + self.assertListEqual( + [_.id for _ in to_compress], + ["1", "2"], + ) + self.assertListEqual( + [_.id for _ in to_reserve], + ["3"], + ) + + # When compress all messages + agent.state.context = [ + UserMsg("User", "".join(["2" for _ in range(30 * 4)]), id="1"), + AssistantMsg( + "Friday", + "".join(["3" for _ in range(15 * 4)]), + id="2", + ), + UserMsg("User", "".join(["3" for _ in range(30 * 4)]), id="3"), + ] + to_compress, to_reserve = await agent._split_context_for_compression( + to_reserved_tokens=80, + tools=[], + ) + self.assertListEqual( + [_.id for _ in to_compress], + ["1", "2", "3"], + ) + self.assertListEqual( + [_.id for _ in to_reserve], + [], + ) + + # When the boundary message has multiple blocks + agent.state.context = [ + UserMsg("User", "".join(["a" for _ in range(30 * 4)]), id="1"), + AssistantMsg( + "Friday", + [ + TextBlock( + text="".join(["b" for _ in range(10 * 4)]), + id="b", + ), + TextBlock( + text="".join(["c" for _ in range(10 * 4)]), + id="c", + ), + ], + id="2", + ), + UserMsg("User", "".join(["d" for _ in range(10 * 4)]), id="3"), + ] + to_compress, to_reserve = await agent._split_context_for_compression( + to_reserved_tokens=80, + tools=[], + ) + self.assertListEqual( + [_.model_dump() for _ in to_compress], + [ + { + "id": "1", + "created_at": AnyString(), + "finished_at": AnyString(), + "name": "User", + "role": "user", + "content": [ + { + "id": AnyString(), + "type": "text", + "text": "a" * 120, + }, + ], + "metadata": {}, + "usage": None, + }, + { + "id": "2", + "created_at": AnyString(), + "finished_at": None, + "name": "Friday", + "role": "assistant", + "content": [ + { + "id": "b", + "text": "b" * 40, + "type": "text", + }, + ], + "metadata": {}, + "usage": None, + }, + ], + ) + self.assertListEqual( + [_.model_dump() for _ in to_reserve], + [ + { + "id": "2", + "created_at": AnyString(), + "finished_at": None, + "name": "Friday", + "role": "assistant", + "content": [ + { + "id": AnyString(), + "text": "c" * 40, + "type": "text", + }, + ], + "metadata": {}, + "usage": None, + }, + { + "id": "3", + "created_at": AnyString(), + "finished_at": AnyString(), + "name": "User", + "role": "user", + "content": [ + { + "id": AnyString(), + "type": "text", + "text": "d" * 40, + }, + ], + "metadata": {}, + "usage": None, + }, + ], + ) + + # When the boundary message has multiple blocks + # Cannot leave any blocks + agent.state.context = [ + UserMsg("User", "".join(["a" for _ in range(30 * 4)]), id="1"), + AssistantMsg( + "Friday", + [ + TextBlock( + text="".join(["b" for _ in range(10 * 4)]), + id="b", + ), + TextBlock( + text="".join(["c" for _ in range(15 * 4)]), + id="c", + ), + ], + id="2", + ), + UserMsg("User", "".join(["d" for _ in range(10 * 4)]), id="3"), + ] + to_compress, to_reserve = await agent._split_context_for_compression( + to_reserved_tokens=80, + tools=[], + ) + self.assertListEqual( + [_.model_dump() for _ in to_compress], + [ + { + "id": "1", + "created_at": AnyString(), + "finished_at": AnyString(), + "name": "User", + "role": "user", + "content": [ + { + "id": AnyString(), + "type": "text", + "text": "a" * 120, + }, + ], + "metadata": {}, + "usage": None, + }, + { + "id": "2", + "created_at": AnyString(), + "finished_at": None, + "name": "Friday", + "role": "assistant", + "content": [ + { + "id": "b", + "text": "b" * 40, + "type": "text", + }, + { + "id": "c", + "text": "c" * 60, + "type": "text", + }, + ], + "metadata": {}, + "usage": None, + }, + ], + ) + self.assertListEqual( + [_.model_dump() for _ in to_reserve], + [ + { + "id": "3", + "created_at": AnyString(), + "finished_at": AnyString(), + "name": "User", + "role": "user", + "content": [ + { + "id": AnyString(), + "type": "text", + "text": "d" * 40, + }, + ], + "metadata": {}, + "usage": None, + }, + ], + ) + + # Leave the last block of the boundary message + agent.state.context = [ + UserMsg("User", "".join(["a" for _ in range(30 * 4)]), id="1"), + AssistantMsg( + "Friday", + [ + TextBlock( + text="".join(["b" for _ in range(10 * 4)]), + id="b", + ), + TextBlock( + text="".join(["c" for _ in range(5 * 4)]), + id="c", + ), + ], + id="2", + ), + UserMsg("User", "".join(["d" for _ in range(10 * 4)]), id="3"), + ] + to_compress, to_reserve = await agent._split_context_for_compression( + to_reserved_tokens=80, + tools=[], + ) + self.assertListEqual( + [_.model_dump() for _ in to_compress], + [ + { + "id": "1", + "created_at": AnyString(), + "finished_at": AnyString(), + "name": "User", + "role": "user", + "content": [ + { + "id": AnyString(), + "type": "text", + "text": "a" * 120, + }, + ], + "metadata": {}, + "usage": None, + }, + { + "id": "2", + "created_at": AnyString(), + "finished_at": None, + "name": "Friday", + "role": "assistant", + "content": [ + { + "id": "b", + "text": "b" * 40, + "type": "text", + }, + ], + "metadata": {}, + "usage": None, + }, + ], + ) + self.assertListEqual( + [_.model_dump() for _ in to_reserve], + [ + { + "id": "2", + "created_at": AnyString(), + "finished_at": None, + "name": "Friday", + "role": "assistant", + "content": [ + { + "id": "c", + "text": "c" * 20, + "type": "text", + }, + ], + "metadata": {}, + "usage": None, + }, + { + "id": "3", + "created_at": AnyString(), + "finished_at": AnyString(), + "name": "User", + "role": "user", + "content": [ + { + "id": AnyString(), + "type": "text", + "text": "d" * 40, + }, + ], + "metadata": {}, + "usage": None, + }, + ], + ) + + # Leave all the blocks + agent.state.context = [ + UserMsg("User", "".join(["a" for _ in range(30 * 4)]), id="1"), + AssistantMsg( + "Friday", + [ + TextBlock( + text="".join(["b" for _ in range(5 * 4)]), + id="b", + ), + TextBlock( + text="".join(["c" for _ in range(5 * 4)]), + id="c", + ), + ], + id="2", + ), + UserMsg("User", "".join(["d" for _ in range(10 * 4)]), id="3"), + ] + to_compress, to_reserve = await agent._split_context_for_compression( + to_reserved_tokens=80, + tools=[], + ) + self.assertListEqual( + [_.model_dump() for _ in to_compress], + [ + { + "id": "1", + "created_at": AnyString(), + "finished_at": AnyString(), + "name": "User", + "role": "user", + "content": [ + { + "id": AnyString(), + "type": "text", + "text": "a" * 120, + }, + ], + "metadata": {}, + "usage": None, + }, + ], + ) + self.assertListEqual( + [_.model_dump() for _ in to_reserve], + [ + { + "id": "2", + "created_at": AnyString(), + "finished_at": None, + "name": "Friday", + "role": "assistant", + "content": [ + { + "id": "b", + "text": "b" * 20, + "type": "text", + }, + { + "id": "c", + "text": "c" * 20, + "type": "text", + }, + ], + "metadata": {}, + "usage": None, + }, + { + "id": "3", + "created_at": AnyString(), + "finished_at": AnyString(), + "name": "User", + "role": "user", + "content": [ + { + "id": AnyString(), + "type": "text", + "text": "d" * 40, + }, + ], + "metadata": {}, + "usage": None, + }, + ], + ) + + # Leave all the messages + agent.state.context = [ + AssistantMsg( + "Friday", + [ + TextBlock( + text="".join(["b" for _ in range(5 * 4)]), + id="b", + ), + TextBlock( + text="".join(["c" for _ in range(5 * 4)]), + id="c", + ), + ], + id="2", + ), + UserMsg("User", "".join(["d" for _ in range(10 * 4)]), id="3"), + ] + to_compress, to_reserve = await agent._split_context_for_compression( + to_reserved_tokens=80, + tools=[], + ) + self.assertListEqual( + [_.model_dump() for _ in to_compress], + [], + ) + self.assertListEqual( + [_.model_dump() for _ in to_reserve], + [ + { + "id": "2", + "created_at": AnyString(), + "finished_at": None, + "name": "Friday", + "role": "assistant", + "content": [ + { + "id": "b", + "text": "b" * 20, + "type": "text", + }, + { + "id": "c", + "text": "c" * 20, + "type": "text", + }, + ], + "metadata": {}, + "usage": None, + }, + { + "id": "3", + "created_at": AnyString(), + "finished_at": AnyString(), + "name": "User", + "role": "user", + "content": [ + { + "id": AnyString(), + "type": "text", + "text": "d" * 40, + }, + ], + "metadata": {}, + "usage": None, + }, + ], + ) + + async def test_context_compression(self) -> None: + """Test the context compression logic.""" + model = MockModel(context_size=100) + agent = Agent( + name="Friday", + system_prompt="".join(["0" for _ in range(20 * 4)]), + model=model, + context_config=ContextConfig( + trigger_ratio=0.7, + reserve_ratio=0.4, + ), + state=AgentState( + session_id="123", + context=[ + UserMsg( + "User", + "".join(["1" for _ in range(30 * 4)]), + id="1", + ), + AssistantMsg( + "Friday", + "".join(["2" for _ in range(10 * 4)]), + id="2", + ), + UserMsg( + "User", + "".join(["3" for _ in range(10 * 4)]), + id="3", + ), + ], + ), + toolkit=Toolkit(), + ) + + model.set_structured_response( + StructuredResponse( + content={ + "task_overview": "1", + "current_state": "2", + "important_discoveries": "3", + "next_steps": "4", + "context_to_preserve": "5", + }, + ), + ) + + await agent.compress_context() + + self.assertEqual( + agent.state.summary, + """Here is a summary of your previous work +# Task Overview +1 + +# Current State +2 + +# Important Discoveries +3 + +# Next Steps +4 + +# Context to Preserve +5""", + ) + + self.assertListEqual( + [_.model_dump() for _ in agent.state.context], + [ + { + "id": "2", + "created_at": AnyString(), + "finished_at": None, + "name": "Friday", + "role": "assistant", + "content": [ + { + "id": AnyString(), + "text": "2" * 40, + "type": "text", + }, + ], + "metadata": {}, + "usage": None, + }, + { + "id": "3", + "created_at": AnyString(), + "finished_at": AnyString(), + "name": "User", + "role": "user", + "content": [ + { + "id": AnyString(), + "text": "3" * 40, + "type": "text", + }, + ], + "metadata": {}, + "usage": None, + }, + ], + ) + + async def test_context_compression_clears_evicted_read_cache(self) -> None: + """Read cache is cleared when its Read block is compressed out.""" + with tempfile.TemporaryDirectory() as temp_dir: + file_path = os.path.join(temp_dir, "test.txt") + with open(file_path, "w", encoding="utf-8") as f: + f.write("content\n") + + model = MockModel(context_size=100) + agent = Agent( + name="Friday", + system_prompt="".join(["0" for _ in range(20 * 4)]), + model=model, + context_config=ContextConfig( + trigger_ratio=0.7, + reserve_ratio=0.4, + ), + state=AgentState( + session_id="123", + context=[ + AssistantMsg( + "Friday", + [ + ToolCallBlock( + id="read-call-1", + name="Read", + input=json.dumps( + {"file_path": file_path}, + ), + ), + ], + id="1", + ), + UserMsg( + "User", + "".join(["2" for _ in range(30 * 4)]), + id="2", + ), + UserMsg( + "User", + "".join(["3" for _ in range(10 * 4)]), + id="3", + ), + ], + ), + toolkit=Toolkit(), + ) + await agent.state.tool_context.cache_file( + file_path=file_path, + lines=["content\n"], + ) + self.assertIsNotNone( + await agent.state.tool_context.get_cache(file_path), + ) + + model.set_structured_response( + StructuredResponse( + content={ + "task_overview": "1", + "current_state": "2", + "important_discoveries": "3", + "next_steps": "4", + "context_to_preserve": "5", + }, + ), + ) + + await agent.compress_context() + + self.assertIsNone( + await agent.state.tool_context.get_cache(file_path), + ) + + async def test_context_compression_keeps_reserved_read_cache( + self, + ) -> None: + """Read cache is kept when the same file is still read in context.""" + with tempfile.TemporaryDirectory() as temp_dir: + file_path = os.path.join(temp_dir, "test.txt") + with open(file_path, "w", encoding="utf-8") as f: + f.write("content\n") + + model = MockModel(context_size=100) + agent = Agent( + name="Friday", + system_prompt="".join(["0" for _ in range(20 * 4)]), + model=model, + context_config=ContextConfig( + trigger_ratio=0.7, + reserve_ratio=0.6, + ), + state=AgentState( + session_id="123", + context=[ + AssistantMsg( + "Friday", + [ + ToolCallBlock( + id="read-call-1", + name="Read", + input=json.dumps( + {"file_path": file_path}, + ), + ), + ], + id="1", + ), + UserMsg( + "User", + "".join(["2" for _ in range(30 * 4)]), + id="2", + ), + AssistantMsg( + "Friday", + [ + ToolCallBlock( + id="read-call-2", + name="Read", + input=json.dumps( + {"file_path": file_path}, + ), + ), + ], + id="3", + ), + ], + ), + toolkit=Toolkit(), + ) + await agent.state.tool_context.cache_file( + file_path=file_path, + lines=["content\n"], + ) + + model.set_structured_response( + StructuredResponse( + content={ + "task_overview": "1", + "current_state": "2", + "important_discoveries": "3", + "next_steps": "4", + "context_to_preserve": "5", + }, + ), + ) + + await agent.compress_context() + + self.assertIsNotNone( + await agent.state.tool_context.get_cache(file_path), + ) + + async def test_context_compression_clears_unreferenced_read_cache( + self, + ) -> None: + """Read cache is cleared when no reserved Read references it.""" + with tempfile.TemporaryDirectory() as temp_dir: + file_path = os.path.join(temp_dir, "test.txt") + with open(file_path, "w", encoding="utf-8") as f: + f.write("content\n") + + model = MockModel(context_size=100) + agent = Agent( + name="Friday", + system_prompt="".join(["0" for _ in range(20 * 4)]), + model=model, + context_config=ContextConfig( + trigger_ratio=0.7, + reserve_ratio=0.4, + ), + state=AgentState( + session_id="123", + context=[ + UserMsg( + "User", + "".join(["2" for _ in range(60 * 4)]), + id="1", + ), + UserMsg( + "User", + "".join(["3" for _ in range(30 * 4)]), + id="2", + ), + ], + ), + toolkit=Toolkit(), + ) + await agent.state.tool_context.cache_file( + file_path=file_path, + lines=["content\n"], + ) + + model.set_structured_response( + StructuredResponse( + content={ + "task_overview": "1", + "current_state": "2", + "important_discoveries": "3", + "next_steps": "4", + "context_to_preserve": "5", + }, + ), + ) + + await agent.compress_context() + + self.assertIsNone( + await agent.state.tool_context.get_cache(file_path), + ) + + async def test_context_compression_injects_instructions_as_hint( + self, + ) -> None: + """Instructions are injected as a HintBlock only for compression.""" + model = RecordingStructuredMockModel(context_size=100) + agent = Agent( + name="Friday", + system_prompt="".join(["0" for _ in range(20 * 4)]), + model=model, + context_config=ContextConfig( + trigger_ratio=0.7, + reserve_ratio=0.4, + ), + state=AgentState( + session_id="123", + context=[ + UserMsg( + "User", + "".join(["1" for _ in range(30 * 4)]), + id="1", + ), + AssistantMsg( + "Friday", + "".join(["2" for _ in range(10 * 4)]), + id="2", + ), + UserMsg( + "User", + "".join(["3" for _ in range(10 * 4)]), + id="3", + ), + ], + ), + toolkit=Toolkit(), + ) + + model.set_structured_response( + StructuredResponse( + content={ + "task_overview": "1", + "current_state": "2", + "important_discoveries": "3", + "next_steps": "4", + "context_to_preserve": "5", + }, + ), + ) + instructions = HintBlock( + hint="Keep user requirements and file paths.", + source="user", + ) + + await agent.compress_context(instructions=instructions) + + self.assertEqual(len(model.recorded_structured_messages), 1) + self.assertTrue( + _has_instruction_hint( + model.recorded_structured_messages[0], + instructions, + ), + ) + self.assertFalse( + any(msg.get_content_blocks("hint") for msg in agent.state.context), + ) + + async def test_context_compression_overflow_retry_keeps_instructions( + self, + ) -> None: + """Overflow retry preserves instructions when rebuilding messages.""" + model = RecordingStructuredMockModel( + context_size=100, + fail_structured_output_times=1, + force_compression_overflow=True, + ) + agent = Agent( + name="Friday", + system_prompt="".join(["0" for _ in range(20 * 4)]), + model=model, + context_config=ContextConfig( + trigger_ratio=0.7, + reserve_ratio=0.4, + ), + state=AgentState( + session_id="123", + context=[ + UserMsg( + "User", + "".join(["1" for _ in range(30 * 4)]), + id="1", + ), + AssistantMsg( + "Friday", + "".join(["2" for _ in range(10 * 4)]), + id="2", + ), + UserMsg( + "User", + "".join(["3" for _ in range(10 * 4)]), + id="3", + ), + ], + ), + toolkit=Toolkit(), + ) + + model.set_structured_response( + StructuredResponse( + content={ + "task_overview": "1", + "current_state": "2", + "important_discoveries": "3", + "next_steps": "4", + "context_to_preserve": "5", + }, + ), + ) + instructions = HintBlock( + hint="Keep the user's original success criteria.", + source="user", + ) + + await agent.compress_context(instructions=instructions) + + self.assertEqual(len(model.recorded_structured_messages), 2) + self.assertTrue( + _has_instruction_hint( + model.recorded_structured_messages[-1], + instructions, + ), + ) + + async def asyncTearDown(self) -> None: + """The async teardown method.""" diff --git a/tests/compress_tool_result_test.py b/tests/compress_tool_result_test.py new file mode 100644 index 0000000000000000000000000000000000000000..73495d4a0b0fbbd1e94359f0973f1359f0351d95 --- /dev/null +++ b/tests/compress_tool_result_test.py @@ -0,0 +1,480 @@ +# -*- coding: utf-8 -*- +"""The unittests for the tool result compression.""" +# pylint: disable=protected-access, unused-argument +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import MockModel + +from agentscope.agent import Agent, ContextConfig +from agentscope.message import ( + ToolResultBlock, + TextBlock, + DataBlock, + Base64Source, +) +from agentscope.state import AgentState +from agentscope.tool import Toolkit + + +class ToolResultCompressionTest(IsolatedAsyncioTestCase): + """Test cases for tool result compression.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.mock_model = MockModel() + self.agent = Agent( + name="TestAgent", + system_prompt="Test system prompt", + model=self.mock_model, + toolkit=Toolkit(), + context_config=ContextConfig( + tool_result_limit=100, + ), + state=AgentState(session_id="test_session"), + ) + + async def test_below_limit(self) -> None: + """Test when tool result is below the token limit.""" + tool_result = ToolResultBlock( + id="test_1", + name="test_tool", + output=[ + TextBlock(text="Short text 1"), + TextBlock(text="Short text 2"), + ], + ) + + async def mock_count_tokens( + messages: list, + tools: list | None = None, + ) -> int: + """Mock token counting function that returns a fixed count.""" + return 50 + + self.mock_model.count_tokens = mock_count_tokens + ( + reserved, + offload, + ) = await self.agent._split_tool_result_for_compression( + tool_result, + ) + + self.assertEqual(reserved, tool_result) + self.assertIsNone(offload) + + async def test_equal_to_limit(self) -> None: + """Test when tool result is exactly at the token limit.""" + tool_result = ToolResultBlock( + id="test_2", + name="test_tool", + output=[TextBlock(text="Text at limit")], + ) + + async def mock_count_tokens( + messages: list, + tools: list | None = None, + ) -> int: + """Mock token counting function that returns a fixed count.""" + return 100 + + self.mock_model.count_tokens = mock_count_tokens + ( + reserved, + offload, + ) = await self.agent._split_tool_result_for_compression( + tool_result, + ) + + self.assertEqual(reserved, tool_result) + self.assertIsNone(offload) + + async def test_boundary_last_block_text(self) -> None: + """Test when boundary is the last block and it is a TextBlock.""" + block1 = TextBlock(text="A" * 20, id="block1") + block2 = TextBlock(text="B" * 20, id="block2") + block3 = TextBlock(text="C" * 100, id="block3") + + tool_result = ToolResultBlock( + id="test_3", + name="test_tool", + output=[block1, block2, block3], + ) + + async def mock_count_tokens( + messages: list, + tools: list | None = None, + ) -> int: + """Mock token counting function that counts text length in + blocks.""" + content = messages[0].content + if isinstance(content, list): + total = sum(len(b.text) for b in content if hasattr(b, "text")) + return total + return 0 + + self.mock_model.count_tokens = mock_count_tokens + ( + reserved, + offload, + ) = await self.agent._split_tool_result_for_compression( + tool_result, + ) + + # Verify results + self.assertIsNotNone(reserved) + self.assertIsNotNone(offload) + + # Verify ToolResultBlock metadata + self.assertEqual(reserved.id, tool_result.id) + self.assertEqual(reserved.name, tool_result.name) + self.assertEqual(reserved.state, tool_result.state) + self.assertEqual(offload.id, tool_result.id) + self.assertEqual(offload.name, tool_result.name) + self.assertEqual(offload.state, tool_result.state) + + # Verify results using assertListEqual + expected_reserved = [ + {"type": "text", "text": "A" * 20, "id": "block1"}, + {"type": "text", "text": "B" * 20 + "C" * 60, "id": "block2"}, + ] + expected_offload = [ + {"type": "text", "text": "C" * 40, "id": "block3"}, + ] + + self.assertListEqual( + [b.model_dump() for b in reserved.output], + expected_reserved, + ) + self.assertListEqual( + [b.model_dump() for b in offload.output], + expected_offload, + ) + + async def test_boundary_last_block_data(self) -> None: + """Test when boundary is the last block and it is a DataBlock.""" + block1 = TextBlock(text="A" * 20, id="block1") + block2 = TextBlock(text="B" * 20, id="block2") + block3 = DataBlock( + source=Base64Source(data="base64data", media_type="image/png"), + id="block3", + ) + + tool_result = ToolResultBlock( + id="test_4", + name="test_tool", + output=[block1, block2, block3], + ) + + async def mock_count_tokens( + messages: list, + tools: list | None = None, + ) -> int: + """Mock token counting function based on content length.""" + content = messages[0].content + if isinstance(content, list): + if len(content) == 3: + return 150 + elif len(content) == 2: + return 80 + return 50 + + self.mock_model.count_tokens = mock_count_tokens + ( + reserved, + offload, + ) = await self.agent._split_tool_result_for_compression( + tool_result, + ) + + # Verify results + self.assertIsNotNone(reserved) + self.assertIsNotNone(offload) + + # Verify ToolResultBlock metadata + self.assertEqual(reserved.id, tool_result.id) + self.assertEqual(offload.id, tool_result.id) + + # Verify results using assertListEqual + expected_reserved = [ + {"type": "text", "text": "A" * 20, "id": "block1"}, + {"type": "text", "text": "B" * 20, "id": "block2"}, + ] + expected_offload = [ + { + "type": "data", + "id": "block3", + "source": { + "type": "base64", + "data": "base64data", + "media_type": "image/png", + }, + "name": None, + }, + ] + + self.assertListEqual( + [b.model_dump() for b in reserved.output], + expected_reserved, + ) + self.assertListEqual( + [b.model_dump() for b in offload.output], + expected_offload, + ) + + async def test_boundary_first_block_text(self) -> None: + """Test when boundary is the first block and it is a TextBlock.""" + block1 = TextBlock(text="A" * 100, id="block1") + block2 = TextBlock(text="B" * 20, id="block2") + block3 = TextBlock(text="C" * 20, id="block3") + + tool_result = ToolResultBlock( + id="test_5", + name="test_tool", + output=[block1, block2, block3], + ) + + async def mock_count_tokens( + messages: list, + tools: list | None = None, + ) -> int: + """Mock token counting function that counts text length in + blocks.""" + content = messages[0].content + if isinstance(content, list): + total = sum(len(b.text) for b in content if hasattr(b, "text")) + return total + return 0 + + self.mock_model.count_tokens = mock_count_tokens + ( + reserved, + offload, + ) = await self.agent._split_tool_result_for_compression( + tool_result, + ) + + # Verify results using assertListEqual + expected_reserved = [ + {"type": "text", "text": "A" * 100, "id": "block1"}, + ] + expected_offload = [ + {"type": "text", "text": "B" * 20, "id": "block2"}, + {"type": "text", "text": "C" * 20, "id": "block3"}, + ] + + self.assertListEqual( + [b.model_dump() for b in reserved.output], + expected_reserved, + ) + print(offload.output) + self.assertListEqual( + [b.model_dump() for b in offload.output], + expected_offload, + ) + + async def test_boundary_first_block_data(self) -> None: + """Test when boundary is the first block and it is a DataBlock.""" + block1 = DataBlock( + source=Base64Source(data="base64data", media_type="image/png"), + id="block1", + ) + block2 = TextBlock(text="B" * 20, id="block2") + block3 = TextBlock(text="C" * 20, id="block3") + + tool_result = ToolResultBlock( + id="test_6", + name="test_tool", + output=[block1, block2, block3], + ) + + async def mock_count_tokens( + messages: list, + tools: list | None = None, + ) -> int: + """Mock token counting function based on content length.""" + content = messages[0].content + if isinstance(content, list): + if len(content) == 3: + return 150 + elif len(content) == 2: + return 80 + elif len(content) == 1: + return 60 + return 50 + + self.mock_model.count_tokens = mock_count_tokens + ( + reserved, + offload, + ) = await self.agent._split_tool_result_for_compression( + tool_result, + ) + + # Verify results + self.assertIsNotNone(reserved) + self.assertIsNotNone(offload) + + # Verify ToolResultBlock metadata + self.assertEqual(reserved.id, tool_result.id) + self.assertEqual(offload.id, tool_result.id) + + # Verify results using assertListEqual + expected_reserved = [ + { + "type": "data", + "id": "block1", + "source": { + "type": "base64", + "data": "base64data", + "media_type": "image/png", + }, + "name": None, + }, + {"type": "text", "text": "B" * 20 + "C" * 5, "id": "block2"}, + ] + expected_offload = [ + {"type": "text", "text": "C" * 15, "id": "block3"}, + ] + + self.assertListEqual( + [b.model_dump() for b in reserved.output], + expected_reserved, + ) + self.assertListEqual( + [b.model_dump() for b in offload.output], + expected_offload, + ) + + async def test_boundary_middle_block_text(self) -> None: + """Test when boundary is a middle block and it is a TextBlock.""" + block1 = TextBlock(text="A" * 20, id="block1") + block2 = TextBlock(text="B" * 100, id="block2") + block3 = TextBlock(text="C" * 20, id="block3") + + tool_result = ToolResultBlock( + id="test_7", + name="test_tool", + output=[block1, block2, block3], + ) + + async def mock_count_tokens( + messages: list, + tools: list | None = None, + ) -> int: + """Mock token counting function that counts text length in + blocks.""" + content = messages[0].content + if isinstance(content, list): + total = sum(len(b.text) for b in content if hasattr(b, "text")) + return total + return 0 + + self.mock_model.count_tokens = mock_count_tokens + ( + reserved, + offload, + ) = await self.agent._split_tool_result_for_compression( + tool_result, + ) + + # Verify results + self.assertIsNotNone(reserved) + self.assertIsNotNone(offload) + + # Verify ToolResultBlock metadata + self.assertEqual(reserved.id, tool_result.id) + self.assertEqual(offload.id, tool_result.id) + + # Verify results using assertListEqual + expected_reserved = [ + {"type": "text", "text": "A" * 20 + "B" * 80, "id": "block1"}, + ] + expected_offload = [ + {"type": "text", "text": "B" * 20 + "C" * 20, "id": "block3"}, + ] + + self.assertListEqual( + [b.model_dump() for b in reserved.output], + expected_reserved, + ) + self.assertListEqual( + [b.model_dump() for b in offload.output], + expected_offload, + ) + + async def test_boundary_middle_block_data(self) -> None: + """Test when boundary is a middle block and it is a DataBlock.""" + block1 = TextBlock(text="A" * 20, id="block1") + block2 = DataBlock( + source=Base64Source(data="base64data", media_type="image/png"), + id="block2", + ) + block3 = TextBlock(text="C" * 20, id="block3") + + tool_result = ToolResultBlock( + id="test_8", + name="test_tool", + output=[block1, block2, block3], + ) + + async def mock_count_tokens( + messages: list, + tools: list | None = None, + ) -> int: + """Mock token counting function based on content length.""" + content = messages[0].content + if isinstance(content, list): + if len(content) == 3: + return 150 + elif len(content) == 2: + return 80 + elif len(content) == 1: + return 40 + return 50 + + self.mock_model.count_tokens = mock_count_tokens + ( + reserved, + offload, + ) = await self.agent._split_tool_result_for_compression( + tool_result, + ) + + # Verify results + self.assertIsNotNone(reserved) + self.assertIsNotNone(offload) + + # Verify ToolResultBlock metadata + self.assertEqual(reserved.id, tool_result.id) + self.assertEqual(offload.id, tool_result.id) + + # Verify results using assertListEqual + expected_reserved = [ + {"type": "text", "text": "A" * 20, "id": "block1"}, + { + "type": "data", + "id": "block2", + "source": { + "type": "base64", + "data": "base64data", + "media_type": "image/png", + }, + "name": None, + }, + {"type": "text", "text": "C" * 5, "id": "block3"}, + ] + expected_offload = [ + {"type": "text", "text": "C" * 15, "id": "block3"}, + ] + + self.assertListEqual( + [b.model_dump() for b in reserved.output], + expected_reserved, + ) + self.assertListEqual( + [b.model_dump() for b in offload.output], + expected_offload, + ) + + async def asyncTearDown(self) -> None: + """The async teardown method.""" diff --git a/tests/embedding_dashscope_test.py b/tests/embedding_dashscope_test.py new file mode 100644 index 0000000000000000000000000000000000000000..9f768f39581b3be21ab9343c0cc87ee088e25e04 --- /dev/null +++ b/tests/embedding_dashscope_test.py @@ -0,0 +1,233 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access,unused-argument +"""Unit tests for DashScopeEmbeddingModel.""" +from dataclasses import asdict +from typing import Any +from unittest import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, MagicMock, patch + +from utils import AnyValue + +from agentscope.credential import DashScopeCredential +from agentscope.embedding import ( + DashScopeEmbeddingModel, + EmbeddingResponse, + EmbeddingUsage, +) +from agentscope.message import DataBlock, Base64Source, URLSource + +A = AnyValue() + + +def _text_resp( + embeddings: list[list[float]], + total_tokens: int = 10, + status_code: int = 200, +) -> MagicMock: + """Build a mock DashScope text embedding response.""" + resp = MagicMock() + resp.status_code = status_code + resp.output = {"embeddings": [{"embedding": e} for e in embeddings]} + resp.usage = {"total_tokens": total_tokens} + return resp + + +def _cred() -> DashScopeCredential: + """Create a test credential.""" + return DashScopeCredential(api_key="k") + + +def _img() -> DataBlock: + """Create a test image DataBlock.""" + return DataBlock( + source=Base64Source(data="aWltYWdl", media_type="image/png"), + ) + + +def _vid() -> DataBlock: + """Create a test video DataBlock.""" + return DataBlock( + source=URLSource(url="https://x.com/v.mp4", media_type="video/mp4"), + ) + + +def _mock_resp(embeddings: list[list[float]]) -> EmbeddingResponse: + """Create a mock EmbeddingResponse.""" + return EmbeddingResponse( + embeddings=embeddings, + usage=EmbeddingUsage(tokens=len(embeddings), time=0.01), + ) + + +class DashScopeListModelsTest(IsolatedAsyncioTestCase): + """Test list_models for DashScope.""" + + async def test_list_models(self) -> None: + """Should list 7 models (text + multimodal).""" + cards = DashScopeEmbeddingModel.list_models() + names = sorted(c.name for c in cards) + self.assertEqual(len(cards), 7) + self.assertIn("text-embedding-v4", names) + self.assertIn("qwen3-vl-embedding", names) + self.assertIn("multimodal-embedding-v1", names) + + async def test_hidden_dimensions(self) -> None: + """multimodal-embedding-v1 declares a fixed dimension (no enum).""" + cards = DashScopeEmbeddingModel.list_models() + v1 = next(c for c in cards if c.name == "multimodal-embedding-v1") + self.assertDictEqual( + v1.model_dump(), + { + "type": "embedding_model", + "name": "multimodal-embedding-v1", + "label": "Multimodal Embedding v1", + "status": "active", + "input_types": [ + "text/plain", + "image/jpeg", + "image/png", + "image/bmp", + ], + "output_types": ["application/x-embedding"], + "dimensions": 1024, + "supported_dimensions": None, + "context_size": 512, + "parameter_schema": { + "type": "object", + "properties": {}, + "required": [], + }, + "parameter_overrides": {}, + }, + ) + + async def test_visible_dimensions(self) -> None: + """qwen3-vl-embedding exposes dimension choices on the card.""" + cards = DashScopeEmbeddingModel.list_models() + qwen = next(c for c in cards if c.name == "qwen3-vl-embedding") + self.assertEqual(qwen.dimensions, 2560) + self.assertEqual( + qwen.supported_dimensions, + [2560, 2048, 1536, 1024, 768, 512, 256], + ) + + +class DashScopeTextCallTest(IsolatedAsyncioTestCase): + """Test DashScope text embedding API calls.""" + + @patch("dashscope.embeddings.TextEmbedding.call") + async def test_text_call(self, mock_api: Any) -> None: + """Text mode returns correct embeddings.""" + mock_api.return_value = _text_resp([[0.1, 0.2], [0.3, 0.4]], 12) + model = DashScopeEmbeddingModel( + credential=_cred(), + model="text-embedding-v4", + dimensions=2, + ) + result = await model(["hello", "world"]) + self.assertDictEqual( + asdict(result), + { + "embeddings": [[0.1, 0.2], [0.3, 0.4]], + "id": A, + "created_at": A, + "type": "embedding", + "usage": {"tokens": 12, "time": A, "type": "embedding"}, + "source": "api", + }, + ) + + @patch("dashscope.embeddings.TextEmbedding.call") + async def test_text_rejects_datablock(self, mock_api: Any) -> None: + """Text mode rejects DataBlock inputs.""" + model = DashScopeEmbeddingModel( + credential=_cred(), + model="text-embedding-v4", + dimensions=1024, + ) + with self.assertRaises(ValueError): + await model([_img()]) + + @patch("dashscope.embeddings.TextEmbedding.call") + async def test_text_api_error_raises(self, mock_api: Any) -> None: + """Non-200 status code raises RuntimeError after retries.""" + mock_api.return_value = _text_resp([], status_code=400) + model = DashScopeEmbeddingModel( + credential=_cred(), + model="text-embedding-v4", + dimensions=1024, + retry_delay=0.0, + ) + with self.assertRaises(RuntimeError): + await model(["hello"]) + + +class DashScopeMultimodalCallTest(IsolatedAsyncioTestCase): + """Test DashScope multimodal embedding via mocked _call_multimodal.""" + + async def test_multimodal_text_and_image(self) -> None: + """Multimodal call with text + image.""" + model = DashScopeEmbeddingModel( + credential=_cred(), + model="qwen3-vl-embedding", + dimensions=2, + ) + model._call_multimodal = AsyncMock( + return_value=_mock_resp([[0.1, 0.2], [0.3, 0.4]]), + ) + result = await model(["describe this", _img()]) + self.assertDictEqual( + asdict(result), + { + "embeddings": [[0.1, 0.2], [0.3, 0.4]], + "id": A, + "created_at": A, + "type": "embedding", + "usage": {"tokens": 2, "time": 0.01, "type": "embedding"}, + "source": "api", + }, + ) + + async def test_multimodal_batching_by_image_limit(self) -> None: + """8 images with max_images=5 produces 2 batches (5+3).""" + model = DashScopeEmbeddingModel( + credential=_cred(), + model="qwen3-vl-embedding", + dimensions=1, + ) + call_count = 0 + + async def _mock(inputs: list, **_kw: Any) -> EmbeddingResponse: + nonlocal call_count + call_count += 1 + return _mock_resp([[0.1]] * len(inputs)) + + model._call_multimodal = _mock # type: ignore[assignment] + result = await model([_img() for _ in range(8)]) + self.assertEqual(result["embeddings"], [[0.1]] * 8) + self.assertEqual(call_count, 2) + + async def test_multimodal_batching_by_video_limit(self) -> None: + """3 videos with max_videos=1 produces 3 batches.""" + model = DashScopeEmbeddingModel( + credential=_cred(), + model="qwen3-vl-embedding", + dimensions=1, + ) + call_count = 0 + + async def _mock(inputs: list, **_kw: Any) -> EmbeddingResponse: + nonlocal call_count + call_count += 1 + return _mock_resp([[0.1]] * len(inputs)) + + model._call_multimodal = _mock # type: ignore[assignment] + result = await model([_vid(), _vid(), _vid()]) + self.assertEqual(result["embeddings"], [[0.1]] * 3) + self.assertEqual(call_count, 3) + + async def test_video_base64_rejected(self) -> None: + """Video with Base64Source raises ValueError.""" + bad = DataBlock(source=Base64Source(data="x", media_type="video/mp4")) + with self.assertRaises(ValueError): + DashScopeEmbeddingModel._format_data_block(bad) diff --git a/tests/embedding_gemini_test.py b/tests/embedding_gemini_test.py new file mode 100644 index 0000000000000000000000000000000000000000..c4e99ea7dbe61ffeab38c35fe9a6c50cc548a120 --- /dev/null +++ b/tests/embedding_gemini_test.py @@ -0,0 +1,179 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for GeminiEmbeddingModel.""" +from dataclasses import asdict +from typing import Any +from unittest import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock + +from utils import AnyValue + +from agentscope.embedding import ( + GeminiEmbeddingModel, + EmbeddingResponse, + EmbeddingUsage, +) +from agentscope.message import DataBlock, Base64Source + +A = AnyValue() + + +def _mock_resp(embeddings: list[list[float]]) -> EmbeddingResponse: + """Create a mock EmbeddingResponse.""" + return EmbeddingResponse( + embeddings=embeddings, + usage=EmbeddingUsage(tokens=len(embeddings), time=0.01), + ) + + +def _img() -> DataBlock: + """Create a test image DataBlock.""" + return DataBlock( + source=Base64Source(data="aWltYWdl", media_type="image/png"), + ) + + +class GeminiListModelsTest(IsolatedAsyncioTestCase): + """Test list_models for Gemini.""" + + async def test_list_models(self) -> None: + """Should list 2 models.""" + cards = GeminiEmbeddingModel.list_models() + names = sorted(c.name for c in cards) + self.assertEqual(names, ["gemini-embedding-001", "gemini-embedding-2"]) + + async def test_text_model_card(self) -> None: + """gemini-embedding-001 is text-only with 2048 context.""" + cards = GeminiEmbeddingModel.list_models() + card = next(c for c in cards if c.name == "gemini-embedding-001") + self.assertDictEqual( + card.model_dump(), + { + "type": "embedding_model", + "name": "gemini-embedding-001", + "label": "Gemini Embedding 001", + "status": "active", + "input_types": ["text/plain"], + "output_types": ["application/x-embedding"], + "dimensions": 3072, + "supported_dimensions": [3072, 1536, 768, 512, 256, 128], + "context_size": 2048, + "parameter_schema": { + "type": "object", + "properties": {}, + "required": [], + }, + "parameter_overrides": {}, + }, + ) + + async def test_multimodal_model_card(self) -> None: + """gemini-embedding-2 is multimodal with 8192 context.""" + cards = GeminiEmbeddingModel.list_models() + card = next(c for c in cards if c.name == "gemini-embedding-2") + self.assertIn("image/png", card.input_types) + self.assertIn("application/pdf", card.input_types) + self.assertEqual(card.context_size, 8192) + self.assertEqual(card.supported_dimensions, [3072, 1536, 768]) + + +class GeminiTextCallTest(IsolatedAsyncioTestCase): + """Test Gemini text embedding via mocked _call_text.""" + + def _make_text_model(self) -> GeminiEmbeddingModel: + """Create a text-mode model bypassing __init__ (no genai).""" + model = GeminiEmbeddingModel.__new__(GeminiEmbeddingModel) + model.model = "gemini-embedding-001" + model.dimensions = 2 + model.context_size = 2048 + model.batch_size = 100 + model.max_retries = 3 + model.retry_delay = 1.0 + model._is_multimodal = False + model.embedding_cache = None + return model + + async def test_text_call(self) -> None: + """Text mode delegates to _call_text.""" + model = self._make_text_model() + model._call_text = AsyncMock( + return_value=_mock_resp([[0.1, 0.2], [0.3, 0.4]]), + ) + result = await model(["hello", "world"]) + self.assertDictEqual( + asdict(result), + { + "embeddings": [[0.1, 0.2], [0.3, 0.4]], + "id": A, + "created_at": A, + "type": "embedding", + "usage": {"tokens": 2, "time": 0.01, "type": "embedding"}, + "source": "api", + }, + ) + + async def test_text_rejects_datablock(self) -> None: + """Text mode rejects DataBlock inputs.""" + model = self._make_text_model() + with self.assertRaises(ValueError): + await GeminiEmbeddingModel._call_text(model, [_img()]) + + +class GeminiMultimodalCallTest(IsolatedAsyncioTestCase): + """Test Gemini multimodal embedding via mocked _call_multimodal.""" + + def _make_multimodal_model(self) -> GeminiEmbeddingModel: + """Create a multimodal-mode model bypassing __init__.""" + model = GeminiEmbeddingModel.__new__(GeminiEmbeddingModel) + model.model = "gemini-embedding-2" + model.dimensions = 1 + model.context_size = 8192 + model.batch_size = 100 + model.max_retries = 3 + model.retry_delay = 1.0 + model._is_multimodal = True + model.embedding_cache = None + from agentscope.embedding._gemini._model import _MultimodalLimits + + model._limits = _MultimodalLimits( + max_elements=20, + max_images=6, + max_videos=1, + max_audios=1, + max_pdfs=1, + ) + return model + + async def test_multimodal_delegates(self) -> None: + """Multimodal mode delegates to _call_multimodal.""" + model = self._make_multimodal_model() + model._call_multimodal = AsyncMock( + return_value=_mock_resp([[0.1], [0.2]]), + ) + result = await model(["hello", "world"]) + self.assertDictEqual( + asdict(result), + { + "embeddings": [[0.1], [0.2]], + "id": A, + "created_at": A, + "type": "embedding", + "usage": {"tokens": 2, "time": 0.01, "type": "embedding"}, + "source": "api", + }, + ) + + async def test_multimodal_batching_by_image_limit(self) -> None: + """8 images with max_images=6 produces 2 batches (6+2).""" + model = self._make_multimodal_model() + call_count = 0 + + async def _mock(inputs: list, **_kw: Any) -> EmbeddingResponse: + nonlocal call_count + call_count += 1 + return _mock_resp([[0.1]] * len(inputs)) + + model._call_multimodal = _mock # type: ignore[assignment] + result = await model([_img() for _ in range(8)]) + self.assertEqual(result["embeddings"], [[0.1]] * 8) + self.assertEqual(call_count, 2) diff --git a/tests/embedding_ollama_test.py b/tests/embedding_ollama_test.py new file mode 100644 index 0000000000000000000000000000000000000000..17c71882bcc72fba72c3a3c2dd22d09abf309e43 --- /dev/null +++ b/tests/embedding_ollama_test.py @@ -0,0 +1,106 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for OllamaEmbeddingModel.""" +from dataclasses import asdict +from typing import Any +from unittest import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock + +from utils import AnyValue + +from agentscope.credential import OllamaCredential +from agentscope.embedding import ( + OllamaEmbeddingModel, + EmbeddingResponse, + EmbeddingUsage, +) + +A = AnyValue() + + +def _cred() -> OllamaCredential: + """Create a test credential.""" + return OllamaCredential(host="http://localhost:11434") + + +def _mock_resp(embeddings: list[list[float]]) -> EmbeddingResponse: + """Create a mock EmbeddingResponse.""" + return EmbeddingResponse( + embeddings=embeddings, + usage=EmbeddingUsage(tokens=len(embeddings), time=0.01), + ) + + +class OllamaListModelsTest(IsolatedAsyncioTestCase): + """Test list_models for Ollama.""" + + async def test_list_models_empty(self) -> None: + """Ollama has no pre-defined YAMLs, returns empty list.""" + self.assertEqual(OllamaEmbeddingModel.list_models(), []) + + +class OllamaEmbeddingCallTest(IsolatedAsyncioTestCase): + """Test Ollama embedding via mocked _call_api.""" + + async def test_basic_call(self) -> None: + """Basic call returns correct embeddings.""" + model = OllamaEmbeddingModel( + credential=_cred(), + model="nomic-embed-text", + dimensions=2, + ) + model._call_api = AsyncMock( + return_value=_mock_resp([[0.1, 0.2], [0.3, 0.4]]), + ) + result = await model(["hello", "world"]) + self.assertDictEqual( + asdict(result), + { + "embeddings": [[0.1, 0.2], [0.3, 0.4]], + "id": A, + "created_at": A, + "type": "embedding", + "usage": {"tokens": 2, "time": 0.01, "type": "embedding"}, + "source": "api", + }, + ) + + async def test_dimensions_and_host(self) -> None: + """Dimensions and host are set correctly from constructor.""" + model = OllamaEmbeddingModel( + credential=OllamaCredential(host="http://gpu:11434"), + model="test", + dimensions=768, + ) + self.assertEqual(model.dimensions, 768) + self.assertEqual(model.host, "http://gpu:11434") + + async def test_multi_batch(self) -> None: + """Batching splits inputs correctly.""" + model = OllamaEmbeddingModel( + credential=_cred(), + model="test", + dimensions=1, + ) + model.batch_size = 2 + call_count = 0 + + async def _mock(inputs: list[str], **_kw: Any) -> EmbeddingResponse: + nonlocal call_count + call_count += 1 + return _mock_resp([[0.1]] * len(inputs)) + + model._call_api = _mock # type: ignore[assignment] + result = await model(["a", "b", "c"]) + self.assertDictEqual( + asdict(result), + { + "embeddings": [[0.1], [0.1], [0.1]], + "id": A, + "created_at": A, + "type": "embedding", + "usage": {"tokens": A, "time": A, "type": "embedding"}, + "source": "api", + }, + ) + self.assertEqual(call_count, 2) diff --git a/tests/embedding_openai_test.py b/tests/embedding_openai_test.py new file mode 100644 index 0000000000000000000000000000000000000000..fbb6a3692b8f3932572eaa608f8e2e822910b8fe --- /dev/null +++ b/tests/embedding_openai_test.py @@ -0,0 +1,183 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for OpenAIEmbeddingModel.""" +from dataclasses import asdict +from typing import Any +from unittest import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, MagicMock, patch + +from utils import AnyValue + +from agentscope.credential import OpenAICredential +from agentscope.embedding import OpenAIEmbeddingModel + +A = AnyValue() + + +def _make_response( + embeddings: list[list[float]], + total_tokens: int = 10, +) -> MagicMock: + """Build a mock ``openai.embeddings.create`` response.""" + resp = MagicMock() + resp.data = [MagicMock(embedding=e) for e in embeddings] + resp.usage = MagicMock(total_tokens=total_tokens) + return resp + + +class OpenAIListModelsTest(IsolatedAsyncioTestCase): + """Test ``list_models()`` for OpenAI.""" + + async def test_list_models(self) -> None: + """Should list 2 models with correct parameter_schema.""" + cards = OpenAIEmbeddingModel.list_models() + names = sorted(c.name for c in cards) + self.assertEqual( + names, + ["text-embedding-3-large", "text-embedding-3-small"], + ) + + card = next(c for c in cards if c.name == "text-embedding-3-small") + self.assertDictEqual( + card.model_dump(), + { + "type": "embedding_model", + "name": "text-embedding-3-small", + "label": "Text Embedding 3 Small", + "status": "active", + "input_types": ["text/plain"], + "output_types": ["application/x-embedding"], + "dimensions": 1536, + "supported_dimensions": [1536, 1024, 768, 512, 256], + "context_size": 8191, + "parameter_schema": { + "type": "object", + "properties": {}, + "required": [], + }, + "parameter_overrides": {}, + }, + ) + + +class OpenAIEmbeddingCallTest(IsolatedAsyncioTestCase): + """Test OpenAI embedding API calls with mocked responses.""" + + @patch("openai.AsyncClient") + async def test_single_batch(self, mock_client_cls: Any) -> None: + """Single batch call returns correct embeddings.""" + mock_client = MagicMock() + mock_client.embeddings.create = AsyncMock( + return_value=_make_response([[0.1, 0.2], [0.3, 0.4]], 8), + ) + mock_client_cls.return_value = mock_client + + model = OpenAIEmbeddingModel( + credential=OpenAICredential(api_key="k"), + model="text-embedding-3-small", + dimensions=2, + ) + result = await model(["hello", "world"]) + + self.assertDictEqual( + asdict(result), + { + "embeddings": [[0.1, 0.2], [0.3, 0.4]], + "id": A, + "created_at": A, + "type": "embedding", + "usage": {"tokens": 8, "time": A, "type": "embedding"}, + "source": "api", + }, + ) + + @patch("openai.AsyncClient") + async def test_multi_batch(self, mock_client_cls: Any) -> None: + """Inputs exceeding batch_size are split and merged.""" + mock_client = MagicMock() + mock_client.embeddings.create = AsyncMock( + side_effect=[ + _make_response([[0.1], [0.2]], 4), + _make_response([[0.3]], 2), + ], + ) + mock_client_cls.return_value = mock_client + + model = OpenAIEmbeddingModel( + credential=OpenAICredential(api_key="k"), + model="text-embedding-3-small", + dimensions=1, + ) + model.batch_size = 2 + + result = await model(["a", "b", "c"]) + + self.assertDictEqual( + asdict(result), + { + "embeddings": [[0.1], [0.2], [0.3]], + "id": A, + "created_at": A, + "type": "embedding", + "usage": {"tokens": 6, "time": A, "type": "embedding"}, + "source": "api", + }, + ) + + @patch("openai.AsyncClient") + async def test_empty_input(self, mock_client_cls: Any) -> None: + """Empty input returns empty response without API call.""" + mock_client = MagicMock() + mock_client_cls.return_value = mock_client + + model = OpenAIEmbeddingModel( + credential=OpenAICredential(api_key="k"), + model="text-embedding-3-small", + dimensions=1536, + ) + result = await model([]) + + self.assertDictEqual( + asdict(result), + { + "embeddings": [], + "id": A, + "created_at": A, + "type": "embedding", + "usage": {"tokens": 0, "time": 0, "type": "embedding"}, + "source": "api", + }, + ) + mock_client.embeddings.create.assert_not_called() + + @patch("openai.AsyncClient") + async def test_retry_on_transient_error( + self, + mock_client_cls: Any, + ) -> None: + """Retryable OpenAI errors are retried.""" + import openai + + mock_client = MagicMock() + mock_client.embeddings.create = AsyncMock( + side_effect=[ + openai.RateLimitError( + message="rate limit", + response=MagicMock(status_code=429), + body=None, + ), + _make_response([[0.1]], 1), + ], + ) + mock_client_cls.return_value = mock_client + + model = OpenAIEmbeddingModel( + credential=OpenAICredential(api_key="k"), + model="text-embedding-3-small", + dimensions=1, + retry_delay=0.0, + ) + result = await model(["hello"]) + + self.assertEqual(result["embeddings"], [[0.1]]) + self.assertEqual(mock_client.embeddings.create.await_count, 2) diff --git a/tests/event_test.py b/tests/event_test.py new file mode 100644 index 0000000000000000000000000000000000000000..c0a9368320ea4dd51b959ec2a23903e23c36b8de --- /dev/null +++ b/tests/event_test.py @@ -0,0 +1,50 @@ +# -*- coding: utf-8 -*- +"""Event test""" +from unittest.async_case import IsolatedAsyncioTestCase +from utils import AnyString +from agentscope.event import ReplyStartEvent + + +class EventTest(IsolatedAsyncioTestCase): + """The event test case.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + + async def test_model_dump(self) -> None: + """Test model dump.""" + event = ReplyStartEvent( + session_id="test_session", + reply_id="test_reply", + name="Friday", + ).model_dump() + self.assertDictEqual( + event, + { + "type": "REPLY_START", + "id": AnyString(), + "created_at": AnyString(), + "metadata": {}, + "session_id": "test_session", + "reply_id": "test_reply", + "name": "Friday", + "role": "assistant", + }, + ) + self.assertIsInstance(event["type"], str) + + async def test_model_validate(self) -> None: + """Test model validate.""" + data = { + "type": "REPLY_START", + "id": "test_id", + "created_at": "2024-01-01T00:00:00", + "session_id": "test_session", + "reply_id": "test_reply", + "name": "Friday", + "role": "assistant", + } + ReplyStartEvent.model_validate(data) + + async def asyncTearDown(self) -> None: + """The async teardown method.""" diff --git a/tests/event_to_message_test.py b/tests/event_to_message_test.py new file mode 100644 index 0000000000000000000000000000000000000000..13fd571329c386e3fdb83d56a34501b34f4aee8e --- /dev/null +++ b/tests/event_to_message_test.py @@ -0,0 +1,923 @@ +# -*- coding: utf-8 -*- +"""Unit tests for Msg.append_event – event-stream-to-Msg accumulation. + +The test drives a single Msg object through a full, realistic streaming +sequence and asserts the complete model_dump() after every individual event. + +Coverage +-------- +* TextBlock : start / delta (×2) / end +* ThinkingBlock : start / delta (×2) / end +* DataBlock (base-64) : start / delta (×2) / end +* ToolCallBlock streaming : start / delta (×2) / end +* RequireUserConfirmEvent → ASKING state +* UserConfirmResultEvent (confirmed=True) → ALLOWED state +* UserConfirmResultEvent (confirmed=False) → FINISHED state +* ToolResultBlock text output : start / text-delta (×2) / end (SUCCESS) +* RequireExternalExecutionEvent → SUBMITTED state +* ExternalExecutionResultEvent → ToolResultBlock appended directly +* ToolResultBlock data output : base-64 delta + URL delta / end (ERROR) +* ModelCallEndEvent (×2) → usage initialized then accumulated +* ReplyEndEvent → finished_at stamped +* Wrong reply_id → event silently skipped +* Missing block → warning, no crash +""" +from typing import Any +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import AnyString + +from agentscope.event import ( + ConfirmResult, + DataBlockDeltaEvent, + DataBlockEndEvent, + DataBlockStartEvent, + ExternalExecutionResultEvent, + ModelCallEndEvent, + ReplyEndEvent, + RequireExternalExecutionEvent, + RequireUserConfirmEvent, + TextBlockDeltaEvent, + TextBlockEndEvent, + TextBlockStartEvent, + ThinkingBlockDeltaEvent, + ThinkingBlockEndEvent, + ThinkingBlockStartEvent, + ToolCallDeltaEvent, + ToolCallEndEvent, + ToolCallStartEvent, + ToolResultDataDeltaEvent, + ToolResultEndEvent, + ToolResultStartEvent, + ToolResultTextDeltaEvent, + UserConfirmResultEvent, +) +from agentscope.message import ( + Msg, + ToolCallBlock, + ToolResultBlock, + ToolResultState, +) + +# --------------------------------------------------------------------------- +# Fixed IDs used throughout – hard-coded so ground-truth dicts are readable. +# --------------------------------------------------------------------------- +_REPLY_ID = "reply_001" +_SESSION_ID = "session_001" + +_B_TEXT = "b_text_001" # TextBlock id +_B_THINK = "b_think_001" # ThinkingBlock id +_B_DATA = "b_data_001" # DataBlock id + +_TC_ALLOW = "tc_allow_001" # tool call that gets confirmed → allowed +_TC_DENY = "tc_deny_001" # tool call that gets denied → finished +_TC_EXT = "tc_ext_001" # tool call for external execution +_TC_IMG = "tc_img_001" # tool call whose result has data blocks + +_RES_DATA_B = "res_data_001" # DataBlock inside tool-result (base-64) +_RES_URL_B = "res_url_001" # DataBlock inside tool-result (URL) + +_FIXED_END_TS = "2026-01-01T12:00:00" # deterministic finished_at + + +# --------------------------------------------------------------------------- +# Block-dict helpers – each call returns a fresh dict to avoid aliasing. +# --------------------------------------------------------------------------- + + +def _tb(block_id: str, text: str) -> dict: + """Text block dict.""" + return {"type": "text", "id": block_id, "text": text} + + +def _thb(block_id: str, thinking: str) -> dict: + """Thinking block dict.""" + return {"type": "thinking", "id": block_id, "thinking": thinking} + + +def _db_b64(block_id: str, data: str, media_type: str) -> dict: + """DataBlock (base-64 source) dict.""" + return { + "type": "data", + "id": block_id, + "source": {"type": "base64", "data": data, "media_type": media_type}, + "name": None, + } + + +def _db_url(block_id: str, url: str, media_type: str) -> dict: + """DataBlock (URL source) dict.""" + return { + "type": "data", + "id": block_id, + "source": {"type": "url", "url": url, "media_type": media_type}, + "name": None, + } + + +def _tcb(tc_id: str, name: str, inp: str, state: str) -> dict: + """ToolCallBlock dict.""" + return { + "type": "tool_call", + "id": tc_id, + "name": name, + "input": inp, + "state": state, + "suggested_rules": [], + } + + +def _trb(tc_id: str, name: str, output: Any, state: str) -> dict: + """ToolResultBlock dict.""" + return { + "type": "tool_result", + "id": tc_id, + "name": name, + "output": output, + "state": state, + "metadata": {}, + } + + +class EventToMessageTest(IsolatedAsyncioTestCase): + """Test Msg.append_event across a full streaming event sequence.""" + + # ------------------------------------------------------------------ + # asyncSetUp: build self.events and self.ground_truths in lock-step. + # ------------------------------------------------------------------ + + async def asyncSetUp(self) -> None: + """Build the Msg, the event list, and the ground-truth list.""" + self.msg = Msg( + id=_REPLY_ID, + name="TestAgent", + role="assistant", + content=[], + ) + _created_at = self.msg.created_at + + def _base( + content: list, + finished_at: str | None = None, + usage: dict | None = None, + ) -> dict: + """Return the expected model_dump() of self.msg.""" + return { + "name": "TestAgent", + "role": "assistant", + "id": _REPLY_ID, + "metadata": {}, + "created_at": _created_at, + "finished_at": finished_at, + "content": content, + "usage": usage, + } + + # ================================================================ + # Stage 1 – Text block streaming + # ================================================================ + ev_text_start = TextBlockStartEvent( + reply_id=_REPLY_ID, + block_id=_B_TEXT, + ) + gt_text_start = _base([_tb(_B_TEXT, "")]) + + ev_text_delta1 = TextBlockDeltaEvent( + reply_id=_REPLY_ID, + block_id=_B_TEXT, + delta="Hello", + ) + gt_text_delta1 = _base([_tb(_B_TEXT, "Hello")]) + + ev_text_delta2 = TextBlockDeltaEvent( + reply_id=_REPLY_ID, + block_id=_B_TEXT, + delta=" World", + ) + gt_text_delta2 = _base([_tb(_B_TEXT, "Hello World")]) + + ev_text_end = TextBlockEndEvent(reply_id=_REPLY_ID, block_id=_B_TEXT) + gt_text_end = _base([_tb(_B_TEXT, "Hello World")]) # unchanged + + # ================================================================ + # Stage 2 – Thinking block streaming + # ================================================================ + ev_think_start = ThinkingBlockStartEvent( + reply_id=_REPLY_ID, + block_id=_B_THINK, + ) + gt_think_start = _base( + [ + _tb(_B_TEXT, "Hello World"), + _thb(_B_THINK, ""), + ], + ) + + ev_think_delta1 = ThinkingBlockDeltaEvent( + reply_id=_REPLY_ID, + block_id=_B_THINK, + delta="Let me", + ) + gt_think_delta1 = _base( + [ + _tb(_B_TEXT, "Hello World"), + _thb(_B_THINK, "Let me"), + ], + ) + + ev_think_delta2 = ThinkingBlockDeltaEvent( + reply_id=_REPLY_ID, + block_id=_B_THINK, + delta=" think", + ) + gt_think_delta2 = _base( + [ + _tb(_B_TEXT, "Hello World"), + _thb(_B_THINK, "Let me think"), + ], + ) + + ev_think_end = ThinkingBlockEndEvent( + reply_id=_REPLY_ID, + block_id=_B_THINK, + ) + gt_think_end = _base( + [ # unchanged + _tb(_B_TEXT, "Hello World"), + _thb(_B_THINK, "Let me think"), + ], + ) + + # ================================================================ + # Stage 3 – Data block streaming (base-64) + # ================================================================ + ev_data_start = DataBlockStartEvent( + reply_id=_REPLY_ID, + block_id=_B_DATA, + media_type="image/png", + ) + gt_data_start = _base( + [ + _tb(_B_TEXT, "Hello World"), + _thb(_B_THINK, "Let me think"), + _db_b64(_B_DATA, "", "image/png"), + ], + ) + + # Each delta carries an independently base64-encoded chunk (with its + # own padding); ``append_event`` decodes -> concats bytes -> re-encodes + # so the message reflects the concatenated underlying bytes, not a + # string-concat of the chunks. Here: b"abc" + b"def" -> "YWJjZGVm". + ev_data_delta1 = DataBlockDeltaEvent( + reply_id=_REPLY_ID, + block_id=_B_DATA, + data="YWJj", # base64(b"abc") + media_type="image/png", + ) + gt_data_delta1 = _base( + [ + _tb(_B_TEXT, "Hello World"), + _thb(_B_THINK, "Let me think"), + _db_b64(_B_DATA, "YWJj", "image/png"), + ], + ) + + ev_data_delta2 = DataBlockDeltaEvent( + reply_id=_REPLY_ID, + block_id=_B_DATA, + data="ZGVm", # base64(b"def") + media_type="image/png", + ) + gt_data_delta2 = _base( + [ + _tb(_B_TEXT, "Hello World"), + _thb(_B_THINK, "Let me think"), + _db_b64(_B_DATA, "YWJjZGVm", "image/png"), + ], + ) + + ev_data_end = DataBlockEndEvent(reply_id=_REPLY_ID, block_id=_B_DATA) + gt_data_end = _base( + [ # unchanged + _tb(_B_TEXT, "Hello World"), + _thb(_B_THINK, "Let me think"), + _db_b64(_B_DATA, "YWJjZGVm", "image/png"), + ], + ) + + # ================================================================ + # Stage 4 – ToolCall (TC_ALLOW): stream → confirm → allowed + # + text tool-result (SUCCESS) + # ================================================================ + ev_tc_allow_start = ToolCallStartEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_ALLOW, + tool_call_name="search", + ) + _s4_prefix = [ + _tb(_B_TEXT, "Hello World"), + _thb(_B_THINK, "Let me think"), + _db_b64(_B_DATA, "YWJjZGVm", "image/png"), + ] + gt_tc_allow_start = _base( + _s4_prefix + [_tcb(_TC_ALLOW, "search", "", "pending")], + ) + + ev_tc_allow_delta1 = ToolCallDeltaEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_ALLOW, + delta='{"q"', + ) + gt_tc_allow_delta1 = _base( + _s4_prefix + [_tcb(_TC_ALLOW, "search", '{"q"', "pending")], + ) + + ev_tc_allow_delta2 = ToolCallDeltaEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_ALLOW, + delta=': "hi"}', + ) + gt_tc_allow_delta2 = _base( + _s4_prefix + [_tcb(_TC_ALLOW, "search", '{"q": "hi"}', "pending")], + ) + + ev_tc_allow_end = ToolCallEndEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_ALLOW, + ) + gt_tc_allow_end = _base( # unchanged + _s4_prefix + [_tcb(_TC_ALLOW, "search", '{"q": "hi"}', "pending")], + ) + + # RequireUserConfirmEvent → state: pending → asking + _tc_allow_block = ToolCallBlock( + id=_TC_ALLOW, + name="search", + input='{"q": "hi"}', + ) + ev_require_confirm = RequireUserConfirmEvent( + reply_id=_REPLY_ID, + tool_calls=[_tc_allow_block], + ) + gt_require_confirm = _base( + _s4_prefix + [_tcb(_TC_ALLOW, "search", '{"q": "hi"}', "asking")], + ) + + # UserConfirmResultEvent (confirmed=True) → state: asking → allowed + ev_user_confirmed = UserConfirmResultEvent( + reply_id=_REPLY_ID, + confirm_results=[ + ConfirmResult(confirmed=True, tool_call=_tc_allow_block), + ], + ) + gt_user_confirmed = _base( + _s4_prefix + [_tcb(_TC_ALLOW, "search", '{"q": "hi"}', "allowed")], + ) + + # ToolResult for _TC_ALLOW – text output + ev_result_start = ToolResultStartEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_ALLOW, + tool_call_name="search", + ) + _s4b_prefix = _s4_prefix + [ + _tcb(_TC_ALLOW, "search", '{"q": "hi"}', "allowed"), + ] + gt_result_start = _base( + _s4b_prefix + [_trb(_TC_ALLOW, "search", [], "running")], + ) + + # First text-delta creates a new TextBlock with auto-generated ID. + ev_result_text1 = ToolResultTextDeltaEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_ALLOW, + delta="Found:", + ) + gt_result_text1 = _base( + _s4b_prefix + + [ + _trb( + _TC_ALLOW, + "search", + [{"type": "text", "id": AnyString(), "text": "Found:"}], + "running", + ), + ], + ) + + # Second text-delta appends to the SAME TextBlock. + ev_result_text2 = ToolResultTextDeltaEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_ALLOW, + delta=" 3 items", + ) + gt_result_text2 = _base( + _s4b_prefix + + [ + _trb( + _TC_ALLOW, + "search", + [ + { + "type": "text", + "id": AnyString(), + "text": "Found: 3 items", + }, + ], + "running", + ), + ], + ) + + ev_result_end_ok = ToolResultEndEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_ALLOW, + state=ToolResultState.SUCCESS, + ) + # TOOL_RESULT_END flips the paired ToolCallBlock to FINISHED, so the + # tool_call state in the prefix changes from "allowed" to "finished" + # from this point onward. + _s4b_done_prefix = _s4_prefix + [ + _tcb(_TC_ALLOW, "search", '{"q": "hi"}', "finished"), + ] + gt_result_end_ok = _base( + _s4b_done_prefix + + [ + _trb( + _TC_ALLOW, + "search", + [ + { + "type": "text", + "id": AnyString(), + "text": "Found: 3 items", + }, + ], + "success", + ), + ], + ) + + # ================================================================ + # Stage 5 – ToolCall (TC_DENY): stream → confirm → denied (finished) + # ================================================================ + _s5_prefix = _s4b_done_prefix + [ + _trb( + _TC_ALLOW, + "search", + [ + { + "type": "text", + "id": AnyString(), + "text": "Found: 3 items", + }, + ], + "success", + ), + ] + + ev_tc_deny_start = ToolCallStartEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_DENY, + tool_call_name="delete", + ) + gt_tc_deny_start = _base( + _s5_prefix + [_tcb(_TC_DENY, "delete", "", "pending")], + ) + + ev_tc_deny_end = ToolCallEndEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_DENY, + ) + gt_tc_deny_end = _base( # unchanged + _s5_prefix + [_tcb(_TC_DENY, "delete", "", "pending")], + ) + + _tc_deny_block = ToolCallBlock(id=_TC_DENY, name="delete", input="") + ev_require_confirm_deny = RequireUserConfirmEvent( + reply_id=_REPLY_ID, + tool_calls=[_tc_deny_block], + ) + gt_require_confirm_deny = _base( + _s5_prefix + [_tcb(_TC_DENY, "delete", "", "asking")], + ) + + # UserConfirmResultEvent (confirmed=False) → state: asking → finished + ev_user_denied = UserConfirmResultEvent( + reply_id=_REPLY_ID, + confirm_results=[ + ConfirmResult(confirmed=False, tool_call=_tc_deny_block), + ], + ) + gt_user_denied = _base( + _s5_prefix + [_tcb(_TC_DENY, "delete", "", "finished")], + ) + + # ================================================================ + # Stage 6 – ToolCall (TC_EXT): external execution flow + # ================================================================ + _s6_prefix = _s5_prefix + [_tcb(_TC_DENY, "delete", "", "finished")] + + ev_tc_ext_start = ToolCallStartEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_EXT, + tool_call_name="run_code", + ) + gt_tc_ext_start = _base( + _s6_prefix + [_tcb(_TC_EXT, "run_code", "", "pending")], + ) + + ev_tc_ext_end = ToolCallEndEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_EXT, + ) + gt_tc_ext_end = _base( # unchanged + _s6_prefix + [_tcb(_TC_EXT, "run_code", "", "pending")], + ) + + _tc_ext_block = ToolCallBlock(id=_TC_EXT, name="run_code", input="") + ev_require_ext = RequireExternalExecutionEvent( + reply_id=_REPLY_ID, + tool_calls=[_tc_ext_block], + ) + gt_require_ext = _base( + _s6_prefix + [_tcb(_TC_EXT, "run_code", "", "submitted")], + ) + + # ExternalExecutionResultEvent – appends a ToolResultBlock directly. + _ext_result_block = ToolResultBlock( + id=_TC_EXT, + name="run_code", + output="output: hello", + state=ToolResultState.SUCCESS, + ) + ev_ext_result = ExternalExecutionResultEvent( + reply_id=_REPLY_ID, + execution_results=[_ext_result_block], + ) + _s6b_prefix = _s6_prefix + [_tcb(_TC_EXT, "run_code", "", "submitted")] + gt_ext_result = _base( + _s6b_prefix + + [_trb(_TC_EXT, "run_code", "output: hello", "success")], + ) + + # ================================================================ + # Stage 7 – ToolResult with data output: base-64 + URL blocks + # ================================================================ + _s7_prefix = _s6b_prefix + [ + _trb(_TC_EXT, "run_code", "output: hello", "success"), + ] + + ev_tc_img_start = ToolCallStartEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_IMG, + tool_call_name="screenshot", + ) + gt_tc_img_start = _base( + _s7_prefix + [_tcb(_TC_IMG, "screenshot", "", "pending")], + ) + + ev_tc_img_end = ToolCallEndEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_IMG, + ) + gt_tc_img_end = _base( # unchanged + _s7_prefix + [_tcb(_TC_IMG, "screenshot", "", "pending")], + ) + + ev_res_img_start = ToolResultStartEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_IMG, + tool_call_name="screenshot", + ) + _s7b_prefix = _s7_prefix + [_tcb(_TC_IMG, "screenshot", "", "pending")] + gt_res_img_start = _base( + _s7b_prefix + [_trb(_TC_IMG, "screenshot", [], "running")], + ) + + # Base-64 data delta → DataBlock(base64) appended to output + ev_res_img_b64 = ToolResultDataDeltaEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_IMG, + block_id=_RES_DATA_B, + media_type="image/png", + data="iVBOR==", + ) + gt_res_img_b64 = _base( + _s7b_prefix + + [ + _trb( + _TC_IMG, + "screenshot", + [_db_b64(_RES_DATA_B, "iVBOR==", "image/png")], + "running", + ), + ], + ) + + # URL data delta → DataBlock(url) appended to output + ev_res_img_url = ToolResultDataDeltaEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_IMG, + block_id=_RES_URL_B, + media_type="image/jpeg", + url="https://example.com/img.jpg", + ) + gt_res_img_url = _base( + _s7b_prefix + + [ + _trb( + _TC_IMG, + "screenshot", + [ + _db_b64(_RES_DATA_B, "iVBOR==", "image/png"), + _db_url( + _RES_URL_B, + "https://example.com/img.jpg", + "image/jpeg", + ), + ], + "running", + ), + ], + ) + + ev_res_img_end = ToolResultEndEvent( + reply_id=_REPLY_ID, + tool_call_id=_TC_IMG, + state=ToolResultState.ERROR, + ) + # TOOL_RESULT_END flips the paired ToolCallBlock to FINISHED. + _s7b_done_prefix = _s7_prefix + [ + _tcb(_TC_IMG, "screenshot", "", "finished"), + ] + gt_res_img_end = _base( + _s7b_done_prefix + + [ + _trb( + _TC_IMG, + "screenshot", + [ + _db_b64(_RES_DATA_B, "iVBOR==", "image/png"), + _db_url( + _RES_URL_B, + "https://example.com/img.jpg", + "image/jpeg", + ), + ], + "error", + ), + ], + ) + + # ================================================================ + # Stage 8 – ModelCallEndEvent (first call: usage initialized; + # second call: usage accumulated) + # ================================================================ + _final_content = _s7b_done_prefix + [ + _trb( + _TC_IMG, + "screenshot", + [ + _db_b64(_RES_DATA_B, "iVBOR==", "image/png"), + _db_url( + _RES_URL_B, + "https://example.com/img.jpg", + "image/jpeg", + ), + ], + "error", + ), + ] + ev_model_call_end_1 = ModelCallEndEvent( + reply_id=_REPLY_ID, + input_tokens=10, + output_tokens=20, + ) + gt_model_call_end_1 = _base( + _final_content, + usage={"input_tokens": 10, "output_tokens": 20}, + ) + + ev_model_call_end_2 = ModelCallEndEvent( + reply_id=_REPLY_ID, + input_tokens=5, + output_tokens=8, + ) + gt_model_call_end_2 = _base( + _final_content, + usage={"input_tokens": 15, "output_tokens": 28}, + ) + + # ================================================================ + # Stage 9 – ReplyEndEvent + # ================================================================ + ev_reply_end = ReplyEndEvent( + reply_id=_REPLY_ID, + session_id=_SESSION_ID, + created_at=_FIXED_END_TS, + ) + gt_reply_end = _base( + _final_content, + finished_at=_FIXED_END_TS, + usage={"input_tokens": 15, "output_tokens": 28}, + ) + + # ================================================================ + # Assemble the two parallel lists + # ================================================================ + self.events = [ + # Stage 1: Text + ev_text_start, + ev_text_delta1, + ev_text_delta2, + ev_text_end, + # Stage 2: Thinking + ev_think_start, + ev_think_delta1, + ev_think_delta2, + ev_think_end, + # Stage 3: Data (base-64) + ev_data_start, + ev_data_delta1, + ev_data_delta2, + ev_data_end, + # Stage 4: ToolCall → confirm (allowed) + text result (success) + ev_tc_allow_start, + ev_tc_allow_delta1, + ev_tc_allow_delta2, + ev_tc_allow_end, + ev_require_confirm, + ev_user_confirmed, + ev_result_start, + ev_result_text1, + ev_result_text2, + ev_result_end_ok, + # Stage 5: ToolCall → confirm (denied) + ev_tc_deny_start, + ev_tc_deny_end, + ev_require_confirm_deny, + ev_user_denied, + # Stage 6: ToolCall → external execution + ev_tc_ext_start, + ev_tc_ext_end, + ev_require_ext, + ev_ext_result, + # Stage 7: ToolResult with data output (base-64 + URL) + ev_tc_img_start, + ev_tc_img_end, + ev_res_img_start, + ev_res_img_b64, + ev_res_img_url, + ev_res_img_end, + # Stage 8: MODEL_CALL_END (init + accumulate) + ev_model_call_end_1, + ev_model_call_end_2, + # Stage 9: REPLY_END + ev_reply_end, + ] + self.ground_truths = [ + # Stage 1 + gt_text_start, + gt_text_delta1, + gt_text_delta2, + gt_text_end, + # Stage 2 + gt_think_start, + gt_think_delta1, + gt_think_delta2, + gt_think_end, + # Stage 3 + gt_data_start, + gt_data_delta1, + gt_data_delta2, + gt_data_end, + # Stage 4 + gt_tc_allow_start, + gt_tc_allow_delta1, + gt_tc_allow_delta2, + gt_tc_allow_end, + gt_require_confirm, + gt_user_confirmed, + gt_result_start, + gt_result_text1, + gt_result_text2, + gt_result_end_ok, + # Stage 5 + gt_tc_deny_start, + gt_tc_deny_end, + gt_require_confirm_deny, + gt_user_denied, + # Stage 6 + gt_tc_ext_start, + gt_tc_ext_end, + gt_require_ext, + gt_ext_result, + # Stage 7 + gt_tc_img_start, + gt_tc_img_end, + gt_res_img_start, + gt_res_img_b64, + gt_res_img_url, + gt_res_img_end, + # Stage 8 + gt_model_call_end_1, + gt_model_call_end_2, + # Stage 9 + gt_reply_end, + ] + + # ------------------------------------------------------------------ + # Tests + # ------------------------------------------------------------------ + + async def test_append_event_stream(self) -> None: + """Apply every event in order and assert the full Msg state after each. + + Uses a zip loop over self.events / self.ground_truths so that the + event index is clear from the assertion failure message. + """ + self.assertEqual( + len(self.events), + len(self.ground_truths), + "events and ground_truths must have equal length", + ) + for idx, (event, expected) in enumerate( + zip(self.events, self.ground_truths), + ): + self.msg.append_event(event) + self.assertDictEqual( + self.msg.model_dump(), + expected, + msg=f"Mismatch after event[{idx}] ({event.type})", + ) + + async def test_wrong_reply_id_is_skipped(self) -> None: + """An event whose reply_id does not match msg.id must be ignored.""" + original_dump = self.msg.model_dump() + wrong_event = TextBlockStartEvent( + reply_id="totally_wrong_id", + block_id="should_not_appear", + ) + self.msg.append_event(wrong_event) + self.assertDictEqual( + self.msg.model_dump(), + original_dump, + msg="Msg must not change when event.reply_id does not match", + ) + + async def test_missing_block_does_not_crash(self) -> None: + """Sending a delta for a non-existent block must log a warning only.""" + original_dump = self.msg.model_dump() + + # Delta events for blocks that were never started + ghost_events = [ + TextBlockDeltaEvent( + reply_id=_REPLY_ID, + block_id="ghost_text", + delta="x", + ), + ThinkingBlockDeltaEvent( + reply_id=_REPLY_ID, + block_id="ghost_think", + delta="x", + ), + DataBlockDeltaEvent( + reply_id=_REPLY_ID, + block_id="ghost_data", + data="x", + media_type="image/png", + ), + ToolCallDeltaEvent( + reply_id=_REPLY_ID, + tool_call_id="ghost_tc", + delta="x", + ), + ToolResultTextDeltaEvent( + reply_id=_REPLY_ID, + tool_call_id="ghost_tr", + delta="x", + ), + ToolResultDataDeltaEvent( + reply_id=_REPLY_ID, + tool_call_id="ghost_tr", + media_type="image/png", + data="x", + ), + ToolResultEndEvent( + reply_id=_REPLY_ID, + tool_call_id="ghost_tr", + state=ToolResultState.SUCCESS, + ), + ] + for ev in ghost_events: + self.msg.append_event(ev) # must not raise + + self.assertDictEqual( + self.msg.model_dump(), + original_dump, + msg="Msg must not change when delta targets a missing block", + ) + + async def asyncTearDown(self) -> None: + """No teardown needed.""" diff --git a/tests/formatter_anthropic_test.py b/tests/formatter_anthropic_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ccd6b0c85487522af0cc93c4a600d84665e3c767 --- /dev/null +++ b/tests/formatter_anthropic_test.py @@ -0,0 +1,958 @@ +# -*- coding: utf-8 -*- +"""Comprehensive formatter unit tests for AnthropicChatFormatter and +AnthropicMultiAgentFormatter, following the reference test style with exact +ground-truth comparisons. +""" +from unittest import IsolatedAsyncioTestCase + +from agentscope.formatter import ( + AnthropicChatFormatter, + AnthropicMultiAgentFormatter, +) +from agentscope.message import ( + UserMsg, + AssistantMsg, + SystemMsg, + TextBlock, + DataBlock, + ToolCallBlock, + ToolResultBlock, + ToolResultState, + Base64Source, + ThinkingBlock, + HintBlock, +) + + +class TestAnthropicFormatter(IsolatedAsyncioTestCase): + """Comprehensive tests for Anthropic Chat and MultiAgent formatters.""" + + async def asyncSetUp(self) -> None: + """Set up shared message fixtures and expected ground-truth dicts.""" + self.image_b64 = "ZmFrZSBpbWFnZSBkYXRh" + + # --------------------------------------------------------------- + # Message fixtures + # (No URL images: Anthropic URL handling downloads from the network) + # --------------------------------------------------------------- + self.msgs_system = [ + SystemMsg( + name="system", + content="You're a helpful assistant.", + ), + ] + + self.msgs_conversation = [ + UserMsg( + name="user", + content="What is the capital of France?", + ), + AssistantMsg( + name="assistant", + content="The capital of France is Paris.", + ), + UserMsg( + name="user", + content="What is the capital of Germany?", + ), + AssistantMsg( + name="assistant", + content="The capital of Germany is Berlin.", + ), + UserMsg( + name="user", + content="What is the capital of Japan?", + ), + ] + + self.msgs_tools = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_1", + name="get_capital", + input='{"country": "Japan"}', + ), + ToolResultBlock( + id="call_1", + name="get_capital", + output=[ + TextBlock(text="The capital of Japan is Tokyo."), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="The capital of Japan is Tokyo."), + ], + ), + ] + + # --------------------------------------------------------------- + # Ground truth: AnthropicChatFormatter + # - No "name" field. + # - Content is always a list of {"type": ..., ...} dicts. + # - ToolResultBlock forces role to "user". + # - ToolCallBlock "input" is a dict (parsed from JSON string). + # --------------------------------------------------------------- + self.gt_chat = [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You're a helpful assistant."}, + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What is the capital of France?", + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of France is Paris.", + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What is the capital of Germany?", + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of Germany is Berlin.", + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What is the capital of Japan?", + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "call_1", + "name": "get_capital", + "input": {"country": "Japan"}, + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_1", + "content": [ + { + "type": "text", + "text": "The capital of Japan is Tokyo.", + }, + ], + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of Japan is Tokyo.", + }, + ], + }, + ] + + # --------------------------------------------------------------- + # Ground truth: AnthropicMultiAgentFormatter + # - System: {"role": "system", "content": [{"type": "text", ...}]} + # - Agent messages (is_first=True): wrapped in hist_prompt + + # .... + # - Agent messages (is_first=False): no wrapping at all. + # --------------------------------------------------------------- + _hist_prompt = ( + AnthropicMultiAgentFormatter().conversation_history_prompt + ) + + _conv_text = ( + "user: What is the capital of France?\n" + "assistant: The capital of France is Paris.\n" + "user: What is the capital of Germany?\n" + "assistant: The capital of Germany is Berlin.\n" + "user: What is the capital of Japan?" + ) + + self._gt_trailing_asst = { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of Japan is Tokyo.", + }, + ], + } + + self._gt_tool_call = { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "call_1", + "name": "get_capital", + "input": {"country": "Japan"}, + }, + ], + } + self._gt_tool_result = { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_1", + "content": [ + { + "type": "text", + "text": "The capital of Japan is Tokyo.", + }, + ], + }, + ], + } + + self.gt_multiagent = [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You're a helpful assistant."}, + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + _hist_prompt + + "\n" + + _conv_text + + "\n" + ), + }, + ], + }, + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ] + + # ------------------------------------------------------------------- + # AnthropicChatFormatter tests + # ------------------------------------------------------------------- + + async def test_chat_formatter(self) -> None: + """Chat formatter produces exact output for various subsets.""" + fmt = AnthropicChatFormatter() + + # Full history + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_chat, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_chat[1:], res) + + # Without conversation + n_tools_gt = len(self.gt_chat) - 1 - len(self.msgs_conversation) + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [self.gt_chat[0]] + self.gt_chat[-n_tools_gt:], + res, + ) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_chat[:-n_tools_gt], res) + + # Empty + res = await fmt.format([]) + self.assertListEqual([], res) + + async def test_chat_formatter_base64_image(self) -> None: + """Base64-encoded image is formatted as Anthropic image source.""" + fmt = AnthropicChatFormatter() + msgs = [ + UserMsg( + name="user", + content=[ + TextBlock(text="What's in this image?"), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": self.image_b64, + }, + }, + ], + }, + ], + res, + ) + + async def test_chat_formatter_thinking_preserved(self) -> None: + """ThinkingBlock with a signature is passed back as a thinking + content block.""" + fmt = AnthropicChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock( + thinking="inner thoughts", + signature="sig_abc", + ), + TextBlock(text="reply"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "inner thoughts", + "signature": "sig_abc", + }, + {"type": "text", "text": "reply"}, + ], + }, + ], + res, + ) + + async def test_chat_formatter_thinking_without_signature_dropped( + self, + ) -> None: + """ThinkingBlock from another provider (no signature) is dropped + rather than forwarded with an empty signature, which Anthropic + rejects with `Invalid signature in thinking block`. An empty-string + signature is treated the same as a missing one.""" + fmt = AnthropicChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="from another provider"), + ThinkingBlock(thinking="empty sig", signature=""), + TextBlock(text="reply"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "reply"}, + ], + }, + ], + res, + ) + + # Thinking-only message produces no output (no empty assistant msg). + res = await fmt.format( + [ + AssistantMsg( + name="assistant", + content=[ThinkingBlock(thinking="only thinking")], + ), + ], + ) + self.assertListEqual([], res) + + async def test_chat_formatter_tool_result_role_forced_to_user( + self, + ) -> None: + """Anthropic forces tool_result messages to role='user'.""" + fmt = AnthropicChatFormatter() + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + tool_result_roles = [ + m["role"] + for m in res + if any( + b.get("type") == "tool_result" + for b in (m.get("content") or []) + ) + ] + self.assertListEqual(tool_result_roles, ["user"]) + + async def test_chat_formatter_tool_result_with_image(self) -> None: + """Tool result containing an image DataBlock inlines the image in the + tool_result content without crashing on TextBlock system-reminders.""" + fmt = AnthropicChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_img", + name="get_chart", + input="{}", + ), + ToolResultBlock( + id="call_img", + name="get_chart", + output=[ + TextBlock(text="Here is the chart."), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="Here is the chart analysis."), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "call_img", + "name": "get_chart", + "input": {}, + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_img", + "content": [ + { + "type": "text", + "text": "Here is the chart.", + }, + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": self.image_b64, + }, + }, + ], + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Here is the chart analysis.", + }, + ], + }, + ], + res, + ) + + # ------------------------------------------------------------------- + # AnthropicMultiAgentFormatter tests + # ------------------------------------------------------------------- + + async def test_multiagent_formatter(self) -> None: + """MultiAgent formatter produces exact output for various subsets.""" + fmt = AnthropicMultiAgentFormatter() + + # Full history + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_multiagent, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_multiagent[1:], res) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_multiagent[:2], res) + + # System only + res = await fmt.format(self.msgs_system) + self.assertListEqual([self.gt_multiagent[0]], res) + + # Conversation only + res = await fmt.format(self.msgs_conversation) + self.assertListEqual([self.gt_multiagent[1]], res) + + # Tools only + res = await fmt.format(self.msgs_tools) + self.assertListEqual( + [ + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # System + tools (no conversation) + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [ + self.gt_multiagent[0], + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # Empty + res = await fmt.format([]) + self.assertListEqual([], res) + + async def test_chat_formatter_complex_multi_step(self) -> None: + """Complex multi-step sequence with interleaved thinking, text, + tool calls, and tool results.""" + fmt = AnthropicChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="thinking_1", signature="sig_1"), + TextBlock(text="text_1"), + ToolCallBlock( + id="call_1", + name="func_1", + input='{"arg": "value1"}', + ), + ToolCallBlock( + id="call_2", + name="func_2", + input='{"arg": "value2"}', + ), + ToolResultBlock( + id="call_1", + name="func_1", + output=[TextBlock(text="result_1")], + state=ToolResultState.SUCCESS, + ), + ToolResultBlock( + id="call_2", + name="func_2", + output=[TextBlock(text="result_2")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_2", signature="sig_2"), + TextBlock(text="text_2"), + ToolCallBlock( + id="call_3", + name="func_3", + input='{"arg": "value3"}', + ), + ToolResultBlock( + id="call_3", + name="func_3", + output=[TextBlock(text="result_3")], + state=ToolResultState.SUCCESS, + ), + ToolCallBlock( + id="call_4", + name="func_4", + input='{"arg": "value4"}', + ), + ToolResultBlock( + id="call_4", + name="func_4", + output=[TextBlock(text="result_4")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_3", signature="sig_3"), + TextBlock(text="text_3"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "thinking_1", + "signature": "sig_1", + }, + {"type": "text", "text": "text_1"}, + { + "type": "tool_use", + "id": "call_1", + "name": "func_1", + "input": {"arg": "value1"}, + }, + { + "type": "tool_use", + "id": "call_2", + "name": "func_2", + "input": {"arg": "value2"}, + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_1", + "content": [ + {"type": "text", "text": "result_1"}, + ], + }, + { + "type": "tool_result", + "tool_use_id": "call_2", + "content": [ + {"type": "text", "text": "result_2"}, + ], + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "thinking_2", + "signature": "sig_2", + }, + {"type": "text", "text": "text_2"}, + { + "type": "tool_use", + "id": "call_3", + "name": "func_3", + "input": {"arg": "value3"}, + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_3", + "content": [ + {"type": "text", "text": "result_3"}, + ], + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "call_4", + "name": "func_4", + "input": {"arg": "value4"}, + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_4", + "content": [ + {"type": "text", "text": "result_4"}, + ], + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "thinking_3", + "signature": "sig_3", + }, + {"type": "text", "text": "text_3"}, + ], + }, + ], + res, + ) + + async def test_chat_formatter_hint_block(self) -> None: + """HintBlock flushes preceding content and becomes a user message.""" + fmt = AnthropicChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + TextBlock(text="Let me think about that."), + HintBlock(hint="Remember to be concise."), + TextBlock(text="Here is my answer."), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me think about that."}, + ], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "Remember to be concise."}, + ], + }, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Here is my answer."}, + ], + }, + ], + res, + ) + + async def test_chat_formatter_hint_block_multimodal(self) -> None: + """Multimodal HintBlock becomes a single user message with text + + image.""" + fmt = AnthropicChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + HintBlock( + hint=[ + TextBlock(text="Inspect this screenshot:"), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + ), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Inspect this screenshot:", + }, + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": self.image_b64, + }, + }, + ], + }, + ], + res, + ) + + async def test_chat_formatter_parallel_tool_results_merged( + self, + ) -> None: + """Parallel tool_results must be grouped into ONE user message. + + Regression test for Issue #1892: AnthropicChatFormatter was flushing + content_blocks on every ToolResultBlock, which split N parallel results + into N separate user messages. Strict endpoints such as DeepSeek's + Anthropic-compatible API reject this with HTTP 400. + """ + fmt = AnthropicChatFormatter() + msgs = [ + UserMsg( + name="user", + content="Get weather for Beijing, Shanghai, and Guangzhou.", + ), + AssistantMsg( + name="assistant", + content=[ + TextBlock(text="Let me check all three cities."), + ToolCallBlock( + id="call_01", + name="get_weather", + input='{"city": "Beijing"}', + ), + ToolCallBlock( + id="call_02", + name="get_weather", + input='{"city": "Shanghai"}', + ), + ToolCallBlock( + id="call_03", + name="get_weather", + input='{"city": "Guangzhou"}', + ), + ToolResultBlock( + id="call_01", + name="get_weather", + output="Sunny, 28\u00b0C", + state=ToolResultState.SUCCESS, + ), + ToolResultBlock( + id="call_02", + name="get_weather", + output="Cloudy, 24\u00b0C", + state=ToolResultState.SUCCESS, + ), + ToolResultBlock( + id="call_03", + name="get_weather", + output="Rainy, 22\u00b0C", + state=ToolResultState.SUCCESS, + ), + ], + ), + ] + res = await fmt.format(msgs) + + # Expected: 3 messages total + # [0] user -> original question + # [1] assistant -> text + 3x tool_use + # [2] user -> ALL 3 tool_results in ONE message (the fix) + self.assertListEqual( + [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Get weather for Beijing, Shanghai, " + "and Guangzhou.", + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Let me check all three cities.", + }, + { + "type": "tool_use", + "id": "call_01", + "name": "get_weather", + "input": {"city": "Beijing"}, + }, + { + "type": "tool_use", + "id": "call_02", + "name": "get_weather", + "input": {"city": "Shanghai"}, + }, + { + "type": "tool_use", + "id": "call_03", + "name": "get_weather", + "input": {"city": "Guangzhou"}, + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_01", + "content": [ + {"type": "text", "text": "Sunny, 28\u00b0C"}, + ], + }, + { + "type": "tool_result", + "tool_use_id": "call_02", + "content": [ + {"type": "text", "text": "Cloudy, 24\u00b0C"}, + ], + }, + { + "type": "tool_result", + "tool_use_id": "call_03", + "content": [ + {"type": "text", "text": "Rainy, 22\u00b0C"}, + ], + }, + ], + }, + ], + res, + ) diff --git a/tests/formatter_dashscope_test.py b/tests/formatter_dashscope_test.py new file mode 100644 index 0000000000000000000000000000000000000000..3193a621d732e2b27125a829efb608325183ae73 --- /dev/null +++ b/tests/formatter_dashscope_test.py @@ -0,0 +1,1035 @@ +# -*- coding: utf-8 -*- +"""Comprehensive formatter unit tests for DashScopeChatFormatter and +DashScopeMultiAgentFormatter, following the reference test style with exact +ground-truth comparisons. +""" +from unittest import IsolatedAsyncioTestCase +from unittest.mock import patch + +from agentscope.formatter import ( + DashScopeChatFormatter, + DashScopeMultiAgentFormatter, +) +from agentscope.message import ( + UserMsg, + AssistantMsg, + SystemMsg, + TextBlock, + DataBlock, + ToolCallBlock, + ToolResultBlock, + ToolResultState, + Base64Source, + URLSource, + ThinkingBlock, + HintBlock, +) + + +# A fixed short-uuid used to make promote-to-multimodal tests deterministic. +_FIXED_ID = "TESTID1234567" + + +class TestDashScopeFormatter(IsolatedAsyncioTestCase): + """Comprehensive tests for DashScope Chat and MultiAgent formatters.""" + + async def asyncSetUp(self) -> None: + """Set up shared message fixtures and expected ground-truth dicts.""" + # --- URL strings --- + # Normalise through URLSource so that pydantic URL normalisation is + # applied consistently to both input and expected values. + _img_src = URLSource( + url="https://example.com/image.png", + media_type="image/png", + ) + _aud_src = URLSource( + url="https://example.com/audio.mp3", + media_type="audio/mpeg", + ) + self.image_url = str(_img_src.url) + self.audio_url = str(_aud_src.url) + + # --- Base64 fixtures --- + self.image_b64 = "ZmFrZSBpbWFnZSBkYXRh" + self.image_data_uri = f"data:image/png;base64,{self.image_b64}" + + # --------------------------------------------------------------- + # Message fixtures + # --------------------------------------------------------------- + self.msgs_system = [ + SystemMsg( + name="system", + content="You're a helpful assistant.", + ), + ] + + self.msgs_conversation = [ + UserMsg( + name="user", + content=[ + TextBlock(text="What is the capital of France?"), + DataBlock( + source=URLSource( + url=self.image_url, + media_type="image/png", + ), + ), + ], + ), + AssistantMsg( + name="assistant", + content="The capital of France is Paris.", + ), + UserMsg( + name="user", + content=[ + TextBlock(text="What is the capital of Germany?"), + DataBlock( + source=URLSource( + url=self.audio_url, + media_type="audio/mpeg", + ), + ), + ], + ), + AssistantMsg( + name="assistant", + content="The capital of Germany is Berlin.", + ), + UserMsg( + name="user", + content="What is the capital of Japan?", + ), + ] + + # Messages with ToolResultBlock must use role="assistant" because the + # system-role validator rejects non-text blocks. + self.msgs_tools = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_1", + name="get_capital", + input='{"country": "Japan"}', + ), + ToolResultBlock( + id="call_1", + name="get_capital", + output=[ + TextBlock(text="The capital of Japan is Tokyo."), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="The capital of Japan is Tokyo."), + ], + ), + ] + + # --------------------------------------------------------------- + # Ground truth: DashScopeChatFormatter (OpenAI-compatible format) + # - Content blocks use {"type": "text", "text": ...} format. + # - Images use {"type": "image_url", "image_url": {"url": ...}}. + # - Audio uses {"type": "input_audio", "input_audio": {...}}. + # - Tool-result content is a plain string. + # --------------------------------------------------------------- + self.gt_chat = [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You're a helpful assistant."}, + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What is the capital of France?", + }, + { + "type": "image_url", + "image_url": {"url": self.image_url}, + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of France is Paris.", + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What is the capital of Germany?", + }, + { + "type": "input_audio", + "input_audio": { + "data": self.audio_url, + "format": "mpeg", + }, + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of Germany is Berlin.", + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What is the capital of Japan?", + }, + ], + }, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_capital", + "arguments": '{"country": "Japan"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "The capital of Japan is Tokyo.", + "name": "get_capital", + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of Japan is Tokyo.", + }, + ], + }, + ] + + # --------------------------------------------------------------- + # Ground truth: DashScopeMultiAgentFormatter + # - System content is a plain string (via get_text_content()). + # - Conversation history is collapsed into a single user message. + # --------------------------------------------------------------- + _hist_prompt = ( + DashScopeMultiAgentFormatter().conversation_history_prompt + ) + + self._gt_trailing_asst = { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of Japan is Tokyo.", + }, + ], + } + + self._gt_tool_call = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_capital", + "arguments": '{"country": "Japan"}', + }, + }, + ], + } + self._gt_tool_result = { + "role": "tool", + "tool_call_id": "call_1", + "content": "The capital of Japan is Tokyo.", + "name": "get_capital", + } + + self.gt_multiagent = [ + { + "role": "system", + "content": "You're a helpful assistant.", + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + _hist_prompt + "\n" + "user: What is the capital of France?\n" + "assistant: The capital of France is Paris.\n" + "user: What is the capital of Germany?\n" + "assistant: The capital of Germany is Berlin.\n" + "user: What is the capital of Japan?\n" + "" + ), + }, + { + "type": "image_url", + "image_url": {"url": self.image_url}, + }, + { + "type": "input_audio", + "input_audio": { + "data": self.audio_url, + "format": "mpeg", + }, + }, + ], + }, + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ] + + # ------------------------------------------------------------------- + # DashScopeChatFormatter tests + # ------------------------------------------------------------------- + + async def test_chat_formatter(self) -> None: + """Chat formatter produces exact output for various subsets.""" + fmt = DashScopeChatFormatter() + + # Full history + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_chat, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_chat[1:], res) + + # Without conversation + n_tools_gt = len(self.gt_chat) - 1 - len(self.msgs_conversation) + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [self.gt_chat[0]] + self.gt_chat[-n_tools_gt:], + res, + ) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_chat[:-n_tools_gt], res) + + # Empty + res = await fmt.format([]) + self.assertListEqual([], res) + + async def test_chat_formatter_base64_image(self) -> None: + """Base64-encoded image is inlined as a data URI.""" + fmt = DashScopeChatFormatter() + msgs = [ + UserMsg( + name="user", + content=[ + TextBlock(text="What's in this image?"), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + { + "type": "image_url", + "image_url": {"url": self.image_data_uri}, + }, + ], + }, + ], + res, + ) + + @patch( + "agentscope.formatter._formatter_base.shortuuid.uuid", + return_value=_FIXED_ID, + ) + async def test_chat_formatter_url_image_in_tool_result( + self, + _mock_uuid: object, + ) -> None: + """URL images in tool results are promoted to a follow-up user message. + + The textual part of the tool result contains a system-reminder with a + unique identifier; the identifier is mocked to be deterministic. + """ + fmt = DashScopeChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_img", + name="get_map", + input='{"city": "Tokyo"}', + ), + ToolResultBlock( + id="call_img", + name="get_map", + output=[ + TextBlock(text="Here is the map."), + DataBlock( + source=URLSource( + url=self.image_url, + media_type="image/png", + ), + ), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="Here is the map of Tokyo."), + ], + ), + ] + res = await fmt.format(msgs) + + expected_tool_content = ( + "Here is the map.\n" + f"A(n) image file is returned " + f"and will be presented to you with the identifier " + f"[{_FIXED_ID}]." + ) + self.assertListEqual( + [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_img", + "type": "function", + "function": { + "name": "get_map", + "arguments": '{"city": "Tokyo"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_img", + "content": expected_tool_content, + "name": "get_map", + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + "The multimodal data " + "and their identifiers are listed as " + "follows:" + ), + }, + { + "type": "text", + "text": f"- {_FIXED_ID} (image file): ", + }, + { + "type": "image_url", + "image_url": {"url": self.image_url}, + }, + { + "type": "text", + "text": "", + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Here is the map of Tokyo.", + }, + ], + }, + ], + res, + ) + + async def test_chat_formatter_thinking_dropped_without_flag(self) -> None: + """ThinkingBlock is silently dropped when application/x-thinking is + absent from input_types.""" + fmt = DashScopeChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="inner thoughts"), + TextBlock(text="reply"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": [{"type": "text", "text": "reply"}], + }, + ], + res, + ) + + async def test_chat_formatter_thinking_becomes_reasoning_content( + self, + ) -> None: + """ThinkingBlock becomes reasoning_content when application/x-thinking + is in input_types.""" + fmt = DashScopeChatFormatter( + input_types=["text/plain", "application/x-thinking"], + ) + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="inner thoughts"), + TextBlock(text="reply"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": [{"type": "text", "text": "reply"}], + "reasoning_content": "inner thoughts", + }, + ], + res, + ) + + async def test_chat_formatter_multiple_thinking_blocks_joined( + self, + ) -> None: + """Multiple ThinkingBlocks are joined with a newline into a single + reasoning_content field.""" + fmt = DashScopeChatFormatter( + input_types=["text/plain", "application/x-thinking"], + ) + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="part one"), + ThinkingBlock(thinking="part two"), + TextBlock(text="answer"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": [{"type": "text", "text": "answer"}], + "reasoning_content": "part one\npart two", + }, + ], + res, + ) + + # ------------------------------------------------------------------- + # DashScopeMultiAgentFormatter tests + # ------------------------------------------------------------------- + + async def test_multiagent_formatter(self) -> None: + """MultiAgent formatter produces exact output for various subsets.""" + fmt = DashScopeMultiAgentFormatter() + + # Full history + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_multiagent, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_multiagent[1:], res) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_multiagent[:2], res) + + # System only + res = await fmt.format(self.msgs_system) + self.assertListEqual([self.gt_multiagent[0]], res) + + # Conversation only + res = await fmt.format(self.msgs_conversation) + self.assertListEqual([self.gt_multiagent[1]], res) + + # Tools only — no prior agent_message group, so the trailing assistant + # is formatted with is_first=True (includes the full history prompt). + res = await fmt.format(self.msgs_tools) + self.assertListEqual( + [ + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # System + tools (no conversation) + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [ + self.gt_multiagent[0], + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # Empty + res = await fmt.format([]) + self.assertListEqual([], res) + + async def test_multiagent_formatter_thinking_in_tool_sequence( + self, + ) -> None: + """ThinkingBlocks inside a tool sequence are forwarded as + reasoning_content when application/x-thinking is in input_types.""" + fmt = DashScopeMultiAgentFormatter( + input_types=["text/plain", "application/x-thinking"], + ) + tc = ToolCallBlock( + id="call_1", + name="get_capital", + input='{"country": "Japan"}', + ) + tr = ToolResultBlock( + id="call_1", + name="get_capital", + output=[TextBlock(text="Tokyo")], + state=ToolResultState.SUCCESS, + ) + msgs = [ + AssistantMsg( + name="assistant", + content=[ThinkingBlock(thinking="Need to check"), tc, tr], + ), + ] + res = await fmt.format(msgs) + asst_msgs = [m for m in res if m.get("role") == "assistant"] + self.assertListEqual( + [m["reasoning_content"] for m in asst_msgs], + ["Need to check"], + ) + + async def test_chat_formatter_complex_multi_step(self) -> None: + """Complex multi-step sequence with interleaved thinking, text, + tool calls, and tool results (thinking dropped without flag).""" + fmt = DashScopeChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="thinking_1"), + TextBlock(text="text_1"), + ToolCallBlock( + id="call_1", + name="func_1", + input='{"arg": "value1"}', + ), + ToolCallBlock( + id="call_2", + name="func_2", + input='{"arg": "value2"}', + ), + ToolResultBlock( + id="call_1", + name="func_1", + output=[TextBlock(text="result_1")], + state=ToolResultState.SUCCESS, + ), + ToolResultBlock( + id="call_2", + name="func_2", + output=[TextBlock(text="result_2")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_2"), + TextBlock(text="text_2"), + ToolCallBlock( + id="call_3", + name="func_3", + input='{"arg": "value3"}', + ), + ToolResultBlock( + id="call_3", + name="func_3", + output=[TextBlock(text="result_3")], + state=ToolResultState.SUCCESS, + ), + ToolCallBlock( + id="call_4", + name="func_4", + input='{"arg": "value4"}', + ), + ToolResultBlock( + id="call_4", + name="func_4", + output=[TextBlock(text="result_4")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_3"), + TextBlock(text="text_3"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": [{"type": "text", "text": "text_1"}], + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "func_1", + "arguments": '{"arg": "value1"}', + }, + }, + { + "id": "call_2", + "type": "function", + "function": { + "name": "func_2", + "arguments": '{"arg": "value2"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "result_1", + "name": "func_1", + }, + { + "role": "tool", + "tool_call_id": "call_2", + "content": "result_2", + "name": "func_2", + }, + { + "role": "assistant", + "content": [{"type": "text", "text": "text_2"}], + "tool_calls": [ + { + "id": "call_3", + "type": "function", + "function": { + "name": "func_3", + "arguments": '{"arg": "value3"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_3", + "content": "result_3", + "name": "func_3", + }, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_4", + "type": "function", + "function": { + "name": "func_4", + "arguments": '{"arg": "value4"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_4", + "content": "result_4", + "name": "func_4", + }, + { + "role": "assistant", + "content": [{"type": "text", "text": "text_3"}], + }, + ], + res, + ) + + async def test_chat_formatter_complex_multi_step_with_thinking( + self, + ) -> None: + """Complex multi-step sequence with thinking preserved via + application/x-thinking flag.""" + fmt = DashScopeChatFormatter( + input_types=["text/plain", "application/x-thinking"], + ) + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="thinking_1"), + TextBlock(text="text_1"), + ToolCallBlock( + id="call_1", + name="func_1", + input='{"arg": "value1"}', + ), + ToolCallBlock( + id="call_2", + name="func_2", + input='{"arg": "value2"}', + ), + ToolResultBlock( + id="call_1", + name="func_1", + output=[TextBlock(text="result_1")], + state=ToolResultState.SUCCESS, + ), + ToolResultBlock( + id="call_2", + name="func_2", + output=[TextBlock(text="result_2")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_2"), + TextBlock(text="text_2"), + ToolCallBlock( + id="call_3", + name="func_3", + input='{"arg": "value3"}', + ), + ToolResultBlock( + id="call_3", + name="func_3", + output=[TextBlock(text="result_3")], + state=ToolResultState.SUCCESS, + ), + ToolCallBlock( + id="call_4", + name="func_4", + input='{"arg": "value4"}', + ), + ToolResultBlock( + id="call_4", + name="func_4", + output=[TextBlock(text="result_4")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_3"), + TextBlock(text="text_3"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": [{"type": "text", "text": "text_1"}], + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "func_1", + "arguments": '{"arg": "value1"}', + }, + }, + { + "id": "call_2", + "type": "function", + "function": { + "name": "func_2", + "arguments": '{"arg": "value2"}', + }, + }, + ], + "reasoning_content": "thinking_1", + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "result_1", + "name": "func_1", + }, + { + "role": "tool", + "tool_call_id": "call_2", + "content": "result_2", + "name": "func_2", + }, + { + "role": "assistant", + "content": [{"type": "text", "text": "text_2"}], + "tool_calls": [ + { + "id": "call_3", + "type": "function", + "function": { + "name": "func_3", + "arguments": '{"arg": "value3"}', + }, + }, + ], + "reasoning_content": "thinking_2", + }, + { + "role": "tool", + "tool_call_id": "call_3", + "content": "result_3", + "name": "func_3", + }, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_4", + "type": "function", + "function": { + "name": "func_4", + "arguments": '{"arg": "value4"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_4", + "content": "result_4", + "name": "func_4", + }, + { + "role": "assistant", + "content": [{"type": "text", "text": "text_3"}], + "reasoning_content": "thinking_3", + }, + ], + res, + ) + + async def test_chat_formatter_hint_block(self) -> None: + """HintBlock flushes preceding content and becomes a user message.""" + fmt = DashScopeChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + TextBlock(text="Let me think about that."), + HintBlock(hint="Remember to be concise."), + TextBlock(text="Here is my answer."), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me think about that."}, + ], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "Remember to be concise."}, + ], + }, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Here is my answer."}, + ], + }, + ], + res, + ) + + async def test_chat_formatter_hint_block_multimodal(self) -> None: + """Multimodal HintBlock becomes a single user message with text + + image.""" + fmt = DashScopeChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + HintBlock( + hint=[ + TextBlock(text="Inspect this screenshot:"), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + ), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Inspect this screenshot:", + }, + { + "type": "image_url", + "image_url": {"url": self.image_data_uri}, + }, + ], + }, + ], + res, + ) diff --git a/tests/formatter_deepseek_test.py b/tests/formatter_deepseek_test.py new file mode 100644 index 0000000000000000000000000000000000000000..fa33a4d163a59c7663baef5aeca24037e2387f77 --- /dev/null +++ b/tests/formatter_deepseek_test.py @@ -0,0 +1,532 @@ +# -*- coding: utf-8 -*- +"""Comprehensive formatter unit tests for DeepSeekChatFormatter and +DeepSeekMultiAgentFormatter, with exact ground-truth comparisons. +""" +from unittest import IsolatedAsyncioTestCase + +from agentscope.formatter import ( + DeepSeekChatFormatter, + DeepSeekMultiAgentFormatter, +) +from agentscope.message import ( + UserMsg, + AssistantMsg, + SystemMsg, + TextBlock, + DataBlock, + Base64Source, + ToolCallBlock, + ToolResultBlock, + ToolResultState, + ThinkingBlock, + HintBlock, +) + + +class TestDeepSeekFormatter(IsolatedAsyncioTestCase): + """Comprehensive tests for DeepSeek Chat and MultiAgent formatters.""" + + async def asyncSetUp(self) -> None: + """Set up shared fixtures and ground-truth dicts.""" + _hist_prompt = ( + DeepSeekMultiAgentFormatter().conversation_history_prompt + ) + + self.msgs_system = [ + SystemMsg( + name="system", + content="You're a helpful assistant.", + ), + ] + self.msgs_conversation = [ + UserMsg( + name="user", + content="What is the capital of France?", + ), + AssistantMsg( + name="assistant", + content="The capital of France is Paris.", + ), + UserMsg( + name="user", + content="What is the capital of Germany?", + ), + AssistantMsg( + name="assistant", + content="The capital of Germany is Berlin.", + ), + UserMsg( + name="user", + content="What is the capital of Japan?", + ), + ] + self.msgs_tools = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_1", + name="get_capital", + input='{"country": "Japan"}', + ), + ToolResultBlock( + id="call_1", + name="get_capital", + output=[ + TextBlock(text="The capital of Japan is Tokyo."), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="The capital of Japan is Tokyo."), + ], + ), + ] + + # --- Chat formatter ground truth --- + # DeepSeek content is a plain string (not a list of blocks). + # All assistant messages include `reasoning_content` (empty string if + # no ThinkingBlock). + self.gt_chat = [ + {"role": "system", "content": "You're a helpful assistant."}, + {"role": "user", "content": "What is the capital of France?"}, + { + "role": "assistant", + "content": "The capital of France is Paris.", + "reasoning_content": "", + }, + {"role": "user", "content": "What is the capital of Germany?"}, + { + "role": "assistant", + "content": "The capital of Germany is Berlin.", + "reasoning_content": "", + }, + {"role": "user", "content": "What is the capital of Japan?"}, + { + "role": "assistant", + "content": None, + "reasoning_content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_capital", + "arguments": '{"country": "Japan"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "The capital of Japan is Tokyo.", + "name": "get_capital", + }, + { + "role": "assistant", + "content": "The capital of Japan is Tokyo.", + "reasoning_content": "", + }, + ] + + # --- MultiAgent formatter ground truth --- + # System content is a plain string. + # History is a plain string (not a list) with tags. + # The trailing assistant message (is_first=False) is wrapped in a + # minimal block without the full prompt prefix. + self._gt_trailing_asst = { + "role": "assistant", + "content": "The capital of Japan is Tokyo.", + "reasoning_content": "", + } + self._gt_tool_call = { + "role": "assistant", + "content": None, + "reasoning_content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_capital", + "arguments": '{"country": "Japan"}', + }, + }, + ], + } + self._gt_tool_result = { + "role": "tool", + "tool_call_id": "call_1", + "content": "The capital of Japan is Tokyo.", + "name": "get_capital", + } + + self.gt_multiagent = [ + {"role": "system", "content": "You're a helpful assistant."}, + { + "role": "user", + "content": ( + _hist_prompt + "\n" + "user: What is the capital of France?\n" + "assistant: The capital of France is Paris.\n" + "user: What is the capital of Germany?\n" + "assistant: The capital of Germany is Berlin.\n" + "user: What is the capital of Japan?\n" + "" + ), + }, + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ] + + # ------------------------------------------------------------------ + # DeepSeekChatFormatter tests + # ------------------------------------------------------------------ + + async def test_chat_formatter(self) -> None: + """Chat formatter produces exact output for various subsets.""" + fmt = DeepSeekChatFormatter() + + # Full history + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_chat, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_chat[1:], res) + + # Without conversation + n_tools_gt = len(self.gt_chat) - 1 - len(self.msgs_conversation) + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [self.gt_chat[0]] + self.gt_chat[-n_tools_gt:], + res, + ) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_chat[:-n_tools_gt], res) + + # Empty + self.assertListEqual([], await fmt.format([])) + + async def test_chat_formatter_reasoning_content_always_present( + self, + ) -> None: + """Every assistant message always has a reasoning_content field.""" + fmt = DeepSeekChatFormatter() + msgs = [AssistantMsg(name="assistant", content="Answer")] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": "Answer", + "reasoning_content": "", + }, + ], + res, + ) + + async def test_chat_formatter_thinking_block(self) -> None: + """ThinkingBlock is placed into reasoning_content.""" + fmt = DeepSeekChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="Let me think..."), + TextBlock(text="Answer"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": "Answer", + "reasoning_content": "Let me think...", + }, + ], + res, + ) + + # ------------------------------------------------------------------ + # DeepSeekMultiAgentFormatter tests + # ------------------------------------------------------------------ + + async def test_multiagent_formatter(self) -> None: + """MultiAgent formatter produces exact output for various subsets.""" + fmt = DeepSeekMultiAgentFormatter() + + # Full + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_multiagent, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_multiagent[1:], res) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_multiagent[:2], res) + + # System only + res = await fmt.format(self.msgs_system) + self.assertListEqual([self.gt_multiagent[0]], res) + + # Conversation only + res = await fmt.format(self.msgs_conversation) + self.assertListEqual([self.gt_multiagent[1]], res) + + # Tools only + res = await fmt.format(self.msgs_tools) + self.assertListEqual( + [ + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # System + tools + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [ + self.gt_multiagent[0], + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # Empty + self.assertListEqual([], await fmt.format([])) + + async def test_chat_formatter_complex_multi_step(self) -> None: + """Complex multi-step sequence with interleaved thinking, text, + tool calls, and tool results.""" + fmt = DeepSeekChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="thinking_1"), + TextBlock(text="text_1"), + ToolCallBlock( + id="call_1", + name="func_1", + input='{"arg": "value1"}', + ), + ToolCallBlock( + id="call_2", + name="func_2", + input='{"arg": "value2"}', + ), + ToolResultBlock( + id="call_1", + name="func_1", + output=[TextBlock(text="result_1")], + state=ToolResultState.SUCCESS, + ), + ToolResultBlock( + id="call_2", + name="func_2", + output=[TextBlock(text="result_2")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_2"), + TextBlock(text="text_2"), + ToolCallBlock( + id="call_3", + name="func_3", + input='{"arg": "value3"}', + ), + ToolResultBlock( + id="call_3", + name="func_3", + output=[TextBlock(text="result_3")], + state=ToolResultState.SUCCESS, + ), + ToolCallBlock( + id="call_4", + name="func_4", + input='{"arg": "value4"}', + ), + ToolResultBlock( + id="call_4", + name="func_4", + output=[TextBlock(text="result_4")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_3"), + TextBlock(text="text_3"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": "text_1", + "reasoning_content": "thinking_1", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "func_1", + "arguments": '{"arg": "value1"}', + }, + }, + { + "id": "call_2", + "type": "function", + "function": { + "name": "func_2", + "arguments": '{"arg": "value2"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "result_1", + "name": "func_1", + }, + { + "role": "tool", + "tool_call_id": "call_2", + "content": "result_2", + "name": "func_2", + }, + { + "role": "assistant", + "content": "text_2", + "reasoning_content": "thinking_2", + "tool_calls": [ + { + "id": "call_3", + "type": "function", + "function": { + "name": "func_3", + "arguments": '{"arg": "value3"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_3", + "content": "result_3", + "name": "func_3", + }, + { + "role": "assistant", + "content": None, + "reasoning_content": "", + "tool_calls": [ + { + "id": "call_4", + "type": "function", + "function": { + "name": "func_4", + "arguments": '{"arg": "value4"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_4", + "content": "result_4", + "name": "func_4", + }, + { + "role": "assistant", + "content": "text_3", + "reasoning_content": "thinking_3", + }, + ], + res, + ) + + async def test_chat_formatter_hint_block(self) -> None: + """HintBlock flushes preceding content and becomes a user message.""" + fmt = DeepSeekChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + TextBlock(text="Let me think about that."), + HintBlock(hint="Remember to be concise."), + TextBlock(text="Here is my answer."), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": "Let me think about that.", + "reasoning_content": "", + }, + { + "role": "user", + "content": "Remember to be concise.", + }, + { + "role": "assistant", + "content": "Here is my answer.", + "reasoning_content": "", + }, + ], + res, + ) + + async def test_chat_formatter_hint_block_multimodal(self) -> None: + """DeepSeek is text-only — DataBlock degrades to a placeholder + string.""" + fmt = DeepSeekChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + HintBlock( + hint=[ + TextBlock(text="Inspect this screenshot:"), + DataBlock( + source=Base64Source( + data="ZmFrZSBpbWFnZSBkYXRh", + media_type="image/png", + ), + ), + ], + ), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "user", + "content": ( + "Inspect this screenshot:\n" + "[image/png attached, not supported by this provider]" + ), + }, + ], + res, + ) diff --git a/tests/formatter_gemini_test.py b/tests/formatter_gemini_test.py new file mode 100644 index 0000000000000000000000000000000000000000..d5d0f6d1fb378fe40ca3d026db9f576db73ee426 --- /dev/null +++ b/tests/formatter_gemini_test.py @@ -0,0 +1,718 @@ +# -*- coding: utf-8 -*- +"""Comprehensive formatter unit tests for GeminiChatFormatter and +GeminiMultiAgentFormatter, following the reference test style with exact +ground-truth comparisons. +""" +from unittest import IsolatedAsyncioTestCase +from unittest.mock import patch + +from agentscope.formatter import ( + GeminiChatFormatter, + GeminiMultiAgentFormatter, +) +from agentscope.message import ( + UserMsg, + AssistantMsg, + SystemMsg, + TextBlock, + DataBlock, + ToolCallBlock, + ToolResultBlock, + ToolResultState, + Base64Source, + ThinkingBlock, + HintBlock, +) + + +_FIXED_ID = "TESTID1234567" + + +class TestGeminiFormatter(IsolatedAsyncioTestCase): + """Comprehensive tests for Gemini Chat and MultiAgent formatters.""" + + async def asyncSetUp(self) -> None: + """Set up shared message fixtures and expected ground-truth dicts.""" + self.image_b64 = "ZmFrZSBpbWFnZSBkYXRh" + + # --------------------------------------------------------------- + # Message fixtures + # (Use base64 images: Gemini URL handling downloads from the network) + # --------------------------------------------------------------- + self.msgs_system = [ + SystemMsg( + name="system", + content="You're a helpful assistant.", + ), + ] + + self.msgs_conversation = [ + UserMsg( + name="user", + content=[ + TextBlock(text="What is the capital of France?"), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + ), + AssistantMsg( + name="assistant", + content="The capital of France is Paris.", + ), + UserMsg( + name="user", + content="What is the capital of Germany?", + ), + AssistantMsg( + name="assistant", + content="The capital of Germany is Berlin.", + ), + UserMsg( + name="user", + content="What is the capital of Japan?", + ), + ] + + self.msgs_tools = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_1", + name="get_capital", + input='{"country": "Japan"}', + ), + ToolResultBlock( + id="call_1", + name="get_capital", + output=[ + TextBlock(text="The capital of Japan is Tokyo."), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="The capital of Japan is Tokyo."), + ], + ), + ] + + _inline_img = { + "inline_data": {"data": self.image_b64, "mime_type": "image/png"}, + } + + # --------------------------------------------------------------- + # Ground truth: GeminiChatFormatter + # - System message becomes role="user" (no special system role). + # - Assistant messages become role="model". + # - Content is in "parts" (not "content") as a list of dicts. + # - ToolCallBlock becomes "function_call" part. + # - ToolResultBlock becomes a separate role="user" message with + # "function_response" part. + # --------------------------------------------------------------- + self.gt_chat = [ + { + "role": "user", + "parts": [{"text": "You're a helpful assistant."}], + }, + { + "role": "user", + "parts": [ + {"text": "What is the capital of France?"}, + _inline_img, + ], + }, + { + "role": "model", + "parts": [{"text": "The capital of France is Paris."}], + }, + { + "role": "user", + "parts": [{"text": "What is the capital of Germany?"}], + }, + { + "role": "model", + "parts": [{"text": "The capital of Germany is Berlin."}], + }, + { + "role": "user", + "parts": [{"text": "What is the capital of Japan?"}], + }, + { + "role": "model", + "parts": [ + { + "function_call": { + "id": "call_1", + "name": "get_capital", + "args": {"country": "Japan"}, + }, + }, + ], + }, + { + "role": "user", + "parts": [ + { + "function_response": { + "id": "call_1", + "name": "get_capital", + "response": { + "output": "The capital of Japan is Tokyo.", + }, + }, + }, + ], + }, + { + "role": "model", + "parts": [{"text": "The capital of Japan is Tokyo."}], + }, + ] + + # --------------------------------------------------------------- + # Ground truth: GeminiMultiAgentFormatter + # - System message: role="user" (same as chat formatter). + # - Agent messages: collapsed into role="user" with parts list. + # - Media blocks interleaved (text flushed before each DataBlock). + # - is_first=False still wraps with (no hist_prompt + # prefix). + # --------------------------------------------------------------- + _hist_prompt = GeminiMultiAgentFormatter().conversation_history_prompt + + self._gt_trailing_asst = { + "role": "model", + "parts": [ + {"text": "The capital of Japan is Tokyo."}, + ], + } + + self._gt_tool_call = { + "role": "model", + "parts": [ + { + "function_call": { + "id": "call_1", + "name": "get_capital", + "args": {"country": "Japan"}, + }, + }, + ], + } + self._gt_tool_result = { + "role": "user", + "parts": [ + { + "function_response": { + "id": "call_1", + "name": "get_capital", + "response": { + "output": "The capital of Japan is Tokyo.", + }, + }, + }, + ], + } + + self.gt_multiagent = [ + { + "role": "user", + "parts": [{"text": "You're a helpful assistant."}], + }, + { + "role": "user", + "parts": [ + { + "text": ( + _hist_prompt + "\n" + "user: What is the capital of France?" + ), + }, + _inline_img, + { + "text": ( + "assistant: The capital of France is Paris.\n" + "user: What is the capital of Germany?\n" + "assistant: The capital of Germany is Berlin.\n" + "user: What is the capital of Japan?\n" + "" + ), + }, + ], + }, + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ] + + # ------------------------------------------------------------------- + # GeminiChatFormatter tests + # ------------------------------------------------------------------- + + async def test_chat_formatter(self) -> None: + """Chat formatter produces exact output for various subsets.""" + fmt = GeminiChatFormatter() + + # Full history + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_chat, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_chat[1:], res) + + # Without conversation + n_tools_gt = len(self.gt_chat) - 1 - len(self.msgs_conversation) + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [self.gt_chat[0]] + self.gt_chat[-n_tools_gt:], + res, + ) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_chat[:-n_tools_gt], res) + + # Empty + res = await fmt.format([]) + self.assertListEqual([], res) + + async def test_chat_formatter_thinking_preserved(self) -> None: + """ThinkingBlock becomes a part with thought=True in Gemini format.""" + fmt = GeminiChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="inner thoughts"), + TextBlock(text="reply"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "model", + "parts": [ + {"thought": True, "text": "inner thoughts"}, + {"text": "reply"}, + ], + }, + ], + res, + ) + + @patch( + "agentscope.formatter._formatter_base.shortuuid.uuid", + return_value=_FIXED_ID, + ) + async def test_chat_formatter_base64_image_in_tool_result( + self, + _mock_uuid: object, + ) -> None: + """Base64 images in tool results are promoted to a follow-up user + message.""" + fmt = GeminiChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_img", + name="get_map", + input='{"city": "Tokyo"}', + ), + ToolResultBlock( + id="call_img", + name="get_map", + output=[ + TextBlock(text="Here is the map."), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="Here is the map of Tokyo."), + ], + ), + ] + res = await fmt.format(msgs) + + expected_tool_output = ( + "Here is the map.\n" + f"A(n) image file is returned " + f"and will be presented to you with the identifier " + f"[{_FIXED_ID}]." + ) + self.assertListEqual( + [ + { + "role": "model", + "parts": [ + { + "function_call": { + "id": "call_img", + "name": "get_map", + "args": {"city": "Tokyo"}, + }, + }, + ], + }, + { + "role": "user", + "parts": [ + { + "function_response": { + "id": "call_img", + "name": "get_map", + "response": {"output": expected_tool_output}, + }, + }, + ], + }, + { + "role": "user", + "parts": [ + { + "text": ( + "The multimodal data " + "and their identifiers are listed as " + "follows:" + ), + }, + { + "text": f"- {_FIXED_ID} (image file): ", + }, + { + "inline_data": { + "data": self.image_b64, + "mime_type": "image/png", + }, + }, + {"text": ""}, + ], + }, + { + "role": "model", + "parts": [ + {"text": "Here is the map of Tokyo."}, + ], + }, + ], + res, + ) + + # ------------------------------------------------------------------- + # GeminiMultiAgentFormatter tests + # ------------------------------------------------------------------- + + async def test_multiagent_formatter(self) -> None: + """MultiAgent formatter produces exact output for various subsets.""" + fmt = GeminiMultiAgentFormatter() + + # Full history + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_multiagent, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_multiagent[1:], res) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_multiagent[:2], res) + + # System only + res = await fmt.format(self.msgs_system) + self.assertListEqual([self.gt_multiagent[0]], res) + + # Conversation only + res = await fmt.format(self.msgs_conversation) + self.assertListEqual([self.gt_multiagent[1]], res) + + # Tools only + res = await fmt.format(self.msgs_tools) + self.assertListEqual( + [ + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # System + tools + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [ + self.gt_multiagent[0], + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # Empty + res = await fmt.format([]) + self.assertListEqual([], res) + + async def test_chat_formatter_complex_multi_step(self) -> None: + """Complex multi-step sequence with interleaved thinking, text, + tool calls, and tool results.""" + fmt = GeminiChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="thinking_1"), + TextBlock(text="text_1"), + ToolCallBlock( + id="call_1", + name="func_1", + input='{"arg": "value1"}', + ), + ToolCallBlock( + id="call_2", + name="func_2", + input='{"arg": "value2"}', + ), + ToolResultBlock( + id="call_1", + name="func_1", + output=[TextBlock(text="result_1")], + state=ToolResultState.SUCCESS, + ), + ToolResultBlock( + id="call_2", + name="func_2", + output=[TextBlock(text="result_2")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_2"), + TextBlock(text="text_2"), + ToolCallBlock( + id="call_3", + name="func_3", + input='{"arg": "value3"}', + ), + ToolResultBlock( + id="call_3", + name="func_3", + output=[TextBlock(text="result_3")], + state=ToolResultState.SUCCESS, + ), + ToolCallBlock( + id="call_4", + name="func_4", + input='{"arg": "value4"}', + ), + ToolResultBlock( + id="call_4", + name="func_4", + output=[TextBlock(text="result_4")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_3"), + TextBlock(text="text_3"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "model", + "parts": [ + {"thought": True, "text": "thinking_1"}, + {"text": "text_1"}, + { + "function_call": { + "id": "call_1", + "name": "func_1", + "args": {"arg": "value1"}, + }, + }, + { + "function_call": { + "id": "call_2", + "name": "func_2", + "args": {"arg": "value2"}, + }, + }, + ], + }, + { + "role": "user", + "parts": [ + { + "function_response": { + "id": "call_1", + "name": "func_1", + "response": {"output": "result_1"}, + }, + }, + ], + }, + { + "role": "user", + "parts": [ + { + "function_response": { + "id": "call_2", + "name": "func_2", + "response": {"output": "result_2"}, + }, + }, + ], + }, + { + "role": "model", + "parts": [ + {"thought": True, "text": "thinking_2"}, + {"text": "text_2"}, + { + "function_call": { + "id": "call_3", + "name": "func_3", + "args": {"arg": "value3"}, + }, + }, + ], + }, + { + "role": "user", + "parts": [ + { + "function_response": { + "id": "call_3", + "name": "func_3", + "response": {"output": "result_3"}, + }, + }, + ], + }, + { + "role": "model", + "parts": [ + { + "function_call": { + "id": "call_4", + "name": "func_4", + "args": {"arg": "value4"}, + }, + }, + ], + }, + { + "role": "user", + "parts": [ + { + "function_response": { + "id": "call_4", + "name": "func_4", + "response": {"output": "result_4"}, + }, + }, + ], + }, + { + "role": "model", + "parts": [ + {"thought": True, "text": "thinking_3"}, + {"text": "text_3"}, + ], + }, + ], + res, + ) + + async def test_chat_formatter_hint_block(self) -> None: + """HintBlock flushes preceding content and becomes a user message.""" + fmt = GeminiChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + TextBlock(text="Let me think about that."), + HintBlock(hint="Remember to be concise."), + TextBlock(text="Here is my answer."), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "model", + "parts": [ + {"text": "Let me think about that."}, + ], + }, + { + "role": "user", + "parts": [ + {"text": "Remember to be concise."}, + ], + }, + { + "role": "model", + "parts": [ + {"text": "Here is my answer."}, + ], + }, + ], + res, + ) + + async def test_chat_formatter_hint_block_multimodal(self) -> None: + """Multimodal HintBlock becomes a single user message with text + + image.""" + fmt = GeminiChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + HintBlock( + hint=[ + TextBlock(text="Inspect this screenshot:"), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + ), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "user", + "parts": [ + {"text": "Inspect this screenshot:"}, + { + "inline_data": { + "data": self.image_b64, + "mime_type": "image/png", + }, + }, + ], + }, + ], + res, + ) diff --git a/tests/formatter_moonshot_test.py b/tests/formatter_moonshot_test.py new file mode 100644 index 0000000000000000000000000000000000000000..f1b966762a0322249a7b0e0cc789a690dd9d57ca --- /dev/null +++ b/tests/formatter_moonshot_test.py @@ -0,0 +1,826 @@ +# -*- coding: utf-8 -*- +"""Comprehensive formatter unit tests for MoonshotChatFormatter and +MoonshotMultiAgentFormatter, following the reference test style with exact +ground-truth comparisons. +""" +import base64 +from unittest import IsolatedAsyncioTestCase +from unittest.mock import patch, MagicMock + +from agentscope.formatter import ( + MoonshotChatFormatter, + MoonshotMultiAgentFormatter, +) +from agentscope.message import ( + UserMsg, + AssistantMsg, + SystemMsg, + TextBlock, + DataBlock, + ToolCallBlock, + ToolResultBlock, + ToolResultState, + Base64Source, + URLSource, + ThinkingBlock, + HintBlock, +) + + +_FIXED_ID = "TESTID1234567" + + +class TestMoonshotFormatter(IsolatedAsyncioTestCase): + """Comprehensive tests for Moonshot Chat and MultiAgent formatters.""" + + async def asyncSetUp(self) -> None: + """Set up shared message fixtures and expected ground-truth dicts.""" + _img_src = URLSource( + url="https://example.com/image.png", + media_type="image/png", + ) + self.image_url = str(_img_src.url) + + # The Moonshot formatter downloads remote image URLs and inlines + # them as base64 data URIs (the Moonshot vision API rejects raw + # HTTPS URLs). Patch `requests.get` so tests don't hit the network + # and produce a deterministic payload that matches `image_b64`. + self.image_b64 = "ZmFrZSBpbWFnZSBkYXRh" + self.image_bytes = base64.b64decode(self.image_b64) + self.image_data_uri = f"data:image/png;base64,{self.image_b64}" + + mock_response = MagicMock() + mock_response.content = self.image_bytes + mock_response.raise_for_status = MagicMock() + self._requests_get_patcher = patch( + "agentscope.formatter._moonshot_formatter.requests.get", + return_value=mock_response, + ) + self._requests_get_patcher.start() + self.addCleanup(self._requests_get_patcher.stop) + + # --------------------------------------------------------------- + # Message fixtures (no audio to avoid downloads) + # --------------------------------------------------------------- + self.msgs_system = [ + SystemMsg( + name="system", + content="You're a helpful assistant.", + ), + ] + + self.msgs_conversation = [ + UserMsg( + name="user", + content=[ + TextBlock(text="What is the capital of France?"), + DataBlock( + source=URLSource( + url=self.image_url, + media_type="image/png", + ), + ), + ], + ), + AssistantMsg( + name="assistant", + content="The capital of France is Paris.", + ), + UserMsg( + name="user", + content="What is the capital of Germany?", + ), + AssistantMsg( + name="assistant", + content="The capital of Germany is Berlin.", + ), + UserMsg( + name="user", + content="What is the capital of Japan?", + ), + ] + + self.msgs_tools = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_1", + name="get_capital", + input='{"country": "Japan"}', + ), + ToolResultBlock( + id="call_1", + name="get_capital", + output=[ + TextBlock(text="The capital of Japan is Tokyo."), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="The capital of Japan is Tokyo."), + ], + ), + ] + + # --------------------------------------------------------------- + # Ground truth: MoonshotChatFormatter + # - Same as OpenAI except ALL assistant messages have an extra + # "reasoning_content" field (empty string when no ThinkingBlock). + # --------------------------------------------------------------- + self.gt_chat = [ + { + "role": "system", + "name": "system", + "content": [ + {"type": "text", "text": "You're a helpful assistant."}, + ], + }, + { + "role": "user", + "name": "user", + "content": [ + {"type": "text", "text": "What is the capital of France?"}, + { + "type": "image_url", + "image_url": {"url": self.image_data_uri}, + }, + ], + }, + { + "role": "assistant", + "name": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of France is Paris.", + }, + ], + "reasoning_content": "", + }, + { + "role": "user", + "name": "user", + "content": [ + { + "type": "text", + "text": "What is the capital of Germany?", + }, + ], + }, + { + "role": "assistant", + "name": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of Germany is Berlin.", + }, + ], + "reasoning_content": "", + }, + { + "role": "user", + "name": "user", + "content": [ + {"type": "text", "text": "What is the capital of Japan?"}, + ], + }, + { + "role": "assistant", + "name": "assistant", + "content": None, + "reasoning_content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_capital", + "arguments": '{"country": "Japan"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "The capital of Japan is Tokyo.", + "name": "get_capital", + }, + { + "role": "assistant", + "name": "assistant", + "content": [ + {"type": "text", "text": "The capital of Japan is Tokyo."}, + ], + "reasoning_content": "", + }, + ] + + # --------------------------------------------------------------- + # Ground truth: MoonshotMultiAgentFormatter + # - Same as OpenAI MultiAgent, but tool-sequence assistant messages + # carry "reasoning_content": "". + # --------------------------------------------------------------- + _hist_prompt = ( + MoonshotMultiAgentFormatter().conversation_history_prompt + ) + + _conv_text = ( + "user: What is the capital of France?\n" + "assistant: The capital of France is Paris.\n" + "user: What is the capital of Germany?\n" + "assistant: The capital of Germany is Berlin.\n" + "user: What is the capital of Japan?" + ) + + self._gt_trailing_asst = { + "role": "assistant", + "name": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of Japan is Tokyo.", + }, + ], + "reasoning_content": "", + } + + self._gt_tool_call = { + "role": "assistant", + "name": "assistant", + "content": None, + "reasoning_content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_capital", + "arguments": '{"country": "Japan"}', + }, + }, + ], + } + self._gt_tool_result = { + "role": "tool", + "tool_call_id": "call_1", + "content": "The capital of Japan is Tokyo.", + "name": "get_capital", + } + + self.gt_multiagent = [ + { + "role": "system", + "content": "You're a helpful assistant.", + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + _hist_prompt + + "\n" + + _conv_text + + "\n" + ), + }, + { + "type": "image_url", + "image_url": {"url": self.image_data_uri}, + }, + ], + }, + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ] + + # ------------------------------------------------------------------- + # MoonshotChatFormatter tests + # ------------------------------------------------------------------- + + async def test_chat_formatter(self) -> None: + """Chat formatter produces exact output for various subsets.""" + fmt = MoonshotChatFormatter() + + # Full history + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_chat, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_chat[1:], res) + + # Without conversation + n_tools_gt = len(self.gt_chat) - 1 - len(self.msgs_conversation) + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [self.gt_chat[0]] + self.gt_chat[-n_tools_gt:], + res, + ) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_chat[:-n_tools_gt], res) + + # Empty + res = await fmt.format([]) + self.assertListEqual([], res) + + async def test_chat_formatter_thinking_to_reasoning_content( + self, + ) -> None: + """ThinkingBlock becomes reasoning_content in Moonshot (Preserved + Thinking).""" + fmt = MoonshotChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="inner thoughts"), + TextBlock(text="reply"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "name": "assistant", + "content": [{"type": "text", "text": "reply"}], + "reasoning_content": "inner thoughts", + }, + ], + res, + ) + + async def test_chat_formatter_assistant_always_has_reasoning_content( + self, + ) -> None: + """All assistant messages always have reasoning_content (even when + empty).""" + fmt = MoonshotChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content="Hello!", + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "name": "assistant", + "content": [{"type": "text", "text": "Hello!"}], + "reasoning_content": "", + }, + ], + res, + ) + + async def test_chat_formatter_base64_image(self) -> None: + """Base64-encoded image is inlined as a data URI.""" + fmt = MoonshotChatFormatter() + msgs = [ + UserMsg( + name="user", + content=[ + TextBlock(text="What's in this image?"), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "user", + "name": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + { + "type": "image_url", + "image_url": {"url": self.image_data_uri}, + }, + ], + }, + ], + res, + ) + + @patch( + "agentscope.formatter._formatter_base.shortuuid.uuid", + return_value=_FIXED_ID, + ) + async def test_chat_formatter_url_image_in_tool_result( + self, + _mock_uuid: object, + ) -> None: + """URL images in tool results are promoted to a follow-up user + message.""" + fmt = MoonshotChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_img", + name="get_map", + input='{"city": "Tokyo"}', + ), + ToolResultBlock( + id="call_img", + name="get_map", + output=[ + TextBlock(text="Here is the map."), + DataBlock( + source=URLSource( + url=self.image_url, + media_type="image/png", + ), + ), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="Here is the map of Tokyo."), + ], + ), + ] + res = await fmt.format(msgs) + + expected_tool_content = ( + "Here is the map.\n" + f"A(n) image file is returned " + f"and will be presented to you with the identifier " + f"[{_FIXED_ID}]." + ) + self.assertListEqual( + [ + { + "role": "assistant", + "name": "assistant", + "content": None, + "reasoning_content": "", + "tool_calls": [ + { + "id": "call_img", + "type": "function", + "function": { + "name": "get_map", + "arguments": '{"city": "Tokyo"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_img", + "content": expected_tool_content, + "name": "get_map", + }, + { + "role": "user", + "name": "system-reminder", + "content": [ + { + "type": "text", + "text": ( + "The multimodal data " + "and their identifiers are listed as " + "follows:" + ), + }, + { + "type": "text", + "text": f"- {_FIXED_ID} (image file): ", + }, + { + "type": "image_url", + "image_url": {"url": self.image_data_uri}, + }, + { + "type": "text", + "text": "", + }, + ], + }, + { + "role": "assistant", + "name": "assistant", + "content": [ + { + "type": "text", + "text": "Here is the map of Tokyo.", + }, + ], + "reasoning_content": "", + }, + ], + res, + ) + + # ------------------------------------------------------------------- + # MoonshotMultiAgentFormatter tests + # ------------------------------------------------------------------- + + async def test_multiagent_formatter(self) -> None: + """MultiAgent formatter produces exact output for various subsets.""" + fmt = MoonshotMultiAgentFormatter() + + # Full history + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_multiagent, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_multiagent[1:], res) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_multiagent[:2], res) + + # System only + res = await fmt.format(self.msgs_system) + self.assertListEqual([self.gt_multiagent[0]], res) + + # Conversation only + res = await fmt.format(self.msgs_conversation) + self.assertListEqual([self.gt_multiagent[1]], res) + + # Tools only + res = await fmt.format(self.msgs_tools) + self.assertListEqual( + [ + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # System + tools + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [ + self.gt_multiagent[0], + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # Empty + res = await fmt.format([]) + self.assertListEqual([], res) + + async def test_chat_formatter_complex_multi_step(self) -> None: + """Complex multi-step sequence with interleaved thinking, text, + tool calls, and tool results.""" + fmt = MoonshotChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="thinking_1"), + TextBlock(text="text_1"), + ToolCallBlock( + id="call_1", + name="func_1", + input='{"arg": "value1"}', + ), + ToolCallBlock( + id="call_2", + name="func_2", + input='{"arg": "value2"}', + ), + ToolResultBlock( + id="call_1", + name="func_1", + output=[TextBlock(text="result_1")], + state=ToolResultState.SUCCESS, + ), + ToolResultBlock( + id="call_2", + name="func_2", + output=[TextBlock(text="result_2")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_2"), + TextBlock(text="text_2"), + ToolCallBlock( + id="call_3", + name="func_3", + input='{"arg": "value3"}', + ), + ToolResultBlock( + id="call_3", + name="func_3", + output=[TextBlock(text="result_3")], + state=ToolResultState.SUCCESS, + ), + ToolCallBlock( + id="call_4", + name="func_4", + input='{"arg": "value4"}', + ), + ToolResultBlock( + id="call_4", + name="func_4", + output=[TextBlock(text="result_4")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_3"), + TextBlock(text="text_3"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "name": "assistant", + "content": [{"type": "text", "text": "text_1"}], + "reasoning_content": "thinking_1", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "func_1", + "arguments": '{"arg": "value1"}', + }, + }, + { + "id": "call_2", + "type": "function", + "function": { + "name": "func_2", + "arguments": '{"arg": "value2"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "result_1", + "name": "func_1", + }, + { + "role": "tool", + "tool_call_id": "call_2", + "content": "result_2", + "name": "func_2", + }, + { + "role": "assistant", + "name": "assistant", + "content": [{"type": "text", "text": "text_2"}], + "reasoning_content": "thinking_2", + "tool_calls": [ + { + "id": "call_3", + "type": "function", + "function": { + "name": "func_3", + "arguments": '{"arg": "value3"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_3", + "content": "result_3", + "name": "func_3", + }, + { + "role": "assistant", + "name": "assistant", + "content": None, + "reasoning_content": "", + "tool_calls": [ + { + "id": "call_4", + "type": "function", + "function": { + "name": "func_4", + "arguments": '{"arg": "value4"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_4", + "content": "result_4", + "name": "func_4", + }, + { + "role": "assistant", + "name": "assistant", + "content": [{"type": "text", "text": "text_3"}], + "reasoning_content": "thinking_3", + }, + ], + res, + ) + + async def test_chat_formatter_hint_block(self) -> None: + """HintBlock flushes preceding content and becomes a user message.""" + fmt = MoonshotChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + TextBlock(text="Let me think about that."), + HintBlock(hint="Remember to be concise."), + TextBlock(text="Here is my answer."), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "name": "assistant", + "content": [ + {"type": "text", "text": "Let me think about that."}, + ], + "reasoning_content": "", + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "Remember to be concise."}, + ], + }, + { + "role": "assistant", + "name": "assistant", + "content": [ + {"type": "text", "text": "Here is my answer."}, + ], + "reasoning_content": "", + }, + ], + res, + ) + + async def test_chat_formatter_hint_block_multimodal(self) -> None: + """Multimodal HintBlock becomes a single user message with text + + image.""" + fmt = MoonshotChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + HintBlock( + hint=[ + TextBlock(text="Inspect this screenshot:"), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + ), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Inspect this screenshot:", + }, + { + "type": "image_url", + "image_url": {"url": self.image_data_uri}, + }, + ], + }, + ], + res, + ) diff --git a/tests/formatter_ollama_test.py b/tests/formatter_ollama_test.py new file mode 100644 index 0000000000000000000000000000000000000000..038271eaac648031818bd94ecb7a301ff626fb24 --- /dev/null +++ b/tests/formatter_ollama_test.py @@ -0,0 +1,576 @@ +# -*- coding: utf-8 -*- +"""Comprehensive formatter unit tests for OllamaChatFormatter and +OllamaMultiAgentFormatter, with exact ground-truth comparisons. +""" +from unittest import IsolatedAsyncioTestCase +from unittest.mock import patch + +from agentscope.formatter import OllamaChatFormatter, OllamaMultiAgentFormatter +from agentscope.message import ( + UserMsg, + AssistantMsg, + SystemMsg, + TextBlock, + DataBlock, + ToolCallBlock, + ToolResultBlock, + ToolResultState, + Base64Source, + ThinkingBlock, + HintBlock, +) + + +_FIXED_ID = "TESTID1234567" + + +class TestOllamaFormatter(IsolatedAsyncioTestCase): + """Comprehensive tests for Ollama Chat and MultiAgent formatters.""" + + async def asyncSetUp(self) -> None: + """Set up shared fixtures and ground-truth dicts.""" + _hist_prompt = OllamaMultiAgentFormatter().conversation_history_prompt + + # Base64 image fixture + self.image_b64 = "ZmFrZSBpbWFnZSBkYXRh" + + self.msgs_system = [ + SystemMsg( + name="system", + content="You're a helpful assistant.", + ), + ] + self.msgs_conversation = [ + UserMsg( + name="user", + content="What is the capital of France?", + ), + AssistantMsg( + name="assistant", + content="The capital of France is Paris.", + ), + UserMsg( + name="user", + content="What is the capital of Germany?", + ), + AssistantMsg( + name="assistant", + content="The capital of Germany is Berlin.", + ), + UserMsg( + name="user", + content="What is the capital of Japan?", + ), + ] + self.msgs_tools = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_1", + name="get_capital", + input='{"country": "Japan"}', + ), + ToolResultBlock( + id="call_1", + name="get_capital", + output=[ + TextBlock(text="The capital of Japan is Tokyo."), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="The capital of Japan is Tokyo."), + ], + ), + ] + + # --- Chat formatter ground truth --- + # Ollama content is always a plain string. + # Tool calls use dict arguments (not JSON string). + self.gt_chat = [ + {"role": "system", "content": "You're a helpful assistant."}, + {"role": "user", "content": "What is the capital of France?"}, + { + "role": "assistant", + "content": "The capital of France is Paris.", + }, + {"role": "user", "content": "What is the capital of Germany?"}, + { + "role": "assistant", + "content": "The capital of Germany is Berlin.", + }, + {"role": "user", "content": "What is the capital of Japan?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "function": { + "name": "get_capital", + "arguments": {"country": "Japan"}, + }, + }, + ], + }, + {"role": "tool", "content": "The capital of Japan is Tokyo."}, + {"role": "assistant", "content": "The capital of Japan is Tokyo."}, + ] + + # --- MultiAgent formatter ground truth --- + # System content is a plain string. + # History is a plain string with format "name:\ntext" per message. + # For is_first=False, there are NO tags (only in + # is_first=True). + _conv_text = ( + "user:\nWhat is the capital of France?\n" + "assistant:\nThe capital of France is Paris.\n" + "user:\nWhat is the capital of Germany?\n" + "assistant:\nThe capital of Germany is Berlin.\n" + "user:\nWhat is the capital of Japan?" + ) + self._gt_trailing_asst = { + "role": "assistant", + "content": "The capital of Japan is Tokyo.", + } + self._gt_tool_call = { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "function": { + "name": "get_capital", + "arguments": {"country": "Japan"}, + }, + }, + ], + } + self._gt_tool_result = { + "role": "tool", + "content": "The capital of Japan is Tokyo.", + } + + self.gt_multiagent = [ + {"role": "system", "content": "You're a helpful assistant."}, + { + "role": "user", + "content": ( + _hist_prompt + "\n" + _conv_text + "\n" + ), + }, + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ] + + # ------------------------------------------------------------------ + # OllamaChatFormatter tests + # ------------------------------------------------------------------ + + async def test_chat_formatter(self) -> None: + """Chat formatter produces exact output for various subsets.""" + fmt = OllamaChatFormatter() + + # Full history + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_chat, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_chat[1:], res) + + # Without conversation + n_tools_gt = len(self.gt_chat) - 1 - len(self.msgs_conversation) + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [self.gt_chat[0]] + self.gt_chat[-n_tools_gt:], + res, + ) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_chat[:-n_tools_gt], res) + + # Empty + self.assertListEqual([], await fmt.format([])) + + async def test_chat_formatter_tool_call_arguments_are_dict(self) -> None: + """Ollama requires tool call arguments as a dict, not a JSON string.""" + fmt = OllamaChatFormatter() + tc = ToolCallBlock(id="c1", name="search", input='{"q": "weather"}') + res = await fmt.format( + [AssistantMsg(name="assistant", content=[tc])], + ) + self.assertListEqual( + [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "function": { + "name": "search", + "arguments": {"q": "weather"}, + }, + }, + ], + }, + ], + res, + ) + + async def test_chat_formatter_base64_image(self) -> None: + """Base64 image is placed in the 'images' list as a raw base64 + string.""" + fmt = OllamaChatFormatter() + msgs = [ + UserMsg( + name="user", + content=[ + TextBlock(text="What is this?"), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "user", + "content": "What is this?", + "images": [self.image_b64], + }, + ], + res, + ) + + @patch( + "agentscope.formatter._formatter_base.shortuuid.uuid", + return_value=_FIXED_ID, + ) + async def test_chat_formatter_base64_image_in_tool_result( + self, + _mock_uuid: object, + ) -> None: + """Base64 images in tool results are promoted to a follow-up user + message with images list.""" + fmt = OllamaChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_img", + name="get_map", + input='{"city": "Tokyo"}', + ), + ToolResultBlock( + id="call_img", + name="get_map", + output=[ + TextBlock(text="Here is the map."), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="Here is the map of Tokyo."), + ], + ), + ] + res = await fmt.format(msgs) + + expected_tool_content = ( + "Here is the map.\n" + f"A(n) image file is returned " + f"and will be presented to you with the identifier " + f"[{_FIXED_ID}]." + ) + self.assertListEqual( + [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "function": { + "name": "get_map", + "arguments": {"city": "Tokyo"}, + }, + }, + ], + }, + { + "role": "tool", + "content": expected_tool_content, + }, + { + "role": "user", + "content": ( + "The multimodal data " + "and their identifiers are listed as follows:\n" + f"- {_FIXED_ID} (image file): \n" + "" + ), + "images": [self.image_b64], + }, + { + "role": "assistant", + "content": "Here is the map of Tokyo.", + }, + ], + res, + ) + + # ------------------------------------------------------------------ + # OllamaMultiAgentFormatter tests + # ------------------------------------------------------------------ + + async def test_multiagent_formatter(self) -> None: + """MultiAgent formatter produces exact output for various subsets.""" + fmt = OllamaMultiAgentFormatter() + + # Full + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_multiagent, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_multiagent[1:], res) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_multiagent[:2], res) + + # System only + res = await fmt.format(self.msgs_system) + self.assertListEqual([self.gt_multiagent[0]], res) + + # Conversation only + res = await fmt.format(self.msgs_conversation) + self.assertListEqual([self.gt_multiagent[1]], res) + + # Tools only + res = await fmt.format(self.msgs_tools) + self.assertListEqual( + [ + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # System + tools + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [ + self.gt_multiagent[0], + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # Empty + self.assertListEqual([], await fmt.format([])) + + async def test_chat_formatter_complex_multi_step(self) -> None: + """Complex multi-step sequence with interleaved thinking, text, + tool calls, and tool results.""" + fmt = OllamaChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="thinking_1"), + TextBlock(text="text_1"), + ToolCallBlock( + id="call_1", + name="func_1", + input='{"arg": "value1"}', + ), + ToolCallBlock( + id="call_2", + name="func_2", + input='{"arg": "value2"}', + ), + ToolResultBlock( + id="call_1", + name="func_1", + output=[TextBlock(text="result_1")], + state=ToolResultState.SUCCESS, + ), + ToolResultBlock( + id="call_2", + name="func_2", + output=[TextBlock(text="result_2")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_2"), + TextBlock(text="text_2"), + ToolCallBlock( + id="call_3", + name="func_3", + input='{"arg": "value3"}', + ), + ToolResultBlock( + id="call_3", + name="func_3", + output=[TextBlock(text="result_3")], + state=ToolResultState.SUCCESS, + ), + ToolCallBlock( + id="call_4", + name="func_4", + input='{"arg": "value4"}', + ), + ToolResultBlock( + id="call_4", + name="func_4", + output=[TextBlock(text="result_4")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_3"), + TextBlock(text="text_3"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": "text_1", + "tool_calls": [ + { + "function": { + "name": "func_1", + "arguments": {"arg": "value1"}, + }, + }, + ], + }, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "function": { + "name": "func_2", + "arguments": {"arg": "value2"}, + }, + }, + ], + }, + {"role": "tool", "content": "result_1"}, + {"role": "tool", "content": "result_2"}, + { + "role": "assistant", + "content": "text_2", + "tool_calls": [ + { + "function": { + "name": "func_3", + "arguments": {"arg": "value3"}, + }, + }, + ], + }, + {"role": "tool", "content": "result_3"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "function": { + "name": "func_4", + "arguments": {"arg": "value4"}, + }, + }, + ], + }, + {"role": "tool", "content": "result_4"}, + {"role": "assistant", "content": "text_3"}, + ], + res, + ) + + async def test_chat_formatter_hint_block(self) -> None: + """HintBlock flushes preceding content and becomes a user message.""" + fmt = OllamaChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + TextBlock(text="Let me think about that."), + HintBlock(hint="Remember to be concise."), + TextBlock(text="Here is my answer."), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": "Let me think about that.", + }, + { + "role": "user", + "content": "Remember to be concise.", + }, + { + "role": "assistant", + "content": "Here is my answer.", + }, + ], + res, + ) + + async def test_chat_formatter_hint_block_multimodal(self) -> None: + """Multimodal HintBlock becomes one user message with text + images + list.""" + fmt = OllamaChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + HintBlock( + hint=[ + TextBlock(text="Inspect this screenshot:"), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + ), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "user", + "content": "Inspect this screenshot:", + "images": [self.image_b64], + }, + ], + res, + ) diff --git a/tests/formatter_openai_chat_test.py b/tests/formatter_openai_chat_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ca929e953fa410882a991cde17949053a1cfb09c --- /dev/null +++ b/tests/formatter_openai_chat_test.py @@ -0,0 +1,772 @@ +# -*- coding: utf-8 -*- +"""Comprehensive formatter unit tests for OpenAIChatFormatter and +OpenAIMultiAgentFormatter, following the reference test style with exact +ground-truth comparisons. +""" +from unittest import IsolatedAsyncioTestCase +from unittest.mock import patch + +from agentscope.formatter import ( + OpenAIChatFormatter, + OpenAIMultiAgentFormatter, +) +from agentscope.message import ( + UserMsg, + AssistantMsg, + SystemMsg, + TextBlock, + DataBlock, + ToolCallBlock, + ToolResultBlock, + ToolResultState, + Base64Source, + URLSource, + ThinkingBlock, + HintBlock, +) + + +_FIXED_ID = "TESTID1234567" + + +class TestOpenAIFormatter(IsolatedAsyncioTestCase): + """Comprehensive tests for OpenAI Chat and MultiAgent formatters.""" + + async def asyncSetUp(self) -> None: + """Set up shared message fixtures and expected ground-truth dicts.""" + _img_src = URLSource( + url="https://example.com/image.png", + media_type="image/png", + ) + self.image_url = str(_img_src.url) + + self.image_b64 = "ZmFrZSBpbWFnZSBkYXRh" + self.image_data_uri = f"data:image/png;base64,{self.image_b64}" + + # --------------------------------------------------------------- + # Message fixtures + # (No audio in conversation: OpenAI URL audio requires a download) + # --------------------------------------------------------------- + self.msgs_system = [ + SystemMsg( + name="system", + content="You're a helpful assistant.", + ), + ] + + self.msgs_conversation = [ + UserMsg( + name="user", + content=[ + TextBlock(text="What is the capital of France?"), + DataBlock( + source=URLSource( + url=self.image_url, + media_type="image/png", + ), + ), + ], + ), + AssistantMsg( + name="assistant", + content="The capital of France is Paris.", + ), + UserMsg( + name="user", + content="What is the capital of Germany?", + ), + AssistantMsg( + name="assistant", + content="The capital of Germany is Berlin.", + ), + UserMsg( + name="user", + content="What is the capital of Japan?", + ), + ] + + self.msgs_tools = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_1", + name="get_capital", + input='{"country": "Japan"}', + ), + ToolResultBlock( + id="call_1", + name="get_capital", + output=[ + TextBlock(text="The capital of Japan is Tokyo."), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="The capital of Japan is Tokyo."), + ], + ), + ] + + # --------------------------------------------------------------- + # Ground truth: OpenAIChatFormatter + # - Content is a list of {"type": ..., ...} dicts. + # - Tool-result content is a plain string. + # - Messages have a "name" field. + # --------------------------------------------------------------- + self.gt_chat = [ + { + "role": "system", + "name": "system", + "content": [ + {"type": "text", "text": "You're a helpful assistant."}, + ], + }, + { + "role": "user", + "name": "user", + "content": [ + {"type": "text", "text": "What is the capital of France?"}, + { + "type": "image_url", + "image_url": {"url": self.image_url}, + }, + ], + }, + { + "role": "assistant", + "name": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of France is Paris.", + }, + ], + }, + { + "role": "user", + "name": "user", + "content": [ + { + "type": "text", + "text": "What is the capital of Germany?", + }, + ], + }, + { + "role": "assistant", + "name": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of Germany is Berlin.", + }, + ], + }, + { + "role": "user", + "name": "user", + "content": [ + {"type": "text", "text": "What is the capital of Japan?"}, + ], + }, + { + "role": "assistant", + "name": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_capital", + "arguments": '{"country": "Japan"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "The capital of Japan is Tokyo.", + "name": "get_capital", + }, + { + "role": "assistant", + "name": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of Japan is Tokyo.", + }, + ], + }, + ] + + # --------------------------------------------------------------- + # Ground truth: OpenAIMultiAgentFormatter + # - System content is a plain string. + # - All conversation text is collapsed into a single text block, + # with media blocks appended after. + # - No "name" field on the user wrapper message. + # --------------------------------------------------------------- + _hist_prompt = OpenAIMultiAgentFormatter().conversation_history_prompt + + _conv_text = ( + "user: What is the capital of France?\n" + "assistant: The capital of France is Paris.\n" + "user: What is the capital of Germany?\n" + "assistant: The capital of Germany is Berlin.\n" + "user: What is the capital of Japan?" + ) + + self._gt_trailing_asst = { + "role": "assistant", + "name": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of Japan is Tokyo.", + }, + ], + } + + self._gt_tool_call = { + "role": "assistant", + "name": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_capital", + "arguments": '{"country": "Japan"}', + }, + }, + ], + } + self._gt_tool_result = { + "role": "tool", + "tool_call_id": "call_1", + "content": "The capital of Japan is Tokyo.", + "name": "get_capital", + } + + self.gt_multiagent = [ + { + "role": "system", + "content": "You're a helpful assistant.", + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + _hist_prompt + + "\n" + + _conv_text + + "\n" + ), + }, + { + "type": "image_url", + "image_url": {"url": self.image_url}, + }, + ], + }, + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ] + + # ------------------------------------------------------------------- + # OpenAIChatFormatter tests + # ------------------------------------------------------------------- + + async def test_chat_formatter(self) -> None: + """Chat formatter produces exact output for various subsets.""" + fmt = OpenAIChatFormatter() + + # Full history + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_chat, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_chat[1:], res) + + # Without conversation + n_tools_gt = len(self.gt_chat) - 1 - len(self.msgs_conversation) + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [self.gt_chat[0]] + self.gt_chat[-n_tools_gt:], + res, + ) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_chat[:-n_tools_gt], res) + + # Empty + res = await fmt.format([]) + self.assertListEqual([], res) + + async def test_chat_formatter_base64_image(self) -> None: + """Base64-encoded image is inlined as a data URI.""" + fmt = OpenAIChatFormatter() + msgs = [ + UserMsg( + name="user", + content=[ + TextBlock(text="What's in this image?"), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "user", + "name": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + { + "type": "image_url", + "image_url": {"url": self.image_data_uri}, + }, + ], + }, + ], + res, + ) + + async def test_chat_formatter_thinking_dropped(self) -> None: + """ThinkingBlock is silently dropped by OpenAI formatter.""" + fmt = OpenAIChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="inner thoughts"), + TextBlock(text="reply"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "name": "assistant", + "content": [{"type": "text", "text": "reply"}], + }, + ], + res, + ) + + @patch( + "agentscope.formatter._formatter_base.shortuuid.uuid", + return_value=_FIXED_ID, + ) + async def test_chat_formatter_url_image_in_tool_result( + self, + _mock_uuid: object, + ) -> None: + """URL images in tool results are promoted to a follow-up user + message.""" + fmt = OpenAIChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_img", + name="get_map", + input='{"city": "Tokyo"}', + ), + ToolResultBlock( + id="call_img", + name="get_map", + output=[ + TextBlock(text="Here is the map."), + DataBlock( + source=URLSource( + url=self.image_url, + media_type="image/png", + ), + ), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="Here is the map of Tokyo."), + ], + ), + ] + res = await fmt.format(msgs) + + expected_tool_content = ( + "Here is the map.\n" + f"A(n) image file is returned " + f"and will be presented to you with the identifier " + f"[{_FIXED_ID}]." + ) + self.assertListEqual( + [ + { + "role": "assistant", + "name": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_img", + "type": "function", + "function": { + "name": "get_map", + "arguments": '{"city": "Tokyo"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_img", + "content": expected_tool_content, + "name": "get_map", + }, + { + "role": "user", + "name": "system-reminder", + "content": [ + { + "type": "text", + "text": ( + "The multimodal data " + "and their identifiers are listed as " + "follows:" + ), + }, + { + "type": "text", + "text": f"- {_FIXED_ID} (image file): ", + }, + { + "type": "image_url", + "image_url": {"url": self.image_url}, + }, + { + "type": "text", + "text": "", + }, + ], + }, + { + "role": "assistant", + "name": "assistant", + "content": [ + { + "type": "text", + "text": "Here is the map of Tokyo.", + }, + ], + }, + ], + res, + ) + + # ------------------------------------------------------------------- + # OpenAIMultiAgentFormatter tests + # ------------------------------------------------------------------- + + async def test_multiagent_formatter(self) -> None: + """MultiAgent formatter produces exact output for various subsets.""" + fmt = OpenAIMultiAgentFormatter() + + # Full history + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_multiagent, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_multiagent[1:], res) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_multiagent[:2], res) + + # System only + res = await fmt.format(self.msgs_system) + self.assertListEqual([self.gt_multiagent[0]], res) + + # Conversation only + res = await fmt.format(self.msgs_conversation) + self.assertListEqual([self.gt_multiagent[1]], res) + + # Tools only + res = await fmt.format(self.msgs_tools) + self.assertListEqual( + [ + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # System + tools (no conversation) + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [ + self.gt_multiagent[0], + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # Empty + res = await fmt.format([]) + self.assertListEqual([], res) + + async def test_chat_formatter_complex_multi_step(self) -> None: + """Complex multi-step sequence with interleaved thinking, text, + tool calls, and tool results.""" + fmt = OpenAIChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="thinking_1"), + TextBlock(text="text_1"), + ToolCallBlock( + id="call_1", + name="func_1", + input='{"arg": "value1"}', + ), + ToolCallBlock( + id="call_2", + name="func_2", + input='{"arg": "value2"}', + ), + ToolResultBlock( + id="call_1", + name="func_1", + output=[TextBlock(text="result_1")], + state=ToolResultState.SUCCESS, + ), + ToolResultBlock( + id="call_2", + name="func_2", + output=[TextBlock(text="result_2")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_2"), + TextBlock(text="text_2"), + ToolCallBlock( + id="call_3", + name="func_3", + input='{"arg": "value3"}', + ), + ToolResultBlock( + id="call_3", + name="func_3", + output=[TextBlock(text="result_3")], + state=ToolResultState.SUCCESS, + ), + ToolCallBlock( + id="call_4", + name="func_4", + input='{"arg": "value4"}', + ), + ToolResultBlock( + id="call_4", + name="func_4", + output=[TextBlock(text="result_4")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_3"), + TextBlock(text="text_3"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "name": "assistant", + "content": [{"type": "text", "text": "text_1"}], + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "func_1", + "arguments": '{"arg": "value1"}', + }, + }, + { + "id": "call_2", + "type": "function", + "function": { + "name": "func_2", + "arguments": '{"arg": "value2"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "result_1", + "name": "func_1", + }, + { + "role": "tool", + "tool_call_id": "call_2", + "content": "result_2", + "name": "func_2", + }, + { + "role": "assistant", + "name": "assistant", + "content": [{"type": "text", "text": "text_2"}], + "tool_calls": [ + { + "id": "call_3", + "type": "function", + "function": { + "name": "func_3", + "arguments": '{"arg": "value3"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_3", + "content": "result_3", + "name": "func_3", + }, + { + "role": "assistant", + "name": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_4", + "type": "function", + "function": { + "name": "func_4", + "arguments": '{"arg": "value4"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_4", + "content": "result_4", + "name": "func_4", + }, + { + "role": "assistant", + "name": "assistant", + "content": [{"type": "text", "text": "text_3"}], + }, + ], + res, + ) + + async def test_chat_formatter_hint_block(self) -> None: + """HintBlock flushes preceding content and becomes a user message.""" + fmt = OpenAIChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + TextBlock(text="Let me think about that."), + HintBlock(hint="Remember to be concise."), + TextBlock(text="Here is my answer."), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "name": "assistant", + "content": [ + {"type": "text", "text": "Let me think about that."}, + ], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "Remember to be concise."}, + ], + }, + { + "role": "assistant", + "name": "assistant", + "content": [ + {"type": "text", "text": "Here is my answer."}, + ], + }, + ], + res, + ) + + async def test_chat_formatter_hint_block_multimodal(self) -> None: + """Multimodal HintBlock becomes a single user message with text + + image.""" + fmt = OpenAIChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + HintBlock( + hint=[ + TextBlock(text="Inspect this screenshot:"), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + ), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Inspect this screenshot:", + }, + { + "type": "image_url", + "image_url": {"url": self.image_data_uri}, + }, + ], + }, + ], + res, + ) diff --git a/tests/formatter_openai_response_test.py b/tests/formatter_openai_response_test.py new file mode 100644 index 0000000000000000000000000000000000000000..b5774582d52f0b994e32e4602e36853ce611d227 --- /dev/null +++ b/tests/formatter_openai_response_test.py @@ -0,0 +1,779 @@ +# -*- coding: utf-8 -*- +"""Comprehensive formatter unit tests for OpenAIResponseFormatter and +OpenAIResponseMultiAgentFormatter, following the reference test style with +exact ground-truth comparisons. + +Key differences from OpenAI Chat formatter: + - Text content type is "input_text" (not "text"). + - Image content type is "input_image" with flat "image_url" string. + - Tool calls become top-level "function_call" items (not nested in a msg). + - Tool results become top-level "function_call_output" items. + - ThinkingBlock: only echoed when it has a "reasoning_item_id" attribute. +""" +from unittest import IsolatedAsyncioTestCase +from unittest.mock import patch + +from agentscope.formatter import ( + OpenAIResponseFormatter, + OpenAIResponseMultiAgentFormatter, +) +from agentscope.message import ( + UserMsg, + AssistantMsg, + SystemMsg, + TextBlock, + DataBlock, + ToolCallBlock, + ToolResultBlock, + ToolResultState, + Base64Source, + URLSource, + ThinkingBlock, + HintBlock, +) + + +_FIXED_ID = "TESTID1234567" + + +class TestOpenAIResponseFormatter(IsolatedAsyncioTestCase): + """Comprehensive tests for OpenAI Responses API formatters.""" + + async def asyncSetUp(self) -> None: + """Set up shared message fixtures and expected ground-truth dicts.""" + _img_src = URLSource( + url="https://example.com/image.png", + media_type="image/png", + ) + self.image_url = str(_img_src.url) + + self.image_b64 = "ZmFrZSBpbWFnZSBkYXRh" + self.image_data_uri = f"data:image/png;base64,{self.image_b64}" + + # --------------------------------------------------------------- + # Message fixtures (no audio to avoid downloads) + # --------------------------------------------------------------- + self.msgs_system = [ + SystemMsg( + name="system", + content="You're a helpful assistant.", + ), + ] + + self.msgs_conversation = [ + UserMsg( + name="user", + content=[ + TextBlock(text="What is the capital of France?"), + DataBlock( + source=URLSource( + url=self.image_url, + media_type="image/png", + ), + ), + ], + ), + AssistantMsg( + name="assistant", + content="The capital of France is Paris.", + ), + UserMsg( + name="user", + content="What is the capital of Germany?", + ), + AssistantMsg( + name="assistant", + content="The capital of Germany is Berlin.", + ), + UserMsg( + name="user", + content="What is the capital of Japan?", + ), + ] + + self.msgs_tools = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_1", + name="get_capital", + input='{"country": "Japan"}', + ), + ToolResultBlock( + id="call_1", + name="get_capital", + output=[ + TextBlock(text="The capital of Japan is Tokyo."), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="The capital of Japan is Tokyo."), + ], + ), + ] + + # --------------------------------------------------------------- + # Ground truth: OpenAIResponseFormatter + # - Text: {"type": "input_text", "text": ...} + # - Image: {"type": "input_image", "image_url": url_string} + # - ToolCallBlock → top-level {"type": "function_call", ...} item + # - ToolResultBlock → top-level {"type": "function_call_output", ...} + # - No "name" field on messages. + # --------------------------------------------------------------- + self.gt_chat = [ + { + "role": "system", + "content": [ + { + "type": "input_text", + "text": "You're a helpful assistant.", + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "input_text", + "text": "What is the capital of France?", + }, + {"type": "input_image", "image_url": self.image_url}, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "input_text", + "text": "The capital of France is Paris.", + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "input_text", + "text": "What is the capital of Germany?", + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "input_text", + "text": "The capital of Germany is Berlin.", + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "input_text", + "text": "What is the capital of Japan?", + }, + ], + }, + { + "type": "function_call", + "id": "call_1", + "call_id": "call_1", + "name": "get_capital", + "arguments": '{"country": "Japan"}', + }, + { + "type": "function_call_output", + "call_id": "call_1", + "output": "The capital of Japan is Tokyo.", + }, + { + "role": "assistant", + "content": [ + { + "type": "input_text", + "text": "The capital of Japan is Tokyo.", + }, + ], + }, + ] + + # --------------------------------------------------------------- + # Ground truth: OpenAIResponseMultiAgentFormatter + # - System: {"role": "system", "content": plain_string} + # - Conversation history: input_text with history wrapping. + # - Tool sequences use OpenAIResponseFormatter (function_call / + # function_call_output top-level items). + # --------------------------------------------------------------- + _hist_prompt = ( + OpenAIResponseMultiAgentFormatter().conversation_history_prompt + ) + + _conv_text = ( + "user: What is the capital of France?\n" + "assistant: The capital of France is Paris.\n" + "user: What is the capital of Germany?\n" + "assistant: The capital of Germany is Berlin.\n" + "user: What is the capital of Japan?" + ) + + self._gt_trailing_asst = { + "role": "assistant", + "content": [ + { + "type": "input_text", + "text": "The capital of Japan is Tokyo.", + }, + ], + } + + self._gt_tool_call = { + "type": "function_call", + "id": "call_1", + "call_id": "call_1", + "name": "get_capital", + "arguments": '{"country": "Japan"}', + } + self._gt_tool_result = { + "type": "function_call_output", + "call_id": "call_1", + "output": "The capital of Japan is Tokyo.", + } + + self.gt_multiagent = [ + { + "role": "system", + "content": "You're a helpful assistant.", + }, + { + "role": "user", + "content": [ + { + "type": "input_text", + "text": ( + _hist_prompt + + "\n" + + _conv_text + + "\n" + ), + }, + {"type": "input_image", "image_url": self.image_url}, + ], + }, + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ] + + # ------------------------------------------------------------------- + # OpenAIResponseFormatter tests + # ------------------------------------------------------------------- + + async def test_chat_formatter(self) -> None: + """Chat formatter produces exact output for various subsets.""" + fmt = OpenAIResponseFormatter() + + # Full history + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_chat, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_chat[1:], res) + + # Without conversation + n_tools_gt = len(self.gt_chat) - 1 - len(self.msgs_conversation) + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [self.gt_chat[0]] + self.gt_chat[-n_tools_gt:], + res, + ) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_chat[:-n_tools_gt], res) + + # Empty + res = await fmt.format([]) + self.assertListEqual([], res) + + async def test_chat_formatter_base64_image(self) -> None: + """Base64-encoded image becomes an input_image item with data URI.""" + fmt = OpenAIResponseFormatter() + msgs = [ + UserMsg( + name="user", + content=[ + TextBlock(text="What's in this image?"), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "user", + "content": [ + { + "type": "input_text", + "text": "What's in this image?", + }, + { + "type": "input_image", + "image_url": self.image_data_uri, + }, + ], + }, + ], + res, + ) + + async def test_chat_formatter_thinking_dropped_without_reasoning_item_id( + self, + ) -> None: + """ThinkingBlock without reasoning_item_id is silently skipped.""" + fmt = OpenAIResponseFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="inner thoughts"), + TextBlock(text="reply"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": [ + {"type": "input_text", "text": "reply"}, + ], + }, + ], + res, + ) + + async def test_chat_formatter_thinking_echoed_with_reasoning_item_id( + self, + ) -> None: + """ThinkingBlock with reasoning_item_id is echoed as a reasoning + item.""" + fmt = OpenAIResponseFormatter() + thinking = ThinkingBlock(thinking="my reasoning") + thinking.reasoning_item_id = "rs_001" + msgs = [ + AssistantMsg( + name="assistant", + content=[thinking, TextBlock(text="reply")], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "type": "reasoning", + "id": "rs_001", + "summary": [ + {"type": "summary_text", "text": "my reasoning"}, + ], + "content": [], + }, + { + "role": "assistant", + "content": [ + {"type": "input_text", "text": "reply"}, + ], + }, + ], + res, + ) + + @patch( + "agentscope.formatter._formatter_base.shortuuid.uuid", + return_value=_FIXED_ID, + ) + async def test_chat_formatter_url_image_in_tool_result( + self, + _mock_uuid: object, + ) -> None: + """URL images in tool results are promoted to a follow-up user + message.""" + fmt = OpenAIResponseFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_img", + name="get_map", + input='{"city": "Tokyo"}', + ), + ToolResultBlock( + id="call_img", + name="get_map", + output=[ + TextBlock(text="Here is the map."), + DataBlock( + source=URLSource( + url=self.image_url, + media_type="image/png", + ), + ), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="Here is the map of Tokyo."), + ], + ), + ] + res = await fmt.format(msgs) + + expected_tool_content = ( + "Here is the map.\n" + f"A(n) image file is returned " + f"and will be presented to you with the identifier " + f"[{_FIXED_ID}]." + ) + self.assertListEqual( + [ + { + "type": "function_call", + "id": "call_img", + "call_id": "call_img", + "name": "get_map", + "arguments": '{"city": "Tokyo"}', + }, + { + "type": "function_call_output", + "call_id": "call_img", + "output": expected_tool_content, + }, + { + "role": "user", + "content": [ + { + "type": "input_text", + "text": ( + "The multimodal data " + "and their identifiers are listed as " + "follows:" + ), + }, + { + "type": "input_text", + "text": f"- {_FIXED_ID} (image file): ", + }, + { + "type": "input_image", + "image_url": self.image_url, + }, + { + "type": "input_text", + "text": "", + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "input_text", + "text": "Here is the map of Tokyo.", + }, + ], + }, + ], + res, + ) + + # ------------------------------------------------------------------- + # OpenAIResponseMultiAgentFormatter tests + # ------------------------------------------------------------------- + + async def test_multiagent_formatter(self) -> None: + """MultiAgent formatter produces exact output for various subsets.""" + fmt = OpenAIResponseMultiAgentFormatter() + + # Full history + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + self.assertListEqual(self.gt_multiagent, res) + + # Without system + res = await fmt.format([*self.msgs_conversation, *self.msgs_tools]) + self.assertListEqual(self.gt_multiagent[1:], res) + + # Without tools + res = await fmt.format([*self.msgs_system, *self.msgs_conversation]) + self.assertListEqual(self.gt_multiagent[:2], res) + + # System only + res = await fmt.format(self.msgs_system) + self.assertListEqual([self.gt_multiagent[0]], res) + + # Conversation only + res = await fmt.format(self.msgs_conversation) + self.assertListEqual([self.gt_multiagent[1]], res) + + # Tools only + res = await fmt.format(self.msgs_tools) + self.assertListEqual( + [ + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # System + tools + res = await fmt.format([*self.msgs_system, *self.msgs_tools]) + self.assertListEqual( + [ + self.gt_multiagent[0], + self._gt_tool_call, + self._gt_tool_result, + self._gt_trailing_asst, + ], + res, + ) + + # Empty + res = await fmt.format([]) + self.assertListEqual([], res) + + async def test_chat_formatter_complex_multi_step(self) -> None: + """Complex multi-step sequence with interleaved thinking, text, + tool calls, and tool results.""" + fmt = OpenAIResponseFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="thinking_1"), + TextBlock(text="text_1"), + ToolCallBlock( + id="call_1", + name="func_1", + input='{"arg": "value1"}', + ), + ToolCallBlock( + id="call_2", + name="func_2", + input='{"arg": "value2"}', + ), + ToolResultBlock( + id="call_1", + name="func_1", + output=[TextBlock(text="result_1")], + state=ToolResultState.SUCCESS, + ), + ToolResultBlock( + id="call_2", + name="func_2", + output=[TextBlock(text="result_2")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_2"), + TextBlock(text="text_2"), + ToolCallBlock( + id="call_3", + name="func_3", + input='{"arg": "value3"}', + ), + ToolResultBlock( + id="call_3", + name="func_3", + output=[TextBlock(text="result_3")], + state=ToolResultState.SUCCESS, + ), + ToolCallBlock( + id="call_4", + name="func_4", + input='{"arg": "value4"}', + ), + ToolResultBlock( + id="call_4", + name="func_4", + output=[TextBlock(text="result_4")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_3"), + TextBlock(text="text_3"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": [ + {"type": "input_text", "text": "text_1"}, + ], + }, + { + "type": "function_call", + "id": "call_1", + "call_id": "call_1", + "name": "func_1", + "arguments": '{"arg": "value1"}', + }, + { + "type": "function_call", + "id": "call_2", + "call_id": "call_2", + "name": "func_2", + "arguments": '{"arg": "value2"}', + }, + { + "type": "function_call_output", + "call_id": "call_1", + "output": "result_1", + }, + { + "type": "function_call_output", + "call_id": "call_2", + "output": "result_2", + }, + { + "role": "assistant", + "content": [ + {"type": "input_text", "text": "text_2"}, + ], + }, + { + "type": "function_call", + "id": "call_3", + "call_id": "call_3", + "name": "func_3", + "arguments": '{"arg": "value3"}', + }, + { + "type": "function_call_output", + "call_id": "call_3", + "output": "result_3", + }, + { + "type": "function_call", + "id": "call_4", + "call_id": "call_4", + "name": "func_4", + "arguments": '{"arg": "value4"}', + }, + { + "type": "function_call_output", + "call_id": "call_4", + "output": "result_4", + }, + { + "role": "assistant", + "content": [ + {"type": "input_text", "text": "text_3"}, + ], + }, + ], + res, + ) + + async def test_chat_formatter_hint_block(self) -> None: + """HintBlock flushes preceding content and becomes a user message.""" + fmt = OpenAIResponseFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + TextBlock(text="Let me think about that."), + HintBlock(hint="Remember to be concise."), + TextBlock(text="Here is my answer."), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "assistant", + "content": [ + { + "type": "input_text", + "text": "Let me think about that.", + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "input_text", + "text": "Remember to be concise.", + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "input_text", + "text": "Here is my answer.", + }, + ], + }, + ], + res, + ) + + async def test_chat_formatter_hint_block_multimodal(self) -> None: + """Multimodal HintBlock becomes a single user message with text + + image.""" + fmt = OpenAIResponseFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + HintBlock( + hint=[ + TextBlock(text="Inspect this screenshot:"), + DataBlock( + source=Base64Source( + data=self.image_b64, + media_type="image/png", + ), + ), + ], + ), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + { + "role": "user", + "content": [ + { + "type": "input_text", + "text": "Inspect this screenshot:", + }, + { + "type": "input_image", + "image_url": self.image_data_uri, + }, + ], + }, + ], + res, + ) diff --git a/tests/formatter_xai_test.py b/tests/formatter_xai_test.py new file mode 100644 index 0000000000000000000000000000000000000000..c610266be974f0fa98616094753f9c91d0e8067e --- /dev/null +++ b/tests/formatter_xai_test.py @@ -0,0 +1,596 @@ +# -*- coding: utf-8 -*- +"""Comprehensive formatter unit tests for XAIChatFormatter and +XAIMultiAgentFormatter (xAI), following the reference test style. + +Because these formatters return xai_sdk protobuf Message objects (not plain +dicts), a lightweight xai_sdk stub is built at module load so that tests run +without the real package. The stub objects support __eq__ and __repr__ so +full assertListEqual comparisons work. +""" +import sys +from typing import Any +from types import ModuleType +from unittest import IsolatedAsyncioTestCase +from unittest.mock import MagicMock + +from agentscope.formatter import XAIChatFormatter, XAIMultiAgentFormatter +from agentscope.message import ( + UserMsg, + AssistantMsg, + SystemMsg, + TextBlock, + DataBlock, + Base64Source, + ToolCallBlock, + ToolResultBlock, + ThinkingBlock, + ToolResultState, + HintBlock, +) + + +# --------------------------------------------------------------------------- +# Comparable stub objects for xai_sdk protobuf messages. +# --------------------------------------------------------------------------- + + +class _StubMessage: + """Comparable stub for xai_sdk.chat.{user,assistant,system,tool_result}.""" + + def __init__( + self, + role: str, + args: tuple = (), + kwargs: dict | None = None, + ) -> None: + self.role = role + self.args = args + self.kwargs = kwargs or {} + + def __eq__(self, other: object) -> bool: + if not isinstance(other, _StubMessage): + return NotImplemented + return ( + self.role == other.role + and self.args == other.args + and self.kwargs == other.kwargs + ) + + def __repr__(self) -> str: + parts = [f"role={self.role!r}", f"args={self.args!r}"] + if self.kwargs: + parts.append(f"kwargs={self.kwargs!r}") + return f"_StubMessage({', '.join(parts)})" + + +class _StubImage: + """Comparable stub for xai_sdk.chat.image().""" + + def __init__(self, url: str) -> None: + self.type = "image" + self.url = url + + def __eq__(self, other: object) -> bool: + if not isinstance(other, _StubImage): + return NotImplemented + return self.url == other.url + + def __repr__(self) -> str: + return f"_StubImage(url={self.url!r})" + + +def user(*args: Any) -> _StubMessage: + """Create a comparable stub user message.""" + return _StubMessage(role="user", args=args) + + +def assistant(*args: Any) -> _StubMessage: + """Create a comparable stub assistant message.""" + return _StubMessage(role="assistant", args=args) + + +def system(*args: Any) -> _StubMessage: + """Create a comparable stub system message.""" + return _StubMessage(role="system", args=args) + + +def tool_result(*args: Any, **kwargs: Any) -> _StubMessage: + """Create a comparable stub tool_result message.""" + return _StubMessage(role="tool", args=args, kwargs=kwargs) + + +def image(url: str) -> _StubImage: + """Create a comparable stub image object.""" + return _StubImage(url=url) + + +# --------------------------------------------------------------------------- +# Build a lightweight xai_sdk stub so tests run without the real package. +# --------------------------------------------------------------------------- + + +def _build_xai_sdk_stub() -> None: + if "xai_sdk" in sys.modules: + return + + chat_pb2 = ModuleType("xai_sdk.chat.chat_pb2") + + class _EnumHelper: + _mapping = { + "ROLE_ASSISTANT": 2, + "TOOL_CALL_TYPE_CLIENT_SIDE_TOOL": 1, + } + + def Value(self, name: str) -> int: + """Return the integer value for the given enum name.""" + return self._mapping.get(name, 0) + + chat_pb2.MessageRole = _EnumHelper() + chat_pb2.ToolCallType = _EnumHelper() + + class _RepeatedField(list): + def __init__(self, factory: Any) -> None: + super().__init__() + self._factory = factory + + def add(self) -> Any: + """Add a new item using the factory and return it.""" + item = self._factory() + self.append(item) + return item + + class _FunctionSpec: + def __init__(self) -> None: + self.name = "" + self.arguments = "" + + class _ToolCallProto: + def __init__(self) -> None: + self.id = "" + self.type = 0 + self.function = _FunctionSpec() + + class _ContentPart: + text: str = "" + + class _MessageProto: + def __init__(self) -> None: + self.role = 0 + self.content = _RepeatedField(_ContentPart) + self.tool_calls = _RepeatedField(_ToolCallProto) + + chat_pb2.Message = _MessageProto + + xai_chat = ModuleType("xai_sdk.chat") + xai_chat.chat_pb2 = chat_pb2 + xai_chat.user = user + xai_chat.assistant = assistant + xai_chat.system = system + xai_chat.tool_result = tool_result + xai_chat.image = image + + xai_sdk = ModuleType("xai_sdk") + xai_sdk.chat = xai_chat + xai_sdk.AsyncClient = MagicMock() + + sys.modules["xai_sdk"] = xai_sdk + sys.modules["xai_sdk.chat"] = xai_chat + sys.modules["xai_sdk.chat.chat_pb2"] = chat_pb2 + + +_build_xai_sdk_stub() + + +class TestXAIFormatter(IsolatedAsyncioTestCase): + """Comprehensive tests for XAI Chat and MultiAgent formatters. + + The stub objects support __eq__, so full assertListEqual works for + user/assistant/system/tool_result messages. Tool-call messages use + _MessageProto which is checked via attribute assertions. + """ + + async def asyncSetUp(self) -> None: + self.msgs_system = [ + SystemMsg( + name="system", + content="You're a helpful assistant.", + ), + ] + + self.msgs_conversation = [ + UserMsg( + name="user", + content="What is the capital of France?", + ), + AssistantMsg( + name="assistant", + content="The capital of France is Paris.", + ), + UserMsg( + name="user", + content="What is the capital of Germany?", + ), + AssistantMsg( + name="assistant", + content="The capital of Germany is Berlin.", + ), + UserMsg( + name="user", + content="What is the capital of Japan?", + ), + ] + + self.msgs_tools = [ + AssistantMsg( + name="assistant", + content=[ + ToolCallBlock( + id="call_1", + name="get_capital", + input='{"country": "Japan"}', + ), + ToolResultBlock( + id="call_1", + name="get_capital", + output=[ + TextBlock(text="The capital of Japan is Tokyo."), + ], + state=ToolResultState.SUCCESS, + ), + TextBlock(text="The capital of Japan is Tokyo."), + ], + ), + ] + + self._hist_prompt = ( + XAIMultiAgentFormatter().conversation_history_prompt + ) + + # ------------------------------------------------------------------- + # XAIChatFormatter tests + # ------------------------------------------------------------------- + + async def test_chat_formatter_system_message(self) -> None: + """System message becomes a system() stub.""" + fmt = XAIChatFormatter() + res = await fmt.format(self.msgs_system) + self.assertListEqual( + [system("You're a helpful assistant.")], + res, + ) + + async def test_chat_formatter_user_assistant(self) -> None: + """User and assistant text messages are passed through correctly.""" + fmt = XAIChatFormatter() + res = await fmt.format(self.msgs_conversation) + self.assertListEqual( + [ + user("What is the capital of France?"), + assistant("The capital of France is Paris."), + user("What is the capital of Germany?"), + assistant("The capital of Germany is Berlin."), + user("What is the capital of Japan?"), + ], + res, + ) + + async def test_chat_formatter_tool_call(self) -> None: + """Assistant tool call becomes a _MessageProto with tool_calls set.""" + fmt = XAIChatFormatter() + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + # Find the tool call proto (has tool_calls list, role=ROLE_ASSISTANT=2) + tool_call_msgs = [ + m + for m in res + if hasattr(m, "tool_calls") and len(m.tool_calls) > 0 + ] + self.assertListEqual( + [len(tool_call_msgs)], + [1], + ) + tc = tool_call_msgs[0].tool_calls[0] + self.assertListEqual( + [tc.id, tc.function.name, tc.function.arguments], + ["call_1", "get_capital", '{"country": "Japan"}'], + ) + + async def test_chat_formatter_tool_result(self) -> None: + """Tool result becomes a tool_result() stub with the right id.""" + fmt = XAIChatFormatter() + res = await fmt.format(self.msgs_tools) + tool_msgs = [m for m in res if m.role == "tool"] + self.assertListEqual( + tool_msgs, + [ + tool_result( + "The capital of Japan is Tokyo.", + tool_call_id="call_1", + ), + ], + ) + + async def test_chat_formatter_thinking_dropped(self) -> None: + """ThinkingBlock is silently ignored in user/assistant xAI + messages.""" + fmt = XAIChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="inner thoughts"), + TextBlock(text="reply"), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual([assistant("reply")], res) + + async def test_chat_formatter_empty(self) -> None: + """Empty input returns empty list.""" + fmt = XAIChatFormatter() + res = await fmt.format([]) + self.assertListEqual([], res) + + # ------------------------------------------------------------------- + # XAIMultiAgentFormatter tests + # ------------------------------------------------------------------- + + async def test_multiagent_formatter_system_message(self) -> None: + """System message is passed through as a system() stub.""" + fmt = XAIMultiAgentFormatter() + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation], + ) + self.assertEqual(res[0], system("You're a helpful assistant.")) + + async def test_multiagent_formatter_conversation_history(self) -> None: + """Non-tool agent messages are collapsed into a user() history + stub.""" + fmt = XAIMultiAgentFormatter() + res = await fmt.format(self.msgs_conversation) + self.assertListEqual( + [ + user( + self._hist_prompt + "\n" + "user: What is the capital of France?\n" + "assistant: The capital of France is Paris.\n" + "user: What is the capital of Germany?\n" + "assistant: The capital of Germany is Berlin.\n" + "user: What is the capital of Japan?\n" + "", + ), + ], + res, + ) + + async def test_multiagent_formatter_first_group_has_hist_prompt( + self, + ) -> None: + """First agent message group includes the conversation history + prompt.""" + fmt = XAIMultiAgentFormatter() + res = await fmt.format(self.msgs_conversation) + hist_text = res[0].args[0] + self.assertListEqual( + [hist_text.startswith(self._hist_prompt)], + [True], + ) + + async def test_multiagent_formatter_full_history(self) -> None: + """Full history produces system + conv_history + tool_call + + tool_result + trailing assistant.""" + fmt = XAIMultiAgentFormatter() + res = await fmt.format( + [*self.msgs_system, *self.msgs_conversation, *self.msgs_tools], + ) + roles = [m.role for m in res] + self.assertListEqual( + sorted(r for r in roles if isinstance(r, str)), + sorted(["system", "user", "tool", "assistant"]), + ) + + async def test_multiagent_formatter_tools_only_is_first(self) -> None: + """When only tools are given, trailing text is formatted as + assistant.""" + fmt = XAIMultiAgentFormatter() + res = await fmt.format(self.msgs_tools) + trailing = [m for m in res if m.role == "assistant"] + self.assertListEqual([len(trailing)], [1]) + + async def test_multiagent_formatter_nonfirst_trailing_is_assistant( + self, + ) -> None: + """Trailing text after a tool sequence is formatted as assistant.""" + fmt = XAIMultiAgentFormatter() + res = await fmt.format( + [*self.msgs_conversation, *self.msgs_tools], + ) + trailing = [m for m in res if m.role == "assistant"] + self.assertListEqual([len(trailing)], [1]) + + async def test_multiagent_formatter_empty(self) -> None: + """Empty input returns empty list.""" + fmt = XAIMultiAgentFormatter() + res = await fmt.format([]) + self.assertListEqual([], res) + + async def test_chat_formatter_complex_multi_step(self) -> None: + """Complex multi-step sequence with interleaved thinking, text, + tool calls, and tool results.""" + fmt = XAIChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + ThinkingBlock(thinking="thinking_1"), + TextBlock(text="text_1"), + ToolCallBlock( + id="call_1", + name="func_1", + input='{"arg": "value1"}', + ), + ToolCallBlock( + id="call_2", + name="func_2", + input='{"arg": "value2"}', + ), + ToolResultBlock( + id="call_1", + name="func_1", + output=[TextBlock(text="result_1")], + state=ToolResultState.SUCCESS, + ), + ToolResultBlock( + id="call_2", + name="func_2", + output=[TextBlock(text="result_2")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_2"), + TextBlock(text="text_2"), + ToolCallBlock( + id="call_3", + name="func_3", + input='{"arg": "value3"}', + ), + ToolResultBlock( + id="call_3", + name="func_3", + output=[TextBlock(text="result_3")], + state=ToolResultState.SUCCESS, + ), + ToolCallBlock( + id="call_4", + name="func_4", + input='{"arg": "value4"}', + ), + ToolResultBlock( + id="call_4", + name="func_4", + output=[TextBlock(text="result_4")], + state=ToolResultState.SUCCESS, + ), + ThinkingBlock(thinking="thinking_3"), + TextBlock(text="text_3"), + ], + ), + ] + res = await fmt.format(msgs) + + # First message: proto with text_1 content + 2 tool_calls + self.assertTrue(hasattr(res[0], "tool_calls")) + self.assertEqual(len(res[0].tool_calls), 2) + self.assertEqual(len(res[0].content), 1) + self.assertEqual(res[0].tool_calls[0].id, "call_1") + self.assertEqual(res[0].tool_calls[0].function.name, "func_1") + self.assertEqual( + res[0].tool_calls[0].function.arguments, + '{"arg": "value1"}', + ) + self.assertEqual(res[0].tool_calls[1].id, "call_2") + self.assertEqual(res[0].tool_calls[1].function.name, "func_2") + self.assertEqual( + res[0].tool_calls[1].function.arguments, + '{"arg": "value2"}', + ) + + # Tool results + self.assertEqual( + res[1], + tool_result("result_1", tool_call_id="call_1"), + ) + self.assertEqual( + res[2], + tool_result("result_2", tool_call_id="call_2"), + ) + + # Second proto: text_2 + 1 tool_call + self.assertTrue(hasattr(res[3], "tool_calls")) + self.assertEqual(len(res[3].tool_calls), 1) + self.assertEqual(len(res[3].content), 1) + self.assertEqual(res[3].tool_calls[0].id, "call_3") + self.assertEqual(res[3].tool_calls[0].function.name, "func_3") + + # Tool result for call_3 + self.assertEqual( + res[4], + tool_result("result_3", tool_call_id="call_3"), + ) + + # Third proto: no content + 1 tool_call + self.assertTrue(hasattr(res[5], "tool_calls")) + self.assertEqual(len(res[5].tool_calls), 1) + self.assertEqual(len(res[5].content), 0) + self.assertEqual(res[5].tool_calls[0].id, "call_4") + self.assertEqual(res[5].tool_calls[0].function.name, "func_4") + + # Tool result for call_4 + self.assertEqual( + res[6], + tool_result("result_4", tool_call_id="call_4"), + ) + + # Final assistant text + self.assertEqual(res[7], assistant("text_3")) + + # Total 8 items + self.assertEqual(len(res), 8) + + async def test_chat_formatter_hint_block(self) -> None: + """HintBlock flushes preceding content and becomes a user message.""" + fmt = XAIChatFormatter() + msgs = [ + AssistantMsg( + name="assistant", + content=[ + TextBlock(text="Let me think about that."), + HintBlock(hint="Remember to be concise."), + TextBlock(text="Here is my answer."), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + assistant("Let me think about that."), + user("Remember to be concise."), + assistant("Here is my answer."), + ], + res, + ) + + async def test_chat_formatter_hint_block_multimodal(self) -> None: + """Multimodal HintBlock becomes one user() stub carrying text + + image.""" + fmt = XAIChatFormatter() + image_b64 = "ZmFrZSBpbWFnZSBkYXRh" + msgs = [ + AssistantMsg( + name="assistant", + content=[ + HintBlock( + hint=[ + TextBlock(text="Inspect this screenshot:"), + DataBlock( + source=Base64Source( + data=image_b64, + media_type="image/png", + ), + ), + ], + ), + ], + ), + ] + res = await fmt.format(msgs) + self.assertListEqual( + [ + user( + "Inspect this screenshot:", + image(f"data:image/png;base64,{image_b64}"), + ), + ], + res, + ) diff --git a/tests/hitl_external_execution_test.py b/tests/hitl_external_execution_test.py new file mode 100644 index 0000000000000000000000000000000000000000..41c4212b49df54a4b112d0f690c6956741989efe --- /dev/null +++ b/tests/hitl_external_execution_test.py @@ -0,0 +1,1344 @@ +# -*- coding: utf-8 -*- +# pylint: disable=redefined-builtin +"""Test the external execution events in the agent class.""" +from typing import Any +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import AnyString, MockModel + +from agentscope.agent import Agent +from agentscope.model import ChatResponse +from agentscope.tool import ( + ToolBase, + Toolkit, + ToolChunk, +) +from agentscope.permission import ( + PermissionDecision, + PermissionBehavior, + PermissionContext, +) +from agentscope.message import ( + TextBlock, + ToolCallBlock, + ToolResultBlock, + UserMsg, + ToolResultState, +) +from agentscope.event import ExternalExecutionResultEvent + + +class MockExternalSequentialTool(ToolBase): + """A mock tool that requires external execution (sequential).""" + + name: str = "mock_external_sequential_tool" + description: str = "A mock external sequential tool for testing" + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "input": {"type": "string", "description": "Input string"}, + }, + "required": ["input"], + } + is_concurrency_safe: bool = False + is_read_only: bool = True + is_external_tool: bool = True + is_mcp: bool = False + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permissions for the tool usage.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + decision_reason="Mock external tool always allows", + message="Mock external tool always allows", + ) + + async def __call__(self, input: str, **kwargs: Any) -> ToolChunk: + """Execute the tool.""" + return ToolChunk( + content=[TextBlock(text=f"External sequential result: {input}")], + ) + + +class MockExternalConcurrentTool(ToolBase): + """A mock tool that requires external execution (concurrent).""" + + name: str = "mock_external_concurrent_tool" + description: str = "A mock external concurrent tool for testing" + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "input": {"type": "string", "description": "Input string"}, + }, + "required": ["input"], + } + is_concurrency_safe: bool = True + is_read_only: bool = True + is_external_tool: bool = True + is_mcp: bool = False + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permissions for the tool usage.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + decision_reason="Mock external tool always allows", + message="Mock external tool always allows", + ) + + async def __call__(self, input: str, **kwargs: Any) -> ToolChunk: + """Execute the tool.""" + return ToolChunk( + content=[TextBlock(text=f"External concurrent result: {input}")], + ) + + +class AgentExternalExecutionTest(IsolatedAsyncioTestCase): + """Test the external execution events in the agent class.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.model = MockModel() + self.agent = Agent( + name="Friday", + system_prompt="You are a helpful assistant.", + model=self.model, + toolkit=Toolkit(), + ) + self.tool_call_id_1 = "tool_call_1" + self.tool_call_id_2 = "tool_call_2" + self.user_input_text = "Test" + self.tool_input_1 = '{"input": "test1"}' + self.tool_input_2 = '{"input": "test2"}' + self.sequential_tool_name = "mock_external_sequential_tool" + self.concurrent_tool_name = "mock_external_concurrent_tool" + self.sequential_result_1 = "External sequential result: test1" + self.sequential_result_2 = "External sequential result: test2" + self.concurrent_result_1 = "External concurrent result: test1" + self.concurrent_result_2 = "External concurrent result: test2" + self.final_response_text = "Final response after external execution" + self.final_text_events = [ + { + "type": "MODEL_CALL_START", + "model_name": "mock-model", + }, + { + "type": "TEXT_BLOCK_START", + "block_id": AnyString(), + }, + { + "type": "TEXT_BLOCK_DELTA", + "block_id": AnyString(), + "delta": self.final_response_text, + }, + { + "type": "TEXT_BLOCK_END", + "block_id": AnyString(), + }, + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + ] + self.final_mock_responses = [ + ChatResponse( + content=[ + TextBlock(text=self.final_response_text), + ], + is_last=False, + usage=None, + ), + ChatResponse( + content=[ + TextBlock(text=self.final_response_text), + ], + is_last=True, + usage=None, + ), + ] + + def _get_event_base(self, reply_id: str) -> dict: + """Get the dict with the basic fields for event assertion.""" + return { + "id": AnyString(), + "created_at": AnyString(), + "metadata": {}, + "reply_id": reply_id, + } + + def _get_msg_base(self) -> dict: + """Get the dict with the basic fields for message assertion.""" + return { + "id": AnyString(), + "created_at": AnyString(), + "finished_at": None, + "metadata": {}, + "name": "Friday", + "role": "assistant", + "usage": None, + } + + def _get_tool_call_events( + self, + id: str, + name: str, + delta: str, + ) -> list[dict]: + """Helper method to get the expected tool call events.""" + return [ + { + "type": "TOOL_CALL_START", + "tool_call_id": id, + "tool_call_name": name, + }, + { + "type": "TOOL_CALL_DELTA", + "tool_call_id": id, + "delta": delta, + }, + { + "type": "TOOL_CALL_END", + "tool_call_id": id, + }, + ] + + def _get_tool_result_events( + self, + id: str, + result: str, + state: str = "success", + ) -> list[dict]: + """Helper method to get the expected tool result events.""" + return [ + { + "type": "TOOL_RESULT_TEXT_DELTA", + "tool_call_id": id, + "delta": result, + }, + { + "type": "TOOL_RESULT_END", + "tool_call_id": id, + "state": state, + }, + ] + + def _get_require_external_execution_events( + self, + reply_id: str, + id: str, + name: str, + tool_input: str, + ) -> list[dict]: + """Helper method to get the expected external execution events.""" + return [ + { + "type": "TOOL_RESULT_START", + "tool_call_id": id, + "tool_call_name": name, + }, + { + "type": "REQUIRE_EXTERNAL_EXECUTION", + "reply_id": reply_id, + "tool_calls": [ + { + "type": "tool_call", + "id": id, + "name": name, + "input": tool_input, + "state": "submitted", + "suggested_rules": [], + }, + ], + }, + ] + + async def test_single_external_execution(self) -> None: + """Test single external execution tool call. + + The agent should: + 1. Generate a tool call that requires external execution + 2. Emit REQUIRE_EXTERNAL_EXECUTION event and pause + 3. Resume when ExternalExecutionResultEvent is provided + 4. Continue execution without calling the model again + """ + # Register external tool + ext_tool = MockExternalSequentialTool() + self.agent.toolkit = Toolkit( + tools=[ext_tool], + ) + + # Set up mock response with tool call (no final text response) + self.model.set_responses( + [ + [ + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.sequential_tool_name, + input=self.tool_input_1, + ), + ], + is_last=False, + usage=None, + ), + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.sequential_tool_name, + input=self.tool_input_1, + ), + ], + is_last=True, + usage=None, + ), + ], + self.final_mock_responses, + ], + ) + + # First call: collect events until REQUIRE_EXTERNAL_EXECUTION + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content=self.user_input_text), + ): + events.append(event.model_dump()) + + # Verify events + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + + expected_events = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + *self._get_tool_call_events( + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ), + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + *self._get_require_external_execution_events( + reply_id, + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ), + ] + + basic_dict = self._get_event_base(reply_id) + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + # Assert context after first call + msg_base = self._get_msg_base() + expected_context = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "submitted", + "suggested_rules": [], + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + # Create external execution result event + external_result_event = ExternalExecutionResultEvent( + reply_id=reply_id, + execution_results=[ + ToolResultBlock( + id=self.tool_call_id_1, + name=self.sequential_tool_name, + output=[ + TextBlock(text=self.sequential_result_1), + ], + state=ToolResultState.SUCCESS, + ), + ], + ) + + # Second call: resume with external execution result + events = [] + async for event in self.agent.reply_stream( + inputs=external_result_event, + ): + events.append(event.model_dump()) + + # Verify events after resumption + expected_events_resume = [ + *self._get_tool_result_events( + self.tool_call_id_1, + self.sequential_result_1, + ), + *self.final_text_events, + {"type": "REPLY_END", "session_id": session_id}, + ] + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events_resume], + ) + + # Assert final context + expected_context_final = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "finished", + "suggested_rules": [], + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.sequential_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.sequential_result_1, + }, + ], + "state": "success", + "metadata": {}, + }, + { + "type": "text", + "id": AnyString(), + "text": self.final_response_text, + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context_final = [ + {**msg_base, **_} for _ in expected_context_final + ] + self.assertListEqual(context_dicts, expected_context_final) + + async def test_sequential_external_execution(self) -> None: + """Test multiple external execution tool calls in sequential execution. + + The agent should: + 1. Generate multiple tool calls that require external execution + 2. All tools have is_concurrent_safe=False (sequential) + 3. Emit REQUIRE_EXTERNAL_EXECUTION event and pause + 4. Resume when ExternalExecutionResultEvent is provided + 5. Continue execution without calling the model again + """ + # Register external sequential tool + ext_tool = MockExternalSequentialTool() + self.agent.toolkit = Toolkit( + tools=[ext_tool], + ) + + # Set up mock response with multiple tool calls + self.model.set_responses( + [ + [ + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.sequential_tool_name, + input=self.tool_input_1, + ), + ToolCallBlock( + id=self.tool_call_id_2, + name=self.sequential_tool_name, + input=self.tool_input_2, + ), + ], + is_last=False, + usage=None, + ), + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.sequential_tool_name, + input=self.tool_input_1, + ), + ToolCallBlock( + id=self.tool_call_id_2, + name=self.sequential_tool_name, + input=self.tool_input_2, + ), + ], + is_last=True, + usage=None, + ), + ], + self.final_mock_responses, + ], + ) + + # First call: collect events until REQUIRE_EXTERNAL_EXECUTION + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content=self.user_input_text), + ): + events.append(event.model_dump()) + + # Verify events + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + tool_call_1_events = self._get_tool_call_events( + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ) + tool_call_2_events = self._get_tool_call_events( + self.tool_call_id_2, + self.sequential_tool_name, + self.tool_input_2, + ) + + expected_events = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + *tool_call_1_events[:2], + *tool_call_2_events[:2], + tool_call_1_events[2], + tool_call_2_events[2], + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + *self._get_require_external_execution_events( + reply_id, + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ), + ] + + basic_dict = self._get_event_base(reply_id) + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + # Assert context after first call + msg_base = self._get_msg_base() + expected_context = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "submitted", + "suggested_rules": [], + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.sequential_tool_name, + "input": self.tool_input_2, + "state": "pending", + "suggested_rules": [], + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + # Create external execution result event + external_result_event = ExternalExecutionResultEvent( + reply_id=reply_id, + execution_results=[ + ToolResultBlock( + id=self.tool_call_id_1, + name=self.sequential_tool_name, + output=[ + TextBlock(text=self.sequential_result_1), + ], + state=ToolResultState.SUCCESS, + ), + ], + ) + + # Second call: resume with external execution result + events = [] + async for event in self.agent.reply_stream( + inputs=external_result_event, + ): + events.append(event.model_dump()) + + # Verify events after resumption (sequential execution) + expected_events_resume = [ + *self._get_tool_result_events( + self.tool_call_id_1, + self.sequential_result_1, + ), + *self._get_require_external_execution_events( + reply_id, + self.tool_call_id_2, + self.sequential_tool_name, + self.tool_input_2, + ), + ] + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events_resume], + ) + + # Given the external execution result of the second tool call + external_result_event = ExternalExecutionResultEvent( + reply_id=reply_id, + execution_results=[ + ToolResultBlock( + id=self.tool_call_id_2, + name=self.sequential_tool_name, + output=[ + TextBlock(text=self.sequential_result_2), + ], + state=ToolResultState.ERROR, + ), + ], + ) + + events = [] + async for evnt in self.agent.reply_stream( + inputs=external_result_event, + ): + events.append(evnt.model_dump()) + + # Assert the events + expected_events_after_second_result = [ + *self._get_tool_result_events( + self.tool_call_id_2, + self.sequential_result_2, + state="error", + ), + *self.final_text_events, + {"type": "REPLY_END", "session_id": session_id}, + ] + + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events_after_second_result], + ) + + # Assert final context + expected_context_final = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "finished", + "suggested_rules": [], + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.sequential_tool_name, + "input": self.tool_input_2, + "state": "finished", + "suggested_rules": [], + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.sequential_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.sequential_result_1, + }, + ], + "state": "success", + "metadata": {}, + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.sequential_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.sequential_result_2, + }, + ], + "state": "error", + "metadata": {}, + }, + { + "type": "text", + "id": AnyString(), + "text": self.final_response_text, + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context_final = [ + {**msg_base, **_} for _ in expected_context_final + ] + self.assertListEqual(context_dicts, expected_context_final) + + async def test_concurrent_external_execution(self) -> None: + """Test multiple external execution tool calls in concurrent execution. + + The agent should: + 1. Generate multiple tool calls that require external execution + 2. All tools have is_concurrent_safe=True (concurrent) + 3. Emit REQUIRE_EXTERNAL_EXECUTION event and pause + 4. Resume when ExternalExecutionResultEvent is provided + 5. Continue execution without calling the model again + """ + # Register external concurrent tool + ext_tool = MockExternalConcurrentTool() + self.agent.toolkit = Toolkit( + tools=[ext_tool], + ) + + # Set up mock response with multiple tool calls + self.model.set_responses( + [ + [ + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.concurrent_tool_name, + input=self.tool_input_1, + ), + ToolCallBlock( + id=self.tool_call_id_2, + name=self.concurrent_tool_name, + input=self.tool_input_2, + ), + ], + is_last=False, + usage=None, + ), + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.concurrent_tool_name, + input=self.tool_input_1, + ), + ToolCallBlock( + id=self.tool_call_id_2, + name=self.concurrent_tool_name, + input=self.tool_input_2, + ), + ], + is_last=True, + usage=None, + ), + ], + self.final_mock_responses, + ], + ) + + # First call: collect events until REQUIRE_EXTERNAL_EXECUTION + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content=self.user_input_text), + ): + events.append(event.model_dump()) + + # Verify events + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + tool_call_1_events = self._get_tool_call_events( + self.tool_call_id_1, + self.concurrent_tool_name, + self.tool_input_1, + ) + tool_call_2_events = self._get_tool_call_events( + self.tool_call_id_2, + self.concurrent_tool_name, + self.tool_input_2, + ) + + expected_events = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + *tool_call_1_events[:2], + *tool_call_2_events[:2], + tool_call_1_events[2], + tool_call_2_events[2], + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + *self._get_require_external_execution_events( + reply_id, + self.tool_call_id_1, + self.concurrent_tool_name, + self.tool_input_1, + ), + *self._get_require_external_execution_events( + reply_id, + self.tool_call_id_2, + self.concurrent_tool_name, + self.tool_input_2, + ), + ] + + basic_dict = self._get_event_base(reply_id) + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + # Assert context after first call + msg_base = self._get_msg_base() + expected_context = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.concurrent_tool_name, + "input": self.tool_input_1, + "state": "submitted", + "suggested_rules": [], + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.concurrent_tool_name, + "input": self.tool_input_2, + "state": "submitted", + "suggested_rules": [], + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + # Create external execution result event + external_result_event = ExternalExecutionResultEvent( + reply_id=reply_id, + execution_results=[ + ToolResultBlock( + id=self.tool_call_id_1, + name=self.concurrent_tool_name, + output=[ + TextBlock(text=self.concurrent_result_1), + ], + state=ToolResultState.SUCCESS, + ), + ], + ) + + # Second call: resume with external execution result + events = [] + async for event in self.agent.reply_stream( + inputs=external_result_event, + ): + events.append(event.model_dump()) + + expected_concurrent = [ + *self._get_tool_result_events( + self.tool_call_id_1, + self.concurrent_result_1, + ), + ] + + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_concurrent], + ) + + # The second tool call result + external_result_event = ExternalExecutionResultEvent( + reply_id=reply_id, + execution_results=[ + ToolResultBlock( + id=self.tool_call_id_2, + name=self.concurrent_tool_name, + output=[ + TextBlock(text=self.concurrent_result_2), + ], + state=ToolResultState.SUCCESS, + ), + ], + ) + + events = [] + async for event in self.agent.reply_stream( + inputs=external_result_event, + ): + events.append(event.model_dump()) + + expected_concurrent = [ + *self._get_tool_result_events( + self.tool_call_id_2, + self.concurrent_result_2, + ), + *self.final_text_events, + { + "type": "REPLY_END", + "session_id": session_id, + }, + ] + + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_concurrent], + ) + + # Assert final context + expected_context_final = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.concurrent_tool_name, + "input": self.tool_input_1, + "state": "finished", + "suggested_rules": [], + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.concurrent_tool_name, + "input": self.tool_input_2, + "state": "finished", + "suggested_rules": [], + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.concurrent_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.concurrent_result_1, + }, + ], + "state": "success", + "metadata": {}, + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.concurrent_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.concurrent_result_2, + }, + ], + "state": "success", + "metadata": {}, + }, + { + "type": "text", + "id": AnyString(), + "text": self.final_response_text, + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context_final = [ + {**msg_base, **_} for _ in expected_context_final + ] + self.assertListEqual(context_dicts, expected_context_final) + + async def test_concurrent_external_execution_in_single_event(self) -> None: + """Test concurrent external execution when two results arrive together. + + The agent should: + 1. Generate multiple external tool calls in concurrent mode + 2. Emit one require event per tool call during the initial run + 3. Resume when a single ExternalExecutionResultEvent carries both + results + 4. Continue reasoning only after both results are applied + """ + ext_tool = MockExternalConcurrentTool() + self.agent.toolkit = Toolkit( + tools=[ext_tool], + ) + + self.model.set_responses( + [ + [ + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.concurrent_tool_name, + input=self.tool_input_1, + ), + ToolCallBlock( + id=self.tool_call_id_2, + name=self.concurrent_tool_name, + input=self.tool_input_2, + ), + ], + is_last=False, + usage=None, + ), + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.concurrent_tool_name, + input=self.tool_input_1, + ), + ToolCallBlock( + id=self.tool_call_id_2, + name=self.concurrent_tool_name, + input=self.tool_input_2, + ), + ], + is_last=True, + usage=None, + ), + ], + self.final_mock_responses, + ], + ) + + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content=self.user_input_text), + ): + events.append(event.model_dump()) + + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + basic_dict = self._get_event_base(reply_id) + msg_base = self._get_msg_base() + tool_call_1_events = self._get_tool_call_events( + self.tool_call_id_1, + self.concurrent_tool_name, + self.tool_input_1, + ) + tool_call_2_events = self._get_tool_call_events( + self.tool_call_id_2, + self.concurrent_tool_name, + self.tool_input_2, + ) + + expected_events = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + *tool_call_1_events[:2], + *tool_call_2_events[:2], + tool_call_1_events[2], + tool_call_2_events[2], + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + *self._get_require_external_execution_events( + reply_id, + self.tool_call_id_1, + self.concurrent_tool_name, + self.tool_input_1, + ), + *self._get_require_external_execution_events( + reply_id, + self.tool_call_id_2, + self.concurrent_tool_name, + self.tool_input_2, + ), + ] + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + expected_context = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.concurrent_tool_name, + "input": self.tool_input_1, + "state": "submitted", + "suggested_rules": [], + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.concurrent_tool_name, + "input": self.tool_input_2, + "state": "submitted", + "suggested_rules": [], + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + external_result_event = ExternalExecutionResultEvent( + reply_id=reply_id, + execution_results=[ + ToolResultBlock( + id=self.tool_call_id_1, + name=self.concurrent_tool_name, + output=[ + TextBlock(text=self.concurrent_result_1), + ], + state=ToolResultState.SUCCESS, + ), + ToolResultBlock( + id=self.tool_call_id_2, + name=self.concurrent_tool_name, + output=[ + TextBlock(text=self.concurrent_result_2), + ], + state=ToolResultState.SUCCESS, + ), + ], + ) + + events = [] + async for event in self.agent.reply_stream( + inputs=external_result_event, + ): + events.append(event.model_dump()) + + expected_events_resume = [ + *self._get_tool_result_events( + self.tool_call_id_1, + self.concurrent_result_1, + ), + *self._get_tool_result_events( + self.tool_call_id_2, + self.concurrent_result_2, + ), + *self.final_text_events, + { + "type": "REPLY_END", + "session_id": session_id, + }, + ] + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events_resume], + ) + + expected_context_final = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.concurrent_tool_name, + "input": self.tool_input_1, + "state": "finished", + "suggested_rules": [], + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.concurrent_tool_name, + "input": self.tool_input_2, + "state": "finished", + "suggested_rules": [], + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.concurrent_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.concurrent_result_1, + }, + ], + "state": "success", + "metadata": {}, + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.concurrent_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.concurrent_result_2, + }, + ], + "state": "success", + "metadata": {}, + }, + { + "type": "text", + "id": AnyString(), + "text": self.final_response_text, + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context_final = [ + {**msg_base, **_} for _ in expected_context_final + ] + self.assertListEqual(context_dicts, expected_context_final) + + async def asyncTearDown(self) -> None: + """The async teardown method.""" diff --git a/tests/hitl_mixed_interrupt.py b/tests/hitl_mixed_interrupt.py new file mode 100644 index 0000000000000000000000000000000000000000..310b889685a99ccb2424bac4095e16a19e1b3070 --- /dev/null +++ b/tests/hitl_mixed_interrupt.py @@ -0,0 +1,1211 @@ +# -*- coding: utf-8 -*- +# pylint: disable=redefined-builtin +"""Test mixed user confirmation and external execution in the agent.""" +from typing import Any +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import AnyString, MockModel + +from agentscope.agent import Agent +from agentscope.model import ChatResponse +from agentscope.tool import ( + ToolBase, + Toolkit, + ToolChunk, + RegisteredTool, +) +from agentscope.permission import ( + PermissionDecision, + PermissionBehavior, + PermissionContext, +) +from agentscope.message import ( + TextBlock, + ToolCallBlock, + ToolResultBlock, + UserMsg, + ToolResultState, +) +from agentscope.event import ( + UserConfirmResultEvent, + ExternalExecutionResultEvent, + ConfirmResult, +) + + +class MockMixedSequentialTool(ToolBase): + """A mock tool that requires confirmation and external execution.""" + + name: str = "mock_mixed_sequential_tool" + description: str = "A mock mixed sequential tool for testing" + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "input": {"type": "string", "description": "Input string"}, + }, + "required": ["input"], + } + is_concurrency_safe: bool = False + is_read_only: bool = False + is_external_tool: bool = True + is_mcp: bool = False + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permissions for the tool usage.""" + return PermissionDecision( + behavior=PermissionBehavior.ASK, + decision_reason="Mock mixed tool requires user confirmation", + message="Mock mixed tool requires user confirmation", + ) + + async def __call__(self, input: str, **kwargs: Any) -> ToolChunk: + """Execute the tool.""" + return ToolChunk( + content=[TextBlock(text=f"Mixed sequential result: {input}")], + ) + + +class MockMixedConcurrentTool(ToolBase): + """A mock tool that requires confirmation and external execution.""" + + name: str = "mock_mixed_concurrent_tool" + description: str = "A mock mixed concurrent tool for testing" + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "input": {"type": "string", "description": "Input string"}, + }, + "required": ["input"], + } + is_concurrency_safe: bool = True + is_read_only: bool = False + is_external_tool: bool = True + is_mcp: bool = False + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permissions for the tool usage.""" + return PermissionDecision( + behavior=PermissionBehavior.ASK, + decision_reason="Mock mixed tool requires user confirmation", + message="Mock mixed tool requires user confirmation", + ) + + async def __call__(self, input: str, **kwargs: Any) -> ToolChunk: + """Execute the tool.""" + return ToolChunk( + content=[TextBlock(text=f"Mixed concurrent result: {input}")], + ) + + +class AgentMixTest(IsolatedAsyncioTestCase): + """Test mixed user confirmation and external execution.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.model = MockModel() + self.agent = Agent( + name="Friday", + system_prompt="You are a helpful assistant.", + model=self.model, + toolkit=Toolkit(), + ) + self.tool_call_id_1 = "tool_call_1" + self.tool_call_id_2 = "tool_call_2" + self.user_input_text = "Test" + self.tool_input_1 = '{"input": "test1"}' + self.tool_input_2 = '{"input": "test2"}' + self.sequential_tool_name = "mock_mixed_sequential_tool" + self.concurrent_tool_name = "mock_mixed_concurrent_tool" + self.sequential_result_1 = "Mixed sequential result: test1" + self.sequential_result_2 = "Mixed sequential result: test2" + self.concurrent_result_1 = "Mixed concurrent result: test1" + self.concurrent_result_2 = "Mixed concurrent result: test2" + self.final_response_text = "Final response after mixed execution" + self.final_text_events = [ + { + "type": "MODEL_CALL_START", + "model_name": "mock-model", + }, + { + "type": "TEXT_BLOCK_START", + "block_id": AnyString(), + }, + { + "type": "TEXT_BLOCK_DELTA", + "block_id": AnyString(), + "delta": self.final_response_text, + }, + { + "type": "TEXT_BLOCK_END", + "block_id": AnyString(), + }, + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + ] + self.final_mock_responses = [ + ChatResponse( + content=[TextBlock(text=self.final_response_text)], + is_last=False, + usage=None, + ), + ChatResponse( + content=[TextBlock(text=self.final_response_text)], + is_last=True, + usage=None, + ), + ] + + def _get_event_base(self, reply_id: str) -> dict: + """Get the dict with the basic fields for event assertion.""" + return { + "id": AnyString(), + "created_at": AnyString(), + "reply_id": reply_id, + } + + def _get_msg_base(self) -> dict: + """Get the dict with the basic fields for message assertion.""" + return { + "id": AnyString(), + "created_at": AnyString(), + "metadata": {}, + "name": "Friday", + "role": "assistant", + } + + def _get_tool_call_events( + self, + id: str, + name: str, + delta: str, + ) -> list[dict]: + """Helper method to get the expected tool call events.""" + return [ + { + "type": "TOOL_CALL_START", + "tool_call_id": id, + "tool_call_name": name, + }, + { + "type": "TOOL_CALL_DELTA", + "tool_call_id": id, + "delta": delta, + }, + { + "type": "TOOL_CALL_END", + "tool_call_id": id, + }, + ] + + def _get_tool_result_events( + self, + id: str, + result: str, + state: str = "success", + ) -> list[dict]: + """Helper method to get the expected tool result events.""" + return [ + { + "type": "TOOL_RESULT_TEXT_DELTA", + "tool_call_id": id, + "delta": result, + }, + { + "type": "TOOL_RESULT_END", + "tool_call_id": id, + "state": state, + }, + ] + + def _get_require_user_confirm_event( + self, + reply_id: str, + id: str, + name: str, + tool_input: str, + ) -> dict: + """Helper method to get the expected user confirmation event.""" + return { + "type": "REQUIRE_USER_CONFIRM", + "reply_id": reply_id, + "tool_calls": [ + { + "type": "tool_call", + "id": id, + "name": name, + "input": tool_input, + "state": "asking", + }, + ], + } + + def _get_require_external_execution_events( + self, + reply_id: str, + id: str, + name: str, + tool_input: str, + ) -> list[dict]: + """Helper method to get the expected external execution events.""" + return [ + { + "type": "TOOL_RESULT_START", + "tool_call_id": id, + "tool_call_name": name, + }, + { + "type": "REQUIRE_EXTERNAL_EXECUTION", + "reply_id": reply_id, + "tool_calls": [ + { + "type": "tool_call", + "id": id, + "name": name, + "input": tool_input, + "state": "submitted", + }, + ], + }, + ] + + def _get_tool_call_block( + self, + id: str, + name: str, + tool_input: str, + ) -> ToolCallBlock: + """Build a tool call block.""" + return ToolCallBlock(id=id, name=name, input=tool_input) + + def _get_tool_result_block( + self, + id: str, + name: str, + result: str, + state: ToolResultState = ToolResultState.SUCCESS, + ) -> ToolResultBlock: + """Build a tool result block.""" + return ToolResultBlock( + id=id, + name=name, + output=[TextBlock(text=result)], + state=state, + ) + + def _get_confirm_result( + self, + id: str, + name: str, + tool_input: str, + ) -> ConfirmResult: + """Build a confirmation result.""" + return ConfirmResult( + confirmed=True, + tool_call=self._get_tool_call_block(id, name, tool_input), + ) + + def _build_tool_calls( + self, + tool_calls: list[tuple[str, str, str]], + ) -> list[ToolCallBlock]: + """Build tool call blocks for mock model responses.""" + return [ + self._get_tool_call_block(id, name, tool_input) + for id, name, tool_input in tool_calls + ] + + def _set_model_tool_call_responses( + self, + tool_calls: list[tuple[str, str, str]], + ) -> None: + """Set mock model responses that emit tool calls then final text.""" + self.model.set_responses( + [ + [ + ChatResponse( + content=self._build_tool_calls(tool_calls), + is_last=False, + usage=None, + ), + ChatResponse( + content=self._build_tool_calls(tool_calls), + is_last=True, + usage=None, + ), + ], + self.final_mock_responses, + ], + ) + + async def test_single_user_confirmation_and_external_execution( + self, + ) -> None: + """Test one tool call that needs confirmation and external + execution.""" + mixed_tool = MockMixedSequentialTool() + self.agent.toolkit.tools[mixed_tool.name] = RegisteredTool( + tool=mixed_tool, + group="basic", + ) + + self._set_model_tool_call_responses( + [ + ( + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ), + ], + ) + + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content=self.user_input_text), + ): + events.append(event.model_dump()) + + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + basic_dict = self._get_event_base(reply_id) + msg_base = self._get_msg_base() + + expected_events = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + *self._get_tool_call_events( + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ), + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + self._get_require_user_confirm_event( + reply_id, + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ), + ] + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + expected_context = [ + { + "name": "user", + "role": "user", + "content": self.user_input_text, + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "asking", + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + user_confirm_event = UserConfirmResultEvent( + reply_id=reply_id, + confirm_results=[ + self._get_confirm_result( + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ), + ], + ) + + events = [] + async for event in self.agent.reply_stream(inputs=user_confirm_event): + events.append(event.model_dump()) + + expected_events_resume = self._get_require_external_execution_events( + reply_id, + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ) + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events_resume], + ) + + expected_context = [ + { + "name": "user", + "role": "user", + "content": self.user_input_text, + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "submitted", + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + external_result_event = ExternalExecutionResultEvent( + reply_id=reply_id, + execution_results=[ + self._get_tool_result_block( + self.tool_call_id_1, + self.sequential_tool_name, + self.sequential_result_1, + ), + ], + ) + + events = [] + async for event in self.agent.reply_stream( + inputs=external_result_event, + ): + events.append(event.model_dump()) + + expected_events_after_result = [ + *self._get_tool_result_events( + self.tool_call_id_1, + self.sequential_result_1, + ), + *self.final_text_events, + {"type": "REPLY_END", "session_id": session_id}, + ] + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events_after_result], + ) + + expected_context_final = [ + { + "name": "user", + "role": "user", + "content": self.user_input_text, + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "finished", + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.sequential_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.sequential_result_1, + }, + ], + "state": "success", + }, + { + "type": "text", + "id": AnyString(), + "text": self.final_response_text, + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context_final = [ + {**msg_base, **_} for _ in expected_context_final + ] + self.assertListEqual(context_dicts, expected_context_final) + + async def test_sequential_user_confirmation_and_external_execution( + self, + ) -> None: + """Test sequential tool calls that need confirmation and external + execution.""" + mixed_tool = MockMixedSequentialTool() + self.agent.toolkit.tools[mixed_tool.name] = RegisteredTool( + tool=mixed_tool, + group="basic", + ) + + self._set_model_tool_call_responses( + [ + ( + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ), + ( + self.tool_call_id_2, + self.sequential_tool_name, + self.tool_input_2, + ), + ], + ) + + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content=self.user_input_text), + ): + events.append(event.model_dump()) + + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + basic_dict = self._get_event_base(reply_id) + msg_base = self._get_msg_base() + tool_call_1_events = self._get_tool_call_events( + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ) + tool_call_2_events = self._get_tool_call_events( + self.tool_call_id_2, + self.sequential_tool_name, + self.tool_input_2, + ) + + expected_events = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + *tool_call_1_events[:2], + *tool_call_2_events[:2], + tool_call_1_events[2], + tool_call_2_events[2], + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + self._get_require_user_confirm_event( + reply_id, + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ), + ] + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + expected_context = [ + { + "name": "user", + "role": "user", + "content": self.user_input_text, + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "asking", + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.sequential_tool_name, + "input": self.tool_input_2, + "state": "pending", + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + user_confirm_event = UserConfirmResultEvent( + reply_id=reply_id, + confirm_results=[ + self._get_confirm_result( + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ), + ], + ) + + events = [] + async for event in self.agent.reply_stream(inputs=user_confirm_event): + events.append(event.model_dump()) + + expected_events_resume = self._get_require_external_execution_events( + reply_id, + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ) + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events_resume], + ) + + expected_context = [ + { + "name": "user", + "role": "user", + "content": self.user_input_text, + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "submitted", + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.sequential_tool_name, + "input": self.tool_input_2, + "state": "pending", + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + external_result_event = ExternalExecutionResultEvent( + reply_id=reply_id, + execution_results=[ + self._get_tool_result_block( + self.tool_call_id_1, + self.sequential_tool_name, + self.sequential_result_1, + ), + ], + ) + + events = [] + async for event in self.agent.reply_stream( + inputs=external_result_event, + ): + events.append(event.model_dump()) + + expected_events_after_first_result = [ + *self._get_tool_result_events( + self.tool_call_id_1, + self.sequential_result_1, + ), + self._get_require_user_confirm_event( + reply_id, + self.tool_call_id_2, + self.sequential_tool_name, + self.tool_input_2, + ), + ] + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events_after_first_result], + ) + + expected_context = [ + { + "name": "user", + "role": "user", + "content": self.user_input_text, + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "finished", + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.sequential_tool_name, + "input": self.tool_input_2, + "state": "asking", + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.sequential_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.sequential_result_1, + }, + ], + "state": "success", + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + user_confirm_event = UserConfirmResultEvent( + reply_id=reply_id, + confirm_results=[ + self._get_confirm_result( + self.tool_call_id_2, + self.sequential_tool_name, + self.tool_input_2, + ), + ], + ) + + events = [] + async for event in self.agent.reply_stream(inputs=user_confirm_event): + events.append(event.model_dump()) + + expected_events_after_second_confirm = ( + self._get_require_external_execution_events( + reply_id, + self.tool_call_id_2, + self.sequential_tool_name, + self.tool_input_2, + ) + ) + self.assertListEqual( + events, + [ + {**basic_dict, **_} + for _ in expected_events_after_second_confirm + ], + ) + + external_result_event = ExternalExecutionResultEvent( + reply_id=reply_id, + execution_results=[ + self._get_tool_result_block( + self.tool_call_id_2, + self.sequential_tool_name, + self.sequential_result_2, + ), + ], + ) + + events = [] + async for event in self.agent.reply_stream( + inputs=external_result_event, + ): + events.append(event.model_dump()) + + expected_events_after_second_result = [ + *self._get_tool_result_events( + self.tool_call_id_2, + self.sequential_result_2, + ), + *self.final_text_events, + {"type": "REPLY_END", "session_id": session_id}, + ] + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events_after_second_result], + ) + + expected_context_final = [ + { + "name": "user", + "role": "user", + "content": self.user_input_text, + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "finished", + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.sequential_tool_name, + "input": self.tool_input_2, + "state": "finished", + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.sequential_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.sequential_result_1, + }, + ], + "state": "success", + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.sequential_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.sequential_result_2, + }, + ], + "state": "success", + }, + { + "type": "text", + "id": AnyString(), + "text": self.final_response_text, + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context_final = [ + {**msg_base, **_} for _ in expected_context_final + ] + self.assertListEqual(context_dicts, expected_context_final) + + async def test_concurrent_user_confirmation_and_external_execution( + self, + ) -> None: + """Test concurrent tool calls that need confirmation and external + execution.""" + mixed_tool = MockMixedConcurrentTool() + self.agent.toolkit.tools[mixed_tool.name] = RegisteredTool( + tool=mixed_tool, + group="basic", + ) + + self._set_model_tool_call_responses( + [ + ( + self.tool_call_id_1, + self.concurrent_tool_name, + self.tool_input_1, + ), + ( + self.tool_call_id_2, + self.concurrent_tool_name, + self.tool_input_2, + ), + ], + ) + + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content=self.user_input_text), + ): + events.append(event.model_dump()) + + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + basic_dict = self._get_event_base(reply_id) + msg_base = self._get_msg_base() + tool_call_1_events = self._get_tool_call_events( + self.tool_call_id_1, + self.concurrent_tool_name, + self.tool_input_1, + ) + tool_call_2_events = self._get_tool_call_events( + self.tool_call_id_2, + self.concurrent_tool_name, + self.tool_input_2, + ) + + expected_events = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + *tool_call_1_events[:2], + *tool_call_2_events[:2], + tool_call_1_events[2], + tool_call_2_events[2], + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + self._get_require_user_confirm_event( + reply_id, + self.tool_call_id_1, + self.concurrent_tool_name, + self.tool_input_1, + ), + self._get_require_user_confirm_event( + reply_id, + self.tool_call_id_2, + self.concurrent_tool_name, + self.tool_input_2, + ), + ] + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + expected_context = [ + { + "name": "user", + "role": "user", + "content": self.user_input_text, + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.concurrent_tool_name, + "input": self.tool_input_1, + "state": "asking", + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.concurrent_tool_name, + "input": self.tool_input_2, + "state": "asking", + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + user_confirm_event = UserConfirmResultEvent( + reply_id=reply_id, + confirm_results=[ + self._get_confirm_result( + self.tool_call_id_1, + self.concurrent_tool_name, + self.tool_input_1, + ), + self._get_confirm_result( + self.tool_call_id_2, + self.concurrent_tool_name, + self.tool_input_2, + ), + ], + ) + + events = [] + async for event in self.agent.reply_stream(inputs=user_confirm_event): + events.append(event.model_dump()) + + self.assertEqual(len(events), 4) + expected_tool_events = { + self.tool_call_id_1: [ + {**basic_dict, **_} + for _ in self._get_require_external_execution_events( + reply_id, + self.tool_call_id_1, + self.concurrent_tool_name, + self.tool_input_1, + ) + ], + self.tool_call_id_2: [ + {**basic_dict, **_} + for _ in self._get_require_external_execution_events( + reply_id, + self.tool_call_id_2, + self.concurrent_tool_name, + self.tool_input_2, + ) + ], + } + for tool_call_id, expected_tool_event in expected_tool_events.items(): + self.assertListEqual( + [ + event + for event in events + if event.get("tool_call_id") == tool_call_id + or ( + event["type"] == "REQUIRE_EXTERNAL_EXECUTION" + and event["tool_calls"][0]["id"] == tool_call_id + ) + ], + expected_tool_event, + ) + + expected_context = [ + { + "name": "user", + "role": "user", + "content": self.user_input_text, + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.concurrent_tool_name, + "input": self.tool_input_1, + "state": "submitted", + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.concurrent_tool_name, + "input": self.tool_input_2, + "state": "submitted", + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + external_result_event = ExternalExecutionResultEvent( + reply_id=reply_id, + execution_results=[ + self._get_tool_result_block( + self.tool_call_id_1, + self.concurrent_tool_name, + self.concurrent_result_1, + ), + self._get_tool_result_block( + self.tool_call_id_2, + self.concurrent_tool_name, + self.concurrent_result_2, + ), + ], + ) + + events = [] + async for event in self.agent.reply_stream( + inputs=external_result_event, + ): + events.append(event.model_dump()) + + expected_events_after_result = [ + *self._get_tool_result_events( + self.tool_call_id_1, + self.concurrent_result_1, + ), + *self._get_tool_result_events( + self.tool_call_id_2, + self.concurrent_result_2, + ), + *self.final_text_events, + {"type": "REPLY_END", "session_id": session_id}, + ] + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events_after_result], + ) + + expected_context_final = [ + { + "name": "user", + "role": "user", + "content": self.user_input_text, + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.concurrent_tool_name, + "input": self.tool_input_1, + "state": "finished", + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.concurrent_tool_name, + "input": self.tool_input_2, + "state": "finished", + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.concurrent_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.concurrent_result_1, + }, + ], + "state": "success", + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.concurrent_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.concurrent_result_2, + }, + ], + "state": "success", + }, + { + "type": "text", + "id": AnyString(), + "text": self.final_response_text, + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context_final = [ + {**msg_base, **_} for _ in expected_context_final + ] + self.assertListEqual(context_dicts, expected_context_final) + + async def asyncTearDown(self) -> None: + """The async teardown method.""" diff --git a/tests/hitl_user_confirmation_test.py b/tests/hitl_user_confirmation_test.py new file mode 100644 index 0000000000000000000000000000000000000000..9bc88085232f82456154078f1279ba490d55d9f3 --- /dev/null +++ b/tests/hitl_user_confirmation_test.py @@ -0,0 +1,1497 @@ +# -*- coding: utf-8 -*- +# pylint: disable=redefined-builtin +"""Test the user confirmation events in the agent class.""" +from typing import Any +from unittest.async_case import IsolatedAsyncioTestCase +from utils import AnyString, MockModel + +from agentscope.agent import Agent +from agentscope.model import ChatResponse +from agentscope.tool import ( + ToolBase, + Toolkit, + ToolChunk, +) +from agentscope.permission import ( + PermissionDecision, + PermissionBehavior, + PermissionContext, +) +from agentscope.message import ( + TextBlock, + ToolCallBlock, + UserMsg, +) +from agentscope.event import UserConfirmResultEvent, ConfirmResult + + +class MockUserConfirmSequentialTool(ToolBase): + """A mock tool that requires user confirmation (sequential).""" + + name: str = "mock_user_confirm_sequential_tool" + description: str = "A mock user confirm sequential tool for testing" + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "input": {"type": "string", "description": "Input string"}, + }, + "required": ["input"], + } + is_concurrency_safe: bool = False + is_read_only: bool = False + is_external_tool: bool = False + is_mcp: bool = False + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permissions for the tool usage.""" + return PermissionDecision( + behavior=PermissionBehavior.ASK, + decision_reason="Mock tool requires user confirmation", + message="Mock tool requires user confirmation", + ) + + async def __call__(self, input: str, **kwargs: Any) -> ToolChunk: + """Execute the tool.""" + return ToolChunk( + content=[ + TextBlock(text=f"User confirm sequential result: {input}"), + ], + ) + + +class MockUserConfirmConcurrentTool(ToolBase): + """A mock tool that requires user confirmation (concurrent).""" + + name: str = "mock_user_confirm_concurrent_tool" + description: str = "A mock user confirm concurrent tool for testing" + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "input": {"type": "string", "description": "Input string"}, + }, + "required": ["input"], + } + is_concurrency_safe: bool = True + is_read_only: bool = False + is_external_tool: bool = False + is_mcp: bool = False + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permissions for the tool usage.""" + return PermissionDecision( + behavior=PermissionBehavior.ASK, + decision_reason="Mock tool requires user confirmation", + message="Mock tool requires user confirmation", + ) + + async def __call__(self, input: str, **kwargs: Any) -> ToolChunk: + """Execute the tool.""" + return ToolChunk( + content=[ + TextBlock(text=f"User confirm concurrent result: {input}"), + ], + ) + + +class AgentUserConfirmationTest(IsolatedAsyncioTestCase): + """Test the user confirmation events in the agent class.""" + + def _get_tool_call_events( + self, + id: str, + name: str, + delta: str, + ) -> list[dict]: + """Helper method to get the expected tool call events.""" + return [ + { + "type": "TOOL_CALL_START", + "tool_call_id": id, + "tool_call_name": name, + }, + { + "type": "TOOL_CALL_DELTA", + "tool_call_id": id, + "delta": delta, + }, + { + "type": "TOOL_CALL_END", + "tool_call_id": id, + }, + ] + + def _get_tool_result_events( + self, + id: str, + name: str, + result: str, + ) -> list[dict]: + """Helper method to get the expected tool result events.""" + return [ + { + "type": "TOOL_RESULT_START", + "tool_call_id": id, + "tool_call_name": name, + }, + { + "type": "TOOL_RESULT_TEXT_DELTA", + "tool_call_id": id, + "delta": result, + }, + { + "type": "TOOL_RESULT_END", + "tool_call_id": id, + "state": "success", + }, + ] + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.model = MockModel() + self.agent = Agent( + name="Friday", + system_prompt="You are a helpful assistant.", + model=self.model, + toolkit=Toolkit(), + ) + self.tool_call_id_1 = "tool_call_1" + self.tool_call_id_2 = "tool_call_2" + self.user_input_text = "Test" + self.tool_input_1 = '{"input": "test1"}' + self.tool_input_2 = '{"input": "test2"}' + self.sequential_tool_name = "mock_user_confirm_sequential_tool" + self.concurrent_tool_name = "mock_user_confirm_concurrent_tool" + self.sequential_result_1 = "User confirm sequential result: test1" + self.sequential_result_2 = "User confirm sequential result: test2" + self.concurrent_result_1 = "User confirm concurrent result: test1" + self.concurrent_result_2 = "User confirm concurrent result: test2" + self.final_response_text = "Result 1" + self.final_text_events = [ + { + "type": "MODEL_CALL_START", + "model_name": "mock-model", + }, + { + "type": "TEXT_BLOCK_START", + "block_id": AnyString(), + }, + { + "type": "TEXT_BLOCK_DELTA", + "block_id": AnyString(), + "delta": self.final_response_text, + }, + { + "type": "TEXT_BLOCK_END", + "block_id": AnyString(), + }, + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + ] + + self.final_mock_responses = [ + ChatResponse( + content=[ + TextBlock(text=self.final_response_text), + ], + is_last=False, + ), + ChatResponse( + content=[ + TextBlock(text=self.final_response_text), + ], + is_last=True, + ), + ] + + def _get_event_base(self, reply_id: str) -> dict: + """Get the dict with the basic fields for event assertion.""" + return { + "id": AnyString(), + "created_at": AnyString(), + "metadata": {}, + "reply_id": reply_id, + } + + def _get_msg_base(self) -> dict: + """Get the dict with the basic fields for message assertion.""" + return { + "id": AnyString(), + "created_at": AnyString(), + "finished_at": None, + "metadata": {}, + "name": "Friday", + "role": "assistant", + "usage": None, + } + + async def test_single_user_confirmation(self) -> None: + """Test single user confirmation tool call. + + The agent should: + 1. Generate a tool call that requires user confirmation + 2. Emit REQUIRE_USER_CONFIRM event and pause + 3. Resume when UserConfirmResultEvent is provided + 4. Execute the tool and continue + """ + # Register user confirm tool + confirm_tool = MockUserConfirmSequentialTool() + self.agent.toolkit = Toolkit( + tools=[confirm_tool], + ) + + # Set up mock response with tool call (no final text response) + self.model.set_responses( + [ + [ + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.sequential_tool_name, + input=self.tool_input_1, + ), + ], + is_last=False, + ), + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.sequential_tool_name, + input=self.tool_input_1, + ), + ], + is_last=True, + ), + ], + self.final_mock_responses, + ], + ) + + # First call: collect events until REQUIRE_USER_CONFIRM + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content=self.user_input_text), + ): + events.append(event.model_dump()) + + # Verify events + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + + expected_events = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + *self._get_tool_call_events( + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ), + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + { + "type": "REQUIRE_USER_CONFIRM", + "reply_id": reply_id, + "tool_calls": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "asking", + "suggested_rules": [ + { + "tool_name": self.sequential_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + ], + }, + ] + + basic_dict = self._get_event_base(reply_id) + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + # Assert context after first call + msg_base = self._get_msg_base() + expected_context = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "asking", + "suggested_rules": [ + { + "tool_name": self.sequential_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + # Create user confirmation result event + user_confirm_event = UserConfirmResultEvent( + reply_id=reply_id, + confirm_results=[ + ConfirmResult( + confirmed=True, + tool_call=ToolCallBlock( + id=self.tool_call_id_1, + name=self.sequential_tool_name, + input=self.tool_input_1, + ), + ), + ], + ) + + # Second call: resume with user confirmation result + events = [] + async for event in self.agent.reply_stream(inputs=user_confirm_event): + events.append(event.model_dump()) + + # Verify events after resumption + expected_events_resume = [ + *self._get_tool_result_events( + self.tool_call_id_1, + self.sequential_tool_name, + self.sequential_result_1, + ), + *self.final_text_events, + {"type": "REPLY_END", "session_id": session_id}, + ] + + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events_resume], + ) + + # Assert final context + expected_context_final = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "finished", + "suggested_rules": [ + { + "tool_name": self.sequential_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.sequential_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.sequential_result_1, + }, + ], + "state": "success", + "metadata": {}, + }, + { + "type": "text", + "id": AnyString(), + "text": self.final_response_text, + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context_final = [ + {**msg_base, **_} for _ in expected_context_final + ] + self.assertListEqual(context_dicts, expected_context_final) + + async def test_sequential_user_confirmation(self) -> None: + """Test multiple user confirmation tool calls in sequential execution. + + The agent should: + 1. Generate multiple tool calls that require user confirmation + 2. All tools have is_concurrent_safe=False (sequential) + 3. Emit REQUIRE_USER_CONFIRM event and pause + 4. Resume when UserConfirmResultEvent is provided + 5. Execute the tools and continue + """ + # Register user confirm sequential tool + confirm_tool = MockUserConfirmSequentialTool() + self.agent.toolkit = Toolkit( + tools=[confirm_tool], + ) + + # Set up mock response with multiple tool calls + self.model.set_responses( + [ + [ + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.sequential_tool_name, + input=self.tool_input_1, + ), + ToolCallBlock( + id=self.tool_call_id_2, + name=self.sequential_tool_name, + input=self.tool_input_2, + ), + ], + is_last=False, + usage=None, + ), + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.sequential_tool_name, + input=self.tool_input_1, + ), + ToolCallBlock( + id=self.tool_call_id_2, + name=self.sequential_tool_name, + input=self.tool_input_2, + ), + ], + is_last=True, + usage=None, + ), + ], + self.final_mock_responses, + ], + ) + + # First call: collect events until REQUIRE_USER_CONFIRM + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content=self.user_input_text), + ): + events.append(event.model_dump()) + + # Verify events + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + + tool_call_1_events = self._get_tool_call_events( + self.tool_call_id_1, + self.sequential_tool_name, + self.tool_input_1, + ) + tool_call_2_events = self._get_tool_call_events( + self.tool_call_id_2, + self.sequential_tool_name, + self.tool_input_2, + ) + + expected_events = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + *tool_call_1_events[:2], + *tool_call_2_events[:2], + tool_call_1_events[2], + tool_call_2_events[2], + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + { + "type": "REQUIRE_USER_CONFIRM", + "reply_id": reply_id, + "tool_calls": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "asking", + "suggested_rules": [ + { + "tool_name": self.sequential_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + ], + }, + ] + + basic_dict = self._get_event_base(reply_id) + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + # Assert context after first call + msg_base = self._get_msg_base() + expected_context = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "asking", + "suggested_rules": [ + { + "tool_name": self.sequential_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.sequential_tool_name, + "input": self.tool_input_2, + "state": "pending", + "suggested_rules": [], + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + # Create user confirmation result event + user_confirm_event = UserConfirmResultEvent( + reply_id=reply_id, + confirm_results=[ + ConfirmResult( + confirmed=True, + tool_call=ToolCallBlock( + id=self.tool_call_id_1, + name=self.sequential_tool_name, + input=self.tool_input_1, + ), + ), + ], + ) + + # resume with user confirmation result + events = [] + async for event in self.agent.reply_stream(inputs=user_confirm_event): + events.append(event.model_dump()) + + # Verify events after resumption (sequential execution) + expected_events_resume = [ + *self._get_tool_result_events( + self.tool_call_id_1, + self.sequential_tool_name, + self.sequential_result_1, + ), + { + "type": "REQUIRE_USER_CONFIRM", + "tool_calls": [ + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.sequential_tool_name, + "input": self.tool_input_2, + "state": "asking", + "suggested_rules": [ + { + "tool_name": self.sequential_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + ], + }, + ] + + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events_resume], + ) + + # Confirm the second tool call + user_confirm_event = UserConfirmResultEvent( + reply_id=reply_id, + confirm_results=[ + ConfirmResult( + confirmed=True, + tool_call=ToolCallBlock( + id=self.tool_call_id_2, + name=self.sequential_tool_name, + input=self.tool_input_2, + ), + ), + ], + ) + + # Second call: resume with user confirmation result + events = [] + async for event in self.agent.reply_stream(inputs=user_confirm_event): + events.append(event.model_dump()) + + expected_events_resume_2 = [ + *self._get_tool_result_events( + self.tool_call_id_2, + self.sequential_tool_name, + self.sequential_result_2, + ), + *self.final_text_events, + {"type": "REPLY_END", "session_id": session_id}, + ] + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events_resume_2], + ) + + # Assert final context + expected_context_final = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.sequential_tool_name, + "input": self.tool_input_1, + "state": "finished", + "suggested_rules": [ + { + "tool_name": self.sequential_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.sequential_tool_name, + "input": self.tool_input_2, + "state": "finished", + "suggested_rules": [ + { + "tool_name": self.sequential_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.sequential_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.sequential_result_1, + }, + ], + "state": "success", + "metadata": {}, + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.sequential_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.sequential_result_2, + }, + ], + "state": "success", + "metadata": {}, + }, + { + "type": "text", + "id": AnyString(), + "text": self.final_response_text, + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context_final = [ + {**msg_base, **_} for _ in expected_context_final + ] + self.assertListEqual(context_dicts, expected_context_final) + + async def test_concurrent_user_confirmation(self) -> None: + """Test multiple user confirmation tool calls in concurrent execution. + + The agent should: + 1. Generate multiple tool calls that require user confirmation + 2. All tools have is_concurrent_safe=True (concurrent) + 3. Emit REQUIRE_USER_CONFIRM event and pause + 4. Resume when UserConfirmResultEvent is provided + 5. Execute the tools and continue + """ + # Register user confirm concurrent tool + confirm_tool = MockUserConfirmConcurrentTool() + self.agent.toolkit = Toolkit( + tools=[confirm_tool], + ) + + # Set up mock response with multiple tool calls + self.model.set_responses( + [ + [ + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.concurrent_tool_name, + input=self.tool_input_1, + ), + ToolCallBlock( + id=self.tool_call_id_2, + name=self.concurrent_tool_name, + input=self.tool_input_2, + ), + ], + is_last=False, + ), + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.concurrent_tool_name, + input=self.tool_input_1, + ), + ToolCallBlock( + id=self.tool_call_id_2, + name=self.concurrent_tool_name, + input=self.tool_input_2, + ), + ], + is_last=True, + ), + ], + self.final_mock_responses, + ], + ) + + # First call: collect events until REQUIRE_USER_CONFIRM + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content=self.user_input_text), + ): + events.append(event.model_dump()) + + # Verify events + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + + tool_call_1_events = self._get_tool_call_events( + self.tool_call_id_1, + self.concurrent_tool_name, + self.tool_input_1, + ) + tool_call_2_events = self._get_tool_call_events( + self.tool_call_id_2, + self.concurrent_tool_name, + self.tool_input_2, + ) + + expected_events = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + *tool_call_1_events[:2], + *tool_call_2_events[:2], + tool_call_1_events[2], + tool_call_2_events[2], + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + { + "type": "REQUIRE_USER_CONFIRM", + "reply_id": reply_id, + "tool_calls": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.concurrent_tool_name, + "input": self.tool_input_1, + "state": "asking", + "suggested_rules": [ + { + "tool_name": self.concurrent_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + ], + }, + { + "type": "REQUIRE_USER_CONFIRM", + "reply_id": reply_id, + "tool_calls": [ + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.concurrent_tool_name, + "input": self.tool_input_2, + "state": "asking", + "suggested_rules": [ + { + "tool_name": self.concurrent_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + ], + }, + ] + + basic_dict = self._get_event_base(reply_id) + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + # Assert context after first call + msg_base = self._get_msg_base() + expected_context = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.concurrent_tool_name, + "input": self.tool_input_1, + "state": "asking", + "suggested_rules": [ + { + "tool_name": self.concurrent_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.concurrent_tool_name, + "input": self.tool_input_2, + "state": "asking", + "suggested_rules": [ + { + "tool_name": self.concurrent_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + # Create user confirmation result event + user_confirm_event = UserConfirmResultEvent( + reply_id=reply_id, + confirm_results=[ + ConfirmResult( + confirmed=True, + tool_call=ToolCallBlock( + id=self.tool_call_id_1, + name=self.concurrent_tool_name, + input=self.tool_input_1, + ), + ), + ], + ) + + # resume with user confirmation result + events = [] + async for event in self.agent.reply_stream(inputs=user_confirm_event): + events.append(event.model_dump()) + + # Verify events for tool call 1 after resumption + expected_events = [ + *self._get_tool_result_events( + self.tool_call_id_1, + self.concurrent_tool_name, + self.concurrent_result_1, + ), + ] + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + # The second tool call + user_confirm_event = UserConfirmResultEvent( + reply_id=reply_id, + confirm_results=[ + ConfirmResult( + confirmed=True, + tool_call=ToolCallBlock( + id=self.tool_call_id_2, + name=self.concurrent_tool_name, + input=self.tool_input_2, + ), + ), + ], + ) + + events = [] + async for event in self.agent.reply_stream(inputs=user_confirm_event): + events.append(event.model_dump()) + + expected_events = [ + *self._get_tool_result_events( + self.tool_call_id_2, + self.concurrent_tool_name, + self.concurrent_result_2, + ), + *self.final_text_events, + {"type": "REPLY_END", "session_id": session_id}, + ] + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + # Assert final context + expected_context_final = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.concurrent_tool_name, + "input": self.tool_input_1, + "state": "finished", + "suggested_rules": [ + { + "tool_name": self.concurrent_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.concurrent_tool_name, + "input": self.tool_input_2, + "state": "finished", + "suggested_rules": [ + { + "tool_name": self.concurrent_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.concurrent_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.concurrent_result_1, + }, + ], + "state": "success", + "metadata": {}, + }, + { + "type": "tool_result", + "id": AnyString(), + "name": self.concurrent_tool_name, + "output": [ + { + "type": "text", + "id": AnyString(), + "text": self.concurrent_result_2, + }, + ], + "state": "success", + "metadata": {}, + }, + { + "type": "text", + "id": AnyString(), + "text": self.final_response_text, + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context_final = [ + {**msg_base, **_} for _ in expected_context_final + ] + self.assertListEqual(context_dicts, expected_context_final) + + async def test_concurrent_user_confirmation_in_single_event(self) -> None: + """Test concurrent user confirmation when two approvals arrive + together. + + The agent should: + 1. Generate multiple tool calls that require user confirmation + 2. Pause in concurrent mode with two asking tool calls + 3. Resume when one UserConfirmResultEvent carries both confirmations + 4. Execute both tools and continue reasoning after both complete + """ + confirm_tool = MockUserConfirmConcurrentTool() + self.agent.toolkit = Toolkit( + tools=[confirm_tool], + ) + + self.model.set_responses( + [ + [ + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.concurrent_tool_name, + input=self.tool_input_1, + ), + ToolCallBlock( + id=self.tool_call_id_2, + name=self.concurrent_tool_name, + input=self.tool_input_2, + ), + ], + is_last=False, + ), + ChatResponse( + content=[ + ToolCallBlock( + id=self.tool_call_id_1, + name=self.concurrent_tool_name, + input=self.tool_input_1, + ), + ToolCallBlock( + id=self.tool_call_id_2, + name=self.concurrent_tool_name, + input=self.tool_input_2, + ), + ], + is_last=True, + ), + ], + self.final_mock_responses, + ], + ) + + events = [] + async for event in self.agent.reply_stream( + UserMsg(name="user", content=self.user_input_text), + ): + events.append(event.model_dump()) + + session_id = self.agent.state.session_id + reply_id = self.agent.state.reply_id + basic_dict = self._get_event_base(reply_id) + msg_base = self._get_msg_base() + + tool_call_1_events = self._get_tool_call_events( + self.tool_call_id_1, + self.concurrent_tool_name, + self.tool_input_1, + ) + tool_call_2_events = self._get_tool_call_events( + self.tool_call_id_2, + self.concurrent_tool_name, + self.tool_input_2, + ) + + expected_events = [ + { + "type": "REPLY_START", + "session_id": session_id, + "name": "Friday", + "role": "assistant", + }, + {"type": "MODEL_CALL_START", "model_name": "mock-model"}, + *tool_call_1_events[:2], + *tool_call_2_events[:2], + tool_call_1_events[2], + tool_call_2_events[2], + { + "type": "MODEL_CALL_END", + "input_tokens": 0, + "output_tokens": 0, + }, + { + "type": "REQUIRE_USER_CONFIRM", + "reply_id": reply_id, + "tool_calls": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.concurrent_tool_name, + "input": self.tool_input_1, + "state": "asking", + "suggested_rules": [ + { + "tool_name": self.concurrent_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + ], + }, + { + "type": "REQUIRE_USER_CONFIRM", + "reply_id": reply_id, + "tool_calls": [ + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.concurrent_tool_name, + "input": self.tool_input_2, + "state": "asking", + "suggested_rules": [ + { + "tool_name": self.concurrent_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + ], + }, + ] + self.assertListEqual( + events, + [{**basic_dict, **_} for _ in expected_events], + ) + + expected_context = [ + { + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "finished_at": AnyString(), + }, + { + "content": [ + { + "type": "tool_call", + "id": self.tool_call_id_1, + "name": self.concurrent_tool_name, + "input": self.tool_input_1, + "state": "asking", + "suggested_rules": [ + { + "tool_name": self.concurrent_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + { + "type": "tool_call", + "id": self.tool_call_id_2, + "name": self.concurrent_tool_name, + "input": self.tool_input_2, + "state": "asking", + "suggested_rules": [ + { + "tool_name": self.concurrent_tool_name, + "rule_content": None, + "behavior": PermissionBehavior.ALLOW, + "source": "suggested", + }, + ], + }, + ], + }, + ] + context_dicts = [msg.model_dump() for msg in self.agent.state.context] + expected_context = [{**msg_base, **_} for _ in expected_context] + self.assertListEqual(context_dicts, expected_context) + + user_confirm_event = UserConfirmResultEvent( + reply_id=reply_id, + confirm_results=[ + ConfirmResult( + confirmed=True, + tool_call=ToolCallBlock( + id=self.tool_call_id_1, + name=self.concurrent_tool_name, + input=self.tool_input_1, + ), + ), + ConfirmResult( + confirmed=True, + tool_call=ToolCallBlock( + id=self.tool_call_id_2, + name=self.concurrent_tool_name, + input=self.tool_input_2, + ), + ), + ], + ) + + events = [] + async for event in self.agent.reply_stream(inputs=user_confirm_event): + events.append(event.model_dump()) + + tool_events = events[:6] + final_events = events[6:] + self.assertEqual(len(tool_events), 6) + + expected_tool_events = { + self.tool_call_id_1: [ + {**basic_dict, **_} + for _ in self._get_tool_result_events( + self.tool_call_id_1, + self.concurrent_tool_name, + self.concurrent_result_1, + ) + ], + self.tool_call_id_2: [ + {**basic_dict, **_} + for _ in self._get_tool_result_events( + self.tool_call_id_2, + self.concurrent_tool_name, + self.concurrent_result_2, + ) + ], + } + for tool_call_id, expected_tool_event in expected_tool_events.items(): + self.assertListEqual( + [ + event + for event in tool_events + if event["tool_call_id"] == tool_call_id + ], + expected_tool_event, + ) + + expected_final_events = [ + *self.final_text_events, + {"type": "REPLY_END", "session_id": session_id}, + ] + self.assertListEqual( + final_events, + [{**basic_dict, **_} for _ in expected_final_events], + ) + + self.assertEqual(len(self.agent.state.context), 2) + self.assertEqual( + self.agent.state.context[0].model_dump(), + { + "id": AnyString(), + "created_at": AnyString(), + "finished_at": AnyString(), + "metadata": {}, + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "id": AnyString(), + "text": self.user_input_text, + }, + ], + "usage": None, + }, + ) + + assistant_msg = self.agent.state.context[-1] + self.assertEqual( + [ + _.model_dump()["state"] + for _ in assistant_msg.get_content_blocks("tool_call") + ], + ["finished", "finished"], + ) + self.assertEqual( + [ + _.model_dump()["id"] + for _ in assistant_msg.get_content_blocks("tool_call") + ], + [self.tool_call_id_1, self.tool_call_id_2], + ) + self.assertEqual( + { + ( + _.model_dump()["name"], + _.model_dump()["state"], + _.output[0].text, + ) + for _ in assistant_msg.get_content_blocks("tool_result") + }, + { + ( + self.concurrent_tool_name, + "success", + self.concurrent_result_1, + ), + ( + self.concurrent_tool_name, + "success", + self.concurrent_result_2, + ), + }, + ) + self.assertEqual( + [_.text for _ in assistant_msg.get_content_blocks("text")], + [self.final_response_text], + ) + + async def asyncTearDown(self) -> None: + """The async teardown method.""" diff --git a/tests/id_factory_test.py b/tests/id_factory_test.py new file mode 100644 index 0000000000000000000000000000000000000000..812ffbfe7635d8b83e6a58a52230b17db36b018f --- /dev/null +++ b/tests/id_factory_test.py @@ -0,0 +1,49 @@ +# -*- coding: utf-8 -*- +"""Tests for the configurable ID factory.""" +import re +from unittest.async_case import IsolatedAsyncioTestCase + +from agentscope import set_id_factory +from agentscope.message import Msg, TextBlock + +_HEX32_RE = re.compile(r"^[0-9a-f]{32}$") + + +class IdFactoryTest(IsolatedAsyncioTestCase): + """Tests for set_id_factory.""" + + async def asyncSetUp(self) -> None: + """Save the current factory before each test.""" + import agentscope._utils._common as common + + # pylint: disable=protected-access + self._saved_factory = common._id_factory + + async def test_default_id_factory_returns_hex32(self) -> None: + """The default ID factory returns uuid.uuid4().hex.""" + msg = Msg( + name="test", + content=[TextBlock(text="hello")], + role="user", + ) + self.assertRegex(msg.id, _HEX32_RE) + self.assertRegex(msg.content[0].id, _HEX32_RE) + + async def test_custom_factory_affects_entities(self) -> None: + """After ``set_id_factory``, entities use the custom factory.""" + set_id_factory(lambda: "custom-entity-id") + + msg = Msg( + name="test", + content=[TextBlock(text="hello")], + role="user", + ) + self.assertEqual(msg.id, "custom-entity-id") + self.assertEqual(msg.content[0].id, "custom-entity-id") + + async def asyncTearDown(self) -> None: + """Restore the original factory after each test.""" + import agentscope._utils._common as common + + # pylint: disable=protected-access + common._id_factory = self._saved_factory diff --git a/tests/in_memory_message_bus_test.py b/tests/in_memory_message_bus_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ead62bd76dcceca3945e24095fe07a79795ab60c --- /dev/null +++ b/tests/in_memory_message_bus_test.py @@ -0,0 +1,503 @@ +# -*- coding: utf-8 -*- +"""Tests for :class:`InMemoryMessageBus`. + +The same abstract surface exercised in ``service_message_bus_test.py`` +(queue / log / pubsub / lock / registry) is tested here against the +pure-Python in-memory backend, plus the domain helpers inherited from +the base :class:`MessageBus` class. + +No external dependencies (no Redis, no fakeredis) — just asyncio. +""" +import asyncio +from contextlib import AsyncExitStack +from unittest import IsolatedAsyncioTestCase + +from agentscope.app.message_bus import InMemoryMessageBus + + +class TestQueuePrimitive(IsolatedAsyncioTestCase): + """Mode A — ``queue_push`` + ``queue_drain`` semantics.""" + + async def asyncSetUp(self) -> None: + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context( + InMemoryMessageBus(), + ) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + + async def test_push_drain_returns_payloads_in_order(self) -> None: + """Entries pushed in order come back out in order, once each.""" + await self.bus.queue_push("k", {"i": 1}) + await self.bus.queue_push("k", {"i": 2}) + entries = await self.bus.queue_drain("k", max_count=10) + self.assertEqual([p for _id, p in entries], [{"i": 1}, {"i": 2}]) + + async def test_drain_is_destructive(self) -> None: + """A drained entry is gone; a second drain yields nothing.""" + await self.bus.queue_push("k", {"x": 1}) + await self.bus.queue_drain("k", max_count=10) + self.assertEqual(await self.bus.queue_drain("k", max_count=10), []) + + async def test_drain_respects_max_count(self) -> None: + """``max_count`` caps the batch size; remaining entries persist.""" + for i in range(5): + await self.bus.queue_push("k", {"i": i}) + first = await self.bus.queue_drain("k", max_count=3) + rest = await self.bus.queue_drain("k", max_count=10) + self.assertEqual([p["i"] for _id, p in first], [0, 1, 2]) + self.assertEqual([p["i"] for _id, p in rest], [3, 4]) + + async def test_drain_empty_queue_returns_empty(self) -> None: + """Draining a key that was never pushed returns an empty list.""" + self.assertEqual(await self.bus.queue_drain("nope"), []) + + async def test_push_returns_unique_ids(self) -> None: + """Each ``queue_push`` returns a distinct entry id.""" + id1 = await self.bus.queue_push("k", {"a": 1}) + id2 = await self.bus.queue_push("k", {"a": 2}) + self.assertNotEqual(id1, id2) + + async def test_queue_delete_removes_all(self) -> None: + """``queue_delete`` drops the entire queue.""" + await self.bus.queue_push("k", {"i": 1}) + await self.bus.queue_push("k", {"i": 2}) + await self.bus.queue_delete("k") + self.assertEqual(await self.bus.queue_drain("k", max_count=10), []) + + async def test_queue_delete_missing_is_noop(self) -> None: + """Deleting a non-existent queue does not raise.""" + await self.bus.queue_delete("never-existed") + + async def test_queue_isolation_between_keys(self) -> None: + """Pushes to different keys are independent.""" + await self.bus.queue_push("a", {"x": 1}) + await self.bus.queue_push("b", {"x": 2}) + a = await self.bus.queue_drain("a", max_count=10) + b = await self.bus.queue_drain("b", max_count=10) + self.assertEqual([p for _id, p in a], [{"x": 1}]) + self.assertEqual([p for _id, p in b], [{"x": 2}]) + + +class TestLogPrimitive(IsolatedAsyncioTestCase): + """Mode C — replay log: append / read with cursor / trim.""" + + async def asyncSetUp(self) -> None: + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context( + InMemoryMessageBus(), + ) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + + async def test_read_returns_everything_when_no_cursor(self) -> None: + """Without a ``since`` cursor, the whole log comes back.""" + await self.bus.log_append("k", {"i": 1}) + await self.bus.log_append("k", {"i": 2}) + entries = await self.bus.log_read("k") + self.assertEqual([p["i"] for _id, p in entries], [1, 2]) + + async def test_read_with_cursor_is_exclusive(self) -> None: + """``since=last_id`` skips that id and returns only newer.""" + await self.bus.log_append("k", {"i": 1}) + await self.bus.log_append("k", {"i": 2}) + await self.bus.log_append("k", {"i": 3}) + all_entries = await self.bus.log_read("k") + cursor = all_entries[1][0] # id of entry 2 + rest = await self.bus.log_read("k", since=cursor) + self.assertEqual([p["i"] for _id, p in rest], [3]) + + async def test_read_respects_max_count(self) -> None: + """``max_count`` caps the batch; remaining entries persist.""" + for i in range(5): + await self.bus.log_append("k", {"i": i}) + first = await self.bus.log_read("k", max_count=3) + self.assertEqual([p["i"] for _id, p in first], [0, 1, 2]) + + async def test_read_is_non_destructive(self) -> None: + """Multiple reads on the same log return the same data.""" + await self.bus.log_append("k", {"i": 1}) + r1 = await self.bus.log_read("k") + r2 = await self.bus.log_read("k") + self.assertEqual( + [p["i"] for _id, p in r1], + [p["i"] for _id, p in r2], + ) + + async def test_read_empty_log(self) -> None: + """Reading a log that never had entries returns ``[]``.""" + self.assertEqual(await self.bus.log_read("nope"), []) + + async def test_read_all_before_cursor(self) -> None: + """When all entries are at or before the cursor, result is ``[]``.""" + id2 = await self.bus.log_append("k", {"i": 2}) + self.assertEqual(await self.bus.log_read("k", since=id2), []) + + async def test_trim_without_before_drops_entire_log(self) -> None: + """``log_trim(key)`` empties the log.""" + await self.bus.log_append("k", {"i": 1}) + await self.bus.log_append("k", {"i": 2}) + await self.bus.log_trim("k") + self.assertEqual(await self.bus.log_read("k"), []) + + async def test_trim_with_before_id_keeps_newer(self) -> None: + """``log_trim(key, before_id)`` drops older entries only.""" + await self.bus.log_append("k", {"i": 1}) + id2 = await self.bus.log_append("k", {"i": 2}) + await self.bus.log_append("k", {"i": 3}) + await self.bus.log_trim("k", before_id=id2) + entries = await self.bus.log_read("k") + self.assertEqual([p["i"] for _id, p in entries], [2, 3]) + + async def test_trim_missing_key_is_noop(self) -> None: + """Trimming a non-existent log does not raise.""" + await self.bus.log_trim("nope") + + async def test_max_len_caps_log_size(self) -> None: + """``max_len`` on ``log_append`` trims older entries when the + log exceeds the cap.""" + for i in range(10): + await self.bus.log_append("k", {"i": i}, max_len=5) + entries = await self.bus.log_read("k", max_count=100) + self.assertLessEqual(len(entries), 5) + # The newest entries must survive. + self.assertEqual(entries[-1][1]["i"], 9) + + +class TestPubSubPrimitive(IsolatedAsyncioTestCase): + """Mode D — transient broadcast: publish / subscribe.""" + + async def asyncSetUp(self) -> None: + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context( + InMemoryMessageBus(), + ) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + + async def test_subscribe_receives_messages_published_after_ready( + self, + ) -> None: + """Subscribers receive payloads published after the subscription + is established.""" + ready = asyncio.Event() + received: list[dict] = [] + + async def _consumer() -> None: + async for payload in self.bus.subscribe( + "ch", + on_ready=ready.set, + ): + received.append(payload) + if len(received) == 2: + break + + task = asyncio.create_task(_consumer()) + await asyncio.wait_for(ready.wait(), timeout=2.0) + + await self.bus.publish("ch", {"i": 1}) + await self.bus.publish("ch", {"i": 2}) + await asyncio.wait_for(task, timeout=2.0) + self.assertEqual([p["i"] for p in received], [1, 2]) + + async def test_publish_without_subscribers_is_noop(self) -> None: + """Publishing when no one is listening does not raise.""" + await self.bus.publish("ch", {"i": 1}) + + async def test_multiple_subscribers_each_receive(self) -> None: + """All active subscribers on a channel receive the payload.""" + ready1 = asyncio.Event() + ready2 = asyncio.Event() + r1: list[dict] = [] + r2: list[dict] = [] + + async def _c1() -> None: + async for payload in self.bus.subscribe( + "ch", + on_ready=ready1.set, + ): + r1.append(payload) + break + + async def _c2() -> None: + async for payload in self.bus.subscribe( + "ch", + on_ready=ready2.set, + ): + r2.append(payload) + break + + t1 = asyncio.create_task(_c1()) + t2 = asyncio.create_task(_c2()) + await asyncio.wait_for(ready1.wait(), timeout=2.0) + await asyncio.wait_for(ready2.wait(), timeout=2.0) + + await self.bus.publish("ch", {"x": 42}) + + await asyncio.wait_for(t1, timeout=2.0) + await asyncio.wait_for(t2, timeout=2.0) + self.assertEqual(r1, [{"x": 42}]) + self.assertEqual(r2, [{"x": 42}]) + + +class TestLockPrimitive(IsolatedAsyncioTestCase): + """Mode E — distributed mutex (process-local asyncio.Lock).""" + + async def asyncSetUp(self) -> None: + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context( + InMemoryMessageBus(), + ) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + + async def test_is_locked_reflects_acquire_release(self) -> None: + """``is_locked`` flips to True while the body runs and back to + False once the context exits.""" + self.assertFalse(await self.bus.is_locked("k")) + async with self.bus.acquire_lock("k", ttl_secs=10): + self.assertTrue(await self.bus.is_locked("k")) + self.assertFalse(await self.bus.is_locked("k")) + + async def test_second_acquirer_waits_until_release(self) -> None: + """A second ``acquire_lock`` on the same key blocks until the + first releases.""" + order: list[str] = [] + + async def _holder() -> None: + async with self.bus.acquire_lock("k", ttl_secs=10): + order.append("first-in") + await asyncio.sleep(0.05) + order.append("first-out") + + async def _challenger() -> None: + await asyncio.sleep(0.005) + async with self.bus.acquire_lock("k", ttl_secs=10): + order.append("second-in") + + await asyncio.gather(_holder(), _challenger()) + self.assertEqual( + order, + ["first-in", "first-out", "second-in"], + ) + + +class TestRegistryPrimitive(IsolatedAsyncioTestCase): + """Mode F — ``registry_*`` hash-keyed namespace operations.""" + + async def asyncSetUp(self) -> None: + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context( + InMemoryMessageBus(), + ) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + + async def test_set_then_exists_and_getall(self) -> None: + """``registry_set`` stores a field; ``exists`` and ``getall`` + round-trip correctly.""" + await self.bus.registry_set("ns", "f1", "v1") + await self.bus.registry_set("ns", "f2", "v2") + + self.assertTrue(await self.bus.registry_exists("ns", "f1")) + self.assertTrue(await self.bus.registry_exists("ns", "f2")) + self.assertFalse(await self.bus.registry_exists("ns", "missing")) + self.assertFalse(await self.bus.registry_exists("other-ns", "f1")) + + self.assertEqual( + await self.bus.registry_getall("ns"), + {"f1": "v1", "f2": "v2"}, + ) + + async def test_set_overwrites_existing_field(self) -> None: + """A second ``registry_set`` for the same field overwrites.""" + await self.bus.registry_set("ns", "f", "v1") + await self.bus.registry_set("ns", "f", "v2") + self.assertEqual( + await self.bus.registry_getall("ns"), + {"f": "v2"}, + ) + + async def test_del_removes_only_the_named_field(self) -> None: + """``registry_del`` removes a single field; siblings survive.""" + await self.bus.registry_set("ns", "keep", "k") + await self.bus.registry_set("ns", "drop", "d") + await self.bus.registry_del("ns", "drop") + self.assertFalse(await self.bus.registry_exists("ns", "drop")) + self.assertTrue(await self.bus.registry_exists("ns", "keep")) + + async def test_del_missing_field_is_noop(self) -> None: + """Deleting a non-existent field does not raise.""" + await self.bus.registry_del("ns", "nope") + + async def test_getall_on_missing_namespace_returns_empty(self) -> None: + """``registry_getall`` for an unknown namespace returns ``{}``.""" + self.assertEqual(await self.bus.registry_getall("ghost"), {}) + + async def test_drop_deletes_entire_namespace(self) -> None: + """``registry_drop`` removes every field under the namespace.""" + await self.bus.registry_set("ns", "f1", "v1") + await self.bus.registry_set("ns", "f2", "v2") + await self.bus.registry_drop("ns") + self.assertEqual(await self.bus.registry_getall("ns"), {}) + + async def test_drop_missing_namespace_is_noop(self) -> None: + """Dropping a namespace that was never written does not raise.""" + await self.bus.registry_drop("never-existed") + + async def test_getall_returns_copy(self) -> None: + """Mutating the returned dict does not affect bus state.""" + await self.bus.registry_set("ns", "f", "v") + out = await self.bus.registry_getall("ns") + out["injected"] = "evil" + self.assertEqual( + await self.bus.registry_getall("ns"), + {"f": "v"}, + ) + + +class TestDomainHelpers(IsolatedAsyncioTestCase): + """Domain helpers inherited from ``MessageBus`` work end-to-end + on the in-memory backend.""" + + async def asyncSetUp(self) -> None: + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context( + InMemoryMessageBus(), + ) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + + async def test_session_run_trims_log_on_exit(self) -> None: + """``session_run`` + ``session_publish_event`` + auto trim.""" + sid = "s-trim" + async with self.bus.session_run(sid): + await self.bus.session_publish_event(sid, {"i": 1}) + await self.bus.session_publish_event(sid, {"i": 2}) + mid = await self.bus.session_read_events(sid) + self.assertEqual([p["i"] for _id, p in mid], [1, 2]) + self.assertEqual(await self.bus.session_read_events(sid), []) + + async def test_session_is_running_reflects_lock(self) -> None: + """``session_is_running`` is True while inside ``session_run``.""" + sid = "s-isrun" + self.assertFalse(await self.bus.session_is_running(sid)) + async with self.bus.session_run(sid): + self.assertTrue(await self.bus.session_is_running(sid)) + self.assertFalse(await self.bus.session_is_running(sid)) + + async def test_inbox_round_trip(self) -> None: + """``inbox_push`` + ``inbox_drain`` FIFO semantics.""" + sid = "s-inbox" + await self.bus.inbox_push(sid, {"hint": "a"}) + await self.bus.inbox_push(sid, {"hint": "b"}) + entries = await self.bus.inbox_drain(sid, max_count=10) + self.assertEqual( + [p["hint"] for _id, p in entries], + ["a", "b"], + ) + self.assertEqual( + await self.bus.inbox_drain(sid, max_count=10), + [], + ) + + async def test_enqueue_wakeup_round_trip(self) -> None: + """``enqueue_wakeup`` → ``dequeue_wakeups`` round-trip.""" + ready = asyncio.Event() + received: list[dict] = [] + + async def _signal_consumer() -> None: + async for payload in self.bus.subscribe_wakeup_signal( + on_ready=ready.set, + ): + received.append(payload) + break + + task = asyncio.create_task(_signal_consumer()) + await asyncio.wait_for(ready.wait(), timeout=2.0) + + await self.bus.enqueue_wakeup( + user_id="u", + session_id="s", + agent_id="a", + ) + await asyncio.wait_for(task, timeout=2.0) + self.assertEqual(len(received), 1) + + entries = await self.bus.dequeue_wakeups(max_count=10) + self.assertEqual(len(entries), 1) + self.assertEqual( + entries[0], + { + "user_id": "u", + "session_id": "s", + "agent_id": "a", + "kind": "wake", + "input": None, + }, + ) + + async def test_bg_task_round_trip(self) -> None: + """``bg_task_register / exists / list / unregister / purge`` + work on the in-memory backend.""" + sid = "s-bg" + self.assertFalse(await self.bus.bg_task_exists(sid, "t1")) + + await self.bus.bg_task_register(sid, "t1", '{"tool":"a"}') + await self.bus.bg_task_register(sid, "t2", '{"tool":"b"}') + + self.assertTrue(await self.bus.bg_task_exists(sid, "t1")) + self.assertEqual( + await self.bus.bg_task_list(sid), + {"t1": '{"tool":"a"}', "t2": '{"tool":"b"}'}, + ) + + await self.bus.bg_task_unregister(sid, "t1") + self.assertFalse(await self.bus.bg_task_exists(sid, "t1")) + + await self.bus.bg_task_purge(sid) + self.assertEqual(await self.bus.bg_task_list(sid), {}) + + async def test_session_purge_clears_all_bus_state(self) -> None: + """``session_purge`` deletes the session's events + inbox + + bg_tasks in one call.""" + sid = "s-purge" + await self.bus.session_publish_event(sid, {"e": 1}) + await self.bus.inbox_push(sid, {"m": 1}) + await self.bus.bg_task_register(sid, "t1", "{}") + + await self.bus.session_purge(sid) + + self.assertEqual(await self.bus.session_read_events(sid), []) + self.assertEqual( + await self.bus.inbox_drain(sid, max_count=10), + [], + ) + self.assertEqual(await self.bus.bg_task_list(sid), {}) + + async def test_task_cancel_pub_sub(self) -> None: + """``task_publish_cancel`` → ``task_subscribe_cancel`` + round-trip.""" + ready = asyncio.Event() + received: list[str] = [] + + async def _consumer() -> None: + async for tid in self.bus.task_subscribe_cancel( + on_ready=ready.set, + ): + received.append(tid) + break + + task = asyncio.create_task(_consumer()) + await asyncio.wait_for(ready.wait(), timeout=2.0) + + await self.bus.task_publish_cancel("task-X") + await asyncio.wait_for(task, timeout=2.0) + + self.assertEqual(received, ["task-X"]) diff --git a/tests/index_worker_lease_test.py b/tests/index_worker_lease_test.py new file mode 100644 index 0000000000000000000000000000000000000000..369fa53b5aa8ebaac05d580543803c80f785dfa2 --- /dev/null +++ b/tests/index_worker_lease_test.py @@ -0,0 +1,196 @@ +# -*- coding: utf-8 -*- +"""Regression tests for :class:`IndexWorker.process` lease handling. + +The pipeline must stop the moment the lease has been stolen by the +sweeper — otherwise the original worker and the worker that just took +over both write the same document into the vector store, producing +duplicate chunks (PR #1926 unresolved review #discussion_r3479544207). +""" +import asyncio +from datetime import timedelta +from typing import Any +from unittest import IsolatedAsyncioTestCase + +from agentscope.app._service._index_worker import IndexWorker + + +class _LeaseStorage: + """Minimal storage stub recording lifecycle calls. + + Driven by a per-document ``renew_results`` queue so tests can stage + a "renew returns True a few times, then False" pattern that mirrors + a sweeper reaping a slow worker. + """ + + def __init__(self) -> None: + self.acquire_returns: bool = True + self.renew_results: list[bool] = [] + self.released: list[dict] = [] + self.status_updates: list[dict] = [] + self.renew_calls = 0 + + async def acquire_knowledge_document_lease( + self, + **kwargs: Any, + ) -> bool: + """Return the staged ``acquire_returns`` flag.""" + del kwargs + return self.acquire_returns + + async def renew_knowledge_document_lease( + self, + **kwargs: Any, + ) -> bool: + """Pop next staged renew result; default to ``True`` once drained.""" + del kwargs + self.renew_calls += 1 + if not self.renew_results: + return True + return self.renew_results.pop(0) + + async def release_knowledge_document_lease( + self, + **kwargs: Any, + ) -> None: + """Record the release call so tests can assert it ran.""" + self.released.append(kwargs) + + async def update_knowledge_document_status( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + status: str, + error: str | None = None, + chunk_count: int | None = None, + ) -> None: + """Record the status transition for later assertion.""" + del user_id, knowledge_base_id, document_id + self.status_updates.append( + { + "status": status, + "error": error, + "chunk_count": chunk_count, + }, + ) + + +class _SlowPipelineWorker(IndexWorker): + """Replaces ``_run_pipeline`` with a long sleep so we can race the + lease timer. + + The whole point of the regression is "what happens if a worker is + *still in_progress* when its lease is taken away" — the only way to + test that deterministically without standing up an embedding model + and a vector store is to make the pipeline trivially long-running. + """ + + def __init__(self, storage: _LeaseStorage, pipeline_seconds: float): + # Skip the real __init__ — we only need a handful of fields. + self._storage = storage # type: ignore[assignment] + self._node_id = "test-node" + self._lease_ttl = timedelta(seconds=10) + self._sem = asyncio.Semaphore(4) + # Renew quickly so a False result is visible within the test. + self._renew_interval = timedelta(seconds=0.05) + self._pipeline_seconds = pipeline_seconds + self.pipeline_started = asyncio.Event() + self.pipeline_cancelled = False + self.pipeline_completed = False + + async def _run_pipeline( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + ) -> None: + """Sleep ``pipeline_seconds`` so the test can race the heartbeat.""" + del user_id, knowledge_base_id, document_id + self.pipeline_started.set() + try: + await asyncio.sleep(self._pipeline_seconds) + self.pipeline_completed = True + except asyncio.CancelledError: + self.pipeline_cancelled = True + raise + + +class IndexWorkerLeaseTest(IsolatedAsyncioTestCase): + """Pipeline-vs-heartbeat race coverage.""" + + async def test_lost_lease_cancels_pipeline_and_marks_error(self) -> None: + """Renew returning False mid-pipeline must abort the pipeline. + + Otherwise the original worker keeps running while the new + worker (that took over the lease) also runs — both end up + inserting the same chunks into the vector store. + """ + storage = _LeaseStorage() + # First renew succeeds, second renew fails (sweeper stole it). + storage.renew_results = [True, False] + + worker = _SlowPipelineWorker(storage, pipeline_seconds=5.0) + # Bound the test so a regression hangs visibly rather than + # silently passing. + await asyncio.wait_for( + worker.process("u", "kb", "doc-1"), + timeout=3.0, + ) + + self.assertTrue( + worker.pipeline_started.is_set(), + "Pipeline never started.", + ) + self.assertTrue( + worker.pipeline_cancelled, + "Pipeline was NOT cancelled after the lease was lost — " + "this is the regression PR #1926 review flagged.", + ) + self.assertFalse( + worker.pipeline_completed, + "Pipeline ran to completion despite the lost lease.", + ) + + # _mark_error must have recorded the lost-lease reason. + errors = [u for u in storage.status_updates if u["status"] == "error"] + self.assertEqual(len(errors), 1) + self.assertIn("Lost lease", errors[0]["error"]) + + # Release is still called (and is a safe no-op server-side). + self.assertEqual(len(storage.released), 1) + + async def test_happy_path_cancels_heartbeat_and_releases(self) -> None: + """Normal completion still tears the heartbeat down cleanly.""" + storage = _LeaseStorage() + # Heartbeat always succeeds. + storage.renew_results = [] + + worker = _SlowPipelineWorker(storage, pipeline_seconds=0.05) + await asyncio.wait_for( + worker.process("u", "kb", "doc-ok"), + timeout=2.0, + ) + + self.assertTrue(worker.pipeline_completed) + self.assertFalse(worker.pipeline_cancelled) + # No error update on the happy path. + self.assertEqual( + [u for u in storage.status_updates if u["status"] == "error"], + [], + ) + self.assertEqual(len(storage.released), 1) + + async def test_not_acquired_short_circuits(self) -> None: + """When the lease is already held by another worker, do nothing.""" + storage = _LeaseStorage() + storage.acquire_returns = False + + worker = _SlowPipelineWorker(storage, pipeline_seconds=5.0) + await asyncio.wait_for( + worker.process("u", "kb", "doc-locked"), + timeout=1.0, + ) + + self.assertFalse(worker.pipeline_started.is_set()) + # No release either — we never held the lease. + self.assertEqual(storage.released, []) diff --git a/tests/mcp_sse_client_test.py b/tests/mcp_sse_client_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ccdd3039210358f104e5b002c7bf77189342a503 --- /dev/null +++ b/tests/mcp_sse_client_test.py @@ -0,0 +1,376 @@ +# -*- coding: utf-8 -*- +"""The MCP client test module in agentscope.""" +import asyncio +import json +from multiprocessing import Process +from unittest.async_case import IsolatedAsyncioTestCase + +from mcp.server import FastMCP +from pydantic import BaseModel + +from agentscope.mcp import MCPClient, HttpMCPConfig +from agentscope.message import ToolCallBlock +from agentscope.tool import ToolResponse, ToolChunk, Toolkit +from agentscope.state import AgentState + + +async def tool_1(arg1: str, arg2: list[int]) -> str: + """A test tool function. + + Args: + arg1 (`str`): + The first argument named arg1. + arg2 (`list[int]`): + The second argument named arg2. + """ + return f"arg1: {arg1}, arg2: {arg2}" + + +def setup_server() -> None: + """Set up the streamable HTTP MCP server.""" + sse_server = FastMCP("SSE", port=8003) + sse_server.tool(description="A test tool function.")(tool_1) + sse_server.run(transport="sse") + + +# --------------------------------------------------------------------------- +# Server / tool definitions for $defs preservation test +# --------------------------------------------------------------------------- + + +class _ItemConfig(BaseModel): + """Config sub-model to generate $defs in the MCP inputSchema.""" + + key: str + count: int + + +async def tool_with_model(name: str, config: _ItemConfig) -> str: + """A tool whose parameter uses a Pydantic sub-model. + + Args: + name: Item name. + config: Item configuration. + """ + return f"name={name}, key={config.key}, count={config.count}" + + +def setup_defs_server() -> None: + """Set up an SSE MCP server that exposes a tool with Pydantic + sub-models.""" + server = FastMCP("DefsSSE", port=8005) + server.tool()(tool_with_model) + server.run(transport="sse") + + +class SseMCPClientTest(IsolatedAsyncioTestCase): + """Test class for MCP server functionality.""" + + async def asyncTearDown(self) -> None: + """Tear down the test environment.""" + del self.toolkit + + while self.process.is_alive(): + self.process.terminate() + await asyncio.sleep(5) + + async def asyncSetUp(self) -> None: + """Set up the test environment.""" + self.port = 8003 + self.process = Process(target=setup_server) + self.process.start() + await asyncio.sleep(10) + + self.toolkit = Toolkit() + self.schemas = [ + { + "type": "function", + "function": { + "name": "mcp__test_sse_client__tool_1", + "description": "A test tool function.", + "parameters": { + "type": "object", + "properties": { + "arg1": { + "type": "string", + }, + "arg2": { + "items": { + "type": "integer", + }, + "type": "array", + }, + }, + "required": [ + "arg1", + "arg2", + ], + }, + }, + }, + ] + + async def test_stateless_client(self) -> None: + """Test the stateless sse MCP client.""" + # Create stateless client (is_stateful=False) + stateless_client = MCPClient( + name="test_sse_client", + is_stateful=False, + mcp_config=HttpMCPConfig( + type="http_mcp", + url=f"http://127.0.0.1:{self.port}/sse", + ), + ) + + mcp_tool_1 = await stateless_client.get_tool("tool_1") + # Repeat to ensure idempotency + res_1: ToolChunk = await mcp_tool_1(arg1="123", arg2=[1, 2, 3]) + res_2: ToolChunk = await mcp_tool_1(arg1="345", arg2=[4, 5, 6]) + res_3: ToolChunk = await mcp_tool_1(arg1="345", arg2=[4, 5, 6]) + + self.assertEqual( + res_1.content[0].text, + "arg1: 123, arg2: [1, 2, 3]", + ) + self.assertEqual( + res_2.content[0].text, + "arg1: 345, arg2: [4, 5, 6]", + ) + self.assertEqual( + res_3.content[0].text, + "arg1: 345, arg2: [4, 5, 6]", + ) + + # Register MCPTool via Toolkit constructor + toolkit_with_mcp = Toolkit(tools=[mcp_tool_1]) + + schemas = await toolkit_with_mcp.get_tool_schemas() + + self.assertListEqual( + schemas, + self.schemas, + ) + + state = AgentState() + res_gen = toolkit_with_mcp.call_tool( + ToolCallBlock( + id="xx", + type="tool_call", + name="mcp__test_sse_client__tool_1", + input=json.dumps( + { + "arg1": "789", + "arg2": [7, 8, 9], + }, + ), + ), + state=state, + ) + + final_response = None + async for chunk in res_gen: + if isinstance(chunk, ToolResponse): + final_response = chunk + else: + self.assertIsInstance(chunk, ToolChunk) + + self.assertIsNotNone(final_response) + self.assertEqual( + final_response.content[0].text, + "arg1: 789, arg2: [7, 8, 9]", + ) + + self.toolkit.clear() + self.assertListEqual(self.toolkit.tool_groups, []) + + # Try to add the mcp client + self.toolkit = Toolkit(mcps=[stateless_client]) + self.assertListEqual( + await self.toolkit.get_tool_schemas(), + self.schemas, + ) + + self.toolkit.clear() + + async def test_stateful_client(self) -> None: + """Test the stateful sse MCP client.""" + + # Test stateful client (is_stateful=True) + stateful_client = MCPClient( + name="test_sse_client", + is_stateful=True, + mcp_config=HttpMCPConfig( + type="http_mcp", + url=f"http://127.0.0.1:{self.port}/sse", + ), + ) + + self.assertFalse(stateful_client.is_connected) + await stateful_client.connect() + + self.assertTrue(stateful_client.is_connected) + + mcp_tool_1 = await stateful_client.get_tool("tool_1") + # Repeat to ensure idempotency + res_1: ToolChunk = await mcp_tool_1(arg1="12", arg2=[1, 2]) + res_2: ToolChunk = await mcp_tool_1(arg1="34", arg2=[4, 5]) + res_3: ToolChunk = await mcp_tool_1(arg1="34", arg2=[4, 5]) + + self.assertEqual( + res_1.content[0].text, + "arg1: 12, arg2: [1, 2]", + ) + self.assertEqual( + res_2.content[0].text, + "arg1: 34, arg2: [4, 5]", + ) + self.assertEqual( + res_3.content[0].text, + "arg1: 34, arg2: [4, 5]", + ) + + # with toolkit - Register MCPTool via Toolkit constructor + toolkit_with_mcp = Toolkit(tools=[mcp_tool_1]) + + self.assertListEqual( + await toolkit_with_mcp.get_tool_schemas(), + self.schemas, + ) + + state = AgentState() + res_gen = toolkit_with_mcp.call_tool( + ToolCallBlock( + id="xx", + type="tool_call", + name="mcp__test_sse_client__tool_1", + input=json.dumps( + { + "arg1": "56", + "arg2": [5, 6], + }, + ), + ), + state=state, + ) + + final_response = None + async for chunk in res_gen: + if isinstance(chunk, ToolResponse): + final_response = chunk + else: + self.assertIsInstance(chunk, ToolChunk) + + self.assertIsNotNone(final_response) + self.assertEqual( + final_response.content[0].text, + "arg1: 56, arg2: [5, 6]", + ) + + # mcp client level test + self.toolkit.clear() + self.assertListEqual(self.toolkit.tool_groups, []) + + self.toolkit = Toolkit(mcps=[stateful_client]) + self.assertListEqual( + await self.toolkit.get_tool_schemas(), + self.schemas, + ) + + await stateful_client.close() + self.assertFalse(stateful_client.is_connected) + + +class SseSchemaDefsPreservationTest(IsolatedAsyncioTestCase): + """End-to-end tests for $defs preservation in MCP tool schemas. + + These tests start a real FastMCP server that exposes a tool whose + parameter is a Pydantic sub-model. FastMCP generates an inputSchema with + ``$defs`` for the sub-model. We verify that the schema returned by + ``await toolkit.get_tool_schemas()`` preserves those ``$defs`` and that + Pydantic-generated ``title`` fields inside ``$defs`` are stripped. + """ + + async def asyncSetUp(self) -> None: + """Start the $defs test server.""" + self.port = 8005 + self.process = Process(target=setup_defs_server) + self.process.start() + await asyncio.sleep(10) + + self.schemas = [ + { + "type": "function", + "function": { + "name": "mcp__test_defs_client__tool_with_model", + "description": "A tool whose parameter uses a " + "Pydantic sub-model.\n\n Args:\n " + "name: Item name.\n " + "config: Item configuration.\n ", + "parameters": { + "$defs": { + "_ItemConfig": { + "description": "Config sub-model to " + "generate $defs in the " + "MCP inputSchema.", + "properties": { + "key": {"type": "string"}, + "count": {"type": "integer"}, + }, + "required": ["key", "count"], + "type": "object", + }, + }, + "properties": { + "name": {"type": "string"}, + "config": {"$ref": "#/$defs/_ItemConfig"}, + }, + "required": ["name", "config"], + "type": "object", + }, + }, + }, + ] + + async def asyncTearDown(self) -> None: + """Stop the $defs test server.""" + while self.process.is_alive(): + self.process.terminate() + await asyncio.sleep(5) + + async def test_defs_preserved_and_titles_stripped(self) -> None: + """$defs from Pydantic sub-model parameters must survive the full + pipeline. + + Failure scenario (before fix): + MCPTool.__init__ only copied ``properties`` and ``required``, + so ``$defs._ItemConfig`` was silently dropped. The LLM would + receive a schema where ``config`` had an unresolvable + ``$ref: "#/$defs/_ItemConfig"``. + + Expected behaviour (after fix): + - ``MCPTool.input_schema`` contains ``$defs._ItemConfig`` + - ``await toolkit.get_tool_schemas()`` output contains ``$defs`` + with the ref resolved and Pydantic titles stripped. + """ + client = MCPClient( + name="test_defs_client", + is_stateful=False, + mcp_config=HttpMCPConfig( + type="http_mcp", + url=f"http://127.0.0.1:{self.port}/sse", + ), + ) + + mcp_tool = await client.get_tool("tool_with_model") + + # 1. input_schema must preserve $defs + self.assertIn( + "$defs", + mcp_tool.input_schema, + "MCPTool.input_schema must preserve $defs from inputSchema", + ) + + # 2. get_tool_schemas() must preserve $defs and strip titles + toolkit = Toolkit(tools=[mcp_tool]) + schemas = await toolkit.get_tool_schemas() + self.assertListEqual(schemas, self.schemas) diff --git a/tests/mcp_streamable_http_client_test.py b/tests/mcp_streamable_http_client_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e4c9913b67e2709cdeaf56713ead3b4dbe405d81 --- /dev/null +++ b/tests/mcp_streamable_http_client_test.py @@ -0,0 +1,147 @@ +# -*- coding: utf-8 -*- +"""The MCP client test module in agentscope.""" +import asyncio +from multiprocessing import Process +from unittest.async_case import IsolatedAsyncioTestCase + +from mcp.server import FastMCP +from mcp.types import EmbeddedResource, TextResourceContents + +from agentscope.mcp import MCPClient, HttpMCPConfig +from agentscope.tool import ToolChunk + + +async def tool_1(arg1: str, arg2: list[int]) -> str: + """A test tool function. + + Args: + arg1 (`str`): + The first argument named arg1. + arg2 (`list[int]`): + The second argument named arg2. + """ + return f"arg1: {arg1}, arg2: {arg2}" + + +async def tool_2() -> list: + """ + A test tool function return the EmbeddedResource type + """ + return [ + EmbeddedResource( + type="resource", + resource=TextResourceContents( + uri="file://tmp.txt", + mimeType="text/plain", + text="test content", + ), + ), + ] + + +def setup_server() -> None: + """Set up the streamable HTTP MCP server.""" + sse_server = FastMCP("StreamableHTTP", port=8002) + sse_server.tool(description="A test tool function.")(tool_1) + sse_server.tool( + description="A test tool function with embedded resource.", + )(tool_2) + sse_server.run(transport="streamable-http") + + +class StreamableHttpMCPClientTest(IsolatedAsyncioTestCase): + """Test class for streamable HTTP MCP client.""" + + async def asyncTearDown(self) -> None: + """Tear down the test environment.""" + while self.process.is_alive(): + self.process.terminate() + await asyncio.sleep(5) + + async def asyncSetUp(self) -> None: + """Set up the test environment.""" + self.port = 8002 + self.process = Process(target=setup_server) + self.process.start() + await asyncio.sleep(10) + + async def test_streamable_http_stateless_client(self) -> None: + """Test the MCP server connection functionality.""" + + # Test stateless client (is_stateful=False) + client = MCPClient( + name="test_streamable_http_stateless_client", + is_stateful=False, + mcp_config=HttpMCPConfig( + type="http_mcp", + url=f"http://127.0.0.1:{self.port}/mcp", + ), + ) + + my_tool_1 = await client.get_tool("tool_1") + res_1: ToolChunk = await my_tool_1(arg1="123", arg2=[1, 2, 3]) + self.assertEqual( + res_1.content[0].text, + "arg1: 123, arg2: [1, 2, 3]", + ) + + res_2: ToolChunk = await my_tool_1(arg1="345", arg2=[4, 5, 6]) + self.assertEqual( + res_2.content[0].text, + "arg1: 345, arg2: [4, 5, 6]", + ) + + # Test stateful client (is_stateful=True) + client = MCPClient( + name="test_streamable_http_stateful_client", + is_stateful=True, + mcp_config=HttpMCPConfig( + type="http_mcp", + url=f"http://127.0.0.1:{self.port}/mcp", + ), + ) + + self.assertFalse(client.is_connected) + await client.connect() + + self.assertTrue(client.is_connected) + + my_tool_1 = await client.get_tool("tool_1") + res_3: ToolChunk = await my_tool_1(arg1="12", arg2=[1, 2]) + self.assertEqual( + res_3.content[0].text, + "arg1: 12, arg2: [1, 2]", + ) + + res_4: ToolChunk = await my_tool_1(arg1="34", arg2=[4, 5]) + self.assertEqual( + res_4.content[0].text, + "arg1: 34, arg2: [4, 5]", + ) + + await client.close() + self.assertFalse(client.is_connected) + + async def test_embedded_content(self) -> None: + """Test the EmbeddedContent functionality.""" + # Test with stateless client (is_stateful=False) + client = MCPClient( + name="test_embedded_content", + is_stateful=False, + mcp_config=HttpMCPConfig( + type="http_mcp", + url=f"http://127.0.0.1:{self.port}/mcp", + ), + ) + + my_tool_2 = await client.get_tool("tool_2") + res: ToolChunk = await my_tool_2() + self.assertEqual( + res.content[0].text, + """{ + "uri": "file://tmp.txt/", + "mimeType": "text/plain", + "meta": null, + "text": "test content" +}""", + ) diff --git a/tests/mem0_agentscope_adapter_test.py b/tests/mem0_agentscope_adapter_test.py new file mode 100644 index 0000000000000000000000000000000000000000..4fad8ca5f320e426700f811b1a4df2e548ae1a08 --- /dev/null +++ b/tests/mem0_agentscope_adapter_test.py @@ -0,0 +1,462 @@ +# -*- coding: utf-8 -*- +# pylint: disable=wrong-import-order +"""Unit tests for the AgentScope ↔ mem0 adapters. + +Verifies that ``AgentScopeLLM`` / ``AgentScopeEmbedding`` correctly +translate between mem0's sync OpenAI-style contract and AgentScope's +async ``Msg`` / ``ContentBlock`` / ``EmbeddingResponse`` shapes, and +that ``register_with_mem0`` plugs them into mem0's factories. +""" +import asyncio +import json +import unittest +from typing import Any + +from agentscope.credential import DashScopeCredential +from agentscope.embedding import EmbeddingModelBase, EmbeddingResponse +from agentscope.message import ( + Msg, + TextBlock, + ThinkingBlock, + ToolCallBlock, +) +from agentscope.middleware._longterm_memory._mem0._agentscope_adapter import ( + AgentScopeEmbedding, + AgentScopeLLM, + _convert_messages_to_agentscope, + _parse_chat_response, + build_mem0_config, +) +from agentscope.model import ChatResponse +from utils import MockModel + + +# ---------------------------------------------------------------------- +# Pure helpers +# ---------------------------------------------------------------------- + + +class TestConvertMessages(unittest.TestCase): + """Tests for converting mem0 dict messages into AgentScope messages.""" + + def test_three_roles_map_to_correct_role(self) -> None: + """System, user, and assistant roles should be preserved.""" + msgs = _convert_messages_to_agentscope( + [ + {"role": "system", "content": "you are helpful"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ], + ) + self.assertEqual( + [m.role for m in msgs], + ["system", "user", "assistant"], + ) + self.assertTrue(all(isinstance(m, Msg) for m in msgs)) + self.assertEqual(msgs[1].get_text_content(), "hi") + + def test_unknown_role_dropped(self) -> None: + """Unsupported message roles should be skipped.""" + msgs = _convert_messages_to_agentscope( + [ + {"role": "tool", "content": "noise"}, + {"role": "user", "content": "real"}, + ], + ) + self.assertEqual(len(msgs), 1) + self.assertEqual(msgs[0].role, "user") + + +class TestParseChatResponse(unittest.TestCase): + """Tests for converting AgentScope chat responses into mem0 output.""" + + def _resp(self, blocks: list) -> ChatResponse: + """Build a final chat response from content blocks.""" + return ChatResponse(content=blocks, is_last=True) + + def test_text_only_returns_string(self) -> None: + """Plain text responses should become plain strings.""" + resp = self._resp([TextBlock(type="text", text="answer")]) + self.assertEqual(_parse_chat_response(resp, has_tool=False), "answer") + + def test_thinking_prefixed_to_text(self) -> None: + """Thinking blocks should be preserved before visible text.""" + resp = self._resp( + [ + ThinkingBlock(type="thinking", thinking="hmm"), + TextBlock(type="text", text="final"), + ], + ) + out = _parse_chat_response(resp, has_tool=False) + self.assertIn("hmm", out) + self.assertIn("final", out) + # thinking comes first to mirror v1 order + self.assertLess(out.index("hmm"), out.index("final")) + + def test_tool_call_with_json_string_input(self) -> None: + """v2 stores ToolCallBlock.input as a JSON string — adapter + must parse it back to a dict for mem0.""" + resp = self._resp( + [ + TextBlock(type="text", text="calling tool"), + ToolCallBlock( + type="tool_call", + id="call_1", + name="lookup", + input=json.dumps({"q": "alice"}), + ), + ], + ) + out = _parse_chat_response(resp, has_tool=True) + self.assertEqual(out["content"], "calling tool") + self.assertEqual( + out["tool_calls"], + [{"name": "lookup", "arguments": {"q": "alice"}}], + ) + + def test_tool_call_with_malformed_input_keeps_raw(self) -> None: + """Malformed JSON tool inputs should remain as raw strings.""" + resp = self._resp( + [ + ToolCallBlock( + type="tool_call", + id="call_2", + name="lookup", + input="not json", + ), + ], + ) + out = _parse_chat_response(resp, has_tool=True) + self.assertEqual(out["tool_calls"][0]["arguments"], "not json") + + def test_empty_content(self) -> None: + """Empty responses should convert to the empty mem0 shapes.""" + resp = self._resp([]) + self.assertEqual(_parse_chat_response(resp, has_tool=False), "") + self.assertEqual( + _parse_chat_response(resp, has_tool=True), + {"content": "", "tool_calls": []}, + ) + + +# ---------------------------------------------------------------------- +# AgentScopeLLM end-to-end (fake AgentScope model on caller event loop) +# ---------------------------------------------------------------------- + + +class _CurrentEventLoopTestCase(unittest.TestCase): + """Provides a current event loop for the adapter's sync bridge.""" + + _event_loop: asyncio.AbstractEventLoop + _previous_event_loop: asyncio.AbstractEventLoop | None + + def setUp(self) -> None: + """Install a fresh event loop for each sync-bridge test.""" + super().setUp() + try: + self._previous_event_loop = asyncio.get_event_loop() + except RuntimeError: + self._previous_event_loop = None + self._event_loop = asyncio.new_event_loop() + asyncio.set_event_loop(self._event_loop) + + def tearDown(self) -> None: + """Restore the previous event loop after each test.""" + asyncio.set_event_loop(self._previous_event_loop) + self._event_loop.close() + super().tearDown() + + +class _RecordingMockChatModel(MockModel): + """Captures the ``messages`` arg so we can assert the v2 Msg + objects mem0's dict messages were converted into.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + """Initialize the recording model.""" + super().__init__(*args, **kwargs) + self.received_messages: list[list[Msg]] = [] + + async def _call_api(self, *args: Any, **kwargs: Any) -> Any: + """Record delivered messages before delegating to MockModel.""" + self.received_messages.append(list(kwargs.get("messages") or [])) + return await super()._call_api(*args, **kwargs) + + +class TestAgentScopeLLM(_CurrentEventLoopTestCase): + """End-to-end tests for the mem0 LLM adapter.""" + + def test_constructor_rejects_non_chatmodel(self) -> None: + """The LLM adapter should reject non-AgentScope chat models.""" + with self.assertRaises(TypeError): + AgentScopeLLM(config={"model": object()}) + + def test_constructor_requires_model(self) -> None: + """The LLM adapter should require a model config entry.""" + with self.assertRaises(ValueError): + AgentScopeLLM(config={}) + + def test_generate_response_routes_through_agentscope_model(self) -> None: + """mem0 generation should call the wrapped AgentScope model.""" + model = _RecordingMockChatModel() + model.set_responses( + [ + ChatResponse( + content=[TextBlock(type="text", text="from agentscope")], + is_last=True, + ), + ], + ) + llm = AgentScopeLLM(config={"model": model}) + + result = llm.generate_response( + [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hello"}, + ], + ) + self.assertEqual(result, "from agentscope") + # The dict messages were converted to Msg objects with the + # correct roles preserved. + delivered = model.received_messages[0] + self.assertEqual([m.role for m in delivered], ["system", "user"]) + + def test_generate_response_with_tools(self) -> None: + """Tool-call responses should be converted to mem0's tool shape.""" + model = _RecordingMockChatModel() + model.set_responses( + [ + ChatResponse( + content=[ + ToolCallBlock( + type="tool_call", + id="call_3", + name="search", + input=json.dumps({"q": "x"}), + ), + ], + is_last=True, + ), + ], + ) + llm = AgentScopeLLM(config={"model": model}) + + result = llm.generate_response( + [{"role": "user", "content": "find x"}], + tools=[{"name": "search"}], + ) + self.assertIsInstance(result, dict) + self.assertEqual( + result["tool_calls"], + [{"name": "search", "arguments": {"q": "x"}}], + ) + + def test_streaming_response_drained_to_last_chunk(self) -> None: + """When the AgentScope model streams, the adapter must + consume all chunks and use the last (which carries the + complete content per AgentScope's streaming contract).""" + model = _RecordingMockChatModel() + # MockModel.set_responses with a list-of-list triggers stream mode + model.set_responses( + [ + [ + ChatResponse( + content=[TextBlock(type="text", text="part 1")], + is_last=False, + ), + ChatResponse( + content=[TextBlock(type="text", text="final")], + is_last=True, + ), + ], + ], + ) + llm = AgentScopeLLM(config={"model": model}) + result = llm.generate_response( + [{"role": "user", "content": "stream"}], + ) + self.assertEqual(result, "final") + + def test_empty_messages_raises(self) -> None: + """Dropping all unsupported messages should raise ValueError.""" + llm = AgentScopeLLM(config={"model": _RecordingMockChatModel()}) + with self.assertRaises(ValueError): + llm.generate_response([{"role": "tool", "content": "ignored"}]) + + +# ---------------------------------------------------------------------- +# AgentScopeEmbedding +# ---------------------------------------------------------------------- + + +class _FakeEmbeddingModel(EmbeddingModelBase): + """Minimal embedding model that records requested texts.""" + + def __init__(self) -> None: + """Initialize the fake embedding model.""" + super().__init__( + credential=DashScopeCredential(api_key="fake"), + model="fake-embed", + parameters=self.Parameters(), + context_size=8192, + batch_size=10, + max_retries=0, + retry_delay=0.0, + dimensions=3, + ) + self.received: list[list[str]] = [] + + async def _call_api( + self, + inputs: list[str], + **kwargs: Any, + ) -> EmbeddingResponse: + """Return a fixed vector for every input text.""" + self.received.append(list(inputs)) + return EmbeddingResponse( + embeddings=[[0.1, 0.2, 0.3] for _ in inputs], + source="api", + ) + + +class TestAgentScopeEmbedding(_CurrentEventLoopTestCase): + """Tests for the mem0 embedding adapter.""" + + def test_constructor_validation(self) -> None: + """The embedding adapter should validate model config.""" + with self.assertRaises(ValueError): + AgentScopeEmbedding(config={}) + with self.assertRaises(TypeError): + AgentScopeEmbedding(config={"model": object()}) + + def test_embed_single_string(self) -> None: + """Single-string inputs should be wrapped before model calls.""" + model = _FakeEmbeddingModel() + emb = AgentScopeEmbedding(config={"model": model}) + result = emb.embed("hello") + self.assertEqual(result, [0.1, 0.2, 0.3]) + # The string was wrapped into a list before reaching the model. + self.assertEqual(model.received, [["hello"]]) + + def test_embed_list_of_strings_returns_first(self) -> None: + """mem0's contract is that ``embed`` returns ONE vector. We + return the first.""" + model = _FakeEmbeddingModel() + emb = AgentScopeEmbedding(config={"model": model}) + result = emb.embed(["a", "b"]) + self.assertEqual(result, [0.1, 0.2, 0.3]) + self.assertEqual(model.received, [["a", "b"]]) + + +# ---------------------------------------------------------------------- +# Factory registration +# ---------------------------------------------------------------------- + + +class TestBuildMem0Config(unittest.TestCase): + """``build_mem0_config`` must bypass mem0's hardcoded provider + whitelist and emit a config that names the AgentScope adapter.""" + + def test_models_only_produces_fresh_config(self) -> None: + """Models-only construction should create AgentScope providers.""" + chat = MockModel() + emb = _FakeEmbeddingModel() + cfg = build_mem0_config(chat_model=chat, embedding_model=emb) + + self.assertEqual(cfg.llm.provider, "agentscope") + self.assertEqual(cfg.embedder.provider, "agentscope") + self.assertIs(cfg.llm.config["model"], chat) + self.assertIs(cfg.embedder.config["model"], emb) + + def test_models_only_requires_both(self) -> None: + """Models-only construction should require chat and embedding.""" + with self.assertRaises(ValueError): + build_mem0_config(chat_model=MockModel()) + with self.assertRaises(ValueError): + build_mem0_config(embedding_model=_FakeEmbeddingModel()) + with self.assertRaises(ValueError): + build_mem0_config() + + def test_base_config_with_both_models_preserves_other_fields( + self, + ) -> None: + """When a ``mem0_config`` base is given, vector_store / + history_db / etc. should survive — only .llm and .embedder + are rewired to the AgentScope adapters.""" + from mem0.configs.base import MemoryConfig + + base = MemoryConfig(history_db_path="/tmp/custom_history.db") + original_vs = base.vector_store + + chat = MockModel() + emb = _FakeEmbeddingModel() + cfg = build_mem0_config( + chat_model=chat, + embedding_model=emb, + mem0_config=base, + ) + + self.assertEqual(cfg.llm.provider, "agentscope") + self.assertIs(cfg.llm.config["model"], chat) + self.assertEqual(cfg.embedder.provider, "agentscope") + self.assertIs(cfg.embedder.config["model"], emb) + # Non-llm/embedder fields preserved. + self.assertEqual(cfg.history_db_path, "/tmp/custom_history.db") + self.assertIs(cfg.vector_store, original_vs) + + def test_base_config_with_only_chat_model_partial_override( + self, + ) -> None: + """Partial override: chat_model alone replaces .llm but + leaves .embedder untouched (whatever the base config had).""" + from mem0.configs.base import MemoryConfig + + base = MemoryConfig() # base has the default openai embedder + original_embedder = base.embedder + + cfg = build_mem0_config( + chat_model=MockModel(), + mem0_config=base, + ) + self.assertEqual(cfg.llm.provider, "agentscope") + # Embedder unchanged from base — still mem0's openai default. + self.assertIs(cfg.embedder, original_embedder) + + def test_base_config_alone_is_pass_through(self) -> None: + """``mem0_config=base`` with no models is just a pass-through — + registration still happens (cheap) but no fields change.""" + from mem0.configs.base import MemoryConfig + + base = MemoryConfig() + cfg = build_mem0_config(mem0_config=base) + self.assertIs(cfg, base) + + def test_registers_adapter_in_factory(self) -> None: + """The helper side-effect: mem0's factory dicts now know how + to resolve provider='agentscope'.""" + from mem0.utils.factory import EmbedderFactory, LlmFactory + + build_mem0_config( + chat_model=MockModel(), + embedding_model=_FakeEmbeddingModel(), + ) + self.assertIn("agentscope", LlmFactory.provider_to_class) + self.assertIn("agentscope", EmbedderFactory.provider_to_class) + + def test_naive_from_config_path_still_rejected(self) -> None: + """Sanity check: confirm WHY the helper exists — calling + ``MemoryConfig`` with ``provider='agentscope'`` directly DOES + raise, which is the failure ``build_mem0_config`` works around.""" + from mem0.configs.base import MemoryConfig + from pydantic import ValidationError + + with self.assertRaises(ValidationError): + MemoryConfig( + llm={ + "provider": "agentscope", + "config": {"model": MockModel()}, + }, + embedder={ + "provider": "agentscope", + "config": {"model": _FakeEmbeddingModel()}, + }, + ) diff --git a/tests/mem0_middleware_test.py b/tests/mem0_middleware_test.py new file mode 100644 index 0000000000000000000000000000000000000000..d587ed90d3ee8b880b2cae88eabee1541fe46196 --- /dev/null +++ b/tests/mem0_middleware_test.py @@ -0,0 +1,891 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access,unused-argument +"""Unit tests for Mem0Middleware. + +The mem0 client itself is mocked — we only exercise the AgentScope hook +wiring (retrieve-before / write-after, system-prompt injection, +list_tools exposure) and the small adapter that translates between +AgentScope and mem0. The OSS ``Memory`` and Platform ``MemoryClient`` +share the same call shape (``search(query, filters=..., top_k=...)`` +and ``add(messages, user_id=..., agent_id=...)``) so one mock covers +both. + +``protected-access`` is disabled because tests legitimately reach into +middleware internals (``mw._client``, ``mw._async_search``, ``mw._top_k``) to +inspect what the public API just did. +""" +from unittest.async_case import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, MagicMock +from typing import Any + +from utils import MockModel + +from agentscope.agent import Agent +from agentscope.event import ( + ExternalExecutionResultEvent, + UserConfirmResultEvent, +) +from agentscope.message import HintBlock, Msg, TextBlock, UserMsg +from agentscope.middleware import Mem0Middleware +from agentscope.middleware._longterm_memory._mem0._utils import ( + _extract_memory_texts, + _extract_query_text, +) +from agentscope.model import ChatResponse +from agentscope.tool import Toolkit + + +# ---------------------------------------------------------------------- +# Test helpers +# ---------------------------------------------------------------------- + + +class RecordingMockModel(MockModel): + """MockModel that captures the ``messages`` of every _call_api call.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + """Initialize the recording mock model.""" + super().__init__(*args, **kwargs) + self.calls: list[list[Msg]] = [] + + async def _call_api( + self, + *args: Any, + **kwargs: Any, + ) -> Any: + """Record messages before delegating to MockModel.""" + messages = kwargs.get("messages") + if messages is None and len(args) >= 2: + messages = args[1] + self.calls.append(list(messages or [])) + return await super()._call_api(*args, **kwargs) + + @property + def last_call_messages(self) -> list[Msg]: + """Return the most recent model-call messages.""" + return self.calls[-1] + + +class _FakeAsyncMem0Client: + """Stands in for any of mem0's four client variants. + + All four expose compatible ``search(query, filters=..., top_k=...)`` + and ``add(messages, user_id=..., agent_id=...)`` shapes, so a single + fake suffices. + """ + + def __init__(self, search_return: Any = None) -> None: + """Initialize the fake mem0 client.""" + self.search_return = search_return or {"results": []} + self.search_calls: list[dict] = [] + self.add_calls: list[dict] = [] + + async def search(self, query: str, **kwargs: Any) -> Any: + """Record a mem0 search call and return the configured result.""" + self.search_calls.append({"query": query, **kwargs}) + return self.search_return + + async def add(self, messages: list[dict], **kwargs: Any) -> None: + """Record a mem0 add call.""" + self.add_calls.append({"messages": messages, **kwargs}) + + +def _single_response(text: str) -> ChatResponse: + """Build a final single-text chat response.""" + return ChatResponse( + content=[TextBlock(type="text", text=text)], + is_last=True, + ) + + +def _all_tool_names(toolkit: Toolkit) -> list[str]: + """Flatten every registered tool's name across every group.""" + return [t.name for g in toolkit.tool_groups for t in g.tools] + + +def _find_tool(toolkit: Toolkit, name: str) -> Any: + """Look a tool up by name across every group on the toolkit.""" + for g in toolkit.tool_groups: + for t in g.tools: + if t.name == name: + return t + raise LookupError(f"tool {name!r} not found in any group") + + +def _find_group(toolkit: Toolkit, name: str) -> Any: + """Look a tool group up by name on the toolkit.""" + for g in toolkit.tool_groups: + if g.name == name: + return g + raise LookupError(f"tool group {name!r} not found") + + +def _chunk_text(chunk: Any) -> str: + """Return the first text block from a tool chunk.""" + return chunk.content[0].text + + +# ---------------------------------------------------------------------- +# Unit tests for module-level helpers +# ---------------------------------------------------------------------- + + +class TestExtractQueryText(IsolatedAsyncioTestCase): + """Tests for extracting the query text from incoming inputs.""" + + def test_none_and_empty(self) -> None: + """None and empty inputs should not produce a query.""" + self.assertIsNone(_extract_query_text(None)) + self.assertIsNone(_extract_query_text([])) + + def test_single_user_msg(self) -> None: + """A single user message should become its text content.""" + msg = UserMsg("user", "hello world") + self.assertEqual(_extract_query_text(msg), "hello world") + + def test_list_of_user_msgs_joined(self) -> None: + """Multiple user messages should be joined by newlines.""" + msgs = [UserMsg("user", "first"), UserMsg("user", "second")] + self.assertEqual(_extract_query_text(msgs), "first\nsecond") + + def test_resumption_events_return_none(self) -> None: + """HITL resumption events should not trigger memory IO.""" + self.assertIsNone( + _extract_query_text( + UserConfirmResultEvent( + reply_id="reply", + confirm_results=[], + ), + ), + ) + self.assertIsNone( + _extract_query_text( + ExternalExecutionResultEvent( + reply_id="reply", + execution_results=[], + ), + ), + ) + + +class TestExtractMemoryTexts(IsolatedAsyncioTestCase): + """Tests for normalizing mem0 search responses into text lists.""" + + def test_dict_with_results(self) -> None: + """Dictionary responses with result dicts should be flattened.""" + raw = {"results": [{"memory": "a"}, {"memory": "b"}]} + self.assertEqual(_extract_memory_texts(raw), ["a", "b"]) + + def test_plain_list_of_dicts(self) -> None: + """Plain list responses should also be supported.""" + self.assertEqual( + _extract_memory_texts([{"memory": "only"}]), + ["only"], + ) + + def test_plain_list_of_strings(self) -> None: + """String results should pass through unchanged.""" + self.assertEqual( + _extract_memory_texts({"results": ["x", "y"]}), + ["x", "y"], + ) + + def test_none_and_garbage(self) -> None: + """Malformed mem0 outputs should normalize to an empty list.""" + self.assertEqual(_extract_memory_texts(None), []) + self.assertEqual(_extract_memory_texts({"results": "nope"}), []) + + +# ---------------------------------------------------------------------- +# Constructor validation +# ---------------------------------------------------------------------- + + +class TestConstructorValidation(IsolatedAsyncioTestCase): + """Tests for Mem0Middleware constructor validation.""" + + def test_missing_user_id_raises(self) -> None: + """Empty user IDs should be rejected.""" + with self.assertRaises(ValueError): + Mem0Middleware(client=MagicMock(), user_id="") + with self.assertRaises(ValueError): + Mem0Middleware(client=MagicMock(), user_id=" ") + + def test_unknown_mode_raises(self) -> None: + """Unknown control modes should be rejected.""" + with self.assertRaises(ValueError): + Mem0Middleware( + client=MagicMock(), + user_id="alice", + mode="garbage", # type: ignore[arg-type] + ) + + def test_neither_client_nor_models_nor_config_raises(self) -> None: + """At least one backend construction path should be provided.""" + with self.assertRaises(ValueError) as ctx: + Mem0Middleware(user_id="alice") + self.assertIn("client", str(ctx.exception)) + self.assertIn("mem0_config", str(ctx.exception)) + self.assertIn("chat_model", str(ctx.exception)) + + def test_client_wins_over_other_backend_kwargs_with_warning( + self, + ) -> None: + """``client`` is the escape hatch — when given, the other + backend kwargs are ignored. A warning is logged so the + mismatch is not invisible.""" + fake_client = _FakeAsyncMem0Client() + with self.assertLogs( + "as", # the project-wide logger name + level="WARNING", + ) as captured: + mw = Mem0Middleware( + user_id="alice", + client=fake_client, + chat_model=MagicMock(), # ignored + embedding_model=MagicMock(), # ignored + mem0_config=MagicMock(), # ignored + ) + + # The fake client wired through unchanged — no AgentScope- + # adapter construction machinery ran. + self.assertIs(mw._client, fake_client) + + # Warning mentions all three ignored kwargs. + joined = "\n".join(captured.output) + self.assertIn("chat_model", joined) + self.assertIn("embedding_model", joined) + self.assertIn("mem0_config", joined) + + def test_client_alone_does_not_warn(self) -> None: + """Sanity: pure client-only path stays quiet.""" + fake_client = _FakeAsyncMem0Client() + with self.assertNoLogs("as", level="WARNING"): + Mem0Middleware(user_id="alice", client=fake_client) + + def test_only_one_of_chat_or_embedding_without_config_raises( + self, + ) -> None: + """Models-only construction requires both chat and embedding.""" + with self.assertRaises(ValueError): + Mem0Middleware(user_id="alice", chat_model=MagicMock()) + with self.assertRaises(ValueError): + Mem0Middleware(user_id="alice", embedding_model=MagicMock()) + + +# ---------------------------------------------------------------------- +# Static-control mode (default) +# ---------------------------------------------------------------------- + + +class TestStaticControlMode(IsolatedAsyncioTestCase): + """Default mode: retrieve-before / inject / write-after, no tools.""" + + async def asyncSetUp(self) -> None: + """Create a fresh recording model and empty toolkit.""" + self.model = RecordingMockModel(context_size=100_000) + self.toolkit = Toolkit() + + def _agent( + self, + middleware: Mem0Middleware, + response_text: str = "ok", + ) -> Agent: + """Build a test agent using ``middleware`` and one response.""" + self.model.set_responses([_single_response(response_text)]) + return Agent( + name="agent_under_test", + system_prompt="base system prompt", + model=self.model, + toolkit=self.toolkit, + middlewares=[middleware], + ) + + async def test_retrieve_inject_write(self) -> None: + """Static control should search, inject, reply, then write.""" + fake = _FakeAsyncMem0Client( + search_return={"results": [{"memory": "alice loves coffee"}]}, + ) + mw = Mem0Middleware( + client=fake, + user_id="alice", + agent_id="agent_under_test", + mode="static_control", + ) + + agent = self._agent(mw, response_text="hi alice") + reply = await agent.reply(UserMsg("user", "remind me what I like")) + + # 1. search called with unified call shape (filters dict, top_k) + self.assertEqual(len(fake.search_calls), 1) + self.assertEqual( + fake.search_calls[0]["query"], + "remind me what I like", + ) + self.assertEqual( + fake.search_calls[0]["filters"], + {"user_id": "alice", "agent_id": "agent_under_test"}, + ) + self.assertEqual(fake.search_calls[0]["top_k"], 5) + + # 2. write called post-turn with user+assistant pair + self.assertEqual(len(fake.add_calls), 1) + self.assertEqual( + fake.add_calls[0]["messages"], + [ + {"role": "user", "content": "remind me what I like"}, + {"role": "assistant", "content": "hi alice"}, + ], + ) + self.assertEqual(fake.add_calls[0]["user_id"], "alice") + self.assertEqual(fake.add_calls[0]["agent_id"], "agent_under_test") + + # 3. system prompt is unchanged — static_control mode does NOT + # add tool instructions (those only fire in agent_control / + # both modes via on_system_prompt). + sys_msg = self.model.last_call_messages[0] + self.assertEqual(sys_msg.role, "system") + self.assertEqual(sys_msg.get_text_content(), "base system prompt") + + # 4. memory appended to the agent's persistent context as an + # assistant-role hint note named "memory". Formatters convert + # HintBlock content into a user message before provider calls. + memory_msgs = [] + for msg in agent.state.context: + if getattr(msg, "name", None) == "memory": + memory_msgs.append(msg) + self.assertEqual(len(memory_msgs), 1) + self.assertEqual(memory_msgs[0].role, "assistant") + memory_hints = memory_msgs[0].get_content_blocks("hint") + self.assertEqual(len(memory_hints), 1) + self.assertIsInstance(memory_hints[0], HintBlock) + memory_text = memory_hints[0].hint + self.assertIn("Relevant memories", memory_text) + self.assertIn("alice loves coffee", memory_text) + # The model saw it on its first call as well. + delivered = [ + m + for m in self.model.last_call_messages + if getattr(m, "name", None) == "memory" + ] + self.assertEqual(len(delivered), 1) + + # 5. no agent-control tools registered (in any group) + all_names = _all_tool_names(self.toolkit) + self.assertNotIn("search_memory", all_names) + self.assertNotIn("add_memory", all_names) + + self.assertEqual(reply.get_text_content(), "hi alice") + + async def test_memory_message_lands_after_user_message(self) -> None: + """Mirroring AgentScope 1.x's ReActAgent placement: the memory + note is inserted right AFTER the new user input lands in the + agent context, not before.""" + fake = _FakeAsyncMem0Client( + search_return={"results": [{"memory": "saved fact"}]}, + ) + mw = Mem0Middleware( + client=fake, + user_id="alice", + mode="static_control", + ) + agent = self._agent(mw, response_text="answer") + await agent.reply(UserMsg("user", "tell me what you know")) + + # context order should be: [...prior history..., USER_MSG, + # MEMORY_NOTE, ASSISTANT_REPLY (added by the agent loop)] + roles_and_names = [ + (m.role, getattr(m, "name", None)) for m in agent.state.context + ] + user_idx = next( + i + for i, (r, n) in enumerate(roles_and_names) + if r == "user" and n != "memory" + ) + memory_idx = next( + i for i, (_, n) in enumerate(roles_and_names) if n == "memory" + ) + self.assertGreater( + memory_idx, + user_idx, + f"memory note at {memory_idx} should come after user msg " + f"at {user_idx}: {roles_and_names}", + ) + + async def test_no_memories_no_injection(self) -> None: + """Empty search results should not add a memory hint message.""" + fake = _FakeAsyncMem0Client(search_return={"results": []}) + mw = Mem0Middleware( + client=fake, + user_id="alice", + mode="static_control", + ) + + agent = self._agent(mw) + await agent.reply(UserMsg("user", "hi")) + + # System prompt unchanged, no synthetic memory msg present. + msgs = self.model.last_call_messages + self.assertEqual(msgs[0].get_text_content(), "base system prompt") + for m in msgs: + self.assertNotEqual(getattr(m, "name", None), "memory") + self.assertEqual(len(fake.search_calls), 1) + + async def test_search_failure_does_not_break_reply(self) -> None: + """Search errors should be logged but not block the reply.""" + fake = _FakeAsyncMem0Client() + fake.search = AsyncMock(side_effect=RuntimeError("mem0 down")) + mw = Mem0Middleware( + client=fake, + user_id="alice", + mode="static_control", + ) + + agent = self._agent(mw, response_text="still works") + reply = await agent.reply(UserMsg("user", "ping")) + + self.assertEqual(reply.get_text_content(), "still works") + # Write still happened — a failed search must not block writes. + self.assertEqual(len(fake.add_calls), 1) + + async def test_scope_search_by_agent_false_drops_agent_id( + self, + ) -> None: + """Unscoped search should omit agent_id from mem0 filters.""" + fake = _FakeAsyncMem0Client() + mw = Mem0Middleware( + client=fake, + user_id="alice", + agent_id="agent_under_test", + mode="static_control", + scope_search_by_agent=False, + ) + agent = self._agent(mw) + await agent.reply(UserMsg("user", "hi")) + + self.assertNotIn("agent_id", fake.search_calls[0]["filters"]) + # Write still tags both for future scoped queries. + self.assertEqual(fake.add_calls[0]["agent_id"], "agent_under_test") + + async def test_sync_client_rejected(self) -> None: + """Sync mem0 clients (``Memory`` / ``MemoryClient``) must be + rejected at construction time — only async clients are + supported.""" + + class SyncFake: + """Fake sync mem0 client used to verify rejection.""" + + def search(self, query: str, **kwargs: Any) -> Any: + """Return an empty synchronous search result.""" + return {"results": []} + + def add(self, messages: list, **kwargs: Any) -> None: + """No-op synchronous add method.""" + return None + + with self.assertRaises(TypeError) as ctx: + Mem0Middleware(client=SyncFake(), user_id="alice") + self.assertIn("AsyncMemory", str(ctx.exception)) + + async def test_async_method_wrapped_by_sync_decorator_accepted( + self, + ) -> None: + """mem0's hosted ``AsyncMemoryClient`` decorates its async + methods with ``@api_error_handler`` — a sync wrapper that + returns the coroutine produced by calling the underlying + ``async def func``. ``inspect.iscoroutinefunction(wrapper)`` + is False, so the naive check would reject the client; we use + ``inspect.unwrap`` to peel through ``@functools.wraps`` and + find the real async function. This test simulates that + pattern and asserts the wrapped-but-async client is accepted. + """ + from functools import wraps + + def sync_wraps_async(func: Any) -> Any: + """Wrap an async function in a sync functools wrapper.""" + + @wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + """Return the coroutine created by the wrapped function.""" + return func(*args, **kwargs) + + return wrapper + + class WrappedFake: + """Fake client whose async methods are sync-wrapped.""" + + @sync_wraps_async + async def search(self, query: str, **kwargs: Any) -> Any: + """Return an empty async search result.""" + return {"results": []} + + @sync_wraps_async + async def add(self, messages: list, **kwargs: Any) -> None: + """No-op async add method.""" + return None + + # Should NOT raise — the unwrap+iscoroutinefunction check sees + # through the sync wrapper. + mw = Mem0Middleware(client=WrappedFake(), user_id="alice") + self.assertIsInstance(mw._client, WrappedFake) + + +# ---------------------------------------------------------------------- +# Agent-control mode +# ---------------------------------------------------------------------- + + +class TestAgentControlMode(IsolatedAsyncioTestCase): + """Tools are listed, but no automatic memory hook behavior.""" + + async def asyncSetUp(self) -> None: + """Create a fresh recording model and empty toolkit.""" + self.model = RecordingMockModel(context_size=100_000) + self.toolkit = Toolkit() + + async def _agent( + self, + middleware: Mem0Middleware, + *, + name: str = "a", + system_prompt: str = "p", + responses: list[str] | None = None, + ) -> Agent: + """Build an agent with tools explicitly listed by middleware.""" + self.model.set_responses( + [_single_response(r) for r in (responses or ["ok"])], + ) + self.toolkit = Toolkit(tools=await middleware.list_tools()) + return Agent( + name=name, + system_prompt=system_prompt, + model=self.model, + toolkit=self.toolkit, + middlewares=[middleware], + ) + + async def test_tools_listed_and_hint_in_prompt(self) -> None: + """Agent-control mode should expose tools and prompt guidance.""" + fake = _FakeAsyncMem0Client( + search_return={"results": [{"memory": "found"}]}, + ) + mw = Mem0Middleware( + client=fake, + user_id="alice", + mode="agent_control", + ) + agent = await self._agent( + mw, + system_prompt="base prompt", + ) + await agent.reply(UserMsg("user", "hello")) + + # Tools are listed by the middleware and explicitly passed into + # the toolkit by the caller. + basic = _find_group(self.toolkit, "basic") + names = [t.name for t in basic.tools] + self.assertIn("search_memory", names) + self.assertIn("add_memory", names) + + # No automatic search/add ran. + self.assertEqual(fake.search_calls, []) + self.assertEqual(fake.add_calls, []) + + # System prompt got a short nudge mentioning the tools — the + # per-tool guidance comes via the standard tool schema, not + # this nudge. + msgs = self.model.last_call_messages + prompt_text = msgs[0].get_text_content() + self.assertIn("base prompt", prompt_text) + self.assertIn("Long-term memory", prompt_text) + self.assertIn("search_memory", prompt_text) + self.assertIn("add_memory", prompt_text) + # No synthetic memory message was injected in agent_control mode. + for m in msgs: + self.assertNotEqual(getattr(m, "name", None), "memory") + + async def test_middleware_does_not_mutate_toolkit(self) -> None: + """Middleware should not register tools unless caller does so.""" + fake = _FakeAsyncMem0Client() + mw = Mem0Middleware( + client=fake, + user_id="alice", + mode="agent_control", + ) + self.model.set_responses([_single_response("ok")]) + toolkit = Toolkit() + agent = Agent( + name="a", + system_prompt="base prompt", + model=self.model, + toolkit=toolkit, + middlewares=[mw], + ) + await agent.reply(UserMsg("user", "hello")) + + self.assertNotIn("search_memory", _all_tool_names(toolkit)) + self.assertNotIn("add_memory", _all_tool_names(toolkit)) + + async def test_search_memory_tool_invokes_mem0(self) -> None: + """The search_memory tool should query mem0 per keyword.""" + fake = _FakeAsyncMem0Client( + search_return={ + "results": [ + {"memory": "first fact"}, + {"memory": "second fact"}, + ], + }, + ) + mw = Mem0Middleware( + client=fake, + user_id="alice", + agent_id="a", + mode="agent_control", + ) + agent = await self._agent( + mw, + system_prompt="base prompt", + ) + # Trigger middleware tool registration by issuing a reply. + await agent.reply(UserMsg("user", "hi")) + + search_tool = _find_tool(self.toolkit, "search_memory") + + # Exercise the tool directly. The signature mirrors + # AgentScope 1.x: a LIST of keywords, each issued as an + # independent parallel search. + result = await search_tool( + keywords=["what does alice like?", "alice preferences"], + limit=3, + ) + result_text = _chunk_text(result) + self.assertIn("first fact", result_text) + self.assertIn("second fact", result_text) + + # Each keyword produces one mem0 call; both share user/agent + # filter and per-keyword limit. + self.assertEqual(len(fake.search_calls), 2) + for call in fake.search_calls: + self.assertEqual( + call["filters"], + {"user_id": "alice", "agent_id": "a"}, + ) + self.assertEqual(call["top_k"], 3) + + async def test_async_search_accepts_per_call_top_k(self) -> None: + """Per-call search limits should not mutate middleware state.""" + fake = _FakeAsyncMem0Client() + mw = Mem0Middleware( + client=fake, + user_id="alice", + mode="agent_control", + top_k=9, + ) + + await mw._async_search( + "q", + user_id="alice", + agent_id="a", + top_k=3, + ) + + self.assertEqual(fake.search_calls[0]["top_k"], 3) + self.assertEqual(mw._top_k, 9) + + async def test_tools_auto_allow_permission(self) -> None: + """Memory tools should be auto-allowed by permission checks.""" + mw = Mem0Middleware( + client=_FakeAsyncMem0Client(), + user_id="alice", + mode="agent_control", + ) + agent = await self._agent(mw) + await agent.reply(UserMsg("user", "hi")) + + from agentscope.permission import PermissionBehavior + + search_tool = _find_tool(self.toolkit, "search_memory") + add_tool = _find_tool(self.toolkit, "add_memory") + decision_search = await search_tool.check_permissions({}, None) + decision_add = await add_tool.check_permissions({}, None) + self.assertEqual(decision_search.behavior, PermissionBehavior.ALLOW) + self.assertEqual(decision_add.behavior, PermissionBehavior.ALLOW) + + async def test_search_memory_dedupes_across_keywords(self) -> None: + """When two keywords return overlapping memories, the tool + merges and dedupes.""" + fake = _FakeAsyncMem0Client( + search_return={ + "results": [{"memory": "shared"}, {"memory": "unique"}], + }, + ) + mw = Mem0Middleware(client=fake, user_id="alice", mode="agent_control") + agent = await self._agent(mw) + await agent.reply(UserMsg("user", "hi")) + search_tool = _find_tool(self.toolkit, "search_memory") + + result = await search_tool( + keywords=["kw1", "kw2"], + limit=5, + ) + result_text = _chunk_text(result) + # "shared" appears once even though both keywords returned it. + self.assertEqual(result_text.count("shared"), 1) + self.assertEqual(result_text.count("unique"), 1) + + async def test_search_memory_failure_returns_error_chunk(self) -> None: + """mem0 raising during search produces a ToolChunk with + state=ERROR so the toolkit marks the call as failed.""" + from agentscope.message import ToolResultState + from agentscope.tool import ToolChunk + + fake = _FakeAsyncMem0Client() + fake.search = AsyncMock(side_effect=RuntimeError("mem0 down")) + mw = Mem0Middleware(client=fake, user_id="alice", mode="agent_control") + agent = await self._agent(mw) + await agent.reply(UserMsg("user", "hi")) + search_tool = _find_tool(self.toolkit, "search_memory") + + result = await search_tool( + keywords=["q"], + limit=5, + ) + self.assertIsInstance(result, ToolChunk) + self.assertEqual(result.state, ToolResultState.ERROR) + self.assertIn("mem0 down", result.content[0].text) + + async def test_add_memory_two_tier_fallback(self) -> None: + """When mem0's extraction LLM returns no memories, the tool + falls back to ``infer=False`` and saves raw text so the + caller's ``add_memory`` invocation isn't silently dropped. + + Two tiers, not three — v1's tier-2 role switch (user → + assistant) was a no-op against current mem0 (v2.x), which + picks the extraction prompt based on the filter dict, not the + message role. See the docstring of + ``Mem0Middleware._async_add_with_fallback`` for the full + rationale.""" + + class CountingFake: + """Fake mem0 client that triggers add fallback once.""" + + def __init__(self) -> None: + """Initialize call tracking and empty search behavior.""" + self.add_calls: list[dict] = [] + self.search = AsyncMock(return_value={"results": []}) + + async def add( + self, + messages: list[dict], + **kwargs: Any, + ) -> dict: + """Return empty extraction first, then a saved memory.""" + self.add_calls.append({"messages": messages, **kwargs}) + # First attempt: extracted nothing → triggers fallback. + # Second attempt (with infer=False): succeeds. + if len(self.add_calls) < 2: + return {"results": []} + return {"results": [{"id": "m1", "memory": "saved"}]} + + fake = CountingFake() + mw = Mem0Middleware(client=fake, user_id="alice", mode="agent_control") + agent = await self._agent(mw) + await agent.reply(UserMsg("user", "hi")) + add_tool = _find_tool(self.toolkit, "add_memory") + + result = await add_tool( + thinking="my reasoning", + content=["fact one"], + ) + self.assertIn("Successfully recorded", _chunk_text(result)) + + # Exactly 2 calls — the role-switch tier from v1 is gone. + self.assertEqual(len(fake.add_calls), 2) + # Both calls use the user role (no more switching). + self.assertEqual(fake.add_calls[0]["messages"][0]["role"], "user") + self.assertEqual(fake.add_calls[1]["messages"][0]["role"], "user") + # First call uses default inference; second explicitly disables it. + self.assertNotIn("infer", fake.add_calls[0]) + self.assertFalse(fake.add_calls[1]["infer"]) + + async def test_add_memory_does_not_persist_thinking(self) -> None: + """``thinking`` is the agent's rationale — it appears in the + tool's return text (auditable in the transcript) but is NOT + sent to mem0, so the stored memory stays clean of agent + self-narration.""" + fake = _FakeAsyncMem0Client() + # First attempt extracts something → no fallback. + fake.add = AsyncMock( + return_value={"results": [{"id": "m1", "memory": "saved"}]}, + ) + mw = Mem0Middleware(client=fake, user_id="alice", mode="agent_control") + agent = await self._agent(mw) + await agent.reply(UserMsg("user", "hi")) + add_tool = _find_tool(self.toolkit, "add_memory") + + thinking = "user said X because Y" + fact = "likes coffee" + result = await add_tool( + thinking=thinking, + content=[fact], + ) + result_text = _chunk_text(result) + + # mem0 only saw the fact, NOT the rationale. + sent = fake.add.call_args.args[0][0]["content"] + self.assertIn(fact, sent) + self.assertNotIn(thinking, sent) + + # ...but the tool return text DOES echo the rationale, so the + # decision is auditable in the agent transcript. + self.assertIn(thinking, result_text) + + +# ---------------------------------------------------------------------- +# Both mode — hooks + tools together +# ---------------------------------------------------------------------- + + +class TestBothMode(IsolatedAsyncioTestCase): + """Tests for combined static-control and agent-control behavior.""" + + async def test_memory_msg_and_tool_hint_both_present(self) -> None: + """Both mode should inject memory and expose memory tools.""" + model = RecordingMockModel(context_size=100_000) + fake = _FakeAsyncMem0Client( + search_return={"results": [{"memory": "auto-injected"}]}, + ) + mw = Mem0Middleware(client=fake, user_id="alice", mode="both") + toolkit = Toolkit(tools=await mw.list_tools()) + model.set_responses([_single_response("ok")]) + agent = Agent( + name="a", + system_prompt="base", + model=model, + toolkit=toolkit, + middlewares=[mw], + ) + await agent.reply(UserMsg("user", "hi")) + + # Static hooks ran. + self.assertEqual(len(fake.search_calls), 1) + self.assertEqual(len(fake.add_calls), 1) + + msgs = model.last_call_messages + + # System prompt has the tool nudge (from on_system_prompt). + prompt_text = msgs[0].get_text_content() + self.assertIn("Long-term memory", prompt_text) + self.assertIn("search_memory", prompt_text) + + # Memory injected as a synthetic HintBlock Msg. + memory_msgs = [m for m in msgs if getattr(m, "name", None) == "memory"] + self.assertEqual(len(memory_msgs), 1) + memory_hints = memory_msgs[0].get_content_blocks("hint") + self.assertEqual(len(memory_hints), 1) + self.assertIn("auto-injected", memory_hints[0].hint) + + # Tools exposed. + names = _all_tool_names(toolkit) + self.assertIn("search_memory", names) + self.assertIn("add_memory", names) diff --git a/tests/message_test.py b/tests/message_test.py new file mode 100644 index 0000000000000000000000000000000000000000..63a9f5146f635174ffed28f87c9d5b4c5bcedf3a --- /dev/null +++ b/tests/message_test.py @@ -0,0 +1,271 @@ +# -*- coding: utf-8 -*- +"""A template test case.""" +from unittest.async_case import IsolatedAsyncioTestCase +from utils import AnyString + +from agentscope.message import ( + UserMsg, + TextBlock, + DataBlock, + URLSource, + Base64Source, + ThinkingBlock, + AssistantMsg, + HintBlock, + Msg, + ToolCallBlock, + ToolResultBlock, + ToolResultState, +) + + +class MessageTest(IsolatedAsyncioTestCase): + """The template test case.""" + + async def test_creating_message(self) -> None: + """The template test.""" + # Test string content + user_msg = UserMsg(name="user", content="hello world") + self.assertDictEqual( + user_msg.model_dump(), + { + "id": AnyString(), + "name": "user", + "role": "user", + "content": [ + { + "id": AnyString(), + "text": "hello world", + "type": "text", + }, + ], + "metadata": {}, + "created_at": AnyString(), + "finished_at": AnyString(), + "usage": None, + }, + ) + + # Test list of content + user_msg = UserMsg( + name="user", + content=[TextBlock(text="1"), TextBlock(text="2")], + ) + self.assertDictEqual( + user_msg.model_dump(), + { + "id": AnyString(), + "name": "user", + "role": "user", + "content": [ + {"type": "text", "text": "1", "id": AnyString()}, + {"type": "text", "text": "2", "id": AnyString()}, + ], + "metadata": {}, + "created_at": AnyString(), + "finished_at": AnyString(), + "usage": None, + }, + ) + + # Test DataBlock content + user_msg = UserMsg( + name="user", + content=[ + TextBlock(text="1"), + DataBlock( + source=URLSource( + url="https://example.com/image.png", + media_type="image/png", + ), + ), + DataBlock( + source=Base64Source( + data="iVBORw0KGgoAAAANSUhEUgAAAAUA", + media_type="image/png", + ), + ), + ], + ) + + self.assertDictEqual( + user_msg.model_dump(), + { + "id": AnyString(), + "name": "user", + "role": "user", + "content": [ + {"type": "text", "text": "1", "id": AnyString()}, + { + "type": "data", + "id": AnyString(), + "source": { + "type": "url", + "url": "https://example.com/image.png", + "media_type": "image/png", + }, + "name": None, + }, + { + "type": "data", + "id": AnyString(), + "source": { + "type": "base64", + "data": "iVBORw0KGgoAAAANSUhEUgAAAAUA", + "media_type": "image/png", + }, + "name": None, + }, + ], + "metadata": {}, + "created_at": AnyString(), + "finished_at": AnyString(), + "usage": None, + }, + ) + + # Test thinking content + msg = AssistantMsg( + name="assistant", + content=[ThinkingBlock(thinking="thinking...")], + ) + self.assertDictEqual( + msg.model_dump(), + { + "id": AnyString(), + "name": "assistant", + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "thinking...", + "id": AnyString(), + }, + ], + "metadata": {}, + "created_at": AnyString(), + "finished_at": None, + "usage": None, + }, + ) + + # Test hint content + msg = AssistantMsg( + name="assistant", + content=[HintBlock(hint="hint...")], + ) + self.assertDictEqual( + msg.model_dump(), + { + "id": AnyString(), + "name": "assistant", + "role": "assistant", + "content": [ + { + "type": "hint", + "hint": "hint...", + "id": AnyString(), + "source": None, + }, + ], + "metadata": {}, + "created_at": AnyString(), + "finished_at": None, + "usage": None, + }, + ) + + async def test_invalid_message(self) -> None: + """Test invalid message creation.""" + # User message with thinking block should raise ValueError + with self.assertRaises(ValueError): + Msg( + name="user", + role="user", + content=[ThinkingBlock(thinking="thinking...")], + ) + + # User message with hint block should raise ValueError + with self.assertRaises(ValueError): + Msg( + name="user", + role="user", + content=[HintBlock(hint="hint...")], + ) + + # User message with tool call block should raise ValueError + with self.assertRaises(ValueError): + Msg( + name="user", + role="user", + content=[ToolCallBlock(id="1", name="tool", input="{}")], + ) + + # User message with tool result block should raise ValueError + with self.assertRaises(ValueError): + Msg( + name="user", + role="user", + content=[ + ToolResultBlock( + id="1", + name="tool", + output="result", + state=ToolResultState.SUCCESS, + ), + ], + ) + + # System message with data block should raise ValueError + with self.assertRaises(ValueError): + Msg( + name="system", + role="system", + content=[ + DataBlock( + source=URLSource( + url="https://example.com/image.png", + media_type="image/png", + ), + ), + ], + ) + + # System message with thinking block should raise ValueError + with self.assertRaises(ValueError): + Msg( + name="system", + role="system", + content=[ThinkingBlock(thinking="thinking...")], + ) + + # System message with hint block should raise ValueError + with self.assertRaises(ValueError): + Msg( + name="system", + role="system", + content=[HintBlock(hint="hint...")], + ) + + # System message with tool call block should raise ValueError + with self.assertRaises(ValueError): + Msg( + name="system", + role="system", + content=[ToolCallBlock(id="1", name="tool", input="{}")], + ) + + # System message with tool result block should raise ValueError + with self.assertRaises(ValueError): + Msg( + name="system", + role="system", + content=[ + ToolResultBlock( + id="1", + name="tool", + output="result", + state=ToolResultState.SUCCESS, + ), + ], + ) diff --git a/tests/middleware_background_acting_test.py b/tests/middleware_background_acting_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tests/middleware_budget_test.py b/tests/middleware_budget_test.py new file mode 100644 index 0000000000000000000000000000000000000000..fe101419e01704288ca69962ec357fe69bbfbd73 --- /dev/null +++ b/tests/middleware_budget_test.py @@ -0,0 +1,476 @@ +# -*- coding: utf-8 -*- +"""Unit tests for BudgetControlMiddleware.""" +from typing import Any +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import MockModel +from agentscope.agent import Agent +from agentscope.message import UserMsg, TextBlock, ToolCallBlock, HintBlock +from agentscope.middleware import ReplyBudgetControlMiddleware +from agentscope.model import ChatResponse, ChatUsage +from agentscope.permission import ( + PermissionBehavior, + PermissionContext, + PermissionDecision, +) +from agentscope.event import UserConfirmResultEvent, ConfirmResult +from agentscope.tool import ToolBase, Toolkit, ToolChunk + + +def _response( + text: str, + input_tokens: int, + output_tokens: int, +) -> ChatResponse: + """Build a non-streaming ChatResponse with usage.""" + return ChatResponse( + content=[TextBlock(text=text)], + is_last=True, + usage=ChatUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + time=0.0, + ), + ) + + +class DummyTool(ToolBase): + """Minimal tool that always allows and returns a fixed result.""" + + name: str = "dummy" + description: str = "A dummy tool for testing" + input_schema: dict[str, Any] = {"type": "object", "properties": {}} + is_concurrency_safe: bool = True + is_read_only: bool = True + is_external_tool: bool = False + is_mcp: bool = False + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Always allow.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + decision_reason="Dummy tool always allows", + message="Dummy tool always allows", + ) + + async def __call__(self, **kwargs: Any) -> ToolChunk: + """Return a fixed result.""" + return ToolChunk(content=[TextBlock(text="ok")]) + + +class ConfirmRequiredTool(ToolBase): + """Minimal tool that always requires user confirmation before running.""" + + name: str = "confirm_required" + description: str = "A tool that requires user confirmation" + input_schema: dict[str, Any] = {"type": "object", "properties": {}} + is_concurrency_safe: bool = False + is_read_only: bool = False + is_external_tool: bool = False + is_mcp: bool = False + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Always require user confirmation.""" + return PermissionDecision( + behavior=PermissionBehavior.ASK, + decision_reason="Confirmation required", + message="Confirmation required", + ) + + async def __call__(self, **kwargs: Any) -> ToolChunk: + """Return a fixed result.""" + return ToolChunk(content=[TextBlock(text="confirmed result")]) + + +def _has_hint_block(msg: Any, hint_message: str) -> bool: + """Return True if *msg* contains a HintBlock with *hint_message*.""" + content = getattr(msg, "content", None) + if not isinstance(content, list): + return False + return any( + isinstance(b, HintBlock) and hint_message in b.hint for b in content + ) + + +class TestBudgetControlMiddleware(IsolatedAsyncioTestCase): + """Test cases for BudgetControlMiddleware.""" + + async def asyncSetUp(self) -> None: + """Set up shared fixtures.""" + self.toolkit = Toolkit() + + async def test_under_budget_no_hint_injected(self) -> None: + """When token usage stays below the budget, no hint is injected.""" + model = MockModel() + model.set_responses( + [_response("done", input_tokens=10, output_tokens=5)], + ) + + middleware = ReplyBudgetControlMiddleware(token_budget=1000) + agent = Agent( + name="test_agent", + system_prompt="you are helpful", + model=model, + toolkit=self.toolkit, + middlewares=[middleware], + ) + + context_before = len(agent.state.context) + await agent.reply(UserMsg("user", "hello")) + + # No HintBlock should have been added to context + hint_msgs = [ + m + for m in agent.state.context[context_before:] + if _has_hint_block(m, middleware.hint_message) + ] + self.assertEqual(len(hint_msgs), 0) + + async def test_budget_exceeded_injects_hint(self) -> None: + """When the budget is exceeded, the hint block is injected. + + Uses token_budget=0 so the budget condition fires on the very first + reasoning call (0 used >= 0 max). + """ + model = MockModel() + model.set_responses( + [_response("wrap up", input_tokens=10, output_tokens=5)], + ) + + middleware = ReplyBudgetControlMiddleware(token_budget=0) + agent = Agent( + name="test_agent", + system_prompt="you are helpful", + model=model, + toolkit=self.toolkit, + middlewares=[middleware], + ) + + context_before = len(agent.state.context) + await agent.reply(UserMsg("user", "hello")) + + hint_msgs = [ + m + for m in agent.state.context[context_before:] + if _has_hint_block(m, middleware.hint_message) + ] + self.assertGreater(len(hint_msgs), 0) + + async def test_budget_exceeded_forces_tool_choice_none(self) -> None: + """When budget is exceeded, tool_choice forwarded to model is none. + + Uses token_budget=0 so the override fires on the first reasoning call. + """ + received_tool_choices: list = [] + + class TrackingModel(MockModel): + """Model that records tool_choice on every call.""" + + async def _call_api( + self, + *args: Any, + **kwargs: Any, + ) -> ChatResponse: + """Record tool_choice and delegate to mock.""" + received_tool_choices.append(kwargs.get("tool_choice")) + return await super()._call_api(*args, **kwargs) + + model = TrackingModel() + model.set_responses( + [_response("wrap up", input_tokens=10, output_tokens=5)], + ) + + middleware = ReplyBudgetControlMiddleware(token_budget=0) + agent = Agent( + name="test_agent", + system_prompt="you are helpful", + model=model, + toolkit=self.toolkit, + middlewares=[middleware], + ) + + await agent.reply(UserMsg("user", "hello")) + + # At least one call must have received tool_choice with mode="none" + self.assertTrue( + any( + getattr(tc, "mode", None) == "none" + for tc in received_tool_choices + if tc is not None + ), + ) + + async def test_token_accumulation_across_steps(self) -> None: + """Tokens accumulate across steps and trigger enforcement correctly. + + Step 1: tool call costs 200+100=300 tokens (token_budget=300 so + step 2 sees used >= max and injects the hint). + """ + toolkit = Toolkit(tools=[DummyTool()]) + + model = MockModel() + model.set_responses( + [ + [ + ChatResponse( + content=[ + ToolCallBlock( + id="tc_1", + name="dummy", + input="{}", + ), + ], + is_last=True, + usage=ChatUsage( + input_tokens=200, + output_tokens=100, + time=0.0, + ), + ), + ], + [ + ChatResponse( + content=[TextBlock(text="done")], + is_last=True, + usage=ChatUsage( + input_tokens=150, + output_tokens=50, + time=0.0, + ), + ), + ], + ], + ) + + middleware = ReplyBudgetControlMiddleware(token_budget=300) + agent = Agent( + name="test_agent", + system_prompt="you are helpful", + model=model, + toolkit=toolkit, + middlewares=[middleware], + ) + + context_before = len(agent.state.context) + await agent.reply(UserMsg("user", "hello")) + + hint_msgs = [ + m + for m in agent.state.context[context_before:] + if _has_hint_block(m, middleware.hint_message) + ] + self.assertGreater(len(hint_msgs), 0) + + async def test_weighted_token_calculation(self) -> None: + """output_token_weight scales output tokens in the budget calculation. + + With input_token_weight=1, output_token_weight=3, token_budget=200: + - Step 1: 50 input * 1 + 50 output * 3 = 200 → budget hit exactly + - Step 2: hint should be injected before the model call + """ + toolkit = Toolkit(tools=[DummyTool()]) + + model = MockModel() + model.set_responses( + [ + [ + ChatResponse( + content=[ + ToolCallBlock( + id="tc_1", + name="dummy", + input="{}", + ), + ], + is_last=True, + usage=ChatUsage( + input_tokens=50, + output_tokens=50, + time=0.0, + ), + ), + ], + [ + ChatResponse( + content=[TextBlock(text="done")], + is_last=True, + usage=ChatUsage( + input_tokens=30, + output_tokens=10, + time=0.0, + ), + ), + ], + ], + ) + + # 50*1 + 50*3 = 200 == token_budget → step 2 triggers enforcement + middleware = ReplyBudgetControlMiddleware( + token_budget=200, + input_token_weight=1, + output_token_weight=3, + ) + agent = Agent( + name="test_agent", + system_prompt="you are helpful", + model=model, + toolkit=toolkit, + middlewares=[middleware], + ) + + context_before = len(agent.state.context) + await agent.reply(UserMsg("user", "hello")) + + hint_msgs = [ + m + for m in agent.state.context[context_before:] + if _has_hint_block(m, middleware.hint_message) + ] + self.assertGreater(len(hint_msgs), 0) + + async def test_state_cleanup_after_reply(self) -> None: + """middle_context entry for the reply is removed after reply ends.""" + model = MockModel() + model.set_responses( + [_response("done", input_tokens=10, output_tokens=5)], + ) + + middleware = ReplyBudgetControlMiddleware(token_budget=1000) + agent = Agent( + name="test_agent", + system_prompt="you are helpful", + model=model, + toolkit=self.toolkit, + middlewares=[middleware], + ) + + await agent.reply(UserMsg("user", "hello")) + + middleware_key = await middleware.get_middleware_key() + bucket = agent.state.middle_context.get(middleware_key, {}) + # All per-reply entries must have been cleaned up + self.assertEqual(len(bucket), 0) + + async def test_token_accumulation_persists_across_hitl(self) -> None: + """Token accumulation in middle_context persists across HITL boundary. + + Scenario: + - token_budget=300, both weights default to 1. + - First reply_stream call: model call costs 200 input + 100 output + = 300 tokens, then pauses at REQUIRE_USER_CONFIRM (no ReplyEndEvent + is emitted). The 300-token count is stored in middle_context. + - Second reply_stream call with UserConfirmResultEvent: the same + reply_id resumes. on_reasoning reads 300 >= 300 from middle_context + and injects the hint + forces tool_choice=none before the final + model call, proving budget state survived the HITL round-trip. + """ + tool_call_id = "tc_hitl" + tool_input = "{}" + toolkit = Toolkit(tools=[ConfirmRequiredTool()]) + + model = MockModel() + model.set_responses( + [ + # Step 1: model produces a tool call that requires confirmation + [ + ChatResponse( + content=[ + ToolCallBlock( + id=tool_call_id, + name="confirm_required", + input=tool_input, + ), + ], + is_last=True, + usage=ChatUsage( + input_tokens=200, + output_tokens=100, + time=0.0, + ), + ), + ], + # Step 2 (after confirmation): final wrap-up text + [ + ChatResponse( + content=[TextBlock(text="wrap up")], + is_last=True, + usage=ChatUsage( + input_tokens=50, + output_tokens=20, + time=0.0, + ), + ), + ], + ], + ) + + # 200*1 + 100*1 = 300 == token_budget → reasoning after resume + # injects hint + middleware = ReplyBudgetControlMiddleware(token_budget=300) + agent = Agent( + name="test_agent", + system_prompt="you are helpful", + model=model, + toolkit=toolkit, + middlewares=[middleware], + ) + + # --- First call: pauses at REQUIRE_USER_CONFIRM --- + events = [] + async for event in agent.reply_stream(UserMsg("user", "hello")): + events.append(event) + + event_types = [e.type for e in events] + self.assertIn("REQUIRE_USER_CONFIRM", event_types) + self.assertNotIn("REPLY_END", event_types) + + reply_id = agent.state.reply_id + + # Token count must be stored in middle_context (survived the pause) + middleware_key = await middleware.get_middleware_key() + stored = agent.state.middle_context.get(middleware_key, {}) + self.assertAlmostEqual(stored.get(reply_id, 0), 300.0) + + # --- Second call: resume with user confirmation --- + user_confirm_event = UserConfirmResultEvent( + reply_id=reply_id, + confirm_results=[ + ConfirmResult( + confirmed=True, + tool_call=ToolCallBlock( + id=tool_call_id, + name="confirm_required", + input=tool_input, + ), + ), + ], + ) + + resume_events = [] + async for event in agent.reply_stream(inputs=user_confirm_event): + resume_events.append(event) + + resume_event_types = [e.type for e in resume_events] + self.assertIn("REPLY_END", resume_event_types) + + # Hint is appended to the existing assistant message (which was created + # in the first call), so we search the full context rather than a + # slice. + hint_msgs = [ + m + for m in agent.state.context + if _has_hint_block(m, middleware.hint_message) + ] + self.assertGreater(len(hint_msgs), 0) + + # middle_context must be cleaned up after reply ends + bucket = agent.state.middle_context.get(middleware_key, {}) + self.assertNotIn(reply_id, bucket) diff --git a/tests/middleware_filesystem_memory_test.py b/tests/middleware_filesystem_memory_test.py new file mode 100644 index 0000000000000000000000000000000000000000..a4233de3ddc1e436513fc2691d542f6a6a83dcf4 --- /dev/null +++ b/tests/middleware_filesystem_memory_test.py @@ -0,0 +1,772 @@ +# -*- coding: utf-8 -*- +"""Unit tests for AgenticMemoryMiddleware with real Agent execution.""" +import os +import shutil +import tempfile +from typing import Any, Type +from unittest.async_case import IsolatedAsyncioTestCase + +from pydantic import BaseModel +from utils import AnyString, AnyValue, MockModel + +from agentscope.agent import Agent +from agentscope.message import ( + HintBlock, + Msg, + TextBlock, + ToolCallBlock, + UserMsg, +) +from agentscope.middleware import AgenticMemoryMiddleware +from agentscope.model import ChatResponse, StructuredResponse +from agentscope.permission import ( + PermissionBehavior, + PermissionContext, + PermissionDecision, +) +from agentscope.tool import ToolBase, ToolChunk, Toolkit + + +class _RecordingMockModel(MockModel): + """A ``MockModel`` that records chat and structured-output calls.""" + + def __init__(self, **kwargs: Any) -> None: + """Initialize the recording mock model. + + Args: + **kwargs (`Any`): + Keyword arguments forwarded to :class:`MockModel`. + """ + kwargs.setdefault("context_size", 100_000) + super().__init__(**kwargs) + self.chat_messages: list[list[Msg]] = [] + self.structured_messages: list[list[Msg]] = [] + + async def _call_api( + self, + *args: Any, + **kwargs: Any, + ) -> ChatResponse: + """Record the chat messages and delegate to ``MockModel``. + + Args: + *args (`Any`): + Positional arguments forwarded to ``MockModel``. + **kwargs (`Any`): + Keyword arguments forwarded to ``MockModel``. + + Returns: + `ChatResponse`: + The configured mock chat response. + """ + self.chat_messages.append(kwargs["messages"]) + return await super()._call_api(*args, **kwargs) + + async def _call_api_with_structured_output( + self, + model_name: str, + messages: list[Msg], + structured_model: Type[BaseModel] | dict, + **kwargs: Any, + ) -> StructuredResponse: + """Record structured-output messages and delegate to ``MockModel``. + + Args: + model_name (`str`): + The model name. + messages (`list[Msg]`): + The structured-output prompt messages. + structured_model (`Type[BaseModel] | dict`): + The expected structured-output schema. + **kwargs (`Any`): + Extra keyword arguments forwarded to ``MockModel``. + + Returns: + `StructuredResponse`: + The configured mock structured response. + """ + self.structured_messages.append(messages) + return await super()._call_api_with_structured_output( + model_name, + messages, + structured_model, + **kwargs, + ) + + +class _DummyTool(ToolBase): + """A minimal tool that forces a second Agent reasoning iteration.""" + + name: str = "dummy" + description: str = "A dummy tool for middleware tests." + input_schema: dict[str, Any] = {"type": "object", "properties": {}} + is_concurrency_safe: bool = True + is_read_only: bool = True + is_external_tool: bool = False + is_mcp: bool = False + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Allow every dummy tool call. + + Args: + tool_input (`dict[str, Any]`): + The tool input. + context (`PermissionContext`): + The permission context. + + Returns: + `PermissionDecision`: + The allow decision. + """ + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + decision_reason="Dummy tool always allows.", + message="Dummy tool always allows.", + ) + + async def __call__(self, **kwargs: Any) -> ToolChunk: + """Return a fixed tool result. + + Args: + **kwargs (`Any`): + Ignored tool arguments. + + Returns: + `ToolChunk`: + The fixed tool output. + """ + return ToolChunk(content=[TextBlock(text="tool result")]) + + +def _text_response(text: str) -> ChatResponse: + """Build a text-only chat response. + + Args: + text (`str`): + The response text. + + Returns: + `ChatResponse`: + A complete chat response with one text block. + """ + return ChatResponse(content=[TextBlock(text=text)], is_last=True) + + +def _tool_response() -> ChatResponse: + """Build a chat response that calls the dummy tool. + + Returns: + `ChatResponse`: + A complete chat response with one tool call block. + """ + return ChatResponse( + content=[ + ToolCallBlock( + id="call_dummy", + name="dummy", + input="{}", + ), + ], + is_last=True, + ) + + +def _structured_response(selected_files: list[str]) -> StructuredResponse: + """Build a structured memory-selection response. + + Args: + selected_files (`list[str]`): + The selected memory filenames. + + Returns: + `StructuredResponse`: + The structured response consumed by the middleware. + """ + return StructuredResponse(content={"selected_files": selected_files}) + + +def _block_to_dict(block: Any) -> dict: + """Convert a message block into a stable assertion dictionary. + + Args: + block (`Any`): + The message block to convert. + + Returns: + `dict`: + The stable block representation. + """ + if isinstance(block, TextBlock): + return { + "type": "text", + "text": block.text, + "id": AnyString(), + } + if isinstance(block, HintBlock): + return { + "type": "hint", + "hint": block.hint, + "id": AnyString(), + "source": block.source, + } + if isinstance(block, ToolCallBlock): + return { + "type": "tool_call", + "id": AnyString(), + "name": block.name, + "input": block.input, + "state": block.state, + "suggested_rules": block.suggested_rules, + } + return block.model_dump() + + +def _message_to_dict(msg: Msg) -> dict: + """Convert a message into a stable assertion dictionary. + + Args: + msg (`Msg`): + The message to convert. + + Returns: + `dict`: + The stable message representation. + """ + return { + "id": AnyString(), + "name": msg.name, + "role": msg.role, + "content": [_block_to_dict(block) for block in msg.content], + "metadata": msg.metadata, + } + + +def _hint_texts(agent: Agent) -> list[str]: + """Collect hint texts from an agent context. + + Args: + agent (`Agent`): + The agent whose context is inspected. + + Returns: + `list[str]`: + The hint texts in context order. + """ + return [ + block.hint + for msg in agent.state.context + for block in msg.content + if isinstance(block, HintBlock) + ] + + +def _write_memory_file( + memory_dir: str, + filename: str, + description: str, + memory_type: str, + body: str, +) -> None: + """Write one Markdown memory file with frontmatter. + + Args: + memory_dir (`str`): + The memory directory. + filename (`str`): + The memory filename relative to ``memory_dir``. + description (`str`): + The frontmatter description. + memory_type (`str`): + The frontmatter memory type. + body (`str`): + The Markdown body. + """ + path = os.path.join(memory_dir, filename) + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w", encoding="utf-8") as f: + f.write( + "---\n" + f"name: {filename}\n" + f"description: {description}\n" + f"type: {memory_type}\n" + "---\n\n" + f"{body}\n", + ) + + +class AgenticMemoryMiddlewareTest(IsolatedAsyncioTestCase): + """Agent-level tests for :class:`AgenticMemoryMiddleware`.""" + + async def asyncSetUp(self) -> None: + """Create a temporary workspace for each test.""" + self.temp_dir = tempfile.mkdtemp() + + async def asyncTearDown(self) -> None: + """Remove the temporary workspace after each test.""" + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def _make_agent( + self, + model: _RecordingMockModel, + middleware: AgenticMemoryMiddleware, + toolkit: Toolkit | None = None, + ) -> Agent: + """Build an Agent with the filesystem memory middleware attached. + + Args: + model (`_RecordingMockModel`): + The mock model used by the agent. + middleware (`AgenticMemoryMiddleware`): + The middleware under test. + toolkit (`Toolkit | None`, optional): + The toolkit for the agent. Defaults to an empty toolkit. + + Returns: + `Agent`: + The configured agent. + """ + return Agent( + name="assistant", + system_prompt="You are helpful.", + model=model, + toolkit=toolkit or Toolkit(), + middlewares=[middleware], + ) + + async def test_agent_reply_creates_layout_and_injects_memory_prompt( + self, + ) -> None: + """Agent reply should create layout and inject memory instructions.""" + model = _RecordingMockModel() + model.set_responses([_text_response("done")]) + middleware = AgenticMemoryMiddleware(workdir=self.temp_dir) + agent = self._make_agent(model, middleware) + + reply = await agent.reply(UserMsg("user", "hello")) + memory_dir = os.path.join(self.temp_dir, "Memory") + system_prompt = model.chat_messages[0][0].get_text_content() + + self.assertDictEqual( + { + "reply": _message_to_dict(reply), + "memory_dir_exists": os.path.isdir(memory_dir), + "memory_md_exists": os.path.isfile( + os.path.join(memory_dir, "MEMORY.md"), + ), + "system_prompt": { + "has_memory_dir": memory_dir in system_prompt, + "has_placeholder": "{memory_dir}" in system_prompt, + "has_memory_header": "## MEMORY.md" in system_prompt, + "has_empty_memory_text": ( + "Your MEMORY.md is currently empty" in system_prompt + ), + }, + }, + { + "reply": { + "id": AnyString(), + "name": "assistant", + "role": "assistant", + "content": [ + { + "type": "text", + "text": "done", + "id": AnyString(), + }, + ], + "metadata": {}, + }, + "memory_dir_exists": True, + "memory_md_exists": True, + "system_prompt": { + "has_memory_dir": True, + "has_placeholder": False, + "has_memory_header": True, + "has_empty_memory_text": True, + }, + }, + ) + + async def test_agent_reasoning_injects_selected_memory_hint( + self, + ) -> None: + """Agent reasoning should inject content selected by structured + output.""" + memory_dir = os.path.join(self.temp_dir, "Memory") + os.makedirs(memory_dir) + with open( + os.path.join(memory_dir, "MEMORY.md"), + "w", + encoding="utf-8", + ) as f: + f.write("- [User profile](user_profile.md) — User profile.\n") + _write_memory_file( + memory_dir, + "user_profile.md", + "User profile details", + "user", + "The user prefers concise Chinese answers.", + ) + + model = _RecordingMockModel() + model.set_structured_response( + _structured_response(["user_profile.md"]), + ) + model.set_responses([_tool_response(), _text_response("final answer")]) + middleware = AgenticMemoryMiddleware(workdir=self.temp_dir) + agent = self._make_agent( + model, + middleware, + toolkit=Toolkit(tools=[_DummyTool()]), + ) + + reply = await agent.reply(UserMsg("user", "what do you remember?")) + hint_texts = _hint_texts(agent) + + self.assertDictEqual( + { + "reply": _message_to_dict(reply), + "hints": [ + { + "has_selected_content": ( + "The user prefers concise Chinese answers." in hint + ), + "has_selected_path": "user_profile.md" in hint, + } + for hint in hint_texts + ], + "context": [ + _message_to_dict(msg) for msg in agent.state.context + ], + "structured_call_count": len(model.structured_messages), + }, + { + "reply": { + "id": AnyString(), + "name": "assistant", + "role": "assistant", + "content": [ + { + "type": "text", + "text": "final answer", + "id": AnyString(), + }, + ], + "metadata": {}, + }, + "hints": [ + { + "has_selected_content": True, + "has_selected_path": True, + }, + ], + "context": [ + { + "id": AnyString(), + "name": "user", + "role": "user", + "content": [ + { + "type": "text", + "text": "what do you remember?", + "id": AnyString(), + }, + ], + "metadata": {}, + }, + { + "id": AnyString(), + "name": "assistant", + "role": "assistant", + "content": [ + { + "type": "tool_call", + "id": AnyString(), + "name": "dummy", + "input": "{}", + "state": "finished", + "suggested_rules": [], + }, + AnyValue(), + { + "type": "hint", + "hint": AnyString(), + "id": AnyString(), + "source": None, + }, + { + "type": "text", + "text": "final answer", + "id": AnyString(), + }, + ], + "metadata": {}, + }, + ], + "structured_call_count": 1, + }, + ) + + async def test_agent_filters_hallucinated_memory_filenames(self) -> None: + """Agent retrieval should ignore filenames not present in memory.""" + memory_dir = os.path.join(self.temp_dir, "Memory") + os.makedirs(memory_dir) + with open( + os.path.join(memory_dir, "MEMORY.md"), + "w", + encoding="utf-8", + ) as f: + f.write("- [User profile](user_profile.md) — User profile.\n") + _write_memory_file( + memory_dir, + "user_profile.md", + "User profile details", + "user", + "Only this real memory should be injected.", + ) + + model = _RecordingMockModel() + model.set_structured_response( + _structured_response(["user_profile.md", "missing.md"]), + ) + model.set_responses([_tool_response(), _text_response("filtered")]) + middleware = AgenticMemoryMiddleware(workdir=self.temp_dir) + agent = self._make_agent( + model, + middleware, + toolkit=Toolkit(tools=[_DummyTool()]), + ) + + await agent.reply(UserMsg("user", "recall memory")) + hint_texts = _hint_texts(agent) + + self.assertListEqual( + [ + { + "has_real_memory": ( + "Only this real memory should be injected." in hint + ), + "has_missing_memory": "missing.md" in hint, + } + for hint in hint_texts + ], + [ + { + "has_real_memory": True, + "has_missing_memory": False, + }, + ], + ) + + async def test_agent_does_not_inject_hint_when_no_file_selected( + self, + ) -> None: + """Agent retrieval should inject no hint when selection is empty.""" + memory_dir = os.path.join(self.temp_dir, "Memory") + os.makedirs(memory_dir) + with open( + os.path.join(memory_dir, "MEMORY.md"), + "w", + encoding="utf-8", + ) as f: + f.write("- [User profile](user_profile.md) — User profile.\n") + _write_memory_file( + memory_dir, + "user_profile.md", + "User profile details", + "user", + "This memory is available but not selected.", + ) + + model = _RecordingMockModel() + model.set_structured_response(_structured_response([])) + model.set_responses([_tool_response(), _text_response("no hint")]) + middleware = AgenticMemoryMiddleware(workdir=self.temp_dir) + agent = self._make_agent( + model, + middleware, + toolkit=Toolkit(tools=[_DummyTool()]), + ) + + reply = await agent.reply(UserMsg("user", "ignore memories")) + + self.assertDictEqual( + { + "reply": _message_to_dict(reply), + "hints": _hint_texts(agent), + "structured_call_count": len(model.structured_messages), + }, + { + "reply": { + "id": AnyString(), + "name": "assistant", + "role": "assistant", + "content": [ + { + "type": "text", + "text": "no hint", + "id": AnyString(), + }, + ], + "metadata": {}, + }, + "hints": [], + "structured_call_count": 1, + }, + ) + + async def test_agent_does_not_retrieve_when_only_memory_index_exists( + self, + ) -> None: + """Agent retrieval should skip structured output without topic + files.""" + model = _RecordingMockModel() + model.set_structured_response(_structured_response(["missing.md"])) + model.set_responses([_tool_response(), _text_response("index only")]) + middleware = AgenticMemoryMiddleware(workdir=self.temp_dir) + agent = self._make_agent( + model, + middleware, + toolkit=Toolkit(tools=[_DummyTool()]), + ) + + reply = await agent.reply(UserMsg("user", "hello")) + + self.assertDictEqual( + { + "reply": _message_to_dict(reply), + "hints": _hint_texts(agent), + "structured_call_count": len(model.structured_messages), + "memory_md_exists": os.path.isfile( + os.path.join(self.temp_dir, "Memory", "MEMORY.md"), + ), + }, + { + "reply": { + "id": AnyString(), + "name": "assistant", + "role": "assistant", + "content": [ + { + "type": "text", + "text": "index only", + "id": AnyString(), + }, + ], + "metadata": {}, + }, + "hints": [], + "structured_call_count": 0, + "memory_md_exists": True, + }, + ) + + async def test_agent_system_prompt_contains_truncation_reminder( + self, + ) -> None: + """Agent system prompt should contain reminder for truncated index.""" + memory_dir = os.path.join(self.temp_dir, "Memory") + os.makedirs(memory_dir) + with open( + os.path.join(memory_dir, "MEMORY.md"), + "w", + encoding="utf-8", + ) as f: + f.write("0123456789" * 80) + + model = _RecordingMockModel() + model.set_responses([_text_response("truncated")]) + middleware = AgenticMemoryMiddleware( + workdir=self.temp_dir, + parameters=AgenticMemoryMiddleware.Parameters( + memory_max_tokens=10, + retrieval_async=False, + ), + ) + agent = self._make_agent(model, middleware) + + await agent.reply(UserMsg("user", "hello")) + system_prompt = model.chat_messages[0][0].get_text_content() + + self.assertDictEqual( + { + "has_truncated_marker": "<<>>" in system_prompt, + "has_offset_reminder": "Use the `Read` tool with offset" + in system_prompt, + "has_memory_path": os.path.join(memory_dir, "MEMORY.md") + in system_prompt, + }, + { + "has_truncated_marker": True, + "has_offset_reminder": True, + "has_memory_path": True, + }, + ) + + async def test_agent_skips_retrieval_when_async_retrieval_disabled( + self, + ) -> None: + """Agent should not run retrieval when ``retrieval_async`` is false.""" + memory_dir = os.path.join(self.temp_dir, "Memory") + os.makedirs(memory_dir) + with open( + os.path.join(memory_dir, "MEMORY.md"), + "w", + encoding="utf-8", + ) as f: + f.write("- [User profile](user_profile.md) — User profile.\n") + _write_memory_file( + memory_dir, + "user_profile.md", + "User profile details", + "user", + "This memory should not be retrieved.", + ) + + model = _RecordingMockModel() + model.set_structured_response( + _structured_response(["user_profile.md"]), + ) + model.set_responses([_tool_response(), _text_response("disabled")]) + middleware = AgenticMemoryMiddleware( + workdir=self.temp_dir, + parameters=AgenticMemoryMiddleware.Parameters( + retrieval_async=False, + ), + ) + agent = self._make_agent( + model, + middleware, + toolkit=Toolkit(tools=[_DummyTool()]), + ) + + reply = await agent.reply(UserMsg("user", "remember?")) + + self.assertDictEqual( + { + "reply": _message_to_dict(reply), + "hints": _hint_texts(agent), + "structured_call_count": len(model.structured_messages), + }, + { + "reply": { + "id": AnyString(), + "name": "assistant", + "role": "assistant", + "content": [ + { + "type": "text", + "text": "disabled", + "id": AnyString(), + }, + ], + "metadata": {}, + }, + "hints": [], + "structured_call_count": 0, + }, + ) diff --git a/tests/middleware_rag_test.py b/tests/middleware_rag_test.py new file mode 100644 index 0000000000000000000000000000000000000000..3d2c844916e30f12129432b513ac9411e41f575a --- /dev/null +++ b/tests/middleware_rag_test.py @@ -0,0 +1,550 @@ +# -*- coding: utf-8 -*- +"""Unit tests for the :class:`RAGMiddleware` class.""" +from contextlib import AsyncExitStack +from types import SimpleNamespace +from typing import Any, AsyncGenerator +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import AnyString + +from agentscope.embedding import EmbeddingResponse +from agentscope.event import EventType, HintBlockEvent +from agentscope.message import ( + Base64Source, + DataBlock, + Msg, + TextBlock, + UserMsg, +) +from agentscope.middleware import RAGMiddleware +from agentscope.rag import Chunk, KnowledgeBase, QdrantStore, VectorRecord + + +_HINT_SOURCE = '{"label": "KnowledgeBase", "sublabel": ""}' + +_EXPECTED_HINT = ( + "The following content is retrieved from the " + "knowledge base(s) and may be helpful for the current " + "request:\n" + "[1] (source: doc-1.txt)\n" + "Paris is in France." +) + + +class _StubEmbeddingModel: + """A stub embedding model returning a fixed vector per input.""" + + supports_multimodal = False + dimensions = 3 + + def __init__(self, vector: list[float]) -> None: + """Initialize the stub. + + Args: + vector (`list[float]`): + The vector returned for every input. + """ + self.vector = vector + self.calls: list[list] = [] + + async def __call__(self, inputs: list) -> EmbeddingResponse: + """Return the fixed vector for each input. + + Args: + inputs (`list`): + The input queries. + + Returns: + `EmbeddingResponse`: + The response with one fixed vector per input. + """ + self.calls.append(inputs) + return EmbeddingResponse(embeddings=[self.vector] * len(inputs)) + + +def _make_record( + text: str, + vector: list[float], + document_id: str, +) -> VectorRecord: + """Build a VectorRecord for testing. + + Args: + text (`str`): + The chunk text content. + vector (`list[float]`): + The embedding vector. + document_id (`str`): + The ID of the source document the record belongs to. + + Returns: + `VectorRecord`: + The constructed record. + """ + return VectorRecord( + vector=vector, + document_id=document_id, + chunk=Chunk( + content=TextBlock(text=text), + source=f"{document_id}.txt", + chunk_index=0, + total_chunks=1, + ), + ) + + +def _make_agent( + context: list[Msg] | None = None, + cur_iter: int = 0, +) -> Any: + """Build a minimal stand-in for an Agent. + + Args: + context (`list[Msg] | None`, optional): + The initial agent context. + cur_iter (`int`, defaults to ``0``): + Value for ``state.cur_iter``; the middleware only searches + on the first reasoning step (``0``). + + Returns: + `Any`: + An object with ``name`` and ``state.context`` / + ``state.reply_id`` / ``state.session_id`` / + ``state.cur_iter`` / ``state.append_context``. + """ + + msgs: list[Msg] = context if context is not None else [] + + def _append_context(name: str, blocks: list) -> None: + # Always append a new assistant carrier message keyed on the + # static reply_id used in these tests. Mirrors the real + # ``AgentState.append_context`` for the purposes of the + # middleware's reverse-scan removal logic. + carrier = Msg(name=name, role="assistant", content=blocks) + carrier.id = "reply-1" + msgs.append(carrier) + + state = SimpleNamespace( + context=msgs, + reply_id="reply-1", + session_id="session-1", + cur_iter=cur_iter, + append_context=_append_context, + ) + return SimpleNamespace(name="assistant", state=state) + + +async def _drain(generator: AsyncGenerator) -> list: + """Exhaust an async generator into a list. + + Args: + generator (`AsyncGenerator`): + The generator to drain. + + Returns: + `list`: + All yielded items. + """ + return [item async for item in generator] + + +class RAGMiddlewareTest(IsolatedAsyncioTestCase): + """The test cases for the :class:`RAGMiddleware` class.""" + + async def asyncSetUp(self) -> None: + """Create an in-memory store seeded with one collection + + one :class:`KnowledgeBase` handle wired to it.""" + self._exit_stack = AsyncExitStack() + self.store = await self._exit_stack.enter_async_context( + QdrantStore(location=":memory:"), + ) + await self.store.create_collection("kb-1", dimensions=3) + await self.store.insert( + "kb-1", + [ + _make_record("Paris is in France.", [1.0, 0.0, 0.0], "doc-1"), + _make_record("Cats are mammals.", [0.0, 1.0, 0.0], "doc-2"), + ], + ) + self.embedding_model = _StubEmbeddingModel([1.0, 0.0, 0.0]) + # Build the KnowledgeBase handle once; tests share it. The + # collection already exists, so ``ensure_collection`` will + # short-circuit on first use. + self.knowledge = KnowledgeBase( + name="paris-kb", + description="Trivia about Paris and cats.", + embedding_model=self.embedding_model, + vector_store=self.store, + collection="kb-1", + ) + + async def asyncTearDown(self) -> None: + """Close the store after each test.""" + await self._exit_stack.aclose() + + def _middleware( + self, + knowledges: list[KnowledgeBase] | None = None, + **kwargs: Any, + ) -> RAGMiddleware: + """Build a middleware bound to ``self.knowledge`` with a + :class:`SearchConfig` assembled from ``kwargs``. + + Args: + knowledges (`list[KnowledgeBase] | None`, optional): + Override the bound knowledge bases. Defaults to + ``[self.knowledge]``. + **kwargs (`Any`): + Forwarded to :class:`SearchConfig` (e.g. ``mode``, + ``top_k``, ``score_threshold``, ``emit_hint_event``, + ``persist_hint``). + + Returns: + `RAGMiddleware`: + The middleware under test. + """ + return RAGMiddleware( + knowledge_bases=knowledges + if knowledges is not None + else [ + self.knowledge, + ], + parameters=RAGMiddleware.Parameters(**kwargs), + ) + + async def _run_with_inputs( + self, + middleware: RAGMiddleware, + agent: Any, + inputs: Msg | list[Msg] | None, + context_during_reasoning: list[dict] | None = None, + ) -> list: + """Drive ``on_reply`` → ``on_reasoning`` end-to-end. + + Mirrors the real agent loop: ``on_reply`` captures the inputs + in the middleware's scratchpad, then ``on_reasoning`` runs + (with ``state.cur_iter == 0``) and may inject a hint. The + reasoning step yields a sentinel ``"reasoning-evt"`` so callers + can assert event order; if ``context_during_reasoning`` is + provided it is filled with a dump of ``agent.state.context`` as + seen by the innermost reasoning callback. + + Args: + middleware (`RAGMiddleware`): + The middleware under test. + agent (`Any`): + The fake agent. + inputs (`Msg | list[Msg] | None`): + The reply inputs to pass through ``on_reply``. + context_during_reasoning (`list[dict] | None`, optional): + When provided, receives a dump of the agent context as + seen by the wrapped (innermost) reasoning call. + + Returns: + `list`: + All events yielded by the on_reply → on_reasoning chain. + """ + + async def reasoning_next(**_kwargs: Any) -> AsyncGenerator: + if context_during_reasoning is not None: + context_during_reasoning.extend( + msg.model_dump() for msg in agent.state.context + ) + yield "reasoning-evt" + + async def reply_next(**_kwargs: Any) -> AsyncGenerator: + # The reply branch drives the reasoning branch — same as + # the real composition. + async for evt in middleware.on_reasoning( + agent=agent, + input_kwargs={"tool_choice": None}, + next_handler=reasoning_next, + ): + yield evt + + return await _drain( + middleware.on_reply( + agent=agent, + input_kwargs={"inputs": inputs}, + next_handler=reply_next, + ), + ) + + # ------------------------------------------------------------------ + # Static mode (auto-injection) + # ------------------------------------------------------------------ + + async def test_static_one_shot_injection(self) -> None: + """The hint participates in one reasoning step and is removed + afterwards (``persist_hint=False``, default).""" + middleware = self._middleware( + mode="static", + top_k=1, + emit_hint_event=False, + ) + agent = _make_agent() + seen_context: list[dict] = [] + + events = await self._run_with_inputs( + middleware, + agent, + UserMsg(name="user", content="Where is Paris?"), + context_during_reasoning=seen_context, + ) + + # No HintBlockEvent (emit_hint_event=False); only downstream + # events. + self.assertEqual(events, ["reasoning-evt"]) + + # The reasoning callback observed exactly one carrier message + # holding the injected hint block. + self.assertEqual(len(seen_context), 1) + carrier = seen_context[0] + self.assertEqual(carrier["role"], "assistant") + self.assertEqual(carrier["id"], "reply-1") + self.assertEqual(len(carrier["content"]), 1) + block = carrier["content"][0] + self.assertEqual(block["type"], "hint") + self.assertEqual(block["source"], _HINT_SOURCE) + self.assertEqual(block["hint"], _EXPECTED_HINT) + + # One-shot: after on_reasoning unwinds, the carrier is emptied. + post = [msg.model_dump() for msg in agent.state.context] + self.assertEqual(len(post), 1) + self.assertEqual(post[0]["content"], []) + + async def test_static_persistent_injection(self) -> None: + """``persist_hint=True`` keeps the hint in the context.""" + middleware = self._middleware( + mode="static", + top_k=1, + persist_hint=True, + emit_hint_event=False, + ) + agent = _make_agent() + seen_context: list[dict] = [] + + await self._run_with_inputs( + middleware, + agent, + UserMsg(name="user", content="Where is Paris?"), + context_during_reasoning=seen_context, + ) + + self.assertEqual( + [msg.model_dump() for msg in agent.state.context], + seen_context, + ) + + async def test_static_event_emission(self) -> None: + """``emit_hint_event=True`` yields one :class:`HintBlockEvent`.""" + middleware = self._middleware( + mode="static", + top_k=1, + emit_hint_event=True, + ) + agent = _make_agent() + + events = await self._run_with_inputs( + middleware, + agent, + UserMsg(name="user", content="Where is Paris?"), + ) + + self.assertEqual(len(events), 2) + self.assertIsInstance(events[0], HintBlockEvent) + self.assertEqual( + events[0].model_dump(), + { + "type": EventType.HINT_BLOCK, + "reply_id": "reply-1", + "block_id": AnyString(), + "source": _HINT_SOURCE, + "hint": _EXPECTED_HINT, + "id": AnyString(), + "created_at": AnyString(), + "metadata": {}, + }, + ) + self.assertEqual(events[1], "reasoning-evt") + + async def test_static_skips_event_inputs(self) -> None: + """Non-message inputs (resumption events / ``None``) skip the + search entirely.""" + middleware = self._middleware(mode="static") + agent = _make_agent() + + events = await self._run_with_inputs(middleware, agent, None) + + self.assertEqual(events, ["reasoning-evt"]) + self.assertEqual(self.embedding_model.calls, []) + self.assertEqual(agent.state.context, []) + + async def test_multimodal_query_extraction(self) -> None: + """DataBlocks reach the embedding model when it declares + ``supports_multimodal``.""" + self.embedding_model.supports_multimodal = True + middleware = self._middleware( + mode="static", + top_k=1, + emit_hint_event=False, + ) + agent = _make_agent() + data_block = DataBlock( + source=Base64Source(data="aGk=", media_type="image/png"), + ) + + await self._run_with_inputs( + middleware, + agent, + UserMsg( + name="user", + content=[TextBlock(text="What is this?"), data_block], + ), + ) + + # The query path prepends ``{name}: `` to the first text + # block; the data block is passed through verbatim. + self.assertEqual(len(self.embedding_model.calls), 1) + query = self.embedding_model.calls[0] + self.assertEqual(len(query), 2) + self.assertEqual(query[0].text, "user: What is this?") + self.assertEqual(query[1], data_block) + + async def test_multimodal_blocks_dropped_for_text_only_model( + self, + ) -> None: + """A text-only embedding model silently drops DataBlock queries + (no exception, no crash).""" + middleware = self._middleware( + mode="static", + top_k=1, + emit_hint_event=False, + ) + agent = _make_agent() + data_block = DataBlock( + source=Base64Source(data="aGk=", media_type="image/png"), + ) + + await self._run_with_inputs( + middleware, + agent, + UserMsg( + name="user", + content=[TextBlock(text="What is this?"), data_block], + ), + ) + + # ``KnowledgeBase.search`` strips the DataBlock when the bound + # embedding model isn't multimodal — the model only saw text. + self.assertEqual(len(self.embedding_model.calls), 1) + for item in self.embedding_model.calls[0]: + self.assertNotIsInstance(item, DataBlock) + + # ------------------------------------------------------------------ + # Agentic mode (tool exposure) + # ------------------------------------------------------------------ + + async def test_agentic_list_tools(self) -> None: + """Agentic mode exposes the search tool; static mode none.""" + agentic_tools = await self._middleware(mode="agentic").list_tools() + static_tools = await self._middleware(mode="static").list_tools() + + self.assertEqual( + [tool.name for tool in agentic_tools], + ["search_knowledge"], + ) + self.assertEqual(static_tools, []) + + async def test_agentic_no_auto_injection(self) -> None: + """Agentic mode never searches or injects automatically.""" + middleware = self._middleware(mode="agentic") + agent = _make_agent() + + events = await self._run_with_inputs( + middleware, + agent, + UserMsg(name="user", content="Where is Paris?"), + ) + + self.assertEqual(events, ["reasoning-evt"]) + self.assertEqual(self.embedding_model.calls, []) + self.assertEqual(agent.state.context, []) + + async def test_search_knowledge_tool_call(self) -> None: + """The tool returns a formatted ``ToolChunk`` for a query. + + ``_SearchKnowledgeTool.call`` is a regular async function (not + an async generator), so ``ToolBase.__call__`` awaits it and + returns the single ``ToolChunk`` directly. + """ + middleware = self._middleware(mode="agentic", top_k=1) + tool = (await middleware.list_tools())[0] + + chunk = await tool(query="Where is Paris?") + + self.assertEqual( + chunk.model_dump(), + { + "content": [ + { + "type": "text", + "text": ( + "[1] (source: doc-1.txt)\nParis is in France." + ), + "id": AnyString(), + }, + ], + "state": "success", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + async def test_search_knowledge_tool_input_schema_enum(self) -> None: + """The tool's ``input_schema`` narrows ``knowledge_bases.items`` + to the equipped KB names.""" + middleware = self._middleware(mode="agentic") + tool = (await middleware.list_tools())[0] + + schema = tool.input_schema + kb_schema = schema["properties"]["knowledge_bases"] + # Pydantic emits Optional[list[str]] as anyOf; pick the array + # branch. + array_variant = next( + v for v in kb_schema["anyOf"] if v.get("type") == "array" + ) + self.assertEqual(array_variant["items"]["enum"], ["paris-kb"]) + + async def test_search_knowledge_tool_filters_by_name(self) -> None: + """Passing ``knowledge_bases=[]`` returns the + ``"No relevant content found."`` notice without touching the + embedding model.""" + middleware = self._middleware(mode="agentic", top_k=1) + tool = (await middleware.list_tools())[0] + + chunk = await tool( + query="Where is Paris?", + knowledge_bases=["does-not-exist"], + ) + + self.assertEqual( + [b["text"] for b in chunk.model_dump()["content"]], + ["No relevant content found."], + ) + self.assertEqual(self.embedding_model.calls, []) + + # ------------------------------------------------------------------ + # Config validation + # ------------------------------------------------------------------ + + async def test_hint_template_must_have_context_placeholder(self) -> None: + """:class:`SearchConfig` rejects a template without exactly one + ``{context}``.""" + with self.assertRaises(ValueError): + RAGMiddleware.Parameters(hint_template="no placeholder here") + with self.assertRaises(ValueError): + RAGMiddleware.Parameters(hint_template="{context} twice {context}") + # Exactly one placeholder is fine. + RAGMiddleware.Parameters(hint_template="wrapped: {context}.") diff --git a/tests/middleware_test.py b/tests/middleware_test.py new file mode 100644 index 0000000000000000000000000000000000000000..7c0c603e74f173ef3690052e54cf3b465ad44adb --- /dev/null +++ b/tests/middleware_test.py @@ -0,0 +1,1060 @@ +# -*- coding: utf-8 -*- +# pylint: disable=abstract-method +"""Unit tests for middleware system.""" +from unittest.async_case import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, patch +from typing import Any, AsyncGenerator, Callable, Union + +from utils import MockModel +from pydantic import BaseModel +from agentscope.event import AgentEvent +from agentscope.agent import Agent, ContextConfig +from agentscope.middleware import MiddlewareBase +from agentscope.model import ChatResponse +from agentscope.message import ( + TextBlock, + HintBlock, + UserMsg, + SystemMsg, + Msg, + ToolCallBlock, +) +from agentscope.tool import Toolkit, ToolBase, ToolChunk +from agentscope.permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, +) + + +class TestMiddleware(IsolatedAsyncioTestCase): + """Test cases for middleware system.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.mock_model = MockModel() + self.toolkit = Toolkit() + self.execution_log = [] + + async def test_on_reply_middleware_pre_post_yield(self) -> None: + """Test on_reply middleware pre, post and yield positions.""" + + class ReplyMiddleware(MiddlewareBase): + """Middleware for testing on_reply hook.""" + + def __init__(self, log: list, name: str) -> None: + """Initialize the reply middleware. + + Args: + log: The execution log list. + name: The middleware name. + """ + self.log = log + self.name = name + + async def on_reply( + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable[[], AsyncGenerator], + ) -> AsyncGenerator: + """The on_reply middleware logic.""" + self.log.append(f"{self.name}_pre") + async for item in next_handler(): + if isinstance(item, AgentEvent): + self.log.append(f"{self.name}_{item.type}") + elif isinstance(item, Msg): + self.log.append(f"{self.name}_msg") + yield item + self.log.append(f"{self.name}_post") + + middleware1 = ReplyMiddleware(self.execution_log, "mw1") + middleware2 = ReplyMiddleware(self.execution_log, "mw2") + + self.mock_model.set_responses( + [ + ChatResponse( + content=[TextBlock(text="test response")], + is_last=True, + ), + ], + ) + + agent = Agent( + name="test_agent", + system_prompt="test prompt", + model=self.mock_model, + toolkit=self.toolkit, + middlewares=[middleware1, middleware2], + ) + + await agent.reply(UserMsg("user", "test message")) + + # Verify execution order + expected = [ + "mw1_pre", + "mw2_pre", + "mw2_REPLY_START", + "mw1_REPLY_START", + "mw2_MODEL_CALL_START", + "mw1_MODEL_CALL_START", + "mw2_TEXT_BLOCK_START", + "mw1_TEXT_BLOCK_START", + "mw2_TEXT_BLOCK_DELTA", + "mw1_TEXT_BLOCK_DELTA", + "mw2_TEXT_BLOCK_END", + "mw1_TEXT_BLOCK_END", + "mw2_MODEL_CALL_END", + "mw1_MODEL_CALL_END", + "mw2_REPLY_END", + "mw1_REPLY_END", + "mw2_msg", + "mw1_msg", + "mw2_post", + "mw1_post", + ] + self.assertListEqual(self.execution_log, expected) + + async def test_on_reasoning_middleware_pre_yield(self) -> None: + """Test on_reasoning middleware pre and yield positions.""" + + class ReasoningMiddleware(MiddlewareBase): + """Middleware for testing on_reasoning hook.""" + + def __init__(self, log: list, name: str) -> None: + """Initialize the reasoning middleware. + + Args: + log: The execution log list. + name: The middleware name. + """ + self.log = log + self.name = name + + async def on_reasoning( + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable[[], AsyncGenerator], + ) -> AsyncGenerator: + """The on_reasoning middleware logic.""" + self.log.append(f"{self.name}_pre") + async for item in next_handler(): + if isinstance(item, AgentEvent): + self.log.append(f"{self.name}_{item.type}") + elif isinstance(item, Msg): + self.log.append(f"{self.name}_msg") + yield item + self.log.append(f"{self.name}_post") + + middleware1 = ReasoningMiddleware(self.execution_log, "mw1") + middleware2 = ReasoningMiddleware(self.execution_log, "mw2") + + self.mock_model.set_responses( + [ + ChatResponse( + content=[TextBlock(text="test response")], + is_last=True, + ), + ], + ) + + agent = Agent( + name="test_agent", + system_prompt="test prompt", + model=self.mock_model, + toolkit=self.toolkit, + middlewares=[middleware1, middleware2], + ) + + await agent.reply(UserMsg("user", "test message")) + + # Verify execution order + expected = [ + "mw1_pre", + "mw2_pre", + "mw2_MODEL_CALL_START", + "mw1_MODEL_CALL_START", + "mw2_TEXT_BLOCK_START", + "mw1_TEXT_BLOCK_START", + "mw2_TEXT_BLOCK_DELTA", + "mw1_TEXT_BLOCK_DELTA", + "mw2_TEXT_BLOCK_END", + "mw1_TEXT_BLOCK_END", + "mw2_MODEL_CALL_END", + "mw1_MODEL_CALL_END", + "mw2_msg", + "mw1_msg", + ] + self.assertListEqual(self.execution_log, expected) + + async def test_on_model_call_middleware_non_streaming(self) -> None: + """Test on_model_call middleware for non-streaming model.""" + + class ModelCallMiddleware(MiddlewareBase): + """Middleware for testing on_model_call hook.""" + + def __init__(self, log: list, name: str) -> None: + """Initialize the model call middleware. + + Args: + log: The execution log list. + name: The middleware name. + """ + self.log = log + self.name = name + + async def on_model_call( + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable, + ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]: + """The on_model_call middleware logic.""" + self.log.append(f"{self.name}_pre") + result = await next_handler() + self.log.append(f"{self.name}_post") + return result + + middleware1 = ModelCallMiddleware(self.execution_log, "mw1") + middleware2 = ModelCallMiddleware(self.execution_log, "mw2") + + self.mock_model.set_responses( + [ + ChatResponse( + content=[TextBlock(text="test response")], + is_last=True, + ), + ], + ) + + agent = Agent( + name="test_agent", + system_prompt="test prompt", + model=self.mock_model, + toolkit=self.toolkit, + middlewares=[middleware1, middleware2], + ) + + await agent.reply(UserMsg("user", "test message")) + + # Verify execution order: mw1_pre -> mw2_pre -> mw2_post -> mw1_post + expected = ["mw1_pre", "mw2_pre", "mw2_post", "mw1_post"] + self.assertListEqual(self.execution_log, expected) + + async def test_on_model_call_middleware_streaming(self) -> None: + """Test on_model_call middleware for streaming model.""" + + class ModelCallMiddleware(MiddlewareBase): + """Middleware for testing on_model_call hook with streaming.""" + + def __init__(self, log: list, name: str) -> None: + """Initialize the model call middleware. + + Args: + log: The execution log list. + name: The middleware name. + """ + self.log = log + self.name = name + + async def on_model_call( + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable, + ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]: + """The on_model_call middleware logic for streaming.""" + self.log.append(f"{self.name}_pre") + result = await next_handler() + + async def wrapped_generator() -> AsyncGenerator[ + ChatResponse, + None, + ]: + """Wrap the generator to log yields.""" + async for chunk in result: + self.log.append(f"{self.name}_chunk") + yield chunk + self.log.append(f"{self.name}_post") + + return wrapped_generator() + + middleware1 = ModelCallMiddleware(self.execution_log, "mw1") + middleware2 = ModelCallMiddleware(self.execution_log, "mw2") + + self.mock_model.set_responses( + [ + [ + ChatResponse( + content=[TextBlock(text="chunk1")], + is_last=False, + ), + ChatResponse( + content=[TextBlock(text="chunk2")], + is_last=False, + ), + ChatResponse( + content=[TextBlock(text="chunk3")], + is_last=True, + ), + ], + ], + ) + + agent = Agent( + name="test_agent", + system_prompt="test prompt", + model=self.mock_model, + toolkit=self.toolkit, + middlewares=[middleware1, middleware2], + ) + + await agent.reply(UserMsg("user", "test message")) + + # Verify execution order + expected = [ + "mw1_pre", + "mw2_pre", + "mw2_chunk", + "mw1_chunk", + "mw2_chunk", + "mw1_chunk", + "mw2_chunk", + "mw1_chunk", + "mw2_post", + "mw1_post", + ] + self.assertListEqual(self.execution_log, expected) + + async def test_on_system_prompt_middleware(self) -> None: + """Test on_system_prompt middleware (transformer pattern).""" + + class SystemPromptMiddleware(MiddlewareBase): + """Middleware for testing on_system_prompt hook.""" + + def __init__(self, log: list, name: str, suffix: str) -> None: + """Initialize the system prompt middleware. + + Args: + log: The execution log list. + name: The middleware name. + suffix: The suffix to append to the prompt. + """ + self.log = log + self.name = name + self.suffix = suffix + + async def on_system_prompt( + self, + agent: Agent, + current_prompt: str, + ) -> str: + """The on_system_prompt middleware logic.""" + self.log.append(f"{self.name}_executed") + return f"{current_prompt} {self.suffix}" + + middleware1 = SystemPromptMiddleware( + self.execution_log, + "mw1", + "[MW1]", + ) + middleware2 = SystemPromptMiddleware( + self.execution_log, + "mw2", + "[MW2]", + ) + + self.mock_model.set_responses( + [ + ChatResponse( + content=[TextBlock(text="test response")], + is_last=True, + ), + ], + ) + + agent = Agent( + name="test_agent", + system_prompt="test prompt", + model=self.mock_model, + toolkit=self.toolkit, + middlewares=[middleware1, middleware2], + ) + + await agent.reply(UserMsg("user", "test message")) + + # Verify execution order: mw1 -> mw2 (sequential transformer pattern) + # Note: system_prompt is called twice (once for initial setup, once + # during reasoning) + expected = [ + "mw1_executed", + "mw2_executed", # First call + "mw1_executed", + "mw2_executed", # Second call during reasoning + ] + self.assertListEqual(self.execution_log, expected) + + async def test_multiple_middleware_types(self) -> None: + """Test multiple middleware types working together.""" + + class MultiMiddleware(MiddlewareBase): + """Middleware implementing multiple hooks.""" + + def __init__(self, log: list, name: str) -> None: + """Initialize the multi middleware. + + Args: + log: The execution log list. + name: The middleware name. + """ + self.log = log + self.name = name + + async def on_reply( + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable[[], AsyncGenerator], + ) -> AsyncGenerator: + """The on_reply middleware logic.""" + self.log.append(f"{self.name}_reply_pre") + async for item in next_handler(): + yield item + self.log.append(f"{self.name}_reply_post") + + async def on_reasoning( + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable[[], AsyncGenerator], + ) -> AsyncGenerator: + """The on_reasoning middleware logic.""" + self.log.append(f"{self.name}_reasoning_pre") + async for item in next_handler(): + yield item + self.log.append(f"{self.name}_reasoning_post") + + async def on_model_call( + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable, + ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]: + """The on_model_call middleware logic.""" + self.log.append(f"{self.name}_model_call_pre") + result = await next_handler() + self.log.append(f"{self.name}_model_call_post") + return result + + async def on_system_prompt( + self, + agent: Agent, + current_prompt: str, + ) -> str: + """The on_system_prompt middleware logic.""" + self.log.append(f"{self.name}_system_prompt") + return current_prompt + + middleware = MultiMiddleware(self.execution_log, "multi") + + self.mock_model.set_responses( + [ + ChatResponse( + content=[TextBlock(text="test response")], + is_last=True, + ), + ], + ) + + agent = Agent( + name="test_agent", + system_prompt="test prompt", + model=self.mock_model, + toolkit=self.toolkit, + middlewares=[middleware], + ) + + await agent.reply(UserMsg("user", "test message")) + + # Verify all middleware hooks were called + expected = [ + "multi_reply_pre", + "multi_system_prompt", + "multi_reasoning_pre", + "multi_system_prompt", + "multi_model_call_pre", + "multi_model_call_post", + "multi_reply_post", + ] + self.assertListEqual(self.execution_log, expected) + + async def test_on_reply_middleware_modify_input(self) -> None: + """Test that on_reply middleware can modify msgs input.""" + + class ModifyMsgsMiddleware(MiddlewareBase): + """Middleware that modifies the msgs input.""" + + async def on_reply( + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Modify inputs before passing to next handler.""" + # Modify the message content + inputs = input_kwargs["inputs"] + if isinstance(inputs, Msg): + modified_msg = UserMsg( + name=inputs.name, + content="MODIFIED: " + inputs.get_text_content(), + ) + async for item in next_handler(inputs=modified_msg): + yield item + else: + async for item in next_handler(**input_kwargs): + yield item + + middleware = ModifyMsgsMiddleware() + + # Track what message the model receives + received_messages = [] + + class TrackingModel(MockModel): + """Model that tracks received messages.""" + + async def _call_api( + self, + *args: Any, + **kwargs: Any, + ) -> ChatResponse: + """Track the messages and return mock response.""" + messages = kwargs.get("messages", []) + received_messages.extend(messages) + return await super()._call_api(*args, **kwargs) + + tracking_model = TrackingModel() + tracking_model.set_responses( + [ + ChatResponse( + content=[TextBlock(text="response")], + is_last=True, + ), + ], + ) + + agent_instance = Agent( + name="test_agent", + system_prompt="test prompt", + model=tracking_model, + toolkit=self.toolkit, + middlewares=[middleware], + ) + + await agent_instance.reply(UserMsg("user", "original message")) + + # Verify the model received the modified message + user_messages = [m for m in received_messages if m.role == "user"] + self.assertTrue(len(user_messages) > 0) + self.assertIn( + "MODIFIED: original message", + user_messages[-1].get_text_content(), + ) + + async def test_on_reasoning_middleware_modify_input(self) -> None: + """Test that on_reasoning middleware can modify tool_choice input.""" + + class ModifyToolChoiceMiddleware(MiddlewareBase): + """Middleware that modifies the tool_choice input.""" + + async def on_reasoning( + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Force tool_choice to 'none' to prevent tool calls.""" + # Override tool_choice to 'none' + async for item in next_handler(tool_choice="none"): + yield item + + middleware = ModifyToolChoiceMiddleware() + + # Track what tool_choice the model receives + received_tool_choices = [] + + class TrackingModel(MockModel): + """Model that tracks received tool_choice.""" + + async def _call_api( + self, + *args: Any, + **kwargs: Any, + ) -> ChatResponse: + """Track the tool_choice and return mock response.""" + tool_choice = kwargs.get("tool_choice") + received_tool_choices.append(tool_choice) + return await super()._call_api(*args, **kwargs) + + tracking_model = TrackingModel() + tracking_model.set_responses( + [ + ChatResponse( + content=[TextBlock(text="response without tools")], + is_last=True, + ), + ], + ) + + agent_instance = Agent( + name="test_agent", + system_prompt="test prompt", + model=tracking_model, + toolkit=self.toolkit, + middlewares=[middleware], + ) + + await agent_instance.reply(UserMsg("user", "test message")) + + # Verify the model received tool_choice='none' + self.assertIn("none", received_tool_choices) + + async def test_on_acting_middleware_intercepts_tool_execution( + self, + ) -> None: + """Test that on_acting middleware intercepts raw tool execution. + + After the refactor, ``on_acting`` wraps only ``_acting_impl`` + (i.e. ``toolkit.call_tool``). Permission checking and context + writes are handled by ``_execute_tool_call`` *outside* the hook. + This test verifies that the middleware can observe and modify the + ``tool_call`` passed to the actual tool function. + """ + + # ------------------------------------------------------------------ # + # A minimal tool that records the raw input it receives. # + # ------------------------------------------------------------------ # + received_inputs: list[str] = [] + + class _EchoParams(BaseModel): + value: str + + class EchoTool(ToolBase): + """Tool that echoes its input and records it.""" + + name: str = "echo" + description: str = "Echo the value." + input_schema: dict = _EchoParams.model_json_schema() + is_concurrency_safe: bool = True + is_read_only: bool = True + is_state_injected: bool = False + is_external_tool: bool = False + is_mcp: bool = False + mcp_name: str | None = None + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Always allow.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message="allowed", + ) + + async def __call__( + self, + value: str, + ) -> ToolChunk: + """Record the value and return it.""" + received_inputs.append(value) + return ToolChunk( + content=[TextBlock(text=f"echo:{value}")], + ) + + toolkit_with_tool = Toolkit(tools=[EchoTool()]) + + # ------------------------------------------------------------------ # + # Middleware that renames the tool_call.input before forwarding. # + # ------------------------------------------------------------------ # + intercepted_tool_calls: list[str] = [] + + class ObserveActingMiddleware(MiddlewareBase): + """Middleware that records the tool_call seen at acting level.""" + + async def on_acting( + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Record the tool_call and forward to next handler.""" + tool_call = input_kwargs["tool_call"] + intercepted_tool_calls.append(tool_call.input) + # Modify the input before execution + import json + + modified = ToolCallBlock( + id=tool_call.id, + name=tool_call.name, + input=json.dumps({"value": "MODIFIED"}), + state=tool_call.state, + ) + async for chunk in next_handler(tool_call=modified): + yield chunk + + middleware = ObserveActingMiddleware() + self.mock_model.set_responses( + [ + ChatResponse( + content=[TextBlock(text="done")], + is_last=True, + ), + ], + ) + + agent_instance = Agent( + name="test_agent", + system_prompt="test prompt", + model=self.mock_model, + toolkit=toolkit_with_tool, + middlewares=[middleware], + ) + + # Call _execute_tool_call with a valid tool call. + tool_call = ToolCallBlock( + id="call_1", + name="echo", + input='{"value": "ORIGINAL"}', + ) + events = [] + # pylint: disable=protected-access + async for evt in agent_instance._execute_tool_call(tool_call): + events.append(evt) + + # Middleware intercepted the tool call at execution level + self.assertEqual(len(intercepted_tool_calls), 1) + self.assertIn("ORIGINAL", intercepted_tool_calls[0]) + + # The tool actually received the MODIFIED value + self.assertEqual(len(received_inputs), 1) + self.assertEqual(received_inputs[0], "MODIFIED") + + async def test_on_model_call_middleware_modify_input(self) -> None: + """Test that on_model_call middleware can modify messages and model.""" + + class ModifyMessagesMiddleware(MiddlewareBase): + """Middleware that modifies messages input.""" + + async def on_model_call( + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable[..., Any], + ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]: + """Prepend a system message to the messages list.""" + messages = input_kwargs["messages"] + modified_messages = [ + SystemMsg( + name="system", + content="INJECTED SYSTEM MESSAGE", + ), + ] + messages + + # Pass modified messages to next handler + return await next_handler( + current_model=input_kwargs["current_model"], + messages=modified_messages, + tools=input_kwargs["tools"], + tool_choice=input_kwargs["tool_choice"], + ) + + middleware = ModifyMessagesMiddleware() + + # Track what messages the model receives + received_messages = [] + + class TrackingModel(MockModel): + """Model that tracks received messages.""" + + async def _call_api( + self, + *args: Any, + **kwargs: Any, + ) -> ChatResponse: + """Track the messages and return mock response.""" + messages = kwargs.get("messages", []) + received_messages.extend(messages) + return await super()._call_api(*args, **kwargs) + + tracking_model = TrackingModel() + tracking_model.set_responses( + [ + ChatResponse( + content=[TextBlock(text="response")], + is_last=True, + ), + ], + ) + + agent_instance = Agent( + name="test_agent", + system_prompt="test prompt", + model=tracking_model, + toolkit=self.toolkit, + middlewares=[middleware], + ) + + await agent_instance.reply(UserMsg("user", "test message")) + + # Verify the injected system message is present + system_messages = [m for m in received_messages if m.role == "system"] + self.assertTrue( + any( + "INJECTED SYSTEM MESSAGE" in m.get_text_content() + for m in system_messages + ), + ) + + async def test_on_compress_context_middleware(self) -> None: + """Test on_compress_context middleware follows the onion chain pattern. + + Verifies that: + - Multiple middlewares are chained in onion order (mw1 wraps mw2). + - ``input_kwargs`` carries the correct ``context_config`` and + ``instructions``. + - The ``next_handler`` ultimately calls ``_compress_context_impl``. + - A middleware can short-circuit and skip the actual implementation. + """ + seen_instructions = [] + + # ------------------------------------------------------------------ # + # Middleware that records pre/post and forwards to next_handler. # + # ------------------------------------------------------------------ # + class CompressContextMiddleware(MiddlewareBase): + """Middleware for testing on_compress_context hook.""" + + def __init__(self, log: list, name: str) -> None: + """Initialize the compress context middleware. + + Args: + log (`list`): + The execution log list. + name (`str`): + The middleware name. + """ + self.log = log + self.name = name + + async def on_compress_context( + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable[..., Any], + ) -> None: + """Forward to next handler, recording pre and post.""" + self.log.append(f"{self.name}_pre") + seen_instructions.append(input_kwargs.get("instructions")) + await next_handler(**input_kwargs) + self.log.append(f"{self.name}_post") + + middleware1 = CompressContextMiddleware(self.execution_log, "mw1") + middleware2 = CompressContextMiddleware(self.execution_log, "mw2") + + self.mock_model.set_responses( + [ + ChatResponse( + content=[TextBlock(text="test response")], + is_last=True, + ), + ], + ) + + context_config = ContextConfig(trigger_ratio=0.8, reserve_ratio=0.1) + agent = Agent( + name="test_agent", + system_prompt="test prompt", + model=self.mock_model, + toolkit=self.toolkit, + middlewares=[middleware1, middleware2], + context_config=context_config, + ) + + # Patch _compress_context_impl to avoid real token counting. + instructions = HintBlock( + hint="Keep user requirements while compressing.", + source="user", + ) + with patch.object( + agent, + "_compress_context_impl", + new_callable=AsyncMock, + ) as mock_impl: + await agent.compress_context( + context_config=context_config, + instructions=instructions, + ) + + # _compress_context_impl must have been called exactly once. + mock_impl.assert_awaited_once_with( + context_config=context_config, + instructions=instructions, + ) + + # Verify onion execution order: mw1_pre -> mw2_pre -> mw2_post -> + # mw1_post + expected = ["mw1_pre", "mw2_pre", "mw2_post", "mw1_post"] + self.assertListEqual(self.execution_log, expected) + self.assertListEqual(seen_instructions, [instructions, instructions]) + + async def test_on_compress_context_middleware_modify_instructions( + self, + ) -> None: + """Test that middleware can replace compress_context instructions.""" + + class ReplaceInstructionsMiddleware(MiddlewareBase): + """Middleware that replaces the compression instructions.""" + + def __init__(self, replacement: HintBlock) -> None: + """Initialize the middleware with replacement instructions.""" + self.replacement = replacement + + async def on_compress_context( + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable[..., Any], + ) -> None: + """Replace instructions before forwarding.""" + input_kwargs["instructions"] = self.replacement + await next_handler(**input_kwargs) + + original = HintBlock( + hint="Keep all requirements.", + source="user", + ) + replacement = HintBlock( + hint="Keep only unresolved requirements.", + source="middleware", + ) + context_config = ContextConfig(trigger_ratio=0.8, reserve_ratio=0.1) + agent = Agent( + name="test_agent", + system_prompt="test prompt", + model=self.mock_model, + toolkit=self.toolkit, + middlewares=[ReplaceInstructionsMiddleware(replacement)], + context_config=context_config, + ) + + with patch.object( + agent, + "_compress_context_impl", + new_callable=AsyncMock, + ) as mock_impl: + await agent.compress_context( + context_config=context_config, + instructions=original, + ) + + mock_impl.assert_awaited_once_with( + context_config=context_config, + instructions=replacement, + ) + + async def test_on_compress_context_middleware_short_circuit( + self, + ) -> None: + """Test that a middleware can skip _compress_context_impl entirely.""" + + class SkipCompressMiddleware(MiddlewareBase): + """Middleware that skips the actual compress_context call.""" + + def __init__(self, log: list) -> None: + """Initialize the skip compress middleware. + + Args: + log (`list`): + The execution log list. + """ + self.log = log + + async def on_compress_context( + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable[..., Any], + ) -> None: + """Record the call and skip forwarding to next_handler.""" + self.log.append("skipped") + # Intentionally NOT calling next_handler. + + middleware = SkipCompressMiddleware(self.execution_log) + + self.mock_model.set_responses( + [ + ChatResponse( + content=[TextBlock(text="test response")], + is_last=True, + ), + ], + ) + + agent = Agent( + name="test_agent", + system_prompt="test prompt", + model=self.mock_model, + toolkit=self.toolkit, + middlewares=[middleware], + ) + + with patch.object( + agent, + "_compress_context_impl", + new_callable=AsyncMock, + ) as mock_impl: + await agent.compress_context() + + # _compress_context_impl should NOT have been called. + mock_impl.assert_not_awaited() + + self.assertListEqual(self.execution_log, ["skipped"]) + + async def test_on_compress_context_no_middleware(self) -> None: + """Test that compress_context calls _compress_context_impl directly + when no middleware is registered.""" + + self.mock_model.set_responses( + [ + ChatResponse( + content=[TextBlock(text="test response")], + is_last=True, + ), + ], + ) + + context_config = ContextConfig(trigger_ratio=0.8, reserve_ratio=0.1) + agent = Agent( + name="test_agent", + system_prompt="test prompt", + model=self.mock_model, + toolkit=self.toolkit, + # No middlewares registered. + context_config=context_config, + ) + + with patch.object( + agent, + "_compress_context_impl", + new_callable=AsyncMock, + ) as mock_impl: + await agent.compress_context(context_config=context_config) + mock_impl.assert_awaited_once_with( + context_config=context_config, + instructions=None, + ) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.execution_log.clear() diff --git a/tests/model_anthropic_test.py b/tests/model_anthropic_test.py new file mode 100644 index 0000000000000000000000000000000000000000..f35fb918f79f4d4eb235072295dca448c3808a1d --- /dev/null +++ b/tests/model_anthropic_test.py @@ -0,0 +1,525 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for AnthropicChatModel with mocked API responses. + +Tests cover both non-streaming and streaming modes. +Anthropic uses event-based streaming (message_start, content_block_start, +content_block_delta, message_delta events). +""" +import json +from typing import Any +import unittest +from unittest import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, MagicMock, patch + +from utils import AnyString + +from agentscope.message import TextBlock, ToolCallBlock, ThinkingBlock +from agentscope.model import AnthropicChatModel +from agentscope.credential import AnthropicCredential +from agentscope.tool import ToolChoice + +A = AnyString() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_model(stream: bool = False) -> Any: + return AnthropicChatModel( + credential=AnthropicCredential(api_key="test"), + model="claude-opus-4-5", + stream=stream, + context_size=200_000, + ) + + +def _mock_completion( + text: Any = None, + tool_calls: Any = None, + thinking: Any = None, + response_id: str = "msg-1", +) -> MagicMock: + """Build a mock non-streaming Anthropic Message response.""" + blocks = [] + if thinking: + b = MagicMock() + b.type = "thinking" + b.thinking = thinking + b.signature = "sig123" + blocks.append(b) + if text: + b = MagicMock() + b.type = "text" + b.text = text + blocks.append(b) + if tool_calls: + for tc in tool_calls: + b = MagicMock() + b.type = "tool_use" + b.id = tc["id"] + b.name = tc["name"] + b.input = tc["input"] + blocks.append(b) + + resp = MagicMock() + resp.id = response_id + resp.content = blocks + resp.usage = MagicMock() + resp.usage.input_tokens = 10 + resp.usage.output_tokens = 5 + resp.usage.cache_creation_input_tokens = 0 + resp.usage.cache_read_input_tokens = 0 + return resp + + +def _make_event(event_type: str, **kwargs: Any) -> MagicMock: + """Build a mock Anthropic streaming event.""" + event = MagicMock() + event.type = event_type + for key, val in kwargs.items(): + setattr(event, key, val) + return event + + +class _MockAsyncEventStream: + """Mock async iterator over Anthropic events.""" + + def __init__(self, events: list) -> None: + self._events = events + self._index = 0 + + def __aiter__(self) -> "_MockAsyncEventStream": + return self + + async def __anext__(self) -> Any: + if self._index >= len(self._events): + raise StopAsyncIteration + event = self._events[self._index] + self._index += 1 + return event + + +# --------------------------------------------------------------------------- +# Non-streaming tests +# --------------------------------------------------------------------------- + + +class TestAnthropicNonStream(IsolatedAsyncioTestCase): + """Tests for AnthropicChatModel in non-streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=False) + + @patch("anthropic.AsyncAnthropic") + async def test_text_response(self, mock_client_cls: MagicMock) -> None: + """Non-stream text response returns a single ChatResponse.""" + mock_create = AsyncMock( + return_value=_mock_completion(text="Hello!"), + ) + mock_client_cls.return_value.messages.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + (True, [TextBlock.model_construct(id=A, text="Hello!")]), + ) + self.assertEqual(result.id, "msg-1") + + @patch("anthropic.AsyncAnthropic") + async def test_tool_call_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream tool call response creates ToolCallBlocks.""" + mock_create = AsyncMock( + return_value=_mock_completion( + tool_calls=[ + { + "id": "toolu_1", + "name": "get_weather", + "input": {"city": "Beijing"}, + }, + ], + ), + ) + mock_client_cls.return_value.messages.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ToolCallBlock( + id="toolu_1", + name="get_weather", + input=json.dumps({"city": "Beijing"}), + ), + ], + ), + ) + + @patch("anthropic.AsyncAnthropic") + async def test_thinking_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream response with reasoning creates ThinkingBlock.""" + mock_create = AsyncMock( + return_value=_mock_completion( + thinking="Deep thought...", + text="Answer", + ), + ) + mock_client_cls.return_value.messages.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Deep thought...", + signature="sig123", + ), + TextBlock.model_construct(id=A, text="Answer"), + ], + ), + ) + + +# --------------------------------------------------------------------------- +# Streaming tests +# --------------------------------------------------------------------------- + + +class TestAnthropicStream(IsolatedAsyncioTestCase): + """Tests for AnthropicChatModel in streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=True) + + @patch("anthropic.AsyncAnthropic") + async def test_stream_text(self, mock_client_cls: MagicMock) -> None: + """Stream text yields n deltas + 1 final with full content.""" + msg_usage = MagicMock() + msg_usage.input_tokens = 10 + msg_usage.output_tokens = 0 + msg_usage.cache_creation_input_tokens = 0 + msg_usage.cache_read_input_tokens = 0 + + message = MagicMock() + message.id = "msg-1" + message.usage = msg_usage + + delta1 = MagicMock() + delta1.type = "text_delta" + delta1.text = "Hello" + + delta2 = MagicMock() + delta2.type = "text_delta" + delta2.text = " world" + + msg_delta_usage = MagicMock() + msg_delta_usage.output_tokens = 5 + + events = [ + _make_event("message_start", message=message), + _make_event("content_block_delta", index=0, delta=delta1), + _make_event("content_block_delta", index=0, delta=delta2), + _make_event( + "message_delta", + usage=msg_delta_usage, + ), + ] + mock_create = AsyncMock( + return_value=_MockAsyncEventStream(events), + ) + mock_client_cls.return_value.messages.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + (False, [TextBlock.model_construct(id=A, text="Hello")]), + (False, [TextBlock.model_construct(id=A, text=" world")]), + (True, [TextBlock.model_construct(id=A, text="Hello world")]), + ], + ) + self.assertEqual(responses[-1].id, "msg-1") + + @patch("anthropic.AsyncAnthropic") + async def test_stream_thinking_and_text( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream thinking + text yields deltas then final with signature.""" + msg_usage = MagicMock() + msg_usage.input_tokens = 10 + msg_usage.output_tokens = 0 + msg_usage.cache_creation_input_tokens = 0 + msg_usage.cache_read_input_tokens = 0 + + message = MagicMock() + message.id = "msg-2" + message.usage = msg_usage + + thinking_delta = MagicMock() + thinking_delta.type = "thinking_delta" + thinking_delta.thinking = "Let me think" + + sig_delta = MagicMock() + sig_delta.type = "signature_delta" + sig_delta.signature = "sig_abc" + + text_delta = MagicMock() + text_delta.type = "text_delta" + text_delta.text = "Result" + + events = [ + _make_event("message_start", message=message), + _make_event("content_block_delta", index=0, delta=thinking_delta), + _make_event("content_block_delta", index=0, delta=sig_delta), + _make_event("content_block_delta", index=1, delta=text_delta), + ] + mock_create = AsyncMock( + return_value=_MockAsyncEventStream(events), + ) + mock_client_cls.return_value.messages.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Let me think", + ), + ], + ), + (False, [TextBlock.model_construct(id=A, text="Result")]), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Let me think", + signature="sig_abc", + ), + TextBlock.model_construct(id=A, text="Result"), + ], + ), + ], + ) + + @patch("anthropic.AsyncAnthropic") + async def test_stream_tool_call( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream tool call yields partial deltas then full accumulated + input.""" + msg_usage = MagicMock() + msg_usage.input_tokens = 10 + msg_usage.output_tokens = 0 + msg_usage.cache_creation_input_tokens = 0 + msg_usage.cache_read_input_tokens = 0 + + message = MagicMock() + message.id = "msg-3" + message.usage = msg_usage + + tool_block = MagicMock() + tool_block.type = "tool_use" + tool_block.id = "toolu_1" + tool_block.name = "get_weather" + + json_delta1 = MagicMock() + json_delta1.type = "input_json_delta" + json_delta1.partial_json = '{"city":' + + json_delta2 = MagicMock() + json_delta2.type = "input_json_delta" + json_delta2.partial_json = '"BJ"}' + + events = [ + _make_event("message_start", message=message), + _make_event( + "content_block_start", + index=0, + content_block=tool_block, + ), + _make_event("content_block_delta", index=0, delta=json_delta1), + _make_event("content_block_delta", index=0, delta=json_delta2), + ] + mock_create = AsyncMock( + return_value=_MockAsyncEventStream(events), + ) + mock_client_cls.return_value.messages.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ + ToolCallBlock( + id="toolu_1", + name="get_weather", + input='{"city":', + ), + ], + ), + ( + False, + [ + ToolCallBlock( + id="toolu_1", + name="get_weather", + input='"BJ"}', + ), + ], + ), + ( + True, + [ + ToolCallBlock( + id="toolu_1", + name="get_weather", + input='{"city":"BJ"}', + ), + ], + ), + ], + ) + + +# --------------------------------------------------------------------------- +# _format_tools tests +# --------------------------------------------------------------------------- + +_FT_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "description": "Get the time", + "parameters": { + "type": "object", + "properties": {"timezone": {"type": "string"}}, + "required": ["timezone"], + }, + }, + }, +] + +_FT_TOOLS_ANTHROPIC = [ + { + "name": "get_weather", + "description": "Get the weather", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + { + "name": "get_time", + "description": "Get the time", + "input_schema": { + "type": "object", + "properties": {"timezone": {"type": "string"}}, + "required": ["timezone"], + }, + }, +] + + +class TestAnthropicFormatTools(unittest.TestCase): + """Tests for AnthropicChatModel._format_tools.""" + + def setUp(self) -> None: + """Set up model instance.""" + self.model = _make_model() + + def test_auto_mode(self) -> None: + """Auto mode returns converted tools and type=auto.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS_ANTHROPIC) + self.assertEqual(fmt_choice, {"type": "auto"}) + + def test_none_mode(self) -> None: + """None mode returns converted tools and type=none.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="none"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS_ANTHROPIC) + self.assertEqual(fmt_choice, {"type": "none"}) + + def test_required_mode(self) -> None: + """Required mode maps to type=any.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="required"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS_ANTHROPIC) + self.assertEqual(fmt_choice, {"type": "any"}) + + def test_str_mode_force_call(self) -> None: + """A specific tool name forces that tool call.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="get_weather"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS_ANTHROPIC) + self.assertEqual(fmt_choice, {"type": "tool", "name": "get_weather"}) + + def test_tools_filtered(self) -> None: + """When tool_choice.tools is set, only those tools are included.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto", tools=["get_weather"]), + ) + self.assertEqual(len(fmt_tools), 1) + self.assertEqual(fmt_tools[0]["name"], "get_weather") + self.assertEqual(fmt_choice, {"type": "auto"}) + + def test_no_tool_choice(self) -> None: + """Without tool_choice, returns converted tools and None.""" + fmt_tools, fmt_choice = self.model._format_tools(_FT_TOOLS, None) + self.assertEqual(fmt_tools, _FT_TOOLS_ANTHROPIC) + self.assertIsNone(fmt_choice) diff --git a/tests/model_dashscope_test.py b/tests/model_dashscope_test.py new file mode 100644 index 0000000000000000000000000000000000000000..4be81954f707ee200c25b6fe58ab24071c639567 --- /dev/null +++ b/tests/model_dashscope_test.py @@ -0,0 +1,589 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for DashScopeChatModel with mocked API responses. + +Tests cover both non-streaming and streaming modes. +""" +import base64 +import io +import wave +from typing import Any +import unittest +from unittest import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, MagicMock, patch + +from utils import AnyString + +from agentscope.message import ( + TextBlock, + ToolCallBlock, + ThinkingBlock, + DataBlock, +) +from agentscope.model import DashScopeChatModel +from agentscope.credential import DashScopeCredential +from agentscope.tool import ToolChoice + +A = AnyString() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_model(stream: bool = False) -> Any: + return DashScopeChatModel( + credential=DashScopeCredential(api_key="test"), + model="qwen3-max", + stream=stream, + max_retries=3, + context_size=131_072, + parameters=DashScopeChatModel.Parameters( + max_tokens=1000, + thinking_enable=True, + thinking_budget=100, + ), + ) + + +def _mock_completion( + text: Any = None, + tool_calls: Any = None, + reasoning: Any = None, + response_id: str = "req-1", +) -> MagicMock: + """Build a mock non-streaming ChatCompletion response.""" + msg = MagicMock() + msg.content = text + msg.reasoning_content = reasoning + msg.tool_calls = None + + if tool_calls: + tc_mocks = [] + for tc in tool_calls: + m = MagicMock() + m.id = tc["id"] + m.function.name = tc["name"] + m.function.arguments = tc["arguments"] + tc_mocks.append(m) + msg.tool_calls = tc_mocks + + choice = MagicMock() + choice.message = msg + + resp = MagicMock() + resp.id = response_id + resp.choices = [choice] + resp.usage = MagicMock() + resp.usage.prompt_tokens = 10 + resp.usage.completion_tokens = 5 + resp.usage.prompt_tokens_details = None + return resp + + +def _make_stream_chunk( + delta_text: str | None = None, + delta_reasoning: str | None = None, + tool_calls: list | None = None, + delta_audio: dict | None = None, + response_id: str = "req-1", + usage: dict | None = None, + has_choices: bool = True, +) -> MagicMock: + """Build a single mock streaming chunk.""" + chunk = MagicMock() + chunk.id = response_id + + if usage: + chunk.usage = MagicMock() + chunk.usage.prompt_tokens = usage.get("prompt_tokens", 0) + chunk.usage.completion_tokens = usage.get("completion_tokens", 0) + chunk.usage.prompt_tokens_details = None + else: + chunk.usage = None + + if has_choices: + delta = MagicMock() + delta.content = delta_text + delta.reasoning_content = delta_reasoning + delta.tool_calls = tool_calls + delta.audio = delta_audio + choice = MagicMock() + choice.delta = delta + chunk.choices = [choice] + else: + chunk.choices = [] + + return chunk + + +def _make_tool_call_delta( + index: int, + tc_id: str | None = None, + name: str | None = None, + arguments: str | None = None, +) -> MagicMock: + tc = MagicMock() + tc.index = index + tc.id = tc_id + tc.function = MagicMock() + tc.function.name = name + tc.function.arguments = arguments + return tc + + +class _MockAsyncStream: + """Mock async stream (context manager + async iterator).""" + + def __init__(self, chunks: list) -> None: + self._chunks = chunks + self._index = 0 + + async def __aenter__(self) -> "_MockAsyncStream": + return self + + async def __aexit__(self, *args: Any) -> None: + pass + + def __aiter__(self) -> "_MockAsyncStream": + return self + + async def __anext__(self) -> Any: + if self._index >= len(self._chunks): + raise StopAsyncIteration + chunk = self._chunks[self._index] + self._index += 1 + return chunk + + +# --------------------------------------------------------------------------- +# Non-streaming tests +# --------------------------------------------------------------------------- + + +class TestDashScopeNonStream(IsolatedAsyncioTestCase): + """Tests for DashScopeChatModel in non-streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=False) + + @patch("openai.AsyncClient") + async def test_text_response(self, mock_client_cls: MagicMock) -> None: + """Non-stream text response returns a single ChatResponse.""" + mock_create = AsyncMock( + return_value=_mock_completion(text="Hello!"), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + (True, [TextBlock.model_construct(id=A, text="Hello!")]), + ) + self.assertEqual(result.id, "req-1") + + @patch("openai.AsyncClient") + async def test_tool_call_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream tool call response creates ToolCallBlocks.""" + mock_create = AsyncMock( + return_value=_mock_completion( + tool_calls=[ + { + "id": "call-1", + "name": "get_weather", + "arguments": '{"city":"Hangzhou"}', + }, + ], + ), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ToolCallBlock( + id="call-1", + name="get_weather", + input='{"city":"Hangzhou"}', + ), + ], + ), + ) + + @patch("openai.AsyncClient") + async def test_thinking_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream response with reasoning creates ThinkingBlock.""" + mock_create = AsyncMock( + return_value=_mock_completion( + text="42", + reasoning="Reasoning step...", + ), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Reasoning step...", + ), + TextBlock.model_construct(id=A, text="42"), + ], + ), + ) + + +# --------------------------------------------------------------------------- +# Streaming tests +# --------------------------------------------------------------------------- + + +class TestDashScopeStream(IsolatedAsyncioTestCase): + """Tests for DashScopeChatModel in streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=True) + + @patch("openai.AsyncClient") + async def test_stream_text(self, mock_client_cls: MagicMock) -> None: + """Stream text yields n deltas + 1 final with full content.""" + chunks = [ + _make_stream_chunk(delta_text="Hello"), + _make_stream_chunk(delta_text=" world"), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 10, "completion_tokens": 2}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + (False, [TextBlock.model_construct(id=A, text="Hello")]), + (False, [TextBlock.model_construct(id=A, text=" world")]), + (True, [TextBlock.model_construct(id=A, text="Hello world")]), + ], + ) + + @patch("openai.AsyncClient") + async def test_stream_thinking_and_text( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream reasoning + text yields deltas then accumulated final.""" + chunks = [ + _make_stream_chunk(delta_reasoning="Think"), + _make_stream_chunk(delta_reasoning="ing"), + _make_stream_chunk(delta_text="Answer"), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 10, "completion_tokens": 5}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ThinkingBlock.model_construct(id=A, thinking="Think")], + ), + (False, [ThinkingBlock.model_construct(id=A, thinking="ing")]), + (False, [TextBlock.model_construct(id=A, text="Answer")]), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Thinking", + ), + TextBlock.model_construct(id=A, text="Answer"), + ], + ), + ], + ) + + @patch("openai.AsyncClient") + async def test_stream_tool_calls( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream tool call chunks accumulate into final ToolCallBlock.""" + chunks = [ + _make_stream_chunk( + tool_calls=[ + _make_tool_call_delta(0, "call-1", "search", '{"q":'), + ], + ), + _make_stream_chunk( + tool_calls=[ + _make_tool_call_delta(0, None, None, '"hello"}'), + ], + ), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 10, "completion_tokens": 5}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ + ToolCallBlock( + id="call-1", + name="search", + input='{"q":', + ), + ], + ), + ( + False, + [ + ToolCallBlock( + id="call-1", + name="search", + input='"hello"}', + ), + ], + ), + ( + True, + [ + ToolCallBlock( + id="call-1", + name="search", + input='{"q":"hello"}', + ), + ], + ), + ], + ) + + @patch("openai.AsyncClient") + async def test_stream_usage(self, mock_client_cls: MagicMock) -> None: + """Stream usage chunk attaches token counts to final response.""" + chunks = [ + _make_stream_chunk(delta_text="X"), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 50, "completion_tokens": 10}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + (False, [TextBlock.model_construct(id=A, text="X")]), + (True, [TextBlock.model_construct(id=A, text="X")]), + ], + ) + self.assertEqual(responses[-1].usage.input_tokens, 50) + self.assertEqual(responses[-1].usage.output_tokens, 10) + + @patch("openai.AsyncClient") + async def test_stream_audio_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream PCM deltas produce per-chunk DataBlocks (first chunk + prefixed with a streaming WAV header) sharing a stable id, plus a + final fixed-size WAV block readable by the ``wave`` module.""" + pcm1 = bytes([1, 2, 3, 4]) + pcm2 = bytes([5, 6, 7, 8]) + pcm3 = bytes([9, 10, 11, 12]) + pcm_full = pcm1 + pcm2 + pcm3 + + chunks = [ + _make_stream_chunk( + delta_audio={"data": base64.b64encode(pcm1).decode()}, + ), + _make_stream_chunk( + delta_audio={"data": base64.b64encode(pcm2).decode()}, + ), + _make_stream_chunk( + delta_audio={"data": base64.b64encode(pcm3).decode()}, + ), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 5, "completion_tokens": 3}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + self.assertEqual(len(responses), 4) + + # All four chunks (3 deltas + 1 final) must share the same audio + # block id so downstream consumers stitch them as one stream. + all_audio_ids = { + block.id + for r in responses + for block in r.content + if isinstance(block, DataBlock) + } + self.assertEqual(len(all_audio_ids), 1) + + # First delta: WAV header (44 bytes, "RIFF"..."WAVE") + pcm1. + first_payload = base64.b64decode(responses[0].content[0].source.data) + self.assertEqual(len(first_payload), 44 + len(pcm1)) + self.assertEqual(first_payload[:4], b"RIFF") + self.assertEqual(first_payload[8:12], b"WAVE") + self.assertEqual(first_payload[44:], pcm1) + + # Subsequent deltas: raw PCM only, no header. + for resp, pcm in zip(responses[1:3], [pcm2, pcm3]): + payload = base64.b64decode(resp.content[0].source.data) + self.assertEqual(payload, pcm) + + # Final ``is_last`` block: a fixed-size WAV the ``wave`` module + # can parse end-to-end at 24kHz / mono / 16-bit. + final = responses[-1] + self.assertTrue(final.is_last) + final_audio = next( + b for b in final.content if isinstance(b, DataBlock) + ) + wav_bytes = base64.b64decode(final_audio.source.data) + with wave.open(io.BytesIO(wav_bytes), "rb") as wav: + self.assertEqual(wav.getnchannels(), 1) + self.assertEqual(wav.getsampwidth(), 2) + self.assertEqual(wav.getframerate(), 24000) + frames = wav.readframes(wav.getnframes()) + self.assertEqual(frames, pcm_full) + + +# --------------------------------------------------------------------------- +# _format_tools tests +# --------------------------------------------------------------------------- + +_FT_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "description": "Get the time", + "parameters": { + "type": "object", + "properties": {"timezone": {"type": "string"}}, + "required": ["timezone"], + }, + }, + }, +] + + +class TestDashScopeFormatTools(unittest.TestCase): + """Tests for DashScopeChatModel._format_tools.""" + + def setUp(self) -> None: + """Set up model instance.""" + self.model = _make_model() + + def test_auto_mode(self) -> None: + """Auto mode returns tools unchanged and string 'auto'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual(fmt_choice, "auto") + + def test_none_mode(self) -> None: + """None mode returns tools unchanged and string 'none'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="none"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual(fmt_choice, "none") + + def test_required_mode_warns(self) -> None: + """Required mode emits a DeprecationWarning and falls back to auto.""" + with self.assertWarns(DeprecationWarning): + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="required"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual(fmt_choice, "auto") + + def test_str_mode_force_call(self) -> None: + """A specific tool name forces that tool call.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="get_weather"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual( + fmt_choice, + {"type": "function", "function": {"name": "get_weather"}}, + ) + + def test_tools_filtered(self) -> None: + """When tool_choice.tools is set, only those tools are included.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto", tools=["get_weather"]), + ) + self.assertEqual(len(fmt_tools), 1) + self.assertEqual(fmt_tools[0]["function"]["name"], "get_weather") + self.assertEqual(fmt_choice, "auto") + + def test_no_tool_choice(self) -> None: + """Without tool_choice, returns tools and None.""" + fmt_tools, fmt_choice = self.model._format_tools(_FT_TOOLS, None) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertIsNone(fmt_choice) diff --git a/tests/model_deepseek_test.py b/tests/model_deepseek_test.py new file mode 100644 index 0000000000000000000000000000000000000000..2d7a36151a428cb62db38f0c0275614cd3864f1c --- /dev/null +++ b/tests/model_deepseek_test.py @@ -0,0 +1,641 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for DeepSeekChatModel with mocked API responses. + +Tests cover both non-streaming and streaming modes, verifying that: +- Non-stream mode returns a single ChatResponse with is_last=True. +- Stream mode yields n delta ChatResponses (is_last=False) followed by + 1 final ChatResponse (is_last=True) with the full accumulated content. +""" +from typing import Any +import unittest +from unittest import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, MagicMock, patch + +from utils import AnyString + +from agentscope.message import TextBlock, ToolCallBlock, ThinkingBlock +from agentscope.model import DeepSeekChatModel +from agentscope.credential import DeepSeekCredential +from agentscope.tool import ToolChoice + +A = AnyString() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_model(stream: bool = False) -> Any: + return DeepSeekChatModel( + credential=DeepSeekCredential(api_key="test"), + model="deepseek-v4-pro", + stream=stream, + context_size=65_536, + ) + + +def _mock_completion( + text: Any = None, + tool_calls: Any = None, + reasoning: Any = None, + response_id: str = "deepseek-1", +) -> MagicMock: + """Build a mock non-streaming ChatCompletion response.""" + msg = MagicMock() + msg.content = text + msg.reasoning_content = reasoning + msg.tool_calls = None + + if tool_calls: + tc_mocks = [] + for tc in tool_calls: + m = MagicMock() + m.id = tc["id"] + m.function.name = tc["name"] + m.function.arguments = tc["arguments"] + tc_mocks.append(m) + msg.tool_calls = tc_mocks + + choice = MagicMock() + choice.message = msg + + resp = MagicMock() + resp.id = response_id + resp.choices = [choice] + resp.usage.prompt_tokens = 10 + resp.usage.completion_tokens = 5 + resp.usage.prompt_cache_hit_tokens = 0 + return resp + + +def _make_stream_chunk( + delta_text: str | None = None, + delta_reasoning: str | None = None, + tool_calls: list | None = None, + response_id: str = "deepseek-1", + usage: dict | None = None, + has_choices: bool = True, +) -> MagicMock: + """Build a single mock streaming chunk.""" + chunk = MagicMock() + chunk.id = response_id + + if usage: + chunk.usage = MagicMock() + chunk.usage.prompt_tokens = usage.get("prompt_tokens", 0) + chunk.usage.completion_tokens = usage.get("completion_tokens", 0) + chunk.usage.prompt_cache_hit_tokens = usage.get( + "prompt_cache_hit_tokens", + 0, + ) + else: + chunk.usage = None + + if has_choices: + delta = MagicMock() + delta.content = delta_text + delta.reasoning_content = delta_reasoning + delta.tool_calls = tool_calls + choice = MagicMock() + choice.delta = delta + chunk.choices = [choice] + else: + chunk.choices = [] + + return chunk + + +def _make_tool_call_delta( + index: int, + tc_id: str | None = None, + name: str | None = None, + arguments: str | None = None, +) -> MagicMock: + """Build a tool_call delta item for streaming.""" + tc = MagicMock() + tc.index = index + tc.id = tc_id + tc.function.name = name + tc.function.arguments = arguments + return tc + + +class _MockAsyncStream: + """Mock async stream that acts as an async context manager + iterator.""" + + def __init__(self, chunks: list) -> None: + self._chunks = chunks + self._index = 0 + + async def __aenter__(self) -> "_MockAsyncStream": + return self + + async def __aexit__(self, *args: Any) -> None: + pass + + def __aiter__(self) -> "_MockAsyncStream": + return self + + async def __anext__(self) -> Any: + if self._index >= len(self._chunks): + raise StopAsyncIteration + chunk = self._chunks[self._index] + self._index += 1 + return chunk + + +# --------------------------------------------------------------------------- +# Non-streaming tests +# --------------------------------------------------------------------------- + + +class TestDeepSeekNonStream(IsolatedAsyncioTestCase): + """Tests for DeepSeekChatModel in non-streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=False) + + @patch("openai.AsyncClient") + async def test_text_response(self, mock_client_cls: MagicMock) -> None: + """Non-stream text response returns a single ChatResponse.""" + mock_create = AsyncMock( + return_value=_mock_completion(text="Hello world!"), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + (True, [TextBlock.model_construct(id=A, text="Hello world!")]), + ) + self.assertEqual(result.id, "deepseek-1") + + @patch("openai.AsyncClient") + async def test_tool_call_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream tool call response creates ToolCallBlocks.""" + mock_create = AsyncMock( + return_value=_mock_completion( + tool_calls=[ + { + "id": "call-1", + "name": "get_weather", + "arguments": '{"city":"Beijing"}', + }, + { + "id": "call-2", + "name": "get_time", + "arguments": '{"tz":"UTC"}', + }, + ], + ), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ToolCallBlock( + id="call-1", + name="get_weather", + input='{"city":"Beijing"}', + ), + ToolCallBlock( + id="call-2", + name="get_time", + input='{"tz":"UTC"}', + ), + ], + ), + ) + + @patch("openai.AsyncClient") + async def test_thinking_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream response with reasoning creates ThinkingBlock.""" + mock_create = AsyncMock( + return_value=_mock_completion( + text="The answer is 42.", + reasoning="Let me think step by step...", + ), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Let me think step by step...", + ), + TextBlock.model_construct(id=A, text="The answer is 42."), + ], + ), + ) + + +# --------------------------------------------------------------------------- +# Streaming tests +# --------------------------------------------------------------------------- + + +class TestDeepSeekStream(IsolatedAsyncioTestCase): + """Tests for DeepSeekChatModel in streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=True) + + @patch("openai.AsyncClient") + async def test_stream_text_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream text yields n deltas (is_last=False) + 1 final + (is_last=True) with full content.""" + chunks = [ + _make_stream_chunk(delta_text="Hello"), + _make_stream_chunk(delta_text=" world"), + _make_stream_chunk(delta_text="!"), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 10, "completion_tokens": 3}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + (False, [TextBlock.model_construct(id=A, text="Hello")]), + (False, [TextBlock.model_construct(id=A, text=" world")]), + (False, [TextBlock.model_construct(id=A, text="!")]), + (True, [TextBlock.model_construct(id=A, text="Hello world!")]), + ], + ) + self.assertEqual(responses[-1].id, "deepseek-1") + + @patch("openai.AsyncClient") + async def test_stream_thinking_and_text( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream with thinking + text yields deltas then final with both.""" + chunks = [ + _make_stream_chunk(delta_reasoning="Think"), + _make_stream_chunk(delta_reasoning="ing..."), + _make_stream_chunk(delta_text="Answer"), + _make_stream_chunk(delta_text=" here."), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 10, "completion_tokens": 8}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ThinkingBlock.model_construct(id=A, thinking="Think")], + ), + ( + False, + [ThinkingBlock.model_construct(id=A, thinking="ing...")], + ), + (False, [TextBlock.model_construct(id=A, text="Answer")]), + (False, [TextBlock.model_construct(id=A, text=" here.")]), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Thinking...", + ), + TextBlock.model_construct(id=A, text="Answer here."), + ], + ), + ], + ) + + @patch("openai.AsyncClient") + async def test_stream_tool_calls( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream tool calls accumulate across chunks into final response.""" + chunks = [ + _make_stream_chunk( + delta_text=None, + tool_calls=[ + _make_tool_call_delta(0, "call-1", "get_weather", '{"ci'), + ], + ), + _make_stream_chunk( + delta_text=None, + tool_calls=[ + _make_tool_call_delta(0, None, None, 'ty":"BJ"}'), + ], + ), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 10, "completion_tokens": 5}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ + ToolCallBlock( + id="call-1", + name="get_weather", + input='{"ci', + ), + ], + ), + ( + False, + [ + ToolCallBlock( + id="call-1", + name="get_weather", + input='ty":"BJ"}', + ), + ], + ), + ( + True, + [ + ToolCallBlock( + id="call-1", + name="get_weather", + input='{"city":"BJ"}', + ), + ], + ), + ], + ) + + @patch("openai.AsyncClient") + async def test_stream_text_then_tool_call( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream with text followed by tool call accumulates both.""" + chunks = [ + _make_stream_chunk(delta_reasoning="Let me check"), + _make_stream_chunk(delta_text="I'll look it up."), + _make_stream_chunk( + delta_text=None, + tool_calls=[ + _make_tool_call_delta( + 0, + "call-1", + "search", + '{"q":"weather"}', + ), + ], + ), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 15, "completion_tokens": 10}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Let me check", + ), + ], + ), + ( + False, + [TextBlock.model_construct(id=A, text="I'll look it up.")], + ), + ( + False, + [ + ToolCallBlock( + id="call-1", + name="search", + input='{"q":"weather"}', + ), + ], + ), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Let me check", + ), + TextBlock.model_construct( + id=A, + text="I'll look it up.", + ), + ToolCallBlock( + id="call-1", + name="search", + input='{"q":"weather"}', + ), + ], + ), + ], + ) + + @patch("openai.AsyncClient") + async def test_stream_usage_in_final( + self, + mock_client_cls: MagicMock, + ) -> None: + """Usage information is captured and present in final response.""" + chunks = [ + _make_stream_chunk(delta_text="Hi"), + _make_stream_chunk( + has_choices=False, + usage={ + "prompt_tokens": 100, + "completion_tokens": 20, + "prompt_cache_hit_tokens": 50, + }, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + (False, [TextBlock.model_construct(id=A, text="Hi")]), + (True, [TextBlock.model_construct(id=A, text="Hi")]), + ], + ) + self.assertEqual(responses[-1].usage.input_tokens, 100) + self.assertEqual(responses[-1].usage.output_tokens, 20) + self.assertEqual(responses[-1].usage.cache_input_tokens, 50) + + +class TestDeepSeekModelParameters(unittest.TestCase): + """Tests for DeepSeekChatModel.Parameters.""" + + def test_thinking_enable_stored_on_model(self) -> None: + """thinking_enable is accessible through model.parameters.""" + model = DeepSeekChatModel( + credential=DeepSeekCredential(api_key="test"), + model="deepseek-reasoner", + stream=False, + context_size=65_536, + parameters=DeepSeekChatModel.Parameters(thinking_enable=True), + ) + self.assertTrue(model.parameters.thinking_enable) + + def test_reasoning_effort_stored_on_model(self) -> None: + """reasoning_effort is accessible through model.parameters.""" + model = DeepSeekChatModel( + credential=DeepSeekCredential(api_key="test"), + model="deepseek-reasoner", + stream=False, + context_size=65_536, + parameters=DeepSeekChatModel.Parameters(reasoning_effort="max"), + ) + self.assertEqual(model.parameters.reasoning_effort, "max") + + +# --------------------------------------------------------------------------- +# Shared _format_tools fixtures +# --------------------------------------------------------------------------- + +_FT_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "description": "Get the time", + "parameters": { + "type": "object", + "properties": {"timezone": {"type": "string"}}, + "required": ["timezone"], + }, + }, + }, +] + + +class TestDeepSeekFormatTools(unittest.TestCase): + """Tests for DeepSeekChatModel._format_tools.""" + + def setUp(self) -> None: + """Set up model instance.""" + self.model = _make_model() + + def test_auto_mode(self) -> None: + """Auto mode returns tools unchanged and string 'auto'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual(fmt_choice, "auto") + + def test_none_mode(self) -> None: + """None mode returns tools unchanged and string 'none'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="none"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual(fmt_choice, "none") + + def test_required_mode(self) -> None: + """Required mode returns tools unchanged and string 'required'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="required"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual(fmt_choice, "required") + + def test_str_mode_force_call(self) -> None: + """A specific tool name returns a type=function dict.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="get_weather"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual( + fmt_choice, + {"type": "function", "function": {"name": "get_weather"}}, + ) + + def test_tools_filtered(self) -> None: + """When tool_choice.tools is set, only those tools are included.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto", tools=["get_weather"]), + ) + self.assertEqual(len(fmt_tools), 1) + self.assertEqual(fmt_tools[0]["function"]["name"], "get_weather") + self.assertEqual(fmt_choice, "auto") + + def test_no_tool_choice(self) -> None: + """Without tool_choice, returns tools and None.""" + fmt_tools, fmt_choice = self.model._format_tools(_FT_TOOLS, None) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertIsNone(fmt_choice) diff --git a/tests/model_gemini_test.py b/tests/model_gemini_test.py new file mode 100644 index 0000000000000000000000000000000000000000..decd6eee0d2ee73f63d8a78c4c90fd62c2aecaa4 --- /dev/null +++ b/tests/model_gemini_test.py @@ -0,0 +1,675 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for GeminiChatModel with mocked API responses. + +Tests cover both non-streaming and streaming modes. +Gemini uses google.genai client with async iterator streaming. +""" +import json +from typing import Any +import unittest +from unittest import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, MagicMock, patch + +from utils import AnyString + +from agentscope.message import TextBlock, ToolCallBlock, ThinkingBlock +from agentscope.model import GeminiChatModel +from agentscope.model._gemini._model import _sanitize_schema_for_gemini +from agentscope._utils._common import _flatten_json_schema +from agentscope.credential import GeminiCredential +from agentscope.tool import ToolChoice + +A = AnyString() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_model(stream: bool = False) -> Any: + return GeminiChatModel( + credential=GeminiCredential(api_key="test"), + model="gemini-2.5-flash", + stream=stream, + context_size=1_048_576, + ) + + +def _make_part( + text: str | None = None, + thought: bool = False, + function_call: dict | None = None, + thought_signature: Any = None, +) -> MagicMock: + """Build a mock Gemini Part.""" + part = MagicMock() + part.text = text + part.thought = thought + part.thought_signature = thought_signature + if function_call: + part.function_call = MagicMock() + part.function_call.name = function_call["name"] + part.function_call.args = function_call.get("args", {}) + part.function_call.id = function_call.get("id", "call-1") + else: + part.function_call = None + return part + + +def _mock_completion( + parts: list, + response_id: str = "resp-gem-1", +) -> MagicMock: + """Build a mock non-streaming Gemini response.""" + resp = MagicMock() + resp.response_id = response_id + resp.candidates = [MagicMock()] + resp.candidates[0].content = MagicMock() + resp.candidates[0].content.parts = parts + resp.usage_metadata = MagicMock() + resp.usage_metadata.prompt_token_count = 10 + resp.usage_metadata.candidates_token_count = 5 + return resp + + +def _make_stream_chunk( + parts: list, + response_id: str = "resp-gem-1", +) -> MagicMock: + """Build a single mock streaming chunk.""" + chunk = MagicMock() + chunk.response_id = response_id + chunk.candidates = [MagicMock()] + chunk.candidates[0].content = MagicMock() + chunk.candidates[0].content.parts = parts + chunk.usage_metadata = MagicMock() + chunk.usage_metadata.prompt_token_count = 10 + chunk.usage_metadata.candidates_token_count = 5 + return chunk + + +class _MockAsyncStream: + """Mock async iterator for Gemini stream.""" + + def __init__(self, chunks: list) -> None: + self._chunks = chunks + self._index = 0 + + def __aiter__(self) -> "_MockAsyncStream": + return self + + async def __anext__(self) -> Any: + if self._index >= len(self._chunks): + raise StopAsyncIteration + chunk = self._chunks[self._index] + self._index += 1 + return chunk + + +# --------------------------------------------------------------------------- +# Non-streaming tests +# --------------------------------------------------------------------------- + + +class TestGeminiNonStream(IsolatedAsyncioTestCase): + """Tests for GeminiChatModel in non-streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=False) + + @patch("google.genai.Client") + async def test_text_response(self, mock_client_cls: MagicMock) -> None: + """Non-stream text response returns a single ChatResponse.""" + parts = [_make_part(text="Hello!")] + mock_client_cls.return_value.aio.models.generate_content = AsyncMock( + return_value=_mock_completion(parts), + ) + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + (True, [TextBlock.model_construct(id=A, text="Hello!")]), + ) + + @patch("google.genai.Client") + async def test_tool_call_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream tool call response creates ToolCallBlocks.""" + parts = [ + _make_part( + function_call={ + "name": "get_weather", + "args": {"city": "Tokyo"}, + "id": "call-1", + }, + ), + ] + mock_client_cls.return_value.aio.models.generate_content = AsyncMock( + return_value=_mock_completion(parts), + ) + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ToolCallBlock( + id="call-1", + name="get_weather", + input=json.dumps( + {"city": "Tokyo"}, + ensure_ascii=False, + ), + ), + ], + ), + ) + + @patch("google.genai.Client") + async def test_tool_call_response_without_id( + self, + mock_client_cls: MagicMock, + ) -> None: + """A function call with no id gets a generated id, not a crash.""" + parts = [ + _make_part( + function_call={ + "name": "get_weather", + "args": {"city": "Tokyo"}, + "id": None, + }, + ), + ] + mock_client_cls.return_value.aio.models.generate_content = AsyncMock( + return_value=_mock_completion(parts), + ) + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ToolCallBlock.model_construct( + id=A, + name="get_weather", + input=json.dumps( + {"city": "Tokyo"}, + ensure_ascii=False, + ), + ), + ], + ), + ) + self.assertIsInstance(result.content[0].id, str) + self.assertTrue(result.content[0].id) + + @patch("google.genai.Client") + async def test_thinking_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream response with reasoning creates ThinkingBlock.""" + parts = [ + _make_part(text="Let me think...", thought=True), + _make_part(text="Answer"), + ] + mock_client_cls.return_value.aio.models.generate_content = AsyncMock( + return_value=_mock_completion(parts), + ) + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Let me think...", + ), + TextBlock.model_construct(id=A, text="Answer"), + ], + ), + ) + + +# --------------------------------------------------------------------------- +# Streaming tests +# --------------------------------------------------------------------------- + + +class TestGeminiStream(IsolatedAsyncioTestCase): + """Tests for GeminiChatModel in streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=True) + + @patch("google.genai.Client") + async def test_stream_text(self, mock_client_cls: MagicMock) -> None: + """Stream text yields n deltas + 1 final with full content.""" + chunks = [ + _make_stream_chunk([_make_part(text="Hello")]), + _make_stream_chunk([_make_part(text=" world")]), + ] + mock_client_cls.return_value.aio.models.generate_content_stream = ( + AsyncMock(return_value=_MockAsyncStream(chunks)) + ) + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + (False, [TextBlock.model_construct(id=A, text="Hello")]), + (False, [TextBlock.model_construct(id=A, text=" world")]), + (True, [TextBlock.model_construct(id=A, text="Hello world")]), + ], + ) + + @patch("google.genai.Client") + async def test_stream_thinking_and_text( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream thinking + text yields deltas then accumulated final.""" + chunks = [ + _make_stream_chunk([_make_part(text="Think", thought=True)]), + _make_stream_chunk([_make_part(text="Answer")]), + ] + mock_client_cls.return_value.aio.models.generate_content_stream = ( + AsyncMock(return_value=_MockAsyncStream(chunks)) + ) + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ThinkingBlock.model_construct(id=A, thinking="Think")], + ), + (False, [TextBlock.model_construct(id=A, text="Answer")]), + ( + True, + [ + ThinkingBlock.model_construct(id=A, thinking="Think"), + TextBlock.model_construct(id=A, text="Answer"), + ], + ), + ], + ) + + @patch("google.genai.Client") + async def test_stream_tool_call( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream tool call yields delta then final with same ToolCallBlock.""" + chunks = [ + _make_stream_chunk( + [ + _make_part( + function_call={ + "name": "search", + "args": {"q": "test"}, + "id": "call-1", + }, + ), + ], + ), + ] + mock_client_cls.return_value.aio.models.generate_content_stream = ( + AsyncMock(return_value=_MockAsyncStream(chunks)) + ) + + gen = await self.model([]) + responses = [r async for r in gen] + + tool_block = ToolCallBlock( + id="call-1", + name="search", + input=json.dumps({"q": "test"}, ensure_ascii=False), + ) + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + (False, [tool_block]), + (True, [tool_block]), + ], + ) + + @patch("google.genai.Client") + async def test_stream_tool_calls_without_id( + self, + mock_client_cls: MagicMock, + ) -> None: + """Two id-less function calls in one chunk get distinct ids.""" + chunks = [ + _make_stream_chunk( + [ + _make_part( + function_call={ + "name": "search", + "args": {"q": "a"}, + "id": None, + }, + ), + _make_part( + function_call={ + "name": "search", + "args": {"q": "b"}, + "id": None, + }, + ), + ], + ), + ] + mock_client_cls.return_value.aio.models.generate_content_stream = ( + AsyncMock(return_value=_MockAsyncStream(chunks)) + ) + + gen = await self.model([]) + responses = [r async for r in gen] + + final = responses[-1] + self.assertTrue(final.is_last) + self.assertEqual(len(final.content), 2) + ids = [block.id for block in final.content] + self.assertTrue(all(isinstance(i, str) and i for i in ids)) + self.assertNotEqual(ids[0], ids[1]) + self.assertEqual( + [block.input for block in final.content], + [ + json.dumps({"q": "a"}, ensure_ascii=False), + json.dumps({"q": "b"}, ensure_ascii=False), + ], + ) + + +# --------------------------------------------------------------------------- +# _format_tools tests +# --------------------------------------------------------------------------- + +_FT_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "description": "Get the time", + "parameters": { + "type": "object", + "properties": {"timezone": {"type": "string"}}, + "required": ["timezone"], + }, + }, + }, +] + +_FT_TOOLS_GEMINI = [ + { + "function_declarations": [ + { + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + { + "name": "get_time", + "description": "Get the time", + "parameters": { + "type": "object", + "properties": {"timezone": {"type": "string"}}, + "required": ["timezone"], + }, + }, + ], + }, +] + + +class TestGeminiFormatTools(unittest.TestCase): + """Tests for GeminiChatModel._format_tools.""" + + def setUp(self) -> None: + """Set up model instance.""" + self.model = _make_model() + + def test_auto_mode(self) -> None: + """Auto mode returns function_declarations and AUTO config.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS_GEMINI) + self.assertEqual( + fmt_choice, + {"function_calling_config": {"mode": "AUTO"}}, + ) + + def test_none_mode(self) -> None: + """None mode returns function_declarations and NONE config.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="none"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS_GEMINI) + self.assertEqual( + fmt_choice, + {"function_calling_config": {"mode": "NONE"}}, + ) + + def test_required_mode(self) -> None: + """Required mode maps to ANY config.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="required"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS_GEMINI) + self.assertEqual( + fmt_choice, + {"function_calling_config": {"mode": "ANY"}}, + ) + + def test_str_mode_force_call(self) -> None: + """A specific tool name restricts via allowed_function_names.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="get_weather"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS_GEMINI) + self.assertEqual( + fmt_choice, + { + "function_calling_config": { + "mode": "ANY", + "allowed_function_names": ["get_weather"], + }, + }, + ) + + def test_tools_filtered(self) -> None: + """When tool_choice.tools is set, only those tools are included.""" + fmt_tools, _ = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto", tools=["get_weather"]), + ) + self.assertEqual(len(fmt_tools[0]["function_declarations"]), 1) + self.assertEqual( + fmt_tools[0]["function_declarations"][0]["name"], + "get_weather", + ) + + def test_no_tool_choice(self) -> None: + """Without tool_choice, returns function_declarations and None.""" + fmt_tools, fmt_choice = self.model._format_tools(_FT_TOOLS, None) + self.assertEqual(fmt_tools, _FT_TOOLS_GEMINI) + self.assertIsNone(fmt_choice) + + +# --------------------------------------------------------------------------- +# Tests for _sanitize_schema_for_gemini and _flatten_json_schema +# --------------------------------------------------------------------------- + + +class TestGeminiSchemaUtils(unittest.TestCase): + """Tests for _sanitize_schema_for_gemini and _flatten_json_schema.""" + + def test_sanitize_removes_additional_properties_and_inlines_optional( + self, + ) -> None: + """additionalProperties removed; anyOf[X, null] inlined to X.""" + self.assertEqual( + _sanitize_schema_for_gemini( + { + "description": "x", + "anyOf": [{"type": "string"}, {"type": "null"}], + }, + ), + {"type": "string", "description": "x"}, + ) + self.assertEqual( + _sanitize_schema_for_gemini( + {"type": "object", "additionalProperties": False}, + ), + {"type": "object"}, + ) + + def test_sanitize_pydantic_optional_list_dict(self) -> None: + """End-to-end: Pydantic Optional[list[dict]] schema is cleaned.""" + result = _sanitize_schema_for_gemini( + { + "type": "object", + "additionalProperties": False, + "properties": { + "actions": { + "description": "List of actions", + "anyOf": [ + { + "type": "array", + "items": { + "type": "object", + "additionalProperties": True, + }, + }, + {"type": "null"}, + ], + }, + }, + }, + ) + self.assertEqual( + result, + { + "type": "object", + "properties": { + "actions": { + "type": "array", + "description": "List of actions", + "items": {"type": "object"}, + }, + }, + }, + ) + + def test_flatten_resolves_ref_and_removes_defs(self) -> None: + """$ref inlined with extra keys merged; $defs removed.""" + self.assertEqual( + _flatten_json_schema( + { + "$defs": {"Name": {"type": "string"}}, + "properties": { + "name": { + "$ref": "#/$defs/Name", + "description": "The name", + }, + }, + }, + ), + { + "properties": { + "name": {"type": "string", "description": "The name"}, + }, + }, + ) + + def test_flatten_circular_ref_returns_placeholder(self) -> None: + """Circular $ref produces a placeholder without infinite recursion.""" + self.assertEqual( + _flatten_json_schema( + { + "$defs": { + "Node": { + "type": "object", + "properties": {"child": {"$ref": "#/$defs/Node"}}, + }, + }, + "properties": {"root": {"$ref": "#/$defs/Node"}}, + }, + ), + { + "properties": { + "root": { + "type": "object", + "properties": { + "child": { + "type": "object", + "description": "(circular: Node)", + }, + }, + }, + }, + }, + ) + + def test_flatten_legacy_definitions_ref(self) -> None: + """Legacy 'definitions' keyword is resolved like '$defs'.""" + self.assertEqual( + _flatten_json_schema( + { + "definitions": {"Address": {"type": "string"}}, + "properties": { + "addr": {"$ref": "#/definitions/Address"}, + }, + }, + ), + { + "properties": { + "addr": {"type": "string"}, + }, + }, + ) + + def test_flatten_no_defs_returns_same_object(self) -> None: + """Schema with no $defs/definitions is returned unchanged ( + identity).""" + schema = {"type": "object", "properties": {"x": {"type": "integer"}}} + self.assertIs(_flatten_json_schema(schema), schema) diff --git a/tests/model_moonshot_test.py b/tests/model_moonshot_test.py new file mode 100644 index 0000000000000000000000000000000000000000..b1e49aeb364c7504190253b3d9e10e68e754de1f --- /dev/null +++ b/tests/model_moonshot_test.py @@ -0,0 +1,491 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for MoonshotChatModel with mocked API responses. + +Tests cover both non-streaming and streaming modes. +""" +from typing import Any +import unittest +from unittest import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, MagicMock, patch + +from utils import AnyString + +from agentscope.message import TextBlock, ToolCallBlock, ThinkingBlock +from agentscope.model import MoonshotChatModel +from agentscope.credential import MoonshotCredential +from agentscope.tool import ToolChoice + +A = AnyString() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_model(stream: bool = False) -> Any: + return MoonshotChatModel( + credential=MoonshotCredential(api_key="test"), + model="kimi-k2-5", + stream=stream, + context_size=131_072, + ) + + +def _mock_completion( + text: Any = None, + tool_calls: Any = None, + reasoning: Any = None, + response_id: str = "kimi-1", +) -> MagicMock: + """Build a mock non-streaming ChatCompletion response.""" + msg = MagicMock() + msg.content = text + msg.reasoning_content = reasoning + msg.tool_calls = None + + if tool_calls: + tc_mocks = [] + for tc in tool_calls: + m = MagicMock() + m.id = tc["id"] + m.function.name = tc["name"] + m.function.arguments = tc["arguments"] + tc_mocks.append(m) + msg.tool_calls = tc_mocks + + choice = MagicMock() + choice.message = msg + + resp = MagicMock() + resp.id = response_id + resp.choices = [choice] + resp.usage.prompt_tokens = 10 + resp.usage.completion_tokens = 5 + resp.usage.cached_tokens = 0 + return resp + + +def _make_stream_chunk( + delta_text: str | None = None, + delta_reasoning: str | None = None, + tool_calls: list | None = None, + response_id: str = "kimi-1", + usage: dict | None = None, + has_choices: bool = True, +) -> MagicMock: + """Build a single mock streaming chunk.""" + chunk = MagicMock() + chunk.id = response_id + + if usage: + chunk.usage = MagicMock() + chunk.usage.prompt_tokens = usage.get("prompt_tokens", 0) + chunk.usage.completion_tokens = usage.get("completion_tokens", 0) + chunk.usage.cached_tokens = usage.get("cached_tokens", 0) + else: + chunk.usage = None + + if has_choices: + delta = MagicMock() + delta.content = delta_text + delta.reasoning_content = delta_reasoning + delta.tool_calls = tool_calls + choice = MagicMock() + choice.delta = delta + chunk.choices = [choice] + else: + chunk.choices = [] + + return chunk + + +def _make_tool_call_delta( + index: int, + tc_id: str | None = None, + name: str | None = None, + arguments: str | None = None, +) -> MagicMock: + tc = MagicMock() + tc.index = index + tc.id = tc_id + tc.function.name = name + tc.function.arguments = arguments + return tc + + +class _MockAsyncStream: + """Mock async stream (context manager + async iterator).""" + + def __init__(self, chunks: list) -> None: + self._chunks = chunks + self._index = 0 + + async def __aenter__(self) -> "_MockAsyncStream": + return self + + async def __aexit__(self, *args: Any) -> None: + pass + + def __aiter__(self) -> "_MockAsyncStream": + return self + + async def __anext__(self) -> Any: + if self._index >= len(self._chunks): + raise StopAsyncIteration + chunk = self._chunks[self._index] + self._index += 1 + return chunk + + +# --------------------------------------------------------------------------- +# Non-streaming tests +# --------------------------------------------------------------------------- + + +class TestMoonshotNonStream(IsolatedAsyncioTestCase): + """Tests for MoonshotChatModel in non-streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=False) + + @patch("openai.AsyncClient") + async def test_text_response(self, mock_client_cls: MagicMock) -> None: + """Non-stream text response returns a single ChatResponse.""" + mock_create = AsyncMock( + return_value=_mock_completion(text="Hello!"), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + (True, [TextBlock.model_construct(id=A, text="Hello!")]), + ) + self.assertEqual(result.id, "kimi-1") + + @patch("openai.AsyncClient") + async def test_tool_call_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream tool call response creates ToolCallBlocks.""" + mock_create = AsyncMock( + return_value=_mock_completion( + tool_calls=[ + { + "id": "call-1", + "name": "get_weather", + "arguments": '{"city":"Shanghai"}', + }, + ], + ), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ToolCallBlock( + id="call-1", + name="get_weather", + input='{"city":"Shanghai"}', + ), + ], + ), + ) + + @patch("openai.AsyncClient") + async def test_thinking_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream response with reasoning creates ThinkingBlock.""" + mock_create = AsyncMock( + return_value=_mock_completion( + text="42", + reasoning="Step by step...", + ), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Step by step...", + ), + TextBlock.model_construct(id=A, text="42"), + ], + ), + ) + + +class TestKimiModelParameters(unittest.TestCase): + """Tests for MoonshotChatModel.Parameters.""" + + def test_thinking_enable_stored_on_model(self) -> None: + """thinking_enable is accessible through model.parameters.""" + model = MoonshotChatModel( + credential=MoonshotCredential(api_key="test"), + model="kimi-k2-5", + stream=False, + context_size=131_072, + parameters=MoonshotChatModel.Parameters(thinking_enable=True), + ) + self.assertTrue(model.parameters.thinking_enable) + + +# --------------------------------------------------------------------------- +# Streaming tests +# --------------------------------------------------------------------------- + + +class TestMoonshotStream(IsolatedAsyncioTestCase): + """Tests for MoonshotChatModel in streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=True) + + @patch("openai.AsyncClient") + async def test_stream_text(self, mock_client_cls: MagicMock) -> None: + """Stream text yields n deltas (is_last=False) + 1 final + (is_last=True) with full content.""" + chunks = [ + _make_stream_chunk(delta_text="Hi"), + _make_stream_chunk(delta_text=" there"), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 10, "completion_tokens": 2}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + (False, [TextBlock.model_construct(id=A, text="Hi")]), + (False, [TextBlock.model_construct(id=A, text=" there")]), + (True, [TextBlock.model_construct(id=A, text="Hi there")]), + ], + ) + + @patch("openai.AsyncClient") + async def test_stream_thinking_then_text( + self, + mock_client_cls: MagicMock, + ) -> None: + """Moonshot yields thinking chunks separately before text.""" + chunks = [ + _make_stream_chunk(delta_reasoning="Let me"), + _make_stream_chunk(delta_reasoning=" think"), + _make_stream_chunk(delta_text="Result"), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 10, "completion_tokens": 5}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ThinkingBlock.model_construct(id=A, thinking="Let me")], + ), + ( + False, + [ThinkingBlock.model_construct(id=A, thinking=" think")], + ), + (False, [TextBlock.model_construct(id=A, text="Result")]), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Let me think", + ), + TextBlock.model_construct(id=A, text="Result"), + ], + ), + ], + ) + + @patch("openai.AsyncClient") + async def test_stream_tool_calls( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream tool calls accumulate across chunks into final response.""" + chunks = [ + _make_stream_chunk( + tool_calls=[ + _make_tool_call_delta(0, "call-1", "search", '{"q":'), + ], + ), + _make_stream_chunk( + tool_calls=[ + _make_tool_call_delta(0, None, None, '"test"}'), + ], + ), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 10, "completion_tokens": 5}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ + ToolCallBlock( + id="call-1", + name="search", + input='{"q":', + ), + ], + ), + ( + False, + [ + ToolCallBlock( + id="call-1", + name="search", + input='"test"}', + ), + ], + ), + ( + True, + [ + ToolCallBlock( + id="call-1", + name="search", + input='{"q":"test"}', + ), + ], + ), + ], + ) + + +# --------------------------------------------------------------------------- +# _format_tools tests +# --------------------------------------------------------------------------- + +_FT_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "description": "Get the time", + "parameters": { + "type": "object", + "properties": {"timezone": {"type": "string"}}, + "required": ["timezone"], + }, + }, + }, +] + + +class TestMoonshotFormatTools(unittest.TestCase): + """Tests for MoonshotChatModel._format_tools.""" + + def setUp(self) -> None: + """Set up model instance.""" + self.model = _make_model() + + def test_auto_mode(self) -> None: + """Auto mode returns tools unchanged and string 'auto'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual(fmt_choice, "auto") + + def test_none_mode(self) -> None: + """None mode returns tools unchanged and string 'none'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="none"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual(fmt_choice, "none") + + def test_required_mode(self) -> None: + """Required mode returns tools unchanged and string 'required'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="required"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual(fmt_choice, "required") + + def test_str_mode_force_call(self) -> None: + """A specific tool name returns a type=function dict.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="get_weather"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual( + fmt_choice, + {"type": "function", "function": {"name": "get_weather"}}, + ) + + def test_tools_filtered(self) -> None: + """When tool_choice.tools is set, only those tools are included.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto", tools=["get_weather"]), + ) + self.assertEqual(len(fmt_tools), 1) + self.assertEqual(fmt_tools[0]["function"]["name"], "get_weather") + self.assertEqual(fmt_choice, "auto") + + def test_no_tool_choice(self) -> None: + """Without tool_choice, returns tools and None.""" + fmt_tools, fmt_choice = self.model._format_tools(_FT_TOOLS, None) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertIsNone(fmt_choice) diff --git a/tests/model_ollama_test.py b/tests/model_ollama_test.py new file mode 100644 index 0000000000000000000000000000000000000000..660f58c2bface98429259d68d533e421f4552ff9 --- /dev/null +++ b/tests/model_ollama_test.py @@ -0,0 +1,365 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for OllamaChatModel with mocked API responses. + +Tests cover both non-streaming and streaming modes. +Ollama uses ollama.AsyncClient with async iterator streaming. +""" +import json +from typing import Any +import unittest +from unittest import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, MagicMock, patch + +from utils import AnyString + +from agentscope.message import TextBlock, ToolCallBlock, ThinkingBlock +from agentscope.model import OllamaChatModel +from agentscope.tool import ToolChoice + +A = AnyString() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_model(stream: bool = False) -> Any: + return OllamaChatModel( + model="qwen3:8b", + stream=stream, + context_size=40_960, + ) + + +def _mock_completion( + content: str = "", + thinking: str | None = None, + tool_calls: list | None = None, +) -> MagicMock: + """Build a mock non-streaming Ollama response.""" + msg = MagicMock() + msg.content = content + msg.thinking = thinking + + if tool_calls: + tc_mocks = [] + for tc in tool_calls: + m = MagicMock() + m.function.name = tc["name"] + m.function.arguments = tc["args"] + tc_mocks.append(m) + msg.tool_calls = tc_mocks + else: + msg.tool_calls = None + + resp = MagicMock() + resp.message = msg + resp.prompt_eval_count = 10 + resp.eval_count = 5 + resp.id = None + return resp + + +def _make_stream_chunk( + content: str = "", + thinking: str | None = None, + tool_calls: list | None = None, +) -> MagicMock: + """Build a single mock Ollama streaming chunk.""" + msg = MagicMock() + msg.content = content + msg.thinking = thinking + + if tool_calls: + tc_mocks = [] + for tc in tool_calls: + m = MagicMock() + m.function.name = tc["name"] + m.function.arguments = tc["args"] + tc_mocks.append(m) + msg.tool_calls = tc_mocks + else: + msg.tool_calls = None + + chunk = MagicMock() + chunk.message = msg + chunk.prompt_eval_count = 10 + chunk.eval_count = 5 + chunk.id = None + return chunk + + +class _MockAsyncStream: + """Mock async iterator for Ollama stream.""" + + def __init__(self, chunks: list) -> None: + self._chunks = chunks + self._index = 0 + + def __aiter__(self) -> "_MockAsyncStream": + return self + + async def __anext__(self) -> Any: + if self._index >= len(self._chunks): + raise StopAsyncIteration + chunk = self._chunks[self._index] + self._index += 1 + return chunk + + +# --------------------------------------------------------------------------- +# Non-streaming tests +# --------------------------------------------------------------------------- + + +class TestOllamaNonStream(IsolatedAsyncioTestCase): + """Tests for OllamaChatModel in non-streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=False) + + @patch("ollama.AsyncClient") + async def test_text_response(self, mock_client_cls: MagicMock) -> None: + """Non-stream text response returns a single ChatResponse.""" + mock_client_cls.return_value.chat = AsyncMock( + return_value=_mock_completion(content="Hello!"), + ) + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + (True, [TextBlock.model_construct(id=A, text="Hello!")]), + ) + + @patch("ollama.AsyncClient") + async def test_tool_call_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Parsing a tool-call response creates a ToolCallBlock.""" + mock_client_cls.return_value.chat = AsyncMock( + return_value=_mock_completion( + tool_calls=[ + {"name": "get_weather", "args": {"city": "SH"}}, + ], + ), + ) + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ToolCallBlock( + id="0_get_weather", + name="get_weather", + input=json.dumps({"city": "SH"}), + ), + ], + ), + ) + + @patch("ollama.AsyncClient") + async def test_thinking_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream thinking plus text returns ThinkingBlock then + TextBlock.""" + mock_client_cls.return_value.chat = AsyncMock( + return_value=_mock_completion( + content="Answer", + thinking="Let me think...", + ), + ) + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Let me think...", + ), + TextBlock.model_construct(id=A, text="Answer"), + ], + ), + ) + + +# --------------------------------------------------------------------------- +# Streaming tests +# --------------------------------------------------------------------------- + + +class TestOllamaStream(IsolatedAsyncioTestCase): + """Tests for OllamaChatModel in streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=True) + + @patch("ollama.AsyncClient") + async def test_stream_text(self, mock_client_cls: MagicMock) -> None: + """Stream text yields deltas then final with full content.""" + chunks = [ + _make_stream_chunk(content="Hi"), + _make_stream_chunk(content=" there"), + ] + mock_client_cls.return_value.chat = AsyncMock( + return_value=_MockAsyncStream(chunks), + ) + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + (False, [TextBlock.model_construct(id=A, text="Hi")]), + (False, [TextBlock.model_construct(id=A, text=" there")]), + (True, [TextBlock.model_construct(id=A, text="Hi there")]), + ], + ) + + @patch("ollama.AsyncClient") + async def test_stream_thinking_and_text( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream thinking and text deltas then final with accumulated + content.""" + chunks = [ + _make_stream_chunk(thinking="Think step"), + _make_stream_chunk(content="Result"), + ] + mock_client_cls.return_value.chat = AsyncMock( + return_value=_MockAsyncStream(chunks), + ) + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Think step", + ), + ], + ), + (False, [TextBlock.model_construct(id=A, text="Result")]), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Think step", + ), + TextBlock.model_construct(id=A, text="Result"), + ], + ), + ], + ) + + @patch("ollama.AsyncClient") + async def test_stream_tool_call( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream tool-call chunk yields delta then final with same + ToolCallBlock.""" + chunks = [ + _make_stream_chunk( + tool_calls=[ + {"name": "search", "args": {"q": "hello"}}, + ], + ), + ] + mock_client_cls.return_value.chat = AsyncMock( + return_value=_MockAsyncStream(chunks), + ) + + gen = await self.model([]) + responses = [r async for r in gen] + + tool_block = ToolCallBlock( + id="0_search", + name="search", + input=json.dumps({"q": "hello"}), + ) + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + (False, [tool_block]), + (True, [tool_block]), + ], + ) + + +# --------------------------------------------------------------------------- +# _format_tools tests +# --------------------------------------------------------------------------- + +_FT_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "description": "Get the time", + "parameters": { + "type": "object", + "properties": {"timezone": {"type": "string"}}, + "required": ["timezone"], + }, + }, + }, +] + + +class TestOllamaFormatTools(unittest.TestCase): + """Tests for OllamaChatModel._format_tools.""" + + def setUp(self) -> None: + self.model = _make_model() + + def test_tools_forwarded_no_choice(self) -> None: + """Tools are forwarded unchanged when tool_choice is None.""" + fmt_tools, fmt_choice = self.model._format_tools(_FT_TOOLS, None) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertIsNone(fmt_choice) + + def test_tools_filtered(self) -> None: + """ToolChoice with tools list filters to matching function names.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto", tools=["get_weather"]), + ) + self.assertIsNotNone(fmt_tools) + assert fmt_tools is not None + self.assertEqual(len(fmt_tools), 1) + self.assertEqual(fmt_tools[0]["function"]["name"], "get_weather") + self.assertIsNone(fmt_choice) diff --git a/tests/model_openai_chat_test.py b/tests/model_openai_chat_test.py new file mode 100644 index 0000000000000000000000000000000000000000..6701885eebcb560de3d9957084a43d591e0e959c --- /dev/null +++ b/tests/model_openai_chat_test.py @@ -0,0 +1,769 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for OpenAIChatModel with mocked API responses. + +Tests cover both non-streaming and streaming modes, verifying that: +- Non-stream mode returns a single ChatResponse with is_last=True. +- Stream mode yields n delta ChatResponses (is_last=False) followed by + 1 final ChatResponse (is_last=True) with the full accumulated content. +""" +import base64 +import io +import wave +from typing import Any +import unittest +from unittest import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, MagicMock, patch + +from utils import AnyString + +from agentscope.message import ( + TextBlock, + ToolCallBlock, + ThinkingBlock, + DataBlock, + Base64Source, +) +from agentscope.model import OpenAIChatModel +from agentscope.credential import OpenAICredential +from agentscope.tool import ToolChoice + +A = AnyString() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_model(stream: bool = False) -> Any: + return OpenAIChatModel( + credential=OpenAICredential(api_key="test"), + model="gpt-4o", + stream=stream, + context_size=128_000, + ) + + +def _mock_completion( + text: Any = None, + tool_calls: Any = None, + reasoning: Any = None, + response_id: str = "resp-1", + audio: dict | None = None, +) -> MagicMock: + """Build a mock non-streaming ChatCompletion response.""" + msg = MagicMock() + msg.content = text + msg.reasoning_content = reasoning + msg.reasoning = None + msg.audio = audio + msg.tool_calls = None + + if tool_calls: + tc_mocks = [] + for tc in tool_calls: + m = MagicMock() + m.id = tc["id"] + m.function.name = tc["name"] + m.function.arguments = tc["arguments"] + tc_mocks.append(m) + msg.tool_calls = tc_mocks + + choice = MagicMock() + choice.message = msg + + resp = MagicMock() + resp.id = response_id + resp.choices = [choice] + resp.usage.prompt_tokens = 10 + resp.usage.completion_tokens = 5 + resp.usage.prompt_tokens_details = None + return resp + + +def _make_stream_chunk( + delta_text: str | None = None, + delta_reasoning: str | None = None, + tool_calls: list | None = None, + response_id: str = "resp-1", + usage: dict | None = None, + has_choices: bool = True, + delta_audio: dict | None = None, +) -> MagicMock: + """Build a single mock streaming chunk.""" + chunk = MagicMock() + chunk.id = response_id + + if usage: + chunk.usage = MagicMock() + chunk.usage.prompt_tokens = usage.get("prompt_tokens", 0) + chunk.usage.completion_tokens = usage.get("completion_tokens", 0) + chunk.usage.prompt_tokens_details = None + else: + chunk.usage = None + + if has_choices: + delta = MagicMock() + delta.content = delta_text + delta.reasoning_content = delta_reasoning + delta.reasoning = None + delta.audio = delta_audio + delta.tool_calls = tool_calls + choice = MagicMock() + choice.delta = delta + chunk.choices = [choice] + else: + chunk.choices = [] + + return chunk + + +def _make_tool_call_delta( + index: int, + tc_id: str | None = None, + name: str | None = None, + arguments: str | None = None, +) -> MagicMock: + """Build a tool_call delta item for streaming.""" + tc = MagicMock() + tc.index = index + tc.id = tc_id + tc.function.name = name + tc.function.arguments = arguments + return tc + + +class _MockAsyncStream: + """Mock async stream that acts as an async context manager + iterator.""" + + def __init__(self, chunks: list) -> None: + self._chunks = chunks + self._index = 0 + + async def __aenter__(self) -> "_MockAsyncStream": + return self + + async def __aexit__(self, *args: Any) -> None: + pass + + def __aiter__(self) -> "_MockAsyncStream": + return self + + async def __anext__(self) -> Any: + if self._index >= len(self._chunks): + raise StopAsyncIteration + chunk = self._chunks[self._index] + self._index += 1 + return chunk + + +# --------------------------------------------------------------------------- +# Non-streaming tests +# --------------------------------------------------------------------------- + + +class TestOpenAIChatNonStream(IsolatedAsyncioTestCase): + """Tests for OpenAIChatModel in non-streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=False) + + @patch("openai.AsyncClient") + async def test_text_response(self, mock_client_cls: MagicMock) -> None: + """Non-stream text response returns a single ChatResponse.""" + mock_create = AsyncMock( + return_value=_mock_completion(text="Hello world!"), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + (True, [TextBlock.model_construct(id=A, text="Hello world!")]), + ) + self.assertEqual(result.id, "resp-1") + + @patch("openai.AsyncClient") + async def test_default_thinking_enable_not_forwarded( + self, + mock_client_cls: MagicMock, + ) -> None: + """Default parameters do not add provider-specific extra_body.""" + mock_create = AsyncMock( + return_value=_mock_completion(text="Hello world!"), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + await self.model([]) + + self.assertNotIn("extra_body", mock_create.call_args.kwargs) + + @patch("openai.AsyncClient") + async def test_constructor_extra_body_forwarded( + self, + mock_client_cls: MagicMock, + ) -> None: + """Custom request fields are forwarded to OpenAI-compatible APIs.""" + model = OpenAIChatModel( + credential=OpenAICredential(api_key="test"), + model="custom-model", + stream=False, + context_size=128_000, + extra_body={"enable_thinking": False}, + ) + mock_create = AsyncMock( + return_value=_mock_completion(text="Hello world!"), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + await model([]) + + self.assertEqual( + mock_create.call_args.kwargs["extra_body"], + {"enable_thinking": False}, + ) + + @patch("openai.AsyncClient") + async def test_generate_kwargs_extra_body_overrides_constructor( + self, + mock_client_cls: MagicMock, + ) -> None: + """Per-call extra_body overrides the constructor default.""" + model = OpenAIChatModel( + credential=OpenAICredential(api_key="test"), + model="custom-model", + stream=False, + context_size=128_000, + extra_body={"enable_thinking": False}, + ) + mock_create = AsyncMock( + return_value=_mock_completion(text="Hello world!"), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + await model([], extra_body={"custom_option": "value"}) + + self.assertEqual( + mock_create.call_args.kwargs["extra_body"], + {"custom_option": "value"}, + ) + + @patch("openai.AsyncClient") + async def test_tool_call_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream tool call response creates ToolCallBlocks.""" + mock_create = AsyncMock( + return_value=_mock_completion( + tool_calls=[ + { + "id": "call-1", + "name": "get_weather", + "arguments": '{"city":"Beijing"}', + }, + ], + ), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ToolCallBlock( + id="call-1", + name="get_weather", + input='{"city":"Beijing"}', + ), + ], + ), + ) + + @patch("openai.AsyncClient") + async def test_audio_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream audio-only output yields transcript TextBlock + audio + DataBlock.""" + mock_create = AsyncMock( + return_value=_mock_completion( + text=None, + audio={ + "data": "QUJDREVG", + "transcript": "Hello from audio.", + }, + ), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + TextBlock.model_construct( + id=A, + text="Hello from audio.", + ), + DataBlock.model_construct( + id=A, + source=Base64Source.model_construct( + type="base64", + media_type="audio/wav", + data="QUJDREVG", + ), + ), + ], + ), + ) + + @patch("openai.AsyncClient") + async def test_thinking_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream response with reasoning creates ThinkingBlock.""" + mock_create = AsyncMock( + return_value=_mock_completion( + text="The answer is 42.", + reasoning="Let me think step by step...", + ), + ) + mock_client_cls.return_value.chat.completions.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Let me think step by step...", + ), + TextBlock.model_construct(id=A, text="The answer is 42."), + ], + ), + ) + + +# --------------------------------------------------------------------------- +# Streaming tests +# --------------------------------------------------------------------------- + + +class TestOpenAIChatStream(IsolatedAsyncioTestCase): + """Tests for OpenAIChatModel in streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=True) + + @patch("openai.AsyncClient") + async def test_stream_text_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream text yields n deltas (is_last=False) + 1 final + (is_last=True) with full content.""" + chunks = [ + _make_stream_chunk(delta_text="Hello"), + _make_stream_chunk(delta_text=" world"), + _make_stream_chunk(delta_text="!"), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 10, "completion_tokens": 3}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + (False, [TextBlock.model_construct(id=A, text="Hello")]), + (False, [TextBlock.model_construct(id=A, text=" world")]), + (False, [TextBlock.model_construct(id=A, text="!")]), + (True, [TextBlock.model_construct(id=A, text="Hello world!")]), + ], + ) + self.assertEqual(responses[-1].id, "resp-1") + + @patch("openai.AsyncClient") + async def test_stream_thinking_and_text( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream with thinking + text yields deltas then final with both.""" + chunks = [ + _make_stream_chunk(delta_reasoning="Think"), + _make_stream_chunk(delta_reasoning="ing..."), + _make_stream_chunk(delta_text="Answer"), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 10, "completion_tokens": 8}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ThinkingBlock.model_construct(id=A, thinking="Think")], + ), + ( + False, + [ThinkingBlock.model_construct(id=A, thinking="ing...")], + ), + (False, [TextBlock.model_construct(id=A, text="Answer")]), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Thinking...", + ), + TextBlock.model_construct(id=A, text="Answer"), + ], + ), + ], + ) + + @patch("openai.AsyncClient") + async def test_stream_tool_calls( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream tool calls accumulate across chunks into final response.""" + chunks = [ + _make_stream_chunk( + tool_calls=[ + _make_tool_call_delta(0, "call-1", "get_weather", '{"ci'), + ], + ), + _make_stream_chunk( + tool_calls=[ + _make_tool_call_delta(0, None, None, 'ty":"BJ"}'), + ], + ), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 10, "completion_tokens": 5}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ + ToolCallBlock( + id="call-1", + name="get_weather", + input='{"ci', + ), + ], + ), + ( + False, + [ + ToolCallBlock( + id="call-1", + name="get_weather", + input='ty":"BJ"}', + ), + ], + ), + ( + True, + [ + ToolCallBlock( + id="call-1", + name="get_weather", + input='{"city":"BJ"}', + ), + ], + ), + ], + ) + + @patch("openai.AsyncClient") + async def test_stream_usage_in_final( + self, + mock_client_cls: MagicMock, + ) -> None: + """Usage information is captured and present in final response.""" + chunks = [ + _make_stream_chunk(delta_text="Hi"), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 100, "completion_tokens": 20}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + (False, [TextBlock.model_construct(id=A, text="Hi")]), + (True, [TextBlock.model_construct(id=A, text="Hi")]), + ], + ) + self.assertEqual(responses[-1].usage.input_tokens, 100) + self.assertEqual(responses[-1].usage.output_tokens, 20) + + @patch("openai.AsyncClient") + async def test_stream_audio_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream PCM deltas produce per-chunk DataBlocks (first chunk + prefixed with a streaming WAV header) sharing a stable id, plus a + final fixed-size WAV block readable by the ``wave`` module. + Transcript chunks ride alongside as TextBlock deltas so the agent + can stream caption text live; the final block carries the full + accumulated transcript.""" + pcm1 = bytes([1, 2, 3, 4]) + pcm2 = bytes([5, 6, 7, 8]) + pcm3 = bytes([9, 10, 11, 12]) + pcm_full = pcm1 + pcm2 + pcm3 + + chunks = [ + _make_stream_chunk( + delta_audio={ + "data": base64.b64encode(pcm1).decode(), + "transcript": "Hello", + }, + ), + _make_stream_chunk( + delta_audio={ + "data": base64.b64encode(pcm2).decode(), + "transcript": " world", + }, + ), + _make_stream_chunk( + delta_audio={ + "data": base64.b64encode(pcm3).decode(), + "transcript": "!", + }, + ), + _make_stream_chunk( + has_choices=False, + usage={"prompt_tokens": 10, "completion_tokens": 6}, + ), + ] + mock_create = AsyncMock(return_value=_MockAsyncStream(chunks)) + mock_client_cls.return_value.chat.completions.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + self.assertEqual(len(responses), 4) + + # All four chunks (3 deltas + 1 final) must share the same audio + # block id so downstream consumers stitch them as one stream. + all_audio_ids = { + block.id + for r in responses + for block in r.content + if isinstance(block, DataBlock) + } + self.assertEqual(len(all_audio_ids), 1) + + # First delta: WAV header (44 bytes, "RIFF"..."WAVE") + pcm1. + first_audio = next( + b for b in responses[0].content if isinstance(b, DataBlock) + ) + first_payload = base64.b64decode(first_audio.source.data) + self.assertEqual(len(first_payload), 44 + len(pcm1)) + self.assertEqual(first_payload[:4], b"RIFF") + self.assertEqual(first_payload[8:12], b"WAVE") + self.assertEqual(first_payload[44:], pcm1) + self.assertEqual(first_audio.source.media_type, "audio/wav") + + # Subsequent deltas: raw PCM only, no header. + for resp, pcm in zip(responses[1:3], [pcm2, pcm3]): + audio_block = next( + b for b in resp.content if isinstance(b, DataBlock) + ) + self.assertEqual(base64.b64decode(audio_block.source.data), pcm) + self.assertEqual(audio_block.source.media_type, "audio/wav") + + # Transcript rides alongside: each delta carries a TextBlock with + # only that chunk's text (so the agent emits TextBlockDeltaEvents + # in real time). + for resp, expected_text in zip( + responses[:3], + ["Hello", " world", "!"], + ): + text_block = next( + b for b in resp.content if isinstance(b, TextBlock) + ) + self.assertEqual(text_block.text, expected_text) + + # Final ``is_last`` block: a fixed-size WAV the ``wave`` module + # can parse end-to-end at 24kHz / mono / 16-bit. + final = responses[-1] + self.assertTrue(final.is_last) + final_audio = next( + b for b in final.content if isinstance(b, DataBlock) + ) + wav_bytes = base64.b64decode(final_audio.source.data) + with wave.open(io.BytesIO(wav_bytes), "rb") as wav: + self.assertEqual(wav.getnchannels(), 1) + self.assertEqual(wav.getsampwidth(), 2) + self.assertEqual(wav.getframerate(), 24000) + frames = wav.readframes(wav.getnframes()) + self.assertEqual(frames, pcm_full) + + # Transcript is accumulated and emitted as a TextBlock alongside. + final_text = next(b for b in final.content if isinstance(b, TextBlock)) + self.assertEqual(final_text.text, "Hello world!") + + +class TestOpenAIChatModelParameters(unittest.TestCase): + """Tests for OpenAIChatModel.Parameters.""" + + def test_reasoning_effort_stored_on_model(self) -> None: + """reasoning_effort is accessible through model.parameters.""" + model = OpenAIChatModel( + credential=OpenAICredential(api_key="test"), + model="o3", + stream=False, + context_size=200_000, + parameters=OpenAIChatModel.Parameters(reasoning_effort="low"), + ) + self.assertEqual(model.parameters.reasoning_effort, "low") + + def test_thinking_enable_stored_on_model(self) -> None: + """thinking_enable is accessible through model.parameters.""" + model = OpenAIChatModel( + credential=OpenAICredential(api_key="test"), + model="o3", + stream=False, + context_size=200_000, + parameters=OpenAIChatModel.Parameters(thinking_enable=True), + ) + self.assertTrue(model.parameters.thinking_enable) + + +# --------------------------------------------------------------------------- +# Shared _format_tools fixtures +# --------------------------------------------------------------------------- + +_FT_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "description": "Get the time", + "parameters": { + "type": "object", + "properties": {"timezone": {"type": "string"}}, + "required": ["timezone"], + }, + }, + }, +] + + +class TestOpenAIChatFormatTools(unittest.TestCase): + """Tests for OpenAIChatModel._format_tools.""" + + def setUp(self) -> None: + """Set up model instance.""" + self.model = _make_model() + + def test_auto_mode(self) -> None: + """Auto mode returns tools unchanged and string 'auto'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual(fmt_choice, "auto") + + def test_none_mode(self) -> None: + """None mode returns tools unchanged and string 'none'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="none"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual(fmt_choice, "none") + + def test_required_mode(self) -> None: + """Required mode returns tools unchanged and string 'required'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="required"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual(fmt_choice, "required") + + def test_str_mode_force_call(self) -> None: + """A specific tool name returns a type=function dict.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="get_weather"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertEqual( + fmt_choice, + {"type": "function", "function": {"name": "get_weather"}}, + ) + + def test_tools_filtered(self) -> None: + """When tool_choice.tools is set, only those tools are included.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto", tools=["get_weather"]), + ) + self.assertEqual(len(fmt_tools), 1) + self.assertEqual(fmt_tools[0]["function"]["name"], "get_weather") + self.assertEqual(fmt_choice, "auto") + + def test_no_tool_choice(self) -> None: + """Without tool_choice, returns tools and None.""" + fmt_tools, fmt_choice = self.model._format_tools(_FT_TOOLS, None) + self.assertEqual(fmt_tools, _FT_TOOLS) + self.assertIsNone(fmt_choice) diff --git a/tests/model_openai_response_test.py b/tests/model_openai_response_test.py new file mode 100644 index 0000000000000000000000000000000000000000..4c407f560cd1a084872886340b84fdc601296016 --- /dev/null +++ b/tests/model_openai_response_test.py @@ -0,0 +1,563 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for OpenAIResponseModel with mocked API responses. + +Tests cover both non-streaming and streaming modes. +OpenAI Responses API uses event-based streaming with response.completed. +""" +from typing import Any +import unittest +from unittest import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, MagicMock, patch + +from utils import AnyString + +from agentscope.message import TextBlock, ToolCallBlock, ThinkingBlock +from agentscope.model import OpenAIResponseModel +from agentscope.credential import OpenAICredential +from agentscope.tool import ToolChoice + +A = AnyString() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_model(stream: bool = False) -> Any: + return OpenAIResponseModel( + credential=OpenAICredential(api_key="test"), + model="o4-mini", + stream=stream, + context_size=200_000, + ) + + +def _mock_completion( + text: Any = None, + function_calls: Any = None, + reasoning_summary: Any = None, + reasoning_id: str = "rs_test123", + response_id: str = "resp-openai-1", +) -> MagicMock: + """Build a mock non-streaming Responses API response.""" + output = [] + + if reasoning_summary: + reasoning_item = MagicMock() + reasoning_item.type = "reasoning" + reasoning_item.id = reasoning_id + summary_mock = MagicMock() + summary_mock.text = reasoning_summary + reasoning_item.summary = [summary_mock] + output.append(reasoning_item) + + if text: + msg_item = MagicMock() + msg_item.type = "message" + part = MagicMock() + part.type = "output_text" + part.text = text + msg_item.content = [part] + output.append(msg_item) + + if function_calls: + for fc in function_calls: + fc_item = MagicMock() + fc_item.type = "function_call" + fc_item.id = fc["id"] + fc_item.call_id = fc["call_id"] + fc_item.name = fc["name"] + fc_item.arguments = fc["arguments"] + output.append(fc_item) + + resp = MagicMock() + resp.id = response_id + resp.output = output + resp.usage = MagicMock() + resp.usage.input_tokens = 10 + resp.usage.output_tokens = 5 + resp.usage.input_tokens_details = None + return resp + + +def _make_event(event_type: str, **kwargs: Any) -> MagicMock: + """Build a mock Responses API streaming event.""" + event = MagicMock() + event.type = event_type + for key, val in kwargs.items(): + setattr(event, key, val) + # Default: no response attribute + if "response" not in kwargs: + event.response = None + return event + + +class _MockAsyncEventStream: + """Mock async iterator over Response events.""" + + def __init__(self, events: list) -> None: + self._events = events + self._index = 0 + + def __aiter__(self) -> "_MockAsyncEventStream": + return self + + async def __anext__(self) -> Any: + if self._index >= len(self._events): + raise StopAsyncIteration + event = self._events[self._index] + self._index += 1 + return event + + +# --------------------------------------------------------------------------- +# Non-streaming tests +# --------------------------------------------------------------------------- + + +class TestOpenAIResponseNonStream(IsolatedAsyncioTestCase): + """Tests for OpenAIResponseModel in non-streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=False) + + @patch("openai.AsyncClient") + async def test_text_response(self, mock_client_cls: MagicMock) -> None: + """Non-stream text response returns a single ChatResponse.""" + mock_create = AsyncMock( + return_value=_mock_completion(text="Hello!"), + ) + mock_client_cls.return_value.responses.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + (True, [TextBlock.model_construct(id=A, text="Hello!")]), + ) + self.assertEqual(result.id, "resp-openai-1") + + @patch("openai.AsyncClient") + async def test_tool_call_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Parsing a tool-call response creates a ToolCallBlock with + call_id.""" + mock_create = AsyncMock( + return_value=_mock_completion( + function_calls=[ + { + "id": "fc_abc", + "call_id": "call-1", + "name": "get_weather", + "arguments": '{"city":"BJ"}', + }, + ], + ), + ) + mock_client_cls.return_value.responses.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ToolCallBlock( + id="fc_abc", + call_id="call-1", + name="get_weather", + input='{"city":"BJ"}', + ), + ], + ), + ) + + @patch("openai.AsyncClient") + async def test_reasoning_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream reasoning summary plus text returns both block types.""" + mock_create = AsyncMock( + return_value=_mock_completion( + reasoning_summary="Thinking step...", + text="Answer", + reasoning_id="rs_abc999", + ), + ) + mock_client_cls.return_value.responses.create = mock_create + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Thinking step...", + reasoning_item_id="rs_abc999", + ), + TextBlock.model_construct(id=A, text="Answer"), + ], + ), + ) + + +class TestOpenAIResponseModelParameters(unittest.TestCase): + """Tests for OpenAIResponseModel.Parameters.""" + + def test_thinking_enable_stored_on_model(self) -> None: + """thinking_enable is accessible through model.parameters.""" + model = OpenAIResponseModel( + credential=OpenAICredential(api_key="test"), + model="o4-mini", + stream=False, + context_size=200_000, + parameters=OpenAIResponseModel.Parameters(thinking_enable=True), + ) + self.assertTrue(model.parameters.thinking_enable) + + def test_reasoning_effort_stored_on_model(self) -> None: + """reasoning_effort is accessible through model.parameters.""" + model = OpenAIResponseModel( + credential=OpenAICredential(api_key="test"), + model="o4-mini", + stream=False, + context_size=200_000, + parameters=OpenAIResponseModel.Parameters( + reasoning_effort="high", + ), + ) + self.assertEqual(model.parameters.reasoning_effort, "high") + + +# --------------------------------------------------------------------------- +# Streaming tests +# --------------------------------------------------------------------------- + + +class TestOpenAIResponseStream(IsolatedAsyncioTestCase): + """Tests for OpenAIResponseModel in streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=True) + + @patch("openai.AsyncClient") + async def test_stream_text(self, mock_client_cls: MagicMock) -> None: + """Stream text yields deltas then final with full content.""" + completed_resp = MagicMock() + completed_resp.id = "resp-1" + completed_resp.output = [] + completed_resp.usage = MagicMock() + completed_resp.usage.input_tokens = 10 + completed_resp.usage.output_tokens = 5 + completed_resp.usage.input_tokens_details = None + + events = [ + _make_event( + "response.output_text.delta", + delta="Hello", + response=MagicMock(id="resp-1"), + ), + _make_event( + "response.output_text.delta", + delta=" world", + ), + _make_event("response.completed", response=completed_resp), + ] + mock_create = AsyncMock( + return_value=_MockAsyncEventStream(events), + ) + mock_client_cls.return_value.responses.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + (False, [TextBlock.model_construct(id=A, text="Hello")]), + (False, [TextBlock.model_construct(id=A, text=" world")]), + (True, [TextBlock.model_construct(id=A, text="Hello world")]), + ], + ) + self.assertEqual(responses[-1].id, "resp-1") + + @patch("openai.AsyncClient") + async def test_stream_reasoning_and_text( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream reasoning and text deltas then final with + reasoning_item_id.""" + reasoning_item = MagicMock() + reasoning_item.type = "reasoning" + reasoning_item.id = "rs_123" + + completed_resp = MagicMock() + completed_resp.id = "resp-2" + completed_resp.output = [reasoning_item] + completed_resp.usage = MagicMock() + completed_resp.usage.input_tokens = 10 + completed_resp.usage.output_tokens = 5 + completed_resp.usage.input_tokens_details = None + + events = [ + _make_event( + "response.reasoning_summary_text.delta", + delta="Thinking", + response=MagicMock(id="resp-2"), + ), + _make_event( + "response.output_text.delta", + delta="Answer", + ), + _make_event("response.completed", response=completed_resp), + ] + mock_create = AsyncMock( + return_value=_MockAsyncEventStream(events), + ) + mock_client_cls.return_value.responses.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ThinkingBlock.model_construct(id=A, thinking="Thinking")], + ), + (False, [TextBlock.model_construct(id=A, text="Answer")]), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Thinking", + reasoning_item_id="rs_123", + ), + TextBlock.model_construct(id=A, text="Answer"), + ], + ), + ], + ) + + @patch("openai.AsyncClient") + async def test_stream_function_call( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream function-call events yield deltas then final + ToolCallBlock.""" + fc_item = MagicMock() + fc_item.type = "function_call" + fc_item.id = "fc_1" + fc_item.call_id = "call-1" + fc_item.name = "search" + + completed_resp = MagicMock() + completed_resp.id = "resp-3" + completed_resp.output = [] + completed_resp.usage = MagicMock() + completed_resp.usage.input_tokens = 10 + completed_resp.usage.output_tokens = 5 + completed_resp.usage.input_tokens_details = None + + events = [ + _make_event( + "response.output_item.added", + item=fc_item, + response=MagicMock(id="resp-3"), + ), + _make_event( + "response.function_call_arguments.delta", + item_id="fc_1", + delta='{"q":', + ), + _make_event( + "response.function_call_arguments.delta", + item_id="fc_1", + delta='"test"}', + ), + _make_event("response.completed", response=completed_resp), + ] + mock_create = AsyncMock( + return_value=_MockAsyncEventStream(events), + ) + mock_client_cls.return_value.responses.create = mock_create + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ + ToolCallBlock( + id="fc_1", + call_id="call-1", + name="search", + input='{"q":', + ), + ], + ), + ( + False, + [ + ToolCallBlock( + id="fc_1", + call_id="call-1", + name="search", + input='"test"}', + ), + ], + ), + ( + True, + [ + ToolCallBlock( + id="fc_1", + call_id="call-1", + name="search", + input='{"q":"test"}', + ), + ], + ), + ], + ) + + +# --------------------------------------------------------------------------- +# _format_tools tests +# --------------------------------------------------------------------------- + +_FT_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "description": "Get the time", + "parameters": { + "type": "object", + "properties": {"timezone": {"type": "string"}}, + "required": ["timezone"], + }, + }, + }, +] + +_FT_TOOLS_RESPONSE = [ + { + "type": "function", + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + { + "type": "function", + "name": "get_time", + "description": "Get the time", + "parameters": { + "type": "object", + "properties": {"timezone": {"type": "string"}}, + "required": ["timezone"], + }, + }, +] + + +class TestOpenAIResponseFormatTools(unittest.TestCase): + """Tests for OpenAIResponseModel._format_tools.""" + + def setUp(self) -> None: + self.model = _make_model() + + def test_auto_mode(self) -> None: + """Auto mode converts tools and sets choice to 'auto'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS_RESPONSE) + self.assertEqual(fmt_choice, "auto") + + def test_none_mode(self) -> None: + """None mode converts tools and sets choice to 'none'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="none"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS_RESPONSE) + self.assertEqual(fmt_choice, "none") + + def test_required_mode(self) -> None: + """Required mode converts tools and sets choice to 'required'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="required"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS_RESPONSE) + self.assertEqual(fmt_choice, "required") + + def test_str_mode_force_call(self) -> None: + """String mode forces a function call for the named tool.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="get_weather"), + ) + self.assertEqual(fmt_tools, _FT_TOOLS_RESPONSE) + self.assertEqual( + fmt_choice, + {"type": "function", "name": "get_weather"}, + ) + + def test_tools_filtered(self) -> None: + """ToolChoice with tools list keeps the full tools schema and + narrows the callable subset via ``allowed_tools`` to preserve + prompt cache hits.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto", tools=["get_weather"]), + ) + self.assertListEqual(fmt_tools, _FT_TOOLS_RESPONSE) + self.assertEqual( + fmt_choice, + { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "function", "name": "get_weather"}], + }, + ) + + def test_no_tool_choice(self) -> None: + """Tools are converted when tool_choice is None.""" + fmt_tools, fmt_choice = self.model._format_tools(_FT_TOOLS, None) + self.assertEqual(fmt_tools, _FT_TOOLS_RESPONSE) + self.assertIsNone(fmt_choice) diff --git a/tests/model_xai_test.py b/tests/model_xai_test.py new file mode 100644 index 0000000000000000000000000000000000000000..a549e922bec8ca6d04b7cc19a65a72910534bf42 --- /dev/null +++ b/tests/model_xai_test.py @@ -0,0 +1,575 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for XAIChatModel with mocked API responses. + +Tests cover both non-streaming and streaming modes. +XAI uses xai_sdk with chat.stream() for streaming. +""" +import sys +from typing import Any +from types import ModuleType +import unittest +from unittest import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, MagicMock, patch + +from utils import AnyString + +from agentscope.message import TextBlock, ToolCallBlock, ThinkingBlock +from agentscope.model import XAIChatModel +from agentscope.credential import XAICredential +from agentscope.tool import ToolChoice + +A = AnyString() + + +# --------------------------------------------------------------------------- +# Build a lightweight xai_sdk stub so tests run without the real package. +# --------------------------------------------------------------------------- + + +def _build_xai_sdk_stub() -> None: + """Register stub modules for xai_sdk so imports don't fail.""" + if "xai_sdk" in sys.modules: + return + + chat_pb2 = ModuleType("xai_sdk.chat.chat_pb2") + + class _EnumHelper: + """Helper that makes .Value() return an integer.""" + + _mapping = { + "ROLE_ASSISTANT": 2, + "TOOL_CALL_TYPE_CLIENT_SIDE_TOOL": 1, + } + + def Value(self, name: str) -> int: + """Return the integer value for the given enum name.""" + return self._mapping.get(name, 0) + + chat_pb2.MessageRole = _EnumHelper() + chat_pb2.ToolCallType = _EnumHelper() + + class _RepeatedField(list): + """Minimal repeated proto field that supports .add().""" + + def __init__(self, factory: Any) -> None: + super().__init__() + self._factory = factory + + def add(self) -> Any: + """Add a new item using the factory and return it.""" + item = self._factory() + self.append(item) + return item + + class _FunctionSpec: + name: str = "" + arguments: str = "" + + class _ToolCallProto: + id: str = "" + type: int = 0 + function = _FunctionSpec() + + class _ContentPart: + text: str = "" + + class _MessageProto: + def __init__(self) -> None: + self.role = 0 + self.content = _RepeatedField(_ContentPart) + self.tool_calls = _RepeatedField(_ToolCallProto) + + chat_pb2.Message = _MessageProto + + xai_chat = ModuleType("xai_sdk.chat") + xai_chat.chat_pb2 = chat_pb2 + xai_chat.user = lambda *args: MagicMock(role="user", args=args) + xai_chat.assistant = lambda *args: MagicMock(role="assistant", args=args) + xai_chat.system = lambda *args: MagicMock(role="system", args=args) + xai_chat.tool_result = lambda *args, **kw: MagicMock( + role="tool", + args=args, + kwargs=kw, + ) + xai_chat.image = lambda url: MagicMock(type="image", url=url) + + xai_sdk = ModuleType("xai_sdk") + xai_sdk.chat = xai_chat + xai_sdk.AsyncClient = MagicMock() + + sys.modules["xai_sdk"] = xai_sdk + sys.modules["xai_sdk.chat"] = xai_chat + sys.modules["xai_sdk.chat.chat_pb2"] = chat_pb2 + + +_build_xai_sdk_stub() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_model(stream: bool = False) -> Any: + return XAIChatModel( + credential=XAICredential(api_key="test"), + model="grok-3", + stream=stream, + context_size=131_072, + ) + + +def _mock_completion( + text: str = "", + reasoning: str = "", + tool_calls: list | None = None, + response_id: str = "xai-resp-1", +) -> MagicMock: + """Build a mock xAI non-streaming response.""" + resp = MagicMock() + resp.id = response_id + resp.content = text + resp.reasoning_content = reasoning + resp.tool_calls = None + resp.usage = None + + if tool_calls: + tc_mocks = [] + for tc in tool_calls: + m = MagicMock() + m.id = tc["id"] + m.function.name = tc["name"] + m.function.arguments = tc["arguments"] + tc_mocks.append(m) + resp.tool_calls = tc_mocks + + return resp + + +class _MockStreamChunk: + """A single chunk from xai chat.stream().""" + + def __init__( + self, + content: str = "", + reasoning_content: str = "", + ) -> None: + self.content = content + self.reasoning_content = reasoning_content + + +class _MockChatStream: + """Mock xai_sdk chat session with stream() and sample().""" + + def __init__( + self, + stream_items: list | None = None, + sample_response: Any = None, + ) -> None: + self._stream_items = stream_items or [] + self._sample_response = sample_response + self._appended: list = [] + + def append(self, msg: Any) -> None: + """Append a message to the conversation.""" + self._appended.append(msg) + + async def sample(self) -> Any: + """Return the pre-configured sample response.""" + return self._sample_response + + def stream(self) -> "_MockStreamIterator": + """Return an async iterator over pre-configured stream items.""" + return _MockStreamIterator(self._stream_items) + + +class _MockStreamIterator: + """Async iterator for (response, chunk) pairs from xai stream.""" + + def __init__(self, items: list) -> None: + self._items = items + self._index = 0 + + def __aiter__(self) -> "_MockStreamIterator": + return self + + async def __anext__(self) -> tuple: + if self._index >= len(self._items): + raise StopAsyncIteration + item = self._items[self._index] + self._index += 1 + return item + + +# --------------------------------------------------------------------------- +# Non-streaming tests +# --------------------------------------------------------------------------- + + +class TestXAINonStream(IsolatedAsyncioTestCase): + """Tests for XAIChatModel in non-streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=False) + + @patch("xai_sdk.AsyncClient") + async def test_text_response(self, mock_client_cls: MagicMock) -> None: + """Non-stream text response returns a single ChatResponse.""" + mock_chat = _MockChatStream( + sample_response=_mock_completion(text="Hello!"), + ) + mock_client_cls.return_value.chat.create.return_value = mock_chat + mock_client_cls.return_value.close = AsyncMock() + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + (True, [TextBlock.model_construct(id=A, text="Hello!")]), + ) + self.assertEqual(result.id, "xai-resp-1") + + @patch("xai_sdk.AsyncClient") + async def test_tool_call_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Parsing a tool-call response creates a ToolCallBlock.""" + mock_chat = _MockChatStream( + sample_response=_mock_completion( + tool_calls=[ + { + "id": "call-1", + "name": "get_weather", + "arguments": '{"city":"NY"}', + }, + ], + ), + ) + mock_client_cls.return_value.chat.create.return_value = mock_chat + mock_client_cls.return_value.close = AsyncMock() + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ToolCallBlock( + id="call-1", + name="get_weather", + input='{"city":"NY"}', + ), + ], + ), + ) + + @patch("xai_sdk.AsyncClient") + async def test_thinking_response( + self, + mock_client_cls: MagicMock, + ) -> None: + """Non-stream reasoning plus text returns ThinkingBlock then + TextBlock.""" + mock_chat = _MockChatStream( + sample_response=_mock_completion( + text="Answer", + reasoning="Deep thinking...", + ), + ) + mock_client_cls.return_value.chat.create.return_value = mock_chat + mock_client_cls.return_value.close = AsyncMock() + + result = await self.model([]) + + self.assertEqual( + (result.is_last, result.content), + ( + True, + [ + ThinkingBlock.model_construct( + id=A, + thinking="Deep thinking...", + ), + TextBlock.model_construct(id=A, text="Answer"), + ], + ), + ) + + +# --------------------------------------------------------------------------- +# Streaming tests +# --------------------------------------------------------------------------- + + +class TestXAIStream(IsolatedAsyncioTestCase): + """Tests for XAIChatModel in streaming mode.""" + + def setUp(self) -> None: + self.model = _make_model(stream=True) + + @patch("xai_sdk.AsyncClient") + async def test_stream_text(self, mock_client_cls: MagicMock) -> None: + """Stream text yields deltas then final with full content.""" + final_response = _mock_completion(text="Hello world") + final_response.tool_calls = None + final_response.usage = MagicMock() + final_response.usage.prompt_tokens = 10 + final_response.usage.completion_tokens = 5 + final_response.usage.cached_prompt_text_tokens = 0 + + stream_items = [ + (final_response, _MockStreamChunk(content="Hello")), + (final_response, _MockStreamChunk(content=" world")), + ] + mock_chat = _MockChatStream(stream_items=stream_items) + mock_client_cls.return_value.chat.create.return_value = mock_chat + mock_client_cls.return_value.close = AsyncMock() + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + (False, [TextBlock.model_construct(id=A, text="Hello")]), + (False, [TextBlock.model_construct(id=A, text=" world")]), + (True, [TextBlock.model_construct(id=A, text="Hello world")]), + ], + ) + + @patch("xai_sdk.AsyncClient") + async def test_stream_thinking_and_text( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream reasoning and text deltas then final with accumulated + content.""" + final_response = _mock_completion(text="") + final_response.tool_calls = None + final_response.usage = MagicMock() + final_response.usage.prompt_tokens = 10 + final_response.usage.completion_tokens = 5 + final_response.usage.cached_prompt_text_tokens = 0 + + stream_items = [ + (final_response, _MockStreamChunk(reasoning_content="Think")), + (final_response, _MockStreamChunk(content="Answer")), + ] + mock_chat = _MockChatStream(stream_items=stream_items) + mock_client_cls.return_value.chat.create.return_value = mock_chat + mock_client_cls.return_value.close = AsyncMock() + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + ( + False, + [ThinkingBlock.model_construct(id=A, thinking="Think")], + ), + (False, [TextBlock.model_construct(id=A, text="Answer")]), + ( + True, + [ + ThinkingBlock.model_construct(id=A, thinking="Think"), + TextBlock.model_construct(id=A, text="Answer"), + ], + ), + ], + ) + + @patch("xai_sdk.AsyncClient") + async def test_stream_tool_calls_in_final( + self, + mock_client_cls: MagicMock, + ) -> None: + """Stream text delta then final adds tool calls from last_response.""" + final_response = _mock_completion( + tool_calls=[ + { + "id": "call-1", + "name": "search", + "arguments": '{"q":"test"}', + }, + ], + ) + final_response.usage = MagicMock() + final_response.usage.prompt_tokens = 10 + final_response.usage.completion_tokens = 5 + final_response.usage.cached_prompt_text_tokens = 0 + + stream_items = [ + (final_response, _MockStreamChunk(content="I'll search")), + ] + mock_chat = _MockChatStream(stream_items=stream_items) + mock_client_cls.return_value.chat.create.return_value = mock_chat + mock_client_cls.return_value.close = AsyncMock() + + gen = await self.model([]) + responses = [r async for r in gen] + + self.assertListEqual( + [(r.is_last, r.content) for r in responses], + [ + (False, [TextBlock.model_construct(id=A, text="I'll search")]), + ( + True, + [ + TextBlock.model_construct(id=A, text="I'll search"), + ToolCallBlock( + id="call-1", + name="search", + input='{"q":"test"}', + ), + ], + ), + ], + ) + + +class TestXAIModelParameters(unittest.TestCase): + """Tests for XAIChatModel.Parameters.""" + + def test_thinking_enable_stored_on_model(self) -> None: + """thinking_enable is accessible through model.parameters.""" + model = XAIChatModel( + credential=XAICredential(api_key="test"), + model="grok-3-mini", + stream=False, + context_size=131_072, + parameters=XAIChatModel.Parameters(thinking_enable=True), + ) + self.assertTrue(model.parameters.thinking_enable) + + def test_reasoning_effort_stored_on_model(self) -> None: + """reasoning_effort is accessible through model.parameters.""" + model = XAIChatModel( + credential=XAICredential(api_key="test"), + model="grok-3-mini", + stream=False, + context_size=131_072, + parameters=XAIChatModel.Parameters(reasoning_effort="high"), + ) + self.assertEqual(model.parameters.reasoning_effort, "high") + + +# --------------------------------------------------------------------------- +# _format_tools tests +# --------------------------------------------------------------------------- + +_FT_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "description": "Get the time", + "parameters": { + "type": "object", + "properties": {"timezone": {"type": "string"}}, + "required": ["timezone"], + }, + }, + }, +] + + +def _extend_xai_stub_for_tools() -> None: + """Add tool and required_tool stubs to the existing xai_sdk stub.""" + xai_chat = sys.modules.get("xai_sdk.chat") + if xai_chat is None or hasattr(xai_chat, "required_tool"): + return + + class _RequiredTool: + def __init__(self, tool_name: str) -> None: + self.tool_name = tool_name + + def __eq__(self, other: object) -> bool: + return ( + isinstance(other, _RequiredTool) + and self.tool_name == other.tool_name + ) + + def _tool_stub( + name: str, + description: str = "", + parameters: Any = None, + ) -> MagicMock: + m = MagicMock() + m.name = name + m.description = description + m.parameters = parameters or {} + return m + + xai_chat.required_tool = _RequiredTool + xai_chat.tool = _tool_stub + + +_extend_xai_stub_for_tools() + + +class TestXAIFormatTools(unittest.TestCase): + """Tests for XAIChatModel._format_tools.""" + + def setUp(self) -> None: + self.model = _make_model() + + def test_no_tool_choice(self) -> None: + """All tools are returned when tool_choice is None.""" + fmt_tools, fmt_choice = self.model._format_tools(_FT_TOOLS, None) + self.assertIsNotNone(fmt_tools) + self.assertEqual(len(fmt_tools), 2) + self.assertIsNone(fmt_choice) + + def test_auto_mode(self) -> None: + """Auto mode passes tools through with choice 'auto'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto"), + ) + self.assertIsNotNone(fmt_tools) + self.assertEqual(fmt_choice, "auto") + + def test_none_mode(self) -> None: + """None mode passes tools through with choice 'none'.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="none"), + ) + self.assertIsNotNone(fmt_tools) + self.assertEqual(fmt_choice, "none") + + def test_str_mode_force_call(self) -> None: + """String mode forces a required_tool for the named function.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="get_weather"), + ) + self.assertIsNotNone(fmt_tools) + self.assertEqual(fmt_choice.tool_name, "get_weather") + + def test_tools_filtered(self) -> None: + """ToolChoice with tools list filters to matching function names.""" + fmt_tools, fmt_choice = self.model._format_tools( + _FT_TOOLS, + ToolChoice(mode="auto", tools=["get_weather"]), + ) + self.assertEqual(len(fmt_tools), 1) + self.assertEqual(fmt_tools[0].name, "get_weather") + self.assertEqual(fmt_choice, "auto") diff --git a/tests/permission_bash_parser_test.py b/tests/permission_bash_parser_test.py new file mode 100644 index 0000000000000000000000000000000000000000..1998f174e618b9593c303c02245c74007f38389d --- /dev/null +++ b/tests/permission_bash_parser_test.py @@ -0,0 +1,1317 @@ +# -*- coding: utf-8 -*- +"""Test cases for BashCommandParser.""" +import sys +import unittest +from unittest.async_case import IsolatedAsyncioTestCase + +from agentscope.tool._builtin._bash_parser import BashCommandParser +from agentscope.tool import Bash + + +@unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", +) +class BashCommandParserTest(IsolatedAsyncioTestCase): + """Test cases for BashCommandParser.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.parser = BashCommandParser() + + async def test_single_command_with_prefix(self) -> None: + """Test single commands that can extract prefixes.""" + test_cases = [ + ("git commit -m 'fix'", ["git commit"]), + ("npm run build", ["npm run"]), + ("docker compose up", ["docker compose"]), + ("cargo build --release", ["cargo build"]), + ("python -m pytest", ["python -m"]), + ] + + for command, expected in test_cases: + with self.subTest(command=command): + result = self.parser.extract_command_prefixes(command) + self.assertEqual(result, expected) + + async def test_single_command_without_prefix(self) -> None: + """Test single commands that cannot extract prefixes.""" + test_cases = [ + ("ls", []), + ("ls -la", []), + ("echo hello", []), + ("cd /tmp", []), + ("pwd", []), + ("cat file.txt", []), + ] + + for command, expected in test_cases: + with self.subTest(command=command): + result = self.parser.extract_command_prefixes(command) + self.assertEqual(result, expected) + + async def test_environment_variables_safe(self) -> None: + """Test commands with safe environment variables.""" + test_cases = [ + ("NODE_ENV=prod npm run build", ["npm run"]), + ("DEBUG=1 npm test", ["npm test"]), + ("PATH=/usr/bin npm run build", ["npm run"]), + ("CI=true npm run test", ["npm run"]), + ("PYTHONUNBUFFERED=1 python -m pytest", ["python -m"]), + ] + + for command, expected in test_cases: + with self.subTest(command=command): + result = self.parser.extract_command_prefixes(command) + self.assertEqual(result, expected) + + async def test_environment_variables_unsafe(self) -> None: + """Test commands with unsafe environment variables.""" + test_cases = [ + ("CUSTOM_VAR=value npm run build", []), + ("MY_SECRET=123 npm test", []), + ("API_KEY=abc npm run deploy", []), + ] + + for command, expected in test_cases: + with self.subTest(command=command): + result = self.parser.extract_command_prefixes(command) + self.assertEqual(result, expected) + + async def test_compound_command_and_operator(self) -> None: + """Test compound commands with && operator.""" + test_cases = [ + ("git add . && git commit", ["git add", "git commit"]), + ( + "npm install && npm run build && npm test", + ["npm install", "npm run", "npm test"], + ), + ("docker build . && docker push", ["docker build", "docker push"]), + ] + + for command, expected in test_cases: + with self.subTest(command=command): + result = self.parser.extract_command_prefixes(command) + self.assertEqual(result, expected) + + async def test_compound_command_or_operator(self) -> None: + """Test compound commands with || operator.""" + test_cases = [ + ("npm run build || echo failed", ["npm run"]), + ("git commit || git status", ["git commit", "git status"]), + ] + + for command, expected in test_cases: + with self.subTest(command=command): + result = self.parser.extract_command_prefixes(command) + self.assertEqual(result, expected) + + async def test_compound_command_semicolon(self) -> None: + """Test compound commands with ; operator.""" + test_cases = [ + ("cd /tmp; ls -la; pwd", []), + ("npm install; npm run build", ["npm install", "npm run"]), + ] + + for command, expected in test_cases: + with self.subTest(command=command): + result = self.parser.extract_command_prefixes(command) + self.assertEqual(result, expected) + + async def test_compound_command_pipe(self) -> None: + """Test compound commands with | operator.""" + test_cases = [ + ("cat file.txt | grep error", []), + ("docker ps | grep nginx", ["docker ps"]), + ("npm run build | tee output.log", ["npm run"]), + ] + + for command, expected in test_cases: + with self.subTest(command=command): + result = self.parser.extract_command_prefixes(command) + self.assertEqual(result, expected) + + async def test_compound_command_mixed(self) -> None: + """Test compound commands with mixed operators.""" + test_cases = [ + ( + "git add . && git commit -m 'fix' || echo failed", + ["git add", "git commit"], + ), + ( + "npm install && npm run build | tee log.txt", + ["npm install", "npm run"], + ), + ] + + for command, expected in test_cases: + with self.subTest(command=command): + result = self.parser.extract_command_prefixes(command) + self.assertEqual(result, expected) + + async def test_edge_cases(self) -> None: + """Test edge cases.""" + test_cases = [ + ("", []), + (" ", []), + ("npm", []), + (" npm run build ", ["npm run"]), + ] + + for command, expected in test_cases: + with self.subTest(command=command): + result = self.parser.extract_command_prefixes(command) + self.assertEqual(result, expected) + + async def test_max_prefixes_limit(self) -> None: + """Test that max_prefixes parameter limits the results.""" + command = ( + "npm install && npm run build && npm test && " + "npm run lint && npm run format && npm run deploy" + ) + + # Default max is 5, but with deduplication we get 3 unique prefixes + result = self.parser.extract_command_prefixes(command) + self.assertEqual(len(result), 3) + self.assertEqual( + result, + ["npm install", "npm run", "npm test"], + ) + + # Custom max + result = self.parser.extract_command_prefixes(command, max_prefixes=3) + self.assertEqual(len(result), 3) + + async def test_deduplication(self) -> None: + """Test that duplicate prefixes are removed.""" + test_cases = [ + ( + "npm run build && npm run test && npm run lint", + ["npm run"], + ), + ( + "git add . && git commit && git push && git status", + ["git add", "git commit", "git push", "git status"], + ), + ] + + for command, expected in test_cases: + with self.subTest(command=command): + result = self.parser.extract_command_prefixes(command) + self.assertEqual(result, expected) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.parser = None + + +@unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", +) +class BashParserReadOnlyTest(IsolatedAsyncioTestCase): + """Test is_read_only_command() method.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.parser = BashCommandParser() + + async def test_single_read_only_git_commands(self) -> None: + """Test single read-only git commands.""" + read_only_commands = [ + "git status", + "git log", + "git diff", + "git show", + "git branch", + "git remote -v", + "git log --oneline", + ] + for cmd in read_only_commands: + with self.subTest(cmd=cmd): + self.assertTrue( + self.parser.is_read_only_command(cmd), + f"Expected '{cmd}' to be read-only", + ) + + async def test_single_read_only_file_commands(self) -> None: + """Test single read-only file commands.""" + read_only_commands = [ + "ls", + "ls -la", + "cat file.txt", + "head -n 10 file.txt", + "tail -f log.txt", + "grep pattern file.txt", + "find . -name '*.py'", + "tree", + "pwd", + "which python", + ] + for cmd in read_only_commands: + with self.subTest(cmd=cmd): + self.assertTrue( + self.parser.is_read_only_command(cmd), + f"Expected '{cmd}' to be read-only", + ) + + async def test_single_read_only_docker_commands(self) -> None: + """Test single read-only docker commands.""" + read_only_commands = [ + "docker ps", + "docker images", + "docker inspect container_id", + "docker logs container_id", + ] + for cmd in read_only_commands: + with self.subTest(cmd=cmd): + self.assertTrue( + self.parser.is_read_only_command(cmd), + f"Expected '{cmd}' to be read-only", + ) + + async def test_single_non_read_only_commands(self) -> None: + """Test single non-read-only commands.""" + non_read_only_commands = [ + "git commit -m 'message'", + "git push", + "git pull", + "rm file.txt", + "mv file1.txt file2.txt", + "cp file1.txt file2.txt", + "chmod +x script.sh", + "mkdir new_dir", + "touch file.txt", + ] + for cmd in non_read_only_commands: + with self.subTest(cmd=cmd): + self.assertFalse( + self.parser.is_read_only_command(cmd), + f"Expected '{cmd}' to be non-read-only", + ) + + async def test_compound_command_all_read_only(self) -> None: + """Test compound commands with all read-only subcommands.""" + compound_commands = [ + "ls -la && cat file.txt", + "git status && git log", + "pwd && ls", + "cat file1.txt || cat file2.txt", + "ls; pwd; cat file.txt", + "git diff | grep pattern", + ] + for cmd in compound_commands: + with self.subTest(cmd=cmd): + self.assertTrue( + self.parser.is_read_only_command(cmd), + f"Expected '{cmd}' to be read-only " + f"(all subcommands are read-only)", + ) + + async def test_compound_command_mixed(self) -> None: + """Test compound commands with mixed read-only and non-read-only.""" + mixed_commands = [ + "ls -la && git commit -m 'message'", + "cat file.txt && rm file.txt", + "git status && git push", + "pwd || mkdir new_dir", + "ls; touch file.txt", + ] + for cmd in mixed_commands: + with self.subTest(cmd=cmd): + self.assertFalse( + self.parser.is_read_only_command(cmd), + f"Expected '{cmd}' to be non-read-only " + f"(contains non-read-only subcommand)", + ) + + async def test_commands_with_output_redirection(self) -> None: + """Test commands with output redirections are not read-only.""" + redirect_commands = [ + "cat file.txt > output.txt", + "ls -la > list.txt", + "git log >> history.txt", + "echo 'hello' > file.txt", + "cat file.txt 2> error.log", + "ls &> output.log", + ] + for cmd in redirect_commands: + with self.subTest(cmd=cmd): + self.assertFalse( + self.parser.is_read_only_command(cmd), + f"Expected '{cmd}' to be non-read-only " + f"(contains output redirection)", + ) + + async def test_commands_with_dangerous_paths(self) -> None: + """Test commands with dangerous paths.""" + # Note: dangerous path check is separate from read-only check + # These commands are still considered read-only if they don't + # modify files + dangerous_read_only = [ + "cat ~/.bashrc", + "ls ~/.ssh", + "cat .git/config", + ] + for cmd in dangerous_read_only: + with self.subTest(cmd=cmd): + self.assertTrue( + self.parser.is_read_only_command(cmd), + f"Expected '{cmd}' to be read-only " + f"(dangerous path doesn't affect read-only status)", + ) + + # These are non-read-only because they modify files + dangerous_non_read_only = [ + "rm ~/.bashrc", + "chmod 600 ~/.ssh/config", + "mv file.txt ~/.ssh/", + ] + for cmd in dangerous_non_read_only: + with self.subTest(cmd=cmd): + self.assertFalse( + self.parser.is_read_only_command(cmd), + f"Expected '{cmd}' to be non-read-only " + f"(modifies files)", + ) + + async def test_empty_and_whitespace_commands(self) -> None: + """Test empty and whitespace-only commands.""" + empty_commands = [ + "", + " ", + "\t", + "\n", + ] + for cmd in empty_commands: + with self.subTest(cmd=repr(cmd)): + # Empty commands return False (not read-only, but + # also not executable) + self.assertFalse( + self.parser.is_read_only_command(cmd), + "Expected empty/whitespace command to return False", + ) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.parser = None + + +@unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", +) +class BashParserFilePathsTest(IsolatedAsyncioTestCase): + """Test extract_file_paths() method.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.parser = BashCommandParser() + + async def test_rm_command(self) -> None: + """Test file path extraction from rm commands.""" + test_cases = [ + ("rm file.txt", [("rm", "file.txt")]), + ("rm -rf /tmp/test", [("rm", "/tmp/test")]), + ( + "rm file1.txt file2.txt", + [("rm", "file1.txt"), ("rm", "file2.txt")], + ), + ("rm -f *.log", [("rm", "*.log")]), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_file_paths(cmd) + self.assertEqual(result, expected) + + async def test_mv_command(self) -> None: + """Test file path extraction from mv commands.""" + test_cases = [ + ("mv old.txt new.txt", [("mv", "old.txt"), ("mv", "new.txt")]), + ( + "mv /tmp/file.txt /home/user/", + [("mv", "/tmp/file.txt"), ("mv", "/home/user/")], + ), + ( + "mv -f file1.txt file2.txt", + [("mv", "file1.txt"), ("mv", "file2.txt")], + ), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_file_paths(cmd) + self.assertEqual(result, expected) + + async def test_cp_command(self) -> None: + """Test file path extraction from cp commands.""" + test_cases = [ + ( + "cp file.txt backup.txt", + [("cp", "file.txt"), ("cp", "backup.txt")], + ), + ("cp -r /src /dest", [("cp", "/src"), ("cp", "/dest")]), + ( + "cp file1.txt file2.txt /tmp/", + [("cp", "file1.txt"), ("cp", "file2.txt"), ("cp", "/tmp/")], + ), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_file_paths(cmd) + self.assertEqual(result, expected) + + async def test_chmod_command(self) -> None: + """Test file path extraction from chmod commands.""" + test_cases = [ + ("chmod +x script.sh", [("chmod", "+x"), ("chmod", "script.sh")]), + ("chmod 755 /usr/bin/tool", [("chmod", "/usr/bin/tool")]), + ("chmod -R 644 /tmp/files", [("chmod", "/tmp/files")]), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_file_paths(cmd) + self.assertEqual(result, expected) + + async def test_chown_command(self) -> None: + """Test file path extraction from chown commands.""" + test_cases = [ + ( + "chown user:group file.txt", + [("chown", "user:group"), ("chown", "file.txt")], + ), + ( + "chown -R user /var/www", + [("chown", "user"), ("chown", "/var/www")], + ), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_file_paths(cmd) + self.assertEqual(result, expected) + + async def test_touch_command(self) -> None: + """Test file path extraction from touch commands.""" + test_cases = [ + ("touch file.txt", [("touch", "file.txt")]), + ( + "touch file1.txt file2.txt", + [("touch", "file1.txt"), ("touch", "file2.txt")], + ), + ("touch /tmp/newfile.log", [("touch", "/tmp/newfile.log")]), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_file_paths(cmd) + self.assertEqual(result, expected) + + async def test_ln_command(self) -> None: + """Test file path extraction from ln commands.""" + test_cases = [ + ("ln -s target link", [("ln", "target"), ("ln", "link")]), + ( + "ln /src/file /dest/file", + [("ln", "/src/file"), ("ln", "/dest/file")], + ), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_file_paths(cmd) + self.assertEqual(result, expected) + + async def test_sed_command(self) -> None: + """Test file path extraction from sed commands.""" + test_cases = [ + ("sed -i 's/old/new/' file.txt", [("sed", "file.txt")]), + ("sed 's/pattern/replacement/' input.txt", [("sed", "input.txt")]), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_file_paths(cmd) + self.assertEqual(result, expected) + + async def test_output_redirections(self) -> None: + """Test file path extraction from output redirections.""" + test_cases = [ + ( + "echo 'hello' > output.txt", + [("redirect", "output.txt")], + ), + ("cat file.txt > backup.txt", [("redirect", "backup.txt")]), + ("ls -la >> list.txt", [("redirect", "list.txt")]), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_file_paths(cmd) + self.assertEqual(result, expected) + + async def test_compound_commands(self) -> None: + """Test file path extraction from compound commands.""" + test_cases = [ + ( + "rm file1.txt && rm file2.txt", + [("rm", "file1.txt"), ("rm", "file2.txt")], + ), + ( + "touch new.txt && chmod +x new.txt", + [("touch", "new.txt"), ("chmod", "+x"), ("chmod", "new.txt")], + ), + ( + "cp src.txt dest.txt || mv src.txt dest.txt", + [ + ("cp", "src.txt"), + ("cp", "dest.txt"), + ("mv", "src.txt"), + ("mv", "dest.txt"), + ], + ), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_file_paths(cmd) + self.assertEqual(result, expected) + + async def test_quoted_paths(self) -> None: + """Test file path extraction with quoted paths.""" + test_cases = [ + ('rm "file with spaces.txt"', []), # Quoted paths not extracted + ("rm 'file.txt'", []), # Quoted paths not extracted + ( + 'mv "old file.txt" "new file.txt"', + [], # Quoted paths not extracted + ), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_file_paths(cmd) + self.assertEqual(result, expected) + + async def test_dangerous_paths(self) -> None: + """Test file path extraction with dangerous paths.""" + test_cases = [ + ("rm ~/.bashrc", [("rm", "~/.bashrc")]), + ("chmod 600 ~/.ssh/config", [("chmod", "~/.ssh/config")]), + ( + "mv file.txt .git/hooks/", + [("mv", "file.txt"), ("mv", ".git/hooks/")], + ), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_file_paths(cmd) + self.assertEqual(result, expected) + + async def test_commands_without_file_operations(self) -> None: + """Test commands that don't operate on files.""" + test_cases = [ + ("ls", []), + ("pwd", []), + ("echo 'hello'", []), + ("git status", []), + ("docker ps", []), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_file_paths(cmd) + self.assertEqual(result, expected) + + async def test_empty_command(self) -> None: + """Test empty command.""" + result = self.parser.extract_file_paths("") + self.assertEqual(result, []) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.parser = None + + +@unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", +) +class BashParserRedirectionsTest(IsolatedAsyncioTestCase): + """Test extract_redirections() method.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.parser = BashCommandParser() + + async def test_simple_output_redirection(self) -> None: + """Test simple output redirection (>).""" + test_cases = [ + ("echo 'hello' > output.txt", ["output.txt"]), + ("cat file.txt > backup.txt", ["backup.txt"]), + ("ls -la > list.txt", ["list.txt"]), + ("git log > history.txt", ["history.txt"]), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_redirections(cmd) + self.assertEqual(result, expected) + + async def test_append_redirection(self) -> None: + """Test append redirection (>>).""" + test_cases = [ + ("echo 'line' >> log.txt", ["log.txt"]), + ("cat file.txt >> combined.txt", ["combined.txt"]), + ("ls >> list.txt", ["list.txt"]), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_redirections(cmd) + self.assertEqual(result, expected) + + async def test_error_redirection(self) -> None: + """Test error redirection (2>).""" + test_cases = [ + ("command 2> error.log", ["error.log"]), + ("python script.py 2> stderr.txt", ["stderr.txt"]), + ("npm run build 2> build_errors.log", ["build_errors.log"]), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_redirections(cmd) + self.assertEqual(result, expected) + + async def test_combined_redirection(self) -> None: + """Test combined stdout/stderr redirection (&>).""" + test_cases = [ + ("command &> output.log", ["output.log"]), + ("python script.py &> all_output.txt", ["all_output.txt"]), + ("npm test &> test_results.log", ["test_results.log"]), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_redirections(cmd) + self.assertEqual(result, expected) + + async def test_multiple_redirections(self) -> None: + """Test commands with multiple redirections.""" + test_cases = [ + ("command > output.txt 2> error.log", ["output.txt", "error.log"]), + ( + "python script.py > stdout.txt 2> stderr.txt", + ["stdout.txt", "stderr.txt"], + ), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_redirections(cmd) + self.assertEqual(result, expected) + + async def test_compound_commands_with_redirections(self) -> None: + """Test compound commands with multiple redirections.""" + test_cases = [ + ( + "echo 'a' > file1.txt && echo 'b' > file2.txt", + ["file1.txt", "file2.txt"], + ), + ( + "cat file.txt > backup.txt || cp file.txt backup.txt", + ["backup.txt"], + ), + ("ls > list1.txt; pwd > list2.txt", ["list1.txt", "list2.txt"]), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_redirections(cmd) + self.assertEqual(result, expected) + + async def test_redirections_with_quoted_paths(self) -> None: + """Test redirections with quoted file paths.""" + test_cases = [ + ( + 'echo "hello" > "output file.txt"', + [], + ), # Quoted redirections not extracted + ( + "cat file.txt > 'backup.txt'", + [], + ), # Quoted redirections not extracted + ( + 'ls > "file with spaces.log"', + [], + ), # Quoted redirections not extracted + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_redirections(cmd) + self.assertEqual(result, expected) + + async def test_redirections_to_dangerous_paths(self) -> None: + """Test redirections to dangerous paths.""" + test_cases = [ + ("echo 'alias' >> ~/.bashrc", ["~/.bashrc"]), + ("cat key > ~/.ssh/authorized_keys", ["~/.ssh/authorized_keys"]), + ("echo 'config' > .git/config", [".git/config"]), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_redirections(cmd) + self.assertEqual(result, expected) + + async def test_redirections_with_absolute_paths(self) -> None: + """Test redirections with absolute paths.""" + test_cases = [ + ("echo 'data' > /tmp/output.txt", ["/tmp/output.txt"]), + ("cat file.txt > /var/log/app.log", ["/var/log/app.log"]), + ("ls > /home/user/list.txt", ["/home/user/list.txt"]), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_redirections(cmd) + self.assertEqual(result, expected) + + async def test_commands_without_redirections(self) -> None: + """Test commands without redirections.""" + test_cases = [ + ("ls -la", []), + ("cat file.txt", []), + ("echo 'hello'", []), + ("git status", []), + ("rm file.txt", []), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_redirections(cmd) + self.assertEqual(result, expected) + + async def test_pipe_not_redirection(self) -> None: + """Test that pipes (|) are not treated as redirections.""" + test_cases = [ + ("cat file.txt | grep pattern", []), + ("ls | wc -l", []), + ("git log | head -n 10", []), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.extract_redirections(cmd) + self.assertEqual(result, expected) + + async def test_empty_command(self) -> None: + """Test empty command.""" + result = self.parser.extract_redirections("") + self.assertEqual(result, []) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.parser = None + + +@unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", +) +class BashParserSedConstraintsTest(IsolatedAsyncioTestCase): + """Test check_sed_constraints() with complex allowlist/denylist.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.parser = BashCommandParser() + self.dangerous_files = [".env", "config.json", "secrets.yaml"] + + async def test_allowlist_line_printing(self) -> None: + """Test allowlist Pattern 1: Line printing with -n flag.""" + test_cases = [ + ("sed -n '5p' file.txt", None), + ("sed -n '10,20p' file.txt", None), + ("sed -n '1p' data.log", None), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_sed_constraints( + cmd, + self.dangerous_files, + ) + self.assertEqual(result, expected) + + async def test_allowlist_substitution(self) -> None: + """Test allowlist Pattern 2: Substitution commands.""" + test_cases = [ + ("sed 's/old/new/' file.txt", None), + ("sed 's/old/new/g' file.txt", None), + ("sed 's|old|new|' file.txt", None), + ("sed 's#pattern#replacement#' file.txt", None), + ] + for cmd, expected in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_sed_constraints( + cmd, + self.dangerous_files, + ) + self.assertEqual(result, expected) + + async def test_denylist_write_operations(self) -> None: + """Test denylist: write operations (w/W).""" + test_cases = [ + ("sed 's/old/new/w output.txt' file.txt", "write operation"), + ("sed 's/old/new/W output.txt' file.txt", "write operation"), + ] + for cmd, expected_substring in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_sed_constraints( + cmd, + self.dangerous_files, + ) + self.assertIsNotNone(result) + self.assertIn(expected_substring, result) + + async def test_denylist_execute_operations(self) -> None: + """Test denylist: execute operations (e/E).""" + test_cases = [ + ("sed 's/old/new/e' file.txt", "execute operation"), + ("sed 's/old/new/E' file.txt", "execute operation"), + ] + for cmd, expected_substring in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_sed_constraints( + cmd, + self.dangerous_files, + ) + self.assertIsNotNone(result) + self.assertIn(expected_substring, result) + + async def test_denylist_dangerous_patterns(self) -> None: + """Test denylist: dangerous patterns (curly braces, negation).""" + test_cases = [ + ("sed '{s/old/new/}' file.txt", "curly braces"), + ("sed '!s/old/new/' file.txt", "negation"), + ] + for cmd, expected_substring in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_sed_constraints( + cmd, + self.dangerous_files, + ) + self.assertIsNotNone(result) + self.assertIn(expected_substring, result) + + async def test_not_in_allowlist(self) -> None: + """Test commands not in allowlist.""" + test_cases = [ + ("sed 'd' file.txt", "not in allowlist"), + ("sed 'a\\text' file.txt", "not in allowlist"), + ("sed 'i\\text' file.txt", "not in allowlist"), + ] + for cmd, expected_substring in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_sed_constraints( + cmd, + self.dangerous_files, + ) + self.assertIsNotNone(result) + self.assertIn(expected_substring, result) + + async def test_inplace_with_dangerous_files(self) -> None: + """Test -i flag with dangerous files.""" + test_cases = [ + ("sed -i 's/old/new/' .env", "dangerous file"), + ("sed -i 's/old/new/' config.json", "dangerous file"), + ("sed --in-place 's/old/new/' secrets.yaml", "dangerous file"), + ] + for cmd, expected_substring in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_sed_constraints( + cmd, + self.dangerous_files, + ) + self.assertIsNotNone(result) + self.assertIn(expected_substring, result) + + async def test_invalid_flags(self) -> None: + """Test invalid flags.""" + test_cases = [ + ("sed -r 's/old/new/' file.txt", "flag -r not allowed"), + ("sed -u 's/old/new/' file.txt", "flag -u not allowed"), + ] + for cmd, expected_substring in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_sed_constraints( + cmd, + self.dangerous_files, + ) + self.assertIsNotNone(result) + self.assertIn(expected_substring, result) + + async def test_non_sed_commands(self) -> None: + """Test non-sed commands return None.""" + test_cases = [ + "ls -la", + "cat file.txt", + "grep pattern file.txt", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_sed_constraints( + cmd, + self.dangerous_files, + ) + self.assertIsNone(result) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.parser = None + + +@unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", +) +class BashWildcardMatchingTest(IsolatedAsyncioTestCase): + """Test match_rule() with complex regex-based wildcard matching.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.bash_tool = Bash() + + async def test_basic_wildcard_matching(self) -> None: + """Test basic wildcard matching with *.""" + test_cases = [ + ("git *", "git status", True), + ("git *", "git", True), # Special optimization + ("git * push", "git origin push", True), + ("npm run *", "npm run build", True), + ("npm run *", "npm run test", True), + ("docker * up", "docker compose up", True), + ] + for pattern, command, expected in test_cases: + with self.subTest(pattern=pattern, command=command): + result = await self.bash_tool.match_rule( + pattern, + {"command": command}, + ) + self.assertEqual(result, expected) + + async def test_wildcard_no_match(self) -> None: + """Test wildcard patterns that don't match.""" + test_cases = [ + ("git *", "github", False), + ("npm run *", "yarn run build", False), + ("docker * up", "docker ps", False), + ] + for pattern, command, expected in test_cases: + with self.subTest(pattern=pattern, command=command): + result = await self.bash_tool.match_rule( + pattern, + {"command": command}, + ) + self.assertEqual(result, expected) + + async def test_escape_sequences(self) -> None: + """Test escape sequences \\* and \\\\.""" + test_cases = [ + ("file\\*.txt", "file*.txt", True), + ("file\\*.txt", "file123.txt", False), + ("path\\\\to\\\\file", "path\\to\\file", True), + ("path\\\\to\\\\file", "path/to/file", False), + ] + for pattern, command, expected in test_cases: + with self.subTest(pattern=pattern, command=command): + result = await self.bash_tool.match_rule( + pattern, + {"command": command}, + ) + self.assertEqual(result, expected) + + async def test_prefix_pattern(self) -> None: + """Test prefix pattern with :* suffix.""" + test_cases = [ + ("git:*", "git status", True), + ("git:*", "git", True), + ("git:*", "github", False), + ("npm:*", "npm install", True), + ("npm:*", "npm", True), + ] + for pattern, command, expected in test_cases: + with self.subTest(pattern=pattern, command=command): + result = await self.bash_tool.match_rule( + pattern, + {"command": command}, + ) + self.assertEqual(result, expected) + + async def test_substring_matching(self) -> None: + """Test substring matching (no wildcards).""" + test_cases = [ + ("npm install", "npm install express", True), + ("npm install", "yarn install", False), + ("git commit", "git commit -m 'fix'", True), + ("git commit", "git push", False), + ] + for pattern, command, expected in test_cases: + with self.subTest(pattern=pattern, command=command): + result = await self.bash_tool.match_rule( + pattern, + {"command": command}, + ) + self.assertEqual(result, expected) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.bash_tool = None + + +@unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", +) +class BashParserInjectionRiskTest(IsolatedAsyncioTestCase): + """Test check_injection_risk() method for detecting dynamic shell + structures.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.parser = BashCommandParser() + + async def test_command_substitution_dollar_paren(self) -> None: + """Test detection of command substitution with $().""" + test_cases = [ + "ls $(pwd)", + "rm $(find . -name '*.tmp')", + "echo $(date)", + "cat $(which python)", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_injection_risk(cmd) + self.assertIsNotNone(result) + self.assertIn("command_substitution", result) + + async def test_command_substitution_backtick(self) -> None: + """Test detection of command substitution with backticks.""" + test_cases = [ + "ls `pwd`", + "rm `find . -name '*.tmp'`", + "echo `date`", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_injection_risk(cmd) + self.assertIsNotNone(result) + self.assertIn("command_substitution", result) + + async def test_process_substitution(self) -> None: + """Test detection of process substitution <().""" + test_cases = [ + "diff <(ls dir1) <(ls dir2)", + "cat <(echo hello)", + "grep pattern <(curl http://example.com)", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_injection_risk(cmd) + self.assertIsNotNone(result) + self.assertIn("process_substitution", result) + + async def test_subshell(self) -> None: + """Test detection of subshells ().""" + test_cases = [ + "(cd /tmp && ls)", + "(export VAR=value; echo $VAR)", + "echo before && (cd /tmp; pwd) && echo after", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_injection_risk(cmd) + self.assertIsNotNone(result) + self.assertIn("subshell", result) + + async def test_for_loop(self) -> None: + """Test detection of for loops.""" + test_cases = [ + "for f in *.txt; do cat $f; done", + "for i in 1 2 3; do echo $i; done", + "for file in $(ls); do rm $file; done", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_injection_risk(cmd) + self.assertIsNotNone(result) + self.assertIn("for_statement", result) + + async def test_while_loop(self) -> None: + """Test detection of while loops.""" + test_cases = [ + "while read line; do echo $line; done < file.txt", + "while true; do sleep 1; done", + "while [ $i -lt 10 ]; do echo $i; i=$((i+1)); done", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_injection_risk(cmd) + self.assertIsNotNone(result) + self.assertIn("while_statement", result) + + async def test_if_statement(self) -> None: + """Test detection of if statements.""" + test_cases = [ + "if [ -f file.txt ]; then cat file.txt; fi", + "if test -d /tmp; then echo exists; fi", + "if [ $? -eq 0 ]; then echo success; else echo fail; fi", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_injection_risk(cmd) + self.assertIsNotNone(result) + self.assertIn("if_statement", result) + + async def test_case_statement(self) -> None: + """Test detection of case statements.""" + test_cases = [ + "case $1 in start) echo starting;; stop) echo stopping;; esac", + "case $var in a) echo A;; b) echo B;; *) echo other;; esac", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_injection_risk(cmd) + self.assertIsNotNone(result) + self.assertIn("case_statement", result) + + async def test_function_definition(self) -> None: + """Test detection of function definitions.""" + test_cases = [ + "function myfunc() { echo hello; }", + "myfunc() { ls -la; }", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_injection_risk(cmd) + self.assertIsNotNone(result) + self.assertIn("function_definition", result) + + async def test_safe_commands(self) -> None: + """Test that safe commands pass injection check.""" + safe_commands = [ + "ls -la", + "cat file.txt", + "git status", + "npm install", + "echo 'hello world'", + "grep pattern file.txt", + "find . -name '*.py'", + "docker ps", + "python script.py", + ] + for cmd in safe_commands: + with self.subTest(cmd=cmd): + result = self.parser.check_injection_risk(cmd) + self.assertIsNone(result, f"Expected '{cmd}' to be safe") + + async def test_compound_commands_safe(self) -> None: + """Test that compound commands without dynamic structures are safe.""" + safe_commands = [ + "ls -la && cat file.txt", + "git add . && git commit -m 'fix'", + "npm install || echo failed", + "cd /tmp; ls; pwd", + "cat file.txt | grep pattern", + ] + for cmd in safe_commands: + with self.subTest(cmd=cmd): + result = self.parser.check_injection_risk(cmd) + self.assertIsNone(result, f"Expected '{cmd}' to be safe") + + async def test_mixed_safe_and_unsafe(self) -> None: + """Test compound commands with both safe and unsafe parts.""" + test_cases = [ + ("ls && rm $(find . -name '*.tmp')", "command_substitution"), + ("git status && (cd /tmp; ls)", "subshell"), + ( + "echo start && for f in *.txt; do cat $f; done", + "for_statement", + ), + ] + for cmd, expected_type in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_injection_risk(cmd) + self.assertIsNotNone(result) + self.assertIn(expected_type, result) + + async def test_empty_command(self) -> None: + """Test empty command.""" + result = self.parser.check_injection_risk("") + # Empty command should be safe (no dangerous nodes) + self.assertIsNone(result) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.parser = None + + +@unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", +) +class BashParserDangerousCommandTest(IsolatedAsyncioTestCase): + """Test check_dangerous_command() method.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.parser = BashCommandParser() + + async def test_rm_rf_pattern(self) -> None: + """Test detection of rm -rf pattern.""" + test_cases = [ + "rm -rf /tmp/test", + "rm -rf .", + "sudo rm -rf /var/log", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_dangerous_command(cmd) + self.assertIsNotNone(result) + self.assertIn("rm -rf", result) + + async def test_sudo_rm_pattern(self) -> None: + """Test detection of sudo rm pattern.""" + test_cases = [ + "sudo rm file.txt", + "sudo rm -f /etc/config", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_dangerous_command(cmd) + self.assertIsNotNone(result) + self.assertIn("sudo rm", result) + + async def test_dd_command(self) -> None: + """Test detection of dd command.""" + test_cases = [ + "dd if=/dev/zero of=/dev/sda", + "dd if=file.iso of=/dev/sdb", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_dangerous_command(cmd) + self.assertIsNotNone(result) + self.assertEqual(result, "dd") + + async def test_chmod_777_pattern(self) -> None: + """Test detection of chmod 777 pattern.""" + test_cases = [ + "chmod 777 file.txt", + "chmod -R 777 /var/www", + ] + for cmd in test_cases: + with self.subTest(cmd=cmd): + result = self.parser.check_dangerous_command(cmd) + self.assertIsNotNone(result) + self.assertIn("chmod", result) + + async def test_safe_commands(self) -> None: + """Test that safe commands are not flagged.""" + safe_commands = [ + "rm file.txt", + "rm -f temp.log", + "chmod +x script.sh", + "chmod 644 file.txt", + "ls -la", + "git status", + ] + for cmd in safe_commands: + with self.subTest(cmd=cmd): + result = self.parser.check_dangerous_command(cmd) + self.assertIsNone(result, f"Expected '{cmd}' to be safe") + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.parser = None diff --git a/tests/permission_engine_test.py b/tests/permission_engine_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ccbc765ea01e8d41afe59405d300fb2dcd78db56 --- /dev/null +++ b/tests/permission_engine_test.py @@ -0,0 +1,813 @@ +# -*- coding: utf-8 -*- +"""Test cases for PermissionEngine. + +Mode-specific tests live in :mod:`tests.permission_mode_test`. This file +covers rule priority, rule pattern matching (Bash / file glob), dangerous +path detection, suggestion generation, Bash read-only command analysis, +and bypass-immune safety checks. +""" +import sys +import unittest +from unittest.async_case import IsolatedAsyncioTestCase + +from agentscope.permission import ( + PermissionEngine, + PermissionMode, + PermissionContext, + PermissionRule, + PermissionBehavior, + PermissionDecision, + AdditionalWorkingDirectory, +) +from agentscope.tool import ( + Bash, + Write, + Read, + Edit, + ToolBase, +) + + +class PermissionEngineRulePriorityTest(IsolatedAsyncioTestCase): + """Test cases for rule priority and decision-making.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.context = PermissionContext(mode=PermissionMode.DEFAULT) + self.engine = PermissionEngine(self.context) + + async def test_deny_rule_priority(self) -> None: + """Test that deny rules have the highest priority.""" + # Add both allow and deny rules for the same tool + self.engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="git:*", + behavior=PermissionBehavior.ALLOW, + source="test", + ), + ) + self.engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="git:*", + behavior=PermissionBehavior.DENY, + source="test", + ), + ) + + decision = await self.engine.check_permission( + Bash(), + {"command": "git status"}, + ) + + # Deny should take precedence over allow + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + async def test_ask_rule_priority(self) -> None: + """Test that ask rules have priority over allow rules.""" + # Add both allow and ask rules for the same tool + self.engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="npm:*", + behavior=PermissionBehavior.ALLOW, + source="test", + ), + ) + self.engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="npm:*", + behavior=PermissionBehavior.ASK, + source="test", + ), + ) + + decision = await self.engine.check_permission( + Bash(), + {"command": "npm install"}, + ) + + # Ask should take precedence over allow + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def test_rule_priority_order(self) -> None: + """Test complete rule priority: deny > ask > allow.""" + # Add all three types of rules + self.engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="test:*", + behavior=PermissionBehavior.ALLOW, + source="test", + ), + ) + self.engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="test:*", + behavior=PermissionBehavior.ASK, + source="test", + ), + ) + self.engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="test:*", + behavior=PermissionBehavior.DENY, + source="test", + ), + ) + + decision = await self.engine.check_permission( + Bash(), + {"command": "test command"}, + ) + + # Deny should win + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.engine = None + self.context = None + + +@unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", +) +class PermissionEngineBashRuleTest(IsolatedAsyncioTestCase): + """Test cases for Bash command rule matching.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.context = PermissionContext(mode=PermissionMode.DEFAULT) + self.engine = PermissionEngine(self.context) + + async def test_bash_prefix_pattern_matching(self) -> None: + """Test bash command prefix pattern matching with :* wildcard.""" + self.engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="git:*", + behavior=PermissionBehavior.ALLOW, + source="test", + ), + ) + + # Test exact command + decision = await self.engine.check_permission( + Bash(), + {"command": "git"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + # Test command with arguments + decision = await self.engine.check_permission( + Bash(), + {"command": "git status"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + # Test command with multiple arguments + decision = await self.engine.check_permission( + Bash(), + {"command": "git add ."}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + # Test non-matching command + decision = await self.engine.check_permission( + Bash(), + {"command": "npm install"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def test_bash_substring_pattern_matching(self) -> None: + """Test bash command substring pattern matching.""" + self.engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="install", + behavior=PermissionBehavior.DENY, + source="test", + ), + ) + + # Test command containing substring + decision = await self.engine.check_permission( + Bash(), + {"command": "npm install package"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + # Test command containing substring in different position + decision = await self.engine.check_permission( + Bash(), + {"command": "pip install requests"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + async def test_bash_multiple_rules(self) -> None: + """Test bash command matching with multiple rules.""" + self.engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="rm:*", + behavior=PermissionBehavior.DENY, + source="test", + ), + ) + self.engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="git:*", + behavior=PermissionBehavior.ALLOW, + source="test", + ), + ) + + # Test deny rule + decision = await self.engine.check_permission( + Bash(), + {"command": "rm -rf /tmp"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + # Test allow rule + decision = await self.engine.check_permission( + Bash(), + {"command": "git status"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + # Test no matching rule + decision = await self.engine.check_permission( + Bash(), + {"command": "npm install"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.engine = None + self.context = None + + +class PermissionEngineFileRuleTest(IsolatedAsyncioTestCase): + """Test cases for file operation rule matching.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.context = PermissionContext(mode=PermissionMode.DEFAULT) + self.engine = PermissionEngine(self.context) + + async def test_file_glob_pattern_matching(self) -> None: + """Test file path glob pattern matching.""" + self.engine.add_rule( + PermissionRule( + tool_name="Read", + rule_content="*.py", + behavior=PermissionBehavior.ALLOW, + source="test", + ), + ) + + # Test matching file + decision = await self.engine.check_permission( + Read(), + {"file_path": "test.py"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + # Test non-matching file + decision = await self.engine.check_permission( + Read(), + {"file_path": "test.txt"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def test_file_directory_pattern_matching(self) -> None: + """Test file path directory pattern matching.""" + self.engine.add_rule( + PermissionRule( + tool_name="Edit", + rule_content="/tmp/**", + behavior=PermissionBehavior.ALLOW, + source="test", + ), + ) + + # Test file in directory - glob pattern /tmp/** should match + decision = await self.engine.check_permission( + Edit(), + {"file_path": "/tmp/test.txt"}, + ) + # Note: The current implementation uses fnmatch/pathlib matching + # /tmp/** may not match /tmp/test.txt depending on implementation + # This test verifies the actual behavior + self.assertIn( + decision.behavior, + [PermissionBehavior.ALLOW, PermissionBehavior.ASK], + ) + + # Test file in subdirectory + decision = await self.engine.check_permission( + Edit(), + {"file_path": "/tmp/subdir/test.txt"}, + ) + self.assertIn( + decision.behavior, + [PermissionBehavior.ALLOW, PermissionBehavior.ASK], + ) + + # Test file outside directory + decision = await self.engine.check_permission( + Edit(), + {"file_path": "/home/user/test.txt"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.engine = None + self.context = None + + +@unittest.skipIf( + sys.platform == "win32", + "Unix-specific paths not supported on Windows", +) +class PermissionEngineDangerousPathTest(IsolatedAsyncioTestCase): + """Test cases for dangerous path detection (bypass-immune safety + checks).""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.context = PermissionContext(mode=PermissionMode.DEFAULT) + self.engine = PermissionEngine(self.context) + + async def test_dangerous_file_blocks_write(self) -> None: + """Test that Write operations on dangerous files require + confirmation.""" + # Try to write a dangerous file + decision = await self.engine.check_permission( + Write(), + {"file_path": "/home/user/.bashrc"}, + ) + + # Should ask for confirmation (safety check) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def test_dangerous_file_blocks_edit(self) -> None: + """Test that Edit operations on dangerous files require + confirmation.""" + # Try to edit a dangerous file + decision = await self.engine.check_permission( + Edit(), + {"file_path": "/home/user/.gitconfig"}, + ) + + # Should ask for confirmation (safety check) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + self.assertIn("safety", decision.decision_reason.lower()) + + async def test_dangerous_directory_blocks_write(self) -> None: + """Test that Write operations in dangerous directories require + confirmation.""" + # Try to write in a dangerous directory + decision = await self.engine.check_permission( + Write(), + {"file_path": "/home/user/.ssh/config"}, + ) + + # Should ask for confirmation (safety check) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + self.assertIn("safety", decision.decision_reason.lower()) + + async def test_dangerous_path_in_bash_command(self) -> None: + """Test that Bash commands on dangerous paths require confirmation.""" + # Try to remove a dangerous file + decision = await self.engine.check_permission( + Bash(), + {"command": "rm /home/user/.bashrc"}, + ) + + # Should ask for confirmation (safety check) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + self.assertIn("safety", decision.decision_reason.lower()) + + async def test_dangerous_path_in_accept_edits_mode(self) -> None: + """Test that dangerous paths require confirmation even in + ACCEPT_EDITS mode.""" + context = PermissionContext( + mode=PermissionMode.ACCEPT_EDITS, + working_directories={ + "/home/user": AdditionalWorkingDirectory( + path="/home/user", + source="test", + ), + }, + ) + engine = PermissionEngine(context) + + # Try to write a dangerous file in working directory + decision = await engine.check_permission( + Write(), + {"file_path": "/home/user/.bashrc"}, + ) + + # Should ask despite being in working directory + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def test_safe_file_allows_write(self) -> None: + """Test that Write operations on safe files work normally.""" + # Try to write a safe file + decision = await self.engine.check_permission( + Write(), + {"file_path": "/home/user/test.py"}, + ) + + # Should ask (no allow rule, but not blocked by safety check) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + # Should NOT be a safety check + self.assertNotIn("safety", decision.decision_reason.lower()) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.engine = None + self.context = None + + +class PermissionEngineSuggestionTest(IsolatedAsyncioTestCase): + """Test cases for permission rule suggestions.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.context = PermissionContext(mode=PermissionMode.DEFAULT) + self.engine = PermissionEngine(self.context) + + async def test_bash_suggestions(self) -> None: + """Test suggestion generation for bash commands.""" + # Use a non-read-only command to test suggestions + decision = await self.engine.check_permission( + Bash(), + {"command": "git commit -m 'test'"}, + ) + + # Should ask (not read-only) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + # Should have suggestions + self.assertIsNotNone(decision.suggested_rules) + self.assertGreater(len(decision.suggested_rules), 0) + + # Check suggestion content - should generate two-word prefix rule + suggestion_contents = [ + s.rule_content for s in decision.suggested_rules + ] + # git commit should generate "git commit:*" + self.assertIn("git commit:*", suggestion_contents) + + async def test_file_suggestions(self) -> None: + """Test suggestion generation for file operations.""" + decision = await self.engine.check_permission( + Read(), + {"file_path": "/tmp/test.py"}, + ) + + # Should have suggestions + self.assertGreater(len(decision.suggested_rules), 0) + + # Check suggestion content - should generate directory rule + suggestion_contents = [ + s.rule_content for s in decision.suggested_rules + ] + self.assertIn("/tmp/**", suggestion_contents) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.engine = None + self.context = None + + +@unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", +) +class PermissionEngineReadOnlyTest(IsolatedAsyncioTestCase): + """Test cases for read-only command auto-allow.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.context = PermissionContext(mode=PermissionMode.DEFAULT) + self.engine = PermissionEngine(self.context) + + async def test_git_status_is_read_only(self) -> None: + """Test that git status is auto-allowed as read-only.""" + decision = await self.engine.check_permission( + Bash(), + {"command": "git status"}, + ) + + # Should be allowed (read-only command) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + self.assertIn("read-only", decision.decision_reason.lower()) + + async def test_ls_is_read_only(self) -> None: + """Test that ls is auto-allowed as read-only.""" + decision = await self.engine.check_permission( + Bash(), + {"command": "ls -la"}, + ) + + # Should be allowed (read-only command) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + async def test_cat_is_read_only(self) -> None: + """Test that cat is auto-allowed as read-only.""" + decision = await self.engine.check_permission( + Bash(), + {"command": "cat file.txt"}, + ) + + # Should be allowed (read-only command) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + async def test_git_commit_is_not_read_only(self) -> None: + """Test that git commit is not auto-allowed (not read-only).""" + decision = await self.engine.check_permission( + Bash(), + {"command": "git commit -m 'test'"}, + ) + + # Should ask (not read-only) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def test_compound_command_with_dangerous_path(self) -> None: + """Test compound command with read-only and dangerous path o + perations.""" + # ls is read-only, but rm ~/.bashrc is dangerous + decision = await self.engine.check_permission( + Bash(), + {"command": "ls -la && rm ~/.bashrc"}, + ) + + # Should ask because of dangerous path (safety check is bypass-immune) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + self.assertIn("safety", decision.decision_reason.lower()) + + async def test_compound_command_all_read_only(self) -> None: + """Test compound command with all read-only operations.""" + # Both ls and cat are read-only + decision = await self.engine.check_permission( + Bash(), + {"command": "ls -la && cat file.txt"}, + ) + + # Should be allowed (all read-only) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + async def test_compound_command_with_write_operation(self) -> None: + """Test compound command with read-only and write operations.""" + # ls is read-only, but git commit is not + decision = await self.engine.check_permission( + Bash(), + {"command": "ls -la && git commit -m 'test'"}, + ) + + # Should ask (contains non-read-only command) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def test_output_redirection_to_dangerous_path(self) -> None: + """Test output redirection to dangerous path.""" + # cat is read-only, but redirecting to ~/.bashrc is dangerous + decision = await self.engine.check_permission( + Bash(), + {"command": "cat file.txt > ~/.bashrc"}, + ) + + # Should ask because of dangerous path in redirection + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + self.assertIn("safety", decision.decision_reason.lower()) + + async def test_output_redirection_to_safe_path(self) -> None: + """Test output redirection to safe path.""" + # cat is read-only, redirecting to safe path + decision = await self.engine.check_permission( + Bash(), + {"command": "cat file.txt > /tmp/output.txt"}, + ) + + # Should ask (redirection is not considered read-only) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.engine = None + self.context = None + + +@unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", +) +class PermissionEngineSafetyCheckAllowRuleImmuneTest( + IsolatedAsyncioTestCase, +): + """Safety checks are immune to allow rules in DEFAULT mode. + + A user-configured allow rule cannot grant permission for an + operation that the tool has flagged as a safety concern + (``bypass_immune=True``). This guarantees that broad rules like + ``Bash:rm:*`` don't accidentally authorize ``rm -rf /``. + + Note: in BYPASS mode the user has explicitly opted out of safety + enforcement, so these tools' safety ASKs are NOT honored — that + behavior is covered by + :class:`tests.permission_mode_test.PermissionEngineBypassModeTest`. + """ + + async def test_injection_check_not_bypassed_by_allow_rule(self) -> None: + """Allow rule for ``ls:*`` does not authorize ``ls $(rm -rf /)``.""" + context = PermissionContext(mode=PermissionMode.DEFAULT) + engine = PermissionEngine(context) + engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="ls:*", + behavior=PermissionBehavior.ALLOW, + source="test", + ), + ) + + decision = await engine.check_permission( + Bash(), + {"command": "ls $(rm -rf /)"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + self.assertIn("command_substitution", decision.message) + + async def test_dangerous_removal_not_bypassed_by_allow_rule(self) -> None: + """Allow rule for ``rm:*`` does not authorize ``rm -rf /``.""" + context = PermissionContext(mode=PermissionMode.DEFAULT) + engine = PermissionEngine(context) + engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="rm:*", + behavior=PermissionBehavior.ALLOW, + source="test", + ), + ) + + decision = await engine.check_permission( + Bash(), + {"command": "rm -rf /"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + # Either dangerous removal path or dangerous command pattern fires + self.assertTrue( + "Dangerous removal operation" in decision.message + or "dangerous pattern" in decision.message, + ) + + +# --------------------------------------------------------------------------- +# bypass_immune field mechanism +# --------------------------------------------------------------------------- + + +class _FakeToolReturningAsk(ToolBase): + """Minimal tool used to test the ``bypass_immune`` field mechanism. + + Returns a single configurable ASK decision from ``check_permissions`` + so we can verify the engine routes bypass-immune vs regular ASKs + correctly regardless of the specific safety trigger. + """ + + name = "FakeAskTool" + description = "Test-only tool" + input_schema = {"type": "object", "properties": {}} + is_concurrency_safe = True + is_read_only = False + is_external_tool = False + is_state_injected = False + is_mcp = False + + def __init__(self, bypass_immune: bool) -> None: + self._bypass_immune = bypass_immune + + async def check_permissions( + self, + tool_input: dict, + context: PermissionContext, + ) -> PermissionDecision: + return PermissionDecision( + behavior=PermissionBehavior.ASK, + message="fake tool requires confirmation", + decision_reason="test-only reason without the s-word", + bypass_immune=self._bypass_immune, + ) + + +class PermissionEngineBypassImmuneFieldTest(IsolatedAsyncioTestCase): + """Tests for the :attr:`PermissionDecision.bypass_immune` field. + + These tests use a minimal fake tool that returns ASK with a + configurable ``bypass_immune`` value — independent of any specific + safety trigger — to verify the engine's mechanism for honoring the + field. Together with + :class:`PermissionEngineSafetyCheckAllowRuleImmuneTest` (which tests + that Bash/Write/Edit's individual safety checks correctly set the + field), this fully covers the bypass-immune contract. + """ + + async def test_bypass_immune_ask_not_overridden_by_allow_rule( + self, + ) -> None: + """An ASK with ``bypass_immune=True`` survives a matching allow + rule that would otherwise grant the operation.""" + context = PermissionContext(mode=PermissionMode.DEFAULT) + engine = PermissionEngine(context) + engine.add_rule( + PermissionRule( + tool_name="FakeAskTool", + rule_content=None, + behavior=PermissionBehavior.ALLOW, + source="test", + ), + ) + + decision = await engine.check_permission( + _FakeToolReturningAsk(bypass_immune=True), + {}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def test_non_bypass_immune_ask_overridden_by_allow_rule( + self, + ) -> None: + """A regular ASK (``bypass_immune=False``) yields to a matching + allow rule — proving the field, not the ``decision_reason`` + text, governs immunity.""" + context = PermissionContext(mode=PermissionMode.DEFAULT) + engine = PermissionEngine(context) + engine.add_rule( + PermissionRule( + tool_name="FakeAskTool", + rule_content=None, + behavior=PermissionBehavior.ALLOW, + source="test", + ), + ) + + decision = await engine.check_permission( + _FakeToolReturningAsk(bypass_immune=False), + {}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + async def test_bypass_mode_ignores_bypass_immune_field(self) -> None: + """BYPASS is intentionally permissive: even an ASK with + ``bypass_immune=True`` is allowed through. The user has opted + out of safety prompts; ``bypass_immune`` only governs allow + rules (in DEFAULT) and DONT_ASK conversion. + """ + context = PermissionContext(mode=PermissionMode.BYPASS) + engine = PermissionEngine(context) + + decision = await engine.check_permission( + _FakeToolReturningAsk(bypass_immune=True), + {}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + async def test_bypass_mode_allows_non_bypass_immune_ask(self) -> None: + """For completeness: a non-bypass-immune ASK is also allowed + through in BYPASS — same outcome, different path.""" + context = PermissionContext(mode=PermissionMode.BYPASS) + engine = PermissionEngine(context) + + decision = await engine.check_permission( + _FakeToolReturningAsk(bypass_immune=False), + {}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + async def test_bypass_immune_ask_converted_to_deny_in_dont_ask( + self, + ) -> None: + """Under DONT_ASK, even a bypass-immune ASK is converted to + DENY (issue #3 contract): no user is available to confirm, + so the only safe action is to refuse.""" + context = PermissionContext(mode=PermissionMode.DONT_ASK) + engine = PermissionEngine(context) + + decision = await engine.check_permission( + _FakeToolReturningAsk(bypass_immune=True), + {}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) diff --git a/tests/permission_mode_test.py b/tests/permission_mode_test.py new file mode 100644 index 0000000000000000000000000000000000000000..fe8ed3afef669d24c9bf63cdfa5228e8cdf042f3 --- /dev/null +++ b/tests/permission_mode_test.py @@ -0,0 +1,895 @@ +# -*- coding: utf-8 -*- +"""Per-mode test cases for ``PermissionEngine``. + +Each :class:`PermissionMode` has its own test class so the policy of +that mode can be verified in isolation. Tests cover: + +- the three rule layers (deny / ask / allow) for that mode +- the ``tool.check_permissions`` return paths (ALLOW / DENY / safety-ASK + / PASSTHROUGH) +- mode-specific behavior (e.g. EXPLORE's read-only resolution, + ACCEPT_EDITS's working-directory auto-allow, BYPASS's safety-ASK + immunity, DONT_ASK's default DENY) +- Bash dynamic read-only / non-read-only / dangerous commands where + relevant + +Specific safety-check triggers (injection, dangerous removal, sed +constraints, dangerous config paths) are tested separately in +``permission_engine_test.py::PermissionEngineSafetyCheckAllowRuleImmuneTest``. +""" +import os +import sys +import tempfile +import unittest +from unittest.async_case import IsolatedAsyncioTestCase + +from agentscope.permission import ( + PermissionEngine, + PermissionMode, + PermissionContext, + PermissionRule, + PermissionBehavior, + AdditionalWorkingDirectory, +) +from agentscope.tool import ( + Bash, + Write, + Read, + Edit, +) + + +# --------------------------------------------------------------------------- +# DEFAULT mode +# --------------------------------------------------------------------------- + + +class PermissionEngineDefaultModeTest(IsolatedAsyncioTestCase): + """Tests for :attr:`PermissionMode.DEFAULT`. + + DEFAULT is the most restrictive non-DONT_ASK mode: every operation + requires explicit permission unless an allow rule matches or the + tool itself returns ALLOW (e.g. Bash auto-allows known read-only + commands). + """ + + async def asyncSetUp(self) -> None: + self.context = PermissionContext(mode=PermissionMode.DEFAULT) + self.engine = PermissionEngine(self.context) + + async def test_default_write_with_no_rules_returns_ask(self) -> None: + """Write to a safe path with no matching rules falls to the + engine's default ASK.""" + decision = await self.engine.check_permission( + Write(), + {"file_path": "/tmp/file.txt"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def test_default_deny_rule_returns_deny(self) -> None: + """Deny rule has the highest priority.""" + self.engine.add_rule( + PermissionRule( + tool_name="Write", + rule_content="*.env", + behavior=PermissionBehavior.DENY, + source="test", + ), + ) + decision = await self.engine.check_permission( + Write(), + {"file_path": "/tmp/secret.env"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + async def test_default_ask_rule_returns_ask_with_suggestions( + self, + ) -> None: + """Ask rule short-circuits before the tool's own check and + attaches ``suggested_rules``.""" + self.engine.add_rule( + PermissionRule( + tool_name="Write", + rule_content="*.py", + behavior=PermissionBehavior.ASK, + source="test", + ), + ) + decision = await self.engine.check_permission( + Write(), + {"file_path": "/tmp/main.py"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + self.assertIsNotNone(decision.suggested_rules) + self.assertGreater(len(decision.suggested_rules), 0) + + async def test_default_allow_rule_returns_allow(self) -> None: + """Allow rule grants permission when no deny/ask rule and no + safety ASK applies.""" + self.engine.add_rule( + PermissionRule( + tool_name="Write", + rule_content="/tmp/**", + behavior=PermissionBehavior.ALLOW, + source="test", + ), + ) + decision = await self.engine.check_permission( + Write(), + {"file_path": "/tmp/file.txt"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_default_bash_read_only_command_auto_allows(self) -> None: + """Bash returns ALLOW for known read-only commands, which the + engine surfaces directly in DEFAULT mode.""" + for command in ("ls -a", "pwd", "git status", "cat README.md"): + decision = await self.engine.check_permission( + Bash(), + {"command": command}, + ) + self.assertEqual( + decision.behavior, + PermissionBehavior.ALLOW, + f"Expected ALLOW for read-only bash command: {command}", + ) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_default_bash_modification_command_returns_ask( + self, + ) -> None: + """Bash modification commands (no matching rule) fall through to + the engine's default ASK.""" + decision = await self.engine.check_permission( + Bash(), + {"command": "npm install"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def test_default_dangerous_path_safety_ask(self) -> None: + """Write to a dangerous path produces a safety ASK from the tool + itself (bypass-immune).""" + decision = await self.engine.check_permission( + Write(), + {"file_path": "/home/user/.bashrc"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + self.assertIn("safety", (decision.decision_reason or "").lower()) + + async def test_default_safety_ask_not_overridden_by_allow_rule( + self, + ) -> None: + """A user-configured allow rule must not override a tool's + safety ASK (bypass-immune).""" + self.engine.add_rule( + PermissionRule( + tool_name="Write", + rule_content="**", + behavior=PermissionBehavior.ALLOW, + source="test", + ), + ) + decision = await self.engine.check_permission( + Write(), + {"file_path": "/home/user/.bashrc"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + +# --------------------------------------------------------------------------- +# EXPLORE mode +# --------------------------------------------------------------------------- + + +class PermissionEngineExploreModeTest(IsolatedAsyncioTestCase): + """Tests for :attr:`PermissionMode.EXPLORE`. + + EXPLORE is the read-only mode: an invocation is ALLOWed iff + :meth:`ToolBase.check_read_only` returns True; everything else is + DENIed outright. ``tool.check_permissions`` is intentionally not + consulted (the read-only verdict is final), and allow rules cannot + grant write access — EXPLORE's read-only invariant is non-negotiable. + """ + + async def asyncSetUp(self) -> None: + self.context = PermissionContext(mode=PermissionMode.EXPLORE) + self.engine = PermissionEngine(self.context) + + async def test_explore_read_tool_allows(self) -> None: + """Statically read-only tools (Read/Glob/Grep) → ALLOW.""" + decision = await self.engine.check_permission( + Read(), + {"file_path": "/tmp/file.txt"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + async def test_explore_write_tool_denies(self) -> None: + """Statically non-read-only tools (Write/Edit) → DENY.""" + decision = await self.engine.check_permission( + Write(), + {"file_path": "/tmp/file.txt"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + async def test_explore_deny_rule_returns_deny(self) -> None: + """Deny rule still has top priority in EXPLORE.""" + self.engine.add_rule( + PermissionRule( + tool_name="Read", + rule_content="/secret/**", + behavior=PermissionBehavior.DENY, + source="test", + ), + ) + decision = await self.engine.check_permission( + Read(), + {"file_path": "/secret/key.txt"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + async def test_explore_ask_rule_on_read_only_tool_returns_ask( + self, + ) -> None: + """Ask rules apply before the read-only fast-path, so a read-only + tool with a matching ask rule still surfaces as ASK. + + Documents current behavior — this is the design point tracked + as issue #7. + """ + self.engine.add_rule( + PermissionRule( + tool_name="Read", + rule_content="**/*.env", + behavior=PermissionBehavior.ASK, + source="test", + ), + ) + decision = await self.engine.check_permission( + Read(), + {"file_path": "/tmp/secret.env"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def test_explore_allow_rule_does_not_override_deny(self) -> None: + """EXPLORE's deny on non-read-only tools is the invariant; an + allow rule must not be able to grant write access.""" + self.engine.add_rule( + PermissionRule( + tool_name="Write", + rule_content="**", + behavior=PermissionBehavior.ALLOW, + source="test", + ), + ) + decision = await self.engine.check_permission( + Write(), + {"file_path": "/tmp/file.txt"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + async def test_explore_dangerous_path_returns_deny_not_safety_ask( + self, + ) -> None: + """EXPLORE never invokes ``tool.check_permissions``, so a dangerous + path on a write tool surfaces as DENY (the read-only verdict) + rather than a safety ASK. DENY is strictly stronger than ASK, so + this is consistent with EXPLORE's "no writes" guarantee. + """ + decision = await self.engine.check_permission( + Write(), + {"file_path": "/home/user/.bashrc"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_explore_bash_read_only_command_allows(self) -> None: + """EXPLORE allows read-only bash commands via + :meth:`Bash.check_read_only` (regression for issue #1).""" + for command in ("ls -a", "pwd", "git status", "cat README.md"): + decision = await self.engine.check_permission( + Bash(), + {"command": command}, + ) + self.assertEqual( + decision.behavior, + PermissionBehavior.ALLOW, + f"Expected ALLOW for read-only bash command: {command}", + ) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_explore_bash_modification_command_denies(self) -> None: + """EXPLORE denies bash commands that are not statically + recognized as read-only.""" + for command in ("cp a b", "mv a b", "touch /tmp/x"): + decision = await self.engine.check_permission( + Bash(), + {"command": command}, + ) + self.assertEqual( + decision.behavior, + PermissionBehavior.DENY, + f"Expected DENY for non-read-only bash command: {command}", + ) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_explore_bash_dangerous_command_denies(self) -> None: + """EXPLORE denies dangerous commands directly — in DEFAULT mode + the same command would surface as a safety ASK, but EXPLORE is + stricter.""" + decision = await self.engine.check_permission( + Bash(), + {"command": "rm -rf /"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + +# --------------------------------------------------------------------------- +# ACCEPT_EDITS mode +# --------------------------------------------------------------------------- + + +class PermissionEngineAcceptEditsModeTest(IsolatedAsyncioTestCase): + """Tests for :attr:`PermissionMode.ACCEPT_EDITS`. + + ACCEPT_EDITS auto-allows file edits within configured working + directories; reads are auto-allowed unconditionally; other + operations follow the normal DEFAULT-like flow. + """ + + async def asyncSetUp(self) -> None: + self.context = PermissionContext( + mode=PermissionMode.ACCEPT_EDITS, + working_directories={ + "/tmp/project": AdditionalWorkingDirectory( + path="/tmp/project", + source="test", + ), + }, + ) + self.engine = PermissionEngine(self.context) + + async def test_accept_edits_within_working_directory(self) -> None: + """Write / Read / Edit within a working directory → ALLOW.""" + for tool in (Write(), Read(), Edit()): + decision = await self.engine.check_permission( + tool, + {"file_path": "/tmp/project/file.txt"}, + ) + self.assertEqual( + decision.behavior, + PermissionBehavior.ALLOW, + f"Expected ALLOW for {tool.name} in working directory", + ) + + @unittest.skipIf( + os.name == "nt", + "os.symlink typically requires admin privileges on Windows", + ) + async def test_accept_edits_resolves_symlinked_working_directory( + self, + ) -> None: + """Working directory comparison must use ``realpath`` so a path + reached through a symlink (e.g. macOS's ``/tmp`` -> + ``/private/tmp``) is recognized. Regression test for the + ``abspath`` → ``realpath`` fix in + :meth:`_path_in_allowed_working_path`. + """ + parent = tempfile.mkdtemp() + try: + real_dir = os.path.join(parent, "real") + os.makedirs(real_dir) + link_dir = os.path.join(parent, "link") + os.symlink(real_dir, link_dir) + + # Case 1: working_dir given as real path, file via link + context = PermissionContext( + mode=PermissionMode.ACCEPT_EDITS, + working_directories={ + real_dir: AdditionalWorkingDirectory( + path=real_dir, + source="test", + ), + }, + ) + engine = PermissionEngine(context) + decision = await engine.check_permission( + Write(), + {"file_path": os.path.join(link_dir, "file.txt")}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + # Case 2: working_dir given as link, file via real path + context = PermissionContext( + mode=PermissionMode.ACCEPT_EDITS, + working_directories={ + link_dir: AdditionalWorkingDirectory( + path=link_dir, + source="test", + ), + }, + ) + engine = PermissionEngine(context) + decision = await engine.check_permission( + Edit(), + {"file_path": os.path.join(real_dir, "file.txt")}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + finally: + import shutil + + shutil.rmtree(parent, ignore_errors=True) + + async def test_accept_edits_outside_working_directory(self) -> None: + """Edit outside a working directory falls to the default ASK.""" + decision = await self.engine.check_permission( + Edit(), + {"file_path": "/home/user/file.txt"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def test_accept_edits_read_operation_auto_allowed(self) -> None: + """Read tool is auto-allowed regardless of path (read-only fast + path).""" + decision = await self.engine.check_permission( + Read(), + {"file_path": "/anywhere/file.txt"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + async def test_accept_edits_deny_rule_returns_deny(self) -> None: + """Deny rule overrides ACCEPT_EDITS's working-directory + auto-allow.""" + self.engine.add_rule( + PermissionRule( + tool_name="Write", + rule_content="**/*.lock", + behavior=PermissionBehavior.DENY, + source="test", + ), + ) + decision = await self.engine.check_permission( + Write(), + {"file_path": "/tmp/project/poetry.lock"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + async def test_accept_edits_ask_rule_returns_ask(self) -> None: + """Ask rule short-circuits before the working-directory check.""" + self.engine.add_rule( + PermissionRule( + tool_name="Edit", + rule_content="**/*.py", + behavior=PermissionBehavior.ASK, + source="test", + ), + ) + decision = await self.engine.check_permission( + Edit(), + {"file_path": "/tmp/project/main.py"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def test_accept_edits_dangerous_path_safety_ask(self) -> None: + """Safety ASK from a dangerous path is bypass-immune even when + the path is inside a working directory. + """ + decision = await self.engine.check_permission( + Write(), + {"file_path": "/home/user/.bashrc"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + self.assertIn("safety", (decision.decision_reason or "").lower()) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_accept_edits_bash_read_only_command_allows(self) -> None: + """Read-only bash commands ALLOW via the read-only fast path.""" + decision = await self.engine.check_permission( + Bash(), + {"command": "ls -a"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_accept_edits_bash_filesystem_command_inside_working_dir( + self, + ) -> None: + """ACCEPT_EDITS auto-allows recognized filesystem commands + (``mkdir``, ``touch``, ``rm``, ``cp``, ``mv``, ``sed``, + ``rmdir``) when **all** target paths are inside the configured + working directories.""" + cases = [ + "touch /tmp/project/new.txt", + "mkdir /tmp/project/newdir", + "rm /tmp/project/old.txt", + "rmdir /tmp/project/olddir", + "cp /tmp/project/a /tmp/project/b", + "mv /tmp/project/a /tmp/project/b", + "sed -i 's/x/y/g' /tmp/project/foo.txt", + ] + for command in cases: + decision = await self.engine.check_permission( + Bash(), + {"command": command}, + ) + self.assertEqual( + decision.behavior, + PermissionBehavior.ALLOW, + f"Expected ALLOW for in-working-dir command: {command}", + ) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_accept_edits_bash_filesystem_command_outside_working_dir( + self, + ) -> None: + """Regression for issue #4: filesystem commands whose targets + escape the working directory must NOT be auto-allowed. + + Without this guard ``cp /etc/hosts /Users/me/other-project/x`` + would silently succeed in ACCEPT_EDITS even though the + equivalent ``Write`` call is correctly denied. The asserted + behavior is ASK (the PASSTHROUGH from Bash + the engine's + default ASK in ACCEPT_EDITS). + """ + cases = [ + # Target entirely outside the working directory + "rm /Users/someone/other/foo", + "touch /Users/someone/other/foo", + "mkdir /Users/someone/other/newdir", + # cp / mv: at least one of (src, dst) outside + "cp /tmp/project/a /Users/someone/other/b", + "cp /Users/someone/other/a /tmp/project/b", + "mv /tmp/project/a /Users/someone/other/b", + # sed in-place modifying a file outside + "sed -i 's/x/y/g' /Users/someone/other/foo.txt", + ] + for command in cases: + decision = await self.engine.check_permission( + Bash(), + {"command": command}, + ) + self.assertEqual( + decision.behavior, + PermissionBehavior.ASK, + f"Expected ASK for outside-working-dir command: {command}", + ) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_accept_edits_bash_filesystem_command_no_args_not_allowed( + self, + ) -> None: + """Conservative behavior: if the parser extracts no target paths + (e.g. a bare ``mkdir`` with no arguments), we do not auto-allow. + The command itself will fail at execution, but it should not + silently pass the permission check.""" + decision = await self.engine.check_permission( + Bash(), + {"command": "mkdir"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + +# --------------------------------------------------------------------------- +# BYPASS mode +# --------------------------------------------------------------------------- + + +class PermissionEngineBypassModeTest(IsolatedAsyncioTestCase): + """Tests for :attr:`PermissionMode.BYPASS`. + + BYPASS is the "fully trusted" mode: the user has opted out of all + safety prompts. The only guardrails left are user-configured deny + / ask rules and tool-emitted DENY. Every bypass-immune safety ASK + (dangerous removal, dangerous paths, sed in-place on sensitive + files, command injection, ...) is intentionally **skipped**. + Use BYPASS only in sandboxed environments or when you fully trust + the agent. + """ + + async def asyncSetUp(self) -> None: + self.context = PermissionContext(mode=PermissionMode.BYPASS) + self.engine = PermissionEngine(self.context) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_bypass_no_rules_allows(self) -> None: + """No rules → ALLOW (BYPASS's default fallback).""" + decision = await self.engine.check_permission( + Bash(), + {"command": "npm install"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_bypass_deny_rule_returns_deny(self) -> None: + """Deny rules are bypass-immune.""" + self.engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="rm:*", + behavior=PermissionBehavior.DENY, + source="test", + ), + ) + decision = await self.engine.check_permission( + Bash(), + {"command": "rm -rf /tmp/foo"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_bypass_ask_rule_returns_ask(self) -> None: + """A user-configured ask rule represents explicit intent to be + prompted; BYPASS must not override it.""" + self.engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="git push:*", + behavior=PermissionBehavior.ASK, + source="test", + ), + ) + decision = await self.engine.check_permission( + Bash(), + {"command": "git push origin main"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ASK) + + async def test_bypass_skips_dangerous_path_safety(self) -> None: + """BYPASS skips the Write tool's dangerous-path safety check — + writing to ``~/.bashrc`` is allowed through.""" + decision = await self.engine.check_permission( + Write(), + {"file_path": "/home/user/.bashrc"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_bypass_skips_dangerous_removal(self) -> None: + """BYPASS skips Bash's dangerous-removal safety check — + ``rm -rf /`` is allowed through.""" + decision = await self.engine.check_permission( + Bash(), + {"command": "rm -rf /"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_bypass_skips_command_injection(self) -> None: + """BYPASS skips Bash's command-injection safety check — + dynamic expansion is allowed through.""" + decision = await self.engine.check_permission( + Bash(), + {"command": "ls $(date +%s)"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_bypass_skips_sed_dangerous_file(self) -> None: + """BYPASS skips Bash's sed-on-dangerous-file safety check.""" + decision = await self.engine.check_permission( + Bash(), + {"command": "sed 's/old/new/e' file.txt"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_bypass_skips_dangerous_config_path_in_bash(self) -> None: + """BYPASS skips Bash's dangerous-config-path safety check — + operating on ``~/.bashrc`` via bash is allowed through.""" + decision = await self.engine.check_permission( + Bash(), + {"command": "rm ~/.bashrc"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_bypass_tool_allow_returns_allow(self) -> None: + """Tool's own ALLOW (e.g. Bash read-only command) is returned + as-is in BYPASS as in any mode.""" + decision = await self.engine.check_permission( + Bash(), + {"command": "ls -a"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + +# --------------------------------------------------------------------------- +# DONT_ASK mode +# --------------------------------------------------------------------------- + + +class PermissionEngineDontAskModeTest(IsolatedAsyncioTestCase): + """Tests for :attr:`PermissionMode.DONT_ASK`. + + DONT_ASK is used when no user is available to answer prompts + (scheduled tasks, background runs). The invariant is "never return + ASK" — every ASK-producing code path (default, ASK rule, safety + ASK) is converted to DENY. + """ + + async def asyncSetUp(self) -> None: + self.context = PermissionContext(mode=PermissionMode.DONT_ASK) + self.engine = PermissionEngine(self.context) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_dont_ask_no_rules_returns_deny(self) -> None: + """No rules → DENY (DONT_ASK's default).""" + decision = await self.engine.check_permission( + Bash(), + {"command": "npm install"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_dont_ask_deny_rule_returns_deny(self) -> None: + """Deny rules apply normally.""" + self.engine.add_rule( + PermissionRule( + tool_name="Bash", + rule_content="rm:*", + behavior=PermissionBehavior.DENY, + source="test", + ), + ) + decision = await self.engine.check_permission( + Bash(), + {"command": "rm -rf /tmp/foo"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + async def test_dont_ask_allow_rule_returns_allow(self) -> None: + """An explicit allow rule still grants permission in DONT_ASK.""" + self.engine.add_rule( + PermissionRule( + tool_name="Write", + rule_content="/tmp/**", + behavior=PermissionBehavior.ALLOW, + source="test", + ), + ) + decision = await self.engine.check_permission( + Write(), + {"file_path": "/tmp/file.txt"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_dont_ask_bash_read_only_command_allows(self) -> None: + """Tool's own ALLOW (Bash read-only command) is still ALLOW + under DONT_ASK — no user prompt is needed.""" + decision = await self.engine.check_permission( + Bash(), + {"command": "ls -a"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.ALLOW) + + async def test_dont_ask_ask_rule_returns_deny(self) -> None: + """An ASK rule hit is converted to DENY (issue #3): the user is + not available to answer the prompt, so the operation cannot + proceed.""" + self.engine.add_rule( + PermissionRule( + tool_name="Write", + rule_content="**/*.py", + behavior=PermissionBehavior.ASK, + source="test", + ), + ) + decision = await self.engine.check_permission( + Write(), + {"file_path": "/tmp/main.py"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + + async def test_dont_ask_ask_rule_conversion_preserves_suggestions( + self, + ) -> None: + """Converted DENY decisions keep the original ASK's + ``suggested_rules`` so callers can surface them to the user + out-of-band (e.g. in a scheduled-task failure report).""" + self.engine.add_rule( + PermissionRule( + tool_name="Write", + rule_content="**/*.py", + behavior=PermissionBehavior.ASK, + source="test", + ), + ) + decision = await self.engine.check_permission( + Write(), + {"file_path": "/tmp/main.py"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + self.assertIsNotNone(decision.suggested_rules) + self.assertGreater(len(decision.suggested_rules), 0) + + async def test_dont_ask_safety_ask_returns_deny(self) -> None: + """A safety ASK from ``tool.check_permissions`` (e.g. Write to a + dangerous path) is converted to DENY (issue #3) — DONT_ASK + respects the safety verdict but cannot ask the user, so the + only safe action is to refuse.""" + decision = await self.engine.check_permission( + Write(), + {"file_path": "/home/user/.bashrc"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) + # Conversion records both the original safety reason and the + # DONT_ASK conversion in the decision_reason. + reason = (decision.decision_reason or "").lower() + self.assertIn("dont_ask", reason) + self.assertIn("safety", reason) + + @unittest.skipIf( + sys.platform == "win32", + "Bash tool is not supported on Windows", + ) + async def test_dont_ask_bash_dangerous_command_returns_deny( + self, + ) -> None: + """``rm -rf /`` triggers a safety ASK from Bash, which DONT_ASK + converts to DENY (issue #3).""" + decision = await self.engine.check_permission( + Bash(), + {"command": "rm -rf /"}, + ) + self.assertEqual(decision.behavior, PermissionBehavior.DENY) diff --git a/tests/rag_chunker_approx_token_test.py b/tests/rag_chunker_approx_token_test.py new file mode 100644 index 0000000000000000000000000000000000000000..71126cdb53829d8590c56d375825a276e0c5d109 --- /dev/null +++ b/tests/rag_chunker_approx_token_test.py @@ -0,0 +1,204 @@ +# -*- coding: utf-8 -*- +"""Unit tests for the ApproxTokenChunker class.""" +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import AnyString + +from agentscope.message import Base64Source, DataBlock, TextBlock +from agentscope.rag import ApproxTokenChunker, Chunk, Section + + +def _dump_chunks(chunks: list[Chunk]) -> list[dict]: + """Convert chunks into plain dicts for whole-structure comparison. + + Args: + chunks (`list[Chunk]`): + The chunks to convert. + + Returns: + `list[dict]`: + The chunks as plain dicts. + """ + return [chunk.model_dump() for chunk in chunks] + + +class ApproxTokenChunkerTest(IsolatedAsyncioTestCase): + """The test cases for the ApproxTokenChunker class.""" + + async def test_short_text_single_chunk(self) -> None: + """Short text should produce a single chunk.""" + chunker = ApproxTokenChunker(chunk_size=512, overlap=50) + sections = [ + Section( + content=TextBlock(text="Hello world!"), + source="a.txt", + metadata={"page": 1}, + ), + ] + + chunks = await chunker.chunk(sections) + + self.assertEqual( + _dump_chunks(chunks), + [ + { + "content": { + "type": "text", + "text": "Hello world!", + "id": AnyString(), + }, + "source": "a.txt", + "chunk_index": 0, + "total_chunks": 1, + "metadata": {"page": 1}, + }, + ], + ) + + async def test_long_text_split_with_overlap(self) -> None: + """Long text should be split into overlapping chunks.""" + # chunk_size=10 tokens -> 40-byte window, overlap=2 -> 8-byte step + # back, so for ASCII text the windows start at 0, 32, 64, ... + chunker = ApproxTokenChunker(chunk_size=10, overlap=2) + text = "abcdefghij" * 20 # 200 ASCII chars => 50 approx tokens + sections = [ + Section( + content=TextBlock(text=text), + source="b.txt", + ), + ] + + chunks = await chunker.chunk(sections) + + self.assertEqual( + _dump_chunks(chunks), + [ + { + "content": { + "type": "text", + "text": text[start : start + 40], + "id": AnyString(), + }, + "source": "b.txt", + "chunk_index": index, + "total_chunks": 6, + "metadata": {}, + } + for index, start in enumerate([0, 32, 64, 96, 128, 160]) + ], + ) + + async def test_data_block_pass_through(self) -> None: + """DataBlock sections should pass through unchanged.""" + chunker = ApproxTokenChunker(chunk_size=10, overlap=2) + data_block = DataBlock( + source=Base64Source(data="aGk=", media_type="image/png"), + ) + sections = [ + Section( + content=TextBlock(text="x" * 100), + source="c.pdf", + ), + Section( + content=data_block, + source="c.pdf", + metadata={"page": 2}, + ), + ] + + chunks = await chunker.chunk(sections) + + self.assertEqual( + _dump_chunks(chunks), + [ + { + "content": { + "type": "text", + "text": "x" * 40, + "id": AnyString(), + }, + "source": "c.pdf", + "chunk_index": 0, + "total_chunks": 4, + "metadata": {}, + }, + { + "content": { + "type": "text", + "text": "x" * 40, + "id": AnyString(), + }, + "source": "c.pdf", + "chunk_index": 1, + "total_chunks": 4, + "metadata": {}, + }, + { + "content": { + "type": "text", + "text": "x" * 36, + "id": AnyString(), + }, + "source": "c.pdf", + "chunk_index": 2, + "total_chunks": 4, + "metadata": {}, + }, + { + "content": { + "type": "data", + "id": AnyString(), + "source": { + "type": "base64", + "data": "aGk=", + "media_type": "image/png", + }, + "name": None, + }, + "source": "c.pdf", + "chunk_index": 3, + "total_chunks": 4, + "metadata": {"page": 2}, + }, + ], + ) + # The DataBlock instance itself is passed through, not copied + self.assertIs(chunks[-1].content, data_block) + + async def test_no_cross_section_merging(self) -> None: + """Chunks must never combine content from two sections.""" + chunker = ApproxTokenChunker(chunk_size=512, overlap=50) + sections = [ + Section(content=TextBlock(text="first"), source="d.txt"), + Section(content=TextBlock(text="second"), source="d.txt"), + ] + + chunks = await chunker.chunk(sections) + + self.assertEqual( + _dump_chunks(chunks), + [ + { + "content": { + "type": "text", + "text": "first", + "id": AnyString(), + }, + "source": "d.txt", + "chunk_index": 0, + "total_chunks": 2, + "metadata": {}, + }, + { + "content": { + "type": "text", + "text": "second", + "id": AnyString(), + }, + "source": "d.txt", + "chunk_index": 1, + "total_chunks": 2, + "metadata": {}, + }, + ], + ) diff --git a/tests/rag_parser_test.py b/tests/rag_parser_test.py new file mode 100644 index 0000000000000000000000000000000000000000..def1a69b1675ff8110b0945e3d41cc46565e798c --- /dev/null +++ b/tests/rag_parser_test.py @@ -0,0 +1,670 @@ +# -*- coding: utf-8 -*- +"""Unit tests for the file parsers in :mod:`agentscope.rag._parser`. + +PDF / PPTX fixtures are produced in-memory via :mod:`reportlab` and +:mod:`python-pptx` so the tests have no on-disk dependencies and can +run anywhere ``agentscope[rag]`` is installed. +""" +import base64 +import io +import os +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import AnyString + +from agentscope.rag import ( + ImageParser, + PDFParser, + PPTParser, + TextParser, +) + + +# Smallest valid PNG ever — a single transparent pixel. Used wherever +# a test needs "some image bytes" without depending on Pillow. +_PNG_PIXEL: bytes = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNgAAIAAAUA" + "AarVyFEAAAAASUVORK5CYII=", +) +_PNG_PIXEL_B64: str = base64.b64encode(_PNG_PIXEL).decode("utf-8") + + +def _make_pdf(pages: list[str]) -> bytes: + """Render the supplied page strings to a PDF in memory.""" + from reportlab.pdfgen import canvas + + buffer = io.BytesIO() + pdf = canvas.Canvas(buffer) + for page in pages: + pdf.drawString(72, 720, page) + pdf.showPage() + pdf.save() + return buffer.getvalue() + + +def _make_pptx_simple(slides: list[str]) -> bytes: + """Build a PPTX in memory with one text shape per slide.""" + from pptx import Presentation + + prs = Presentation() + blank_layout = prs.slide_layouts[5] # title-only layout + for text in slides: + slide = prs.slides.add_slide(blank_layout) + slide.shapes.title.text = text + buffer = io.BytesIO() + prs.save(buffer) + return buffer.getvalue() + + +def _make_pptx_rich() -> bytes: + """Build a richer PPTX with text, a table, and an embedded image. + + Slide 1: a single title text shape. + Slide 2: a title text shape, a 2x2 table, and a trailing text shape. + Slide 3: an embedded PNG image only (blank title). + """ + from pptx import Presentation + from pptx.util import Inches + + prs = Presentation() + blank_layout = prs.slide_layouts[5] + + slide1 = prs.slides.add_slide(blank_layout) + slide1.shapes.title.text = "Hello" + + slide2 = prs.slides.add_slide(blank_layout) + slide2.shapes.title.text = "Header" + table = slide2.shapes.add_table( + rows=2, + cols=2, + left=Inches(1), + top=Inches(2), + width=Inches(4), + height=Inches(1), + ).table + table.cell(0, 0).text = "A" + table.cell(0, 1).text = "B" + table.cell(1, 0).text = "1" + table.cell(1, 1).text = "2" + trailing = slide2.shapes.add_textbox( + Inches(1), + Inches(4), + Inches(4), + Inches(1), + ) + trailing.text_frame.text = "Footer" + + slide3 = prs.slides.add_slide(blank_layout) + slide3.shapes.title.text = "" # blank title + slide3.shapes.add_picture( + io.BytesIO(_PNG_PIXEL), + Inches(1), + Inches(1), + Inches(1), + Inches(1), + ) + + buffer = io.BytesIO() + prs.save(buffer) + return buffer.getvalue() + + +class TextParserTest(IsolatedAsyncioTestCase): + """Behavioural coverage for :class:`TextParser`.""" + + async def test_decode_bytes(self) -> None: + """Bytes are decoded with the configured encoding.""" + parser = TextParser() + sections = await parser.parse(b"hello", "x.txt") + self.assertEqual( + [s.model_dump() for s in sections], + [ + { + "content": { + "type": "text", + "text": "hello", + "id": AnyString(), + }, + "source": "x.txt", + "metadata": {}, + }, + ], + ) + + async def test_pre_decoded_string_round_trips(self) -> None: + """``str`` inputs skip the decode step when no such file exists.""" + parser = TextParser() + sections = await parser.parse("preset", "x.md") + self.assertEqual( + [s.model_dump() for s in sections], + [ + { + "content": { + "type": "text", + "text": "preset", + "id": AnyString(), + }, + "source": "x.md", + "metadata": {}, + }, + ], + ) + + async def test_string_input_treated_as_path(self) -> None: + """A ``str`` that names an existing file is read as a path.""" + import tempfile + + with tempfile.NamedTemporaryFile( + mode="w", + suffix=".txt", + delete=False, + ) as f: + f.write("from disk") + path = f.name + try: + parser = TextParser() + sections = await parser.parse(path, "x.txt") + self.assertEqual(len(sections), 1) + self.assertEqual(sections[0].content.text, "from disk") + finally: + os.unlink(path) + + async def test_bad_encoding_raises_value_error(self) -> None: + """An undecodable byte sequence surfaces a clear error.""" + parser = TextParser(encoding="ascii") + with self.assertRaises(ValueError): + await parser.parse(b"\xff\xfe", "bad.txt") + + +class PDFParserTest(IsolatedAsyncioTestCase): + """Behavioural coverage for :class:`PDFParser`.""" + + async def test_one_section_per_page(self) -> None: + """Each page in the PDF round-trips to its own Section.""" + pdf_bytes = _make_pdf(["First page text", "Second page text"]) + parser = PDFParser() + sections = await parser.parse(pdf_bytes, "demo.pdf") + + self.assertEqual( + [s.model_dump() for s in sections], + [ + { + "content": { + "type": "text", + "text": "First page text\n", + "id": AnyString(), + }, + "source": "demo.pdf", + "metadata": {"page": 1}, + }, + { + "content": { + "type": "text", + "text": "Second page text\n", + "id": AnyString(), + }, + "source": "demo.pdf", + "metadata": {"page": 2}, + }, + ], + ) + + async def test_string_input_treated_as_path(self) -> None: + """``str`` is interpreted as a filesystem path to the PDF.""" + import tempfile + + pdf_bytes = _make_pdf(["Hello"]) + with tempfile.NamedTemporaryFile( + suffix=".pdf", + delete=False, + ) as f: + f.write(pdf_bytes) + path = f.name + try: + parser = PDFParser() + sections = await parser.parse(path, "demo.pdf") + self.assertEqual(len(sections), 1) + finally: + os.unlink(path) + + async def test_missing_path_raises_file_not_found(self) -> None: + """A ``str`` pointing at a missing path raises FileNotFoundError.""" + parser = PDFParser() + with self.assertRaises(FileNotFoundError): + await parser.parse("/no/such/file.pdf", "x.pdf") + + async def test_invalid_bytes_raise_value_error(self) -> None: + """Garbage bytes surface as :class:`ValueError`.""" + parser = PDFParser() + with self.assertRaises(ValueError): + await parser.parse(b"not a pdf", "broken.pdf") + + async def test_supported_extensions(self) -> None: + """``.pdf`` is the only extension exposed to the file picker.""" + self.assertEqual(PDFParser.supported_extensions(), [".pdf"]) + + +class ImageParserTest(IsolatedAsyncioTestCase): + """Behavioural coverage for :class:`ImageParser`.""" + + async def test_wraps_bytes_in_data_block(self) -> None: + """A single Section is emitted with the image as a DataBlock.""" + parser = ImageParser() + sections = await parser.parse(_PNG_PIXEL, "pixel.png") + + self.assertEqual( + [s.model_dump() for s in sections], + [ + { + "content": { + "type": "data", + "id": AnyString(), + "source": { + "type": "base64", + "data": _PNG_PIXEL_B64, + "media_type": "image/png", + }, + "name": "pixel.png", + }, + "source": "pixel.png", + "metadata": {"media_type": "image/png"}, + }, + ], + ) + + async def test_media_type_sniffed_from_jpeg(self) -> None: + """JPEG magic bytes yield ``image/jpeg``.""" + jpeg_bytes = b"\xff\xd8\xff\xe0rest-of-jpeg" + parser = ImageParser() + sections = await parser.parse(jpeg_bytes, "x.jpg") + self.assertEqual( + [s.model_dump() for s in sections], + [ + { + "content": { + "type": "data", + "id": AnyString(), + "source": { + "type": "base64", + "data": base64.b64encode(jpeg_bytes).decode( + "utf-8", + ), + "media_type": "image/jpeg", + }, + "name": "x.jpg", + }, + "source": "x.jpg", + "metadata": {"media_type": "image/jpeg"}, + }, + ], + ) + + async def test_string_input_treated_as_path(self) -> None: + """``str`` is interpreted as a filesystem path to the image.""" + import tempfile + + with tempfile.NamedTemporaryFile( + suffix=".png", + delete=False, + ) as f: + f.write(_PNG_PIXEL) + path = f.name + try: + parser = ImageParser() + sections = await parser.parse(path, "pixel.png") + self.assertEqual(len(sections), 1) + finally: + os.unlink(path) + + async def test_missing_path_raises_file_not_found(self) -> None: + """A ``str`` pointing at a missing path raises FileNotFoundError.""" + parser = ImageParser() + with self.assertRaises(FileNotFoundError): + await parser.parse("/no/such/file.png", "x.png") + + +class PPTParserTest(IsolatedAsyncioTestCase): + """Behavioural coverage for :class:`PPTParser`.""" + + async def test_simple_deck_text_only(self) -> None: + """A simple text-only deck round-trips through wrapping tags.""" + pptx_bytes = _make_pptx_simple(["Alpha", "Beta"]) + parser = PPTParser(include_image=False) + sections = await parser.parse(pptx_bytes, "demo.pptx") + + self.assertEqual( + [s.model_dump() for s in sections], + [ + { + "content": { + "type": "text", + "text": "\nAlpha\n", + "id": AnyString(), + }, + "source": "demo.pptx", + "metadata": {"slide": 1}, + }, + { + "content": { + "type": "text", + "text": "\nBeta\n", + "id": AnyString(), + }, + "source": "demo.pptx", + "metadata": {"slide": 2}, + }, + ], + ) + + async def test_without_slide_tags(self) -> None: + """Disabling prefix/suffix removes the slide tags entirely.""" + pptx_bytes = _make_pptx_simple(["Alpha"]) + parser = PPTParser( + include_image=False, + slide_prefix=None, + slide_suffix=None, + ) + sections = await parser.parse(pptx_bytes, "demo.pptx") + self.assertEqual( + [s.model_dump() for s in sections], + [ + { + "content": { + "type": "text", + "text": "Alpha", + "id": AnyString(), + }, + "source": "demo.pptx", + "metadata": {"slide": 1}, + }, + ], + ) + + async def test_table_merges_with_surrounding_text_by_default( + self, + ) -> None: + """``separate_table=False`` merges the table into the running + text section.""" + pptx_bytes = _make_pptx_rich() + parser = PPTParser(include_image=False, separate_table=False) + sections = await parser.parse(pptx_bytes, "rich.pptx") + + self.assertEqual( + [s.model_dump() for s in sections], + [ + { + "content": { + "type": "text", + "text": "\nHello\n", + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 1}, + }, + { + "content": { + "type": "text", + "text": ( + "\n" + "Header\n" + "| A | B |\n" + "| --- | --- |\n" + "| 1 | 2 |\n\n" + "Footer\n" + "" + ), + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 2}, + }, + { + "content": { + "type": "text", + "text": "\n", + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 3}, + }, + ], + ) + + async def test_table_separated_when_separate_table_true(self) -> None: + """``separate_table=True`` flushes the running text around the + table.""" + pptx_bytes = _make_pptx_rich() + parser = PPTParser(include_image=False, separate_table=True) + sections = await parser.parse(pptx_bytes, "rich.pptx") + + self.assertEqual( + [s.model_dump() for s in sections], + [ + { + "content": { + "type": "text", + "text": "\nHello\n", + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 1}, + }, + { + "content": { + "type": "text", + "text": "\nHeader", + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 2}, + }, + { + "content": { + "type": "text", + "text": ( + "| A | B |\n" "| --- | --- |\n" "| 1 | 2 |\n" + ), + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 2}, + }, + { + "content": { + "type": "text", + "text": "Footer\n", + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 2}, + }, + { + "content": { + "type": "text", + "text": "\n", + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 3}, + }, + ], + ) + + async def test_image_emits_data_block(self) -> None: + """An embedded picture becomes its own DataBlock section.""" + pptx_bytes = _make_pptx_rich() + parser = PPTParser(include_image=True, separate_table=True) + sections = await parser.parse(pptx_bytes, "rich.pptx") + + self.assertEqual( + [s.model_dump() for s in sections], + [ + { + "content": { + "type": "text", + "text": "\nHello\n", + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 1}, + }, + { + "content": { + "type": "text", + "text": "\nHeader", + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 2}, + }, + { + "content": { + "type": "text", + "text": ( + "| A | B |\n" "| --- | --- |\n" "| 1 | 2 |\n" + ), + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 2}, + }, + { + "content": { + "type": "text", + "text": "Footer\n", + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 2}, + }, + { + "content": { + "type": "text", + "text": "", + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 3}, + }, + { + "content": { + "type": "data", + "id": AnyString(), + "source": { + "type": "base64", + "data": _PNG_PIXEL_B64, + "media_type": "image/png", + }, + "name": "rich.pptx", + }, + "source": "rich.pptx", + "metadata": { + "slide": 3, + "media_type": "image/png", + }, + }, + { + "content": { + "type": "text", + "text": "", + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 3}, + }, + ], + ) + + async def test_table_json_format(self) -> None: + """``table_format="json"`` emits the JSON marker payload.""" + pptx_bytes = _make_pptx_rich() + parser = PPTParser( + include_image=False, + separate_table=True, + table_format="json", + ) + sections = await parser.parse(pptx_bytes, "rich.pptx") + + self.assertEqual( + [s.model_dump() for s in sections], + [ + { + "content": { + "type": "text", + "text": "\nHello\n", + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 1}, + }, + { + "content": { + "type": "text", + "text": "\nHeader", + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 2}, + }, + { + "content": { + "type": "text", + "text": ( + "A table loaded as a JSON " + "array:\n" + '[["A", "B"], ["1", "2"]]' + ), + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 2}, + }, + { + "content": { + "type": "text", + "text": "Footer\n", + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 2}, + }, + { + "content": { + "type": "text", + "text": "\n", + "id": AnyString(), + }, + "source": "rich.pptx", + "metadata": {"slide": 3}, + }, + ], + ) + + async def test_table_format_validation(self) -> None: + """Unknown ``table_format`` raises :class:`ValueError`.""" + with self.assertRaises(ValueError): + PPTParser(table_format="csv") # type: ignore[arg-type] + + async def test_string_input_treated_as_path(self) -> None: + """``str`` is interpreted as a filesystem path to the PPTX.""" + import tempfile + + pptx_bytes = _make_pptx_simple(["Alpha"]) + with tempfile.NamedTemporaryFile( + suffix=".pptx", + delete=False, + ) as f: + f.write(pptx_bytes) + path = f.name + try: + parser = PPTParser(include_image=False) + sections = await parser.parse(path, "demo.pptx") + self.assertEqual(len(sections), 1) + finally: + os.unlink(path) + + async def test_missing_path_raises_file_not_found(self) -> None: + """A ``str`` pointing at a missing path raises FileNotFoundError.""" + parser = PPTParser() + with self.assertRaises(FileNotFoundError): + await parser.parse("/no/such/file.pptx", "x.pptx") diff --git a/tests/rag_vdb_qdrant_test.py b/tests/rag_vdb_qdrant_test.py new file mode 100644 index 0000000000000000000000000000000000000000..3d15424d46d8d0ceabc5e36e6245d008b319085b --- /dev/null +++ b/tests/rag_vdb_qdrant_test.py @@ -0,0 +1,376 @@ +# -*- coding: utf-8 -*- +"""Unit tests for the QdrantStore class.""" +from contextlib import AsyncExitStack +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import AnyString + +from agentscope.message import TextBlock +from agentscope.rag import ( + Chunk, + QdrantStore, + VectorRecord, + VectorSearchResult, +) + + +def _dump_results(results: list[VectorSearchResult]) -> list[dict]: + """Convert search results into plain dicts for whole-structure + comparison. + + Args: + results (`list[VectorSearchResult]`): + The search results to convert. + + Returns: + `list[dict]`: + The results as plain dicts. + """ + return [result.model_dump() for result in results] + + +def _make_record( + text: str, + vector: list[float], + document_id: str, + chunk_index: int = 0, + total_chunks: int = 1, +) -> VectorRecord: + """Build a VectorRecord for testing. + + Args: + text (`str`): + The chunk text content. + vector (`list[float]`): + The embedding vector. + document_id (`str`): + The ID of the source document the record belongs to. + chunk_index (`int`, defaults to ``0``): + The chunk index within the document. + total_chunks (`int`, defaults to ``1``): + The total number of chunks in the document. + + Returns: + `VectorRecord`: + The constructed record. + """ + return VectorRecord( + vector=vector, + document_id=document_id, + chunk=Chunk( + content=TextBlock(text=text), + source=f"{document_id}.txt", + chunk_index=chunk_index, + total_chunks=total_chunks, + ), + ) + + +class QdrantStoreTest(IsolatedAsyncioTestCase): + """The test cases for the QdrantStore class.""" + + async def asyncSetUp(self) -> None: + """Create an in-memory Qdrant store before each test.""" + self._exit_stack = AsyncExitStack() + self.store = await self._exit_stack.enter_async_context( + QdrantStore(location=":memory:"), + ) + + async def asyncTearDown(self) -> None: + """Close the store after each test.""" + await self._exit_stack.aclose() + + async def test_collection_lifecycle(self) -> None: + """Collections can be created, checked, and deleted.""" + self.assertEqual(await self.store.has_collection("kb-1"), False) + + await self.store.create_collection("kb-1", dimensions=3) + self.assertEqual(await self.store.has_collection("kb-1"), True) + + # Creating an existing collection is a no-op + await self.store.create_collection("kb-1", dimensions=3) + self.assertEqual(await self.store.has_collection("kb-1"), True) + + await self.store.delete_collection("kb-1") + self.assertEqual(await self.store.has_collection("kb-1"), False) + + async def test_insert_and_search(self) -> None: + """Inserted records are searchable, ordered by similarity.""" + await self.store.create_collection("kb-1", dimensions=3) + await self.store.insert( + "kb-1", + [ + _make_record( + "Hello world!", + [1.0, 0.0, 0.0], + document_id="doc-1", + chunk_index=0, + total_chunks=2, + ), + _make_record( + "Goodbye world!", + [0.0, 1.0, 0.0], + document_id="doc-1", + chunk_index=1, + total_chunks=2, + ), + ], + ) + + results = await self.store.search( + "kb-1", + query_vector=[1.0, 0.0, 0.0], + top_k=2, + ) + + self.assertEqual( + _dump_results(results), + [ + { + "score": 1.0, + "document_id": "doc-1", + "chunk": { + "content": { + "type": "text", + "text": "Hello world!", + "id": AnyString(), + }, + "source": "doc-1.txt", + "chunk_index": 0, + "total_chunks": 2, + "metadata": {}, + }, + }, + { + "score": 0.0, + "document_id": "doc-1", + "chunk": { + "content": { + "type": "text", + "text": "Goodbye world!", + "id": AnyString(), + }, + "source": "doc-1.txt", + "chunk_index": 1, + "total_chunks": 2, + "metadata": {}, + }, + }, + ], + ) + + async def test_search_top_k(self) -> None: + """top_k limits the number of returned results.""" + await self.store.create_collection("kb-1", dimensions=3) + await self.store.insert( + "kb-1", + [ + _make_record("A", [1.0, 0.0, 0.0], document_id="doc-1"), + _make_record("B", [0.9, 0.1, 0.0], document_id="doc-2"), + _make_record("C", [0.0, 0.0, 1.0], document_id="doc-3"), + ], + ) + + results = await self.store.search( + "kb-1", + query_vector=[1.0, 0.0, 0.0], + top_k=1, + ) + + self.assertEqual( + _dump_results(results), + [ + { + "score": 1.0, + "document_id": "doc-1", + "chunk": { + "content": { + "type": "text", + "text": "A", + "id": AnyString(), + }, + "source": "doc-1.txt", + "chunk_index": 0, + "total_chunks": 1, + "metadata": {}, + }, + }, + ], + ) + + async def test_delete_by_document_id(self) -> None: + """delete removes all records of one document only.""" + await self.store.create_collection("kb-1", dimensions=3) + await self.store.insert( + "kb-1", + [ + _make_record( + "doc1-chunk0", + [1.0, 0.0, 0.0], + document_id="doc-1", + chunk_index=0, + total_chunks=2, + ), + _make_record( + "doc1-chunk1", + [0.9, 0.1, 0.0], + document_id="doc-1", + chunk_index=1, + total_chunks=2, + ), + _make_record( + "doc2-chunk0", + [0.0, 1.0, 0.0], + document_id="doc-2", + ), + ], + ) + + await self.store.delete("kb-1", document_id="doc-1") + + results = await self.store.search( + "kb-1", + query_vector=[1.0, 0.0, 0.0], + top_k=5, + ) + + self.assertEqual( + _dump_results(results), + [ + { + "score": 0.0, + "document_id": "doc-2", + "chunk": { + "content": { + "type": "text", + "text": "doc2-chunk0", + "id": AnyString(), + }, + "source": "doc-2.txt", + "chunk_index": 0, + "total_chunks": 1, + "metadata": {}, + }, + }, + ], + ) + + async def test_insert_empty_records(self) -> None: + """Inserting an empty record list is a no-op.""" + await self.store.create_collection("kb-1", dimensions=3) + await self.store.insert("kb-1", []) + + results = await self.store.search( + "kb-1", + query_vector=[1.0, 0.0, 0.0], + ) + + self.assertEqual(_dump_results(results), []) + + async def test_list_documents_aggregates_by_document_id(self) -> None: + """list_documents groups chunks by document_id.""" + await self.store.create_collection("kb-1", dimensions=3) + + def _record_with_metadata( + text: str, + document_id: str, + metadata: dict, + chunk_index: int = 0, + total_chunks: int = 1, + ) -> VectorRecord: + return VectorRecord( + vector=[1.0, 0.0, 0.0], + document_id=document_id, + chunk=Chunk( + content=TextBlock(text=text), + source=metadata.get("filename", f"{document_id}.txt"), + chunk_index=chunk_index, + total_chunks=total_chunks, + metadata=metadata, + ), + ) + + await self.store.insert( + "kb-1", + [ + _record_with_metadata( + "A", + "doc-1", + {"filename": "alpha.txt", "media_type": "text/plain"}, + 0, + 2, + ), + _record_with_metadata( + "B", + "doc-1", + {"filename": "alpha.txt", "media_type": "text/plain"}, + 1, + 2, + ), + _record_with_metadata( + "C", + "doc-2", + {"filename": "beta.md", "media_type": "text/markdown"}, + 0, + 1, + ), + ], + ) + + summaries = await self.store.list_documents("kb-1") + summaries_by_id = {s.document_id: s for s in summaries} + + self.assertEqual(set(summaries_by_id), {"doc-1", "doc-2"}) + self.assertEqual(summaries_by_id["doc-1"].chunk_count, 2) + self.assertEqual(summaries_by_id["doc-1"].source, "alpha.txt") + self.assertEqual( + summaries_by_id["doc-1"].metadata, + {"filename": "alpha.txt", "media_type": "text/plain"}, + ) + self.assertEqual(summaries_by_id["doc-2"].chunk_count, 1) + self.assertEqual(summaries_by_id["doc-2"].source, "beta.md") + + async def test_search_metadata_filter(self) -> None: + """search applies the metadata_filter as a payload predicate.""" + await self.store.create_collection("kb-1", dimensions=3) + + def _record( + text: str, + document_id: str, + kb_scope: str, + ) -> VectorRecord: + return VectorRecord( + vector=[1.0, 0.0, 0.0], + document_id=document_id, + chunk=Chunk( + content=TextBlock(text=text), + source=f"{document_id}.txt", + chunk_index=0, + total_chunks=1, + metadata={"kb_scope": kb_scope}, + ), + ) + + await self.store.insert( + "kb-1", + [ + _record("A", "doc-1", "kb-a"), + _record("B", "doc-2", "kb-b"), + ], + ) + + results = await self.store.search( + "kb-1", + query_vector=[1.0, 0.0, 0.0], + top_k=5, + metadata_filter={"kb_scope": "kb-a"}, + ) + self.assertEqual([r.document_id for r in results], ["doc-1"]) + + results = await self.store.search( + "kb-1", + query_vector=[1.0, 0.0, 0.0], + top_k=5, + metadata_filter={"kb_scope": "kb-b"}, + ) + self.assertEqual([r.document_id for r in results], ["doc-2"]) diff --git a/tests/service_cancel_dispatcher_test.py b/tests/service_cancel_dispatcher_test.py new file mode 100644 index 0000000000000000000000000000000000000000..1924cf1d0d0889abc0f30b0f4365b7d5da2d606f --- /dev/null +++ b/tests/service_cancel_dispatcher_test.py @@ -0,0 +1,463 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Tests for :class:`CancelDispatcher` — one-per-process consumer of the +shared session-cancel broadcast channel. + +Verifies that on each incoming ``session_id`` the dispatcher: + +- Cancels the chat-run task in :class:`ChatRunRegistry` when it owns one + locally. +- Asks :class:`BackgroundTaskManager` to cancel local BG tasks for the + same session. +- Silently does nothing for sessions whose state lives on other + processes. +""" +import asyncio +from contextlib import asynccontextmanager +from typing import AsyncGenerator, Callable +from unittest import IsolatedAsyncioTestCase + +from agentscope.app._manager import ( + BackgroundTaskManager, + CancelDispatcher, + ChatRunRegistry, +) +from agentscope.app._manager._background_task_manager import ToolStop +from agentscope.app.message_bus import MessageBus +from agentscope.message import ToolResultState + + +class _FakeBus(MessageBus): + """In-memory bus with just enough behaviour for the dispatcher. + + Only the cancel-broadcast channel is exercised here; the other + primitives are stubbed. + """ + + def __init__(self) -> None: + self._channels: dict[str, asyncio.Queue] = {} + self._locks: set[str] = set() + self._registries: dict[str, dict[str, str]] = {} + + def _channel(self, key: str) -> asyncio.Queue: + return self._channels.setdefault(key, asyncio.Queue()) + + # Mode A — queue (unused) + async def queue_push( + self, + key: str, + payload: dict, + *, + ttl_secs: int | None = None, + ) -> str: + return "n/a" + + async def queue_drain( + self, + key: str, + max_count: int = 100, + ) -> list[tuple[str, dict]]: + return [] + + async def queue_delete(self, key: str) -> None: + return None + + # Mode C — log (unused) + async def log_append( + self, + key: str, + payload: dict, + *, + max_len: int | None = None, + ttl_secs: int | None = None, + ) -> str: + return "n/a" + + async def log_read( + self, + key: str, + since: str | None = None, + max_count: int = 100, + ) -> list[tuple[str, dict]]: + return [] + + async def log_trim( + self, + key: str, + before_id: str | None = None, + ) -> None: + return None + + # Mode D — pub/sub + async def publish(self, key: str, payload: dict) -> None: + await self._channel(key).put(payload) + + async def subscribe( + self, + key: str, + *, + on_ready: Callable[[], None] | None = None, + ) -> AsyncGenerator[dict, None]: + if on_ready is not None: + on_ready() + while True: + yield await self._channel(key).get() + + # Mode E — lock (unused) + @asynccontextmanager + async def acquire_lock( + self, + key: str, + *, + ttl_secs: int = 600, + ) -> AsyncGenerator[None, None]: + self._locks.add(key) + try: + yield + finally: + self._locks.discard(key) + + async def is_locked(self, key: str) -> bool: + return key in self._locks + + # Mode F — registry (in-memory dict) + async def registry_set( + self, + namespace: str, + field: str, + value: str, + *, + ttl_secs: int | None = None, + ) -> None: + self._registries.setdefault(namespace, {})[field] = value + + async def registry_del(self, namespace: str, field: str) -> None: + if namespace in self._registries: + self._registries[namespace].pop(field, None) + + async def registry_exists(self, namespace: str, field: str) -> bool: + return field in self._registries.get(namespace, {}) + + async def registry_getall(self, namespace: str) -> dict[str, str]: + return dict(self._registries.get(namespace, {})) + + async def registry_drop(self, namespace: str) -> None: + self._registries.pop(namespace, None) + + +async def _yield_a_few_times(ticks: int = 8) -> None: + """Yield the event loop a few times so spawned tasks make progress.""" + for _ in range(ticks): + await asyncio.sleep(0) + + +class _NeverEndingCoro: + """Helper: yields a fresh coroutine that sleeps forever.""" + + @staticmethod + async def run() -> None: + """The fake coroutine.""" + await asyncio.Event().wait() + + +class TestCancelDispatcher(IsolatedAsyncioTestCase): + """Verifies the cross-process cancel fan-out.""" + + async def test_cancel_signal_cancels_local_chat_run(self) -> None: + """Broadcast for a session whose chat run is registered locally + cancels the registered asyncio task.""" + bus = _FakeBus() + registry = ChatRunRegistry() + bg_manager = BackgroundTaskManager(message_bus=bus) + + async with bg_manager, registry, CancelDispatcher( + message_bus=bus, + registry=registry, + bg_manager=bg_manager, + ): + chat_task = registry.spawn( + _NeverEndingCoro.run(), + session_id="sess-A", + ) + await bus.session_publish_cancel("sess-A") + + # Wait until the cancel actually propagates. + for _ in range(50): + if chat_task.cancelled() or chat_task.done(): + break + await asyncio.sleep(0.01) + + self.assertTrue(chat_task.cancelled() or chat_task.done()) + + async def test_cancel_signal_cancels_local_bg_tasks(self) -> None: + """Broadcast for a session with locally-registered BG tasks + cancels each of them; tasks for other sessions are untouched.""" + bus = _FakeBus() + registry = ChatRunRegistry() + bg_manager = BackgroundTaskManager(message_bus=bus) + + async with bg_manager, registry, CancelDispatcher( + message_bus=bus, + registry=registry, + bg_manager=bg_manager, + ): + bg_task_a1 = asyncio.create_task(_NeverEndingCoro.run()) + bg_task_a2 = asyncio.create_task(_NeverEndingCoro.run()) + bg_task_b = asyncio.create_task(_NeverEndingCoro.run()) + + await bg_manager.register_task( + bg_task_a1, + session_id="sess-A", + agent_id="agent-A", + user_id="u", + ) + await bg_manager.register_task( + bg_task_a2, + session_id="sess-A", + agent_id="agent-A", + user_id="u", + ) + await bg_manager.register_task( + bg_task_b, + session_id="sess-B", + agent_id="agent-B", + user_id="u", + ) + + await bus.session_publish_cancel("sess-A") + + for _ in range(50): + if bg_task_a1.cancelled() and bg_task_a2.cancelled(): + break + await asyncio.sleep(0.01) + + self.assertTrue(bg_task_a1.cancelled() or bg_task_a1.done()) + self.assertTrue(bg_task_a2.cancelled() or bg_task_a2.done()) + # sess-B BG task is left running until shutdown cancels it. + self.assertFalse(bg_task_b.cancelled()) + bg_task_b.cancel() + + async def test_cancel_signal_for_remote_session_is_noop(self) -> None: + """Broadcast for a session held on another process is silently + ignored — no exception, no spurious cancel.""" + bus = _FakeBus() + registry = ChatRunRegistry() + bg_manager = BackgroundTaskManager(message_bus=bus) + + async with bg_manager, registry, CancelDispatcher( + message_bus=bus, + registry=registry, + bg_manager=bg_manager, + ): + # Register an unrelated chat run + unrelated BG task so we + # can verify the unrelated work survives the broadcast. + unrelated_chat = registry.spawn( + _NeverEndingCoro.run(), + session_id="other", + ) + unrelated_bg = asyncio.create_task(_NeverEndingCoro.run()) + await bg_manager.register_task( + unrelated_bg, + session_id="other", + agent_id="agent", + user_id="u", + ) + + await bus.session_publish_cancel("not-on-this-process") + await _yield_a_few_times() + + self.assertFalse(unrelated_chat.cancelled()) + self.assertFalse(unrelated_bg.cancelled()) + + # __aexit__ of ChatRunRegistry + BackgroundTaskManager cancels + # the unrelated tasks on shutdown. + + async def test_cancel_fans_out_to_both_chat_and_bg_in_one_signal( + self, + ) -> None: + """A single cancel broadcast cancels both the local chat run + and the local BG task(s) for the session, not just one.""" + bus = _FakeBus() + registry = ChatRunRegistry() + bg_manager = BackgroundTaskManager(message_bus=bus) + + async with bg_manager, registry, CancelDispatcher( + message_bus=bus, + registry=registry, + bg_manager=bg_manager, + ): + chat_task = registry.spawn( + _NeverEndingCoro.run(), + session_id="sess-X", + ) + bg_task = asyncio.create_task(_NeverEndingCoro.run()) + await bg_manager.register_task( + bg_task, + session_id="sess-X", + agent_id="agent-X", + user_id="u", + ) + + await bus.session_publish_cancel("sess-X") + + for _ in range(50): + if chat_task.cancelled() and bg_task.cancelled(): + break + await asyncio.sleep(0.01) + + self.assertTrue(chat_task.cancelled() or chat_task.done()) + self.assertTrue(bg_task.cancelled() or bg_task.done()) + + +class TestBackgroundTaskManagerCancelSessionTasks(IsolatedAsyncioTestCase): + """Verifies :meth:`BackgroundTaskManager.cancel_session_tasks`.""" + + async def test_cancels_only_matching_session(self) -> None: + """Only tasks whose ``session_id`` matches are cancelled; the + return value reports the local count.""" + bg_manager = BackgroundTaskManager(message_bus=_FakeBus()) + async with bg_manager: + task_a = asyncio.create_task(_NeverEndingCoro.run()) + task_b = asyncio.create_task(_NeverEndingCoro.run()) + await bg_manager.register_task( + task_a, + session_id="match", + agent_id="a", + user_id="u", + ) + await bg_manager.register_task( + task_b, + session_id="other", + agent_id="b", + user_id="u", + ) + + count = bg_manager.cancel_session_tasks("match") + self.assertEqual(count, 1) + + for _ in range(50): + if task_a.cancelled(): + break + await asyncio.sleep(0.01) + + self.assertTrue(task_a.cancelled()) + self.assertFalse(task_b.cancelled()) + + async def test_no_matches_returns_zero(self) -> None: + """A session with no locally-registered tasks returns 0 and + does no work.""" + bg_manager = BackgroundTaskManager(message_bus=_FakeBus()) + async with bg_manager: + task = asyncio.create_task(_NeverEndingCoro.run()) + await bg_manager.register_task( + task, + session_id="other", + agent_id="a", + user_id="u", + ) + + self.assertEqual( + bg_manager.cancel_session_tasks("ghost"), + 0, + ) + self.assertFalse(task.cancelled()) + + +class TestToolStopRemoteCancel(IsolatedAsyncioTestCase): + """Verifies the cross-worker cancel path of :class:`ToolStop`. + + A "worker A" registers a BG task in the shared bus registry, and a + "worker B" — which has the task only in the global registry, not in + its local cache — issues ``ToolStop``. The dispatcher on worker A + must receive the broadcast and cancel the task locally. + """ + + async def test_remote_cancel_via_toolstop_broadcast(self) -> None: + """ToolStop on a worker without the task publishes a task-level + cancel; the owning worker's CancelDispatcher cancels the task.""" + bus = _FakeBus() + + # Worker A — owns the task and runs CancelDispatcher. + bg_manager_owner = BackgroundTaskManager(message_bus=bus) + registry_owner = ChatRunRegistry() + + # Worker B — only sees the task via the shared registry. + bg_manager_caller = BackgroundTaskManager(message_bus=bus) + + async with bg_manager_owner, registry_owner, CancelDispatcher( + message_bus=bus, + registry=registry_owner, + bg_manager=bg_manager_owner, + ), bg_manager_caller: + owned_task = asyncio.create_task(_NeverEndingCoro.run()) + task_id = await bg_manager_owner.register_task( + owned_task, + session_id="sess-shared", + agent_id="agent", + user_id="u", + tool_name="LongRunningTool", + ) + + # Worker B's ToolStop: task_id is in the global registry but + # not in worker B's local cache, so the remote-cancel path + # is taken. + tool_stop = ToolStop( + background_tasks=bg_manager_caller.tasks, + message_bus=bus, + session_id="sess-shared", + ) + chunk = await tool_stop(task_id=task_id) + + self.assertEqual(chunk.state, ToolResultState.SUCCESS) + self.assertIn( + "Cancel request sent", + chunk.content[0].text, + ) + + for _ in range(50): + if owned_task.cancelled() or owned_task.done(): + break + await asyncio.sleep(0.01) + + self.assertTrue(owned_task.cancelled() or owned_task.done()) + + async def test_toolstop_does_not_cancel_other_session_locally( + self, + ) -> None: + """A ToolStop instance bound to session A must not cancel a + locally-tracked task that belongs to session B, even if the + guessed task_id is correct.""" + bus = _FakeBus() + bg_manager = BackgroundTaskManager(message_bus=bus) + + async with bg_manager: + victim_task = asyncio.create_task(_NeverEndingCoro.run()) + victim_task_id = await bg_manager.register_task( + victim_task, + session_id="sess-victim", + agent_id="agent-v", + user_id="u", + ) + + # ToolStop is bound to a *different* session; it should not + # cancel ``victim_task`` directly. The shared registry is + # also keyed by the bound session id, so the lookup misses + # and we fall through to "not found". + tool_stop = ToolStop( + background_tasks=bg_manager.tasks, + message_bus=bus, + session_id="sess-attacker", + ) + chunk = await tool_stop(task_id=victim_task_id) + + await _yield_a_few_times() + + self.assertEqual(chunk.state, ToolResultState.ERROR) + self.assertIn( + "TaskNotFoundError", + chunk.content[0].text, + ) + self.assertFalse(victim_task.cancelled()) + self.assertIn(victim_task_id, bg_manager.tasks) + + victim_task.cancel() diff --git a/tests/service_enqueue_index_task_test.py b/tests/service_enqueue_index_task_test.py new file mode 100644 index 0000000000000000000000000000000000000000..2f0b951be0c93298730d5b5a5858cada15bdc536 --- /dev/null +++ b/tests/service_enqueue_index_task_test.py @@ -0,0 +1,154 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Tests for :func:`agentscope.app._bus_ops.enqueue_index_task`. + +The helper is a three-line composition over two bus primitives — +``queue_push`` and ``publish`` — so the tests verify only that: + +- the queued payload carries the three fields the + :class:`~agentscope.app._service.IndexTaskConsumer` reads; +- a signal is published exactly once per call; +- both primitives reach the production + :class:`~agentscope.app.message_bus.RedisMessageBus` backend + (proxied by ``fakeredis``), not just an in-memory mock — keeping + the test honest against the wire-level contract the worker side + relies on. + +If a future change adds metadata to the payload, this test is the +contract gate: callers are expected to keep ``user_id`` / +``knowledge_base_id`` / ``document_id`` as the primary keys. +""" +import asyncio +from contextlib import AsyncExitStack +from unittest import IsolatedAsyncioTestCase + +import fakeredis.aioredis + +from agentscope.app._bus_ops import enqueue_index_task +from agentscope.app.message_bus import MessageBusKeys, RedisMessageBus + + +def _make_bus(fr: fakeredis.aioredis.FakeRedis) -> RedisMessageBus: + """Construct a :class:`RedisMessageBus` bound to *fr*. + + Args: + fr (`fakeredis.aioredis.FakeRedis`): + The fake Redis client to inject into the bus. + + Returns: + `RedisMessageBus`: + A bus subclass whose ``__aenter__`` skips the real + connection setup and binds *fr* as the client. + """ + + class _B(RedisMessageBus): + """Test-only bus that reuses the supplied fakeredis client.""" + + async def __aenter__( + self, + ) -> "RedisMessageBus": # type: ignore[override] + """Bind the pre-supplied fakeredis client and return self. + + Returns: + `RedisMessageBus`: + This bus, ready for the with-block body. + """ + self._client = fr + return self + + async def aclose(self) -> None: + """Drop the client reference without touching the network.""" + self._client = None + + return _B() + + +class TestEnqueueIndexTask(IsolatedAsyncioTestCase): + """Verifies the queue_push + publish composition.""" + + async def asyncSetUp(self) -> None: + """Wire a fakeredis-backed bus into an async exit stack.""" + self.fr = fakeredis.aioredis.FakeRedis(decode_responses=True) + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context(_make_bus(self.fr)) + + async def asyncTearDown(self) -> None: + """Tear the exit stack and the fakeredis client down.""" + await self._stack.aclose() + await self.fr.aclose() + + async def test_enqueue_pushes_payload_and_publishes_signal( + self, + ) -> None: + """One enqueue puts one structured entry on the durable queue + and one opaque payload on the signal channel. + + Order matters: the queue push must precede the publish so a + worker woken by the signal is guaranteed to find the entry + when it drains. We do not assert order directly (would + require instrumenting the bus), but we assert both primitives + landed. + """ + ready = asyncio.Event() + received: list[dict] = [] + + async def _signal_consumer() -> None: + """Consume one signal payload and exit.""" + async for payload in self.bus.subscribe( + MessageBusKeys.index_tasks_signal(), + on_ready=ready.set, + ): + received.append(payload) + break + + task = asyncio.create_task(_signal_consumer()) + await asyncio.wait_for(ready.wait(), timeout=2.0) + + await enqueue_index_task( + self.bus, + user_id="u", + knowledge_base_id="kb", + document_id="doc", + ) + await asyncio.wait_for(task, timeout=2.0) + + # Signal was published. + self.assertEqual(len(received), 1) + + # Queue holds the structured entry under the well-known key. + entries = await self.bus.queue_drain( + MessageBusKeys.index_tasks_queue(), + max_count=10, + ) + self.assertEqual(len(entries), 1) + _entry_id, payload = entries[0] + self.assertEqual( + payload, + { + "user_id": "u", + "knowledge_base_id": "kb", + "document_id": "doc", + }, + ) + + async def test_double_enqueue_queues_twice(self) -> None: + """Two enqueues for the same document leave two entries on + the queue — the helper is intentionally not deduplicating, + the worker's lease CAS is the dedup contract.""" + await enqueue_index_task( + self.bus, + user_id="u", + knowledge_base_id="kb", + document_id="doc", + ) + await enqueue_index_task( + self.bus, + user_id="u", + knowledge_base_id="kb", + document_id="doc", + ) + entries = await self.bus.queue_drain( + MessageBusKeys.index_tasks_queue(), + max_count=10, + ) + self.assertEqual(len(entries), 2) diff --git a/tests/service_inbox_middleware_test.py b/tests/service_inbox_middleware_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ef2f6e40027db482a18d21ae97691560625024f3 --- /dev/null +++ b/tests/service_inbox_middleware_test.py @@ -0,0 +1,477 @@ +# -*- coding: utf-8 -*- +# pylint: disable=abstract-method,protected-access +"""Tests for :class:`InboxMiddleware`. + +Every cross-session message delivery in the framework (team messages +via ``TeamSay`` / ``AgentCreate``, background-tool completion results, +scheduler fires) goes through this middleware on the consumer side, so +the four branches below cover the whole inbox→context pipeline: + +- empty inbox → no injection, no event yield, downstream unchanged; +- last context msg already an assistant msg from this agent → hints are + *extended* into its ``content``; +- last context msg is something else (system / different agent) → a + fresh ``AssistantMsg`` is appended; +- empty context → a fresh ``AssistantMsg`` is appended. + +Each non-empty drain must also yield one ``HintBlockEvent`` per hint so +the SSE stream renders them. +""" +import uuid +from contextlib import asynccontextmanager +from types import SimpleNamespace +from typing import Any, AsyncGenerator, Callable +from unittest import IsolatedAsyncioTestCase + +from utils import AnyString + +from agentscope.app.message_bus import MessageBus +from agentscope.app.middleware import InboxMiddleware +from agentscope.message import ( + AssistantMsg, + HintBlock, + SystemMsg, + TextBlock, +) + + +class _FakeBus(MessageBus): + """In-memory bus that only implements the inbox API needed by + :class:`InboxMiddleware`. All other primitives raise ``NotImplemented`` + to make accidental dependencies obvious.""" + + def __init__(self) -> None: + self._queues: dict[str, list[tuple[str, dict]]] = {} + self._next_id = 0 + + def _alloc_id(self) -> str: + """Allocate a monotonically increasing entry id.""" + self._next_id += 1 + return f"id-{self._next_id}" + + # Mode A — drain queue + async def queue_push( + self, + key: str, + payload: dict, + *, + ttl_secs: int | None = None, + ) -> str: + entry_id = self._alloc_id() + self._queues.setdefault(key, []).append((entry_id, payload)) + return entry_id + + async def queue_drain( + self, + key: str, + *, + max_count: int, + ) -> list[tuple[str, dict]]: + entries = self._queues.get(key, [])[:max_count] + self._queues[key] = self._queues.get(key, [])[max_count:] + return entries + + async def queue_delete(self, key: str) -> None: + self._queues.pop(key, None) + + # Mode C — log (unused) + async def log_append( + self, + key: str, + payload: dict, + *, + max_len: int | None = None, + ttl_secs: int | None = None, + ) -> str: + raise NotImplementedError + + async def log_read( + self, + key: str, + since: str | None = None, + max_count: int = 100, + ) -> list[tuple[str, dict]]: + raise NotImplementedError + + async def log_trim( + self, + key: str, + before_id: str | None = None, + ) -> None: + raise NotImplementedError + + # Mode D — pub/sub (unused) + async def publish(self, key: str, payload: dict) -> None: + raise NotImplementedError + + async def subscribe( + self, + key: str, + *, + on_ready: Callable[[], None] | None = None, + ) -> AsyncGenerator[dict, None]: + raise NotImplementedError + yield # pragma: no cover # pylint: disable=unreachable + + # Mode E — lock (unused) + @asynccontextmanager + async def acquire_lock( + self, + key: str, + *, + ttl_secs: int = 600, + ) -> AsyncGenerator[None, None]: + raise NotImplementedError + yield # pragma: no cover # pylint: disable=unreachable + + async def is_locked(self, key: str) -> bool: + raise NotImplementedError + + # Mode F — registry (unused) + async def registry_set( + self, + namespace: str, + field: str, + value: str, + *, + ttl_secs: int | None = None, + ) -> None: + raise NotImplementedError + + async def registry_del(self, namespace: str, field: str) -> None: + raise NotImplementedError + + async def registry_exists(self, namespace: str, field: str) -> bool: + raise NotImplementedError + + async def registry_getall(self, namespace: str) -> dict[str, str]: + raise NotImplementedError + + async def registry_drop(self, namespace: str) -> None: + raise NotImplementedError + + +def _make_agent( + *, + name: str, + session_id: str, + reply_id: str, + context: list, +) -> Any: + """Build the smallest object that satisfies what + :class:`InboxMiddleware.on_reasoning` reads off ``agent``. + + Args: + name: Agent display name. + session_id: Identifier for the inbox key. + reply_id: Reply id stamped onto freshly-appended AssistantMsg. + context: Mutable list of messages — the middleware mutates this + in place to inject hints. + + Returns: + A ``SimpleNamespace`` shaped like the real :class:`Agent` for + the fields the middleware touches. + """ + return SimpleNamespace( + name=name, + state=SimpleNamespace( + session_id=session_id, + reply_id=reply_id, + context=context, + ), + ) + + +async def _noop_next_handler(**_kwargs: Any) -> AsyncGenerator: + """Stand-in for the downstream reasoning chain. Yields nothing — + InboxMiddleware should run its inbox-drain step first, then exit + the iteration cleanly.""" + return + yield # pragma: no cover # pylint: disable=unreachable + + +async def _drain(gen: AsyncGenerator) -> list: + """Collect every value yielded by an async generator.""" + out: list = [] + async for item in gen: + out.append(item) + return out + + +def _push_hint( + bus: _FakeBus, + sid: str, + hint: HintBlock, +) -> None: + """Push a :class:`HintBlock` into the per-session inbox key the + middleware will drain.""" + key = MessageBus._INBOX_KEY.format(sid=sid) + # asyncio.run on the bus would be overkill — `_queues` is a plain + # dict, mutate it directly. + bus._queues.setdefault(key, []).append( + (bus._alloc_id(), hint.model_dump(mode="json")), + ) + + +class TestInboxMiddlewareEmptyInbox(IsolatedAsyncioTestCase): + """When the inbox is empty, ``on_reasoning`` injects nothing, yields + nothing, and delegates straight to ``next_handler``.""" + + async def test_empty_inbox_is_noop(self) -> None: + """An empty inbox produces no events and leaves context untouched.""" + bus = _FakeBus() + agent = _make_agent( + name="A", + session_id="s", + reply_id=uuid.uuid4().hex, + context=[], + ) + mw = InboxMiddleware(bus) + + out = await _drain( + mw.on_reasoning(agent, {}, _noop_next_handler), + ) + + self.assertEqual(out, []) + self.assertEqual(agent.state.context, []) + + +class TestInboxMiddlewareInjection(IsolatedAsyncioTestCase): + """Branch coverage for the three injection cases.""" + + async def test_extends_into_last_assistant_msg_from_same_agent( + self, + ) -> None: + """When the last context msg is already an assistant msg from + this agent, hints are extended into its ``content`` and a new + msg is NOT appended.""" + bus = _FakeBus() + existing = AssistantMsg( + name="A", + content=[TextBlock(text="hello")], + ) + agent = _make_agent( + name="A", + session_id="s", + reply_id=uuid.uuid4().hex, + context=[existing], + ) + hint = HintBlock(hint="poke", source="tester") + _push_hint(bus, "s", hint) + mw = InboxMiddleware(bus) + + events = await _drain( + mw.on_reasoning(agent, {}, _noop_next_handler), + ) + + # No new message appended; hint extended into existing assistant msg. + self.assertEqual(len(agent.state.context), 1) + self.assertIs(agent.state.context[0], existing) + self.assertDictEqual( + existing.model_dump(), + { + "id": AnyString(), + "name": "A", + "role": "assistant", + "content": [ + {"type": "text", "text": "hello", "id": AnyString()}, + { + "type": "hint", + "hint": "poke", + "id": hint.id, + "source": "tester", + }, + ], + "metadata": {}, + "created_at": AnyString(), + "finished_at": None, + "usage": None, + }, + ) + # One HintBlockEvent yielded. + self.assertEqual( + [e.model_dump(mode="json") for e in events], + [ + { + "type": "HINT_BLOCK", + "id": AnyString(), + "created_at": AnyString(), + "metadata": {}, + "reply_id": agent.state.reply_id, + "block_id": hint.id, + "source": "tester", + "hint": "poke", + }, + ], + ) + + async def test_appends_new_msg_when_last_is_different_agent( + self, + ) -> None: + """When the last context msg is from a different agent (or a + system msg), a fresh :class:`AssistantMsg` is appended with the + hints as its content.""" + bus = _FakeBus() + system_msg = SystemMsg(name="system", content="boot") + agent = _make_agent( + name="A", + session_id="s", + reply_id="rid-1", + context=[system_msg], + ) + hint = HintBlock(hint="hi", source="x") + _push_hint(bus, "s", hint) + mw = InboxMiddleware(bus) + + await _drain(mw.on_reasoning(agent, {}, _noop_next_handler)) + + self.assertEqual(len(agent.state.context), 2) + self.assertDictEqual( + agent.state.context[-1].model_dump(), + { + "id": "rid-1", + "name": "A", + "role": "assistant", + "content": [ + { + "type": "hint", + "hint": "hi", + "id": hint.id, + "source": "x", + }, + ], + "metadata": {}, + "created_at": AnyString(), + "finished_at": None, + "usage": None, + }, + ) + + async def test_appends_new_msg_when_context_empty(self) -> None: + """When ``agent.state.context`` is empty, a fresh + :class:`AssistantMsg` containing the hints is the first entry.""" + bus = _FakeBus() + agent = _make_agent( + name="A", + session_id="s", + reply_id="rid-empty", + context=[], + ) + hint = HintBlock(hint="hi", source="x") + _push_hint(bus, "s", hint) + mw = InboxMiddleware(bus) + + await _drain(mw.on_reasoning(agent, {}, _noop_next_handler)) + + self.assertEqual(len(agent.state.context), 1) + self.assertDictEqual( + agent.state.context[0].model_dump(), + { + "id": "rid-empty", + "name": "A", + "role": "assistant", + "content": [ + { + "type": "hint", + "hint": "hi", + "id": hint.id, + "source": "x", + }, + ], + "metadata": {}, + "created_at": AnyString(), + "finished_at": None, + "usage": None, + }, + ) + + +class TestInboxMiddlewareYieldsHintBlockEvents(IsolatedAsyncioTestCase): + """One ``HintBlockEvent`` is yielded per injected hint, in order, + each carrying the hint's own block_id / source / content.""" + + async def test_event_count_and_payload(self) -> None: + """One HintBlockEvent is yielded per inbox hint, in arrival order.""" + bus = _FakeBus() + agent = _make_agent( + name="A", + session_id="s", + reply_id="rid-evt", + context=[], + ) + h1 = HintBlock(hint="a", source="alice") + h2 = HintBlock(hint="b", source="bob") + _push_hint(bus, "s", h1) + _push_hint(bus, "s", h2) + mw = InboxMiddleware(bus) + + events = await _drain( + mw.on_reasoning(agent, {}, _noop_next_handler), + ) + + self.assertEqual( + [e.model_dump(mode="json") for e in events], + [ + { + "type": "HINT_BLOCK", + "id": AnyString(), + "created_at": AnyString(), + "metadata": {}, + "reply_id": "rid-evt", + "block_id": h1.id, + "source": "alice", + "hint": "a", + }, + { + "type": "HINT_BLOCK", + "id": AnyString(), + "created_at": AnyString(), + "metadata": {}, + "reply_id": "rid-evt", + "block_id": h2.id, + "source": "bob", + "hint": "b", + }, + ], + ) + + +class TestInboxMiddlewareDelegatesDownstream(IsolatedAsyncioTestCase): + """After the inbox is drained, events from ``next_handler`` come + out of the middleware in order.""" + + async def test_downstream_events_pass_through(self) -> None: + """Downstream events appear after any HintBlockEvents.""" + bus = _FakeBus() + agent = _make_agent( + name="A", + session_id="s", + reply_id="rid", + context=[], + ) + hint = HintBlock(hint="hi", source="x") + _push_hint(bus, "s", hint) + mw = InboxMiddleware(bus) + + async def downstream(**_k: Any) -> AsyncGenerator[str, None]: + yield "ds-1" + yield "ds-2" + + out = await _drain(mw.on_reasoning(agent, {}, downstream)) + + # 1 HintBlockEvent followed by 2 downstream items. + self.assertEqual(len(out), 3) + self.assertDictEqual( + out[0].model_dump(mode="json"), + { + "type": "HINT_BLOCK", + "id": AnyString(), + "created_at": AnyString(), + "metadata": {}, + "reply_id": "rid", + "block_id": hint.id, + "source": "x", + "hint": "hi", + }, + ) + self.assertEqual(out[1:], ["ds-1", "ds-2"]) diff --git a/tests/service_index_task_consumer_test.py b/tests/service_index_task_consumer_test.py new file mode 100644 index 0000000000000000000000000000000000000000..c6ed78c8c1acbf8e9933aa36da768a190b6b61f4 --- /dev/null +++ b/tests/service_index_task_consumer_test.py @@ -0,0 +1,399 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Tests for :class:`IndexTaskConsumer` — the worker-process side of +the message-bus index dispatch flow. + +Verifies the four behaviours callers rely on: + +- Lifecycle is purely ACM: ``__aenter__`` starts the loop and performs + an initial drain; ``__aexit__`` cancels the loop and any in-flight + ``worker.process`` task cleanly. +- A signal triggers a queue drain; each entry is dispatched as a + ``worker.process`` call. +- Entries left on the queue from before ``__aenter__`` are picked up + on the initial drain without waiting for a fresh signal. +- Malformed entries are logged and skipped, not raised; later valid + entries still dispatch. +- An exception inside ``worker.process`` is logged but does not crash + the consumer loop — subsequent signals still drain. + +Mirrors :mod:`service_wakeup_dispatcher_test` deliberately; the +flows are identical at the primitive level. +""" +import asyncio +from contextlib import asynccontextmanager +from typing import AsyncGenerator, Callable +from unittest import IsolatedAsyncioTestCase + +from agentscope.app._service import IndexTaskConsumer +from agentscope.app.message_bus import MessageBus, MessageBusKeys + + +class _FakeBus(MessageBus): + """In-memory bus with just enough behaviour for the consumer. + + Implements the primitives the consumer actually uses + (``queue_push`` / ``queue_drain`` / ``subscribe`` / ``publish``) + and stubs the others. The consumer reaches for the bus through + the well-known channel/key constants — no domain methods exist + on the bus. + """ + + def __init__(self) -> None: + self.queues: dict[str, list[tuple[str, dict]]] = {} + self._channels: dict[str, asyncio.Queue] = {} + self._next = 0 + self._locks: set[str] = set() + + def _channel(self, key: str) -> asyncio.Queue: + return self._channels.setdefault(key, asyncio.Queue()) + + # Mode A — queue + async def queue_push( + self, + key: str, + payload: dict, + *, + ttl_secs: int | None = None, + ) -> str: + self._next += 1 + entry_id = str(self._next) + self.queues.setdefault(key, []).append((entry_id, payload)) + return entry_id + + async def queue_drain( + self, + key: str, + *, + max_count: int, + ) -> list[tuple[str, dict]]: + entries = self.queues.get(key, [])[:max_count] + self.queues[key] = self.queues.get(key, [])[max_count:] + return entries + + async def queue_delete(self, key: str) -> None: + self.queues.pop(key, None) + + # Mode C — log (unused) + async def log_append( + self, + key: str, + payload: dict, + *, + max_len: int | None = None, + ttl_secs: int | None = None, + ) -> str: + return "n/a" + + async def log_read( + self, + key: str, + since: str | None = None, + max_count: int = 100, + ) -> list[tuple[str, dict]]: + return [] + + async def log_trim( + self, + key: str, + before_id: str | None = None, + ) -> None: + return None + + # Mode D — pub/sub + async def publish(self, key: str, payload: dict) -> None: + await self._channel(key).put(payload) + + async def subscribe( + self, + key: str, + *, + on_ready: Callable[[], None] | None = None, + ) -> AsyncGenerator[dict, None]: + if on_ready is not None: + on_ready() + while True: + yield await self._channel(key).get() + + # Mode E — lock + @asynccontextmanager + async def acquire_lock( + self, + key: str, + *, + ttl_secs: int = 600, + ) -> AsyncGenerator[None, None]: + self._locks.add(key) + try: + yield + finally: + self._locks.discard(key) + + async def is_locked(self, key: str) -> bool: + return key in self._locks + + # Mode F — registry (unused by IndexTaskConsumer; raise so any + # accidental dependency surfaces immediately rather than silently + # passing through a stub). + async def registry_set( + self, + namespace: str, + field: str, + value: str, + *, + ttl_secs: int | None = None, + ) -> None: + raise NotImplementedError + + async def registry_del(self, namespace: str, field: str) -> None: + raise NotImplementedError + + async def registry_exists(self, namespace: str, field: str) -> bool: + raise NotImplementedError + + async def registry_getall(self, namespace: str) -> dict[str, str]: + raise NotImplementedError + + async def registry_drop(self, namespace: str) -> None: + raise NotImplementedError + + +class _RecordingWorker: + """Records calls to :meth:`process` so tests can assert dispatch. + + The real ``IndexWorker`` has many more methods, but the consumer + only ever invokes :meth:`process` — so the test surface stays + narrow on purpose. + """ + + def __init__(self) -> None: + self.calls: list[dict] = [] + self.notify = asyncio.Event() + self.fail_on_doc: str | None = None + self.process_delay = 0.0 + + async def process( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + ) -> None: + """Record the call and signal the test, optionally failing or + sleeping first to exercise the consumer's error and timing + paths. + + Args: + user_id (`str`): + The owning user id forwarded by the consumer. + knowledge_base_id (`str`): + The knowledge base id forwarded by the consumer. + document_id (`str`): + The document id forwarded by the consumer. + + Raises: + `RuntimeError`: + When ``fail_on_doc`` matches ``document_id`` — used by + the consumer-resilience tests. + """ + if self.process_delay: + await asyncio.sleep(self.process_delay) + self.calls.append( + { + "user_id": user_id, + "knowledge_base_id": knowledge_base_id, + "document_id": document_id, + }, + ) + self.notify.set() + if self.fail_on_doc and document_id == self.fail_on_doc: + raise RuntimeError(f"boom for {document_id}") + + +async def _yield_a_few_times(ticks: int = 8) -> None: + """Yield the event loop a few times so spawned tasks make progress.""" + for _ in range(ticks): + await asyncio.sleep(0) + + +class TestIndexTaskConsumerDispatch(IsolatedAsyncioTestCase): + """Verifies the signal-driven dispatch path.""" + + async def test_signal_drives_dispatch(self) -> None: + """A signal causes the queue to be drained and each entry + dispatched as ``worker.process``.""" + bus = _FakeBus() + worker = _RecordingWorker() + async with IndexTaskConsumer(message_bus=bus, worker=worker): + await bus.queue_push( + MessageBusKeys.index_tasks_queue(), + { + "user_id": "u", + "knowledge_base_id": "kb", + "document_id": "d1", + }, + ) + await bus.publish(MessageBusKeys.index_tasks_signal(), {}) + + await asyncio.wait_for(worker.notify.wait(), timeout=2.0) + + self.assertEqual( + worker.calls, + [ + { + "user_id": "u", + "knowledge_base_id": "kb", + "document_id": "d1", + }, + ], + ) + + async def test_initial_drain_picks_up_pending_entries(self) -> None: + """Entries on the queue from before ``__aenter__`` are picked + up without waiting for a fresh signal.""" + bus = _FakeBus() + worker = _RecordingWorker() + await bus.queue_push( + MessageBusKeys.index_tasks_queue(), + { + "user_id": "u", + "knowledge_base_id": "kb", + "document_id": "pre", + }, + ) + + async with IndexTaskConsumer(message_bus=bus, worker=worker): + await asyncio.wait_for(worker.notify.wait(), timeout=2.0) + + self.assertEqual( + worker.calls, + [ + { + "user_id": "u", + "knowledge_base_id": "kb", + "document_id": "pre", + }, + ], + ) + + async def test_malformed_entry_skipped(self) -> None: + """A queue entry missing required fields is logged and + skipped, not raised; later valid entries still dispatch.""" + bus = _FakeBus() + worker = _RecordingWorker() + + async with IndexTaskConsumer(message_bus=bus, worker=worker): + await bus.queue_push( + MessageBusKeys.index_tasks_queue(), + {"oops": True}, + ) + await bus.queue_push( + MessageBusKeys.index_tasks_queue(), + { + "user_id": "u", + "knowledge_base_id": "kb", + "document_id": "d2", + }, + ) + await bus.publish(MessageBusKeys.index_tasks_signal(), {}) + await asyncio.wait_for(worker.notify.wait(), timeout=2.0) + + self.assertEqual( + worker.calls, + [ + { + "user_id": "u", + "knowledge_base_id": "kb", + "document_id": "d2", + }, + ], + ) + + async def test_worker_exception_is_isolated(self) -> None: + """An exception inside ``worker.process`` for one entry does + not stop subsequent entries from dispatching.""" + bus = _FakeBus() + worker = _RecordingWorker() + worker.fail_on_doc = "boom" + + async with IndexTaskConsumer(message_bus=bus, worker=worker): + await bus.queue_push( + MessageBusKeys.index_tasks_queue(), + { + "user_id": "u", + "knowledge_base_id": "kb", + "document_id": "boom", + }, + ) + await bus.publish(MessageBusKeys.index_tasks_signal(), {}) + # Reset notify so we can wait for the *second* dispatch. + await _yield_a_few_times() + worker.notify.clear() + + await bus.queue_push( + MessageBusKeys.index_tasks_queue(), + { + "user_id": "u", + "knowledge_base_id": "kb", + "document_id": "ok", + }, + ) + await bus.publish(MessageBusKeys.index_tasks_signal(), {}) + await asyncio.wait_for(worker.notify.wait(), timeout=2.0) + + doc_ids = [c["document_id"] for c in worker.calls] + self.assertEqual(doc_ids, ["boom", "ok"]) + + +class TestIndexTaskConsumerLifecycle(IsolatedAsyncioTestCase): + """Tests covering the ``__aenter__`` / ``__aexit__`` ACM behaviour.""" + + async def test_exit_cancels_loop_cleanly(self) -> None: + """``__aexit__`` cancels the consumer's loop task and returns + without re-raising the cancellation.""" + bus = _FakeBus() + worker = _RecordingWorker() + consumer = IndexTaskConsumer(message_bus=bus, worker=worker) + + # pylint: disable=unnecessary-dunder-call + await consumer.__aenter__() + loop_task = consumer._task + self.assertIsNotNone(loop_task) + + await consumer.__aexit__(None, None, None) + + self.assertIsNone(consumer._task) + self.assertTrue(loop_task.cancelled() or loop_task.done()) + + async def test_exit_drains_inflight_tasks(self) -> None: + """``__aexit__`` cancels in-flight ``worker.process`` tasks + and waits for them to settle so exceptions surface in logs + and the event loop does not close on top of running + coroutines.""" + bus = _FakeBus() + worker = _RecordingWorker() + # Make worker.process slow so the consumer exits while one + # process call is still mid-flight. + worker.process_delay = 0.2 + + consumer = IndexTaskConsumer(message_bus=bus, worker=worker) + # pylint: disable=unnecessary-dunder-call + await consumer.__aenter__() + + await bus.queue_push( + MessageBusKeys.index_tasks_queue(), + { + "user_id": "u", + "knowledge_base_id": "kb", + "document_id": "slow", + }, + ) + await bus.publish(MessageBusKeys.index_tasks_signal(), {}) + # Give the loop a tick to start the worker task. + await asyncio.sleep(0) + self.assertGreaterEqual(len(consumer._inflight), 1) + + await consumer.__aexit__(None, None, None) + + # The in-flight set is cleared and the consumer task is gone. + self.assertEqual(consumer._inflight, set()) + self.assertIsNone(consumer._task) diff --git a/tests/service_knowledge_base_upload_test.py b/tests/service_knowledge_base_upload_test.py new file mode 100644 index 0000000000000000000000000000000000000000..2b03940b3b25ef97d79a0b9c545f6e488a5dc66e --- /dev/null +++ b/tests/service_knowledge_base_upload_test.py @@ -0,0 +1,444 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""End-to-end wiring test for the knowledge-base upload pipeline. + +Boots the full FastAPI app via :func:`create_app` against fakeredis + +in-memory KB-side fakes, then drives the upload → status → list → +delete flow through ``TestClient``. Verifies that: + +* ``create_app`` wires the new ``blob_store`` / dispatcher / sweeper / + service into ``app.state`` correctly; +* the upload endpoint streams bytes into the blob store, persists a + ``pending`` record, and dispatches the worker; +* the in-process worker drives the record to ``ready`` through the + parse → chunk → index phases; +* the status / list endpoints surface the lifecycle correctly; +* delete tears down the vector-store and storage records together. +""" +import asyncio +import tempfile +from typing import Any +from unittest.async_case import IsolatedAsyncioTestCase + +import fakeredis.aioredis +from fastapi.testclient import TestClient + +from agentscope.app import create_app +from agentscope.app.rag.blob_store import LocalBlobStore +from agentscope.app.rag.knowledge_base_manager import ( + KnowledgeBaseManagerBase, + KnowledgeBaseNotFoundError, +) +from agentscope.app.rag.knowledge_base_manager._dimension_policy import ( + DimensionPolicy, + DimensionPolicyKind, +) +from agentscope.app.message_bus import RedisMessageBus +from agentscope.app.storage import ( + EmbeddingModelConfig, + KnowledgeBaseRecord, + RedisStorage, +) +from agentscope.app.workspace_manager._base import WorkspaceManagerBase +from agentscope.rag import VectorStoreBase +from agentscope.rag._vdb._vector_store import ( + DocumentSummary, + VectorRecord, + VectorSearchResult, +) + + +# ---------------------------------------------------------------------- +# Test doubles +# ---------------------------------------------------------------------- + + +class _FakeVectorStore(VectorStoreBase): + """In-memory vector store — records are kept by collection. + + The worker never reads back from it, so there is no need for real + similarity math; we only have to honour the ``insert`` / ``delete`` + / ``has_collection`` contract. + """ + + def __init__(self) -> None: + self._collections: dict[str, list[VectorRecord]] = {} + + async def create_collection(self, name: str, dimensions: int) -> None: + self._collections.setdefault(name, []) + + async def delete_collection(self, name: str) -> None: + self._collections.pop(name, None) + + async def has_collection(self, name: str) -> bool: + return name in self._collections + + async def insert( + self, + collection: str, + records: list[VectorRecord], + ) -> None: + self._collections.setdefault(collection, []).extend(records) + + async def delete(self, collection: str, document_id: str) -> None: + bucket = self._collections.get(collection) + if bucket is None: + return + self._collections[collection] = [ + r for r in bucket if r.document_id != document_id + ] + + async def search( + self, + collection: str, + query_vector: list[float], + top_k: int = 5, + metadata_filter: dict[str, Any] | None = None, + ) -> list[VectorSearchResult]: + return [] + + async def list_documents( + self, + collection: str, + metadata_filter: dict[str, Any] | None = None, + ) -> list[DocumentSummary]: + return [] + + +class _FakeKnowledge: + """Minimal stand-in for :class:`KnowledgeBase` used by the worker. + + Bypasses embedding-model construction — instead just funnels the + chunks into the bound :class:`_FakeVectorStore` with a fixed + zero-vector so ``insert_document`` succeeds end-to-end. + """ + + def __init__( + self, + vector_store: _FakeVectorStore, + collection_name: str, + ) -> None: + self._vector_store = vector_store + self._collection_name = collection_name + + async def insert_document( + self, + chunks: list, + document_id: str | None = None, + document_metadata: dict | None = None, + ) -> str: + """Pretend to embed and insert ``chunks`` into the bound store. + + The fake skips the real embedding step — it stamps a single + scalar vector on every record so the upload pipeline can be + exercised without an embedding model. + + Args: + chunks (`list`): + The parsed and chunked document content. + document_id (`str | None`, optional): + Caller-supplied document id; the fake just echoes it + back rather than generating a UUID. + document_metadata (`dict | None`, optional): + Document-level metadata; ignored — the upload tests + don't assert on metadata propagation. + + Returns: + `str`: + The (caller-supplied) document id, or ``""`` when + none was passed. + """ + del document_metadata # unused — see docstring + records = [ + VectorRecord( + vector=[0.0], + document_id=document_id or "", + chunk=chunk, + ) + for chunk in chunks + ] + await self._vector_store.insert(self._collection_name, records) + return document_id or "" + + async def delete_document(self, document_id: str) -> None: + """Remove every record for ``document_id`` from the bound store. + + Args: + document_id (`str`): + The document whose records should be deleted. + """ + await self._vector_store.delete(self._collection_name, document_id) + + async def search(self, queries: list, top_k: int = 5) -> list: + """Return an empty result list — search is out of scope here. + + Args: + queries (`list`): + The query inputs; ignored. + top_k (`int`, defaults to ``5``): + The maximum result count; ignored. + + Returns: + `list`: + Always empty — the upload tests do not exercise + retrieval. + """ + del queries, top_k # unused — see docstring + return [] + + +class _FakeKbManager(KnowledgeBaseManagerBase): + """KB manager that uses the storage + a fake vector store directly. + + Skips the real ``CollectionPerKbManager`` so we don't need a live + embedding model — the indexing pipeline only requires + ``insert_document`` / ``delete_document``, which the + :class:`_FakeKnowledge` returned here implements directly. + """ + + async def get_dimension_policy(self) -> DimensionPolicy: + return DimensionPolicy(kind=DimensionPolicyKind.ANY, dimension=None) + + async def create_knowledge_base( + self, + user_id: str, + name: str, + description: str, + embedding_model_config: EmbeddingModelConfig, + ) -> KnowledgeBaseRecord: + record = KnowledgeBaseRecord( + user_id=user_id, + name=name, + description=description, + embedding_model_config=embedding_model_config, + collection_name="", + ) + record.collection_name = f"kb_{record.id}" + await self._vector_store.create_collection( + name=record.collection_name, + dimensions=embedding_model_config.dimensions, + ) + return await self._storage.upsert_knowledge_base(user_id, record) + + async def delete_knowledge_base( + self, + user_id: str, + knowledge_base_id: str, + ) -> bool: + record = await self._storage.get_knowledge_base( + user_id, + knowledge_base_id, + ) + if record is None: + return False + await self._vector_store.delete_collection(record.collection_name) + return await self._storage.delete_knowledge_base( + user_id, + knowledge_base_id, + ) + + async def get_knowledge( + self, + user_id: str, + knowledge_base_id: str, + ) -> _FakeKnowledge: + record = await self._storage.get_knowledge_base( + user_id, + knowledge_base_id, + ) + if record is None: + raise KnowledgeBaseNotFoundError( + f"Knowledge base {knowledge_base_id!r} not found.", + ) + return _FakeKnowledge( + vector_store=self._vector_store, + collection_name=record.collection_name, + ) + + +class _NoopWorkspaceManager(WorkspaceManagerBase): + """Workspace manager that does nothing — the KB pipeline never + touches it, but ``create_app`` requires one to be wired in.""" + + async def get_workspace(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError + + async def create_workspace(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError + + async def close(self, workspace_id: str) -> None: + return None + + async def close_all(self) -> None: + return None + + +def _make_storage(fr: fakeredis.aioredis.FakeRedis) -> RedisStorage: + """Build a RedisStorage already bound to *fr*. + + Pre-populates ``_client`` so the lifespan's ``__aenter__`` no-ops on + the connection-pool side and just reuses our fakeredis handle. + """ + + class _FakeStorage(RedisStorage): + async def __aenter__(self) -> "_FakeStorage": + self._client = fr + return self + + async def aclose(self) -> None: + self._client = None + + return _FakeStorage() + + +def _make_bus(fr: fakeredis.aioredis.FakeRedis) -> RedisMessageBus: + """Build a RedisMessageBus bound to *fr* (same trick as in the bus + tests).""" + + class _FakeBus(RedisMessageBus): + async def __aenter__(self) -> "_FakeBus": + self._client = fr + return self + + async def aclose(self) -> None: + self._client = None + + return _FakeBus() + + +# ---------------------------------------------------------------------- +# Tests +# ---------------------------------------------------------------------- + + +class KnowledgeBaseUploadFlowTest(IsolatedAsyncioTestCase): + """End-to-end wiring of the upload pipeline through ``TestClient``.""" + + async def asyncSetUp(self) -> None: + self._tmp = tempfile.TemporaryDirectory() + self._fr = fakeredis.aioredis.FakeRedis(decode_responses=True) + self._vector_store = _FakeVectorStore() + storage = _make_storage(self._fr) + message_bus = _make_bus(self._fr) + + self._app = create_app( + storage=storage, + message_bus=message_bus, + workspace_manager=_NoopWorkspaceManager(), + knowledge_base_manager=_FakeKbManager( + storage=storage, + vector_store=self._vector_store, + ), + blob_store=LocalBlobStore(root_dir=self._tmp.name), + ) + # Seed a knowledge base directly through storage so we don't + # have to mock the manager's create flow over HTTP. + kb_record = KnowledgeBaseRecord( + user_id="user-1", + name="kb", + description="", + embedding_model_config=EmbeddingModelConfig( + type="openai_credential", + credential_id="cred-1", + model="text-embedding-3-small", + dimensions=1, + ), + collection_name="", + ) + kb_record.collection_name = f"kb_{kb_record.id}" + await self._vector_store.create_collection( + kb_record.collection_name, + 1, + ) + + # Drop the seed record into fakeredis directly — we need the KB + # in place before lifespan starts the sweeper. + storage._client = self._fr + await storage.upsert_knowledge_base("user-1", kb_record) + storage._client = None + self._kb_id = kb_record.id + + async def asyncTearDown(self) -> None: + await self._fr.aclose() + self._tmp.cleanup() + + async def test_upload_drives_document_to_ready(self) -> None: + """Upload a small text file and observe the lifecycle.""" + headers = {"X-User-ID": "user-1"} + with TestClient(self._app) as client: + files = { + "file": ( + "hello.txt", + b"hello world\n" * 16, + "text/plain", + ), + } + resp = client.post( + f"/knowledge_bases/{self._kb_id}/documents", + files=files, + headers=headers, + ) + self.assertEqual(resp.status_code, 201, resp.text) + body = resp.json() + document_id = body["document_id"] + self.assertEqual(body["filename"], "hello.txt") + self.assertIn(body["status"], ("pending", "ready")) + + # Wait for the in-process worker to drive the record to + # ``ready``. We poll with a generous overall timeout — the + # actual work is < 100 ms but CI machines can be slow. + deadline = 5.0 + poll = 0.05 + elapsed = 0.0 + final_status = body["status"] + while elapsed < deadline: + resp = client.get( + f"/knowledge_bases/{self._kb_id}/documents/status", + params={"ids": document_id}, + headers=headers, + ) + self.assertEqual(resp.status_code, 200, resp.text) + items = resp.json()["items"] + self.assertEqual(len(items), 1, items) + final_status = items[0]["status"] + if final_status in ("ready", "error"): + break + await asyncio.sleep(poll) + elapsed += poll + self.assertEqual(final_status, "ready", items) + + # Listing shows the document with the same ``ready`` state. + resp = client.get( + f"/knowledge_bases/{self._kb_id}/documents", + headers=headers, + ) + self.assertEqual(resp.status_code, 200, resp.text) + documents = resp.json()["documents"] + self.assertEqual(len(documents), 1) + self.assertEqual(documents[0]["id"], document_id) + self.assertEqual(documents[0]["status"], "ready") + self.assertGreaterEqual(documents[0]["chunk_count"], 1) + + # Delete and confirm the listing comes back empty. + resp = client.delete( + f"/knowledge_bases/{self._kb_id}/documents/{document_id}", + headers=headers, + ) + self.assertEqual(resp.status_code, 204, resp.text) + resp = client.get( + f"/knowledge_bases/{self._kb_id}/documents", + headers=headers, + ) + self.assertEqual(resp.json()["documents"], []) + + async def test_status_for_unknown_id_is_silently_skipped(self) -> None: + """Asking for a non-existent doc returns an empty items list.""" + headers = {"X-User-ID": "user-1"} + with TestClient(self._app) as client: + resp = client.get( + f"/knowledge_bases/{self._kb_id}/documents/status", + params={"ids": "does-not-exist"}, + headers=headers, + ) + self.assertEqual(resp.status_code, 200, resp.text) + self.assertEqual(resp.json()["items"], []) diff --git a/tests/service_message_bus_test.py b/tests/service_message_bus_test.py new file mode 100644 index 0000000000000000000000000000000000000000..6edea09dca111be9fad86dd2c1ba3e9c4f95c716 --- /dev/null +++ b/tests/service_message_bus_test.py @@ -0,0 +1,591 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Tests for :class:`RedisMessageBus` and the domain helpers on the base +:class:`MessageBus` class. + +The Redis backend is exercised against ``fakeredis`` so tests cover both +the abstract surface (queue / log / pubsub / lock) and the domain helpers +(``session_run`` / ``session_publish_event`` / ``inbox_*`` / ``wakeup_*``) +that are layered on top. +""" +import asyncio +from contextlib import AsyncExitStack +from unittest import IsolatedAsyncioTestCase + +import fakeredis.aioredis + +from agentscope.app.message_bus import MessageBus, RedisMessageBus + + +def _make_bus( + fake_redis: fakeredis.aioredis.FakeRedis, +) -> RedisMessageBus: + """Construct a :class:`RedisMessageBus` that uses *fake_redis*. + + The bus subclass overrides ``__aenter__`` so it talks to fakeredis + instead of opening a real connection pool. + + Args: + fake_redis (`fakeredis.aioredis.FakeRedis`): + A fakeredis client whose pubsub / streams APIs are async. + + Returns: + `RedisMessageBus`: + A bus instance ready to be used as an async context manager. + """ + + class _FakeBus(RedisMessageBus): + """Bus subclass that returns the supplied fakeredis client on + context entry instead of building a real one.""" + + async def __aenter__(self) -> "RedisMessageBus": + self._client = fake_redis + return self + + async def aclose(self) -> None: + # The fakeredis client is owned by the test, not the bus. + self._client = None + + return _FakeBus() + + +class TestQueuePrimitive(IsolatedAsyncioTestCase): + """Mode A — ``queue_push`` + ``queue_drain`` semantics.""" + + async def asyncSetUp(self) -> None: + self.fr = fakeredis.aioredis.FakeRedis(decode_responses=True) + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context(_make_bus(self.fr)) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + await self.fr.aclose() + + async def test_push_drain_returns_payloads_in_order(self) -> None: + """Entries pushed in order come back out in order, once each.""" + await self.bus.queue_push("k", {"i": 1}) + await self.bus.queue_push("k", {"i": 2}) + entries = await self.bus.queue_drain("k", max_count=10) + self.assertEqual([p for _id, p in entries], [{"i": 1}, {"i": 2}]) + + async def test_drain_is_destructive(self) -> None: + """A drained entry is gone; a second drain yields nothing.""" + await self.bus.queue_push("k", {"x": 1}) + await self.bus.queue_drain("k", max_count=10) + self.assertEqual(await self.bus.queue_drain("k", max_count=10), []) + + async def test_drain_respects_max_count(self) -> None: + """``max_count`` caps the batch size; remaining entries persist.""" + for i in range(5): + await self.bus.queue_push("k", {"i": i}) + first = await self.bus.queue_drain("k", max_count=3) + rest = await self.bus.queue_drain("k", max_count=10) + self.assertEqual([p["i"] for _id, p in first], [0, 1, 2]) + self.assertEqual([p["i"] for _id, p in rest], [3, 4]) + + +class TestLogPrimitive(IsolatedAsyncioTestCase): + """Mode C — replay log: append / read with cursor / trim.""" + + async def asyncSetUp(self) -> None: + self.fr = fakeredis.aioredis.FakeRedis(decode_responses=True) + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context(_make_bus(self.fr)) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + await self.fr.aclose() + + async def test_read_returns_everything_when_no_cursor(self) -> None: + """Without a ``since`` cursor, the whole log comes back.""" + await self.bus.log_append("k", {"i": 1}) + await self.bus.log_append("k", {"i": 2}) + entries = await self.bus.log_read("k") + self.assertEqual([p["i"] for _id, p in entries], [1, 2]) + + async def test_read_with_cursor_is_exclusive(self) -> None: + """``since=last_id`` skips that id and returns only newer entries.""" + await self.bus.log_append("k", {"i": 1}) + await self.bus.log_append("k", {"i": 2}) + await self.bus.log_append("k", {"i": 3}) + all_entries = await self.bus.log_read("k") + cursor = all_entries[1][0] # id of entry 2 + rest = await self.bus.log_read("k", since=cursor) + self.assertEqual([p["i"] for _id, p in rest], [3]) + + async def test_trim_without_before_drops_entire_log(self) -> None: + """``log_trim(key)`` empties the log; subsequent read is empty.""" + await self.bus.log_append("k", {"i": 1}) + await self.bus.log_append("k", {"i": 2}) + await self.bus.log_trim("k") + self.assertEqual(await self.bus.log_read("k"), []) + + +class TestPubSubPrimitive(IsolatedAsyncioTestCase): + """Mode D — transient broadcast: publish / subscribe.""" + + async def asyncSetUp(self) -> None: + self.fr = fakeredis.aioredis.FakeRedis(decode_responses=True) + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context(_make_bus(self.fr)) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + await self.fr.aclose() + + async def test_subscribe_receives_messages_published_after_ready( + self, + ) -> None: + """Subscribers receive payloads published after the subscription + is established. The ``on_ready`` hook fires once before any + payload is yielded.""" + ready = asyncio.Event() + received: list[dict] = [] + + async def _consumer() -> None: + async for payload in self.bus.subscribe( + "ch", + on_ready=ready.set, + ): + received.append(payload) + if len(received) == 2: + break + + task = asyncio.create_task(_consumer()) + await asyncio.wait_for(ready.wait(), timeout=2.0) + + await self.bus.publish("ch", {"i": 1}) + await self.bus.publish("ch", {"i": 2}) + await asyncio.wait_for(task, timeout=2.0) + self.assertEqual([p["i"] for p in received], [1, 2]) + + +class TestLockPrimitive(IsolatedAsyncioTestCase): + """Mode E — distributed mutex semantics.""" + + async def asyncSetUp(self) -> None: + self.fr = fakeredis.aioredis.FakeRedis(decode_responses=True) + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context(_make_bus(self.fr)) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + await self.fr.aclose() + + async def test_is_locked_reflects_acquire_release(self) -> None: + """``is_locked`` flips to True while the body runs and back to + False once the context exits.""" + self.assertFalse(await self.bus.is_locked("k")) + async with self.bus.acquire_lock("k", ttl_secs=10): + self.assertTrue(await self.bus.is_locked("k")) + self.assertFalse(await self.bus.is_locked("k")) + + async def test_second_acquirer_waits_until_release(self) -> None: + """A second ``acquire_lock`` on the same key blocks until the + first releases.""" + order: list[str] = [] + + async def _holder() -> None: + async with self.bus.acquire_lock("k", ttl_secs=10): + order.append("first-in") + await asyncio.sleep(0.05) + order.append("first-out") + + async def _challenger() -> None: + # Tiny delay so the holder grabs the lock first. + await asyncio.sleep(0.005) + async with self.bus.acquire_lock("k", ttl_secs=10): + order.append("second-in") + + await asyncio.gather(_holder(), _challenger()) + self.assertEqual( + order, + ["first-in", "first-out", "second-in"], + ) + + +class TestSessionRunAutoTrimsLog(IsolatedAsyncioTestCase): + """``session_run.__aexit__`` must trim the session's replay log + *before* releasing the distributed lock, so any subscriber that + connects between two runs sees a clean slate.""" + + async def asyncSetUp(self) -> None: + self.fr = fakeredis.aioredis.FakeRedis(decode_responses=True) + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context(_make_bus(self.fr)) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + await self.fr.aclose() + + async def test_log_trim_happens_on_session_run_exit(self) -> None: + """Events published inside a ``session_run`` block are trimmed + once the block exits — a fresh ``session_read_events`` returns + an empty list.""" + sid = "s-trim" + async with self.bus.session_run(sid): + await self.bus.session_publish_event(sid, {"i": 1}) + await self.bus.session_publish_event(sid, {"i": 2}) + mid = await self.bus.session_read_events(sid) + self.assertEqual([p["i"] for _id, p in mid], [1, 2]) + self.assertEqual(await self.bus.session_read_events(sid), []) + + async def test_log_trim_runs_even_when_body_raises(self) -> None: + """If the body raises, the log is still trimmed before the lock + releases — the next run starts clean.""" + sid = "s-raise" + + class _Boom(RuntimeError): + """Marker exception raised inside the run body.""" + + with self.assertRaises(_Boom): + async with self.bus.session_run(sid): + await self.bus.session_publish_event(sid, {"i": 1}) + raise _Boom() + + self.assertEqual(await self.bus.session_read_events(sid), []) + + +class TestSessionDomainHelpers(IsolatedAsyncioTestCase): + """``session_publish_event`` + ``session_subscribe_events`` + + ``session_is_running`` round-trip behaviour.""" + + async def asyncSetUp(self) -> None: + self.fr = fakeredis.aioredis.FakeRedis(decode_responses=True) + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context(_make_bus(self.fr)) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + await self.fr.aclose() + + async def test_publish_event_writes_to_log_and_pubsub(self) -> None: + """``session_publish_event`` writes one entry to the replay log + AND fans it out on the live channel; subscribers receive the + payload with the ``_entry_id`` field stripped.""" + sid = "s-pub" + ready = asyncio.Event() + received: list[dict] = [] + + async def _consumer() -> None: + async for payload in self.bus.session_subscribe_events( + sid, + on_ready=ready.set, + ): + received.append(payload) + break + + task = asyncio.create_task(_consumer()) + await asyncio.wait_for(ready.wait(), timeout=2.0) + + await self.bus.session_publish_event(sid, {"hello": "world"}) + await asyncio.wait_for(task, timeout=2.0) + + # Replay log captured the entry too. + log_entries = await self.bus.session_read_events(sid) + self.assertEqual(len(log_entries), 1) + self.assertEqual(log_entries[0][1], {"hello": "world"}) + + # Live subscriber saw it without the internal _entry_id key. + self.assertEqual(received, [{"hello": "world"}]) + + async def test_session_is_running_reflects_session_run(self) -> None: + """``session_is_running`` returns True while inside + ``session_run`` and False after.""" + sid = "s-isrun" + self.assertFalse(await self.bus.session_is_running(sid)) + async with self.bus.session_run(sid): + self.assertTrue(await self.bus.session_is_running(sid)) + self.assertFalse(await self.bus.session_is_running(sid)) + + +class TestInboxAndWakeupHelpers(IsolatedAsyncioTestCase): + """Inbox + wakeup domain helpers used by team / tool-offload / + scheduler to deliver work to idle sessions.""" + + async def asyncSetUp(self) -> None: + self.fr = fakeredis.aioredis.FakeRedis(decode_responses=True) + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context(_make_bus(self.fr)) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + await self.fr.aclose() + + async def test_inbox_push_drain_round_trip(self) -> None: + """``inbox_push`` payloads are returned by ``inbox_drain`` in + push order, exactly once.""" + sid = "s-inbox" + await self.bus.inbox_push(sid, {"hint": "a"}) + await self.bus.inbox_push(sid, {"hint": "b"}) + entries = await self.bus.inbox_drain(sid, max_count=10) + self.assertEqual( + [p["hint"] for _id, p in entries], + ["a", "b"], + ) + self.assertEqual( + await self.bus.inbox_drain(sid, max_count=10), + [], + ) + + async def test_enqueue_wakeup_signals_and_queues(self) -> None: + """``enqueue_wakeup`` puts the payload on the durable queue and + fires the signal channel; a subscriber and a ``dequeue_wakeups`` + call both see it.""" + ready = asyncio.Event() + received: list[dict] = [] + + async def _signal_consumer() -> None: + async for payload in self.bus.subscribe_wakeup_signal( + on_ready=ready.set, + ): + received.append(payload) + break + + task = asyncio.create_task(_signal_consumer()) + await asyncio.wait_for(ready.wait(), timeout=2.0) + + await self.bus.enqueue_wakeup( + user_id="u", + session_id="s", + agent_id="a", + ) + await asyncio.wait_for(task, timeout=2.0) + + # Signal fired. + self.assertEqual(len(received), 1) + + # Queue holds the structured entry. ``enqueue_wakeup`` is the + # idle-wake shortcut, so the entry carries ``kind="wake"`` and a + # null input alongside the routing fields. + entries = await self.bus.dequeue_wakeups(max_count=10) + self.assertEqual(len(entries), 1) + self.assertEqual( + entries[0], + { + "user_id": "u", + "session_id": "s", + "agent_id": "a", + "kind": "wake", + "input": None, + }, + ) + + +class TestRegistryPrimitive(IsolatedAsyncioTestCase): + """Mode F — ``registry_*`` hash-keyed namespace operations.""" + + async def asyncSetUp(self) -> None: + self.fr = fakeredis.aioredis.FakeRedis(decode_responses=True) + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context(_make_bus(self.fr)) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + await self.fr.aclose() + + async def test_set_then_exists_and_getall(self) -> None: + """``registry_set`` stores a field under a namespace; ``exists`` + is True, ``getall`` returns the full mapping.""" + await self.bus.registry_set("ns", "f1", "v1") + await self.bus.registry_set("ns", "f2", "v2") + + self.assertTrue(await self.bus.registry_exists("ns", "f1")) + self.assertTrue(await self.bus.registry_exists("ns", "f2")) + self.assertFalse(await self.bus.registry_exists("ns", "missing")) + self.assertFalse(await self.bus.registry_exists("other-ns", "f1")) + + self.assertEqual( + await self.bus.registry_getall("ns"), + {"f1": "v1", "f2": "v2"}, + ) + + async def test_set_overwrites_existing_field(self) -> None: + """A second ``registry_set`` for the same field overwrites the + previous value (``HSET`` semantics).""" + await self.bus.registry_set("ns", "f", "v1") + await self.bus.registry_set("ns", "f", "v2") + self.assertEqual( + await self.bus.registry_getall("ns"), + {"f": "v2"}, + ) + + async def test_del_removes_only_the_named_field(self) -> None: + """``registry_del`` removes a single field; siblings survive.""" + await self.bus.registry_set("ns", "keep", "k") + await self.bus.registry_set("ns", "drop", "d") + + await self.bus.registry_del("ns", "drop") + + self.assertFalse(await self.bus.registry_exists("ns", "drop")) + self.assertTrue(await self.bus.registry_exists("ns", "keep")) + self.assertEqual( + await self.bus.registry_getall("ns"), + {"keep": "k"}, + ) + + async def test_del_missing_field_is_noop(self) -> None: + """Deleting a non-existent field does not raise.""" + await self.bus.registry_del("ns", "nope") + await self.bus.registry_set("ns", "keep", "k") + await self.bus.registry_del("ns", "still-missing") + self.assertEqual( + await self.bus.registry_getall("ns"), + {"keep": "k"}, + ) + + async def test_getall_on_missing_namespace_returns_empty_dict( + self, + ) -> None: + """``registry_getall`` for an unknown namespace returns ``{}`` + rather than ``None``.""" + self.assertEqual(await self.bus.registry_getall("ghost"), {}) + + async def test_drop_deletes_entire_namespace(self) -> None: + """``registry_drop`` removes every field under the namespace.""" + await self.bus.registry_set("ns", "f1", "v1") + await self.bus.registry_set("ns", "f2", "v2") + + await self.bus.registry_drop("ns") + + self.assertFalse(await self.bus.registry_exists("ns", "f1")) + self.assertFalse(await self.bus.registry_exists("ns", "f2")) + self.assertEqual(await self.bus.registry_getall("ns"), {}) + + async def test_drop_missing_namespace_is_noop(self) -> None: + """Dropping a namespace that was never written does not raise.""" + await self.bus.registry_drop("never-existed") + + async def test_set_with_ttl_applies_expire_and_refreshes(self) -> None: + """``registry_set`` with ``ttl_secs`` sets a TTL on the hash key; + a subsequent set with a longer ``ttl_secs`` refreshes it.""" + await self.bus.registry_set("ns", "f", "v", ttl_secs=60) + ttl_first = await self.fr.ttl("ns") + self.assertGreater(ttl_first, 0) + self.assertLessEqual(ttl_first, 60) + + # Refresh with a much larger TTL — must overwrite the old one. + await self.bus.registry_set("ns", "f", "v", ttl_secs=3600) + ttl_refreshed = await self.fr.ttl("ns") + self.assertGreater(ttl_refreshed, 60) + + async def test_set_without_ttl_leaves_namespace_persistent( + self, + ) -> None: + """Without ``ttl_secs`` the namespace has no expiry (TTL == -1).""" + await self.bus.registry_set("ns", "f", "v") + self.assertEqual(await self.fr.ttl("ns"), -1) + + +class TestBackgroundTaskRegistryHelpers(IsolatedAsyncioTestCase): + """Domain helpers built on Mode F: ``bg_task_register / unregister / + exists / list / purge`` plus ``task_publish_cancel / + task_subscribe_cancel``.""" + + async def asyncSetUp(self) -> None: + self.fr = fakeredis.aioredis.FakeRedis(decode_responses=True) + self._stack = AsyncExitStack() + self.bus = await self._stack.enter_async_context(_make_bus(self.fr)) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + await self.fr.aclose() + + async def test_register_then_exists_list_unregister(self) -> None: + """End-to-end registry round-trip for a single session.""" + sid = "s-bg" + + self.assertFalse(await self.bus.bg_task_exists(sid, "t1")) + + await self.bus.bg_task_register(sid, "t1", '{"tool":"a"}') + await self.bus.bg_task_register(sid, "t2", '{"tool":"b"}') + + self.assertTrue(await self.bus.bg_task_exists(sid, "t1")) + self.assertTrue(await self.bus.bg_task_exists(sid, "t2")) + self.assertEqual( + await self.bus.bg_task_list(sid), + {"t1": '{"tool":"a"}', "t2": '{"tool":"b"}'}, + ) + + await self.bus.bg_task_unregister(sid, "t1") + self.assertFalse(await self.bus.bg_task_exists(sid, "t1")) + self.assertTrue(await self.bus.bg_task_exists(sid, "t2")) + self.assertEqual( + await self.bus.bg_task_list(sid), + {"t2": '{"tool":"b"}'}, + ) + + async def test_register_isolates_sessions(self) -> None: + """Tasks registered under one session id are invisible to other + session ids.""" + await self.bus.bg_task_register("s1", "t", "{}") + + self.assertTrue(await self.bus.bg_task_exists("s1", "t")) + self.assertFalse(await self.bus.bg_task_exists("s2", "t")) + self.assertEqual(await self.bus.bg_task_list("s2"), {}) + + async def test_register_applies_fallback_ttl(self) -> None: + """``bg_task_register`` sets the per-session fallback TTL on the + hash key so abandoned entries can't accumulate forever.""" + sid = "s-ttl" + await self.bus.bg_task_register(sid, "t", "{}") + + ttl = await self.fr.ttl(self.bus._BG_TASKS_KEY.format(sid=sid)) + self.assertGreater(ttl, 0) + self.assertLessEqual(ttl, self.bus._BG_TASKS_TTL_SECS) + + # A second register-call refreshes the TTL back near the cap. + await asyncio.sleep(0) # let any fakeredis internals settle + await self.bus.bg_task_register(sid, "t2", "{}") + ttl_refreshed = await self.fr.ttl( + self.bus._BG_TASKS_KEY.format(sid=sid), + ) + self.assertGreater(ttl_refreshed, 0) + + async def test_purge_clears_all_session_entries(self) -> None: + """``bg_task_purge`` deletes every task entry for a session in + a single call (used during session deletion).""" + sid = "s-purge" + await self.bus.bg_task_register(sid, "t1", "{}") + await self.bus.bg_task_register(sid, "t2", "{}") + # Other session must survive. + await self.bus.bg_task_register("s-keep", "t", "{}") + + await self.bus.bg_task_purge(sid) + + self.assertEqual(await self.bus.bg_task_list(sid), {}) + self.assertFalse(await self.bus.bg_task_exists(sid, "t1")) + self.assertTrue(await self.bus.bg_task_exists("s-keep", "t")) + + async def test_task_publish_cancel_reaches_subscriber(self) -> None: + """``task_publish_cancel`` fans out to every active + ``task_subscribe_cancel`` listener; the yielded value is the + ``task_id`` from the payload.""" + ready = asyncio.Event() + received: list[str] = [] + + async def _consumer() -> None: + async for tid in self.bus.task_subscribe_cancel( + on_ready=ready.set, + ): + received.append(tid) + break + + task = asyncio.create_task(_consumer()) + await asyncio.wait_for(ready.wait(), timeout=2.0) + + await self.bus.task_publish_cancel("task-X") + await asyncio.wait_for(task, timeout=2.0) + + self.assertEqual(received, ["task-X"]) + + +class TestBaseClassIsAbstract(IsolatedAsyncioTestCase): + """``MessageBus`` itself cannot be instantiated.""" + + def test_instantiating_base_class_raises(self) -> None: + """The abstract methods are not implemented on the base — direct + instantiation must fail.""" + with self.assertRaises(TypeError): + # pylint: disable=abstract-class-instantiated + MessageBus() # type: ignore[abstract] diff --git a/tests/service_scheduler_test.py b/tests/service_scheduler_test.py new file mode 100644 index 0000000000000000000000000000000000000000..dbe31a80afdcde547385899fabc92482455e8bc5 --- /dev/null +++ b/tests/service_scheduler_test.py @@ -0,0 +1,267 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Tests for :meth:`SchedulerManager._build_trigger`. + +We don't drive APScheduler here — we ask the manager to build a trigger +coroutine for a record and invoke it directly. The trigger's contract is: + +- when ``ScheduleData.enabled`` is False → no side effects; +- when enabled → resolve / create a target session, push a + ````-wrapped :class:`HintBlock` to the session inbox, + and enqueue one wakeup pointing at that session. + +In stateful mode the session id is deterministic (``{record_id}_stateful``) +and reused across fires; in non-stateful mode a fresh session id is +created every fire. +""" +import json +from contextlib import AsyncExitStack +from datetime import datetime +from unittest import IsolatedAsyncioTestCase + +import fakeredis.aioredis + +from utils import AnyString + +from agentscope.app._manager import SchedulerManager +from agentscope.app.message_bus import RedisMessageBus +from agentscope.app.storage import ( + ChatModelConfig, + RedisStorage, + ScheduleData, + ScheduleRecord, + SessionSource, +) +from agentscope.permission import PermissionMode + + +def _make_storage( + fr: fakeredis.aioredis.FakeRedis, +) -> RedisStorage: + """Construct a :class:`RedisStorage` bound to *fr*.""" + + class _S(RedisStorage): + async def __aenter__(self) -> "RedisStorage": # type: ignore[override] + self._client = fr + return self + + async def aclose(self) -> None: + self._client = None + + return _S() + + +def _make_bus( + fr: fakeredis.aioredis.FakeRedis, +) -> RedisMessageBus: + """Construct a :class:`RedisMessageBus` bound to *fr*.""" + + class _B(RedisMessageBus): + async def __aenter__( # type: ignore[override] + self, + ) -> "RedisMessageBus": + self._client = fr + return self + + async def aclose(self) -> None: + self._client = None + + return _B() + + +def _make_record( + *, + user_id: str = "u", + agent_id: str = "a", + enabled: bool = True, + stateful: bool = False, + description: str = "run nightly summary", +) -> ScheduleRecord: + """Build a minimal :class:`ScheduleRecord` for the trigger test.""" + return ScheduleRecord( + user_id=user_id, + agent_id=agent_id, + data=ScheduleData( + name="sched-a", + description=description, + enabled=enabled, + cron_expression="0 0 * * *", + started_at=datetime(2025, 1, 1), + chat_model_config=ChatModelConfig( + type="dashscope_credential", + credential_id="c", + model="m", + parameters={}, + ), + stateful=stateful, + permission_mode=PermissionMode.DONT_ASK, + ), + ) + + +class _SchedulerFireTestBase(IsolatedAsyncioTestCase): + """Shared fakeredis + storage + bus + manager fixture.""" + + async def asyncSetUp(self) -> None: + self.fr = fakeredis.aioredis.FakeRedis(decode_responses=True) + self._stack = AsyncExitStack() + self.storage = await self._stack.enter_async_context( + _make_storage(self.fr), + ) + self.bus = await self._stack.enter_async_context(_make_bus(self.fr)) + # Do NOT enter the SchedulerManager context — that would start + # APScheduler. We only need ``_build_trigger`` from the + # un-started manager. + self.manager = SchedulerManager( + storage=self.storage, + message_bus=self.bus, + ) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + await self.fr.aclose() + + +class TestSchedulerFireDelivery(_SchedulerFireTestBase): + """A fire delivers the prompt as a HintBlock + wakeup.""" + + async def test_fire_pushes_hint_and_wakeup(self) -> None: + """A fire creates a session, pushes the wrapped HintBlock to its + inbox, and enqueues one wakeup pointing at that session.""" + record = _make_record(description="please summarise the news") + trigger = self.manager._build_trigger(record) + await trigger() + + # A session was created. + sessions = await self.storage.list_sessions( + record.user_id, + record.agent_id, + ) + self.assertEqual(len(sessions), 1) + session = sessions[0] + self.assertEqual( + { + "source": session.source, + "source_schedule_id": session.source_schedule_id, + }, + { + "source": SessionSource.SCHEDULE, + "source_schedule_id": record.id, + }, + ) + + # Inbox has the wrapped HintBlock. + inbox = await self.bus.inbox_drain(session.id, max_count=10) + self.assertEqual(len(inbox), 1) + hint = inbox[0][1] + self.assertDictEqual( + hint, + { + "type": "hint", + "id": AnyString(), + "hint": AnyString(), + "source": json.dumps( + {"label": "schedule", "sublabel": record.data.name}, + ), + }, + ) + self.assertIn("", hint["hint"]) + self.assertIn("please summarise the news", hint["hint"]) + + # A wakeup is enqueued for that session. + wakeups = await self.bus.dequeue_wakeups(max_count=10) + self.assertEqual(len(wakeups), 1) + self.assertEqual( + wakeups[0], + { + "session_id": session.id, + "agent_id": record.agent_id, + "user_id": record.user_id, + "kind": "wake", + "input": None, + }, + ) + + +class TestSchedulerFireDisabled(_SchedulerFireTestBase): + """Disabled schedules are a no-op.""" + + async def test_disabled_fire_does_nothing(self) -> None: + """A fire on a disabled schedule creates no session and no wakeup.""" + record = _make_record(enabled=False) + trigger = self.manager._build_trigger(record) + await trigger() + + # No session created, no wakeup enqueued. + sessions = await self.storage.list_sessions( + record.user_id, + record.agent_id, + ) + self.assertEqual(sessions, []) + wakeups = await self.bus.dequeue_wakeups(max_count=10) + self.assertEqual(wakeups, []) + + +class TestSchedulerFireStatefulMode(_SchedulerFireTestBase): + """Stateful schedules reuse the same session id across fires.""" + + async def test_stateful_fires_share_one_session(self) -> None: + """Two fires of a stateful schedule reuse the same session id.""" + record = _make_record(stateful=True) + trigger = self.manager._build_trigger(record) + await trigger() + await trigger() + + sessions = await self.storage.list_sessions( + record.user_id, + record.agent_id, + ) + # Exactly ONE session reused. + self.assertEqual(len(sessions), 1) + self.assertEqual([s.id for s in sessions], [f"{record.id}_stateful"]) + + # That single session has two HintBlocks in its inbox. + inbox = await self.bus.inbox_drain(sessions[0].id, max_count=10) + self.assertEqual(len(inbox), 2) + + # Two wakeups, both pointing at the same session. + wakeups = await self.bus.dequeue_wakeups(max_count=10) + self.assertEqual( + wakeups, + [ + { + "session_id": sessions[0].id, + "agent_id": record.agent_id, + "user_id": record.user_id, + "kind": "wake", + "input": None, + }, + { + "session_id": sessions[0].id, + "agent_id": record.agent_id, + "user_id": record.user_id, + "kind": "wake", + "input": None, + }, + ], + ) + + +class TestSchedulerFireNonStatefulMode(_SchedulerFireTestBase): + """Non-stateful schedules create a fresh session every fire.""" + + async def test_non_stateful_fires_create_distinct_sessions( + self, + ) -> None: + """Two fires of a non-stateful schedule create distinct sessions.""" + record = _make_record(stateful=False) + trigger = self.manager._build_trigger(record) + await trigger() + await trigger() + + sessions = await self.storage.list_sessions( + record.user_id, + record.agent_id, + ) + self.assertEqual(len(sessions), 2) + self.assertNotEqual(sessions[0].id, sessions[1].id) diff --git a/tests/service_subagent_hitl_projector_test.py b/tests/service_subagent_hitl_projector_test.py new file mode 100644 index 0000000000000000000000000000000000000000..1b7e56ccd15ea257fd5a56aebca1595e4aae32d9 --- /dev/null +++ b/tests/service_subagent_hitl_projector_test.py @@ -0,0 +1,350 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Tests for :class:`SubagentHitlProjector` — the strategy that mirrors a +team *member* (worker) session's HITL request onto its *leader* session. + +Covers the projection policy: + +- A worker ``Require*`` event upserts a pending card on the leader and + publishes the live ``EVT_REQUIRE`` notification. +- A worker ``*Result`` / ``ReplyEnd`` event deletes the card and + publishes the ``EVT_RESULT`` clear notification. +- No-op for non-team agents, for sessions with no ``team_id``, and for + the leader session itself (its own HITL reaches its client directly). +- ``resolve`` finds the worker entry behind a leader+reply_id, and + returns ``None`` when none matches. +""" +from typing import Any +from unittest import IsolatedAsyncioTestCase + +from agentscope.app._service import SubagentHitlProjector +from agentscope.app.storage._model._agent import AgentRecord, AgentData +from agentscope.app.storage._model._session import SessionRecord +from agentscope.app.storage._model._team import TeamRecord +from agentscope.event import ( + RequireUserConfirmEvent, + RequireExternalExecutionEvent, + UserConfirmResultEvent, + ExternalExecutionResultEvent, + ReplyEndEvent, +) + + +class _FakeProjection: + """Records every mutation the projector makes, no Redis required. + + Mimics :class:`SessionProjection`'s surface (``upsert`` / ``delete`` + / ``publish`` / ``list``) over an in-memory ``{(sid, kind): {eid: + payload}}`` store so tests can assert exactly what was projected. + """ + + def __init__(self) -> None: + self.store: dict[tuple[str, str], dict[str, dict]] = {} + self.published: list[tuple[str, str, dict]] = [] + + async def upsert( + self, + target_sid: str, + kind: str, + entry_id: str, + payload: dict, + ) -> None: + """Record an upsert in the in-memory store.""" + self.store.setdefault((target_sid, kind), {})[entry_id] = payload + + async def delete( + self, + target_sid: str, + kind: str, + entry_id: str, + ) -> None: + """Remove an entry from the in-memory store.""" + self.store.get((target_sid, kind), {}).pop(entry_id, None) + + async def list(self, target_sid: str, kind: str) -> list[dict]: + """Return all entries for the given session and kind.""" + return list(self.store.get((target_sid, kind), {}).values()) + + async def publish( + self, + target_sid: str, + event_name: str, + value: dict, + ) -> None: + """Record a published event for later assertion.""" + self.published.append((target_sid, event_name, value)) + + +class _FakeStorage: + """Returns a fixed team for ``get_team``; ``None`` to simulate a + missing / non-team lookup.""" + + def __init__(self, team: TeamRecord | None) -> None: + self._team = team + + async def get_team( + self, + _user_id: str, + _team_id: str, + ) -> TeamRecord | None: + """Return the pre-configured team record.""" + return self._team + + +_LEADER_SID = "leader-sid" +_WORKER_SID = "worker-sid" +_TEAM_ID = "team-1" + + +def _team() -> TeamRecord: + """A team whose leader is ``_LEADER_SID``.""" + return TeamRecord.model_construct( + id=_TEAM_ID, + user_id="u", + session_id=_LEADER_SID, + ) + + +def _session(sid: str, team_id: str | None = _TEAM_ID) -> SessionRecord: + """A session record with just the fields the projector reads.""" + return SessionRecord.model_construct( + id=sid, + user_id="u", + agent_id="wa1", + team_id=team_id, + ) + + +def _agent(source: str = "team") -> AgentRecord: + """An agent record with just the fields the projector reads.""" + return AgentRecord.model_construct( + id="wa1", + user_id="u", + source=source, + data=AgentData.model_construct(name="researcher"), + ) + + +def _projector(team: TeamRecord | None = None) -> SubagentHitlProjector: + return SubagentHitlProjector(_FakeStorage(team if team else _team())) + + +def _entry_id() -> str: + return SubagentHitlProjector.entry_id(_WORKER_SID, "r1") + + +class TestSubagentHitlProjectorRequire(IsolatedAsyncioTestCase): + """The require → upsert + publish path.""" + + async def _run_require(self, event: Any) -> _FakeProjection: + projection = _FakeProjection() + await _projector().maybe_project( + "u", + _session(_WORKER_SID), + _agent(), + event, + projection, + ) + return projection + + async def test_require_user_confirm_upserts_and_publishes(self) -> None: + """A worker ``RequireUserConfirmEvent`` writes a leader card and + fires the live require notification.""" + event = RequireUserConfirmEvent.model_construct( + reply_id="r1", + tool_calls=[], + ) + projection = await self._run_require(event) + + card = projection.store[(_LEADER_SID, SubagentHitlProjector.KIND)] + self.assertIn(_entry_id(), card) + payload = card[_entry_id()] + self.assertEqual(payload["worker_session_id"], _WORKER_SID) + self.assertEqual(payload["reply_id"], "r1") + self.assertEqual(payload["event_type"], "require_user_confirm") + + self.assertEqual(len(projection.published), 1) + sid, name, _ = projection.published[0] + self.assertEqual(sid, _LEADER_SID) + self.assertEqual(name, SubagentHitlProjector.EVT_REQUIRE) + + async def test_require_external_execution_marks_event_type(self) -> None: + """A ``RequireExternalExecutionEvent`` carries the matching + ``event_type`` discriminator.""" + event = RequireExternalExecutionEvent.model_construct( + reply_id="r1", + tool_calls=[], + ) + projection = await self._run_require(event) + + card = projection.store[(_LEADER_SID, SubagentHitlProjector.KIND)] + self.assertEqual( + card[_entry_id()]["event_type"], + "require_external_execution", + ) + + +class TestSubagentHitlProjectorClear(IsolatedAsyncioTestCase): + """The result / reply-end → delete + publish path.""" + + async def _seed_then(self, clear_event: Any) -> _FakeProjection: + projection = _FakeProjection() + projector = _projector() + # Seed a pending card first. + await projector.maybe_project( + "u", + _session(_WORKER_SID), + _agent(), + RequireUserConfirmEvent.model_construct( + reply_id="r1", + tool_calls=[], + ), + projection, + ) + projection.published.clear() + # Now clear it. + await projector.maybe_project( + "u", + _session(_WORKER_SID), + _agent(), + clear_event, + projection, + ) + return projection + + async def test_user_confirm_result_clears_card(self) -> None: + """A worker ``UserConfirmResultEvent`` deletes the leader card and + publishes the clear notification.""" + projection = await self._seed_then( + UserConfirmResultEvent.model_construct( + reply_id="r1", + confirm_results=[], + ), + ) + card = projection.store[(_LEADER_SID, SubagentHitlProjector.KIND)] + self.assertNotIn(_entry_id(), card) + + self.assertEqual(len(projection.published), 1) + sid, name, value = projection.published[0] + self.assertEqual(sid, _LEADER_SID) + self.assertEqual(name, SubagentHitlProjector.EVT_RESULT) + self.assertEqual(value["worker_session_id"], _WORKER_SID) + self.assertEqual(value["reply_id"], "r1") + + async def test_external_execution_result_clears_card(self) -> None: + """An ``ExternalExecutionResultEvent`` also clears the card.""" + projection = await self._seed_then( + ExternalExecutionResultEvent.model_construct( + reply_id="r1", + execution_results=[], + ), + ) + card = projection.store[(_LEADER_SID, SubagentHitlProjector.KIND)] + self.assertNotIn(_entry_id(), card) + self.assertEqual( + projection.published[0][1], + SubagentHitlProjector.EVT_RESULT, + ) + + async def test_reply_end_clears_card(self) -> None: + """``ReplyEndEvent`` is the primary clear signal (the resume's + continuation events are not republished through the stream).""" + projection = await self._seed_then( + ReplyEndEvent.model_construct(reply_id="r1"), + ) + card = projection.store[(_LEADER_SID, SubagentHitlProjector.KIND)] + self.assertNotIn(_entry_id(), card) + self.assertEqual( + projection.published[0][1], + SubagentHitlProjector.EVT_RESULT, + ) + + +class TestSubagentHitlProjectorNoOp(IsolatedAsyncioTestCase): + """Cases where nothing should be projected.""" + + async def _assert_noop( + self, + *, + agent_source: str = "team", + team_id: str | None = _TEAM_ID, + team: TeamRecord | None = None, + session_id: str = _WORKER_SID, + ) -> None: + projection = _FakeProjection() + projector = SubagentHitlProjector( + _FakeStorage(team if team is not None else _team()), + ) + await projector.maybe_project( + "u", + _session(session_id, team_id=team_id), + _agent(source=agent_source), + RequireUserConfirmEvent.model_construct( + reply_id="r1", + tool_calls=[], + ), + projection, + ) + self.assertEqual(projection.store, {}) + self.assertEqual(projection.published, []) + + async def test_non_team_agent_is_noop(self) -> None: + """A ``source="user"`` agent never projects.""" + await self._assert_noop(agent_source="user") + + async def test_session_without_team_id_is_noop(self) -> None: + """A session with no ``team_id`` never projects.""" + await self._assert_noop(team_id=None) + + async def test_leader_session_is_noop(self) -> None: + """The leader's own HITL reaches its client directly — no + self-projection.""" + await self._assert_noop(session_id=_LEADER_SID) + + async def test_unrelated_event_is_noop(self) -> None: + """Events outside the HITL set are ignored.""" + projection = _FakeProjection() + await _projector().maybe_project( + "u", + _session(_WORKER_SID), + _agent(), + ReplyEndEvent.model_construct(reply_id="r1"), + projection, + ) + # ReplyEnd with no seeded card: delete is a no-op, but it DOES + # publish a (harmless, idempotent) clear. Assert no card written. + self.assertEqual(projection.store, {}) + + +class TestSubagentHitlProjectorResolve(IsolatedAsyncioTestCase): + """The router-side ``resolve`` helper.""" + + async def test_resolve_finds_and_misses(self) -> None: + """``resolve`` returns the worker entry for a known reply_id and + ``None`` for an unknown one.""" + projection = _FakeProjection() + await _projector().maybe_project( + "u", + _session(_WORKER_SID), + _agent(), + RequireUserConfirmEvent.model_construct( + reply_id="r1", + tool_calls=[], + ), + projection, + ) + + hit = await SubagentHitlProjector.resolve( + projection, + _LEADER_SID, + "r1", + ) + self.assertIsNotNone(hit) + self.assertEqual(hit["worker_session_id"], _WORKER_SID) + + miss = await SubagentHitlProjector.resolve( + projection, + _LEADER_SID, + "nope", + ) + self.assertIsNone(miss) diff --git a/tests/service_team_tools_test.py b/tests/service_team_tools_test.py new file mode 100644 index 0000000000000000000000000000000000000000..33f535fe1c76d16eca5aa6d373d52fbb2bf9fdd7 --- /dev/null +++ b/tests/service_team_tools_test.py @@ -0,0 +1,1181 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Tests for the four framework-builtin team tools — :class:`TeamCreate`, +:class:`AgentCreate`, :class:`TeamSay`, :class:`TeamDelete`. + +Each tool's business logic now lives inline in its ``__call__`` (the +old ``TeamService`` orchestration layer is gone), so unit tests run the +tools directly against a real :class:`RedisStorage` + :class:`RedisMessageBus` +backed by ``fakeredis``. They assert: + +- the success path: storage rows updated, inbox + wakeup pushed where + expected; +- the failure paths: each precondition (not in a team / not the leader / + recipient missing / etc.) returns an ``ERROR`` ``ToolChunk`` instead + of raising. +""" +from contextlib import AsyncExitStack +from unittest import IsolatedAsyncioTestCase + +import fakeredis.aioredis + +from utils import AnyString + +from agentscope.agent import ContextConfig, ReActConfig +from agentscope.app._tool import ( + AgentCreate, + DEFAULT_SUB_AGENT_TEMPLATE, + TeamCreate, + TeamDelete, + TeamSay, +) +from agentscope.app._types import SubAgentTemplate +from agentscope.app.message_bus import RedisMessageBus +from agentscope.app.storage import ( + AgentData, + AgentRecord, + RedisStorage, + SessionConfig, +) +from agentscope.permission import ( + AdditionalWorkingDirectory, + PermissionBehavior, + PermissionContext, + PermissionMode, + PermissionRule, +) +from agentscope.state import AgentState + + +def _make_storage( + fr: fakeredis.aioredis.FakeRedis, +) -> RedisStorage: + """Construct a :class:`RedisStorage` that talks to *fr*.""" + + class _S(RedisStorage): + async def __aenter__(self) -> "RedisStorage": # type: ignore[override] + self._client = fr + return self + + async def aclose(self) -> None: + self._client = None + + return _S() + + +def _make_bus( + fr: fakeredis.aioredis.FakeRedis, +) -> RedisMessageBus: + """Construct a :class:`RedisMessageBus` that talks to *fr*.""" + + class _B(RedisMessageBus): + async def __aenter__( # type: ignore[override] + self, + ) -> "RedisMessageBus": + self._client = fr + return self + + async def aclose(self) -> None: + self._client = None + + return _B() + + +def _make_agent_record( + user_id: str, + name: str, + source: str = "user", +) -> AgentRecord: + """Build a minimal :class:`AgentRecord`.""" + return AgentRecord( + user_id=user_id, + source=source, + data=AgentData( + name=name, + system_prompt=f"You are {name}.", + context_config=ContextConfig(), + react_config=ReActConfig(), + ), + ) + + +class _TeamToolsTestBase(IsolatedAsyncioTestCase): + """Shared fixture: a fakeredis-backed storage + bus, a leader + agent record, and a leader session. + + Sub-classes set up their own teams / workers on top. + """ + + user_id = "u" + + async def asyncSetUp(self) -> None: + self.fr = fakeredis.aioredis.FakeRedis(decode_responses=True) + self._stack = AsyncExitStack() + self.storage = await self._stack.enter_async_context( + _make_storage(self.fr), + ) + self.bus = await self._stack.enter_async_context(_make_bus(self.fr)) + + # Leader agent + its session. + self.leader_agent = _make_agent_record(self.user_id, "leader") + await self.storage.upsert_agent(self.user_id, self.leader_agent) + self.leader_session = await self.storage.upsert_session( + user_id=self.user_id, + agent_id=self.leader_agent.id, + config=SessionConfig(workspace_id="ws"), + ) + + async def asyncTearDown(self) -> None: + await self._stack.aclose() + await self.fr.aclose() + + +class TestTeamCreate(_TeamToolsTestBase): + """``TeamCreate`` creates a TeamRecord and stamps ``team_id`` on + the calling session.""" + + async def test_creates_team_and_stamps_session(self) -> None: + """A successful ``TeamCreate`` writes a TeamRecord and stamps the + leader's session with the new ``team_id``.""" + tool = TeamCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + ) + chunk = await tool(name="alpha", description="t-desc") + + self.assertDictEqual( + chunk.model_dump(), + { + "content": [ + {"type": "text", "text": AnyString(), "id": AnyString()}, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + # A team exists and the leader's session now points at it. + sess = await self.storage.get_session( + self.user_id, + self.leader_agent.id, + self.leader_session.id, + ) + self.assertIsNotNone(sess.team_id) + team = await self.storage.get_team(self.user_id, sess.team_id) + self.assertIsNotNone(team) + self.assertDictEqual( + { + "name": team.data.name, + "session_id": team.session_id, + "member_ids": team.data.member_ids, + }, + { + "name": "alpha", + "session_id": self.leader_session.id, + "member_ids": [], + }, + ) + + async def test_rejects_when_session_already_in_team(self) -> None: + """A session can lead at most one team.""" + tool = TeamCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + ) + await tool(name="first", description="d") + chunk = await tool(name="second", description="d") + self.assertDictEqual( + chunk.model_dump(), + { + "content": [ + {"type": "text", "text": AnyString(), "id": AnyString()}, + ], + "state": "error", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + +class TestAgentCreate(_TeamToolsTestBase): + """``AgentCreate`` spawns a worker agent + session, appends it to + the team, and delivers the initial prompt via inbox + wakeup.""" + + async def asyncSetUp(self) -> None: + await super().asyncSetUp() + # Pre-create a team so the leader has one to add to. + await TeamCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + )(name="team", description="team desc") + + async def test_spawns_worker_and_delivers_initial_prompt( + self, + ) -> None: + """A successful ``AgentCreate`` adds the worker agent + session to + the team and delivers the initial prompt via inbox + wakeup.""" + tool = AgentCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + ) + chunk = await tool( + name="worker", + description="does research", + prompt="please look up X", + ) + self.assertDictEqual( + chunk.model_dump(), + { + "content": [ + {"type": "text", "text": AnyString(), "id": AnyString()}, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + # Team has one member. + sess = await self.storage.get_session( + self.user_id, + self.leader_agent.id, + self.leader_session.id, + ) + team = await self.storage.get_team(self.user_id, sess.team_id) + self.assertEqual(len(team.data.member_ids), 1) + worker_agent_id = team.data.member_ids[0] + + # Worker agent exists with source=team. + worker_agent = await self.storage.get_agent( + self.user_id, + worker_agent_id, + ) + self.assertEqual( + {"source": worker_agent.source, "name": worker_agent.data.name}, + {"source": "team", "name": "worker"}, + ) + + # Worker has exactly one session, marked with team_id. + worker_sessions = await self.storage.list_sessions( + self.user_id, + worker_agent_id, + ) + self.assertEqual(len(worker_sessions), 1) + self.assertEqual(worker_sessions[0].team_id, sess.team_id) + + # The initial prompt is in the worker's inbox as a HintBlock + # wrapped in a tag. + inbox = await self.bus.inbox_drain( + worker_sessions[0].id, + max_count=10, + ) + self.assertEqual(len(inbox), 1) + hint_payload = inbox[0][1] + self.assertDictEqual( + hint_payload, + { + "type": "hint", + "id": AnyString(), + "hint": AnyString(), + "source": '{"label": "team_message", "sublabel": "leader"}', + }, + ) + self.assertIn("please look up X", hint_payload["hint"]) + self.assertIn(" PermissionContext: + """Run ``AgentCreate`` with the given template + leader state + and return the worker session's :class:`PermissionContext`.""" + tool = AgentCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + sub_agent_templates={template.type: template}, + ) + chunk = await tool( + name=worker_name, + description="does work", + prompt="work in the repo", + subagent_type=template.type, + _agent_state=leader_state, + ) + self.assertEqual(chunk.state.value, "running") + sess = await self.storage.get_session( + self.user_id, + self.leader_agent.id, + self.leader_session.id, + ) + team = await self.storage.get_team(self.user_id, sess.team_id) + # The worker is whichever member was added last. + worker_agent_id = team.data.member_ids[-1] + worker_sessions = await self.storage.list_sessions( + self.user_id, + worker_agent_id, + ) + return worker_sessions[0].state.permission_context + + def _make_leader_state(self) -> AgentState: + """Build a leader :class:`AgentState` with one of every kind of + permission entry, so tests can assert which pieces leak through + each flag.""" + return AgentState( + permission_context=PermissionContext( + mode=PermissionMode.ACCEPT_EDITS, + working_directories={ + "/tmp/as-workspace": AdditionalWorkingDirectory( + path="/tmp/as-workspace", + source="session", + ), + }, + allow_rules={ + "Bash": [ + PermissionRule( + tool_name="Bash", + rule_content="git status", + behavior=PermissionBehavior.ALLOW, + source="session", + ), + ], + }, + ), + ) + + async def test_default_template_follows_leader_completely(self) -> None: + """The built-in default template inherits the leader's mode, + working directories, and rules — its own + :attr:`permission_context` is empty, so the worker effectively + mirrors the leader.""" + leader_state = self._make_leader_state() + worker_context = await self._spawn_worker_with_template( + leader_state, + DEFAULT_SUB_AGENT_TEMPLATE, + ) + self.assertEqual(worker_context.mode, PermissionMode.ACCEPT_EDITS) + self.assertIn("/tmp/as-workspace", worker_context.working_directories) + self.assertEqual( + worker_context.allow_rules["Bash"][0].rule_content, + "git status", + ) + + async def test_override_leader_mode_pins_template_mode(self) -> None: + """When ``override_leader_mode=True`` the template's mode wins.""" + leader_state = self._make_leader_state() + explorer = SubAgentTemplate( + type="explorer", + description="Read-only worker.", + system_prompt_template=( + DEFAULT_SUB_AGENT_TEMPLATE.system_prompt_template + ), + permission_context=PermissionContext( + mode=PermissionMode.EXPLORE, + ), + override_leader_mode=True, + ) + worker_context = await self._spawn_worker_with_template( + leader_state, + explorer, + ) + self.assertEqual(worker_context.mode, PermissionMode.EXPLORE) + # Rules and dirs still inherited (defaults). + self.assertIn("/tmp/as-workspace", worker_context.working_directories) + self.assertIn("Bash", worker_context.allow_rules) + + async def test_extend_flags_off_isolate_template(self) -> None: + """``extend_*=False`` keeps the leader's rules and dirs out of + the worker; the template's own entries are the worker's + complete set.""" + leader_state = self._make_leader_state() + sandbox = SubAgentTemplate( + type="sandbox", + description="Fully isolated worker.", + system_prompt_template=( + DEFAULT_SUB_AGENT_TEMPLATE.system_prompt_template + ), + permission_context=PermissionContext( + mode=PermissionMode.BYPASS, + deny_rules={ + "Write": [ + PermissionRule( + tool_name="Write", + rule_content=None, + behavior=PermissionBehavior.DENY, + source="template", + ), + ], + }, + ), + override_leader_mode=True, + extend_leader_permission_rules=False, + extend_leader_working_directories=False, + ) + worker_context = await self._spawn_worker_with_template( + leader_state, + sandbox, + ) + self.assertEqual(worker_context.mode, PermissionMode.BYPASS) + self.assertEqual(worker_context.working_directories, {}) + self.assertNotIn("Bash", worker_context.allow_rules) + # Template's own deny rule is preserved. + self.assertEqual( + worker_context.deny_rules["Write"][0].source, + "template", + ) + + async def test_extend_rules_keeps_template_rules_first(self) -> None: + """When merging rules for the same tool, the template's rules + appear first in the list so the engine evaluates them before + the leader's.""" + leader_state = AgentState( + permission_context=PermissionContext( + mode=PermissionMode.DEFAULT, + allow_rules={ + "Bash": [ + PermissionRule( + tool_name="Bash", + rule_content="git status", + behavior=PermissionBehavior.ALLOW, + source="session", + ), + ], + }, + ), + ) + template = SubAgentTemplate( + type="custom", + description="Custom worker.", + system_prompt_template=( + DEFAULT_SUB_AGENT_TEMPLATE.system_prompt_template + ), + permission_context=PermissionContext( + allow_rules={ + "Bash": [ + PermissionRule( + tool_name="Bash", + rule_content="ls", + behavior=PermissionBehavior.ALLOW, + source="template", + ), + ], + }, + ), + ) + worker_context = await self._spawn_worker_with_template( + leader_state, + template, + ) + bash_rules = worker_context.allow_rules["Bash"] + self.assertEqual( + [r.source for r in bash_rules], + ["template", "session"], + ) + + async def test_extend_dirs_template_wins_on_collision(self) -> None: + """When the template and leader both declare the same working + directory path, the template's entry is kept.""" + leader_state = AgentState( + permission_context=PermissionContext( + working_directories={ + "/tmp/shared": AdditionalWorkingDirectory( + path="/tmp/shared", + source="session", + ), + }, + ), + ) + template = SubAgentTemplate( + type="custom", + description="Custom worker.", + system_prompt_template=( + DEFAULT_SUB_AGENT_TEMPLATE.system_prompt_template + ), + permission_context=PermissionContext( + working_directories={ + "/tmp/shared": AdditionalWorkingDirectory( + path="/tmp/shared", + source="template", + ), + }, + ), + ) + worker_context = await self._spawn_worker_with_template( + leader_state, + template, + ) + self.assertEqual( + worker_context.working_directories["/tmp/shared"].source, + "template", + ) + + async def test_rejects_when_not_in_team(self) -> None: + """Calling ``AgentCreate`` from a session that hasn't run + ``TeamCreate`` yet returns an error chunk.""" + # Build a second leader session that is NOT in any team. + loner_session = await self.storage.upsert_session( + user_id=self.user_id, + agent_id=self.leader_agent.id, + config=SessionConfig(workspace_id="ws2"), + ) + tool = AgentCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=loner_session.id, + agent_id=self.leader_agent.id, + ) + chunk = await tool( + name="worker", + description="d", + prompt="p", + ) + self.assertDictEqual( + chunk.model_dump(), + { + "content": [ + {"type": "text", "text": AnyString(), "id": AnyString()}, + ], + "state": "error", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + async def test_rejects_unknown_subagent_type(self) -> None: + """An unrecognised ``subagent_type`` returns an error chunk.""" + tool = AgentCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + ) + chunk = await tool( + name="w", + description="d", + prompt="p", + subagent_type="not-a-type", + ) + self.assertDictEqual( + chunk.model_dump(), + { + "content": [ + {"type": "text", "text": AnyString(), "id": AnyString()}, + ], + "state": "error", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + async def test_rejects_duplicate_member_name(self) -> None: + """Two ``AgentCreate`` calls with the same ``name`` is rejected: + TeamSay routes by name, so duplicates would be ambiguous.""" + tool = AgentCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + ) + first = await tool( + name="worker", + description="d", + prompt="p", + ) + self.assertEqual(first.state.value, "running") + + second = await tool( + name="worker", + description="d", + prompt="p", + ) + self.assertEqual(second.state.value, "error") + + # The team still only has one member — the second call did not + # persist anything. + sess = await self.storage.get_session( + self.user_id, + self.leader_agent.id, + self.leader_session.id, + ) + team = await self.storage.get_team(self.user_id, sess.team_id) + self.assertEqual(len(team.data.member_ids), 1) + + async def test_rejects_member_name_colliding_with_leader(self) -> None: + """A worker name that matches the leader's name is rejected — + ``TeamSay(to=)`` must remain unambiguous.""" + tool = AgentCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + ) + chunk = await tool( + name=self.leader_agent.data.name, # "leader" + description="d", + prompt="p", + ) + self.assertEqual(chunk.state.value, "error") + + # No worker was added. + sess = await self.storage.get_session( + self.user_id, + self.leader_agent.id, + self.leader_session.id, + ) + team = await self.storage.get_team(self.user_id, sess.team_id) + self.assertEqual(team.data.member_ids, []) + + +class TestAgentCreateTemplates(_TeamToolsTestBase): + """Template-aware ``AgentCreate`` behaviour: schema dynamics, + template routing, and config isolation.""" + + _explorer_template = SubAgentTemplate( + type="explorer", + description="Read-only exploration agent.", + system_prompt_template=( + "You are {member_name}, an explorer in team " + "'{team_name}' led by {leader_name}.\n\n" + "Team purpose: {team_description}\n\n" + "Your role: {member_description}" + ), + ) + + async def asyncSetUp(self) -> None: + await super().asyncSetUp() + await TeamCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + )(name="team", description="team desc") + + async def test_schema_omits_subagent_type_when_no_custom_templates( + self, + ) -> None: + """When only the built-in ``"default"`` template exists, the + ``input_schema`` must NOT contain a ``subagent_type`` field.""" + tool = AgentCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + ) + self.assertNotIn("subagent_type", tool.input_schema["properties"]) + + async def test_schema_includes_subagent_type_with_custom_templates( + self, + ) -> None: + """When custom templates are registered, ``subagent_type`` appears + in the schema with the correct enum values.""" + templates = {"explorer": self._explorer_template} + tool = AgentCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + sub_agent_templates=templates, + ) + self.assertIn("subagent_type", tool.input_schema["properties"]) + enum_values = tool.input_schema["properties"]["subagent_type"]["enum"] + self.assertIn("default", enum_values) + self.assertIn("explorer", enum_values) + + async def test_default_template_injected_when_missing(self) -> None: + """The built-in default template is always available even when + only custom templates are provided.""" + templates = {"explorer": self._explorer_template} + tool = AgentCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + sub_agent_templates=templates, + ) + self.assertIn("default", tool._sub_agent_templates) + self.assertIs( + tool._sub_agent_templates["default"], + DEFAULT_SUB_AGENT_TEMPLATE, + ) + + async def test_custom_template_applies_system_prompt(self) -> None: + """A worker created with a custom template gets the template's + system prompt (not the default one).""" + templates = {"explorer": self._explorer_template} + tool = AgentCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + sub_agent_templates=templates, + ) + chunk = await tool( + name="scout", + description="explores code", + prompt="look around", + subagent_type="explorer", + ) + self.assertEqual(chunk.state.value, "running") + + sess = await self.storage.get_session( + self.user_id, + self.leader_agent.id, + self.leader_session.id, + ) + team = await self.storage.get_team(self.user_id, sess.team_id) + worker_agent = await self.storage.get_agent( + self.user_id, + team.data.member_ids[0], + ) + self.assertIn("an explorer in team", worker_agent.data.system_prompt) + self.assertNotIn( + "You communicate with the team leader", + worker_agent.data.system_prompt, + ) + + async def test_configs_are_deep_copied(self) -> None: + """Each spawned worker receives its own copy of the template's + config objects — mutations must not leak across agents.""" + templates = {"explorer": self._explorer_template} + tool = AgentCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + sub_agent_templates=templates, + ) + await tool( + name="w1", + description="d", + prompt="p", + subagent_type="explorer", + ) + await tool( + name="w2", + description="d", + prompt="p", + subagent_type="explorer", + ) + + sess = await self.storage.get_session( + self.user_id, + self.leader_agent.id, + self.leader_session.id, + ) + team = await self.storage.get_team(self.user_id, sess.team_id) + a1 = await self.storage.get_agent( + self.user_id, + team.data.member_ids[0], + ) + a2 = await self.storage.get_agent( + self.user_id, + team.data.member_ids[1], + ) + self.assertIsNot( + a1.data.context_config, + a2.data.context_config, + ) + self.assertIsNot( + a1.data.react_config, + a2.data.react_config, + ) + + +class TestTeamSay(_TeamToolsTestBase): + """``TeamSay`` delivers a HintBlock + wakeup to each addressed + teammate.""" + + async def asyncSetUp(self) -> None: + await super().asyncSetUp() + await TeamCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + )(name="team", description="team desc") + # Add 2 workers. + agent_create = AgentCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + ) + await agent_create( + name="w1", + description="d", + prompt="p1", + ) + await agent_create( + name="w2", + description="d", + prompt="p2", + ) + # Drain the pre-existing wakeups so subsequent assertions only + # see what TeamSay enqueues. + await self.bus.dequeue_wakeups(max_count=100) + # Resolve worker IDs. + sess = await self.storage.get_session( + self.user_id, + self.leader_agent.id, + self.leader_session.id, + ) + team = await self.storage.get_team(self.user_id, sess.team_id) + self.worker_ids = team.data.member_ids + self.worker_sessions = {} + for aid in self.worker_ids: + ss = await self.storage.list_sessions(self.user_id, aid) + self.worker_sessions[aid] = ss[0].id + # Drain initial-prompt hints too. + await self.bus.inbox_drain(ss[0].id, max_count=100) + + async def test_targeted_message_delivers_to_one(self) -> None: + """``to=`` delivers a HintBlock + wakeup to that worker + only; other team members see nothing.""" + tool = TeamSay( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + role="leader", + ) + target_aid = self.worker_ids[0] + target_agent = await self.storage.get_agent( + self.user_id, + target_aid, + ) + chunk = await tool(content="hi w1", to=target_agent.data.name) + self.assertDictEqual( + chunk.model_dump(), + { + "content": [ + {"type": "text", "text": AnyString(), "id": AnyString()}, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + # Inbox of target has the hint; the other worker's inbox stays + # empty. + target_sid = self.worker_sessions[target_aid] + other_sid = self.worker_sessions[self.worker_ids[1]] + target_inbox = await self.bus.inbox_drain(target_sid, max_count=10) + other_inbox = await self.bus.inbox_drain(other_sid, max_count=10) + self.assertEqual(len(target_inbox), 1) + self.assertEqual(len(other_inbox), 0) + self.assertDictEqual( + target_inbox[0][1], + { + "type": "hint", + "id": AnyString(), + "hint": AnyString(), + "source": AnyString(), + }, + ) + self.assertIn("hi w1", target_inbox[0][1]["hint"]) + + # Wakeup for the target only. + wakeups = await self.bus.dequeue_wakeups(max_count=10) + self.assertEqual( + wakeups, + [ + { + "session_id": target_sid, + "agent_id": target_aid, + "user_id": self.user_id, + "kind": "wake", + "input": None, + }, + ], + ) + + async def test_broadcast_delivers_to_all_others(self) -> None: + """``to=None`` broadcasts to everyone in the team except the + sender.""" + tool = TeamSay( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + role="leader", + ) + chunk = await tool(content="all hands", to=None) + self.assertDictEqual( + chunk.model_dump(), + { + "content": [ + {"type": "text", "text": AnyString(), "id": AnyString()}, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + # Both workers receive; leader doesn't loopback to itself. + for aid, sid in self.worker_sessions.items(): + inbox = await self.bus.inbox_drain(sid, max_count=10) + self.assertEqual( + len(inbox), + 1, + f"worker {aid} missed broadcast", + ) + leader_inbox = await self.bus.inbox_drain( + self.leader_session.id, + max_count=10, + ) + self.assertEqual(leader_inbox, []) + + wakeups = await self.bus.dequeue_wakeups(max_count=10) + self.assertEqual(len(wakeups), 2) + + async def test_rejects_when_session_not_in_team(self) -> None: + """A session without a team can't TeamSay.""" + loner_session = await self.storage.upsert_session( + user_id=self.user_id, + agent_id=self.leader_agent.id, + config=SessionConfig(workspace_id="ws-lone"), + ) + tool = TeamSay( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=loner_session.id, + agent_id=self.leader_agent.id, + role="leader", + ) + chunk = await tool(content="hi", to=None) + self.assertDictEqual( + chunk.model_dump(), + { + "content": [ + {"type": "text", "text": AnyString(), "id": AnyString()}, + ], + "state": "error", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + async def test_rejects_unknown_recipient(self) -> None: + """A ``to=`` that doesn't resolve to a team member name returns + an error chunk.""" + tool = TeamSay( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + role="leader", + ) + chunk = await tool(content="hi", to="ghost-name-not-in-team") + self.assertDictEqual( + chunk.model_dump(), + { + "content": [ + {"type": "text", "text": AnyString(), "id": AnyString()}, + ], + "state": "error", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + async def test_rejects_self_target(self) -> None: + """``to=`` is rejected — talk to yourself in + reasoning, not via TeamSay.""" + tool = TeamSay( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + role="leader", + ) + chunk = await tool(content="hi", to=self.leader_agent.data.name) + self.assertDictEqual( + chunk.model_dump(), + { + "content": [ + {"type": "text", "text": AnyString(), "id": AnyString()}, + ], + "state": "error", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + async def test_worker_can_address_leader_by_name(self) -> None: + """A worker constructed with ``role="worker"`` can address the + leader using the leader's name — the only identifier a worker + ever has for the leader (received via the ``from=`` attribute + of the initial ```` hint).""" + worker_aid = self.worker_ids[0] + worker_sid = self.worker_sessions[worker_aid] + + tool = TeamSay( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=worker_sid, + agent_id=worker_aid, + role="worker", + ) + chunk = await tool( + content="task done", + to=self.leader_agent.data.name, + ) + self.assertEqual(chunk.state.value, "running") + + # Leader's inbox received the worker's reply. + leader_inbox = await self.bus.inbox_drain( + self.leader_session.id, + max_count=10, + ) + self.assertEqual(len(leader_inbox), 1) + self.assertIn("task done", leader_inbox[0][1]["hint"]) + + # Wakeup was enqueued for the leader. + wakeups = await self.bus.dequeue_wakeups(max_count=10) + self.assertEqual( + wakeups, + [ + { + "session_id": self.leader_session.id, + "agent_id": self.leader_agent.id, + "user_id": self.user_id, + "kind": "wake", + "input": None, + }, + ], + ) + + +class TestTeamDelete(_TeamToolsTestBase): + """``TeamDelete`` dissolves the team and clears the leader's + ``team_id``.""" + + async def asyncSetUp(self) -> None: + await super().asyncSetUp() + await TeamCreate( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + )(name="team", description="d") + + async def test_dissolves_team_from_leader(self) -> None: + """``TeamDelete`` removes the team and clears the leader's + ``team_id``.""" + tool = TeamDelete( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=self.leader_session.id, + agent_id=self.leader_agent.id, + ) + chunk = await tool() + self.assertDictEqual( + chunk.model_dump(), + { + "content": [ + {"type": "text", "text": AnyString(), "id": AnyString()}, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + sess = await self.storage.get_session( + self.user_id, + self.leader_agent.id, + self.leader_session.id, + ) + self.assertIsNone(sess.team_id) + + async def test_rejects_when_not_in_team(self) -> None: + """Calling ``TeamDelete`` from a session that isn't in any team + returns an error chunk.""" + loner_session = await self.storage.upsert_session( + user_id=self.user_id, + agent_id=self.leader_agent.id, + config=SessionConfig(workspace_id="ws-lone"), + ) + tool = TeamDelete( + storage=self.storage, + message_bus=self.bus, + user_id=self.user_id, + session_id=loner_session.id, + agent_id=self.leader_agent.id, + ) + chunk = await tool() + self.assertDictEqual( + chunk.model_dump(), + { + "content": [ + {"type": "text", "text": AnyString(), "id": AnyString()}, + ], + "state": "error", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) diff --git a/tests/service_toolkit_test.py b/tests/service_toolkit_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ba4a39bbc39bacd817228c48d0eab18162c7dd3b --- /dev/null +++ b/tests/service_toolkit_test.py @@ -0,0 +1,331 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Tests for :func:`get_toolkit` — the single entry point that assembles +the per-chat-turn :class:`Toolkit` from every tool source the framework +manages. + +Verifies the assembly rules: + +- workspace builtins are always included; +- the four ``Task*`` planning tools are always included; +- :class:`ToolStop` (from ``BackgroundTaskManager``) is always included; +- the four ``Schedule*`` tools only when ``session.config.chat_model_config`` + is set (they need a model to fire new runs with); +- team tools are role-gated by ``agent_record.source``: ``"team"`` → + one ``TeamSay`` (worker variant); anything else → the full + leader-side toolset of four; +- caller-supplied ``extra_factory`` results land at the end. +""" +from typing import Any +from unittest import IsolatedAsyncioTestCase + +from agentscope.agent import ContextConfig, ReActConfig +from agentscope.app._manager import ( + BackgroundTaskManager, + SchedulerManager, +) +from agentscope.app._service import get_toolkit +from agentscope.app.storage import ( + AgentData, + AgentRecord, + ChatModelConfig, + SessionConfig, + SessionRecord, +) +from agentscope.tool import ToolBase + + +class _FakeWorkspace: + """Stand-in for a resolved :class:`WorkspaceBase` — only the three + methods :func:`get_toolkit` calls are implemented.""" + + def __init__( + self, + tools: list[ToolBase] | None = None, + skills: list | None = None, + mcps: list | None = None, + ) -> None: + self._tools = tools or [] + self._skills = skills or [] + self._mcps = mcps or [] + + async def list_tools(self) -> list[ToolBase]: + """Return the configured workspace tools.""" + return list(self._tools) + + async def list_skills(self) -> list: + """Return the configured workspace skills.""" + return list(self._skills) + + async def list_mcps(self) -> list: + """Return the configured workspace MCP descriptors.""" + return list(self._mcps) + + +class _NullBus: + """``MessageBus`` placeholder. Team tools only carry the reference; + nothing in :func:`get_toolkit` actually awaits it.""" + + +def _make_agent(*, source: str = "user", name: str = "A") -> AgentRecord: + """Build a minimal :class:`AgentRecord`.""" + return AgentRecord( + user_id="u", + source=source, + data=AgentData( + name=name, + system_prompt=f"You are {name}.", + context_config=ContextConfig(), + react_config=ReActConfig(), + ), + ) + + +def _make_session( + *, + user_id: str, + agent_id: str, + with_model: bool, +) -> SessionRecord: + """Build a minimal :class:`SessionRecord`, optionally with a chat + model config.""" + cfg = SessionConfig( + workspace_id="ws", + chat_model_config=( + ChatModelConfig( + type="dashscope_credential", + credential_id="c", + model="m", + parameters={}, + ) + if with_model + else None + ), + ) + return SessionRecord(user_id=user_id, agent_id=agent_id, config=cfg) + + +class _NoOpStorage: + """Storage placeholder. ``get_toolkit`` itself does not call any + storage method — the team tools bind a reference for later use.""" + + +def _tool_names(toolkit: Any) -> list[str]: + """Extract every registered tool name from a :class:`Toolkit`, + walking its tool groups.""" + return [t.name for group in toolkit.tool_groups for t in group.tools] + + +class _StubTool(ToolBase): + """Minimal :class:`ToolBase` subclass that satisfies the abstract + methods so :func:`get_toolkit` can register the instance. + + Sub-classes override ``name`` and ``description``; nothing in the + tests actually calls the tool, so the implementations are no-ops. + """ + + name: str = "stub" + description: str = "stub tool" + input_schema: dict = {} + is_concurrency_safe: bool = True + is_read_only: bool = False + is_state_injected: bool = False + is_external_tool: bool = False + is_mcp: bool = False + mcp_name: str | None = None + + async def check_permissions(self, *args: Any, **kwargs: Any) -> None: + """No-op permission check — the tests do not exercise it.""" + + async def __call__(self, *args: Any, **kwargs: Any) -> None: + """No-op invocation — the tests never execute the tool.""" + + +class TestGetToolkitBaseAssembly(IsolatedAsyncioTestCase): + """User-owned agent (``source="user"``) gets the full set.""" + + async def test_user_agent_gets_all_sources(self) -> None: + """A user-owned agent receives workspace, planning, scheduling, + ToolStop, and the four leader-side team tools.""" + agent = _make_agent(source="user") + session = _make_session( + user_id="u", + agent_id=agent.id, + with_model=True, + ) + + class _WsTool(_StubTool): + """Stub workspace tool registered through ``_FakeWorkspace``.""" + + name: str = "ws-bash" + description: str = "stub workspace tool" + + ws_tool = _WsTool() + workspace = _FakeWorkspace(tools=[ws_tool]) + + toolkit = await get_toolkit( + storage=_NoOpStorage(), # type: ignore[arg-type] + workspace=workspace, # type: ignore[arg-type] + scheduler_manager=SchedulerManager( + storage=_NoOpStorage(), # type: ignore[arg-type] + message_bus=_NullBus(), # type: ignore[arg-type] + ), + background_task_manager=BackgroundTaskManager( + message_bus=_NullBus(), # type: ignore[arg-type] + ), + message_bus=_NullBus(), # type: ignore[arg-type] + user_id="u", + agent_record=agent, + session_record=session, + extra_factory=None, + middlewares=[], + ) + + names = set(_tool_names(toolkit)) + # Workspace tool present. + self.assertIn("ws-bash", names) + # Planning tools present. + self.assertTrue( + {"TaskCreate", "TaskList", "TaskGet", "TaskUpdate"} <= names, + ) + # Background task control present. + self.assertIn("ToolStop", names) + # Schedule control present (model_config is set). + self.assertTrue( + { + "ScheduleCreate", + "ScheduleView", + "ScheduleDelete", + "ScheduleList", + } + <= names, + ) + # Leader-side team tools (4 of them). + self.assertTrue( + {"TeamCreate", "AgentCreate", "TeamSay", "TeamDelete"} <= names, + ) + + +class TestGetToolkitWorkerVariant(IsolatedAsyncioTestCase): + """Worker agent (``source="team"``) only gets ``TeamSay``.""" + + async def test_worker_only_gets_team_say(self) -> None: + """A worker agent (``source="team"``) receives only ``TeamSay`` + from the team toolset.""" + agent = _make_agent(source="team", name="worker") + session = _make_session( + user_id="u", + agent_id=agent.id, + with_model=True, + ) + toolkit = await get_toolkit( + storage=_NoOpStorage(), # type: ignore[arg-type] + workspace=_FakeWorkspace(), # type: ignore[arg-type] + scheduler_manager=SchedulerManager( + storage=_NoOpStorage(), # type: ignore[arg-type] + message_bus=_NullBus(), # type: ignore[arg-type] + ), + background_task_manager=BackgroundTaskManager( + message_bus=_NullBus(), # type: ignore[arg-type] + ), + message_bus=_NullBus(), # type: ignore[arg-type] + user_id="u", + agent_record=agent, + session_record=session, + extra_factory=None, + middlewares=[], + ) + names = set(_tool_names(toolkit)) + # Only TeamSay from the team toolset. + self.assertIn("TeamSay", names) + for missing in ("TeamCreate", "AgentCreate", "TeamDelete"): + self.assertNotIn(missing, names) + + +class TestGetToolkitSchedulingGuard(IsolatedAsyncioTestCase): + """``Schedule*`` tools are only attached when the session has a + model configured.""" + + async def test_no_schedule_tools_without_model_config(self) -> None: + """Without a ``chat_model_config`` on the session, the four + ``Schedule*`` tools are omitted from the toolkit.""" + agent = _make_agent(source="user") + session = _make_session( + user_id="u", + agent_id=agent.id, + with_model=False, + ) + toolkit = await get_toolkit( + storage=_NoOpStorage(), # type: ignore[arg-type] + workspace=_FakeWorkspace(), # type: ignore[arg-type] + scheduler_manager=SchedulerManager( + storage=_NoOpStorage(), # type: ignore[arg-type] + message_bus=_NullBus(), # type: ignore[arg-type] + ), + background_task_manager=BackgroundTaskManager( + message_bus=_NullBus(), # type: ignore[arg-type] + ), + message_bus=_NullBus(), # type: ignore[arg-type] + user_id="u", + agent_record=agent, + session_record=session, + extra_factory=None, + middlewares=[], + ) + names = set(_tool_names(toolkit)) + for missing in ( + "ScheduleCreate", + "ScheduleView", + "ScheduleDelete", + "ScheduleList", + ): + self.assertNotIn(missing, names) + + +class TestGetToolkitExtraFactory(IsolatedAsyncioTestCase): + """Tools returned by ``extra_factory`` end up in the final toolkit.""" + + async def test_extra_factory_tools_are_attached(self) -> None: + """Tools returned by ``extra_factory`` end up in the final + toolkit alongside the framework-builtin ones.""" + + class _ExtraTool(_StubTool): + """Stub tool emitted by the ``extra_factory`` callback.""" + + name: str = "my-extra" + description: str = "stub extra tool" + + extra_tool = _ExtraTool() + + async def factory( + _user_id: str, + _agent_id: str, + _session_id: str, + ) -> list[ToolBase]: + """Stub extra-factory that always returns ``extra_tool``.""" + return [extra_tool] + + agent = _make_agent() + session = _make_session( + user_id="u", + agent_id=agent.id, + with_model=True, + ) + toolkit = await get_toolkit( + storage=_NoOpStorage(), # type: ignore[arg-type] + workspace=_FakeWorkspace(), # type: ignore[arg-type] + scheduler_manager=SchedulerManager( + storage=_NoOpStorage(), # type: ignore[arg-type] + message_bus=_NullBus(), # type: ignore[arg-type] + ), + background_task_manager=BackgroundTaskManager( + message_bus=_NullBus(), # type: ignore[arg-type] + ), + message_bus=_NullBus(), # type: ignore[arg-type] + user_id="u", + agent_record=agent, + session_record=session, + extra_factory=factory, + middlewares=[], + ) + self.assertIn("my-extra", _tool_names(toolkit)) diff --git a/tests/service_wakeup_dispatcher_test.py b/tests/service_wakeup_dispatcher_test.py new file mode 100644 index 0000000000000000000000000000000000000000..75996a8160e9523b9031f8c954a561bfcc0c808d --- /dev/null +++ b/tests/service_wakeup_dispatcher_test.py @@ -0,0 +1,479 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Tests for :class:`WakeupDispatcher` — one-per-process consumer of the +shared wake-up queue + signal channel. + +Verifies the four behaviours that callers rely on: + +- Lifecycle is purely ACM: ``__aenter__`` starts the loop and performs + an initial drain; ``__aexit__`` cancels the loop cleanly. +- A wake-up signal triggers a queue drain; each entry is dispatched as + a fire-and-forget ``ChatService.run`` call. +- Entries left on the queue from before startup are picked up on + ``__aenter__`` without waiting for a fresh signal. +- Sessions that are already running are skipped (no duplicate run). +- Malformed entries are logged and skipped, not raised. +""" +import asyncio +from contextlib import asynccontextmanager +from typing import Any, AsyncGenerator, Callable +from unittest import IsolatedAsyncioTestCase + +from agentscope.app._manager import ChatRunRegistry, WakeupDispatcher +from agentscope.app.message_bus import MessageBus, MessageBusKeys + + +class _FakeStorage: + """Minimal storage stand-in for the dispatcher's orphan-guard check. + + ``get_session`` returns a truthy sentinel for every session id by + default; tests that exercise the orphan path mutate + ``missing_session_ids``. + """ + + def __init__(self) -> None: + self.missing_session_ids: set[str] = set() + + async def get_session( + self, + _user_id: str, + _agent_id: str, + session_id: str, + ) -> object | None: + """Get a session id from the orphan guard.""" + if session_id in self.missing_session_ids: + return None + return object() + + +class _FakeBus(MessageBus): + """In-memory bus with just enough behaviour for the dispatcher. + + Implements the four primitives the dispatcher uses + (``queue_push`` / ``dequeue_wakeups`` indirectly via the parent's + domain helper / ``subscribe_wakeup_signal`` / ``is_locked`` / + ``publish``) and stubs the others. + """ + + def __init__(self) -> None: + self.queues: dict[str, list[tuple[str, dict]]] = {} + self._channels: dict[str, asyncio.Queue] = {} + self._next = 0 + self._locks: set[str] = set() + + def _channel(self, key: str) -> asyncio.Queue: + return self._channels.setdefault(key, asyncio.Queue()) + + # Mode A — queue + async def queue_push( + self, + key: str, + payload: dict, + *, + ttl_secs: int | None = None, + ) -> str: + self._next += 1 + entry_id = str(self._next) + self.queues.setdefault(key, []).append((entry_id, payload)) + return entry_id + + async def queue_drain( + self, + key: str, + *, + max_count: int, + ) -> list[tuple[str, dict]]: + entries = self.queues.get(key, [])[:max_count] + self.queues[key] = self.queues.get(key, [])[max_count:] + return entries + + async def queue_delete(self, key: str) -> None: + self.queues.pop(key, None) + + # Mode C — log (unused here) + async def log_append( + self, + key: str, + payload: dict, + *, + max_len: int | None = None, + ttl_secs: int | None = None, + ) -> str: + return "n/a" + + async def log_read( + self, + key: str, + since: str | None = None, + max_count: int = 100, + ) -> list[tuple[str, dict]]: + return [] + + async def log_trim( + self, + key: str, + before_id: str | None = None, + ) -> None: + return None + + # Mode D — pub/sub + async def publish(self, key: str, payload: dict) -> None: + await self._channel(key).put(payload) + + async def subscribe( + self, + key: str, + *, + on_ready: Callable[[], None] | None = None, + ) -> AsyncGenerator[dict, None]: + if on_ready is not None: + on_ready() + while True: + yield await self._channel(key).get() + + # Mode E — lock + @asynccontextmanager + async def acquire_lock( + self, + key: str, + *, + ttl_secs: int = 600, + ) -> AsyncGenerator[None, None]: + self._locks.add(key) + try: + yield + finally: + self._locks.discard(key) + + async def is_locked(self, key: str) -> bool: + return key in self._locks + + # Mode F — registry (unused by WakeupDispatcher; raise so any + # accidental dependency surfaces immediately rather than silently + # passing through a stub). + async def registry_set( + self, + namespace: str, + field: str, + value: str, + *, + ttl_secs: int | None = None, + ) -> None: + raise NotImplementedError + + async def registry_del(self, namespace: str, field: str) -> None: + raise NotImplementedError + + async def registry_exists(self, namespace: str, field: str) -> bool: + raise NotImplementedError + + async def registry_getall(self, namespace: str) -> dict[str, str]: + raise NotImplementedError + + async def registry_drop(self, namespace: str) -> None: + raise NotImplementedError + + +class _FakeChatService: + """Records calls to :meth:`run` so tests can assert dispatch.""" + + def __init__(self) -> None: + self.calls: list[dict] = [] + self.notify = asyncio.Event() + + async def run( + self, + user_id: str, + session_id: str, + agent_id: str, + input_msg: Any = None, + ) -> None: + """Record the call and signal a waiter.""" + self.calls.append( + { + "user_id": user_id, + "session_id": session_id, + "agent_id": agent_id, + "input_msg": input_msg, + }, + ) + self.notify.set() + + +async def _yield_a_few_times(ticks: int = 8) -> None: + """Yield the event loop a few times so spawned tasks make progress.""" + for _ in range(ticks): + await asyncio.sleep(0) + + +class TestWakeupDispatcherDispatch(IsolatedAsyncioTestCase): + """Verifies the signal-driven dispatch path.""" + + async def test_signal_drives_dispatch(self) -> None: + """A wake-up signal causes the queue to be drained and each + entry dispatched as a chat run.""" + bus = _FakeBus() + chat = _FakeChatService() + async with WakeupDispatcher( + message_bus=bus, + storage=_FakeStorage(), + chat_service=chat, + chat_run_registry=ChatRunRegistry(), + ): + await bus.queue_push( + MessageBusKeys.wakeup_queue(), + {"user_id": "u", "session_id": "s1", "agent_id": "a1"}, + ) + await bus.publish(MessageBusKeys.wakeup_signal(), {}) + + await asyncio.wait_for(chat.notify.wait(), timeout=2.0) + + self.assertEqual( + chat.calls, + [ + { + "user_id": "u", + "session_id": "s1", + "agent_id": "a1", + "input_msg": None, + }, + ], + ) + + async def test_initial_drain_picks_up_pending_entries(self) -> None: + """Entries on the queue from before ``__aenter__`` are picked up + without waiting for a fresh signal.""" + bus = _FakeBus() + chat = _FakeChatService() + await bus.queue_push( + MessageBusKeys.wakeup_queue(), + {"user_id": "u", "session_id": "pre", "agent_id": "a"}, + ) + + async with WakeupDispatcher( + message_bus=bus, + storage=_FakeStorage(), + chat_service=chat, + chat_run_registry=ChatRunRegistry(), + ): + await _yield_a_few_times() + + self.assertEqual( + chat.calls, + [ + { + "user_id": "u", + "session_id": "pre", + "agent_id": "a", + "input_msg": None, + }, + ], + ) + + async def test_active_session_skipped(self) -> None: + """If the target session is already running, no chat run is + spawned for it.""" + bus = _FakeBus() + chat = _FakeChatService() + bus._locks.add(MessageBus._SESSION_LOCK_KEY.format(sid="busy")) + + async with WakeupDispatcher( + message_bus=bus, + storage=_FakeStorage(), + chat_service=chat, + chat_run_registry=ChatRunRegistry(), + ): + await bus.queue_push( + MessageBusKeys.wakeup_queue(), + {"user_id": "u", "session_id": "busy", "agent_id": "a"}, + ) + await bus.publish(MessageBusKeys.wakeup_signal(), {}) + await asyncio.sleep(0.05) + + self.assertEqual(chat.calls, []) + + async def test_malformed_entry_skipped(self) -> None: + """A wake-up entry missing required fields is logged and skipped, + not raised; later valid entries still dispatch.""" + bus = _FakeBus() + chat = _FakeChatService() + + async with WakeupDispatcher( + message_bus=bus, + storage=_FakeStorage(), + chat_service=chat, + chat_run_registry=ChatRunRegistry(), + ): + await bus.queue_push( + MessageBusKeys.wakeup_queue(), + {"oops": True}, + ) + await bus.queue_push( + MessageBusKeys.wakeup_queue(), + {"user_id": "u", "session_id": "s2", "agent_id": "a"}, + ) + await bus.publish(MessageBusKeys.wakeup_signal(), {}) + await asyncio.wait_for(chat.notify.wait(), timeout=2.0) + + # Only the valid entry made it through. + self.assertEqual( + chat.calls, + [ + { + "user_id": "u", + "session_id": "s2", + "agent_id": "a", + "input_msg": None, + }, + ], + ) + + async def test_deleted_session_skipped(self) -> None: + """A wake-up whose target session no longer exists in storage + is dropped without spawning a chat run; later wake-ups for live + sessions still dispatch.""" + bus = _FakeBus() + chat = _FakeChatService() + storage = _FakeStorage() + storage.missing_session_ids.add("ghost") + + async with WakeupDispatcher( + message_bus=bus, + storage=storage, + chat_service=chat, + chat_run_registry=ChatRunRegistry(), + ): + await bus.queue_push( + MessageBusKeys.wakeup_queue(), + {"user_id": "u", "session_id": "ghost", "agent_id": "a"}, + ) + await bus.queue_push( + MessageBusKeys.wakeup_queue(), + {"user_id": "u", "session_id": "live", "agent_id": "a"}, + ) + await bus.publish(MessageBusKeys.wakeup_signal(), {}) + await asyncio.wait_for(chat.notify.wait(), timeout=2.0) + + self.assertEqual( + chat.calls, + [ + { + "user_id": "u", + "session_id": "live", + "agent_id": "a", + "input_msg": None, + }, + ], + ) + + async def test_resume_idle_spawns_with_parsed_event(self) -> None: + """A ``resume`` trigger for an idle session spawns a run whose + ``input_msg`` is the carried HITL event, rebuilt from its dump.""" + from agentscope.event import UserConfirmResultEvent + + bus = _FakeBus() + chat = _FakeChatService() + event = UserConfirmResultEvent.model_construct( + reply_id="r1", + confirm_results=[], + ) + + async with WakeupDispatcher( + message_bus=bus, + storage=_FakeStorage(), + chat_service=chat, + chat_run_registry=ChatRunRegistry(), + ): + await bus.queue_push( + MessageBusKeys.wakeup_queue(), + { + "user_id": "u", + "session_id": "w1", + "agent_id": "wa1", + "kind": MessageBusKeys.WAKEUP_KIND_RESUME, + "input": event.model_dump(mode="json"), + }, + ) + await bus.publish(MessageBusKeys.wakeup_signal(), {}) + await asyncio.wait_for(chat.notify.wait(), timeout=2.0) + + self.assertEqual(len(chat.calls), 1) + call = chat.calls[0] + self.assertEqual(call["session_id"], "w1") + self.assertIsInstance(call["input_msg"], UserConfirmResultEvent) + self.assertEqual(call["input_msg"].reply_id, "r1") + + async def test_resume_running_session_requeues_until_free(self) -> None: + """A ``resume`` whose target is still running is NOT dropped: it + is re-queued (with backoff) and dispatched once the session lock + releases. This is the structural fix for the parked-run 409 race. + """ + from agentscope.event import UserConfirmResultEvent + + bus = _FakeBus() + chat = _FakeChatService() + lock_key = MessageBus._SESSION_LOCK_KEY.format(sid="w1") + bus._locks.add(lock_key) # session is busy finishing its park tail + event = UserConfirmResultEvent.model_construct( + reply_id="r1", + confirm_results=[], + ) + + async with WakeupDispatcher( + message_bus=bus, + storage=_FakeStorage(), + chat_service=chat, + chat_run_registry=ChatRunRegistry(), + ): + await bus.queue_push( + MessageBusKeys.wakeup_queue(), + { + "user_id": "u", + "session_id": "w1", + "agent_id": "wa1", + "kind": MessageBusKeys.WAKEUP_KIND_RESUME, + "input": event.model_dump(mode="json"), + }, + ) + await bus.publish(MessageBusKeys.wakeup_signal(), {}) + + # While locked, the resume must keep deferring — no run yet. + await asyncio.sleep(0.25) + self.assertEqual(chat.calls, []) + + # Release the lock; the re-queued resume now lands. + bus._locks.discard(lock_key) + await asyncio.wait_for(chat.notify.wait(), timeout=2.0) + + self.assertEqual(len(chat.calls), 1) + self.assertEqual(chat.calls[0]["session_id"], "w1") + self.assertIsInstance( + chat.calls[0]["input_msg"], + UserConfirmResultEvent, + ) + + +class TestWakeupDispatcherLifecycle(IsolatedAsyncioTestCase): + """Tests covering the ``__aenter__`` / ``__aexit__`` ACM behaviour.""" + + async def test_exit_cancels_loop_cleanly(self) -> None: + """``__aexit__`` cancels the dispatcher's loop task and returns + without re-raising the cancellation.""" + bus = _FakeBus() + chat = _FakeChatService() + dispatcher = WakeupDispatcher( + message_bus=bus, + storage=_FakeStorage(), + chat_service=chat, + chat_run_registry=ChatRunRegistry(), + ) + + # pylint: disable=unnecessary-dunder-call + await dispatcher.__aenter__() + loop_task = dispatcher._task + self.assertIsNotNone(loop_task) + + await dispatcher.__aexit__(None, None, None) + + self.assertIsNone(dispatcher._task) + self.assertTrue(loop_task.cancelled() or loop_task.done()) diff --git a/tests/skill_loader_test.py b/tests/skill_loader_test.py new file mode 100644 index 0000000000000000000000000000000000000000..c838cfb133bc6e6594bbc015906af56163347507 --- /dev/null +++ b/tests/skill_loader_test.py @@ -0,0 +1,212 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Test cases for LocalSkillLoader.""" +import os +import tempfile +import shutil +import time +from unittest.async_case import IsolatedAsyncioTestCase + +from agentscope.skill import LocalSkillLoader + + +class SkillLoaderTest(IsolatedAsyncioTestCase): + """Test cases for LocalSkillLoader.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + # Create a temporary directory for testing + self.test_dir = tempfile.mkdtemp() + + # Create test SKILL.md files + # 1. Root level skill + self.root_skill_content = """--- +name: root_skill +description: A skill in the root directory +--- + +This is the root skill content. +""" + with open( + os.path.join(self.test_dir, "SKILL.md"), + "w", + encoding="utf-8", + ) as f: + f.write(self.root_skill_content) + + # 2. Subdirectory skill + self.subdir1 = os.path.join(self.test_dir, "subdir1") + os.makedirs(self.subdir1) + self.subdir1_skill_content = """--- +name: subdir1_skill +description: A skill in subdirectory 1 +--- + +This is the subdir1 skill content. +""" + with open( + os.path.join(self.subdir1, "SKILL.md"), + "w", + encoding="utf-8", + ) as f: + f.write(self.subdir1_skill_content) + + # 3. Nested subdirectory skill + self.subdir2 = os.path.join(self.test_dir, "subdir1", "subdir2") + os.makedirs(self.subdir2) + self.subdir2_skill_content = """--- +name: subdir2_skill +description: A skill in nested subdirectory 2 +--- + +This is the subdir2 skill content. +""" + with open( + os.path.join(self.subdir2, "SKILL.md"), + "w", + encoding="utf-8", + ) as f: + f.write(self.subdir2_skill_content) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + # Remove the temporary directory + if os.path.exists(self.test_dir): + shutil.rmtree(self.test_dir) + + async def test_nonexistent_skill(self) -> None: + """Test loading from a directory without SKILL.md.""" + # Create a directory without SKILL.md + empty_dir = tempfile.mkdtemp() + try: + loader = LocalSkillLoader(empty_dir, scan_subdir=False) + skills = await loader.list_skills() + self.assertEqual(len(skills), 0) + finally: + shutil.rmtree(empty_dir) + + async def test_scan_subdir_flag_controls_subdirectory_scanning( + self, + ) -> None: + """Test that scan_subdir only controls subdirectory scanning. + + The root directory is always scanned regardless of scan_subdir. + When scan_subdir=False, only root SKILL.md is loaded. + When scan_subdir=True, root + all subdirectory SKILL.md files + are loaded. + """ + loader_no_scan = LocalSkillLoader(self.test_dir, scan_subdir=False) + skills_no_scan = await loader_no_scan.list_skills() + + loader_with_scan = LocalSkillLoader(self.test_dir, scan_subdir=True) + skills_with_scan = await loader_with_scan.list_skills() + + # scan_subdir=False: only root skill + self.assertEqual(len(skills_no_scan), 1) + self.assertEqual(skills_no_scan[0].name, "root_skill") + + # scan_subdir=True: root skill + both subdirectory skills + self.assertEqual(len(skills_with_scan), 3) + skill_names = {s.name for s in skills_with_scan} + self.assertIn("root_skill", skill_names) + self.assertIn("subdir1_skill", skill_names) + self.assertIn("subdir2_skill", skill_names) + + async def test_load_skill_without_scan_subdir(self) -> None: + """Test loading skill from root directory only (scan_subdir=False).""" + loader = LocalSkillLoader(self.test_dir, scan_subdir=False) + skills = await loader.list_skills() + + # Should only load the root skill + self.assertEqual(len(skills), 1) + self.assertEqual(skills[0].name, "root_skill") + self.assertEqual( + skills[0].description, + "A skill in the root directory", + ) + self.assertEqual(skills[0].dir, self.test_dir) + self.assertIn("This is the root skill content.", skills[0].markdown) + self.assertIsInstance(skills[0].updated_at, float) + self.assertGreater(skills[0].updated_at, 0) + + async def test_load_skill_with_scan_subdir(self) -> None: + """Test loading skills from subdirectories (scan_subdir=True).""" + loader = LocalSkillLoader(self.test_dir, scan_subdir=True) + skills = await loader.list_skills() + + # Should load all three skills + self.assertEqual(len(skills), 3) + + # Check skill names + skill_names = {skill.name for skill in skills} + self.assertEqual( + skill_names, + {"root_skill", "subdir1_skill", "subdir2_skill"}, + ) + + # Verify each skill has correct attributes + for skill in skills: + self.assertIsInstance(skill.name, str) + self.assertIsInstance(skill.description, str) + self.assertIsInstance(skill.dir, str) + self.assertIsInstance(skill.markdown, str) + self.assertIsInstance(skill.updated_at, float) + self.assertGreater(skill.updated_at, 0) + + async def test_cache_mechanism(self) -> None: + """Test that cache is used when file is not modified.""" + loader = LocalSkillLoader(self.test_dir, scan_subdir=False) + + # First load - should read from file + skills_first = await loader.list_skills() + self.assertEqual(len(skills_first), 1) + first_skill = skills_first[0] + + # Verify cache is populated + self.assertIn(self.test_dir, loader._cache) + cached_skill = loader._cache[self.test_dir] + self.assertEqual(cached_skill.name, first_skill.name) + self.assertEqual(cached_skill.updated_at, first_skill.updated_at) + + # Second load - should use cache (file not modified) + skills_second = await loader.list_skills() + self.assertEqual(len(skills_second), 1) + second_skill = skills_second[0] + + # Should return the same cached object + self.assertIs(second_skill, cached_skill) + self.assertEqual(second_skill.updated_at, first_skill.updated_at) + + # Modify the file + time.sleep(0.01) # Ensure different mtime + modified_content = """--- +name: modified_root_skill +description: Modified skill description +--- + +This is the modified content. +""" + with open( + os.path.join(self.test_dir, "SKILL.md"), + "w", + encoding="utf-8", + ) as f: + f.write(modified_content) + + # Third load - should detect change and reload + skills_third = await loader.list_skills() + self.assertEqual(len(skills_third), 1) + third_skill = skills_third[0] + + # Should have new content + self.assertEqual(third_skill.name, "modified_root_skill") + self.assertEqual(third_skill.description, "Modified skill description") + self.assertIn("This is the modified content.", third_skill.markdown) + self.assertNotEqual(third_skill.updated_at, first_skill.updated_at) + + # Cache should be updated + self.assertIn(self.test_dir, loader._cache) + self.assertEqual( + loader._cache[self.test_dir].name, + "modified_root_skill", + ) diff --git a/tests/storage_redis_knowledge_base_test.py b/tests/storage_redis_knowledge_base_test.py new file mode 100644 index 0000000000000000000000000000000000000000..fb3cc95c446185e13faddbd818b36e373d7ac2a3 --- /dev/null +++ b/tests/storage_redis_knowledge_base_test.py @@ -0,0 +1,110 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for the knowledge base persistence layer of RedisStorage.""" +from unittest.async_case import IsolatedAsyncioTestCase + +import fakeredis.aioredis + +from agentscope.app.storage import ( + EmbeddingModelConfig, + KnowledgeBaseRecord, + RedisStorage, +) + + +def make_storage() -> RedisStorage: + """Create a RedisStorage instance backed by fakeredis.""" + storage = RedisStorage.__new__(RedisStorage) + storage._client = fakeredis.aioredis.FakeRedis(decode_responses=True) + storage.key_ttl = None + storage.key_config = RedisStorage.KeyConfig() + return storage + + +def make_record(user_id: str, name: str = "kb") -> KnowledgeBaseRecord: + """Build a KnowledgeBaseRecord with a default embedding config.""" + return KnowledgeBaseRecord( + user_id=user_id, + name=name, + description="desc", + embedding_model_config=EmbeddingModelConfig( + type="openai_credential", + credential_id="cred-1", + model="text-embedding-3-small", + dimensions=1536, + ), + collection_name="kb_abc", + ) + + +class KnowledgeBaseStorageTest(IsolatedAsyncioTestCase): + """Tests for the KnowledgeBaseRecord CRUD methods on RedisStorage.""" + + async def test_upsert_get_list_delete(self) -> None: + """Round-trip and isolation across users.""" + storage = make_storage() + + rec_a = make_record("user-1", "first") + rec_b = make_record("user-1", "second") + rec_c = make_record("user-2", "other") + + stored_a = await storage.upsert_knowledge_base("user-1", rec_a) + stored_b = await storage.upsert_knowledge_base("user-1", rec_b) + stored_c = await storage.upsert_knowledge_base("user-2", rec_c) + + # get returns the persisted record for the right owner + fetched = await storage.get_knowledge_base("user-1", stored_a.id) + self.assertEqual(fetched.id, stored_a.id) + self.assertEqual(fetched.name, "first") + + # cross-user lookups return None + self.assertIsNone( + await storage.get_knowledge_base("user-2", stored_a.id), + ) + + # list scoped to user-1 returns only its records + listed = await storage.list_knowledge_bases("user-1") + self.assertEqual( + sorted(r.id for r in listed), + sorted([stored_a.id, stored_b.id]), + ) + + # delete only the requested record; other users are untouched + deleted = await storage.delete_knowledge_base("user-1", stored_a.id) + self.assertTrue(deleted) + self.assertIsNone( + await storage.get_knowledge_base("user-1", stored_a.id), + ) + self.assertEqual( + [r.id for r in await storage.list_knowledge_bases("user-1")], + [stored_b.id], + ) + self.assertEqual( + [r.id for r in await storage.list_knowledge_bases("user-2")], + [stored_c.id], + ) + + # double-delete returns False + self.assertFalse( + await storage.delete_knowledge_base("user-1", stored_a.id), + ) + + async def test_upsert_rejects_user_id_mismatch(self) -> None: + """upsert refuses records whose user_id does not match arg.""" + storage = make_storage() + rec = make_record("user-1") + with self.assertRaises(ValueError): + await storage.upsert_knowledge_base("user-2", rec) + + async def test_upsert_overwrites_and_preserves_created_at(self) -> None: + """Re-upsert with same id keeps created_at, refreshes updated_at.""" + storage = make_storage() + rec = make_record("user-1") + first = await storage.upsert_knowledge_base("user-1", rec) + + rec.name = "renamed" + second = await storage.upsert_knowledge_base("user-1", rec) + + self.assertEqual(second.id, first.id) + self.assertEqual(second.created_at, first.created_at) + self.assertEqual(second.name, "renamed") diff --git a/tests/storage_redis_test.py b/tests/storage_redis_test.py new file mode 100644 index 0000000000000000000000000000000000000000..6d0369c66b5f63ea82e9423ed89fa993c605aa22 --- /dev/null +++ b/tests/storage_redis_test.py @@ -0,0 +1,1110 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for RedisStorage using fakeredis.""" + +from unittest.async_case import IsolatedAsyncioTestCase + +import fakeredis.aioredis + +from agentscope.app.storage import ( + RedisStorage, + AgentRecord, + SessionConfig, + SessionRecord, + ChatModelConfig, + ScheduleRecord, + ScheduleData, + SessionSource, + TeamData, + TeamRecord, +) +from agentscope.credential import OllamaCredential +from agentscope.app.storage import AgentData +from agentscope.agent import ContextConfig, ReActConfig +from agentscope.message import UserMsg, AssistantMsg, TextBlock +from agentscope.state import AgentState + + +def make_storage() -> RedisStorage: + """Create a RedisStorage instance backed by fakeredis.""" + storage = RedisStorage.__new__(RedisStorage) + # pylint: disable=protected-access + storage._client = fakeredis.aioredis.FakeRedis(decode_responses=True) + storage.key_ttl = None + storage.key_config = RedisStorage.KeyConfig() + return storage + + +def make_agent_record(user_id: str) -> AgentRecord: + """Create a test AgentRecord with all-default sub-configs.""" + return AgentRecord( + user_id=user_id, + data=AgentData( + id="agent-data-id", + name="test-agent", + system_prompt="You are a helpful assistant.", + context_config=ContextConfig(), + react_config=ReActConfig(), + ), + ) + + +def make_session_config(workspace_id: str = "ws-1") -> SessionConfig: + """Create a test SessionConfig with a chat model config.""" + return SessionConfig( + workspace_id=workspace_id, + chat_model_config=ChatModelConfig( + type="openai", + credential_id="cred-1", + model="gpt-4", + parameters={}, + ), + ) + + +class TestCredential(IsolatedAsyncioTestCase): + """Tests for credential CRUD and cascading operations.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.storage = make_storage() + self.user_id = "user-1" + + async def test_create(self) -> None: + """Create a credential and verify it is retrievable via list.""" + cred_id = await self.storage.upsert_credential( + self.user_id, + OllamaCredential(host="http://localhost:11434"), + ) + records = await self.storage.list_credentials(self.user_id) + self.assertEqual(len(records), 1) + self.assertEqual(records[0].id, cred_id) + self.assertEqual(records[0].data.get("type"), "ollama_credential") + self.assertEqual( + records[0].data.get("host"), + "http://localhost:11434", + ) + + async def test_list_empty(self) -> None: + """Verify list returns empty when no records exist.""" + records = await self.storage.list_credentials(self.user_id) + self.assertEqual(records, []) + + async def test_update_in_place(self) -> None: + """Update a credential and verify data changed without adding + a new record.""" + cred_id = await self.storage.upsert_credential( + self.user_id, + OllamaCredential(host="http://old-host:11434"), + ) + await self.storage.upsert_credential( + self.user_id, + OllamaCredential(id=cred_id, host="http://new-host:11434"), + ) + records = await self.storage.list_credentials(self.user_id) + self.assertEqual(len(records), 1) + self.assertEqual(records[0].data.get("host"), "http://new-host:11434") + + async def test_delete(self) -> None: + """Delete a credential and verify it is gone from Redis.""" + cred_id = await self.storage.upsert_credential( + self.user_id, + OllamaCredential(host="http://localhost:11434"), + ) + result = await self.storage.delete_credential(self.user_id, cred_id) + self.assertTrue(result) + records = await self.storage.list_credentials(self.user_id) + self.assertEqual(records, []) + + async def test_delete_nonexistent(self) -> None: + """Verify delete returns False for non-existent record.""" + result = await self.storage.delete_credential( + self.user_id, + "no-such-id", + ) + self.assertFalse(result) + + async def test_user_isolation(self) -> None: + """Verify different users cannot see each other's records.""" + await self.storage.upsert_credential( + "user-A", + OllamaCredential(host="http://localhost:11434"), + ) + records = await self.storage.list_credentials("user-B") + self.assertEqual(records, []) + + +class TestAgent(IsolatedAsyncioTestCase): + """Tests for agent CRUD and cascading operations.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.storage = make_storage() + self.user_id = "user-1" + + async def test_create(self) -> None: + """Create an agent and verify it is retrievable via list.""" + record = make_agent_record(self.user_id) + agent_id = await self.storage.upsert_agent(self.user_id, record) + records = await self.storage.list_agents(self.user_id) + self.assertEqual(len(records), 1) + self.assertEqual(records[0].id, agent_id) + self.assertEqual(records[0].data.name, "test-agent") + + async def test_list_empty(self) -> None: + """Verify list returns empty when no records exist.""" + records = await self.storage.list_agents(self.user_id) + self.assertEqual(records, []) + + async def test_delete(self) -> None: + """Delete an agent and verify it is gone from Redis.""" + record = make_agent_record(self.user_id) + await self.storage.upsert_agent(self.user_id, record) + result = await self.storage.delete_agent(self.user_id, record.id) + self.assertTrue(result) + records = await self.storage.list_agents(self.user_id) + self.assertEqual(records, []) + + async def test_delete_nonexistent(self) -> None: + """Verify delete returns False for non-existent record.""" + result = await self.storage.delete_agent(self.user_id, "no-such-id") + self.assertFalse(result) + + async def test_user_isolation(self) -> None: + """Verify different users cannot see each other's records.""" + await self.storage.upsert_agent("user-A", make_agent_record("user-A")) + records = await self.storage.list_agents("user-B") + self.assertEqual(records, []) + + +class TestSession(IsolatedAsyncioTestCase): + """Tests for session CRUD and cascading operations.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.storage = make_storage() + self.user_id = "user-1" + self.agent_id = "agent-1" + self.workspace_id = "ws-1" + + async def test_create(self) -> None: + """Create a session and verify it is retrievable via list.""" + await self.storage.upsert_session( + self.user_id, + self.agent_id, + make_session_config(self.workspace_id), + ) + records = await self.storage.list_sessions(self.user_id, self.agent_id) + self.assertEqual(len(records), 1) + self.assertEqual(records[0].config.workspace_id, self.workspace_id) + self.assertEqual(records[0].agent_id, self.agent_id) + + async def test_list_empty(self) -> None: + """Verify list returns empty when no records exist.""" + records = await self.storage.list_sessions(self.user_id, self.agent_id) + self.assertEqual(records, []) + + async def test_upsert_same_triple_updates_in_place(self) -> None: + """Second upsert with the same session_id must update the existing + record, not create a second one.""" + session = await self.storage.upsert_session( + self.user_id, + self.agent_id, + make_session_config(self.workspace_id), + ) + first_id = session.id + + await self.storage.upsert_session( + self.user_id, + self.agent_id, + make_session_config(self.workspace_id), + session_id=first_id, + ) + records_after = await self.storage.list_sessions( + self.user_id, + self.agent_id, + ) + self.assertEqual(len(records_after), 1) + self.assertEqual(records_after[0].id, first_id) + + async def test_create_with_explicit_session_id_uses_that_key(self) -> None: + """A caller-provided session_id should be the stored record id.""" + session_id = "session-from-router" + + session = await self.storage.upsert_session( + self.user_id, + self.agent_id, + make_session_config(self.workspace_id), + session_id=session_id, + ) + + self.assertEqual(session.id, session_id) + fetched = await self.storage.get_session( + self.user_id, + self.agent_id, + session_id, + ) + self.assertIsNotNone(fetched) + self.assertEqual(fetched.id, session_id) + + await self.storage.update_session_state( + self.user_id, + self.agent_id, + session_id, + AgentState(), + ) + records = await self.storage.list_sessions(self.user_id, self.agent_id) + self.assertEqual([record.id for record in records], [session_id]) + + async def test_delete(self) -> None: + """Delete a session and verify it is gone from Redis.""" + await self.storage.upsert_session( + self.user_id, + self.agent_id, + make_session_config(self.workspace_id), + ) + records = await self.storage.list_sessions(self.user_id, self.agent_id) + result = await self.storage.delete_session( + self.user_id, + self.agent_id, + records[0].id, + ) + self.assertTrue(result) + remaining = await self.storage.list_sessions( + self.user_id, + self.agent_id, + ) + self.assertEqual(remaining, []) + + async def test_delete_cascades_lookup_key(self) -> None: + """Deleting a session must remove the lookup key so a subsequent upsert + for the same (user, agent) pair creates a fresh session with a new + id.""" + await self.storage.upsert_session( + self.user_id, + self.agent_id, + make_session_config(self.workspace_id), + ) + records = await self.storage.list_sessions(self.user_id, self.agent_id) + old_id = records[0].id + + await self.storage.delete_session(self.user_id, self.agent_id, old_id) + + await self.storage.upsert_session( + self.user_id, + self.agent_id, + make_session_config(self.workspace_id), + ) + new_records = await self.storage.list_sessions( + self.user_id, + self.agent_id, + ) + self.assertEqual(len(new_records), 1) + self.assertNotEqual(new_records[0].id, old_id) + + async def test_delete_nonexistent(self) -> None: + """Verify delete returns False for non-existent record.""" + result = await self.storage.delete_session( + self.user_id, + self.agent_id, + "no-such-id", + ) + self.assertFalse(result) + + async def test_agent_isolation(self) -> None: + """Verify different agents cannot see each other's sessions.""" + await self.storage.upsert_session( + self.user_id, + "agent-A", + make_session_config(self.workspace_id), + ) + records = await self.storage.list_sessions(self.user_id, "agent-B") + self.assertEqual(records, []) + + +class TestMessage(IsolatedAsyncioTestCase): + """Tests for message persistence: upsert_message, get_message and + list_messages.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.storage = make_storage() + self.user_id = "user-1" + self.session_id = "session-1" + + async def test_upsert_appends_new_message(self) -> None: + """Upserting a new message appends it to the session list.""" + msg = UserMsg(name="alice", content="hello") + await self.storage.upsert_message(self.user_id, self.session_id, msg) + messages = await self.storage.list_messages( + self.user_id, + self.session_id, + ) + self.assertListEqual( + [m.model_dump() for m in messages], + [msg.model_dump()], + ) + + async def test_upsert_refreshes_message_list_ttl(self) -> None: + """Message list keys expire with the session storage TTL.""" + self.storage.key_ttl = 60 + msg = UserMsg(name="alice", content="hello") + + await self.storage.upsert_message(self.user_id, self.session_id, msg) + + ttl = await self.storage._client.ttl( + self.storage._message_key(self.user_id, self.session_id), + ) + self.assertGreater(ttl, 0) + + async def test_upsert_replaces_last_message_with_same_id(self) -> None: + """Upserting a message whose id matches the last entry replaces it + in-place (streaming overwrite), rather than creating a duplicate.""" + msg = AssistantMsg(name="bot", content="v1") + await self.storage.upsert_message(self.user_id, self.session_id, msg) + + # Keep the same id but replace content — simulates a streaming update. + updated = msg.model_copy( + update={"content": [TextBlock(text="v2")]}, + ) + await self.storage.upsert_message( + self.user_id, + self.session_id, + updated, + ) + + messages = await self.storage.list_messages( + self.user_id, + self.session_id, + ) + self.assertListEqual( + [m.model_dump() for m in messages], + [updated.model_dump()], + "Duplicate must not be created; existing entry must be replaced.", + ) + + async def test_replace_refreshes_message_list_ttl(self) -> None: + """Streaming message replacement keeps the message key expiring.""" + self.storage.key_ttl = 60 + msg = AssistantMsg(name="bot", content="v1") + await self.storage.upsert_message(self.user_id, self.session_id, msg) + await self.storage._client.persist( + self.storage._message_key(self.user_id, self.session_id), + ) + updated = msg.model_copy( + update={"content": [TextBlock(text="v2")]}, + ) + + await self.storage.upsert_message( + self.user_id, + self.session_id, + updated, + ) + + ttl = await self.storage._client.ttl( + self.storage._message_key(self.user_id, self.session_id), + ) + self.assertGreater(ttl, 0) + + async def test_upsert_appends_when_id_differs_from_last(self) -> None: + """Upserting a message with a different id than the last always + appends, even if an earlier message shares the same id.""" + msg1 = UserMsg(name="alice", content="first") + msg2 = UserMsg(name="alice", content="second") + await self.storage.upsert_message(self.user_id, self.session_id, msg1) + await self.storage.upsert_message(self.user_id, self.session_id, msg2) + messages = await self.storage.list_messages( + self.user_id, + self.session_id, + ) + self.assertListEqual( + [m.model_dump() for m in messages], + [msg1.model_dump(), msg2.model_dump()], + ) + + async def test_get_message_returns_correct_message(self) -> None: + """get_message fetches the message matching the given id.""" + msg1 = UserMsg(name="alice", content="first") + msg2 = UserMsg(name="alice", content="second") + await self.storage.upsert_message(self.user_id, self.session_id, msg1) + await self.storage.upsert_message(self.user_id, self.session_id, msg2) + + fetched = await self.storage.get_message( + self.user_id, + self.session_id, + msg1.id, + ) + self.assertIsNotNone(fetched) + self.assertDictEqual(fetched.model_dump(), msg1.model_dump()) + + async def test_get_message_nonexistent_returns_none(self) -> None: + """get_message returns None when the message id does not exist.""" + result = await self.storage.get_message( + self.user_id, + self.session_id, + "no-such-id", + ) + self.assertIsNone(result) + + async def test_list_messages_empty_session(self) -> None: + """list_messages returns an empty list for a session with no + messages.""" + messages = await self.storage.list_messages( + self.user_id, + self.session_id, + ) + self.assertListEqual(messages, []) + + async def test_list_messages_pagination(self) -> None: + """list_messages respects offset and limit parameters.""" + msgs = [UserMsg(name="alice", content=f"msg-{i}") for i in range(5)] + for m in msgs: + await self.storage.upsert_message( + self.user_id, + self.session_id, + m, + ) + + # Fetch the middle slice: offset=1, limit=3 → msgs[1], msgs[2], msgs[3] + page = await self.storage.list_messages( + self.user_id, + self.session_id, + offset=1, + limit=3, + ) + self.assertListEqual( + [m.model_dump() for m in page], + [m.model_dump() for m in msgs[1:4]], + ) + + async def test_list_messages_order_preserved(self) -> None: + """Messages are returned in the insertion order (chronological).""" + msgs = [ + UserMsg(name="alice", content=text) + for text in ["alpha", "beta", "gamma"] + ] + for m in msgs: + await self.storage.upsert_message( + self.user_id, + self.session_id, + m, + ) + messages = await self.storage.list_messages( + self.user_id, + self.session_id, + ) + self.assertListEqual( + [m.model_dump() for m in messages], + [m.model_dump() for m in msgs], + ) + + async def test_session_isolation(self) -> None: + """Messages belonging to different sessions do not interfere.""" + await self.storage.upsert_message( + self.user_id, + "session-A", + UserMsg(name="alice", content="in A"), + ) + messages = await self.storage.list_messages( + self.user_id, + "session-B", + ) + self.assertListEqual(messages, []) + + +def make_schedule_record(user_id: str, agent_id: str) -> ScheduleRecord: + """Create a test ScheduleRecord.""" + return ScheduleRecord( + user_id=user_id, + agent_id=agent_id, + data=ScheduleData( + name="test-schedule", + cron_expression="0 9 * * *", + started_at="2026-01-01T00:00:00", + chat_model_config=ChatModelConfig( + type="openai", + credential_id="cred-1", + model="gpt-4", + parameters={}, + ), + ), + ) + + +class TestScheduleSession(IsolatedAsyncioTestCase): + """Tests for schedule-session index and cascade deletion.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.storage = make_storage() + self.user_id = "user-1" + self.agent_id = "agent-1" + + async def test_list_sessions_by_schedule(self) -> None: + """Sessions created with source_schedule_id are queryable by + schedule.""" + schedule = make_schedule_record(self.user_id, self.agent_id) + await self.storage.upsert_schedule(self.user_id, schedule) + + session = await self.storage.upsert_session( + self.user_id, + self.agent_id, + make_session_config(), + source=SessionSource.SCHEDULE, + source_schedule_id=schedule.id, + ) + + results = await self.storage.list_sessions_by_schedule( + self.user_id, + schedule.id, + ) + self.assertEqual(len(results), 1) + self.assertEqual(results[0].id, session.id) + self.assertEqual(results[0].source_schedule_id, schedule.id) + + async def test_list_sessions_by_schedule_empty(self) -> None: + """Returns empty list when no sessions exist for a schedule.""" + results = await self.storage.list_sessions_by_schedule( + self.user_id, + "nonexistent-schedule", + ) + self.assertEqual(results, []) + + async def test_schedule_session_also_in_agent_index(self) -> None: + """A schedule-created session appears in both the schedule and agent + session indexes.""" + schedule = make_schedule_record(self.user_id, self.agent_id) + await self.storage.upsert_schedule(self.user_id, schedule) + + session = await self.storage.upsert_session( + self.user_id, + self.agent_id, + make_session_config(), + source=SessionSource.SCHEDULE, + source_schedule_id=schedule.id, + ) + + agent_sessions = await self.storage.list_sessions( + self.user_id, + self.agent_id, + ) + self.assertEqual(len(agent_sessions), 1) + self.assertEqual(agent_sessions[0].id, session.id) + + async def test_delete_schedule_cascades_sessions(self) -> None: + """Deleting a schedule removes all its execution sessions.""" + schedule = make_schedule_record(self.user_id, self.agent_id) + await self.storage.upsert_schedule(self.user_id, schedule) + + await self.storage.upsert_session( + self.user_id, + self.agent_id, + make_session_config(), + source=SessionSource.SCHEDULE, + source_schedule_id=schedule.id, + ) + await self.storage.upsert_session( + self.user_id, + self.agent_id, + make_session_config(), + source=SessionSource.SCHEDULE, + source_schedule_id=schedule.id, + ) + + await self.storage.delete_schedule(self.user_id, schedule.id) + + schedule_sessions = await self.storage.list_sessions_by_schedule( + self.user_id, + schedule.id, + ) + self.assertEqual(schedule_sessions, []) + + agent_sessions = await self.storage.list_sessions( + self.user_id, + self.agent_id, + ) + self.assertEqual(agent_sessions, []) + + async def test_delete_session_cleans_schedule_index(self) -> None: + """Deleting a session removes it from the schedule session index.""" + schedule = make_schedule_record(self.user_id, self.agent_id) + await self.storage.upsert_schedule(self.user_id, schedule) + + session = await self.storage.upsert_session( + self.user_id, + self.agent_id, + make_session_config(), + source=SessionSource.SCHEDULE, + source_schedule_id=schedule.id, + ) + + await self.storage.delete_session( + self.user_id, + self.agent_id, + session.id, + ) + + results = await self.storage.list_sessions_by_schedule( + self.user_id, + schedule.id, + ) + self.assertEqual(results, []) + + +def make_team_record( + user_id: str, + session_id: str = "leader-session-1", + name: str = "test-team", + member_ids: list[str] | None = None, +) -> TeamRecord: + """Create a test TeamRecord.""" + return TeamRecord( + user_id=user_id, + session_id=session_id, + data=TeamData( + name=name, + member_ids=member_ids if member_ids is not None else [], + ), + ) + + +class TestAgentSource(IsolatedAsyncioTestCase): + """Tests for the ``source`` field on AgentRecord.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.storage = make_storage() + self.user_id = "user-1" + + async def test_default_source_is_user(self) -> None: + """An AgentRecord built without ``source`` defaults to ``"user"``.""" + record = make_agent_record(self.user_id) + self.assertEqual(record.source, "user") + await self.storage.upsert_agent(self.user_id, record) + loaded = await self.storage.get_agent(self.user_id, record.id) + self.assertIsNotNone(loaded) + self.assertEqual(loaded.source, "user") + + async def test_team_source_round_trips(self) -> None: + """A worker AgentRecord persists ``source='team'`` through Redis.""" + record = AgentRecord( + user_id=self.user_id, + source="team", + data=AgentData( + id="worker-1", + name="worker", + system_prompt="You are a team worker.", + context_config=ContextConfig(), + react_config=ReActConfig(), + ), + ) + await self.storage.upsert_agent(self.user_id, record) + loaded = await self.storage.get_agent(self.user_id, record.id) + self.assertIsNotNone(loaded) + self.assertEqual(loaded.source, "team") + + async def test_list_agents_filters_out_team_workers(self) -> None: + """``storage.list_agents`` returns only ``source='user'`` agents. + Workers (``source='team'``) are scoped to a team and addressable + only via direct id lookup.""" + user_agent = make_agent_record(self.user_id) + worker = AgentRecord( + id="worker-1", + user_id=self.user_id, + source="team", + data=AgentData( + id="data-w", + name="worker", + system_prompt="", + context_config=ContextConfig(), + react_config=ReActConfig(), + ), + ) + await self.storage.upsert_agent(self.user_id, user_agent) + await self.storage.upsert_agent(self.user_id, worker) + + listed = await self.storage.list_agents(self.user_id) + listed_ids = {a.id for a in listed} + self.assertIn(user_agent.id, listed_ids) + self.assertNotIn(worker.id, listed_ids) + + # But direct lookup still works for the worker. + loaded_worker = await self.storage.get_agent( + self.user_id, + worker.id, + ) + self.assertIsNotNone(loaded_worker) + self.assertEqual(loaded_worker.source, "team") + + async def test_legacy_record_without_source_deserializes(self) -> None: + """JSON written before the ``source`` field was added still loads + and falls back to the default ``"user"``.""" + legacy_record = make_agent_record(self.user_id) + legacy_json = legacy_record.model_dump_json() + # Strip the new field to simulate pre-migration data. + import json + + payload = json.loads(legacy_json) + payload.pop("source", None) + # pylint: disable=protected-access + key = self.storage._key( + self.storage.key_config.agent, + user_id=self.user_id, + agent_id=legacy_record.id, + ) + await self.storage._client.set(key, json.dumps(payload)) + await self.storage._client.sadd( + self.storage._key( + self.storage.key_config.agent_index, + user_id=self.user_id, + ), + legacy_record.id, + ) + + loaded = await self.storage.get_agent(self.user_id, legacy_record.id) + self.assertIsNotNone(loaded) + self.assertEqual(loaded.source, "user") + + +class TestSessionTeamId(IsolatedAsyncioTestCase): + """Tests for the ``team_id`` field on SessionRecord.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.storage = make_storage() + self.user_id = "user-1" + self.agent_id = "agent-1" + + async def test_default_team_id_is_none(self) -> None: + """A SessionRecord built without ``team_id`` defaults to ``None``.""" + session = await self.storage.upsert_session( + self.user_id, + self.agent_id, + make_session_config(), + ) + self.assertIsNone(session.team_id) + + loaded = await self.storage.get_session( + self.user_id, + self.agent_id, + session.id, + ) + self.assertIsNotNone(loaded) + self.assertIsNone(loaded.team_id) + + async def test_legacy_session_without_team_id_deserializes(self) -> None: + """JSON written before ``team_id`` existed still loads with default + ``None``.""" + session = await self.storage.upsert_session( + self.user_id, + self.agent_id, + make_session_config(), + ) + + # Strip the new field from the persisted JSON to simulate pre-migration + # data, then write it back at the same key. + import json + + # pylint: disable=protected-access + key = self.storage._key( + self.storage.key_config.session, + user_id=self.user_id, + session_id=session.id, + ) + raw = await self.storage._client.get(key) + payload = json.loads(raw) + payload.pop("team_id", None) + await self.storage._client.set(key, json.dumps(payload)) + + loaded = await self.storage.get_session( + self.user_id, + self.agent_id, + session.id, + ) + self.assertIsNotNone(loaded) + self.assertIsNone(loaded.team_id) + + +class TestTeam(IsolatedAsyncioTestCase): + """Tests for team CRUD.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.storage = make_storage() + self.user_id = "user-1" + + async def test_create(self) -> None: + """Create a team and verify it is retrievable via list.""" + record = make_team_record(self.user_id) + stored = await self.storage.upsert_team(self.user_id, record) + records = await self.storage.list_teams(self.user_id) + self.assertEqual(len(records), 1) + self.assertEqual(records[0].id, stored.id) + self.assertEqual(records[0].session_id, "leader-session-1") + self.assertEqual(records[0].data.name, "test-team") + self.assertEqual(records[0].data.member_ids, []) + + async def test_list_empty(self) -> None: + """Verify list returns empty when no teams exist.""" + records = await self.storage.list_teams(self.user_id) + self.assertEqual(records, []) + + async def test_get_returns_record(self) -> None: + """get_team returns the persisted record by id.""" + record = make_team_record( + self.user_id, + member_ids=["worker-a", "worker-b"], + ) + await self.storage.upsert_team(self.user_id, record) + loaded = await self.storage.get_team(self.user_id, record.id) + self.assertIsNotNone(loaded) + self.assertEqual(loaded.id, record.id) + self.assertEqual(loaded.data.member_ids, ["worker-a", "worker-b"]) + + async def test_get_nonexistent_returns_none(self) -> None: + """get_team returns None when the id does not exist.""" + loaded = await self.storage.get_team(self.user_id, "no-such-id") + self.assertIsNone(loaded) + + async def test_update_in_place(self) -> None: + """Upsert with the same id overwrites the existing record.""" + record = make_team_record(self.user_id, name="original") + await self.storage.upsert_team(self.user_id, record) + + record.data.name = "renamed" + record.data.member_ids = ["new-worker"] + await self.storage.upsert_team(self.user_id, record) + + records = await self.storage.list_teams(self.user_id) + self.assertEqual(len(records), 1) + self.assertEqual(records[0].data.name, "renamed") + self.assertEqual(records[0].data.member_ids, ["new-worker"]) + + async def test_upsert_refreshes_updated_at(self) -> None: + """upsert_team bumps ``updated_at`` on each write.""" + record = make_team_record(self.user_id) + first = await self.storage.upsert_team(self.user_id, record) + first_updated = first.updated_at + + # Second write must have a strictly non-decreasing updated_at. + second = await self.storage.upsert_team(self.user_id, record) + self.assertGreaterEqual(second.updated_at, first_updated) + + async def test_delete(self) -> None: + """Delete a team and verify it is gone from Redis and the index.""" + record = make_team_record(self.user_id) + await self.storage.upsert_team(self.user_id, record) + + result = await self.storage.delete_team(self.user_id, record.id) + self.assertTrue(result) + + loaded = await self.storage.get_team(self.user_id, record.id) + self.assertIsNone(loaded) + self.assertEqual(await self.storage.list_teams(self.user_id), []) + + async def test_delete_nonexistent(self) -> None: + """delete_team returns False for an unknown id.""" + result = await self.storage.delete_team(self.user_id, "no-such-id") + self.assertFalse(result) + + async def test_user_isolation(self) -> None: + """Teams from one user are invisible to another.""" + await self.storage.upsert_team( + "user-A", + make_team_record("user-A"), + ) + records = await self.storage.list_teams("user-B") + self.assertEqual(records, []) + + +def make_worker_agent(user_id: str, agent_id: str) -> AgentRecord: + """Create a source='team' worker AgentRecord.""" + return AgentRecord( + id=agent_id, + user_id=user_id, + source="team", + data=AgentData( + id=f"data-{agent_id}", + name=f"name-{agent_id}", + system_prompt="worker", + context_config=ContextConfig(), + react_config=ReActConfig(), + ), + ) + + +class TestTeamCascade(IsolatedAsyncioTestCase): + """End-to-end cascade tests for the team relationship graph.""" + + async def asyncSetUp(self) -> None: + """Build a small team fixture: 1 leader user-agent + 2 workers.""" + self.storage = make_storage() + self.user_id = "user-1" + + # Leader agent (user-created) and its session + self.leader_agent = make_agent_record(self.user_id) + await self.storage.upsert_agent(self.user_id, self.leader_agent) + leader_session = await self.storage.upsert_session( + self.user_id, + self.leader_agent.id, + make_session_config(), + ) + self.leader_session_id = leader_session.id + + # Two worker agents and their sessions + self.worker_a = make_worker_agent(self.user_id, "worker-a") + self.worker_b = make_worker_agent(self.user_id, "worker-b") + await self.storage.upsert_agent(self.user_id, self.worker_a) + await self.storage.upsert_agent(self.user_id, self.worker_b) + worker_a_session = await self.storage.upsert_session( + self.user_id, + self.worker_a.id, + make_session_config(), + ) + worker_b_session = await self.storage.upsert_session( + self.user_id, + self.worker_b.id, + make_session_config(), + ) + self.worker_a_session_id = worker_a_session.id + self.worker_b_session_id = worker_b_session.id + + # The team itself + leader's team_id back-reference + self.team = make_team_record( + self.user_id, + session_id=self.leader_session_id, + member_ids=[self.worker_a.id, self.worker_b.id], + ) + await self.storage.upsert_team(self.user_id, self.team) + + # Stamp team_id on every team-participating session + # pylint: disable=protected-access + for sid in [ + self.leader_session_id, + self.worker_a_session_id, + self.worker_b_session_id, + ]: + session_key = self.storage._key( + self.storage.key_config.session, + user_id=self.user_id, + session_id=sid, + ) + raw = await self.storage._client.get(session_key) + rec = SessionRecord.model_validate_json(raw) + rec.team_id = self.team.id + await self.storage._client.set(session_key, rec.model_dump_json()) + + async def test_delete_team_cascades_workers_and_clears_leader( + self, + ) -> None: + """delete_team removes all workers + sessions + clears leader.""" + result = await self.storage.delete_team(self.user_id, self.team.id) + self.assertTrue(result) + + # Team record gone + self.assertIsNone( + await self.storage.get_team(self.user_id, self.team.id), + ) + # Worker agents gone + self.assertIsNone( + await self.storage.get_agent(self.user_id, self.worker_a.id), + ) + self.assertIsNone( + await self.storage.get_agent(self.user_id, self.worker_b.id), + ) + # Worker sessions gone + self.assertIsNone( + await self.storage.get_session( + self.user_id, + self.worker_a.id, + self.worker_a_session_id, + ), + ) + # Leader session still exists, team_id cleared + leader = await self.storage.get_session( + self.user_id, + self.leader_agent.id, + self.leader_session_id, + ) + self.assertIsNotNone(leader) + self.assertIsNone(leader.team_id) + + async def test_delete_leader_session_dissolves_team(self) -> None: + """Deleting a leader session auto-dissolves its team.""" + await self.storage.delete_session( + self.user_id, + self.leader_agent.id, + self.leader_session_id, + ) + + # Team gone + self.assertIsNone( + await self.storage.get_team(self.user_id, self.team.id), + ) + # Workers gone + self.assertIsNone( + await self.storage.get_agent(self.user_id, self.worker_a.id), + ) + # Leader agent itself still exists (only the session was deleted) + self.assertIsNotNone( + await self.storage.get_agent(self.user_id, self.leader_agent.id), + ) + + async def test_delete_leader_agent_dissolves_all_its_teams(self) -> None: + """Deleting the leader agent dissolves every team it leads.""" + await self.storage.delete_agent(self.user_id, self.leader_agent.id) + + self.assertIsNone( + await self.storage.get_team(self.user_id, self.team.id), + ) + self.assertIsNone( + await self.storage.get_agent(self.user_id, self.worker_a.id), + ) + self.assertIsNone( + await self.storage.get_agent(self.user_id, self.leader_agent.id), + ) + + async def test_direct_delete_worker_agent_scrubs_member_ids(self) -> None: + """Bypassing delete_team still keeps team.member_ids consistent.""" + await self.storage.delete_agent(self.user_id, self.worker_a.id) + + team = await self.storage.get_team(self.user_id, self.team.id) + self.assertIsNotNone(team) + self.assertEqual(team.data.member_ids, [self.worker_b.id]) + # The other worker is untouched + self.assertIsNotNone( + await self.storage.get_agent(self.user_id, self.worker_b.id), + ) + + async def test_direct_delete_worker_session_does_not_dissolve_team( + self, + ) -> None: + """Deleting a worker's session leaves the team intact (asymmetric + with leader-session deletion — there's no FK from session to its + owning agent).""" + await self.storage.delete_session( + self.user_id, + self.worker_a.id, + self.worker_a_session_id, + ) + + # Team and worker agent still exist + self.assertIsNotNone( + await self.storage.get_team(self.user_id, self.team.id), + ) + self.assertIsNotNone( + await self.storage.get_agent(self.user_id, self.worker_a.id), + ) + # The session is gone though + self.assertIsNone( + await self.storage.get_session( + self.user_id, + self.worker_a.id, + self.worker_a_session_id, + ), + ) + + async def test_delete_team_idempotent_on_nonexistent(self) -> None: + """delete_team on a missing team returns False without crashing.""" + result = await self.storage.delete_team(self.user_id, "no-such-id") + self.assertFalse(result) diff --git a/tests/task_tool_test.py b/tests/task_tool_test.py new file mode 100644 index 0000000000000000000000000000000000000000..bda9608a79214fc09c8c5952d028723d77a4a9ef --- /dev/null +++ b/tests/task_tool_test.py @@ -0,0 +1,960 @@ +# -*- coding: utf-8 -*- +"""Unit tests for task tools.""" +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import AnyString + +from agentscope.state import AgentState +from agentscope.tool import TaskCreate, TaskGet, TaskList, TaskUpdate + + +class TestTaskCreate(IsolatedAsyncioTestCase): + """Test cases for TaskCreate tool.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.agent_state = AgentState() + self.task_create = TaskCreate() + + async def test_create_single_task(self) -> None: + """Test creating a single task.""" + result = await self.task_create( + subject="Test Task 1", + description="This is a test task", + _agent_state=self.agent_state, + ) + + # Check result using model_dump + task_id = self.agent_state.tasks_context.tasks[0].id + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": f"Task (id={task_id}) created successfully: " + f"Test Task 1", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) + + # Check task was added to agent state using model_dump + tasks_dump = [ + task.model_dump() for task in self.agent_state.tasks_context.tasks + ] + expected = [ + { + "subject": "Test Task 1", + "description": "This is a test task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ] + self.assertEqual(tasks_dump, expected) + + async def test_create_multiple_tasks(self) -> None: + """Test creating multiple tasks.""" + # Create first task + result1 = await self.task_create( + subject="Task 1", + description="First task", + _agent_state=self.agent_state, + ) + task1_id = self.agent_state.tasks_context.tasks[0].id + expected_result1 = { + "content": [ + { + "text": f"Task (id={task1_id}) created successfully: " + f"Task 1", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result1.model_dump(mode="json"), expected_result1) + + # Create second task + result2 = await self.task_create( + subject="Task 2", + description="Second task", + _agent_state=self.agent_state, + ) + task2_id = self.agent_state.tasks_context.tasks[1].id + expected_result2 = { + "content": [ + { + "text": f"Task (id={task2_id}) created successfully: " + f"Task 2", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result2.model_dump(mode="json"), expected_result2) + + # Create third task with metadata + result3 = await self.task_create( + subject="Task 3", + description="Third task", + metadata={"priority": "high"}, + _agent_state=self.agent_state, + ) + task3_id = self.agent_state.tasks_context.tasks[2].id + expected_result3 = { + "content": [ + { + "text": f"Task (id={task3_id}) created successfully: " + f"Task 3", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result3.model_dump(mode="json"), expected_result3) + + # Check all tasks using model_dump + tasks_dump = [ + task.model_dump() for task in self.agent_state.tasks_context.tasks + ] + expected = [ + { + "subject": "Task 1", + "description": "First task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": AnyString(), + "owner": None, + "blocks": [], + "blocked_by": [], + }, + { + "subject": "Task 2", + "description": "Second task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": AnyString(), + "owner": None, + "blocks": [], + "blocked_by": [], + }, + { + "subject": "Task 3", + "description": "Third task", + "metadata": {"priority": "high"}, + "created_at": AnyString(), + "state": "pending", + "id": AnyString(), + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ] + self.assertEqual(tasks_dump, expected) + + async def test_create_task_with_metadata(self) -> None: + """Test creating a task with metadata.""" + metadata = {"priority": "high", "tags": ["urgent", "bug"]} + result = await self.task_create( + subject="Bug Fix", + description="Fix critical bug", + metadata=metadata, + _agent_state=self.agent_state, + ) + + # Check result using model_dump + task_id = self.agent_state.tasks_context.tasks[0].id + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": f"Task (id={task_id}) created successfully: " + f"Bug Fix", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) + + # Check task using model_dump + tasks_dump = [ + task.model_dump() for task in self.agent_state.tasks_context.tasks + ] + expected = [ + { + "subject": "Bug Fix", + "description": "Fix critical bug", + "metadata": {"priority": "high", "tags": ["urgent", "bug"]}, + "created_at": AnyString(), + "state": "pending", + "id": AnyString(), + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ] + self.assertEqual(tasks_dump, expected) + + +class TestTaskList(IsolatedAsyncioTestCase): + """Test cases for TaskList tool.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.agent_state = AgentState() + self.task_list = TaskList() + self.task_create = TaskCreate() + + async def test_list_no_tasks(self) -> None: + """Test listing when there are no tasks.""" + result = await self.task_list(_agent_state=self.agent_state) + + self.assertEqual(len(result.content), 1) + self.assertEqual(result.content[0].text, "No tasks available.") + + async def test_list_with_tasks(self) -> None: + """Test listing when there are tasks.""" + # Create multiple tasks + await self.task_create( + subject="Task 1", + description="First task", + _agent_state=self.agent_state, + ) + await self.task_create( + subject="Task 2", + description="Second task", + _agent_state=self.agent_state, + ) + await self.task_create( + subject="Task 3", + description="Third task", + _agent_state=self.agent_state, + ) + + task1_id = self.agent_state.tasks_context.tasks[0].id + task2_id = self.agent_state.tasks_context.tasks[1].id + task3_id = self.agent_state.tasks_context.tasks[2].id + + # List tasks + result = await self.task_list(_agent_state=self.agent_state) + + # Check result using model_dump + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": f"{task1_id} [pending] Task 1\n" + f"{task2_id} [pending] Task 2\n" + f"{task3_id} [pending] Task 3", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) + + +class TestTaskGet(IsolatedAsyncioTestCase): + """Test cases for TaskGet tool.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.agent_state = AgentState() + self.task_get = TaskGet() + self.task_create = TaskCreate() + + async def test_get_existing_task(self) -> None: + """Test getting an existing task.""" + # Create a task + await self.task_create( + subject="Test Task", + description="This is a test task with details", + metadata={"priority": "high"}, + _agent_state=self.agent_state, + ) + + task_id = self.agent_state.tasks_context.tasks[0].id + + # Get the task + result = await self.task_get( + task_id=task_id, + _agent_state=self.agent_state, + ) + + # Check result using model_dump + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": f"Task (id={task_id}): Test Task\n" + f"Status: pending\n" + f"Description: This is a test task with details\n" + f"Metadata: {{'priority': 'high'}}", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) + + async def test_get_nonexistent_task(self) -> None: + """Test getting a task that doesn't exist.""" + result = await self.task_get( + task_id="nonexistent-id", + _agent_state=self.agent_state, + ) + + # Check result using model_dump + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": "Task not found", + "type": "text", + "id": AnyString(), + }, + ], + "state": "error", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) + + +class TestTaskUpdate(IsolatedAsyncioTestCase): + """Test cases for TaskUpdate tool.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.agent_state = AgentState() + self.task_update = TaskUpdate() + self.task_create = TaskCreate() + + async def test_update_subject(self) -> None: + """Test updating task subject.""" + # Create a task + await self.task_create( + subject="Original Subject", + description="Test description", + _agent_state=self.agent_state, + ) + task_id = self.agent_state.tasks_context.tasks[0].id + + # Update subject + result = await self.task_update( + task_id=task_id, + subject="Updated Subject", + _agent_state=self.agent_state, + ) + + # Check result using model_dump + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": f"Update task (id={task_id}) subject.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) + + # Check task using model_dump + tasks_dump = [ + task.model_dump() for task in self.agent_state.tasks_context.tasks + ] + expected = [ + { + "subject": "Updated Subject", + "description": "Test description", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ] + self.assertEqual(tasks_dump, expected) + + async def test_update_description(self) -> None: + """Test updating task description.""" + # Create a task + await self.task_create( + subject="Test Task", + description="Original description", + _agent_state=self.agent_state, + ) + task_id = self.agent_state.tasks_context.tasks[0].id + + # Update description + result = await self.task_update( + task_id=task_id, + description="Updated description", + _agent_state=self.agent_state, + ) + + # Check result using model_dump + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": f"Update task (id={task_id}) description.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) + + # Check task using model_dump + tasks_dump = [ + task.model_dump() for task in self.agent_state.tasks_context.tasks + ] + expected = [ + { + "subject": "Test Task", + "description": "Updated description", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ] + self.assertEqual(tasks_dump, expected) + + async def test_update_status(self) -> None: + """Test updating task status.""" + # Create a task + await self.task_create( + subject="Test Task", + description="Test description", + _agent_state=self.agent_state, + ) + task_id = self.agent_state.tasks_context.tasks[0].id + + # Update status to in_progress + result = await self.task_update( + task_id=task_id, + status="in_progress", + _agent_state=self.agent_state, + ) + + # Check result using model_dump + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": f"Update task (id={task_id}) status.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) + + # Check task using model_dump + tasks_dump = [ + task.model_dump() for task in self.agent_state.tasks_context.tasks + ] + expected = [ + { + "subject": "Test Task", + "description": "Test description", + "metadata": {}, + "created_at": AnyString(), + "state": "in_progress", + "id": task_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ] + self.assertEqual(tasks_dump, expected) + + # Update status to completed + result = await self.task_update( + task_id=task_id, + status="completed", + _agent_state=self.agent_state, + ) + + # Check result using model_dump + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": f"Update task (id={task_id}) status.\n\n" + f"Task completed. " + f"Call TaskList now to find your next available " + f"task or see if your work unblocked others.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) + + # Check task using model_dump + tasks_dump = [ + task.model_dump() for task in self.agent_state.tasks_context.tasks + ] + expected[0]["state"] = "completed" + self.assertEqual(tasks_dump, expected) + + # Test deleted status - create a new task first + await self.task_create( + subject="Task to Delete", + description="This task will be deleted", + _agent_state=self.agent_state, + ) + task_to_delete_id = self.agent_state.tasks_context.tasks[1].id + + # Update status to deleted + result = await self.task_update( + task_id=task_to_delete_id, + status="deleted", + _agent_state=self.agent_state, + ) + + # Check result using model_dump + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": f"Task (id={task_to_delete_id}) has been deleted.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) + + # Check that only the first task remains + self.assertEqual(len(self.agent_state.tasks_context.tasks), 1) + self.assertEqual(self.agent_state.tasks_context.tasks[0].id, task_id) + + async def test_update_owner(self) -> None: + """Test updating task owner.""" + # Create a task + await self.task_create( + subject="Test Task", + description="Test description", + _agent_state=self.agent_state, + ) + task_id = self.agent_state.tasks_context.tasks[0].id + + # Update owner + result = await self.task_update( + task_id=task_id, + owner="agent-1", + _agent_state=self.agent_state, + ) + + # Check result using model_dump + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": f"Update task (id={task_id}) owner.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) + + # Check task using model_dump + tasks_dump = [ + task.model_dump() for task in self.agent_state.tasks_context.tasks + ] + expected = [ + { + "subject": "Test Task", + "description": "Test description", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_id, + "owner": "agent-1", + "blocks": [], + "blocked_by": [], + }, + ] + self.assertEqual(tasks_dump, expected) + + async def test_update_metadata(self) -> None: + """Test updating task metadata.""" + # Create a task with initial metadata + await self.task_create( + subject="Test Task", + description="Test description", + metadata={"priority": "low", "tags": ["test"]}, + _agent_state=self.agent_state, + ) + task_id = self.agent_state.tasks_context.tasks[0].id + + # Update metadata + result = await self.task_update( + task_id=task_id, + metadata={"priority": "high", "new_field": "value"}, + _agent_state=self.agent_state, + ) + + # Check result using model_dump + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": f"Update task (id={task_id}) metadata.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) + + # Check task using model_dump + tasks_dump = [ + task.model_dump() for task in self.agent_state.tasks_context.tasks + ] + expected = [ + { + "subject": "Test Task", + "description": "Test description", + "metadata": { + "priority": "high", + "tags": ["test"], + "new_field": "value", + }, + "created_at": AnyString(), + "state": "pending", + "id": task_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ] + self.assertEqual(tasks_dump, expected) + + async def test_update_add_blocks(self) -> None: + """Test updating task blocks relationship.""" + # Create two tasks + await self.task_create( + subject="Task 1", + description="First task", + _agent_state=self.agent_state, + ) + await self.task_create( + subject="Task 2", + description="Second task", + _agent_state=self.agent_state, + ) + + task1_id = self.agent_state.tasks_context.tasks[0].id + task2_id = self.agent_state.tasks_context.tasks[1].id + + # Update task1 to block task2 + result = await self.task_update( + task_id=task1_id, + add_blocks=[task2_id], + _agent_state=self.agent_state, + ) + + # Check result using model_dump + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": f"Update task (id={task1_id}) add_blocks.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) + + # Check tasks using model_dump + tasks_dump = [ + task.model_dump() for task in self.agent_state.tasks_context.tasks + ] + expected = [ + { + "subject": "Task 1", + "description": "First task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task1_id, + "owner": None, + "blocks": [task2_id], + "blocked_by": [], + }, + { + "subject": "Task 2", + "description": "Second task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task2_id, + "owner": None, + "blocks": [], + "blocked_by": [task1_id], + }, + ] + self.assertEqual(tasks_dump, expected) + + async def test_update_add_blocked_by(self) -> None: + """Test updating task blockedBy relationship.""" + # Create two tasks + await self.task_create( + subject="Task 1", + description="First task", + _agent_state=self.agent_state, + ) + await self.task_create( + subject="Task 2", + description="Second task", + _agent_state=self.agent_state, + ) + + task1_id = self.agent_state.tasks_context.tasks[0].id + task2_id = self.agent_state.tasks_context.tasks[1].id + + # Update task2 to be blocked by task1 + result = await self.task_update( + task_id=task2_id, + add_blocked_by=[task1_id], + _agent_state=self.agent_state, + ) + + # Check result using model_dump + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": f"Update task (id={task2_id}) add_blocked_by.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) + + # Check tasks using model_dump + tasks_dump = [ + task.model_dump() for task in self.agent_state.tasks_context.tasks + ] + expected = [ + { + "subject": "Task 1", + "description": "First task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task1_id, + "owner": None, + "blocks": [task2_id], + "blocked_by": [], + }, + { + "subject": "Task 2", + "description": "Second task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task2_id, + "owner": None, + "blocks": [], + "blocked_by": [task1_id], + }, + ] + self.assertEqual(tasks_dump, expected) + + async def test_update_delete_task(self) -> None: + """Test deleting a task and removing it from blocks/blocked_by.""" + # Create three tasks with dependencies + await self.task_create( + subject="Task 1", + description="First task", + _agent_state=self.agent_state, + ) + await self.task_create( + subject="Task 2", + description="Second task", + _agent_state=self.agent_state, + ) + await self.task_create( + subject="Task 3", + description="Third task", + _agent_state=self.agent_state, + ) + + task1_id = self.agent_state.tasks_context.tasks[0].id + task2_id = self.agent_state.tasks_context.tasks[1].id + task3_id = self.agent_state.tasks_context.tasks[2].id + + # Set up dependencies: Task 1 blocks Task 2, Task 2 blocks Task 3 + await self.task_update( + task_id=task1_id, + add_blocks=[task2_id], + _agent_state=self.agent_state, + ) + await self.task_update( + task_id=task2_id, + add_blocks=[task3_id], + _agent_state=self.agent_state, + ) + + # Verify dependencies are set up correctly + tasks_before = [ + task.model_dump() for task in self.agent_state.tasks_context.tasks + ] + self.assertEqual(tasks_before[0]["blocks"], [task2_id]) + self.assertEqual(tasks_before[1]["blocked_by"], [task1_id]) + self.assertEqual(tasks_before[1]["blocks"], [task3_id]) + self.assertEqual(tasks_before[2]["blocked_by"], [task2_id]) + + # Delete task 2 (which is in the middle of the dependency chain) + result = await self.task_update( + task_id=task2_id, + status="deleted", + _agent_state=self.agent_state, + ) + + # Check result using model_dump + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": f"Task (id={task2_id}) has been deleted.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) + + # Check task was removed and task2_id was removed from + # blocks/blocked_by + self.assertEqual(len(self.agent_state.tasks_context.tasks), 2) + + tasks_after = [ + task.model_dump() for task in self.agent_state.tasks_context.tasks + ] + # Task 1 should no longer have task2_id in blocks + self.assertEqual(tasks_after[0]["id"], task1_id) + self.assertEqual(tasks_after[0]["blocks"], []) + + # Task 3 should no longer have task2_id in blocked_by + self.assertEqual(tasks_after[1]["id"], task3_id) + self.assertEqual(tasks_after[1]["blocked_by"], []) + + async def test_update_nonexistent_task(self) -> None: + """Test updating a task that doesn't exist.""" + result = await self.task_update( + task_id="nonexistent-id", + subject="New Subject", + description=None, + _agent_state=self.agent_state, + ) + + # Check result using model_dump + result_dump = result.model_dump(mode="json") + expected_result = { + "content": [ + { + "text": "TaskNotFoundError: " + "The task (id=nonexistent-id) does not exist.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "error", + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + self.assertDictEqual(result_dump, expected_result) diff --git a/tests/test_template.py b/tests/test_template.py new file mode 100644 index 0000000000000000000000000000000000000000..26ff54ae8c3baf99c5f207835181e08b8287ed7a --- /dev/null +++ b/tests/test_template.py @@ -0,0 +1,16 @@ +# -*- coding: utf-8 -*- +"""A template test case.""" +from unittest.async_case import IsolatedAsyncioTestCase + + +class TemplateTest(IsolatedAsyncioTestCase): + """The template test case.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + + async def test_template(self) -> None: + """The template test.""" + + async def asyncTearDown(self) -> None: + """The async teardown method.""" diff --git a/tests/tool_middleware_test.py b/tests/tool_middleware_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ba3d4b56a927594e5364bf0c3ed9e7a8bb3e3e05 --- /dev/null +++ b/tests/tool_middleware_test.py @@ -0,0 +1,344 @@ +# -*- coding: utf-8 -*- +# pylint: disable=unused-argument +"""Test cases for the tool-level onion middleware mechanism.""" +from typing import Any, AsyncGenerator +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import AnyString + +from agentscope.tool import ( + ToolBase, + ToolMiddlewareBase, + ToolChunk, +) +from agentscope.message import TextBlock, ToolResultState +from agentscope.permission import ( + PermissionDecision, + PermissionBehavior, +) + + +def _expected_chunk(text: str) -> dict: + """Build the full expected ``ToolChunk.model_dump()`` dict for the given + text, with random ``id`` fields matched by :class:`AnyString`.""" + return { + "content": [{"type": "text", "text": text, "id": AnyString()}], + "state": ToolResultState.RUNNING, + "is_last": True, + "metadata": {}, + "id": AnyString(), + } + + +class _NonStreamingTool(ToolBase): + """A tool whose ``call`` returns a single ToolChunk (coroutine).""" + + name: str = "non_streaming_tool" + description: str = "A tool that returns a single chunk." + input_schema: dict = {"type": "object", "properties": {}} + is_concurrency_safe: bool = True + is_read_only: bool = True + is_mcp: bool = False + + async def check_permissions( + self, + *args: Any, + **kwargs: Any, + ) -> PermissionDecision: + """Check permissions for the tool.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message="allowed", + ) + + async def call(self, **kwargs: Any) -> ToolChunk: + """Run the tool, echoing back the received kwargs.""" + return ToolChunk( + content=[TextBlock(text=f"call() kwargs={kwargs}")], + ) + + +class _StreamingTool(ToolBase): + """A tool whose ``call`` is an async generator function.""" + + name: str = "streaming_tool" + description: str = "A tool that yields several chunks." + input_schema: dict = {"type": "object", "properties": {}} + is_concurrency_safe: bool = True + is_read_only: bool = True + is_mcp: bool = False + + async def check_permissions( + self, + *args: Any, + **kwargs: Any, + ) -> PermissionDecision: + """Check permissions for the tool.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message="allowed", + ) + + async def call( + self, + n: int = 2, + **kwargs: Any, + ) -> AsyncGenerator[ToolChunk, None]: + """Yield ``n`` chunks.""" + for i in range(n): + yield ToolChunk(content=[TextBlock(text=f"chunk-{i}")]) + + +class _NoSuperInitTool(ToolBase): + """A tool that overrides ``__init__`` without calling ``super().__init__``. + + Used to verify the no-middleware path still works (via the ``getattr`` + fallback for ``_middlewares``). + """ + + name: str = "no_super_init_tool" + description: str = "A tool that skips super().__init__()." + input_schema: dict = {"type": "object", "properties": {}} + is_concurrency_safe: bool = True + is_read_only: bool = True + is_mcp: bool = False + + def __init__(self) -> None: # pylint: disable=super-init-not-called + """Intentionally does not call ``super().__init__()``.""" + self.marker = "initialized" + + async def check_permissions( + self, + *args: Any, + **kwargs: Any, + ) -> PermissionDecision: + """Check permissions for the tool.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message="allowed", + ) + + async def call(self, **kwargs: Any) -> ToolChunk: + """Run the tool.""" + return ToolChunk(content=[TextBlock(text="no super init")]) + + +def _make_recording_middleware( + label: str, + execution_order: list[str], +) -> ToolMiddlewareBase: + """Build a middleware that records pre/post markers in ``execution_order`` + and transparently forwards the chunks.""" + + class _Middleware(ToolMiddlewareBase): + async def on_tool_call( + self, + tool: Any, + input_kwargs: dict, + next_handler: Any, + ) -> AsyncGenerator: + execution_order.append(f"{label}-pre") + async for chunk in next_handler(**input_kwargs): + yield chunk + execution_order.append(f"{label}-post") + + return _Middleware() + + +class ToolMiddlewareTest(IsolatedAsyncioTestCase): + """Tests for the tool-level onion middleware mechanism.""" + + async def _drain(self, result: Any) -> list[dict]: + """Collect all chunks from a middleware-wrapped ``__call__`` result, + returning their full ``model_dump()`` dicts for whole-object asserts. + """ + chunks = [] + async for chunk in result: + chunks.append(chunk.model_dump()) + return chunks + + async def test_no_middleware_non_streaming(self) -> None: + """A non-streaming tool with no middleware returns a single chunk.""" + tool = _NonStreamingTool() + result = await tool() + self.assertIsInstance(result, ToolChunk) + self.assertEqual( + result.model_dump(), + _expected_chunk("call() kwargs={}"), + ) + + async def test_no_middleware_streaming(self) -> None: + """A streaming tool with no middleware yields its chunks directly.""" + tool = _StreamingTool() + result = await tool(n=3) + chunks = await self._drain(result) + self.assertEqual( + chunks, + [ + _expected_chunk("chunk-0"), + _expected_chunk("chunk-1"), + _expected_chunk("chunk-2"), + ], + ) + + async def test_middleware_wraps_call_in_onion_order(self) -> None: + """Two middlewares wrap call() in the correct onion order. + + Execution order must be: + outer-pre -> inner-pre -> call() -> inner-post -> outer-post. + """ + execution_order: list[str] = [] + tool = _NonStreamingTool( + middlewares=[ + _make_recording_middleware("outer", execution_order), + _make_recording_middleware("inner", execution_order), + ], + ) + + chunks = await self._drain(await tool()) + + self.assertEqual(chunks, [_expected_chunk("call() kwargs={}")]) + self.assertEqual( + execution_order, + ["outer-pre", "inner-pre", "inner-post", "outer-post"], + ) + + async def test_single_middleware(self) -> None: + """A single middleware fires pre and post around call().""" + execution_order: list[str] = [] + tool = _NonStreamingTool( + middlewares=[_make_recording_middleware("only", execution_order)], + ) + + chunks = await self._drain(await tool()) + + self.assertEqual(chunks, [_expected_chunk("call() kwargs={}")]) + self.assertEqual(execution_order, ["only-pre", "only-post"]) + + async def test_middleware_with_streaming_tool(self) -> None: + """Middleware transparently forwards each chunk of a streaming tool.""" + execution_order: list[str] = [] + tool = _StreamingTool( + middlewares=[_make_recording_middleware("mw", execution_order)], + ) + + chunks = await self._drain(await tool(n=3)) + + self.assertEqual( + chunks, + [ + _expected_chunk("chunk-0"), + _expected_chunk("chunk-1"), + _expected_chunk("chunk-2"), + ], + ) + self.assertEqual(execution_order, ["mw-pre", "mw-post"]) + + async def test_middleware_can_rewrite_input_kwargs(self) -> None: + """A middleware can mutate input_kwargs so the tool sees new args.""" + + class _RewriteMiddleware(ToolMiddlewareBase): + async def on_tool_call( + self, + tool: Any, + input_kwargs: dict, + next_handler: Any, + ) -> AsyncGenerator: + input_kwargs["injected"] = "value" + async for chunk in next_handler(**input_kwargs): + yield chunk + + tool = _NonStreamingTool(middlewares=[_RewriteMiddleware()]) + chunks = await self._drain(await tool(original="x")) + + self.assertEqual( + chunks, + [ + _expected_chunk( + "call() kwargs={'original': 'x', 'injected': 'value'}", + ), + ], + ) + + async def test_middleware_can_transform_chunks(self) -> None: + """A middleware can transform the chunks yielded by the tool.""" + + class _UppercaseMiddleware(ToolMiddlewareBase): + async def on_tool_call( + self, + tool: Any, + input_kwargs: dict, + next_handler: Any, + ) -> AsyncGenerator: + async for chunk in next_handler(**input_kwargs): + yield ToolChunk( + content=[ + TextBlock(text=chunk.content[0].text.upper()), + ], + ) + + tool = _StreamingTool(middlewares=[_UppercaseMiddleware()]) + chunks = await self._drain(await tool(n=2)) + + self.assertEqual( + chunks, + [_expected_chunk("CHUNK-0"), _expected_chunk("CHUNK-1")], + ) + + async def test_middleware_can_short_circuit(self) -> None: + """A middleware that never calls next_handler skips the tool.""" + called = {"value": False} + + class _ShortCircuitTool(_NonStreamingTool): + async def call(self, **kwargs: Any) -> ToolChunk: + called["value"] = True + return await super().call(**kwargs) + + class _ShortCircuitMiddleware(ToolMiddlewareBase): + async def on_tool_call( + self, + tool: Any, + input_kwargs: dict, + next_handler: Any, + ) -> AsyncGenerator: + yield ToolChunk(content=[TextBlock(text="short-circuited")]) + + tool = _ShortCircuitTool(middlewares=[_ShortCircuitMiddleware()]) + chunks = await self._drain(await tool()) + + self.assertFalse(called["value"]) + self.assertEqual(chunks, [_expected_chunk("short-circuited")]) + + async def test_middleware_exception_propagates(self) -> None: + """An exception raised in a middleware propagates to the caller.""" + + class _BoomMiddleware(ToolMiddlewareBase): + async def on_tool_call( + self, + tool: Any, + input_kwargs: dict, + next_handler: Any, + ) -> AsyncGenerator: + raise ValueError("boom") + yield # pylint: disable=unreachable + + tool = _NonStreamingTool(middlewares=[_BoomMiddleware()]) + + with self.assertRaises(ValueError): + await self._drain(await tool()) + + async def test_positional_args_rejected(self) -> None: + """Calling a tool with positional args raises TypeError instead of + silently dropping them.""" + tool = _NonStreamingTool() + with self.assertRaises(TypeError): + await tool("positional") # type: ignore[call-arg] + + async def test_missing_super_init_no_middleware(self) -> None: + """A tool that skips super().__init__() still works on the + no-middleware path via the getattr fallback for _middlewares.""" + tool = _NoSuperInitTool() + self.assertFalse(hasattr(tool, "_middlewares")) + result = await tool() + self.assertIsInstance(result, ToolChunk) + self.assertEqual(result.model_dump(), _expected_chunk("no super init")) diff --git a/tests/tool_offload_middleware_test.py b/tests/tool_offload_middleware_test.py new file mode 100644 index 0000000000000000000000000000000000000000..23d02851ab0658cac57b589160fec9cf56add6fb --- /dev/null +++ b/tests/tool_offload_middleware_test.py @@ -0,0 +1,400 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for ToolOffloadMiddleware.""" +import asyncio +import json +from typing import Any +from unittest.async_case import IsolatedAsyncioTestCase +from unittest.mock import MagicMock + +from pydantic import BaseModel + + +from utils import AnyString, MockModel + +from agentscope.agent import Agent +from agentscope.app.message_bus import MessageBus, MessageBusKeys +from agentscope.app.middleware import ToolOffloadMiddleware +from agentscope.app._manager import BackgroundTaskManager +from agentscope.message import TextBlock, ToolCallBlock +from agentscope.permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, +) +from agentscope.tool import ToolBase, ToolChunk, Toolkit, ToolResponse + + +class _SlowToolParams(BaseModel): + """Parameters for the slow test tool.""" + + delay: float + + +class SlowTool(ToolBase): + """A tool that sleeps for ``delay`` seconds before returning.""" + + name: str = "slow_tool" + description: str = "A slow tool for testing background offload." + input_schema: dict = _SlowToolParams.model_json_schema() + is_concurrency_safe: bool = True + is_read_only: bool = True + is_state_injected: bool = False + is_external_tool: bool = False + is_mcp: bool = False + mcp_name: str | None = None + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Always allow. + + Args: + tool_input (`dict[str, Any]`): + The tool input parameters. + context (`PermissionContext`): + The permission context. + + Returns: + `PermissionDecision`: + Always ALLOW. + """ + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message="allowed", + ) + + async def __call__( # type: ignore[override] + self, + delay: float, + ) -> ToolChunk: + """Sleep for *delay* seconds then return a result. + + Args: + delay (`float`): + Seconds to sleep. + + Returns: + `ToolChunk`: + A chunk containing the result text. + """ + await asyncio.sleep(delay) + return ToolChunk( + content=[TextBlock(text=f"SlowTool finished after {delay}s")], + ) + + +class _FastToolParams(BaseModel): + """Parameters for the fast test tool.""" + + value: str + + +class FastTool(ToolBase): + """A tool that returns immediately.""" + + name: str = "fast_tool" + description: str = "A fast tool for testing normal execution." + input_schema: dict = _FastToolParams.model_json_schema() + is_concurrency_safe: bool = True + is_read_only: bool = True + is_state_injected: bool = False + is_external_tool: bool = False + is_mcp: bool = False + mcp_name: str | None = None + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Always allow. + + Args: + tool_input (`dict[str, Any]`): + The tool input parameters. + context (`PermissionContext`): + The permission context. + + Returns: + `PermissionDecision`: + Always ALLOW. + """ + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message="allowed", + ) + + async def __call__( # type: ignore[override] + self, + value: str, + ) -> ToolChunk: + """Return a chunk with *value*. + + Args: + value (`str`): + The value to echo. + + Returns: + `ToolChunk`: + A chunk containing the value. + """ + return ToolChunk( + content=[TextBlock(text=f"FastTool: {value}")], + ) + + +class ToolOffloadMiddlewareTest(IsolatedAsyncioTestCase): + """Test cases for the ToolOffloadMiddleware.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + self.mock_model = MockModel() + self.bg_manager = BackgroundTaskManager( + message_bus=MagicMock(spec=MessageBus), + ) + + # ------------------------------------------------------------------ + # Helper + # ------------------------------------------------------------------ + + def _make_agent( + self, + toolkit: Toolkit, + timeout_secs: float, + ) -> tuple[Agent, ToolOffloadMiddleware]: + """Create an agent with ToolOffloadMiddleware attached. + + Args: + toolkit (`Toolkit`): + The toolkit to attach to the agent. + timeout_secs (`float`): + The middleware timeout. + + Returns: + `tuple[Agent, ToolOffloadMiddleware]`: + The configured agent and the middleware instance. + """ + middleware = ToolOffloadMiddleware( + bg_manager=self.bg_manager, + message_bus=MagicMock(spec=MessageBus), + user_id="u", + agent_id="a", + timeout_secs=timeout_secs, + ) + agent = Agent( + name="test_agent", + system_prompt="test prompt", + model=self.mock_model, + toolkit=toolkit, + middlewares=[middleware], + ) + return agent, middleware + + # ------------------------------------------------------------------ + # Tests + # ------------------------------------------------------------------ + + async def test_fast_tool_completes_normally(self) -> None: + """A tool that finishes within the timeout yields its real result.""" + + toolkit = Toolkit(tools=[FastTool()]) + agent, _ = self._make_agent(toolkit, timeout_secs=5.0) + + tool_call = ToolCallBlock( + id="call_fast", + name="fast_tool", + input=json.dumps({"value": "hello"}), + ) + + results: list = [] + # pylint: disable=protected-access + async for item in agent._acting(tool_call): + results.append(item) + + # Should yield real ToolResponse (not synthetic) + responses = [r for r in results if isinstance(r, ToolResponse)] + self.assertEqual(len(responses), 1) + text = responses[0].content[0].text # type: ignore[union-attr] + self.assertIn("FastTool: hello", text) + # No background tasks registered + self.assertEqual(len(self.bg_manager.tasks), 0) + + async def test_slow_tool_offloaded_to_background(self) -> None: + """A tool that exceeds timeout returns a synthetic result.""" + + toolkit = Toolkit(tools=[SlowTool()]) + # Set a very short timeout so the 0.5s tool is always offloaded + agent, _ = self._make_agent(toolkit, timeout_secs=0.05) + + tool_call = ToolCallBlock( + id="call_slow", + name="slow_tool", + input=json.dumps({"delay": 0.5}), + ) + + results: list = [] + # pylint: disable=protected-access + async for item in agent._acting(tool_call): + results.append(item) + + # Should yield a synthetic ToolResponse immediately + responses = [r for r in results if isinstance(r, ToolResponse)] + self.assertEqual(len(responses), 1) + self.assertDictEqual( + responses[0].model_dump(), + { + "content": [ + { + "type": "text", + "text": AnyString(), + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "call_slow", + }, + ) + text = responses[0].content[0].text # type: ignore[union-attr] + self.assertIn("background", text) + self.assertIn("id=", text) + + # Background task should be registered + self.assertEqual(len(self.bg_manager.tasks), 1) + + async def test_background_task_result_injected_into_context( + self, + ) -> None: + """After the background tool finishes, the result is pushed to the + session inbox on the message bus as a serialised HintBlock.""" + + toolkit = Toolkit(tools=[SlowTool()]) + agent, middleware = self._make_agent(toolkit, timeout_secs=0.05) + + tool_call = ToolCallBlock( + id="call_bg", + name="slow_tool", + input=json.dumps({"delay": 0.2}), + ) + + # Trigger offload + # pylint: disable=protected-access + async for _ in agent._acting(tool_call): + pass + + # Wait long enough for the background tool (0.2s) to finish + await asyncio.sleep(0.4) + + # The completed result should now have been pushed to the message bus + # inbox as a model-dumped HintBlock. + mock_bus = middleware._message_bus + # The middleware uses queue_push(inbox_key, payload) rather + # than the deprecated inbox_push(session_id, payload). + inbox_calls = [ + c + for c in mock_bus.queue_push.call_args_list + if c.args[0] == MessageBusKeys.inbox(agent.state.session_id) + ] + self.assertEqual(len(inbox_calls), 1) + _, hint_dict = inbox_calls[0].args + session_id_called = agent.state.session_id + self.assertEqual(session_id_called, agent.state.session_id) + self.maxDiff = None + self.assertDictEqual( + hint_dict, + { + "type": "hint", + "id": AnyString(), + "source": '{"label": "tool_output", "sublabel": "slow_tool · ' + 'call_bg"}', + "hint": [ + { + "type": "text", + "text": AnyString(), + "id": AnyString(), + }, + ], + }, + ) + hint_text = hint_dict["hint"][0]["text"] + self.assertIn("SlowTool finished", hint_text) + self.assertIn("", hint_text) + + async def test_background_task_triggers_wakeup_on_completion( + self, + ) -> None: + """After a background tool finishes, a wakeup is enqueued on the + message bus so that an idle session can be restarted automatically.""" + + toolkit = Toolkit(tools=[SlowTool()]) + agent, middleware = self._make_agent(toolkit, timeout_secs=0.05) + + tool_call = ToolCallBlock( + id="call_wakeup", + name="slow_tool", + input=json.dumps({"delay": 0.2}), + ) + + # Trigger offload + # pylint: disable=protected-access + async for _ in agent._acting(tool_call): + pass + + # Wait long enough for the background tool (0.2s) to finish + await asyncio.sleep(0.4) + + # enqueue_wakeup must be called exactly once with the correct ids so + # WakeupDispatcher can re-invoke ChatService.run for this session. + mock_bus = middleware._message_bus + # enqueue_run_trigger is a standalone function that calls + # bus.queue_push + bus.publish under the hood. + wakeup_calls = [ + c + for c in mock_bus.queue_push.call_args_list + if c.args[0] == MessageBusKeys.wakeup_queue() + ] + self.assertEqual(len(wakeup_calls), 1) + payload = wakeup_calls[0].args[1] + self.assertEqual(payload["user_id"], "u") + self.assertEqual(payload["session_id"], agent.state.session_id) + self.assertEqual(payload["agent_id"], "a") + + async def test_tool_stop_cancels_background_task(self) -> None: + """ToolStop tool cancels the running background asyncio task.""" + + toolkit = Toolkit(tools=[SlowTool()]) + agent, _ = self._make_agent(toolkit, timeout_secs=0.05) + + tool_call = ToolCallBlock( + id="call_cancel", + name="slow_tool", + input=json.dumps({"delay": 10.0}), + ) + + # Offload the slow tool + # pylint: disable=protected-access + async for _ in agent._acting(tool_call): + pass + + self.assertEqual(len(self.bg_manager.tasks), 1) + task_id = next(iter(self.bg_manager.tasks)) + bg_task = self.bg_manager.tasks[task_id] + asyncio_task = bg_task.asyncio_task + + # Call ToolStop bound to the same session as the registered + # background task, so the local cancel path matches. + tool_stop_tools = await self.bg_manager.list_tools( + session_id=bg_task.session_id, + ) + tool_stop = tool_stop_tools[0] + result = await tool_stop(task_id=task_id) + text = result.content[0].text # type: ignore[union-attr] + self.assertIn("stopped successfully", text) + + # The asyncio task should be cancelling + self.assertTrue(asyncio_task.cancelled() or asyncio_task.cancelling()) + # Removed from manager + self.assertEqual(len(self.bg_manager.tasks), 0) diff --git a/tests/toolkit_skill_test.py b/tests/toolkit_skill_test.py new file mode 100644 index 0000000000000000000000000000000000000000..af67843054e0351b59abe42d574f709baf78caf1 --- /dev/null +++ b/tests/toolkit_skill_test.py @@ -0,0 +1,408 @@ +# -*- coding: utf-8 -*- +"""Test cases for Toolkit skill-related functionality.""" +import json +import os +import tempfile +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import AnyString + +from agentscope.skill import SkillLoaderBase, Skill +from agentscope.tool import Toolkit, ToolChunk, ToolResponse, ToolGroup +from agentscope.message import ToolCallBlock +from agentscope.state import AgentState + + +def _make_skill( + name: str, + description: str = "desc", + dir_: str = "/tmp", +) -> Skill: + """Helper function to create a Skill object for testing.""" + return Skill( + name=name, + description=description, + dir=dir_, + markdown="", + updated_at=0.0, + ) + + +class MockSkillLoader(SkillLoaderBase): + """A mock skill loader for testing.""" + + def __init__(self, skills: list[Skill]) -> None: + self._skills = skills + + async def list_skills(self) -> list[Skill]: + """Return a list of all skills.""" + return self._skills + + +class ToolkitSkillTest(IsolatedAsyncioTestCase): + """Test cases for Toolkit skill functionality.""" + + async def test_init_with_various_types(self) -> None: + """Test Toolkit initialization with str path, SkillLoaderBase, and + direct Skill objects, then assert get_skill_instructions output.""" + with tempfile.TemporaryDirectory() as skill_dir: + # Create a minimal SKILL.md so LocalSkillLoader can load it + skill_md_path = os.path.join(skill_dir, "SKILL.md") + with open(skill_md_path, "w", encoding="utf-8") as f: + f.write( + "---\n" + "name: path_skill\n" + "description: A skill loaded from a path\n" + "---\n\n" + "# Path Skill\n", + ) + + loader_skill = _make_skill( + "loader_skill", + description="A skill from loader", + dir_="/loader/dir", + ) + direct_skill = _make_skill( + "direct_skill", + description="A directly provided skill", + dir_="/direct/dir", + ) + + toolkit = Toolkit( + skills_or_loaders=[ + skill_dir, # str -> LocalSkillLoader + MockSkillLoader([loader_skill]), # SkillLoaderBase + direct_skill, # Skill directly + ], + ) + + result = await toolkit.get_skill_instructions() + + self.assertEqual( + result, + # pylint: disable=line-too-long + f""" +Skills are a collection of instructions, scripts, and resources to extend your capabilities. + +**IMPORTANT**: Skills are NOT tools, and you cannot call a skill directly. To use a skill, you MUST use the `Skill` tool to read the skill's full instructions, and then follow those instructions to use the tools and resources provided by the skill. + +# Available Skills: + +path_skill +A skill loaded from a path +{skill_dir} + + +loader_skill +A skill from loader +/loader/dir + + +direct_skill +A directly provided skill +/direct/dir + +""", # noqa: E501 + ) + + async def test_get_skill_instructions_no_skills(self) -> None: + """Test that get_skill_instructions returns None when no skills + registered.""" + toolkit = Toolkit() + result = await toolkit.get_skill_instructions() + self.assertIsNone(result) + + async def test_get_skill_instructions_multiple_loaders(self) -> None: + """Test that get_skill_instructions aggregates skills from multiple + loaders.""" + loader1 = MockSkillLoader([_make_skill("skill_x")]) + loader2 = MockSkillLoader([_make_skill("skill_y")]) + toolkit = Toolkit(skills_or_loaders=[loader1, loader2]) + + result = await toolkit.get_skill_instructions() + + self.assertEqual( + result, + # pylint: disable=line-too-long + """ +Skills are a collection of instructions, scripts, and resources to extend your capabilities. + +**IMPORTANT**: Skills are NOT tools, and you cannot call a skill directly. To use a skill, you MUST use the `Skill` tool to read the skill's full instructions, and then follow those instructions to use the tools and resources provided by the skill. + +# Available Skills: + +skill_x +desc +/tmp + + +skill_y +desc +/tmp + +""", # noqa: E501 + ) + + async def test_get_skill_instructions_empty_loader(self) -> None: + """Test that an empty loader contributes no skills.""" + loader = MockSkillLoader([]) + toolkit = Toolkit(skills_or_loaders=[loader]) + + result = await toolkit.get_skill_instructions() + self.assertIsNone(result) + + async def test_get_skill_instructions_includes_tool_group_skills( + self, + ) -> None: + """Test that skill instructions include skills from custom groups.""" + toolkit = Toolkit( + tool_groups=[ + ToolGroup( + name="repair", + description="Repair tools", + skills_or_loaders=[ + MockSkillLoader([_make_skill("repair_skill")]), + ], + ), + ], + ) + + result = await toolkit.get_skill_instructions() + + self.assertIsNotNone(result) + assert result is not None + self.assertIn("repair_skill", result) + + async def test_get_skill_instructions_filters_inactive_group_skills( + self, + ) -> None: + """Test that inactive group skills are hidden from the prompt.""" + toolkit = Toolkit( + skills_or_loaders=[MockSkillLoader([_make_skill("basic_skill")])], + tool_groups=[ + ToolGroup( + name="repair", + description="Repair tools", + skills_or_loaders=[ + MockSkillLoader([_make_skill("repair_skill")]), + ], + ), + ], + ) + + inactive_result = await toolkit.get_skill_instructions( + activated_groups=[], + ) + active_result = await toolkit.get_skill_instructions( + activated_groups=["repair"], + ) + + self.assertIsNotNone(inactive_result) + self.assertIsNotNone(active_result) + assert inactive_result is not None + assert active_result is not None + self.assertIn("basic_skill", inactive_result) + self.assertNotIn("repair_skill", inactive_result) + self.assertIn("basic_skill", active_result) + self.assertIn("repair_skill", active_result) + + +class ToolkitSkillViewerTest(IsolatedAsyncioTestCase): + """Test cases for Toolkit SkillViewer functionality.""" + + async def test_register_skill_and_get_function_schemas(self) -> None: + """Test that registering skills makes SkillViewer available in + function schemas.""" + skill = _make_skill("test_skill", description="A test skill") + loader = MockSkillLoader([skill]) + toolkit = Toolkit(skills_or_loaders=[loader]) + + schemas = await toolkit.get_tool_schemas() + + self.assertListEqual( + schemas, + [ + { + "type": "function", + "function": { + "name": "Skill", + "description": ( + "Retrieve a skill within the conversation. " + "When users asks you to perform tasks, check if " + "any of the available skills match. " + "Skills provide specialized capabilities and " + "domain knowledge." + ), + "parameters": { + "type": "object", + "properties": { + "skill": { + "type": "string", + "description": "The exact name of the " + "skill to view. ", + }, + }, + "required": ["skill"], + }, + }, + }, + ], + ) + + async def test_call_skill_viewer_success(self) -> None: + """Test calling SkillViewer with an existing skill.""" + skill = _make_skill( + "my_skill", + description="My test skill", + dir_="/test/dir", + ) + skill.markdown = "# My Skill\nThis is the skill content." + loader = MockSkillLoader([skill]) + toolkit = Toolkit(skills_or_loaders=[loader]) + + tool_call = ToolCallBlock( + id="test_call_1", + name="Skill", + input=json.dumps({"skill": "my_skill"}), + ) + state = AgentState() + + chunks = [] + response = None + async for result in toolkit.call_tool(tool_call, state): + if isinstance(result, ToolChunk): + chunks.append(result) + elif isinstance(result, ToolResponse): + response = result + + self.assertEqual(len(chunks), 1) + self.assertDictEqual( + chunks[0].model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "# My Skill\nThis is the skill content.", + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + self.assertIsNotNone(response) + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "# My Skill\nThis is the skill content.", + }, + ], + "state": "success", + "metadata": {}, + "id": "test_call_1", + }, + ) + + async def test_call_skill_viewer_not_found(self) -> None: + """Test calling SkillViewer with a non-existent skill.""" + skill = _make_skill("existing_skill") + loader = MockSkillLoader([skill]) + toolkit = Toolkit(skills_or_loaders=[loader]) + + tool_call = ToolCallBlock( + id="test_call_2", + name="Skill", + input=json.dumps({"skill": "non_existent_skill"}), + ) + state = AgentState() + + chunks = [] + response = None + async for result in toolkit.call_tool(tool_call, state): + if isinstance(result, ToolChunk): + chunks.append(result) + elif isinstance(result, ToolResponse): + response = result + + self.assertEqual(len(chunks), 1) + self.assertDictEqual( + chunks[0].model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "SkillNotFoundError: " + "Skill 'non_existent_skill' not found.", + }, + ], + "state": "error", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + self.assertIsNotNone(response) + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "SkillNotFoundError: " + "Skill 'non_existent_skill' not found.", + }, + ], + "state": "error", + "metadata": {}, + "id": "test_call_2", + }, + ) + + async def test_skill_viewer_uses_active_tool_group_skills(self) -> None: + """Test that activated custom-group skills expose SkillViewer.""" + skill = _make_skill("repair_skill") + skill.markdown = "# Repair Skill\nUse this for repair tasks." + toolkit = Toolkit( + tool_groups=[ + ToolGroup( + name="repair", + description="Repair tools", + skills_or_loaders=[MockSkillLoader([skill])], + ), + ], + ) + + schema_names = [ + schema["function"]["name"] + for schema in await toolkit.get_tool_schemas(groups=["repair"]) + ] + self.assertIn("Skill", schema_names) + + tool_call = ToolCallBlock( + id="test_call_group_skill", + name="Skill", + input=json.dumps({"skill": "repair_skill"}), + ) + state = AgentState() + state.tool_context.activated_groups.append("repair") + + chunks = [] + async for result in toolkit.call_tool(tool_call, state): + if isinstance(result, ToolChunk): + chunks.append(result) + + self.assertEqual(len(chunks), 1) + self.assertEqual( + chunks[0].content[0].text, + "# Repair Skill\nUse this for repair tasks.", + ) diff --git a/tests/toolkit_task_test.py b/tests/toolkit_task_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e47c9950e9915d441dada0a57da79f1133f94484 --- /dev/null +++ b/tests/toolkit_task_test.py @@ -0,0 +1,1187 @@ +# -*- coding: utf-8 -*- +"""Unit tests for task tools executed through toolkit.""" +import json +from typing import Any +from unittest.async_case import IsolatedAsyncioTestCase + +from utils import AnyString + +from agentscope.message import ToolCallBlock +from agentscope.state import AgentState +from agentscope.tool import ( + Toolkit, + ToolResponse, + TaskCreate, + TaskGet, + TaskList, + TaskUpdate, +) + + +class _ToolkitTaskTestBase(IsolatedAsyncioTestCase): + """Shared helpers for toolkit task tool tests.""" + + async def asyncSetUp(self) -> None: + """Set up shared test fixtures.""" + self.agent_state = AgentState() + self.toolkit = Toolkit( + tools=[ + TaskCreate(), + TaskList(), + TaskGet(), + TaskUpdate(), + ], + ) + + async def _call_tool( + self, + name: str, + tool_input: dict[str, Any], + tool_call_id: str, + ) -> ToolResponse: + """Call a task tool through toolkit and return the final response.""" + response = None + async for result in self.toolkit.call_tool( + ToolCallBlock( + id=tool_call_id, + name=name, + input=json.dumps(tool_input), + ), + self.agent_state, + ): + if isinstance(result, ToolResponse): + response = result + + self.assertIsNotNone(response) + return response + + def _dump_tasks(self) -> list[dict[str, Any]]: + """Dump all tasks from the agent state for assertions.""" + return [ + task.model_dump() for task in self.agent_state.tasks_context.tasks + ] + + +class TestToolkitTaskCreate(_ToolkitTaskTestBase): + """Test cases for TaskCreate through toolkit.""" + + async def test_create_single_task(self) -> None: + """Test creating a single task.""" + response = await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Test Task 1", + "description": "This is a test task", + }, + tool_call_id="task-create-single", + ) + + task_id = self.agent_state.tasks_context.tasks[0].id + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": f"Task (id={task_id}) created successfully: " + "Test Task 1", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-create-single", + }, + ) + + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Test Task 1", + "description": "This is a test task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ], + ) + + async def test_create_multiple_tasks(self) -> None: + """Test creating multiple tasks.""" + response_1 = await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Task 1", + "description": "First task", + }, + tool_call_id="task-create-1", + ) + task_1_id = self.agent_state.tasks_context.tasks[0].id + self.assertDictEqual( + response_1.model_dump(), + { + "content": [ + { + "text": f"Task (id={task_1_id}) created successfully: " + f"Task 1", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-create-1", + }, + ) + + response_2 = await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Task 2", + "description": "Second task", + }, + tool_call_id="task-create-2", + ) + task_2_id = self.agent_state.tasks_context.tasks[1].id + self.assertDictEqual( + response_2.model_dump(), + { + "content": [ + { + "text": f"Task (id={task_2_id}) created successfully: " + f"Task 2", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-create-2", + }, + ) + + response_3 = await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Task 3", + "description": "Third task", + "metadata": {"priority": "high"}, + }, + tool_call_id="task-create-3", + ) + task_3_id = self.agent_state.tasks_context.tasks[2].id + self.assertDictEqual( + response_3.model_dump(), + { + "content": [ + { + "text": f"Task (id={task_3_id}) created successfully: " + f"Task 3", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-create-3", + }, + ) + + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Task 1", + "description": "First task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_1_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + { + "subject": "Task 2", + "description": "Second task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_2_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + { + "subject": "Task 3", + "description": "Third task", + "metadata": {"priority": "high"}, + "created_at": AnyString(), + "state": "pending", + "id": task_3_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ], + ) + + async def test_create_task_with_metadata(self) -> None: + """Test creating a task with metadata.""" + response = await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Bug Fix", + "description": "Fix critical bug", + "metadata": { + "priority": "high", + "tags": ["urgent", "bug"], + }, + }, + tool_call_id="task-create-metadata", + ) + + task_id = self.agent_state.tasks_context.tasks[0].id + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": f"Task (id={task_id}) created successfully: " + f"Bug Fix", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-create-metadata", + }, + ) + + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Bug Fix", + "description": "Fix critical bug", + "metadata": { + "priority": "high", + "tags": ["urgent", "bug"], + }, + "created_at": AnyString(), + "state": "pending", + "id": task_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ], + ) + + +class TestToolkitTaskList(_ToolkitTaskTestBase): + """Test cases for TaskList through toolkit.""" + + async def test_list_no_tasks(self) -> None: + """Test listing when there are no tasks.""" + response = await self._call_tool( + name="TaskList", + tool_input={}, + tool_call_id="task-list-empty", + ) + + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": "No tasks available.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-list-empty", + }, + ) + self.assertEqual(self._dump_tasks(), []) + + async def test_list_with_tasks(self) -> None: + """Test listing when there are tasks.""" + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Task 1", + "description": "First task", + }, + tool_call_id="task-list-create-1", + ) + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Task 2", + "description": "Second task", + }, + tool_call_id="task-list-create-2", + ) + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Task 3", + "description": "Third task", + }, + tool_call_id="task-list-create-3", + ) + + task_1_id = self.agent_state.tasks_context.tasks[0].id + task_2_id = self.agent_state.tasks_context.tasks[1].id + task_3_id = self.agent_state.tasks_context.tasks[2].id + + response = await self._call_tool( + name="TaskList", + tool_input={}, + tool_call_id="task-list-with-tasks", + ) + + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": f"{task_1_id} [pending] Task 1\n" + f"{task_2_id} [pending] Task 2\n" + f"{task_3_id} [pending] Task 3", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-list-with-tasks", + }, + ) + + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Task 1", + "description": "First task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_1_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + { + "subject": "Task 2", + "description": "Second task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_2_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + { + "subject": "Task 3", + "description": "Third task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_3_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ], + ) + + +class TestToolkitTaskGet(_ToolkitTaskTestBase): + """Test cases for TaskGet through toolkit.""" + + async def test_get_existing_task(self) -> None: + """Test getting an existing task.""" + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Test Task", + "description": "This is a test task with details", + "metadata": {"priority": "high"}, + }, + tool_call_id="task-get-create", + ) + + task_id = self.agent_state.tasks_context.tasks[0].id + response = await self._call_tool( + name="TaskGet", + tool_input={"task_id": task_id}, + tool_call_id="task-get-existing", + ) + + self.assertDictEqual( + response.model_dump(mode="json"), + { + "content": [ + { + "text": f"Task (id={task_id}): Test Task\n" + "Status: pending\n" + "Description: This is a test task with details\n" + "Metadata: {'priority': 'high'}", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-get-existing", + }, + ) + + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Test Task", + "description": "This is a test task with details", + "metadata": {"priority": "high"}, + "created_at": AnyString(), + "state": "pending", + "id": task_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ], + ) + + async def test_get_nonexistent_task(self) -> None: + """Test getting a task that does not exist.""" + response = await self._call_tool( + name="TaskGet", + tool_input={"task_id": "nonexistent-id"}, + tool_call_id="task-get-missing", + ) + + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": "Task not found", + "type": "text", + "id": AnyString(), + }, + ], + "state": "error", + "metadata": {}, + "id": "task-get-missing", + }, + ) + self.assertEqual(self._dump_tasks(), []) + + +class TestToolkitTaskUpdate(_ToolkitTaskTestBase): + """Test cases for TaskUpdate through toolkit.""" + + async def test_update_subject(self) -> None: + """Test updating task subject.""" + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Original Subject", + "description": "Test description", + }, + tool_call_id="task-update-subject-create", + ) + task_id = self.agent_state.tasks_context.tasks[0].id + + response = await self._call_tool( + name="TaskUpdate", + tool_input={ + "task_id": task_id, + "subject": "Updated Subject", + }, + tool_call_id="task-update-subject", + ) + + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": f"Update task (id={task_id}) subject.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-update-subject", + }, + ) + + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Updated Subject", + "description": "Test description", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ], + ) + + async def test_update_description(self) -> None: + """Test updating task description.""" + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Test Task", + "description": "Original description", + }, + tool_call_id="task-update-description-create", + ) + task_id = self.agent_state.tasks_context.tasks[0].id + + response = await self._call_tool( + name="TaskUpdate", + tool_input={ + "task_id": task_id, + "description": "Updated description", + }, + tool_call_id="task-update-description", + ) + + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": f"Update task (id={task_id}) description.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-update-description", + }, + ) + + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Test Task", + "description": "Updated description", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ], + ) + + async def test_update_status(self) -> None: + """Test updating task status.""" + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Test Task", + "description": "Test description", + }, + tool_call_id="task-update-status-create", + ) + task_id = self.agent_state.tasks_context.tasks[0].id + + response = await self._call_tool( + name="TaskUpdate", + tool_input={ + "task_id": task_id, + "status": "in_progress", + }, + tool_call_id="task-update-status-in-progress", + ) + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": f"Update task (id={task_id}) status.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-update-status-in-progress", + }, + ) + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Test Task", + "description": "Test description", + "metadata": {}, + "created_at": AnyString(), + "state": "in_progress", + "id": task_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ], + ) + + response = await self._call_tool( + name="TaskUpdate", + tool_input={ + "task_id": task_id, + "status": "completed", + }, + tool_call_id="task-update-status-completed", + ) + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": f"Update task (id={task_id}) status.\n\n" + "Task completed. Call TaskList now to find your " + "next available task or see if your work " + "unblocked others.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-update-status-completed", + }, + ) + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Test Task", + "description": "Test description", + "metadata": {}, + "created_at": AnyString(), + "state": "completed", + "id": task_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ], + ) + + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Task to Delete", + "description": "This task will be deleted", + }, + tool_call_id="task-update-status-create-delete", + ) + task_to_delete_id = self.agent_state.tasks_context.tasks[1].id + + response = await self._call_tool( + name="TaskUpdate", + tool_input={ + "task_id": task_to_delete_id, + "status": "deleted", + }, + tool_call_id="task-update-status-deleted", + ) + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": f"Task (id={task_to_delete_id}) has been " + f"deleted.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-update-status-deleted", + }, + ) + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Test Task", + "description": "Test description", + "metadata": {}, + "created_at": AnyString(), + "state": "completed", + "id": task_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ], + ) + + async def test_update_owner(self) -> None: + """Test updating task owner.""" + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Test Task", + "description": "Test description", + }, + tool_call_id="task-update-owner-create", + ) + task_id = self.agent_state.tasks_context.tasks[0].id + + response = await self._call_tool( + name="TaskUpdate", + tool_input={ + "task_id": task_id, + "owner": "agent-1", + }, + tool_call_id="task-update-owner", + ) + + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": f"Update task (id={task_id}) owner.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-update-owner", + }, + ) + + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Test Task", + "description": "Test description", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_id, + "owner": "agent-1", + "blocks": [], + "blocked_by": [], + }, + ], + ) + + async def test_update_metadata(self) -> None: + """Test updating task metadata.""" + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Test Task", + "description": "Test description", + "metadata": { + "priority": "low", + "tags": ["test"], + }, + }, + tool_call_id="task-update-metadata-create", + ) + task_id = self.agent_state.tasks_context.tasks[0].id + + response = await self._call_tool( + name="TaskUpdate", + tool_input={ + "task_id": task_id, + "metadata": { + "priority": "high", + "new_field": "value", + }, + }, + tool_call_id="task-update-metadata", + ) + + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": f"Update task (id={task_id}) metadata.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-update-metadata", + }, + ) + + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Test Task", + "description": "Test description", + "metadata": { + "priority": "high", + "tags": ["test"], + "new_field": "value", + }, + "created_at": AnyString(), + "state": "pending", + "id": task_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ], + ) + + async def test_update_add_blocks(self) -> None: + """Test updating task blocks relationship.""" + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Task 1", + "description": "First task", + }, + tool_call_id="task-update-add-blocks-create-1", + ) + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Task 2", + "description": "Second task", + }, + tool_call_id="task-update-add-blocks-create-2", + ) + + task_1_id = self.agent_state.tasks_context.tasks[0].id + task_2_id = self.agent_state.tasks_context.tasks[1].id + response = await self._call_tool( + name="TaskUpdate", + tool_input={ + "task_id": task_1_id, + "add_blocks": [task_2_id], + }, + tool_call_id="task-update-add-blocks", + ) + + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": f"Update task (id={task_1_id}) add_blocks.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-update-add-blocks", + }, + ) + + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Task 1", + "description": "First task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_1_id, + "owner": None, + "blocks": [task_2_id], + "blocked_by": [], + }, + { + "subject": "Task 2", + "description": "Second task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_2_id, + "owner": None, + "blocks": [], + "blocked_by": [task_1_id], + }, + ], + ) + + async def test_update_add_blocked_by(self) -> None: + """Test updating task blocked_by relationship.""" + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Task 1", + "description": "First task", + }, + tool_call_id="task-update-add-blocked-by-create-1", + ) + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Task 2", + "description": "Second task", + }, + tool_call_id="task-update-add-blocked-by-create-2", + ) + + task_1_id = self.agent_state.tasks_context.tasks[0].id + task_2_id = self.agent_state.tasks_context.tasks[1].id + response = await self._call_tool( + name="TaskUpdate", + tool_input={ + "task_id": task_2_id, + "add_blocked_by": [task_1_id], + }, + tool_call_id="task-update-add-blocked-by", + ) + + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": f"Update task (id={task_2_id}) " + f"add_blocked_by.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-update-add-blocked-by", + }, + ) + + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Task 1", + "description": "First task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_1_id, + "owner": None, + "blocks": [task_2_id], + "blocked_by": [], + }, + { + "subject": "Task 2", + "description": "Second task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_2_id, + "owner": None, + "blocks": [], + "blocked_by": [task_1_id], + }, + ], + ) + + async def test_update_delete_task(self) -> None: + """Test deleting a task and cleaning dependency relations.""" + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Task 1", + "description": "First task", + }, + tool_call_id="task-update-delete-create-1", + ) + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Task 2", + "description": "Second task", + }, + tool_call_id="task-update-delete-create-2", + ) + await self._call_tool( + name="TaskCreate", + tool_input={ + "subject": "Task 3", + "description": "Third task", + }, + tool_call_id="task-update-delete-create-3", + ) + + task_1_id = self.agent_state.tasks_context.tasks[0].id + task_2_id = self.agent_state.tasks_context.tasks[1].id + task_3_id = self.agent_state.tasks_context.tasks[2].id + + await self._call_tool( + name="TaskUpdate", + tool_input={ + "task_id": task_1_id, + "add_blocks": [task_2_id], + }, + tool_call_id="task-update-delete-link-1", + ) + await self._call_tool( + name="TaskUpdate", + tool_input={ + "task_id": task_2_id, + "add_blocks": [task_3_id], + }, + tool_call_id="task-update-delete-link-2", + ) + + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Task 1", + "description": "First task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_1_id, + "owner": None, + "blocks": [task_2_id], + "blocked_by": [], + }, + { + "subject": "Task 2", + "description": "Second task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_2_id, + "owner": None, + "blocks": [task_3_id], + "blocked_by": [task_1_id], + }, + { + "subject": "Task 3", + "description": "Third task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_3_id, + "owner": None, + "blocks": [], + "blocked_by": [task_2_id], + }, + ], + ) + + response = await self._call_tool( + name="TaskUpdate", + tool_input={ + "task_id": task_2_id, + "status": "deleted", + }, + tool_call_id="task-update-delete", + ) + + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": f"Task (id={task_2_id}) has been deleted.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + "id": "task-update-delete", + }, + ) + self.assertEqual( + self._dump_tasks(), + [ + { + "subject": "Task 1", + "description": "First task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_1_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + { + "subject": "Task 3", + "description": "Third task", + "metadata": {}, + "created_at": AnyString(), + "state": "pending", + "id": task_3_id, + "owner": None, + "blocks": [], + "blocked_by": [], + }, + ], + ) + + async def test_update_nonexistent_task(self) -> None: + """Test updating a task that does not exist.""" + response = await self._call_tool( + name="TaskUpdate", + tool_input={ + "task_id": "nonexistent-id", + "subject": "New Subject", + "description": None, + }, + tool_call_id="task-update-missing", + ) + + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "text": "TaskNotFoundError: " + "The task (id=nonexistent-id) does not exist.", + "type": "text", + "id": AnyString(), + }, + ], + "state": "error", + "metadata": {}, + "id": "task-update-missing", + }, + ) + self.assertEqual(self._dump_tasks(), []) diff --git a/tests/toolkit_test.py b/tests/toolkit_test.py new file mode 100644 index 0000000000000000000000000000000000000000..410dbcb2abdfb435264511e5088f909785312031 --- /dev/null +++ b/tests/toolkit_test.py @@ -0,0 +1,1387 @@ +# -*- coding: utf-8 -*- +# pylint: disable=unused-argument +"""Toolkit test case.""" +import base64 +import json +from typing import Any, AsyncGenerator, Generator +from unittest import TestCase +from unittest.async_case import IsolatedAsyncioTestCase + + +from utils import AnyString + +from agentscope.state import AgentState +from agentscope.message import ( + TextBlock, + ToolCallBlock, + DataBlock, + Base64Source, +) +from agentscope.tool import ( + Toolkit, + ToolBase, + ToolChunk, + ToolResponse, + ToolGroup, + FunctionTool, +) +from agentscope.permission import ( + PermissionDecision, + PermissionBehavior, +) + + +class Tool1(ToolBase): + """A simple tool for testing.""" + + name: str = "tool_1" + description: str = "A simple tool for testing." + input_schema: dict = { + "type": "object", + "properties": {}, + } + is_concurrency_safe: bool = True + is_read_only: bool = True + is_mcp: bool = False + + async def check_permissions( + self, + *args: Any, + **kwargs: Any, + ) -> PermissionDecision: + """Check permissions for the tool.""" + return PermissionDecision( + behavior=PermissionBehavior.ASK, + message="Do you want to use my_tool?", + ) + + async def call(self, **kwargs: Any) -> ToolChunk: + """Run the tool.""" + return ToolChunk( + content=[TextBlock(text="Hello, world!")], + ) + + +class Tool2(ToolBase): + """Test tool 2""" + + name: str = "tool_2" + description: str = "Test tool 2." + input_schema: dict = { + "type": "object", + "properties": {}, + } + is_concurrency_safe: bool = True + is_read_only: bool = True + is_mcp: bool = False + + async def check_permissions( + self, + *args: Any, + **kwargs: Any, + ) -> PermissionDecision: + """Check permissions for the tool.""" + return PermissionDecision( + behavior=PermissionBehavior.ASK, + message="Do you want to use my_tool?", + ) + + async def call(self, **kwargs: Any) -> AsyncGenerator[ToolChunk, None]: + """Run the tool.""" + yield ToolChunk( + content=[TextBlock(text="123", id="a")], + ) + yield ToolChunk( + content=[TextBlock(text="456", id="b")], + ) + yield ToolChunk( + content=[TextBlock(text="789", id="b")], + ) + yield ToolChunk( + content=[ + DataBlock( + id="1", + source=Base64Source( + data="abc", + media_type="image/jpeg", + ), + ), + ], + ) + yield ToolChunk( + content=[ + DataBlock( + id="2", + source=Base64Source( + data="***", + media_type="image/jpeg", + ), + ), + ], + ) + yield ToolChunk( + content=[ + DataBlock( + id="1", + source=Base64Source( + data="def", + media_type="image/jpeg", + ), + ), + ], + ) + + +class ToolkitTest(IsolatedAsyncioTestCase): + """The toolkit test case.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + + async def test_initialize(self) -> None: + """The template test.""" + + # Initialize the toolkit + toolkit = Toolkit() + + # No tools + schemas = await toolkit.get_tool_schemas() + self.assertEqual(len(schemas), 0) + + # Initialize the toolkit with tools + toolkit = Toolkit(tools=[Tool1(), Tool2()]) + schemas = await toolkit.get_tool_schemas() + self.assertListEqual( + schemas, + [ + { + "type": "function", + "function": { + "name": "tool_1", + "description": "A simple tool for testing.", + "parameters": { + "type": "object", + "properties": {}, + }, + }, + }, + { + "type": "function", + "function": { + "name": "tool_2", + "description": "Test tool 2.", + "parameters": { + "type": "object", + "properties": {}, + }, + }, + }, + ], + ) + + async def test_tool(self) -> None: + """Test executing a tool.""" + toolkit = Toolkit(tools=[Tool1(), Tool2()]) + state = AgentState() + + # Test Tool1 (returns single ToolChunk) + tool_call_1 = ToolCallBlock( + id="test_1", + name="tool_1", + input=json.dumps({}), + ) + + chunks_1 = [] + response_1 = None + async for result in toolkit.call_tool(tool_call_1, state): + if isinstance(result, ToolChunk): + chunks_1.append(result) + elif isinstance(result, ToolResponse): + response_1 = result + + # Verify Tool1 chunks + self.assertEqual(len(chunks_1), 1) + self.assertDictEqual( + chunks_1[0].model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hello, world!", + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + # Verify Tool1 response + self.assertIsNotNone(response_1) + self.assertDictEqual( + response_1.model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Hello, world!", + }, + ], + "state": "success", + "metadata": {}, + "id": "test_1", + }, + ) + + # Test Tool2 (returns async generator of ToolChunks) + tool_call_2 = ToolCallBlock( + id="test_2", + name="tool_2", + input=json.dumps({}), + ) + + chunks_2 = [] + response_2 = None + async for result in toolkit.call_tool(tool_call_2, state): + if isinstance(result, ToolChunk): + chunks_2.append(result) + elif isinstance(result, ToolResponse): + response_2 = result + + # Verify Tool2 chunks + self.assertEqual(len(chunks_2), 6) + + # First chunk - TextBlock id="a" + self.assertDictEqual( + chunks_2[0].model_dump(), + { + "content": [ + { + "type": "text", + "id": "a", + "text": "123", + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + # Second chunk - TextBlock id="b" + self.assertDictEqual( + chunks_2[1].model_dump(), + { + "content": [ + { + "type": "text", + "id": "b", + "text": "456", + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + # Third chunk - TextBlock id="b" (same id as second) + self.assertDictEqual( + chunks_2[2].model_dump(), + { + "content": [ + { + "type": "text", + "id": "b", + "text": "789", + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + # Fourth chunk - DataBlock id="1" + self.assertDictEqual( + chunks_2[3].model_dump(), + { + "content": [ + { + "type": "data", + "id": "1", + "name": None, + "source": { + "type": "base64", + "media_type": "image/jpeg", + "data": "abc", + }, + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + # Fifth chunk - DataBlock id="2" + self.assertDictEqual( + chunks_2[4].model_dump(), + { + "content": [ + { + "type": "data", + "id": "2", + "name": None, + "source": { + "type": "base64", + "media_type": "image/jpeg", + "data": "***", + }, + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + # Sixth chunk - DataBlock id="1" (same id as fourth) + self.assertDictEqual( + chunks_2[5].model_dump(), + { + "content": [ + { + "type": "data", + "id": "1", + "name": None, + "source": { + "type": "base64", + "media_type": "image/jpeg", + "data": "def", + }, + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + # Verify Tool2 response - blocks with same id are merged, + # and consecutive TextBlocks are also merged + # TextBlock id="a" (123) + id="b" (456) + id="b" (789) -> merged to + # "123456789" with id="a" + # DataBlock id="1" appears twice, should be merged to "abcdef" + self.assertIsNotNone(response_2) + self.assertDictEqual( + response_2.model_dump(), + { + "content": [ + { + "type": "text", + "id": "a", + # All consecutive TextBlocks merged + "text": "123456789", + }, + { + "type": "data", + "id": "1", + "name": None, + "source": { + "type": "base64", + "media_type": "image/jpeg", + "data": "abcdef", # Merged: "abc" + "def" + }, + }, + { + "type": "data", + "id": "2", + "name": None, + "source": { + "type": "base64", + "media_type": "image/jpeg", + "data": "***", + }, + }, + ], + "state": "success", + "metadata": {}, + "id": "test_2", + }, + ) + + async def test_tool_response_merges_base64_chunks_by_bytes(self) -> None: + """Base64 chunks with padding should merge as bytes, not strings.""" + response = ToolResponse() + first = base64.b64encode(b"hello").decode("ascii") + second = base64.b64encode(b"world").decode("ascii") + + response.append_chunk( + ToolChunk( + content=[ + DataBlock( + id="image", + source=Base64Source( + data=first, + media_type="image/png", + ), + ), + ], + ), + ) + response.append_chunk( + ToolChunk( + content=[ + DataBlock( + id="image", + source=Base64Source( + data=second, + media_type="image/png", + ), + ), + ], + ), + ) + + self.assertEqual(len(response.content), 1) + merged = response.content[0] + self.assertIsInstance(merged, DataBlock) + self.assertEqual( + base64.b64decode(merged.source.data), + b"helloworld", + ) + + +class RegisterFunctionTest(IsolatedAsyncioTestCase): + """Test registering different functions in the toolkit.""" + + async def test_sync_non_streaming_function(self) -> None: + """Test registering a synchronous non-streaming function.""" + + def add_numbers(a: int, b: int) -> ToolChunk: + """Add two numbers together. + + Args: + a: The first number + b: The second number + """ + result = a + b + return ToolChunk( + content=[TextBlock(text=f"Result: {result}")], + ) + + toolkit = Toolkit( + tools=[FunctionTool(add_numbers)], + ) + + # Test schema + schemas = await toolkit.get_tool_schemas() + self.assertEqual(len(schemas), 1) + self.assertDictEqual( + schemas[0], + { + "type": "function", + "function": { + "name": "add_numbers", + "description": "Add two numbers together.", + "parameters": { + "type": "object", + "properties": { + "a": { + "type": "integer", + "description": "The first number", + }, + "b": { + "type": "integer", + "description": "The second number", + }, + }, + "required": ["a", "b"], + }, + }, + }, + ) + + # Test execution + state = AgentState() + tool_call = ToolCallBlock( + id="test_add", + name="add_numbers", + input=json.dumps({"a": 3, "b": 5}), + ) + + chunks = [] + response = None + async for result in toolkit.call_tool(tool_call, state): + if isinstance(result, ToolChunk): + chunks.append(result) + elif isinstance(result, ToolResponse): + response = result + + # Verify chunks + self.assertEqual(len(chunks), 1) + self.assertDictEqual( + chunks[0].model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Result: 8", + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + # Verify response + self.assertIsNotNone(response) + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Result: 8", + }, + ], + "state": "success", + "metadata": {}, + "id": "test_add", + }, + ) + + async def test_sync_function_returning_plain_string(self) -> None: + """Test wrapping a sync function that returns a plain string.""" + + def get_weather(location: str) -> str: + """Get weather information. + + Args: + location: The location to get weather for + """ + return f"The weather in {location} is sunny." + + toolkit = Toolkit( + tools=[FunctionTool(get_weather)], + ) + + state = AgentState() + tool_call = ToolCallBlock( + id="test_weather", + name="get_weather", + input=json.dumps({"location": "Chengdu"}), + ) + + chunks = [] + response = None + async for result in toolkit.call_tool(tool_call, state): + if isinstance(result, ToolChunk): + chunks.append(result) + elif isinstance(result, ToolResponse): + response = result + + self.assertEqual(len(chunks), 1) + self.assertDictEqual( + chunks[0].model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "The weather in Chengdu is sunny.", + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + self.assertIsNotNone(response) + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "The weather in Chengdu is sunny.", + }, + ], + "state": "success", + "metadata": {}, + "id": "test_weather", + }, + ) + + async def test_async_function_returning_json_serializable_value( + self, + ) -> None: + """Test wrapping an async function that returns a JSON value.""" + + async def get_weather(location: str) -> dict: + """Get weather information. + + Args: + location: The location to get weather for + """ + return { + "location": location, + "weather": { + "condition": "sunny", + "temperature": 22, + }, + } + + toolkit = Toolkit( + tools=[FunctionTool(get_weather)], + ) + + state = AgentState() + tool_call = ToolCallBlock( + id="test_weather_json", + name="get_weather", + input=json.dumps({"location": "成都"}), + ) + expected_text = json.dumps( + { + "location": "成都", + "weather": { + "condition": "sunny", + "temperature": 22, + }, + }, + ensure_ascii=False, + ) + + chunks = [] + response = None + async for result in toolkit.call_tool(tool_call, state): + if isinstance(result, ToolChunk): + chunks.append(result) + elif isinstance(result, ToolResponse): + response = result + + self.assertEqual(len(chunks), 1) + self.assertEqual(chunks[0].content[0].text, expected_text) + + self.assertIsNotNone(response) + self.assertEqual(response.content[0].text, expected_text) + + async def test_sync_streaming_function(self) -> None: + """Test registering a synchronous streaming function.""" + + def count_to_n(n: int) -> Generator[ToolChunk, None, None]: + """Count from 1 to n. + + Args: + n: The number to count to + """ + for i in range(1, n + 1): + yield ToolChunk( + content=[TextBlock(text=str(i))], + ) + + toolkit = Toolkit( + tools=[ + FunctionTool( + func=count_to_n, + ), + ], + ) + + # Test schema + schemas = await toolkit.get_tool_schemas() + self.assertEqual(len(schemas), 1) + self.assertDictEqual( + schemas[0], + { + "type": "function", + "function": { + "name": "count_to_n", + "description": "Count from 1 to n.", + "parameters": { + "type": "object", + "properties": { + "n": { + "type": "integer", + "description": "The number to count to", + }, + }, + "required": ["n"], + }, + }, + }, + ) + + # Test execution + state = AgentState() + tool_call = ToolCallBlock( + id="test_count", + name="count_to_n", + input=json.dumps({"n": 3}), + ) + + chunks = [] + response = None + async for result in toolkit.call_tool(tool_call, state): + if isinstance(result, ToolChunk): + chunks.append(result) + elif isinstance(result, ToolResponse): + response = result + + # Verify chunks + self.assertEqual(len(chunks), 3) + for i, chunk in enumerate(chunks, 1): + self.assertDictEqual( + chunk.model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": str(i), + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + # Verify response - consecutive TextBlocks are merged + self.assertIsNotNone(response) + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "123", # All consecutive TextBlocks merged + }, + ], + "state": "success", + "metadata": {}, + "id": "test_count", + }, + ) + + async def test_sync_streaming_function_returning_plain_values( + self, + ) -> None: + """Test wrapping a sync generator that yields plain values.""" + + def stream_weather() -> Generator[Any, None, None]: + """Stream weather information.""" + yield "checking" + yield { + "condition": "sunny", + "temperature": 22, + } + + toolkit = Toolkit( + tools=[FunctionTool(stream_weather)], + ) + + state = AgentState() + tool_call = ToolCallBlock( + id="test_weather_stream", + name="stream_weather", + input=json.dumps({}), + ) + expected_dict_text = json.dumps( + { + "condition": "sunny", + "temperature": 22, + }, + ensure_ascii=False, + ) + + chunks = [] + response = None + async for result in toolkit.call_tool(tool_call, state): + if isinstance(result, ToolChunk): + chunks.append(result) + elif isinstance(result, ToolResponse): + response = result + + self.assertEqual( + [chunk.content[0].text for chunk in chunks], + ["checking", expected_dict_text], + ) + + self.assertIsNotNone(response) + self.assertEqual( + response.content[0].text, + f"checking{expected_dict_text}", + ) + + async def test_async_non_streaming_function(self) -> None: + """Test registering an asynchronous non-streaming function.""" + + async def multiply_numbers(x: float, y: float) -> ToolChunk: + """Multiply two numbers. + + Args: + x: The first number + y: The second number + """ + result = x * y + return ToolChunk( + content=[TextBlock(text=f"Product: {result}")], + ) + + toolkit = Toolkit( + tools=[FunctionTool(multiply_numbers)], + ) + + # Test schema + schemas = await toolkit.get_tool_schemas() + self.assertEqual(len(schemas), 1) + self.assertDictEqual( + schemas[0], + { + "type": "function", + "function": { + "name": "multiply_numbers", + "description": "Multiply two numbers.", + "parameters": { + "type": "object", + "properties": { + "x": { + "type": "number", + "description": "The first number", + }, + "y": { + "type": "number", + "description": "The second number", + }, + }, + "required": ["x", "y"], + }, + }, + }, + ) + + # Test execution + state = AgentState() + tool_call = ToolCallBlock( + id="test_multiply", + name="multiply_numbers", + input=json.dumps({"x": 2.5, "y": 4.0}), + ) + + chunks = [] + response = None + async for result in toolkit.call_tool(tool_call, state): + if isinstance(result, ToolChunk): + chunks.append(result) + elif isinstance(result, ToolResponse): + response = result + + # Verify chunks + self.assertEqual(len(chunks), 1) + self.assertDictEqual( + chunks[0].model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Product: 10.0", + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + # Verify response + self.assertIsNotNone(response) + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "Product: 10.0", + }, + ], + "state": "success", + "metadata": {}, + "id": "test_multiply", + }, + ) + + async def test_async_streaming_function(self) -> None: + """Test registering an asynchronous streaming function.""" + + async def generate_sequence( + start: int, + end: int, + ) -> AsyncGenerator[ToolChunk, None]: + """Generate a sequence of numbers. + + Args: + start: The starting number + end: The ending number + """ + for i in range(start, end + 1): + yield ToolChunk( + content=[TextBlock(text=f"Number: {i}")], + ) + + toolkit = Toolkit( + tools=[FunctionTool(generate_sequence)], + ) + + # Test schema + schemas = await toolkit.get_tool_schemas() + self.assertEqual(len(schemas), 1) + self.assertDictEqual( + schemas[0], + { + "type": "function", + "function": { + "name": "generate_sequence", + "description": "Generate a sequence of numbers.", + "parameters": { + "type": "object", + "properties": { + "start": { + "type": "integer", + "description": "The starting number", + }, + "end": { + "type": "integer", + "description": "The ending number", + }, + }, + "required": ["start", "end"], + }, + }, + }, + ) + + # Test execution + state = AgentState() + tool_call = ToolCallBlock( + id="test_sequence", + name="generate_sequence", + input=json.dumps({"start": 5, "end": 7}), + ) + + chunks = [] + response = None + async for result in toolkit.call_tool(tool_call, state): + if isinstance(result, ToolChunk): + chunks.append(result) + elif isinstance(result, ToolResponse): + response = result + + # Verify chunks + self.assertEqual(len(chunks), 3) + for chunk, num in zip(chunks, [5, 6, 7]): + self.assertDictEqual( + chunk.model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + "text": f"Number: {num}", + }, + ], + "state": "running", + "is_last": True, + "metadata": {}, + "id": AnyString(), + }, + ) + + # Verify response - consecutive TextBlocks are merged + self.assertIsNotNone(response) + self.assertDictEqual( + response.model_dump(), + { + "content": [ + { + "type": "text", + "id": AnyString(), + # All consecutive TextBlocks merged + "text": "Number: 5Number: 6Number: 7", + }, + ], + "state": "success", + "metadata": {}, + "id": "test_sequence", + }, + ) + + async def test_async_streaming_function_returning_plain_values( + self, + ) -> None: + """Test wrapping an async generator that yields plain values.""" + + async def stream_status() -> AsyncGenerator[Any, None]: + """Stream status information.""" + yield "started" + yield { + "done": True, + } + + toolkit = Toolkit( + tools=[FunctionTool(stream_status)], + ) + + state = AgentState() + tool_call = ToolCallBlock( + id="test_status_stream", + name="stream_status", + input=json.dumps({}), + ) + expected_dict_text = json.dumps( + { + "done": True, + }, + ensure_ascii=False, + ) + + chunks = [] + response = None + async for result in toolkit.call_tool(tool_call, state): + if isinstance(result, ToolChunk): + chunks.append(result) + elif isinstance(result, ToolResponse): + response = result + + self.assertEqual( + [chunk.content[0].text for chunk in chunks], + ["started", expected_dict_text], + ) + + self.assertIsNotNone(response) + self.assertEqual( + response.content[0].text, + f"started{expected_dict_text}", + ) + + +class ToolGroupTest(IsolatedAsyncioTestCase): + """The tool group test case.""" + + async def asyncSetUp(self) -> None: + """The async setup method.""" + self.meta_tool_description = ( + "This tool allows you to reset your equipped tools based on your " + "current task requirements. These tools are organized into " + "different groups, and you can activate/deactivate them " + "by specifying the boolean values for each group in the input.\n\n" + "**Important: The input booleans are the final state of their " + "corresponding tool groups, not incremental changes.** Any " + "group not explicitly set to True will be deactivated, " + "regardless of its previous state.\n\n" + "**Best practice**: Actively manage your tool groups——activate " + "only what you need for the current task, and promptly " + "deactivate groups as soon as they are no longer needed to" + " conserve context space.\n\n" + "This tool will return the usage instructions for the activated " + "tool groups, which you **MUST pay attention to and follow**. " + "You can also reuse this tool to re-check the instructions." + ) + + async def test_meta_tool(self) -> None: + """Test creating a tool group.""" + toolkit = Toolkit() + + # Test meta tool when no groups exist + schemas = await toolkit.get_tool_schemas() + self.assertEqual(len(schemas), 0) + + toolkit = Toolkit( + tool_groups=[ + ToolGroup( + name="group_1", + description="Group 1", + ), + ], + ) + + # The group is created successfully + self.assertEqual(len(toolkit.tool_groups), 2) + # The builtin meta tool is activated + schemas = await toolkit.get_tool_schemas() + + self.assertListEqual( + schemas, + [ + { + "type": "function", + "function": { + "name": "reset_tools", + "description": self.meta_tool_description, + "parameters": { + "properties": { + "group_1": { + "default": False, + "description": "Group 1", + "type": "boolean", + }, + }, + "type": "object", + }, + }, + }, + ], + ) + + # Name conflict + with self.assertRaises(ValueError): + Toolkit( + tool_groups=[ + ToolGroup( + name="group_2", + description="Group 2", + ), + ToolGroup( + name="group_2", + description="Group 2", + ), + ], + ) + + # A new group with tools + toolkit = Toolkit( + tool_groups=[ + ToolGroup( + name="group_1", + description="Group 1", + ), + ToolGroup( + name="group_2", + description="Group 2", + tools=[Tool1(), Tool2()], + instructions="This is group 2.", + ), + ], + ) + + self.assertEqual(len(toolkit.tool_groups), 3) + schemas = await toolkit.get_tool_schemas() + self.assertListEqual( + schemas, + [ + { + "type": "function", + "function": { + "name": "reset_tools", + "description": self.meta_tool_description, + "parameters": { + "properties": { + "group_1": { + "default": False, + "description": "Group 1", + "type": "boolean", + }, + "group_2": { + "default": False, + "description": "Group 2", + "type": "boolean", + }, + }, + "type": "object", + }, + }, + }, + ], + ) + + # Active one group + state = AgentState() + res = toolkit.call_tool( + ToolCallBlock( + id="xxx", + name="reset_tools", + input=json.dumps({"group_2": True}), + ), + state, + ) + + chunk = await anext(res) + self.assertIsInstance(chunk, ToolChunk) + + chunk = await anext(res) + self.assertIsInstance(chunk, ToolResponse) + self.assertDictEqual( + chunk.model_dump(), + { + "id": AnyString(), + "content": [ + { + "type": "text", + "id": AnyString(), + "text": """The currently activated tool group(s): group_2. + +The tool instructions are a collection of suggestions, rules and notifications about how to use the tools in the activated groups. +This is group 2. +""", # noqa: E501 + }, + ], + "metadata": {}, + "state": "success", + }, + ) + + # Activate both groups + res = toolkit.call_tool( + ToolCallBlock( + id="xxx", + name="reset_tools", + input=json.dumps({"group_2": True, "group_1": True}), + ), + state, + ) + + last_chunk = None + async for chunk in res: + last_chunk = chunk + + self.assertDictEqual( + last_chunk.model_dump(), + { + "id": AnyString(), + "content": [ + { + "type": "text", + "id": AnyString(), + "text": """The currently activated tool group(s): group_1, group_2. + +The tool instructions are a collection of suggestions, rules and notifications about how to use the tools in the activated groups. +This is group 2. +""", # noqa: E501 + }, + ], + "metadata": {}, + "state": "success", + }, + ) + + # deactivate all groups + res = toolkit.call_tool( + ToolCallBlock( + id="xxx", + name="reset_tools", + input=json.dumps({}), + ), + state, + ) + + last_chunk = None + async for chunk in res: + last_chunk = chunk + + self.assertDictEqual( + last_chunk.model_dump(), + { + "id": AnyString(), + "content": [ + { + "type": "text", + "id": AnyString(), + "text": "All tool groups are currently deactivated.", + }, + ], + "metadata": {}, + "state": "success", + }, + ) + + +class RemoveTitleFieldTest(TestCase): + """Unit tests for _remove_title_field.""" + + def setUp(self) -> None: + from agentscope.tool._utils import _remove_title_field + + self.fn = _remove_title_field + + def test_removes_top_level_title(self) -> None: + """The top level title field must be removed.""" + schema = {"title": "Root", "type": "object", "properties": {}} + self.fn(schema) + self.assertNotIn("title", schema) + + def test_removes_property_titles(self) -> None: + """Titles inside properties must be removed.""" + schema: dict[str, Any] = { + "type": "object", + "properties": { + "name": {"title": "Name", "type": "string"}, + }, + } + self.fn(schema) + self.assertNotIn("title", schema["properties"]["name"]) + + def test_removes_defs_titles(self) -> None: + """Titles inside $defs must be recursively stripped.""" + schema: dict[str, Any] = { + "type": "object", + "properties": {"x": {"$ref": "#/$defs/MyModel"}}, + "$defs": { + "MyModel": { + "title": "MyModel", + "type": "object", + "properties": { + "val": {"title": "Val", "type": "string"}, + }, + }, + }, + } + self.fn(schema) + + self.assertDictEqual( + schema, + { + "type": "object", + "properties": {"x": {"$ref": "#/$defs/MyModel"}}, + "$defs": { + "MyModel": { + "type": "object", + "properties": { + "val": {"type": "string"}, + }, + }, + }, + }, + ) + + def test_does_not_mutate_non_dict_defs(self) -> None: + """Boolean schema values inside $defs should not raise.""" + schema: dict[str, Any] = { + "type": "object", + "properties": {}, + "$defs": {"AlwaysTrue": True}, + } + self.fn(schema) + self.assertEqual(schema["$defs"]["AlwaysTrue"], True) diff --git a/tests/tracing_test.py b/tests/tracing_test.py new file mode 100644 index 0000000000000000000000000000000000000000..96f14c2b02487f2dc8a55eedbeef487aec671a00 --- /dev/null +++ b/tests/tracing_test.py @@ -0,0 +1,846 @@ +# -*- coding: utf-8 -*- +"""Unit tests for the tracing module using an in-memory OTel exporter.""" +import json +from typing import Any +from unittest.async_case import IsolatedAsyncioTestCase + +from opentelemetry import trace as otel_trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) + +from utils import MockModel + +from agentscope.agent import Agent +from agentscope.event import ( + ConfirmResult, + ExternalExecutionResultEvent, + RequireExternalExecutionEvent, + RequireUserConfirmEvent, + UserConfirmResultEvent, +) +from agentscope.message import ( + TextBlock, + ToolCallBlock, + ToolResultBlock, + ToolResultState, + UserMsg, +) +from agentscope.model import ChatResponse, ChatUsage +from agentscope.permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, +) +from agentscope.tool import Toolkit, ToolBase +from agentscope.middleware import TracingMiddleware + + +# --------------------------------------------------------------------------- +# Shared test fixtures +# --------------------------------------------------------------------------- + + +class WeatherTool(ToolBase): + """Stub weather tool for tracing tests.""" + + name: str = "get_weather" + description: str = "Return stub weather for a city." + input_schema: dict = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + } + is_concurrency_safe: bool = True + is_read_only: bool = True + is_external_tool: bool = False + is_mcp: bool = False + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message="always allowed", + ) + + async def execute(self, city: str) -> str: + """Stub weather tool for tracing tests.""" + return f"{city}: sunny, 25°C." + + +class HitlWeatherTool(ToolBase): + """Weather tool that always asks user for confirmation (HITL).""" + + name: str = "get_weather" + description: str = "Return weather for a city, requires user confirmation." + input_schema: dict = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + } + is_concurrency_safe: bool = True + is_read_only: bool = True + is_external_tool: bool = False + is_mcp: bool = False + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + return PermissionDecision( + behavior=PermissionBehavior.ASK, + message="user must confirm this tool call", + ) + + async def execute(self, city: str) -> str: + """Stub weather tool for tracing tests.""" + return f"{city}: sunny, 25°C." + + +class ExternalWeatherTool(ToolBase): + """Weather tool that is always executed externally.""" + + name: str = "get_weather" + description: str = "Return weather for a city via external execution." + input_schema: dict = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + } + is_concurrency_safe: bool = True + is_read_only: bool = True + is_external_tool: bool = True # Mark as external + is_mcp: bool = False + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message="always allowed", + ) + + async def execute(self, city: str) -> str: + """Stub weather tool for tracing tests.""" + return f"{city}: sunny, 25°C." + + +def _make_tool_call_response(tool_id: str, city: str) -> ChatResponse: + return ChatResponse( + content=[ + ToolCallBlock( + id=tool_id, + name="get_weather", + input=json.dumps({"city": city}), + ), + ], + is_last=True, + usage=ChatUsage(input_tokens=10, output_tokens=5, time=0.05), + ) + + +def _make_text_response(text: str) -> ChatResponse: + return ChatResponse( + content=[TextBlock(text=text)], + is_last=True, + usage=ChatUsage(input_tokens=15, output_tokens=8, time=0.05), + ) + + +class TracingTest(IsolatedAsyncioTestCase): + """Tests that OTel spans are emitted with correct attributes. + + The in-memory exporter is set up once per class (setUpClass) because + the OTel global TracerProvider cannot be replaced once installed. + setUp only creates fresh model/agent and clears the exporter. + """ + + exporter: InMemorySpanExporter + + @classmethod + def setUpClass(cls) -> None: + """Configure an in-memory OTel provider once for the whole class.""" + cls.exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(cls.exporter)) + otel_trace.set_tracer_provider(provider) + + def setUp(self) -> None: + """Create a fresh agent and clear accumulated spans before each + test.""" + self.exporter.clear() + self.model = MockModel() + self.agent = Agent( + name="test-agent", + system_prompt="You are a test assistant.", + model=self.model, + toolkit=Toolkit(tools=[WeatherTool()]), + middlewares=[TracingMiddleware()], + ) + + # ----------------------------------------------------------------------- + # Helpers + # ----------------------------------------------------------------------- + + def _spans_by_name(self, fragment: str) -> list: + return [ + s + for s in self.exporter.get_finished_spans() + if fragment in (s.name or "") + ] + + def _all_conv_ids(self) -> set: + return { + dict(s.attributes or {}).get("gen_ai.conversation.id") + for s in self.exporter.get_finished_spans() + if "gen_ai.conversation.id" in (s.attributes or {}) + } + + # ----------------------------------------------------------------------- + # Tests: Agent.reply + # ----------------------------------------------------------------------- + + async def test_reply_spans_share_conversation_id(self) -> None: + """All spans from a single reply must share the same + conversation_id.""" + self.model.set_responses( + [ + _make_tool_call_response("c3", "Guangzhou"), + _make_text_response("Guangzhou is sunny."), + ], + ) + msg = UserMsg(name="user", content="Weather in Guangzhou?") + await self.agent.reply(msg) + + conv_ids = self._all_conv_ids() + self.assertEqual( + len(conv_ids), + 1, + f"All spans must share exactly one conversation_id, " + f"got: {conv_ids}", + ) + + async def test_invoke_agent_span_has_response_attributes(self) -> None: + """invoke_agent span must carry gen_ai.output.messages attribute.""" + self.model.set_responses( + [ + _make_tool_call_response("c4", "Wuhan"), + _make_text_response("Wuhan weather: clear sky."), + ], + ) + msg = UserMsg(name="user", content="Weather in Wuhan?") + await self.agent.reply(msg) + + agent_spans = self._spans_by_name("invoke_agent") + self.assertEqual( + len(agent_spans), + 1, + "Expected exactly one invoke_agent span", + ) + span_attrs = dict(agent_spans[0].attributes or {}) + output_raw = span_attrs.get("gen_ai.output.messages") + assert isinstance( + output_raw, + str, + ), "invoke_agent span gen_ai.output.messages should be a string" + output = json.loads(output_raw) + self.assertEqual( + output, + [ + { + "role": "assistant", + "parts": [ + { + "type": "text", + "content": "Wuhan weather: clear sky.", + }, + ], + "name": "test-agent", + "finish_reason": "stop", + }, + ], + ) + + async def test_invoke_agent_span_has_input_attributes(self) -> None: + """invoke_agent span must carry gen_ai.input.messages attribute.""" + self.model.set_responses( + [ + _make_text_response("Simple answer."), + ], + ) + msg = UserMsg(name="user", content="Simple question?") + await self.agent.reply(msg) + + agent_spans = self._spans_by_name("invoke_agent") + self.assertEqual( + len(agent_spans), + 1, + "Expected exactly one invoke_agent span", + ) + span_attrs = dict(agent_spans[0].attributes or {}) + input_raw = span_attrs.get("gen_ai.input.messages") + assert isinstance( + input_raw, + str, + ), "invoke_agent span gen_ai.input.messages should be a string" + input_msgs = json.loads(input_raw) + self.assertEqual( + input_msgs, + [ + { + "role": "user", + "parts": [ + {"type": "text", "content": "Simple question?"}, + ], + "name": "user", + "finish_reason": "stop", + }, + ], + ) + + async def test_execute_tool_span_has_tool_name_attribute(self) -> None: + """execute_tool span must have the correct gen_ai.tool.name + attribute.""" + self.model.set_responses( + [ + _make_tool_call_response("c8", "Nanjing"), + _make_text_response("Nanjing result."), + ], + ) + msg = UserMsg(name="user", content="Weather in Nanjing?") + await self.agent.reply(msg) + + tool_spans = self._spans_by_name("execute_tool") + self.assertEqual( + len(tool_spans), + 1, + "Expected exactly one execute_tool span", + ) + span_attrs = dict(tool_spans[0].attributes or {}) + self.assertEqual( + span_attrs.get("gen_ai.tool.name"), + "get_weather", + "execute_tool span should have gen_ai.tool.name = get_weather", + ) + + # ----------------------------------------------------------------------- + # Tests: chat (LLM) span + # ----------------------------------------------------------------------- + + async def test_chat_span_has_model_and_provider(self) -> None: + """chat span must carry gen_ai.request.model and + gen_ai.provider.name.""" + self.model.set_responses( + [_make_text_response("Hello from model.")], + ) + msg = UserMsg(name="user", content="Hello?") + await self.agent.reply(msg) + + chat_spans = self._spans_by_name("chat") + self.assertEqual(len(chat_spans), 1, "Expected exactly one chat span") + span_attrs = dict(chat_spans[0].attributes or {}) + self.assertEqual( + span_attrs.get("gen_ai.request.model"), + "mock-model", + "chat span gen_ai.request.model should equal mock-model", + ) + self.assertEqual( + span_attrs.get("gen_ai.operation.name"), + "chat", + "chat span gen_ai.operation.name should equal chat", + ) + + async def test_chat_span_has_output_messages(self) -> None: + """chat span must carry gen_ai.output.messages with response + content.""" + self.model.set_responses( + [_make_text_response("Weather is fine.")], + ) + msg = UserMsg(name="user", content="How is the weather?") + await self.agent.reply(msg) + + chat_spans = self._spans_by_name("chat") + self.assertEqual(len(chat_spans), 1, "Expected exactly one chat span") + span_attrs = dict(chat_spans[0].attributes or {}) + output_raw = span_attrs.get("gen_ai.output.messages") + assert isinstance( + output_raw, + str, + ), "chat span gen_ai.output.messages should be a string" + output = json.loads(output_raw) + self.assertEqual( + output, + [ + { + "role": "assistant", + "parts": [ + {"type": "text", "content": "Weather is fine."}, + ], + "finish_reason": "stop", + }, + ], + ) + + async def test_chat_span_has_usage_tokens(self) -> None: + """chat span must carry input/output token counts from usage.""" + self.model.set_responses( + [_make_text_response("Token test.")], + ) + msg = UserMsg(name="user", content="Count tokens?") + await self.agent.reply(msg) + + chat_spans = self._spans_by_name("chat") + self.assertEqual(len(chat_spans), 1, "Expected exactly one chat span") + span_attrs = dict(chat_spans[0].attributes or {}) + self.assertEqual( + span_attrs.get("gen_ai.usage.input_tokens"), + 15, + "chat span gen_ai.usage.input_tokens should equal 15", + ) + self.assertEqual( + span_attrs.get("gen_ai.usage.output_tokens"), + 8, + "chat span gen_ai.usage.output_tokens should equal 8", + ) + + # ----------------------------------------------------------------------- + # Tests: reply_id attribute + # ----------------------------------------------------------------------- + + async def test_invoke_agent_span_has_reply_id(self) -> None: + """invoke_agent span from a normal reply must carry reply_id. + + reply() delegates to _reply(), which is the only decorated entry + point, so there is exactly one invoke_agent span. We search by + attribute rather than relying on list order for robustness. + """ + self.model.set_responses( + [ + _make_tool_call_response("r1", "Wuhan"), + _make_text_response("Wuhan: clear sky."), + ], + ) + msg = UserMsg(name="user", content="Weather in Wuhan?") + await self.agent.reply(msg) + + agent_spans = self._spans_by_name("invoke_agent") + self.assertEqual( + len(agent_spans), + 1, + "Expected exactly one invoke_agent span", + ) + span_attrs = dict(agent_spans[0].attributes or {}) + self.assertIn( + "agentscope.agent.reply_id", + span_attrs, + "invoke_agent span should have agentscope.agent.reply_id", + ) + + # ----------------------------------------------------------------------- + # Tests: HITL (Human-in-the-loop) + # ----------------------------------------------------------------------- + + async def test_hitl_first_call_has_hitl_pending_attribute(self) -> None: + """First call in HITL flow must have + agentscope.agent.hitl_pending_tools.""" + hitl_agent = Agent( + name="hitl-agent", + system_prompt="You are a test assistant.", + model=self.model, + toolkit=Toolkit(tools=[HitlWeatherTool()]), + middlewares=[TracingMiddleware()], + ) + # The HITL tool returns ASK; + # first call ends with RequireUserConfirmEvent + self.model.set_responses( + [_make_tool_call_response("h1", "Beijing")], + ) + self.exporter.clear() + + msg = UserMsg(name="user", content="Weather in Beijing?") + async for _ in hitl_agent.reply_stream(msg): + pass + + first_spans = self._spans_by_name("invoke_agent") + self.assertEqual( + len(first_spans), + 1, + "Expected exactly one invoke_agent span from first HITL call", + ) + span_attrs = dict(first_spans[0].attributes or {}) + self.assertIn( + "agentscope.agent.hitl_pending_tools", + span_attrs, + "First HITL span should carry agentscope.agent.hitl_pending_tools", + ) + pending = json.loads(span_attrs["agentscope.agent.hitl_pending_tools"]) + self.assertEqual(pending, ["get_weather"]) + + async def test_hitl_spans_share_reply_id(self) -> None: + """Both HITL calls must share the same agentscope.agent.reply_id.""" + hitl_agent = Agent( + name="hitl-agent2", + system_prompt="You are a test assistant.", + model=self.model, + toolkit=Toolkit(tools=[HitlWeatherTool()]), + middlewares=[TracingMiddleware()], + ) + + # First call: agent asks for confirmation + self.model.set_responses( + [_make_tool_call_response("h2", "Shanghai")], + ) + self.exporter.clear() + + require_confirm_event = None + async for evt in hitl_agent.reply_stream( + UserMsg(name="user", content="Weather in Shanghai?"), + ): + if isinstance(evt, RequireUserConfirmEvent): + require_confirm_event = evt + + self.assertIsNotNone( + require_confirm_event, + "Expected RequireUserConfirmEvent", + ) + + first_spans = self._spans_by_name("invoke_agent") + first_reply_ids = { + dict(s.attributes or {}).get("agentscope.agent.reply_id") + for s in first_spans + if "agentscope.agent.reply_id" in (s.attributes or {}) + } + self.assertEqual(len(first_reply_ids), 1) + reply_id_first = next(iter(first_reply_ids)) + + # Second call: user confirms + self.model.set_responses( + [_make_text_response("Shanghai: 18°C, raining.")], + ) + self.exporter.clear() + + confirm_event = UserConfirmResultEvent( + reply_id=require_confirm_event.reply_id, + confirm_results=[ + ConfirmResult( + confirmed=True, + tool_call=require_confirm_event.tool_calls[0], + ), + ], + ) + await hitl_agent.reply(inputs=confirm_event) + + second_spans = self._spans_by_name("invoke_agent") + second_reply_ids = { + dict(s.attributes or {}).get("agentscope.agent.reply_id") + for s in second_spans + if "agentscope.agent.reply_id" in (s.attributes or {}) + } + self.assertEqual(len(second_reply_ids), 1) + reply_id_second = next(iter(second_reply_ids)) + + self.assertEqual( + reply_id_first, + reply_id_second, + "Both HITL calls must share the same reply_id", + ) + + async def test_hitl_second_call_has_incoming_event_type(self) -> None: + """Second call in HITL flow must carry + incoming_event_type=user_confirm_result.""" + hitl_agent = Agent( + name="hitl-agent3", + system_prompt="You are a test assistant.", + model=self.model, + toolkit=Toolkit(tools=[HitlWeatherTool()]), + middlewares=[TracingMiddleware()], + ) + + # First call: agent asks for confirmation + self.model.set_responses( + [_make_tool_call_response("h3", "Tianjin")], + ) + self.exporter.clear() + + require_confirm_event = None + async for evt in hitl_agent.reply_stream( + UserMsg(name="user", content="Weather in Tianjin?"), + ): + if isinstance(evt, RequireUserConfirmEvent): + require_confirm_event = evt + + self.assertIsNotNone( + require_confirm_event, + "Expected RequireUserConfirmEvent", + ) + + # Second call: user confirms + self.model.set_responses( + [_make_text_response("Tianjin: windy, 12°C.")], + ) + self.exporter.clear() + + confirm_event = UserConfirmResultEvent( + reply_id=require_confirm_event.reply_id, + confirm_results=[ + ConfirmResult( + confirmed=True, + tool_call=require_confirm_event.tool_calls[0], + ), + ], + ) + await hitl_agent.reply(inputs=confirm_event) + + agent_spans = self._spans_by_name("invoke_agent") + self.assertEqual( + len(agent_spans), + 1, + "Expected exactly one invoke_agent span from second HITL call", + ) + span_attrs = dict(agent_spans[0].attributes or {}) + self.assertEqual( + span_attrs.get("agentscope.agent.incoming_event_type"), + "user_confirm_result", + "Second HITL invoke_agent span should have " + "incoming_event_type=user_confirm_result", + ) + + # ----------------------------------------------------------------------- + # Tests: External execution + # ----------------------------------------------------------------------- + + async def test_external_execution_first_call_has_pending_attribute( + self, + ) -> None: + """First call must have + agentscope.agent.external_execution_pending_tools.""" + ext_agent = Agent( + name="ext-agent", + system_prompt="You are a test assistant.", + model=self.model, + toolkit=Toolkit(tools=[ExternalWeatherTool()]), + middlewares=[TracingMiddleware()], + ) + self.model.set_responses( + [_make_tool_call_response("e1", "Guangzhou")], + ) + self.exporter.clear() + + async for _ in ext_agent.reply_stream( + UserMsg(name="user", content="Weather in Guangzhou?"), + ): + pass + + first_spans = self._spans_by_name("invoke_agent") + self.assertEqual( + len(first_spans), + 1, + "Expected exactly one invoke_agent span", + ) + span_attrs = dict(first_spans[0].attributes or {}) + self.assertIn( + "agentscope.agent.external_execution_pending_tools", + span_attrs, + "First external-execution span should carry " + "agentscope.agent.external_execution_pending_tools", + ) + pending = json.loads( + span_attrs["agentscope.agent.external_execution_pending_tools"], + ) + self.assertEqual(pending, ["get_weather"]) + + async def test_external_execution_spans_share_reply_id(self) -> None: + """Both external-execution calls must share the same reply_id.""" + ext_agent = Agent( + name="ext-agent-rid", + system_prompt="You are a test assistant.", + model=self.model, + toolkit=Toolkit(tools=[ExternalWeatherTool()]), + middlewares=[TracingMiddleware()], + ) + + # First call: agent requires external execution + self.model.set_responses( + [_make_tool_call_response("ex-r1", "Nanjing")], + ) + self.exporter.clear() + + require_ext_event = None + async for evt in ext_agent.reply_stream( + UserMsg(name="user", content="Weather in Nanjing?"), + ): + if isinstance(evt, RequireExternalExecutionEvent): + require_ext_event = evt + + self.assertIsNotNone( + require_ext_event, + "Expected RequireExternalExecutionEvent", + ) + + first_spans = self._spans_by_name("invoke_agent") + first_reply_ids = { + dict(s.attributes or {}).get("agentscope.agent.reply_id") + for s in first_spans + if "agentscope.agent.reply_id" in (s.attributes or {}) + } + self.assertEqual(len(first_reply_ids), 1) + reply_id_first = next(iter(first_reply_ids)) + + # Second call: inject external result + self.model.set_responses( + [_make_text_response("Nanjing: clear, 18°C.")], + ) + self.exporter.clear() + + ext_result = ExternalExecutionResultEvent( + reply_id=require_ext_event.reply_id, + execution_results=[ + ToolResultBlock( + id=require_ext_event.tool_calls[0].id, + name="get_weather", + output="Nanjing: clear, 18°C.", + state=ToolResultState.SUCCESS, + ), + ], + ) + await ext_agent.reply(inputs=ext_result) + + second_spans = self._spans_by_name("invoke_agent") + second_reply_ids = { + dict(s.attributes or {}).get("agentscope.agent.reply_id") + for s in second_spans + if "agentscope.agent.reply_id" in (s.attributes or {}) + } + self.assertEqual(len(second_reply_ids), 1) + reply_id_second = next(iter(second_reply_ids)) + + self.assertEqual( + reply_id_first, + reply_id_second, + "Both external-execution calls must share the same reply_id", + ) + + async def test_external_execution_second_call_has_synthetic_tool_span( + self, + ) -> None: + """Second call with ExternalExecutionResultEvent must produce + execute_tool span.""" + ext_agent = Agent( + name="ext-agent2", + system_prompt="You are a test assistant.", + model=self.model, + toolkit=Toolkit(tools=[ExternalWeatherTool()]), + middlewares=[TracingMiddleware()], + ) + + # First call + self.model.set_responses( + [_make_tool_call_response("e2", "Shenzhen")], + ) + self.exporter.clear() + + require_ext_event = None + async for evt in ext_agent.reply_stream( + UserMsg(name="user", content="Weather in Shenzhen?"), + ): + if isinstance(evt, RequireExternalExecutionEvent): + require_ext_event = evt + + self.assertIsNotNone(require_ext_event) + + # Second call: inject external result + self.model.set_responses( + [_make_text_response("Shenzhen: warm, 28°C.")], + ) + self.exporter.clear() + + ext_result = ExternalExecutionResultEvent( + reply_id=require_ext_event.reply_id, + execution_results=[ + ToolResultBlock( + id=require_ext_event.tool_calls[0].id, + name="get_weather", + output="Shenzhen: warm, 28°C.", + state=ToolResultState.SUCCESS, + ), + ], + ) + await ext_agent.reply(inputs=ext_result) + + tool_spans = self._spans_by_name("execute_tool") + self.assertEqual( + len(tool_spans), + 1, + "Expected exactly one synthetic execute_tool span", + ) + span_attrs = dict(tool_spans[0].attributes or {}) + self.assertEqual( + span_attrs.get("agentscope.agent.is_external_execution"), + True, + "Synthetic execute_tool span should have " + "is_external_execution=True", + ) + self.assertEqual(span_attrs.get("gen_ai.tool.name"), "get_weather") + + async def test_external_execution_second_call_has_incoming_event_type( + self, + ) -> None: + """Second call span must have + incoming_event_type=external_execution_result.""" + ext_agent = Agent( + name="ext-agent3", + system_prompt="You are a test assistant.", + model=self.model, + toolkit=Toolkit(tools=[ExternalWeatherTool()]), + middlewares=[TracingMiddleware()], + ) + + # First call + self.model.set_responses( + [_make_tool_call_response("e3", "Chengdu")], + ) + self.exporter.clear() + + require_ext_event = None + async for evt in ext_agent.reply_stream( + UserMsg(name="user", content="Weather in Chengdu?"), + ): + if isinstance(evt, RequireExternalExecutionEvent): + require_ext_event = evt + + # Second call + self.model.set_responses([_make_text_response("Chengdu: cloudy.")]) + self.exporter.clear() + + ext_result = ExternalExecutionResultEvent( + reply_id=require_ext_event.reply_id, + execution_results=[ + ToolResultBlock( + id=require_ext_event.tool_calls[0].id, + name="get_weather", + output="Chengdu: cloudy, 15°C.", + state=ToolResultState.SUCCESS, + ), + ], + ) + await ext_agent.reply(inputs=ext_result) + + agent_spans = self._spans_by_name("invoke_agent") + self.assertEqual( + len(agent_spans), + 1, + "Expected exactly one invoke_agent span", + ) + span_attrs = dict(agent_spans[0].attributes or {}) + self.assertEqual( + span_attrs.get("agentscope.agent.incoming_event_type"), + "external_execution_result", + ) diff --git a/tests/tts_dashscope_test.py b/tests/tts_dashscope_test.py new file mode 100644 index 0000000000000000000000000000000000000000..23e28d5714c48c50a02215eaf159add5073805e7 --- /dev/null +++ b/tests/tts_dashscope_test.py @@ -0,0 +1,1260 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Unit tests for the TTS module. + +Covers: + * ``TTSModelBase`` default no-op behaviour for ``connect`` / ``close`` / + ``push`` (so non-realtime subclasses needn't override them). + * ``DashScopeTTSModel`` non-streaming aggregation. + * ``DashScopeTTSModel`` streaming: incremental deltas and ``is_last`` + placement at the final chunk only. + * ``DashScopeRealtimeTTSModel`` connect / close / push / synthesize + lifecycle over a mocked WebSocket. + * ``DashScopeCosyVoiceRealtimeTTSModel`` connect / close / push / + synthesize lifecycle over a mocked SpeechSynthesizer. +""" +import base64 +import io +import wave +from typing import Any, AsyncGenerator +from unittest import IsolatedAsyncioTestCase +from unittest.mock import MagicMock, Mock, patch + +from agentscope.credential import DashScopeCredential +from agentscope.tts import ( + DashScopeCosyVoiceRealtimeTTSModel, + DashScopeTTSModel, + DashScopeRealtimeTTSModel, + TTSModelBase, + TTSResponse, +) + + +_MEDIA_TYPE = "audio/wav" +# DashScope TTS emits 24kHz / mono / 16-bit PCM; the WAV wrapping in the +# model layer uses the same parameters. +_TTS_SAMPLE_RATE = 24000 +_TTS_CHANNELS = 1 +_TTS_SAMPLE_WIDTH = 2 # bytes (= 16 bit) +_WAV_HEADER_LEN = 44 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_api_chunk( + data_bytes: bytes | None, + usage: Any = None, +) -> MagicMock: + """Build a chunk shaped like what dashscope.MultiModalConversation + yields. ``data_bytes=None`` represents a chunk with no output.""" + chunk = MagicMock() + chunk.usage = usage + if data_bytes is None: + chunk.output = None + return chunk + chunk.output = MagicMock() + chunk.output.audio = MagicMock() + chunk.output.audio.data = base64.b64encode(data_bytes).decode("ascii") + return chunk + + +def _make_usage( + input_tokens: int = 0, + output_tokens: int = 0, + characters: int = 0, +) -> MagicMock: + """Build a usage object shaped like what the DashScope API returns.""" + usage = MagicMock() + usage.input_tokens = input_tokens + usage.output_tokens = output_tokens + usage.characters = characters + return usage + + +def _make_api_generator(chunks: list[bytes | None]) -> Any: + """Build a sync generator like ``MultiModalConversation.call`` returns. + The last chunk carries a usage object.""" + + def _gen() -> Any: + for i, data in enumerate(chunks): + is_last = i == len(chunks) - 1 + usage = _make_usage(characters=10) if is_last else None + yield _make_api_chunk(data, usage=usage) + + return _gen() + + +# --------------------------------------------------------------------------- +# TTSModelBase — default no-op surface for non-realtime subclasses +# --------------------------------------------------------------------------- + + +class _DummyTTS(TTSModelBase): + """Minimal subclass that implements only ``synthesize`` — exercises the + base class's no-op ``connect`` / ``close`` / ``push`` defaults.""" + + async def synthesize( + self, + text: str | None = None, + **kwargs: Any, + ) -> TTSResponse | AsyncGenerator[TTSResponse, None]: + del text, kwargs + return TTSResponse(content=None) + + +class _RealtimeDummyTTS(_DummyTTS): + """Realtime-flavoured dummy to assert ``__aenter__`` drives the lifecycle + hooks when ``realtime`` is True.""" + + realtime = True + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.connect_calls = 0 + self.close_calls = 0 + + async def connect(self) -> None: + self.connect_calls += 1 + + async def close(self) -> None: + self.close_calls += 1 + + +def _make_dummy(cls: type = _DummyTTS) -> TTSModelBase: + return cls( + credential=DashScopeCredential(api_key="test"), + model="x", + stream=False, + ) + + +class TestTTSModelBaseDefaults(IsolatedAsyncioTestCase): + """The base class supplies safe no-op defaults so non-realtime subclasses + don't need to implement realtime-only hooks.""" + + async def test_default_connect_close_noop(self) -> None: + """Default connect/close return without raising.""" + model = _make_dummy() + await model.connect() + await model.close() + + async def test_default_push_returns_empty(self) -> None: + """Default push returns an empty TTSResponse rather than raising, + so a misuse on a non-realtime model degrades gracefully.""" + model = _make_dummy() + resp = await model.push("ignored") + self.assertIsNone(resp.content) + + async def test_aenter_skips_hooks_for_non_realtime(self) -> None: + """``async with`` on a non-realtime model must not invoke connect/ + close (gated by ``realtime``).""" + model = _make_dummy(_DummyTTS) + async with model as m: + self.assertIs(m, model) + + async def test_aenter_invokes_hooks_for_realtime(self) -> None: + """For ``realtime=True`` subclasses, connect/close fire on + enter/exit exactly once.""" + model = _make_dummy(_RealtimeDummyTTS) + async with model: + self.assertEqual(model.connect_calls, 1) + self.assertEqual(model.close_calls, 0) + self.assertEqual(model.close_calls, 1) + + +# --------------------------------------------------------------------------- +# DashScopeTTSModel — non-streaming and streaming +# --------------------------------------------------------------------------- + + +def _parse_wav_payload(wav_bytes: bytes) -> bytes: + """Decode a full WAV file and return its raw PCM frames.""" + with wave.open(io.BytesIO(wav_bytes), "rb") as wav: + return wav.readframes(wav.getnframes()) + + +class TestDashScopeTTSModel(IsolatedAsyncioTestCase): + """The unittests for DashScope TTS model (non-realtime).""" + + def setUp(self) -> None: + """Set up the test case.""" + self.patcher = patch("dashscope.MultiModalConversation") + self.mock_mmc = self.patcher.start() + + def tearDown(self) -> None: + """Tear down the test case.""" + self.patcher.stop() + + def _make_model(self, stream: bool = False) -> DashScopeTTSModel: + """Create a DashScopeTTSModel with test credentials.""" + return DashScopeTTSModel( + credential=DashScopeCredential(api_key="test"), + model="qwen3-tts-flash", + parameters=DashScopeTTSModel.Parameters(voice="Cherry"), + stream=stream, + ) + + # -- non-streaming -- + + async def test_aggregates_chunks(self) -> None: + """All API chunks are aggregated into one self-contained WAV.""" + self.mock_mmc.call.return_value = _make_api_generator( + [b"AAAA", b"BBBB", b"CCCC"], + ) + model = self._make_model(stream=False) + + result = await model.synthesize("Hello world") + + self.assertIsInstance(result, TTSResponse) + self.assertEqual(result.content.source.media_type, _MEDIA_TYPE) + wav_bytes = base64.b64decode(result.content.source.data) + with wave.open(io.BytesIO(wav_bytes), "rb") as wav: + self.assertEqual( + { + "framerate": wav.getframerate(), + "channels": wav.getnchannels(), + "sampwidth": wav.getsampwidth(), + "frames": wav.readframes(wav.getnframes()), + }, + { + "framerate": _TTS_SAMPLE_RATE, + "channels": _TTS_CHANNELS, + "sampwidth": _TTS_SAMPLE_WIDTH, + "frames": b"AAAABBBBCCCC", + }, + ) + self.assertTrue(result.is_last) + + async def test_none_short_circuits(self) -> None: + """``synthesize(None)`` returns an empty response without touching + the API.""" + model = self._make_model(stream=False) + + result = await model.synthesize(None) + + self.assertIsNone(result.content) + self.mock_mmc.call.assert_not_called() + + async def test_empty_string_short_circuits(self) -> None: + """``synthesize("")`` returns an empty response without touching + the API.""" + model = self._make_model(stream=False) + + result = await model.synthesize("") + + self.assertIsNone(result.content) + self.mock_mmc.call.assert_not_called() + + async def test_skips_empty_chunks(self) -> None: + """Chunks without ``output`` are ignored during aggregation.""" + self.mock_mmc.call.return_value = _make_api_generator( + [None, b"AAAA", None, b"BBBB"], + ) + model = self._make_model(stream=False) + + result = await model.synthesize("Hello world") + + wav_bytes = base64.b64decode(result.content.source.data) + self.assertEqual(_parse_wav_payload(wav_bytes), b"AAAABBBB") + + # -- streaming -- + + async def test_incremental_deltas(self) -> None: + """Each API chunk yields one TTSResponse with incremental PCM.""" + self.mock_mmc.call.return_value = _make_api_generator( + [b"AAAA", b"BBBB", b"CCCC"], + ) + model = self._make_model(stream=True) + + gen = await model.synthesize("Hello world") + chunks = [c async for c in gen] + + payloads = [base64.b64decode(c.content.source.data) for c in chunks] + + self.assertTrue(payloads[0].startswith(b"RIFF")) + self.assertEqual(payloads[0][8:12], b"WAVE") + self.assertEqual(payloads[0][_WAV_HEADER_LEN:], b"AAAA") + self.assertEqual(payloads[1], b"BBBB") + self.assertEqual(payloads[2], b"CCCC") + + self.assertEqual( + [c.is_last for c in chunks], + [False, False, True], + ) + self.assertEqual( + [c.content.source.media_type for c in chunks], + [_MEDIA_TYPE, _MEDIA_TYPE, _MEDIA_TYPE], + ) + + async def test_single_chunk_marked_last(self) -> None: + """A lone audio chunk is flagged ``is_last=True`` with a streaming + WAV header.""" + self.mock_mmc.call.return_value = _make_api_generator([b"ONLYCHUNK"]) + model = self._make_model(stream=True) + + gen = await model.synthesize("Hello world") + chunks = [c async for c in gen] + + self.assertEqual(len(chunks), 1) + self.assertTrue(chunks[0].is_last) + payload = base64.b64decode(chunks[0].content.source.data) + self.assertTrue(payload.startswith(b"RIFF")) + self.assertEqual(payload[_WAV_HEADER_LEN:], b"ONLYCHUNK") + + async def test_empty_stream_yields_terminal(self) -> None: + """When the API yields no audio, the generator emits a terminal + sentinel so consumers can detect EOS.""" + self.mock_mmc.call.return_value = _make_api_generator([None, None]) + model = self._make_model(stream=True) + + gen = await model.synthesize("Hello world") + chunks = [c async for c in gen] + + self.assertEqual(len(chunks), 1) + self.assertIsNone(chunks[0].content) + self.assertTrue(chunks[0].is_last) + + +# --------------------------------------------------------------------------- +# DashScopeRealtimeTTSModel — realtime push / synthesize lifecycle +# --------------------------------------------------------------------------- + + +class TestDashScopeRealtimeTTSModel( # pylint: disable=too-many-public-methods + IsolatedAsyncioTestCase, +): + """The unittests for DashScope Realtime TTS model.""" + + def setUp(self) -> None: + self.mock_modules = self._create_mock_dashscope_modules() + self.mock_client = self._create_mock_tts_client() + mock_tts_class = Mock(return_value=self.mock_client) + self.mock_modules[ + "dashscope.audio.qwen_tts_realtime" + ].QwenTtsRealtime = mock_tts_class + + @staticmethod + def _create_mock_dashscope_modules() -> dict: + mock_qwen_tts_realtime = MagicMock() + mock_qwen_tts_realtime.QwenTtsRealtime = Mock + mock_qwen_tts_realtime.QwenTtsRealtimeCallback = Mock + + mock_audio = MagicMock() + mock_audio.qwen_tts_realtime = mock_qwen_tts_realtime + + mock_dashscope = MagicMock() + mock_dashscope.api_key = None + mock_dashscope.audio = mock_audio + + return { + "dashscope": mock_dashscope, + "dashscope.audio": mock_audio, + "dashscope.audio.qwen_tts_realtime": mock_qwen_tts_realtime, + } + + @staticmethod + def _create_mock_tts_client() -> Mock: + client = Mock() + client.connect = Mock() + client.close = Mock() + client.finish = Mock() + client.commit = Mock() + client.update_session = Mock() + client.append_text = Mock() + return client + + def _make_model(self, **kwargs: Any) -> DashScopeRealtimeTTSModel: + defaults: dict[str, Any] = { + "credential": DashScopeCredential(api_key="test"), + "model": "qwen3-tts-flash-realtime", + "stream": True, + "max_retries": 1, + "retry_delay": 0.0, + } + defaults.update(kwargs) + return DashScopeRealtimeTTSModel(**defaults) + + def _mock_synthesize_callback(self, model: Any) -> None: + """Set up callback mocks so ``synthesize()`` doesn't block.""" + model._callback.finish_event = Mock() + model._callback.finish_event.wait = Mock() + model._callback.has_audio_data = Mock(return_value=True) + model._callback.get_audio_response = Mock( + return_value=TTSResponse(content=None), + ) + + # -- connect / close -- + + async def test_async_context_manager(self) -> None: + """``async with`` triggers connect on enter and close on exit.""" + with patch.dict("sys.modules", self.mock_modules): + model = self._make_model() + async with model: + self.assertTrue(model._connected) + self.assertFalse(model._connected) + + # -- push -- + + async def test_push_incremental_deltas(self) -> None: + """Consecutive deltas are each forwarded verbatim.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model() as model: + await model.push("Hello") + await model.push(" world") + + self.assertEqual( + self.mock_client.append_text.call_count, + 2, + ) + self.mock_client.append_text.assert_any_call("Hello") + self.mock_client.append_text.assert_any_call(" world") + self.assertEqual(model._accumulated_text, "Hello world") + + async def test_push_returns_audio_when_available(self) -> None: + """push() returns audio data from the callback when available.""" + mock_audio_data = base64.b64encode(b"PCMDATA").decode("ascii") + + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model() as model: + model._callback.get_audio_response = Mock( + return_value=TTSResponse( + content=MagicMock( + source=MagicMock( + data=mock_audio_data, + media_type=_MEDIA_TYPE, + ), + ), + ), + ) + res = await model.push("Hello") + + self.assertIsNotNone(res.content) + self.assertEqual(res.content.source.data, mock_audio_data) + + async def test_push_cold_start_buffers_across_deltas(self) -> None: + """Multiple small deltas are buffered until cold_start_length is + met, then flushed as a single ``append_text`` call.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model(cold_start_length=10) as model: + await model.push("Hi") + await model.push(" there") + self.mock_client.append_text.assert_not_called() + self.assertFalse(model._cold_start_done) + + await model.push(" friend!") + self.mock_client.append_text.assert_called_once_with( + "Hi there friend!", + ) + self.assertTrue(model._cold_start_done) + + async def test_push_after_cold_start_forwards_directly(self) -> None: + """Once cold start is done, subsequent deltas bypass the buffer.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model(cold_start_length=3) as model: + await model.push("Hello") + self.mock_client.append_text.assert_called_with("Hello") + + await model.push(" world") + self.mock_client.append_text.assert_called_with(" world") + self.assertEqual( + self.mock_client.append_text.call_count, + 2, + ) + + async def test_push_cold_start_words_buffers(self) -> None: + """Deltas below cold_start_words are buffered.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model(cold_start_words=3) as model: + await model.push("Hello") + await model.push(" world") + + self.mock_client.append_text.assert_not_called() + self.assertFalse(model._cold_start_done) + + # -- synthesize -- + + async def test_synthesize_commits_and_finishes(self) -> None: + """synthesize(text=...) appends text, commits, and finishes.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model() as model: + self._mock_synthesize_callback(model) + + await model.synthesize(text="Hello") + + self.mock_client.append_text.assert_called_once_with("Hello") + self.mock_client.commit.assert_called_once() + self.mock_client.finish.assert_called_once() + + async def test_synthesize_appends_extra_text(self) -> None: + """synthesize(text=...) after push appends the extra text.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model() as model: + model._callback.get_audio_response = Mock( + return_value=TTSResponse(content=None), + ) + await model.push("Hello") + self.mock_client.append_text.reset_mock() + + self._mock_synthesize_callback(model) + await model.synthesize(text=" world") + + self.mock_client.append_text.assert_called_once_with( + " world", + ) + self.mock_client.commit.assert_called_once() + + async def test_synthesize_no_text_drain(self) -> None: + """synthesize(None) after push commits without extra append.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model() as model: + model._callback.get_audio_response = Mock( + return_value=TTSResponse(content=None), + ) + await model.push("Hello") + self.mock_client.append_text.reset_mock() + + self._mock_synthesize_callback(model) + await model.synthesize() + + self.mock_client.append_text.assert_not_called() + self.mock_client.commit.assert_called_once() + + async def test_synthesize_stream_returns_generator(self) -> None: + """stream=True returns an async generator with is_last.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model(stream=True) as model: + self._mock_synthesize_callback(model) + + async def mock_chunks() -> AsyncGenerator[TTSResponse, None]: + yield TTSResponse(content=None, is_last=True) + + model._callback.get_audio_chunks = mock_chunks + + gen = await model.synthesize(text="Hello") + chunks = [c async for c in gen] + + self.assertTrue(len(chunks) >= 1) + self.assertTrue(chunks[-1].is_last) + + async def test_synthesize_flushes_cold_start_buffer(self) -> None: + """If cold start was never met during push, synthesize flushes the + buffered text before committing.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model(cold_start_length=100) as model: + await model.push("Hi") + await model.push(" there") + self.mock_client.append_text.assert_not_called() + + self._mock_synthesize_callback(model) + await model.synthesize() + + self.mock_client.append_text.assert_called_once_with( + "Hi there", + ) + self.mock_client.commit.assert_called_once() + + async def test_synthesize_no_audio_raises_after_retries(self) -> None: + """RuntimeError after max_retries with no audio received.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model( + max_retries=1, + retry_delay=0.0, + ) as model: + model._callback.finish_event = Mock() + model._callback.finish_event.wait = Mock() + model._callback.has_audio_data = Mock(return_value=False) + + with self.assertRaises(RuntimeError): + await model.synthesize(text="Hello") + + # -- callback -- + + async def test_get_audio_chunks_prepends_header_without_prior_push( + self, + ) -> None: + """get_audio_chunks() prepends a WAV header when push() has not + yet consumed any bytes (_consumed == 0).""" + with patch.dict("sys.modules", self.mock_modules): + from agentscope.tts._dashscope._realtime_model import ( + _make_callback_class, + ) + + callback_cls = _make_callback_class() + cb = callback_cls() + + audio_b64 = base64.b64encode(b"ALLPCM").decode() + cb.on_event({"type": "response.audio.delta", "delta": audio_b64}) + cb.on_event({"type": "session.finished"}) + + chunks = [c async for c in cb.get_audio_chunks()] + raw_payloads = [ + base64.b64decode(c.content.source.data) + for c in chunks + if c.content is not None + ] + + self.assertTrue(len(raw_payloads) > 0) + self.assertTrue(raw_payloads[0].startswith(b"RIFF")) + self.assertTrue(raw_payloads[0].endswith(b"ALLPCM")) + + async def test_get_audio_chunks_skips_header_after_partial_push( + self, + ) -> None: + """get_audio_chunks() does NOT prepend a second WAV header when + push() already consumed and sent the first chunk with a header + (_consumed > 0).""" + with patch.dict("sys.modules", self.mock_modules): + from agentscope.tts._dashscope._realtime_model import ( + _make_callback_class, + ) + + callback_cls = _make_callback_class() + cb = callback_cls() + + # Simulate push() consuming the first chunk (with WAV header) + first_b64 = base64.b64encode(b"FIRSTPCM").decode() + cb.on_event( + {"type": "response.audio.delta", "delta": first_b64}, + ) + first_resp = cb.get_audio_response(block=False) + self.assertIsNotNone(first_resp.content) + self.assertGreater(cb._consumed, 0) + + # More data arrives and session finishes + more_b64 = base64.b64encode(b"MOREPCM").decode() + cb.on_event( + {"type": "response.audio.delta", "delta": more_b64}, + ) + cb.on_event({"type": "session.finished"}) + + chunks = [c async for c in cb.get_audio_chunks()] + raw_payloads = [ + base64.b64decode(c.content.source.data) + for c in chunks + if c.content is not None + ] + + self.assertTrue(len(raw_payloads) > 0) + # No second RIFF header should appear mid-stream + self.assertFalse(raw_payloads[0].startswith(b"RIFF")) + self.assertEqual(raw_payloads[0], b"MOREPCM") + + async def test_get_audio_response_includes_header_without_prior_push( + self, + ) -> None: + """get_audio_response() prepends a WAV header when no data has been + consumed yet (_consumed == 0).""" + with patch.dict("sys.modules", self.mock_modules): + from agentscope.tts._dashscope._realtime_model import ( + _make_callback_class, + ) + + callback_cls = _make_callback_class() + cb = callback_cls() + + audio_b64 = base64.b64encode(b"ALLPCM").decode() + cb.on_event({"type": "response.audio.delta", "delta": audio_b64}) + cb.on_event({"type": "session.finished"}) + + resp = cb.get_audio_response(block=True) + raw = base64.b64decode(resp.content.source.data) + + self.assertTrue(raw.startswith(b"RIFF")) + self.assertTrue(raw.endswith(b"ALLPCM")) + + async def test_get_audio_response_skips_header_after_partial_push( + self, + ) -> None: + """get_audio_response() does NOT prepend a second WAV header when + push() already consumed and sent the first chunk with a header + (_consumed > 0).""" + with patch.dict("sys.modules", self.mock_modules): + from agentscope.tts._dashscope._realtime_model import ( + _make_callback_class, + ) + + callback_cls = _make_callback_class() + cb = callback_cls() + + first_b64 = base64.b64encode(b"FIRSTPCM").decode() + cb.on_event( + {"type": "response.audio.delta", "delta": first_b64}, + ) + first_resp = cb.get_audio_response(block=False) + self.assertIsNotNone(first_resp.content) + self.assertGreater(cb._consumed, 0) + + more_b64 = base64.b64encode(b"MOREPCM").decode() + cb.on_event( + {"type": "response.audio.delta", "delta": more_b64}, + ) + cb.on_event({"type": "session.finished"}) + + resp = cb.get_audio_response(block=True) + raw = base64.b64decode(resp.content.source.data) + + self.assertFalse(raw.startswith(b"RIFF")) + self.assertEqual(raw, b"MOREPCM") + + +# --------------------------------------------------------------------------- +# DashScopeCosyVoiceRealtimeTTSModel — realtime push / synthesize lifecycle +# --------------------------------------------------------------------------- + + +# pylint: disable=too-many-public-methods +class TestDashScopeCosyVoiceRealtimeTTSModel( + IsolatedAsyncioTestCase, +): + """Unit tests for the CosyVoice Realtime TTS model.""" + + def setUp(self) -> None: + self.mock_modules = self._create_mock_cosyvoice_modules() + self.mock_synthesizer = self._create_mock_synthesizer() + self.mock_modules["dashscope.audio.tts_v2"].SpeechSynthesizer = Mock( + return_value=self.mock_synthesizer, + ) + + @staticmethod + def _create_mock_cosyvoice_modules() -> dict: + mock_result_callback = Mock + mock_audio_format = MagicMock() + mock_audio_format.PCM_24000HZ_MONO_16BIT = "pcm_24000hz_mono_16bit" + + mock_tts_v2 = MagicMock() + mock_tts_v2.ResultCallback = mock_result_callback + mock_tts_v2.AudioFormat = mock_audio_format + mock_tts_v2.SpeechSynthesizer = Mock + + mock_audio = MagicMock() + mock_audio.tts_v2 = mock_tts_v2 + + mock_dashscope = MagicMock() + mock_dashscope.api_key = None + mock_dashscope.audio = mock_audio + + return { + "dashscope": mock_dashscope, + "dashscope.audio": mock_audio, + "dashscope.audio.tts_v2": mock_tts_v2, + } + + @staticmethod + def _create_mock_synthesizer() -> Mock: + synth = Mock() + synth.streaming_call = Mock() + synth.streaming_complete = Mock() + synth.close = Mock() + return synth + + def _make_model( + self, + **kwargs: Any, + ) -> DashScopeCosyVoiceRealtimeTTSModel: + defaults: dict[str, Any] = { + "credential": DashScopeCredential(api_key="test"), + "model": "cosyvoice-v3-plus", + "stream": True, + "max_retries": 1, + "retry_delay": 0.0, + } + defaults.update(kwargs) + return DashScopeCosyVoiceRealtimeTTSModel(**defaults) + + def _mock_synthesize_callback(self, model: Any) -> None: + """Set up callback mocks so synthesize() doesn't block.""" + model._callback.finish_event = Mock() + model._callback.finish_event.wait = Mock(return_value=True) + model._callback.has_audio_data = Mock(return_value=True) + model._callback.get_audio_response = Mock( + return_value=TTSResponse(content=None), + ) + + # -- connect / close -- + + async def test_async_context_manager(self) -> None: + """async with triggers connect on enter and close on exit.""" + with patch.dict("sys.modules", self.mock_modules): + model = self._make_model() + async with model: + self.assertTrue(model._connected) + self.assertFalse(model._connected) + + # -- push -- + + async def test_push_incremental_deltas(self) -> None: + """Consecutive deltas are each forwarded via streaming_call.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model() as model: + model._callback.get_audio_response = Mock( + return_value=TTSResponse(content=None), + ) + await model.push("Hello") + await model.push(" world") + + self.assertEqual( + self.mock_synthesizer.streaming_call.call_count, + 2, + ) + self.mock_synthesizer.streaming_call.assert_any_call("Hello") + self.mock_synthesizer.streaming_call.assert_any_call(" world") + + async def test_push_returns_audio_when_available(self) -> None: + """push() returns audio data from the callback when available.""" + mock_audio_data = base64.b64encode(b"PCMDATA").decode("ascii") + + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model() as model: + model._callback.get_audio_response = Mock( + return_value=TTSResponse( + content=MagicMock( + source=MagicMock( + data=mock_audio_data, + media_type=_MEDIA_TYPE, + ), + ), + ), + ) + res = await model.push("Hello") + + self.assertIsNotNone(res.content) + self.assertEqual(res.content.source.data, mock_audio_data) + + async def test_push_cold_start_buffers_across_deltas(self) -> None: + """Multiple small deltas are buffered until cold_start_length is + met, then flushed as a single streaming_call.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model(cold_start_length=10) as model: + model._callback.get_audio_response = Mock( + return_value=TTSResponse(content=None), + ) + await model.push("Hi") + await model.push(" there") + self.mock_synthesizer.streaming_call.assert_not_called() + self.assertFalse(model._cold_start_done) + + await model.push(" friend!") + self.mock_synthesizer.streaming_call.assert_called_once_with( + "Hi there friend!", + ) + self.assertTrue(model._cold_start_done) + + async def test_push_after_cold_start_forwards_directly(self) -> None: + """Once cold start is done, subsequent deltas bypass the buffer.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model(cold_start_length=3) as model: + model._callback.get_audio_response = Mock( + return_value=TTSResponse(content=None), + ) + await model.push("Hello") + self.mock_synthesizer.streaming_call.assert_called_with( + "Hello", + ) + + await model.push(" world") + self.mock_synthesizer.streaming_call.assert_called_with( + " world", + ) + self.assertEqual( + self.mock_synthesizer.streaming_call.call_count, + 2, + ) + + async def test_push_cold_start_words_buffers(self) -> None: + """Deltas below cold_start_words are buffered.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model(cold_start_words=3) as model: + model._callback.get_audio_response = Mock( + return_value=TTSResponse(content=None), + ) + await model.push("Hello") + await model.push(" world") + + self.mock_synthesizer.streaming_call.assert_not_called() + self.assertFalse(model._cold_start_done) + + async def test_push_exception_returns_empty(self) -> None: + """push() returns empty response if streaming_call raises.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model() as model: + self.mock_synthesizer.streaming_call.side_effect = Exception( + "connection error", + ) + res = await model.push("Hello") + + self.assertIsNone(res.content) + + # -- synthesize -- + + async def test_synthesize_calls_streaming_complete(self) -> None: + """synthesize() calls streaming_call and streaming_complete.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model() as model: + self._mock_synthesize_callback(model) + + await model.synthesize(text="Hello") + + self.mock_synthesizer.streaming_call.assert_called_once_with( + "Hello", + ) + self.mock_synthesizer.streaming_complete.assert_called_once() + + async def test_synthesize_appends_extra_text(self) -> None: + """synthesize(text=...) after push appends the extra text.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model() as model: + model._callback.get_audio_response = Mock( + return_value=TTSResponse(content=None), + ) + await model.push("Hello") + self.mock_synthesizer.streaming_call.reset_mock() + + self._mock_synthesize_callback(model) + await model.synthesize(text=" world") + + self.mock_synthesizer.streaming_call.assert_called_once_with( + " world", + ) + self.mock_synthesizer.streaming_complete.assert_called_once() + + async def test_synthesize_no_text_drain(self) -> None: + """synthesize(None) after push calls streaming_complete without + extra streaming_call.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model() as model: + model._callback.get_audio_response = Mock( + return_value=TTSResponse(content=None), + ) + await model.push("Hello") + self.mock_synthesizer.streaming_call.reset_mock() + + self._mock_synthesize_callback(model) + await model.synthesize() + + self.mock_synthesizer.streaming_call.assert_not_called() + self.mock_synthesizer.streaming_complete.assert_called_once() + + async def test_synthesize_stream_returns_generator(self) -> None: + """stream=True returns an async generator with is_last.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model(stream=True) as model: + self._mock_synthesize_callback(model) + + async def mock_chunks() -> AsyncGenerator[TTSResponse, None]: + yield TTSResponse(content=None, is_last=True) + + model._callback.get_audio_chunks = mock_chunks + + gen = await model.synthesize(text="Hello") + chunks = [c async for c in gen] + + self.assertTrue(len(chunks) >= 1) + self.assertTrue(chunks[-1].is_last) + + async def test_synthesize_non_stream_returns_single(self) -> None: + """stream=False returns a single TTSResponse.""" + mock_audio_data = base64.b64encode(b"PCMDATA").decode("ascii") + + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model(stream=False) as model: + self._mock_synthesize_callback(model) + model._callback.get_audio_response = Mock( + return_value=TTSResponse( + content=MagicMock( + source=MagicMock( + data=mock_audio_data, + media_type=_MEDIA_TYPE, + ), + ), + ), + ) + + res = await model.synthesize(text="Hello") + + self.assertIsInstance(res, TTSResponse) + self.assertIsNotNone(res.content) + + async def test_synthesize_flushes_cold_start_buffer(self) -> None: + """If cold start was never met during push, synthesize flushes the + buffered text before calling streaming_complete.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model(cold_start_length=100) as model: + model._callback.get_audio_response = Mock( + return_value=TTSResponse(content=None), + ) + await model.push("Hi") + await model.push(" there") + self.mock_synthesizer.streaming_call.assert_not_called() + + self._mock_synthesize_callback(model) + await model.synthesize() + + self.mock_synthesizer.streaming_call.assert_called_once_with( + "Hi there", + ) + self.mock_synthesizer.streaming_complete.assert_called_once() + + async def test_synthesize_empty_text_short_circuits(self) -> None: + """synthesize() with no accumulated text returns empty.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model(stream=True) as model: + gen = await model.synthesize() + chunks = [c async for c in gen] + + self.assertEqual(len(chunks), 1) + self.assertIsNone(chunks[0].content) + self.assertTrue(chunks[0].is_last) + self.mock_synthesizer.streaming_complete.assert_not_called() + + async def test_synthesize_no_audio_raises_after_retries(self) -> None: + """RuntimeError after max_retries with no audio received.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model( + max_retries=1, + retry_delay=0.0, + ) as model: + model._callback.finish_event = Mock() + model._callback.finish_event.wait = Mock(return_value=True) + model._callback.has_audio_data = Mock(return_value=False) + + with self.assertRaises(RuntimeError): + await model.synthesize(text="Hello") + + async def test_synthesize_timeout_raises(self) -> None: + """RuntimeError when finish_event.wait times out.""" + with patch.dict("sys.modules", self.mock_modules): + async with self._make_model( + max_retries=1, + retry_delay=0.0, + ) as model: + model._callback.finish_event = Mock() + model._callback.finish_event.wait = Mock(return_value=False) + + with self.assertRaises(RuntimeError) as ctx: + await model.synthesize(text="Hello") + + self.assertIn("timed out", str(ctx.exception)) + + # -- callback -- + + async def test_callback_on_complete_sets_finish_event(self) -> None: + """on_complete() correctly sets finish_event and chunk_event.""" + with patch.dict("sys.modules", self.mock_modules): + from agentscope.tts._dashscope._cosyvoice_realtime_model import ( + _make_cosyvoice_callback_class, + ) + + callback_cls = _make_cosyvoice_callback_class() + cb = callback_cls() + + self.assertFalse(cb.finish_event.is_set()) + self.assertFalse(cb.chunk_event.is_set()) + + cb.on_complete() + + self.assertTrue(cb.finish_event.is_set()) + self.assertTrue(cb.chunk_event.is_set()) + + async def test_callback_on_data_accumulates(self) -> None: + """on_data() accumulates PCM bytes and signals chunk_event.""" + with patch.dict("sys.modules", self.mock_modules): + from agentscope.tts._dashscope._cosyvoice_realtime_model import ( + _make_cosyvoice_callback_class, + ) + + callback_cls = _make_cosyvoice_callback_class() + cb = callback_cls() + + cb.on_data(b"AAAA") + cb.on_data(b"BBBB") + + self.assertEqual(bytes(cb._pcm_bytes), b"AAAABBBB") + self.assertTrue(cb.chunk_event.is_set()) + + async def test_callback_take_delta_incremental(self) -> None: + """_take_delta() returns only new bytes since last call.""" + with patch.dict("sys.modules", self.mock_modules): + from agentscope.tts._dashscope._cosyvoice_realtime_model import ( + _make_cosyvoice_callback_class, + ) + + callback_cls = _make_cosyvoice_callback_class() + cb = callback_cls() + + cb.on_data(b"AAAA") + delta1 = cb._take_delta(header=False) + self.assertEqual(delta1, b"AAAA") + + cb.on_data(b"BBBB") + delta2 = cb._take_delta(header=False) + self.assertEqual(delta2, b"BBBB") + + delta3 = cb._take_delta(header=False) + self.assertIsNone(delta3) + + async def test_callback_take_delta_with_header(self) -> None: + """_take_delta(header=True) prepends WAV header to first chunk.""" + with patch.dict("sys.modules", self.mock_modules): + from agentscope.tts._dashscope._cosyvoice_realtime_model import ( + _make_cosyvoice_callback_class, + ) + + callback_cls = _make_cosyvoice_callback_class() + cb = callback_cls() + + cb.on_data(b"PCMPCM") + delta = cb._take_delta(header=True) + + self.assertTrue(delta.startswith(b"RIFF")) + self.assertIn(b"WAVE", delta[:12]) + self.assertTrue(delta.endswith(b"PCMPCM")) + + async def test_get_audio_chunks_prepends_header_without_prior_push( + self, + ) -> None: + """get_audio_chunks() prepends a WAV header when push() has not + yet consumed any bytes (_consumed == 0).""" + with patch.dict("sys.modules", self.mock_modules): + from agentscope.tts._dashscope._cosyvoice_realtime_model import ( + _make_cosyvoice_callback_class, + ) + + callback_cls = _make_cosyvoice_callback_class() + cb = callback_cls() + + cb.on_data(b"ALLPCM") + cb.on_complete() + + chunks = [c async for c in cb.get_audio_chunks()] + raw_payloads = [ + base64.b64decode(c.content.source.data) + for c in chunks + if c.content is not None + ] + + self.assertTrue(len(raw_payloads) > 0) + self.assertTrue(raw_payloads[0].startswith(b"RIFF")) + self.assertTrue(raw_payloads[0].endswith(b"ALLPCM")) + + async def test_get_audio_chunks_skips_header_after_partial_push( + self, + ) -> None: + """get_audio_chunks() does NOT prepend a second WAV header when + push() already consumed and sent the first chunk with a header + (_consumed > 0).""" + with patch.dict("sys.modules", self.mock_modules): + from agentscope.tts._dashscope._cosyvoice_realtime_model import ( + _make_cosyvoice_callback_class, + ) + + callback_cls = _make_cosyvoice_callback_class() + cb = callback_cls() + + # Simulate push() consuming the first chunk (with WAV header) + cb.on_data(b"FIRSTPCM") + first_resp = cb.get_audio_response(block=False) + self.assertIsNotNone(first_resp.content) + self.assertGreater(cb._consumed, 0) + + # More data arrives and synthesis completes + cb.on_data(b"MOREPCM") + cb.on_complete() + + chunks = [c async for c in cb.get_audio_chunks()] + raw_payloads = [ + base64.b64decode(c.content.source.data) + for c in chunks + if c.content is not None + ] + + self.assertTrue(len(raw_payloads) > 0) + # No second RIFF header should appear mid-stream + self.assertFalse(raw_payloads[0].startswith(b"RIFF")) + self.assertEqual(raw_payloads[0], b"MOREPCM") + + async def test_get_audio_response_includes_header_without_prior_push( + self, + ) -> None: + """get_audio_response() prepends a WAV header when no data has been + consumed yet (_consumed == 0).""" + with patch.dict("sys.modules", self.mock_modules): + from agentscope.tts._dashscope._cosyvoice_realtime_model import ( + _make_cosyvoice_callback_class, + ) + + callback_cls = _make_cosyvoice_callback_class() + cb = callback_cls() + + cb.on_data(b"ALLPCM") + cb.on_complete() + + resp = cb.get_audio_response(block=True) + raw = base64.b64decode(resp.content.source.data) + + self.assertTrue(raw.startswith(b"RIFF")) + self.assertTrue(raw.endswith(b"ALLPCM")) + + async def test_get_audio_response_skips_header_after_partial_push( + self, + ) -> None: + """get_audio_response() does NOT prepend a second WAV header when + push() already consumed and sent the first chunk with a header + (_consumed > 0).""" + with patch.dict("sys.modules", self.mock_modules): + from agentscope.tts._dashscope._cosyvoice_realtime_model import ( + _make_cosyvoice_callback_class, + ) + + callback_cls = _make_cosyvoice_callback_class() + cb = callback_cls() + + cb.on_data(b"FIRSTPCM") + first_resp = cb.get_audio_response(block=False) + self.assertIsNotNone(first_resp.content) + self.assertGreater(cb._consumed, 0) + + cb.on_data(b"MOREPCM") + cb.on_complete() + + resp = cb.get_audio_response(block=True) + raw = base64.b64decode(resp.content.source.data) + + self.assertFalse(raw.startswith(b"RIFF")) + self.assertEqual(raw, b"MOREPCM") + + async def test_callback_reset(self) -> None: + """reset() clears all state.""" + with patch.dict("sys.modules", self.mock_modules): + from agentscope.tts._dashscope._cosyvoice_realtime_model import ( + _make_cosyvoice_callback_class, + ) + + callback_cls = _make_cosyvoice_callback_class() + cb = callback_cls() + + cb.on_data(b"AAAA") + cb.on_complete() + cb.reset() + + self.assertFalse(cb.finish_event.is_set()) + self.assertFalse(cb.chunk_event.is_set()) + self.assertEqual(bytes(cb._pcm_bytes), b"") + self.assertEqual(cb._consumed, 0) + + # -- reconnect -- + + async def test_reconnect_recreates_synthesizer(self) -> None: + """_reconnect() closes old synthesizer and creates a new one.""" + with patch.dict("sys.modules", self.mock_modules): + model = self._make_model() + await model.connect() + old_callback = model._callback + + await model._reconnect() + + self.mock_synthesizer.close.assert_called_once() + self.assertTrue(model._connected) + self.assertIsNot(model._callback, old_callback) diff --git a/tests/tts_middleware_test.py b/tests/tts_middleware_test.py new file mode 100644 index 0000000000000000000000000000000000000000..219744303febd227d44577c26f4e4c309b7d9171 --- /dev/null +++ b/tests/tts_middleware_test.py @@ -0,0 +1,426 @@ +# -*- coding: utf-8 -*- +"""Unit tests for TTSMiddleware.""" +import base64 +from typing import Any, AsyncGenerator +from unittest import IsolatedAsyncioTestCase +from unittest.mock import AsyncMock, MagicMock + +from agentscope import set_id_factory +from agentscope.event import ( + DataBlockDeltaEvent, + DataBlockEndEvent, + DataBlockStartEvent, + TextBlockDeltaEvent, + TextBlockEndEvent, +) +from agentscope.message import Base64Source, DataBlock +from agentscope.middleware import TTSMiddleware +from agentscope.tts import TTSModelBase, TTSResponse + +_EXCLUDE = {"id", "created_at", "metadata"} + + +def _dump(evt: Any) -> dict: + """Dump an event to a dict excluding auto-generated fields.""" + return evt.model_dump(exclude=_EXCLUDE) + + +def _make_tts_response( + data: str, + media_type: str = "audio/wav", +) -> TTSResponse: + return TTSResponse( + content=DataBlock( + source=Base64Source(data=data, media_type=media_type), + ), + ) + + +def _make_agent_stub( + reply_id: str = "reply-1", + name: str = "agent", +) -> MagicMock: + agent = MagicMock() + agent.name = name + agent.state.reply_id = reply_id + return agent + + +class TestTTSMiddlewareNonRealtime(IsolatedAsyncioTestCase): + """Tests for non-realtime (non-streaming-input) TTS path.""" + + async def test_synthesize_on_text_block_end(self) -> None: + """After TextBlockEnd, synthesize() is called with accumulated text + and DATA_BLOCK_START + DATA_BLOCK_DELTA + DATA_BLOCK_END are emitted. + """ + audio_b64 = base64.b64encode(b"\x00\x01\x02").decode() + + tts = MagicMock(spec=TTSModelBase) + tts.realtime = False + tts.__aenter__ = AsyncMock(return_value=tts) + tts.__aexit__ = AsyncMock(return_value=None) + tts.synthesize = AsyncMock( + return_value=_make_tts_response(audio_b64), + ) + + middleware = TTSMiddleware(tts_model=tts) + agent = _make_agent_stub() + + upstream_events = [ + TextBlockDeltaEvent( + reply_id="reply-1", + block_id="blk-1", + delta="Hello ", + ), + TextBlockDeltaEvent( + reply_id="reply-1", + block_id="blk-1", + delta="world", + ), + TextBlockEndEvent(reply_id="reply-1", block_id="blk-1"), + ] + + async def next_handler(**_kwargs: Any) -> AsyncGenerator: + for evt in upstream_events: + yield evt + + emitted = [] + async for evt in middleware.on_reply(agent, {}, next_handler): + emitted.append(evt) + + # Upstream events are passed through + self.assertEqual( + [_dump(e) for e in emitted[:3]], + [_dump(e) for e in upstream_events], + ) + + # After TextBlockEnd: START + DELTA + END + block_id = emitted[3].block_id + self.assertEqual( + [_dump(e) for e in emitted[3:]], + [ + { + "type": "DATA_BLOCK_START", + "reply_id": "reply-1", + "block_id": block_id, + "media_type": "audio/wav", + }, + { + "type": "DATA_BLOCK_DELTA", + "reply_id": "reply-1", + "block_id": block_id, + "data": audio_b64, + "media_type": "audio/wav", + }, + { + "type": "DATA_BLOCK_END", + "reply_id": "reply-1", + "block_id": block_id, + }, + ], + ) + + # synthesize was called with the full text + tts.synthesize.assert_called_once_with("Hello world") + + async def test_empty_text_skips_synthesize(self) -> None: + """When accumulated text is whitespace-only, synthesize is not called + and no DATA_BLOCK events are emitted.""" + tts = MagicMock(spec=TTSModelBase) + tts.realtime = False + tts.__aenter__ = AsyncMock(return_value=tts) + tts.__aexit__ = AsyncMock(return_value=None) + tts.synthesize = AsyncMock() + + middleware = TTSMiddleware(tts_model=tts) + agent = _make_agent_stub() + + upstream_events = [ + TextBlockDeltaEvent( + reply_id="reply-1", + block_id="blk-1", + delta=" ", + ), + TextBlockEndEvent(reply_id="reply-1", block_id="blk-1"), + ] + + async def next_handler(**_kwargs: Any) -> AsyncGenerator: + for evt in upstream_events: + yield evt + + emitted = [] + async for evt in middleware.on_reply(agent, {}, next_handler): + emitted.append(evt) + + # Only the upstream events are passed through + self.assertEqual( + [_dump(e) for e in emitted], + [_dump(e) for e in upstream_events], + ) + tts.synthesize.assert_not_called() + + async def test_audio_block_id_uses_configured_id_factory(self) -> None: + """TTS audio DATA_BLOCK events use the configured ID factory.""" + import agentscope._utils._common as common + + # pylint: disable=protected-access + saved_factory = common._id_factory + self.addCleanup(setattr, common, "_id_factory", saved_factory) + set_id_factory(lambda: "custom-audio-id") + + tts = MagicMock(spec=TTSModelBase) + tts.realtime = False + tts.__aenter__ = AsyncMock(return_value=tts) + tts.__aexit__ = AsyncMock(return_value=None) + tts.synthesize = AsyncMock( + return_value=_make_tts_response("AAAA"), + ) + + middleware = TTSMiddleware(tts_model=tts) + agent = _make_agent_stub() + + async def next_handler(**_kwargs: Any) -> AsyncGenerator: + yield TextBlockDeltaEvent( + reply_id="reply-1", + block_id="blk-1", + delta="hi", + ) + yield TextBlockEndEvent(reply_id="reply-1", block_id="blk-1") + + emitted = [ + evt async for evt in middleware.on_reply(agent, {}, next_handler) + ] + + self.assertEqual( + [evt.block_id for evt in emitted[2:]], + ["custom-audio-id"] * 3, + ) + + async def test_streaming_output_multiple_chunks(self) -> None: + """When synthesize() returns an async generator, each chunk produces + a DATA_BLOCK_DELTA under the same block_id.""" + chunk1 = _make_tts_response("AAAA") + chunk2 = _make_tts_response("BBBB") + + async def synth_gen(*_args: Any, **_kwargs: Any) -> AsyncGenerator: + yield chunk1 + yield chunk2 + + tts = MagicMock(spec=TTSModelBase) + tts.realtime = False + tts.__aenter__ = AsyncMock(return_value=tts) + tts.__aexit__ = AsyncMock(return_value=None) + tts.synthesize = AsyncMock(return_value=synth_gen()) + + middleware = TTSMiddleware(tts_model=tts) + agent = _make_agent_stub() + + upstream_events = [ + TextBlockDeltaEvent( + reply_id="reply-1", + block_id="blk-1", + delta="hi", + ), + TextBlockEndEvent(reply_id="reply-1", block_id="blk-1"), + ] + + async def next_handler(**_kwargs: Any) -> AsyncGenerator: + for evt in upstream_events: + yield evt + + emitted = [] + async for evt in middleware.on_reply(agent, {}, next_handler): + emitted.append(evt) + + # Upstream pass-through + START + 2x DELTA + END + data_events = emitted[2:] + block_id = data_events[0].block_id + self.assertEqual( + [_dump(e) for e in data_events], + [ + { + "type": "DATA_BLOCK_START", + "reply_id": "reply-1", + "block_id": block_id, + "media_type": "audio/wav", + }, + { + "type": "DATA_BLOCK_DELTA", + "reply_id": "reply-1", + "block_id": block_id, + "data": "AAAA", + "media_type": "audio/wav", + }, + { + "type": "DATA_BLOCK_DELTA", + "reply_id": "reply-1", + "block_id": block_id, + "data": "BBBB", + "media_type": "audio/wav", + }, + { + "type": "DATA_BLOCK_END", + "reply_id": "reply-1", + "block_id": block_id, + }, + ], + ) + + +class TestTTSMiddlewareRealtime(IsolatedAsyncioTestCase): + """Tests for realtime (streaming-input) TTS path.""" + + async def test_push_on_delta_and_drain_on_end(self) -> None: + """In realtime mode, push() is called on each TextBlockDelta and + synthesize() drains on TextBlockEnd.""" + push_audio = _make_tts_response("PUSH1") + drain_audio = _make_tts_response("DRAIN") + + tts = MagicMock(spec=TTSModelBase) + tts.realtime = True + tts.__aenter__ = AsyncMock(return_value=tts) + tts.__aexit__ = AsyncMock(return_value=None) + tts.push = AsyncMock(return_value=push_audio) + tts.synthesize = AsyncMock(return_value=drain_audio) + + middleware = TTSMiddleware(tts_model=tts) + agent = _make_agent_stub() + + upstream_events = [ + TextBlockDeltaEvent( + reply_id="reply-1", + block_id="blk-1", + delta="Hello", + ), + TextBlockEndEvent(reply_id="reply-1", block_id="blk-1"), + ] + + async def next_handler(**_kwargs: Any) -> AsyncGenerator: + for evt in upstream_events: + yield evt + + emitted = [] + async for evt in middleware.on_reply(agent, {}, next_handler): + emitted.append(evt) + + # push() called with the delta text + tts.push.assert_called_once_with("Hello") + + # synthesize() called to drain + tts.synthesize.assert_called_once() + + data_events = [ + e + for e in emitted + if isinstance( + e, + (DataBlockStartEvent, DataBlockDeltaEvent, DataBlockEndEvent), + ) + ] + # START + DELTA(push) + DELTA(drain) + END + block_id = data_events[0].block_id + self.assertEqual( + [_dump(e) for e in data_events], + [ + { + "type": "DATA_BLOCK_START", + "reply_id": "reply-1", + "block_id": block_id, + "media_type": "audio/wav", + }, + { + "type": "DATA_BLOCK_DELTA", + "reply_id": "reply-1", + "block_id": block_id, + "data": "PUSH1", + "media_type": "audio/wav", + }, + { + "type": "DATA_BLOCK_DELTA", + "reply_id": "reply-1", + "block_id": block_id, + "data": "DRAIN", + "media_type": "audio/wav", + }, + { + "type": "DATA_BLOCK_END", + "reply_id": "reply-1", + "block_id": block_id, + }, + ], + ) + + async def test_push_returns_none_no_audio_emitted(self) -> None: + """When push() returns empty content, no DATA_BLOCK events are emitted + until synthesize() drains.""" + empty_response = TTSResponse(content=None) + drain_audio = _make_tts_response("FINAL") + + tts = MagicMock(spec=TTSModelBase) + tts.realtime = True + tts.__aenter__ = AsyncMock(return_value=tts) + tts.__aexit__ = AsyncMock(return_value=None) + tts.push = AsyncMock(return_value=empty_response) + tts.synthesize = AsyncMock(return_value=drain_audio) + + middleware = TTSMiddleware(tts_model=tts) + agent = _make_agent_stub() + + upstream_events = [ + TextBlockDeltaEvent( + reply_id="reply-1", + block_id="blk-1", + delta="Hi", + ), + TextBlockDeltaEvent( + reply_id="reply-1", + block_id="blk-1", + delta=" there", + ), + TextBlockEndEvent(reply_id="reply-1", block_id="blk-1"), + ] + + async def next_handler(**_kwargs: Any) -> AsyncGenerator: + for evt in upstream_events: + yield evt + + emitted = [] + async for evt in middleware.on_reply(agent, {}, next_handler): + emitted.append(evt) + + # push called twice but produced no audio + self.assertEqual(tts.push.call_count, 2) + + data_events = [ + e + for e in emitted + if isinstance( + e, + (DataBlockStartEvent, DataBlockDeltaEvent, DataBlockEndEvent), + ) + ] + # Only drain produces: START + DELTA + END + block_id = data_events[0].block_id + self.assertEqual( + [_dump(e) for e in data_events], + [ + { + "type": "DATA_BLOCK_START", + "reply_id": "reply-1", + "block_id": block_id, + "media_type": "audio/wav", + }, + { + "type": "DATA_BLOCK_DELTA", + "reply_id": "reply-1", + "block_id": block_id, + "data": "FINAL", + "media_type": "audio/wav", + }, + { + "type": "DATA_BLOCK_END", + "reply_id": "reply-1", + "block_id": block_id, + }, + ], + ) diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..56de5f31802350c6b5cec17bd8377201535c0c3e --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,128 @@ +# -*- coding: utf-8 -*- +"""The utility module for unit tests in agentscope.""" +import json +from typing import Any, AsyncGenerator, Type + +from pydantic import BaseModel + +from agentscope.credential import CredentialBase +from agentscope.message import Msg +from agentscope.model import ChatModelBase, ChatResponse, StructuredResponse + + +class AnyString(str): + """A helper class for asserting any string value in unit tests.""" + + def __eq__(self, other: object) -> bool: + """Override equality check to match any string.""" + return isinstance(other, str) + + def __repr__(self) -> str: + """Return a string representation for debugging purposes.""" + return "" + + +class AnyValue: + """A helper class for asserting any value (str, int, float, etc.) + in unit tests. Useful for dynamic fields like timestamps, tokens, + or durations where the exact value is unpredictable.""" + + def __eq__(self, other: object) -> bool: + """Override equality check to match any value.""" + return True + + def __repr__(self) -> str: + """Return a string representation for debugging purposes.""" + return "" + + +class MockCredential(CredentialBase): + """The mock credential class.""" + + @classmethod + def get_chat_model_class(cls) -> Type["ChatModelBase"]: + """Return the mock model class.""" + return MockModel + + +class MockModel(ChatModelBase): + """A mock model for testing.""" + + class Parameters(BaseModel): + """The parameters.""" + + def __init__( + self, + model: str = "mock-model", + stream: bool = True, + context_size: int = 1000, + mock_chat_responses: list | None = None, + mock_structured_response: Any = None, + ) -> None: + """Initialize the mock model.""" + super().__init__( + credential=MockCredential(), + model=model, + stream=stream, + parameters=MockModel.Parameters(), + context_size=context_size, + ) + self.mock_chat_responses = mock_chat_responses or [] + self.mock_structured_response = mock_structured_response + self.cnt = 0 + + def set_responses( + self, + mock_responses: list[ChatResponse | list[ChatResponse]], + ) -> None: + """Set the mock responses.""" + self.mock_chat_responses = mock_responses + if all(isinstance(_, ChatResponse) for _ in mock_responses): + self.stream = False + else: + self.stream = True + self.cnt = 0 + + async def _call_api( + self, # pylint: disable=unused-argument + *args: Any, + **kwargs: Any, + ) -> ChatResponse | AsyncGenerator[ChatResponse, None]: + """Mock the API call.""" + mock_responses = self.mock_chat_responses[self.cnt] + self.cnt += 1 + if isinstance(mock_responses, list): + + async def _stream() -> AsyncGenerator[ChatResponse, None]: + for response in mock_responses: + yield response + + return _stream() + + if isinstance(mock_responses, ChatResponse): + return mock_responses + + raise AssertionError + + def set_structured_response( + self, + mock_response: StructuredResponse, + ) -> None: + """Set the mock structured responses.""" + self.mock_structured_response = mock_response + + async def _call_api_with_structured_output( + self, + model_name: str, + messages: list[Msg], + structured_model: Type[BaseModel] | dict, + **kwargs: Any, + ) -> StructuredResponse: + """Mock the API call with structured output.""" + return self.mock_structured_response + + +def compare_by_printing(a: Any, b: Any) -> None: + """Compare the expected output with the actual output by printing them.""" + print(json.dumps(a, indent=4)) + print(json.dumps(b, indent=4)) diff --git a/tests/workspace_docker_test.py b/tests/workspace_docker_test.py new file mode 100644 index 0000000000000000000000000000000000000000..3326e20bcbf3acc57d370c1b5d95e9631dfae6a1 --- /dev/null +++ b/tests/workspace_docker_test.py @@ -0,0 +1,821 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Test cases for DockerWorkspace. + +Mirrors the structure of ``workspace_local_test.py`` so that the two +backends are validated against the same surface contract: + +* offload_context (text + DataBlock + dedup) +* offload_tool_result (string + multimodal blocks) +* skill seeding / listing +* lifecycle (initialize idempotent, close, workspace_id stability) + +Two practical differences vs. the local tests: + +1. Paths *returned* by DockerWorkspace are container-side + (``/workspace/sessions/...``). Host-side file checks therefore + compose paths against ``self.temp_dir.name``, which is bind-mounted + to ``CONTAINER_WORKDIR`` inside the container. + +2. ``DataBlock`` URLs persisted by DockerWorkspace use the container + path (``file:///workspace/data/.png``). We assert the URL + string against the container path and verify the file's *bytes* via + the corresponding host-side path. + +The whole module is skipped when no Docker daemon is reachable. +""" + +import base64 +import hashlib +import os +import re +import shutil +import subprocess +import tempfile +import unittest +import uuid +from typing import Any +from unittest.async_case import IsolatedAsyncioTestCase +from urllib.parse import urlparse + +import aiofiles + +from agentscope.message import ( + AssistantMsg, + Base64Source, + DataBlock, + Msg, + TextBlock, + ToolResultBlock, + ToolResultState, + URLSource, + UserMsg, +) +from agentscope.permission import PermissionBehavior, PermissionDecision +from agentscope.tool import ToolBase, ToolChunk +from agentscope.workspace import DockerWorkspace +from agentscope.workspace._docker._make_dockerfile import ( + CONTAINER_SESSIONS_DIR, + CONTAINER_SKILLS_DIR, +) + +# ── docker daemon detection ──────────────────────────────────────── + + +def _docker_available() -> bool: + """Return ``True`` iff the Docker daemon is reachable. + + Probes via the ``docker`` CLI rather than the aiodocker async client + so the check is cheap and synchronous (runs at module import time). + """ + if shutil.which("docker") is None: + return False + try: + result = subprocess.run( + ["docker", "info"], + capture_output=True, + timeout=5, + check=False, + ) + return result.returncode == 0 + except (subprocess.TimeoutExpired, FileNotFoundError, OSError): + return False + + +_DOCKER_OK = _docker_available() +_SKIP_REASON = "Docker daemon not available" + + +# ── helper tool for agent-integration test ──────────────────────── + + +class _LongResultTool(ToolBase): + """Mock tool that returns a long string + base64 DataBlock.""" + + name: str = "long_result_tool" + description: str = "A tool that returns a long string." + input_schema: dict[str, Any] = { + "type": "object", + "properties": {}, + "required": [], + } + is_concurrency_safe: bool = True + is_read_only: bool = True + is_external_tool: bool = False + is_mcp: bool = False + + async def check_permissions( + self, + *_args: Any, + **_kwargs: Any, + ) -> PermissionDecision: + """Always allow.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + decision_reason="Mock tool always allows", + message="Mock tool always allows", + ) + + async def __call__(self, **_kwargs: Any) -> ToolChunk: + """Long text + base64 DataBlock so we exercise both offload paths.""" + return ToolChunk( + content=[ + TextBlock(text="0" * 30000), + DataBlock( + name="fake_image.png", + source=Base64Source( + data="AAECAwQF", + media_type="image/png", + ), + ), + ], + state=ToolResultState.SUCCESS, + ) + + +# ── offload tests ───────────────────────────────────────────────── + + +@unittest.skipUnless(_DOCKER_OK, _SKIP_REASON) +class TestDockerWorkspaceOffload(IsolatedAsyncioTestCase): + """Test cases for DockerWorkspace offload functionality.""" + + async def asyncSetUp(self) -> None: + """Build a fresh workspace bound to a temp host dir. + + The workspace_id is randomised so test runs do not collide on + the deterministic ``as_ws_`` container name. + """ + # pylint: disable=consider-using-with + self.temp_dir = tempfile.TemporaryDirectory() + self.workspace = DockerWorkspace( + workspace_id=f"test-{uuid.uuid4().hex[:8]}", + host_workdir=self.temp_dir.name, + ) + await self.workspace.initialize() + + async def asyncTearDown(self) -> None: + """Stop the container and clean up the temp host dir.""" + try: + await self.workspace.close() + finally: + self.temp_dir.cleanup() + + async def test_offload_context_pure_text(self) -> None: + """Text-only offload: returned path is container-side, lines match. + + Verifies: + 1. The returned path is the *container-side* JSONL path. + 2. The host-mirror file (under the bind-mounted workdir) contains + one ``model_dump_json()`` line per input message. + """ + session_id = "test_session_pure_text" + msgs = [ + UserMsg(name="user", content="Hello, world!"), + AssistantMsg(name="assistant", content="Hi there!"), + ] + + file_path = await self.workspace.offload_context(session_id, msgs) + + # Returned path is container-side (/workspace/sessions/...). + expected_container_path = ( + f"{CONTAINER_SESSIONS_DIR}/{session_id}/context.jsonl" + ) + self.assertEqual(file_path, expected_container_path) + + # Host-side mirror of that path (bind mount). + host_path = os.path.join( + self.temp_dir.name, + "sessions", + session_id, + "context.jsonl", + ) + self.assertTrue(os.path.exists(host_path)) + + async with aiofiles.open(host_path, "r") as f: + content = await f.read() + lines = content.strip().split("\n") + self.assertEqual(len(lines), 2) + self.assertListEqual(lines, [m.model_dump_json() for m in msgs]) + + async def test_offload_context_multiple_calls(self) -> None: + """Repeated calls append to the same file (proper JSONL). + + Verifies: + 1. Both calls return the same container-side path. + 2. The host-mirror file ends up with one line per message, + preserving order. + """ + session_id = "test_session_multiple" + msgs1 = [ + UserMsg(name="user", content="First message"), + AssistantMsg(name="assistant", content="First response"), + ] + msgs2 = [ + UserMsg(name="user", content="Second message"), + AssistantMsg(name="assistant", content="Second response"), + ] + + path1 = await self.workspace.offload_context(session_id, msgs1) + path2 = await self.workspace.offload_context(session_id, msgs2) + self.assertEqual(path1, path2) + + host_path = os.path.join( + self.temp_dir.name, + "sessions", + session_id, + "context.jsonl", + ) + async with aiofiles.open(host_path, "r") as f: + content = await f.read() + lines = content.strip().split("\n") + self.assertEqual(len(lines), 4) + self.assertListEqual( + lines, + [m.model_dump_json() for m in msgs1 + msgs2], + ) + for line in lines: + self.assertIsNotNone(Msg.model_validate_json(line)) + + async def test_offload_context_with_datablock(self) -> None: + """DataBlock(Base64Source) is rewritten to URLSource on a + separate file. + + Verifies: + 1. The decoded payload lands on disk under ``/data/`` + (container-side ``/workspace/data/``). + 2. The offloaded JSONL line carries a ``URLSource`` whose URL is + the *container* path (``file:///workspace/data/.``). + 3. Decoded bytes match the original. + """ + session_id = "test_session_datablock" + b64_data = base64.b64encode( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01" + b"\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde", + ).decode() + data_block = DataBlock( + source=Base64Source(data=b64_data, media_type="image/png"), + name="test_image", + ) + msgs = [ + UserMsg( + name="user", + content=[TextBlock(text="Check this image:"), data_block], + ), + ] + + file_path = await self.workspace.offload_context(session_id, msgs) + host_path = os.path.join( + self.temp_dir.name, + "sessions", + session_id, + "context.jsonl", + ) + self.assertTrue(os.path.exists(host_path)) + + async with aiofiles.open(host_path, "r") as f: + content = await f.read() + loaded_msg = Msg.model_validate_json(content.strip()) + + # The DataBlock's source should now be URLSource pointing to a + # container-side file://... + self.assertEqual(len(loaded_msg.content), 2) + data_url = str(loaded_msg.content[1].source.url) + self.assertTrue(data_url.startswith("file:///workspace/data/")) + + # Verify the bytes via the host-side mirror. + container_path = urlparse(data_url).path # /workspace/data/.png + rel = os.path.relpath(container_path, "/workspace") + host_data_path = os.path.join(self.temp_dir.name, rel) + self.assertTrue(os.path.exists(host_data_path)) + async with aiofiles.open(host_data_path, "rb") as f: + saved = await f.read() + self.assertEqual(saved, base64.b64decode(b64_data)) + + # Sanity-check the returned (container) JSONL path too. + self.assertEqual( + file_path, + f"{CONTAINER_SESSIONS_DIR}/{session_id}/context.jsonl", + ) + + async def test_offload_data_block_deduplication(self) -> None: + """Two identical DataBlocks share a single persisted file. + + Verifies: + 1. Both ``_offload_data_block`` calls return DataBlocks pointing + at the same URL. + 2. Only one file ends up in ``/data/``. + """ + b64_data = base64.b64encode(b"test content").decode() + block1 = DataBlock( + source=Base64Source(data=b64_data, media_type="text/plain"), + name="file1", + ) + block2 = DataBlock( + source=Base64Source(data=b64_data, media_type="text/plain"), + name="file2", + ) + + result1 = await self.workspace._offload_data_block(block1) + result2 = await self.workspace._offload_data_block(block2) + self.assertEqual(str(result1.source.url), str(result2.source.url)) + + # Check the host-mirror data dir has exactly one file. + host_data_dir = os.path.join(self.temp_dir.name, "data") + self.assertTrue(os.path.isdir(host_data_dir)) + files = os.listdir(host_data_dir) + self.assertEqual(len(files), 1) + + # Hash key is sha256 of the *base64 string* (matches LocalWorkspace). + expected_name = hashlib.sha256(b64_data.encode()).hexdigest() + ".txt" + self.assertIn(expected_name, files) + + async def test_offload_data_block_url_source(self) -> None: + """A DataBlock that already has a URLSource is returned unchanged. + + Verifies: + 1. ``_offload_data_block`` is a no-op for URL-sourced blocks. + 2. No file lands in ``/data/``. + """ + from pydantic import AnyUrl + + block = DataBlock( + source=URLSource( + url=AnyUrl("https://example.com/image.png"), + media_type="image/png", + ), + name="remote_image", + ) + + result = await self.workspace._offload_data_block(block) + self.assertDictEqual(result.model_dump(), block.model_dump()) + + host_data_dir = os.path.join(self.temp_dir.name, "data") + if os.path.isdir(host_data_dir): + self.assertEqual(len(os.listdir(host_data_dir)), 0) + + async def test_offload_tool_result_string(self) -> None: + """String-output tool result writes the raw string to disk. + + Verifies: + 1. The returned path is the container-side ``tool_result-.txt``. + 2. The host-mirror file contains exactly the output string. + """ + session_id = "test_session_tool_result" + tool_result = ToolResultBlock( + id="tool_123", + name="test_tool", + output="Tool execution successful!", + state=ToolResultState.SUCCESS, + ) + + file_path = await self.workspace.offload_tool_result( + session_id, + tool_result, + ) + self.assertEqual( + file_path, + f"{CONTAINER_SESSIONS_DIR}/{session_id}/" + f"tool_result-{tool_result.id}.txt", + ) + + host_path = os.path.join( + self.temp_dir.name, + "sessions", + session_id, + f"tool_result-{tool_result.id}.txt", + ) + self.assertTrue(os.path.exists(host_path)) + async with aiofiles.open(host_path, "r") as f: + self.assertEqual(await f.read(), "Tool execution successful!") + + async def test_offload_tool_result_with_blocks(self) -> None: + """Multimodal tool result: text concatenated, DataBlock placeholder. + + Verifies: + 1. The text portion is preserved verbatim. + 2. The DataBlock is rendered as + ```` with a + container-side URL. + 3. The decoded payload is reachable via the host mirror. + """ + session_id = "test_session_tool_result_blocks" + b64_data = base64.b64encode(b"test file content").decode() + data_block = DataBlock( + source=Base64Source(data=b64_data, media_type="text/plain"), + name="output.txt", + ) + tool_result = ToolResultBlock( + id="tool_456", + name="file_tool", + output=[ + TextBlock(text="File created successfully: "), + data_block, + ], + state=ToolResultState.SUCCESS, + ) + + file_path = await self.workspace.offload_tool_result( + session_id, + tool_result, + ) + host_path = os.path.join( + self.temp_dir.name, + "sessions", + session_id, + f"tool_result-{tool_result.id}.txt", + ) + self.assertTrue(os.path.exists(host_path)) + async with aiofiles.open(host_path, "r") as f: + content = await f.read() + + self.assertTrue(content.startswith("File created successfully: ")) + self.assertIn("")) + + # The data file is reachable via the host mirror. + match = re.search(r"url='([^']+)'", content) + self.assertIsNotNone(match) + container_path = urlparse(match.group(1)).path + rel = os.path.relpath(container_path, "/workspace") + host_data_path = os.path.join(self.temp_dir.name, rel) + self.assertTrue(os.path.exists(host_data_path)) + + # And the returned path is container-side. + self.assertEqual( + file_path, + f"{CONTAINER_SESSIONS_DIR}/{session_id}/" + f"tool_result-{tool_result.id}.txt", + ) + + +# ── skill tests ──────────────────────────────────────────────────── + + +@unittest.skipUnless(_DOCKER_OK, _SKIP_REASON) +class TestDockerWorkspaceSkills(IsolatedAsyncioTestCase): + """Test cases for DockerWorkspace skill management.""" + + async def asyncSetUp(self) -> None: + """Build separate temp dirs for the workspace and the skill source.""" + # pylint: disable=consider-using-with + self.temp_dir = tempfile.TemporaryDirectory() + self.test_skills_dir = tempfile.TemporaryDirectory() + self.workspace: DockerWorkspace | None = None + + async def asyncTearDown(self) -> None: + """Stop the workspace if started, then drop both temp dirs.""" + try: + if self.workspace is not None: + await self.workspace.close() + finally: + self.temp_dir.cleanup() + self.test_skills_dir.cleanup() + + def _create_test_skill( + self, + skill_name: str, + description: str, + additional_files: dict[str, str] | None = None, + ) -> str: + """Create a host-side skill directory containing SKILL.md. + + Args: + skill_name: Name written to the SKILL.md front matter and + used as the directory basename. + description: ``description`` field for the front matter. + additional_files: Optional ``{filename: content}`` map of + supplementary files written alongside SKILL.md. + + Returns: + Absolute path to the created directory. + """ + skill_dir = os.path.join(self.test_skills_dir.name, skill_name) + os.makedirs(skill_dir, exist_ok=True) + skill_md = ( + f"---\nname: {skill_name}\ndescription: {description}\n---\n\n" + f"# {skill_name}\n\n{description}\n" + ) + with open( + os.path.join(skill_dir, "SKILL.md"), + "w", + encoding="utf-8", + ) as f: + f.write(skill_md) + if additional_files: + for filename, content in additional_files.items(): + with open( + os.path.join(skill_dir, filename), + "w", + encoding="utf-8", + ) as f: + f.write(content) + return skill_dir + + async def test_initialize_copy_skills(self) -> None: + """``skill_paths`` are copied into the container's ``skills/``. + + Verifies: + 1. Both seed skills appear under ``/skills/`` (host + mirror) — directory + SKILL.md + supplementary files. + """ + skill1 = self._create_test_skill( + "test_skill_1", + "First test skill", + {"tool.py": "def t():\n pass\n"}, + ) + skill2 = self._create_test_skill( + "test_skill_2", + "Second test skill", + {"helper.py": "def h():\n return 42\n"}, + ) + + self.workspace = DockerWorkspace( + workspace_id=f"test-{uuid.uuid4().hex[:8]}", + host_workdir=self.temp_dir.name, + skill_paths=[skill1, skill2], + ) + await self.workspace.initialize() + + skills_host = os.path.join(self.temp_dir.name, "skills") + self.assertTrue(os.path.isdir(skills_host)) + for name, extra in ( + ("test_skill_1", "tool.py"), + ("test_skill_2", "helper.py"), + ): + target = os.path.join(skills_host, name) + self.assertTrue(os.path.isdir(target)) + self.assertTrue( + os.path.isfile(os.path.join(target, "SKILL.md")), + ) + self.assertTrue(os.path.isfile(os.path.join(target, extra))) + + async def test_list_skills(self) -> None: + """``list_skills`` enumerates skills via in-container ``find``. + + Verifies: + 1. Every seeded skill is returned with the right name + dir. + 2. The dir field uses the *container-side* path. + """ + skill1 = self._create_test_skill( + "list_skill_1", + "First skill for listing", + ) + skill2 = self._create_test_skill( + "list_skill_2", + "Second skill for listing", + ) + + self.workspace = DockerWorkspace( + workspace_id=f"test-{uuid.uuid4().hex[:8]}", + host_workdir=self.temp_dir.name, + skill_paths=[skill1, skill2], + ) + await self.workspace.initialize() + + skills = await self.workspace.list_skills() + self.assertEqual(len(skills), 2) + skills_sorted = sorted(skills, key=lambda s: s.name) + + self.assertEqual(skills_sorted[0].name, "list_skill_1") + self.assertEqual( + skills_sorted[0].description, + "First skill for listing", + ) + self.assertEqual( + skills_sorted[0].dir, + f"{CONTAINER_SKILLS_DIR}/list_skill_1", + ) + + self.assertEqual(skills_sorted[1].name, "list_skill_2") + self.assertEqual( + skills_sorted[1].description, + "Second skill for listing", + ) + self.assertEqual( + skills_sorted[1].dir, + f"{CONTAINER_SKILLS_DIR}/list_skill_2", + ) + + async def test_list_skills_empty(self) -> None: + """Empty workspace → empty skill list (no errors).""" + self.workspace = DockerWorkspace( + workspace_id=f"test-{uuid.uuid4().hex[:8]}", + host_workdir=self.temp_dir.name, + ) + await self.workspace.initialize() + skills = await self.workspace.list_skills() + self.assertListEqual(skills, []) + + async def test_add_skill_then_list(self) -> None: + """``add_skill`` after init shows up in subsequent ``list_skills``.""" + self.workspace = DockerWorkspace( + workspace_id=f"test-{uuid.uuid4().hex[:8]}", + host_workdir=self.temp_dir.name, + ) + await self.workspace.initialize() + self.assertListEqual(await self.workspace.list_skills(), []) + + new_skill = self._create_test_skill( + "added_skill", + "Added at runtime", + ) + await self.workspace.add_skill(new_skill) + + skills = await self.workspace.list_skills() + self.assertEqual(len(skills), 1) + self.assertEqual(skills[0].name, "added_skill") + self.assertEqual(skills[0].description, "Added at runtime") + + async def test_add_skill_invalid_no_md(self) -> None: + """``add_skill`` rejects directories without SKILL.md.""" + self.workspace = DockerWorkspace( + workspace_id=f"test-{uuid.uuid4().hex[:8]}", + host_workdir=self.temp_dir.name, + ) + await self.workspace.initialize() + + invalid = os.path.join(self.test_skills_dir.name, "no_md") + os.makedirs(invalid, exist_ok=True) + with open( + os.path.join(invalid, "tool.py"), + "w", + encoding="utf-8", + ) as f: + f.write("def t():\n pass\n") + + with self.assertRaises(ValueError): + await self.workspace.add_skill(invalid) + + +# ── lifecycle tests ─────────────────────────────────────────────── + + +@unittest.skipUnless(_DOCKER_OK, _SKIP_REASON) +class TestDockerWorkspaceLifecycle(IsolatedAsyncioTestCase): + """Test cases for DockerWorkspace lifecycle (init / close / restart).""" + + async def asyncSetUp(self) -> None: + """Per-test temp workdir; workspace built lazily inside each test.""" + # pylint: disable=consider-using-with + self.temp_dir = tempfile.TemporaryDirectory() + self.workspace_id = f"test-{uuid.uuid4().hex[:8]}" + + async def asyncTearDown(self) -> None: + """Best-effort cleanup of the temp workdir.""" + self.temp_dir.cleanup() + + async def test_initialize_idempotent(self) -> None: + """Calling ``initialize`` on a live workspace is a no-op.""" + ws = DockerWorkspace( + workspace_id=self.workspace_id, + host_workdir=self.temp_dir.name, + ) + try: + await ws.initialize() + self.assertTrue(ws.is_alive) + container = ws._container + + await ws.initialize() # second call — must be a no-op + self.assertTrue(ws.is_alive) + self.assertIs(ws._container, container) + finally: + await ws.close() + + async def test_close_marks_inactive(self) -> None: + """``close`` flips ``is_alive`` and tears the container down.""" + ws = DockerWorkspace( + workspace_id=self.workspace_id, + host_workdir=self.temp_dir.name, + ) + await ws.initialize() + self.assertTrue(ws.is_alive) + await ws.close() + self.assertFalse(ws.is_alive) + self.assertIsNone(ws._container) + self.assertIsNone(ws._client) + + async def test_list_mcps_empty(self) -> None: + """No MCPs registered → ``list_mcps`` returns an empty list.""" + ws = DockerWorkspace( + workspace_id=self.workspace_id, + host_workdir=self.temp_dir.name, + ) + try: + await ws.initialize() + self.assertListEqual(await ws.list_mcps(), []) + finally: + await ws.close() + + async def test_list_tools_builtin(self) -> None: + """returns the six builtin tools backed by DockerBackend.""" + from agentscope.tool._builtin import ( + Bash, + Edit, + Glob, + Grep, + Read, + Write, + ) + from agentscope.workspace._docker._docker_backend import DockerBackend + + ws = DockerWorkspace( + workspace_id=self.workspace_id, + host_workdir=self.temp_dir.name, + ) + try: + await ws.initialize() + tools = await ws.list_tools() + self.assertEqual(len(tools), 6) + + expected_types = {Bash, Edit, Glob, Grep, Read, Write} + actual_types = {type(t) for t in tools} + self.assertSetEqual(actual_types, expected_types) + + for tool in tools: + self.assertIsInstance(tool._backend, DockerBackend) + finally: + await ws.close() + + async def test_workdir_persistence_across_restart(self) -> None: + """Same ``workspace_id`` + same ``workdir`` → state survives close. + + Specifically: an offloaded session file written before ``close`` + is still readable from the host mirror after a fresh + ``initialize`` (the new container re-bind-mounts the same host + workdir). + """ + session_id = "persisted_session" + msg = UserMsg(name="user", content="durable line") + + ws1 = DockerWorkspace( + workspace_id=self.workspace_id, + host_workdir=self.temp_dir.name, + ) + await ws1.initialize() + await ws1.offload_context(session_id, [msg]) + await ws1.close() + + host_path = os.path.join( + self.temp_dir.name, + "sessions", + session_id, + "context.jsonl", + ) + self.assertTrue(os.path.exists(host_path)) + + ws2 = DockerWorkspace( + workspace_id=self.workspace_id, + host_workdir=self.temp_dir.name, + ) + try: + await ws2.initialize() + # The container mounts the same host dir; the file should + # still be visible via the host mirror. + self.assertTrue(os.path.exists(host_path)) + async with aiofiles.open(host_path, "r") as f: + content = await f.read() + self.assertEqual(content.strip(), msg.model_dump_json()) + finally: + await ws2.close() + + async def test_reset_clears_sessions_and_data(self) -> None: + """``reset`` removes ``sessions/`` and ``data/`` (host-visible).""" + ws = DockerWorkspace( + workspace_id=self.workspace_id, + host_workdir=self.temp_dir.name, + ) + try: + await ws.initialize() + + # Seed both directories via offload_*. + session_id = "reset_session" + await ws.offload_context( + session_id, + [UserMsg(name="user", content="hi")], + ) + await ws._offload_data_block( + DataBlock( + source=Base64Source( + data=base64.b64encode(b"x").decode(), + media_type="text/plain", + ), + name="x", + ), + ) + sessions_host = os.path.join(self.temp_dir.name, "sessions") + data_host = os.path.join(self.temp_dir.name, "data") + self.assertTrue(os.path.isdir(sessions_host)) + self.assertTrue(os.path.isdir(data_host)) + + await ws.reset() + self.assertFalse(os.path.exists(sessions_host)) + self.assertFalse(os.path.exists(data_host)) + finally: + await ws.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/workspace_e2b_test.py b/tests/workspace_e2b_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e0e95ed86af52c6b6d145c4a410521b2f4749c25 --- /dev/null +++ b/tests/workspace_e2b_test.py @@ -0,0 +1,72 @@ +# -*- coding: utf-8 -*- +"""Test cases for E2BWorkspace. + +The whole module is skipped when the ``E2B_API_KEY`` environment variable is +not set, because every test requires a live E2B cloud sandbox. +""" +import os +import unittest +from unittest.async_case import IsolatedAsyncioTestCase + +from agentscope.mcp import MCPClient, StdioMCPConfig +from agentscope.workspace import E2BWorkspace + + +# ── E2B availability check ───────────────────────────────────────── + +_E2B_API_KEY = os.getenv("E2B_API_KEY", "") +_SKIP_REASON = "E2B_API_KEY environment variable is not set" + + +# ── lifecycle tests ──────────────────────────────────────────────── + + +@unittest.skipUnless(_E2B_API_KEY, _SKIP_REASON) +class TestE2BWorkspaceLifecycle(IsolatedAsyncioTestCase): + """Test cases for E2BWorkspace lifecycle and MCP integration. + + Each test creates a real E2B cloud sandbox and tears it down + (``pause``) afterwards. The suite is skipped entirely when + ``E2B_API_KEY`` is absent so that CI runs without E2B credentials + are unaffected. + """ + + async def asyncSetUp(self) -> None: + """No shared setup — each test manages its own workspace.""" + + async def asyncTearDown(self) -> None: + """No shared teardown — each test closes its own workspace.""" + + async def test_initialize_and_list_mcps(self) -> None: + """``initialize`` starts the sandbox and ``list_mcps`` enumerates MCPs. + + Verifies: + 1. The workspace initializes without raising. + 2. ``list_mcps`` returns at least the seeded MCP (browser-use). + 3. Each MCP exposes at least one tool via ``list_raw_tools``. + 4. ``close`` (sandbox pause) completes without raising. + """ + workspace = E2BWorkspace( + api_key=_E2B_API_KEY, + default_mcps=[ + MCPClient( + name="browser-use", + mcp_config=StdioMCPConfig( + command="npx", + args=["@playwright/mcp@latest"], + ), + is_stateful=True, + ), + ], + ) + + await workspace.initialize() + + mcps = await workspace.list_mcps() + self.assertGreater(len(mcps), 0) + + for mcp in mcps: + tools = await mcp.list_raw_tools() + self.assertGreater(len(tools), 0) + + await workspace.close() diff --git a/tests/workspace_local_test.py b/tests/workspace_local_test.py new file mode 100644 index 0000000000000000000000000000000000000000..4bd52c73638d7d72b4c8d188fa7529ad15682796 --- /dev/null +++ b/tests/workspace_local_test.py @@ -0,0 +1,1300 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Test cases for LocalWorkspace.""" +import os +import json +import base64 +import hashlib +import tempfile +from pathlib import Path +from typing import Any +from unittest.async_case import IsolatedAsyncioTestCase +from dataclasses import asdict +from urllib.parse import urlparse +from urllib.request import url2pathname + +import aiofiles +from utils import AnyString, MockModel +from agentscope.agent import Agent, ContextConfig +from agentscope.model import ChatResponse, StructuredResponse +from agentscope.state import AgentState +from agentscope.tool import Toolkit, ToolBase, ToolChunk +from agentscope.permission import PermissionDecision, PermissionBehavior +from agentscope.workspace import LocalWorkspace +from agentscope.mcp import MCPClient, StdioMCPConfig +from agentscope.message import ( + Msg, + UserMsg, + AssistantMsg, + DataBlock, + Base64Source, + URLSource, + TextBlock, + ToolResultBlock, + ToolResultState, + ToolCallBlock, +) + + +class _LongResultTool(ToolBase): + """A mock tool that returns a long string result for offload testing.""" + + name: str = "long_result_tool" + description: str = "A tool that returns a long string." + input_schema: dict[str, Any] = { + "type": "object", + "properties": {}, + "required": [], + } + is_concurrency_safe: bool = True + is_read_only: bool = True + is_external_tool: bool = False + is_mcp: bool = False + + async def check_permissions( + self, + *_args: Any, + **_kwargs: Any, + ) -> PermissionDecision: + """Always allow.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + decision_reason="Mock tool always allows", + message="Mock tool always allows", + ) + + async def __call__(self, **_kwargs: Any) -> ToolChunk: + """Return a long string result followed by a base64 data block, so we + can also verify base64 data offloading.""" + return ToolChunk( + content=[ + TextBlock(text="0" * 30000), + DataBlock( + name="fake_image.png", + source=Base64Source( + data="AAECAwQF", + media_type="image/png", + ), + ), + ], + state=ToolResultState.SUCCESS, + ) + + +class TestLocalWorkspaceOffload(IsolatedAsyncioTestCase): + """Test cases for LocalWorkspace offload functionality.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + # pylint: disable=consider-using-with + self.temp_dir = tempfile.TemporaryDirectory() + self.workspace = LocalWorkspace(workdir=self.temp_dir.name) + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.temp_dir.cleanup() + + async def test_offload_context_pure_text(self) -> None: + """Test offloading messages with pure text content. + + This test verifies that: + 1. Messages with string content are correctly offloaded + 2. The offloaded file is created at the expected path + 3. The file contains valid JSONL with all message fields preserved + """ + session_id = "test_session_pure_text" + msgs = [ + UserMsg(name="user", content="Hello, world!"), + AssistantMsg(name="assistant", content="Hi there!"), + ] + + # Offload the messages + file_path = await self.workspace.offload_context(session_id, msgs) + + # Verify the file was created at the expected path + expected_path = os.path.join( + self.temp_dir.name, + "sessions", + session_id, + "context.jsonl", + ) + self.assertEqual(file_path, expected_path) + self.assertTrue(os.path.exists(file_path)) + + # Read and verify the offloaded messages + async with aiofiles.open(file_path, "r") as f: + content = await f.read() + + lines = content.strip().split("\n") + self.assertEqual(len(lines), 2) + + # Compare with expected JSON strings + expected_lines = [msg.model_dump_json() for msg in msgs] + self.assertListEqual(lines, expected_lines) + + async def test_offload_context_multiple_calls(self) -> None: + """Test multiple calls to offload_context for the same session. + + This test verifies that: + 1. Multiple calls to offload_context append correctly to the file + 2. Each message is on its own line (proper JSONL format) + 3. No lines are concatenated together + """ + session_id = "test_session_multiple" + + # First batch of messages + msgs1 = [ + UserMsg(name="user", content="First message"), + AssistantMsg(name="assistant", content="First response"), + ] + + # Second batch of messages + msgs2 = [ + UserMsg(name="user", content="Second message"), + AssistantMsg(name="assistant", content="Second response"), + ] + + # Offload first batch + file_path = await self.workspace.offload_context(session_id, msgs1) + + # Offload second batch + file_path2 = await self.workspace.offload_context(session_id, msgs2) + + # Verify both calls return the same path + self.assertEqual(file_path, file_path2) + + # Read and verify the offloaded messages + async with aiofiles.open(file_path, "r") as f: + content = await f.read() + + lines = content.strip().split("\n") + self.assertEqual(len(lines), 4) + + # Compare with expected JSON strings + expected_lines = [msg.model_dump_json() for msg in msgs1 + msgs2] + self.assertListEqual(lines, expected_lines) + + # Verify each line is valid JSON + for line in lines: + msg = Msg.model_validate_json(line) + self.assertIsNotNone(msg) + + async def test_offload_context_with_datablock(self) -> None: + """Test offloading messages with DataBlock content. + + This test verifies that: + 1. Messages with DataBlock (Base64Source) are correctly offloaded + 2. DataBlock data is persisted to separate files + 3. DataBlock source is converted from Base64Source to URLSource + 4. The offloaded message file contains the updated DataBlock + """ + session_id = "test_session_datablock" + + # Create a test image data (1x1 red pixel PNG) + test_data = base64.b64encode( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01" + b"\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde", + ).decode() + + data_block = DataBlock( + source=Base64Source(data=test_data, media_type="image/png"), + name="test_image", + ) + + msgs = [ + UserMsg( + name="user", + content=[TextBlock(text="Check this image:"), data_block], + ), + ] + + # Offload the messages + file_path = await self.workspace.offload_context(session_id, msgs) + + # Verify the message file was created + self.assertTrue(os.path.exists(file_path)) + + # Read and verify the offloaded message + async with aiofiles.open(file_path, "r") as f: + content = await f.read() + + loaded_msg = Msg.model_validate_json(content.strip()) + + # Verify the data file was created and extract the URL + self.assertIsInstance(loaded_msg.content, list) + self.assertEqual(len(loaded_msg.content), 2) + data_url = str(loaded_msg.content[1].source.url) + self.assertTrue(data_url.startswith("file://")) + # Convert file URL to local path (works on both Windows and Unix) + data_file_path = url2pathname(urlparse(data_url).path) + self.assertTrue(os.path.exists(data_file_path)) + + # Verify the data file contains the correct content + async with aiofiles.open(data_file_path, "rb") as f: + saved_data = await f.read() + self.assertEqual(saved_data, base64.b64decode(test_data)) + + # Build expected message with URLSource for comparison + # Use the actual IDs from loaded message to avoid UUID mismatch + expected_msg = UserMsg( + name="user", + content=[ + TextBlock( + text="Check this image:", + id=loaded_msg.content[0].id, + ), + DataBlock( + id=loaded_msg.content[1].id, + source=loaded_msg.content[1].source, + name="test_image", + ), + ], + id=loaded_msg.id, + created_at=loaded_msg.created_at, + ) + self.assertEqual( + loaded_msg.model_dump_json(), + expected_msg.model_dump_json(), + ) + + async def test_offload_data_block_deduplication(self) -> None: + """Test that duplicate DataBlocks are deduplicated. + + This test verifies that: + 1. Multiple DataBlocks with the same content share the same file + 2. Only one file is created for duplicate data + 3. Both DataBlocks point to the same file path + """ + # Create two DataBlocks with identical data + test_data = base64.b64encode(b"test content").decode() + + data_block1 = DataBlock( + source=Base64Source(data=test_data, media_type="text/plain"), + name="file1", + ) + data_block2 = DataBlock( + source=Base64Source(data=test_data, media_type="text/plain"), + name="file2", + ) + + # Offload both data blocks + result1 = await self.workspace._offload_data_block(data_block1) + result2 = await self.workspace._offload_data_block(data_block2) + + # Verify both point to the same file by comparing source URLs + self.assertEqual(str(result1.source.url), str(result2.source.url)) + + # Verify the file exists + data_url = str(result1.source.url) + # Convert file URL to local path (works on both Windows and Unix) + data_file_path = url2pathname(urlparse(data_url).path) + self.assertTrue(os.path.exists(data_file_path)) + + # Verify only one file was created in the data directory + data_dir = os.path.join(self.temp_dir.name, "data") + files = os.listdir(data_dir) + self.assertEqual(len(files), 1) + + async def test_offload_data_block_url_source(self) -> None: + """Test offloading DataBlock with URLSource. + + This test verifies that: + 1. DataBlock with URLSource is returned as-is + 2. No file is created for URLSource DataBlocks + """ + from pydantic import AnyUrl + + data_block = DataBlock( + source=URLSource( + url=AnyUrl("https://example.com/image.png"), + media_type="image/png", + ), + name="remote_image", + ) + + # Offload the data block + result = await self.workspace._offload_data_block(data_block) + + # Verify the data block is returned as-is by comparing full objects + self.assertDictEqual(result.model_dump(), data_block.model_dump()) + + # Verify no file was created in the data directory + data_dir = os.path.join(self.temp_dir.name, "data") + if os.path.exists(data_dir): + files = os.listdir(data_dir) + self.assertEqual(len(files), 0) + + async def test_offload_tool_result_string(self) -> None: + """Test offloading tool result with string output. + + This test verifies that: + 1. Tool result with string output is correctly offloaded + 2. The offloaded file is created at the expected path + 3. The file contains the correct string content + """ + session_id = "test_session_tool_result" + tool_result = ToolResultBlock( + id="tool_123", + name="test_tool", + output="Tool execution successful!", + state=ToolResultState.SUCCESS, + ) + + # Offload the tool result + file_path = await self.workspace.offload_tool_result( + session_id, + tool_result, + ) + + # Verify the file was created at the expected path + expected_path = os.path.join( + self.temp_dir.name, + "sessions", + session_id, + f"tool_result-{tool_result.id}.txt", + ) + self.assertEqual(file_path, expected_path) + self.assertTrue(os.path.exists(file_path)) + + # Read and verify the content + async with aiofiles.open(file_path, "r") as f: + content = await f.read() + + expected_content = "Tool execution successful!" + self.assertEqual(content, expected_content) + + async def test_offload_tool_result_with_blocks(self) -> None: + """Test offloading tool result with TextBlock and DataBlock output. + + This test verifies that: + 1. Tool result with list of blocks is correctly offloaded + 2. TextBlock content is extracted and written to file + 3. DataBlock is offloaded and referenced in the output file + 4. The output file contains the correct format + """ + session_id = "test_session_tool_result_blocks" + + # Create test data + test_data = base64.b64encode(b"test file content").decode() + data_block = DataBlock( + source=Base64Source(data=test_data, media_type="text/plain"), + name="output.txt", + ) + + tool_result = ToolResultBlock( + id="tool_456", + name="file_tool", + output=[ + TextBlock(text="File created successfully: "), + data_block, + ], + state=ToolResultState.SUCCESS, + ) + + # Offload the tool result + file_path = await self.workspace.offload_tool_result( + session_id, + tool_result, + ) + + # Verify the file was created + self.assertTrue(os.path.exists(file_path)) + + # Read and verify the content + async with aiofiles.open(file_path, "r") as f: + content = await f.read() + + # Verify the content structure (URL format varies by platform) + self.assertTrue(content.startswith("File created successfully: ")) + self.assertIn("")) + + # Extract and verify the data file exists + # Parse the URL from the content + import re + + url_match = re.search(r"url='([^']+)'", content) + self.assertIsNotNone(url_match) + data_url = url_match.group(1) + # Convert file URL to local path (works on both Windows and Unix) + data_file_path = url2pathname(urlparse(data_url).path) + self.assertTrue(os.path.exists(data_file_path)) + + +class TestLocalWorkspaceSkills(IsolatedAsyncioTestCase): + """Test cases for LocalWorkspace skill management functionality.""" + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + # pylint: disable=consider-using-with + self.temp_dir = tempfile.TemporaryDirectory() + # pylint: disable=consider-using-with + self.test_skills_dir = tempfile.TemporaryDirectory() + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.temp_dir.cleanup() + self.test_skills_dir.cleanup() + + def _create_test_skill( + self, + skill_name: str, + description: str, + additional_files: dict[str, str] | None = None, + ) -> str: + """Create a test skill directory with SKILL.md. + + Args: + skill_name (`str`): + The name of the skill. + description (`str`): + The description of the skill. + additional_files (`dict[str, str] | None`, optional): + Additional files to create in the skill directory. + Keys are file names, values are file contents. + + Returns: + `str`: + The path to the created skill directory. + """ + skill_dir = os.path.join(self.test_skills_dir.name, skill_name) + os.makedirs(skill_dir, exist_ok=True) + + # Create SKILL.md with frontmatter + skill_md_content = f"""--- +name: {skill_name} +description: {description} +--- + +# {skill_name} + +{description} +""" + with open( + os.path.join(skill_dir, "SKILL.md"), + "w", + encoding="utf-8", + ) as f: + f.write(skill_md_content) + + # Create additional files if provided + if additional_files: + for filename, content in additional_files.items(): + with open( + os.path.join(skill_dir, filename), + "w", + encoding="utf-8", + ) as f: + f.write(content) + + return skill_dir + + async def test_initialize_copy_skills(self) -> None: + """Test copying skills to workspace. + + This test verifies that: + 1. Skills are correctly copied from source paths to workspace + 2. The .skills file is created with correct hash mappings + 3. All skill files are preserved during copying + """ + # Create test skills + skill1_dir = self._create_test_skill( + "test_skill_1", + "A test skill for testing", + {"tool.py": "def test_tool():\n pass\n"}, + ) + skill2_dir = self._create_test_skill( + "test_skill_2", + "Another test skill", + {"helper.py": "def helper():\n return 42\n"}, + ) + + # Create workspace with skill paths + workspace = LocalWorkspace( + workdir=self.temp_dir.name, + skill_paths=[skill1_dir, skill2_dir], + ) + + # Initialize the workspace + await workspace.initialize() + + # Verify skills were copied + skills_dir = os.path.join(self.temp_dir.name, "skills") + self.assertTrue(os.path.exists(skills_dir)) + + # Verify skill directories exist + skill1_target = os.path.join(skills_dir, "test_skill_1") + skill2_target = os.path.join(skills_dir, "test_skill_2") + self.assertTrue(os.path.exists(skill1_target)) + self.assertTrue(os.path.exists(skill2_target)) + + # Verify SKILL.md files exist + self.assertTrue( + os.path.exists(os.path.join(skill1_target, "SKILL.md")), + ) + self.assertTrue( + os.path.exists(os.path.join(skill2_target, "SKILL.md")), + ) + + # Verify additional files were copied + self.assertTrue(os.path.exists(os.path.join(skill1_target, "tool.py"))) + self.assertTrue( + os.path.exists(os.path.join(skill2_target, "helper.py")), + ) + + # Verify .skills file was created with correct new structure + skills_hash_file = os.path.join(skills_dir, ".skills") + self.assertTrue(os.path.exists(skills_hash_file)) + + async with aiofiles.open(skills_hash_file, "r") as f: + skills_data = json.loads(await f.read()) + + # Verify top-level structure + self.assertIn("skills_dir_mtime", skills_data) + self.assertIn("skills", skills_data) + + skills_index = skills_data["skills"] + self.assertEqual(len(skills_index), 2) + + # Verify each entry has the correct structure + self.assertIn("test_skill_1", skills_index) + self.assertIn("test_skill_2", skills_index) + self.assertDictEqual( + {k: v["skill_name"] for k, v in skills_index.items()}, + {"test_skill_1": "test_skill_1", "test_skill_2": "test_skill_2"}, + ) + + async def test_initialize_skip_duplicate_skills(self) -> None: + """Test that duplicate skills are not copied again. + + This test verifies that: + 1. Skills are copied on first initialization + 2. Running initialize again does not copy duplicate skills + 3. The .skills file is not modified on second initialization + """ + # Create test skill + skill_dir = self._create_test_skill( + "test_skill_dup", + "A test skill for duplication testing", + ) + + # Create workspace and initialize + workspace = LocalWorkspace( + workdir=self.temp_dir.name, + skill_paths=[skill_dir], + ) + await workspace.initialize() + + # Get the .skills file content after first initialization + skills_hash_file = os.path.join( + self.temp_dir.name, + "skills", + ".skills", + ) + async with aiofiles.open(skills_hash_file, "r") as f: + hash_data_first = await f.read() + + # Get modification time of the skill directory + skill_target = os.path.join( + self.temp_dir.name, + "skills", + "test_skill_dup", + ) + mtime_first = os.path.getmtime(skill_target) + + # Initialize again + await workspace.initialize() + + # Verify .skills file is unchanged + async with aiofiles.open(skills_hash_file, "r") as f: + hash_data_second = await f.read() + self.assertEqual(hash_data_first, hash_data_second) + + # Verify skill directory was not modified + mtime_second = os.path.getmtime(skill_target) + self.assertEqual(mtime_first, mtime_second) + + async def test_initialize_deduplicate_skills(self) -> None: + """Test that duplicate skills in skill_paths are deduplicated. + + This test verifies that: + 1. When skill_paths contains duplicates (same hash), only one is copied + 2. No concurrent copy conflicts occur + 3. The .skills file contains only one entry for the duplicated skill + """ + # Create a test skill + skill_dir = self._create_test_skill( + "test_skill_dedup", + "A test skill for deduplication testing", + ) + + # Create workspace with the same skill path listed multiple times + workspace = LocalWorkspace( + workdir=self.temp_dir.name, + skill_paths=[skill_dir, skill_dir, skill_dir], # Same path 3 times + ) + + # Initialize the workspace + await workspace.initialize() + + # Verify only one skill was copied + skills_dir = os.path.join(self.temp_dir.name, "skills") + skill_target = os.path.join(skills_dir, "test_skill_dedup") + self.assertTrue(os.path.exists(skill_target)) + + # Verify .skills file contains only one entry + skills_hash_file = os.path.join(skills_dir, ".skills") + self.assertTrue(os.path.exists(skills_hash_file)) + + async with aiofiles.open(skills_hash_file, "r") as f: + skills_data = json.loads(await f.read()) + + # Should have exactly one entry in the skills index + skills_index = skills_data["skills"] + self.assertEqual(len(skills_index), 1) + self.assertIn("test_skill_dedup", skills_index) + self.assertEqual( + skills_index["test_skill_dedup"]["skill_name"], + "test_skill_dedup", + ) + + async def test_initialize_invalid_skill(self) -> None: + """Test handling of invalid skills. + + This test verifies that: + 1. Skills without SKILL.md are not copied + 2. Skills with invalid frontmatter are not copied + 3. Valid skills are still copied correctly + """ + # Create a valid skill + valid_skill_dir = self._create_test_skill( + "valid_skill", + "A valid test skill", + ) + + # Create an invalid skill without SKILL.md + invalid_skill_no_md = os.path.join( + self.test_skills_dir.name, + "invalid_no_md", + ) + os.makedirs(invalid_skill_no_md, exist_ok=True) + with open( + os.path.join(invalid_skill_no_md, "tool.py"), + "w", + encoding="utf-8", + ) as f: + f.write("def tool():\n pass\n") + + # Create an invalid skill with malformed frontmatter + invalid_skill_bad_fm = os.path.join( + self.test_skills_dir.name, + "invalid_bad_fm", + ) + os.makedirs(invalid_skill_bad_fm, exist_ok=True) + with open( + os.path.join(invalid_skill_bad_fm, "SKILL.md"), + "w", + encoding="utf-8", + ) as f: + f.write( + "---\nname: missing_description\n---\n\nNo description field!", + ) + + # Create workspace with all skill paths + workspace = LocalWorkspace( + workdir=self.temp_dir.name, + skill_paths=[ + valid_skill_dir, + invalid_skill_no_md, + invalid_skill_bad_fm, + ], + ) + + # Initialize the workspace + await workspace.initialize() + + # Verify only the valid skill was copied + skills_dir = os.path.join(self.temp_dir.name, "skills") + self.assertTrue(os.path.exists(skills_dir)) + + # Verify valid skill exists + valid_target = os.path.join(skills_dir, "valid_skill") + self.assertTrue(os.path.exists(valid_target)) + + # Verify invalid skills do not exist + invalid_target_no_md = os.path.join(skills_dir, "invalid_no_md") + invalid_target_bad_fm = os.path.join(skills_dir, "invalid_bad_fm") + self.assertFalse(os.path.exists(invalid_target_no_md)) + self.assertFalse(os.path.exists(invalid_target_bad_fm)) + + async def test_list_skills(self) -> None: + """Test listing skills from workspace. + + This test verifies that: + 1. All skills in the workspace are correctly listed + 2. Each skill has the correct name, description, and directory + 3. The returned list matches the expected skills + """ + # Create test skills + skill1_dir = self._create_test_skill( + "list_skill_1", + "First skill for listing", + ) + skill2_dir = self._create_test_skill( + "list_skill_2", + "Second skill for listing", + ) + + # Create workspace and initialize + workspace = LocalWorkspace( + workdir=self.temp_dir.name, + skill_paths=[skill1_dir, skill2_dir], + ) + await workspace.initialize() + + # List skills + skills = await workspace.list_skills() + + # Verify the number of skills + self.assertEqual(len(skills), 2) + + # Sort skills by name for consistent comparison + skills_sorted = sorted(skills, key=lambda s: s.name) + + # Build expected skills for comparison + expected_skills = [ + { + "name": "list_skill_1", + "description": "First skill for listing", + "dir": skills_sorted[0].dir, # Use actual dir path + "markdown": skills_sorted[0].markdown, # Use actual markdown + "updated_at": skills_sorted[ + 0 + ].updated_at, # Use actual timestamp + }, + { + "name": "list_skill_2", + "description": "Second skill for listing", + "dir": skills_sorted[1].dir, # Use actual dir path + "markdown": skills_sorted[1].markdown, # Use actual markdown + "updated_at": skills_sorted[ + 1 + ].updated_at, # Use actual timestamp + }, + ] + + # Compare full skill objects using dataclasses.asdict + actual_skills = [asdict(skill) for skill in skills_sorted] + self.assertListEqual(actual_skills, expected_skills) + + async def test_list_skills_empty(self) -> None: + """Test listing skills when no skills exist. + + This test verifies that: + 1. An empty list is returned when no skills are in the workspace + 2. No errors are raised when the skills directory doesn't exist + """ + # Create workspace without initializing + workspace = LocalWorkspace(workdir=self.temp_dir.name) + + # List skills (should return empty list) + skills = await workspace.list_skills() + + # Verify empty list is returned + self.assertListEqual(skills, []) + + +class TestLocalWorkspaceWithAgent(IsolatedAsyncioTestCase): + """Test the local workspace class offloading with the agent.""" + + async def test_offload_tool_result(self) -> None: + """Test integration with the agent when offloading tool result. + + This test verifies that: + 1. A long tool result is split into a reserved part (kept in context) + and an offloaded part (written to disk). + 2. The reserved tool result block in the context is truncated and + contains a system reminder pointing to the offload file. + 3. The offloaded file contains the truncated remainder. + 4. A second reply with a fresh tool call produces a new offload file. + """ + with tempfile.TemporaryDirectory() as workdir: + session_id = "test_session" + model = MockModel(stream=False) + agent = Agent( + name="Friday", + system_prompt="You're a helpful assistant named Friday.", + model=model, + toolkit=Toolkit( + tools=[_LongResultTool()], + ), + context_config=ContextConfig( + tool_result_limit=50, + ), + offloader=LocalWorkspace( + workdir=workdir, + ), + state=AgentState(session_id=session_id), + ) + + model.set_responses( + mock_responses=[ + [ + ChatResponse( + content=[ + ToolCallBlock( + id="1", + name="long_result_tool", + input="{}", + ), + ], + is_last=True, + ), + ], + [ + ChatResponse( + content=[ + TextBlock(text="End_1."), + ], + is_last=True, + ), + ], + ], + ) + + await agent.reply() + + # === Assert offload file content === + offload_path_1 = os.path.join( + workdir, + "sessions", + session_id, + "tool_result-1.txt", + ) + self.assertTrue(os.path.exists(offload_path_1)) + async with aiofiles.open(offload_path_1, "r") as f: + offload_content = await f.read() + + # The base64 payload is hashed with sha256 and persisted under + # `{workdir}/data/{hash}.{ext}` with the decoded bytes. + b64_data = "AAECAwQF" + data_hash = hashlib.sha256(b64_data.encode()).hexdigest() + data_file_path = os.path.join( + workdir, + "data", + f"{data_hash}.png", + ) + self.assertTrue(os.path.exists(data_file_path)) + async with aiofiles.open(data_file_path, "rb") as f: + self.assertEqual(await f.read(), base64.b64decode(b64_data)) + + # The full text is "0" * 30000 followed by a base64 DataBlock; + # tool_result_limit=50 reserves ~200 chars of text in context, the + # remaining 29800 chars + the DataBlock placeholder are offloaded. + data_url = Path(data_file_path).as_uri() + expected_offload_content = ( + "0" * 29800 + f"" + ) + self.assertEqual(offload_content, expected_offload_content) + + # === Assert context content === + reminder_1 = ( + "\n<<>>\nThe remaining content " + "has been omitted for limited context. You can refer to the " + f"file in '{offload_path_1}' for the truncated content if " + "needed." + ) + expected_first_msg = { + "id": AnyString(), + "name": "Friday", + "role": "assistant", + "content": [ + { + "type": "tool_call", + "id": "1", + "name": "long_result_tool", + "input": "{}", + "state": "finished", + "suggested_rules": [], + }, + { + "type": "tool_result", + "id": "1", + "name": "long_result_tool", + "output": [ + { + "type": "text", + "text": "0" * 200 + reminder_1, + "id": AnyString(), + }, + ], + "state": "success", + "metadata": {}, + }, + { + "type": "text", + "text": "End_1.", + "id": AnyString(), + }, + ], + "metadata": {}, + "created_at": AnyString(), + "finished_at": None, + "usage": None, + } + self.assertListEqual( + [_.model_dump() for _ in agent.state.context], + [expected_first_msg], + ) + + async def test_offload_context(self) -> None: + """Test integration with the agent when offloading context. + + This test triggers context compression twice in the same session and + verifies that: + 1. The offload file ``context.jsonl`` is appended to (not + overwritten) across the two compressions. + 2. When the compressed context contains a base64-encoded + ``DataBlock``, the binary payload is persisted to a separate + data file and the offloaded JSON line references that file via a + ``URLSource`` instead of embedding the base64 inline. + 3. ``agent.state.summary`` is rewritten on every compression and + ends with a system-reminder pointing to the offload file. + 4. ``agent.state.context`` only retains the latest assistant reply. + + Note: compression triggers based on ``model.context_size`` together + with the default ``ContextConfig.trigger_ratio`` (0.8) — we set + ``context_size=100`` here so the threshold is just 80 tokens, and a + ~500-byte user message (~125 tokens) is enough to trigger + compression on each reply. The default ``ContextConfig`` is used. + """ + with tempfile.TemporaryDirectory() as workdir: + session_id = "test_session_ctx" + model = MockModel(stream=False, context_size=100) + agent = Agent( + name="Friday", + system_prompt="You're Friday.", + model=model, + toolkit=Toolkit(), + offloader=LocalWorkspace(workdir=workdir), + state=AgentState(session_id=session_id), + ) + + # The mock structured response is reused across both compression + # calls (same summary fields each time). + model.set_structured_response( + StructuredResponse( + content={ + "task_overview": "TASK", + "current_state": "STATE", + "important_discoveries": "DISCOVERIES", + "next_steps": "NEXT", + "context_to_preserve": "PRESERVE", + }, + ), + ) + + # Each reply yields a single final-text response (no tool calls). + model.set_responses( + mock_responses=[ + ChatResponse( + content=[TextBlock(text="End_1.")], + is_last=True, + ), + ChatResponse( + content=[TextBlock(text="End_2.")], + is_last=True, + ), + ], + ) + + offload_path = os.path.join( + workdir, + "sessions", + session_id, + "context.jsonl", + ) + + # ===== First reply ===== + # Build user_msg_a with **fixed** random fields (msg id, the + # content-block ids, timestamps) so the offloaded JSONL is a + # fully deterministic string we can assert against literally. + # + # Putting the DataBlock FIRST (before the long TextBlock) makes + # the boundary split include both blocks on the compress side, + # so the very first compression offloads a multimodal message — + # the DataBlock is rewritten to a URLSource alongside the + # original TextBlock, exercising the multimodal offload path. + b64_data = "AAECAwQF" + user_msg_a = UserMsg( + name="user", + content=[ + DataBlock( + id="data_block_a", + name="fake_image_a.png", + source=Base64Source( + data=b64_data, + media_type="image/png", + ), + ), + TextBlock(id="text_block_a", text="A" * 500), + ], + id="msg_a", + created_at="2026-01-01T00:00:00", + finished_at="2026-01-01T00:00:00", + ) + await agent.reply(user_msg_a) + + self.assertTrue(os.path.exists(offload_path)) + async with aiofiles.open(offload_path, "r") as f: + content_after_first = await f.read() + + # The DataBlock is persisted to ``{workdir}/data/`` as soon as + # it is included in an offloaded line — this happens during the + # first compression because both blocks land on the compress + # side of the boundary split. + data_hash = hashlib.sha256(b64_data.encode()).hexdigest() + data_file_path = os.path.join( + workdir, + "data", + f"{data_hash}.png", + ) + self.assertTrue(os.path.exists(data_file_path)) + async with aiofiles.open(data_file_path, "rb") as f: + self.assertEqual(await f.read(), base64.b64decode(b64_data)) + + # The single offloaded line carries user_msg_a with the + # DataBlock's source rewritten from ``Base64Source`` to + # ``URLSource`` (pointing at the persisted data file) while the + # TextBlock is preserved as-is. The expected JSONL is written + # literally so a developer can read off exactly what gets + # persisted; only the temp-dir-dependent file URL is + # interpolated via ``data_url``. + data_url = Path(data_file_path).as_uri() + expected_user_msg_a_offloaded_json = ( + '{"name":"user","content":[' + '{"type":"data","id":"data_block_a","source":' + '{"type":"url","url":"' + data_url + '",' + '"media_type":"image/png"},"name":"fake_image_a.png"},' + '{"type":"text","text":"' + "A" * 500 + '",' + '"id":"text_block_a"}' + '],"role":"user","id":"msg_a","metadata":{},' + '"created_at":"2026-01-01T00:00:00",' + '"finished_at":"2026-01-01T00:00:00","usage":null}' + ) + self.assertEqual( + content_after_first, + expected_user_msg_a_offloaded_json + "\n", + ) + + # ``state.context`` after the first compression is empty + # (msgs_to_reserve is empty since both content blocks of + # user_msg_a went to the compress side). After reasoning, + # ``state.context[0]`` is the assistant's "End_1." reply. The + # assistant fields (msg id, text-block id, timestamps) are + # generated by the agent — we capture them here and substitute + # them into the expected string. + assistant_1 = agent.state.context[0] + + # ===== Second reply ===== + user_msg_b = UserMsg( + name="user", + content=[ + TextBlock(id="text_block_b", text="B" * 500), + ], + id="msg_b", + created_at="2026-01-02T00:00:00", + finished_at="2026-01-02T00:00:00", + ) + await agent.reply(user_msg_b) + + async with aiofiles.open(offload_path, "r") as f: + content_after_second = await f.read() + + # The second compression offloads ``assistant_1`` and + # ``user_msg_b``. The file is appended to (mode="a"), so it + # now contains 3 lines: the multimodal user_msg_a from the + # first compression, plus assistant_1 and user_msg_b from the + # second. + expected_assistant_1_json = ( + '{"name":"Friday","content":[' + '{"type":"text","text":"End_1.","id":"' + + assistant_1.content[0].id + + '"}' + '],"role":"assistant","id":"' + assistant_1.id + '",' + '"metadata":{},"created_at":"' + assistant_1.created_at + '",' + '"finished_at":null,"usage":null}' + ) + expected_user_msg_b_json = ( + '{"name":"user","content":[' + '{"type":"text","text":"' + "B" * 500 + '",' + '"id":"text_block_b"}' + '],"role":"user","id":"msg_b","metadata":{},' + '"created_at":"2026-01-02T00:00:00",' + '"finished_at":"2026-01-02T00:00:00","usage":null}' + ) + self.assertEqual( + content_after_second, + expected_user_msg_a_offloaded_json + + "\n" + + expected_assistant_1_json + + "\n" + + expected_user_msg_b_json + + "\n", + ) + + # ``state.summary`` is rewritten on every compression, so the + # final value is just one rendering of the summary template plus + # one offload pointer (both compressions wrote to the same file). + expected_summary = ( + "Here is a summary of your previous work\n" + "# Task Overview\n" + "TASK\n\n" + "# Current State\n" + "STATE\n\n" + "# Important Discoveries\n" + "DISCOVERIES\n\n" + "# Next Steps\n" + "NEXT\n\n" + "# Context to Preserve\n" + "PRESERVE\n" + f"The compressed context is offloaded " + f"to '{offload_path}', you can refer to it when needed." + f"" + ) + self.assertEqual(agent.state.summary, expected_summary) + + # ``state.context`` only retains the latest assistant text; + # everything else has been offloaded. + expected_second_assistant = { + "id": AnyString(), + "name": "Friday", + "role": "assistant", + "content": [ + { + "type": "text", + "text": "End_2.", + "id": AnyString(), + }, + ], + "metadata": {}, + "created_at": AnyString(), + "finished_at": None, + "usage": None, + } + self.assertListEqual( + [_.model_dump() for _ in agent.state.context], + [expected_second_assistant], + ) + + +class TestLocalWorkspaceMCPInit(IsolatedAsyncioTestCase): + """Test MCP loading error handling in LocalWorkspace.initialize(). + + Covers: + - Invalid entries in persisted .mcp are skipped (not crashing) + - Stateful MCP connection failures are skipped (not crashing) + - Valid MCPs still load despite invalid neighbours + """ + + async def asyncSetUp(self) -> None: + """Set up test fixtures.""" + # pylint: disable=consider-using-with + self.temp_dir = tempfile.TemporaryDirectory() + + async def asyncTearDown(self) -> None: + """Clean up test fixtures.""" + self.temp_dir.cleanup() + + async def _write_mcp_file(self, entries: list[dict]) -> str: + """Write a list of MCP config dicts to ``/.mcp``. + + Args: + entries: List of raw MCP config dicts. + + Returns: + The path to the written file. + """ + mcp_file = os.path.join(self.temp_dir.name, ".mcp") + async with aiofiles.open(mcp_file, "w", encoding="utf-8") as f: + await f.write(json.dumps(entries, indent=2, ensure_ascii=False)) + return mcp_file + + @staticmethod + def _make_http_mcp(name: str) -> dict: + """Return a valid stateless HTTP MCP entry.""" + return { + "name": name, + "is_stateful": False, + "mcp_config": { + "type": "http_mcp", + "url": "http://localhost:19999/nonexistent", + }, + "enable_tools": None, + "disable_tools": None, + "execution_timeout": None, + } + + @staticmethod + def _make_bad_stdio_mcp(name: str) -> dict: + """Return an invalid STDIO MCP entry (is_stateful=False).""" + return { + "name": name, + "is_stateful": False, + "mcp_config": { + "type": "stdio_mcp", + "command": "nonexistent_cmd", + }, + "enable_tools": None, + "disable_tools": None, + "execution_timeout": None, + } + + # ----------------------------------------------------------------- + # persisted .mcp + # ----------------------------------------------------------------- + + async def test_initialize_skips_bad_entry_keeps_good(self) -> None: + """A persisted .mcp with one bad entry should skip it and still + load the valid entry.""" + await self._write_mcp_file( + [ + self._make_bad_stdio_mcp("bad_one"), + self._make_http_mcp("good_one"), + ], + ) + + ws = LocalWorkspace(workdir=self.temp_dir.name) + await ws.initialize() + + mcps = await ws.list_mcps() + names = [m.name for m in mcps] + self.assertIn("good_one", names) + self.assertNotIn("bad_one", names) + + # ----------------------------------------------------------------- + # default_mcps + connect failure + # ----------------------------------------------------------------- + + async def test_initialize_connect_failure_removes_mcp(self) -> (None): + """A stateful MCP whose connect() raises should not crash + initialize() and should be removed from the MCP list.""" + ws = LocalWorkspace( + workdir=self.temp_dir.name, + default_mcps=[ + MCPClient( + name="will_fail_connect", + is_stateful=True, + mcp_config=StdioMCPConfig( + command="nonexistent_command_xyz", + ), + ), + ], + ) + await ws.initialize() + self.assertTrue(ws.is_alive) + names = [m.name for m in await ws.list_mcps()] + self.assertNotIn("will_fail_connect", names)