Spaces:
Paused
Paused
Upload folder using huggingface_hub (part 3)
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- src/agentscope/model/_openai_response/_models/gpt-5.4.yaml +22 -0
- src/agentscope/model/_openai_response/_models/gpt-5.5.yaml +26 -0
- src/agentscope/model/_openai_response/_models/o3.yaml +24 -0
- src/agentscope/model/_openai_response/_models/o4-mini.yaml +24 -0
- src/agentscope/model/_xai/__init__.py +9 -0
- src/agentscope/model/_xai/_model.py +455 -0
- src/agentscope/model/_xai/_models/grok-3-fast.yaml +19 -0
- src/agentscope/model/_xai/_models/grok-3-mini.yaml +17 -0
- src/agentscope/model/_xai/_models/grok-3.yaml +19 -0
- src/agentscope/model/_xai/_models/grok-4.3.yaml +19 -0
- src/agentscope/permission/__init__.py +18 -0
- src/agentscope/permission/_context.py +46 -0
- src/agentscope/permission/_decision.py +68 -0
- src/agentscope/permission/_engine.py +729 -0
- src/agentscope/permission/_rule.py +36 -0
- src/agentscope/permission/_types.py +102 -0
- src/agentscope/py.typed +0 -0
- src/agentscope/rag/__init__.py +35 -0
- src/agentscope/rag/_chunker/__init__.py +10 -0
- src/agentscope/rag/_chunker/_approx_token_chunker.py +172 -0
- src/agentscope/rag/_chunker/_base.py +62 -0
- src/agentscope/rag/_document.py +102 -0
- src/agentscope/rag/_knowledge.py +384 -0
- src/agentscope/rag/_parser/__init__.py +16 -0
- src/agentscope/rag/_parser/_base.py +116 -0
- src/agentscope/rag/_parser/_image.py +92 -0
- src/agentscope/rag/_parser/_pdf.py +92 -0
- src/agentscope/rag/_parser/_ppt.py +341 -0
- src/agentscope/rag/_parser/_text.py +125 -0
- src/agentscope/rag/_parser/_utils.py +99 -0
- src/agentscope/rag/_vdb/__init__.py +18 -0
- src/agentscope/rag/_vdb/_qdrant.py +392 -0
- src/agentscope/rag/_vdb/_vector_store.py +291 -0
- src/agentscope/skill/__init__.py +11 -0
- src/agentscope/skill/_base.py +29 -0
- src/agentscope/skill/_local_loader.py +171 -0
- src/agentscope/state/__init__.py +11 -0
- src/agentscope/state/_state.py +221 -0
- src/agentscope/state/_task.py +39 -0
- src/agentscope/tool/__init__.py +58 -0
- src/agentscope/tool/_adapters.py +394 -0
- src/agentscope/tool/_base.py +451 -0
- src/agentscope/tool/_builtin/__init__.py +26 -0
- src/agentscope/tool/_builtin/_backend.py +733 -0
- src/agentscope/tool/_builtin/_bash.py +778 -0
- src/agentscope/tool/_builtin/_bash_parser.py +884 -0
- src/agentscope/tool/_builtin/_edit.py +416 -0
- src/agentscope/tool/_builtin/_glob.py +305 -0
- src/agentscope/tool/_builtin/_grep.py +488 -0
- src/agentscope/tool/_builtin/_meta.py +130 -0
src/agentscope/model/_openai_response/_models/gpt-5.4.yaml
ADDED
|
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
name: gpt-5.4
|
| 2 |
+
label: GPT-5.4 (Responses API)
|
| 3 |
+
status: active
|
| 4 |
+
|
| 5 |
+
input_types:
|
| 6 |
+
- text/plain
|
| 7 |
+
- image/jpeg
|
| 8 |
+
- image/png
|
| 9 |
+
- image/gif
|
| 10 |
+
- image/webp
|
| 11 |
+
|
| 12 |
+
output_types:
|
| 13 |
+
- text/plain
|
| 14 |
+
|
| 15 |
+
context_size: 1050000
|
| 16 |
+
output_size: 128000
|
| 17 |
+
|
| 18 |
+
parameter_overrides:
|
| 19 |
+
max_tokens:
|
| 20 |
+
maximum: 128000
|
| 21 |
+
reasoning_effort:
|
| 22 |
+
hidden: true
|
src/agentscope/model/_openai_response/_models/gpt-5.5.yaml
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
name: gpt-5.5
|
| 2 |
+
label: GPT-5.5
|
| 3 |
+
status: active
|
| 4 |
+
|
| 5 |
+
input_types:
|
| 6 |
+
- text/plain
|
| 7 |
+
- image/jpeg
|
| 8 |
+
- image/png
|
| 9 |
+
- image/gif
|
| 10 |
+
- image/webp
|
| 11 |
+
- audio/mp3
|
| 12 |
+
- audio/wav
|
| 13 |
+
|
| 14 |
+
output_types:
|
| 15 |
+
- text/plain
|
| 16 |
+
|
| 17 |
+
context_size: 1050000
|
| 18 |
+
output_size: 128000
|
| 19 |
+
|
| 20 |
+
parameter_overrides:
|
| 21 |
+
max_tokens:
|
| 22 |
+
maximum: 128000
|
| 23 |
+
thinking_enable:
|
| 24 |
+
hidden: true
|
| 25 |
+
reasoning_effort:
|
| 26 |
+
hidden: true
|
src/agentscope/model/_openai_response/_models/o3.yaml
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
name: o3
|
| 2 |
+
label: o3 (Responses API)
|
| 3 |
+
status: active
|
| 4 |
+
|
| 5 |
+
input_types:
|
| 6 |
+
- text/plain
|
| 7 |
+
- application/x-thinking
|
| 8 |
+
- image/jpeg
|
| 9 |
+
- image/png
|
| 10 |
+
- image/gif
|
| 11 |
+
- image/webp
|
| 12 |
+
|
| 13 |
+
output_types:
|
| 14 |
+
- text/plain
|
| 15 |
+
- application/x-thinking
|
| 16 |
+
|
| 17 |
+
context_size: 200000
|
| 18 |
+
output_size: 100000
|
| 19 |
+
|
| 20 |
+
parameter_overrides:
|
| 21 |
+
max_tokens:
|
| 22 |
+
maximum: 100000
|
| 23 |
+
temperature:
|
| 24 |
+
hidden: true
|
src/agentscope/model/_openai_response/_models/o4-mini.yaml
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
name: o4-mini
|
| 2 |
+
label: o4-mini (Responses API)
|
| 3 |
+
status: active
|
| 4 |
+
|
| 5 |
+
input_types:
|
| 6 |
+
- text/plain
|
| 7 |
+
- application/x-thinking
|
| 8 |
+
- image/jpeg
|
| 9 |
+
- image/png
|
| 10 |
+
- image/gif
|
| 11 |
+
- image/webp
|
| 12 |
+
|
| 13 |
+
output_types:
|
| 14 |
+
- text/plain
|
| 15 |
+
- application/x-thinking
|
| 16 |
+
|
| 17 |
+
context_size: 200000
|
| 18 |
+
output_size: 100000
|
| 19 |
+
|
| 20 |
+
parameter_overrides:
|
| 21 |
+
max_tokens:
|
| 22 |
+
maximum: 100000
|
| 23 |
+
temperature:
|
| 24 |
+
hidden: true
|
src/agentscope/model/_xai/__init__.py
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The xAI LLM API modules."""
|
| 3 |
+
|
| 4 |
+
from ._model import XAICredential, XAIChatModel
|
| 5 |
+
|
| 6 |
+
__all__ = [
|
| 7 |
+
"XAICredential",
|
| 8 |
+
"XAIChatModel",
|
| 9 |
+
]
|
src/agentscope/model/_xai/_model.py
ADDED
|
@@ -0,0 +1,455 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The xAI chat model implementation using the official xai_sdk."""
|
| 3 |
+
from datetime import datetime
|
| 4 |
+
from typing import Any, AsyncGenerator, List, Literal, TYPE_CHECKING, Type
|
| 5 |
+
|
| 6 |
+
from pydantic import BaseModel, Field
|
| 7 |
+
|
| 8 |
+
from .._base import ChatModelBase, _TOOL_CHOICE_LITERAL_MODES
|
| 9 |
+
from .._model_response import ChatResponse
|
| 10 |
+
from .._model_usage import ChatUsage
|
| 11 |
+
from ...credential import XAICredential
|
| 12 |
+
from ...formatter import XAIChatFormatter
|
| 13 |
+
from ...message import (
|
| 14 |
+
Msg,
|
| 15 |
+
TextBlock,
|
| 16 |
+
ThinkingBlock,
|
| 17 |
+
ToolCallBlock,
|
| 18 |
+
)
|
| 19 |
+
from ...tool import ToolChoice
|
| 20 |
+
|
| 21 |
+
if TYPE_CHECKING:
|
| 22 |
+
from xai_sdk import AsyncClient
|
| 23 |
+
from xai_sdk.chat import Response
|
| 24 |
+
else:
|
| 25 |
+
AsyncClient = Any
|
| 26 |
+
Response = Any
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
class XAIChatModel(ChatModelBase):
|
| 30 |
+
"""The xAI chat model using the official ``xai_sdk`` gRPC client.
|
| 31 |
+
|
| 32 |
+
This model provides native access to xAI-specific features such as
|
| 33 |
+
server-side agentic tools (web search, X search, code execution) and
|
| 34 |
+
reasoning effort control, which are not available through the
|
| 35 |
+
OpenAI-compatible REST endpoint.
|
| 36 |
+
"""
|
| 37 |
+
|
| 38 |
+
class Parameters(BaseModel):
|
| 39 |
+
"""The parameters for the xAI chat model."""
|
| 40 |
+
|
| 41 |
+
max_tokens: int | None = Field(
|
| 42 |
+
default=None,
|
| 43 |
+
title="Max Tokens",
|
| 44 |
+
description="The maximum number of tokens for the LLM output.",
|
| 45 |
+
gt=0,
|
| 46 |
+
)
|
| 47 |
+
"""The maximum number of tokens to generate."""
|
| 48 |
+
|
| 49 |
+
thinking_enable: bool = Field(
|
| 50 |
+
default=False,
|
| 51 |
+
title="Thinking",
|
| 52 |
+
description=(
|
| 53 |
+
"Whether to enable reasoning for models that support "
|
| 54 |
+
"extended thinking (e.g. ``grok-3-mini``). Use "
|
| 55 |
+
"reasoning_effort to control the depth of reasoning."
|
| 56 |
+
),
|
| 57 |
+
)
|
| 58 |
+
|
| 59 |
+
reasoning_effort: Literal["low", "medium", "high"] | None = Field(
|
| 60 |
+
default=None,
|
| 61 |
+
title="Reasoning Effort",
|
| 62 |
+
description=(
|
| 63 |
+
"Controls the depth of reasoning for models that support "
|
| 64 |
+
"extended thinking (e.g. ``grok-3-mini``). Set to "
|
| 65 |
+
"``'low'``, ``'medium'``, or ``'high'`` to enable reasoning "
|
| 66 |
+
"with the corresponding effort level. ``None`` disables "
|
| 67 |
+
"reasoning."
|
| 68 |
+
),
|
| 69 |
+
)
|
| 70 |
+
|
| 71 |
+
temperature: float | None = Field(
|
| 72 |
+
default=None,
|
| 73 |
+
title="Temperature",
|
| 74 |
+
description="The temperature for the LLM output.",
|
| 75 |
+
ge=0,
|
| 76 |
+
le=2,
|
| 77 |
+
)
|
| 78 |
+
"""The sampling temperature."""
|
| 79 |
+
|
| 80 |
+
top_p: float | None = Field(
|
| 81 |
+
default=None,
|
| 82 |
+
title="Top P",
|
| 83 |
+
description="The top-p nucleus sampling value.",
|
| 84 |
+
gt=0,
|
| 85 |
+
le=1,
|
| 86 |
+
)
|
| 87 |
+
"""The top-p sampling parameter."""
|
| 88 |
+
|
| 89 |
+
type: Literal["xai_chat"] = "xai_chat"
|
| 90 |
+
"""The type of the chat model."""
|
| 91 |
+
|
| 92 |
+
def __init__(
|
| 93 |
+
self,
|
| 94 |
+
credential: XAICredential,
|
| 95 |
+
model: str,
|
| 96 |
+
parameters: "XAIChatModel.Parameters | None" = None,
|
| 97 |
+
stream: bool = True,
|
| 98 |
+
max_retries: int = 3,
|
| 99 |
+
retry_delay: float = 1.0,
|
| 100 |
+
context_size: int = 131072,
|
| 101 |
+
formatter: XAIChatFormatter | None = None,
|
| 102 |
+
client_kwargs: dict[str, Any] | None = None,
|
| 103 |
+
) -> None:
|
| 104 |
+
"""Initialize the xAI chat model.
|
| 105 |
+
|
| 106 |
+
Args:
|
| 107 |
+
credential (`XAICredential`):
|
| 108 |
+
The xAI credential used to authenticate API calls.
|
| 109 |
+
model (`str`):
|
| 110 |
+
The xAI model name, e.g. ``grok-3`` or ``grok-3-mini``.
|
| 111 |
+
parameters (`XAIChatModel.Parameters | None`, defaults to \
|
| 112 |
+
`None`):
|
| 113 |
+
The xAI API parameters. When ``None``, the default
|
| 114 |
+
parameters will be used.
|
| 115 |
+
stream (`bool`, defaults to `True`):
|
| 116 |
+
Whether to enable streaming output.
|
| 117 |
+
max_retries (`int`, defaults to `3`):
|
| 118 |
+
The maximum number of retries for the xAI API.
|
| 119 |
+
retry_delay (`float`, defaults to `1.0`):
|
| 120 |
+
Seconds to sleep between retry attempts.
|
| 121 |
+
context_size (`int`, defaults to `131072`):
|
| 122 |
+
The model context size used for context compression.
|
| 123 |
+
formatter (`XAIChatFormatter | None`, defaults to `None`):
|
| 124 |
+
The formatter that converts ``Msg`` objects to xai_sdk
|
| 125 |
+
proto messages. When ``None``, an ``XAIChatFormatter``
|
| 126 |
+
instance will be used.
|
| 127 |
+
client_kwargs (`dict[str, Any] | None`, defaults to `None`):
|
| 128 |
+
Extra keyword arguments forwarded to ``xai_sdk.AsyncClient``
|
| 129 |
+
(e.g. ``timeout``, ``metadata``, ``channel_options``). Keys
|
| 130 |
+
that overlap with credential-derived arguments (such as
|
| 131 |
+
``api_key`` or ``api_host``) take precedence over the
|
| 132 |
+
credential values.
|
| 133 |
+
"""
|
| 134 |
+
super().__init__(
|
| 135 |
+
credential=credential,
|
| 136 |
+
model=model,
|
| 137 |
+
parameters=parameters or self.Parameters(),
|
| 138 |
+
stream=stream,
|
| 139 |
+
max_retries=max_retries,
|
| 140 |
+
retry_delay=retry_delay,
|
| 141 |
+
context_size=context_size,
|
| 142 |
+
)
|
| 143 |
+
self.formatter = formatter or XAIChatFormatter()
|
| 144 |
+
self.client_kwargs = client_kwargs or {}
|
| 145 |
+
|
| 146 |
+
@classmethod
|
| 147 |
+
def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]:
|
| 148 |
+
import grpc
|
| 149 |
+
|
| 150 |
+
# xai_sdk uses grpc.aio under the hood; transport/API failures surface
|
| 151 |
+
# as grpc.aio.AioRpcError (subclass of grpc.RpcError). We retry the
|
| 152 |
+
# whole class because the retry mechanism here only filters by type,
|
| 153 |
+
# not by status code — so 4xx-equivalents (UNAUTHENTICATED, INVALID
|
| 154 |
+
# _ARGUMENT) will also be retried a few times before failing, which
|
| 155 |
+
# we accept. Note xai_sdk additionally retries UNAVAILABLE 5 times at
|
| 156 |
+
# the gRPC layer with exponential backoff before raising to us.
|
| 157 |
+
return (grpc.RpcError,)
|
| 158 |
+
|
| 159 |
+
async def _call_api(
|
| 160 |
+
self,
|
| 161 |
+
model_name: str,
|
| 162 |
+
messages: list[Msg],
|
| 163 |
+
tools: list[dict] | None = None,
|
| 164 |
+
tool_choice: ToolChoice | None = None,
|
| 165 |
+
**generate_kwargs: Any,
|
| 166 |
+
) -> ChatResponse | AsyncGenerator[ChatResponse, None]:
|
| 167 |
+
"""Call the xAI API using the official ``xai_sdk`` gRPC client.
|
| 168 |
+
|
| 169 |
+
Args:
|
| 170 |
+
model_name (`str`):
|
| 171 |
+
The model name to use for this call.
|
| 172 |
+
messages (`list`):
|
| 173 |
+
A list of ``Msg`` objects representing the conversation.
|
| 174 |
+
tools (`list[dict]`, default `None`):
|
| 175 |
+
The tools JSON schemas.
|
| 176 |
+
tool_choice (`ToolChoice | None`, optional):
|
| 177 |
+
Controls which (if any) tool is called by the model.
|
| 178 |
+
**generate_kwargs (`Any`):
|
| 179 |
+
Extra keyword arguments forwarded to the API.
|
| 180 |
+
|
| 181 |
+
Returns:
|
| 182 |
+
`ChatResponse | AsyncGenerator[ChatResponse, None]`:
|
| 183 |
+
A ``ChatResponse`` when streaming is disabled, or an async
|
| 184 |
+
generator of ``ChatResponse`` objects when streaming is
|
| 185 |
+
enabled.
|
| 186 |
+
"""
|
| 187 |
+
from xai_sdk import AsyncClient
|
| 188 |
+
|
| 189 |
+
client = AsyncClient(
|
| 190 |
+
**{
|
| 191 |
+
"api_key": self.credential.api_key.get_secret_value(),
|
| 192 |
+
"api_host": self.credential.api_host,
|
| 193 |
+
**self.client_kwargs,
|
| 194 |
+
},
|
| 195 |
+
)
|
| 196 |
+
|
| 197 |
+
xai_messages = await self.formatter.format(messages)
|
| 198 |
+
|
| 199 |
+
xai_tools, xai_tool_choice = self._format_tools(tools, tool_choice)
|
| 200 |
+
|
| 201 |
+
create_kwargs: dict[str, Any] = {"model": model_name}
|
| 202 |
+
if self.parameters.max_tokens is not None:
|
| 203 |
+
create_kwargs["max_tokens"] = self.parameters.max_tokens
|
| 204 |
+
if self.parameters.temperature is not None:
|
| 205 |
+
create_kwargs["temperature"] = self.parameters.temperature
|
| 206 |
+
if self.parameters.top_p is not None:
|
| 207 |
+
create_kwargs["top_p"] = self.parameters.top_p
|
| 208 |
+
if (
|
| 209 |
+
self.parameters.thinking_enable
|
| 210 |
+
and self.parameters.reasoning_effort
|
| 211 |
+
):
|
| 212 |
+
create_kwargs[
|
| 213 |
+
"reasoning_effort"
|
| 214 |
+
] = self.parameters.reasoning_effort
|
| 215 |
+
if xai_tools:
|
| 216 |
+
create_kwargs["tools"] = xai_tools
|
| 217 |
+
if xai_tool_choice is not None:
|
| 218 |
+
create_kwargs["tool_choice"] = xai_tool_choice
|
| 219 |
+
|
| 220 |
+
create_kwargs.update(generate_kwargs)
|
| 221 |
+
|
| 222 |
+
chat = client.chat.create(**create_kwargs)
|
| 223 |
+
for xai_msg in xai_messages:
|
| 224 |
+
chat.append(xai_msg)
|
| 225 |
+
|
| 226 |
+
start_datetime = datetime.now()
|
| 227 |
+
|
| 228 |
+
if self.stream:
|
| 229 |
+
return self._parse_stream_response(start_datetime, chat, client)
|
| 230 |
+
|
| 231 |
+
try:
|
| 232 |
+
response = await chat.sample()
|
| 233 |
+
finally:
|
| 234 |
+
await client.close()
|
| 235 |
+
|
| 236 |
+
return self._parse_completion_response(start_datetime, response)
|
| 237 |
+
|
| 238 |
+
def _format_tools(
|
| 239 |
+
self,
|
| 240 |
+
tools: list[dict] | None,
|
| 241 |
+
tool_choice: ToolChoice | None,
|
| 242 |
+
) -> tuple[list | None, Any]:
|
| 243 |
+
"""Validate, filter, and format tools and tool_choice for the xAI API.
|
| 244 |
+
|
| 245 |
+
When ``tool_choice.tools`` is specified the schemas list is filtered
|
| 246 |
+
to only those tools. When ``tool_choice.mode`` is a specific tool name
|
| 247 |
+
(str) the model is forced to call exactly that tool without needing to
|
| 248 |
+
filter the list, preserving prompt-cache efficiency.
|
| 249 |
+
|
| 250 |
+
Args:
|
| 251 |
+
tools (`list[dict] | None`, optional):
|
| 252 |
+
The raw tool schemas.
|
| 253 |
+
tool_choice (`ToolChoice | None`, optional):
|
| 254 |
+
The tool choice configuration.
|
| 255 |
+
|
| 256 |
+
Returns:
|
| 257 |
+
`tuple[list | None, Any]`:
|
| 258 |
+
A tuple of (xai_tools, xai_tool_choice) ready for the
|
| 259 |
+
``xai_sdk`` client.
|
| 260 |
+
"""
|
| 261 |
+
from xai_sdk.chat import required_tool, tool
|
| 262 |
+
|
| 263 |
+
if tool_choice and tools:
|
| 264 |
+
self._validate_tool_choice(tool_choice, tools)
|
| 265 |
+
if tool_choice.tools:
|
| 266 |
+
allowed = set(tool_choice.tools)
|
| 267 |
+
tools = [t for t in tools if t["function"]["name"] in allowed]
|
| 268 |
+
|
| 269 |
+
xai_tools = None
|
| 270 |
+
if tools:
|
| 271 |
+
xai_tools = []
|
| 272 |
+
for t in tools:
|
| 273 |
+
if t.get("type") == "function" and "function" in t:
|
| 274 |
+
fn = t["function"]
|
| 275 |
+
xai_tools.append(
|
| 276 |
+
tool(
|
| 277 |
+
name=fn["name"],
|
| 278 |
+
description=fn.get("description", ""),
|
| 279 |
+
parameters=fn.get("parameters", {}),
|
| 280 |
+
),
|
| 281 |
+
)
|
| 282 |
+
|
| 283 |
+
if not tool_choice:
|
| 284 |
+
return xai_tools, None
|
| 285 |
+
|
| 286 |
+
mode = tool_choice.mode
|
| 287 |
+
|
| 288 |
+
if mode in _TOOL_CHOICE_LITERAL_MODES:
|
| 289 |
+
return xai_tools, mode
|
| 290 |
+
|
| 291 |
+
return xai_tools, required_tool(mode)
|
| 292 |
+
|
| 293 |
+
async def _parse_stream_response(
|
| 294 |
+
self,
|
| 295 |
+
start_datetime: datetime,
|
| 296 |
+
chat: Any,
|
| 297 |
+
client: AsyncClient,
|
| 298 |
+
) -> AsyncGenerator[ChatResponse, None]:
|
| 299 |
+
"""Parse the xAI streaming response from ``xai_sdk``.
|
| 300 |
+
|
| 301 |
+
Args:
|
| 302 |
+
start_datetime (`datetime`):
|
| 303 |
+
The start datetime of the response generation.
|
| 304 |
+
chat (`Any`):
|
| 305 |
+
The ``xai_sdk`` chat session object.
|
| 306 |
+
client (`Any`):
|
| 307 |
+
The ``xai_sdk.AsyncClient`` instance; closed when the
|
| 308 |
+
generator is exhausted or abandoned.
|
| 309 |
+
|
| 310 |
+
Yields:
|
| 311 |
+
`ChatResponse`:
|
| 312 |
+
Incremental ``ChatResponse`` objects with ``is_last=False``
|
| 313 |
+
followed by a final one with ``is_last=True``.
|
| 314 |
+
"""
|
| 315 |
+
acc_text = TextBlock(text="")
|
| 316 |
+
acc_thinking = ThinkingBlock(thinking="")
|
| 317 |
+
last_response = None
|
| 318 |
+
response_id: str | None = None
|
| 319 |
+
|
| 320 |
+
try:
|
| 321 |
+
async for response, chunk in chat.stream():
|
| 322 |
+
if response_id is None:
|
| 323 |
+
response_id = getattr(response, "id", None) or None
|
| 324 |
+
|
| 325 |
+
delta_text: str = chunk.content or ""
|
| 326 |
+
delta_thinking: str = chunk.reasoning_content or ""
|
| 327 |
+
|
| 328 |
+
delta_contents: List[TextBlock | ThinkingBlock] = []
|
| 329 |
+
|
| 330 |
+
if delta_thinking:
|
| 331 |
+
acc_thinking.thinking += delta_thinking
|
| 332 |
+
delta_contents.append(
|
| 333 |
+
ThinkingBlock(
|
| 334 |
+
id=acc_thinking.id,
|
| 335 |
+
thinking=delta_thinking,
|
| 336 |
+
),
|
| 337 |
+
)
|
| 338 |
+
if delta_text:
|
| 339 |
+
acc_text.text += delta_text
|
| 340 |
+
delta_contents.append(
|
| 341 |
+
TextBlock(id=acc_text.id, text=delta_text),
|
| 342 |
+
)
|
| 343 |
+
|
| 344 |
+
if delta_contents:
|
| 345 |
+
_kwargs: dict[str, Any] = {
|
| 346 |
+
"content": delta_contents,
|
| 347 |
+
"is_last": False,
|
| 348 |
+
}
|
| 349 |
+
if response_id:
|
| 350 |
+
_kwargs["id"] = response_id
|
| 351 |
+
yield ChatResponse(**_kwargs)
|
| 352 |
+
|
| 353 |
+
last_response = response
|
| 354 |
+
|
| 355 |
+
finally:
|
| 356 |
+
await client.close()
|
| 357 |
+
|
| 358 |
+
final_contents: List[TextBlock | ToolCallBlock | ThinkingBlock] = []
|
| 359 |
+
if acc_thinking.thinking:
|
| 360 |
+
final_contents.append(acc_thinking)
|
| 361 |
+
if acc_text.text:
|
| 362 |
+
final_contents.append(acc_text)
|
| 363 |
+
|
| 364 |
+
if last_response is not None:
|
| 365 |
+
for tc in last_response.tool_calls or []:
|
| 366 |
+
final_contents.append(
|
| 367 |
+
ToolCallBlock(
|
| 368 |
+
id=tc.id,
|
| 369 |
+
name=tc.function.name,
|
| 370 |
+
input=tc.function.arguments,
|
| 371 |
+
),
|
| 372 |
+
)
|
| 373 |
+
|
| 374 |
+
usage = None
|
| 375 |
+
if last_response is not None and last_response.usage is not None:
|
| 376 |
+
u = last_response.usage
|
| 377 |
+
usage = ChatUsage(
|
| 378 |
+
input_tokens=u.prompt_tokens,
|
| 379 |
+
output_tokens=u.completion_tokens,
|
| 380 |
+
time=(datetime.now() - start_datetime).total_seconds(),
|
| 381 |
+
cache_input_tokens=getattr(
|
| 382 |
+
u,
|
| 383 |
+
"cached_prompt_text_tokens",
|
| 384 |
+
0,
|
| 385 |
+
),
|
| 386 |
+
)
|
| 387 |
+
|
| 388 |
+
final_kwargs: dict[str, Any] = {
|
| 389 |
+
"content": final_contents,
|
| 390 |
+
"usage": usage,
|
| 391 |
+
"is_last": True,
|
| 392 |
+
}
|
| 393 |
+
if response_id:
|
| 394 |
+
final_kwargs["id"] = response_id
|
| 395 |
+
yield ChatResponse(**final_kwargs)
|
| 396 |
+
|
| 397 |
+
def _parse_completion_response(
|
| 398 |
+
self,
|
| 399 |
+
start_datetime: datetime,
|
| 400 |
+
response: Response,
|
| 401 |
+
) -> ChatResponse:
|
| 402 |
+
"""Parse the xAI non-streaming response from ``xai_sdk``.
|
| 403 |
+
|
| 404 |
+
Args:
|
| 405 |
+
start_datetime (`datetime`):
|
| 406 |
+
The start datetime of the response generation.
|
| 407 |
+
response (`Any`):
|
| 408 |
+
The ``xai_sdk`` ``Response`` object.
|
| 409 |
+
|
| 410 |
+
Returns:
|
| 411 |
+
`ChatResponse`:
|
| 412 |
+
A single ``ChatResponse`` with ``is_last=True``.
|
| 413 |
+
"""
|
| 414 |
+
content_blocks: List[TextBlock | ToolCallBlock | ThinkingBlock] = []
|
| 415 |
+
|
| 416 |
+
if response.reasoning_content:
|
| 417 |
+
content_blocks.append(
|
| 418 |
+
ThinkingBlock(thinking=response.reasoning_content),
|
| 419 |
+
)
|
| 420 |
+
if response.content:
|
| 421 |
+
content_blocks.append(TextBlock(text=response.content))
|
| 422 |
+
|
| 423 |
+
for tc in response.tool_calls or []:
|
| 424 |
+
content_blocks.append(
|
| 425 |
+
ToolCallBlock(
|
| 426 |
+
id=tc.id,
|
| 427 |
+
name=tc.function.name,
|
| 428 |
+
input=tc.function.arguments,
|
| 429 |
+
),
|
| 430 |
+
)
|
| 431 |
+
|
| 432 |
+
usage = None
|
| 433 |
+
if response.usage is not None:
|
| 434 |
+
u = response.usage
|
| 435 |
+
usage = ChatUsage(
|
| 436 |
+
input_tokens=u.prompt_tokens,
|
| 437 |
+
output_tokens=u.completion_tokens,
|
| 438 |
+
time=(datetime.now() - start_datetime).total_seconds(),
|
| 439 |
+
cache_input_tokens=getattr(
|
| 440 |
+
u,
|
| 441 |
+
"cached_prompt_text_tokens",
|
| 442 |
+
0,
|
| 443 |
+
),
|
| 444 |
+
)
|
| 445 |
+
|
| 446 |
+
resp_kwargs: dict[str, Any] = {
|
| 447 |
+
"content": content_blocks,
|
| 448 |
+
"is_last": True,
|
| 449 |
+
"usage": usage,
|
| 450 |
+
}
|
| 451 |
+
response_id = getattr(response, "id", None)
|
| 452 |
+
if response_id:
|
| 453 |
+
resp_kwargs["id"] = response_id
|
| 454 |
+
|
| 455 |
+
return ChatResponse(**resp_kwargs)
|
src/agentscope/model/_xai/_models/grok-3-fast.yaml
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
name: grok-3-fast
|
| 2 |
+
label: Grok 3 Fast
|
| 3 |
+
status: active
|
| 4 |
+
|
| 5 |
+
input_types:
|
| 6 |
+
- text/plain
|
| 7 |
+
- image/jpeg
|
| 8 |
+
- image/png
|
| 9 |
+
|
| 10 |
+
output_types:
|
| 11 |
+
- text/plain
|
| 12 |
+
- application/x-thinking
|
| 13 |
+
|
| 14 |
+
context_size: 131072
|
| 15 |
+
output_size: 16000
|
| 16 |
+
|
| 17 |
+
parameter_overrides:
|
| 18 |
+
max_tokens:
|
| 19 |
+
maximum: 16000
|
src/agentscope/model/_xai/_models/grok-3-mini.yaml
ADDED
|
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
name: grok-3-mini
|
| 2 |
+
label: Grok 3 Mini
|
| 3 |
+
status: active
|
| 4 |
+
|
| 5 |
+
input_types:
|
| 6 |
+
- text/plain
|
| 7 |
+
|
| 8 |
+
output_types:
|
| 9 |
+
- text/plain
|
| 10 |
+
- application/x-thinking
|
| 11 |
+
|
| 12 |
+
context_size: 131072
|
| 13 |
+
output_size: 16000
|
| 14 |
+
|
| 15 |
+
parameter_overrides:
|
| 16 |
+
max_tokens:
|
| 17 |
+
maximum: 16000
|
src/agentscope/model/_xai/_models/grok-3.yaml
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
name: grok-3
|
| 2 |
+
label: Grok 3
|
| 3 |
+
status: active
|
| 4 |
+
|
| 5 |
+
input_types:
|
| 6 |
+
- text/plain
|
| 7 |
+
- image/jpeg
|
| 8 |
+
- image/png
|
| 9 |
+
|
| 10 |
+
output_types:
|
| 11 |
+
- text/plain
|
| 12 |
+
- application/x-thinking
|
| 13 |
+
|
| 14 |
+
context_size: 131072
|
| 15 |
+
output_size: 16000
|
| 16 |
+
|
| 17 |
+
parameter_overrides:
|
| 18 |
+
max_tokens:
|
| 19 |
+
maximum: 16000
|
src/agentscope/model/_xai/_models/grok-4.3.yaml
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
name: grok-4.3
|
| 2 |
+
label: Grok 4.3
|
| 3 |
+
status: active
|
| 4 |
+
|
| 5 |
+
input_types:
|
| 6 |
+
- text/plain
|
| 7 |
+
- image/jpeg
|
| 8 |
+
- image/png
|
| 9 |
+
|
| 10 |
+
output_types:
|
| 11 |
+
- text/plain
|
| 12 |
+
- application/x-thinking
|
| 13 |
+
|
| 14 |
+
context_size: 1000000
|
| 15 |
+
output_size: 30000
|
| 16 |
+
|
| 17 |
+
parameter_overrides:
|
| 18 |
+
max_tokens:
|
| 19 |
+
maximum: 30000
|
src/agentscope/permission/__init__.py
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The tool permission related types and functions."""
|
| 3 |
+
|
| 4 |
+
from ._context import PermissionContext, AdditionalWorkingDirectory
|
| 5 |
+
from ._decision import PermissionDecision
|
| 6 |
+
from ._engine import PermissionEngine
|
| 7 |
+
from ._rule import PermissionRule
|
| 8 |
+
from ._types import PermissionMode, PermissionBehavior
|
| 9 |
+
|
| 10 |
+
__all__ = [
|
| 11 |
+
"PermissionContext",
|
| 12 |
+
"AdditionalWorkingDirectory",
|
| 13 |
+
"PermissionDecision",
|
| 14 |
+
"PermissionEngine",
|
| 15 |
+
"PermissionRule",
|
| 16 |
+
"PermissionMode",
|
| 17 |
+
"PermissionBehavior",
|
| 18 |
+
]
|
src/agentscope/permission/_context.py
ADDED
|
@@ -0,0 +1,46 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The permission context module."""
|
| 3 |
+
from pydantic import BaseModel, Field
|
| 4 |
+
|
| 5 |
+
from ._rule import PermissionRule
|
| 6 |
+
from ._types import PermissionMode
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
class AdditionalWorkingDirectory(BaseModel):
|
| 10 |
+
"""An additional directory included in permission scope.
|
| 11 |
+
|
| 12 |
+
Working directories are used to determine which file paths should be
|
| 13 |
+
automatically allowed in ACCEPT_EDITS mode.
|
| 14 |
+
"""
|
| 15 |
+
|
| 16 |
+
path: str
|
| 17 |
+
"""Absolute path to the directory."""
|
| 18 |
+
|
| 19 |
+
source: str
|
| 20 |
+
"""Where this directory permission originated from
|
| 21 |
+
(e.g., 'userSettings', 'session')."""
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
class PermissionContext(BaseModel):
|
| 25 |
+
"""Context for permission checking.
|
| 26 |
+
|
| 27 |
+
Contains the permission mode, working directories, and all configured
|
| 28 |
+
permission rules organized by behavior type (allow, deny, ask).
|
| 29 |
+
"""
|
| 30 |
+
|
| 31 |
+
mode: PermissionMode = PermissionMode.DEFAULT
|
| 32 |
+
"""The current permission mode."""
|
| 33 |
+
|
| 34 |
+
working_directories: dict[str, AdditionalWorkingDirectory] = Field(
|
| 35 |
+
default_factory=dict,
|
| 36 |
+
)
|
| 37 |
+
"""Additional directories allowed for file operations, keyed by path."""
|
| 38 |
+
|
| 39 |
+
allow_rules: dict[str, list[PermissionRule]] = Field(default_factory=dict)
|
| 40 |
+
"""Rules that allow tool execution, keyed by tool name."""
|
| 41 |
+
|
| 42 |
+
deny_rules: dict[str, list[PermissionRule]] = Field(default_factory=dict)
|
| 43 |
+
"""Rules that deny tool execution, keyed by tool name."""
|
| 44 |
+
|
| 45 |
+
ask_rules: dict[str, list[PermissionRule]] = Field(default_factory=dict)
|
| 46 |
+
"""Rules that require user confirmation, keyed by tool name."""
|
src/agentscope/permission/_decision.py
ADDED
|
@@ -0,0 +1,68 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The permission decision result."""
|
| 3 |
+
from dataclasses import dataclass
|
| 4 |
+
from typing import Any
|
| 5 |
+
|
| 6 |
+
from ._rule import PermissionRule
|
| 7 |
+
from ._types import PermissionBehavior
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
@dataclass
|
| 11 |
+
class PermissionDecision:
|
| 12 |
+
"""Decision result from permission checking.
|
| 13 |
+
|
| 14 |
+
Represents the outcome of a permission check, including whether
|
| 15 |
+
the action should be allowed, denied, or require user confirmation.
|
| 16 |
+
"""
|
| 17 |
+
|
| 18 |
+
behavior: PermissionBehavior
|
| 19 |
+
"""The permission behavior decision."""
|
| 20 |
+
|
| 21 |
+
message: str
|
| 22 |
+
"""Human-readable message describing the decision."""
|
| 23 |
+
|
| 24 |
+
decision_reason: str | None = None
|
| 25 |
+
"""Optional explanation for why this decision was made."""
|
| 26 |
+
|
| 27 |
+
updated_input: dict[str, Any] | None = None
|
| 28 |
+
"""Optional modified input data (e.g., sanitized paths)."""
|
| 29 |
+
|
| 30 |
+
suggested_rules: list[PermissionRule] | None = None
|
| 31 |
+
"""Optional list of suggested permission rules for user to apply."""
|
| 32 |
+
|
| 33 |
+
bypass_immune: bool = False
|
| 34 |
+
"""Whether this decision is immune to being silenced by allow rules
|
| 35 |
+
("bypass-immune").
|
| 36 |
+
|
| 37 |
+
Only meaningful when :attr:`behavior` is :attr:`PermissionBehavior.ASK`.
|
| 38 |
+
A tool sets this to ``True`` to signal that the operation is
|
| 39 |
+
dangerous enough that **no allow rule** may convert the ASK into an
|
| 40 |
+
ALLOW — the user must explicitly confirm in-the-moment. In
|
| 41 |
+
:attr:`PermissionMode.DONT_ASK` where no user is available, a
|
| 42 |
+
bypass-immune ASK is converted to DENY rather than silently allowed.
|
| 43 |
+
|
| 44 |
+
Per-mode handling of a ``bypass_immune=True`` ASK:
|
| 45 |
+
|
| 46 |
+
- ``DEFAULT`` / ``ACCEPT_EDITS``: honored — allow rules cannot
|
| 47 |
+
override.
|
| 48 |
+
- ``EXPLORE``: not applicable (the engine resolves EXPLORE via
|
| 49 |
+
:meth:`ToolBase.check_read_only` and does not invoke
|
| 50 |
+
:meth:`ToolBase.check_permissions`).
|
| 51 |
+
- ``BYPASS``: **intentionally ignored** — BYPASS's contract is
|
| 52 |
+
"the user has opted out of safety prompts; only deny / ask
|
| 53 |
+
rules remain as guardrails." Use deny rules in BYPASS to
|
| 54 |
+
enforce specific protections.
|
| 55 |
+
- ``DONT_ASK``: converted to DENY (no user available).
|
| 56 |
+
|
| 57 |
+
Default is ``False``: a regular ASK that may be overridden by an
|
| 58 |
+
allow rule in DEFAULT / ACCEPT_EDITS, and is silently allowed by
|
| 59 |
+
BYPASS's fallback. Tools should set this only for genuine safety
|
| 60 |
+
checks (e.g. writes to dangerous paths, ``rm -rf /``, command
|
| 61 |
+
injection patterns) — not for "I'd prefer user input" cases.
|
| 62 |
+
|
| 63 |
+
Note: this field is internal metadata for the permission engine.
|
| 64 |
+
Callers handling the decision (agent loop, HITL backend, UI) treat
|
| 65 |
+
a bypass-immune ASK the same as a regular ASK — both prompt the
|
| 66 |
+
user. The distinction only governs whether engine-level rules /
|
| 67 |
+
modes may override it before reaching the caller.
|
| 68 |
+
"""
|
src/agentscope/permission/_engine.py
ADDED
|
@@ -0,0 +1,729 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The permission engine for checking and enforcing permission rules."""
|
| 3 |
+
from typing import Any, List, TYPE_CHECKING
|
| 4 |
+
|
| 5 |
+
from ._context import PermissionContext
|
| 6 |
+
from ._rule import PermissionRule
|
| 7 |
+
from ._decision import PermissionDecision, PermissionBehavior
|
| 8 |
+
from ._types import PermissionMode
|
| 9 |
+
from .._utils._common import _execute_async_or_sync_func
|
| 10 |
+
|
| 11 |
+
if TYPE_CHECKING:
|
| 12 |
+
from ..tool import ToolBase
|
| 13 |
+
else:
|
| 14 |
+
ToolBase = "ToolBase"
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
class PermissionEngine:
|
| 18 |
+
"""Engine for checking and enforcing permission rules.
|
| 19 |
+
|
| 20 |
+
Evaluates tool execution requests against configured permission rules.
|
| 21 |
+
Matching strategy is delegated to each tool's :meth:`ToolBase.match_rule`:
|
| 22 |
+
|
| 23 |
+
- Bash tools: substring / prefix wildcard matching against the command
|
| 24 |
+
- Write/Read/Edit tools: glob matching against file paths
|
| 25 |
+
- Other tools: generic pattern matching (or tool-name-level only)
|
| 26 |
+
|
| 27 |
+
Each :class:`PermissionMode` has its own ``_check_<mode>`` method so
|
| 28 |
+
that mode policies are self-contained and readable in isolation. See
|
| 29 |
+
:meth:`check_permission` for the dispatcher and the individual methods
|
| 30 |
+
for per-mode evaluation order.
|
| 31 |
+
"""
|
| 32 |
+
|
| 33 |
+
def __init__(
|
| 34 |
+
self,
|
| 35 |
+
context: PermissionContext,
|
| 36 |
+
) -> None:
|
| 37 |
+
"""Initialize the permission engine.
|
| 38 |
+
|
| 39 |
+
Args:
|
| 40 |
+
context (`PermissionContext`):
|
| 41 |
+
The permission context containing rules and mode
|
| 42 |
+
|
| 43 |
+
Example:
|
| 44 |
+
>>> context = PermissionContext(mode=PermissionMode.ACCEPT_EDITS)
|
| 45 |
+
>>> engine = PermissionEngine(context)
|
| 46 |
+
"""
|
| 47 |
+
self.context = context
|
| 48 |
+
|
| 49 |
+
def add_rule(self, rule: PermissionRule) -> None:
|
| 50 |
+
"""Add a permission rule to the context.
|
| 51 |
+
|
| 52 |
+
Args:
|
| 53 |
+
rule (`PermissionRule`):
|
| 54 |
+
The permission rule to add
|
| 55 |
+
|
| 56 |
+
Example:
|
| 57 |
+
>>> engine.add_rule(PermissionRule(
|
| 58 |
+
... tool_name="Bash",
|
| 59 |
+
... rule_content="git:*",
|
| 60 |
+
... behavior=PermissionBehavior.ALLOW,
|
| 61 |
+
... ))
|
| 62 |
+
"""
|
| 63 |
+
|
| 64 |
+
if rule.behavior == PermissionBehavior.ALLOW:
|
| 65 |
+
if rule.tool_name not in self.context.allow_rules:
|
| 66 |
+
self.context.allow_rules[rule.tool_name] = []
|
| 67 |
+
self.context.allow_rules[rule.tool_name].append(rule)
|
| 68 |
+
elif rule.behavior == PermissionBehavior.DENY:
|
| 69 |
+
if rule.tool_name not in self.context.deny_rules:
|
| 70 |
+
self.context.deny_rules[rule.tool_name] = []
|
| 71 |
+
self.context.deny_rules[rule.tool_name].append(rule)
|
| 72 |
+
elif rule.behavior == PermissionBehavior.ASK:
|
| 73 |
+
if rule.tool_name not in self.context.ask_rules:
|
| 74 |
+
self.context.ask_rules[rule.tool_name] = []
|
| 75 |
+
self.context.ask_rules[rule.tool_name].append(rule)
|
| 76 |
+
|
| 77 |
+
async def check_permission(
|
| 78 |
+
self,
|
| 79 |
+
tool: ToolBase,
|
| 80 |
+
tool_input: dict[str, Any],
|
| 81 |
+
) -> PermissionDecision:
|
| 82 |
+
"""Check permission for a tool execution request.
|
| 83 |
+
|
| 84 |
+
Dispatches to a per-mode private method so each mode's policy
|
| 85 |
+
is self-contained and readable in isolation:
|
| 86 |
+
|
| 87 |
+
- DEFAULT → :meth:`_check_default`
|
| 88 |
+
- EXPLORE → :meth:`_check_explore`
|
| 89 |
+
- ACCEPT_EDITS → :meth:`_check_accept_edits`
|
| 90 |
+
- BYPASS → :meth:`_check_bypass`
|
| 91 |
+
- DONT_ASK → :meth:`_check_dont_ask`
|
| 92 |
+
|
| 93 |
+
Args:
|
| 94 |
+
tool (`ToolBase`):
|
| 95 |
+
The tool instance being called.
|
| 96 |
+
tool_input (`dict[str, Any]`):
|
| 97 |
+
The tool input data, used for rule matching and
|
| 98 |
+
tool-specific checks.
|
| 99 |
+
|
| 100 |
+
Returns:
|
| 101 |
+
`PermissionDecision`:
|
| 102 |
+
Decision indicating whether to allow, deny, or ask.
|
| 103 |
+
"""
|
| 104 |
+
mode = self.context.mode
|
| 105 |
+
if mode == PermissionMode.DEFAULT:
|
| 106 |
+
return await self._check_default(tool, tool_input)
|
| 107 |
+
if mode == PermissionMode.EXPLORE:
|
| 108 |
+
return await self._check_explore(tool, tool_input)
|
| 109 |
+
if mode == PermissionMode.ACCEPT_EDITS:
|
| 110 |
+
return await self._check_accept_edits(tool, tool_input)
|
| 111 |
+
if mode == PermissionMode.BYPASS:
|
| 112 |
+
return await self._check_bypass(tool, tool_input)
|
| 113 |
+
if mode == PermissionMode.DONT_ASK:
|
| 114 |
+
return await self._check_dont_ask(tool, tool_input)
|
| 115 |
+
raise ValueError(f"Unknown permission mode: {mode}")
|
| 116 |
+
|
| 117 |
+
async def _check_default(
|
| 118 |
+
self,
|
| 119 |
+
tool: ToolBase,
|
| 120 |
+
tool_input: dict[str, Any],
|
| 121 |
+
) -> PermissionDecision:
|
| 122 |
+
"""Permission check for :attr:`PermissionMode.DEFAULT`.
|
| 123 |
+
|
| 124 |
+
Every operation requires explicit permission unless either an
|
| 125 |
+
allow rule matches or the tool's own ``check_permissions``
|
| 126 |
+
explicitly returns ALLOW (e.g. ``Bash`` auto-allows recognized
|
| 127 |
+
read-only commands like ``ls``/``git status``). Evaluation order:
|
| 128 |
+
|
| 129 |
+
1. Deny rules → DENY
|
| 130 |
+
2. Ask rules → ASK (with suggestions)
|
| 131 |
+
3. ``tool.check_permissions``:
|
| 132 |
+
- ALLOW / DENY → returned as-is
|
| 133 |
+
- Safety ASK (bypass-immune) → returned with suggestions; cannot
|
| 134 |
+
be overridden by allow rules
|
| 135 |
+
- Non-safety ASK / PASSTHROUGH → continue
|
| 136 |
+
4. Allow rules → ALLOW
|
| 137 |
+
5. Default → ASK (with suggestions)
|
| 138 |
+
|
| 139 |
+
Args:
|
| 140 |
+
tool (`ToolBase`):
|
| 141 |
+
The tool instance being called.
|
| 142 |
+
tool_input (`dict[str, Any]`):
|
| 143 |
+
The tool input data.
|
| 144 |
+
|
| 145 |
+
Returns:
|
| 146 |
+
`PermissionDecision`:
|
| 147 |
+
The final decision.
|
| 148 |
+
"""
|
| 149 |
+
# step 1: deny rules — highest priority
|
| 150 |
+
deny = await self._check_deny_rules(tool, tool_input)
|
| 151 |
+
if deny:
|
| 152 |
+
return deny
|
| 153 |
+
|
| 154 |
+
# step 2: ask rules
|
| 155 |
+
ask = await self._check_ask_rules(tool, tool_input)
|
| 156 |
+
if ask:
|
| 157 |
+
ask.suggested_rules = await self._generate_suggestions(
|
| 158 |
+
tool,
|
| 159 |
+
tool_input,
|
| 160 |
+
)
|
| 161 |
+
return ask
|
| 162 |
+
|
| 163 |
+
# step 3: tool's own check_permissions
|
| 164 |
+
tool_decision = await tool.check_permissions(tool_input, self.context)
|
| 165 |
+
# step 3a: tool ALLOW / DENY returned as-is
|
| 166 |
+
if tool_decision.behavior in (
|
| 167 |
+
PermissionBehavior.ALLOW,
|
| 168 |
+
PermissionBehavior.DENY,
|
| 169 |
+
):
|
| 170 |
+
return tool_decision
|
| 171 |
+
# step 3b: safety ASK is bypass-immune — allow rules can't override
|
| 172 |
+
if self._is_safety_ask(tool_decision):
|
| 173 |
+
tool_decision.suggested_rules = await self._generate_suggestions(
|
| 174 |
+
tool,
|
| 175 |
+
tool_input,
|
| 176 |
+
)
|
| 177 |
+
return tool_decision
|
| 178 |
+
|
| 179 |
+
# step 4: allow rules
|
| 180 |
+
allow = await self._check_allow_rules(tool, tool_input)
|
| 181 |
+
if allow:
|
| 182 |
+
return allow
|
| 183 |
+
|
| 184 |
+
# step 5: default — ASK the user
|
| 185 |
+
default = PermissionDecision(
|
| 186 |
+
behavior=PermissionBehavior.ASK,
|
| 187 |
+
message=f"Permission required for {tool.name}",
|
| 188 |
+
decision_reason=f"Mode: {self.context.mode.value}",
|
| 189 |
+
)
|
| 190 |
+
default.suggested_rules = await self._generate_suggestions(
|
| 191 |
+
tool,
|
| 192 |
+
tool_input,
|
| 193 |
+
)
|
| 194 |
+
return default
|
| 195 |
+
|
| 196 |
+
async def _check_explore(
|
| 197 |
+
self,
|
| 198 |
+
tool: ToolBase,
|
| 199 |
+
tool_input: dict[str, Any],
|
| 200 |
+
) -> PermissionDecision:
|
| 201 |
+
"""Permission check for :attr:`PermissionMode.EXPLORE`.
|
| 202 |
+
|
| 203 |
+
Read-only mode — modifications are categorically denied. Evaluation
|
| 204 |
+
order:
|
| 205 |
+
|
| 206 |
+
1. Deny rules → DENY
|
| 207 |
+
2. Ask rules → ASK (with suggestions)
|
| 208 |
+
3. :meth:`ToolBase.check_read_only` (input-aware):
|
| 209 |
+
- True → ALLOW
|
| 210 |
+
- False → DENY
|
| 211 |
+
|
| 212 |
+
``tool.check_permissions`` is not invoked: EXPLORE is fully
|
| 213 |
+
resolved by the read-only verdict, so safety ASK paths (e.g.
|
| 214 |
+
``rm -rf /``) are subsumed into the broader DENY. Allow rules are
|
| 215 |
+
intentionally not consulted — EXPLORE's read-only guarantee
|
| 216 |
+
cannot be granted away by a user-configured rule.
|
| 217 |
+
|
| 218 |
+
Args:
|
| 219 |
+
tool (`ToolBase`):
|
| 220 |
+
The tool instance being called.
|
| 221 |
+
tool_input (`dict[str, Any]`):
|
| 222 |
+
The tool input data.
|
| 223 |
+
|
| 224 |
+
Returns:
|
| 225 |
+
`PermissionDecision`:
|
| 226 |
+
ALLOW for read-only invocations, DENY otherwise.
|
| 227 |
+
"""
|
| 228 |
+
# step 1: deny rules
|
| 229 |
+
deny = await self._check_deny_rules(tool, tool_input)
|
| 230 |
+
if deny:
|
| 231 |
+
return deny
|
| 232 |
+
|
| 233 |
+
# step 2: ask rules
|
| 234 |
+
ask = await self._check_ask_rules(tool, tool_input)
|
| 235 |
+
if ask:
|
| 236 |
+
ask.suggested_rules = await self._generate_suggestions(
|
| 237 |
+
tool,
|
| 238 |
+
tool_input,
|
| 239 |
+
)
|
| 240 |
+
return ask
|
| 241 |
+
|
| 242 |
+
# step 3: read-only verdict decides everything (ALLOW or DENY)
|
| 243 |
+
if await tool.check_read_only(tool_input):
|
| 244 |
+
return PermissionDecision(
|
| 245 |
+
behavior=PermissionBehavior.ALLOW,
|
| 246 |
+
message=(
|
| 247 |
+
f"Permission granted for {tool.name} "
|
| 248 |
+
f"(explore mode - read-only invocation)"
|
| 249 |
+
),
|
| 250 |
+
decision_reason="Explore mode allows read-only operations",
|
| 251 |
+
)
|
| 252 |
+
return PermissionDecision(
|
| 253 |
+
behavior=PermissionBehavior.DENY,
|
| 254 |
+
message=(
|
| 255 |
+
f"Permission denied for {tool.name} "
|
| 256 |
+
f"(explore mode is read-only)"
|
| 257 |
+
),
|
| 258 |
+
decision_reason="Explore mode does not allow modifications",
|
| 259 |
+
)
|
| 260 |
+
|
| 261 |
+
async def _check_accept_edits(
|
| 262 |
+
self,
|
| 263 |
+
tool: ToolBase,
|
| 264 |
+
tool_input: dict[str, Any],
|
| 265 |
+
) -> PermissionDecision:
|
| 266 |
+
"""Permission check for :attr:`PermissionMode.ACCEPT_EDITS`.
|
| 267 |
+
|
| 268 |
+
Edits within working directories are auto-allowed by each tool's
|
| 269 |
+
own ``check_permissions``; other operations follow the normal
|
| 270 |
+
flow. Evaluation order:
|
| 271 |
+
|
| 272 |
+
1. Deny rules → DENY
|
| 273 |
+
2. Ask rules → ASK (with suggestions)
|
| 274 |
+
3. :meth:`ToolBase.check_read_only` → True → ALLOW (fast path)
|
| 275 |
+
4. ``tool.check_permissions``:
|
| 276 |
+
- ALLOW (e.g. ``Write`` to a file in the working directory) /
|
| 277 |
+
DENY → returned as-is
|
| 278 |
+
- Safety ASK (bypass-immune) → returned with suggestions
|
| 279 |
+
- Non-safety ASK / PASSTHROUGH → continue
|
| 280 |
+
5. Allow rules → ALLOW
|
| 281 |
+
6. Default → ASK (with suggestions)
|
| 282 |
+
|
| 283 |
+
Args:
|
| 284 |
+
tool (`ToolBase`):
|
| 285 |
+
The tool instance being called.
|
| 286 |
+
tool_input (`dict[str, Any]`):
|
| 287 |
+
The tool input data.
|
| 288 |
+
|
| 289 |
+
Returns:
|
| 290 |
+
`PermissionDecision`:
|
| 291 |
+
The final decision.
|
| 292 |
+
"""
|
| 293 |
+
# step 1: deny rules
|
| 294 |
+
deny = await self._check_deny_rules(tool, tool_input)
|
| 295 |
+
if deny:
|
| 296 |
+
return deny
|
| 297 |
+
|
| 298 |
+
# step 2: ask rules
|
| 299 |
+
ask = await self._check_ask_rules(tool, tool_input)
|
| 300 |
+
if ask:
|
| 301 |
+
ask.suggested_rules = await self._generate_suggestions(
|
| 302 |
+
tool,
|
| 303 |
+
tool_input,
|
| 304 |
+
)
|
| 305 |
+
return ask
|
| 306 |
+
|
| 307 |
+
# step 3: read-only fast path — ALLOW without invoking the tool
|
| 308 |
+
if await tool.check_read_only(tool_input):
|
| 309 |
+
return PermissionDecision(
|
| 310 |
+
behavior=PermissionBehavior.ALLOW,
|
| 311 |
+
message=(
|
| 312 |
+
f"Permission granted for {tool.name} "
|
| 313 |
+
f"(accept edits mode - read-only invocation)"
|
| 314 |
+
),
|
| 315 |
+
decision_reason="Accept edits mode allows read-only "
|
| 316 |
+
"operations",
|
| 317 |
+
)
|
| 318 |
+
|
| 319 |
+
# step 4: tool's own check_permissions (working-directory check
|
| 320 |
+
# for Write/Edit, path-checked auto-allow for Bash, ...)
|
| 321 |
+
tool_decision = await tool.check_permissions(tool_input, self.context)
|
| 322 |
+
# step 4a: tool ALLOW / DENY returned as-is
|
| 323 |
+
if tool_decision.behavior in (
|
| 324 |
+
PermissionBehavior.ALLOW,
|
| 325 |
+
PermissionBehavior.DENY,
|
| 326 |
+
):
|
| 327 |
+
return tool_decision
|
| 328 |
+
# step 4b: safety ASK is bypass-immune
|
| 329 |
+
if self._is_safety_ask(tool_decision):
|
| 330 |
+
tool_decision.suggested_rules = await self._generate_suggestions(
|
| 331 |
+
tool,
|
| 332 |
+
tool_input,
|
| 333 |
+
)
|
| 334 |
+
return tool_decision
|
| 335 |
+
|
| 336 |
+
# step 5: allow rules
|
| 337 |
+
allow = await self._check_allow_rules(tool, tool_input)
|
| 338 |
+
if allow:
|
| 339 |
+
return allow
|
| 340 |
+
|
| 341 |
+
# step 6: default — ASK the user
|
| 342 |
+
default = PermissionDecision(
|
| 343 |
+
behavior=PermissionBehavior.ASK,
|
| 344 |
+
message=f"Permission required for {tool.name}",
|
| 345 |
+
decision_reason=f"Mode: {self.context.mode.value}",
|
| 346 |
+
)
|
| 347 |
+
default.suggested_rules = await self._generate_suggestions(
|
| 348 |
+
tool,
|
| 349 |
+
tool_input,
|
| 350 |
+
)
|
| 351 |
+
return default
|
| 352 |
+
|
| 353 |
+
async def _check_bypass(
|
| 354 |
+
self,
|
| 355 |
+
tool: ToolBase,
|
| 356 |
+
tool_input: dict[str, Any],
|
| 357 |
+
) -> PermissionDecision:
|
| 358 |
+
"""Permission check for :attr:`PermissionMode.BYPASS`.
|
| 359 |
+
|
| 360 |
+
BYPASS is the "fully trusted" mode: the user has explicitly
|
| 361 |
+
opted out of safety prompts. All tool-emitted safety ASKs
|
| 362 |
+
(``rm -rf /``, write to ``~/.bashrc``, command-injection
|
| 363 |
+
patterns, dangerous sed, etc.) are **skipped** — only
|
| 364 |
+
user-configured deny / ask rules and tool-emitted DENY remain
|
| 365 |
+
as guardrails. The :attr:`PermissionDecision.bypass_immune`
|
| 366 |
+
field has no effect in BYPASS by design.
|
| 367 |
+
|
| 368 |
+
Use BYPASS only for sandboxed / containerized environments or
|
| 369 |
+
when you fully trust the agent's behavior. For unattended
|
| 370 |
+
execution where safety still matters, use
|
| 371 |
+
:attr:`PermissionMode.DONT_ASK` instead — it converts safety
|
| 372 |
+
ASKs to DENY rather than skipping them.
|
| 373 |
+
|
| 374 |
+
Evaluation order:
|
| 375 |
+
|
| 376 |
+
1. Deny rules → DENY
|
| 377 |
+
2. Ask rules → ASK (with suggestions; honors explicit user intent)
|
| 378 |
+
3. ``tool.check_permissions``:
|
| 379 |
+
- ALLOW / DENY → returned as-is
|
| 380 |
+
- ASK (including bypass-immune safety ASKs) → falls through
|
| 381 |
+
- PASSTHROUGH → falls through
|
| 382 |
+
4. Allow rules → ALLOW
|
| 383 |
+
5. Fallback → ALLOW (BYPASS)
|
| 384 |
+
|
| 385 |
+
Args:
|
| 386 |
+
tool (`ToolBase`):
|
| 387 |
+
The tool instance being called.
|
| 388 |
+
tool_input (`dict[str, Any]`):
|
| 389 |
+
The tool input data.
|
| 390 |
+
|
| 391 |
+
Returns:
|
| 392 |
+
`PermissionDecision`:
|
| 393 |
+
The final decision.
|
| 394 |
+
"""
|
| 395 |
+
# step 1: deny rules
|
| 396 |
+
deny = await self._check_deny_rules(tool, tool_input)
|
| 397 |
+
if deny:
|
| 398 |
+
return deny
|
| 399 |
+
|
| 400 |
+
# step 2: ask rules (honor explicit user intent to be prompted)
|
| 401 |
+
ask = await self._check_ask_rules(tool, tool_input)
|
| 402 |
+
if ask:
|
| 403 |
+
ask.suggested_rules = await self._generate_suggestions(
|
| 404 |
+
tool,
|
| 405 |
+
tool_input,
|
| 406 |
+
)
|
| 407 |
+
return ask
|
| 408 |
+
|
| 409 |
+
# step 3: tool's own check_permissions — ALLOW / DENY returned;
|
| 410 |
+
# any ASK (including bypass-immune safety ASK) is intentionally
|
| 411 |
+
# NOT honored here, per BYPASS's "skip safety prompts" contract.
|
| 412 |
+
tool_decision = await tool.check_permissions(tool_input, self.context)
|
| 413 |
+
if tool_decision.behavior in (
|
| 414 |
+
PermissionBehavior.ALLOW,
|
| 415 |
+
PermissionBehavior.DENY,
|
| 416 |
+
):
|
| 417 |
+
return tool_decision
|
| 418 |
+
|
| 419 |
+
# step 4: allow rules
|
| 420 |
+
allow = await self._check_allow_rules(tool, tool_input)
|
| 421 |
+
if allow:
|
| 422 |
+
return allow
|
| 423 |
+
|
| 424 |
+
# step 5: bypass fallback — ALLOW everything else
|
| 425 |
+
return PermissionDecision(
|
| 426 |
+
behavior=PermissionBehavior.ALLOW,
|
| 427 |
+
message=f"Permission granted for {tool.name} (bypass mode)",
|
| 428 |
+
decision_reason="Bypass mode allows all operations",
|
| 429 |
+
)
|
| 430 |
+
|
| 431 |
+
async def _check_dont_ask(
|
| 432 |
+
self,
|
| 433 |
+
tool: ToolBase,
|
| 434 |
+
tool_input: dict[str, Any],
|
| 435 |
+
) -> PermissionDecision:
|
| 436 |
+
"""Permission check for :attr:`PermissionMode.DONT_ASK`.
|
| 437 |
+
|
| 438 |
+
Used when no user is available to answer prompts (scheduled
|
| 439 |
+
tasks, background runs). Invariant: this method must never
|
| 440 |
+
return :attr:`PermissionBehavior.ASK` — every code path that
|
| 441 |
+
would otherwise ASK is converted to DENY via
|
| 442 |
+
:meth:`_convert_ask_to_deny`. Evaluation order:
|
| 443 |
+
|
| 444 |
+
1. Deny rules → DENY
|
| 445 |
+
2. Ask rules → DENY (converted, with suggestions preserved)
|
| 446 |
+
3. ``tool.check_permissions``:
|
| 447 |
+
- ALLOW / DENY → returned as-is
|
| 448 |
+
- Safety ASK → DENY (converted, with suggestions preserved)
|
| 449 |
+
- Non-safety ASK / PASSTHROUGH → continue
|
| 450 |
+
4. Allow rules → ALLOW
|
| 451 |
+
5. Default → DENY (user not available to answer)
|
| 452 |
+
|
| 453 |
+
Args:
|
| 454 |
+
tool (`ToolBase`):
|
| 455 |
+
The tool instance being called.
|
| 456 |
+
tool_input (`dict[str, Any]`):
|
| 457 |
+
The tool input data.
|
| 458 |
+
|
| 459 |
+
Returns:
|
| 460 |
+
`PermissionDecision`:
|
| 461 |
+
The final decision (never ASK).
|
| 462 |
+
"""
|
| 463 |
+
# step 1: deny rules
|
| 464 |
+
deny = await self._check_deny_rules(tool, tool_input)
|
| 465 |
+
if deny:
|
| 466 |
+
return deny
|
| 467 |
+
|
| 468 |
+
# step 2: ask rules — converted to DENY (no user available)
|
| 469 |
+
ask = await self._check_ask_rules(tool, tool_input)
|
| 470 |
+
if ask:
|
| 471 |
+
ask.suggested_rules = await self._generate_suggestions(
|
| 472 |
+
tool,
|
| 473 |
+
tool_input,
|
| 474 |
+
)
|
| 475 |
+
return self._convert_ask_to_deny(tool, ask)
|
| 476 |
+
|
| 477 |
+
# step 3: tool's own check_permissions
|
| 478 |
+
tool_decision = await tool.check_permissions(tool_input, self.context)
|
| 479 |
+
# step 3a: tool ALLOW / DENY returned as-is
|
| 480 |
+
if tool_decision.behavior in (
|
| 481 |
+
PermissionBehavior.ALLOW,
|
| 482 |
+
PermissionBehavior.DENY,
|
| 483 |
+
):
|
| 484 |
+
return tool_decision
|
| 485 |
+
# step 3b: safety ASK converted to DENY (no user available)
|
| 486 |
+
if self._is_safety_ask(tool_decision):
|
| 487 |
+
tool_decision.suggested_rules = await self._generate_suggestions(
|
| 488 |
+
tool,
|
| 489 |
+
tool_input,
|
| 490 |
+
)
|
| 491 |
+
return self._convert_ask_to_deny(tool, tool_decision)
|
| 492 |
+
|
| 493 |
+
# step 4: allow rules
|
| 494 |
+
allow = await self._check_allow_rules(tool, tool_input)
|
| 495 |
+
if allow:
|
| 496 |
+
return allow
|
| 497 |
+
|
| 498 |
+
# step 5: default — DENY (no user available to confirm)
|
| 499 |
+
return PermissionDecision(
|
| 500 |
+
behavior=PermissionBehavior.DENY,
|
| 501 |
+
message=(
|
| 502 |
+
f"Permission denied for {tool.name} "
|
| 503 |
+
f"(dont_ask mode - user not available)"
|
| 504 |
+
),
|
| 505 |
+
decision_reason="User is not available to answer permission "
|
| 506 |
+
"prompts",
|
| 507 |
+
)
|
| 508 |
+
|
| 509 |
+
@staticmethod
|
| 510 |
+
def _convert_ask_to_deny(
|
| 511 |
+
tool: ToolBase,
|
| 512 |
+
ask_decision: PermissionDecision,
|
| 513 |
+
) -> PermissionDecision:
|
| 514 |
+
"""Convert an ASK decision into a DENY for DONT_ASK mode.
|
| 515 |
+
|
| 516 |
+
DONT_ASK's invariant is "never return ASK" — the user is not
|
| 517 |
+
available to answer prompts. This helper turns whatever produced
|
| 518 |
+
the ASK (an ASK rule, a safety check) into a DENY while
|
| 519 |
+
preserving traceability by carrying the original reason and
|
| 520 |
+
``suggested_rules`` forward; callers (e.g. a UI surfacing the
|
| 521 |
+
scheduled-task failure) can still show the user what rule they
|
| 522 |
+
could add to unblock the operation in the future.
|
| 523 |
+
|
| 524 |
+
Args:
|
| 525 |
+
tool (`ToolBase`):
|
| 526 |
+
The tool whose invocation is being denied.
|
| 527 |
+
ask_decision (`PermissionDecision`):
|
| 528 |
+
The original ASK decision to convert.
|
| 529 |
+
|
| 530 |
+
Returns:
|
| 531 |
+
`PermissionDecision`:
|
| 532 |
+
A DENY decision with the original ASK's reason and
|
| 533 |
+
suggestions attached.
|
| 534 |
+
"""
|
| 535 |
+
return PermissionDecision(
|
| 536 |
+
behavior=PermissionBehavior.DENY,
|
| 537 |
+
message=(
|
| 538 |
+
f"Permission denied for {tool.name} "
|
| 539 |
+
f"(dont_ask mode - ASK converted to DENY, "
|
| 540 |
+
f"user not available)"
|
| 541 |
+
),
|
| 542 |
+
decision_reason=(
|
| 543 |
+
f"DONT_ASK mode converted ASK to DENY. "
|
| 544 |
+
f"Original reason: {ask_decision.decision_reason}"
|
| 545 |
+
),
|
| 546 |
+
suggested_rules=ask_decision.suggested_rules,
|
| 547 |
+
)
|
| 548 |
+
|
| 549 |
+
@staticmethod
|
| 550 |
+
def _is_safety_ask(decision: PermissionDecision) -> bool:
|
| 551 |
+
"""Whether a decision is a bypass-immune safety ASK.
|
| 552 |
+
|
| 553 |
+
A safety ASK is an ASK that a tool has explicitly marked with
|
| 554 |
+
:attr:`PermissionDecision.bypass_immune` ``= True``. Tools emit
|
| 555 |
+
these for dangerous operations (e.g. write to ``~/.bashrc``,
|
| 556 |
+
``rm -rf /``, command-injection patterns) that must be surfaced
|
| 557 |
+
to the user regardless of allow rules in
|
| 558 |
+
``DEFAULT``/``ACCEPT_EDITS``. ``BYPASS`` mode intentionally
|
| 559 |
+
skips this check (see :meth:`_check_bypass`); ``DONT_ASK``
|
| 560 |
+
converts the ASK to DENY (see :meth:`_check_dont_ask`).
|
| 561 |
+
|
| 562 |
+
Args:
|
| 563 |
+
decision (`PermissionDecision`):
|
| 564 |
+
The decision returned by a tool's ``check_permissions``.
|
| 565 |
+
|
| 566 |
+
Returns:
|
| 567 |
+
`bool`:
|
| 568 |
+
True iff ``behavior == ASK`` and ``bypass_immune`` is set.
|
| 569 |
+
"""
|
| 570 |
+
return (
|
| 571 |
+
decision.behavior == PermissionBehavior.ASK
|
| 572 |
+
and decision.bypass_immune
|
| 573 |
+
)
|
| 574 |
+
|
| 575 |
+
async def _check_deny_rules(
|
| 576 |
+
self,
|
| 577 |
+
tool: ToolBase,
|
| 578 |
+
input_data: dict[str, Any],
|
| 579 |
+
) -> PermissionDecision | None:
|
| 580 |
+
"""Check if any deny rules match the request.
|
| 581 |
+
|
| 582 |
+
Args:
|
| 583 |
+
tool (`ToolBase`):
|
| 584 |
+
The tool instance being called
|
| 585 |
+
input_data (`dict[str, Any]`):
|
| 586 |
+
The tool input data
|
| 587 |
+
|
| 588 |
+
Returns:
|
| 589 |
+
`PermissionDecision | None`:
|
| 590 |
+
DENY decision if a rule matches, None otherwise
|
| 591 |
+
"""
|
| 592 |
+
rules = self.context.deny_rules.get(tool.name, [])
|
| 593 |
+
for rule in rules:
|
| 594 |
+
if await self._rule_matches(tool, rule, input_data):
|
| 595 |
+
return PermissionDecision(
|
| 596 |
+
behavior=PermissionBehavior.DENY,
|
| 597 |
+
message=f"Permission to use {tool.name} has been denied",
|
| 598 |
+
decision_reason=f"Rule: {rule.rule_content}",
|
| 599 |
+
)
|
| 600 |
+
return None
|
| 601 |
+
|
| 602 |
+
async def _check_ask_rules(
|
| 603 |
+
self,
|
| 604 |
+
tool: ToolBase,
|
| 605 |
+
input_data: dict[str, Any],
|
| 606 |
+
) -> PermissionDecision | None:
|
| 607 |
+
"""Check if any ask rules match the request.
|
| 608 |
+
|
| 609 |
+
Args:
|
| 610 |
+
tool (`ToolBase`):
|
| 611 |
+
The tool instance being called (used for tool-specific checks)
|
| 612 |
+
input_data (`dict[str, Any]`):
|
| 613 |
+
The tool input data
|
| 614 |
+
|
| 615 |
+
Returns:
|
| 616 |
+
`PermissionDecision | None`:
|
| 617 |
+
ASK decision if a rule matches, None otherwise
|
| 618 |
+
"""
|
| 619 |
+
rules = self.context.ask_rules.get(tool.name, [])
|
| 620 |
+
for rule in rules:
|
| 621 |
+
if await self._rule_matches(tool, rule, input_data):
|
| 622 |
+
return PermissionDecision(
|
| 623 |
+
behavior=PermissionBehavior.ASK,
|
| 624 |
+
message=f"Permission required for {tool.name}",
|
| 625 |
+
decision_reason=f"Rule: {rule.rule_content}",
|
| 626 |
+
)
|
| 627 |
+
return None
|
| 628 |
+
|
| 629 |
+
async def _check_allow_rules(
|
| 630 |
+
self,
|
| 631 |
+
tool: ToolBase,
|
| 632 |
+
input_data: dict[str, Any],
|
| 633 |
+
) -> PermissionDecision | None:
|
| 634 |
+
"""Check if any allow rules match the request.
|
| 635 |
+
|
| 636 |
+
Args:
|
| 637 |
+
tool (`ToolBase`):
|
| 638 |
+
The tool instance being called (used for tool-specific checks)
|
| 639 |
+
input_data (`dict[str, Any]`):
|
| 640 |
+
The tool input data
|
| 641 |
+
|
| 642 |
+
Returns:
|
| 643 |
+
`PermissionDecision | None`:
|
| 644 |
+
ALLOW decision if a rule matches, None otherwise
|
| 645 |
+
"""
|
| 646 |
+
rules = self.context.allow_rules.get(tool.name, [])
|
| 647 |
+
for rule in rules:
|
| 648 |
+
if await self._rule_matches(tool, rule, input_data):
|
| 649 |
+
return PermissionDecision(
|
| 650 |
+
behavior=PermissionBehavior.ALLOW,
|
| 651 |
+
message=f"Permission granted for {tool.name}",
|
| 652 |
+
updated_input=input_data,
|
| 653 |
+
)
|
| 654 |
+
return None
|
| 655 |
+
|
| 656 |
+
async def _rule_matches(
|
| 657 |
+
self,
|
| 658 |
+
tool: ToolBase,
|
| 659 |
+
rule: PermissionRule,
|
| 660 |
+
input_data: dict[str, Any],
|
| 661 |
+
) -> bool:
|
| 662 |
+
"""Check if a rule matches the input data.
|
| 663 |
+
|
| 664 |
+
The matching strategy depends on the tool type:
|
| 665 |
+
- Bash: Substring matching against the command
|
| 666 |
+
- Write/Read: Glob pattern matching against file paths
|
| 667 |
+
- Other: Generic pattern matching
|
| 668 |
+
|
| 669 |
+
Args:
|
| 670 |
+
rule (`PermissionRule`):
|
| 671 |
+
The permission rule to check
|
| 672 |
+
input_data (`dict[str, Any]`):
|
| 673 |
+
The tool input data
|
| 674 |
+
|
| 675 |
+
Returns:
|
| 676 |
+
`bool`:
|
| 677 |
+
True if the rule matches, False otherwise
|
| 678 |
+
"""
|
| 679 |
+
# Empty rule_content matches everything
|
| 680 |
+
if not rule.rule_content:
|
| 681 |
+
return True
|
| 682 |
+
|
| 683 |
+
# Try to use tool's match_rule method if available.
|
| 684 |
+
# ``_execute_async_or_sync_func`` keeps backward compatibility
|
| 685 |
+
# with third-party tools that still override match_rule with a
|
| 686 |
+
# sync ``def`` (the framework's signature is now ``async def``).
|
| 687 |
+
return await _execute_async_or_sync_func(
|
| 688 |
+
tool.match_rule,
|
| 689 |
+
rule.rule_content,
|
| 690 |
+
input_data,
|
| 691 |
+
)
|
| 692 |
+
|
| 693 |
+
async def _generate_suggestions(
|
| 694 |
+
self,
|
| 695 |
+
tool: ToolBase,
|
| 696 |
+
tool_input: dict[str, Any],
|
| 697 |
+
) -> List[PermissionRule]:
|
| 698 |
+
"""Generate suggested permission rules from a tool call.
|
| 699 |
+
|
| 700 |
+
This method analyzes the tool call and generates broader permission
|
| 701 |
+
suggestions that the user can apply to avoid future confirmations.
|
| 702 |
+
|
| 703 |
+
Strategy:
|
| 704 |
+
- For Bash: Extract command prefix (e.g., "npm run" -> "npm run:*")
|
| 705 |
+
- For File operations: Extract directory (e.g.,
|
| 706 |
+
"src/file.py" -> "src/**")
|
| 707 |
+
- For other tools: Generate exact match rule
|
| 708 |
+
|
| 709 |
+
Args:
|
| 710 |
+
tool (`ToolBase`):
|
| 711 |
+
The tool instance being called (used for tool-specific
|
| 712 |
+
suggestions)
|
| 713 |
+
tool_input (`dict[str, Any]`):
|
| 714 |
+
The tool input data (used for generating suggestions)
|
| 715 |
+
|
| 716 |
+
Returns:
|
| 717 |
+
`List[PermissionRule]`:
|
| 718 |
+
List of suggested permission rules (usually 1, max 5 for
|
| 719 |
+
compound commands)
|
| 720 |
+
"""
|
| 721 |
+
|
| 722 |
+
# Try to use tool's generate_suggestions method if available.
|
| 723 |
+
# ``_execute_async_or_sync_func`` keeps backward compatibility
|
| 724 |
+
# with third-party tools that still override this method with
|
| 725 |
+
# a sync ``def`` (the framework's signature is now ``async def``).
|
| 726 |
+
return await _execute_async_or_sync_func(
|
| 727 |
+
tool.generate_suggestions,
|
| 728 |
+
tool_input,
|
| 729 |
+
)
|
src/agentscope/permission/_rule.py
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""Permission rule model for tool usage."""
|
| 3 |
+
from pydantic import BaseModel
|
| 4 |
+
|
| 5 |
+
from ._types import PermissionBehavior
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
class PermissionRule(BaseModel):
|
| 9 |
+
"""Permission rule for tool usage.
|
| 10 |
+
|
| 11 |
+
A permission rule defines whether a specific tool or tool operation
|
| 12 |
+
should be allowed, denied, or require user confirmation. The
|
| 13 |
+
rule_content field has different semantics depending on the tool_name:
|
| 14 |
+
|
| 15 |
+
- For "Bash": rule_content is a substring pattern matched against the
|
| 16 |
+
command Example: rule_content="npm install" matches "npm install express"
|
| 17 |
+
|
| 18 |
+
- For "Write"/"Read": rule_content is a glob pattern matched against file
|
| 19 |
+
paths Example: rule_content="src/**" matches "src/main.py"
|
| 20 |
+
|
| 21 |
+
- For other tools: rule_content is a tool-specific filter pattern
|
| 22 |
+
"""
|
| 23 |
+
|
| 24 |
+
tool_name: str
|
| 25 |
+
"""The name of the tool this rule applies to (e.g., "Bash",
|
| 26 |
+
"Write", "Read")."""
|
| 27 |
+
|
| 28 |
+
rule_content: str | None
|
| 29 |
+
"""Optional filter pattern - semantics depend on tool_name."""
|
| 30 |
+
|
| 31 |
+
behavior: PermissionBehavior
|
| 32 |
+
"""The permission behavior ("allow", "deny", or "ask")."""
|
| 33 |
+
|
| 34 |
+
source: str
|
| 35 |
+
"""Where this rule originated from (e.g., "userSettings",
|
| 36 |
+
"projectSettings")."""
|
src/agentscope/permission/_types.py
ADDED
|
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
# pylint: disable=line-too-long
|
| 3 |
+
"""Permission system types and engine for tool usage control.
|
| 4 |
+
|
| 5 |
+
This module implements a permission system that controls tool execution based
|
| 6 |
+
on configurable rules. The permission system supports different matching
|
| 7 |
+
strategies depending on the tool type:
|
| 8 |
+
|
| 9 |
+
- For Bash tools: rule_content is a substring pattern matched against commands
|
| 10 |
+
- For Write/Read tools: rule_content is a glob pattern matched against file
|
| 11 |
+
paths
|
| 12 |
+
- For other tools: rule_content uses generic matching logic
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
from enum import Enum
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
class PermissionMode(Enum):
|
| 19 |
+
"""The mode of permission.
|
| 20 |
+
|
| 21 |
+
Permission modes control how the system handles tool execution requests.
|
| 22 |
+
Different modes are suitable for different scenarios:
|
| 23 |
+
|
| 24 |
+
+---------------+--------------------------------------------------+--------------------------------+
|
| 25 |
+
| Mode | Behavior | Use Case |
|
| 26 |
+
+===============+==================================================+================================+
|
| 27 |
+
| DEFAULT | Every operation asks for permission unless: | Default mode, most secure |
|
| 28 |
+
| | - an allow rule matches, OR | |
|
| 29 |
+
| | - the tool's ``check_permissions`` explicitly | |
|
| 30 |
+
| | returns ALLOW for the invocation (currently | |
|
| 31 |
+
| | only ``Bash`` auto-allows recognized | |
|
| 32 |
+
| | read-only commands such as ``ls``/``git | |
|
| 33 |
+
| | status``). Read/Glob/Grep return PASSTHROUGH | |
|
| 34 |
+
| | and fall through to the default ASK unless | |
|
| 35 |
+
| | an allow rule matches. | |
|
| 36 |
+
+---------------+--------------------------------------------------+--------------------------------+
|
| 37 |
+
| ACCEPT_EDITS | - Auto-allow file writes in working directories | User present, rapid iteration |
|
| 38 |
+
| | - Auto-allow file reads in working directories | development |
|
| 39 |
+
| | - Auto-allow filesystem commands (mkdir, rm, | |
|
| 40 |
+
| | mv, cp, ...) **only when all target paths** | |
|
| 41 |
+
| | **resolve inside a working directory** | |
|
| 42 |
+
| | - Other operations follow normal rules | |
|
| 43 |
+
+---------------+--------------------------------------------------+--------------------------------+
|
| 44 |
+
| EXPLORE | Read-only mode: | Exploring codebase, planning |
|
| 45 |
+
| | - Allow: read-only tools (``Read``/``Grep``/ | implementation |
|
| 46 |
+
| | ``Glob``) and read-only bash commands | |
|
| 47 |
+
| | (e.g. ``ls``, ``git status``) | |
|
| 48 |
+
| | - Deny: any modification tool / command | |
|
| 49 |
+
| | - User-configured DENY or ASK rules take | |
|
| 50 |
+
| | precedence over the read-only auto-allow | |
|
| 51 |
+
+---------------+--------------------------------------------------+--------------------------------+
|
| 52 |
+
| BYPASS | Skip all permission checks except explicit | Sandboxed environments |
|
| 53 |
+
| | user-configured deny / ask rules and tool | (container, VM), unattended |
|
| 54 |
+
| | DENY. **Safety ASKs from tools are NOT** | runs where you fully trust |
|
| 55 |
+
| | **enforced** — including ``rm -rf /``, writes | the agent |
|
| 56 |
+
| | to ``~/.bashrc``, command-injection patterns, | |
|
| 57 |
+
| | etc. Use deny rules to protect specific paths. | |
|
| 58 |
+
| | For unattended runs that still need safety, | |
|
| 59 |
+
| | prefer DONT_ASK. | |
|
| 60 |
+
+---------------+--------------------------------------------------+--------------------------------+
|
| 61 |
+
| DONT_ASK | Convert every ASK (including safety ASKs and | Scheduled tasks, background |
|
| 62 |
+
| | ASK-rule hits) to DENY. Safe-by-default for | execution when user is away |
|
| 63 |
+
| | unattended execution. | |
|
| 64 |
+
+---------------+--------------------------------------------------+--------------------------------+
|
| 65 |
+
|
| 66 |
+
Attributes:
|
| 67 |
+
DEFAULT: Default mode - explicit permission per action. The
|
| 68 |
+
only auto-allow path is the tool's own ``check_permissions``
|
| 69 |
+
returning ALLOW (currently just ``Bash`` for recognized
|
| 70 |
+
read-only commands like ``ls``/``git status``).
|
| 71 |
+
ACCEPT_EDITS: Accept edits mode - automatically allows file
|
| 72 |
+
edits within working directories (including filesystem
|
| 73 |
+
bash commands whose every target is in a working dir).
|
| 74 |
+
EXPLORE: Explore mode - read-only; modifications are denied.
|
| 75 |
+
BYPASS: Bypass mode - skips safety checks; relies on user
|
| 76 |
+
deny / ask rules as the only guardrail.
|
| 77 |
+
DONT_ASK: Don't ask mode - converts all ASK decisions to DENY
|
| 78 |
+
(for unattended execution).
|
| 79 |
+
""" # noqa: E501
|
| 80 |
+
|
| 81 |
+
DEFAULT = "default"
|
| 82 |
+
ACCEPT_EDITS = "accept_edits"
|
| 83 |
+
EXPLORE = "explore"
|
| 84 |
+
BYPASS = "bypass"
|
| 85 |
+
DONT_ASK = "dont_ask"
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
class PermissionBehavior(Enum):
|
| 89 |
+
"""The behavior of permission.
|
| 90 |
+
|
| 91 |
+
Attributes:
|
| 92 |
+
ALLOW: Allow the operation
|
| 93 |
+
DENY: Deny the operation
|
| 94 |
+
ASK: Ask the user for permission
|
| 95 |
+
PASSTHROUGH: Let the permission engine continue with rule matching
|
| 96 |
+
(used by tools to defer decision to the engine)
|
| 97 |
+
"""
|
| 98 |
+
|
| 99 |
+
ALLOW = "allow"
|
| 100 |
+
DENY = "deny"
|
| 101 |
+
ASK = "ask"
|
| 102 |
+
PASSTHROUGH = "passthrough"
|
src/agentscope/py.typed
ADDED
|
File without changes
|
src/agentscope/rag/__init__.py
ADDED
|
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The retrieval-augmented generation (RAG) module in AgentScope."""
|
| 3 |
+
|
| 4 |
+
from ._chunker import ApproxTokenChunker, ChunkerBase
|
| 5 |
+
from ._document import (
|
| 6 |
+
Section,
|
| 7 |
+
Chunk,
|
| 8 |
+
)
|
| 9 |
+
from ._parser import ImageParser, ParserBase, PDFParser, PPTParser, TextParser
|
| 10 |
+
from ._vdb import (
|
| 11 |
+
DocumentSummary,
|
| 12 |
+
VectorStoreBase,
|
| 13 |
+
VectorRecord,
|
| 14 |
+
VectorSearchResult,
|
| 15 |
+
QdrantStore,
|
| 16 |
+
)
|
| 17 |
+
from ._knowledge import KnowledgeBase
|
| 18 |
+
|
| 19 |
+
__all__ = [
|
| 20 |
+
"ApproxTokenChunker",
|
| 21 |
+
"ChunkerBase",
|
| 22 |
+
"Chunk",
|
| 23 |
+
"DocumentSummary",
|
| 24 |
+
"ImageParser",
|
| 25 |
+
"ParserBase",
|
| 26 |
+
"PDFParser",
|
| 27 |
+
"PPTParser",
|
| 28 |
+
"TextParser",
|
| 29 |
+
"Section",
|
| 30 |
+
"VectorStoreBase",
|
| 31 |
+
"VectorRecord",
|
| 32 |
+
"VectorSearchResult",
|
| 33 |
+
"QdrantStore",
|
| 34 |
+
"KnowledgeBase",
|
| 35 |
+
]
|
src/agentscope/rag/_chunker/__init__.py
ADDED
|
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""Chunker implementations for the RAG indexing pipeline."""
|
| 3 |
+
|
| 4 |
+
from ._approx_token_chunker import ApproxTokenChunker
|
| 5 |
+
from ._base import ChunkerBase
|
| 6 |
+
|
| 7 |
+
__all__ = [
|
| 8 |
+
"ApproxTokenChunker",
|
| 9 |
+
"ChunkerBase",
|
| 10 |
+
]
|
src/agentscope/rag/_chunker/_approx_token_chunker.py
ADDED
|
@@ -0,0 +1,172 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""A chunker that splits text by an approximate token count.
|
| 3 |
+
|
| 4 |
+
The token count is approximated as ``len(text.encode("utf-8")) // 4``,
|
| 5 |
+
which avoids a hard dependency on any tokenizer library while staying
|
| 6 |
+
within the right order of magnitude for most LLM tokenizers.
|
| 7 |
+
"""
|
| 8 |
+
from bisect import bisect_right
|
| 9 |
+
from itertools import accumulate
|
| 10 |
+
|
| 11 |
+
from ._base import ChunkerBase
|
| 12 |
+
from .._document import Chunk, Section
|
| 13 |
+
from ...message import TextBlock, DataBlock
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
class ApproxTokenChunker(ChunkerBase):
|
| 17 |
+
"""A chunker based on an approximate token counting strategy.
|
| 18 |
+
|
| 19 |
+
Text sections are sliced into pieces of at most ``chunk_size``
|
| 20 |
+
approximate tokens, with ``overlap`` approximate tokens shared
|
| 21 |
+
between two consecutive pieces. The token count of a string is
|
| 22 |
+
approximated as ``len(text.encode("utf-8")) // 4``, so no
|
| 23 |
+
tokenizer dependency is required.
|
| 24 |
+
|
| 25 |
+
Sections carrying a :class:`~agentscope.message.DataBlock`
|
| 26 |
+
(images, video, etc.) are passed through unchanged as a single
|
| 27 |
+
chunk.
|
| 28 |
+
|
| 29 |
+
.. note:: Chunks never span across two input Sections, as
|
| 30 |
+
required by :class:`ChunkerBase`.
|
| 31 |
+
"""
|
| 32 |
+
|
| 33 |
+
def __init__(self, chunk_size: int = 512, overlap: int = 50) -> None:
|
| 34 |
+
"""Initialize the approx token chunker.
|
| 35 |
+
|
| 36 |
+
Args:
|
| 37 |
+
chunk_size (`int`, defaults to `512`):
|
| 38 |
+
The maximum number of approximate tokens per chunk.
|
| 39 |
+
Must be a positive integer.
|
| 40 |
+
overlap (`int`, defaults to `50`):
|
| 41 |
+
The number of approximate tokens shared between two
|
| 42 |
+
consecutive chunks. Must be non-negative and smaller
|
| 43 |
+
than ``chunk_size``.
|
| 44 |
+
|
| 45 |
+
Raises:
|
| 46 |
+
`ValueError`:
|
| 47 |
+
If ``chunk_size`` is not positive, or ``overlap`` is
|
| 48 |
+
negative or not smaller than ``chunk_size``.
|
| 49 |
+
"""
|
| 50 |
+
if chunk_size <= 0:
|
| 51 |
+
raise ValueError(
|
| 52 |
+
f"chunk_size must be positive, got {chunk_size}.",
|
| 53 |
+
)
|
| 54 |
+
if overlap < 0 or overlap >= chunk_size:
|
| 55 |
+
raise ValueError(
|
| 56 |
+
"overlap must satisfy 0 <= overlap < chunk_size, "
|
| 57 |
+
f"got overlap={overlap}, chunk_size={chunk_size}.",
|
| 58 |
+
)
|
| 59 |
+
|
| 60 |
+
self.chunk_size = chunk_size
|
| 61 |
+
self.overlap = overlap
|
| 62 |
+
|
| 63 |
+
async def chunk(self, sections: list[Section]) -> list[Chunk]:
|
| 64 |
+
"""Chunk the input sections into smaller chunks based on an approx
|
| 65 |
+
token counting strategy.
|
| 66 |
+
|
| 67 |
+
Args:
|
| 68 |
+
sections (`list[Section]`):
|
| 69 |
+
A list of sections to chunk.
|
| 70 |
+
|
| 71 |
+
Returns:
|
| 72 |
+
`list[Chunk]`:
|
| 73 |
+
A list of chunks, with ``chunk_index`` numbered
|
| 74 |
+
``0..N-1`` and ``total_chunks`` set to ``N`` on every
|
| 75 |
+
chunk.
|
| 76 |
+
"""
|
| 77 |
+
chunks: list[Chunk] = []
|
| 78 |
+
for section in sections:
|
| 79 |
+
contents: list[TextBlock | DataBlock]
|
| 80 |
+
if isinstance(section.content, TextBlock):
|
| 81 |
+
contents = [
|
| 82 |
+
TextBlock(text=piece)
|
| 83 |
+
for piece in self._split_text(section.content.text)
|
| 84 |
+
]
|
| 85 |
+
else:
|
| 86 |
+
# DataBlock pass-through: never slice multimodal data
|
| 87 |
+
contents = [section.content]
|
| 88 |
+
|
| 89 |
+
chunks.extend(
|
| 90 |
+
Chunk(
|
| 91 |
+
content=content,
|
| 92 |
+
source=section.source,
|
| 93 |
+
chunk_index=0, # renumbered below
|
| 94 |
+
total_chunks=0, # renumbered below
|
| 95 |
+
metadata=dict(section.metadata),
|
| 96 |
+
)
|
| 97 |
+
for content in contents
|
| 98 |
+
)
|
| 99 |
+
|
| 100 |
+
for index, chunk in enumerate(chunks):
|
| 101 |
+
chunk.chunk_index = index
|
| 102 |
+
chunk.total_chunks = len(chunks)
|
| 103 |
+
|
| 104 |
+
return chunks
|
| 105 |
+
|
| 106 |
+
def _split_text(self, text: str) -> list[str]:
|
| 107 |
+
"""Split text into pieces of at most ``chunk_size`` approx tokens.
|
| 108 |
+
|
| 109 |
+
Consecutive pieces share approximately ``overlap`` tokens.
|
| 110 |
+
|
| 111 |
+
Args:
|
| 112 |
+
text (`str`):
|
| 113 |
+
The text to split.
|
| 114 |
+
|
| 115 |
+
Returns:
|
| 116 |
+
`list[str]`:
|
| 117 |
+
The text pieces, in document order.
|
| 118 |
+
"""
|
| 119 |
+
if self._approx_count_tokens(text) <= self.chunk_size:
|
| 120 |
+
return [text]
|
| 121 |
+
|
| 122 |
+
# Cumulative UTF-8 byte length after each character, so that
|
| 123 |
+
# the byte length of text[i:j] == byte_offsets[j] - byte_offsets[i]
|
| 124 |
+
byte_offsets = [0, *accumulate(len(c.encode("utf-8")) for c in text)]
|
| 125 |
+
|
| 126 |
+
chunk_bytes = self.chunk_size * 4
|
| 127 |
+
overlap_bytes = self.overlap * 4
|
| 128 |
+
|
| 129 |
+
pieces: list[str] = []
|
| 130 |
+
start = 0
|
| 131 |
+
while start < len(text):
|
| 132 |
+
# The largest end such that the slice fits the byte budget
|
| 133 |
+
end = (
|
| 134 |
+
bisect_right(
|
| 135 |
+
byte_offsets,
|
| 136 |
+
byte_offsets[start] + chunk_bytes,
|
| 137 |
+
)
|
| 138 |
+
- 1
|
| 139 |
+
)
|
| 140 |
+
# Always make progress, even for characters whose UTF-8
|
| 141 |
+
# encoding exceeds the budget on their own
|
| 142 |
+
end = max(end, start + 1)
|
| 143 |
+
pieces.append(text[start:end])
|
| 144 |
+
|
| 145 |
+
if end >= len(text):
|
| 146 |
+
break
|
| 147 |
+
|
| 148 |
+
# Step back by the overlap budget, ensuring forward progress
|
| 149 |
+
next_start = (
|
| 150 |
+
bisect_right(
|
| 151 |
+
byte_offsets,
|
| 152 |
+
byte_offsets[end] - overlap_bytes,
|
| 153 |
+
)
|
| 154 |
+
- 1
|
| 155 |
+
)
|
| 156 |
+
start = max(next_start, start + 1)
|
| 157 |
+
|
| 158 |
+
return pieces
|
| 159 |
+
|
| 160 |
+
@staticmethod
|
| 161 |
+
def _approx_count_tokens(text: str) -> int:
|
| 162 |
+
"""The approx count of tokens.
|
| 163 |
+
|
| 164 |
+
Args:
|
| 165 |
+
text (`str`):
|
| 166 |
+
The text to be counted.
|
| 167 |
+
|
| 168 |
+
Returns:
|
| 169 |
+
`int`:
|
| 170 |
+
The approx count of tokens.
|
| 171 |
+
"""
|
| 172 |
+
return len(text.encode("utf-8")) // 4
|
src/agentscope/rag/_chunker/_base.py
ADDED
|
@@ -0,0 +1,62 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""Abstract base class for chunkers.
|
| 3 |
+
|
| 4 |
+
A :class:`ChunkerBase` subclass takes the :class:`Section` list
|
| 5 |
+
produced by a :class:`~agentscope.rag.ParserBase` and splits the
|
| 6 |
+
content into final :class:`Chunk` objects suitable for embedding and
|
| 7 |
+
storage in a vector database.
|
| 8 |
+
|
| 9 |
+
Chunkers are **format-agnostic** — they operate on the unified
|
| 10 |
+
``TextBlock | DataBlock`` content carried in each Section. Long
|
| 11 |
+
:class:`TextBlock` content is sliced according to a chunking strategy
|
| 12 |
+
(by character count, by tokens, by semantic boundaries, etc.); short
|
| 13 |
+
text and :class:`DataBlock` content are passed through unchanged.
|
| 14 |
+
|
| 15 |
+
Chunkers **never combine content across Section boundaries**. This
|
| 16 |
+
guarantee preserves the structural metadata attached by the Parser
|
| 17 |
+
(page numbers, slide indices, embedded-image isolation, etc.).
|
| 18 |
+
"""
|
| 19 |
+
from abc import ABC, abstractmethod
|
| 20 |
+
|
| 21 |
+
from .._document import Chunk, Section
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
class ChunkerBase(ABC):
|
| 25 |
+
"""Abstract base class for chunkers.
|
| 26 |
+
|
| 27 |
+
Subclasses implement a specific chunking strategy (by character
|
| 28 |
+
count, by token count, by semantic boundary, etc.). The
|
| 29 |
+
chunker is configured once at construction time and reused
|
| 30 |
+
across many ``chunk()`` calls within the same knowledge base.
|
| 31 |
+
|
| 32 |
+
Subclasses must guarantee:
|
| 33 |
+
|
| 34 |
+
- **No cross-Section merging**: every output :class:`Chunk` is
|
| 35 |
+
derived from exactly one input :class:`Section`.
|
| 36 |
+
- **DataBlock pass-through**: a Section whose content is a
|
| 37 |
+
:class:`~agentscope.message.DataBlock` becomes a single Chunk
|
| 38 |
+
with the same content; multimodal data is never sliced.
|
| 39 |
+
- **Continuous indexing**: ``chunk_index`` runs from ``0`` to
|
| 40 |
+
``total_chunks - 1`` across the entire output list, even
|
| 41 |
+
when the input contains many Sections.
|
| 42 |
+
- **Consistent total_chunks**: every output Chunk carries the
|
| 43 |
+
same ``total_chunks`` value (the length of the output list).
|
| 44 |
+
- **Metadata inheritance**: each output Chunk's ``source`` and
|
| 45 |
+
``metadata`` are copied from its parent Section.
|
| 46 |
+
"""
|
| 47 |
+
|
| 48 |
+
@abstractmethod
|
| 49 |
+
async def chunk(self, sections: list[Section]) -> list[Chunk]:
|
| 50 |
+
"""Split a list of Sections into Chunks.
|
| 51 |
+
|
| 52 |
+
Args:
|
| 53 |
+
sections (`list[Section]`):
|
| 54 |
+
The Sections produced by a :class:`ParserBase`, in
|
| 55 |
+
document order.
|
| 56 |
+
|
| 57 |
+
Returns:
|
| 58 |
+
`list[Chunk]`:
|
| 59 |
+
The final chunks, in document order, with
|
| 60 |
+
``chunk_index`` numbered ``0..N-1`` and
|
| 61 |
+
``total_chunks`` set to ``N`` on every chunk.
|
| 62 |
+
"""
|
src/agentscope/rag/_document.py
ADDED
|
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""Data structures used in the RAG indexing pipeline.
|
| 3 |
+
|
| 4 |
+
The indexing pipeline has two stages, each producing its own
|
| 5 |
+
structured output:
|
| 6 |
+
|
| 7 |
+
1. :class:`Section` — produced by a :class:`ParserBase` from a raw
|
| 8 |
+
file. Each ``Section`` represents one "natural boundary" of the
|
| 9 |
+
source (a PDF page, a PPTX slide, an embedded image, a Markdown
|
| 10 |
+
heading section, etc.). A ``Chunker`` never combines content
|
| 11 |
+
across two ``Section`` instances, so ``Section`` is also a hard
|
| 12 |
+
boundary that prevents leakage of format-specific structure into
|
| 13 |
+
downstream chunks.
|
| 14 |
+
|
| 15 |
+
2. :class:`Chunk` — produced by a :class:`ChunkerBase` from one or
|
| 16 |
+
more ``Section`` instances. Each ``Chunk`` is the final unit
|
| 17 |
+
that gets embedded and inserted into the vector store.
|
| 18 |
+
|
| 19 |
+
Neither structure is persisted on its own — they are transient
|
| 20 |
+
in-memory carriers between pipeline stages. Persistence happens at
|
| 21 |
+
the :class:`~agentscope.rag.VectorRecord` and
|
| 22 |
+
``KnowledgeDocumentRecord`` layers.
|
| 23 |
+
"""
|
| 24 |
+
from typing import Any
|
| 25 |
+
|
| 26 |
+
from pydantic import BaseModel, Field
|
| 27 |
+
|
| 28 |
+
from ..message import TextBlock, DataBlock
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
class Section(BaseModel):
|
| 32 |
+
"""A single natural section produced by a :class:`ParserBase`.
|
| 33 |
+
|
| 34 |
+
A ``Section`` represents one logical region of the source file.
|
| 35 |
+
The :class:`ChunkerBase` guarantees that no resulting
|
| 36 |
+
:class:`Chunk` ever spans content from two different sections.
|
| 37 |
+
|
| 38 |
+
The granularity of a ``Section`` is format-specific:
|
| 39 |
+
|
| 40 |
+
- **PDF**: one section per page (plus separate sections for
|
| 41 |
+
embedded images).
|
| 42 |
+
- **PPTX**: one section per slide.
|
| 43 |
+
- **Markdown**: one section per top-level heading, or the entire
|
| 44 |
+
file if unstructured.
|
| 45 |
+
- **TXT / image / video**: one section for the whole file.
|
| 46 |
+
"""
|
| 47 |
+
|
| 48 |
+
content: TextBlock | DataBlock
|
| 49 |
+
"""The section content. Text sections use :class:`TextBlock`;
|
| 50 |
+
multimodal sections (images, video, etc.) use :class:`DataBlock`."""
|
| 51 |
+
|
| 52 |
+
source: str
|
| 53 |
+
"""The source filename (e.g. ``"report.pdf"``). Carried through
|
| 54 |
+
to every downstream :class:`Chunk` and into the vector store
|
| 55 |
+
metadata for citation / display."""
|
| 56 |
+
|
| 57 |
+
metadata: dict[str, Any] = Field(default_factory=dict)
|
| 58 |
+
"""Format-specific metadata written by the parser. Examples:
|
| 59 |
+
|
| 60 |
+
- PDFParser: ``{"page": 3}``
|
| 61 |
+
- PPTXParser: ``{"slide": 2}``
|
| 62 |
+
- ExcelParser: ``{"sheet": "Q3 Sales"}``
|
| 63 |
+
|
| 64 |
+
These keys are not part of any retrieval / pipeline contract —
|
| 65 |
+
they are passed through verbatim to the vector store metadata for
|
| 66 |
+
later citation. Each Chunk inherits this dict from its parent
|
| 67 |
+
Section.
|
| 68 |
+
"""
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
class Chunk(BaseModel):
|
| 72 |
+
"""A final indexable chunk produced by a :class:`ChunkerBase`.
|
| 73 |
+
|
| 74 |
+
Each ``Chunk`` corresponds to one record in the vector store.
|
| 75 |
+
The required structural fields (``source``, ``chunk_index``,
|
| 76 |
+
``total_chunks``) enable downstream features such as "expand
|
| 77 |
+
context around a hit" during retrieval.
|
| 78 |
+
"""
|
| 79 |
+
|
| 80 |
+
content: TextBlock | DataBlock
|
| 81 |
+
"""The chunk content (sliced from a text :class:`Section`, or a
|
| 82 |
+
multimodal :class:`DataBlock` passed through unchanged)."""
|
| 83 |
+
|
| 84 |
+
source: str
|
| 85 |
+
"""The source filename — inherited from the parent
|
| 86 |
+
:class:`Section`. Used for display / citation."""
|
| 87 |
+
|
| 88 |
+
chunk_index: int
|
| 89 |
+
"""The 0-based index of this chunk **within the document**.
|
| 90 |
+
Sequential across all sections of the same source file. Used to
|
| 91 |
+
locate neighbouring chunks for "context expansion" at query time.
|
| 92 |
+
"""
|
| 93 |
+
|
| 94 |
+
total_chunks: int
|
| 95 |
+
"""The total number of chunks produced from the same source file.
|
| 96 |
+
Together with :attr:`chunk_index` lets callers know whether a hit
|
| 97 |
+
is near the start / end of the document, and bounds the
|
| 98 |
+
expansion range."""
|
| 99 |
+
|
| 100 |
+
metadata: dict[str, Any] = Field(default_factory=dict)
|
| 101 |
+
"""Format-specific metadata inherited from the parent
|
| 102 |
+
:class:`Section`. See :attr:`Section.metadata`."""
|
src/agentscope/rag/_knowledge.py
ADDED
|
@@ -0,0 +1,384 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""Runtime handle for a single knowledge base.
|
| 3 |
+
|
| 4 |
+
A :class:`KnowledgeBase` instance is the **single algorithmic source of
|
| 5 |
+
truth** for talking to one knowledge base: it pairs an embedding model
|
| 6 |
+
with a vector-store collection (optionally scoped by a payload
|
| 7 |
+
``metadata_filter``) and exposes the four operations a caller ever
|
| 8 |
+
needs — :meth:`search`, :meth:`insert_document`,
|
| 9 |
+
:meth:`delete_document`, :meth:`list_documents`.
|
| 10 |
+
|
| 11 |
+
The handle is *narrow on purpose* — it carries the resolved runtime
|
| 12 |
+
state (embedding model + vector store + scope) and delegates every
|
| 13 |
+
operation to the bound :class:`VectorStoreBase`. Document parsing,
|
| 14 |
+
chunking, credential resolution, dimension policy validation, and
|
| 15 |
+
persistence of knowledge-base records all belong one layer up
|
| 16 |
+
(service-side :class:`KnowledgeBaseManagerBase` for hosted
|
| 17 |
+
deployments; the caller directly otherwise).
|
| 18 |
+
|
| 19 |
+
The backing collection is created on first use — each operation
|
| 20 |
+
transparently calls :meth:`ensure_collection`, which is itself
|
| 21 |
+
idempotent and memoised after the first success, so the only cost is
|
| 22 |
+
one extra round-trip on the very first call against a fresh
|
| 23 |
+
deployment.
|
| 24 |
+
|
| 25 |
+
``metadata_filter`` is the defense-in-depth scoping mechanism for
|
| 26 |
+
co-locating multiple logical knowledge bases inside the same physical
|
| 27 |
+
collection — typically multi-tenant deployments where every record
|
| 28 |
+
carries a ``{"tenant_id": "..."}`` payload. It is set once at
|
| 29 |
+
construction time and **always** applied: search/list never escape
|
| 30 |
+
it, and insert forces it onto every chunk's metadata so a malicious or
|
| 31 |
+
buggy parser cannot rebind a record into another scope.
|
| 32 |
+
"""
|
| 33 |
+
|
| 34 |
+
import asyncio
|
| 35 |
+
|
| 36 |
+
from ._document import Chunk
|
| 37 |
+
from ._vdb import VectorRecord, VectorSearchResult, VectorStoreBase
|
| 38 |
+
from .._utils._common import _generate_id
|
| 39 |
+
from ..embedding import EmbeddingModelBase
|
| 40 |
+
from ..message import DataBlock, TextBlock
|
| 41 |
+
from ._vdb import DocumentSummary
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
class KnowledgeBase:
|
| 45 |
+
"""Runtime handle for one knowledge base.
|
| 46 |
+
|
| 47 |
+
Binds an embedding model and a vector-store collection together so
|
| 48 |
+
callers can retrieve / insert / delete / list documents without
|
| 49 |
+
repeating the wiring. Cheap to construct (no I/O); the collection
|
| 50 |
+
itself is created lazily on the first operation, so a fresh
|
| 51 |
+
deployment "just works" without an explicit setup step.
|
| 52 |
+
|
| 53 |
+
.. code-block:: python
|
| 54 |
+
|
| 55 |
+
kb = KnowledgeBase(
|
| 56 |
+
name="company-handbook",
|
| 57 |
+
description="Internal HR and onboarding documents.",
|
| 58 |
+
embedding_model=embedding_model,
|
| 59 |
+
vector_store=vector_store,
|
| 60 |
+
collection="handbook",
|
| 61 |
+
)
|
| 62 |
+
await kb.insert_document(chunks)
|
| 63 |
+
results = await kb.search(["What is the PTO policy?"])
|
| 64 |
+
"""
|
| 65 |
+
|
| 66 |
+
name: str
|
| 67 |
+
"""Agent-oriented knowledge base name — used by tool descriptions
|
| 68 |
+
and frontend rendering."""
|
| 69 |
+
|
| 70 |
+
description: str
|
| 71 |
+
"""Agent-oriented knowledge base description — what this knowledge
|
| 72 |
+
base contains and when to retrieve from it."""
|
| 73 |
+
|
| 74 |
+
def __init__(
|
| 75 |
+
self,
|
| 76 |
+
name: str,
|
| 77 |
+
description: str,
|
| 78 |
+
embedding_model: EmbeddingModelBase,
|
| 79 |
+
vector_store: VectorStoreBase,
|
| 80 |
+
collection: str,
|
| 81 |
+
metadata_filter: dict | None = None,
|
| 82 |
+
) -> None:
|
| 83 |
+
"""Initialize the runtime handle.
|
| 84 |
+
|
| 85 |
+
Args:
|
| 86 |
+
name (`str`):
|
| 87 |
+
Agent-oriented knowledge base name. Surfaced to the
|
| 88 |
+
LLM (via tool descriptions) and to the front-end.
|
| 89 |
+
description (`str`):
|
| 90 |
+
Agent-oriented description. Should answer "what is in
|
| 91 |
+
this knowledge base and when should I search it?" — the
|
| 92 |
+
LLM uses it to decide whether to call the search tool
|
| 93 |
+
in agentic mode.
|
| 94 |
+
embedding_model (`EmbeddingModelBase`):
|
| 95 |
+
The embedding model used to embed both queries and
|
| 96 |
+
inserted chunks. Must be the same model used at
|
| 97 |
+
indexing time and at retrieval time, otherwise vectors
|
| 98 |
+
will not be comparable.
|
| 99 |
+
vector_store (`VectorStoreBase`):
|
| 100 |
+
The shared vector-store connection. The store must
|
| 101 |
+
already be entered (its own ``__aenter__`` already
|
| 102 |
+
called) before any operation on this handle runs.
|
| 103 |
+
collection (`str`):
|
| 104 |
+
The physical collection backing this knowledge base.
|
| 105 |
+
Created lazily on the first operation; see
|
| 106 |
+
:meth:`ensure_collection`.
|
| 107 |
+
metadata_filter (`dict | None`, optional):
|
| 108 |
+
Defense-in-depth payload filter. When set:
|
| 109 |
+
|
| 110 |
+
- :meth:`search` and :meth:`list_documents` restrict
|
| 111 |
+
results to records whose payload matches every
|
| 112 |
+
``key == value`` pair;
|
| 113 |
+
- :meth:`insert_document` forces these keys onto every
|
| 114 |
+
inserted chunk's metadata, overriding caller-supplied
|
| 115 |
+
values, so records cannot leak into another scope.
|
| 116 |
+
|
| 117 |
+
``None`` disables filtering — the default for
|
| 118 |
+
deployments where every knowledge base owns its
|
| 119 |
+
collection outright.
|
| 120 |
+
"""
|
| 121 |
+
self.name = name
|
| 122 |
+
self.description = description
|
| 123 |
+
self._embedding_model = embedding_model
|
| 124 |
+
self._vector_store = vector_store
|
| 125 |
+
self._collection = collection
|
| 126 |
+
self._metadata_filter = metadata_filter
|
| 127 |
+
# Memoise the "collection exists" check after the first
|
| 128 |
+
# successful ensure_collection so subsequent operations avoid
|
| 129 |
+
# the extra round-trip.
|
| 130 |
+
self._collection_ready = False
|
| 131 |
+
|
| 132 |
+
# ------------------------------------------------------------------
|
| 133 |
+
# Read-only accessors
|
| 134 |
+
# ------------------------------------------------------------------
|
| 135 |
+
|
| 136 |
+
@property
|
| 137 |
+
def embedding_model(self) -> EmbeddingModelBase:
|
| 138 |
+
"""The bound embedding model."""
|
| 139 |
+
return self._embedding_model
|
| 140 |
+
|
| 141 |
+
@property
|
| 142 |
+
def vector_store(self) -> VectorStoreBase:
|
| 143 |
+
"""The bound vector store."""
|
| 144 |
+
return self._vector_store
|
| 145 |
+
|
| 146 |
+
@property
|
| 147 |
+
def collection(self) -> str:
|
| 148 |
+
"""The physical collection backing this knowledge base."""
|
| 149 |
+
return self._collection
|
| 150 |
+
|
| 151 |
+
@property
|
| 152 |
+
def metadata_filter(self) -> dict | None:
|
| 153 |
+
"""The defense-in-depth payload filter, or ``None``."""
|
| 154 |
+
return self._metadata_filter
|
| 155 |
+
|
| 156 |
+
# ------------------------------------------------------------------
|
| 157 |
+
# Lifecycle
|
| 158 |
+
# ------------------------------------------------------------------
|
| 159 |
+
|
| 160 |
+
async def ensure_collection(self) -> None:
|
| 161 |
+
"""Idempotently create the backing collection if missing.
|
| 162 |
+
|
| 163 |
+
Called transparently at the top of every public operation —
|
| 164 |
+
callers should not need to invoke it themselves. Memoised on
|
| 165 |
+
the instance after the first success, so subsequent calls are
|
| 166 |
+
a single ``if`` check.
|
| 167 |
+
|
| 168 |
+
Looks up the collection via
|
| 169 |
+
:meth:`VectorStoreBase.has_collection` and creates it with the
|
| 170 |
+
embedding model's :attr:`~EmbeddingModelBase.dimensions` when
|
| 171 |
+
absent.
|
| 172 |
+
|
| 173 |
+
Raises whatever the backend raises if the collection exists at
|
| 174 |
+
an incompatible dimension (the backend is the authority on
|
| 175 |
+
that; we do not double-check here).
|
| 176 |
+
"""
|
| 177 |
+
if self._collection_ready:
|
| 178 |
+
return
|
| 179 |
+
if not await self._vector_store.has_collection(self._collection):
|
| 180 |
+
await self._vector_store.create_collection(
|
| 181 |
+
self._collection,
|
| 182 |
+
dimensions=self._embedding_model.dimensions,
|
| 183 |
+
)
|
| 184 |
+
self._collection_ready = True
|
| 185 |
+
|
| 186 |
+
# ------------------------------------------------------------------
|
| 187 |
+
# Search
|
| 188 |
+
# ------------------------------------------------------------------
|
| 189 |
+
|
| 190 |
+
async def search(
|
| 191 |
+
self,
|
| 192 |
+
queries: list[str | TextBlock | DataBlock],
|
| 193 |
+
top_k: int = 5,
|
| 194 |
+
score_threshold: float | None = None,
|
| 195 |
+
) -> list[VectorSearchResult]:
|
| 196 |
+
"""Search the knowledge base with one or more queries.
|
| 197 |
+
|
| 198 |
+
All queries are embedded in a single batch, then searched
|
| 199 |
+
concurrently against the bound collection (with
|
| 200 |
+
:attr:`metadata_filter` applied). Hits are deduplicated by
|
| 201 |
+
``(document_id, chunk_index)`` keeping the best score,
|
| 202 |
+
optionally filtered by ``score_threshold``, sorted by
|
| 203 |
+
descending score, and truncated to ``top_k``.
|
| 204 |
+
|
| 205 |
+
Args:
|
| 206 |
+
queries (`list[str | TextBlock | DataBlock]`):
|
| 207 |
+
Query inputs. Text may be either bare ``str`` or
|
| 208 |
+
:class:`TextBlock`; :class:`DataBlock` items are
|
| 209 |
+
**silently dropped** when the bound embedding model
|
| 210 |
+
does not declare ``supports_multimodal`` — text-only
|
| 211 |
+
models would otherwise reject them. Callers can
|
| 212 |
+
therefore pass a mixed list without per-KB filtering.
|
| 213 |
+
top_k (`int`, defaults to ``5``):
|
| 214 |
+
Maximum number of results returned across all queries
|
| 215 |
+
(after dedup).
|
| 216 |
+
score_threshold (`float | None`, optional):
|
| 217 |
+
Minimum similarity score for a hit to be retained.
|
| 218 |
+
Only meaningful for similarity metrics where higher is
|
| 219 |
+
better (cosine / dot-product). ``None`` disables
|
| 220 |
+
filtering.
|
| 221 |
+
|
| 222 |
+
Returns:
|
| 223 |
+
`list[VectorSearchResult]`:
|
| 224 |
+
At most ``top_k`` deduplicated hits ordered by
|
| 225 |
+
descending similarity score. Empty when there are no
|
| 226 |
+
queries the bound embedding model can consume.
|
| 227 |
+
"""
|
| 228 |
+
if not queries:
|
| 229 |
+
return []
|
| 230 |
+
|
| 231 |
+
if not self._embedding_model.supports_multimodal:
|
| 232 |
+
queries = [q for q in queries if not isinstance(q, DataBlock)]
|
| 233 |
+
if not queries:
|
| 234 |
+
return []
|
| 235 |
+
|
| 236 |
+
await self.ensure_collection()
|
| 237 |
+
response = await self._embedding_model(queries)
|
| 238 |
+
|
| 239 |
+
results_per_query = await asyncio.gather(
|
| 240 |
+
*(
|
| 241 |
+
self._vector_store.search(
|
| 242 |
+
collection=self._collection,
|
| 243 |
+
query_vector=vector,
|
| 244 |
+
top_k=top_k,
|
| 245 |
+
metadata_filter=self._metadata_filter,
|
| 246 |
+
)
|
| 247 |
+
for vector in response.embeddings
|
| 248 |
+
),
|
| 249 |
+
)
|
| 250 |
+
|
| 251 |
+
best: dict[tuple[str, int], VectorSearchResult] = {}
|
| 252 |
+
for results in results_per_query:
|
| 253 |
+
for result in results:
|
| 254 |
+
if (
|
| 255 |
+
score_threshold is not None
|
| 256 |
+
and result.score < score_threshold
|
| 257 |
+
):
|
| 258 |
+
continue
|
| 259 |
+
# ``(document_id, chunk_index)`` is the stable identity
|
| 260 |
+
# of a chunk: it survives reindex (block UUIDs do not)
|
| 261 |
+
# and uniquely names "this slice of that document"
|
| 262 |
+
# regardless of which query surfaced it.
|
| 263 |
+
key = (result.document_id, result.chunk.chunk_index)
|
| 264 |
+
if key not in best or result.score > best[key].score:
|
| 265 |
+
best[key] = result
|
| 266 |
+
|
| 267 |
+
merged = sorted(
|
| 268 |
+
best.values(),
|
| 269 |
+
key=lambda result: result.score,
|
| 270 |
+
reverse=True,
|
| 271 |
+
)
|
| 272 |
+
return merged[:top_k]
|
| 273 |
+
|
| 274 |
+
# ------------------------------------------------------------------
|
| 275 |
+
# Document management
|
| 276 |
+
# ------------------------------------------------------------------
|
| 277 |
+
|
| 278 |
+
async def insert_document(
|
| 279 |
+
self,
|
| 280 |
+
chunks: list[Chunk],
|
| 281 |
+
document_id: str | None = None,
|
| 282 |
+
document_metadata: dict | None = None,
|
| 283 |
+
) -> str:
|
| 284 |
+
"""Embed and insert a list of chunks as a single source document.
|
| 285 |
+
|
| 286 |
+
All chunks share the resolved ``document_id``;
|
| 287 |
+
:meth:`delete_document` later removes them as a unit. Each
|
| 288 |
+
chunk's metadata is merged in this precedence (highest wins):
|
| 289 |
+
|
| 290 |
+
1. :attr:`metadata_filter` keys — defense-in-depth scoping, so
|
| 291 |
+
a chunk can never be inserted with a payload that escapes
|
| 292 |
+
the filter (any escape would silently disappear at retrieve
|
| 293 |
+
time anyway, but failing closed at insert is clearer).
|
| 294 |
+
2. The chunk's pre-existing ``metadata`` — parser-supplied.
|
| 295 |
+
3. ``document_metadata`` — document-level fields propagated
|
| 296 |
+
down (filename, media type, upload time, ...).
|
| 297 |
+
|
| 298 |
+
Args:
|
| 299 |
+
chunks (`list[Chunk]`):
|
| 300 |
+
The pre-chunked document content (already produced by
|
| 301 |
+
a parser + chunker pipeline). An empty list is a
|
| 302 |
+
no-op.
|
| 303 |
+
document_id (`str | None`, optional):
|
| 304 |
+
The document identifier. When ``None`` a fresh UUID
|
| 305 |
+
hex is generated and returned so the caller can record
|
| 306 |
+
it for future :meth:`delete_document` calls.
|
| 307 |
+
document_metadata (`dict | None`, optional):
|
| 308 |
+
Document-level metadata (filename, media type, size,
|
| 309 |
+
upload time, ...). Merged into each chunk's
|
| 310 |
+
``metadata``.
|
| 311 |
+
|
| 312 |
+
Returns:
|
| 313 |
+
`str`:
|
| 314 |
+
The (possibly generated) document id.
|
| 315 |
+
|
| 316 |
+
Raises:
|
| 317 |
+
`RuntimeError`:
|
| 318 |
+
If the embedding model returns a number of vectors
|
| 319 |
+
that does not match the number of chunks.
|
| 320 |
+
"""
|
| 321 |
+
if not chunks:
|
| 322 |
+
return document_id or _generate_id()
|
| 323 |
+
document_id = document_id or _generate_id()
|
| 324 |
+
|
| 325 |
+
await self.ensure_collection()
|
| 326 |
+
|
| 327 |
+
# Precedence: metadata_filter wins (security boundary), then
|
| 328 |
+
# chunk metadata, then document_metadata. See docstring.
|
| 329 |
+
for chunk in chunks:
|
| 330 |
+
chunk.metadata = {
|
| 331 |
+
**(document_metadata or {}),
|
| 332 |
+
**chunk.metadata,
|
| 333 |
+
**(self._metadata_filter or {}),
|
| 334 |
+
}
|
| 335 |
+
|
| 336 |
+
response = await self._embedding_model(
|
| 337 |
+
[chunk.content for chunk in chunks],
|
| 338 |
+
)
|
| 339 |
+
|
| 340 |
+
if len(response.embeddings) != len(chunks):
|
| 341 |
+
raise RuntimeError(
|
| 342 |
+
f"Embedding model returned {len(response.embeddings)} "
|
| 343 |
+
f"vectors for {len(chunks)} chunks.",
|
| 344 |
+
)
|
| 345 |
+
|
| 346 |
+
records = [
|
| 347 |
+
VectorRecord(
|
| 348 |
+
vector=vector,
|
| 349 |
+
document_id=document_id,
|
| 350 |
+
chunk=chunk,
|
| 351 |
+
)
|
| 352 |
+
for vector, chunk in zip(response.embeddings, chunks)
|
| 353 |
+
]
|
| 354 |
+
await self._vector_store.insert(self._collection, records)
|
| 355 |
+
return document_id
|
| 356 |
+
|
| 357 |
+
async def delete_document(self, document_id: str) -> None:
|
| 358 |
+
"""Remove every record for one source document.
|
| 359 |
+
|
| 360 |
+
Args:
|
| 361 |
+
document_id (`str`):
|
| 362 |
+
The source document id whose records should be removed.
|
| 363 |
+
"""
|
| 364 |
+
await self.ensure_collection()
|
| 365 |
+
await self._vector_store.delete(
|
| 366 |
+
self._collection,
|
| 367 |
+
document_id,
|
| 368 |
+
)
|
| 369 |
+
|
| 370 |
+
async def list_documents(self) -> list["DocumentSummary"]:
|
| 371 |
+
"""List all distinct source documents in this knowledge base.
|
| 372 |
+
|
| 373 |
+
Filtered by :attr:`metadata_filter` when set, so callers only
|
| 374 |
+
ever see documents within their own scope.
|
| 375 |
+
|
| 376 |
+
Returns:
|
| 377 |
+
`list[DocumentSummary]`:
|
| 378 |
+
One summary per indexed document, in unspecified order.
|
| 379 |
+
"""
|
| 380 |
+
await self.ensure_collection()
|
| 381 |
+
return await self._vector_store.list_documents(
|
| 382 |
+
self._collection,
|
| 383 |
+
metadata_filter=self._metadata_filter,
|
| 384 |
+
)
|
src/agentscope/rag/_parser/__init__.py
ADDED
|
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""File parser implementations for the RAG indexing pipeline."""
|
| 3 |
+
|
| 4 |
+
from ._base import ParserBase
|
| 5 |
+
from ._image import ImageParser
|
| 6 |
+
from ._pdf import PDFParser
|
| 7 |
+
from ._ppt import PPTParser
|
| 8 |
+
from ._text import TextParser
|
| 9 |
+
|
| 10 |
+
__all__ = [
|
| 11 |
+
"ParserBase",
|
| 12 |
+
"PDFParser",
|
| 13 |
+
"PPTParser",
|
| 14 |
+
"ImageParser",
|
| 15 |
+
"TextParser",
|
| 16 |
+
]
|
src/agentscope/rag/_parser/_base.py
ADDED
|
@@ -0,0 +1,116 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""Abstract base class for file parsers.
|
| 3 |
+
|
| 4 |
+
A :class:`ParserBase` subclass handles **one file format**. Its job
|
| 5 |
+
is to read a file's raw bytes and produce a list of
|
| 6 |
+
:class:`~agentscope.rag.Section` objects, each representing a natural
|
| 7 |
+
boundary of the source (e.g. one PDF page, one PPTX slide, one
|
| 8 |
+
embedded image).
|
| 9 |
+
|
| 10 |
+
Parsers **do not chunk text**. Long text is left intact inside the
|
| 11 |
+
Section; splitting happens later in a
|
| 12 |
+
:class:`~agentscope.rag.ChunkerBase`. Parsers also do not need to
|
| 13 |
+
worry about output size — only about preserving the structural
|
| 14 |
+
boundaries that downstream consumers must not cross.
|
| 15 |
+
"""
|
| 16 |
+
import mimetypes
|
| 17 |
+
from abc import ABC, abstractmethod
|
| 18 |
+
|
| 19 |
+
from .._document import Section
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
class ParserBase(ABC):
|
| 23 |
+
"""Abstract base class for file-format parsers.
|
| 24 |
+
|
| 25 |
+
Each subclass handles a single file format (or a related family,
|
| 26 |
+
e.g. all plain-text MIME types). Subclasses are typically
|
| 27 |
+
instantiated once and reused across many ``parse()`` calls.
|
| 28 |
+
|
| 29 |
+
Subclasses should be stateless or thread-safe — a single
|
| 30 |
+
instance may be invoked concurrently from multiple agent runs.
|
| 31 |
+
|
| 32 |
+
Subclasses must declare :attr:`supported_media_types` so that the
|
| 33 |
+
KnowledgeBaseManager can route uploaded files to the right parser
|
| 34 |
+
based on standard IANA media types (RFC 6838).
|
| 35 |
+
"""
|
| 36 |
+
|
| 37 |
+
supported_media_types: list[str]
|
| 38 |
+
"""Standard IANA media types (RFC 6838) this parser handles,
|
| 39 |
+
e.g. ``["application/pdf"]`` or
|
| 40 |
+
``["text/plain", "text/markdown"]``. Used by the
|
| 41 |
+
KnowledgeBaseManager to select a parser for an uploaded file."""
|
| 42 |
+
|
| 43 |
+
@classmethod
|
| 44 |
+
def supported_extensions(cls) -> list[str]:
|
| 45 |
+
"""Filename extensions (including the leading ``.``) this parser
|
| 46 |
+
can produce uploads for.
|
| 47 |
+
|
| 48 |
+
The base implementation derives extensions from
|
| 49 |
+
:attr:`supported_media_types` via
|
| 50 |
+
:func:`mimetypes.guess_all_extensions` — good enough for clean
|
| 51 |
+
IANA types like ``application/pdf``. Subclasses **should
|
| 52 |
+
override** this when the default reverse-lookup is noisy
|
| 53 |
+
(``text/plain`` resolves to ``.bat`` / ``.c`` / ``.pl`` and a
|
| 54 |
+
dozen other developer extensions no KB user wants in the file
|
| 55 |
+
picker) or when a media type has no registered extension at all
|
| 56 |
+
(``application/x-yaml`` returns the empty list).
|
| 57 |
+
|
| 58 |
+
The result is consumed by the front-end's ``<input accept>`` and
|
| 59 |
+
by the client-side filename guard; it is **not** consulted for
|
| 60 |
+
media-type routing — that always goes through
|
| 61 |
+
:attr:`supported_media_types`.
|
| 62 |
+
|
| 63 |
+
Returns:
|
| 64 |
+
`list[str]`:
|
| 65 |
+
Deduplicated, sorted extensions (each starting with
|
| 66 |
+
``.``). May be empty when no media type resolves.
|
| 67 |
+
"""
|
| 68 |
+
out: set[str] = set()
|
| 69 |
+
for media_type in cls.supported_media_types:
|
| 70 |
+
out.update(mimetypes.guess_all_extensions(media_type))
|
| 71 |
+
return sorted(out)
|
| 72 |
+
|
| 73 |
+
@abstractmethod
|
| 74 |
+
async def parse(
|
| 75 |
+
self,
|
| 76 |
+
file: bytes | str,
|
| 77 |
+
filename: str,
|
| 78 |
+
) -> list[Section]:
|
| 79 |
+
"""Parse a file into a list of :class:`Section` objects.
|
| 80 |
+
|
| 81 |
+
The ``file`` argument is a union covering the three call sites
|
| 82 |
+
a parser sees in practice:
|
| 83 |
+
|
| 84 |
+
- ``bytes`` — the raw payload, as handed in by HTTP uploads
|
| 85 |
+
and blob-store reads.
|
| 86 |
+
- ``str`` for binary parsers (PDF, PPT, image, …) — a
|
| 87 |
+
**filesystem path** to the file to read. The parser opens
|
| 88 |
+
the path itself; callers do not need to read the bytes first.
|
| 89 |
+
- ``str`` for :class:`TextParser` — disambiguated at runtime:
|
| 90 |
+
if the string names an existing file on disk it is treated
|
| 91 |
+
as a path and the file is decoded with the configured
|
| 92 |
+
encoding; otherwise it is treated as pre-decoded text.
|
| 93 |
+
|
| 94 |
+
Args:
|
| 95 |
+
file (`bytes | str`):
|
| 96 |
+
The file content or a path to it (see above).
|
| 97 |
+
filename (`str`):
|
| 98 |
+
The original filename (e.g. ``"report.pdf"``). Used
|
| 99 |
+
for error messages and copied into each Section's
|
| 100 |
+
:attr:`Section.source` field for downstream display
|
| 101 |
+
/ citation.
|
| 102 |
+
|
| 103 |
+
Returns:
|
| 104 |
+
`list[Section]`:
|
| 105 |
+
One Section per natural boundary in the source file.
|
| 106 |
+
For unstructured formats (plain text, image, video),
|
| 107 |
+
a single Section may cover the whole file. Sections
|
| 108 |
+
are returned in document order.
|
| 109 |
+
|
| 110 |
+
Raises:
|
| 111 |
+
`TypeError`: If the subclass does not accept the supplied
|
| 112 |
+
``file`` form.
|
| 113 |
+
`FileNotFoundError`: If a binary parser is handed a
|
| 114 |
+
``str`` that does not name an existing file.
|
| 115 |
+
`ValueError`: If the file cannot be parsed.
|
| 116 |
+
"""
|
src/agentscope/rag/_parser/_image.py
ADDED
|
@@ -0,0 +1,92 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""Image file parser.
|
| 3 |
+
|
| 4 |
+
A single :class:`Section` carrying the raw image bytes as a
|
| 5 |
+
base64-encoded :class:`DataBlock`. No OCR, no captioning — the
|
| 6 |
+
section is the image, ready to flow through to a multimodal
|
| 7 |
+
embedding model unchanged.
|
| 8 |
+
"""
|
| 9 |
+
import base64
|
| 10 |
+
|
| 11 |
+
from ...message import Base64Source, DataBlock
|
| 12 |
+
from .._document import Section
|
| 13 |
+
from ._base import ParserBase
|
| 14 |
+
from ._utils import _guess_image_media_type
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
class ImageParser(ParserBase):
|
| 18 |
+
"""Parser for image files.
|
| 19 |
+
|
| 20 |
+
Wraps the entire file as a single :class:`Section` holding a
|
| 21 |
+
:class:`DataBlock` with the image's base64-encoded bytes. This is
|
| 22 |
+
the input shape a multimodal embedding model expects; the
|
| 23 |
+
surrounding pipeline (chunker, vector store) treats the section
|
| 24 |
+
opaquely.
|
| 25 |
+
|
| 26 |
+
The IANA media type is sniffed from the bytes' magic number so
|
| 27 |
+
callers do not need to pass it explicitly.
|
| 28 |
+
"""
|
| 29 |
+
|
| 30 |
+
supported_media_types: list[str] = [
|
| 31 |
+
"image/png",
|
| 32 |
+
"image/jpeg",
|
| 33 |
+
"image/gif",
|
| 34 |
+
"image/bmp",
|
| 35 |
+
"image/webp",
|
| 36 |
+
]
|
| 37 |
+
|
| 38 |
+
@classmethod
|
| 39 |
+
def supported_extensions(cls) -> list[str]:
|
| 40 |
+
"""Return the canonical image extensions."""
|
| 41 |
+
return [
|
| 42 |
+
".bmp",
|
| 43 |
+
".gif",
|
| 44 |
+
".jpeg",
|
| 45 |
+
".jpg",
|
| 46 |
+
".png",
|
| 47 |
+
".webp",
|
| 48 |
+
]
|
| 49 |
+
|
| 50 |
+
async def parse(
|
| 51 |
+
self,
|
| 52 |
+
file: bytes | str,
|
| 53 |
+
filename: str,
|
| 54 |
+
) -> list[Section]:
|
| 55 |
+
"""Wrap the image bytes in a single :class:`Section`.
|
| 56 |
+
|
| 57 |
+
Args:
|
| 58 |
+
file (`bytes | str`):
|
| 59 |
+
Either the raw image bytes, or a filesystem path to
|
| 60 |
+
the image file.
|
| 61 |
+
filename (`str`):
|
| 62 |
+
The source filename, copied into
|
| 63 |
+
:attr:`Section.source`.
|
| 64 |
+
|
| 65 |
+
Returns:
|
| 66 |
+
`list[Section]`:
|
| 67 |
+
A one-element list whose section's ``content`` is a
|
| 68 |
+
:class:`DataBlock` with the base64-encoded image data.
|
| 69 |
+
|
| 70 |
+
Raises:
|
| 71 |
+
`FileNotFoundError`: If ``file`` is a ``str`` pointing to
|
| 72 |
+
a path that does not exist.
|
| 73 |
+
"""
|
| 74 |
+
if isinstance(file, str):
|
| 75 |
+
with open(file, "rb") as fp:
|
| 76 |
+
file = fp.read()
|
| 77 |
+
|
| 78 |
+
media_type = _guess_image_media_type(file)
|
| 79 |
+
data = base64.b64encode(file).decode("utf-8")
|
| 80 |
+
return [
|
| 81 |
+
Section(
|
| 82 |
+
content=DataBlock(
|
| 83 |
+
source=Base64Source(
|
| 84 |
+
media_type=media_type,
|
| 85 |
+
data=data,
|
| 86 |
+
),
|
| 87 |
+
name=filename,
|
| 88 |
+
),
|
| 89 |
+
source=filename,
|
| 90 |
+
metadata={"media_type": media_type},
|
| 91 |
+
),
|
| 92 |
+
]
|
src/agentscope/rag/_parser/_pdf.py
ADDED
|
@@ -0,0 +1,92 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""PDF file parser.
|
| 3 |
+
|
| 4 |
+
One :class:`Section` per page so a downstream
|
| 5 |
+
:class:`~agentscope.rag.ChunkerBase` never combines text across page
|
| 6 |
+
boundaries. Each section's :attr:`Section.metadata` carries the
|
| 7 |
+
page number (starting at 1) for later citation.
|
| 8 |
+
"""
|
| 9 |
+
import io
|
| 10 |
+
|
| 11 |
+
from ...message import TextBlock
|
| 12 |
+
from .._document import Section
|
| 13 |
+
from ._base import ParserBase
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
class PDFParser(ParserBase):
|
| 17 |
+
"""Parser for PDF files.
|
| 18 |
+
|
| 19 |
+
Returns one :class:`Section` per page. Empty / image-only pages
|
| 20 |
+
still produce a Section (with empty text) so the page → section
|
| 21 |
+
correspondence stays exact — downstream chunkers naturally drop
|
| 22 |
+
empty content.
|
| 23 |
+
|
| 24 |
+
Requires :mod:`pypdf`; install with ``pip install pypdf`` (or via
|
| 25 |
+
the ``agentscope[rag]`` extra).
|
| 26 |
+
"""
|
| 27 |
+
|
| 28 |
+
supported_media_types: list[str] = ["application/pdf"]
|
| 29 |
+
|
| 30 |
+
@classmethod
|
| 31 |
+
def supported_extensions(cls) -> list[str]:
|
| 32 |
+
"""Return the canonical ``.pdf`` extension."""
|
| 33 |
+
return [".pdf"]
|
| 34 |
+
|
| 35 |
+
async def parse(
|
| 36 |
+
self,
|
| 37 |
+
file: bytes | str,
|
| 38 |
+
filename: str,
|
| 39 |
+
) -> list[Section]:
|
| 40 |
+
"""Read the PDF bytes and return one Section per page.
|
| 41 |
+
|
| 42 |
+
Args:
|
| 43 |
+
file (`bytes | str`):
|
| 44 |
+
Either the raw PDF bytes, or a filesystem path to
|
| 45 |
+
the PDF file.
|
| 46 |
+
filename (`str`):
|
| 47 |
+
The source filename, copied verbatim into each
|
| 48 |
+
Section's :attr:`Section.source` field.
|
| 49 |
+
|
| 50 |
+
Returns:
|
| 51 |
+
`list[Section]`:
|
| 52 |
+
One Section per page, in document order. Each
|
| 53 |
+
section's metadata holds ``{"page": <starting at 1>}``.
|
| 54 |
+
|
| 55 |
+
Raises:
|
| 56 |
+
`FileNotFoundError`: If ``file`` is a ``str`` pointing to
|
| 57 |
+
a path that does not exist.
|
| 58 |
+
`ImportError`: If :mod:`pypdf` is not installed.
|
| 59 |
+
`ValueError`: If the bytes cannot be parsed as PDF.
|
| 60 |
+
"""
|
| 61 |
+
if isinstance(file, str):
|
| 62 |
+
with open(file, "rb") as fp:
|
| 63 |
+
file = fp.read()
|
| 64 |
+
|
| 65 |
+
try:
|
| 66 |
+
from pypdf import PdfReader
|
| 67 |
+
from pypdf.errors import PdfReadError
|
| 68 |
+
except ImportError as e:
|
| 69 |
+
raise ImportError(
|
| 70 |
+
"Please install pypdf to use the PDF parser. "
|
| 71 |
+
"You can install it by `pip install pypdf` (or "
|
| 72 |
+
"`pip install agentscope[rag]`).",
|
| 73 |
+
) from e
|
| 74 |
+
|
| 75 |
+
try:
|
| 76 |
+
reader = PdfReader(io.BytesIO(file))
|
| 77 |
+
except PdfReadError as e:
|
| 78 |
+
raise ValueError(
|
| 79 |
+
f"Failed to parse {filename!r} as PDF: {e}",
|
| 80 |
+
) from e
|
| 81 |
+
|
| 82 |
+
sections: list[Section] = []
|
| 83 |
+
for page_idx, page in enumerate(reader.pages, start=1):
|
| 84 |
+
text = page.extract_text() or ""
|
| 85 |
+
sections.append(
|
| 86 |
+
Section(
|
| 87 |
+
content=TextBlock(text=text),
|
| 88 |
+
source=filename,
|
| 89 |
+
metadata={"page": page_idx},
|
| 90 |
+
),
|
| 91 |
+
)
|
| 92 |
+
return sections
|
src/agentscope/rag/_parser/_ppt.py
ADDED
|
@@ -0,0 +1,341 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""PowerPoint (.pptx) file parser.
|
| 3 |
+
|
| 4 |
+
Walks the deck slide-by-slide and emits one :class:`Section` per
|
| 5 |
+
contiguous content block — adjacent text runs (and, by default,
|
| 6 |
+
tables) are merged into a single text section; embedded images are
|
| 7 |
+
emitted as their own :class:`DataBlock` sections. Each section's
|
| 8 |
+
metadata carries the slide index (starting at 1) for later citation.
|
| 9 |
+
|
| 10 |
+
Mirrors the knob set of the v1 ``PowerPointReader``:
|
| 11 |
+
``include_image``, ``separate_table``, ``table_format``,
|
| 12 |
+
``slide_prefix``, ``slide_suffix``. Chunking is **not** done here —
|
| 13 |
+
long text stays intact inside a section and is split downstream by a
|
| 14 |
+
:class:`~agentscope.rag.ChunkerBase`.
|
| 15 |
+
"""
|
| 16 |
+
import base64
|
| 17 |
+
import io
|
| 18 |
+
from typing import Any, Literal
|
| 19 |
+
|
| 20 |
+
from ..._logging import logger
|
| 21 |
+
from ...message import Base64Source, DataBlock, TextBlock
|
| 22 |
+
from .._document import Section
|
| 23 |
+
from ._base import ParserBase
|
| 24 |
+
from ._utils import (
|
| 25 |
+
_guess_image_media_type,
|
| 26 |
+
_table_to_json,
|
| 27 |
+
_table_to_markdown,
|
| 28 |
+
)
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def _extract_table_rows(table: Any) -> list[list[str]]:
|
| 32 |
+
"""Read a python-pptx table into a 2-D ``list[list[str]]``.
|
| 33 |
+
|
| 34 |
+
Args:
|
| 35 |
+
table (`Any`):
|
| 36 |
+
The python-pptx ``Table`` object.
|
| 37 |
+
|
| 38 |
+
Returns:
|
| 39 |
+
`list[list[str]]`:
|
| 40 |
+
One inner list per row; per-cell line breaks are
|
| 41 |
+
normalised to ``\\n``.
|
| 42 |
+
"""
|
| 43 |
+
rows: list[list[str]] = []
|
| 44 |
+
for row in table.rows:
|
| 45 |
+
cells: list[str] = []
|
| 46 |
+
for cell in row.cells:
|
| 47 |
+
text = cell.text.strip()
|
| 48 |
+
text = text.replace("\r\n", "\n").replace("\r", "\n")
|
| 49 |
+
cells.append(text)
|
| 50 |
+
rows.append(cells)
|
| 51 |
+
return rows
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def _extract_image_bytes(shape: Any) -> bytes | None:
|
| 55 |
+
"""Return the embedded image bytes for a picture shape, or ``None``.
|
| 56 |
+
|
| 57 |
+
Args:
|
| 58 |
+
shape (`Any`):
|
| 59 |
+
A python-pptx shape. Non-picture shapes return ``None``.
|
| 60 |
+
|
| 61 |
+
Returns:
|
| 62 |
+
`bytes | None`:
|
| 63 |
+
The raw image bytes, or ``None`` when ``shape`` is not a
|
| 64 |
+
picture / the bytes are unreadable.
|
| 65 |
+
"""
|
| 66 |
+
try:
|
| 67 |
+
from pptx.enum.shapes import MSO_SHAPE_TYPE
|
| 68 |
+
|
| 69 |
+
picture_type = MSO_SHAPE_TYPE.PICTURE
|
| 70 |
+
except ImportError:
|
| 71 |
+
# MSO_SHAPE_TYPE.PICTURE numeric value used as the fallback
|
| 72 |
+
# so the parser still works against pptx builds where the
|
| 73 |
+
# enum import path has moved.
|
| 74 |
+
picture_type = 13
|
| 75 |
+
|
| 76 |
+
if shape.shape_type != picture_type:
|
| 77 |
+
return None
|
| 78 |
+
try:
|
| 79 |
+
return shape.image.blob
|
| 80 |
+
except Exception as e: # pylint: disable=broad-except
|
| 81 |
+
logger.warning("Failed to extract image from PPT shape: %s", e)
|
| 82 |
+
return None
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
class PPTParser(ParserBase):
|
| 86 |
+
"""Parser for PowerPoint ``.pptx`` files.
|
| 87 |
+
|
| 88 |
+
Slide order is preserved. Within a slide, shapes are visited in
|
| 89 |
+
document order and grouped into a minimum number of sections:
|
| 90 |
+
|
| 91 |
+
- **Text** shapes and **table** shapes contribute to one running
|
| 92 |
+
text section. When ``separate_table=True`` a table closes the
|
| 93 |
+
running section and starts a new one of its own.
|
| 94 |
+
- **Picture** shapes emit a standalone :class:`Section` whose
|
| 95 |
+
``content`` is a :class:`DataBlock` holding the base64-encoded
|
| 96 |
+
image bytes.
|
| 97 |
+
|
| 98 |
+
``slide_prefix`` / ``slide_suffix`` are wrapped around each
|
| 99 |
+
slide's text content (the prefix is prepended to the slide's
|
| 100 |
+
first text section, the suffix is appended to its last text
|
| 101 |
+
section). Use ``None`` on either to disable wrapping.
|
| 102 |
+
"""
|
| 103 |
+
|
| 104 |
+
supported_media_types: list[str] = [
|
| 105 |
+
"application/vnd.openxmlformats-officedocument.presentationml"
|
| 106 |
+
".presentation",
|
| 107 |
+
]
|
| 108 |
+
|
| 109 |
+
@classmethod
|
| 110 |
+
def supported_extensions(cls) -> list[str]:
|
| 111 |
+
"""Return ``[".pptx"]`` — the only format ``python-pptx``
|
| 112 |
+
reads."""
|
| 113 |
+
return [".pptx"]
|
| 114 |
+
|
| 115 |
+
def __init__(
|
| 116 |
+
self,
|
| 117 |
+
include_image: bool = True,
|
| 118 |
+
separate_table: bool = False,
|
| 119 |
+
table_format: Literal["markdown", "json"] = "markdown",
|
| 120 |
+
slide_prefix: str | None = "<slide index={index}>",
|
| 121 |
+
slide_suffix: str | None = "</slide>",
|
| 122 |
+
) -> None:
|
| 123 |
+
"""Initialize the PowerPoint parser.
|
| 124 |
+
|
| 125 |
+
Args:
|
| 126 |
+
include_image (`bool`, defaults to ``True``):
|
| 127 |
+
When ``True``, picture shapes are emitted as
|
| 128 |
+
:class:`DataBlock` sections. Set to ``False`` to keep
|
| 129 |
+
a text-only index.
|
| 130 |
+
separate_table (`bool`, defaults to ``False``):
|
| 131 |
+
When ``True``, each table becomes its own text
|
| 132 |
+
section, never merged with surrounding text.
|
| 133 |
+
table_format (`Literal["markdown", "json"]`, defaults to
|
| 134 |
+
``"markdown"``):
|
| 135 |
+
How to render tables. ``"markdown"`` uses pipe-table
|
| 136 |
+
syntax; ``"json"`` emits a JSON array prefixed with a
|
| 137 |
+
``<system-info>`` marker — choose JSON when cells
|
| 138 |
+
contain newlines that would corrupt Markdown layout.
|
| 139 |
+
slide_prefix (`str | None`, defaults to
|
| 140 |
+
``"<slide index={index}>"``):
|
| 141 |
+
Prepended to the first text section of each slide.
|
| 142 |
+
Supports the ``{index}`` placeholder (starting at 1). Use
|
| 143 |
+
``None`` to disable.
|
| 144 |
+
slide_suffix (`str | None`, defaults to ``"</slide>"``):
|
| 145 |
+
Appended to the last text section of each slide. Use
|
| 146 |
+
``None`` to disable.
|
| 147 |
+
|
| 148 |
+
Raises:
|
| 149 |
+
`ValueError`: If ``table_format`` is not one of
|
| 150 |
+
``"markdown"`` / ``"json"``.
|
| 151 |
+
"""
|
| 152 |
+
if table_format not in ("markdown", "json"):
|
| 153 |
+
raise ValueError(
|
| 154 |
+
"The table_format must be one of 'markdown' or 'json', "
|
| 155 |
+
f"got {table_format!r}.",
|
| 156 |
+
)
|
| 157 |
+
self.include_image = include_image
|
| 158 |
+
self.separate_table = separate_table
|
| 159 |
+
self.table_format = table_format
|
| 160 |
+
self.slide_prefix = slide_prefix
|
| 161 |
+
self.slide_suffix = slide_suffix
|
| 162 |
+
|
| 163 |
+
async def parse(
|
| 164 |
+
self,
|
| 165 |
+
file: bytes | str,
|
| 166 |
+
filename: str,
|
| 167 |
+
) -> list[Section]:
|
| 168 |
+
"""Parse a PPTX file into a list of :class:`Section` objects.
|
| 169 |
+
|
| 170 |
+
Args:
|
| 171 |
+
file (`bytes | str`):
|
| 172 |
+
Either the raw PPTX bytes, or a filesystem path to
|
| 173 |
+
the PPTX file.
|
| 174 |
+
filename (`str`):
|
| 175 |
+
The source filename, copied into each Section's
|
| 176 |
+
:attr:`Section.source`.
|
| 177 |
+
|
| 178 |
+
Returns:
|
| 179 |
+
`list[Section]`:
|
| 180 |
+
Sections in deck order. Text sections carry
|
| 181 |
+
``{"slide": <starting at 1>}``; image sections add
|
| 182 |
+
``{"media_type": "image/..."}``.
|
| 183 |
+
|
| 184 |
+
Raises:
|
| 185 |
+
`FileNotFoundError`: If ``file`` is a ``str`` pointing to
|
| 186 |
+
a path that does not exist.
|
| 187 |
+
`ImportError`: If :mod:`python-pptx` is not installed.
|
| 188 |
+
`ValueError`: If the bytes cannot be parsed.
|
| 189 |
+
"""
|
| 190 |
+
if isinstance(file, str):
|
| 191 |
+
with open(file, "rb") as fp:
|
| 192 |
+
file = fp.read()
|
| 193 |
+
|
| 194 |
+
try:
|
| 195 |
+
from pptx import Presentation
|
| 196 |
+
except ImportError as e:
|
| 197 |
+
raise ImportError(
|
| 198 |
+
"Please install python-pptx to use the PowerPoint "
|
| 199 |
+
"parser. You can install it by "
|
| 200 |
+
"`pip install python-pptx` (or "
|
| 201 |
+
"`pip install agentscope[rag]`).",
|
| 202 |
+
) from e
|
| 203 |
+
|
| 204 |
+
try:
|
| 205 |
+
prs = Presentation(io.BytesIO(file))
|
| 206 |
+
except Exception as e: # pylint: disable=broad-except
|
| 207 |
+
raise ValueError(
|
| 208 |
+
f"Failed to parse {filename!r} as PPTX: {e}",
|
| 209 |
+
) from e
|
| 210 |
+
|
| 211 |
+
sections: list[Section] = []
|
| 212 |
+
for slide_idx, slide in enumerate(prs.slides):
|
| 213 |
+
sections.extend(
|
| 214 |
+
self._parse_slide(slide, slide_idx, filename),
|
| 215 |
+
)
|
| 216 |
+
return sections
|
| 217 |
+
|
| 218 |
+
# ------------------------------------------------------------------
|
| 219 |
+
# Slide-level parsing
|
| 220 |
+
# ------------------------------------------------------------------
|
| 221 |
+
|
| 222 |
+
def _parse_slide(
|
| 223 |
+
self,
|
| 224 |
+
slide: Any,
|
| 225 |
+
slide_idx: int,
|
| 226 |
+
filename: str,
|
| 227 |
+
) -> list[Section]:
|
| 228 |
+
"""Walk one slide and return its ordered sections."""
|
| 229 |
+
slide_no = slide_idx + 1
|
| 230 |
+
prefix = (
|
| 231 |
+
self.slide_prefix.format(index=slide_no)
|
| 232 |
+
if self.slide_prefix is not None
|
| 233 |
+
else ""
|
| 234 |
+
)
|
| 235 |
+
|
| 236 |
+
slide_sections: list[Section] = []
|
| 237 |
+
# ``text_buffer`` accumulates the running text section; it is
|
| 238 |
+
# flushed whenever an image arrives, the slide ends, or
|
| 239 |
+
# ``separate_table`` forces a break around a table shape.
|
| 240 |
+
text_buffer: list[str] = []
|
| 241 |
+
|
| 242 |
+
def flush_text() -> None:
|
| 243 |
+
if not text_buffer:
|
| 244 |
+
return
|
| 245 |
+
slide_sections.append(
|
| 246 |
+
Section(
|
| 247 |
+
content=TextBlock(text="\n".join(text_buffer)),
|
| 248 |
+
source=filename,
|
| 249 |
+
metadata={"slide": slide_no},
|
| 250 |
+
),
|
| 251 |
+
)
|
| 252 |
+
text_buffer.clear()
|
| 253 |
+
|
| 254 |
+
# Slide prefix lives at the very top of the first text section.
|
| 255 |
+
if prefix:
|
| 256 |
+
text_buffer.append(prefix)
|
| 257 |
+
|
| 258 |
+
for shape in slide.shapes:
|
| 259 |
+
# 1. Pictures — flush running text, emit a DataBlock section.
|
| 260 |
+
if self.include_image:
|
| 261 |
+
image_bytes = _extract_image_bytes(shape)
|
| 262 |
+
if image_bytes is not None:
|
| 263 |
+
flush_text()
|
| 264 |
+
media_type = _guess_image_media_type(image_bytes)
|
| 265 |
+
data = base64.b64encode(image_bytes).decode("utf-8")
|
| 266 |
+
slide_sections.append(
|
| 267 |
+
Section(
|
| 268 |
+
content=DataBlock(
|
| 269 |
+
source=Base64Source(
|
| 270 |
+
media_type=media_type,
|
| 271 |
+
data=data,
|
| 272 |
+
),
|
| 273 |
+
name=filename,
|
| 274 |
+
),
|
| 275 |
+
source=filename,
|
| 276 |
+
metadata={
|
| 277 |
+
"slide": slide_no,
|
| 278 |
+
"media_type": media_type,
|
| 279 |
+
},
|
| 280 |
+
),
|
| 281 |
+
)
|
| 282 |
+
continue
|
| 283 |
+
|
| 284 |
+
# 2. Tables — render to text; optionally flush around them.
|
| 285 |
+
if getattr(shape, "has_table", False):
|
| 286 |
+
try:
|
| 287 |
+
rows = _extract_table_rows(shape.table)
|
| 288 |
+
except Exception as e: # pylint: disable=broad-except
|
| 289 |
+
logger.warning(
|
| 290 |
+
"Failed to extract table from slide %d: %s",
|
| 291 |
+
slide_no,
|
| 292 |
+
e,
|
| 293 |
+
)
|
| 294 |
+
continue
|
| 295 |
+
rendered = (
|
| 296 |
+
_table_to_markdown(rows)
|
| 297 |
+
if self.table_format == "markdown"
|
| 298 |
+
else _table_to_json(rows)
|
| 299 |
+
)
|
| 300 |
+
if not rendered:
|
| 301 |
+
continue
|
| 302 |
+
if self.separate_table:
|
| 303 |
+
flush_text()
|
| 304 |
+
slide_sections.append(
|
| 305 |
+
Section(
|
| 306 |
+
content=TextBlock(text=rendered),
|
| 307 |
+
source=filename,
|
| 308 |
+
metadata={"slide": slide_no},
|
| 309 |
+
),
|
| 310 |
+
)
|
| 311 |
+
else:
|
| 312 |
+
text_buffer.append(rendered)
|
| 313 |
+
continue
|
| 314 |
+
|
| 315 |
+
# 3. Text frames.
|
| 316 |
+
if getattr(shape, "has_text_frame", False):
|
| 317 |
+
try:
|
| 318 |
+
parts = [
|
| 319 |
+
para.text.strip()
|
| 320 |
+
for para in shape.text_frame.paragraphs
|
| 321 |
+
if para.text.strip()
|
| 322 |
+
]
|
| 323 |
+
except Exception as e: # pylint: disable=broad-except
|
| 324 |
+
logger.warning(
|
| 325 |
+
"Failed to extract text from shape in slide %d: %s",
|
| 326 |
+
slide_no,
|
| 327 |
+
e,
|
| 328 |
+
)
|
| 329 |
+
continue
|
| 330 |
+
if parts:
|
| 331 |
+
text_buffer.append("\n".join(parts))
|
| 332 |
+
|
| 333 |
+
# Suffix goes onto the last text section of this slide. If the
|
| 334 |
+
# slide ends on an image (text_buffer empty, no prior text
|
| 335 |
+
# section in this slide), create a small text-only carrier so
|
| 336 |
+
# the suffix is preserved.
|
| 337 |
+
if self.slide_suffix is not None:
|
| 338 |
+
text_buffer.append(self.slide_suffix)
|
| 339 |
+
|
| 340 |
+
flush_text()
|
| 341 |
+
return slide_sections
|
src/agentscope/rag/_parser/_text.py
ADDED
|
@@ -0,0 +1,125 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""Plain-text file parser."""
|
| 3 |
+
import os
|
| 4 |
+
|
| 5 |
+
from ...message import TextBlock
|
| 6 |
+
from .._document import Section
|
| 7 |
+
from ._base import ParserBase
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
class TextParser(ParserBase):
|
| 11 |
+
"""Parser for plain-text file formats.
|
| 12 |
+
|
| 13 |
+
Reads the entire file as UTF-8 text and returns a single
|
| 14 |
+
:class:`Section`. No internal boundaries are inferred — the file
|
| 15 |
+
is treated as one unstructured blob, leaving all splitting to a
|
| 16 |
+
downstream :class:`~agentscope.rag.ChunkerBase`.
|
| 17 |
+
|
| 18 |
+
Supports a fixed set of standard text-based IANA media types
|
| 19 |
+
(``text/plain``, ``text/markdown``, ``text/csv``, …). Use
|
| 20 |
+
``TextParser.supported_media_types`` to enumerate them.
|
| 21 |
+
"""
|
| 22 |
+
|
| 23 |
+
supported_media_types: list[str] = [
|
| 24 |
+
"text/plain",
|
| 25 |
+
"text/markdown",
|
| 26 |
+
"text/csv",
|
| 27 |
+
"text/html",
|
| 28 |
+
"text/x-rst",
|
| 29 |
+
"application/json",
|
| 30 |
+
"application/xml",
|
| 31 |
+
"application/x-yaml",
|
| 32 |
+
]
|
| 33 |
+
"""Standard IANA media types this parser handles."""
|
| 34 |
+
|
| 35 |
+
@classmethod
|
| 36 |
+
def supported_extensions(cls) -> list[str]:
|
| 37 |
+
"""Return the human-friendly text extensions.
|
| 38 |
+
|
| 39 |
+
Override the base reverse-lookup because
|
| 40 |
+
:func:`mimetypes.guess_all_extensions` returns a long tail of
|
| 41 |
+
developer-tool extensions for ``text/plain`` (``.bat`` /
|
| 42 |
+
``.c`` / ``.pl`` / ``.ksh`` / …) that have no place in a KB
|
| 43 |
+
file picker, and returns the empty list for
|
| 44 |
+
``application/x-yaml``.
|
| 45 |
+
"""
|
| 46 |
+
return [
|
| 47 |
+
".csv",
|
| 48 |
+
".htm",
|
| 49 |
+
".html",
|
| 50 |
+
".json",
|
| 51 |
+
".markdown",
|
| 52 |
+
".md",
|
| 53 |
+
".rst",
|
| 54 |
+
".txt",
|
| 55 |
+
".xml",
|
| 56 |
+
".yaml",
|
| 57 |
+
".yml",
|
| 58 |
+
]
|
| 59 |
+
|
| 60 |
+
def __init__(self, encoding: str = "utf-8") -> None:
|
| 61 |
+
"""Initialize the text parser.
|
| 62 |
+
|
| 63 |
+
Args:
|
| 64 |
+
encoding (`str`, defaults to ``"utf-8"``):
|
| 65 |
+
The text encoding used to decode the file bytes.
|
| 66 |
+
"""
|
| 67 |
+
self.encoding = encoding
|
| 68 |
+
|
| 69 |
+
async def parse(
|
| 70 |
+
self,
|
| 71 |
+
file: bytes | str,
|
| 72 |
+
filename: str,
|
| 73 |
+
) -> list[Section]:
|
| 74 |
+
"""Read the file as text and return a single :class:`Section`.
|
| 75 |
+
|
| 76 |
+
Args:
|
| 77 |
+
file (`bytes | str`):
|
| 78 |
+
The file content. ``bytes`` is decoded with the
|
| 79 |
+
configured encoding. ``str`` is disambiguated at
|
| 80 |
+
runtime: if it names an existing file on disk the
|
| 81 |
+
file is read and decoded; otherwise it is used
|
| 82 |
+
verbatim as pre-decoded text — letting local-mode
|
| 83 |
+
callers skip the encode → decode round trip.
|
| 84 |
+
filename (`str`):
|
| 85 |
+
The source filename, copied verbatim into
|
| 86 |
+
:attr:`Section.source`.
|
| 87 |
+
|
| 88 |
+
Returns:
|
| 89 |
+
`list[Section]`:
|
| 90 |
+
Always a one-element list containing the entire file
|
| 91 |
+
contents.
|
| 92 |
+
|
| 93 |
+
Raises:
|
| 94 |
+
`ValueError`: If the bytes cannot be decoded with the
|
| 95 |
+
configured encoding.
|
| 96 |
+
"""
|
| 97 |
+
if isinstance(file, str):
|
| 98 |
+
if os.path.isfile(file):
|
| 99 |
+
with open(file, "rb") as fp:
|
| 100 |
+
raw = fp.read()
|
| 101 |
+
try:
|
| 102 |
+
text = raw.decode(self.encoding)
|
| 103 |
+
except UnicodeDecodeError as e:
|
| 104 |
+
raise ValueError(
|
| 105 |
+
f"Failed to decode {filename!r} as "
|
| 106 |
+
f"{self.encoding!r}: {e}",
|
| 107 |
+
) from e
|
| 108 |
+
else:
|
| 109 |
+
text = file
|
| 110 |
+
else:
|
| 111 |
+
try:
|
| 112 |
+
text = file.decode(self.encoding)
|
| 113 |
+
except UnicodeDecodeError as e:
|
| 114 |
+
raise ValueError(
|
| 115 |
+
f"Failed to decode {filename!r} as "
|
| 116 |
+
f"{self.encoding!r}: {e}",
|
| 117 |
+
) from e
|
| 118 |
+
|
| 119 |
+
return [
|
| 120 |
+
Section(
|
| 121 |
+
content=TextBlock(text=text),
|
| 122 |
+
source=filename,
|
| 123 |
+
metadata={},
|
| 124 |
+
),
|
| 125 |
+
]
|
src/agentscope/rag/_parser/_utils.py
ADDED
|
@@ -0,0 +1,99 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""Shared helpers for binary parsers.
|
| 3 |
+
|
| 4 |
+
Three small utilities used by :class:`PDFParser`, :class:`ImageParser`,
|
| 5 |
+
and :class:`PPTParser`:
|
| 6 |
+
|
| 7 |
+
- :func:`_guess_image_media_type` — sniff the IANA media type from raw
|
| 8 |
+
image bytes by looking at the magic number. Used to populate the
|
| 9 |
+
``media_type`` field of :class:`Base64Source` for embedded images.
|
| 10 |
+
- :func:`_table_to_markdown` — render a 2-D ``list[list[str]]`` as a
|
| 11 |
+
Markdown pipe-table; the default rendering for table content.
|
| 12 |
+
- :func:`_table_to_json` — render the same shape as a JSON array
|
| 13 |
+
prefixed with a one-line ``<system-info>`` marker; used when the
|
| 14 |
+
caller picks ``table_format="json"`` to avoid Markdown's
|
| 15 |
+
multi-line-cell ambiguity.
|
| 16 |
+
"""
|
| 17 |
+
import json
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def _guess_image_media_type(data: bytes) -> str:
|
| 21 |
+
"""Sniff the IANA media type of an image from its magic number.
|
| 22 |
+
|
| 23 |
+
Args:
|
| 24 |
+
data (`bytes`):
|
| 25 |
+
The raw image bytes.
|
| 26 |
+
|
| 27 |
+
Returns:
|
| 28 |
+
`str`:
|
| 29 |
+
The IANA media type (e.g. ``"image/png"``). Falls back to
|
| 30 |
+
``"image/jpeg"`` when no signature matches — JPEG is the
|
| 31 |
+
most permissive default and matches what every consumer can
|
| 32 |
+
decode.
|
| 33 |
+
"""
|
| 34 |
+
signatures = {
|
| 35 |
+
b"\x89PNG\r\n\x1a\n": "image/png",
|
| 36 |
+
b"\xff\xd8": "image/jpeg",
|
| 37 |
+
b"GIF87a": "image/gif",
|
| 38 |
+
b"GIF89a": "image/gif",
|
| 39 |
+
b"BM": "image/bmp",
|
| 40 |
+
}
|
| 41 |
+
for signature, media_type in signatures.items():
|
| 42 |
+
if data.startswith(signature):
|
| 43 |
+
return media_type
|
| 44 |
+
# WebP: ``RIFF`` at offset 0 + ``WEBP`` at offset 8.
|
| 45 |
+
if len(data) > 12 and data[:4] == b"RIFF" and data[8:12] == b"WEBP":
|
| 46 |
+
return "image/webp"
|
| 47 |
+
return "image/jpeg"
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def _table_to_markdown(table_data: list[list[str]]) -> str:
|
| 51 |
+
"""Render a 2-D table as a Markdown pipe-table.
|
| 52 |
+
|
| 53 |
+
Args:
|
| 54 |
+
table_data (`list[list[str]]`):
|
| 55 |
+
The table data; ``table_data[0]`` is the header row.
|
| 56 |
+
|
| 57 |
+
Returns:
|
| 58 |
+
`str`:
|
| 59 |
+
The Markdown rendering, or the empty string when
|
| 60 |
+
``table_data`` is empty / column-less.
|
| 61 |
+
"""
|
| 62 |
+
if not table_data:
|
| 63 |
+
return ""
|
| 64 |
+
|
| 65 |
+
num_cols = len(table_data[0])
|
| 66 |
+
if num_cols == 0:
|
| 67 |
+
return ""
|
| 68 |
+
|
| 69 |
+
lines = [
|
| 70 |
+
"| " + " | ".join(table_data[0]) + " |",
|
| 71 |
+
"| " + " | ".join(["---"] * num_cols) + " |",
|
| 72 |
+
]
|
| 73 |
+
for row in table_data[1:]:
|
| 74 |
+
# Pad short rows so column counts match the header.
|
| 75 |
+
padded = list(row) + [""] * max(0, num_cols - len(row))
|
| 76 |
+
lines.append("| " + " | ".join(padded[:num_cols]) + " |")
|
| 77 |
+
return "\n".join(lines) + "\n"
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
def _table_to_json(table_data: list[list[str]]) -> str:
|
| 81 |
+
"""Render a 2-D table as JSON prefixed by a one-line marker.
|
| 82 |
+
|
| 83 |
+
The ``<system-info>`` marker lets the LLM (and any
|
| 84 |
+
structure-aware downstream renderer) tell at a glance that this
|
| 85 |
+
block is tabular data, not free text — which is otherwise
|
| 86 |
+
indistinguishable from a raw JSON dump in the middle of a chunk.
|
| 87 |
+
|
| 88 |
+
Args:
|
| 89 |
+
table_data (`list[list[str]]`):
|
| 90 |
+
The table data.
|
| 91 |
+
|
| 92 |
+
Returns:
|
| 93 |
+
`str`:
|
| 94 |
+
``"<system-info>...</system-info>\\n<json>"``.
|
| 95 |
+
"""
|
| 96 |
+
return (
|
| 97 |
+
"<system-info>A table loaded as a JSON array:</system-info>\n"
|
| 98 |
+
+ json.dumps(table_data, ensure_ascii=False)
|
| 99 |
+
)
|
src/agentscope/rag/_vdb/__init__.py
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The vector store classes in AgentScope."""
|
| 3 |
+
|
| 4 |
+
from ._vector_store import (
|
| 5 |
+
DocumentSummary,
|
| 6 |
+
VectorRecord,
|
| 7 |
+
VectorSearchResult,
|
| 8 |
+
VectorStoreBase,
|
| 9 |
+
)
|
| 10 |
+
from ._qdrant import QdrantStore
|
| 11 |
+
|
| 12 |
+
__all__ = [
|
| 13 |
+
"DocumentSummary",
|
| 14 |
+
"VectorStoreBase",
|
| 15 |
+
"VectorRecord",
|
| 16 |
+
"VectorSearchResult",
|
| 17 |
+
"QdrantStore",
|
| 18 |
+
]
|
src/agentscope/rag/_vdb/_qdrant.py
ADDED
|
@@ -0,0 +1,392 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""Qdrant implementation of the vector store backend.
|
| 3 |
+
|
| 4 |
+
Built on the official ``qdrant-client`` SDK using its fully
|
| 5 |
+
asynchronous client (:class:`~qdrant_client.AsyncQdrantClient`), so all
|
| 6 |
+
operations are non-blocking and safe to call from the application's
|
| 7 |
+
event loop.
|
| 8 |
+
|
| 9 |
+
The same class supports all Qdrant deployment modes through the
|
| 10 |
+
constructor arguments:
|
| 11 |
+
|
| 12 |
+
- ``location=":memory:"`` — in-process, ephemeral (ideal for tests)
|
| 13 |
+
- ``path="/path/to/db"`` — in-process, persisted to local disk
|
| 14 |
+
- ``url="http://localhost:6333"`` — remote Qdrant server / cloud
|
| 15 |
+
"""
|
| 16 |
+
import uuid
|
| 17 |
+
from typing import TYPE_CHECKING, Any, Literal
|
| 18 |
+
|
| 19 |
+
from ._vector_store import (
|
| 20 |
+
DocumentSummary,
|
| 21 |
+
VectorRecord,
|
| 22 |
+
VectorSearchResult,
|
| 23 |
+
VectorStoreBase,
|
| 24 |
+
)
|
| 25 |
+
from .._document import Chunk
|
| 26 |
+
|
| 27 |
+
if TYPE_CHECKING:
|
| 28 |
+
from qdrant_client import AsyncQdrantClient
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
class QdrantStore(VectorStoreBase):
|
| 32 |
+
"""Vector store backend backed by `Qdrant <https://qdrant.tech>`_.
|
| 33 |
+
|
| 34 |
+
Each knowledge base maps to one Qdrant collection. Every point
|
| 35 |
+
payload stores the owning ``document_id`` plus the serialized
|
| 36 |
+
:class:`~agentscope.rag.Chunk`, which is reconstructed on
|
| 37 |
+
retrieval.
|
| 38 |
+
|
| 39 |
+
.. note:: The ``qdrant-client`` package is required. Install it
|
| 40 |
+
with ``pip install qdrant-client``.
|
| 41 |
+
|
| 42 |
+
.. code-block:: python
|
| 43 |
+
|
| 44 |
+
# In-memory (tests / prototyping)
|
| 45 |
+
store = QdrantStore(location=":memory:")
|
| 46 |
+
|
| 47 |
+
# Remote server
|
| 48 |
+
store = QdrantStore(
|
| 49 |
+
url="http://localhost:6333",
|
| 50 |
+
api_key="...",
|
| 51 |
+
)
|
| 52 |
+
|
| 53 |
+
async with store:
|
| 54 |
+
await store.create_collection("kb-1", dimensions=768)
|
| 55 |
+
|
| 56 |
+
"""
|
| 57 |
+
|
| 58 |
+
def __init__(
|
| 59 |
+
self,
|
| 60 |
+
location: str | None = None,
|
| 61 |
+
url: str | None = None,
|
| 62 |
+
path: str | None = None,
|
| 63 |
+
api_key: str | None = None,
|
| 64 |
+
distance: Literal["Cosine", "Dot", "Euclid", "Manhattan"] = "Cosine",
|
| 65 |
+
client_kwargs: dict[str, Any] | None = None,
|
| 66 |
+
) -> None:
|
| 67 |
+
"""Initialize the Qdrant vector store.
|
| 68 |
+
|
| 69 |
+
Args:
|
| 70 |
+
location (`str | None`, optional):
|
| 71 |
+
Pass ``":memory:"`` for an ephemeral in-process
|
| 72 |
+
instance. Mutually exclusive with ``url`` and ``path``.
|
| 73 |
+
url (`str | None`, optional):
|
| 74 |
+
The URL of a remote Qdrant server, e.g.
|
| 75 |
+
``"http://localhost:6333"``.
|
| 76 |
+
path (`str | None`, optional):
|
| 77 |
+
A local directory for an in-process, on-disk instance.
|
| 78 |
+
api_key (`str | None`, optional):
|
| 79 |
+
The API key for Qdrant Cloud or a secured server.
|
| 80 |
+
distance (`Literal["Cosine", "Dot", "Euclid", "Manhattan"]`, \
|
| 81 |
+
defaults to ``"Cosine"``):
|
| 82 |
+
The distance metric used when creating collections.
|
| 83 |
+
client_kwargs (`dict[str, Any] | None`, optional):
|
| 84 |
+
Extra keyword arguments forwarded to
|
| 85 |
+
:class:`~qdrant_client.AsyncQdrantClient`.
|
| 86 |
+
"""
|
| 87 |
+
self._location = location
|
| 88 |
+
self._url = url
|
| 89 |
+
self._path = path
|
| 90 |
+
self._api_key = api_key
|
| 91 |
+
self._distance = distance
|
| 92 |
+
self._client_kwargs = client_kwargs or {}
|
| 93 |
+
self._client: "AsyncQdrantClient | None" = None
|
| 94 |
+
|
| 95 |
+
def get_client(self) -> "AsyncQdrantClient":
|
| 96 |
+
"""Lazily create and cache the async Qdrant client.
|
| 97 |
+
|
| 98 |
+
Returns:
|
| 99 |
+
`AsyncQdrantClient`:
|
| 100 |
+
The shared async client instance.
|
| 101 |
+
"""
|
| 102 |
+
if self._client is None:
|
| 103 |
+
from qdrant_client import AsyncQdrantClient
|
| 104 |
+
|
| 105 |
+
self._client = AsyncQdrantClient(
|
| 106 |
+
location=self._location,
|
| 107 |
+
url=self._url,
|
| 108 |
+
path=self._path,
|
| 109 |
+
api_key=self._api_key,
|
| 110 |
+
**self._client_kwargs,
|
| 111 |
+
)
|
| 112 |
+
return self._client
|
| 113 |
+
|
| 114 |
+
async def __aexit__(
|
| 115 |
+
self,
|
| 116 |
+
exc_type: type[BaseException] | None,
|
| 117 |
+
exc_val: BaseException | None,
|
| 118 |
+
exc_tb: Any,
|
| 119 |
+
) -> None:
|
| 120 |
+
"""Exit the async context — close the underlying client."""
|
| 121 |
+
if self._client is not None:
|
| 122 |
+
await self._client.close()
|
| 123 |
+
self._client = None
|
| 124 |
+
|
| 125 |
+
# ------------------------------------------------------------------
|
| 126 |
+
# Collection management
|
| 127 |
+
# ------------------------------------------------------------------
|
| 128 |
+
|
| 129 |
+
async def create_collection(
|
| 130 |
+
self,
|
| 131 |
+
name: str,
|
| 132 |
+
dimensions: int,
|
| 133 |
+
) -> None:
|
| 134 |
+
"""Create a new Qdrant collection.
|
| 135 |
+
|
| 136 |
+
No-op if the collection already exists.
|
| 137 |
+
|
| 138 |
+
Args:
|
| 139 |
+
name (`str`):
|
| 140 |
+
The collection name. Typically, the knowledge base ID.
|
| 141 |
+
dimensions (`int`):
|
| 142 |
+
The fixed vector dimensionality for this collection.
|
| 143 |
+
"""
|
| 144 |
+
from qdrant_client import models
|
| 145 |
+
|
| 146 |
+
client = self.get_client()
|
| 147 |
+
if await client.collection_exists(name):
|
| 148 |
+
return
|
| 149 |
+
await client.create_collection(
|
| 150 |
+
collection_name=name,
|
| 151 |
+
vectors_config=models.VectorParams(
|
| 152 |
+
size=dimensions,
|
| 153 |
+
distance=models.Distance(self._distance),
|
| 154 |
+
),
|
| 155 |
+
)
|
| 156 |
+
|
| 157 |
+
async def delete_collection(self, name: str) -> None:
|
| 158 |
+
"""Delete a collection and all its data.
|
| 159 |
+
|
| 160 |
+
Args:
|
| 161 |
+
name (`str`):
|
| 162 |
+
The collection name to delete.
|
| 163 |
+
"""
|
| 164 |
+
await self.get_client().delete_collection(name)
|
| 165 |
+
|
| 166 |
+
async def has_collection(self, name: str) -> bool:
|
| 167 |
+
"""Check whether a collection exists.
|
| 168 |
+
|
| 169 |
+
Args:
|
| 170 |
+
name (`str`):
|
| 171 |
+
The collection name to check.
|
| 172 |
+
|
| 173 |
+
Returns:
|
| 174 |
+
`bool`: ``True`` if the collection exists.
|
| 175 |
+
"""
|
| 176 |
+
return await self.get_client().collection_exists(name)
|
| 177 |
+
|
| 178 |
+
# ------------------------------------------------------------------
|
| 179 |
+
# Data operations
|
| 180 |
+
# ------------------------------------------------------------------
|
| 181 |
+
|
| 182 |
+
async def insert(
|
| 183 |
+
self,
|
| 184 |
+
collection: str,
|
| 185 |
+
records: list[VectorRecord],
|
| 186 |
+
) -> None:
|
| 187 |
+
"""Insert records into a collection.
|
| 188 |
+
|
| 189 |
+
Each point payload stores the :attr:`VectorRecord.document_id`
|
| 190 |
+
under the ``document_id`` key and the serialized
|
| 191 |
+
:class:`Chunk` under the ``chunk`` key, so that :meth:`delete`
|
| 192 |
+
can remove all records of one document.
|
| 193 |
+
|
| 194 |
+
Args:
|
| 195 |
+
collection (`str`):
|
| 196 |
+
The target collection name.
|
| 197 |
+
records (`list[VectorRecord]`):
|
| 198 |
+
The records to insert (each carrying a
|
| 199 |
+
:class:`Chunk` and its embedding vector).
|
| 200 |
+
"""
|
| 201 |
+
|
| 202 |
+
from qdrant_client import models
|
| 203 |
+
|
| 204 |
+
if not records:
|
| 205 |
+
return
|
| 206 |
+
await self.get_client().upsert(
|
| 207 |
+
collection_name=collection,
|
| 208 |
+
points=[
|
| 209 |
+
models.PointStruct(
|
| 210 |
+
id=str(uuid.uuid4()),
|
| 211 |
+
vector=record.vector,
|
| 212 |
+
payload={
|
| 213 |
+
"document_id": record.document_id,
|
| 214 |
+
"chunk": record.chunk.model_dump(mode="json"),
|
| 215 |
+
},
|
| 216 |
+
)
|
| 217 |
+
for record in records
|
| 218 |
+
],
|
| 219 |
+
)
|
| 220 |
+
|
| 221 |
+
async def delete(
|
| 222 |
+
self,
|
| 223 |
+
collection: str,
|
| 224 |
+
document_id: str,
|
| 225 |
+
) -> None:
|
| 226 |
+
"""Delete all records belonging to one source document.
|
| 227 |
+
|
| 228 |
+
Matches the ``document_id`` payload key written by
|
| 229 |
+
:meth:`insert` from :attr:`VectorRecord.document_id`.
|
| 230 |
+
|
| 231 |
+
Args:
|
| 232 |
+
collection (`str`):
|
| 233 |
+
The target collection name.
|
| 234 |
+
document_id (`str`):
|
| 235 |
+
The source document ID whose records should be
|
| 236 |
+
removed.
|
| 237 |
+
"""
|
| 238 |
+
from qdrant_client import models
|
| 239 |
+
|
| 240 |
+
await self.get_client().delete(
|
| 241 |
+
collection_name=collection,
|
| 242 |
+
points_selector=models.FilterSelector(
|
| 243 |
+
filter=models.Filter(
|
| 244 |
+
must=[
|
| 245 |
+
models.FieldCondition(
|
| 246 |
+
key="document_id",
|
| 247 |
+
match=models.MatchValue(value=document_id),
|
| 248 |
+
),
|
| 249 |
+
],
|
| 250 |
+
),
|
| 251 |
+
),
|
| 252 |
+
)
|
| 253 |
+
|
| 254 |
+
# ------------------------------------------------------------------
|
| 255 |
+
# Search
|
| 256 |
+
# ------------------------------------------------------------------
|
| 257 |
+
|
| 258 |
+
async def search(
|
| 259 |
+
self,
|
| 260 |
+
collection: str,
|
| 261 |
+
query_vector: list[float],
|
| 262 |
+
top_k: int = 5,
|
| 263 |
+
metadata_filter: dict[str, Any] | None = None,
|
| 264 |
+
) -> list[VectorSearchResult]:
|
| 265 |
+
"""Find the most similar records to a query vector.
|
| 266 |
+
|
| 267 |
+
Args:
|
| 268 |
+
collection (`str`):
|
| 269 |
+
The collection to search.
|
| 270 |
+
query_vector (`list[float]`):
|
| 271 |
+
The query embedding vector.
|
| 272 |
+
top_k (`int`, defaults to ``5``):
|
| 273 |
+
Maximum number of results to return.
|
| 274 |
+
metadata_filter (`dict[str, Any] | None`, optional):
|
| 275 |
+
If provided, restrict the search to records whose
|
| 276 |
+
``chunk.metadata`` matches every ``key == value`` pair
|
| 277 |
+
in this dict (translated into a Qdrant ``must`` payload
|
| 278 |
+
filter against ``chunk.metadata.<key>``).
|
| 279 |
+
|
| 280 |
+
Returns:
|
| 281 |
+
`list[VectorSearchResult]`:
|
| 282 |
+
Results ordered by descending similarity score.
|
| 283 |
+
"""
|
| 284 |
+
response = await self.get_client().query_points(
|
| 285 |
+
collection_name=collection,
|
| 286 |
+
query=query_vector,
|
| 287 |
+
limit=top_k,
|
| 288 |
+
with_payload=True,
|
| 289 |
+
query_filter=self._build_metadata_filter(metadata_filter),
|
| 290 |
+
)
|
| 291 |
+
return [
|
| 292 |
+
VectorSearchResult(
|
| 293 |
+
score=point.score,
|
| 294 |
+
document_id=point.payload["document_id"],
|
| 295 |
+
chunk=Chunk.model_validate(point.payload["chunk"]),
|
| 296 |
+
)
|
| 297 |
+
for point in response.points
|
| 298 |
+
]
|
| 299 |
+
|
| 300 |
+
# ------------------------------------------------------------------
|
| 301 |
+
# Document listing
|
| 302 |
+
# ------------------------------------------------------------------
|
| 303 |
+
|
| 304 |
+
async def list_documents(
|
| 305 |
+
self,
|
| 306 |
+
collection: str,
|
| 307 |
+
metadata_filter: dict[str, Any] | None = None,
|
| 308 |
+
) -> list[DocumentSummary]:
|
| 309 |
+
"""List all distinct source documents indexed in a collection.
|
| 310 |
+
|
| 311 |
+
Scrolls the collection in payload-only mode (vectors disabled)
|
| 312 |
+
and aggregates by ``document_id``. The first chunk encountered
|
| 313 |
+
for each document supplies the ``source`` filename and the
|
| 314 |
+
document-level ``metadata``.
|
| 315 |
+
|
| 316 |
+
Args:
|
| 317 |
+
collection (`str`):
|
| 318 |
+
The target collection name.
|
| 319 |
+
metadata_filter (`dict[str, Any] | None`, optional):
|
| 320 |
+
If provided, restrict aggregation to records whose
|
| 321 |
+
``chunk.metadata`` matches every ``key == value`` pair.
|
| 322 |
+
|
| 323 |
+
Returns:
|
| 324 |
+
`list[DocumentSummary]`:
|
| 325 |
+
One summary per distinct ``document_id``.
|
| 326 |
+
"""
|
| 327 |
+
client = self.get_client()
|
| 328 |
+
query_filter = self._build_metadata_filter(metadata_filter)
|
| 329 |
+
summaries: dict[str, DocumentSummary] = {}
|
| 330 |
+
offset: Any = None
|
| 331 |
+
|
| 332 |
+
while True:
|
| 333 |
+
points, next_offset = await client.scroll(
|
| 334 |
+
collection_name=collection,
|
| 335 |
+
scroll_filter=query_filter,
|
| 336 |
+
limit=256,
|
| 337 |
+
offset=offset,
|
| 338 |
+
with_payload=True,
|
| 339 |
+
with_vectors=False,
|
| 340 |
+
)
|
| 341 |
+
for point in points:
|
| 342 |
+
doc_id = point.payload["document_id"]
|
| 343 |
+
summary = summaries.get(doc_id)
|
| 344 |
+
if summary is None:
|
| 345 |
+
chunk_payload = point.payload["chunk"]
|
| 346 |
+
summaries[doc_id] = DocumentSummary(
|
| 347 |
+
document_id=doc_id,
|
| 348 |
+
source=chunk_payload.get("source", ""),
|
| 349 |
+
chunk_count=1,
|
| 350 |
+
metadata=dict(chunk_payload.get("metadata", {})),
|
| 351 |
+
)
|
| 352 |
+
else:
|
| 353 |
+
summary.chunk_count += 1
|
| 354 |
+
if next_offset is None:
|
| 355 |
+
break
|
| 356 |
+
offset = next_offset
|
| 357 |
+
|
| 358 |
+
return list(summaries.values())
|
| 359 |
+
|
| 360 |
+
@staticmethod
|
| 361 |
+
def _build_metadata_filter(
|
| 362 |
+
metadata_filter: dict[str, Any] | None,
|
| 363 |
+
) -> Any:
|
| 364 |
+
"""Translate a flat ``{key: value}`` filter into a Qdrant filter.
|
| 365 |
+
|
| 366 |
+
Each ``key`` is matched against the corresponding nested path
|
| 367 |
+
``chunk.metadata.<key>`` written by :meth:`insert`. Returns
|
| 368 |
+
``None`` when ``metadata_filter`` is empty so that callers
|
| 369 |
+
skip the filter argument entirely.
|
| 370 |
+
|
| 371 |
+
Args:
|
| 372 |
+
metadata_filter (`dict[str, Any] | None`):
|
| 373 |
+
The flat filter, or ``None`` for no filter.
|
| 374 |
+
|
| 375 |
+
Returns:
|
| 376 |
+
`qdrant_client.models.Filter | None`:
|
| 377 |
+
A Qdrant ``Filter`` object, or ``None``.
|
| 378 |
+
"""
|
| 379 |
+
if not metadata_filter:
|
| 380 |
+
return None
|
| 381 |
+
|
| 382 |
+
from qdrant_client import models
|
| 383 |
+
|
| 384 |
+
return models.Filter(
|
| 385 |
+
must=[
|
| 386 |
+
models.FieldCondition(
|
| 387 |
+
key=f"chunk.metadata.{key}",
|
| 388 |
+
match=models.MatchValue(value=value),
|
| 389 |
+
)
|
| 390 |
+
for key, value in metadata_filter.items()
|
| 391 |
+
],
|
| 392 |
+
)
|
src/agentscope/rag/_vdb/_vector_store.py
ADDED
|
@@ -0,0 +1,291 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""Abstract base class for vector store backends.
|
| 3 |
+
|
| 4 |
+
A :class:`VectorStoreBase` instance is the single connection point to
|
| 5 |
+
one vector database deployment. It is created once at application
|
| 6 |
+
startup, passed into ``create_app(vector_store=...)``, and shared
|
| 7 |
+
across all requests for the lifetime of the process — similar to
|
| 8 |
+
:class:`~agentscope.app.storage.StorageBase` and
|
| 9 |
+
:class:`~agentscope.app.message_bus.MessageBus`.
|
| 10 |
+
|
| 11 |
+
Each **knowledge base** maps to one **collection** inside the vector
|
| 12 |
+
store. Collections are isolated: different knowledge bases never
|
| 13 |
+
share a collection, so retrieval is always scoped to a single
|
| 14 |
+
collection without cross-collection filtering.
|
| 15 |
+
|
| 16 |
+
Lifecycle is managed via the async context manager protocol
|
| 17 |
+
(``__aenter__`` / ``__aexit__``), which the app lifespan calls
|
| 18 |
+
automatically.
|
| 19 |
+
"""
|
| 20 |
+
from abc import ABC, abstractmethod
|
| 21 |
+
from typing import Any, Self
|
| 22 |
+
|
| 23 |
+
from pydantic import BaseModel, Field
|
| 24 |
+
|
| 25 |
+
from .._document import Chunk
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
class VectorRecord(BaseModel):
|
| 29 |
+
"""A single record to insert into a vector store collection.
|
| 30 |
+
|
| 31 |
+
Pairs a :class:`Chunk` (the business payload — content, source,
|
| 32 |
+
structural metadata) with its dense embedding vector. ``Chunk``
|
| 33 |
+
is intentionally not extended with an ``embedding`` field so its
|
| 34 |
+
semantics stay stable across the indexing pipeline; instead the
|
| 35 |
+
vector lives in this wrapper whose only purpose is "I am about
|
| 36 |
+
to be inserted into a vector database."
|
| 37 |
+
"""
|
| 38 |
+
|
| 39 |
+
vector: list[float]
|
| 40 |
+
"""The dense embedding vector for :attr:`chunk`."""
|
| 41 |
+
|
| 42 |
+
document_id: str
|
| 43 |
+
"""The ID of the source document this record belongs to.
|
| 44 |
+
Assigned by the knowledge base layer when the document is
|
| 45 |
+
registered. Backends must persist it at insertion time so that
|
| 46 |
+
:meth:`VectorStoreBase.delete` can remove all records of one
|
| 47 |
+
document as a unit."""
|
| 48 |
+
|
| 49 |
+
chunk: Chunk
|
| 50 |
+
"""The business payload — content, source, structural metadata."""
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
class VectorSearchResult(BaseModel):
|
| 54 |
+
"""A single result returned by a similarity search.
|
| 55 |
+
|
| 56 |
+
Pairs the matched :class:`Chunk` with its similarity score.
|
| 57 |
+
``Chunk`` is intentionally not extended with a ``score`` field so
|
| 58 |
+
its semantics stay stable; instead the score lives in this
|
| 59 |
+
wrapper whose only purpose is "I am a query hit."
|
| 60 |
+
"""
|
| 61 |
+
|
| 62 |
+
score: float
|
| 63 |
+
"""Similarity score. Higher = more similar for cosine /
|
| 64 |
+
dot-product; lower = more similar for L2 distance."""
|
| 65 |
+
|
| 66 |
+
document_id: str
|
| 67 |
+
"""The ID of the source document the matched chunk belongs to —
|
| 68 |
+
the same value carried by :attr:`VectorRecord.document_id` at
|
| 69 |
+
insertion time. Lets callers cite, group, or delete the source
|
| 70 |
+
document of a hit."""
|
| 71 |
+
|
| 72 |
+
chunk: Chunk
|
| 73 |
+
"""The matched business payload."""
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
class DocumentSummary(BaseModel):
|
| 77 |
+
"""A lightweight description of one source document inside a collection.
|
| 78 |
+
|
| 79 |
+
Aggregated by :meth:`VectorStoreBase.list_documents` from the
|
| 80 |
+
records of each ``document_id`` — the vector store is the single
|
| 81 |
+
source of truth for "what documents exist in a knowledge base".
|
| 82 |
+
"""
|
| 83 |
+
|
| 84 |
+
document_id: str
|
| 85 |
+
"""The source document identifier — the same value carried by
|
| 86 |
+
:attr:`VectorRecord.document_id` at insertion time."""
|
| 87 |
+
|
| 88 |
+
source: str
|
| 89 |
+
"""The original filename, taken from the first chunk encountered.
|
| 90 |
+
All chunks of the same document share the same filename so any
|
| 91 |
+
chunk yields the same value."""
|
| 92 |
+
|
| 93 |
+
chunk_count: int
|
| 94 |
+
"""The total number of chunks indexed for this document."""
|
| 95 |
+
|
| 96 |
+
metadata: dict[str, Any] = Field(default_factory=dict)
|
| 97 |
+
"""Document-level metadata propagated from the parser / uploader
|
| 98 |
+
(media type, size, upload time, ...). Taken from the first chunk
|
| 99 |
+
encountered."""
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
class VectorStoreBase(ABC):
|
| 103 |
+
"""Abstract base class for vector store backends.
|
| 104 |
+
|
| 105 |
+
Subclasses implement the concrete connection and query logic for a
|
| 106 |
+
specific vector database (Chroma, Milvus, Qdrant, FAISS, etc.).
|
| 107 |
+
|
| 108 |
+
A single instance is shared across the entire application. The
|
| 109 |
+
underlying client SDK is expected to handle connection pooling and
|
| 110 |
+
thread safety internally.
|
| 111 |
+
"""
|
| 112 |
+
|
| 113 |
+
# ------------------------------------------------------------------
|
| 114 |
+
# Lifecycle
|
| 115 |
+
# ------------------------------------------------------------------
|
| 116 |
+
|
| 117 |
+
async def __aenter__(self) -> Self:
|
| 118 |
+
"""Enter the async context — open connections if needed.
|
| 119 |
+
|
| 120 |
+
The default implementation is a no-op. Subclasses that need
|
| 121 |
+
explicit connection setup should override this.
|
| 122 |
+
|
| 123 |
+
Returns:
|
| 124 |
+
`VectorStoreBase`: ``self``.
|
| 125 |
+
"""
|
| 126 |
+
return self
|
| 127 |
+
|
| 128 |
+
async def __aexit__(
|
| 129 |
+
self,
|
| 130 |
+
exc_type: type[BaseException] | None,
|
| 131 |
+
exc_val: BaseException | None,
|
| 132 |
+
exc_tb: Any,
|
| 133 |
+
) -> None:
|
| 134 |
+
"""Exit the async context — close connections if needed.
|
| 135 |
+
|
| 136 |
+
The default implementation is a no-op.
|
| 137 |
+
"""
|
| 138 |
+
|
| 139 |
+
# ------------------------------------------------------------------
|
| 140 |
+
# Collection management
|
| 141 |
+
# ------------------------------------------------------------------
|
| 142 |
+
|
| 143 |
+
@abstractmethod
|
| 144 |
+
async def create_collection(
|
| 145 |
+
self,
|
| 146 |
+
name: str,
|
| 147 |
+
dimensions: int,
|
| 148 |
+
) -> None:
|
| 149 |
+
"""Create a new collection (vector index).
|
| 150 |
+
|
| 151 |
+
If the collection already exists, implementations should raise
|
| 152 |
+
or silently no-op depending on the backend's semantics.
|
| 153 |
+
|
| 154 |
+
Args:
|
| 155 |
+
name (`str`):
|
| 156 |
+
The collection name. Typically, the knowledge base ID.
|
| 157 |
+
dimensions (`int`):
|
| 158 |
+
The fixed vector dimensionality for this collection.
|
| 159 |
+
All vectors inserted later must have this many elements.
|
| 160 |
+
"""
|
| 161 |
+
|
| 162 |
+
@abstractmethod
|
| 163 |
+
async def delete_collection(self, name: str) -> None:
|
| 164 |
+
"""Delete a collection and all its data.
|
| 165 |
+
|
| 166 |
+
Args:
|
| 167 |
+
name (`str`):
|
| 168 |
+
The collection name to delete.
|
| 169 |
+
"""
|
| 170 |
+
|
| 171 |
+
@abstractmethod
|
| 172 |
+
async def has_collection(self, name: str) -> bool:
|
| 173 |
+
"""Check whether a collection exists.
|
| 174 |
+
|
| 175 |
+
Args:
|
| 176 |
+
name (`str`):
|
| 177 |
+
The collection name to check.
|
| 178 |
+
|
| 179 |
+
Returns:
|
| 180 |
+
`bool`: ``True`` if the collection exists.
|
| 181 |
+
"""
|
| 182 |
+
|
| 183 |
+
# ------------------------------------------------------------------
|
| 184 |
+
# Data operations
|
| 185 |
+
# ------------------------------------------------------------------
|
| 186 |
+
|
| 187 |
+
@abstractmethod
|
| 188 |
+
async def insert(
|
| 189 |
+
self,
|
| 190 |
+
collection: str,
|
| 191 |
+
records: list[VectorRecord],
|
| 192 |
+
) -> None:
|
| 193 |
+
"""Insert records into a collection.
|
| 194 |
+
|
| 195 |
+
Args:
|
| 196 |
+
collection (`str`):
|
| 197 |
+
The target collection name.
|
| 198 |
+
records (`list[VectorRecord]`):
|
| 199 |
+
The records to insert (each carrying a
|
| 200 |
+
:class:`Chunk` and its embedding vector).
|
| 201 |
+
"""
|
| 202 |
+
|
| 203 |
+
@abstractmethod
|
| 204 |
+
async def delete(
|
| 205 |
+
self,
|
| 206 |
+
collection: str,
|
| 207 |
+
document_id: str,
|
| 208 |
+
) -> None:
|
| 209 |
+
"""Delete all records belonging to one source document.
|
| 210 |
+
|
| 211 |
+
Identifies records by the :attr:`VectorRecord.document_id`
|
| 212 |
+
field that backends persist at insertion time. This matches
|
| 213 |
+
the typical RAG workflow where a user uploads or removes a
|
| 214 |
+
file as a unit.
|
| 215 |
+
|
| 216 |
+
Args:
|
| 217 |
+
collection (`str`):
|
| 218 |
+
The target collection name.
|
| 219 |
+
document_id (`str`):
|
| 220 |
+
The source document ID whose records should be
|
| 221 |
+
removed.
|
| 222 |
+
"""
|
| 223 |
+
|
| 224 |
+
# ------------------------------------------------------------------
|
| 225 |
+
# Search
|
| 226 |
+
# ------------------------------------------------------------------
|
| 227 |
+
|
| 228 |
+
@abstractmethod
|
| 229 |
+
async def search(
|
| 230 |
+
self,
|
| 231 |
+
collection: str,
|
| 232 |
+
query_vector: list[float],
|
| 233 |
+
top_k: int = 5,
|
| 234 |
+
metadata_filter: dict[str, Any] | None = None,
|
| 235 |
+
) -> list[VectorSearchResult]:
|
| 236 |
+
"""Find the most similar records to a query vector.
|
| 237 |
+
|
| 238 |
+
Args:
|
| 239 |
+
collection (`str`):
|
| 240 |
+
The collection to search.
|
| 241 |
+
query_vector (`list[float]`):
|
| 242 |
+
The query embedding vector.
|
| 243 |
+
top_k (`int`, defaults to ``5``):
|
| 244 |
+
Maximum number of results to return.
|
| 245 |
+
metadata_filter (`dict[str, Any] | None`, optional):
|
| 246 |
+
If provided, restrict the search to records whose
|
| 247 |
+
``chunk.metadata`` matches every ``key == value`` pair
|
| 248 |
+
in this dict. Backends translate this into a native
|
| 249 |
+
payload filter. Used for defense-in-depth
|
| 250 |
+
cross-tenant scoping when an isolation strategy
|
| 251 |
+
co-locates multiple knowledge bases inside the same
|
| 252 |
+
collection.
|
| 253 |
+
|
| 254 |
+
Returns:
|
| 255 |
+
`list[VectorSearchResult]`:
|
| 256 |
+
Results ordered by descending similarity score.
|
| 257 |
+
"""
|
| 258 |
+
|
| 259 |
+
# ------------------------------------------------------------------
|
| 260 |
+
# Document listing
|
| 261 |
+
# ------------------------------------------------------------------
|
| 262 |
+
|
| 263 |
+
@abstractmethod
|
| 264 |
+
async def list_documents(
|
| 265 |
+
self,
|
| 266 |
+
collection: str,
|
| 267 |
+
metadata_filter: dict[str, Any] | None = None,
|
| 268 |
+
) -> list[DocumentSummary]:
|
| 269 |
+
"""List all distinct source documents indexed in a collection.
|
| 270 |
+
|
| 271 |
+
Aggregates records by :attr:`VectorRecord.document_id` and
|
| 272 |
+
returns one :class:`DocumentSummary` per document. Backends
|
| 273 |
+
are free to use whatever scrolling / aggregation primitive
|
| 274 |
+
they expose; this method is expected to be O(documents) not
|
| 275 |
+
O(chunks) on backends that support payload-only scans.
|
| 276 |
+
|
| 277 |
+
Args:
|
| 278 |
+
collection (`str`):
|
| 279 |
+
The target collection name.
|
| 280 |
+
metadata_filter (`dict[str, Any] | None`, optional):
|
| 281 |
+
If provided, restrict aggregation to records whose
|
| 282 |
+
``chunk.metadata`` matches every ``key == value`` pair
|
| 283 |
+
in this dict. Used together with the search-time
|
| 284 |
+
filter when an isolation strategy co-locates multiple
|
| 285 |
+
knowledge bases inside the same collection.
|
| 286 |
+
|
| 287 |
+
Returns:
|
| 288 |
+
`list[DocumentSummary]`:
|
| 289 |
+
One summary per distinct ``document_id``, in
|
| 290 |
+
unspecified order.
|
| 291 |
+
"""
|
src/agentscope/skill/__init__.py
ADDED
|
@@ -0,0 +1,11 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The skill related classes and functions."""
|
| 3 |
+
|
| 4 |
+
from ._base import SkillLoaderBase, Skill
|
| 5 |
+
from ._local_loader import LocalSkillLoader
|
| 6 |
+
|
| 7 |
+
__all__ = [
|
| 8 |
+
"Skill",
|
| 9 |
+
"SkillLoaderBase",
|
| 10 |
+
"LocalSkillLoader",
|
| 11 |
+
]
|
src/agentscope/skill/_base.py
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The skill loader base class."""
|
| 3 |
+
from abc import abstractmethod, ABC
|
| 4 |
+
from dataclasses import dataclass
|
| 5 |
+
|
| 6 |
+
|
| 7 |
+
@dataclass
|
| 8 |
+
class Skill:
|
| 9 |
+
"""The agent skill class"""
|
| 10 |
+
|
| 11 |
+
name: str
|
| 12 |
+
"""The name of the skill."""
|
| 13 |
+
description: str
|
| 14 |
+
"""The description of the skill."""
|
| 15 |
+
dir: str
|
| 16 |
+
"""The directory of the agent skill."""
|
| 17 |
+
markdown: str
|
| 18 |
+
"""The markdown content of the agent skill."""
|
| 19 |
+
updated_at: float
|
| 20 |
+
"""The last updated time of the skill."""
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
class SkillLoaderBase(ABC):
|
| 24 |
+
"""The base class for skill loader."""
|
| 25 |
+
|
| 26 |
+
@abstractmethod
|
| 27 |
+
async def list_skills(self) -> list[Skill]:
|
| 28 |
+
"""List all the skills that can be loaded by this loader."""
|
| 29 |
+
raise NotImplementedError
|
src/agentscope/skill/_local_loader.py
ADDED
|
@@ -0,0 +1,171 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The local skill loader class."""
|
| 3 |
+
import asyncio
|
| 4 |
+
import os
|
| 5 |
+
|
| 6 |
+
import aiofiles
|
| 7 |
+
import aiofiles.ospath
|
| 8 |
+
import frontmatter
|
| 9 |
+
|
| 10 |
+
from ._base import SkillLoaderBase
|
| 11 |
+
from .._logging import logger
|
| 12 |
+
from ..skill import Skill
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
class LocalSkillLoader(SkillLoaderBase):
|
| 16 |
+
"""The skill loader that loads skills from a local directory."""
|
| 17 |
+
|
| 18 |
+
def __init__(self, directory: str, scan_subdir: bool = False) -> None:
|
| 19 |
+
"""Initialize the loader with the directory.
|
| 20 |
+
|
| 21 |
+
Args:
|
| 22 |
+
directory (`str`):
|
| 23 |
+
The directory to load skills from.
|
| 24 |
+
scan_subdir (`bool`, defaults to False):
|
| 25 |
+
Whether to scan subdirectories. Defaults to False (only
|
| 26 |
+
scan current directory).
|
| 27 |
+
"""
|
| 28 |
+
self.directory = os.path.abspath(directory)
|
| 29 |
+
self.scan_subdir = scan_subdir
|
| 30 |
+
self._cache: dict[str, Skill] = {}
|
| 31 |
+
|
| 32 |
+
async def _load_single_skill(self, skill_root: str) -> Skill | None:
|
| 33 |
+
"""Load a single skill from a skill root directory.
|
| 34 |
+
|
| 35 |
+
Args:
|
| 36 |
+
skill_root (`str`): The skill root directory containing SKILL.md.
|
| 37 |
+
|
| 38 |
+
Returns:
|
| 39 |
+
`Skill | None`: A Skill object or None if loading failed.
|
| 40 |
+
"""
|
| 41 |
+
skill_md_path = os.path.join(skill_root, "SKILL.md")
|
| 42 |
+
|
| 43 |
+
try:
|
| 44 |
+
# Check if SKILL.md exists
|
| 45 |
+
if not await aiofiles.ospath.isfile(skill_md_path):
|
| 46 |
+
return None
|
| 47 |
+
|
| 48 |
+
# Get file modification time
|
| 49 |
+
updated_at = await aiofiles.ospath.getmtime(skill_md_path)
|
| 50 |
+
|
| 51 |
+
# Check cache: if cached skill exists and updated_at matches,
|
| 52 |
+
# return cached
|
| 53 |
+
if skill_root in self._cache:
|
| 54 |
+
cached_skill = self._cache[skill_root]
|
| 55 |
+
if cached_skill.updated_at == updated_at:
|
| 56 |
+
return cached_skill
|
| 57 |
+
|
| 58 |
+
# Read and parse SKILL.md
|
| 59 |
+
async with aiofiles.open(
|
| 60 |
+
skill_md_path,
|
| 61 |
+
"r",
|
| 62 |
+
encoding="utf-8",
|
| 63 |
+
) as f:
|
| 64 |
+
content_str = await f.read()
|
| 65 |
+
content = frontmatter.loads(content_str)
|
| 66 |
+
|
| 67 |
+
name = content.get("name")
|
| 68 |
+
description = content.get("description")
|
| 69 |
+
|
| 70 |
+
if not name or not description:
|
| 71 |
+
logger.warning(
|
| 72 |
+
"SKILL.md in %s is missing required fields "
|
| 73 |
+
"(name or description). Skipping.",
|
| 74 |
+
skill_root,
|
| 75 |
+
)
|
| 76 |
+
return None
|
| 77 |
+
|
| 78 |
+
skill = Skill(
|
| 79 |
+
name=str(name),
|
| 80 |
+
description=str(description),
|
| 81 |
+
dir=skill_root,
|
| 82 |
+
markdown=content.content,
|
| 83 |
+
updated_at=updated_at,
|
| 84 |
+
)
|
| 85 |
+
|
| 86 |
+
# Update cache
|
| 87 |
+
self._cache[skill_root] = skill
|
| 88 |
+
|
| 89 |
+
return skill
|
| 90 |
+
|
| 91 |
+
except Exception as e:
|
| 92 |
+
logger.warning(
|
| 93 |
+
"Failed to load skill from %s: %s",
|
| 94 |
+
skill_root,
|
| 95 |
+
str(e),
|
| 96 |
+
)
|
| 97 |
+
return None
|
| 98 |
+
|
| 99 |
+
async def list_skills(self) -> list[Skill]:
|
| 100 |
+
"""List all the available skills from the directory.
|
| 101 |
+
|
| 102 |
+
This method will:
|
| 103 |
+
1. Search for SKILL.md in the current directory
|
| 104 |
+
2. If scan_subdir is True, search for SKILL.md in all subdirectories
|
| 105 |
+
3. Load all SKILL.md files concurrently
|
| 106 |
+
|
| 107 |
+
Returns:
|
| 108 |
+
`list[Skill]`: A list of Skill objects.
|
| 109 |
+
"""
|
| 110 |
+
try:
|
| 111 |
+
# Check if directory exists
|
| 112 |
+
if not await aiofiles.ospath.isdir(self.directory):
|
| 113 |
+
logger.warning(
|
| 114 |
+
"Skill directory %s does not exist.",
|
| 115 |
+
self.directory,
|
| 116 |
+
)
|
| 117 |
+
return []
|
| 118 |
+
|
| 119 |
+
# Find all directories containing SKILL.md
|
| 120 |
+
def _find_skill_dirs() -> list[str]:
|
| 121 |
+
"""Find all directories containing SKILL.md file."""
|
| 122 |
+
dirs = []
|
| 123 |
+
|
| 124 |
+
if os.path.isfile(os.path.join(self.directory, "SKILL.md")):
|
| 125 |
+
dirs.append(self.directory)
|
| 126 |
+
|
| 127 |
+
if self.scan_subdir:
|
| 128 |
+
for root, _, filenames in os.walk(self.directory):
|
| 129 |
+
if root == self.directory:
|
| 130 |
+
continue
|
| 131 |
+
if "SKILL.md" in filenames:
|
| 132 |
+
dirs.append(root)
|
| 133 |
+
|
| 134 |
+
return dirs
|
| 135 |
+
|
| 136 |
+
skill_dirs = await asyncio.to_thread(_find_skill_dirs)
|
| 137 |
+
|
| 138 |
+
if not skill_dirs:
|
| 139 |
+
logger.info(
|
| 140 |
+
"No SKILL.md files found in %s",
|
| 141 |
+
self.directory,
|
| 142 |
+
)
|
| 143 |
+
return []
|
| 144 |
+
|
| 145 |
+
# Load all skills concurrently
|
| 146 |
+
tasks = [
|
| 147 |
+
self._load_single_skill(skill_dir) for skill_dir in skill_dirs
|
| 148 |
+
]
|
| 149 |
+
results = await asyncio.gather(*tasks, return_exceptions=True)
|
| 150 |
+
|
| 151 |
+
# Filter out None results and exceptions
|
| 152 |
+
skills: list = []
|
| 153 |
+
for i, result in enumerate(results):
|
| 154 |
+
if isinstance(result, Exception):
|
| 155 |
+
logger.warning(
|
| 156 |
+
"Failed to load skill from %s: %s",
|
| 157 |
+
skill_dirs[i],
|
| 158 |
+
str(result),
|
| 159 |
+
)
|
| 160 |
+
elif result is not None:
|
| 161 |
+
skills.append(result)
|
| 162 |
+
|
| 163 |
+
return skills
|
| 164 |
+
|
| 165 |
+
except Exception as e:
|
| 166 |
+
logger.warning(
|
| 167 |
+
"Failed to list skills from directory %s: %s",
|
| 168 |
+
self.directory,
|
| 169 |
+
str(e),
|
| 170 |
+
)
|
| 171 |
+
return []
|
src/agentscope/state/__init__.py
ADDED
|
@@ -0,0 +1,11 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The agent state module in agentscope."""
|
| 3 |
+
|
| 4 |
+
from ._state import AgentState, TaskContext
|
| 5 |
+
from ._task import Task
|
| 6 |
+
|
| 7 |
+
__all__ = [
|
| 8 |
+
"Task",
|
| 9 |
+
"TaskContext",
|
| 10 |
+
"AgentState",
|
| 11 |
+
]
|
src/agentscope/state/_state.py
ADDED
|
@@ -0,0 +1,221 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The agent state class."""
|
| 3 |
+
from typing import Any
|
| 4 |
+
|
| 5 |
+
from pydantic import BaseModel, Field
|
| 6 |
+
|
| 7 |
+
import aiofiles.os
|
| 8 |
+
|
| 9 |
+
from .._utils._common import _generate_id
|
| 10 |
+
from ._task import Task
|
| 11 |
+
from ..message import (
|
| 12 |
+
TextBlock,
|
| 13 |
+
DataBlock,
|
| 14 |
+
Msg,
|
| 15 |
+
ToolCallBlock,
|
| 16 |
+
ToolResultBlock,
|
| 17 |
+
HintBlock,
|
| 18 |
+
)
|
| 19 |
+
from ..permission import PermissionContext
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
class ReadCacheEntry(BaseModel):
|
| 23 |
+
"""The read file cache."""
|
| 24 |
+
|
| 25 |
+
lines: list[str]
|
| 26 |
+
updated_at: float
|
| 27 |
+
bytes: float
|
| 28 |
+
file_path: str
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
class ToolContext(BaseModel):
|
| 32 |
+
"""The tool context, e.g. tool cache"""
|
| 33 |
+
|
| 34 |
+
max_cache_files: int = Field(default=100, gt=1)
|
| 35 |
+
"""The maximum number of cached files."""
|
| 36 |
+
max_cache_bytes: float = Field(default=25000, gt=10000)
|
| 37 |
+
"""The maximum size of the accumulated read file cache."""
|
| 38 |
+
read_file_cache: list[ReadCacheEntry] = Field(default_factory=list)
|
| 39 |
+
"""The cache for Read/Write/Edit file tools."""
|
| 40 |
+
|
| 41 |
+
activated_groups: list[str] = Field(default_factory=list)
|
| 42 |
+
"""The names of the activated tool groups, each group contains a set of
|
| 43 |
+
tools."""
|
| 44 |
+
|
| 45 |
+
async def get_cache(self, file_path: str) -> ReadCacheEntry | None:
|
| 46 |
+
"""Get cached file content if still valid.
|
| 47 |
+
|
| 48 |
+
Args:
|
| 49 |
+
file_path: The absolute path of the file.
|
| 50 |
+
|
| 51 |
+
Returns:
|
| 52 |
+
The cached entry if valid, otherwise None.
|
| 53 |
+
"""
|
| 54 |
+
|
| 55 |
+
# Find the cache entry
|
| 56 |
+
for entry in self.read_file_cache:
|
| 57 |
+
if entry.file_path == file_path:
|
| 58 |
+
# Check if cache is still valid
|
| 59 |
+
try:
|
| 60 |
+
updated_at = await aiofiles.os.path.getmtime(file_path)
|
| 61 |
+
if updated_at == entry.updated_at:
|
| 62 |
+
return entry
|
| 63 |
+
else:
|
| 64 |
+
# Cache is outdated, remove it
|
| 65 |
+
self.read_file_cache.remove(entry)
|
| 66 |
+
return None
|
| 67 |
+
except Exception:
|
| 68 |
+
# File might not exist anymore
|
| 69 |
+
self.read_file_cache.remove(entry)
|
| 70 |
+
return None
|
| 71 |
+
return None
|
| 72 |
+
|
| 73 |
+
async def cache_file(self, file_path: str, lines: list[str]) -> None:
|
| 74 |
+
"""Cache file content with LRU eviction.
|
| 75 |
+
|
| 76 |
+
Args:
|
| 77 |
+
file_path: The absolute path of the file.
|
| 78 |
+
lines: The lines of the file content.
|
| 79 |
+
"""
|
| 80 |
+
try:
|
| 81 |
+
updated_at = await aiofiles.os.path.getmtime(file_path)
|
| 82 |
+
except Exception:
|
| 83 |
+
# Cannot get mtime, skip caching
|
| 84 |
+
return
|
| 85 |
+
|
| 86 |
+
# Calculate size in KB
|
| 87 |
+
new_entry_bytes = (
|
| 88 |
+
sum(len(line.encode("utf-8")) for line in lines) / 1024
|
| 89 |
+
)
|
| 90 |
+
|
| 91 |
+
# Remove existing cache for this file if present
|
| 92 |
+
self.read_file_cache = [
|
| 93 |
+
entry
|
| 94 |
+
for entry in self.read_file_cache
|
| 95 |
+
if entry.file_path != file_path
|
| 96 |
+
]
|
| 97 |
+
|
| 98 |
+
# Evict the oldest entries if exceeding max_cache_files
|
| 99 |
+
while len(self.read_file_cache) >= self.max_cache_files:
|
| 100 |
+
self.read_file_cache.pop(0)
|
| 101 |
+
|
| 102 |
+
# Evict the oldest entries if exceeding max_cache_bytes
|
| 103 |
+
current_size = sum(entry.bytes for entry in self.read_file_cache)
|
| 104 |
+
while (
|
| 105 |
+
self.read_file_cache
|
| 106 |
+
and current_size + new_entry_bytes > self.max_cache_bytes
|
| 107 |
+
):
|
| 108 |
+
removed = self.read_file_cache.pop(0)
|
| 109 |
+
current_size -= removed.bytes
|
| 110 |
+
|
| 111 |
+
# Add new entry to the end (most recent)
|
| 112 |
+
self.read_file_cache.append(
|
| 113 |
+
ReadCacheEntry(
|
| 114 |
+
lines=lines,
|
| 115 |
+
updated_at=updated_at,
|
| 116 |
+
bytes=new_entry_bytes,
|
| 117 |
+
file_path=file_path,
|
| 118 |
+
),
|
| 119 |
+
)
|
| 120 |
+
|
| 121 |
+
async def clean_file_cache(
|
| 122 |
+
self,
|
| 123 |
+
reserved_file_paths: set[str] | None = None,
|
| 124 |
+
) -> None:
|
| 125 |
+
"""Drop read caches whose paths are not in ``reserved_file_paths``.
|
| 126 |
+
|
| 127 |
+
Args:
|
| 128 |
+
reserved_file_paths: File paths from Read calls that remain in the
|
| 129 |
+
context. Caches for these files are kept; all others are
|
| 130 |
+
evicted.
|
| 131 |
+
"""
|
| 132 |
+
reserved_file_paths = reserved_file_paths or set()
|
| 133 |
+
|
| 134 |
+
self.read_file_cache = [
|
| 135 |
+
entry
|
| 136 |
+
for entry in self.read_file_cache
|
| 137 |
+
if entry.file_path in reserved_file_paths
|
| 138 |
+
]
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
class TaskContext(BaseModel):
|
| 142 |
+
"""The task context."""
|
| 143 |
+
|
| 144 |
+
tasks: list[Task] = Field(default_factory=lambda: [])
|
| 145 |
+
"""The task context."""
|
| 146 |
+
|
| 147 |
+
|
| 148 |
+
class AgentState(BaseModel):
|
| 149 |
+
"""The agent state that should be saved and loaded from storage."""
|
| 150 |
+
|
| 151 |
+
session_id: str = Field(default_factory=_generate_id)
|
| 152 |
+
"""The session id of the agent. Normally, each session will maintain one
|
| 153 |
+
independent agent state for each agent."""
|
| 154 |
+
|
| 155 |
+
summary: str | list[TextBlock | DataBlock] = ""
|
| 156 |
+
"""The compressed summary of the context, which will be prepended to the
|
| 157 |
+
context when feed into the LLM."""
|
| 158 |
+
context: list[Msg] = Field(default_factory=list)
|
| 159 |
+
"""The uncompressed conversation context, that will be feed into the LLM"""
|
| 160 |
+
reply_id: str = Field(default_factory=_generate_id)
|
| 161 |
+
"""The id of the current reply, which is also used as the id of the
|
| 162 |
+
final message of the reply."""
|
| 163 |
+
cur_iter: int = 0
|
| 164 |
+
"""The current iteration of the agent's reasoning-acting loop."""
|
| 165 |
+
|
| 166 |
+
# =================================================================
|
| 167 |
+
# The permission context
|
| 168 |
+
# =================================================================
|
| 169 |
+
permission_context: PermissionContext = Field(
|
| 170 |
+
default_factory=PermissionContext,
|
| 171 |
+
)
|
| 172 |
+
"""The permission context that will be passed to the toolkit to determine
|
| 173 |
+
the tool permissions."""
|
| 174 |
+
|
| 175 |
+
# =================================================================
|
| 176 |
+
# The tool context
|
| 177 |
+
# =================================================================
|
| 178 |
+
tool_context: ToolContext = Field(default_factory=ToolContext)
|
| 179 |
+
|
| 180 |
+
# =================================================================
|
| 181 |
+
# The tasks context
|
| 182 |
+
# =================================================================
|
| 183 |
+
tasks_context: TaskContext = Field(default_factory=TaskContext)
|
| 184 |
+
"""The task context that records the agent tasks."""
|
| 185 |
+
|
| 186 |
+
# =================================================================
|
| 187 |
+
# The middleware context
|
| 188 |
+
# =================================================================
|
| 189 |
+
middle_context: dict[str, Any] = Field(default_factory=dict)
|
| 190 |
+
"""The context that allow the middlewares to store/get data across
|
| 191 |
+
different replies."""
|
| 192 |
+
|
| 193 |
+
def append_context(
|
| 194 |
+
self,
|
| 195 |
+
name: str,
|
| 196 |
+
blocks: list[
|
| 197 |
+
TextBlock | DataBlock | HintBlock | ToolCallBlock | ToolResultBlock
|
| 198 |
+
],
|
| 199 |
+
) -> None:
|
| 200 |
+
"""Append the given blocks to the agent's own message with the current
|
| 201 |
+
`reply_id`. If such message doesn't exist, a new assistant message
|
| 202 |
+
with agent's name and current reply ID will be created.
|
| 203 |
+
"""
|
| 204 |
+
# If append to the latest message
|
| 205 |
+
if (
|
| 206 |
+
self.context
|
| 207 |
+
and self.context[-1].role == "assistant"
|
| 208 |
+
and self.context[-1].name == name
|
| 209 |
+
and self.context[-1].id == self.reply_id
|
| 210 |
+
):
|
| 211 |
+
self.context[-1].content.extend(blocks)
|
| 212 |
+
else:
|
| 213 |
+
# Create a new assistant message with the current reply ID
|
| 214 |
+
self.context.append(
|
| 215 |
+
Msg(
|
| 216 |
+
id=self.reply_id,
|
| 217 |
+
role="assistant",
|
| 218 |
+
name=name,
|
| 219 |
+
content=blocks,
|
| 220 |
+
),
|
| 221 |
+
)
|
src/agentscope/state/_task.py
ADDED
|
@@ -0,0 +1,39 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The task class."""
|
| 3 |
+
from datetime import datetime
|
| 4 |
+
from typing import Any, Literal
|
| 5 |
+
|
| 6 |
+
from pydantic import BaseModel, Field
|
| 7 |
+
|
| 8 |
+
from .._utils._common import _generate_id
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
class Task(BaseModel):
|
| 12 |
+
"""The agent task."""
|
| 13 |
+
|
| 14 |
+
subject: str
|
| 15 |
+
"""The subject of the task."""
|
| 16 |
+
|
| 17 |
+
description: str
|
| 18 |
+
"""The task description."""
|
| 19 |
+
|
| 20 |
+
metadata: dict[str, Any]
|
| 21 |
+
"""The additional metadata of the task."""
|
| 22 |
+
|
| 23 |
+
created_at: str = Field(default_factory=lambda: datetime.now().isoformat())
|
| 24 |
+
"""The created timestamp."""
|
| 25 |
+
|
| 26 |
+
state: Literal["pending", "in_progress", "completed"] = "pending"
|
| 27 |
+
"""The task state."""
|
| 28 |
+
|
| 29 |
+
id: str = Field(default_factory=_generate_id)
|
| 30 |
+
"""The task identifier."""
|
| 31 |
+
|
| 32 |
+
owner: str | None = None
|
| 33 |
+
"""The owner of the task."""
|
| 34 |
+
|
| 35 |
+
blocks: list[str] = Field(default_factory=lambda: [])
|
| 36 |
+
"""The task ids blocked by this task."""
|
| 37 |
+
|
| 38 |
+
blocked_by: list[str] = Field(default_factory=lambda: [])
|
| 39 |
+
"""The task ids blocking this task."""
|
src/agentscope/tool/__init__.py
ADDED
|
@@ -0,0 +1,58 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The tool module in agentscope."""
|
| 3 |
+
|
| 4 |
+
from ._types import ToolChoice, Function, RegisteredTool
|
| 5 |
+
from ._response import ToolResponse, ToolChunk
|
| 6 |
+
from ._toolkit import Toolkit
|
| 7 |
+
from ._base import ToolBase, ParamsBase, ToolMiddlewareBase
|
| 8 |
+
from ._adapters import MCPTool, FunctionTool
|
| 9 |
+
from ._builtin import (
|
| 10 |
+
ResetTools,
|
| 11 |
+
Bash,
|
| 12 |
+
Edit,
|
| 13 |
+
Glob,
|
| 14 |
+
Grep,
|
| 15 |
+
Read,
|
| 16 |
+
Write,
|
| 17 |
+
BackendBase,
|
| 18 |
+
ExecResult,
|
| 19 |
+
LocalBackend,
|
| 20 |
+
)
|
| 21 |
+
from ._task import (
|
| 22 |
+
TaskUpdate,
|
| 23 |
+
TaskGet,
|
| 24 |
+
TaskList,
|
| 25 |
+
TaskCreate,
|
| 26 |
+
)
|
| 27 |
+
from ._tool_group import ToolGroup
|
| 28 |
+
|
| 29 |
+
__all__ = [
|
| 30 |
+
# Basic tool related types and functions
|
| 31 |
+
"ToolChoice",
|
| 32 |
+
"Function",
|
| 33 |
+
"ToolBase",
|
| 34 |
+
"ParamsBase",
|
| 35 |
+
"ToolMiddlewareBase",
|
| 36 |
+
"MCPTool",
|
| 37 |
+
"FunctionTool",
|
| 38 |
+
"ToolGroup",
|
| 39 |
+
"Toolkit",
|
| 40 |
+
"ToolChunk",
|
| 41 |
+
"ToolResponse",
|
| 42 |
+
"RegisteredTool",
|
| 43 |
+
# Builtin tools
|
| 44 |
+
"BackendBase",
|
| 45 |
+
"LocalBackend",
|
| 46 |
+
"ExecResult",
|
| 47 |
+
"ResetTools",
|
| 48 |
+
"Bash",
|
| 49 |
+
"Edit",
|
| 50 |
+
"Glob",
|
| 51 |
+
"Grep",
|
| 52 |
+
"Read",
|
| 53 |
+
"Write",
|
| 54 |
+
"TaskUpdate",
|
| 55 |
+
"TaskGet",
|
| 56 |
+
"TaskList",
|
| 57 |
+
"TaskCreate",
|
| 58 |
+
]
|
src/agentscope/tool/_adapters.py
ADDED
|
@@ -0,0 +1,394 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""Adapters to convert functions and MCP tools to ToolProtocol."""
|
| 3 |
+
import inspect
|
| 4 |
+
import json
|
| 5 |
+
import re
|
| 6 |
+
from contextlib import _AsyncGeneratorContextManager
|
| 7 |
+
from datetime import timedelta
|
| 8 |
+
from typing import Callable, Any, AsyncGenerator, Generator
|
| 9 |
+
|
| 10 |
+
from mcp import ClientSession
|
| 11 |
+
import mcp
|
| 12 |
+
|
| 13 |
+
from ._types import Function
|
| 14 |
+
from ._base import ToolBase, ToolMiddlewareBase
|
| 15 |
+
from ..permission import (
|
| 16 |
+
PermissionBehavior,
|
| 17 |
+
PermissionDecision,
|
| 18 |
+
)
|
| 19 |
+
from ._response import ToolChunk
|
| 20 |
+
from ._utils import _extract_func_description, _extract_input_schema
|
| 21 |
+
from .._logging import logger
|
| 22 |
+
from ..message import (
|
| 23 |
+
TextBlock,
|
| 24 |
+
DataBlock,
|
| 25 |
+
Base64Source,
|
| 26 |
+
URLSource,
|
| 27 |
+
ToolResultState,
|
| 28 |
+
)
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
class FunctionTool(ToolBase):
|
| 32 |
+
"""Adapter to convert a Python function to ToolProtocol.
|
| 33 |
+
|
| 34 |
+
This class wraps a regular Python function and makes it compatible with
|
| 35 |
+
the ToolProtocol interface. It automatically extracts metadata from the
|
| 36 |
+
function's signature and docstring, and normalizes the return value to
|
| 37 |
+
ToolChunk or AsyncGenerator[ToolChunk, None].
|
| 38 |
+
"""
|
| 39 |
+
|
| 40 |
+
is_external_tool: bool = False
|
| 41 |
+
"""If this tool is an external tool, which doesn't need to implement the
|
| 42 |
+
__call__ method and the agent will yield the external tool call event."""
|
| 43 |
+
is_mcp: bool = False
|
| 44 |
+
"""If this tool is an MCP tool, which will be used in the permission"""
|
| 45 |
+
mcp_name: str | None = None
|
| 46 |
+
"""The name of the MCP server this tool belongs to, which is required if
|
| 47 |
+
this tool is an MCP tool."""
|
| 48 |
+
|
| 49 |
+
def __init__(
|
| 50 |
+
self,
|
| 51 |
+
func: Function,
|
| 52 |
+
name: str | None = None,
|
| 53 |
+
description: str | None = None,
|
| 54 |
+
is_concurrency_safe: bool = True,
|
| 55 |
+
is_read_only: bool = False,
|
| 56 |
+
is_state_injected: bool = False,
|
| 57 |
+
middlewares: list[ToolMiddlewareBase] | None = None,
|
| 58 |
+
) -> None:
|
| 59 |
+
"""Initialize the FunctionTool.
|
| 60 |
+
|
| 61 |
+
Args:
|
| 62 |
+
func (`Callable`):
|
| 63 |
+
The Python function to wrap.
|
| 64 |
+
name (`str | None`, optional):
|
| 65 |
+
Custom tool name. If None, uses the function name.
|
| 66 |
+
description (`str | None`, optional):
|
| 67 |
+
Custom tool description. If None, extracts from docstring.
|
| 68 |
+
is_concurrency_safe (`bool`, optional):
|
| 69 |
+
Whether this tool is safe to call concurrently.
|
| 70 |
+
is_read_only (`bool`, optional):
|
| 71 |
+
Whether this tool only reads data without side effects.
|
| 72 |
+
is_state_injected (`bool`, optional):
|
| 73 |
+
Whether this tool requires agent state injection.
|
| 74 |
+
middlewares (`list[ToolMiddlewareBase] | None`, optional):
|
| 75 |
+
Tool middlewares wrapping the tool execution.
|
| 76 |
+
"""
|
| 77 |
+
super().__init__(middlewares=middlewares)
|
| 78 |
+
self.name = name or func.__name__
|
| 79 |
+
self.description = description or _extract_func_description(
|
| 80 |
+
func.__doc__ or "",
|
| 81 |
+
)
|
| 82 |
+
self.input_schema = _extract_input_schema(func)
|
| 83 |
+
self.is_concurrency_safe = is_concurrency_safe
|
| 84 |
+
self.is_read_only = is_read_only
|
| 85 |
+
self.is_state_injected = is_state_injected
|
| 86 |
+
self.is_external_tool = False
|
| 87 |
+
self.is_mcp = False
|
| 88 |
+
self._func = func
|
| 89 |
+
|
| 90 |
+
async def check_permissions(
|
| 91 |
+
self,
|
| 92 |
+
*_args: Any,
|
| 93 |
+
**_kwargs: Any,
|
| 94 |
+
) -> PermissionDecision:
|
| 95 |
+
"""Check permissions for the tool usage.
|
| 96 |
+
|
| 97 |
+
Default implementation allows all operations.
|
| 98 |
+
|
| 99 |
+
Returns:
|
| 100 |
+
`PermissionDecision`:
|
| 101 |
+
Permission decision (default: allowed).
|
| 102 |
+
"""
|
| 103 |
+
return PermissionDecision(
|
| 104 |
+
behavior=PermissionBehavior.ASK,
|
| 105 |
+
message="Custom function tools must be explicitly allowed "
|
| 106 |
+
"by the user.",
|
| 107 |
+
)
|
| 108 |
+
|
| 109 |
+
async def call(
|
| 110 |
+
self,
|
| 111 |
+
**kwargs: Any,
|
| 112 |
+
) -> ToolChunk | AsyncGenerator[ToolChunk, None]:
|
| 113 |
+
"""Invoke the wrapped function in an async style.
|
| 114 |
+
|
| 115 |
+
Returns:
|
| 116 |
+
`ToolChunk` or `AsyncGenerator[ToolChunk, None]`:
|
| 117 |
+
The normalized result of the function execution.
|
| 118 |
+
"""
|
| 119 |
+
if inspect.iscoroutinefunction(self._func):
|
| 120 |
+
result = await self._func(**kwargs)
|
| 121 |
+
else:
|
| 122 |
+
result = self._func(**kwargs)
|
| 123 |
+
|
| 124 |
+
if isinstance(result, AsyncGenerator):
|
| 125 |
+
|
| 126 |
+
async def _stream() -> AsyncGenerator[ToolChunk, None]:
|
| 127 |
+
async for chunk in result:
|
| 128 |
+
if isinstance(chunk, ToolChunk):
|
| 129 |
+
yield chunk
|
| 130 |
+
else:
|
| 131 |
+
yield self._convert_func_result_to_chunk(chunk)
|
| 132 |
+
|
| 133 |
+
return _stream()
|
| 134 |
+
|
| 135 |
+
if isinstance(result, Generator):
|
| 136 |
+
|
| 137 |
+
async def _stream() -> AsyncGenerator[ToolChunk, None]:
|
| 138 |
+
for chunk in result:
|
| 139 |
+
if isinstance(chunk, ToolChunk):
|
| 140 |
+
yield chunk
|
| 141 |
+
else:
|
| 142 |
+
yield self._convert_func_result_to_chunk(chunk)
|
| 143 |
+
|
| 144 |
+
return _stream()
|
| 145 |
+
|
| 146 |
+
return self._convert_func_result_to_chunk(result)
|
| 147 |
+
|
| 148 |
+
@staticmethod
|
| 149 |
+
def _convert_func_result_to_chunk(
|
| 150 |
+
result: Any,
|
| 151 |
+
) -> ToolChunk:
|
| 152 |
+
if isinstance(result, ToolChunk):
|
| 153 |
+
return result
|
| 154 |
+
if isinstance(result, str):
|
| 155 |
+
text = result
|
| 156 |
+
else:
|
| 157 |
+
try:
|
| 158 |
+
text = json.dumps(result, ensure_ascii=False)
|
| 159 |
+
except (TypeError, ValueError):
|
| 160 |
+
text = str(result)
|
| 161 |
+
return ToolChunk(
|
| 162 |
+
content=[TextBlock(text=text)],
|
| 163 |
+
state=ToolResultState.RUNNING,
|
| 164 |
+
)
|
| 165 |
+
|
| 166 |
+
|
| 167 |
+
class MCPTool(ToolBase):
|
| 168 |
+
"""Adapter to convert an MCP tool to ToolProtocol.
|
| 169 |
+
|
| 170 |
+
This class wraps an MCP tool and makes it compatible with the ToolProtocol
|
| 171 |
+
interface. It handles the conversion between MCP's result format and
|
| 172 |
+
AgentScope's ToolChunk format.
|
| 173 |
+
"""
|
| 174 |
+
|
| 175 |
+
is_mcp: bool = True
|
| 176 |
+
"""Whether this tool is an MCP tool."""
|
| 177 |
+
is_state_injected: bool = False
|
| 178 |
+
"""The mcp tools is prohibited state injection for safety reason."""
|
| 179 |
+
|
| 180 |
+
def __init__(
|
| 181 |
+
self,
|
| 182 |
+
mcp_name: str,
|
| 183 |
+
tool: mcp.types.Tool,
|
| 184 |
+
client_gen: Callable[..., _AsyncGeneratorContextManager[Any]]
|
| 185 |
+
| None = None,
|
| 186 |
+
session: Any | None = None,
|
| 187 |
+
timeout: float | None = None,
|
| 188 |
+
middlewares: list[ToolMiddlewareBase] | None = None,
|
| 189 |
+
) -> None:
|
| 190 |
+
"""Initialize the MCPTool.
|
| 191 |
+
|
| 192 |
+
Args:
|
| 193 |
+
mcp_name (`str`):
|
| 194 |
+
The name of the MCP server instance.
|
| 195 |
+
tool (`mcp.types.Tool`):
|
| 196 |
+
The MCP tool definition.
|
| 197 |
+
client_gen (`Callable[..., _AsyncGeneratorContextManager[Any]] \
|
| 198 |
+
| None`, optional):
|
| 199 |
+
The MCP client generator function for stateless clients.
|
| 200 |
+
Either this or ``session`` must be provided.
|
| 201 |
+
session (`mcp.ClientSession | None`, optional):
|
| 202 |
+
The MCP client session for stateful clients.
|
| 203 |
+
Either this or ``client_gen`` must be provided.
|
| 204 |
+
timeout (`float | None`, optional):
|
| 205 |
+
The timeout in seconds for tool execution.
|
| 206 |
+
middlewares (`list[ToolMiddlewareBase] | None`, optional):
|
| 207 |
+
Tool middlewares wrapping the tool execution.
|
| 208 |
+
"""
|
| 209 |
+
super().__init__(middlewares=middlewares)
|
| 210 |
+
self.mcp_name = mcp_name
|
| 211 |
+
|
| 212 |
+
# LLM providers enforce ^[a-zA-Z0-9_-]+$ on tool names.
|
| 213 |
+
# mcp_name is validated in MCPClient.model_post_init;
|
| 214 |
+
# tool.name comes from the MCP server and may contain dots,
|
| 215 |
+
# colons, etc. — replace illegal chars with "x" (not "_")
|
| 216 |
+
# to avoid collisions with the "__" separator.
|
| 217 |
+
# self._tool.name retains the original for server-side calls.
|
| 218 |
+
sanitized_tool = re.sub(r"[^a-zA-Z0-9_-]", "x", tool.name)
|
| 219 |
+
self.name = f"mcp__{mcp_name}__{sanitized_tool}"
|
| 220 |
+
if sanitized_tool != tool.name:
|
| 221 |
+
logger.debug(
|
| 222 |
+
"MCP tool name sanitized: '%s' -> '%s'.",
|
| 223 |
+
tool.name,
|
| 224 |
+
self.name,
|
| 225 |
+
)
|
| 226 |
+
|
| 227 |
+
self.description = tool.description or ""
|
| 228 |
+
|
| 229 |
+
# Preserve the full inputSchema (including $defs, anyOf, oneOf, etc.)
|
| 230 |
+
# rather than only copying "properties" and "required", which would
|
| 231 |
+
# silently drop any nested type definitions that the LLM needs to
|
| 232 |
+
# resolve $ref pointers.
|
| 233 |
+
_schema = dict(tool.inputSchema) if tool.inputSchema else {}
|
| 234 |
+
_schema.setdefault("type", "object")
|
| 235 |
+
_schema.setdefault("properties", {})
|
| 236 |
+
_schema.setdefault("required", [])
|
| 237 |
+
self.input_schema = _schema
|
| 238 |
+
|
| 239 |
+
# By default
|
| 240 |
+
self.is_concurrency_safe = False
|
| 241 |
+
self.is_external_tool = False
|
| 242 |
+
|
| 243 |
+
# Extract is_read_only from MCP tool annotations
|
| 244 |
+
self.is_read_only = False
|
| 245 |
+
if tool.annotations and hasattr(tool.annotations, "readOnlyHint"):
|
| 246 |
+
self.is_read_only = tool.annotations.readOnlyHint or False
|
| 247 |
+
|
| 248 |
+
# Store MCP tool and connection info
|
| 249 |
+
self._tool = tool
|
| 250 |
+
self._client_gen = client_gen
|
| 251 |
+
self._session = session
|
| 252 |
+
|
| 253 |
+
if timeout:
|
| 254 |
+
self._timeout = timedelta(seconds=timeout)
|
| 255 |
+
else:
|
| 256 |
+
self._timeout = None
|
| 257 |
+
|
| 258 |
+
# Validate that either client_gen or session is provided
|
| 259 |
+
if (client_gen is None and session is None) or (
|
| 260 |
+
client_gen is not None and session is not None
|
| 261 |
+
):
|
| 262 |
+
raise ValueError(
|
| 263 |
+
"Either client_gen or session must be provided, but not both.",
|
| 264 |
+
)
|
| 265 |
+
|
| 266 |
+
async def check_permissions(
|
| 267 |
+
self,
|
| 268 |
+
*_args: Any,
|
| 269 |
+
**_kwargs: Any,
|
| 270 |
+
) -> PermissionDecision:
|
| 271 |
+
"""Check permissions for the MCP tool usage.
|
| 272 |
+
|
| 273 |
+
Default implementation allows all operations.
|
| 274 |
+
|
| 275 |
+
Returns:
|
| 276 |
+
`PermissionDecision`:
|
| 277 |
+
Permission decision (default: ask for confirmation).
|
| 278 |
+
"""
|
| 279 |
+
if self.is_read_only:
|
| 280 |
+
return PermissionDecision(
|
| 281 |
+
behavior=PermissionBehavior.ALLOW,
|
| 282 |
+
message="This is a read-only MCP tool. Allowing execution.",
|
| 283 |
+
)
|
| 284 |
+
return PermissionDecision(
|
| 285 |
+
behavior=PermissionBehavior.ASK,
|
| 286 |
+
message="MCP tools must be explicitly allowed by the user.",
|
| 287 |
+
)
|
| 288 |
+
|
| 289 |
+
async def call(
|
| 290 |
+
self,
|
| 291 |
+
**kwargs: Any,
|
| 292 |
+
) -> ToolChunk:
|
| 293 |
+
"""Invoke the MCP tool and convert the result to ToolChunk.
|
| 294 |
+
|
| 295 |
+
Args:
|
| 296 |
+
**kwargs: Arguments to pass to the MCP tool.
|
| 297 |
+
|
| 298 |
+
Returns:
|
| 299 |
+
`ToolChunk`: The converted tool execution result.
|
| 300 |
+
"""
|
| 301 |
+
|
| 302 |
+
# Call the MCP tool
|
| 303 |
+
if self._client_gen:
|
| 304 |
+
# Stateless client: create temporary session
|
| 305 |
+
async with self._client_gen() as cli:
|
| 306 |
+
read_stream, write_stream = cli[0], cli[1]
|
| 307 |
+
async with ClientSession(read_stream, write_stream) as session:
|
| 308 |
+
await session.initialize()
|
| 309 |
+
result = await session.call_tool(
|
| 310 |
+
self._tool.name,
|
| 311 |
+
arguments=kwargs,
|
| 312 |
+
read_timeout_seconds=self._timeout,
|
| 313 |
+
)
|
| 314 |
+
else:
|
| 315 |
+
# Stateful client: use existing session
|
| 316 |
+
result = await self._session.call_tool(
|
| 317 |
+
self._tool.name,
|
| 318 |
+
arguments=kwargs,
|
| 319 |
+
read_timeout_seconds=self._timeout,
|
| 320 |
+
)
|
| 321 |
+
|
| 322 |
+
# Convert MCP result to AgentScope blocks
|
| 323 |
+
return ToolChunk(
|
| 324 |
+
content=self._convert_mcp_content_to_blocks(result.content),
|
| 325 |
+
state=ToolResultState.ERROR
|
| 326 |
+
if result.isError
|
| 327 |
+
else ToolResultState.RUNNING,
|
| 328 |
+
)
|
| 329 |
+
|
| 330 |
+
@staticmethod
|
| 331 |
+
def _convert_mcp_content_to_blocks(
|
| 332 |
+
mcp_content_blocks: list,
|
| 333 |
+
) -> list[TextBlock | DataBlock]:
|
| 334 |
+
"""Convert MCP content to AgentScope blocks.
|
| 335 |
+
|
| 336 |
+
Args:
|
| 337 |
+
mcp_content_blocks (`list`):
|
| 338 |
+
The MCP content blocks to convert.
|
| 339 |
+
|
| 340 |
+
Returns:
|
| 341 |
+
`list[TextBlock | DataBlock]`: Converted AgentScope blocks.
|
| 342 |
+
"""
|
| 343 |
+
|
| 344 |
+
as_content = []
|
| 345 |
+
for content in mcp_content_blocks:
|
| 346 |
+
if isinstance(content, mcp.types.TextContent):
|
| 347 |
+
as_content.append(TextBlock(text=content.text))
|
| 348 |
+
elif isinstance(
|
| 349 |
+
content,
|
| 350 |
+
(mcp.types.ImageContent, mcp.types.AudioContent),
|
| 351 |
+
):
|
| 352 |
+
as_content.append(
|
| 353 |
+
DataBlock(
|
| 354 |
+
source=Base64Source(
|
| 355 |
+
type="base64",
|
| 356 |
+
media_type=content.mimeType,
|
| 357 |
+
data=content.data,
|
| 358 |
+
),
|
| 359 |
+
),
|
| 360 |
+
)
|
| 361 |
+
|
| 362 |
+
elif isinstance(content, mcp.types.EmbeddedResource):
|
| 363 |
+
if isinstance(
|
| 364 |
+
content.resource,
|
| 365 |
+
mcp.types.TextResourceContents,
|
| 366 |
+
):
|
| 367 |
+
as_content.append(
|
| 368 |
+
TextBlock(
|
| 369 |
+
text=content.resource.model_dump_json(indent=2),
|
| 370 |
+
),
|
| 371 |
+
)
|
| 372 |
+
else:
|
| 373 |
+
logger.error(
|
| 374 |
+
"Unsupported EmbeddedResource content type: %s. "
|
| 375 |
+
"Skipping this content.",
|
| 376 |
+
type(content.resource),
|
| 377 |
+
)
|
| 378 |
+
|
| 379 |
+
elif isinstance(content, mcp.types.ResourceContents):
|
| 380 |
+
as_content.append(
|
| 381 |
+
DataBlock(
|
| 382 |
+
source=URLSource(
|
| 383 |
+
media_type=content.mimeType,
|
| 384 |
+
url=content.uri,
|
| 385 |
+
),
|
| 386 |
+
),
|
| 387 |
+
)
|
| 388 |
+
|
| 389 |
+
else:
|
| 390 |
+
logger.warning(
|
| 391 |
+
"Unsupported content type: %s. Skipping this content.",
|
| 392 |
+
type(content),
|
| 393 |
+
)
|
| 394 |
+
return as_content
|
src/agentscope/tool/_base.py
ADDED
|
@@ -0,0 +1,451 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
# pylint: disable=unused-argument
|
| 3 |
+
"""The tool protocol in agentscope."""
|
| 4 |
+
import inspect
|
| 5 |
+
import os
|
| 6 |
+
from abc import abstractmethod, ABC
|
| 7 |
+
from pathlib import Path
|
| 8 |
+
from typing import AsyncGenerator, Any, Callable, List
|
| 9 |
+
|
| 10 |
+
from pydantic import BaseModel
|
| 11 |
+
|
| 12 |
+
from ._constants import DEFAULT_DANGEROUS_FILES, DEFAULT_DANGEROUS_DIRECTORIES
|
| 13 |
+
from ..permission import (
|
| 14 |
+
PermissionContext,
|
| 15 |
+
PermissionDecision,
|
| 16 |
+
PermissionRule,
|
| 17 |
+
PermissionBehavior,
|
| 18 |
+
)
|
| 19 |
+
from ._response import ToolChunk
|
| 20 |
+
from ._utils import _remove_title_field
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
class ParamsBase(BaseModel):
|
| 24 |
+
"""A base class for tool parameters that remove the title field from the
|
| 25 |
+
exported JSON schema.
|
| 26 |
+
"""
|
| 27 |
+
|
| 28 |
+
@classmethod
|
| 29 |
+
def model_json_schema(cls, *args: Any, **kwargs: Any) -> dict:
|
| 30 |
+
"""An override implementation to remove the title field from the
|
| 31 |
+
exported schema.
|
| 32 |
+
"""
|
| 33 |
+
return _remove_title_field(super().model_json_schema(*args, **kwargs))
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
class ToolMiddlewareBase(ABC):
|
| 37 |
+
"""Base class for tool middlewares.
|
| 38 |
+
|
| 39 |
+
A tool middleware wraps the execution of a tool in an onion fashion: the
|
| 40 |
+
first registered middleware is the outermost layer and runs its pre-logic
|
| 41 |
+
before any inner layer, then its post-logic after all inner layers have
|
| 42 |
+
completed. Subclass this and implement :meth:`on_tool_call` — the signature
|
| 43 |
+
is already spelled out, so second-party developers only need to fill in the
|
| 44 |
+
body without reasoning about the wrapping protocol.
|
| 45 |
+
|
| 46 |
+
Streaming and non-streaming tools are unified: ``next_handler`` always
|
| 47 |
+
returns an async generator, so a middleware never needs to know whether the
|
| 48 |
+
underlying tool yields a stream of chunks or returns a single chunk.
|
| 49 |
+
|
| 50 |
+
Example:
|
| 51 |
+
```python
|
| 52 |
+
class LoggingMiddleware(ToolMiddlewareBase):
|
| 53 |
+
async def on_tool_call(self, tool, input_kwargs, next_handler):
|
| 54 |
+
print(f"Calling {tool.name} with {input_kwargs}")
|
| 55 |
+
async for chunk in next_handler(**input_kwargs):
|
| 56 |
+
yield chunk
|
| 57 |
+
print(f"Finished {tool.name}")
|
| 58 |
+
|
| 59 |
+
tool = MyTool(middlewares=[LoggingMiddleware()])
|
| 60 |
+
```
|
| 61 |
+
"""
|
| 62 |
+
|
| 63 |
+
@abstractmethod
|
| 64 |
+
async def on_tool_call(
|
| 65 |
+
self,
|
| 66 |
+
tool: "ToolBase",
|
| 67 |
+
input_kwargs: dict[str, Any],
|
| 68 |
+
next_handler: Callable[..., AsyncGenerator[ToolChunk, None]],
|
| 69 |
+
) -> AsyncGenerator[ToolChunk, None]:
|
| 70 |
+
"""Intercept a single tool invocation.
|
| 71 |
+
|
| 72 |
+
Add pre-/post-logic around ``next_handler``, rewrite the tool inputs by
|
| 73 |
+
passing modified keyword arguments to ``next_handler``, or transform
|
| 74 |
+
the yielded chunks.
|
| 75 |
+
|
| 76 |
+
Args:
|
| 77 |
+
tool (`ToolBase`):
|
| 78 |
+
The tool instance being invoked.
|
| 79 |
+
input_kwargs (`dict[str, Any]`):
|
| 80 |
+
The tool's input arguments for this invocation. Pass them on
|
| 81 |
+
via ``next_handler(**input_kwargs)``; mutate or replace them to
|
| 82 |
+
change what the inner layers and the tool itself receive.
|
| 83 |
+
next_handler (`Callable[..., AsyncGenerator[ToolChunk, None]]`):
|
| 84 |
+
Call it as ``next_handler(**input_kwargs)`` to run the next
|
| 85 |
+
layer. It always returns an async generator, regardless of
|
| 86 |
+
whether the underlying tool is streaming or not.
|
| 87 |
+
|
| 88 |
+
Yields:
|
| 89 |
+
`ToolChunk`:
|
| 90 |
+
The chunks produced by this tool invocation.
|
| 91 |
+
"""
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
class ToolBase(ABC):
|
| 95 |
+
"""The tool protocol."""
|
| 96 |
+
|
| 97 |
+
name: str
|
| 98 |
+
"""The name presented to the agent."""
|
| 99 |
+
description: str
|
| 100 |
+
"""The agent-oriented tool description."""
|
| 101 |
+
input_schema: dict[str, Any]
|
| 102 |
+
"""The input schema of the tool, following JSON schema format."""
|
| 103 |
+
is_concurrency_safe: bool
|
| 104 |
+
"""If this tool is concurrency safe."""
|
| 105 |
+
is_read_only: bool
|
| 106 |
+
"""If this tool is read-only, which will be used in the permission
|
| 107 |
+
checking."""
|
| 108 |
+
is_external_tool: bool = False
|
| 109 |
+
"""If this tool is an external tool, which doesn't need to implement the
|
| 110 |
+
__call__ method and the agent will yield the external tool call event."""
|
| 111 |
+
is_state_injected: bool = False
|
| 112 |
+
"""If this tool requires agent state to be injected when called. If `True`,
|
| 113 |
+
the state will be injected by an argument named `_agent_state`. Note your
|
| 114 |
+
tool should be able to accept such argument.
|
| 115 |
+
"""
|
| 116 |
+
is_mcp: bool = False
|
| 117 |
+
"""If this tool is an MCP tool, which will be used in the permission"""
|
| 118 |
+
mcp_name: str | None = None
|
| 119 |
+
"""The name of the MCP server this tool belongs to, which is required if
|
| 120 |
+
this tool is an MCP tool."""
|
| 121 |
+
|
| 122 |
+
# Class attributes for dangerous path checking
|
| 123 |
+
dangerous_files: list[str] = DEFAULT_DANGEROUS_FILES
|
| 124 |
+
"""List of dangerous files that should be protected from auto-editing."""
|
| 125 |
+
dangerous_directories: list[str] = DEFAULT_DANGEROUS_DIRECTORIES
|
| 126 |
+
"""List of dangerous directories that should be protected from
|
| 127 |
+
auto-editing."""
|
| 128 |
+
|
| 129 |
+
def __init__(
|
| 130 |
+
self,
|
| 131 |
+
middlewares: List["ToolMiddlewareBase"] | None = None,
|
| 132 |
+
) -> None:
|
| 133 |
+
"""Initialize the tool with optional middlewares.
|
| 134 |
+
|
| 135 |
+
Args:
|
| 136 |
+
middlewares (`List[ToolMiddlewareBase] | None`, optional):
|
| 137 |
+
A list of :class:`ToolMiddlewareBase` instances wrapping the
|
| 138 |
+
tool execution in an onion fashion. Defaults to an empty list.
|
| 139 |
+
"""
|
| 140 |
+
self._middlewares: List["ToolMiddlewareBase"] = (
|
| 141 |
+
middlewares if middlewares is not None else []
|
| 142 |
+
)
|
| 143 |
+
|
| 144 |
+
async def call(
|
| 145 |
+
self,
|
| 146 |
+
*args: Any,
|
| 147 |
+
**kwargs: Any,
|
| 148 |
+
) -> ToolChunk | AsyncGenerator[ToolChunk, None]:
|
| 149 |
+
"""Execute the tool logic.
|
| 150 |
+
|
| 151 |
+
This is the new override point for tool implementations.
|
| 152 |
+
Subclasses should override this method instead of
|
| 153 |
+
:meth:`__call__`. The base implementation raises
|
| 154 |
+
:exc:`NotImplementedError` for non-external tools and
|
| 155 |
+
:exc:`RuntimeError` for external tools.
|
| 156 |
+
|
| 157 |
+
Args:
|
| 158 |
+
**kwargs: Tool input arguments.
|
| 159 |
+
|
| 160 |
+
Returns:
|
| 161 |
+
`ToolChunk | AsyncGenerator[ToolChunk, None]`:
|
| 162 |
+
A single :class:`~agentscope.tool.ToolChunk` or an
|
| 163 |
+
async generator that yields them.
|
| 164 |
+
"""
|
| 165 |
+
if not self.is_external_tool:
|
| 166 |
+
raise NotImplementedError(
|
| 167 |
+
f"{self.__class__.__name__} does not implement call",
|
| 168 |
+
)
|
| 169 |
+
|
| 170 |
+
raise RuntimeError(
|
| 171 |
+
f"{self.__class__.__name__} is an external tool and should not "
|
| 172 |
+
f"be called directly",
|
| 173 |
+
)
|
| 174 |
+
|
| 175 |
+
async def __call__(
|
| 176 |
+
self,
|
| 177 |
+
*args: Any,
|
| 178 |
+
**kwargs: Any,
|
| 179 |
+
) -> ToolChunk | AsyncGenerator[ToolChunk, None]:
|
| 180 |
+
"""Invoke the tool, layering any registered middlewares around
|
| 181 |
+
:meth:`call`.
|
| 182 |
+
|
| 183 |
+
Tools are always invoked with keyword arguments only. ``*args`` is
|
| 184 |
+
accepted in the signature solely to stay Liskov-compatible with
|
| 185 |
+
subclasses that override ``__call__`` with their own positional
|
| 186 |
+
parameters; any positional argument actually passed here is rejected
|
| 187 |
+
(raising :exc:`TypeError`) so it fails loudly instead of being silently
|
| 188 |
+
dropped.
|
| 189 |
+
|
| 190 |
+
Middlewares are applied in an onion fashion: the first registered
|
| 191 |
+
middleware is the outermost layer and runs its pre-logic before
|
| 192 |
+
any inner layers, then its post-logic after all inner layers
|
| 193 |
+
have completed.
|
| 194 |
+
"""
|
| 195 |
+
if args:
|
| 196 |
+
raise TypeError(
|
| 197 |
+
f"{type(self).__name__} must be called with keyword arguments "
|
| 198 |
+
f"only, but got {len(args)} positional argument(s).",
|
| 199 |
+
)
|
| 200 |
+
# ``getattr`` with a default so the no-middleware path keeps working
|
| 201 |
+
# even if a subclass overrides ``__init__`` without calling
|
| 202 |
+
# ``super().__init__()``.
|
| 203 |
+
middlewares = getattr(self, "_middlewares", [])
|
| 204 |
+
if not middlewares:
|
| 205 |
+
if inspect.isasyncgenfunction(self.call):
|
| 206 |
+
return self.call(**kwargs)
|
| 207 |
+
return await self.call(**kwargs)
|
| 208 |
+
|
| 209 |
+
async def execute_chain(
|
| 210 |
+
index: int = 0,
|
| 211 |
+
**chain_kwargs: Any,
|
| 212 |
+
) -> AsyncGenerator[ToolChunk, None]:
|
| 213 |
+
"""Execute the tool middleware chain."""
|
| 214 |
+
if index >= len(middlewares):
|
| 215 |
+
# Innermost layer: run the tool's own ``call``. ``call`` is
|
| 216 |
+
# always async but comes in two shapes — an async generator
|
| 217 |
+
# function (e.g. ``Bash``) or a coroutine returning a single
|
| 218 |
+
# ``ToolChunk`` / an async generator (e.g. ``FunctionTool``).
|
| 219 |
+
# Normalize both into a single stream so middlewares never have
|
| 220 |
+
# to distinguish them.
|
| 221 |
+
if inspect.isasyncgenfunction(self.call):
|
| 222 |
+
async for chunk in self.call(**chain_kwargs):
|
| 223 |
+
yield chunk
|
| 224 |
+
else:
|
| 225 |
+
result = await self.call(**chain_kwargs)
|
| 226 |
+
if isinstance(result, AsyncGenerator):
|
| 227 |
+
async for chunk in result:
|
| 228 |
+
yield chunk
|
| 229 |
+
else:
|
| 230 |
+
yield result
|
| 231 |
+
else:
|
| 232 |
+
mw = middlewares[index]
|
| 233 |
+
input_kwargs = dict(chain_kwargs)
|
| 234 |
+
|
| 235 |
+
async def next_handler(
|
| 236 |
+
**kw: Any,
|
| 237 |
+
) -> AsyncGenerator[ToolChunk, None]:
|
| 238 |
+
async for chunk in execute_chain(index + 1, **kw):
|
| 239 |
+
yield chunk
|
| 240 |
+
|
| 241 |
+
async for chunk in mw.on_tool_call(
|
| 242 |
+
tool=self,
|
| 243 |
+
input_kwargs=input_kwargs,
|
| 244 |
+
next_handler=next_handler,
|
| 245 |
+
):
|
| 246 |
+
yield chunk
|
| 247 |
+
|
| 248 |
+
return execute_chain(**kwargs)
|
| 249 |
+
|
| 250 |
+
@abstractmethod
|
| 251 |
+
async def check_permissions(
|
| 252 |
+
self,
|
| 253 |
+
tool_input: dict[str, Any],
|
| 254 |
+
context: PermissionContext,
|
| 255 |
+
) -> PermissionDecision:
|
| 256 |
+
"""Check permissions for the tool usage."""
|
| 257 |
+
|
| 258 |
+
async def check_read_only(
|
| 259 |
+
self,
|
| 260 |
+
tool_input: dict[str, Any],
|
| 261 |
+
) -> bool:
|
| 262 |
+
"""Decide whether this specific invocation is read-only.
|
| 263 |
+
|
| 264 |
+
Returns the static :attr:`is_read_only` attribute by default.
|
| 265 |
+
Subclasses with input-dependent semantics (e.g. ``Bash``) should
|
| 266 |
+
override this to inspect ``tool_input`` — for example, ``Bash`` is
|
| 267 |
+
statically marked as not read-only but ``ls -a`` is in fact read-only.
|
| 268 |
+
|
| 269 |
+
Should be cheap — the permission engine may call this before the
|
| 270 |
+
full :meth:`check_permissions` flow.
|
| 271 |
+
|
| 272 |
+
Args:
|
| 273 |
+
tool_input (`dict[str, Any]`):
|
| 274 |
+
The tool input data for this invocation.
|
| 275 |
+
|
| 276 |
+
Returns:
|
| 277 |
+
`bool`:
|
| 278 |
+
``True`` if this invocation is read-only, ``False`` otherwise.
|
| 279 |
+
"""
|
| 280 |
+
return self.is_read_only
|
| 281 |
+
|
| 282 |
+
async def match_rule(
|
| 283 |
+
self,
|
| 284 |
+
rule_content: str | None,
|
| 285 |
+
tool_input: dict[str, Any],
|
| 286 |
+
) -> bool:
|
| 287 |
+
"""Check if a permission rule matches the tool input.
|
| 288 |
+
|
| 289 |
+
.. note:: This is an optional method. A rule with no content (``None``)
|
| 290 |
+
is a tool-name-level rule that matches every invocation; a rule
|
| 291 |
+
with content requires the tool to override this method with its
|
| 292 |
+
own matching logic, otherwise it returns ``False``.
|
| 293 |
+
|
| 294 |
+
This means:
|
| 295 |
+
- ``_FunctionTool`` and ``MCPTool`` (which do not override this)
|
| 296 |
+
can still be controlled at the tool-name level via rules like
|
| 297 |
+
``{"tool_name": "my_tool", "rule_content": None}``.
|
| 298 |
+
- Specific tools (Bash, Read, Write, Edit, Glob, Grep) override
|
| 299 |
+
this method to support fine-grained pattern matching.
|
| 300 |
+
|
| 301 |
+
Args:
|
| 302 |
+
rule_content (`str | None`):
|
| 303 |
+
The rule pattern to match. ``None`` means "match all
|
| 304 |
+
invocations of this tool" (tool-name-level rule).
|
| 305 |
+
tool_input (`dict[str, Any]`):
|
| 306 |
+
The tool input data
|
| 307 |
+
|
| 308 |
+
Returns:
|
| 309 |
+
`bool`:
|
| 310 |
+
True if the rule matches, False otherwise
|
| 311 |
+
"""
|
| 312 |
+
# None rule_content = tool-name-level rule, matches everything
|
| 313 |
+
return rule_content is None
|
| 314 |
+
|
| 315 |
+
async def generate_suggestions(
|
| 316 |
+
self,
|
| 317 |
+
tool_input: dict[str, Any],
|
| 318 |
+
) -> List[PermissionRule]:
|
| 319 |
+
"""Generate suggested permission rules for the tool input.
|
| 320 |
+
|
| 321 |
+
.. note:: Suggest a single tool-name-level rule (``rule_content=None``)
|
| 322 |
+
that allows all invocations of this tool. Tools can override this to
|
| 323 |
+
provide finer-grained suggestions.
|
| 324 |
+
|
| 325 |
+
For example:
|
| 326 |
+
- File tools (Read/Write/Edit): suggest a glob pattern covering the
|
| 327 |
+
parent directory (e.g., "src/main.py" -> "src/**")
|
| 328 |
+
- Bash: suggest command prefix patterns (e.g., "git commit -m 'xxx'"
|
| 329 |
+
-> "git commit:*")
|
| 330 |
+
- Grep/Glob: suggest patterns based on search paths
|
| 331 |
+
|
| 332 |
+
Args:
|
| 333 |
+
tool_input (`dict[str, Any]`):
|
| 334 |
+
The tool input data
|
| 335 |
+
|
| 336 |
+
Returns:
|
| 337 |
+
`List[PermissionRule]`:
|
| 338 |
+
List of suggested permission rules (usually 1, max 5 for
|
| 339 |
+
compound operations)
|
| 340 |
+
"""
|
| 341 |
+
return [
|
| 342 |
+
PermissionRule(
|
| 343 |
+
tool_name=self.name,
|
| 344 |
+
rule_content=None,
|
| 345 |
+
behavior=PermissionBehavior.ALLOW,
|
| 346 |
+
source="suggested",
|
| 347 |
+
),
|
| 348 |
+
]
|
| 349 |
+
|
| 350 |
+
def _path_in_allowed_working_path(
|
| 351 |
+
self,
|
| 352 |
+
file_path: str,
|
| 353 |
+
context: PermissionContext,
|
| 354 |
+
) -> bool:
|
| 355 |
+
"""Check if a file path is within any allowed working directory.
|
| 356 |
+
|
| 357 |
+
A "working directory" is the process's current directory plus any
|
| 358 |
+
entries in :attr:`PermissionContext.working_directories`. Paths
|
| 359 |
+
are compared via :func:`os.path.realpath` so that aliases like
|
| 360 |
+
macOS's ``/tmp`` → ``/private/tmp`` and symlinked working
|
| 361 |
+
directories compare equal on both sides.
|
| 362 |
+
|
| 363 |
+
Used by tools that conditionally auto-allow file operations in
|
| 364 |
+
:attr:`PermissionMode.ACCEPT_EDITS` (e.g. Write, Edit, and the
|
| 365 |
+
filesystem-command branch of Bash).
|
| 366 |
+
|
| 367 |
+
Args:
|
| 368 |
+
file_path (`str`):
|
| 369 |
+
The file path to check.
|
| 370 |
+
context (`PermissionContext`):
|
| 371 |
+
The permission context containing the working directories.
|
| 372 |
+
|
| 373 |
+
Returns:
|
| 374 |
+
`bool`:
|
| 375 |
+
True if ``file_path`` is within any allowed working
|
| 376 |
+
directory.
|
| 377 |
+
"""
|
| 378 |
+
current_dir = os.getcwd()
|
| 379 |
+
additional_dirs = list(context.working_directories.keys())
|
| 380 |
+
all_working_dirs = [current_dir] + additional_dirs
|
| 381 |
+
|
| 382 |
+
abs_file_path = os.path.realpath(os.path.expanduser(file_path))
|
| 383 |
+
|
| 384 |
+
for working_dir in all_working_dirs:
|
| 385 |
+
abs_working_dir = os.path.realpath(
|
| 386 |
+
os.path.expanduser(working_dir),
|
| 387 |
+
)
|
| 388 |
+
try:
|
| 389 |
+
os.path.relpath(abs_file_path, abs_working_dir)
|
| 390 |
+
if (
|
| 391 |
+
abs_file_path.startswith(abs_working_dir + os.sep)
|
| 392 |
+
or abs_file_path == abs_working_dir
|
| 393 |
+
):
|
| 394 |
+
return True
|
| 395 |
+
except ValueError:
|
| 396 |
+
# On Windows, relpath raises ValueError if paths are on
|
| 397 |
+
# different drives.
|
| 398 |
+
continue
|
| 399 |
+
|
| 400 |
+
return False
|
| 401 |
+
|
| 402 |
+
def _is_dangerous_path(self, file_path: str) -> bool:
|
| 403 |
+
"""Check if a file path is dangerous (sensitive file or directory).
|
| 404 |
+
|
| 405 |
+
A path is considered dangerous if:
|
| 406 |
+
1. The filename matches a dangerous file (e.g., .bashrc, .gitconfig)
|
| 407 |
+
2. Any path segment matches a dangerous directory (e.g., .git, .ssh)
|
| 408 |
+
|
| 409 |
+
Case-insensitive matching is used to prevent bypasses on
|
| 410 |
+
case-insensitive filesystems (macOS, Windows).
|
| 411 |
+
|
| 412 |
+
Args:
|
| 413 |
+
file_path (`str`):
|
| 414 |
+
The file path to check
|
| 415 |
+
|
| 416 |
+
Returns:
|
| 417 |
+
`bool`:
|
| 418 |
+
True if the path is dangerous and should require explicit
|
| 419 |
+
permission
|
| 420 |
+
|
| 421 |
+
Example:
|
| 422 |
+
>>> self._is_dangerous_path("/home/user/.bashrc")
|
| 423 |
+
True
|
| 424 |
+
>>> self._is_dangerous_path("/home/user/.git/config")
|
| 425 |
+
True
|
| 426 |
+
>>> self._is_dangerous_path("/home/user/project/main.py")
|
| 427 |
+
False
|
| 428 |
+
"""
|
| 429 |
+
|
| 430 |
+
# Normalize path
|
| 431 |
+
abs_path = os.path.abspath(os.path.expanduser(file_path))
|
| 432 |
+
|
| 433 |
+
# Split path into segments
|
| 434 |
+
path_parts = Path(abs_path).parts
|
| 435 |
+
path_parts_lower = [p.lower() for p in path_parts]
|
| 436 |
+
|
| 437 |
+
# Check if filename matches dangerous files (case-insensitive)
|
| 438 |
+
filename = os.path.basename(abs_path)
|
| 439 |
+
filename_lower = filename.lower()
|
| 440 |
+
for dangerous_file in self.dangerous_files:
|
| 441 |
+
if filename_lower == dangerous_file.lower():
|
| 442 |
+
return True
|
| 443 |
+
|
| 444 |
+
# Check if any path segment matches dangerous directories
|
| 445 |
+
# (case-insensitive)
|
| 446 |
+
for dangerous_dir in self.dangerous_directories:
|
| 447 |
+
dangerous_dir_lower = dangerous_dir.lower()
|
| 448 |
+
if dangerous_dir_lower in path_parts_lower:
|
| 449 |
+
return True
|
| 450 |
+
|
| 451 |
+
return False
|
src/agentscope/tool/_builtin/__init__.py
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The builtin tools in agentscope."""
|
| 3 |
+
|
| 4 |
+
from ._backend import BackendBase, ExecResult, LocalBackend
|
| 5 |
+
from ._bash import Bash
|
| 6 |
+
from ._edit import Edit
|
| 7 |
+
from ._glob import Glob
|
| 8 |
+
from ._grep import Grep
|
| 9 |
+
from ._meta import ResetTools
|
| 10 |
+
from ._read import Read
|
| 11 |
+
from ._skill import SkillViewer
|
| 12 |
+
from ._write import Write
|
| 13 |
+
|
| 14 |
+
__all__ = [
|
| 15 |
+
"ResetTools",
|
| 16 |
+
"SkillViewer",
|
| 17 |
+
"Bash",
|
| 18 |
+
"Edit",
|
| 19 |
+
"Glob",
|
| 20 |
+
"Grep",
|
| 21 |
+
"Read",
|
| 22 |
+
"Write",
|
| 23 |
+
"BackendBase",
|
| 24 |
+
"LocalBackend",
|
| 25 |
+
"ExecResult",
|
| 26 |
+
]
|
src/agentscope/tool/_builtin/_backend.py
ADDED
|
@@ -0,0 +1,733 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""Backend abstraction for builtin tools.
|
| 3 |
+
|
| 4 |
+
Provides a :class:`BackendBase` abstract base class that captures the
|
| 5 |
+
core I/O primitives shared across all six builtin tools (Bash, Read,
|
| 6 |
+
Write, Edit, Grep, Glob).
|
| 7 |
+
|
| 8 |
+
Every backend implements exactly **three** abstract primitives whose
|
| 9 |
+
mechanism genuinely differs per environment:
|
| 10 |
+
|
| 11 |
+
* :meth:`BackendBase.exec_shell` — run a program from an argv list
|
| 12 |
+
(no shell; callers needing shell features wrap with ``sh -c``).
|
| 13 |
+
* :meth:`BackendBase.read_file` — read raw bytes.
|
| 14 |
+
* :meth:`BackendBase.write_file` — write raw bytes.
|
| 15 |
+
|
| 16 |
+
All remaining filesystem operations (``file_exists``, ``is_dir``,
|
| 17 |
+
``list_dir``, ``stat_mtime``, ``delete_path``) are derived on the base
|
| 18 |
+
class from ``exec_shell`` and work out-of-the-box for any remote
|
| 19 |
+
backend. A backend that has a cheaper native path (e.g.
|
| 20 |
+
:class:`LocalBackend` using ``os.*``) simply overrides them.
|
| 21 |
+
|
| 22 |
+
Concrete implementations:
|
| 23 |
+
|
| 24 |
+
* :class:`LocalBackend` — default; uses ``asyncio`` subprocesses,
|
| 25 |
+
``aiofiles``, and ``os.*`` for host-local I/O. Injected automatically
|
| 26 |
+
when no explicit backend is given.
|
| 27 |
+
* ``DockerBackend`` — uses ``aiodocker`` exec / archive APIs.
|
| 28 |
+
* ``E2BBackend`` — uses the E2B SDK ``commands`` / ``files`` APIs.
|
| 29 |
+
|
| 30 |
+
By accepting a ``BackendBase`` parameter, each builtin tool can
|
| 31 |
+
operate identically in local, Docker, and E2B workspaces without any
|
| 32 |
+
workspace-specific branching inside the tool code itself.
|
| 33 |
+
"""
|
| 34 |
+
|
| 35 |
+
from __future__ import annotations
|
| 36 |
+
|
| 37 |
+
import asyncio
|
| 38 |
+
import os
|
| 39 |
+
import posixpath
|
| 40 |
+
import shlex
|
| 41 |
+
import shutil
|
| 42 |
+
from abc import ABC, abstractmethod
|
| 43 |
+
from dataclasses import dataclass
|
| 44 |
+
from types import ModuleType
|
| 45 |
+
from typing import Any
|
| 46 |
+
|
| 47 |
+
import aiofiles
|
| 48 |
+
|
| 49 |
+
# ── data class ─────────────────────────────────────────────────────────
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
@dataclass(frozen=True, slots=True)
|
| 53 |
+
class ExecResult:
|
| 54 |
+
"""Result of running a shell command via a backend.
|
| 55 |
+
|
| 56 |
+
Attributes:
|
| 57 |
+
exit_code: Process exit code. ``-1`` conventionally indicates
|
| 58 |
+
an internal failure (timeout, connection error, …).
|
| 59 |
+
stdout: Raw bytes captured from standard output.
|
| 60 |
+
stderr: Raw bytes captured from standard error.
|
| 61 |
+
"""
|
| 62 |
+
|
| 63 |
+
exit_code: int
|
| 64 |
+
stdout: bytes
|
| 65 |
+
stderr: bytes
|
| 66 |
+
|
| 67 |
+
def ok(self) -> bool:
|
| 68 |
+
"""Whether the command exited successfully.
|
| 69 |
+
|
| 70 |
+
Returns:
|
| 71 |
+
`bool`:
|
| 72 |
+
``True`` iff the command exited with code ``0``.
|
| 73 |
+
"""
|
| 74 |
+
return self.exit_code == 0
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
# ── helpers ────────────────────────────────────────────────────────────
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
def _normalize_newlines(text: str) -> str:
|
| 81 |
+
"""Normalize Windows/old-Mac line endings to ``\\n``.
|
| 82 |
+
|
| 83 |
+
Converts ``\\r\\n`` (Windows) and lone ``\\r`` (classic Mac) to a
|
| 84 |
+
single ``\\n``. Builtin tools read files as raw bytes (so binary
|
| 85 |
+
payloads survive intact); when the bytes are decoded as text for
|
| 86 |
+
line-based caching, editing, or matching, the line endings must be
|
| 87 |
+
normalized so that content written on Windows behaves identically to
|
| 88 |
+
content written on POSIX.
|
| 89 |
+
|
| 90 |
+
Args:
|
| 91 |
+
text (`str`):
|
| 92 |
+
Decoded file contents.
|
| 93 |
+
|
| 94 |
+
Returns:
|
| 95 |
+
`str`:
|
| 96 |
+
The text with all line endings collapsed to ``\\n``.
|
| 97 |
+
"""
|
| 98 |
+
return text.replace("\r\n", "\n").replace("\r", "\n")
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
# ── base class ─────────────────────────────────────────────────────────
|
| 102 |
+
|
| 103 |
+
|
| 104 |
+
class BackendBase(ABC):
|
| 105 |
+
"""Filesystem + subprocess interface consumed by builtin tools.
|
| 106 |
+
|
| 107 |
+
Subclasses must implement three abstract primitives — ``exec_shell``,
|
| 108 |
+
``read_file``, and ``write_file`` — which are the only operations
|
| 109 |
+
whose mechanism genuinely differs per environment. The remaining
|
| 110 |
+
filesystem helpers are implemented here on top of ``exec_shell`` and
|
| 111 |
+
work for any backend whose shell is POSIX-like; a backend with a
|
| 112 |
+
cheaper native path may override them (see :class:`LocalBackend`).
|
| 113 |
+
"""
|
| 114 |
+
|
| 115 |
+
#: Path-manipulation module whose semantics match the backend's
|
| 116 |
+
#: environment. Used by :meth:`join_path`, :meth:`dirname`,
|
| 117 |
+
#: :meth:`isabs`, :meth:`normpath`, and :meth:`abspath` to ensure
|
| 118 |
+
#: correct behavior when the host OS and the backend OS differ
|
| 119 |
+
#: (e.g. a Windows host driving a Linux Docker container).
|
| 120 |
+
#:
|
| 121 |
+
#: Defaults to :mod:`posixpath`, which is correct for any backend
|
| 122 |
+
#: whose environment is Linux/macOS (Docker, E2B, …). Subclasses
|
| 123 |
+
#: targeting a different environment override this attribute, e.g.
|
| 124 |
+
#: :class:`LocalBackend` sets it to :mod:`os.path`, and a future
|
| 125 |
+
#: Windows-container backend would set it to :mod:`ntpath`.
|
| 126 |
+
#:
|
| 127 |
+
#: .. important::
|
| 128 |
+
#:
|
| 129 |
+
#: Only **pure string operations** on this module are safe to
|
| 130 |
+
#: call (``join``, ``split``, ``dirname``, ``basename``,
|
| 131 |
+
#: ``normpath``, ``isabs``, ``splitext``, ``splitdrive``, …).
|
| 132 |
+
#:
|
| 133 |
+
#: Do **not** call functions that touch the filesystem or
|
| 134 |
+
#: environment variables — ``exists``, ``isfile``, ``isdir``,
|
| 135 |
+
#: ``getmtime``, ``realpath``, ``expanduser``, ``expandvars``,
|
| 136 |
+
#: and parameterless ``abspath`` — because those read the
|
| 137 |
+
#: **host** process's filesystem / ``$HOME`` / ``cwd``, which
|
| 138 |
+
#: is meaningless (and a silent bug) for remote backends. Use
|
| 139 |
+
#: the async I/O methods on the backend instead
|
| 140 |
+
#: (:meth:`file_exists`, :meth:`is_dir`, :meth:`stat_mtime`,
|
| 141 |
+
#: …) or the :meth:`abspath` wrapper below, which requires an
|
| 142 |
+
#: explicit ``cwd`` argument.
|
| 143 |
+
_path_module: ModuleType = posixpath
|
| 144 |
+
|
| 145 |
+
# ── path manipulation helpers (pure string ops) ────────────────
|
| 146 |
+
|
| 147 |
+
def join_path(self, path: str, *paths: str) -> str:
|
| 148 |
+
"""Join one or more path components using the backend's separator.
|
| 149 |
+
|
| 150 |
+
Args:
|
| 151 |
+
path (`str`):
|
| 152 |
+
The first path component.
|
| 153 |
+
*paths (`str`):
|
| 154 |
+
Additional components to join onto ``path``.
|
| 155 |
+
|
| 156 |
+
Returns:
|
| 157 |
+
`str`:
|
| 158 |
+
The joined path, using the backend environment's path
|
| 159 |
+
separator.
|
| 160 |
+
"""
|
| 161 |
+
return self._path_module.join(path, *paths)
|
| 162 |
+
|
| 163 |
+
def dirname(self, path: str) -> str:
|
| 164 |
+
"""Return the directory component of ``path``.
|
| 165 |
+
|
| 166 |
+
Args:
|
| 167 |
+
path (`str`):
|
| 168 |
+
A path inside the backend's environment.
|
| 169 |
+
|
| 170 |
+
Returns:
|
| 171 |
+
`str`:
|
| 172 |
+
Everything up to (but not including) the last path
|
| 173 |
+
separator. Empty string if ``path`` has no separator.
|
| 174 |
+
"""
|
| 175 |
+
return self._path_module.dirname(path)
|
| 176 |
+
|
| 177 |
+
def isabs(self, path: str) -> bool:
|
| 178 |
+
"""Return ``True`` if ``path`` is absolute in the backend.
|
| 179 |
+
|
| 180 |
+
Args:
|
| 181 |
+
path (`str`):
|
| 182 |
+
A path inside the backend's environment.
|
| 183 |
+
|
| 184 |
+
Returns:
|
| 185 |
+
`bool`:
|
| 186 |
+
``True`` iff ``path`` is absolute under the backend
|
| 187 |
+
environment's path semantics.
|
| 188 |
+
"""
|
| 189 |
+
return self._path_module.isabs(path)
|
| 190 |
+
|
| 191 |
+
def normpath(self, path: str) -> str:
|
| 192 |
+
"""Normalize ``path`` (collapse ``..``, ``.``, duplicate seps).
|
| 193 |
+
|
| 194 |
+
Pure string operation — does not touch the filesystem.
|
| 195 |
+
|
| 196 |
+
Args:
|
| 197 |
+
path (`str`):
|
| 198 |
+
A path inside the backend's environment.
|
| 199 |
+
|
| 200 |
+
Returns:
|
| 201 |
+
`str`:
|
| 202 |
+
The normalized path.
|
| 203 |
+
"""
|
| 204 |
+
return self._path_module.normpath(path)
|
| 205 |
+
|
| 206 |
+
def abspath(self, path: str, *, cwd: str) -> str:
|
| 207 |
+
"""Return an absolute, normalized version of ``path``.
|
| 208 |
+
|
| 209 |
+
Unlike :func:`os.path.abspath`, this helper **never** reads
|
| 210 |
+
the host process's working directory: when ``path`` is
|
| 211 |
+
relative it is joined with the explicitly supplied ``cwd``,
|
| 212 |
+
which must itself be a path that is meaningful inside the
|
| 213 |
+
backend's environment. This avoids the silent bug where the
|
| 214 |
+
host's ``os.getcwd()`` leaks into paths that will actually be
|
| 215 |
+
used on a remote backend.
|
| 216 |
+
|
| 217 |
+
Args:
|
| 218 |
+
path (`str`):
|
| 219 |
+
A path inside the backend's environment.
|
| 220 |
+
cwd (`str`):
|
| 221 |
+
Directory to resolve a relative ``path`` against.
|
| 222 |
+
Ignored when ``path`` is already absolute.
|
| 223 |
+
|
| 224 |
+
Returns:
|
| 225 |
+
`str`:
|
| 226 |
+
An absolute, normalized path.
|
| 227 |
+
"""
|
| 228 |
+
if self._path_module.isabs(path):
|
| 229 |
+
return self._path_module.normpath(path)
|
| 230 |
+
return self._path_module.normpath(
|
| 231 |
+
self._path_module.join(cwd, path),
|
| 232 |
+
)
|
| 233 |
+
|
| 234 |
+
# ── abstract primitives ────────────────────────────────────────
|
| 235 |
+
|
| 236 |
+
@abstractmethod
|
| 237 |
+
async def exec_shell(
|
| 238 |
+
self,
|
| 239 |
+
command: list[str],
|
| 240 |
+
*,
|
| 241 |
+
cwd: str | None = None,
|
| 242 |
+
timeout: float | None = None,
|
| 243 |
+
) -> ExecResult:
|
| 244 |
+
"""Run a program directly from an argument vector.
|
| 245 |
+
|
| 246 |
+
*command* is an executable followed by its arguments — it is
|
| 247 |
+
**not** passed through a shell, so callers never have to quote
|
| 248 |
+
or escape arguments and there is no platform-specific quoting
|
| 249 |
+
bug. This makes the primitive portable to Windows, where POSIX
|
| 250 |
+
single-quote escaping (``shlex.quote``) is not understood by
|
| 251 |
+
``cmd.exe``.
|
| 252 |
+
|
| 253 |
+
Callers that genuinely need shell features (pipes, redirects,
|
| 254 |
+
``&&``) must wrap their command line explicitly, e.g.
|
| 255 |
+
``["/bin/sh", "-c", command_line]``.
|
| 256 |
+
|
| 257 |
+
Args:
|
| 258 |
+
command (`list[str]`):
|
| 259 |
+
Executable path/name followed by its arguments.
|
| 260 |
+
cwd (`str | None`, optional):
|
| 261 |
+
Working directory to run the command in. When ``None``
|
| 262 |
+
the backend's default working directory is used.
|
| 263 |
+
timeout (`float | None`, optional):
|
| 264 |
+
Maximum number of seconds to wait. When ``None`` the
|
| 265 |
+
call waits indefinitely. On timeout the result carries
|
| 266 |
+
an ``exit_code`` of ``-1``.
|
| 267 |
+
|
| 268 |
+
Returns:
|
| 269 |
+
`ExecResult`:
|
| 270 |
+
The captured exit code, stdout, and stderr.
|
| 271 |
+
"""
|
| 272 |
+
|
| 273 |
+
@abstractmethod
|
| 274 |
+
async def read_file(self, path: str) -> bytes:
|
| 275 |
+
"""Read the full contents of ``path`` as raw bytes.
|
| 276 |
+
|
| 277 |
+
Args:
|
| 278 |
+
path (`str`):
|
| 279 |
+
Path to the file inside the backend's environment.
|
| 280 |
+
|
| 281 |
+
Returns:
|
| 282 |
+
`bytes`:
|
| 283 |
+
The raw file contents.
|
| 284 |
+
"""
|
| 285 |
+
|
| 286 |
+
@abstractmethod
|
| 287 |
+
async def write_file(self, path: str, data: bytes) -> None:
|
| 288 |
+
"""Write ``data`` to ``path``, creating parent directories.
|
| 289 |
+
|
| 290 |
+
Args:
|
| 291 |
+
path (`str`):
|
| 292 |
+
Destination path inside the backend's environment.
|
| 293 |
+
data (`bytes`):
|
| 294 |
+
The raw bytes to write.
|
| 295 |
+
"""
|
| 296 |
+
|
| 297 |
+
# ── derived filesystem ops (shell-based defaults) ──────────────
|
| 298 |
+
|
| 299 |
+
async def getcwd(self) -> str:
|
| 300 |
+
"""Return the backend environment's current working directory.
|
| 301 |
+
|
| 302 |
+
This is the directory that bare ``exec_shell`` invocations
|
| 303 |
+
(those with ``cwd=None``) execute in. Tools should call this
|
| 304 |
+
— instead of :func:`os.getcwd` — whenever they need a default
|
| 305 |
+
path that is meaningful inside the backend, because the host
|
| 306 |
+
process's cwd is meaningless for remote backends
|
| 307 |
+
(Docker / E2B).
|
| 308 |
+
|
| 309 |
+
The default implementation runs ``pwd`` via :meth:`exec_shell`,
|
| 310 |
+
which works for any POSIX-like backend. Backends with cheaper
|
| 311 |
+
native access (e.g. :class:`LocalBackend`, or remote backends
|
| 312 |
+
that already track their workdir) should override it.
|
| 313 |
+
|
| 314 |
+
Returns:
|
| 315 |
+
`str`:
|
| 316 |
+
The backend's current working directory.
|
| 317 |
+
"""
|
| 318 |
+
result = await self.exec_shell(["pwd"])
|
| 319 |
+
return result.stdout.decode("utf-8", errors="replace").strip()
|
| 320 |
+
|
| 321 |
+
async def expanduser(self, path: str) -> str:
|
| 322 |
+
"""Expand a leading ``~`` / ``~/`` to the backend's home directory.
|
| 323 |
+
|
| 324 |
+
Tools should call this — instead of :func:`os.path.expanduser`
|
| 325 |
+
— whenever they need to expand ``~`` in a path that lives
|
| 326 |
+
inside the backend's environment, because the host process's
|
| 327 |
+
``$HOME`` is meaningless for remote backends.
|
| 328 |
+
|
| 329 |
+
The default implementation queries ``$HOME`` via
|
| 330 |
+
:meth:`exec_shell` (POSIX-only). Only the leading ``~`` /
|
| 331 |
+
``~/foo`` form is expanded; ``~user/...`` is not supported by
|
| 332 |
+
the default and is returned unchanged. Backends with cheaper
|
| 333 |
+
native access should override (e.g. :class:`LocalBackend`).
|
| 334 |
+
|
| 335 |
+
Args:
|
| 336 |
+
path (`str`):
|
| 337 |
+
A path inside the backend's environment, possibly
|
| 338 |
+
starting with ``~``.
|
| 339 |
+
|
| 340 |
+
Returns:
|
| 341 |
+
`str`:
|
| 342 |
+
``path`` with a leading ``~`` / ``~/`` expanded. If
|
| 343 |
+
``path`` does not start with ``~``, or starts with
|
| 344 |
+
``~user`` (unsupported), it is returned unchanged.
|
| 345 |
+
"""
|
| 346 |
+
if not path or path[0] != "~":
|
| 347 |
+
return path
|
| 348 |
+
# ``~user/...`` form — not supported by the default impl.
|
| 349 |
+
if len(path) > 1 and path[1] not in ("/", self._path_module.sep):
|
| 350 |
+
return path
|
| 351 |
+
result = await self.exec_shell(["printenv", "HOME"])
|
| 352 |
+
home = result.stdout.decode("utf-8", errors="replace").strip()
|
| 353 |
+
if not home:
|
| 354 |
+
return path
|
| 355 |
+
return home + path[1:]
|
| 356 |
+
|
| 357 |
+
async def file_exists(self, path: str) -> bool:
|
| 358 |
+
"""Return ``True`` if ``path`` exists (file or directory).
|
| 359 |
+
|
| 360 |
+
Args:
|
| 361 |
+
path (`str`):
|
| 362 |
+
Path to test inside the backend's environment.
|
| 363 |
+
|
| 364 |
+
Returns:
|
| 365 |
+
`bool`:
|
| 366 |
+
``True`` if the path exists, ``False`` otherwise.
|
| 367 |
+
"""
|
| 368 |
+
result = await self.exec_shell(["test", "-e", path])
|
| 369 |
+
return result.ok()
|
| 370 |
+
|
| 371 |
+
async def is_dir(self, path: str) -> bool:
|
| 372 |
+
"""Return ``True`` if ``path`` is an existing directory.
|
| 373 |
+
|
| 374 |
+
Args:
|
| 375 |
+
path (`str`):
|
| 376 |
+
Path to test inside the backend's environment.
|
| 377 |
+
|
| 378 |
+
Returns:
|
| 379 |
+
`bool`:
|
| 380 |
+
``True`` if the path is an existing directory.
|
| 381 |
+
"""
|
| 382 |
+
result = await self.exec_shell(["test", "-d", path])
|
| 383 |
+
return result.ok()
|
| 384 |
+
|
| 385 |
+
async def list_dir(
|
| 386 |
+
self,
|
| 387 |
+
path: str,
|
| 388 |
+
*,
|
| 389 |
+
recursive: bool = False,
|
| 390 |
+
) -> list[str]:
|
| 391 |
+
"""List entries under ``path``.
|
| 392 |
+
|
| 393 |
+
Output is NUL-delimited (``find -print0`` / ``-printf '%f\\0'``)
|
| 394 |
+
and split on ``\\0`` so that file names containing spaces or
|
| 395 |
+
newlines are handled correctly. ``find -printf`` is a GNU
|
| 396 |
+
extension; backends running on non-GNU userlands should override
|
| 397 |
+
this method.
|
| 398 |
+
|
| 399 |
+
Args:
|
| 400 |
+
path (`str`):
|
| 401 |
+
Directory to list inside the backend's environment.
|
| 402 |
+
recursive (`bool`, optional):
|
| 403 |
+
When ``True``, return all files underneath ``path`` as
|
| 404 |
+
paths (like ``find path -type f``). When ``False``
|
| 405 |
+
(default), return the immediate children's base names
|
| 406 |
+
(like ``os.listdir``).
|
| 407 |
+
|
| 408 |
+
Returns:
|
| 409 |
+
`list[str]`:
|
| 410 |
+
The matched entries, or an empty list if ``path`` does
|
| 411 |
+
not exist or cannot be listed.
|
| 412 |
+
"""
|
| 413 |
+
if recursive:
|
| 414 |
+
command = ["find", path, "-type", "f", "-print0"]
|
| 415 |
+
else:
|
| 416 |
+
command = [
|
| 417 |
+
"find",
|
| 418 |
+
path,
|
| 419 |
+
"-mindepth",
|
| 420 |
+
"1",
|
| 421 |
+
"-maxdepth",
|
| 422 |
+
"1",
|
| 423 |
+
"-printf",
|
| 424 |
+
"%f\\0",
|
| 425 |
+
]
|
| 426 |
+
result = await self.exec_shell(command)
|
| 427 |
+
if not result.ok():
|
| 428 |
+
return []
|
| 429 |
+
return [
|
| 430 |
+
part.decode("utf-8", errors="surrogateescape")
|
| 431 |
+
for part in result.stdout.split(b"\0")
|
| 432 |
+
if part
|
| 433 |
+
]
|
| 434 |
+
|
| 435 |
+
async def stat_mtime(self, path: str) -> float | None:
|
| 436 |
+
"""Return the modification time of ``path``, or ``None``.
|
| 437 |
+
|
| 438 |
+
Tries GNU ``stat -c %Y`` first and falls back to BSD
|
| 439 |
+
``stat -f %m`` so the same call works across coreutils and
|
| 440 |
+
BSD/macOS userlands. The two attempts are combined with ``||``,
|
| 441 |
+
so this default wraps a ``sh -c`` script; backends without a
|
| 442 |
+
POSIX shell (e.g. :class:`LocalBackend`) override it.
|
| 443 |
+
|
| 444 |
+
Args:
|
| 445 |
+
path (`str`):
|
| 446 |
+
Path to stat inside the backend's environment.
|
| 447 |
+
|
| 448 |
+
Returns:
|
| 449 |
+
`float | None`:
|
| 450 |
+
The modification time as a POSIX timestamp, or ``None``
|
| 451 |
+
if the path does not exist or cannot be stat'd.
|
| 452 |
+
"""
|
| 453 |
+
quoted = shlex.quote(path)
|
| 454 |
+
script = (
|
| 455 |
+
f"stat -c %Y {quoted} 2>/dev/null || "
|
| 456 |
+
f"stat -f %m {quoted} 2>/dev/null"
|
| 457 |
+
)
|
| 458 |
+
result = await self.exec_shell(["sh", "-c", script])
|
| 459 |
+
if not result.ok():
|
| 460 |
+
return None
|
| 461 |
+
try:
|
| 462 |
+
return float(
|
| 463 |
+
result.stdout.decode("utf-8", errors="replace").strip(),
|
| 464 |
+
)
|
| 465 |
+
except ValueError:
|
| 466 |
+
return None
|
| 467 |
+
|
| 468 |
+
async def delete_path(self, path: str) -> None:
|
| 469 |
+
"""Delete ``path`` (file or directory tree).
|
| 470 |
+
|
| 471 |
+
If ``path`` does not exist the call is a silent no-op (like
|
| 472 |
+
``rm -rf``). Handles both files and directories (recursively).
|
| 473 |
+
|
| 474 |
+
Args:
|
| 475 |
+
path (`str`):
|
| 476 |
+
Path to delete inside the backend's environment.
|
| 477 |
+
"""
|
| 478 |
+
await self.exec_shell(["rm", "-rf", path])
|
| 479 |
+
|
| 480 |
+
|
| 481 |
+
# ── local backend ──────────────────────────────────────────────────────
|
| 482 |
+
|
| 483 |
+
|
| 484 |
+
def _subprocess_creation_kwargs() -> dict[str, Any]:
|
| 485 |
+
"""Return platform-specific subprocess creation options.
|
| 486 |
+
|
| 487 |
+
Returns:
|
| 488 |
+
`dict[str, Any]`:
|
| 489 |
+
Extra keyword arguments for ``create_subprocess_shell``.
|
| 490 |
+
Empty on POSIX; on Windows it sets ``creationflags`` to
|
| 491 |
+
suppress a console window.
|
| 492 |
+
"""
|
| 493 |
+
if os.name != "nt":
|
| 494 |
+
return {}
|
| 495 |
+
|
| 496 |
+
import subprocess
|
| 497 |
+
|
| 498 |
+
return {
|
| 499 |
+
"creationflags": getattr(
|
| 500 |
+
subprocess,
|
| 501 |
+
"CREATE_NO_WINDOW",
|
| 502 |
+
0x08000000,
|
| 503 |
+
),
|
| 504 |
+
}
|
| 505 |
+
|
| 506 |
+
|
| 507 |
+
class LocalBackend(BackendBase):
|
| 508 |
+
"""Host-local :class:`BackendBase` implementation.
|
| 509 |
+
|
| 510 |
+
Uses ``asyncio.create_subprocess_exec``, ``aiofiles``, and the
|
| 511 |
+
``os`` module. This is the default backend injected when no
|
| 512 |
+
explicit one is given to a builtin tool. Commands are spawned
|
| 513 |
+
directly from their argument vector (no shell), which avoids the
|
| 514 |
+
POSIX-vs-``cmd.exe`` quoting mismatch and makes the backend work on
|
| 515 |
+
Windows. The derived filesystem helpers are overridden with native
|
| 516 |
+
``os.*`` calls — faster and more robust than shelling out, and
|
| 517 |
+
portable to Windows where ``test`` / ``find`` / ``stat`` are
|
| 518 |
+
unavailable.
|
| 519 |
+
"""
|
| 520 |
+
|
| 521 |
+
# Use the host OS's path semantics (Windows or POSIX) instead of
|
| 522 |
+
# the base class default (``posixpath``), so path helpers behave
|
| 523 |
+
# correctly when running on a Windows host.
|
| 524 |
+
_path_module = os.path
|
| 525 |
+
|
| 526 |
+
async def exec_shell(
|
| 527 |
+
self,
|
| 528 |
+
command: list[str],
|
| 529 |
+
*,
|
| 530 |
+
cwd: str | None = None,
|
| 531 |
+
timeout: float | None = None,
|
| 532 |
+
) -> ExecResult:
|
| 533 |
+
"""Run a program via ``asyncio.create_subprocess_exec``.
|
| 534 |
+
|
| 535 |
+
The program is spawned directly from *command* without an
|
| 536 |
+
intervening shell, so no argument quoting is required and the
|
| 537 |
+
same code path works on POSIX and Windows.
|
| 538 |
+
|
| 539 |
+
Args:
|
| 540 |
+
command (`list[str]`):
|
| 541 |
+
Executable path/name followed by its arguments.
|
| 542 |
+
cwd (`str | None`, optional):
|
| 543 |
+
Working directory for the subprocess. When ``None`` the
|
| 544 |
+
current process working directory is used.
|
| 545 |
+
timeout (`float | None`, optional):
|
| 546 |
+
Maximum number of seconds to wait before the process is
|
| 547 |
+
killed and an ``exit_code`` of ``-1`` is returned.
|
| 548 |
+
|
| 549 |
+
Returns:
|
| 550 |
+
`ExecResult`:
|
| 551 |
+
The captured exit code, stdout, and stderr. If the
|
| 552 |
+
executable cannot be found or spawned, ``exit_code`` is
|
| 553 |
+
``127`` (matching a shell's "command not found"), with
|
| 554 |
+
the OS error message on stderr.
|
| 555 |
+
"""
|
| 556 |
+
kwargs = _subprocess_creation_kwargs()
|
| 557 |
+
if cwd is not None:
|
| 558 |
+
kwargs["cwd"] = cwd
|
| 559 |
+
|
| 560 |
+
try:
|
| 561 |
+
process = await asyncio.create_subprocess_exec(
|
| 562 |
+
*command,
|
| 563 |
+
stdout=asyncio.subprocess.PIPE,
|
| 564 |
+
stderr=asyncio.subprocess.PIPE,
|
| 565 |
+
**kwargs,
|
| 566 |
+
)
|
| 567 |
+
except (FileNotFoundError, NotADirectoryError, OSError) as exc:
|
| 568 |
+
# The executable could not be found or spawned. A shell would
|
| 569 |
+
# have returned 127 ("command not found"); mirror that so
|
| 570 |
+
# callers see a normal non-zero ExecResult instead of an
|
| 571 |
+
# exception.
|
| 572 |
+
return ExecResult(
|
| 573 |
+
exit_code=127,
|
| 574 |
+
stdout=b"",
|
| 575 |
+
stderr=str(exc).encode("utf-8"),
|
| 576 |
+
)
|
| 577 |
+
|
| 578 |
+
try:
|
| 579 |
+
stdout, stderr = await asyncio.wait_for(
|
| 580 |
+
process.communicate(),
|
| 581 |
+
timeout=timeout,
|
| 582 |
+
)
|
| 583 |
+
except asyncio.TimeoutError:
|
| 584 |
+
process.kill()
|
| 585 |
+
await process.communicate()
|
| 586 |
+
return ExecResult(exit_code=-1, stdout=b"", stderr=b"timed out")
|
| 587 |
+
|
| 588 |
+
return ExecResult(
|
| 589 |
+
exit_code=process.returncode or 0,
|
| 590 |
+
stdout=stdout,
|
| 591 |
+
stderr=stderr,
|
| 592 |
+
)
|
| 593 |
+
|
| 594 |
+
async def read_file(self, path: str) -> bytes:
|
| 595 |
+
"""Read a local file as raw bytes.
|
| 596 |
+
|
| 597 |
+
Args:
|
| 598 |
+
path (`str`):
|
| 599 |
+
Path to the local file.
|
| 600 |
+
|
| 601 |
+
Returns:
|
| 602 |
+
`bytes`:
|
| 603 |
+
The raw file contents.
|
| 604 |
+
"""
|
| 605 |
+
async with aiofiles.open(path, mode="rb") as f:
|
| 606 |
+
return await f.read()
|
| 607 |
+
|
| 608 |
+
async def write_file(self, path: str, data: bytes) -> None:
|
| 609 |
+
"""Write *data* to a local file, creating parent dirs.
|
| 610 |
+
|
| 611 |
+
Args:
|
| 612 |
+
path (`str`):
|
| 613 |
+
Destination path on the local filesystem.
|
| 614 |
+
data (`bytes`):
|
| 615 |
+
The raw bytes to write.
|
| 616 |
+
"""
|
| 617 |
+
parent = os.path.dirname(path)
|
| 618 |
+
if parent:
|
| 619 |
+
os.makedirs(parent, exist_ok=True)
|
| 620 |
+
async with aiofiles.open(path, mode="wb") as f:
|
| 621 |
+
await f.write(data)
|
| 622 |
+
|
| 623 |
+
async def getcwd(self) -> str:
|
| 624 |
+
"""Return the host process's current working directory.
|
| 625 |
+
|
| 626 |
+
Returns:
|
| 627 |
+
`str`:
|
| 628 |
+
``os.getcwd()`` — avoids spawning a ``pwd`` subprocess.
|
| 629 |
+
"""
|
| 630 |
+
return os.getcwd()
|
| 631 |
+
|
| 632 |
+
async def expanduser(self, path: str) -> str:
|
| 633 |
+
"""Expand ``~`` using the host process's ``$HOME``.
|
| 634 |
+
|
| 635 |
+
Args:
|
| 636 |
+
path (`str`):
|
| 637 |
+
A local path, possibly starting with ``~``.
|
| 638 |
+
|
| 639 |
+
Returns:
|
| 640 |
+
`str`:
|
| 641 |
+
``os.path.expanduser(path)`` — avoids spawning a
|
| 642 |
+
subprocess.
|
| 643 |
+
"""
|
| 644 |
+
return os.path.expanduser(path)
|
| 645 |
+
|
| 646 |
+
async def file_exists(self, path: str) -> bool:
|
| 647 |
+
"""Check if a local path exists.
|
| 648 |
+
|
| 649 |
+
Args:
|
| 650 |
+
path (`str`):
|
| 651 |
+
Path to test.
|
| 652 |
+
|
| 653 |
+
Returns:
|
| 654 |
+
`bool`:
|
| 655 |
+
``True`` if the path exists.
|
| 656 |
+
"""
|
| 657 |
+
return os.path.exists(path)
|
| 658 |
+
|
| 659 |
+
async def is_dir(self, path: str) -> bool:
|
| 660 |
+
"""Check if a local path is a directory.
|
| 661 |
+
|
| 662 |
+
Args:
|
| 663 |
+
path (`str`):
|
| 664 |
+
Path to test.
|
| 665 |
+
|
| 666 |
+
Returns:
|
| 667 |
+
`bool`:
|
| 668 |
+
``True`` if the path is an existing directory.
|
| 669 |
+
"""
|
| 670 |
+
return os.path.isdir(path)
|
| 671 |
+
|
| 672 |
+
async def list_dir(
|
| 673 |
+
self,
|
| 674 |
+
path: str,
|
| 675 |
+
*,
|
| 676 |
+
recursive: bool = False,
|
| 677 |
+
) -> list[str]:
|
| 678 |
+
"""List local directory entries.
|
| 679 |
+
|
| 680 |
+
Mirrors the base contract using native ``os`` calls.
|
| 681 |
+
|
| 682 |
+
Args:
|
| 683 |
+
path (`str`):
|
| 684 |
+
Directory to list.
|
| 685 |
+
recursive (`bool`, optional):
|
| 686 |
+
When ``True``, return file paths underneath ``path``
|
| 687 |
+
(``os.walk``). When ``False`` (default), return the
|
| 688 |
+
immediate children's base names (``os.listdir``).
|
| 689 |
+
|
| 690 |
+
Returns:
|
| 691 |
+
`list[str]`:
|
| 692 |
+
The matched entries.
|
| 693 |
+
"""
|
| 694 |
+
if recursive:
|
| 695 |
+
results: list[str] = []
|
| 696 |
+
for root, _dirs, files in os.walk(path):
|
| 697 |
+
for f in files:
|
| 698 |
+
results.append(os.path.join(root, f))
|
| 699 |
+
return results
|
| 700 |
+
return os.listdir(path)
|
| 701 |
+
|
| 702 |
+
async def stat_mtime(self, path: str) -> float | None:
|
| 703 |
+
"""Return the modification time of a local file.
|
| 704 |
+
|
| 705 |
+
Args:
|
| 706 |
+
path (`str`):
|
| 707 |
+
Path to stat.
|
| 708 |
+
|
| 709 |
+
Returns:
|
| 710 |
+
`float | None`:
|
| 711 |
+
The modification time as a POSIX timestamp, or ``None``
|
| 712 |
+
if the path does not exist or cannot be stat'd.
|
| 713 |
+
"""
|
| 714 |
+
try:
|
| 715 |
+
return os.stat(path).st_mtime
|
| 716 |
+
except (OSError, FileNotFoundError):
|
| 717 |
+
return None
|
| 718 |
+
|
| 719 |
+
async def delete_path(self, path: str) -> None:
|
| 720 |
+
"""Delete a local file or directory tree.
|
| 721 |
+
|
| 722 |
+
No-op if *path* does not exist.
|
| 723 |
+
|
| 724 |
+
Args:
|
| 725 |
+
path (`str`):
|
| 726 |
+
Path to delete.
|
| 727 |
+
"""
|
| 728 |
+
if not os.path.exists(path):
|
| 729 |
+
return
|
| 730 |
+
if os.path.isdir(path):
|
| 731 |
+
shutil.rmtree(path)
|
| 732 |
+
else:
|
| 733 |
+
os.remove(path)
|
src/agentscope/tool/_builtin/_bash.py
ADDED
|
@@ -0,0 +1,778 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The bash tool in agentscope."""
|
| 3 |
+
import os
|
| 4 |
+
from typing import AsyncGenerator, Any, List
|
| 5 |
+
import re
|
| 6 |
+
|
| 7 |
+
from ._bash_parser import BashCommandParser
|
| 8 |
+
from .._base import ToolBase, ToolMiddlewareBase
|
| 9 |
+
from .._constants import (
|
| 10 |
+
DEFAULT_DANGEROUS_FILES,
|
| 11 |
+
DEFAULT_DANGEROUS_DIRECTORIES,
|
| 12 |
+
)
|
| 13 |
+
from ...permission import (
|
| 14 |
+
PermissionContext,
|
| 15 |
+
PermissionDecision,
|
| 16 |
+
PermissionBehavior,
|
| 17 |
+
PermissionMode,
|
| 18 |
+
PermissionRule,
|
| 19 |
+
)
|
| 20 |
+
from ...message import TextBlock, ToolResultState
|
| 21 |
+
from .._response import ToolChunk
|
| 22 |
+
from ._backend import BackendBase
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
class Bash(ToolBase):
|
| 26 |
+
"""The bash tool."""
|
| 27 |
+
|
| 28 |
+
name: str = "Bash"
|
| 29 |
+
"""The tool name presented to the agent."""
|
| 30 |
+
|
| 31 |
+
description: str = """Executes a bash command and returns its output.
|
| 32 |
+
|
| 33 |
+
The working directory persists between commands, but shell state does
|
| 34 |
+
not. The shell environment is initialized from the user's profile
|
| 35 |
+
(bash or zsh).
|
| 36 |
+
|
| 37 |
+
IMPORTANT: Avoid using this tool to run `find`, `grep`, `cat`, `head`,
|
| 38 |
+
`tail`, `sed`, `awk`, or `echo` commands, unless explicitly instructed
|
| 39 |
+
or after you have verified that a dedicated tool cannot accomplish your
|
| 40 |
+
task. Instead, use the appropriate dedicated tool as this will provide
|
| 41 |
+
a much better experience for the user:
|
| 42 |
+
|
| 43 |
+
- File search: Use Glob (NOT find or ls)
|
| 44 |
+
- Content search: Use Grep (NOT grep or rg)
|
| 45 |
+
- Read files: Use Read (NOT cat/head/tail)
|
| 46 |
+
- Edit files: Use Edit (NOT sed/awk)
|
| 47 |
+
- Write files: Use Write (NOT echo >/cat <<EOF)
|
| 48 |
+
- Communication: Output text directly (NOT echo/printf)
|
| 49 |
+
|
| 50 |
+
While the Bash tool can do similar things, it's better to use the
|
| 51 |
+
built-in tools as they provide a better user experience and make it
|
| 52 |
+
easier to review tool calls and give permission.
|
| 53 |
+
|
| 54 |
+
# Instructions
|
| 55 |
+
- If your command will create new directories or files, first use
|
| 56 |
+
this tool to run `ls` to verify the parent directory exists and is
|
| 57 |
+
the correct location.
|
| 58 |
+
- Always quote file paths that contain spaces with double quotes in
|
| 59 |
+
your command (e.g., cd "path with spaces/file.txt")
|
| 60 |
+
- Try to maintain your current working directory throughout the
|
| 61 |
+
session by using absolute paths and avoiding usage of `cd`. You may
|
| 62 |
+
use `cd` if the User explicitly requests it.
|
| 63 |
+
- You may specify an optional timeout in milliseconds (up to 600000ms
|
| 64 |
+
/ 10 minutes). By default, your command will timeout after 120000ms
|
| 65 |
+
(2 minutes).
|
| 66 |
+
- Write a clear, concise description of what your command does. For
|
| 67 |
+
simple commands, keep it brief (5-10 words). For complex commands
|
| 68 |
+
(piped commands, obscure flags, or anything hard to understand at a
|
| 69 |
+
glance), include enough context so that the user can understand what
|
| 70 |
+
your command will do.
|
| 71 |
+
- When issuing multiple commands:
|
| 72 |
+
- If the commands are independent and can run in parallel, make
|
| 73 |
+
multiple Bash tool calls in a single message. Example: if you need
|
| 74 |
+
to run "git status" and "git diff", send a single message with two
|
| 75 |
+
Bash tool calls in parallel.
|
| 76 |
+
- If the commands depend on each other and must run sequentially,
|
| 77 |
+
use a single Bash call with '&&' to chain them together.
|
| 78 |
+
- Use ';' only when you need to run commands sequentially but don't
|
| 79 |
+
care if earlier commands fail.
|
| 80 |
+
- DO NOT use newlines to separate commands (newlines are ok in
|
| 81 |
+
quoted strings).
|
| 82 |
+
- For git commands:
|
| 83 |
+
- Prefer to create a new commit rather than amending an existing
|
| 84 |
+
commit.
|
| 85 |
+
- Before running destructive operations (e.g., git reset --hard, git
|
| 86 |
+
push --force, git checkout --), consider whether there is a safer
|
| 87 |
+
alternative that achieves the same goal. Only use destructive
|
| 88 |
+
operations when they are truly the best approach.
|
| 89 |
+
- Never skip hooks (--no-verify) or bypass signing (--no-gpg-sign,
|
| 90 |
+
-c commit.gpgsign=false) unless the user has explicitly asked for
|
| 91 |
+
it. If a hook fails, investigate and fix the underlying issue.
|
| 92 |
+
- Avoid unnecessary `sleep` commands:
|
| 93 |
+
- Do not sleep between commands that can run immediately — just run
|
| 94 |
+
them.
|
| 95 |
+
- Do not retry failing commands in a sleep loop — diagnose the root
|
| 96 |
+
cause or consider an alternative approach.
|
| 97 |
+
- If you must sleep, keep the duration short (1-5 seconds) to avoid
|
| 98 |
+
blocking the user."""
|
| 99 |
+
"""The description presented to the agent."""
|
| 100 |
+
|
| 101 |
+
input_schema: dict[str, Any] = {
|
| 102 |
+
"type": "object",
|
| 103 |
+
"properties": {
|
| 104 |
+
"command": {
|
| 105 |
+
"type": "string",
|
| 106 |
+
"description": "The bash command to execute.",
|
| 107 |
+
},
|
| 108 |
+
"description": {
|
| 109 |
+
"type": "string",
|
| 110 |
+
"description": (
|
| 111 |
+
"Clear, concise description of what this command "
|
| 112 |
+
"does. For simple commands, keep it brief (5-10 "
|
| 113 |
+
"words). For complex commands, include enough "
|
| 114 |
+
"context."
|
| 115 |
+
),
|
| 116 |
+
},
|
| 117 |
+
"timeout": {
|
| 118 |
+
"type": "integer",
|
| 119 |
+
"description": (
|
| 120 |
+
"Optional timeout in milliseconds "
|
| 121 |
+
"(default: 120000, max: 600000)"
|
| 122 |
+
),
|
| 123 |
+
"default": 120000,
|
| 124 |
+
"maximum": 600000,
|
| 125 |
+
"minimum": 0,
|
| 126 |
+
},
|
| 127 |
+
},
|
| 128 |
+
"required": ["command"],
|
| 129 |
+
}
|
| 130 |
+
|
| 131 |
+
is_mcp: bool = False
|
| 132 |
+
is_read_only: bool = False
|
| 133 |
+
is_concurrency_safe: bool = False
|
| 134 |
+
is_external_tool: bool = False
|
| 135 |
+
is_state_injected: bool = False
|
| 136 |
+
|
| 137 |
+
def __init__( # pylint: disable=dangerous-default-value
|
| 138 |
+
self,
|
| 139 |
+
dangerous_files: list[str] = DEFAULT_DANGEROUS_FILES,
|
| 140 |
+
dangerous_directories: list[str] = DEFAULT_DANGEROUS_DIRECTORIES,
|
| 141 |
+
cwd: str | os.PathLike[str] | None = None,
|
| 142 |
+
middlewares: List[ToolMiddlewareBase] | None = None,
|
| 143 |
+
backend: BackendBase | None = None,
|
| 144 |
+
) -> None:
|
| 145 |
+
"""Initialize the bash tool.
|
| 146 |
+
|
| 147 |
+
Args:
|
| 148 |
+
dangerous_files (`list[str]`, optional):
|
| 149 |
+
Sensitive files that require explicit user confirmation,
|
| 150 |
+
even in BYPASS mode. Matched by basename
|
| 151 |
+
(case-insensitive). Defaults to `DEFAULT_DANGEROUS_FILES`.
|
| 152 |
+
Pass a custom list to fully replace the defaults, or `[]`
|
| 153 |
+
to disable the filename check.
|
| 154 |
+
dangerous_directories (`list[str]`, optional):
|
| 155 |
+
Sensitive directories that require explicit user
|
| 156 |
+
confirmation. Matched when any path segment equals an
|
| 157 |
+
entry (case-insensitive). Defaults to
|
| 158 |
+
`DEFAULT_DANGEROUS_DIRECTORIES`. Pass a custom list to
|
| 159 |
+
fully replace the defaults, or `[]` to disable the
|
| 160 |
+
directory check.
|
| 161 |
+
cwd (`str | os.PathLike[str] | None`, optional):
|
| 162 |
+
The working directory used when executing bash commands.
|
| 163 |
+
middlewares (`List[ToolMiddlewareBase] | None`, optional):
|
| 164 |
+
Tool middlewares wrapping the tool execution.
|
| 165 |
+
backend (`BackendBase | None`, optional):
|
| 166 |
+
The sandbox backend to use for shell execution. When
|
| 167 |
+
``None``, a :class:`LocalBackend` is created.
|
| 168 |
+
"""
|
| 169 |
+
from ._backend import LocalBackend
|
| 170 |
+
|
| 171 |
+
super().__init__(middlewares=middlewares)
|
| 172 |
+
self._bash_parser = BashCommandParser()
|
| 173 |
+
|
| 174 |
+
self.dangerous_files = list(dangerous_files)
|
| 175 |
+
self.dangerous_directories = list(dangerous_directories)
|
| 176 |
+
self._cwd = os.fspath(cwd) if cwd is not None else None
|
| 177 |
+
self._backend = backend or LocalBackend()
|
| 178 |
+
|
| 179 |
+
async def check_read_only(
|
| 180 |
+
self,
|
| 181 |
+
tool_input: dict[str, Any],
|
| 182 |
+
) -> bool:
|
| 183 |
+
"""Decide whether this specific bash invocation is read-only.
|
| 184 |
+
|
| 185 |
+
Inspects the command and returns ``True`` for known-safe read-only
|
| 186 |
+
commands (e.g. ``ls``, ``cat``, ``grep``, ``git status``). The
|
| 187 |
+
static :attr:`is_read_only` class attribute is ``False`` because
|
| 188 |
+
Bash can execute arbitrary commands; this method overrides that
|
| 189 |
+
with a per-invocation answer.
|
| 190 |
+
"""
|
| 191 |
+
command = tool_input.get("command", "")
|
| 192 |
+
if not command:
|
| 193 |
+
return self.is_read_only
|
| 194 |
+
return self._bash_parser.is_read_only_command(command)
|
| 195 |
+
|
| 196 |
+
async def check_permissions(
|
| 197 |
+
self,
|
| 198 |
+
tool_input: dict[str, Any],
|
| 199 |
+
context: PermissionContext,
|
| 200 |
+
) -> PermissionDecision:
|
| 201 |
+
"""Check permissions for bash command execution.
|
| 202 |
+
|
| 203 |
+
This method implements Bash-specific permission checks:
|
| 204 |
+
|
| 205 |
+
0. Injection risk check (bypass-immune safety ASK if command
|
| 206 |
+
contains dynamic expansion like ``$(...)`` or ``<(...)``)
|
| 207 |
+
1. Read-only command check — auto-ALLOW in **every mode**
|
| 208 |
+
(including DEFAULT) for known-safe read-only commands
|
| 209 |
+
(``ls``, ``pwd``, ``git status``, ``cat``, etc.). This is
|
| 210 |
+
the static counterpart to :meth:`check_read_only`.
|
| 211 |
+
2. Dangerous command pattern check (bypass-immune safety ASK)
|
| 212 |
+
3. Sed in-place constraint check (bypass-immune safety ASK)
|
| 213 |
+
4. Dangerous path check for config files (bypass-immune safety
|
| 214 |
+
ASK)
|
| 215 |
+
5. Dangerous removal path check for system dirs (bypass-immune
|
| 216 |
+
safety ASK)
|
| 217 |
+
6. ACCEPT_EDITS auto-allow for ``mkdir``/``touch``/``rm``/
|
| 218 |
+
``rmdir``/``mv``/``cp``/``sed`` — only when **every**
|
| 219 |
+
target path resolves inside a working directory
|
| 220 |
+
7. PASSTHROUGH (engine continues with rule matching)
|
| 221 |
+
|
| 222 |
+
"Bypass-immune" decisions set
|
| 223 |
+
:attr:`PermissionDecision.bypass_immune` so they cannot be
|
| 224 |
+
silenced by allow rules in DEFAULT mode. In BYPASS mode all
|
| 225 |
+
bypass-immune ASKs are intentionally skipped — see
|
| 226 |
+
:attr:`PermissionMode.BYPASS`.
|
| 227 |
+
|
| 228 |
+
Args:
|
| 229 |
+
tool_input (`dict[str, Any]`):
|
| 230 |
+
The tool input containing "command" key
|
| 231 |
+
context (`PermissionContext`):
|
| 232 |
+
The permission context with mode and rules
|
| 233 |
+
|
| 234 |
+
Returns:
|
| 235 |
+
`PermissionDecision`:
|
| 236 |
+
ALLOW for safe operations, ASK for dangerous operations,
|
| 237 |
+
PASSTHROUGH to let Engine continue with rule matching
|
| 238 |
+
"""
|
| 239 |
+
|
| 240 |
+
command = tool_input.get("command", "")
|
| 241 |
+
if not command:
|
| 242 |
+
return PermissionDecision(
|
| 243 |
+
behavior=PermissionBehavior.PASSTHROUGH,
|
| 244 |
+
message="Empty command",
|
| 245 |
+
)
|
| 246 |
+
|
| 247 |
+
# 0. Injection check: detect dynamic shell structures that cannot be
|
| 248 |
+
# statically analyzed (command substitution, process substitution,
|
| 249 |
+
# control flow, etc.). Must run before read-only check so that
|
| 250 |
+
# `$(rm -rf /)` inside an otherwise-safe command is caught.
|
| 251 |
+
injection_reason = self._bash_parser.check_injection_risk(command)
|
| 252 |
+
if injection_reason:
|
| 253 |
+
return PermissionDecision(
|
| 254 |
+
behavior=PermissionBehavior.ASK,
|
| 255 |
+
message=f"Permission required: {injection_reason}",
|
| 256 |
+
decision_reason="Safety check: command contains dynamic "
|
| 257 |
+
"expansion that cannot be statically analyzed",
|
| 258 |
+
bypass_immune=True,
|
| 259 |
+
)
|
| 260 |
+
|
| 261 |
+
# 1. Check if command is read-only (auto-allow)
|
| 262 |
+
if self._bash_parser.is_read_only_command(command):
|
| 263 |
+
return PermissionDecision(
|
| 264 |
+
behavior=PermissionBehavior.ALLOW,
|
| 265 |
+
message="Permission granted for read-only command",
|
| 266 |
+
decision_reason="Read-only command is allowed",
|
| 267 |
+
)
|
| 268 |
+
|
| 269 |
+
# 2. Check for dangerous commands (safety check, bypass-immune)
|
| 270 |
+
dangerous_pattern = self._bash_parser.check_dangerous_command(command)
|
| 271 |
+
if dangerous_pattern:
|
| 272 |
+
return PermissionDecision(
|
| 273 |
+
behavior=PermissionBehavior.ASK,
|
| 274 |
+
message=f"Permission required: Command contains dangerous "
|
| 275 |
+
f"pattern: {dangerous_pattern}",
|
| 276 |
+
decision_reason="Safety check: dangerous command pattern "
|
| 277 |
+
"detected",
|
| 278 |
+
bypass_immune=True,
|
| 279 |
+
)
|
| 280 |
+
|
| 281 |
+
# 3. Check for sed constraints (safety check, bypass-immune)
|
| 282 |
+
sed_error = self._bash_parser.check_sed_constraints(
|
| 283 |
+
command,
|
| 284 |
+
self.dangerous_files,
|
| 285 |
+
)
|
| 286 |
+
if sed_error:
|
| 287 |
+
return PermissionDecision(
|
| 288 |
+
behavior=PermissionBehavior.ASK,
|
| 289 |
+
message=f"Permission required: {sed_error}",
|
| 290 |
+
decision_reason="Safety check: sed in-place modification "
|
| 291 |
+
"of dangerous file",
|
| 292 |
+
bypass_immune=True,
|
| 293 |
+
)
|
| 294 |
+
|
| 295 |
+
# 4. Check for dangerous paths in sensitive config files/dirs
|
| 296 |
+
# (safety check, bypass-immune)
|
| 297 |
+
dangerous_paths = self._extract_dangerous_paths_from_bash(command)
|
| 298 |
+
if dangerous_paths:
|
| 299 |
+
paths_str = ", ".join(dangerous_paths)
|
| 300 |
+
return PermissionDecision(
|
| 301 |
+
behavior=PermissionBehavior.ASK,
|
| 302 |
+
message=f"Permission required: Bash command operates on "
|
| 303 |
+
f"sensitive paths: {paths_str}",
|
| 304 |
+
decision_reason="Safety check: dangerous file or "
|
| 305 |
+
"directory in bash command",
|
| 306 |
+
bypass_immune=True,
|
| 307 |
+
)
|
| 308 |
+
|
| 309 |
+
# 5. Check for dangerous removal paths: rm/rmdir targeting system
|
| 310 |
+
# critical directories like /, /usr, /etc, ~ (bypass-immune).
|
| 311 |
+
# Checked separately from step 4 because these paths are not in the
|
| 312 |
+
# dangerous_files/directories lists — they are system-level paths
|
| 313 |
+
# that should never be removed regardless of user configuration.
|
| 314 |
+
removal_path = await self._check_dangerous_removal_path(command)
|
| 315 |
+
if removal_path:
|
| 316 |
+
return PermissionDecision(
|
| 317 |
+
behavior=PermissionBehavior.ASK,
|
| 318 |
+
message=f"Dangerous removal operation detected: "
|
| 319 |
+
f"'{removal_path}'\n\nThis command would remove a critical "
|
| 320 |
+
f"system directory. This requires explicit approval and "
|
| 321 |
+
f"cannot be auto-allowed by permission rules.",
|
| 322 |
+
decision_reason="Safety check: dangerous removal of "
|
| 323 |
+
"critical system path",
|
| 324 |
+
bypass_immune=True,
|
| 325 |
+
)
|
| 326 |
+
|
| 327 |
+
# 6. ACCEPT_EDITS auto-allow for filesystem commands whose targets
|
| 328 |
+
# all live inside a working directory. Mirrors Write/Edit's strict
|
| 329 |
+
# working-directory check — we never auto-allow a bash command that
|
| 330 |
+
# would touch a path outside the configured working set (e.g.
|
| 331 |
+
# ``cp /etc/hosts /tmp/x`` must not pass even though ``cp`` is in
|
| 332 |
+
# the auto-allow list).
|
| 333 |
+
if context.mode == PermissionMode.ACCEPT_EDITS:
|
| 334 |
+
filesystem_commands = {
|
| 335 |
+
"mkdir",
|
| 336 |
+
"touch",
|
| 337 |
+
"rm",
|
| 338 |
+
"rmdir",
|
| 339 |
+
"mv",
|
| 340 |
+
"cp",
|
| 341 |
+
"sed",
|
| 342 |
+
}
|
| 343 |
+
base_command = (
|
| 344 |
+
command.strip().split()[0] if command.strip() else ""
|
| 345 |
+
)
|
| 346 |
+
|
| 347 |
+
if base_command in filesystem_commands:
|
| 348 |
+
# Collect every target path: file arguments AND output
|
| 349 |
+
# redirections. ``extract_file_paths`` includes both.
|
| 350 |
+
target_paths = [
|
| 351 |
+
path
|
| 352 |
+
for _cmd, path in self._bash_parser.extract_file_paths(
|
| 353 |
+
command,
|
| 354 |
+
)
|
| 355 |
+
]
|
| 356 |
+
# Conservative: only auto-allow when we extracted at least
|
| 357 |
+
# one target AND every target resolves inside a working
|
| 358 |
+
# directory. An empty list means the parser found nothing
|
| 359 |
+
# actionable (or the command has no args) — in that case
|
| 360 |
+
# we fall through to PASSTHROUGH rather than blindly
|
| 361 |
+
# allowing.
|
| 362 |
+
if target_paths and all(
|
| 363 |
+
self._path_in_allowed_working_path(path, context)
|
| 364 |
+
for path in target_paths
|
| 365 |
+
):
|
| 366 |
+
return PermissionDecision(
|
| 367 |
+
behavior=PermissionBehavior.ALLOW,
|
| 368 |
+
message=f"Permission granted for '{base_command}' "
|
| 369 |
+
f"command (accept edits mode - filesystem command, "
|
| 370 |
+
f"all targets in working directory)",
|
| 371 |
+
decision_reason=(
|
| 372 |
+
f"Filesystem command '{base_command}' is "
|
| 373 |
+
f"auto-allowed in accept edits mode because "
|
| 374 |
+
f"all target paths are within a working "
|
| 375 |
+
f"directory"
|
| 376 |
+
),
|
| 377 |
+
)
|
| 378 |
+
|
| 379 |
+
# 7. Passthrough to let Engine continue with rule matching
|
| 380 |
+
return PermissionDecision(
|
| 381 |
+
behavior=PermissionBehavior.PASSTHROUGH,
|
| 382 |
+
message=f"Execute bash command: {command}",
|
| 383 |
+
)
|
| 384 |
+
|
| 385 |
+
async def match_rule(
|
| 386 |
+
self,
|
| 387 |
+
rule_content: str | None,
|
| 388 |
+
tool_input: dict[str, Any],
|
| 389 |
+
) -> bool:
|
| 390 |
+
r"""Match Bash command using regex-based wildcard matching.
|
| 391 |
+
|
| 392 |
+
Implements wildcard matching with escape sequences:
|
| 393 |
+
- Supports \* for literal asterisk and \\ for literal backslash
|
| 394 |
+
- Special optimization: "git *" matches both "git" and "git add"
|
| 395 |
+
- Prefix pattern (e.g., "git:*"): matches commands starting with "git "
|
| 396 |
+
- Wildcard pattern: converts to regex with proper escape handling
|
| 397 |
+
- Substring pattern: exact substring matching
|
| 398 |
+
- If rule_content is None, matches all invocations
|
| 399 |
+
(tool-name-level rule)
|
| 400 |
+
|
| 401 |
+
Args:
|
| 402 |
+
rule_content: The command pattern to match, or None to match all
|
| 403 |
+
tool_input: Must contain a "command" key with the command string
|
| 404 |
+
|
| 405 |
+
Returns:
|
| 406 |
+
True if pattern matches the command
|
| 407 |
+
"""
|
| 408 |
+
# None = tool-name-level rule, matches everything
|
| 409 |
+
if rule_content is None:
|
| 410 |
+
return True
|
| 411 |
+
|
| 412 |
+
command = tool_input.get("command", "")
|
| 413 |
+
|
| 414 |
+
# Check if pattern is a prefix pattern (ends with :*)
|
| 415 |
+
if rule_content.endswith(":*"):
|
| 416 |
+
prefix = rule_content[:-2].strip()
|
| 417 |
+
return command.startswith(prefix + " ") or command == prefix
|
| 418 |
+
|
| 419 |
+
# Check if pattern has unescaped wildcards
|
| 420 |
+
def has_wildcards(pattern: str) -> bool:
|
| 421 |
+
"""Check if pattern contains unescaped * wildcards."""
|
| 422 |
+
i = 0
|
| 423 |
+
while i < len(pattern):
|
| 424 |
+
if pattern[i] == "\\":
|
| 425 |
+
i += 2 # Skip escaped character
|
| 426 |
+
elif pattern[i] == "*":
|
| 427 |
+
return True
|
| 428 |
+
else:
|
| 429 |
+
i += 1
|
| 430 |
+
return False
|
| 431 |
+
|
| 432 |
+
if not has_wildcards(rule_content):
|
| 433 |
+
# No wildcards, but may have escape sequences
|
| 434 |
+
# Convert escape sequences for matching
|
| 435 |
+
pattern = rule_content
|
| 436 |
+
pattern = pattern.replace("\\\\", "\x00BACKSLASH\x00")
|
| 437 |
+
pattern = pattern.replace("\\*", "*")
|
| 438 |
+
pattern = pattern.replace("\x00BACKSLASH\x00", "\\")
|
| 439 |
+
# Use substring matching with converted pattern
|
| 440 |
+
return pattern in command
|
| 441 |
+
|
| 442 |
+
# Convert wildcard pattern to regex with escape handling
|
| 443 |
+
# Use placeholders for escaped sequences
|
| 444 |
+
ESCAPED_STAR = "\x00ESCAPED_STAR\x00"
|
| 445 |
+
ESCAPED_BACKSLASH = "\x00ESCAPED_BACKSLASH\x00"
|
| 446 |
+
|
| 447 |
+
pattern = rule_content
|
| 448 |
+
# Replace \\ with placeholder
|
| 449 |
+
pattern = pattern.replace("\\\\", ESCAPED_BACKSLASH)
|
| 450 |
+
# Replace \* with placeholder
|
| 451 |
+
pattern = pattern.replace("\\*", ESCAPED_STAR)
|
| 452 |
+
|
| 453 |
+
# Manually escape regex special characters (except *)
|
| 454 |
+
# Don't use re.escape() as it escapes spaces too
|
| 455 |
+
special_chars = r".^$+?{}[]|()"
|
| 456 |
+
for char in special_chars:
|
| 457 |
+
pattern = pattern.replace(char, "\\" + char)
|
| 458 |
+
|
| 459 |
+
# Convert * to regex .* (match any characters)
|
| 460 |
+
pattern = pattern.replace("*", ".*")
|
| 461 |
+
|
| 462 |
+
# Restore escaped sequences
|
| 463 |
+
pattern = pattern.replace(ESCAPED_STAR, r"\*")
|
| 464 |
+
pattern = pattern.replace(ESCAPED_BACKSLASH, r"\\")
|
| 465 |
+
|
| 466 |
+
# Special optimization: "git *" should match both "git" and "git add"
|
| 467 |
+
# Pattern: if ends with .*, make it optional
|
| 468 |
+
if pattern.endswith(".*"):
|
| 469 |
+
base_pattern = pattern[:-2] # Remove .*
|
| 470 |
+
# Try exact match first (handles trailing space)
|
| 471 |
+
base_pattern = base_pattern.rstrip()
|
| 472 |
+
if re.fullmatch(base_pattern, command):
|
| 473 |
+
return True
|
| 474 |
+
|
| 475 |
+
# Full regex match
|
| 476 |
+
try:
|
| 477 |
+
return bool(re.fullmatch(pattern, command))
|
| 478 |
+
except re.error:
|
| 479 |
+
# Invalid regex, fall back to substring matching
|
| 480 |
+
return rule_content.replace("*", "") in command
|
| 481 |
+
|
| 482 |
+
async def generate_suggestions(
|
| 483 |
+
self,
|
| 484 |
+
tool_input: dict[str, Any],
|
| 485 |
+
) -> List["PermissionRule"]:
|
| 486 |
+
"""Generate suggested permission rules for Bash commands.
|
| 487 |
+
|
| 488 |
+
Generates prefix rules based on command + subcommand (two words).
|
| 489 |
+
For example, "git commit -m 'xxx'" generates "git commit:*".
|
| 490 |
+
|
| 491 |
+
Args:
|
| 492 |
+
tool_input (`dict[str, Any]`):
|
| 493 |
+
The tool input data containing "command" key
|
| 494 |
+
|
| 495 |
+
Returns:
|
| 496 |
+
`List[PermissionRule]`:
|
| 497 |
+
List of suggested permission rules based on command prefixes
|
| 498 |
+
"""
|
| 499 |
+
|
| 500 |
+
command = tool_input.get("command", "")
|
| 501 |
+
if not command:
|
| 502 |
+
return []
|
| 503 |
+
|
| 504 |
+
# Use bash parser to extract command prefixes
|
| 505 |
+
prefixes = self._bash_parser.extract_command_prefixes(
|
| 506 |
+
command,
|
| 507 |
+
max_prefixes=5,
|
| 508 |
+
)
|
| 509 |
+
|
| 510 |
+
if not prefixes:
|
| 511 |
+
# Cannot extract any prefix, return empty
|
| 512 |
+
return []
|
| 513 |
+
|
| 514 |
+
# Generate rules for each prefix
|
| 515 |
+
rules = []
|
| 516 |
+
for prefix in prefixes:
|
| 517 |
+
rules.append(
|
| 518 |
+
PermissionRule(
|
| 519 |
+
tool_name="Bash",
|
| 520 |
+
rule_content=f"{prefix}:*",
|
| 521 |
+
behavior=PermissionBehavior.ALLOW,
|
| 522 |
+
source="suggested",
|
| 523 |
+
),
|
| 524 |
+
)
|
| 525 |
+
|
| 526 |
+
return rules
|
| 527 |
+
|
| 528 |
+
def _extract_dangerous_paths_from_bash(
|
| 529 |
+
self,
|
| 530 |
+
command: str,
|
| 531 |
+
) -> list[str]:
|
| 532 |
+
"""Extract dangerous paths from a bash command using tree-sitter.
|
| 533 |
+
|
| 534 |
+
Checks for dangerous paths in:
|
| 535 |
+
- File-manipulating commands (rm, mv, cp, chmod, chown, sed, touch)
|
| 536 |
+
- Output redirections (>, >>)
|
| 537 |
+
|
| 538 |
+
Args:
|
| 539 |
+
command (`str`):
|
| 540 |
+
The bash command string
|
| 541 |
+
|
| 542 |
+
Returns:
|
| 543 |
+
`list[str]`:
|
| 544 |
+
List of dangerous paths found in the command
|
| 545 |
+
"""
|
| 546 |
+
dangerous_paths = []
|
| 547 |
+
|
| 548 |
+
# Use tree-sitter to extract file paths
|
| 549 |
+
file_paths = self._bash_parser.extract_file_paths(command)
|
| 550 |
+
|
| 551 |
+
for _cmd_name, path in file_paths:
|
| 552 |
+
if self._is_dangerous_path(path):
|
| 553 |
+
dangerous_paths.append(path)
|
| 554 |
+
|
| 555 |
+
return dangerous_paths
|
| 556 |
+
|
| 557 |
+
async def _check_dangerous_removal_path(self, command: str) -> str | None:
|
| 558 |
+
"""Check if a rm/rmdir command targets a critical system path.
|
| 559 |
+
|
| 560 |
+
Detects commands like `rm -rf /`, `rm -rf /usr`, `rmdir ~` that
|
| 561 |
+
would destroy critical system directories. Unlike _is_dangerous_path
|
| 562 |
+
(which checks against a configurable list of sensitive config files),
|
| 563 |
+
this checks against a fixed set of system-level paths that must
|
| 564 |
+
never be removed regardless of user configuration.
|
| 565 |
+
|
| 566 |
+
Dangerous paths are:
|
| 567 |
+
- Root directory (/)
|
| 568 |
+
- Home directory (~)
|
| 569 |
+
- Wildcard alone (*) or as dir/* (removes everything)
|
| 570 |
+
- Direct children of root (/usr, /etc, /tmp, /var, etc.)
|
| 571 |
+
|
| 572 |
+
Args:
|
| 573 |
+
command (`str`):
|
| 574 |
+
The bash command string
|
| 575 |
+
|
| 576 |
+
Returns:
|
| 577 |
+
`str | None`:
|
| 578 |
+
The dangerous path if found, None otherwise
|
| 579 |
+
"""
|
| 580 |
+
tokens = command.strip().split()
|
| 581 |
+
if not tokens:
|
| 582 |
+
return None
|
| 583 |
+
|
| 584 |
+
# Find rm or rmdir subcommands (handle compound commands)
|
| 585 |
+
try:
|
| 586 |
+
tree = self._bash_parser.parser.parse(bytes(command, "utf8"))
|
| 587 |
+
subcommands = self._bash_parser.split_compound_command(
|
| 588 |
+
tree.root_node,
|
| 589 |
+
command,
|
| 590 |
+
)
|
| 591 |
+
except Exception:
|
| 592 |
+
subcommands = [command]
|
| 593 |
+
|
| 594 |
+
# Check each subcommand for rm/rmdir
|
| 595 |
+
for subcmd in subcommands:
|
| 596 |
+
subcmd_tokens = subcmd.strip().split()
|
| 597 |
+
if not subcmd_tokens:
|
| 598 |
+
continue
|
| 599 |
+
base = subcmd_tokens[0]
|
| 600 |
+
if base not in ("rm", "rmdir"):
|
| 601 |
+
continue
|
| 602 |
+
|
| 603 |
+
# Collect non-flag arguments as potential paths
|
| 604 |
+
i = 1
|
| 605 |
+
while i < len(subcmd_tokens):
|
| 606 |
+
tok = subcmd_tokens[i]
|
| 607 |
+
# Skip flags
|
| 608 |
+
if tok.startswith("-"):
|
| 609 |
+
i += 1
|
| 610 |
+
continue
|
| 611 |
+
path = tok.strip("'\"")
|
| 612 |
+
if await self._is_dangerous_removal_path(path):
|
| 613 |
+
return path
|
| 614 |
+
i += 1
|
| 615 |
+
|
| 616 |
+
return None
|
| 617 |
+
|
| 618 |
+
async def _is_dangerous_removal_path(self, path: str) -> bool:
|
| 619 |
+
"""Check if a path is a critical system directory that must not be
|
| 620 |
+
removed.
|
| 621 |
+
|
| 622 |
+
All path resolution is performed via the backend so that the
|
| 623 |
+
check operates on the **backend environment's** ``$HOME`` /
|
| 624 |
+
``cwd`` / path semantics, not the host process's.
|
| 625 |
+
|
| 626 |
+
Args:
|
| 627 |
+
path (`str`):
|
| 628 |
+
The path to check (may be relative, absolute, or contain ~)
|
| 629 |
+
|
| 630 |
+
Returns:
|
| 631 |
+
`bool`:
|
| 632 |
+
True if removing this path would be catastrophic
|
| 633 |
+
"""
|
| 634 |
+
|
| 635 |
+
# Bare wildcard
|
| 636 |
+
if path in ("*", "./*", "/"):
|
| 637 |
+
return True
|
| 638 |
+
# Ends with /* — removes everything in a directory
|
| 639 |
+
if path.endswith("/*") or path.endswith("\\*"):
|
| 640 |
+
return True
|
| 641 |
+
|
| 642 |
+
# Expand tilde and resolve to an absolute path inside the
|
| 643 |
+
# backend environment. Don't resolve symlinks — ``/tmp`` is a
|
| 644 |
+
# symlink on macOS but is still a root-child and should be
|
| 645 |
+
# flagged.
|
| 646 |
+
expanded = await self._backend.expanduser(path)
|
| 647 |
+
backend_cwd = await self._backend.getcwd()
|
| 648 |
+
abs_path = self._backend.abspath(expanded, cwd=backend_cwd)
|
| 649 |
+
|
| 650 |
+
# Home directory
|
| 651 |
+
home = await self._backend.expanduser("~")
|
| 652 |
+
if abs_path == home:
|
| 653 |
+
return True
|
| 654 |
+
|
| 655 |
+
# Root itself: ``dirname(root) == root`` on both POSIX
|
| 656 |
+
# (``"/"``) and Windows (``"C:\\"``), so this check is
|
| 657 |
+
# path-flavor agnostic.
|
| 658 |
+
parent = self._backend.dirname(abs_path)
|
| 659 |
+
if abs_path == parent:
|
| 660 |
+
return True
|
| 661 |
+
|
| 662 |
+
# Direct children of root (e.g. ``/usr``, ``/etc``, ``/tmp``):
|
| 663 |
+
# the *parent* of these is the root, where
|
| 664 |
+
# ``dirname(parent) == parent``.
|
| 665 |
+
if self._backend.dirname(parent) == parent:
|
| 666 |
+
return True
|
| 667 |
+
|
| 668 |
+
return False
|
| 669 |
+
|
| 670 |
+
async def call( # type: ignore[override] # pylint: disable=unused-argument
|
| 671 |
+
self,
|
| 672 |
+
command: str,
|
| 673 |
+
description: str = "",
|
| 674 |
+
timeout: int = 120000,
|
| 675 |
+
) -> AsyncGenerator[ToolChunk, None]:
|
| 676 |
+
"""Execute the bash and return the output.
|
| 677 |
+
|
| 678 |
+
Args:
|
| 679 |
+
command: The bash command to execute.
|
| 680 |
+
description: Optional description of what the command does.
|
| 681 |
+
timeout: Timeout in milliseconds (default: 120000, max: 600000).
|
| 682 |
+
|
| 683 |
+
Yields:
|
| 684 |
+
ToolChunk: The tool execution result with stdout/stderr content.
|
| 685 |
+
"""
|
| 686 |
+
|
| 687 |
+
# Clamp timeout to max 600000ms and convert to seconds
|
| 688 |
+
timeout_ms = min(timeout, 600000)
|
| 689 |
+
timeout_sec = timeout_ms / 1000.0
|
| 690 |
+
|
| 691 |
+
try:
|
| 692 |
+
# ``command`` is a full shell command line (it may contain
|
| 693 |
+
# pipes, redirects, ``&&``, …), so wrap it in a shell — the
|
| 694 |
+
# backend primitive runs the argv directly without one. Pick
|
| 695 |
+
# the platform's native shell so the Windows experience that
|
| 696 |
+
# ``main`` had (commands interpreted by ``cmd.exe``) is
|
| 697 |
+
# preserved; POSIX hosts use ``/bin/sh``.
|
| 698 |
+
if os.name == "nt":
|
| 699 |
+
shell_command = ["cmd", "/c", command]
|
| 700 |
+
else:
|
| 701 |
+
shell_command = ["/bin/sh", "-c", command]
|
| 702 |
+
result = await self._backend.exec_shell(
|
| 703 |
+
shell_command,
|
| 704 |
+
cwd=self._cwd,
|
| 705 |
+
timeout=timeout_sec,
|
| 706 |
+
)
|
| 707 |
+
|
| 708 |
+
# Decode and normalize line endings
|
| 709 |
+
stdout = result.stdout.decode(
|
| 710 |
+
"utf-8",
|
| 711 |
+
errors="replace",
|
| 712 |
+
).replace("\r\n", "\n")
|
| 713 |
+
stderr = result.stderr.decode(
|
| 714 |
+
"utf-8",
|
| 715 |
+
errors="replace",
|
| 716 |
+
).replace("\r\n", "\n")
|
| 717 |
+
|
| 718 |
+
# Check for timeout (backend returns exit_code=-1,
|
| 719 |
+
# stderr=b"timed out")
|
| 720 |
+
if result.exit_code == -1 and result.stderr == b"timed out":
|
| 721 |
+
error_msg = (
|
| 722 |
+
f"Command timed out after {timeout_ms}ms: {command}"
|
| 723 |
+
)
|
| 724 |
+
yield ToolChunk(
|
| 725 |
+
content=[TextBlock(text=error_msg)],
|
| 726 |
+
state=ToolResultState.ERROR,
|
| 727 |
+
is_last=True,
|
| 728 |
+
)
|
| 729 |
+
return
|
| 730 |
+
|
| 731 |
+
# Combine output
|
| 732 |
+
output = stdout
|
| 733 |
+
if stderr:
|
| 734 |
+
if output:
|
| 735 |
+
output += "\n"
|
| 736 |
+
output += stderr
|
| 737 |
+
|
| 738 |
+
# Truncate if exceeds 30000 characters
|
| 739 |
+
if len(output) > 30000:
|
| 740 |
+
output = output[:30000] + "\n... (output truncated)"
|
| 741 |
+
|
| 742 |
+
# Check exit code
|
| 743 |
+
if not result.ok():
|
| 744 |
+
# Command failed
|
| 745 |
+
error_result = f"Command failed: {command}\n"
|
| 746 |
+
if stdout:
|
| 747 |
+
error_result += f"\nStdout:\n{stdout}"
|
| 748 |
+
if stderr:
|
| 749 |
+
error_result += f"\nStderr:\n{stderr}"
|
| 750 |
+
|
| 751 |
+
# Truncate error message if needed
|
| 752 |
+
if len(error_result) > 30000:
|
| 753 |
+
error_result = (
|
| 754 |
+
error_result[:30000] + "\n... (output truncated)"
|
| 755 |
+
)
|
| 756 |
+
|
| 757 |
+
yield ToolChunk(
|
| 758 |
+
content=[TextBlock(text=error_result)],
|
| 759 |
+
state=ToolResultState.ERROR,
|
| 760 |
+
is_last=True,
|
| 761 |
+
)
|
| 762 |
+
else:
|
| 763 |
+
# Command succeeded - note: ToolChunk uses "running" state
|
| 764 |
+
# which will be converted to "finished" in ToolResponse
|
| 765 |
+
yield ToolChunk(
|
| 766 |
+
content=[TextBlock(text=output)],
|
| 767 |
+
state=ToolResultState.RUNNING,
|
| 768 |
+
is_last=True,
|
| 769 |
+
)
|
| 770 |
+
|
| 771 |
+
except Exception as e:
|
| 772 |
+
# Other errors
|
| 773 |
+
error_msg = f"Command failed: {command}\nError: {str(e)}"
|
| 774 |
+
yield ToolChunk(
|
| 775 |
+
content=[TextBlock(text=error_msg)],
|
| 776 |
+
state=ToolResultState.ERROR,
|
| 777 |
+
is_last=True,
|
| 778 |
+
)
|
src/agentscope/tool/_builtin/_bash_parser.py
ADDED
|
@@ -0,0 +1,884 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""Bash command parser using tree-sitter for precise syntax analysis.
|
| 3 |
+
|
| 4 |
+
This module provides utilities to parse Bash commands and extract meaningful
|
| 5 |
+
information for permission rule generation, including:
|
| 6 |
+
- Splitting compound commands (&&, ||, ;, |)
|
| 7 |
+
- Extracting command prefixes (e.g., "npm run" from "npm run build")
|
| 8 |
+
- Extracting file paths from commands for dangerous path detection
|
| 9 |
+
- Extracting output redirections
|
| 10 |
+
- Checking if commands are read-only
|
| 11 |
+
"""
|
| 12 |
+
|
| 13 |
+
from typing import List, Optional, Set, Tuple
|
| 14 |
+
|
| 15 |
+
import re
|
| 16 |
+
import shlex
|
| 17 |
+
|
| 18 |
+
import tree_sitter_bash as tsbash
|
| 19 |
+
from tree_sitter import Language, Parser, Node
|
| 20 |
+
|
| 21 |
+
from .._constants import DANGEROUS_NODE_TYPES, DANGEROUS_COMMANDS
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
# Commands that are considered safe and don't require permission rules
|
| 25 |
+
SAFE_COMMANDS: Set[str] = {
|
| 26 |
+
"echo",
|
| 27 |
+
"cat",
|
| 28 |
+
"ls",
|
| 29 |
+
"pwd",
|
| 30 |
+
"cd",
|
| 31 |
+
"true",
|
| 32 |
+
"false",
|
| 33 |
+
"printf",
|
| 34 |
+
"grep",
|
| 35 |
+
"tee",
|
| 36 |
+
}
|
| 37 |
+
|
| 38 |
+
# Safe environment variables that can be skipped when extracting command prefix
|
| 39 |
+
SAFE_ENV_VARS = {
|
| 40 |
+
"NODE_ENV",
|
| 41 |
+
"PYTHONUNBUFFERED",
|
| 42 |
+
"RUST_LOG",
|
| 43 |
+
"LANG",
|
| 44 |
+
"TERM",
|
| 45 |
+
"NO_COLOR",
|
| 46 |
+
"FORCE_COLOR",
|
| 47 |
+
"DEBUG",
|
| 48 |
+
"VERBOSE",
|
| 49 |
+
"CI",
|
| 50 |
+
"PATH",
|
| 51 |
+
"HOME",
|
| 52 |
+
"USER",
|
| 53 |
+
"SHELL",
|
| 54 |
+
"EDITOR",
|
| 55 |
+
"PAGER",
|
| 56 |
+
"TZ",
|
| 57 |
+
"LC_ALL",
|
| 58 |
+
"LC_CTYPE",
|
| 59 |
+
"COLUMNS",
|
| 60 |
+
"LINES",
|
| 61 |
+
"CLICOLOR",
|
| 62 |
+
"CLICOLOR_FORCE",
|
| 63 |
+
}
|
| 64 |
+
|
| 65 |
+
# Read-only git commands
|
| 66 |
+
GIT_READ_ONLY_COMMANDS = {
|
| 67 |
+
"git status",
|
| 68 |
+
"git log",
|
| 69 |
+
"git diff",
|
| 70 |
+
"git show",
|
| 71 |
+
"git branch",
|
| 72 |
+
"git tag",
|
| 73 |
+
"git remote",
|
| 74 |
+
"git ls-files",
|
| 75 |
+
"git ls-tree",
|
| 76 |
+
"git cat-file",
|
| 77 |
+
"git rev-parse",
|
| 78 |
+
"git rev-list",
|
| 79 |
+
"git describe",
|
| 80 |
+
"git shortlog",
|
| 81 |
+
"git blame",
|
| 82 |
+
"git grep",
|
| 83 |
+
"git reflog",
|
| 84 |
+
"git config --get",
|
| 85 |
+
"git config --list",
|
| 86 |
+
}
|
| 87 |
+
|
| 88 |
+
# Read-only commands for various tools
|
| 89 |
+
READ_ONLY_COMMANDS = {
|
| 90 |
+
# Basic file operations
|
| 91 |
+
"ls",
|
| 92 |
+
"cat",
|
| 93 |
+
"head",
|
| 94 |
+
"tail",
|
| 95 |
+
"less",
|
| 96 |
+
"more",
|
| 97 |
+
"file",
|
| 98 |
+
"stat",
|
| 99 |
+
"wc",
|
| 100 |
+
"grep",
|
| 101 |
+
"rg",
|
| 102 |
+
"ag",
|
| 103 |
+
"ack",
|
| 104 |
+
"find",
|
| 105 |
+
"tree",
|
| 106 |
+
"pwd",
|
| 107 |
+
"which",
|
| 108 |
+
"whereis",
|
| 109 |
+
"type",
|
| 110 |
+
# Git commands
|
| 111 |
+
*GIT_READ_ONLY_COMMANDS,
|
| 112 |
+
# Docker read-only
|
| 113 |
+
"docker ps",
|
| 114 |
+
"docker images",
|
| 115 |
+
"docker inspect",
|
| 116 |
+
"docker logs",
|
| 117 |
+
"docker version",
|
| 118 |
+
"docker info",
|
| 119 |
+
# GitHub CLI read-only
|
| 120 |
+
"gh repo view",
|
| 121 |
+
"gh issue list",
|
| 122 |
+
"gh pr list",
|
| 123 |
+
"gh status",
|
| 124 |
+
# Python/Node tools
|
| 125 |
+
"python --version",
|
| 126 |
+
"python -V",
|
| 127 |
+
"node --version",
|
| 128 |
+
"node -v",
|
| 129 |
+
"npm list",
|
| 130 |
+
"npm ls",
|
| 131 |
+
"pip list",
|
| 132 |
+
"pip show",
|
| 133 |
+
}
|
| 134 |
+
|
| 135 |
+
|
| 136 |
+
class BashCommandParser:
|
| 137 |
+
"""Parse Bash commands using tree-sitter for accurate syntax analysis."""
|
| 138 |
+
|
| 139 |
+
def __init__(self) -> None:
|
| 140 |
+
"""Initialize the parser with tree-sitter-bash language."""
|
| 141 |
+
self.parser = Parser(Language(tsbash.language()))
|
| 142 |
+
|
| 143 |
+
def is_read_only_command(self, command: str) -> bool:
|
| 144 |
+
"""Check if a command is read-only (safe to auto-allow).
|
| 145 |
+
|
| 146 |
+
For compound commands (&&, ||, ;, |), ALL subcommands must be
|
| 147 |
+
read-only for the entire command to be considered read-only.
|
| 148 |
+
|
| 149 |
+
Commands with output redirections (>, >>) are NOT considered read-only.
|
| 150 |
+
|
| 151 |
+
Args:
|
| 152 |
+
command (`str`):
|
| 153 |
+
The bash command string
|
| 154 |
+
|
| 155 |
+
Returns:
|
| 156 |
+
`bool`:
|
| 157 |
+
True if the command (and all subcommands) are read-only,
|
| 158 |
+
False otherwise
|
| 159 |
+
"""
|
| 160 |
+
# Normalize command (strip leading/trailing whitespace)
|
| 161 |
+
cmd = command.strip()
|
| 162 |
+
|
| 163 |
+
# Check for output redirections - these are NOT read-only
|
| 164 |
+
if ">" in cmd:
|
| 165 |
+
return False
|
| 166 |
+
|
| 167 |
+
# Check if it's a compound command
|
| 168 |
+
if any(op in cmd for op in ["&&", "||", ";", "|"]):
|
| 169 |
+
# Split into subcommands and check each one
|
| 170 |
+
try:
|
| 171 |
+
tree = self.parser.parse(bytes(cmd, "utf8"))
|
| 172 |
+
root = tree.root_node
|
| 173 |
+
subcommands = self.split_compound_command(root, cmd)
|
| 174 |
+
|
| 175 |
+
# All subcommands must be read-only
|
| 176 |
+
for subcmd in subcommands:
|
| 177 |
+
if not self._is_single_command_read_only(subcmd.strip()):
|
| 178 |
+
return False
|
| 179 |
+
return True
|
| 180 |
+
except Exception:
|
| 181 |
+
# If parsing fails, be conservative
|
| 182 |
+
return False
|
| 183 |
+
|
| 184 |
+
# Single command - check directly
|
| 185 |
+
return self._is_single_command_read_only(cmd)
|
| 186 |
+
|
| 187 |
+
def _is_single_command_read_only(self, cmd: str) -> bool:
|
| 188 |
+
"""Check if a single (non-compound) command is read-only.
|
| 189 |
+
|
| 190 |
+
Args:
|
| 191 |
+
cmd (`str`):
|
| 192 |
+
A single command string (no &&, ||, ;, |)
|
| 193 |
+
|
| 194 |
+
Returns:
|
| 195 |
+
`bool`:
|
| 196 |
+
True if the command is read-only, False otherwise
|
| 197 |
+
"""
|
| 198 |
+
# Check exact match in read-only commands
|
| 199 |
+
if cmd in READ_ONLY_COMMANDS:
|
| 200 |
+
return True
|
| 201 |
+
|
| 202 |
+
# Check if it starts with a read-only prefix
|
| 203 |
+
for readonly_cmd in READ_ONLY_COMMANDS:
|
| 204 |
+
if cmd == readonly_cmd or cmd.startswith(readonly_cmd + " "):
|
| 205 |
+
return True
|
| 206 |
+
|
| 207 |
+
# Check base command for simple read-only operations
|
| 208 |
+
tokens = cmd.split()
|
| 209 |
+
if tokens:
|
| 210 |
+
base_cmd = tokens[0]
|
| 211 |
+
# Skip environment variables
|
| 212 |
+
i = 0
|
| 213 |
+
while i < len(tokens) and "=" in tokens[i]:
|
| 214 |
+
i += 1
|
| 215 |
+
if i < len(tokens):
|
| 216 |
+
base_cmd = tokens[i]
|
| 217 |
+
|
| 218 |
+
# Check if base command is in safe commands
|
| 219 |
+
if base_cmd in SAFE_COMMANDS:
|
| 220 |
+
return True
|
| 221 |
+
|
| 222 |
+
return False
|
| 223 |
+
|
| 224 |
+
def extract_file_paths(
|
| 225 |
+
self,
|
| 226 |
+
command: str,
|
| 227 |
+
) -> List[Tuple[str, str]]:
|
| 228 |
+
"""Extract file paths from a bash command using tree-sitter.
|
| 229 |
+
|
| 230 |
+
Returns paths that are arguments to file-manipulating commands
|
| 231 |
+
(rm, mv, cp, chmod, chown, etc.) and output redirection targets.
|
| 232 |
+
|
| 233 |
+
Args:
|
| 234 |
+
command (`str`):
|
| 235 |
+
The bash command string
|
| 236 |
+
|
| 237 |
+
Returns:
|
| 238 |
+
`List[Tuple[str, str]]`:
|
| 239 |
+
List of tuples (command_name, file_path)
|
| 240 |
+
"""
|
| 241 |
+
paths = []
|
| 242 |
+
|
| 243 |
+
try:
|
| 244 |
+
# Parse command to AST
|
| 245 |
+
tree = self.parser.parse(bytes(command, "utf8"))
|
| 246 |
+
root = tree.root_node
|
| 247 |
+
|
| 248 |
+
# Extract paths from commands
|
| 249 |
+
self._extract_paths_from_node(root, command, paths)
|
| 250 |
+
|
| 251 |
+
except Exception:
|
| 252 |
+
# Fallback to simple token-based extraction
|
| 253 |
+
paths = self._extract_paths_fallback(command)
|
| 254 |
+
|
| 255 |
+
return paths
|
| 256 |
+
|
| 257 |
+
def _extract_paths_from_node(
|
| 258 |
+
self,
|
| 259 |
+
node: Node,
|
| 260 |
+
command: str,
|
| 261 |
+
paths: List[Tuple[str, str]],
|
| 262 |
+
) -> None:
|
| 263 |
+
"""Recursively extract file paths from AST nodes.
|
| 264 |
+
|
| 265 |
+
Args:
|
| 266 |
+
node (`Node`):
|
| 267 |
+
The AST node to process
|
| 268 |
+
command (`str`):
|
| 269 |
+
The original command string
|
| 270 |
+
paths (`List[Tuple[str, str]]`):
|
| 271 |
+
List to append (command_name, path) tuples to
|
| 272 |
+
"""
|
| 273 |
+
# Check for redirections
|
| 274 |
+
if node.type == "file_redirect":
|
| 275 |
+
# Extract the target file
|
| 276 |
+
for child in node.children:
|
| 277 |
+
if child.type == "word":
|
| 278 |
+
path = command[child.start_byte : child.end_byte]
|
| 279 |
+
paths.append(("redirect", path.strip("'\"")))
|
| 280 |
+
# Check for commands
|
| 281 |
+
if node.type == "command":
|
| 282 |
+
# Extract command name and arguments
|
| 283 |
+
cmd_name = None
|
| 284 |
+
args = []
|
| 285 |
+
|
| 286 |
+
for child in node.children:
|
| 287 |
+
if child.type == "command_name":
|
| 288 |
+
cmd_name = command[child.start_byte : child.end_byte]
|
| 289 |
+
elif child.type == "word" and cmd_name:
|
| 290 |
+
arg = command[child.start_byte : child.end_byte]
|
| 291 |
+
args.append(arg.strip("'\""))
|
| 292 |
+
|
| 293 |
+
# Check if this is a file-manipulating command
|
| 294 |
+
if cmd_name in [
|
| 295 |
+
"rm",
|
| 296 |
+
"mv",
|
| 297 |
+
"cp",
|
| 298 |
+
"chmod",
|
| 299 |
+
"chown",
|
| 300 |
+
"chgrp",
|
| 301 |
+
"touch",
|
| 302 |
+
"ln",
|
| 303 |
+
"sed",
|
| 304 |
+
"mkdir",
|
| 305 |
+
"rmdir",
|
| 306 |
+
]:
|
| 307 |
+
# Extract file arguments (skip flags)
|
| 308 |
+
for arg in args:
|
| 309 |
+
if not arg.startswith("-"):
|
| 310 |
+
paths.append((cmd_name, arg))
|
| 311 |
+
|
| 312 |
+
# Recursively process children
|
| 313 |
+
for child in node.children:
|
| 314 |
+
self._extract_paths_from_node(child, command, paths)
|
| 315 |
+
|
| 316 |
+
def _extract_paths_fallback(
|
| 317 |
+
self,
|
| 318 |
+
command: str,
|
| 319 |
+
) -> List[Tuple[str, str]]:
|
| 320 |
+
"""Fallback path extraction using simple token parsing.
|
| 321 |
+
|
| 322 |
+
Args:
|
| 323 |
+
command (`str`):
|
| 324 |
+
The bash command string
|
| 325 |
+
|
| 326 |
+
Returns:
|
| 327 |
+
`List[Tuple[str, str]]`:
|
| 328 |
+
List of tuples (command_name, file_path)
|
| 329 |
+
"""
|
| 330 |
+
paths = []
|
| 331 |
+
tokens = command.split()
|
| 332 |
+
i = 0
|
| 333 |
+
|
| 334 |
+
while i < len(tokens):
|
| 335 |
+
token = tokens[i]
|
| 336 |
+
|
| 337 |
+
# Check for output redirections
|
| 338 |
+
if token in [">", ">>", "2>", "&>"]:
|
| 339 |
+
if i + 1 < len(tokens):
|
| 340 |
+
path = tokens[i + 1].strip("'\"")
|
| 341 |
+
paths.append(("redirect", path))
|
| 342 |
+
i += 2
|
| 343 |
+
continue
|
| 344 |
+
|
| 345 |
+
# Check for file-manipulating commands
|
| 346 |
+
if token in [
|
| 347 |
+
"rm",
|
| 348 |
+
"mv",
|
| 349 |
+
"cp",
|
| 350 |
+
"chmod",
|
| 351 |
+
"chown",
|
| 352 |
+
"sed",
|
| 353 |
+
"touch",
|
| 354 |
+
"mkdir",
|
| 355 |
+
"rmdir",
|
| 356 |
+
]:
|
| 357 |
+
cmd_name = token
|
| 358 |
+
# Look for file arguments after this command
|
| 359 |
+
j = i + 1
|
| 360 |
+
while j < len(tokens):
|
| 361 |
+
arg = tokens[j].strip("'\"")
|
| 362 |
+
# Skip flags
|
| 363 |
+
if arg.startswith("-"):
|
| 364 |
+
j += 1
|
| 365 |
+
continue
|
| 366 |
+
# This is a file argument
|
| 367 |
+
paths.append((cmd_name, arg))
|
| 368 |
+
j += 1
|
| 369 |
+
break
|
| 370 |
+
|
| 371 |
+
i += 1
|
| 372 |
+
|
| 373 |
+
return paths
|
| 374 |
+
|
| 375 |
+
def extract_redirections(self, command: str) -> List[str]:
|
| 376 |
+
"""Extract output redirection targets from a bash command.
|
| 377 |
+
|
| 378 |
+
Args:
|
| 379 |
+
command (`str`):
|
| 380 |
+
The bash command string
|
| 381 |
+
|
| 382 |
+
Returns:
|
| 383 |
+
`List[str]`:
|
| 384 |
+
List of file paths that are redirection targets
|
| 385 |
+
"""
|
| 386 |
+
redirections = []
|
| 387 |
+
|
| 388 |
+
try:
|
| 389 |
+
# Parse command to AST
|
| 390 |
+
tree = self.parser.parse(bytes(command, "utf8"))
|
| 391 |
+
root = tree.root_node
|
| 392 |
+
|
| 393 |
+
# Extract redirections
|
| 394 |
+
self._extract_redirections_from_node(root, command, redirections)
|
| 395 |
+
|
| 396 |
+
except Exception:
|
| 397 |
+
# Fallback to simple extraction
|
| 398 |
+
tokens = command.split()
|
| 399 |
+
for i, token in enumerate(tokens):
|
| 400 |
+
if token in [">", ">>", "2>", "&>"] and i + 1 < len(tokens):
|
| 401 |
+
path = tokens[i + 1].strip("'\"")
|
| 402 |
+
redirections.append(path)
|
| 403 |
+
|
| 404 |
+
return redirections
|
| 405 |
+
|
| 406 |
+
def _extract_redirections_from_node(
|
| 407 |
+
self,
|
| 408 |
+
node: Node,
|
| 409 |
+
command: str,
|
| 410 |
+
redirections: List[str],
|
| 411 |
+
) -> None:
|
| 412 |
+
"""Recursively extract redirections from AST nodes.
|
| 413 |
+
|
| 414 |
+
Args:
|
| 415 |
+
node (`Node`):
|
| 416 |
+
The AST node to process
|
| 417 |
+
command (`str`):
|
| 418 |
+
The original command string
|
| 419 |
+
redirections (`List[str]`):
|
| 420 |
+
List to append redirection targets to
|
| 421 |
+
"""
|
| 422 |
+
if node.type == "file_redirect":
|
| 423 |
+
# Extract the target file
|
| 424 |
+
for child in node.children:
|
| 425 |
+
if child.type == "word":
|
| 426 |
+
path = command[child.start_byte : child.end_byte]
|
| 427 |
+
redirections.append(path.strip("'\""))
|
| 428 |
+
|
| 429 |
+
# Recursively process children
|
| 430 |
+
for child in node.children:
|
| 431 |
+
self._extract_redirections_from_node(child, command, redirections)
|
| 432 |
+
|
| 433 |
+
def extract_command_prefixes(
|
| 434 |
+
self,
|
| 435 |
+
command: str,
|
| 436 |
+
max_prefixes: int = 5,
|
| 437 |
+
) -> List[str]:
|
| 438 |
+
"""Extract command prefixes from a bash command.
|
| 439 |
+
|
| 440 |
+
Automatically handles compound commands (&&, ||, ;, |) and extracts
|
| 441 |
+
prefixes from each subcommand. Returns deduplicated list of prefixes.
|
| 442 |
+
|
| 443 |
+
Args:
|
| 444 |
+
command (`str`):
|
| 445 |
+
The bash command string (may be compound)
|
| 446 |
+
max_prefixes (`int`):
|
| 447 |
+
Maximum number of prefixes to return (default: 5)
|
| 448 |
+
|
| 449 |
+
Returns:
|
| 450 |
+
`List[str]`:
|
| 451 |
+
List of command prefixes (deduplicated), e.g., ["npm run",
|
| 452 |
+
"git commit"]
|
| 453 |
+
|
| 454 |
+
Examples:
|
| 455 |
+
>>> parser.extract_command_prefixes("git add . && git commit")
|
| 456 |
+
['git add', 'git commit']
|
| 457 |
+
>>> parser.extract_command_prefixes("npm run build")
|
| 458 |
+
['npm run']
|
| 459 |
+
>>> parser.extract_command_prefixes("ls -la")
|
| 460 |
+
[]
|
| 461 |
+
"""
|
| 462 |
+
if not command or not command.strip():
|
| 463 |
+
return []
|
| 464 |
+
|
| 465 |
+
# Parse command to AST
|
| 466 |
+
tree = self.parser.parse(bytes(command, "utf8"))
|
| 467 |
+
root = tree.root_node
|
| 468 |
+
|
| 469 |
+
# Split compound commands
|
| 470 |
+
subcommands = self.split_compound_command(root, command)
|
| 471 |
+
|
| 472 |
+
# Extract prefixes from each subcommand
|
| 473 |
+
prefixes = []
|
| 474 |
+
seen = set()
|
| 475 |
+
|
| 476 |
+
for subcmd in subcommands[:max_prefixes]:
|
| 477 |
+
prefix = self._extract_command_prefix(subcmd)
|
| 478 |
+
if prefix and prefix not in seen:
|
| 479 |
+
prefixes.append(prefix)
|
| 480 |
+
seen.add(prefix)
|
| 481 |
+
|
| 482 |
+
if len(prefixes) >= max_prefixes:
|
| 483 |
+
break
|
| 484 |
+
|
| 485 |
+
return prefixes
|
| 486 |
+
|
| 487 |
+
def split_compound_command(self, root: Node, command: str) -> List[str]:
|
| 488 |
+
"""Split compound commands using tree-sitter for precise parsing.
|
| 489 |
+
|
| 490 |
+
Recognizes: &&, ||, ;, |
|
| 491 |
+
|
| 492 |
+
Args:
|
| 493 |
+
root (`Node`):
|
| 494 |
+
The root AST node
|
| 495 |
+
command (`str`):
|
| 496 |
+
The original command string
|
| 497 |
+
|
| 498 |
+
Returns:
|
| 499 |
+
`List[str]`:
|
| 500 |
+
List of individual subcommands
|
| 501 |
+
"""
|
| 502 |
+
subcommands = []
|
| 503 |
+
|
| 504 |
+
def extract_commands(node: Node) -> None:
|
| 505 |
+
"""Recursively extract commands from AST."""
|
| 506 |
+
if node.type == "command":
|
| 507 |
+
# Extract command text
|
| 508 |
+
cmd_text = command[node.start_byte : node.end_byte]
|
| 509 |
+
subcommands.append(cmd_text)
|
| 510 |
+
elif node.type in ["list", "pipeline", "command_list"]:
|
| 511 |
+
# Recursively process compound structures
|
| 512 |
+
for child in node.children:
|
| 513 |
+
if child.type not in ["&&", "||", ";", "|", "|&"]:
|
| 514 |
+
extract_commands(child)
|
| 515 |
+
else:
|
| 516 |
+
# Continue traversing
|
| 517 |
+
for child in node.children:
|
| 518 |
+
extract_commands(child)
|
| 519 |
+
|
| 520 |
+
extract_commands(root)
|
| 521 |
+
return subcommands if subcommands else [command]
|
| 522 |
+
|
| 523 |
+
def _extract_command_prefix(
|
| 524 |
+
self,
|
| 525 |
+
subcmd: str,
|
| 526 |
+
) -> Optional[str]:
|
| 527 |
+
"""Extract command prefix (first two words) from a subcommand.
|
| 528 |
+
|
| 529 |
+
Logic:
|
| 530 |
+
1. Skip safe environment variable assignments
|
| 531 |
+
2. Extract command name and first subcommand
|
| 532 |
+
3. Verify the second word looks like a subcommand (not a flag)
|
| 533 |
+
|
| 534 |
+
Args:
|
| 535 |
+
subcmd (`str`):
|
| 536 |
+
The subcommand string to extract prefix from
|
| 537 |
+
|
| 538 |
+
Returns:
|
| 539 |
+
`Optional[str]`:
|
| 540 |
+
Command prefix (e.g., "npm run") or None if cannot extract
|
| 541 |
+
"""
|
| 542 |
+
# Parse the subcommand
|
| 543 |
+
tree = self.parser.parse(bytes(subcmd, "utf8"))
|
| 544 |
+
root = tree.root_node
|
| 545 |
+
|
| 546 |
+
# Find the first simple_command node
|
| 547 |
+
simple_cmd = self._find_first_simple_command(root)
|
| 548 |
+
if not simple_cmd:
|
| 549 |
+
return None
|
| 550 |
+
|
| 551 |
+
# Extract command parts
|
| 552 |
+
parts = []
|
| 553 |
+
env_vars = []
|
| 554 |
+
|
| 555 |
+
for child in simple_cmd.children:
|
| 556 |
+
if child.type == "variable_assignment":
|
| 557 |
+
# Environment variable assignment
|
| 558 |
+
var_name = subcmd[child.start_byte : child.end_byte].split(
|
| 559 |
+
"=",
|
| 560 |
+
)[0]
|
| 561 |
+
env_vars.append(var_name)
|
| 562 |
+
elif child.type == "command_name":
|
| 563 |
+
# Command name
|
| 564 |
+
parts.append(subcmd[child.start_byte : child.end_byte])
|
| 565 |
+
elif child.type == "word" and len(parts) >= 1:
|
| 566 |
+
# Argument (might be a flag or subcommand)
|
| 567 |
+
word = subcmd[child.start_byte : child.end_byte]
|
| 568 |
+
parts.append(word)
|
| 569 |
+
# Stop after we have command + first argument
|
| 570 |
+
if len(parts) >= 2:
|
| 571 |
+
break
|
| 572 |
+
|
| 573 |
+
# Check if environment variables are safe
|
| 574 |
+
if env_vars and not all(v in SAFE_ENV_VARS for v in env_vars):
|
| 575 |
+
return None
|
| 576 |
+
|
| 577 |
+
# Check if the command is a safe command that doesn't need permission
|
| 578 |
+
if parts and parts[0].lower() in SAFE_COMMANDS:
|
| 579 |
+
return None
|
| 580 |
+
|
| 581 |
+
# Return first two words
|
| 582 |
+
if len(parts) >= 2:
|
| 583 |
+
return " ".join(parts[:2])
|
| 584 |
+
|
| 585 |
+
return None
|
| 586 |
+
|
| 587 |
+
def _find_first_simple_command(self, node: Node) -> Optional[Node]:
|
| 588 |
+
"""Recursively find the first command node in AST.
|
| 589 |
+
|
| 590 |
+
Args:
|
| 591 |
+
node (`Node`):
|
| 592 |
+
The AST node to search from
|
| 593 |
+
|
| 594 |
+
Returns:
|
| 595 |
+
`Optional[Node]`:
|
| 596 |
+
The first command node found, or None
|
| 597 |
+
"""
|
| 598 |
+
if node.type == "command":
|
| 599 |
+
return node
|
| 600 |
+
|
| 601 |
+
for child in node.children:
|
| 602 |
+
result = self._find_first_simple_command(child)
|
| 603 |
+
if result:
|
| 604 |
+
return result
|
| 605 |
+
|
| 606 |
+
return None
|
| 607 |
+
|
| 608 |
+
def check_dangerous_command(self, command: str) -> Optional[str]:
|
| 609 |
+
"""Check if command contains dangerous patterns.
|
| 610 |
+
|
| 611 |
+
Uses word-boundary aware matching to avoid false positives like
|
| 612 |
+
'git add' matching 'dd' pattern.
|
| 613 |
+
|
| 614 |
+
Args:
|
| 615 |
+
command (`str`):
|
| 616 |
+
The bash command to check
|
| 617 |
+
|
| 618 |
+
Returns:
|
| 619 |
+
`Optional[str]`:
|
| 620 |
+
The matched dangerous pattern if found, None otherwise
|
| 621 |
+
"""
|
| 622 |
+
|
| 623 |
+
# Normalize command for matching
|
| 624 |
+
normalized = " ".join(command.split())
|
| 625 |
+
|
| 626 |
+
# Check each dangerous pattern
|
| 627 |
+
for pattern in DANGEROUS_COMMANDS:
|
| 628 |
+
# For single-word patterns like "dd", use word boundary matching
|
| 629 |
+
# to avoid false positives (e.g., "git add" shouldn't match "dd")
|
| 630 |
+
if " " not in pattern and len(pattern) <= 4:
|
| 631 |
+
# Single word pattern - use word boundaries
|
| 632 |
+
regex = r"\b" + re.escape(pattern) + r"\b"
|
| 633 |
+
if re.search(regex, normalized):
|
| 634 |
+
return pattern
|
| 635 |
+
else:
|
| 636 |
+
# Multi-word pattern or longer pattern - use substring match
|
| 637 |
+
if pattern in normalized:
|
| 638 |
+
return pattern
|
| 639 |
+
|
| 640 |
+
return None
|
| 641 |
+
|
| 642 |
+
# pylint: disable=too-many-return-statements, too-many-branches
|
| 643 |
+
def check_sed_constraints(
|
| 644 |
+
self,
|
| 645 |
+
command: str,
|
| 646 |
+
dangerous_files: List[str],
|
| 647 |
+
) -> str | None:
|
| 648 |
+
"""Check if sed command violates safety constraints.
|
| 649 |
+
|
| 650 |
+
Implements allowlist/denylist system:
|
| 651 |
+
- Allowlist: Line printing (sed -n 'Np') and substitution (sed 's///')
|
| 652 |
+
- Denylist: Dangerous operations (w/W/e/E), file writes, command
|
| 653 |
+
execution
|
| 654 |
+
|
| 655 |
+
Args:
|
| 656 |
+
command: The bash command to check
|
| 657 |
+
dangerous_files: List of dangerous file patterns
|
| 658 |
+
|
| 659 |
+
Returns:
|
| 660 |
+
Error message if dangerous sed operation found, None otherwise
|
| 661 |
+
"""
|
| 662 |
+
|
| 663 |
+
if "sed" not in command:
|
| 664 |
+
return None
|
| 665 |
+
|
| 666 |
+
# Parse command using shlex
|
| 667 |
+
try:
|
| 668 |
+
tokens = shlex.split(command)
|
| 669 |
+
except ValueError:
|
| 670 |
+
return "sed command has invalid shell syntax"
|
| 671 |
+
|
| 672 |
+
# Find sed command position
|
| 673 |
+
sed_idx = None
|
| 674 |
+
for i, token in enumerate(tokens):
|
| 675 |
+
if token == "sed" or token.endswith("/sed"):
|
| 676 |
+
sed_idx = i
|
| 677 |
+
break
|
| 678 |
+
|
| 679 |
+
if sed_idx is None:
|
| 680 |
+
return None
|
| 681 |
+
|
| 682 |
+
# Extract flags and expressions
|
| 683 |
+
args = tokens[sed_idx + 1 :]
|
| 684 |
+
flags = []
|
| 685 |
+
expressions = []
|
| 686 |
+
file_args = []
|
| 687 |
+
i = 0
|
| 688 |
+
found_first_expr = False
|
| 689 |
+
|
| 690 |
+
while i < len(args):
|
| 691 |
+
arg = args[i]
|
| 692 |
+
|
| 693 |
+
# Handle flags
|
| 694 |
+
if arg.startswith("-") and not arg.startswith("--"):
|
| 695 |
+
# Combined flags like -nE
|
| 696 |
+
flag_chars = arg[1:]
|
| 697 |
+
for char in flag_chars:
|
| 698 |
+
flags.append(char)
|
| 699 |
+
# -i flag may have optional backup extension argument
|
| 700 |
+
# But only skip if next arg doesn't look like an expression
|
| 701 |
+
if "i" in flag_chars and i + 1 < len(args):
|
| 702 |
+
next_arg = args[i + 1]
|
| 703 |
+
# Skip backup extension only if it's not an expression
|
| 704 |
+
# or file
|
| 705 |
+
if (
|
| 706 |
+
not next_arg.startswith("-")
|
| 707 |
+
and not next_arg.startswith("s")
|
| 708 |
+
and "." not in next_arg
|
| 709 |
+
):
|
| 710 |
+
i += 1 # Skip backup extension
|
| 711 |
+
elif arg == "--in-place":
|
| 712 |
+
flags.append("i")
|
| 713 |
+
if i + 1 < len(args):
|
| 714 |
+
next_arg = args[i + 1]
|
| 715 |
+
if (
|
| 716 |
+
not next_arg.startswith("-")
|
| 717 |
+
and not next_arg.startswith("s")
|
| 718 |
+
and "." not in next_arg
|
| 719 |
+
):
|
| 720 |
+
i += 1
|
| 721 |
+
elif arg in ["-e", "--expression"]:
|
| 722 |
+
if i + 1 < len(args):
|
| 723 |
+
expressions.append(args[i + 1])
|
| 724 |
+
i += 1
|
| 725 |
+
elif not arg.startswith("-"):
|
| 726 |
+
# First non-flag, non-option arg is expression (if no -e used)
|
| 727 |
+
if not found_first_expr:
|
| 728 |
+
expressions.append(arg)
|
| 729 |
+
found_first_expr = True
|
| 730 |
+
else:
|
| 731 |
+
file_args.append(arg)
|
| 732 |
+
|
| 733 |
+
i += 1
|
| 734 |
+
|
| 735 |
+
# If no expressions found, command is invalid
|
| 736 |
+
if not expressions:
|
| 737 |
+
return "sed command missing expression"
|
| 738 |
+
|
| 739 |
+
# Validate flags - only allow specific flags
|
| 740 |
+
allowed_flags = {"n", "E", "e", "i"}
|
| 741 |
+
for flag in flags:
|
| 742 |
+
if flag not in allowed_flags:
|
| 743 |
+
return f"sed flag -{flag} not allowed"
|
| 744 |
+
|
| 745 |
+
# Check allowlist patterns
|
| 746 |
+
has_n_flag = "n" in flags
|
| 747 |
+
has_i_flag = "i" in flags
|
| 748 |
+
|
| 749 |
+
for expr in expressions:
|
| 750 |
+
# Denylist checks first - dangerous operations
|
| 751 |
+
# Check for write operations (w, W) - must be at end or followed
|
| 752 |
+
# by space/filename
|
| 753 |
+
if (
|
| 754 |
+
re.search(r"/[wW]\s+\S+", expr)
|
| 755 |
+
or expr.endswith("/w")
|
| 756 |
+
or expr.endswith("/W")
|
| 757 |
+
):
|
| 758 |
+
return "sed write operation (w/W) not allowed"
|
| 759 |
+
|
| 760 |
+
# Check for execute operations (e, E) - must be at end or
|
| 761 |
+
# followed by space
|
| 762 |
+
if (
|
| 763 |
+
re.search(r"/[eE](?:\s|$)", expr)
|
| 764 |
+
or expr.endswith("/e")
|
| 765 |
+
or expr.endswith("/E")
|
| 766 |
+
):
|
| 767 |
+
return "sed execute operation (e/E) not allowed"
|
| 768 |
+
|
| 769 |
+
# Check for dangerous patterns
|
| 770 |
+
if "{" in expr or "}" in expr:
|
| 771 |
+
return "sed curly braces not allowed"
|
| 772 |
+
if expr.startswith("!"):
|
| 773 |
+
return "sed negation (!) not allowed"
|
| 774 |
+
if "#" in expr and not expr.startswith("s#"):
|
| 775 |
+
return "sed comments not allowed"
|
| 776 |
+
|
| 777 |
+
# Pattern 1: Line printing with -n flag (sed -n 'Np' or 'N,Mp')
|
| 778 |
+
if has_n_flag:
|
| 779 |
+
# Match: number followed by 'p', or range 'N,Mp'
|
| 780 |
+
if re.match(
|
| 781 |
+
r"^\d+p$",
|
| 782 |
+
expr,
|
| 783 |
+
) or re.match(
|
| 784 |
+
r"^\d+,\d+p$",
|
| 785 |
+
expr,
|
| 786 |
+
):
|
| 787 |
+
continue
|
| 788 |
+
|
| 789 |
+
# Pattern 2: Substitution command
|
| 790 |
+
# (sed 's/pattern/replacement/flags')
|
| 791 |
+
if (
|
| 792 |
+
expr.startswith("s/")
|
| 793 |
+
or expr.startswith("s|")
|
| 794 |
+
or expr.startswith("s#")
|
| 795 |
+
):
|
| 796 |
+
delimiter = expr[1]
|
| 797 |
+
parts = expr[2:].split(delimiter)
|
| 798 |
+
if len(parts) >= 2:
|
| 799 |
+
# Valid substitution
|
| 800 |
+
# Check substitution flags (g, p, number, etc.)
|
| 801 |
+
if len(parts) > 2:
|
| 802 |
+
sub_flags = parts[2]
|
| 803 |
+
# Allow common substitution flags
|
| 804 |
+
if all(c in "gp0123456789" for c in sub_flags):
|
| 805 |
+
continue
|
| 806 |
+
else:
|
| 807 |
+
continue
|
| 808 |
+
|
| 809 |
+
# If we reach here, expression doesn't match allowlist
|
| 810 |
+
return f"sed expression '{expr}' not in allowlist"
|
| 811 |
+
|
| 812 |
+
# Check -i flag with dangerous files
|
| 813 |
+
if has_i_flag and file_args:
|
| 814 |
+
for file_path in file_args:
|
| 815 |
+
for dangerous_file in dangerous_files:
|
| 816 |
+
if dangerous_file in file_path or file_path.endswith(
|
| 817 |
+
dangerous_file,
|
| 818 |
+
):
|
| 819 |
+
return f"sed -i modifying dangerous file: {file_path}"
|
| 820 |
+
|
| 821 |
+
return None
|
| 822 |
+
|
| 823 |
+
def check_injection_risk(self, command: str) -> Optional[str]:
|
| 824 |
+
"""Check if command contains structures that cannot be statically
|
| 825 |
+
analyzed.
|
| 826 |
+
|
| 827 |
+
This detects command substitution, process substitution, complex
|
| 828 |
+
expansions, control flow, and other dynamic shell features that
|
| 829 |
+
make it impossible to determine the command's behavior without
|
| 830 |
+
execution.
|
| 831 |
+
|
| 832 |
+
Args:
|
| 833 |
+
command (`str`):
|
| 834 |
+
The bash command to check
|
| 835 |
+
|
| 836 |
+
Returns:
|
| 837 |
+
`Optional[str]`:
|
| 838 |
+
Reason string if command is too complex, None if it can be
|
| 839 |
+
statically analyzed
|
| 840 |
+
|
| 841 |
+
Examples:
|
| 842 |
+
>>> parser.check_injection_risk("ls -la")
|
| 843 |
+
None
|
| 844 |
+
>>> parser.check_injection_risk("rm $(find . -name '*.tmp')")
|
| 845 |
+
"Command contains command_substitution which cannot be statically
|
| 846 |
+
analyzed"
|
| 847 |
+
>>> parser.check_injection_risk("for f in *.txt; do cat $f; done")
|
| 848 |
+
"Command contains for_statement which cannot be statically
|
| 849 |
+
analyzed"
|
| 850 |
+
"""
|
| 851 |
+
|
| 852 |
+
try:
|
| 853 |
+
tree = self.parser.parse(bytes(command, "utf8"))
|
| 854 |
+
return self._walk_for_dangerous_nodes(tree.root_node)
|
| 855 |
+
except Exception:
|
| 856 |
+
# If parsing fails, be conservative and require review
|
| 857 |
+
return "Command parsing failed, cannot verify safety"
|
| 858 |
+
|
| 859 |
+
def _walk_for_dangerous_nodes(self, node: Node) -> Optional[str]:
|
| 860 |
+
"""Recursively walk AST to find dangerous node types.
|
| 861 |
+
|
| 862 |
+
Args:
|
| 863 |
+
node (`Node`):
|
| 864 |
+
The AST node to check
|
| 865 |
+
|
| 866 |
+
Returns:
|
| 867 |
+
`Optional[str]`:
|
| 868 |
+
Reason string if dangerous node found, None otherwise
|
| 869 |
+
"""
|
| 870 |
+
|
| 871 |
+
# Check if this node is a dangerous type
|
| 872 |
+
if node.type in DANGEROUS_NODE_TYPES:
|
| 873 |
+
return (
|
| 874 |
+
f"Command contains {node.type} which cannot be "
|
| 875 |
+
f"statically analyzed"
|
| 876 |
+
)
|
| 877 |
+
|
| 878 |
+
# Recursively check children
|
| 879 |
+
for child in node.children:
|
| 880 |
+
result = self._walk_for_dangerous_nodes(child)
|
| 881 |
+
if result:
|
| 882 |
+
return result
|
| 883 |
+
|
| 884 |
+
return None
|
src/agentscope/tool/_builtin/_edit.py
ADDED
|
@@ -0,0 +1,416 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The edit tool in agentscope."""
|
| 3 |
+
import difflib
|
| 4 |
+
import fnmatch
|
| 5 |
+
from typing import Any, List
|
| 6 |
+
|
| 7 |
+
from .._base import ToolBase, ToolMiddlewareBase
|
| 8 |
+
from .._constants import (
|
| 9 |
+
DEFAULT_DANGEROUS_FILES,
|
| 10 |
+
DEFAULT_DANGEROUS_DIRECTORIES,
|
| 11 |
+
)
|
| 12 |
+
from ...permission import (
|
| 13 |
+
PermissionContext,
|
| 14 |
+
PermissionDecision,
|
| 15 |
+
PermissionBehavior,
|
| 16 |
+
PermissionMode,
|
| 17 |
+
PermissionRule,
|
| 18 |
+
)
|
| 19 |
+
from .._response import ToolChunk
|
| 20 |
+
from ...message import TextBlock, ToolResultState
|
| 21 |
+
from ...state import AgentState
|
| 22 |
+
from ._backend import BackendBase, _normalize_newlines
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
class Edit(ToolBase):
|
| 26 |
+
"""The edit tool for performing exact string replacements in files."""
|
| 27 |
+
|
| 28 |
+
name: str = "Edit"
|
| 29 |
+
"""The tool name presented to the agent."""
|
| 30 |
+
|
| 31 |
+
description: str = """Performs exact string replacements in files.
|
| 32 |
+
|
| 33 |
+
Usage:
|
| 34 |
+
- You must use your `Read` tool at least once in the conversation
|
| 35 |
+
before editing. This tool will error if you attempt an edit without
|
| 36 |
+
reading the file.
|
| 37 |
+
- When editing text from Read tool output, ensure you preserve the
|
| 38 |
+
exact indentation (tabs/spaces) as it appears AFTER the line number
|
| 39 |
+
prefix. The line number prefix format is: line number + tab.
|
| 40 |
+
Everything after that is the actual file content to match. Never
|
| 41 |
+
include any part of the line number prefix in the old_string or
|
| 42 |
+
new_string.
|
| 43 |
+
- ALWAYS prefer editing existing files in the codebase. NEVER write
|
| 44 |
+
new files unless explicitly required.
|
| 45 |
+
- Only use emojis if the user explicitly requests it. Avoid adding
|
| 46 |
+
emojis to files unless asked.
|
| 47 |
+
- The edit will FAIL if `old_string` is not unique in the file.""" # noqa: E501
|
| 48 |
+
"""The description presented to the agent."""
|
| 49 |
+
|
| 50 |
+
input_schema: dict[str, Any] = {
|
| 51 |
+
"type": "object",
|
| 52 |
+
"properties": {
|
| 53 |
+
"file_path": {
|
| 54 |
+
"type": "string",
|
| 55 |
+
"description": "The absolute path to the file to edit.",
|
| 56 |
+
},
|
| 57 |
+
"old_string": {
|
| 58 |
+
"type": "string",
|
| 59 |
+
"description": (
|
| 60 |
+
"The exact string to replace. Must match exactly "
|
| 61 |
+
"including whitespace and indentation."
|
| 62 |
+
),
|
| 63 |
+
},
|
| 64 |
+
"new_string": {
|
| 65 |
+
"type": "string",
|
| 66 |
+
"description": "The string to replace old_string with.",
|
| 67 |
+
},
|
| 68 |
+
"replace_all": {
|
| 69 |
+
"type": "boolean",
|
| 70 |
+
"description": (
|
| 71 |
+
"If true, replace all occurrences. If false "
|
| 72 |
+
"(default), only replace if there is exactly one "
|
| 73 |
+
"occurrence."
|
| 74 |
+
),
|
| 75 |
+
"default": False,
|
| 76 |
+
},
|
| 77 |
+
},
|
| 78 |
+
"required": ["file_path", "old_string", "new_string"],
|
| 79 |
+
}
|
| 80 |
+
|
| 81 |
+
is_mcp: bool = False
|
| 82 |
+
is_read_only: bool = False
|
| 83 |
+
is_concurrency_safe: bool = False
|
| 84 |
+
is_external_tool: bool = False
|
| 85 |
+
is_state_injected: bool = True
|
| 86 |
+
|
| 87 |
+
def __init__( # pylint: disable=dangerous-default-value
|
| 88 |
+
self,
|
| 89 |
+
dangerous_files: list[str] = DEFAULT_DANGEROUS_FILES,
|
| 90 |
+
dangerous_directories: list[str] = DEFAULT_DANGEROUS_DIRECTORIES,
|
| 91 |
+
middlewares: List[ToolMiddlewareBase] | None = None,
|
| 92 |
+
backend: BackendBase | None = None,
|
| 93 |
+
) -> None:
|
| 94 |
+
"""Initialize the edit tool.
|
| 95 |
+
|
| 96 |
+
Args:
|
| 97 |
+
dangerous_files (`list[str]`, optional):
|
| 98 |
+
Sensitive files that require explicit user confirmation,
|
| 99 |
+
even in BYPASS mode. Matched by basename
|
| 100 |
+
(case-insensitive). Defaults to `DEFAULT_DANGEROUS_FILES`.
|
| 101 |
+
Pass a custom list to fully replace the defaults, or `[]`
|
| 102 |
+
to disable the filename check.
|
| 103 |
+
dangerous_directories (`list[str]`, optional):
|
| 104 |
+
Sensitive directories that require explicit user
|
| 105 |
+
confirmation. Matched when any path segment equals an
|
| 106 |
+
entry (case-insensitive). Defaults to
|
| 107 |
+
`DEFAULT_DANGEROUS_DIRECTORIES`. Pass a custom list to
|
| 108 |
+
fully replace the defaults, or `[]` to disable the
|
| 109 |
+
directory check.
|
| 110 |
+
middlewares (`List[ToolMiddlewareBase] | None`, optional):
|
| 111 |
+
Tool middlewares wrapping the tool execution.
|
| 112 |
+
backend (`BackendBase | None`, optional):
|
| 113 |
+
The sandbox backend to use for file I/O. When ``None``,
|
| 114 |
+
a :class:`LocalBackend` is created.
|
| 115 |
+
"""
|
| 116 |
+
from ._backend import LocalBackend
|
| 117 |
+
|
| 118 |
+
super().__init__(middlewares=middlewares)
|
| 119 |
+
self.dangerous_files = list(dangerous_files)
|
| 120 |
+
self.dangerous_directories = list(dangerous_directories)
|
| 121 |
+
self._backend = backend or LocalBackend()
|
| 122 |
+
|
| 123 |
+
async def check_permissions(
|
| 124 |
+
self,
|
| 125 |
+
tool_input: dict[str, Any],
|
| 126 |
+
context: PermissionContext,
|
| 127 |
+
) -> PermissionDecision:
|
| 128 |
+
"""Check permissions for file editing.
|
| 129 |
+
|
| 130 |
+
This method implements Edit-specific permission checks:
|
| 131 |
+
1. Dangerous path check (safety check, bypass-immune)
|
| 132 |
+
2. ACCEPT_EDITS mode check for files in working directories
|
| 133 |
+
|
| 134 |
+
Args:
|
| 135 |
+
tool_input (`dict[str, Any]`):
|
| 136 |
+
The tool input containing "file_path" key
|
| 137 |
+
context (`PermissionContext`):
|
| 138 |
+
The permission context with mode and rules
|
| 139 |
+
|
| 140 |
+
Returns:
|
| 141 |
+
`PermissionDecision`:
|
| 142 |
+
ASK for dangerous paths, ALLOW for safe operations in
|
| 143 |
+
ACCEPT_EDITS mode, PASSTHROUGH otherwise
|
| 144 |
+
"""
|
| 145 |
+
|
| 146 |
+
file_path = tool_input.get("file_path")
|
| 147 |
+
if not file_path:
|
| 148 |
+
return PermissionDecision(
|
| 149 |
+
behavior=PermissionBehavior.PASSTHROUGH,
|
| 150 |
+
message="No file path provided",
|
| 151 |
+
)
|
| 152 |
+
|
| 153 |
+
# 1. Check for dangerous paths (safety check, bypass-immune)
|
| 154 |
+
if self._is_dangerous_path(file_path):
|
| 155 |
+
return PermissionDecision(
|
| 156 |
+
behavior=PermissionBehavior.ASK,
|
| 157 |
+
message=f"Permission required: Edit operation on "
|
| 158 |
+
f"sensitive file {file_path}",
|
| 159 |
+
decision_reason="Safety check: dangerous file or directory",
|
| 160 |
+
bypass_immune=True,
|
| 161 |
+
)
|
| 162 |
+
|
| 163 |
+
# 2. Check ACCEPT_EDITS mode for files in working directories
|
| 164 |
+
if context.mode == PermissionMode.ACCEPT_EDITS:
|
| 165 |
+
if self._path_in_allowed_working_path(file_path, context):
|
| 166 |
+
return PermissionDecision(
|
| 167 |
+
behavior=PermissionBehavior.ALLOW,
|
| 168 |
+
message=f"Permission granted for editing {file_path} "
|
| 169 |
+
f"(accept edits mode - in working directory)",
|
| 170 |
+
decision_reason="File is in working directory and not "
|
| 171 |
+
"a dangerous path",
|
| 172 |
+
)
|
| 173 |
+
|
| 174 |
+
# 3. Return PASSTHROUGH to let PermissionEngine check allow rules
|
| 175 |
+
# This ensures allow rules can grant Edit permissions
|
| 176 |
+
return PermissionDecision(
|
| 177 |
+
behavior=PermissionBehavior.PASSTHROUGH,
|
| 178 |
+
message="",
|
| 179 |
+
)
|
| 180 |
+
|
| 181 |
+
async def match_rule(
|
| 182 |
+
self,
|
| 183 |
+
rule_content: str | None,
|
| 184 |
+
tool_input: dict[str, Any],
|
| 185 |
+
) -> bool:
|
| 186 |
+
"""Check if a permission rule matches the file path.
|
| 187 |
+
|
| 188 |
+
Matches rule_content as a glob pattern against the "file_path"
|
| 189 |
+
parameter using fnmatch. If rule_content is None, matches all
|
| 190 |
+
invocations (tool-name-level rule).
|
| 191 |
+
|
| 192 |
+
Args:
|
| 193 |
+
rule_content (`str | None`):
|
| 194 |
+
Glob pattern to match against the file path (e.g., "src/**"),
|
| 195 |
+
or None to match all invocations
|
| 196 |
+
tool_input (`dict[str, Any]`):
|
| 197 |
+
The tool input data containing "file_path" key
|
| 198 |
+
|
| 199 |
+
Returns:
|
| 200 |
+
`bool`:
|
| 201 |
+
True if the glob pattern matches the file path, False otherwise
|
| 202 |
+
"""
|
| 203 |
+
if rule_content is None:
|
| 204 |
+
return True
|
| 205 |
+
|
| 206 |
+
file_path = tool_input.get("file_path", "")
|
| 207 |
+
if not file_path:
|
| 208 |
+
return False
|
| 209 |
+
return fnmatch.fnmatch(file_path, rule_content)
|
| 210 |
+
|
| 211 |
+
async def generate_suggestions(
|
| 212 |
+
self,
|
| 213 |
+
tool_input: dict[str, Any],
|
| 214 |
+
) -> List[PermissionRule]:
|
| 215 |
+
"""Generate suggested permission rules for the file path.
|
| 216 |
+
|
| 217 |
+
Suggests a glob pattern covering the parent directory of the file,
|
| 218 |
+
allowing the user to grant permission for the entire directory at once.
|
| 219 |
+
|
| 220 |
+
Args:
|
| 221 |
+
tool_input (`dict[str, Any]`):
|
| 222 |
+
The tool input data containing "file_path" key
|
| 223 |
+
|
| 224 |
+
Returns:
|
| 225 |
+
`List[PermissionRule]`:
|
| 226 |
+
A single suggested rule covering the parent directory
|
| 227 |
+
(e.g., file "/src/main.py" -> rule "src/**")
|
| 228 |
+
"""
|
| 229 |
+
file_path = tool_input.get("file_path", "")
|
| 230 |
+
if not file_path:
|
| 231 |
+
return []
|
| 232 |
+
|
| 233 |
+
parent = self._backend.dirname(file_path)
|
| 234 |
+
# Glob patterns are POSIX-style strings (matched by fnmatch),
|
| 235 |
+
# not real filesystem paths — do NOT use backend.join_path here.
|
| 236 |
+
pattern = (parent.rstrip("/\\") + "/**") if parent else "**"
|
| 237 |
+
|
| 238 |
+
return [
|
| 239 |
+
PermissionRule(
|
| 240 |
+
tool_name=self.name,
|
| 241 |
+
rule_content=pattern,
|
| 242 |
+
behavior=PermissionBehavior.ALLOW,
|
| 243 |
+
source="suggested",
|
| 244 |
+
),
|
| 245 |
+
]
|
| 246 |
+
|
| 247 |
+
async def call( # type: ignore[override]
|
| 248 |
+
self,
|
| 249 |
+
file_path: str,
|
| 250 |
+
old_string: str,
|
| 251 |
+
new_string: str,
|
| 252 |
+
replace_all: bool = False,
|
| 253 |
+
_agent_state: AgentState | None = None,
|
| 254 |
+
) -> ToolChunk:
|
| 255 |
+
"""Execute the edit and return the result."""
|
| 256 |
+
# Validate file_path is absolute
|
| 257 |
+
if not self._backend.isabs(file_path):
|
| 258 |
+
return ToolChunk(
|
| 259 |
+
content=[
|
| 260 |
+
TextBlock(
|
| 261 |
+
text=(
|
| 262 |
+
f"Error: file_path must be an absolute "
|
| 263 |
+
f"path, got: {file_path}"
|
| 264 |
+
),
|
| 265 |
+
),
|
| 266 |
+
],
|
| 267 |
+
state=ToolResultState.ERROR,
|
| 268 |
+
is_last=True,
|
| 269 |
+
)
|
| 270 |
+
|
| 271 |
+
# Check file exists
|
| 272 |
+
if not await self._backend.file_exists(file_path):
|
| 273 |
+
return ToolChunk(
|
| 274 |
+
content=[
|
| 275 |
+
TextBlock(text=f"Error: File not found: {file_path}"),
|
| 276 |
+
],
|
| 277 |
+
state=ToolResultState.ERROR,
|
| 278 |
+
is_last=True,
|
| 279 |
+
)
|
| 280 |
+
|
| 281 |
+
# Check old_string != new_string
|
| 282 |
+
if old_string == new_string:
|
| 283 |
+
return ToolChunk(
|
| 284 |
+
content=[
|
| 285 |
+
TextBlock(
|
| 286 |
+
text=(
|
| 287 |
+
"Error: old_string and new_string are "
|
| 288 |
+
"identical. No changes to make."
|
| 289 |
+
),
|
| 290 |
+
),
|
| 291 |
+
],
|
| 292 |
+
state=ToolResultState.ERROR,
|
| 293 |
+
is_last=True,
|
| 294 |
+
)
|
| 295 |
+
|
| 296 |
+
content = None
|
| 297 |
+
if _agent_state is not None:
|
| 298 |
+
cache = await _agent_state.tool_context.get_cache(file_path)
|
| 299 |
+
if cache is None:
|
| 300 |
+
# Haven't read this file before
|
| 301 |
+
return ToolChunk(
|
| 302 |
+
content=[
|
| 303 |
+
TextBlock(
|
| 304 |
+
text="Error: To edit a file, you must first read "
|
| 305 |
+
"it using the Read tool.",
|
| 306 |
+
),
|
| 307 |
+
],
|
| 308 |
+
state=ToolResultState.ERROR,
|
| 309 |
+
is_last=True,
|
| 310 |
+
)
|
| 311 |
+
content = "".join(cache.lines)
|
| 312 |
+
else:
|
| 313 |
+
# No state provided, read from backend
|
| 314 |
+
try:
|
| 315 |
+
raw = await self._backend.read_file(file_path)
|
| 316 |
+
# Normalize CRLF/CR to match the cached-content path and
|
| 317 |
+
# the LF-based old_string the caller supplies.
|
| 318 |
+
content = _normalize_newlines(
|
| 319 |
+
raw.decode("utf-8", errors="replace"),
|
| 320 |
+
)
|
| 321 |
+
except Exception as e:
|
| 322 |
+
return ToolChunk(
|
| 323 |
+
content=[TextBlock(text=f"Error reading file: {str(e)}")],
|
| 324 |
+
state=ToolResultState.ERROR,
|
| 325 |
+
is_last=True,
|
| 326 |
+
)
|
| 327 |
+
|
| 328 |
+
# Count occurrences
|
| 329 |
+
occurrences = content.count(old_string)
|
| 330 |
+
|
| 331 |
+
# If occurrences == 0, raise error
|
| 332 |
+
if occurrences == 0:
|
| 333 |
+
return ToolChunk(
|
| 334 |
+
content=[
|
| 335 |
+
TextBlock(
|
| 336 |
+
text=f"Error: old_string not found in {file_path}",
|
| 337 |
+
),
|
| 338 |
+
],
|
| 339 |
+
state=ToolResultState.ERROR,
|
| 340 |
+
is_last=True,
|
| 341 |
+
)
|
| 342 |
+
|
| 343 |
+
# If occurrences > 1 and not replace_all, raise error
|
| 344 |
+
if occurrences > 1 and not replace_all:
|
| 345 |
+
return ToolChunk(
|
| 346 |
+
content=[
|
| 347 |
+
TextBlock(
|
| 348 |
+
text=(
|
| 349 |
+
f"Error: old_string appears {occurrences} "
|
| 350 |
+
f"times in {file_path}. Set replace_all=true "
|
| 351 |
+
f"to replace all occurrences, or make "
|
| 352 |
+
f"old_string more specific."
|
| 353 |
+
),
|
| 354 |
+
),
|
| 355 |
+
],
|
| 356 |
+
state=ToolResultState.ERROR,
|
| 357 |
+
is_last=True,
|
| 358 |
+
)
|
| 359 |
+
|
| 360 |
+
# Perform replacement
|
| 361 |
+
if replace_all:
|
| 362 |
+
updated_content = content.replace(old_string, new_string)
|
| 363 |
+
else:
|
| 364 |
+
updated_content = content.replace(
|
| 365 |
+
old_string,
|
| 366 |
+
new_string,
|
| 367 |
+
1,
|
| 368 |
+
)
|
| 369 |
+
|
| 370 |
+
# Write updated content back to file via backend
|
| 371 |
+
try:
|
| 372 |
+
await self._backend.write_file(
|
| 373 |
+
file_path,
|
| 374 |
+
updated_content.encode("utf-8"),
|
| 375 |
+
)
|
| 376 |
+
except Exception as e:
|
| 377 |
+
return ToolChunk(
|
| 378 |
+
content=[TextBlock(text=f"Error writing file: {str(e)}")],
|
| 379 |
+
state=ToolResultState.ERROR,
|
| 380 |
+
is_last=True,
|
| 381 |
+
)
|
| 382 |
+
|
| 383 |
+
# Return success message
|
| 384 |
+
replacement_msg = (
|
| 385 |
+
f"all {occurrences} occurrences" if replace_all else "1 occurrence"
|
| 386 |
+
)
|
| 387 |
+
|
| 388 |
+
# Build a unified diff of the change with absolute line numbers so the
|
| 389 |
+
# web UI can render it with real line numbers and proper inter-hunk
|
| 390 |
+
# gaps. The diff is kept in ``metadata`` only (not in the textual
|
| 391 |
+
# output) so it does not bloat the LLM context.
|
| 392 |
+
diff_text = "".join(
|
| 393 |
+
difflib.unified_diff(
|
| 394 |
+
content.splitlines(keepends=True),
|
| 395 |
+
updated_content.splitlines(keepends=True),
|
| 396 |
+
fromfile=f"a/{file_path}",
|
| 397 |
+
tofile=f"b/{file_path}",
|
| 398 |
+
n=3,
|
| 399 |
+
),
|
| 400 |
+
)
|
| 401 |
+
|
| 402 |
+
return ToolChunk(
|
| 403 |
+
content=[
|
| 404 |
+
TextBlock(
|
| 405 |
+
text=f"Successfully replaced {replacement_msg} "
|
| 406 |
+
f"in {file_path}",
|
| 407 |
+
),
|
| 408 |
+
],
|
| 409 |
+
state=ToolResultState.RUNNING,
|
| 410 |
+
is_last=True,
|
| 411 |
+
metadata={
|
| 412 |
+
"diff": diff_text,
|
| 413 |
+
"file_path": file_path,
|
| 414 |
+
"occurrences": occurrences if replace_all else 1,
|
| 415 |
+
},
|
| 416 |
+
)
|
src/agentscope/tool/_builtin/_glob.py
ADDED
|
@@ -0,0 +1,305 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The glob tool in agentscope."""
|
| 3 |
+
|
| 4 |
+
from __future__ import annotations
|
| 5 |
+
|
| 6 |
+
import fnmatch
|
| 7 |
+
import json
|
| 8 |
+
import sys
|
| 9 |
+
from typing import TYPE_CHECKING, Any, List
|
| 10 |
+
|
| 11 |
+
from ...message import TextBlock, ToolResultState
|
| 12 |
+
from ...permission import (
|
| 13 |
+
PermissionBehavior,
|
| 14 |
+
PermissionContext,
|
| 15 |
+
PermissionDecision,
|
| 16 |
+
PermissionRule,
|
| 17 |
+
)
|
| 18 |
+
from .._base import ToolBase, ToolMiddlewareBase
|
| 19 |
+
from .._response import ToolChunk
|
| 20 |
+
|
| 21 |
+
if TYPE_CHECKING:
|
| 22 |
+
from ._backend import BackendBase
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def _default_glob_helper_path() -> str:
|
| 26 |
+
"""Resolve the on-disk path of the bundled ``_glob_helper.py`` script.
|
| 27 |
+
|
| 28 |
+
Used by :class:`Glob` when no explicit ``glob_helper_path`` is
|
| 29 |
+
provided (i.e. the local-workspace case). The path is obtained via
|
| 30 |
+
:mod:`importlib.resources` so it works for both editable and
|
| 31 |
+
installed packages.
|
| 32 |
+
"""
|
| 33 |
+
import importlib.resources as _res
|
| 34 |
+
|
| 35 |
+
ref = _res.files("agentscope.tool._builtin._scripts").joinpath(
|
| 36 |
+
"_glob_helper.py",
|
| 37 |
+
)
|
| 38 |
+
# as_posix() on a MultiplexedPath / PosixPath gives a str path
|
| 39 |
+
return str(ref)
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
class Glob(ToolBase):
|
| 43 |
+
"""The glob tool for fast file pattern matching."""
|
| 44 |
+
|
| 45 |
+
name: str = "Glob"
|
| 46 |
+
"""The tool name presented to the agent."""
|
| 47 |
+
|
| 48 |
+
description: str = """Fast file pattern matching tool that works with
|
| 49 |
+
any codebase size.
|
| 50 |
+
|
| 51 |
+
Supports glob patterns like "**/*.js" or "src/**/*.ts" and returns
|
| 52 |
+
matching file paths sorted by modification time (newest first).
|
| 53 |
+
|
| 54 |
+
Use this tool when you need to find files by pattern across the
|
| 55 |
+
codebase.""" # ignore: E501
|
| 56 |
+
"""The description presented to the agent."""
|
| 57 |
+
|
| 58 |
+
input_schema: dict[str, Any] = {
|
| 59 |
+
"type": "object",
|
| 60 |
+
"properties": {
|
| 61 |
+
"pattern": {
|
| 62 |
+
"type": "string",
|
| 63 |
+
"description": "The glob pattern to match against "
|
| 64 |
+
"(e.g., '**/*.py', 'src/**/*.ts')",
|
| 65 |
+
},
|
| 66 |
+
"path": {
|
| 67 |
+
"type": "string",
|
| 68 |
+
"description": "The base directory to search from "
|
| 69 |
+
"(defaults to current working directory)",
|
| 70 |
+
},
|
| 71 |
+
},
|
| 72 |
+
"required": ["pattern"],
|
| 73 |
+
}
|
| 74 |
+
|
| 75 |
+
is_mcp: bool = False
|
| 76 |
+
is_read_only: bool = True
|
| 77 |
+
is_concurrency_safe: bool = True
|
| 78 |
+
is_external_tool: bool = False
|
| 79 |
+
is_state_injected: bool = False
|
| 80 |
+
|
| 81 |
+
def __init__(
|
| 82 |
+
self,
|
| 83 |
+
backend: BackendBase | None = None,
|
| 84 |
+
glob_helper_path: str | None = None,
|
| 85 |
+
middlewares: List[ToolMiddlewareBase] | None = None,
|
| 86 |
+
) -> None:
|
| 87 |
+
"""Initialize the glob tool.
|
| 88 |
+
|
| 89 |
+
Args:
|
| 90 |
+
middlewares (`List[ToolMiddlewareBase] | None`, optional):
|
| 91 |
+
Tool middlewares wrapping the tool execution.
|
| 92 |
+
backend (`BackendBase | None`, optional):
|
| 93 |
+
The sandbox backend to use. When ``None``, a
|
| 94 |
+
:class:`LocalBackend` is created automatically.
|
| 95 |
+
glob_helper_path (`str | None`, optional):
|
| 96 |
+
Filesystem path (inside the backend's environment) to
|
| 97 |
+
the ``_glob_helper.py`` script. When ``None``, the
|
| 98 |
+
path is resolved from the installed package resources
|
| 99 |
+
(suitable for :class:`LocalBackend`). Remote backends
|
| 100 |
+
(Docker, E2B) should pass the path where the script
|
| 101 |
+
was deployed during workspace initialization.
|
| 102 |
+
"""
|
| 103 |
+
from ._backend import LocalBackend
|
| 104 |
+
|
| 105 |
+
super().__init__(middlewares=middlewares)
|
| 106 |
+
self._backend = backend or LocalBackend()
|
| 107 |
+
# When running against the host, invoke the helper with the
|
| 108 |
+
# current interpreter (``sys.executable``) rather than assuming
|
| 109 |
+
# ``python3`` is on PATH.
|
| 110 |
+
self._is_local = isinstance(self._backend, LocalBackend)
|
| 111 |
+
self._glob_helper_path = (
|
| 112 |
+
glob_helper_path
|
| 113 |
+
if glob_helper_path is not None
|
| 114 |
+
else _default_glob_helper_path()
|
| 115 |
+
)
|
| 116 |
+
|
| 117 |
+
async def check_permissions(
|
| 118 |
+
self,
|
| 119 |
+
tool_input: dict[str, Any],
|
| 120 |
+
context: PermissionContext,
|
| 121 |
+
) -> PermissionDecision:
|
| 122 |
+
"""Check permissions for glob pattern matching.
|
| 123 |
+
|
| 124 |
+
Glob is a read-only tool. Return PASSTHROUGH to let the engine
|
| 125 |
+
handle EXPLORE mode and rule matching.
|
| 126 |
+
"""
|
| 127 |
+
return PermissionDecision(
|
| 128 |
+
behavior=PermissionBehavior.PASSTHROUGH,
|
| 129 |
+
message="Glob pattern matching is read-only.",
|
| 130 |
+
)
|
| 131 |
+
|
| 132 |
+
async def match_rule(
|
| 133 |
+
self,
|
| 134 |
+
rule_content: str | None,
|
| 135 |
+
tool_input: dict[str, Any],
|
| 136 |
+
) -> bool:
|
| 137 |
+
"""Check if a permission rule matches the glob pattern or path.
|
| 138 |
+
|
| 139 |
+
Matches rule_content as a glob pattern against the "pattern" or "path"
|
| 140 |
+
parameters. This allows rules to match either the search pattern itself
|
| 141 |
+
or the directory being searched. If rule_content is None, matches all
|
| 142 |
+
invocations (tool-name-level rule).
|
| 143 |
+
|
| 144 |
+
Args:
|
| 145 |
+
rule_content (`str | None`):
|
| 146 |
+
Glob pattern to match (e.g., "src/**" to match searches in
|
| 147 |
+
src), or None to match all invocations
|
| 148 |
+
tool_input (`dict[str, Any]`):
|
| 149 |
+
The tool input data containing "pattern" and optional "path"
|
| 150 |
+
|
| 151 |
+
Returns:
|
| 152 |
+
`bool`:
|
| 153 |
+
True if the rule matches the pattern or path, False otherwise
|
| 154 |
+
"""
|
| 155 |
+
# None = tool-name-level rule, matches everything
|
| 156 |
+
if rule_content is None:
|
| 157 |
+
return True
|
| 158 |
+
|
| 159 |
+
# Try matching against the search path first
|
| 160 |
+
path = tool_input.get("path", "")
|
| 161 |
+
if path and fnmatch.fnmatch(path, rule_content):
|
| 162 |
+
return True
|
| 163 |
+
|
| 164 |
+
# Fall back to matching against the pattern itself
|
| 165 |
+
pattern = tool_input.get("pattern", "")
|
| 166 |
+
if pattern and fnmatch.fnmatch(pattern, rule_content):
|
| 167 |
+
return True
|
| 168 |
+
|
| 169 |
+
return False
|
| 170 |
+
|
| 171 |
+
async def generate_suggestions(
|
| 172 |
+
self,
|
| 173 |
+
tool_input: dict[str, Any],
|
| 174 |
+
) -> List[PermissionRule]:
|
| 175 |
+
"""Generate suggested permission rules for the glob search.
|
| 176 |
+
|
| 177 |
+
Suggests a rule based on the search path. If no path is provided,
|
| 178 |
+
suggests a rule for the current directory.
|
| 179 |
+
|
| 180 |
+
Args:
|
| 181 |
+
tool_input (`dict[str, Any]`):
|
| 182 |
+
The tool input data containing optional "path" key
|
| 183 |
+
|
| 184 |
+
Returns:
|
| 185 |
+
`List[PermissionRule]`:
|
| 186 |
+
A single suggested rule covering the search directory
|
| 187 |
+
"""
|
| 188 |
+
backend_cwd = await self._backend.getcwd()
|
| 189 |
+
path = tool_input.get("path") or backend_cwd
|
| 190 |
+
|
| 191 |
+
# Normalize path and build a glob pattern. Glob patterns are
|
| 192 |
+
# POSIX-style strings (matched by fnmatch), not real filesystem
|
| 193 |
+
# paths — do NOT use backend.join_path here.
|
| 194 |
+
abs_path = self._backend.abspath(path, cwd=backend_cwd)
|
| 195 |
+
pattern = abs_path.rstrip("/\\") + "/**"
|
| 196 |
+
|
| 197 |
+
return [
|
| 198 |
+
PermissionRule(
|
| 199 |
+
tool_name=self.name,
|
| 200 |
+
rule_content=pattern,
|
| 201 |
+
behavior=PermissionBehavior.ALLOW,
|
| 202 |
+
source="suggested",
|
| 203 |
+
),
|
| 204 |
+
]
|
| 205 |
+
|
| 206 |
+
async def call( # type: ignore[override]
|
| 207 |
+
self,
|
| 208 |
+
pattern: str,
|
| 209 |
+
path: str | None = None,
|
| 210 |
+
) -> ToolChunk:
|
| 211 |
+
"""Execute the glob pattern matching and return the results.
|
| 212 |
+
|
| 213 |
+
Invokes the standalone ``_glob_helper.py`` script via
|
| 214 |
+
``exec_shell``. The script performs high-performance
|
| 215 |
+
``os.walk`` + ``os.scandir`` matching and returns results
|
| 216 |
+
sorted by modification time (newest first) as JSON.
|
| 217 |
+
|
| 218 |
+
This unified path works identically across Local, Docker,
|
| 219 |
+
and E2B backends.
|
| 220 |
+
|
| 221 |
+
Args:
|
| 222 |
+
pattern (`str`):
|
| 223 |
+
The glob pattern to match against (e.g. ``**/*.py``).
|
| 224 |
+
path (`str | None`, optional):
|
| 225 |
+
Base directory to search from. Defaults to the current
|
| 226 |
+
working directory when ``None``.
|
| 227 |
+
|
| 228 |
+
Returns:
|
| 229 |
+
`ToolChunk`:
|
| 230 |
+
On success, the matched file paths joined by newlines
|
| 231 |
+
(or a "no files found" message). If the base directory
|
| 232 |
+
is missing or the helper fails, an error chunk with
|
| 233 |
+
``ToolResultState.ERROR``.
|
| 234 |
+
"""
|
| 235 |
+
base_dir = path if path else await self._backend.getcwd()
|
| 236 |
+
|
| 237 |
+
# The base must be an existing directory; a regular file would
|
| 238 |
+
# otherwise be accepted here and fail later with a confusing
|
| 239 |
+
# error from the helper.
|
| 240 |
+
if not await self._backend.is_dir(base_dir):
|
| 241 |
+
return ToolChunk(
|
| 242 |
+
content=[
|
| 243 |
+
TextBlock(text=f"Directory not found: {base_dir}"),
|
| 244 |
+
],
|
| 245 |
+
state=ToolResultState.ERROR,
|
| 246 |
+
is_last=True,
|
| 247 |
+
)
|
| 248 |
+
|
| 249 |
+
# Invoke the glob helper script via exec_shell as an argv list
|
| 250 |
+
# (run directly, without a shell, so no platform-specific
|
| 251 |
+
# quoting is needed). Use the current interpreter locally
|
| 252 |
+
# (``python3`` may be absent, e.g. on Windows or venvs exposing
|
| 253 |
+
# only ``python``); remote backends run inside Linux images
|
| 254 |
+
# where ``python3`` is the safe choice.
|
| 255 |
+
python = sys.executable if self._is_local else "python3"
|
| 256 |
+
command = [
|
| 257 |
+
python,
|
| 258 |
+
self._glob_helper_path,
|
| 259 |
+
"--pattern",
|
| 260 |
+
pattern,
|
| 261 |
+
"--base-dir",
|
| 262 |
+
base_dir,
|
| 263 |
+
]
|
| 264 |
+
result = await self._backend.exec_shell(command, timeout=30.0)
|
| 265 |
+
|
| 266 |
+
# A non-zero exit means the helper itself failed (missing
|
| 267 |
+
# interpreter/script, permission error, …) — surface it rather
|
| 268 |
+
# than masking it as an empty match.
|
| 269 |
+
if not result.ok():
|
| 270 |
+
stderr = result.stderr.decode("utf-8", errors="replace").strip()
|
| 271 |
+
return ToolChunk(
|
| 272 |
+
content=[
|
| 273 |
+
TextBlock(
|
| 274 |
+
text=f"Glob helper failed: {stderr}"
|
| 275 |
+
if stderr
|
| 276 |
+
else "Glob helper failed with no error output.",
|
| 277 |
+
),
|
| 278 |
+
],
|
| 279 |
+
state=ToolResultState.ERROR,
|
| 280 |
+
is_last=True,
|
| 281 |
+
)
|
| 282 |
+
|
| 283 |
+
try:
|
| 284 |
+
matches = json.loads(
|
| 285 |
+
result.stdout.decode("utf-8", errors="replace"),
|
| 286 |
+
)
|
| 287 |
+
except (json.JSONDecodeError, ValueError):
|
| 288 |
+
matches = []
|
| 289 |
+
|
| 290 |
+
if len(matches) == 0:
|
| 291 |
+
return ToolChunk(
|
| 292 |
+
content=[
|
| 293 |
+
TextBlock(
|
| 294 |
+
text=f"No files found matching pattern: {pattern}",
|
| 295 |
+
),
|
| 296 |
+
],
|
| 297 |
+
state=ToolResultState.RUNNING,
|
| 298 |
+
is_last=True,
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
return ToolChunk(
|
| 302 |
+
content=[TextBlock(text="\n".join(matches))],
|
| 303 |
+
state=ToolResultState.RUNNING,
|
| 304 |
+
is_last=True,
|
| 305 |
+
)
|
src/agentscope/tool/_builtin/_grep.py
ADDED
|
@@ -0,0 +1,488 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The grep tool in agentscope."""
|
| 3 |
+
import fnmatch
|
| 4 |
+
from typing import Any, List, Literal
|
| 5 |
+
|
| 6 |
+
from .._base import ToolBase, ToolMiddlewareBase
|
| 7 |
+
from ...permission import (
|
| 8 |
+
PermissionContext,
|
| 9 |
+
PermissionDecision,
|
| 10 |
+
PermissionBehavior,
|
| 11 |
+
PermissionRule,
|
| 12 |
+
)
|
| 13 |
+
from .._response import ToolChunk
|
| 14 |
+
from ...message import TextBlock, ToolResultState
|
| 15 |
+
from ._backend import BackendBase
|
| 16 |
+
|
| 17 |
+
# Version control system directories to exclude from searches
|
| 18 |
+
VCS_DIRECTORIES_TO_EXCLUDE = [
|
| 19 |
+
".git",
|
| 20 |
+
".svn",
|
| 21 |
+
".hg",
|
| 22 |
+
".bzr",
|
| 23 |
+
".jj",
|
| 24 |
+
".sl",
|
| 25 |
+
]
|
| 26 |
+
|
| 27 |
+
# Default cap on grep results when head_limit is unspecified
|
| 28 |
+
DEFAULT_HEAD_LIMIT = 250
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
class RipgrepTimeoutError(Exception):
|
| 32 |
+
"""Custom error class for ripgrep timeouts."""
|
| 33 |
+
|
| 34 |
+
def __init__(self, message: str, partial_results: list[str]):
|
| 35 |
+
super().__init__(message)
|
| 36 |
+
self.partial_results = partial_results
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
class Grep(ToolBase):
|
| 40 |
+
"""The grep tool for searching file contents using ripgrep."""
|
| 41 |
+
|
| 42 |
+
name: str = "Grep"
|
| 43 |
+
"""The tool name presented to the agent."""
|
| 44 |
+
|
| 45 |
+
description: str = """A powerful search tool built on ripgrep
|
| 46 |
+
|
| 47 |
+
Usage:
|
| 48 |
+
- 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.
|
| 49 |
+
- Supports full regex syntax (e.g., "log.*Error", "function\\s+\\w+")
|
| 50 |
+
- Filter files with glob parameter (e.g., "*.js", "**/*.tsx") or type parameter (e.g., "js", "py", "rust")
|
| 51 |
+
- Output modes: "content" shows matching lines, "files_with_matches" shows only file paths (default), "count" shows match counts per file
|
| 52 |
+
- Context lines: use context parameter or -A/-B/-C for lines after/before/around matches
|
| 53 |
+
- Case-insensitive search: set i to true
|
| 54 |
+
- Multiline regex: set multiline to true for patterns spanning multiple lines
|
| 55 |
+
- Limit results: use head_limit to cap the number of results returned""" # noqa: E501
|
| 56 |
+
"""The description presented to the agent."""
|
| 57 |
+
|
| 58 |
+
input_schema: dict[str, Any] = {
|
| 59 |
+
"type": "object",
|
| 60 |
+
"properties": {
|
| 61 |
+
"pattern": {
|
| 62 |
+
"type": "string",
|
| 63 |
+
"description": "The regular expression pattern to search "
|
| 64 |
+
"for in file contents.",
|
| 65 |
+
},
|
| 66 |
+
"path": {
|
| 67 |
+
"type": "string",
|
| 68 |
+
"description": "File or directory to search in. Defaults "
|
| 69 |
+
"to current working directory.",
|
| 70 |
+
},
|
| 71 |
+
"output_mode": {
|
| 72 |
+
"type": "string",
|
| 73 |
+
"enum": ["content", "files_with_matches", "count"],
|
| 74 |
+
"description": "Output mode: 'content' shows matching lines "
|
| 75 |
+
"(supports -A/-B/-C context, -n line numbers, "
|
| 76 |
+
"head_limit), 'files_with_matches' shows file "
|
| 77 |
+
"paths (supports head_limit), 'count' shows "
|
| 78 |
+
"match counts (supports head_limit). "
|
| 79 |
+
"Defaults to 'files_with_matches'.",
|
| 80 |
+
"default": "files_with_matches",
|
| 81 |
+
},
|
| 82 |
+
"glob": {
|
| 83 |
+
"type": "string",
|
| 84 |
+
"description": "Glob pattern to filter files (e.g., '*.js', "
|
| 85 |
+
"'*.{ts,tsx}').",
|
| 86 |
+
},
|
| 87 |
+
"type": {
|
| 88 |
+
"type": "string",
|
| 89 |
+
"description": "File type to search (rg --type). "
|
| 90 |
+
"Common types: js, py, rust, go, java, etc.",
|
| 91 |
+
},
|
| 92 |
+
"-A": {
|
| 93 |
+
"type": "integer",
|
| 94 |
+
"description": "Number of lines to show after each match. "
|
| 95 |
+
"Requires output_mode: 'content'.",
|
| 96 |
+
},
|
| 97 |
+
"-B": {
|
| 98 |
+
"type": "integer",
|
| 99 |
+
"description": "Number of lines to show before each match. "
|
| 100 |
+
"Requires output_mode: 'content'.",
|
| 101 |
+
},
|
| 102 |
+
"-C": {
|
| 103 |
+
"type": "integer",
|
| 104 |
+
"description": "Alias for context.",
|
| 105 |
+
},
|
| 106 |
+
"context": {
|
| 107 |
+
"type": "integer",
|
| 108 |
+
"description": "Number of context lines to show before and "
|
| 109 |
+
"after matches. Requires output_mode: "
|
| 110 |
+
"'content'.",
|
| 111 |
+
},
|
| 112 |
+
"n": {
|
| 113 |
+
"type": "boolean",
|
| 114 |
+
"description": "Show line numbers in output. Requires "
|
| 115 |
+
"output_mode: 'content'. Defaults to true.",
|
| 116 |
+
"default": True,
|
| 117 |
+
},
|
| 118 |
+
"i": {
|
| 119 |
+
"type": "boolean",
|
| 120 |
+
"description": "Case insensitive search.",
|
| 121 |
+
"default": False,
|
| 122 |
+
},
|
| 123 |
+
"case_insensitive": {
|
| 124 |
+
"type": "boolean",
|
| 125 |
+
"description": "Case insensitive search (alias for i).",
|
| 126 |
+
"default": False,
|
| 127 |
+
},
|
| 128 |
+
"multiline": {
|
| 129 |
+
"type": "boolean",
|
| 130 |
+
"description": "Enable multiline mode where . matches "
|
| 131 |
+
"newlines and patterns can span lines. "
|
| 132 |
+
"Default: false.",
|
| 133 |
+
"default": False,
|
| 134 |
+
},
|
| 135 |
+
"head_limit": {
|
| 136 |
+
"type": "integer",
|
| 137 |
+
"description": "Limit output to first N lines/entries. "
|
| 138 |
+
"Defaults to 250 when unspecified. "
|
| 139 |
+
"Pass 0 for unlimited.",
|
| 140 |
+
"minimum": 0,
|
| 141 |
+
},
|
| 142 |
+
"offset": {
|
| 143 |
+
"type": "integer",
|
| 144 |
+
"description": "Skip first N lines/entries before applying "
|
| 145 |
+
"head_limit. Defaults to 0.",
|
| 146 |
+
"default": 0,
|
| 147 |
+
"minimum": 0,
|
| 148 |
+
},
|
| 149 |
+
},
|
| 150 |
+
"required": ["pattern"],
|
| 151 |
+
}
|
| 152 |
+
|
| 153 |
+
is_mcp: bool = False
|
| 154 |
+
is_read_only: bool = True
|
| 155 |
+
is_concurrency_safe: bool = True
|
| 156 |
+
is_external_tool: bool = False
|
| 157 |
+
is_state_injected: bool = False
|
| 158 |
+
|
| 159 |
+
def __init__(
|
| 160 |
+
self,
|
| 161 |
+
middlewares: List[ToolMiddlewareBase] | None = None,
|
| 162 |
+
backend: BackendBase | None = None,
|
| 163 |
+
) -> None:
|
| 164 |
+
"""Initialize the grep tool.
|
| 165 |
+
|
| 166 |
+
Args:
|
| 167 |
+
middlewares (`List[ToolMiddlewareBase] | None`, optional):
|
| 168 |
+
Tool middlewares wrapping the tool execution.
|
| 169 |
+
backend (`BackendBase | None`, optional):
|
| 170 |
+
The sandbox backend to use for shell execution. When
|
| 171 |
+
``None``, a :class:`LocalBackend` is created.
|
| 172 |
+
Ripgrep is always invoked via ``exec_shell`` so that
|
| 173 |
+
the same code path works for local, Docker, and E2B
|
| 174 |
+
backends.
|
| 175 |
+
"""
|
| 176 |
+
from ._backend import LocalBackend
|
| 177 |
+
|
| 178 |
+
super().__init__(middlewares=middlewares)
|
| 179 |
+
self._backend = backend or LocalBackend()
|
| 180 |
+
|
| 181 |
+
async def check_permissions(
|
| 182 |
+
self,
|
| 183 |
+
tool_input: dict[str, Any],
|
| 184 |
+
context: PermissionContext,
|
| 185 |
+
) -> PermissionDecision:
|
| 186 |
+
"""Check permissions for grep search."""
|
| 187 |
+
return PermissionDecision(
|
| 188 |
+
behavior=PermissionBehavior.PASSTHROUGH,
|
| 189 |
+
message="Grep search is read-only.",
|
| 190 |
+
)
|
| 191 |
+
|
| 192 |
+
async def match_rule(
|
| 193 |
+
self,
|
| 194 |
+
rule_content: str | None,
|
| 195 |
+
tool_input: dict[str, Any],
|
| 196 |
+
) -> bool:
|
| 197 |
+
"""Check if a permission rule matches the grep search path.
|
| 198 |
+
|
| 199 |
+
Matches rule_content as a glob pattern against the "path" parameter.
|
| 200 |
+
If no path is given, falls back to the current working directory.
|
| 201 |
+
If rule_content is None, matches all invocations (tool-name-level
|
| 202 |
+
rule).
|
| 203 |
+
|
| 204 |
+
Args:
|
| 205 |
+
rule_content (`str | None`):
|
| 206 |
+
Glob pattern to match against the search path (e.g., "src/**"),
|
| 207 |
+
or None to match all invocations
|
| 208 |
+
tool_input (`dict[str, Any]`):
|
| 209 |
+
The tool input data containing optional "path" key
|
| 210 |
+
|
| 211 |
+
Returns:
|
| 212 |
+
`bool`:
|
| 213 |
+
True if the glob pattern matches the search path, False
|
| 214 |
+
otherwise
|
| 215 |
+
"""
|
| 216 |
+
# None = tool-name-level rule, matches everything
|
| 217 |
+
if rule_content is None:
|
| 218 |
+
return True
|
| 219 |
+
|
| 220 |
+
path = tool_input.get("path", "")
|
| 221 |
+
if not path:
|
| 222 |
+
path = await self._backend.getcwd()
|
| 223 |
+
return fnmatch.fnmatch(path, rule_content)
|
| 224 |
+
|
| 225 |
+
async def generate_suggestions(
|
| 226 |
+
self,
|
| 227 |
+
tool_input: dict[str, Any],
|
| 228 |
+
) -> List[PermissionRule]:
|
| 229 |
+
"""Generate suggested permission rules for the grep search path.
|
| 230 |
+
|
| 231 |
+
Suggests a rule based on the search path. If no path is provided,
|
| 232 |
+
suggests a rule for the current directory.
|
| 233 |
+
|
| 234 |
+
Args:
|
| 235 |
+
tool_input (`dict[str, Any]`):
|
| 236 |
+
The tool input data containing optional "path" key
|
| 237 |
+
|
| 238 |
+
Returns:
|
| 239 |
+
`List[PermissionRule]`:
|
| 240 |
+
A single suggested rule covering the search directory
|
| 241 |
+
"""
|
| 242 |
+
backend_cwd = await self._backend.getcwd()
|
| 243 |
+
path = tool_input.get("path") or backend_cwd
|
| 244 |
+
|
| 245 |
+
abs_path = self._backend.abspath(path, cwd=backend_cwd)
|
| 246 |
+
# Glob patterns are POSIX-style strings (matched by fnmatch),
|
| 247 |
+
# not real filesystem paths — do NOT use backend.join_path here.
|
| 248 |
+
pattern = abs_path.rstrip("/\\") + "/**"
|
| 249 |
+
|
| 250 |
+
return [
|
| 251 |
+
PermissionRule(
|
| 252 |
+
tool_name=self.name,
|
| 253 |
+
rule_content=pattern,
|
| 254 |
+
behavior=PermissionBehavior.ALLOW,
|
| 255 |
+
source="suggested",
|
| 256 |
+
),
|
| 257 |
+
]
|
| 258 |
+
|
| 259 |
+
def _apply_head_limit(
|
| 260 |
+
self,
|
| 261 |
+
items: list[str],
|
| 262 |
+
limit: int | None,
|
| 263 |
+
offset: int = 0,
|
| 264 |
+
) -> tuple[list[str], int | None]:
|
| 265 |
+
"""Apply head_limit and offset to a list of items.
|
| 266 |
+
|
| 267 |
+
Returns (sliced_items, applied_limit_if_truncated).
|
| 268 |
+
"""
|
| 269 |
+
if limit == 0:
|
| 270 |
+
return items[offset:], None
|
| 271 |
+
effective_limit = limit if limit is not None else DEFAULT_HEAD_LIMIT
|
| 272 |
+
sliced = items[offset : offset + effective_limit]
|
| 273 |
+
was_truncated = len(items) - offset > effective_limit
|
| 274 |
+
return sliced, (effective_limit if was_truncated else None)
|
| 275 |
+
|
| 276 |
+
async def _run_ripgrep(
|
| 277 |
+
self,
|
| 278 |
+
args: list[str],
|
| 279 |
+
search_path: str,
|
| 280 |
+
timeout: int = 30,
|
| 281 |
+
) -> list[str]:
|
| 282 |
+
"""Run ripgrep and return output lines.
|
| 283 |
+
|
| 284 |
+
Builds an argument vector and dispatches it through
|
| 285 |
+
``backend.exec_shell`` (which runs the program directly, without
|
| 286 |
+
a shell), so the same code path works for local, Docker, and E2B
|
| 287 |
+
backends and needs no platform-specific argument quoting.
|
| 288 |
+
"""
|
| 289 |
+
command = ["rg", *args, search_path]
|
| 290 |
+
|
| 291 |
+
result = await self._backend.exec_shell(
|
| 292 |
+
command,
|
| 293 |
+
timeout=float(timeout),
|
| 294 |
+
)
|
| 295 |
+
|
| 296 |
+
if result.exit_code == -1 and result.stderr == b"timed out":
|
| 297 |
+
raise RipgrepTimeoutError(
|
| 298 |
+
f"Ripgrep search timed out after {timeout} seconds. "
|
| 299 |
+
"Try searching a more specific path or pattern.",
|
| 300 |
+
[],
|
| 301 |
+
)
|
| 302 |
+
|
| 303 |
+
# returncode 0 = matches found, 1 = no matches
|
| 304 |
+
if result.exit_code not in (0, 1):
|
| 305 |
+
error_msg = result.stderr.decode(
|
| 306 |
+
"utf-8",
|
| 307 |
+
errors="ignore",
|
| 308 |
+
).strip()
|
| 309 |
+
raise RuntimeError(
|
| 310 |
+
f"ripgrep error (code {result.exit_code}): {error_msg}",
|
| 311 |
+
)
|
| 312 |
+
|
| 313 |
+
raw = result.stdout.decode("utf-8", errors="ignore")
|
| 314 |
+
|
| 315 |
+
lines = [
|
| 316 |
+
line.rstrip("\r") for line in raw.split("\n") if line.rstrip("\r")
|
| 317 |
+
]
|
| 318 |
+
return lines
|
| 319 |
+
|
| 320 |
+
async def call( # type: ignore[override]
|
| 321 |
+
self,
|
| 322 |
+
pattern: str,
|
| 323 |
+
path: str | None = None,
|
| 324 |
+
output_mode: Literal[
|
| 325 |
+
"content",
|
| 326 |
+
"files_with_matches",
|
| 327 |
+
"count",
|
| 328 |
+
] = "files_with_matches",
|
| 329 |
+
glob: str | None = None,
|
| 330 |
+
type: str | None = None, # pylint: disable=redefined-builtin
|
| 331 |
+
i: bool = False,
|
| 332 |
+
case_insensitive: bool = False,
|
| 333 |
+
context: int | None = None,
|
| 334 |
+
multiline: bool = False,
|
| 335 |
+
head_limit: int | None = None,
|
| 336 |
+
offset: int = 0,
|
| 337 |
+
n: bool = True,
|
| 338 |
+
**kwargs: Any,
|
| 339 |
+
) -> ToolChunk:
|
| 340 |
+
"""Execute the grep search using ripgrep.
|
| 341 |
+
|
| 342 |
+
Args:
|
| 343 |
+
pattern: The regex pattern to search for
|
| 344 |
+
path: The directory or file path to search in
|
| 345 |
+
output_mode: Output mode ('content', 'files_with_matches', 'count')
|
| 346 |
+
glob: Glob pattern to filter files
|
| 347 |
+
type: File type to filter by (rg --type)
|
| 348 |
+
i: Case-insensitive search (rg -i)
|
| 349 |
+
case_insensitive: Alias for i (backward compatibility)
|
| 350 |
+
context: Number of context lines around matches
|
| 351 |
+
multiline: Enable multiline regex matching
|
| 352 |
+
head_limit: Maximum number of results to return
|
| 353 |
+
(default 250, 0=unlimited)
|
| 354 |
+
offset: Skip first N results
|
| 355 |
+
n: Show line numbers (content mode only, default True)
|
| 356 |
+
**kwargs: Additional parameters (-A, -B, -C)
|
| 357 |
+
"""
|
| 358 |
+
search_path = path or await self._backend.getcwd()
|
| 359 |
+
|
| 360 |
+
if head_limit is not None and head_limit < 0:
|
| 361 |
+
return ToolChunk(
|
| 362 |
+
content=[
|
| 363 |
+
TextBlock(text="Error: head_limit must be non-negative."),
|
| 364 |
+
],
|
| 365 |
+
state=ToolResultState.ERROR,
|
| 366 |
+
is_last=True,
|
| 367 |
+
)
|
| 368 |
+
|
| 369 |
+
if offset < 0:
|
| 370 |
+
return ToolChunk(
|
| 371 |
+
content=[
|
| 372 |
+
TextBlock(text="Error: offset must be non-negative."),
|
| 373 |
+
],
|
| 374 |
+
state=ToolResultState.ERROR,
|
| 375 |
+
is_last=True,
|
| 376 |
+
)
|
| 377 |
+
|
| 378 |
+
args: list[str] = ["--hidden"]
|
| 379 |
+
|
| 380 |
+
# Exclude VCS directories
|
| 381 |
+
for vcs_dir in VCS_DIRECTORIES_TO_EXCLUDE:
|
| 382 |
+
args.extend(["--glob", f"!{vcs_dir}"])
|
| 383 |
+
|
| 384 |
+
# Limit line length to prevent base64/minified content
|
| 385 |
+
args.extend(["--max-columns", "500"])
|
| 386 |
+
|
| 387 |
+
# Multiline mode
|
| 388 |
+
if multiline:
|
| 389 |
+
args.extend(["-U", "--multiline-dotall"])
|
| 390 |
+
|
| 391 |
+
# Case-insensitive (support both i and case_insensitive
|
| 392 |
+
# for compatibility)
|
| 393 |
+
if i or case_insensitive:
|
| 394 |
+
args.append("-i")
|
| 395 |
+
|
| 396 |
+
# Output mode flags
|
| 397 |
+
if output_mode == "files_with_matches":
|
| 398 |
+
args.append("-l")
|
| 399 |
+
elif output_mode == "count":
|
| 400 |
+
args.append("-c")
|
| 401 |
+
|
| 402 |
+
# Line numbers (content mode only)
|
| 403 |
+
if n and output_mode == "content":
|
| 404 |
+
args.append("-n")
|
| 405 |
+
|
| 406 |
+
# Context flags (content mode only)
|
| 407 |
+
if output_mode == "content":
|
| 408 |
+
A = kwargs.get("-A")
|
| 409 |
+
B = kwargs.get("-B")
|
| 410 |
+
C = kwargs.get("-C")
|
| 411 |
+
|
| 412 |
+
if context is not None:
|
| 413 |
+
args.extend(["-C", str(context)])
|
| 414 |
+
elif C is not None:
|
| 415 |
+
args.extend(["-C", str(C)])
|
| 416 |
+
else:
|
| 417 |
+
if B is not None:
|
| 418 |
+
args.extend(["-B", str(B)])
|
| 419 |
+
if A is not None:
|
| 420 |
+
args.extend(["-A", str(A)])
|
| 421 |
+
|
| 422 |
+
# Pattern — use -e if it starts with a dash
|
| 423 |
+
if pattern.startswith("-"):
|
| 424 |
+
args.extend(["-e", pattern])
|
| 425 |
+
else:
|
| 426 |
+
args.append(pattern)
|
| 427 |
+
|
| 428 |
+
# File type filter
|
| 429 |
+
if type is not None:
|
| 430 |
+
args.extend(["--type", type])
|
| 431 |
+
|
| 432 |
+
# Glob filter
|
| 433 |
+
if glob is not None:
|
| 434 |
+
raw_patterns = glob.split()
|
| 435 |
+
glob_patterns: list[str] = []
|
| 436 |
+
for raw in raw_patterns:
|
| 437 |
+
if "{" in raw and "}" in raw:
|
| 438 |
+
glob_patterns.append(raw)
|
| 439 |
+
else:
|
| 440 |
+
glob_patterns.extend(p for p in raw.split(",") if p)
|
| 441 |
+
for gp in glob_patterns:
|
| 442 |
+
args.extend(["--glob", gp])
|
| 443 |
+
|
| 444 |
+
try:
|
| 445 |
+
results = await self._run_ripgrep(args, search_path)
|
| 446 |
+
except RipgrepTimeoutError as e:
|
| 447 |
+
return ToolChunk(
|
| 448 |
+
content=[TextBlock(text=str(e))],
|
| 449 |
+
state=ToolResultState.ERROR,
|
| 450 |
+
is_last=True,
|
| 451 |
+
)
|
| 452 |
+
except RuntimeError as e:
|
| 453 |
+
return ToolChunk(
|
| 454 |
+
content=[TextBlock(text=str(e))],
|
| 455 |
+
state=ToolResultState.ERROR,
|
| 456 |
+
is_last=True,
|
| 457 |
+
)
|
| 458 |
+
|
| 459 |
+
if not results:
|
| 460 |
+
return ToolChunk(
|
| 461 |
+
content=[
|
| 462 |
+
TextBlock(text=f"No matches found for pattern: {pattern}"),
|
| 463 |
+
],
|
| 464 |
+
state=ToolResultState.SUCCESS,
|
| 465 |
+
is_last=True,
|
| 466 |
+
)
|
| 467 |
+
|
| 468 |
+
limited, applied_limit = self._apply_head_limit(
|
| 469 |
+
results,
|
| 470 |
+
head_limit,
|
| 471 |
+
offset,
|
| 472 |
+
)
|
| 473 |
+
|
| 474 |
+
suffix = ""
|
| 475 |
+
if applied_limit is not None:
|
| 476 |
+
suffix = (
|
| 477 |
+
f"\n\n[Showing results with pagination = "
|
| 478 |
+
f"limit: {applied_limit}"
|
| 479 |
+
)
|
| 480 |
+
if offset:
|
| 481 |
+
suffix += f", offset: {offset}"
|
| 482 |
+
suffix += "]"
|
| 483 |
+
|
| 484 |
+
return ToolChunk(
|
| 485 |
+
content=[TextBlock(text="\n".join(limited) + suffix)],
|
| 486 |
+
state=ToolResultState.SUCCESS,
|
| 487 |
+
is_last=True,
|
| 488 |
+
)
|
src/agentscope/tool/_builtin/_meta.py
ADDED
|
@@ -0,0 +1,130 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""The meta tool class."""
|
| 3 |
+
from typing import Any, List
|
| 4 |
+
|
| 5 |
+
from pydantic import Field, create_model
|
| 6 |
+
from jinja2 import Template
|
| 7 |
+
|
| 8 |
+
from .._tool_group import ToolGroup
|
| 9 |
+
from ...permission import (
|
| 10 |
+
PermissionContext,
|
| 11 |
+
PermissionDecision,
|
| 12 |
+
PermissionBehavior,
|
| 13 |
+
)
|
| 14 |
+
from .._response import ToolChunk
|
| 15 |
+
from .._base import ToolBase, ToolMiddlewareBase
|
| 16 |
+
from ...exception import DeveloperOrientedException
|
| 17 |
+
from ...message import TextBlock
|
| 18 |
+
from ...state import AgentState
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
class ResetTools(ToolBase):
|
| 22 |
+
"""A meta tool allows agent to self-manage its equipped tools by
|
| 23 |
+
activating or deactivating tool groups dynamically."""
|
| 24 |
+
|
| 25 |
+
name: str = "reset_tools"
|
| 26 |
+
description: str = (
|
| 27 |
+
"This tool allows you to reset your equipped tools based on your "
|
| 28 |
+
"current task requirements. These tools are organized into different "
|
| 29 |
+
"groups, and you can activate/deactivate them by specifying the "
|
| 30 |
+
"boolean values for each group in the input.\n\n"
|
| 31 |
+
"**Important: The input booleans are the final state of their "
|
| 32 |
+
"corresponding tool groups, not incremental changes.** Any group not "
|
| 33 |
+
"explicitly set to True will be deactivated, regardless of its "
|
| 34 |
+
"previous state.\n\n"
|
| 35 |
+
"**Best practice**: Actively manage your tool groups——activate only "
|
| 36 |
+
"what you need for the current task, and promptly deactivate groups "
|
| 37 |
+
"as soon as they are no longer needed to conserve context space.\n\n"
|
| 38 |
+
"This tool will return the usage instructions for the activated tool "
|
| 39 |
+
"groups, which you **MUST pay attention to and follow**. You can "
|
| 40 |
+
"also reuse this tool to re-check the instructions."
|
| 41 |
+
)
|
| 42 |
+
is_mcp: bool = False
|
| 43 |
+
is_read_only: bool = False
|
| 44 |
+
is_concurrency_safe: bool = True
|
| 45 |
+
is_external_tool: bool = False
|
| 46 |
+
is_state_injected: bool = True
|
| 47 |
+
|
| 48 |
+
def __init__(
|
| 49 |
+
self,
|
| 50 |
+
groups: list[ToolGroup],
|
| 51 |
+
response_template: str,
|
| 52 |
+
middlewares: List[ToolMiddlewareBase] | None = None,
|
| 53 |
+
) -> None:
|
| 54 |
+
"""Initialize the meta tool with the current tool groups."""
|
| 55 |
+
super().__init__(middlewares=middlewares)
|
| 56 |
+
self.groups = groups
|
| 57 |
+
self.response_template = response_template
|
| 58 |
+
|
| 59 |
+
@property
|
| 60 |
+
def input_schema(self) -> dict[str, Any]: # type: ignore[override]
|
| 61 |
+
"""Dynamically generate the input schema based on the current
|
| 62 |
+
available tool groups."""
|
| 63 |
+
fields = {}
|
| 64 |
+
for group in self.groups:
|
| 65 |
+
if group.name == "basic":
|
| 66 |
+
continue
|
| 67 |
+
fields[group.name] = (
|
| 68 |
+
bool,
|
| 69 |
+
Field(
|
| 70 |
+
default=False,
|
| 71 |
+
description=group.description,
|
| 72 |
+
),
|
| 73 |
+
)
|
| 74 |
+
|
| 75 |
+
model = create_model("_DynamicModel", **fields)
|
| 76 |
+
schema = model.model_json_schema()
|
| 77 |
+
return schema
|
| 78 |
+
|
| 79 |
+
async def check_permissions(
|
| 80 |
+
self,
|
| 81 |
+
tool_input: dict[str, Any],
|
| 82 |
+
context: PermissionContext,
|
| 83 |
+
) -> PermissionDecision:
|
| 84 |
+
"""The meta tool is always allowed to be called."""
|
| 85 |
+
return PermissionDecision(
|
| 86 |
+
behavior=PermissionBehavior.ALLOW,
|
| 87 |
+
message="The meta tool is always allowed to be called.",
|
| 88 |
+
)
|
| 89 |
+
|
| 90 |
+
async def call(
|
| 91 |
+
self,
|
| 92 |
+
_agent_state: AgentState,
|
| 93 |
+
**kwargs: Any,
|
| 94 |
+
) -> ToolChunk:
|
| 95 |
+
"""Activate or deactivate tool groups based on the input arguments,
|
| 96 |
+
and return their usage instructions."""
|
| 97 |
+
if _agent_state is None:
|
| 98 |
+
raise DeveloperOrientedException(
|
| 99 |
+
"Error: ResetTools requires state to be provided.",
|
| 100 |
+
)
|
| 101 |
+
|
| 102 |
+
# Deactivate all tool groups first
|
| 103 |
+
_agent_state.tool_context.activated_groups.clear()
|
| 104 |
+
|
| 105 |
+
to_activate = []
|
| 106 |
+
for key, value in kwargs.items():
|
| 107 |
+
if not isinstance(value, bool):
|
| 108 |
+
return ToolChunk(
|
| 109 |
+
content=[
|
| 110 |
+
TextBlock(
|
| 111 |
+
text=f"Invalid arguments: the argument {key} "
|
| 112 |
+
f"should be a bool value, but got {type(value)}.",
|
| 113 |
+
),
|
| 114 |
+
],
|
| 115 |
+
)
|
| 116 |
+
|
| 117 |
+
if value:
|
| 118 |
+
to_activate.append(key)
|
| 119 |
+
|
| 120 |
+
_agent_state.tool_context.activated_groups.extend(to_activate)
|
| 121 |
+
|
| 122 |
+
template = Template(self.response_template)
|
| 123 |
+
activated_groups = [_ for _ in self.groups if _.name in to_activate]
|
| 124 |
+
return ToolChunk(
|
| 125 |
+
content=[
|
| 126 |
+
TextBlock(
|
| 127 |
+
text=template.render(groups=activated_groups),
|
| 128 |
+
),
|
| 129 |
+
],
|
| 130 |
+
)
|