agent_os_kernel 3.1.0__py3-none-any.whl

agent_os_observability/tracer.py

@@ -0,0 +1,232 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+OpenTelemetry Tracing for Agent OS Kernel.
+
+Every kernel operation emits traces for debugging and compliance.
+"""
+
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.trace import Status, StatusCode
+from functools import wraps
+from typing import Any, Callable, Optional
+import time
+
+
+class KernelTracer:
+    """
+    OpenTelemetry tracer for Agent OS kernel.
+    
+    Usage:
+        tracer = KernelTracer(service_name="agent-os")
+        
+        # Trace an operation
+        with tracer.span("policy_check", {"agent_id": "001", "action": "query"}):
+            result = policy_engine.check(action)
+        
+        # Or use decorator
+        @tracer.trace("execute_action")
+        async def execute(action, params):
+            ...
+    """
+    
+    def __init__(
+        self,
+        service_name: str = "agent-os-kernel",
+        exporter=None,
+        attributes: Optional[dict] = None
+    ):
+        """
+        Initialize tracer.
+        
+        Args:
+            service_name: Name of the service
+            exporter: OpenTelemetry exporter (default: console)
+            attributes: Additional resource attributes
+        """
+        resource_attrs = {
+            "service.name": service_name,
+            "service.version": "0.4.0",
+            "deployment.environment": "production"
+        }
+        if attributes:
+            resource_attrs.update(attributes)
+        
+        resource = Resource.create(resource_attrs)
+        provider = TracerProvider(resource=resource)
+        
+        if exporter:
+            provider.add_span_processor(BatchSpanProcessor(exporter))
+        
+        trace.set_tracer_provider(provider)
+        self.tracer = trace.get_tracer(__name__)
+    
+    def span(self, name: str, attributes: Optional[dict] = None):
+        """
+        Create a span context manager.
+        
+        Usage:
+            with tracer.span("operation", {"key": "value"}):
+                do_work()
+        """
+        span = self.tracer.start_span(name)
+        if attributes:
+            for key, value in attributes.items():
+                span.set_attribute(key, value)
+        return SpanContext(span)
+    
+    def trace(self, name: str, extract_attributes: Optional[Callable] = None):
+        """
+        Decorator to trace a function.
+        
+        Usage:
+            @tracer.trace("my_function")
+            def my_function(x, y):
+                return x + y
+        """
+        def decorator(func):
+            @wraps(func)
+            async def async_wrapper(*args, **kwargs):
+                attrs = {}
+                if extract_attributes:
+                    attrs = extract_attributes(*args, **kwargs)
+                
+                with self.span(name, attrs) as span:
+                    try:
+                        result = await func(*args, **kwargs)
+                        span.set_attribute("status", "success")
+                        return result
+                    except Exception as e:
+                        span.record_exception(e)
+                        span.set_status(Status(StatusCode.ERROR, str(e)))
+                        raise
+            
+            @wraps(func)
+            def sync_wrapper(*args, **kwargs):
+                attrs = {}
+                if extract_attributes:
+                    attrs = extract_attributes(*args, **kwargs)
+                
+                with self.span(name, attrs) as span:
+                    try:
+                        result = func(*args, **kwargs)
+                        span.set_attribute("status", "success")
+                        return result
+                    except Exception as e:
+                        span.record_exception(e)
+                        span.set_status(Status(StatusCode.ERROR, str(e)))
+                        raise
+            
+            import asyncio
+            if asyncio.iscoroutinefunction(func):
+                return async_wrapper
+            return sync_wrapper
+        
+        return decorator
+    
+    # =========================================================================
+    # Pre-built Kernel Traces
+    # =========================================================================
+    
+    def trace_policy_check(self, agent_id: str, action: str, policies: list):
+        """Start a policy check trace."""
+        return self.span("kernel.policy_check", {
+            "agent.id": agent_id,
+            "action": action,
+            "policies": ",".join(policies),
+            "kernel.component": "policy_engine"
+        })
+    
+    def trace_execution(self, agent_id: str, action: str):
+        """Start an execution trace."""
+        return self.span("kernel.execute", {
+            "agent.id": agent_id,
+            "action": action,
+            "kernel.component": "dispatcher"
+        })
+    
+    def trace_signal(self, agent_id: str, signal: str, reason: str):
+        """Start a signal trace."""
+        return self.span("kernel.signal", {
+            "agent.id": agent_id,
+            "signal": signal,
+            "reason": reason,
+            "kernel.component": "signal_dispatcher"
+        })
+    
+    def trace_violation(self, agent_id: str, action: str, policy: str, reason: str):
+        """Start a violation trace."""
+        return self.span("kernel.violation", {
+            "agent.id": agent_id,
+            "action": action,
+            "policy": policy,
+            "violation.reason": reason,
+            "kernel.component": "policy_engine",
+            "severity": "high"
+        })
+
+
+class SpanContext:
+    """Context manager for spans."""
+    
+    def __init__(self, span):
+        self.span = span
+        self.start_time = None
+    
+    def __enter__(self):
+        self.start_time = time.time()
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        duration = time.time() - self.start_time
+        self.span.set_attribute("duration_ms", duration * 1000)
+        
+        if exc_type:
+            self.span.record_exception(exc_val)
+            self.span.set_status(Status(StatusCode.ERROR, str(exc_val)))
+        else:
+            self.span.set_status(Status(StatusCode.OK))
+        
+        self.span.end()
+        return False
+    
+    def set_attribute(self, key: str, value: Any):
+        """Set an attribute on the span."""
+        self.span.set_attribute(key, value)
+    
+    def add_event(self, name: str, attributes: Optional[dict] = None):
+        """Add an event to the span."""
+        self.span.add_event(name, attributes or {})
+    
+    def record_exception(self, exception: Exception):
+        """Record an exception."""
+        self.span.record_exception(exception)
+
+    def set_status(self, status):
+        """Set the span status."""
+        self.span.set_status(status)
+
+
+def trace_operation(
+    tracer: KernelTracer,
+    name: str,
+    attributes: Optional[dict] = None
+):
+    """
+    Decorator to trace an operation.
+    
+    Usage:
+        @trace_operation(tracer, "my_operation", {"key": "value"})
+        def my_function():
+            pass
+    """
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            with tracer.span(name, attributes):
+                return func(*args, **kwargs)
+        return wrapper
+    return decorator

agent_primitives/__init__.py

@@ -0,0 +1,24 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Agent Primitives - Shared data models for Agent OS.
+
+This is a Layer 1 primitive package providing foundational models
+used across the Agent OS stack.
+"""
+
+__version__ = "3.1.0"
+
+from .failures import (
+    FailureType,
+    FailureSeverity,
+    FailureTrace,
+    AgentFailure,
+)
+
+__all__ = [
+    "FailureType",
+    "FailureSeverity",
+    "FailureTrace",
+    "AgentFailure",
+]

agent_primitives/failures.py

@@ -0,0 +1,84 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Failure models for Agent OS.
+
+These are the foundational failure tracking primitives used across the stack.
+Extracted to Layer 1 to allow proper dependency layering.
+"""
+
+from typing import Dict, List, Optional, Any
+from datetime import datetime
+from enum import Enum
+from pydantic import BaseModel, Field, ConfigDict
+
+
+class FailureType(str, Enum):
+    """Types of agent failures."""
+    BLOCKED_BY_CONTROL_PLANE = "blocked_by_control_plane"
+    TIMEOUT = "timeout"
+    INVALID_ACTION = "invalid_action"
+    RESOURCE_EXHAUSTED = "resource_exhausted"
+    LOGIC_ERROR = "logic_error"
+    UNKNOWN = "unknown"
+
+
+class FailureSeverity(str, Enum):
+    """Severity levels for failures."""
+    LOW = "low"
+    MEDIUM = "medium"
+    HIGH = "high"
+    CRITICAL = "critical"
+
+
+class FailureTrace(BaseModel):
+    """Full trace of an agent failure including reasoning chain."""
+    
+    user_prompt: str = Field(..., description="Original user prompt that led to failure")
+    chain_of_thought: List[str] = Field(default_factory=list, description="Agent's reasoning steps")
+    failed_action: Dict[str, Any] = Field(..., description="The action that failed")
+    error_details: str = Field(..., description="Detailed error information")
+    timestamp: datetime = Field(default_factory=datetime.utcnow)
+    
+    model_config = ConfigDict(
+        json_schema_extra={
+            "example": {
+                "user_prompt": "Delete the recent user records",
+                "chain_of_thought": [
+                    "User wants to delete records",
+                    "I need to identify which records are 'recent'",
+                    "I'll delete from users table"
+                ],
+                "failed_action": {
+                    "action": "execute_sql",
+                    "query": "DELETE FROM users WHERE created_at > '2024-01-01'"
+                },
+                "error_details": "Action blocked by control plane: Dangerous SQL query"
+            }
+        }
+    )
+
+
+class AgentFailure(BaseModel):
+    """Represents a failure detected in an agent."""
+    
+    agent_id: str = Field(..., description="Unique identifier for the agent")
+    failure_type: FailureType = Field(..., description="Type of failure")
+    severity: FailureSeverity = Field(default=FailureSeverity.MEDIUM)
+    timestamp: datetime = Field(default_factory=datetime.utcnow)
+    error_message: str = Field(..., description="Error message from the failure")
+    context: Dict[str, Any] = Field(default_factory=dict, description="Additional context")
+    stack_trace: Optional[str] = Field(None, description="Stack trace if available")
+    failure_trace: Optional[FailureTrace] = Field(None, description="Full failure trace if available")
+    
+    model_config = ConfigDict(
+        json_schema_extra={
+            "example": {
+                "agent_id": "agent-123",
+                "failure_type": "blocked_by_control_plane",
+                "severity": "high",
+                "error_message": "Agent action blocked by control plane policy",
+                "context": {"action": "delete_file", "resource": "/etc/passwd"}
+            }
+        }
+    )

amb_core/__init__.py

@@ -0,0 +1,177 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+AMB Core - A lightweight, broker-agnostic message bus for AI Agents.
+
+AMB (Agent Message Bus) provides a decoupled communication layer that allows
+AI agents to emit signals, broadcast intentions, and coordinate without tight
+coupling between senders and receivers.
+
+Key Features:
+    - Broker-agnostic: Works with Redis, RabbitMQ, Kafka, or in-memory
+    - Async-first: Built on asyncio/anyio for non-blocking operation
+    - Multiple patterns: Fire-and-forget, acknowledgment, request-response
+    - Type-safe: Full type hints with Pydantic validation
+    
+New in v0.2.0:
+    - Message persistence for replay capability (AMB-001)
+    - Dead Letter Queue for failed messages (AMB-002)
+    - Schema validation via SchemaRegistry (AMB-003)
+    - Distributed tracing support (AMB-004)
+    - Message prioritization (AMB-005)
+    - Message TTL/expiration (AMB-007)
+
+Quick Start:
+    >>> import asyncio
+    >>> from amb_core import MessageBus, Message
+    >>>
+    >>> async def main():
+    ...     async with MessageBus() as bus:
+    ...         await bus.publish("agent.thoughts", {"thought": "Hello!"})
+    >>>
+    >>> asyncio.run(main())
+
+Advanced Usage:
+    >>> from amb_core import MessageBus, SchemaRegistry, Priority
+    >>> from pydantic import BaseModel
+    >>>
+    >>> class FraudAlert(BaseModel):
+    ...     transaction_id: str
+    ...     risk_score: float
+    >>>
+    >>> schemas = SchemaRegistry()
+    >>> schemas.register("fraud.alerts", FraudAlert)
+    >>>
+    >>> async def main():
+    ...     bus = MessageBus(
+    ...         persistence=True,
+    ...         schema_registry=schemas,
+    ...         dlq_enabled=True
+    ...     )
+    ...     async with bus:
+    ...         await bus.publish(
+    ...             "fraud.alerts",
+    ...             {"transaction_id": "tx123", "risk_score": 0.95},
+    ...             priority=Priority.CRITICAL,
+    ...             ttl_seconds=300
+    ...         )
+
+For more information, see: https://github.com/microsoft/agent-governance-toolkit
+"""
+
+from __future__ import annotations
+
+__version__ = "3.1.0"
+__author__ = "Microsoft Corporation"
+__license__ = "MIT"
+
+# Core models
+from amb_core.models import Message, MessagePriority, Priority, MessageStatus
+
+# Main bus
+from amb_core.bus import MessageBus
+
+# Broker interface
+from amb_core.broker import BrokerAdapter, MessageHandler
+from amb_core.memory_broker import InMemoryBroker
+
+# Schema validation (AMB-003)
+from amb_core.schema import (
+    SchemaRegistry,
+    Schema,
+    PydanticSchema,
+    DictSchema,
+    CallableSchema,
+    SchemaValidationError,
+)
+
+# Dead Letter Queue (AMB-002)
+from amb_core.dlq import (
+    DeadLetterQueue,
+    DLQEntry,
+    DLQReason,
+    DLQHandler,
+)
+
+# Persistence (AMB-001)
+from amb_core.persistence import (
+    MessageStore,
+    InMemoryMessageStore,
+    FileMessageStore,
+    PersistedMessage,
+    MessageStatus as PersistenceMessageStatus,
+)
+
+# Tracing (AMB-004)
+from amb_core.tracing import (
+    TraceContext,
+    TraceSpan,
+    get_current_trace,
+    inject_trace,
+    extract_trace,
+)
+
+# CloudEvents support (AMB-008)
+from amb_core.cloudevents import (
+    CloudEvent,
+    CloudEventBatch,
+    to_cloudevent,
+    from_cloudevent,
+    to_http_headers,
+    from_http_headers,
+    topic_to_type,
+    type_to_topic,
+    CLOUDEVENTS_SPEC_VERSION,
+    CLOUDEVENTS_CONTENT_TYPE,
+)
+
+__all__ = [
+    # Version info
+    "__version__",
+    "__author__",
+    "__license__",
+    # Core classes
+    "Message",
+    "MessagePriority",
+    "Priority",
+    "MessageStatus",
+    "MessageBus",
+    # Broker interface
+    "BrokerAdapter",
+    "MessageHandler",
+    "InMemoryBroker",
+    # Schema validation
+    "SchemaRegistry",
+    "Schema",
+    "PydanticSchema",
+    "DictSchema",
+    "CallableSchema",
+    "SchemaValidationError",
+    # DLQ
+    "DeadLetterQueue",
+    "DLQEntry",
+    "DLQReason",
+    "DLQHandler",
+    # Persistence
+    "MessageStore",
+    "InMemoryMessageStore",
+    "FileMessageStore",
+    "PersistedMessage",
+    # Tracing
+    "TraceContext",
+    "TraceSpan",
+    "get_current_trace",
+    "inject_trace",
+    "extract_trace",
+    # CloudEvents
+    "CloudEvent",
+    "CloudEventBatch",
+    "to_cloudevent",
+    "from_cloudevent",
+    "to_http_headers",
+    "from_http_headers",
+    "topic_to_type",
+    "type_to_topic",
+    "CLOUDEVENTS_SPEC_VERSION",
+    "CLOUDEVENTS_CONTENT_TYPE",
+]

amb_core/adapters/__init__.py

@@ -0,0 +1,57 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Broker adapters package for AMB.
+
+This module provides optional adapters for different message brokers.
+Each adapter requires its corresponding extra dependencies.
+
+Available Adapters:
+    - RedisBroker: pip install amb-core[redis]
+    - RabbitMQBroker: pip install amb-core[rabbitmq]
+    - KafkaBroker: pip install amb-core[kafka]
+    - NATSBroker: pip install amb-core[nats]
+    - AzureServiceBusBroker: pip install amb-core[azure]
+    - AWSSQSBroker: pip install amb-core[aws]
+
+Example:
+    >>> from amb_core.adapters.redis_broker import RedisBroker
+    >>> broker = RedisBroker(url="redis://localhost:6379/0")
+    
+    >>> from amb_core.adapters.nats_broker import NATSBroker
+    >>> broker = NATSBroker(servers=["nats://localhost:4222"])
+"""
+
+from typing import TYPE_CHECKING, List
+
+__all__: List[str] = [
+    "RedisBroker",
+    "RabbitMQBroker",
+    "KafkaBroker",
+    "NATSBroker",
+    "AzureServiceBusBroker",
+    "AWSSQSBroker",
+]
+
+
+def __getattr__(name: str):
+    """Lazy import adapters to avoid import errors when dependencies missing."""
+    if name == "RedisBroker":
+        from amb_core.adapters.redis_broker import RedisBroker
+        return RedisBroker
+    elif name == "RabbitMQBroker":
+        from amb_core.adapters.rabbitmq_broker import RabbitMQBroker
+        return RabbitMQBroker
+    elif name == "KafkaBroker":
+        from amb_core.adapters.kafka_broker import KafkaBroker
+        return KafkaBroker
+    elif name == "NATSBroker":
+        from amb_core.adapters.nats_broker import NATSBroker
+        return NATSBroker
+    elif name == "AzureServiceBusBroker":
+        from amb_core.adapters.azure_servicebus_broker import AzureServiceBusBroker
+        return AzureServiceBusBroker
+    elif name == "AWSSQSBroker":
+        from amb_core.adapters.aws_sqs_broker import AWSSQSBroker
+        return AWSSQSBroker
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")