agentpool 2.1.9__py3-none-any.whl

agentpool/tools/manager.py

@@ -0,0 +1,255 @@
+"""Tool management for AgentPool."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable, Sequence
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Literal, assert_never
+
+from agentpool.log import get_logger
+from agentpool.resource_providers import StaticResourceProvider
+from agentpool.tools.base import Tool
+from agentpool.utils.baseregistry import AgentPoolError
+
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from agentpool.common_types import ToolType
+    from agentpool.prompts.prompts import MCPClientPrompt
+    from agentpool.resource_providers import ResourceProvider
+    from agentpool.resource_providers.codemode.provider import CodeModeResourceProvider
+
+
+logger = get_logger(__name__)
+
+MAX_LEN_DESCRIPTION = 2000
+ToolState = Literal["all", "enabled", "disabled"]
+ProviderName = str
+OwnerType = Literal["pool", "team", "node"]
+ToolMode = Literal["codemode"]
+
+
+class ToolError(AgentPoolError):
+    """Base exception for tool-related errors."""
+
+
+class ToolManager:
+    """Manages tool registration, enabling/disabling and access."""
+
+    def __init__(
+        self,
+        tools: Sequence[ToolType] | None = None,
+        tool_mode: ToolMode | None = None,
+    ) -> None:
+        """Initialize tool manager.
+
+        Args:
+            tools: Initial tools to register
+            tool_mode: Tool execution mode (None or "codemode")
+        """
+        super().__init__()
+        self.external_providers: list[ResourceProvider] = []
+        self.worker_provider = StaticResourceProvider(name="workers")
+        self.builtin_provider = StaticResourceProvider(name="builtin")
+        self.tool_mode = tool_mode
+
+        # CodeModeResourceProvider gets populated with providers in providers property
+        from agentpool.resource_providers.codemode.provider import CodeModeResourceProvider
+
+        self._codemode_provider: CodeModeResourceProvider = CodeModeResourceProvider([])
+
+        # Forward to provider methods
+        self.tool = self.builtin_provider.tool
+        self.register_tool = self.builtin_provider.register_tool
+        self.register_worker = self.worker_provider.register_worker
+
+        # Register initial tools
+        for tool in tools or []:
+            t = self._validate_item(tool)
+            self.builtin_provider.add_tool(t)
+
+    @property
+    def providers(self) -> list[ResourceProvider]:
+        """Get all providers: external + worker + builtin providers."""
+        if self.tool_mode == "codemode":
+            # Update the providers list with current providers
+            self._codemode_provider.providers[:] = [
+                *self.external_providers,
+                self.worker_provider,
+                self.builtin_provider,
+            ]
+            return [self._codemode_provider]
+
+        return [*self.external_providers, self.worker_provider, self.builtin_provider]
+
+    async def __prompt__(self) -> str:
+        enabled_tools = [t.name for t in await self.get_tools() if t.enabled]
+        if not enabled_tools:
+            return "No tools available"
+        return f"Available tools: {', '.join(enabled_tools)}"
+
+    def add_provider(self, provider: ResourceProvider, owner: str | None = None) -> None:
+        """Add an external resource provider.
+
+        Args:
+            provider: ResourceProvider instance (e.g., MCP server, custom provider)
+            owner: Optional owner for the provider
+        """
+        if owner:
+            provider.owner = owner
+        self.external_providers.append(provider)
+
+    def remove_provider(self, provider: ResourceProvider | ProviderName) -> None:
+        """Remove an external resource provider."""
+        from agentpool.resource_providers import ResourceProvider
+
+        match provider:
+            case ResourceProvider():
+                self.external_providers.remove(provider)
+            case str():
+                for p in self.external_providers:
+                    if p.name == provider:
+                        self.external_providers.remove(p)
+            case _ as unreachable:
+                assert_never(unreachable)
+
+    async def reset_states(self) -> None:
+        """Reset all tools to their default enabled states."""
+        for info in await self.get_tools():
+            info.enabled = True
+
+    def _validate_item(self, item: ToolType) -> Tool:
+        """Validate and convert items before registration."""
+        match item:
+            case Tool():
+                return item
+            case Callable() | str():
+                return Tool.from_callable(item)
+            case _:
+                typ = type(item)
+                msg = f"Item must be Tool or callable. Got {typ}"
+                raise ToolError(msg)
+
+    async def enable_tool(self, tool_name: str) -> None:
+        """Enable a previously disabled tool."""
+        tool_info = await self.get_tool(tool_name)
+        tool_info.enabled = True
+        logger.debug("Enabled tool", tool_name=tool_name)
+
+    async def disable_tool(self, tool_name: str) -> None:
+        """Disable a tool."""
+        tool_info = await self.get_tool(tool_name)
+        tool_info.enabled = False
+        logger.debug("Disabled tool", tool_name=tool_name)
+
+    async def list_tools(self) -> dict[str, bool]:
+        """Get a mapping of all tools and their enabled status."""
+        return {tool.name: tool.enabled for tool in await self.get_tools()}
+
+    async def get_tools(
+        self,
+        state: ToolState = "all",
+        names: str | list[str] | None = None,
+    ) -> list[Tool]:
+        """Get tool objects based on filters."""
+        tools: list[Tool] = []
+        # Get tools from providers concurrently
+        provider_coroutines = [provider.get_tools() for provider in self.providers]
+        results = await asyncio.gather(*provider_coroutines, return_exceptions=True)
+        for provider, result in zip(self.providers, results, strict=False):
+            if isinstance(result, BaseException):
+                msg = "Failed to get tools from provider"
+                logger.warning(msg, provider=provider, result=result)
+                continue
+            tools.extend(t for t in result if t.matches_filter(state))
+
+        match names:
+            case str():
+                tools = [t for t in tools if t.name == names]
+            case list():
+                tools = [t for t in tools if t.name in names]
+        return tools
+
+    async def get_tool(self, name: str) -> Tool:
+        """Get a specific tool by name.
+
+        First checks local tools, then uses concurrent provider fetching.
+        """
+        tool = next((tool for tool in await self.get_tools() if tool.name == name), None)
+        if not tool:
+            msg = f"Tool not found: {tool}"
+            raise ToolError(msg)
+        return tool
+
+    async def list_prompts(self) -> list[MCPClientPrompt]:
+        """Get all prompts from all providers."""
+        from agentpool.mcp_server.manager import MCPManager
+
+        all_prompts: list[MCPClientPrompt] = []
+        # Get prompts from all external providers (check if they're MCP providers)
+        for provider in self.external_providers:
+            if isinstance(provider, MCPManager):
+                try:
+                    # Get prompts from MCP providers via the aggregating provider
+                    agg_provider = provider.get_aggregating_provider()
+                    prompts = await agg_provider.get_prompts()
+                    all_prompts.extend(prompts)
+                except Exception:
+                    logger.exception("Failed to get prompts from provider", provider=provider)
+
+        return all_prompts
+
+    @asynccontextmanager
+    async def temporary_tools(
+        self,
+        tools: ToolType | Sequence[ToolType],
+        *,
+        exclusive: bool = False,
+    ) -> AsyncIterator[list[Tool]]:
+        """Temporarily register tools.
+
+        Args:
+            tools: Tool(s) to register
+            exclusive: Whether to temporarily disable all other tools
+
+        Yields:
+            List of registered tool infos
+        """
+        # Normalize inputs to lists
+        tools_list: list[ToolType] = [tools] if not isinstance(tools, Sequence) else list(tools)
+        # Store original tool states if exclusive
+        tools = await self.get_tools()
+        original_states: dict[str, bool] = {}
+        if exclusive:
+            original_states = {t.name: t.enabled for t in tools}
+            # Disable all existing tools
+            for t in tools:
+                t.enabled = False
+        # Register all tools
+        registered_tools: list[Tool] = []
+        try:
+            for tool in tools_list:
+                tool_info = self.register_tool(tool)
+                registered_tools.append(tool_info)
+            yield registered_tools
+
+        finally:
+            # Remove temporary tools
+            for tool_info in registered_tools:
+                self.builtin_provider.remove_tool(tool_info.name)
+            # Restore original tool states if exclusive
+            if exclusive:
+                for name_, was_enabled in original_states.items():
+                    t_ = await self.get_tool(name_)
+                    t_.enabled = was_enabled
+
+
+if __name__ == "__main__":
+    manager = ToolManager()
+
+    @manager.tool(name="custom_name", description="Custom description")
+    def with_params(query: str) -> str:
+        """With parameters."""
+        return query.upper()

agentpool/tools/tool_call_info.py

@@ -0,0 +1,87 @@
+"""Tool management for AgentPool."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, Literal
+from uuid import uuid4
+
+from pydantic import Field
+from schemez import Schema
+
+from agentpool.text_templates import get_tool_call_template
+from agentpool.utils.now import get_now
+
+
+if TYPE_CHECKING:
+    from agentpool.text_templates import FormatStyle
+
+
+class ToolCallInfo(Schema):
+    """Information about an executed tool call."""
+
+    tool_name: str
+    """Name of the tool that was called."""
+
+    args: dict[str, Any]
+    """Arguments passed to the tool."""
+
+    result: Any
+    """Result returned by the tool."""
+
+    agent_name: str
+    """Name of the calling agent."""
+
+    tool_call_id: str = Field(default_factory=lambda: str(uuid4()))
+    """ID provided by the model (e.g. OpenAI function call ID)."""
+
+    timestamp: datetime = Field(default_factory=get_now)
+    """When the tool was called."""
+
+    message_id: str | None = None
+    """ID of the message that triggered this tool call."""
+
+    error: str | None = None
+    """Error message if the tool call failed."""
+
+    timing: float | None = None
+    """Time taken for this specific tool call in seconds."""
+
+    agent_tool_name: str | None = None
+    """If this tool is agent-based, the name of that agent."""
+
+    def format(
+        self,
+        style: FormatStyle | Literal["custom"] = "simple",
+        *,
+        template: str | None = None,
+        variables: dict[str, Any] | None = None,
+        show_timing: bool = True,
+    ) -> str:
+        """Format tool call information with configurable style.
+
+        Args:
+            style: Predefined style to use:
+                - simple: Compact single-line format
+                - detailed: Multi-line with all details
+                - markdown: Formatted markdown with syntax highlighting
+            template: Optional custom template (required if style="custom")
+            variables: Additional variables for template rendering
+            show_timing: Whether to include execution timing
+
+        Raises:
+            ValueError: If style is invalid or custom template is missing
+        """
+        from jinjarope import Environment
+
+        template_str = template if style == "custom" else get_tool_call_template(style)
+        if not template_str:
+            raise ValueError("Custom template is required for style='custom'")
+        env = Environment(trim_blocks=True, lstrip_blocks=True)
+        env.filters["repr"] = repr
+        template_obj = env.from_string(template_str)
+        vars_ = self.model_dump()
+        vars_["show_timing"] = show_timing
+        if variables:
+            vars_.update(variables)
+        return template_obj.render(**vars_)

agentpool/ui/__init__.py

@@ -0,0 +1,2 @@
+# src/agentpool/interfaces/__init__.py
+"""Interface definitions for AgentPool."""

agentpool/ui/base.py

@@ -0,0 +1,89 @@
+"""Base input provider class."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+
+if TYPE_CHECKING:
+    from collections.abc import Coroutine
+
+    from mcp import types
+    from pydantic import BaseModel
+
+    from agentpool.agents.context import ConfirmationResult
+    from agentpool.messaging import ChatMessage
+    from agentpool.messaging.context import NodeContext
+    from agentpool.tools.base import Tool
+
+
+class InputProvider(ABC):
+    """Base class for handling all UI interactions."""
+
+    async def get_input(
+        self,
+        context: NodeContext,
+        prompt: str,
+        output_type: type | None = None,
+        message_history: list[ChatMessage[Any]] | None = None,
+    ) -> Any:
+        """Get normal input (used by HumanProvider).
+
+        Args:
+            context: Current agent context
+            prompt: The prompt to show to the user
+            output_type: Optional type for structured responses
+            message_history: Optional conversation history
+        """
+        if output_type:
+            return await self.get_structured_input(context, prompt, output_type, message_history)
+        return await self.get_text_input(context, prompt, message_history)
+
+    async def get_text_input(
+        self,
+        context: NodeContext[Any],
+        prompt: str,
+        message_history: list[ChatMessage[Any]] | None = None,
+    ) -> str:
+        """Get normal text input."""
+        raise NotImplementedError
+
+    async def get_structured_input(
+        self,
+        context: NodeContext[Any],
+        prompt: str,
+        output_type: type[BaseModel],
+        message_history: list[ChatMessage[Any]] | None = None,
+    ) -> BaseModel:
+        """Get structured input."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_tool_confirmation(
+        self,
+        context: NodeContext[Any],
+        tool: Tool,
+        args: dict[str, Any],
+        message_history: list[ChatMessage[Any]] | None = None,
+    ) -> Coroutine[Any, Any, ConfirmationResult]:
+        """Get tool execution confirmation.
+
+        Args:
+            context: Current node context
+            tool: Information about the tool to be executed
+            args: Tool arguments
+            message_history: Optional conversation history
+        """
+
+    @abstractmethod
+    def get_elicitation(
+        self,
+        params: types.ElicitRequestParams,
+    ) -> Coroutine[Any, Any, types.ElicitResult | types.ErrorData]:
+        """Get user response to elicitation request.
+
+        Args:
+            context: Current agent context
+            params: MCP elicit request parameters
+        """

agentpool/ui/mock_provider.py

@@ -0,0 +1,81 @@
+"""Mock input provider implementation."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal
+
+from agentpool.ui.base import InputProvider
+
+
+if TYPE_CHECKING:
+    from mcp import types
+
+    from agentpool.agents.context import ConfirmationResult
+    from agentpool.messaging import ChatMessage
+    from agentpool.messaging.context import NodeContext
+    from agentpool.tools.base import Tool
+
+InputMethod = Literal["get_input", "get_tool_confirmation", "get_elicitation"]
+
+
+@dataclass
+class InputCall:
+    """Record of an input provider call."""
+
+    method: InputMethod
+    args: tuple[Any, ...]
+    kwargs: dict[str, Any]
+    result: Any
+
+
+class MockInputProvider(InputProvider):
+    """Provider that records calls and returns pre-configured responses."""
+
+    def __init__(
+        self,
+        *,
+        input_response: str = "mock response",
+        tool_confirmation: ConfirmationResult = "allow",
+        elicitation_response: dict[str, Any] | None = None,
+    ) -> None:
+        self.input_response = input_response
+        self.tool_confirmation = tool_confirmation
+        self.elicitation_response = elicitation_response or {"response": "mock response"}
+        self.calls: list[InputCall] = []
+
+    async def get_input(
+        self,
+        context: NodeContext,
+        prompt: str,
+        output_type: type | None = None,
+        message_history: list[ChatMessage[Any]] | None = None,
+    ) -> Any:
+        kwargs = {"output_type": output_type, "message_history": message_history}
+        call = InputCall("get_input", (context, prompt), kwargs, result=self.input_response)
+        self.calls.append(call)
+        return self.input_response
+
+    async def get_tool_confirmation(
+        self,
+        context: NodeContext,
+        tool: Tool,
+        args: dict[str, Any],
+        message_history: list[ChatMessage[Any]] | None = None,
+    ) -> ConfirmationResult:
+        kwargs = {"message_history": message_history}
+        result = self.tool_confirmation
+        call = InputCall("get_tool_confirmation", (context, tool, args), kwargs, result=result)
+        self.calls.append(call)
+        return result  # pyright: ignore
+
+    async def get_elicitation(
+        self,
+        params: types.ElicitRequestParams,
+    ) -> types.ElicitResult | types.ErrorData:
+        from mcp import types
+
+        result = types.ElicitResult(action="accept", content=self.elicitation_response)
+        call = InputCall("get_elicitation", (params,), {}, result=result)
+        self.calls.append(call)
+        return result

agentpool/ui/stdlib_provider.py

@@ -0,0 +1,150 @@
+"""Stdlib input provider."""
+
+from __future__ import annotations
+
+import sys
+from textwrap import dedent
+from typing import TYPE_CHECKING, Any
+
+import anyenv
+from mcp import types
+
+from agentpool.log import get_logger
+from agentpool.tools.exceptions import ToolError
+from agentpool.ui.base import InputProvider
+
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+    from agentpool.agents.context import ConfirmationResult
+    from agentpool.messaging import ChatMessage
+    from agentpool.messaging.context import NodeContext
+    from agentpool.tools.base import Tool
+
+
+logger = get_logger(__name__)
+
+
+class StdlibInputProvider(InputProvider):
+    """Input provider using only Python stdlib functionality."""
+
+    async def get_text_input(
+        self,
+        context: NodeContext,
+        prompt: str,
+        message_history: list[ChatMessage[Any]] | None = None,
+    ) -> str:
+        print(f"{prompt}\n> ", end="", file=sys.stderr, flush=True)
+        return input()
+
+    async def get_structured_input(
+        self,
+        context: NodeContext,
+        prompt: str,
+        output_type: type[BaseModel],
+        message_history: list[ChatMessage[Any]] | None = None,
+    ) -> BaseModel:
+        """Get structured input, with promptantic and fallback handling."""
+        if result := await _get_promptantic_result(output_type):
+            return result
+
+        # Fallback: Get raw input and validate
+        prompt = f"{prompt}\n(Please provide response as {output_type.__name__})"
+        raw_input = await self.get_input(context, prompt, message_history=message_history)
+        try:
+            return output_type.model_validate_json(raw_input)
+        except Exception as e:
+            msg = f"Invalid response format: {e}"
+            raise ToolError(msg) from e
+
+    async def get_tool_confirmation(
+        self,
+        context: NodeContext,
+        tool: Tool,
+        args: dict[str, Any],
+        message_history: list[ChatMessage[Any]] | None = None,
+    ) -> ConfirmationResult:
+        agent_name = context.node_name
+        prompt = dedent(f"""
+            Tool Execution Confirmation
+            -------------------------
+            Tool: {tool.name}
+            Description: {tool.description or "No description"}
+            Agent: {agent_name}
+
+            Arguments:
+            {anyenv.dump_json(args, indent=True)}
+
+            Options:
+            - y: allow execution
+            - n/skip: skip this tool
+            - abort: abort current run
+            - quit: abort entire chain
+            """).strip()
+
+        print(f"{prompt}\nChoice [y/n/abort/quit]: ", end="", file=sys.stderr, flush=True)
+        response = input().lower()
+        match response:
+            case "y" | "yes":
+                return "allow"
+            case "abort":
+                return "abort_run"
+            case "quit":
+                return "abort_chain"
+            case _:
+                return "skip"
+
+    async def get_elicitation(
+        self,
+        params: types.ElicitRequestParams,
+    ) -> types.ElicitResult | types.ErrorData:
+        """Get user response to elicitation request using stdlib input."""
+        try:
+            print(f"\n{params.message}", file=sys.stderr)
+
+            # URL mode: prompt user to open external URL
+            if isinstance(params, types.ElicitRequestURLParams):
+                print(f"URL: {params.url}", file=sys.stderr)
+                print("Open this URL? [y/n]: ", end="", file=sys.stderr, flush=True)
+                response = input().strip().lower()
+                action = (
+                    "accept"
+                    if response in ("y", "yes")
+                    else ("decline" if response in ("n", "no") else "cancel")
+                )
+                return types.ElicitResult(action=action)
+
+            # Form mode: collect structured JSON input
+            print("Please provide response as JSON:", file=sys.stderr)
+            if params.requestedSchema:
+                schema_json = anyenv.dump_json(params.requestedSchema, indent=True)
+                print(f"Expected schema:\n{schema_json}", file=sys.stderr)
+            print("> ", end="", file=sys.stderr, flush=True)
+            response = input()
+            try:
+                content = anyenv.load_json(response, return_type=dict)
+                return types.ElicitResult(action="accept", content=content)
+            except anyenv.JsonLoadError as e:
+                return types.ErrorData(code=types.INVALID_REQUEST, message=f"Invalid JSON: {e}")
+
+        except KeyboardInterrupt:
+            return types.ElicitResult(action="cancel")
+        except Exception as e:  # noqa: BLE001
+            return types.ErrorData(code=types.INVALID_REQUEST, message=f"Elicitation failed: {e}")
+
+
+async def _get_promptantic_result(output_type: type[BaseModel]) -> BaseModel | None:
+    """Helper to get structured input via promptantic.
+
+    Returns None if promptantic is not available or fails.
+    """
+    try:
+        from promptantic import ModelGenerator
+
+        return await ModelGenerator().apopulate(output_type)
+    except ImportError:
+        return None
+    except Exception as e:  # noqa: BLE001
+        logger.warning("Promptantic failed", error=e)
+        return None