agent_os_kernel 3.1.0__py3-none-any.whl

amb_core/schema.py

@@ -0,0 +1,294 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Schema validation for AMB messages.
+
+This module provides schema registry and validation capabilities
+to ensure message payloads conform to expected schemas.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional, Type, Union, Callable
+from pydantic import BaseModel, ValidationError
+import json
+
+
+class SchemaValidationError(Exception):
+    """Raised when message schema validation fails."""
+    
+    def __init__(self, topic: str, errors: list, payload: Dict[str, Any]):
+        self.topic = topic
+        self.errors = errors
+        self.payload = payload
+        super().__init__(f"Schema validation failed for topic '{topic}': {errors}")
+
+
+class Schema(ABC):
+    """Abstract base class for message schemas."""
+    
+    @abstractmethod
+    def validate(self, payload: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Validate a payload against the schema.
+        
+        Args:
+            payload: The payload to validate
+            
+        Returns:
+            The validated payload (potentially coerced/normalized)
+            
+        Raises:
+            SchemaValidationError: If validation fails
+        """
+        pass
+
+
+class PydanticSchema(Schema):
+    """Schema backed by a Pydantic model."""
+    
+    def __init__(self, model: Type[BaseModel]):
+        """
+        Initialize with a Pydantic model.
+        
+        Args:
+            model: Pydantic model class to use for validation
+        """
+        self._model = model
+    
+    def validate(self, payload: Dict[str, Any]) -> Dict[str, Any]:
+        """Validate payload using Pydantic model."""
+        try:
+            validated = self._model.model_validate(payload)
+            return validated.model_dump()
+        except ValidationError as e:
+            raise SchemaValidationError(
+                topic="unknown",
+                errors=e.errors(),
+                payload=payload
+            )
+
+
+class DictSchema(Schema):
+    """Schema based on a dictionary specification with type checking."""
+    
+    def __init__(self, spec: Dict[str, type], strict: bool = False):
+        """
+        Initialize with a dictionary specification.
+        
+        Args:
+            spec: Dict mapping field names to expected types
+            strict: If True, reject payloads with extra fields
+            
+        Example:
+            schema = DictSchema({
+                "user_id": str,
+                "amount": float,
+                "timestamp": str
+            })
+        """
+        self._spec = spec
+        self._strict = strict
+    
+    def validate(self, payload: Dict[str, Any]) -> Dict[str, Any]:
+        """Validate payload against dictionary specification."""
+        errors = []
+        
+        # Check required fields
+        for field, expected_type in self._spec.items():
+            if field not in payload:
+                errors.append({
+                    "type": "missing",
+                    "loc": [field],
+                    "msg": f"Field '{field}' is required"
+                })
+            elif not isinstance(payload[field], expected_type):
+                errors.append({
+                    "type": "type_error",
+                    "loc": [field],
+                    "msg": f"Expected {expected_type.__name__}, got {type(payload[field]).__name__}"
+                })
+        
+        # Check for extra fields if strict
+        if self._strict:
+            extra_fields = set(payload.keys()) - set(self._spec.keys())
+            for field in extra_fields:
+                errors.append({
+                    "type": "extra_forbidden",
+                    "loc": [field],
+                    "msg": f"Extra field '{field}' not allowed"
+                })
+        
+        if errors:
+            raise SchemaValidationError(
+                topic="unknown",
+                errors=errors,
+                payload=payload
+            )
+        
+        return payload
+
+
+class CallableSchema(Schema):
+    """Schema using a custom validation function."""
+    
+    def __init__(self, validator: Callable[[Dict[str, Any]], Dict[str, Any]]):
+        """
+        Initialize with a validation function.
+        
+        Args:
+            validator: Function that takes payload and returns validated payload,
+                      or raises an exception on validation failure
+        """
+        self._validator = validator
+    
+    def validate(self, payload: Dict[str, Any]) -> Dict[str, Any]:
+        """Validate using custom function."""
+        try:
+            return self._validator(payload)
+        except Exception as e:
+            raise SchemaValidationError(
+                topic="unknown",
+                errors=[{"type": "validation_error", "msg": str(e)}],
+                payload=payload
+            )
+
+
+class SchemaRegistry:
+    """
+    Registry for message schemas.
+    
+    Provides centralized schema management for topic validation.
+    Supports multiple schema types: Pydantic models, dict specs, and custom validators.
+    
+    Example:
+        from pydantic import BaseModel
+        
+        class FraudAlertPayload(BaseModel):
+            transaction_id: str
+            amount: float
+            risk_score: float
+        
+        registry = SchemaRegistry()
+        registry.register("fraud.alerts", FraudAlertPayload)
+        
+        # Or with dict schema
+        registry.register("user.events", {
+            "user_id": str,
+            "event_type": str
+        })
+    """
+    
+    def __init__(self, strict: bool = True):
+        """
+        Initialize the schema registry.
+        
+        Args:
+            strict: If True, require schema for all topics when validating
+        """
+        self._schemas: Dict[str, Schema] = {}
+        self._strict = strict
+    
+    def register(
+        self,
+        topic: str,
+        schema: Union[Type[BaseModel], Dict[str, type], Schema, Callable]
+    ) -> None:
+        """
+        Register a schema for a topic.
+        
+        Args:
+            topic: Topic pattern to register schema for
+            schema: Schema specification (Pydantic model, dict, Schema instance, or callable)
+            
+        Example:
+            # Pydantic model
+            registry.register("fraud.alerts", FraudAlertSchema)
+            
+            # Dict specification
+            registry.register("user.events", {"user_id": str, "event": str})
+            
+            # Custom Schema instance
+            registry.register("custom.topic", MyCustomSchema())
+            
+            # Callable validator
+            registry.register("validated.topic", lambda p: p if p.get("valid") else raise_error())
+        """
+        if isinstance(schema, type) and issubclass(schema, BaseModel):
+            self._schemas[topic] = PydanticSchema(schema)
+        elif isinstance(schema, dict):
+            self._schemas[topic] = DictSchema(schema)
+        elif isinstance(schema, Schema):
+            self._schemas[topic] = schema
+        elif callable(schema):
+            self._schemas[topic] = CallableSchema(schema)
+        else:
+            raise TypeError(
+                f"Schema must be a Pydantic model, dict, Schema instance, or callable. "
+                f"Got {type(schema)}"
+            )
+    
+    def unregister(self, topic: str) -> bool:
+        """
+        Unregister a schema for a topic.
+        
+        Args:
+            topic: Topic to unregister
+            
+        Returns:
+            True if schema was removed, False if topic wasn't registered
+        """
+        if topic in self._schemas:
+            del self._schemas[topic]
+            return True
+        return False
+    
+    def has_schema(self, topic: str) -> bool:
+        """Check if a topic has a registered schema."""
+        return topic in self._schemas
+    
+    def get_schema(self, topic: str) -> Optional[Schema]:
+        """Get the schema for a topic."""
+        return self._schemas.get(topic)
+    
+    def validate(self, topic: str, payload: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Validate a payload for a given topic.
+        
+        Args:
+            topic: Topic the message is for
+            payload: Message payload to validate
+            
+        Returns:
+            Validated payload (potentially normalized)
+            
+        Raises:
+            SchemaValidationError: If validation fails
+            ValueError: If strict mode and no schema registered for topic
+        """
+        schema = self._schemas.get(topic)
+        
+        if schema is None:
+            if self._strict:
+                raise ValueError(f"No schema registered for topic '{topic}'")
+            return payload
+        
+        try:
+            return schema.validate(payload)
+        except SchemaValidationError as e:
+            # Update error with correct topic
+            raise SchemaValidationError(
+                topic=topic,
+                errors=e.errors,
+                payload=e.payload
+            )
+    
+    def list_topics(self) -> list:
+        """Get list of topics with registered schemas."""
+        return list(self._schemas.keys())
+    
+    def __len__(self) -> int:
+        """Get number of registered schemas."""
+        return len(self._schemas)
+    
+    def __contains__(self, topic: str) -> bool:
+        """Check if topic has a schema registered."""
+        return topic in self._schemas

amb_core/tracing.py

@@ -0,0 +1,358 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Distributed tracing for AMB.
+
+This module provides distributed tracing capabilities for tracking
+message flow across agents and services.
+"""
+
+from typing import Any, Dict, Optional, List
+from datetime import datetime, timezone
+from dataclasses import dataclass, field
+from uuid import uuid4
+from contextvars import ContextVar
+import json
+
+
+# Context variable for current trace
+_current_trace: ContextVar[Optional["TraceContext"]] = ContextVar(
+    "current_trace", default=None
+)
+
+
+@dataclass
+class TraceSpan:
+    """
+    A span within a trace representing a single operation.
+    """
+    span_id: str
+    operation_name: str
+    trace_id: str
+    parent_span_id: Optional[str] = None
+    start_time: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    end_time: Optional[datetime] = None
+    tags: Dict[str, Any] = field(default_factory=dict)
+    logs: List[Dict[str, Any]] = field(default_factory=list)
+    status: str = "in_progress"  # in_progress, success, error
+    
+    def finish(self, status: str = "success") -> None:
+        """Mark span as finished."""
+        self.end_time = datetime.now(timezone.utc)
+        self.status = status
+    
+    def log(self, event: str, **kwargs) -> None:
+        """Add a log entry to the span."""
+        self.logs.append({
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "event": event,
+            **kwargs
+        })
+    
+    def set_tag(self, key: str, value: Any) -> None:
+        """Set a tag on the span."""
+        self.tags[key] = value
+    
+    @property
+    def duration_ms(self) -> Optional[float]:
+        """Get duration in milliseconds."""
+        if not self.end_time:
+            return None
+        delta = self.end_time - self.start_time
+        return delta.total_seconds() * 1000
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "span_id": self.span_id,
+            "operation_name": self.operation_name,
+            "trace_id": self.trace_id,
+            "parent_span_id": self.parent_span_id,
+            "start_time": self.start_time.isoformat(),
+            "end_time": self.end_time.isoformat() if self.end_time else None,
+            "duration_ms": self.duration_ms,
+            "tags": self.tags,
+            "logs": self.logs,
+            "status": self.status
+        }
+
+
+@dataclass
+class TraceContext:
+    """
+    Context for distributed tracing across message flows.
+    
+    TraceContext maintains trace and span IDs that propagate with messages,
+    allowing you to track the full journey of a message through the system.
+    
+    Example:
+        # Start a new trace
+        with TraceContext.start("process_order") as ctx:
+            await bus.publish("orders.new", payload, trace_id=ctx.trace_id)
+            
+        # Continue an existing trace
+        with TraceContext.from_message(message) as ctx:
+            ctx.log("Processing started")
+            # ... process message
+            ctx.log("Processing complete")
+    """
+    trace_id: str
+    span_id: str
+    parent_span_id: Optional[str] = None
+    baggage: Dict[str, str] = field(default_factory=dict)
+    spans: List[TraceSpan] = field(default_factory=list)
+    _current_span: Optional[TraceSpan] = field(default=None, repr=False)
+    
+    @classmethod
+    def new(cls, operation_name: str = "root") -> "TraceContext":
+        """
+        Create a new trace context.
+        
+        Args:
+            operation_name: Name of the root operation
+            
+        Returns:
+            New TraceContext
+        """
+        trace_id = str(uuid4())
+        span_id = str(uuid4())
+        
+        ctx = cls(
+            trace_id=trace_id,
+            span_id=span_id,
+            parent_span_id=None
+        )
+        
+        # Create root span
+        root_span = TraceSpan(
+            span_id=span_id,
+            operation_name=operation_name,
+            trace_id=trace_id
+        )
+        ctx.spans.append(root_span)
+        ctx._current_span = root_span
+        
+        return ctx
+    
+    @classmethod
+    def from_headers(cls, headers: Dict[str, str]) -> "TraceContext":
+        """
+        Create trace context from message headers.
+        
+        Args:
+            headers: Headers containing trace information
+            
+        Returns:
+            TraceContext (new or continued)
+        """
+        trace_id = headers.get("x-trace-id")
+        parent_span_id = headers.get("x-span-id")
+        baggage_json = headers.get("x-trace-baggage", "{}")
+        
+        if not trace_id:
+            return cls.new()
+        
+        try:
+            baggage = json.loads(baggage_json)
+        except json.JSONDecodeError:
+            baggage = {}
+        
+        span_id = str(uuid4())
+        
+        return cls(
+            trace_id=trace_id,
+            span_id=span_id,
+            parent_span_id=parent_span_id,
+            baggage=baggage
+        )
+    
+    @classmethod
+    def from_message(cls, message: "Message") -> "TraceContext":  # noqa: F821
+        """
+        Extract trace context from a message.
+        
+        Args:
+            message: Message with trace metadata
+            
+        Returns:
+            TraceContext (new or continued)
+        """
+        trace_id = message.metadata.get("trace_id")
+        parent_span_id = message.metadata.get("span_id")
+        baggage = message.metadata.get("trace_baggage", {})
+        
+        if not trace_id:
+            return cls.new()
+        
+        span_id = str(uuid4())
+        
+        return cls(
+            trace_id=trace_id,
+            span_id=span_id,
+            parent_span_id=parent_span_id,
+            baggage=baggage if isinstance(baggage, dict) else {}
+        )
+    
+    @classmethod
+    def current(cls) -> Optional["TraceContext"]:
+        """Get current trace context from context var."""
+        return _current_trace.get()
+    
+    @classmethod
+    def start(cls, operation_name: str = "root") -> "TraceContext":
+        """
+        Start a new trace and set as current.
+        
+        Args:
+            operation_name: Name of the root operation
+            
+        Returns:
+            New TraceContext
+        """
+        ctx = cls.new(operation_name)
+        _current_trace.set(ctx)
+        return ctx
+    
+    def start_span(self, operation_name: str) -> TraceSpan:
+        """
+        Start a new child span.
+        
+        Args:
+            operation_name: Name of the operation
+            
+        Returns:
+            New TraceSpan
+        """
+        parent_id = self._current_span.span_id if self._current_span else self.span_id
+        
+        span = TraceSpan(
+            span_id=str(uuid4()),
+            operation_name=operation_name,
+            trace_id=self.trace_id,
+            parent_span_id=parent_id
+        )
+        
+        self.spans.append(span)
+        self._current_span = span
+        return span
+    
+    def finish_span(self, status: str = "success") -> None:
+        """Finish the current span."""
+        if self._current_span:
+            self._current_span.finish(status)
+            
+            # Find parent span
+            if self._current_span.parent_span_id:
+                for span in self.spans:
+                    if span.span_id == self._current_span.parent_span_id:
+                        self._current_span = span
+                        return
+            
+            self._current_span = None
+    
+    def log(self, event: str, **kwargs) -> None:
+        """Log an event to the current span."""
+        if self._current_span:
+            self._current_span.log(event, **kwargs)
+    
+    def set_tag(self, key: str, value: Any) -> None:
+        """Set a tag on the current span."""
+        if self._current_span:
+            self._current_span.set_tag(key, value)
+    
+    def set_baggage(self, key: str, value: str) -> None:
+        """
+        Set baggage item that propagates with the trace.
+        
+        Baggage items are key-value pairs that travel with the trace
+        across all services.
+        """
+        self.baggage[key] = value
+    
+    def get_baggage(self, key: str) -> Optional[str]:
+        """Get a baggage item."""
+        return self.baggage.get(key)
+    
+    def to_headers(self) -> Dict[str, str]:
+        """
+        Convert trace context to headers for propagation.
+        
+        Returns:
+            Headers dict
+        """
+        return {
+            "x-trace-id": self.trace_id,
+            "x-span-id": self.span_id,
+            "x-parent-span-id": self.parent_span_id or "",
+            "x-trace-baggage": json.dumps(self.baggage)
+        }
+    
+    def to_message_metadata(self) -> Dict[str, Any]:
+        """
+        Convert trace context to message metadata.
+        
+        Returns:
+            Metadata dict to add to message
+        """
+        return {
+            "trace_id": self.trace_id,
+            "span_id": self.span_id,
+            "parent_span_id": self.parent_span_id,
+            "trace_baggage": self.baggage
+        }
+    
+    def __enter__(self) -> "TraceContext":
+        """Context manager entry."""
+        _current_trace.set(self)
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Context manager exit."""
+        if exc_type:
+            self.finish_span(status="error")
+        else:
+            self.finish_span(status="success")
+        _current_trace.set(None)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "trace_id": self.trace_id,
+            "span_id": self.span_id,
+            "parent_span_id": self.parent_span_id,
+            "baggage": self.baggage,
+            "spans": [s.to_dict() for s in self.spans]
+        }
+
+
+def get_current_trace() -> Optional[TraceContext]:
+    """Get the current trace context."""
+    return _current_trace.get()
+
+
+def inject_trace(message: "Message") -> "Message":  # noqa: F821
+    """
+    Inject current trace context into a message.
+    
+    Args:
+        message: Message to inject trace into
+        
+    Returns:
+        Message with trace metadata
+    """
+    ctx = get_current_trace()
+    if ctx:
+        message.metadata.update(ctx.to_message_metadata())
+    return message
+
+
+def extract_trace(message: "Message") -> TraceContext:  # noqa: F821
+    """
+    Extract trace context from a message.
+    
+    Args:
+        message: Message to extract trace from
+        
+    Returns:
+        TraceContext (new or continued)
+    """
+    return TraceContext.from_message(message)