ai-lib-python 0.5.0__py3-none-any.whl

ai_lib_python/plugins/middleware.py

@@ -0,0 +1,285 @@
+"""
+Middleware system for request/response processing.
+
+Provides a chain-of-responsibility pattern for processing.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+
+@dataclass
+class MiddlewareContext:
+    """Context for middleware execution.
+
+    Attributes:
+        request: Current request
+        response: Current response (if available)
+        model: Model identifier
+        provider: Provider identifier
+        metadata: Additional metadata
+        aborted: Whether processing was aborted
+    """
+
+    request: dict[str, Any] = field(default_factory=dict)
+    response: dict[str, Any] | None = None
+    model: str = ""
+    provider: str = ""
+    metadata: dict[str, Any] = field(default_factory=dict)
+    aborted: bool = False
+
+    def abort(self, response: dict[str, Any] | None = None) -> None:
+        """Abort middleware chain and return response.
+
+        Args:
+            response: Response to return
+        """
+        self.aborted = True
+        if response:
+            self.response = response
+
+
+class Middleware(ABC):
+    """Base class for middleware.
+
+    Middleware processes requests and responses in a chain.
+
+    Example:
+        >>> class LoggingMiddleware(Middleware):
+        ...     async def process(self, ctx, next):
+        ...         print(f"Request: {ctx.request}")
+        ...         response = await next(ctx)
+        ...         print(f"Response: {response}")
+        ...         return response
+    """
+
+    @property
+    def name(self) -> str:
+        """Get middleware name.
+
+        Returns:
+            Middleware name
+        """
+        return self.__class__.__name__
+
+    @abstractmethod
+    async def process(
+        self,
+        ctx: MiddlewareContext,
+        next: Callable[[MiddlewareContext], Awaitable[dict[str, Any] | None]],
+    ) -> dict[str, Any] | None:
+        """Process the request/response.
+
+        Args:
+            ctx: Middleware context
+            next: Next middleware in chain
+
+        Returns:
+            Response data
+        """
+        raise NotImplementedError
+
+
+class MiddlewareChain:
+    """Chain of middleware for processing.
+
+    Example:
+        >>> chain = MiddlewareChain()
+        >>> chain.use(LoggingMiddleware())
+        >>> chain.use(CachingMiddleware())
+        >>>
+        >>> async def handler(ctx):
+        ...     # Make actual request
+        ...     return response
+        >>>
+        >>> result = await chain.execute(ctx, handler)
+    """
+
+    def __init__(self) -> None:
+        """Initialize middleware chain."""
+        self._middleware: list[Middleware] = []
+
+    def use(self, middleware: Middleware) -> MiddlewareChain:
+        """Add middleware to the chain.
+
+        Args:
+            middleware: Middleware to add
+
+        Returns:
+            Self for chaining
+        """
+        self._middleware.append(middleware)
+        return self
+
+    def use_before(
+        self, middleware: Middleware, before: type | str
+    ) -> MiddlewareChain:
+        """Add middleware before another.
+
+        Args:
+            middleware: Middleware to add
+            before: Middleware class or name to insert before
+
+        Returns:
+            Self for chaining
+        """
+        for i, m in enumerate(self._middleware):
+            if (isinstance(before, type) and isinstance(m, before)) or (
+                isinstance(before, str) and m.name == before
+            ):
+                self._middleware.insert(i, middleware)
+                return self
+
+        # If not found, add at start
+        self._middleware.insert(0, middleware)
+        return self
+
+    def use_after(
+        self, middleware: Middleware, after: type | str
+    ) -> MiddlewareChain:
+        """Add middleware after another.
+
+        Args:
+            middleware: Middleware to add
+            after: Middleware class or name to insert after
+
+        Returns:
+            Self for chaining
+        """
+        for i, m in enumerate(self._middleware):
+            if (isinstance(after, type) and isinstance(m, after)) or (
+                isinstance(after, str) and m.name == after
+            ):
+                self._middleware.insert(i + 1, middleware)
+                return self
+
+        # If not found, add at end
+        self._middleware.append(middleware)
+        return self
+
+    def remove(self, middleware: type | str | Middleware) -> bool:
+        """Remove middleware from the chain.
+
+        Args:
+            middleware: Middleware to remove (class, name, or instance)
+
+        Returns:
+            True if removed, False if not found
+        """
+        for i, m in enumerate(self._middleware):
+            if (
+                m is middleware
+                or (isinstance(middleware, type) and isinstance(m, middleware))
+                or (isinstance(middleware, str) and m.name == middleware)
+            ):
+                self._middleware.pop(i)
+                return True
+        return False
+
+    async def execute(
+        self,
+        ctx: MiddlewareContext,
+        handler: Callable[[MiddlewareContext], Awaitable[dict[str, Any] | None]],
+    ) -> dict[str, Any] | None:
+        """Execute the middleware chain.
+
+        Args:
+            ctx: Middleware context
+            handler: Final handler to execute
+
+        Returns:
+            Response data
+        """
+        if not self._middleware:
+            return await handler(ctx)
+
+        # Build chain from end to start
+        chain = handler
+
+        for middleware in reversed(self._middleware):
+            # Capture middleware in closure
+            chain = self._create_next(middleware, chain)
+
+        return await chain(ctx)
+
+    def _create_next(
+        self,
+        middleware: Middleware,
+        next_handler: Callable[[MiddlewareContext], Awaitable[dict[str, Any] | None]],
+    ) -> Callable[[MiddlewareContext], Awaitable[dict[str, Any] | None]]:
+        """Create next handler in chain.
+
+        Args:
+            middleware: Current middleware
+            next_handler: Next handler
+
+        Returns:
+            Wrapped handler
+        """
+
+        async def handler(ctx: MiddlewareContext) -> dict[str, Any] | None:
+            if ctx.aborted:
+                return ctx.response
+            return await middleware.process(ctx, next_handler)
+
+        return handler
+
+    @property
+    def middleware_names(self) -> list[str]:
+        """Get list of middleware names."""
+        return [m.name for m in self._middleware]
+
+    def __len__(self) -> int:
+        """Get number of middleware."""
+        return len(self._middleware)
+
+
+class FunctionMiddleware(Middleware):
+    """Middleware created from a function.
+
+    Example:
+        >>> async def log_request(ctx, next):
+        ...     print(f"Request: {ctx.request}")
+        ...     return await next(ctx)
+        >>>
+        >>> middleware = FunctionMiddleware("logging", log_request)
+    """
+
+    def __init__(
+        self,
+        name: str,
+        func: Callable[
+            [
+                MiddlewareContext,
+                Callable[[MiddlewareContext], Awaitable[dict[str, Any] | None]],
+            ],
+            Awaitable[dict[str, Any] | None],
+        ],
+    ) -> None:
+        """Initialize function middleware.
+
+        Args:
+            name: Middleware name
+            func: Middleware function
+        """
+        self._name = name
+        self._func = func
+
+    @property
+    def name(self) -> str:
+        """Get middleware name."""
+        return self._name
+
+    async def process(
+        self,
+        ctx: MiddlewareContext,
+        next: Callable[[MiddlewareContext], Awaitable[dict[str, Any] | None]],
+    ) -> dict[str, Any] | None:
+        """Process using the function."""
+        return await self._func(ctx, next)

ai_lib_python/plugins/registry.py

@@ -0,0 +1,294 @@
+"""
+Plugin registry for managing plugins.
+
+Provides plugin registration, discovery, and lifecycle management.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from ai_lib_python.plugins.base import Plugin, PluginContext
+
+
+class PluginRegistry:
+    """Registry for managing plugins.
+
+    Handles plugin registration, lookup, and lifecycle.
+
+    Example:
+        >>> registry = PluginRegistry()
+        >>> registry.register(LoggingPlugin())
+        >>> registry.register(CachingPlugin())
+        >>>
+        >>> # Initialize all plugins
+        >>> await registry.init_all(ctx)
+        >>>
+        >>> # Process request through plugins
+        >>> request = await registry.process_request(ctx, request)
+    """
+
+    def __init__(self) -> None:
+        """Initialize plugin registry."""
+        self._plugins: dict[str, Plugin] = {}
+        self._ordered: list[Plugin] = []
+        self._initialized = False
+
+    def register(self, plugin: Plugin) -> PluginRegistry:
+        """Register a plugin.
+
+        Args:
+            plugin: Plugin to register
+
+        Returns:
+            Self for chaining
+        """
+        if plugin.name in self._plugins:
+            raise ValueError(f"Plugin already registered: {plugin.name}")
+
+        self._plugins[plugin.name] = plugin
+        self._update_order()
+        return self
+
+    def unregister(self, name: str) -> bool:
+        """Unregister a plugin.
+
+        Args:
+            name: Plugin name
+
+        Returns:
+            True if unregistered, False if not found
+        """
+        if name in self._plugins:
+            del self._plugins[name]
+            self._update_order()
+            return True
+        return False
+
+    def get(self, name: str) -> Plugin | None:
+        """Get a plugin by name.
+
+        Args:
+            name: Plugin name
+
+        Returns:
+            Plugin or None
+        """
+        return self._plugins.get(name)
+
+    def has(self, name: str) -> bool:
+        """Check if a plugin is registered.
+
+        Args:
+            name: Plugin name
+
+        Returns:
+            True if registered
+        """
+        return name in self._plugins
+
+    def enable(self, name: str) -> bool:
+        """Enable a plugin (no-op if already enabled).
+
+        Args:
+            name: Plugin name
+
+        Returns:
+            True if plugin exists
+        """
+        return name in self._plugins
+
+    def disable(self, name: str) -> bool:
+        """Disable a plugin by unregistering it.
+
+        Args:
+            name: Plugin name
+
+        Returns:
+            True if disabled
+        """
+        return self.unregister(name)
+
+    async def init_all(self, ctx: PluginContext | None = None) -> None:
+        """Initialize all registered plugins.
+
+        Args:
+            ctx: Plugin context
+        """
+        if self._initialized:
+            return
+
+        ctx = ctx or PluginContext()
+
+        for plugin in self._ordered:
+            if plugin.enabled:
+                await plugin.on_init(ctx)
+
+        self._initialized = True
+
+    async def shutdown_all(self, ctx: PluginContext | None = None) -> None:
+        """Shutdown all registered plugins.
+
+        Args:
+            ctx: Plugin context
+        """
+        if not self._initialized:
+            return
+
+        ctx = ctx or PluginContext()
+
+        for plugin in reversed(self._ordered):
+            if plugin.enabled:
+                await plugin.on_shutdown(ctx)
+
+        self._initialized = False
+
+    async def process_request(
+        self,
+        ctx: PluginContext,
+        request: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Process request through all plugins.
+
+        Args:
+            ctx: Plugin context
+            request: Request payload
+
+        Returns:
+            Modified request
+        """
+        for plugin in self._ordered:
+            if plugin.enabled:
+                request = await plugin.on_request(ctx, request)
+        return request
+
+    async def process_response(
+        self,
+        ctx: PluginContext,
+        response: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Process response through all plugins.
+
+        Args:
+            ctx: Plugin context
+            response: Response data
+
+        Returns:
+            Modified response
+        """
+        for plugin in reversed(self._ordered):
+            if plugin.enabled:
+                response = await plugin.on_response(ctx, response)
+        return response
+
+    async def process_error(
+        self,
+        ctx: PluginContext,
+        error: Exception,
+    ) -> Exception | None:
+        """Process error through all plugins.
+
+        Args:
+            ctx: Plugin context
+            error: The error
+
+        Returns:
+            Modified error or None to suppress
+        """
+        for plugin in self._ordered:
+            if plugin.enabled:
+                result = await plugin.on_error(ctx, error)
+                if result is None:
+                    return None
+                error = result
+        return error
+
+    async def on_stream_start(self, ctx: PluginContext) -> None:
+        """Notify plugins of stream start.
+
+        Args:
+            ctx: Plugin context
+        """
+        for plugin in self._ordered:
+            if plugin.enabled:
+                await plugin.on_stream_start(ctx)
+
+    async def process_stream_chunk(
+        self,
+        ctx: PluginContext,
+        chunk: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Process stream chunk through all plugins.
+
+        Args:
+            ctx: Plugin context
+            chunk: Stream chunk
+
+        Returns:
+            Modified chunk
+        """
+        for plugin in self._ordered:
+            if plugin.enabled:
+                chunk = await plugin.on_stream_chunk(ctx, chunk)
+        return chunk
+
+    async def on_stream_end(self, ctx: PluginContext) -> None:
+        """Notify plugins of stream end.
+
+        Args:
+            ctx: Plugin context
+        """
+        for plugin in reversed(self._ordered):
+            if plugin.enabled:
+                await plugin.on_stream_end(ctx)
+
+    def _update_order(self) -> None:
+        """Update the plugin execution order."""
+        self._ordered = sorted(
+            self._plugins.values(),
+            key=lambda p: p.priority,
+        )
+
+    @property
+    def plugins(self) -> list[Plugin]:
+        """Get all plugins in execution order."""
+        return list(self._ordered)
+
+    @property
+    def plugin_names(self) -> list[str]:
+        """Get all plugin names."""
+        return [p.name for p in self._ordered]
+
+    def __len__(self) -> int:
+        """Get number of registered plugins."""
+        return len(self._plugins)
+
+    def __contains__(self, name: str) -> bool:
+        """Check if plugin is registered."""
+        return name in self._plugins
+
+
+# Global plugin registry
+_global_registry: PluginRegistry | None = None
+
+
+def get_plugin_registry() -> PluginRegistry:
+    """Get the global plugin registry.
+
+    Returns:
+        Global PluginRegistry instance
+    """
+    global _global_registry
+    if _global_registry is None:
+        _global_registry = PluginRegistry()
+    return _global_registry
+
+
+def set_plugin_registry(registry: PluginRegistry) -> None:
+    """Set the global plugin registry.
+
+    Args:
+        registry: PluginRegistry instance
+    """
+    global _global_registry
+    _global_registry = registry

ai_lib_python/protocol/__init__.py

@@ -0,0 +1,71 @@
+"""
+Protocol layer - Protocol loading, validation, and manifest models.
+
+This module handles:
+- Loading protocol manifests from various sources
+- Validating manifests against JSON Schema
+- Protocol version validation
+- Strict streaming validation
+- Typed manifest models for runtime use
+"""
+
+from ai_lib_python.protocol.loader import ProtocolLoader
+from ai_lib_python.protocol.manifest import (
+    AccumulatorConfig,
+    AuthConfig,
+    AvailabilityConfig,
+    CandidateConfig,
+    CapabilitiesConfig,
+    DecoderConfig,
+    EndpointConfig,
+    ErrorClassification,
+    EventMapRule,
+    HealthCheckConfig,
+    ProtocolManifest,
+    RateLimitHeaders,
+    RetryPolicy,
+    ServiceConfig,
+    StreamingConfig,
+    ToolingConfig,
+    ToolUseConfig,
+)
+from ai_lib_python.protocol.validator import (
+    SUPPORTED_PROTOCOL_VERSIONS,
+    ProtocolValidator,
+    ValidationResult,
+    validate_manifest,
+    validate_manifest_or_raise,
+    validate_protocol_version,
+    validate_streaming_config,
+)
+
+__all__ = [
+    "AccumulatorConfig",
+    "AuthConfig",
+    "AvailabilityConfig",
+    "CandidateConfig",
+    "CapabilitiesConfig",
+    "DecoderConfig",
+    "EndpointConfig",
+    "ErrorClassification",
+    "EventMapRule",
+    "HealthCheckConfig",
+    # Loader and validator
+    "ProtocolLoader",
+    # Manifest models
+    "ProtocolManifest",
+    "ProtocolValidator",
+    "RateLimitHeaders",
+    "RetryPolicy",
+    "SUPPORTED_PROTOCOL_VERSIONS",
+    "ServiceConfig",
+    "StreamingConfig",
+    "ToolUseConfig",
+    "ToolingConfig",
+    # Validation functions
+    "ValidationResult",
+    "validate_manifest",
+    "validate_manifest_or_raise",
+    "validate_protocol_version",
+    "validate_streaming_config",
+]