agent_os_kernel 3.1.0__py3-none-any.whl

agent_os/integrations/profiling.py

@@ -0,0 +1,144 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Performance profiling for governance checks and adapter operations.
+
+Provides a decorator and context manager for measuring execution time,
+call counts, and memory usage of adapter methods.
+"""
+
+import functools
+import time
+import tracemalloc
+from dataclasses import dataclass, field
+from typing import Any, Callable, Optional
+
+
+@dataclass
+class MethodStats:
+    """Statistics for a single profiled method."""
+    name: str
+    call_count: int = 0
+    total_time_ms: float = 0.0
+    min_time_ms: float = float("inf")
+    max_time_ms: float = 0.0
+    total_memory_delta: int = 0
+
+    @property
+    def avg_time_ms(self) -> float:
+        return self.total_time_ms / self.call_count if self.call_count else 0.0
+
+
+@dataclass
+class ProfilingReport:
+    """Aggregated profiling results."""
+    methods: dict[str, MethodStats] = field(default_factory=dict)
+
+    @property
+    def total_calls(self) -> int:
+        return sum(m.call_count for m in self.methods.values())
+
+    @property
+    def total_time_ms(self) -> float:
+        return sum(m.total_time_ms for m in self.methods.values())
+
+    def format_report(self) -> str:
+        """Return a human-readable table of profiling results."""
+        if not self.methods:
+            return "No profiling data collected."
+        header = f"{'Method':<30} {'Calls':>6} {'Total ms':>10} {'Avg ms':>10} {'Min ms':>10} {'Max ms':>10}"
+        sep = "-" * len(header)
+        lines = [header, sep]
+        for stats in sorted(self.methods.values(), key=lambda s: s.total_time_ms, reverse=True):
+            lines.append(
+                f"{stats.name:<30} {stats.call_count:>6} "
+                f"{stats.total_time_ms:>10.2f} {stats.avg_time_ms:>10.2f} "
+                f"{stats.min_time_ms:>10.2f} {stats.max_time_ms:>10.2f}"
+            )
+        lines.append(sep)
+        lines.append(f"{'TOTAL':<30} {self.total_calls:>6} {self.total_time_ms:>10.2f}")
+        return "\n".join(lines)
+
+
+# Global report used by the decorator
+_global_report = ProfilingReport()
+
+
+def get_report() -> ProfilingReport:
+    """Retrieve the global profiling report."""
+    return _global_report
+
+
+def reset_report() -> None:
+    """Reset the global profiling report."""
+    _global_report.methods.clear()
+
+
+def profile_governance(func: Optional[Callable] = None, *, track_memory: bool = False):
+    """Decorator that profiles execution time (and optionally memory) of a method.
+
+    Usage:
+        @profile_governance
+        def my_method(self, ...): ...
+
+        @profile_governance(track_memory=True)
+        def my_method(self, ...): ...
+    """
+    def decorator(fn: Callable) -> Callable:
+        @functools.wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            key = fn.__name__
+            if key not in _global_report.methods:
+                _global_report.methods[key] = MethodStats(name=key)
+            stats = _global_report.methods[key]
+
+            mem_before = 0
+            if track_memory:
+                if not tracemalloc.is_tracing():
+                    tracemalloc.start()
+                _, mem_before = tracemalloc.get_traced_memory()
+
+            start = time.perf_counter()
+            try:
+                return fn(*args, **kwargs)
+            finally:
+                elapsed_ms = (time.perf_counter() - start) * 1000
+                stats.call_count += 1
+                stats.total_time_ms += elapsed_ms
+                stats.min_time_ms = min(stats.min_time_ms, elapsed_ms)
+                stats.max_time_ms = max(stats.max_time_ms, elapsed_ms)
+
+                if track_memory:
+                    _, mem_after = tracemalloc.get_traced_memory()
+                    stats.total_memory_delta += mem_after - mem_before
+
+        return wrapper
+
+    if func is not None:
+        return decorator(func)
+    return decorator
+
+
+class ProfileGovernanceContext:
+    """Context manager for scoped profiling.
+
+    Usage:
+        with ProfileGovernanceContext() as report:
+            # run profiled code
+        print(report.format_report())
+    """
+
+    def __init__(self, track_memory: bool = False):
+        self.track_memory = track_memory
+        self.report = ProfilingReport()
+        self._previous_report: Optional[ProfilingReport] = None
+
+    def __enter__(self) -> ProfilingReport:
+        global _global_report
+        self._previous_report = _global_report
+        _global_report = self.report
+        return self.report
+
+    def __exit__(self, *exc: Any) -> None:
+        global _global_report
+        _global_report = self._previous_report  # type: ignore

agent_os/integrations/pydantic_ai_adapter.py

@@ -0,0 +1,420 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+PydanticAI Integration for Agent-OS
+====================================
+
+Provides kernel-level governance for PydanticAI agent workflows.
+
+Features:
+- Policy enforcement for agent tool calls
+- Tool call interception via PydanticAI's tool system
+- Human approval workflows for sensitive operations
+- Call budget enforcement (max_tool_calls)
+- Audit logging for all tool executions
+- Blocked pattern detection in tool arguments
+- Graceful degradation when pydantic-ai is not installed
+
+Example:
+    >>> from agent_os.integrations.pydantic_ai_adapter import PydanticAIKernel
+    >>> from agent_os.integrations.base import GovernancePolicy
+    >>> from pydantic_ai import Agent
+    >>>
+    >>> policy = GovernancePolicy(
+    ...     max_tool_calls=10,
+    ...     allowed_tools=["search", "read_file"],
+    ...     blocked_patterns=["rm -rf", "DROP TABLE"],
+    ... )
+    >>> kernel = PydanticAIKernel(policy=policy)
+    >>>
+    >>> agent = Agent("openai:gpt-4o", system_prompt="You are helpful.")
+    >>> governed = kernel.wrap(agent)
+    >>>
+    >>> result = await governed.run("Analyze this data")
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from datetime import datetime, timezone
+from functools import wraps
+from typing import Any, Callable
+
+from .base import (
+    BaseIntegration,
+    ExecutionContext,
+    GovernancePolicy,
+    PolicyInterceptor,
+    PolicyViolationError,
+    ToolCallRequest,
+    ToolCallResult,
+)
+
+logger = logging.getLogger(__name__)
+
+# Graceful import handling for pydantic-ai
+try:
+    import pydantic_ai  # noqa: F401
+    HAS_PYDANTIC_AI = True
+except ImportError:
+    HAS_PYDANTIC_AI = False
+
+
+class HumanApprovalRequired(PolicyViolationError):
+    """Raised when a tool call requires human approval."""
+
+    def __init__(self, tool_name: str, arguments: dict[str, Any]):
+        self.tool_name = tool_name
+        self.arguments = arguments
+        super().__init__(
+            f"Tool '{tool_name}' requires human approval before execution"
+        )
+
+
+class PydanticAIKernel(BaseIntegration):
+    """
+    PydanticAI adapter for Agent OS.
+
+    Supports:
+    - Agent wrapping with governance (run / run_sync)
+    - Individual tool call interception (allowed_tools, blocked_patterns)
+    - Human approval workflows for sensitive tools
+    - Call budget enforcement (max_tool_calls)
+    - Audit logging of all tool executions
+    """
+
+    def __init__(
+        self,
+        policy: GovernancePolicy | None = None,
+        approval_callback: Callable[[str, dict[str, Any]], bool] | None = None,
+    ) -> None:
+        super().__init__(policy)
+        self._wrapped_agents: dict[int, Any] = {}
+        self._audit_log: list[dict[str, Any]] = []
+        self._approval_callback = approval_callback
+        self._start_time: float = time.monotonic()
+        self._last_error: str | None = None
+        logger.debug("PydanticAIKernel initialized with policy=%s", policy)
+
+    @property
+    def audit_log(self) -> list[dict[str, Any]]:
+        """Return the full audit log."""
+        return list(self._audit_log)
+
+    def _record_audit(
+        self,
+        event_type: str,
+        tool_name: str = "",
+        allowed: bool = True,
+        reason: str = "",
+        arguments: dict[str, Any] | None = None,
+        agent_id: str = "",
+    ) -> dict[str, Any]:
+        """Record an audit entry and return it."""
+        entry = {
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "event_type": event_type,
+            "tool_name": tool_name,
+            "allowed": allowed,
+            "reason": reason,
+            "arguments": arguments or {},
+            "agent_id": agent_id,
+        }
+        if self.policy.log_all_calls:
+            self._audit_log.append(entry)
+        return entry
+
+    def wrap(self, agent: Any) -> Any:
+        """
+        Wrap a PydanticAI Agent with governance.
+
+        Intercepts:
+        - agent.run() / agent.run_sync()
+        - All registered tool calls
+        - Result validation
+
+        Args:
+            agent: A pydantic_ai.Agent instance (or mock).
+
+        Returns:
+            A governed wrapper around the agent.
+        """
+        agent_id = getattr(agent, "name", None) or f"agent-{id(agent)}"
+        ctx = self.create_context(agent_id)
+        self._wrapped_agents[id(agent)] = agent
+
+        logger.info(
+            "Wrapping PydanticAI agent with governance: agent_id=%s", agent_id
+        )
+
+        original = agent
+        kernel = self
+
+        class GovernedPydanticAIAgent:
+            """PydanticAI agent wrapped with Agent OS governance."""
+
+            def __init__(self_inner):
+                self_inner._original = original
+                self_inner._ctx = ctx
+                self_inner._kernel = kernel
+                self_inner._agent_id = agent_id
+                self_inner._wrap_tools()
+
+            def _wrap_tools(self_inner):
+                """Intercept all tools registered on the agent."""
+                tools = _get_agent_tools(self_inner._original)
+                for tool_entry in tools:
+                    _wrap_single_tool(tool_entry, self_inner, kernel, ctx)
+
+            async def run(self_inner, prompt: str, **kwargs) -> Any:
+                """Governed async run."""
+                allowed, reason = kernel.pre_execute(ctx, prompt)
+                if not allowed:
+                    kernel._last_error = reason
+                    kernel._record_audit(
+                        "run_blocked",
+                        reason=reason or "",
+                        agent_id=agent_id,
+                    )
+                    raise PolicyViolationError(reason or "Pre-execution check failed")
+
+                kernel._record_audit(
+                    "run_start",
+                    agent_id=agent_id,
+                    reason=f"prompt_length={len(prompt)}",
+                )
+
+                try:
+                    result = await self_inner._original.run(prompt, **kwargs)
+                    kernel._record_audit("run_complete", agent_id=agent_id)
+                    return result
+                except PolicyViolationError:
+                    raise
+                except Exception as exc:
+                    kernel._last_error = str(exc)
+                    kernel._record_audit(
+                        "run_error",
+                        agent_id=agent_id,
+                        reason=str(exc),
+                        allowed=False,
+                    )
+                    raise
+
+            def run_sync(self_inner, prompt: str, **kwargs) -> Any:
+                """Governed sync run."""
+                allowed, reason = kernel.pre_execute(ctx, prompt)
+                if not allowed:
+                    kernel._last_error = reason
+                    kernel._record_audit(
+                        "run_blocked",
+                        reason=reason or "",
+                        agent_id=agent_id,
+                    )
+                    raise PolicyViolationError(reason or "Pre-execution check failed")
+
+                kernel._record_audit(
+                    "run_start",
+                    agent_id=agent_id,
+                    reason=f"prompt_length={len(prompt)}",
+                )
+
+                try:
+                    result = self_inner._original.run_sync(prompt, **kwargs)
+                    kernel._record_audit("run_complete", agent_id=agent_id)
+                    return result
+                except PolicyViolationError:
+                    raise
+                except Exception as exc:
+                    kernel._last_error = str(exc)
+                    kernel._record_audit(
+                        "run_error",
+                        agent_id=agent_id,
+                        reason=str(exc),
+                        allowed=False,
+                    )
+                    raise
+
+            @property
+            def original(self_inner) -> Any:
+                """Return the original unwrapped agent before governance wrapping."""
+                return self_inner._original
+
+            @property
+            def context(self_inner) -> ExecutionContext:
+                """Return the ExecutionContext tracking call counts and session state."""
+                return self_inner._ctx
+
+            def __getattr__(self_inner, name: str) -> Any:
+                return getattr(self_inner._original, name)
+
+        return GovernedPydanticAIAgent()
+
+    def unwrap(self, governed_agent: Any) -> Any:
+        """Remove governance wrapper and return original agent."""
+        if hasattr(governed_agent, "_original"):
+            return governed_agent._original
+        return governed_agent
+
+    def intercept_tool_call(
+        self,
+        ctx: ExecutionContext,
+        tool_name: str,
+        arguments: dict[str, Any],
+    ) -> ToolCallResult:
+        """
+        Evaluate a tool call against the governance policy.
+
+        Returns a ToolCallResult indicating whether the call is allowed.
+        """
+        # Handle human approval callback before the interceptor
+        if self.policy.require_human_approval:
+            if self._approval_callback:
+                approved = self._approval_callback(tool_name, arguments)
+                if not approved:
+                    return ToolCallResult(
+                        allowed=False,
+                        reason=f"Human approval denied for tool '{tool_name}'",
+                    )
+                # Approved — skip the interceptor's require_human_approval check
+                # by using a policy copy without the flag
+                from dataclasses import replace
+                policy_for_interceptor = replace(self.policy, require_human_approval=False)
+            else:
+                return ToolCallResult(
+                    allowed=False,
+                    reason=f"Tool '{tool_name}' requires human approval",
+                )
+        else:
+            policy_for_interceptor = self.policy
+
+        interceptor = PolicyInterceptor(policy_for_interceptor, ctx)
+        request = ToolCallRequest(
+            tool_name=tool_name,
+            arguments=arguments,
+            agent_id=ctx.agent_id,
+        )
+        return interceptor.intercept(request)
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get governance statistics."""
+        total_calls = sum(c.call_count for c in self.contexts.values())
+        return {
+            "total_sessions": len(self.contexts),
+            "wrapped_agents": len(self._wrapped_agents),
+            "total_tool_calls": total_calls,
+            "audit_entries": len(self._audit_log),
+            "policy": {
+                "max_tool_calls": self.policy.max_tool_calls,
+                "allowed_tools": self.policy.allowed_tools,
+                "blocked_patterns": [
+                    p if isinstance(p, str) else p[0]
+                    for p in self.policy.blocked_patterns
+                ],
+                "require_human_approval": self.policy.require_human_approval,
+            },
+        }
+
+    def health_check(self) -> dict[str, Any]:
+        """Return adapter health status."""
+        uptime = time.monotonic() - self._start_time
+        status = "degraded" if self._last_error else "healthy"
+        return {
+            "status": status,
+            "backend": "pydantic_ai",
+            "backend_available": HAS_PYDANTIC_AI,
+            "backend_connected": bool(self._wrapped_agents),
+            "last_error": self._last_error,
+            "uptime_seconds": round(uptime, 2),
+        }
+
+
+# ── Helper functions ──────────────────────────────────────────
+
+
+def _get_agent_tools(agent: Any) -> list:
+    """Extract the list of tool entries from a PydanticAI agent."""
+    # PydanticAI stores tools in _function_tools (list of Tool objects)
+    if hasattr(agent, "_function_tools"):
+        return list(agent._function_tools)
+    # Fallback for mocks or alternative structures
+    if hasattr(agent, "tools"):
+        tools = agent.tools
+        return list(tools) if tools else []
+    return []
+
+
+def _wrap_single_tool(
+    tool_entry: Any,
+    governed: Any,
+    kernel: PydanticAIKernel,
+    ctx: ExecutionContext,
+) -> None:
+    """Wrap a single tool's function with governance interception."""
+    if getattr(tool_entry, "_governed", False):
+        return
+
+    # Determine the tool name and callable
+    tool_name = getattr(tool_entry, "name", None) or getattr(
+        tool_entry, "__name__", str(tool_entry)
+    )
+    original_fn = getattr(tool_entry, "function", None) or getattr(
+        tool_entry, "_run", None
+    )
+    if original_fn is None:
+        return
+
+    @wraps(original_fn)
+    def governed_fn(*args: Any, **kwargs: Any) -> Any:
+        """Governed wrapper that validates and delegates PydanticAI tool calls."""
+        # Build arguments dict for policy check
+        call_args: dict[str, Any] = kwargs.copy()
+        if args:
+            call_args["_positional"] = list(args)
+
+        result = kernel.intercept_tool_call(ctx, tool_name, call_args)
+
+        if not result.allowed:
+            kernel._record_audit(
+                "tool_blocked",
+                tool_name=tool_name,
+                allowed=False,
+                reason=result.reason or "",
+                arguments=call_args,
+                agent_id=ctx.agent_id,
+            )
+            raise PolicyViolationError(
+                result.reason or f"Tool '{tool_name}' blocked by policy"
+            )
+
+        ctx.call_count += 1
+        kernel._record_audit(
+            "tool_executed",
+            tool_name=tool_name,
+            allowed=True,
+            arguments=call_args,
+            agent_id=ctx.agent_id,
+        )
+        return original_fn(*args, **kwargs)
+
+    # Patch the tool entry
+    if hasattr(tool_entry, "function"):
+        tool_entry.function = governed_fn
+    elif hasattr(tool_entry, "_run"):
+        tool_entry._run = governed_fn
+
+    tool_entry._governed = True
+
+
+# Convenience function
+def wrap(agent: Any, policy: GovernancePolicy | None = None, **kwargs) -> Any:
+    """Quick wrapper for PydanticAI agents."""
+    return PydanticAIKernel(policy, **kwargs).wrap(agent)
+
+
+__all__ = [
+    "PydanticAIKernel",
+    "HumanApprovalRequired",
+    "HAS_PYDANTIC_AI",
+    "wrap",
+]

agent_os/integrations/rate_limiter.py

@@ -0,0 +1,130 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Tool-call rate limiting tied to governance policy.
+
+This module enforces token-bucket limits for tool invocations, optionally scoped
+per agent and governed by ``GovernancePolicy.max_tool_calls``.
+
+See also:
+    - hypervisor.security.rate_limiter: runtime-layer per-agent/per-ring limits.
+    - agentmesh.services.rate_limiter: service/proxy-level limits in Agent Mesh.
+    - agentmesh.services.rate_limit_middleware: HTTP edge middleware in Agent Mesh.
+    - agent_os.policies.rate_limiting: shared token-bucket primitives.
+"""
+
+import threading
+import time
+from dataclasses import dataclass
+from typing import Optional
+
+from .base import GovernancePolicy
+
+
+@dataclass(frozen=True)
+class RateLimitStatus:
+    """Snapshot of an agent's rate-limit state."""
+    allowed: bool
+    remaining_calls: int
+    reset_at: float
+    wait_seconds: float
+
+
+class RateLimiter:
+    """Thread-safe token-bucket rate limiter for tool calls.
+
+    Args:
+        max_calls: Maximum number of calls allowed per time window (bucket size).
+        time_window: Duration of the time window in seconds.
+        per_agent: If ``True``, limits are tracked independently per agent.
+            If ``False``, a single global bucket is used for all agents.
+        policy: Optional GovernancePolicy whose ``max_tool_calls`` overrides
+            *max_calls*.
+    """
+
+    _GLOBAL_KEY = "__global__"
+
+    def __init__(
+        self,
+        max_calls: int = 10,
+        time_window: float = 60.0,
+        per_agent: bool = True,
+        policy: Optional[GovernancePolicy] = None,
+    ) -> None:
+        if max_calls <= 0:
+            raise ValueError("max_calls must be positive")
+        if time_window <= 0:
+            raise ValueError("time_window must be positive")
+
+        self._max_calls = policy.max_tool_calls if policy is not None else max_calls
+        self._time_window = float(time_window)
+        self._per_agent = per_agent
+        self._lock = threading.Lock()
+        # Each bucket: (tokens: float, last_refill: float)
+        self._buckets: dict[str, list] = {}
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _key(self, agent_id: str) -> str:
+        return agent_id if self._per_agent else self._GLOBAL_KEY
+
+    def _refill(self, bucket: list, now: float) -> None:
+        """Add tokens accrued since the last refill."""
+        elapsed = now - bucket[1]
+        if elapsed > 0:
+            rate = self._max_calls / self._time_window
+            bucket[0] = min(self._max_calls, bucket[0] + elapsed * rate)
+            bucket[1] = now
+
+    def _get_bucket(self, key: str, now: float) -> list:
+        bucket = self._buckets.get(key)
+        if bucket is None:
+            bucket = [float(self._max_calls), now]
+            self._buckets[key] = bucket
+        else:
+            self._refill(bucket, now)
+        return bucket
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def allow(self, agent_id: str) -> bool:
+        """Try to consume one token. Returns ``True`` if the call is allowed."""
+        now = time.monotonic()
+        with self._lock:
+            bucket = self._get_bucket(self._key(agent_id), now)
+            if bucket[0] >= 1.0:
+                bucket[0] -= 1.0
+                return True
+            return False
+
+    def check(self, agent_id: str) -> RateLimitStatus:
+        """Return current rate-limit status without consuming a token."""
+        now = time.monotonic()
+        with self._lock:
+            bucket = self._get_bucket(self._key(agent_id), now)
+            remaining = int(bucket[0])
+            allowed = remaining >= 1
+            if allowed:
+                wait = 0.0
+            else:
+                rate = self._max_calls / self._time_window
+                wait = (1.0 - bucket[0]) / rate if rate > 0 else 0.0
+            reset_at = now + self._time_window
+        return RateLimitStatus(
+            allowed=allowed,
+            remaining_calls=remaining,
+            reset_at=reset_at,
+            wait_seconds=wait,
+        )
+
+    def wait_time(self, agent_id: str) -> float:
+        """Return seconds until at least one token is available (0.0 if available now)."""
+        return self.check(agent_id).wait_seconds
+
+    def reset(self, agent_id: str) -> None:
+        """Reset the bucket for *agent_id* (or the global bucket if ``per_agent=False``)."""
+        with self._lock:
+            self._buckets.pop(self._key(agent_id), None)