RouteKitAI 0.1.0__py3-none-any.whl

routekitai/core/agent.py

@@ -0,0 +1,325 @@
+"""Agent primitive for RouteKit."""
+
+from __future__ import annotations
+
+import asyncio
+import uuid
+from collections.abc import AsyncIterator
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from routekitai.core.hooks import ToolFilter
+from routekitai.core.memory import Memory
+from routekitai.core.message import Message
+from routekitai.core.model import Model
+from routekitai.core.policy import Policy
+from routekitai.core.policy_adapter import PolicyAdapter
+from routekitai.core.runtime import Runtime
+from routekitai.core.tool import Tool
+from routekitai.graphs.graph import Graph
+from routekitai.observability.trace import Trace, TraceEvent
+
+
+class RunResult(BaseModel):
+    """Result of an agent run."""
+
+    output: Message = Field(..., description="Final output message")
+    trace_id: str = Field(..., description="Trace ID for this run")
+    final_state: dict[str, Any] = Field(default_factory=dict, description="Final agent state")
+    messages: list[Message] = Field(default_factory=list, description="All messages in the run")
+
+
+class Agent(BaseModel):
+    """Agent with model, tools, policy, and memory.
+
+    Provides a clean API for creating and running agents.
+    """
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    name: str = Field(..., description="Agent name")
+    model: Model = Field(..., description="Model used by the agent")
+    tools: list[Tool] = Field(default_factory=list, description="Tools available to the agent")
+    policy: Policy | dict[str, Any] | None = Field(
+        default=None, description="Agent policy or policy configuration"
+    )
+    memory: Memory | None = Field(default=None, description="Agent memory system")
+    tool_filter: ToolFilter | None = Field(
+        default=None, description="Agent-level tool allow/deny list"
+    )
+    trace_dir: Path | None = Field(default=None, description="Directory for trace files")
+
+    def __init__(self, **data: Any) -> None:
+        """Initialize agent with optional runtime."""
+        super().__init__(**data)
+        # Create internal runtime if not provided
+        if not hasattr(self, "_runtime"):
+            trace_dir = data.get("trace_dir")
+            if trace_dir is None:
+                trace_dir = Path(".routekit") / "traces"
+            self._runtime = Runtime(trace_dir=trace_dir)
+            self._runtime.register_agent(self)
+
+    async def run(self, prompt: str, **kwargs: Any) -> RunResult:
+        """Run the agent with a prompt.
+
+        Args:
+            prompt: User prompt
+            **kwargs: Additional run parameters
+
+        Returns:
+            RunResult with output, trace_id, final_state, and messages
+
+        Raises:
+            PolicyError: If agent policy fails
+        """
+        # Convert policy to PolicyAdapter if needed
+        policy_adapter = None
+        if self.policy:
+            if isinstance(self.policy, Policy):
+                policy_adapter = PolicyAdapter(self.policy)
+            elif isinstance(self.policy, dict):
+                # Legacy dict-based policy config - convert to appropriate policy
+                from routekitai.core.policies import ReActPolicy
+
+                policy_adapter = PolicyAdapter(ReActPolicy(**self.policy))
+
+        # Run via runtime
+        return await self._runtime.run(self.name, prompt, policy=policy_adapter, **kwargs)
+
+    async def run_stream(self, prompt: str, **kwargs: Any) -> AsyncIterator[dict[str, Any]]:
+        """Run the agent with streaming updates.
+
+        Yields trace events and progress updates in real-time.
+
+        Args:
+            prompt: User prompt
+            **kwargs: Additional run parameters
+
+        Yields:
+            Dicts with event information:
+            - type: Event type (trace_event, progress_update, result)
+            - data: Event data
+            - result: Final RunResult (only in last event)
+
+        Examples:
+            >>> async for event in agent.run_stream("Hello"):
+            ...     if event["type"] == "trace_event":
+            ...         print(f"Event: {event['data']['type']}")
+            ...     elif event["type"] == "progress_update":
+            ...         print(f"Progress: {event['data']['progress_percent']}%")
+        """
+        # Use queues to collect events and progress updates
+        trace_queue: asyncio.Queue[TraceEvent] = asyncio.Queue()
+        progress_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
+        result_queue: asyncio.Queue[RunResult] = asyncio.Queue()
+
+        def trace_callback(event: TraceEvent) -> None:
+            trace_queue.put_nowait(event)
+
+        def progress_callback(progress: dict[str, Any]) -> None:
+            progress_queue.put_nowait(progress)
+
+        # Convert policy to PolicyAdapter if needed
+        policy_adapter = None
+        if self.policy:
+            if isinstance(self.policy, Policy):
+                policy_adapter = PolicyAdapter(self.policy)
+            elif isinstance(self.policy, dict):
+                from routekitai.core.policies import ReActPolicy
+
+                policy_adapter = PolicyAdapter(ReActPolicy(**self.policy))
+
+        # Add callbacks to runtime
+        self._runtime.add_progress_callback(progress_callback)
+
+        # Create a trace to capture events
+        trace_id = str(uuid.uuid4())
+        trace = Trace(trace_id=trace_id, metadata={"agent": self.name, "prompt": prompt})
+        trace.add_event_callback(trace_callback)
+
+        # Run agent in background task
+        async def run_agent() -> None:
+            try:
+                result = await self._runtime.run(self.name, prompt, policy=policy_adapter, **kwargs)
+                result_queue.put_nowait(result)
+            except Exception as e:
+                # Put error in result queue
+                result_queue.put_nowait(None)  # type: ignore[arg-type]
+                trace_queue.put_nowait(
+                    TraceEvent(
+                        type="error",
+                        timestamp=asyncio.get_event_loop().time(),
+                        data={"error": str(e), "error_type": type(e).__name__},
+                    )
+                )
+
+        run_task = asyncio.create_task(run_agent())
+
+        try:
+            # Stream events as they arrive
+            while True:
+                # Check all queues with timeout
+                done, pending = await asyncio.wait(
+                    [
+                        asyncio.create_task(trace_queue.get()),
+                        asyncio.create_task(progress_queue.get()),
+                        asyncio.create_task(result_queue.get()),
+                        run_task,
+                    ],
+                    return_when=asyncio.FIRST_COMPLETED,
+                    timeout=0.1,
+                )
+
+                if not done:
+                    continue
+
+                for task in done:
+                    try:
+                        result = await task
+                        if task == run_task:
+                            # Agent run completed
+                            if result is not None:
+                                # Type guard: result should be RunResult when task is run_task
+                                from routekitai.core.agent import RunResult
+
+                                assert isinstance(result, RunResult), (
+                                    "Expected RunResult from run_task"
+                                )
+                                yield {
+                                    "type": "result",
+                                    "result": {
+                                        "output": result.output.model_dump(mode="json"),
+                                        "trace_id": result.trace_id,
+                                        "final_state": result.final_state,
+                                        "messages": [
+                                            msg.model_dump(mode="json") for msg in result.messages
+                                        ],
+                                    },
+                                }
+                            return
+                        elif isinstance(result, TraceEvent):
+                            yield {
+                                "type": "trace_event",
+                                "data": {
+                                    "type": result.type,
+                                    "timestamp": result.timestamp,
+                                    "data": result.data,
+                                },
+                            }
+                        elif isinstance(result, dict):
+                            yield {"type": "progress_update", "data": result}
+                    except asyncio.CancelledError:
+                        return
+                    except Exception:
+                        pass
+
+                # Cancel pending tasks
+                for task in pending:
+                    task.cancel()
+
+        finally:
+            # Clean up callbacks
+            self._runtime.remove_progress_callback(progress_callback)
+            if not run_task.done():
+                run_task.cancel()
+                try:
+                    await run_task
+                except asyncio.CancelledError:
+                    pass
+
+    def run_sync(self, prompt: str, **kwargs: Any) -> RunResult:
+        """Synchronous wrapper for agent.run().
+
+        Args:
+            prompt: User prompt
+            **kwargs: Additional run parameters
+
+        Returns:
+            RunResult with output, trace_id, final_state, and messages
+
+        Raises:
+            RuntimeError: If called from within an async context
+        """
+        try:
+            asyncio.get_running_loop()
+            # We're in an async context - need to use a different approach
+            # Try to use nest_asyncio if available, otherwise raise error
+            try:
+                import nest_asyncio
+
+                nest_asyncio.apply()
+                return asyncio.run(self.run(prompt, **kwargs))
+            except ImportError:
+                raise RuntimeError(
+                    "Cannot use run_sync() in an async context. "
+                    "Either use 'await agent.run()' or install nest_asyncio: pip install nest-asyncio"
+                ) from None
+        except RuntimeError:
+            # No event loop exists, safe to use asyncio.run
+            return asyncio.run(self.run(prompt, **kwargs))
+
+    @classmethod
+    def with_fake_model(
+        cls,
+        name: str,
+        responses: list[str | dict[str, Any]],
+        tools: list[Tool] | None = None,
+    ) -> Agent:
+        """Create agent with FakeModel for testing.
+
+        Args:
+            name: Agent name
+            responses: List of responses to queue in FakeModel
+            tools: Optional list of tools
+
+        Returns:
+            Agent instance with FakeModel
+
+        Examples:
+            >>> agent = Agent.with_fake_model("test", ["Hello!", "World!"])
+            >>> result = await agent.run("Hi")
+        """
+        from routekitai.providers.local import FakeModel
+
+        model = FakeModel(name=f"{name}_model")
+        for response in responses:
+            model.add_response(response)
+        return cls(name=name, model=model, tools=tools or [])
+
+    @classmethod
+    def with_graph_policy(
+        cls,
+        name: str,
+        graph: Graph,
+        model: Model,
+        runtime: Runtime | None = None,
+        **kwargs: Any,
+    ) -> Agent:
+        """Create agent with graph policy.
+
+        Args:
+            name: Agent name
+            graph: Graph to execute
+            model: Model instance
+            runtime: Optional runtime (will create one if not provided)
+            **kwargs: Additional agent parameters
+
+        Returns:
+            Agent instance with GraphPolicy
+
+        Examples:
+            >>> from routekitai.graphs import Graph, GraphNode, NodeType
+            >>> graph = Graph(name="test", entry_node="start", nodes=[...])
+            >>> agent = Agent.with_graph_policy("graph_agent", graph, model)
+        """
+
+        from routekitai.core.policies import GraphPolicy
+
+        if runtime is None:
+            runtime = Runtime(trace_dir=kwargs.get("trace_dir"))
+
+        policy = GraphPolicy(graph=graph, runtime=runtime)
+        return cls(name=name, model=model, policy=policy, **kwargs)

routekitai/core/errors.py

@@ -0,0 +1,49 @@
+"""Well-typed exceptions for RouteKit."""
+
+from typing import Any
+
+
+class RouteKitError(Exception):
+    """Base exception for all routkitai errors."""
+
+    def __init__(self, message: str, context: dict[str, Any] | None = None) -> None:
+        """Initialize error with message and optional context.
+
+        Args:
+            message: Error message
+            context: Optional context dictionary (trace_id, agent_name, step_id, etc.)
+        """
+        super().__init__(message)
+        self.message = message
+        self.context = context or {}
+
+    def __str__(self) -> str:
+        """Return error message with context if available."""
+        if self.context:
+            context_str = ", ".join(f"{k}={v}" for k, v in self.context.items())
+            return f"{self.message} (context: {context_str})"
+        return self.message
+
+
+class ModelError(RouteKitError):
+    """Error raised when model operations fail."""
+
+    pass
+
+
+class ToolError(RouteKitError):
+    """Error raised when tool execution fails."""
+
+    pass
+
+
+class PolicyError(RouteKitError):
+    """Error raised when agent policy fails."""
+
+    pass
+
+
+class RuntimeError(RouteKitError):
+    """Error raised when runtime operations fail."""
+
+    pass

routekitai/core/hooks.py

@@ -0,0 +1,174 @@
+"""Policy hooks for routkitai runtime."""
+
+import re
+from collections.abc import Callable
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class PIIRedactionHook(BaseModel):
+    """PII redaction hook for redacting sensitive information.
+
+    Uses regex patterns to identify and redact PII like emails and phone numbers.
+    """
+
+    redact_emails: bool = Field(default=True, description="Redact email addresses")
+    redact_phones: bool = Field(default=True, description="Redact phone numbers")
+    redact_patterns: list[tuple[str, str]] = Field(
+        default_factory=list, description="Custom (pattern, replacement) tuples"
+    )
+    replacement: str = Field(
+        default="[REDACTED]", description="Replacement string for redacted content"
+    )
+
+    # Compiled regex patterns
+    _email_pattern: re.Pattern[str] | None = None
+    _phone_pattern: re.Pattern[str] | None = None
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize PII redaction hook."""
+        super().__init__(**kwargs)
+        if self.redact_emails:
+            self._email_pattern = re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b")
+        if self.redact_phones:
+            # Matches various phone formats
+            self._phone_pattern = re.compile(
+                r"(\+?\d{1,3}[-.\s]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}"
+            )
+
+    def redact(self, text: str) -> str:
+        """Redact PII from text.
+
+        Args:
+            text: Text to redact
+
+        Returns:
+            Text with PII redacted
+        """
+        result = text
+
+        if self._email_pattern:
+            result = self._email_pattern.sub(self.replacement, result)
+
+        if self._phone_pattern:
+            result = self._phone_pattern.sub(self.replacement, result)
+
+        # Apply custom patterns
+        for pattern, replacement in self.redact_patterns:
+            result = re.sub(pattern, replacement, result)
+
+        return result
+
+    def redact_dict(self, data: dict[str, Any]) -> dict[str, Any]:
+        """Recursively redact PII from a dictionary.
+
+        Args:
+            data: Dictionary to redact
+
+        Returns:
+            Dictionary with PII redacted
+        """
+        redacted: dict[str, Any] = {}
+        for key, value in data.items():
+            if isinstance(value, str):
+                redacted[key] = self.redact(value)
+            elif isinstance(value, dict):
+                redacted[key] = self.redact_dict(value)
+            elif isinstance(value, list):
+                redacted[key] = [
+                    self.redact(item) if isinstance(item, str) else item for item in value
+                ]
+            else:
+                redacted[key] = value
+        return redacted
+
+
+class ToolFilter(BaseModel):
+    """Tool allow/deny list filter."""
+
+    allowed_tools: list[str] | None = Field(
+        default=None, description="List of allowed tool names (None = allow all)"
+    )
+    denied_tools: list[str] = Field(default_factory=list, description="List of denied tool names")
+
+    def is_allowed(self, tool_name: str) -> bool:
+        """Check if a tool is allowed.
+
+        Args:
+            tool_name: Name of the tool
+
+        Returns:
+            True if tool is allowed, False otherwise
+        """
+        # Deny list takes precedence
+        if tool_name in self.denied_tools:
+            return False
+
+        # If allow list is set, tool must be in it
+        if self.allowed_tools is not None:
+            return tool_name in self.allowed_tools
+
+        # Default: allow if not in deny list
+        return True
+
+
+class ApprovalGate(BaseModel):
+    """Approval gate for blocking tools until approved.
+
+    Uses a callback function to determine if a tool should be approved.
+    """
+
+    approval_callback: Callable[[str, dict[str, Any]], bool] | None = Field(
+        default=None, description="Callback(tool_name, tool_args) -> bool"
+    )
+    require_approval_for_permissions: list[str] = Field(
+        default_factory=list, description="Permissions that require approval"
+    )
+
+    def requires_approval(
+        self, tool_name: str, tool_args: dict[str, Any], tool_permissions: list[str]
+    ) -> bool:
+        """Check if a tool requires approval.
+
+        Args:
+            tool_name: Name of the tool
+            tool_args: Tool arguments
+            tool_permissions: Tool permissions
+
+        Returns:
+            True if approval is required
+        """
+        # Check if tool has permissions that require approval
+        if any(perm in self.require_approval_for_permissions for perm in tool_permissions):
+            return True
+
+        # Check callback if provided
+        if self.approval_callback:
+            return not self.approval_callback(tool_name, tool_args)
+
+        return False
+
+    def is_approved(self, tool_name: str, tool_args: dict[str, Any]) -> bool:
+        """Check if a tool is approved.
+
+        Args:
+            tool_name: Name of the tool
+            tool_args: Tool arguments
+
+        Returns:
+            True if approved, False otherwise
+        """
+        if self.approval_callback:
+            return self.approval_callback(tool_name, tool_args)
+        return True  # Default: approved if no callback
+
+
+class PolicyHooks(BaseModel):
+    """Collection of policy hooks for runtime."""
+
+    pii_redaction: PIIRedactionHook | None = Field(
+        default=None, description="PII redaction hook for messages and tool arguments"
+    )
+    tool_filter: ToolFilter | None = Field(default=None, description="Tool allow/deny list")
+    approval_gate: ApprovalGate | None = Field(default=None, description="Approval gate for tools")

routekitai/core/memory.py

@@ -0,0 +1,54 @@
+"""Memory interface for routkitai agents."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class Memory(ABC):
+    """Base interface for agent memory systems."""
+
+    @abstractmethod
+    async def get(self, key: str) -> Any:
+        """Get value by key.
+
+        Args:
+            key: Key to retrieve
+
+        Returns:
+            Stored value or None if not found
+        """
+        raise NotImplementedError("Subclasses must implement get")
+
+    @abstractmethod
+    async def set(self, key: str, value: Any) -> None:
+        """Set value by key.
+
+        Args:
+            key: Key to set
+            value: Value to store
+        """
+        raise NotImplementedError("Subclasses must implement set")
+
+    @abstractmethod
+    async def append(self, event: dict[str, Any]) -> None:
+        """Append an event to memory.
+
+        Args:
+            event: Event dictionary to append
+        """
+        raise NotImplementedError("Subclasses must implement append")
+
+    async def search(self, query: str, k: int = 5) -> list[dict[str, Any]]:
+        """Search memory (optional for retrieval memory).
+
+        Args:
+            query: Search query
+            k: Number of results to return
+
+        Returns:
+            List of matching results
+
+        Raises:
+            NotImplementedError: If search is not supported
+        """
+        raise NotImplementedError("Search not supported by this memory type")