daita-agents 0.1.0__py3-none-any.whl

daita/core/tools.py

@@ -0,0 +1,402 @@
+"""
+Universal tool abstraction for Daita agents.
+
+Tools are LLM-callable functions that can come from:
+- Plugins (database queries, S3 operations, API calls)
+- MCP servers (external tools via Model Context Protocol)
+- Custom functions (user-defined Python functions)
+
+This abstraction is provider-agnostic and supports both prompt-based
+and native function calling (OpenAI, Anthropic, etc).
+"""
+
+from dataclasses import dataclass
+from typing import Callable, Dict, Any, List, Optional, Awaitable, Union
+import asyncio
+import inspect
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AgentTool:
+    """
+    Universal tool definition for agent-LLM integration.
+
+    Represents any callable function that an agent can use, regardless of source.
+    Designed to be compatible with industry standards (LangChain, AutoGen, etc).
+
+    Example:
+        ```python
+        tool = AgentTool(
+            name="search_database",
+            description="Search for records in the database",
+            parameters={
+                "query": {
+                    "type": "string",
+                    "description": "SQL query to execute",
+                    "required": True
+                }
+            },
+            handler=async_search_function
+        )
+        ```
+    """
+
+    # Core fields (required)
+    name: str
+    description: str
+    parameters: Dict[str, Any]  # JSON Schema format
+    handler: Callable[[Dict[str, Any]], Awaitable[Any]]  # async function
+
+    # Optional metadata
+    category: Optional[str] = None  # "database", "storage", "api", etc
+    source: str = "custom"  # "plugin", "mcp", "custom"
+    plugin_name: Optional[str] = None  # Which plugin provides this tool
+
+    # Safety features
+    timeout_seconds: Optional[int] = None  # Execution timeout
+
+    def to_openai_function(self) -> Dict[str, Any]:
+        """
+        Convert to OpenAI function calling format.
+
+        Returns:
+            OpenAI function definition dict
+
+        Reference:
+            https://platform.openai.com/docs/guides/function-calling
+        """
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": {
+                    "type": "object",
+                    "properties": self.parameters,
+                    "required": [
+                        k for k, v in self.parameters.items()
+                        if v.get("required", False)
+                    ]
+                }
+            }
+        }
+
+    def to_anthropic_tool(self) -> Dict[str, Any]:
+        """
+        Convert to Anthropic tool format.
+
+        Returns:
+            Anthropic tool definition dict
+
+        Reference:
+            https://docs.anthropic.com/claude/docs/tool-use
+        """
+        return {
+            "name": self.name,
+            "description": self.description,
+            "input_schema": {
+                "type": "object",
+                "properties": self.parameters,
+                "required": [
+                    k for k, v in self.parameters.items()
+                    if v.get("required", False)
+                ]
+            }
+        }
+
+    def to_llm_function(self) -> Dict[str, Any]:
+        """
+        Generic LLM function format (works for most providers).
+
+        For provider-specific formats, use to_openai_function() or to_anthropic_tool().
+        This format is used for prompt-based tool calling.
+        """
+        return {
+            "name": self.name,
+            "description": self.description,
+            "parameters": {
+                "type": "object",
+                "properties": self.parameters,
+                "required": [
+                    k for k, v in self.parameters.items()
+                    if v.get("required", False)
+                ]
+            }
+        }
+
+    def to_prompt_description(self) -> str:
+        """
+        Generate human-readable tool description for prompt injection.
+
+        Used in prompt-based tool calling (non-native function calling).
+
+        Returns:
+            Formatted tool description string
+        """
+        params_desc = []
+        for param_name, param_info in self.parameters.items():
+            required = " (required)" if param_info.get("required") else " (optional)"
+            param_type = param_info.get("type", "any")
+            param_desc = param_info.get("description", "")
+            params_desc.append(
+                f"  - {param_name} ({param_type}){required}: {param_desc}"
+            )
+
+        params_str = "\n".join(params_desc) if params_desc else "  (no parameters)"
+
+        return f"{self.name}: {self.description}\nParameters:\n{params_str}"
+
+    async def execute(self, arguments: Dict[str, Any]) -> Any:
+        """
+        Execute the tool with given arguments.
+
+        Args:
+            arguments: Tool arguments matching the parameter schema
+
+        Returns:
+            Tool execution result
+
+        Raises:
+            RuntimeError: If tool execution fails or times out
+        """
+        if not callable(self.handler):
+            raise RuntimeError(f"Tool '{self.name}' has non-callable handler")
+
+        # Execute with timeout if specified
+        if self.timeout_seconds:
+            try:
+                result = await asyncio.wait_for(
+                    self.handler(arguments),
+                    timeout=self.timeout_seconds
+                )
+                return result
+            except asyncio.TimeoutError:
+                raise RuntimeError(
+                    f"Tool '{self.name}' execution timed out after {self.timeout_seconds}s"
+                )
+        else:
+            # No timeout
+            return await self.handler(arguments)
+
+    @classmethod
+    def from_function(
+        cls,
+        func: Callable,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        parameters: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> 'AgentTool':
+        """
+        Create AgentTool from a Python function.
+
+        Automatically extracts name, docstring, and parameter info if not provided.
+
+        Args:
+            func: Python function (sync or async)
+            name: Tool name (defaults to function name)
+            description: Tool description (defaults to docstring)
+            parameters: Parameter schema (auto-detected from type hints if not provided)
+            **kwargs: Additional AgentTool fields (category, timeout_seconds, etc)
+
+        Returns:
+            AgentTool instance
+
+        Example:
+            ```python
+            async def search_docs(query: str, limit: int = 10):
+                '''Search documentation with a query'''
+                return search_results
+
+            # Auto-detection from function signature
+            tool = AgentTool.from_function(search_docs)
+
+            # Or with explicit parameters
+            tool = AgentTool.from_function(
+                search_docs,
+                parameters={
+                    "query": {
+                        "type": "string",
+                        "description": "Search query",
+                        "required": True
+                    },
+                    "limit": {
+                        "type": "integer",
+                        "description": "Max results",
+                        "required": False
+                    }
+                }
+            )
+            ```
+        """
+        # Use function name if not provided
+        tool_name = name or func.__name__
+
+        # Use docstring if no description
+        tool_description = description or (func.__doc__ or f"Execute {tool_name}").strip()
+
+        # Auto-detect parameters from type hints if not provided
+        if parameters is None:
+            parameters = {}
+            sig = inspect.signature(func)
+
+            for param_name, param in sig.parameters.items():
+                # Skip self/cls parameters
+                if param_name in ('self', 'cls'):
+                    continue
+
+                # Map Python types to JSON Schema types
+                param_type = "string"  # Default type
+                if param.annotation != inspect.Parameter.empty:
+                    if param.annotation in (int, 'int'):
+                        param_type = "integer"
+                    elif param.annotation in (float, 'float'):
+                        param_type = "number"
+                    elif param.annotation in (bool, 'bool'):
+                        param_type = "boolean"
+                    elif param.annotation in (list, List):
+                        param_type = "array"
+                    elif param.annotation in (dict, Dict):
+                        param_type = "object"
+
+                # Determine if required (no default value)
+                is_required = param.default == inspect.Parameter.empty
+
+                parameters[param_name] = {
+                    "type": param_type,
+                    "description": f"Parameter: {param_name}",
+                    "required": is_required
+                }
+
+        # Wrap sync functions to make them async
+        if asyncio.iscoroutinefunction(func):
+            handler = func
+        else:
+            async def async_wrapper(args: Dict[str, Any]) -> Any:
+                return func(**args)
+            handler = async_wrapper
+
+        return cls(
+            name=tool_name,
+            description=tool_description,
+            parameters=parameters,
+            handler=handler,
+            source="custom",
+            **kwargs
+        )
+
+    @classmethod
+    def from_mcp_tool(cls, mcp_tool, mcp_registry) -> 'AgentTool':
+        """
+        Create AgentTool from an MCP tool.
+
+        Args:
+            mcp_tool: MCPTool instance from MCP plugin
+            mcp_registry: MCPToolRegistry for routing calls
+
+        Returns:
+            AgentTool instance that wraps the MCP tool
+        """
+        # Create handler that routes to MCP registry
+        async def mcp_handler(arguments: Dict[str, Any]) -> Any:
+            return await mcp_registry.call_tool(mcp_tool.name, arguments)
+
+        return cls(
+            name=mcp_tool.name,
+            description=mcp_tool.description,
+            parameters=mcp_tool.input_schema.get("properties", {}),
+            handler=mcp_handler,
+            source="mcp",
+            category="mcp"
+        )
+
+
+class ToolRegistry:
+    """
+    Registry for managing tools from multiple sources.
+
+    Used internally by agents to aggregate tools from plugins, MCP servers,
+    and custom functions.
+    """
+
+    def __init__(self):
+        """Initialize empty tool registry"""
+        self.tools: List[AgentTool] = []
+        self._tool_map: Dict[str, AgentTool] = {}
+
+    def register(self, tool: AgentTool) -> None:
+        """
+        Register a tool.
+
+        Args:
+            tool: AgentTool to register
+        """
+        if tool.name in self._tool_map:
+            logger.warning(
+                f"Tool '{tool.name}' already registered. "
+                f"Overwriting (old source: {self._tool_map[tool.name].source}, "
+                f"new source: {tool.source})"
+            )
+
+        self.tools.append(tool)
+        self._tool_map[tool.name] = tool
+
+        logger.debug(f"Registered tool: {tool.name} (source: {tool.source})")
+
+    def register_many(self, tools: List[AgentTool]) -> None:
+        """
+        Register multiple tools at once.
+
+        Args:
+            tools: List of AgentTool instances
+        """
+        for tool in tools:
+            self.register(tool)
+
+    def get(self, name: str) -> Optional[AgentTool]:
+        """
+        Get tool by name.
+
+        Args:
+            name: Tool name
+
+        Returns:
+            AgentTool instance or None if not found
+        """
+        return self._tool_map.get(name)
+
+    async def execute(self, name: str, arguments: Dict[str, Any]) -> Any:
+        """
+        Execute a tool by name.
+
+        Args:
+            name: Tool name
+            arguments: Tool arguments
+
+        Returns:
+            Tool execution result
+
+        Raises:
+            RuntimeError: If tool not found or execution fails
+        """
+        tool = self.get(name)
+        if not tool:
+            available = [t.name for t in self.tools]
+            raise RuntimeError(
+                f"Tool '{name}' not found. Available tools: {', '.join(available)}"
+            )
+
+        return await tool.execute(arguments)
+
+    @property
+    def tool_count(self) -> int:
+        """Total number of registered tools"""
+        return len(self.tools)
+
+    @property
+    def tool_names(self) -> List[str]:
+        """List of all tool names"""
+        return list(self._tool_map.keys())