aury-agent 0.0.4__py3-none-any.whl

aury/agents/tool/builtin/yield_result.py

@@ -0,0 +1,130 @@
+"""Yield result tool - return control to parent agent.
+
+Only available in DELEGATED mode sub-agents.
+Dynamically injected when control is transferred.
+"""
+from __future__ import annotations
+
+from typing import Any
+
+from ...core.logging import tool_logger as logger
+from ...core.types.tool import BaseTool, ToolContext, ToolResult
+from ...core.types.block import BlockEvent, BlockKind, BlockOp
+
+
+class YieldResultTool(BaseTool):
+    """Return control to parent agent with results.
+    
+    This tool is only available to sub-agents running in DELEGATED mode.
+    It is dynamically injected when control is transferred via delegate().
+    
+    When called:
+    1. Pops current frame from control_stack
+    2. Passes result to parent agent
+    3. Resumes parent invocation
+    """
+    
+    _name = "yield_result"
+    
+    def __init__(self, parent_invocation_id: str):
+        """Initialize with parent context.
+        
+        Args:
+            parent_invocation_id: ID of the parent invocation to return to
+        """
+        self.parent_invocation_id = parent_invocation_id
+    
+    @property
+    def name(self) -> str:
+        return self._name
+    
+    @property
+    def description(self) -> str:
+        return """Return control to the parent agent with task results.
+
+Use this when the delegated task is complete and you want to
+return the results to the agent that delegated to you.
+
+This will end your current session and resume the parent."""
+    
+    @property
+    def parameters(self) -> dict[str, Any]:
+        return {
+            "type": "object",
+            "properties": {
+                "result": {
+                    "type": "string",
+                    "description": "Summary of what was accomplished",
+                },
+                "data": {
+                    "type": "object",
+                    "description": "Structured data to pass back to parent",
+                },
+                "status": {
+                    "type": "string",
+                    "enum": ["completed", "failed", "cancelled"],
+                    "description": "Status of the delegated task",
+                    "default": "completed",
+                },
+            },
+            "required": ["result"],
+        }
+    
+    async def execute(
+        self,
+        params: dict[str, Any],
+        ctx: ToolContext,
+    ) -> ToolResult:
+        """Execute yield - return to parent."""
+        result_text = params.get("result", "")
+        data = params.get("data", {})
+        status = params.get("status", "completed")
+        
+        logger.info(
+            "Yielding result to parent",
+            extra={
+                "parent_invocation_id": self.parent_invocation_id,
+                "status": status,
+            },
+        )
+        
+        # Emit YIELD block
+        await self._emit_yield_block(ctx, result_text, data, status)
+        
+        # Note: Real implementation would:
+        # 1. Pop current frame from session.control_stack
+        # 2. Serialize result for parent
+        # 3. Update state to trigger parent resumption
+        # 4. End current invocation
+        
+        return ToolResult(output=f"Returning to parent agent with status: {status}\nResult: {result_text}")
+    
+    async def _emit_yield_block(
+        self,
+        ctx: ToolContext,
+        result: str,
+        data: dict[str, Any],
+        status: str,
+    ) -> None:
+        """Emit YIELD block."""
+        emit = getattr(ctx, 'emit', None)
+        if emit is None:
+            return
+        
+        block = BlockEvent(
+            kind=BlockKind.YIELD,
+            op=BlockOp.APPLY,
+            data={
+                "result": result,
+                "data": data,
+                "status": status,
+                "parent_invocation_id": self.parent_invocation_id,
+            },
+            session_id=ctx.session_id,
+            invocation_id=ctx.invocation_id,
+        )
+        
+        await emit(block)
+
+
+__all__ = ["YieldResultTool"]

aury/agents/tool/decorator.py

@@ -0,0 +1,252 @@
+"""Tool decorators for creating tools from functions."""
+from __future__ import annotations
+
+import asyncio
+import inspect
+from functools import wraps
+from typing import Any, Callable, get_type_hints, get_origin, get_args, Union
+
+from ..core.types.tool import BaseTool, ToolContext, ToolResult, ToolInfo
+
+
+def _type_to_schema(t: type) -> dict[str, Any]:
+    """Convert Python type to JSON Schema."""
+    # Handle None/NoneType
+    if t is type(None):
+        return {"type": "null"}
+    
+    # Handle Optional (Union with None)
+    origin = get_origin(t)
+    if origin is Union:
+        args = get_args(t)
+        # Filter out None
+        non_none = [a for a in args if a is not type(None)]
+        if len(non_none) == 1:
+            # Optional[X] -> X with nullable
+            schema = _type_to_schema(non_none[0])
+            return schema
+        # Union of multiple types
+        return {"oneOf": [_type_to_schema(a) for a in non_none]}
+    
+    # Handle list
+    if origin is list:
+        args = get_args(t)
+        if args:
+            return {"type": "array", "items": _type_to_schema(args[0])}
+        return {"type": "array"}
+    
+    # Handle dict
+    if origin is dict:
+        args = get_args(t)
+        if len(args) == 2:
+            return {
+                "type": "object",
+                "additionalProperties": _type_to_schema(args[1]),
+            }
+        return {"type": "object"}
+    
+    # Basic types
+    type_map = {
+        str: {"type": "string"},
+        int: {"type": "integer"},
+        float: {"type": "number"},
+        bool: {"type": "boolean"},
+        list: {"type": "array"},
+        dict: {"type": "object"},
+        Any: {},
+    }
+    
+    return type_map.get(t, {"type": "string"})
+
+
+def _parse_docstring(docstring: str | None) -> tuple[str, dict[str, str]]:
+    """Parse docstring to extract description and param descriptions.
+    
+    Returns:
+        Tuple of (main description, {param_name: param_description})
+    """
+    if not docstring:
+        return "", {}
+    
+    lines = docstring.strip().split("\n")
+    description_lines = []
+    param_docs: dict[str, str] = {}
+    
+    in_params = False
+    current_param = None
+    
+    for line in lines:
+        stripped = line.strip()
+        
+        # Check for param section
+        if stripped.lower().startswith(("args:", "arguments:", "parameters:", "params:")):
+            in_params = True
+            continue
+        
+        if stripped.lower().startswith(("returns:", "return:", "raises:", "yields:")):
+            in_params = False
+            current_param = None
+            continue
+        
+        if in_params:
+            # Check for param definition: "param_name: description" or "param_name (type): description"
+            if ":" in stripped and not stripped.startswith(" "):
+                parts = stripped.split(":", 1)
+                param_name = parts[0].strip()
+                # Remove type annotation if present
+                if "(" in param_name:
+                    param_name = param_name.split("(")[0].strip()
+                param_docs[param_name] = parts[1].strip() if len(parts) > 1 else ""
+                current_param = param_name
+            elif current_param and stripped:
+                # Continuation of previous param description
+                param_docs[current_param] += " " + stripped
+        else:
+            description_lines.append(stripped)
+    
+    description = " ".join(description_lines).strip()
+    return description, param_docs
+
+
+def tool(
+    func: Callable | None = None,
+    *,
+    name: str | None = None,
+    description: str | None = None,
+    parameters: dict[str, Any] | None = None,
+) -> Callable[[Callable], BaseTool] | BaseTool:
+    """Decorator to create a Tool from a function.
+    
+    The function can be sync or async. Parameters are automatically converted
+    to JSON Schema based on type hints, unless manually specified.
+    
+    Args:
+        name: Tool name (defaults to function name)
+        description: Tool description (defaults to docstring)
+        parameters: Manual JSON Schema for parameters (optional)
+        
+    Returns:
+        Decorated function as a Tool
+        
+    Example:
+        @tool
+        def greet(name: str) -> str:
+            return f"Hello, {name}!"
+        
+        @tool(name="calculator", description="Perform calculations")
+        def calculate(expression: str) -> str:
+            return str(eval(expression))
+        
+        @tool(parameters={"type": "object", "properties": {...}})
+        def custom_tool(arg: str) -> str:
+            return arg
+    """
+    def decorator(fn: Callable) -> BaseTool:
+        tool_name = name or fn.__name__
+        
+        # Parse docstring for descriptions
+        doc_desc, param_docs = _parse_docstring(fn.__doc__)
+        tool_desc = description or doc_desc or f"Tool: {tool_name}"
+        
+        # Use manual parameters or auto-generate
+        if parameters is not None:
+            parameters_schema = parameters
+        else:
+            # Get type hints and signature
+            try:
+                hints = get_type_hints(fn)
+            except Exception:
+                hints = {}
+            
+            sig = inspect.signature(fn)
+            
+            # Build parameters schema
+            properties: dict[str, Any] = {}
+            required: list[str] = []
+            
+            for param_name, param in sig.parameters.items():
+                # Skip special parameters
+                if param_name in ("self", "cls", "ctx", "context"):
+                    continue
+                
+                # Get type
+                param_type = hints.get(param_name, str)
+                schema = _type_to_schema(param_type)
+                
+                # Add description from docstring
+                if param_name in param_docs:
+                    schema["description"] = param_docs[param_name]
+                
+                properties[param_name] = schema
+                
+                # Check if required
+                if param.default is inspect.Parameter.empty:
+                    required.append(param_name)
+            
+            parameters_schema = {
+                "type": "object",
+                "properties": properties,
+                "required": required,
+            }
+        
+        sig = inspect.signature(fn)
+        
+        # Check if function accepts ctx
+        accepts_ctx = "ctx" in sig.parameters or "context" in sig.parameters
+        ctx_param_name = "ctx" if "ctx" in sig.parameters else "context"
+        
+        class FunctionTool:
+            """Tool created from decorated function."""
+            
+            @property
+            def name(self) -> str:
+                return tool_name
+            
+            @property
+            def description(self) -> str:
+                return tool_desc
+            
+            @property
+            def parameters(self) -> dict[str, Any]:
+                return parameters_schema
+            
+            async def execute(
+                self,
+                params: dict[str, Any],
+                ctx: ToolContext,
+            ) -> ToolResult:
+                """Execute the wrapped function."""
+                try:
+                    # Add context if function accepts it
+                    if accepts_ctx:
+                        params = {**params, ctx_param_name: ctx}
+                    
+                    # Call function
+                    result = fn(**params)
+                    
+                    # Await if coroutine
+                    if asyncio.iscoroutine(result):
+                        result = await result
+                    
+                    # Convert result to ToolResult
+                    if isinstance(result, ToolResult):
+                        return result
+                    
+                    return ToolResult(output=str(result))
+                
+                except Exception as e:
+                    return ToolResult.error(str(e))
+            
+            def get_info(self) -> ToolInfo:
+                return ToolInfo(
+                    name=tool_name,
+                    description=tool_desc,
+                    parameters=parameters_schema,
+                )
+        
+        return FunctionTool()
+    
+    # Support both @tool and @tool(...) usage
+    if func is not None:
+        return decorator(func)
+    return decorator

aury/agents/tool/set.py

@@ -0,0 +1,204 @@
+"""ToolSet - A collection of tools."""
+from __future__ import annotations
+
+from typing import Any, Callable
+
+from ..core.types.tool import BaseTool, ToolInfo
+
+
+class ToolSet:
+    """A collection of tools with lookup, filtering, and merging capabilities.
+    
+    ToolSet provides:
+    - Fast lookup by tool ID
+    - Lazy instantiation with caching
+    - Filtering (include/exclude)
+    - Merging multiple sets
+    
+    Usage:
+        # From list
+        tools = ToolSet.from_list([tool1, tool2])
+        
+        # Manual construction
+        tools = ToolSet()
+        tools.add(my_tool)
+        
+        # Filtering
+        subset = tools.filtered(include=["read_file", "write_file"])
+        
+        # Merging
+        combined = tools.merge(other_tools)
+    """
+    
+    def __init__(self) -> None:
+        self._tools: dict[str, Callable[[], BaseTool]] = {}
+        self._instances: dict[str, BaseTool] = {}  # Cached instances
+    
+    @classmethod
+    def from_list(cls, tools: list[BaseTool]) -> "ToolSet":
+        """Create ToolSet from a list of tools.
+        
+        Args:
+            tools: List of tool instances
+            
+        Returns:
+            New ToolSet containing all tools
+        """
+        ts = cls()
+        for tool in tools:
+            ts.add(tool)
+        return ts
+    
+    def add(
+        self,
+        tool: BaseTool | Callable[[], BaseTool],
+        replace: bool = False,
+    ) -> None:
+        """Add a tool to the set.
+        
+        Args:
+            tool: Tool instance or factory function
+            replace: Allow replacing existing tool
+        """
+        # If it's a BaseTool instance, wrap it in a lambda
+        if isinstance(tool, BaseTool):
+            tool_name = tool.name
+            factory: Callable[[], BaseTool] = lambda t=tool: t
+        else:
+            # It's a factory, call it once to get the name
+            instance = tool()
+            tool_name = instance.name
+            factory = tool
+        
+        if tool_name in self._tools and not replace:
+            raise ValueError(f"Tool already registered: {tool_name}")
+        
+        self._tools[tool_name] = factory
+        # Clear cached instance
+        if tool_name in self._instances:
+            del self._instances[tool_name]
+    
+    def remove(self, tool_id: str) -> bool:
+        """Remove a tool from the set.
+        
+        Args:
+            tool_id: Tool identifier to remove
+            
+        Returns:
+            True if tool was removed, False if not found
+        """
+        if tool_id in self._tools:
+            del self._tools[tool_id]
+            if tool_id in self._instances:
+                del self._instances[tool_id]
+            return True
+        return False
+    
+    def get(self, tool_id: str, cached: bool = True) -> BaseTool:
+        """Get a tool instance.
+        
+        Args:
+            tool_id: Tool identifier
+            cached: Whether to use cached instance
+            
+        Returns:
+            BaseTool instance
+            
+        Raises:
+            KeyError: If tool not found
+        """
+        if tool_id not in self._tools:
+            raise KeyError(f"Tool not found: {tool_id}")
+        
+        if cached and tool_id in self._instances:
+            return self._instances[tool_id]
+        
+        instance = self._tools[tool_id]()
+        if cached:
+            self._instances[tool_id] = instance
+        
+        return instance
+    
+    def has(self, tool_id: str) -> bool:
+        """Check if tool is registered."""
+        return tool_id in self._tools
+    
+    def all(self, cached: bool = True) -> list[BaseTool]:
+        """Get all registered tools.
+        
+        Args:
+            cached: Whether to use cached instances
+        """
+        return [self.get(tid, cached=cached) for tid in self._tools]
+    
+    def all_info(self) -> list[ToolInfo]:
+        """Get info for all registered tools."""
+        return [tool.get_info() for tool in self.all()]
+    
+    def ids(self) -> list[str]:
+        """Get all registered tool IDs."""
+        return list(self._tools.keys())
+    
+    def filtered(
+        self,
+        include: list[str] | None = None,
+        exclude: list[str] | None = None,
+    ) -> "ToolSet":
+        """Create a filtered copy of this set.
+        
+        Args:
+            include: Only include these tools (None = all)
+            exclude: Exclude these tools
+            
+        Returns:
+            New ToolSet with filtered tools
+        """
+        filtered = ToolSet()
+        
+        for tool_id, factory in self._tools.items():
+            # Check include filter
+            if include is not None and tool_id not in include:
+                continue
+            
+            # Check exclude filter
+            if exclude is not None and tool_id in exclude:
+                continue
+            
+            filtered._tools[tool_id] = factory
+        
+        return filtered
+    
+    def merge(self, other: "ToolSet", replace: bool = False) -> "ToolSet":
+        """Merge another ToolSet into this one.
+        
+        Args:
+            other: ToolSet to merge
+            replace: Allow replacing existing tools
+            
+        Returns:
+            Self for chaining
+        """
+        for tool_id, factory in other._tools.items():
+            if tool_id not in self._tools or replace:
+                self._tools[tool_id] = factory
+        return self
+    
+    def copy(self) -> "ToolSet":
+        """Create a copy of this set."""
+        new = ToolSet()
+        new._tools = dict(self._tools)
+        return new
+    
+    def clear(self) -> None:
+        """Clear all tools."""
+        self._tools.clear()
+        self._instances.clear()
+    
+    def __len__(self) -> int:
+        return len(self._tools)
+    
+    def __contains__(self, tool_id: str) -> bool:
+        return tool_id in self._tools
+    
+    def __iter__(self):
+        return iter(self.all())

aury/agents/usage/__init__.py

@@ -0,0 +1,12 @@
+"""Usage tracking for cost management."""
+from .tracker import (
+    UsageType,
+    UsageEntry,
+    UsageTracker,
+)
+
+__all__ = [
+    "UsageType",
+    "UsageEntry",
+    "UsageTracker",
+]