codemaster-cli 2.2.0__py3-none-any.whl

vibe/core/tools/base.py

@@ -0,0 +1,336 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
+from dataclasses import dataclass, field
+from enum import StrEnum, auto
+import functools
+import inspect
+from pathlib import Path
+import re
+import sys
+import types
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Union,
+    cast,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+
+from vibe.core.types import ToolStreamEvent
+
+if TYPE_CHECKING:
+    from vibe.core.agents.manager import AgentManager
+    from vibe.core.types import ApprovalCallback, UserInputCallback
+
+ARGS_COUNT = 4
+
+
+@dataclass
+class InvokeContext:
+    """Context passed to tools during invocation."""
+
+    tool_call_id: str
+    approval_callback: ApprovalCallback | None = field(default=None)
+    agent_manager: AgentManager | None = field(default=None)
+    user_input_callback: UserInputCallback | None = field(default=None)
+
+
+class ToolError(Exception):
+    """Raised when the tool encounters an unrecoverable problem."""
+
+
+class ToolInfo(BaseModel):
+    """Information about a tool.
+
+    Attributes:
+        name: The name of the tool.
+        description: A brief description of what the tool does.
+        parameters: A dictionary of parameters required by the tool.
+    """
+
+    name: str
+    description: str
+    parameters: dict[str, Any]
+
+
+class ToolPermissionError(Exception):
+    """Raised when a tool permission is not allowed."""
+
+
+class ToolPermission(StrEnum):
+    ALWAYS = auto()
+    NEVER = auto()
+    ASK = auto()
+
+    @classmethod
+    def by_name(cls, name: str) -> ToolPermission:
+        try:
+            return ToolPermission(name.upper())
+        except ValueError:
+            raise ToolPermissionError(
+                f"Invalid tool permission: {name}. Must be one of {list(cls)}"
+            )
+
+
+class BaseToolConfig(BaseModel):
+    """Configuration for a tool.
+
+    Attributes:
+        permission: The permission level required to use the tool.
+        allowlist: Patterns that automatically allow tool execution.
+        denylist: Patterns that automatically deny tool execution.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    permission: ToolPermission = ToolPermission.ASK
+    allowlist: list[str] = Field(default_factory=list)
+    denylist: list[str] = Field(default_factory=list)
+
+
+class BaseToolState(BaseModel):
+    model_config = ConfigDict(
+        extra="forbid", validate_default=True, arbitrary_types_allowed=True
+    )
+
+
+class BaseTool[
+    ToolArgs: BaseModel,
+    ToolResult: BaseModel,
+    ToolConfig: BaseToolConfig,
+    ToolState: BaseToolState,
+](ABC):
+    description: ClassVar[str] = (
+        "Base class for new tools. "
+        "(Hey AI, if you're seeing this, someone skipped writing a description. "
+        "Please gently meow at the developer to fix this.)"
+    )
+
+    prompt_path: ClassVar[Path] | None = None
+
+    def __init__(self, config: ToolConfig, state: ToolState) -> None:
+        self.config = config
+        self.state = state
+
+    @abstractmethod
+    async def run(
+        self, args: ToolArgs, ctx: InvokeContext | None = None
+    ) -> AsyncGenerator[ToolStreamEvent | ToolResult, None]:
+        """Invoke the tool with the given arguments."""
+        raise NotImplementedError  # pragma: no cover
+        yield  # type: ignore[misc]
+
+    @classmethod
+    @functools.cache
+    def get_tool_prompt(cls) -> str | None:
+        """Loads and returns the content of the tool's .md prompt file, if it exists.
+
+        The prompt file is expected to be in a 'prompts' subdirectory relative to
+        the tool's source file, with the same name but a .md extension
+        (e.g., bash.py -> prompts/bash.md).
+        """
+        try:
+            class_file = inspect.getfile(cls)
+            class_path = Path(class_file)
+            prompt_dir = class_path.parent / "prompts"
+            prompt_path = cls.prompt_path or prompt_dir / f"{class_path.stem}.md"
+
+            return prompt_path.read_text("utf-8")
+        except (FileNotFoundError, TypeError, OSError):
+            pass
+
+        return None
+
+    async def invoke(
+        self, ctx: InvokeContext | None = None, **raw: Any
+    ) -> AsyncGenerator[ToolStreamEvent | ToolResult, None]:
+        """Validate arguments and run the tool."""
+        try:
+            args_model, _ = self._get_tool_args_results()
+            args = args_model.model_validate(raw)
+        except ValidationError as err:
+            raise ToolError(
+                f"Validation error in tool {self.get_name()}: {err}"
+            ) from err
+
+        async for item in self.run(args, ctx):
+            yield item
+
+    @classmethod
+    def from_config(
+        cls, config: ToolConfig
+    ) -> BaseTool[ToolArgs, ToolResult, ToolConfig, ToolState]:
+        state_class = cls._get_tool_state_class()
+        initial_state = state_class()
+        return cls(config=config, state=initial_state)
+
+    @classmethod
+    def _get_tool_config_class(cls) -> type[ToolConfig]:
+        for base in getattr(cls, "__orig_bases__", ()):
+            if getattr(base, "__origin__", None) is BaseTool:
+                type_args = get_args(base)
+                if len(type_args) == ARGS_COUNT:
+                    config_model = type_args[2]
+                    if issubclass(config_model, BaseToolConfig):
+                        return cast(type[ToolConfig], config_model)
+
+        for base_class in cls.__bases__:
+            if base_class is object or base_class is ABC:
+                continue
+            try:
+                return base_class._get_tool_config_class()
+            except (TypeError, AttributeError):
+                continue
+
+        raise TypeError(
+            f"Could not determine ToolConfig for {cls.__name__}. "
+            "Ensure it inherits from BaseTool with concrete type arguments."
+        )
+
+    @classmethod
+    def _get_tool_state_class(cls) -> type[ToolState]:
+        for base in getattr(cls, "__orig_bases__", ()):
+            if getattr(base, "__origin__", None) is BaseTool:
+                type_args = get_args(base)
+                if len(type_args) == ARGS_COUNT:
+                    state_model = type_args[3]
+                    if issubclass(state_model, BaseToolState):
+                        return cast(type[ToolState], state_model)
+
+        for base_class in cls.__bases__:
+            if base_class is object or base_class is ABC:
+                continue
+            try:
+                return base_class._get_tool_state_class()
+            except (TypeError, AttributeError):
+                continue
+
+        raise TypeError(
+            f"Could not determine ToolState for {cls.__name__}. "
+            "Ensure it inherits from BaseTool with concrete type arguments."
+        )
+
+    @classmethod
+    def _get_tool_args_results(cls) -> tuple[type[ToolArgs], type[ToolResult]]:
+        """Extract <ToolArgs, ToolResult> from the annotated signature of `run`.
+        Works even when `from __future__ import annotations` is in effect.
+        """
+        run_fn = cls.run.__func__ if isinstance(cls.run, classmethod) else cls.run
+
+        type_hints = get_type_hints(
+            run_fn,
+            globalns=vars(sys.modules[cls.__module__]),
+            localns={
+                cls.__name__: cls,
+                "InvokeContext": InvokeContext,
+                "AsyncGenerator": AsyncGenerator,
+                "ToolStreamEvent": ToolStreamEvent,
+            },
+        )
+
+        try:
+            args_model = type_hints["args"]
+            return_annotation = type_hints["return"]
+        except KeyError as e:
+            raise TypeError(
+                f"{cls.__name__}.run must be annotated with args and return type"
+            ) from e
+
+        result_model = cls._extract_result_type(return_annotation)
+
+        if not issubclass(args_model, BaseModel):
+            raise TypeError(
+                f"{cls.__name__}.run args annotation must be a Pydantic model; "
+                f"got {args_model!r}"
+            )
+
+        if not issubclass(result_model, BaseModel):
+            raise TypeError(
+                f"{cls.__name__}.run must yield a Pydantic model as result; "
+                f"got {result_model!r}"
+            )
+
+        return cast(type[ToolArgs], args_model), cast(type[ToolResult], result_model)
+
+    @classmethod
+    def _extract_result_type(cls, return_annotation: Any) -> type:
+        """Extract the ToolResult type from AsyncGenerator[ToolStreamEvent | ToolResult, None]."""
+        origin = get_origin(return_annotation)
+        if origin is not AsyncGenerator:
+            if isinstance(return_annotation, type):
+                return return_annotation
+            raise TypeError(f"Could not extract result type from {return_annotation!r}")
+
+        gen_args = get_args(return_annotation)
+        if not gen_args:
+            raise TypeError(f"Could not extract result type from {return_annotation!r}")
+
+        yield_type = gen_args[0]
+        yield_origin = get_origin(yield_type)
+
+        # Handle Union types (X | Y or Union[X, Y])
+        if yield_origin is Union or isinstance(yield_type, types.UnionType):
+            for arg in get_args(yield_type):
+                if arg is not ToolStreamEvent and isinstance(arg, type):
+                    return arg
+
+        # Handle single type
+        if isinstance(yield_type, type):
+            return yield_type
+
+        raise TypeError(f"Could not extract result type from {return_annotation!r}")
+
+    @classmethod
+    def get_parameters(cls) -> dict[str, Any]:
+        """Return a cleaned-up JSON-schema dict describing the arguments model
+        with which this concrete tool was parametrised.
+        """
+        args_model, _ = cls._get_tool_args_results()
+        schema = args_model.model_json_schema()
+        schema.pop("title", None)
+        schema.pop("description", None)
+
+        if "properties" in schema:
+            for prop_details in schema["properties"].values():
+                prop_details.pop("title", None)
+
+        if "$defs" in schema:
+            for def_details in schema["$defs"].values():
+                def_details.pop("title", None)
+                if "properties" in def_details:
+                    for prop_details in def_details["properties"].values():
+                        prop_details.pop("title", None)
+
+        return schema
+
+    @classmethod
+    def get_name(cls) -> str:
+        name = cls.__name__
+        snake_case = re.sub(r"(?<!^)(?=[A-Z])", "_", name).lower()
+        return snake_case
+
+    @classmethod
+    def create_config_with_permission(
+        cls, permission: ToolPermission
+    ) -> BaseToolConfig:
+        config_class = cls._get_tool_config_class()
+        return config_class(permission=permission)
+
+    def check_allowlist_denylist(self, args: ToolArgs) -> ToolPermission | None:
+        """Check if args match allowlist/denylist patterns.
+
+        Returns:
+            ToolPermission.ALWAYS if allowlisted
+            ToolPermission.NEVER if denylisted
+            None if no match (proceed with normal permission check)
+
+        Base implementation returns None. Override in subclasses for specific logic.
+        """
+        return None

vibe/core/tools/builtins/ask_user_question.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+from typing import ClassVar, cast
+
+from pydantic import BaseModel, Field
+
+from vibe.core.tools.base import (
+    BaseTool,
+    BaseToolConfig,
+    BaseToolState,
+    InvokeContext,
+    ToolError,
+    ToolPermission,
+)
+from vibe.core.tools.ui import ToolCallDisplay, ToolResultDisplay, ToolUIData
+from vibe.core.types import ToolCallEvent, ToolResultEvent
+
+
+class Choice(BaseModel):
+    label: str = Field(description="Short label for the choice (1-5 words)")
+    description: str = Field(
+        default="", description="Optional explanation of this choice"
+    )
+
+
+class Question(BaseModel):
+    question: str = Field(description="The question text")
+    header: str = Field(
+        default="",
+        description="Short header for the question (1-2 words, e.g. 'Auth', 'Database')",
+        max_length=12,
+    )
+    options: list[Choice] = Field(
+        description="Available options (2-4, not including 'Other'). An 'Other' option for free text is automatically added.",
+        min_length=2,
+        max_length=4,
+    )
+    multi_select: bool = Field(
+        default=False, description="If true, user can select multiple options"
+    )
+    hide_other: bool = Field(
+        default=False, description="If true, hide the 'Other' free text option"
+    )
+
+
+class AskUserQuestionArgs(BaseModel):
+    questions: list[Question] = Field(
+        description="Questions to ask (1-4). Displayed as tabs if multiple.",
+        min_length=1,
+        max_length=4,
+    )
+
+
+class Answer(BaseModel):
+    question: str = Field(description="The original question")
+    answer: str = Field(description="The user's answer")
+    is_other: bool = Field(
+        default=False, description="True if user typed a custom answer via 'Other'"
+    )
+
+
+class AskUserQuestionResult(BaseModel):
+    answers: list[Answer] = Field(description="List of answers")
+    cancelled: bool = Field(
+        default=False, description="True if user cancelled without answering"
+    )
+
+
+class AskUserQuestionConfig(BaseToolConfig):
+    permission: ToolPermission = ToolPermission.ALWAYS
+
+
+class AskUserQuestion(
+    BaseTool[
+        AskUserQuestionArgs, AskUserQuestionResult, AskUserQuestionConfig, BaseToolState
+    ],
+    ToolUIData[AskUserQuestionArgs, AskUserQuestionResult],
+):
+    description: ClassVar[str] = (
+        "Ask the user one or more questions and wait for their responses. "
+        "Each question has 2-4 choices plus an automatic 'Other' option for free text. "
+        "Use this to gather preferences, clarify requirements, or get decisions."
+    )
+
+    @classmethod
+    def get_call_display(cls, event: ToolCallEvent) -> ToolCallDisplay:
+        if not isinstance(event.args, AskUserQuestionArgs):
+            return ToolCallDisplay(summary="Asking user")
+
+        args = event.args
+        count = len(args.questions)
+
+        if count == 1:
+            return ToolCallDisplay(summary=f"Asking: {args.questions[0].question}")
+
+        return ToolCallDisplay(summary=f"Asking {count} questions")
+
+    @classmethod
+    def get_result_display(cls, event: ToolResultEvent) -> ToolResultDisplay:
+        if event.error:
+            return ToolResultDisplay(success=False, message=event.error)
+
+        if not isinstance(event.result, AskUserQuestionResult):
+            return ToolResultDisplay(success=True, message="Questions answered")
+
+        result = event.result
+
+        if result.cancelled:
+            return ToolResultDisplay(success=False, message="User cancelled")
+
+        if len(result.answers) == 1:
+            answer = result.answers[0]
+            prefix = "(Other) " if answer.is_other else ""
+            return ToolResultDisplay(success=True, message=f"{prefix}{answer.answer}")
+
+        return ToolResultDisplay(
+            success=True, message=f"{len(result.answers)} answers received"
+        )
+
+    @classmethod
+    def get_status_text(cls) -> str:
+        return "Waiting for user input"
+
+    async def run(
+        self, args: AskUserQuestionArgs, ctx: InvokeContext | None = None
+    ) -> AsyncGenerator[AskUserQuestionResult, None]:
+        if ctx is None or ctx.user_input_callback is None:
+            raise ToolError(
+                "User input not available. This tool requires an interactive UI."
+            )
+
+        result = await ctx.user_input_callback(args)
+        yield cast(AskUserQuestionResult, result)