foundry-mcp 0.8.22__py3-none-any.whl

foundry_mcp/core/providers/base.py

@@ -0,0 +1,515 @@
+"""
+Base provider abstractions for foundry-mcp.
+
+This module provides the core provider contracts adapted from sdd-toolkit,
+enabling pluggable LLM backends for CLI operations. The abstractions support
+capability negotiation, request/response normalization, and lifecycle hooks.
+
+Design principles:
+- Frozen dataclasses for immutability
+- Enum-based capabilities for type-safe routing
+- Status codes aligned with existing ProviderStatus patterns
+- Error hierarchy for granular exception handling
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Sequence, Set
+
+if TYPE_CHECKING:
+    from foundry_mcp.core.responses import ErrorType
+
+
+class ProviderCapability(Enum):
+    """
+    Feature flags a provider can expose to routing heuristics.
+
+    These capabilities enable callers to select providers based on
+    required features (vision, streaming, etc.) and allow registries
+    to route requests to appropriate backends.
+
+    Values:
+        TEXT: Basic text generation capability
+        VISION: Image/vision input processing
+        FUNCTION_CALLING: Tool/function invocation support
+        STREAMING: Incremental response streaming
+        THINKING: Extended reasoning/chain-of-thought support
+    """
+
+    TEXT = "text_generation"
+    VISION = "vision"
+    FUNCTION_CALLING = "function_calling"
+    STREAMING = "streaming"
+    THINKING = "thinking"
+
+
+class ProviderStatus(Enum):
+    """
+    Normalized execution outcomes emitted by providers.
+
+    These status codes provide a consistent interface for handling provider
+    responses across different backends, enabling uniform error handling
+    and retry logic.
+
+    Values:
+        SUCCESS: Operation completed successfully
+        TIMEOUT: Operation exceeded time limit (retryable)
+        NOT_FOUND: Provider/resource not available (not retryable)
+        INVALID_OUTPUT: Provider returned malformed response (not retryable)
+        ERROR: Generic error during execution (retryable)
+        CANCELED: Operation was explicitly canceled (not retryable)
+    """
+
+    SUCCESS = "success"
+    TIMEOUT = "timeout"
+    NOT_FOUND = "not_found"
+    INVALID_OUTPUT = "invalid_output"
+    ERROR = "error"
+    CANCELED = "canceled"
+
+    def is_retryable(self) -> bool:
+        """
+        Check if this status represents a transient failure that may succeed on retry.
+
+        Retryable statuses:
+            - TIMEOUT: May succeed with more time or if temporary resource contention resolves
+            - ERROR: Generic errors may be transient (network issues, rate limits, etc.)
+
+        Non-retryable statuses:
+            - SUCCESS: Already succeeded
+            - NOT_FOUND: Provider/tool not available (won't change on retry)
+            - INVALID_OUTPUT: Provider responded but output was malformed
+            - CANCELED: Explicitly canceled (shouldn't retry)
+
+        Returns:
+            True if the status is retryable, False otherwise
+        """
+        return self in (ProviderStatus.TIMEOUT, ProviderStatus.ERROR)
+
+    def to_error_type(self) -> "ErrorType":
+        """
+        Map provider status to foundry-mcp ErrorType for response envelopes.
+
+        This enables consistent error categorization across MCP and CLI surfaces,
+        allowing callers to handle provider errors using the standard error taxonomy.
+
+        Mapping:
+            - SUCCESS: raises ValueError (not an error state)
+            - TIMEOUT: UNAVAILABLE (503 analog, retryable)
+            - NOT_FOUND: NOT_FOUND (404 analog, not retryable)
+            - INVALID_OUTPUT: VALIDATION (400 analog, not retryable)
+            - ERROR: INTERNAL (500 analog, retryable)
+            - CANCELED: INTERNAL (operation aborted)
+
+        Returns:
+            ErrorType enum value corresponding to this status
+
+        Raises:
+            ValueError: If called on SUCCESS status (not an error)
+        """
+        from foundry_mcp.core.responses import ErrorType
+
+        if self == ProviderStatus.SUCCESS:
+            raise ValueError("SUCCESS status cannot be mapped to an error type")
+
+        mapping = {
+            ProviderStatus.TIMEOUT: ErrorType.UNAVAILABLE,
+            ProviderStatus.NOT_FOUND: ErrorType.NOT_FOUND,
+            ProviderStatus.INVALID_OUTPUT: ErrorType.VALIDATION,
+            ProviderStatus.ERROR: ErrorType.INTERNAL,
+            ProviderStatus.CANCELED: ErrorType.INTERNAL,
+        }
+        return mapping[self]
+
+
+@dataclass(frozen=True)
+class ProviderRequest:
+    """
+    Normalized request payload for provider execution.
+
+    This dataclass encapsulates all parameters needed to make a generation
+    request to any provider backend. Fields follow common LLM API conventions
+    to ensure portability across different providers.
+
+    Attributes:
+        prompt: The user's input prompt/message
+        system_prompt: Optional system/instruction prompt
+        model: Model identifier (provider-specific, e.g., "pro", "flash")
+        timeout: Request timeout in seconds (None = provider default)
+        temperature: Sampling temperature (0.0-2.0, None = provider default)
+        max_tokens: Maximum output tokens (None = provider default)
+        metadata: Arbitrary request metadata (tracing IDs, feature flags, etc.)
+        stream: Whether to request streaming response
+        attachments: File paths or URIs for multimodal inputs
+    """
+
+    prompt: str
+    system_prompt: Optional[str] = None
+    model: Optional[str] = None
+    timeout: Optional[float] = None
+    temperature: Optional[float] = None
+    max_tokens: Optional[int] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    stream: bool = False
+    attachments: Sequence[str] = field(default_factory=list)
+
+
+@dataclass(frozen=True)
+class TokenUsage:
+    """
+    Token accounting information reported by providers.
+
+    Tracks input, output, and cached tokens for cost estimation
+    and usage monitoring.
+
+    Attributes:
+        input_tokens: Tokens consumed by the prompt
+        output_tokens: Tokens generated in the response
+        cached_input_tokens: Tokens served from cache (if supported)
+        total_tokens: Sum of all token counts
+    """
+
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cached_input_tokens: int = 0
+    total_tokens: int = 0
+
+
+@dataclass(frozen=True)
+class ProviderResult:
+    """
+    Normalized provider response.
+
+    This dataclass encapsulates all data returned from a provider execution,
+    providing a consistent interface regardless of the backend used.
+
+    Attributes:
+        content: Final text output (aggregated if streaming was used)
+        provider_id: Canonical provider identifier (e.g., "gemini", "codex")
+        model_used: Fully-qualified model identifier (e.g., "gemini:pro")
+        status: ProviderStatus describing execution outcome
+        tokens: Token usage data (if reported by provider)
+        duration_ms: Execution duration in milliseconds
+        stderr: Captured stderr/log output for debugging
+        raw_payload: Provider-specific metadata (traces, debug info, etc.)
+    """
+
+    content: str
+    provider_id: str
+    model_used: str
+    status: ProviderStatus
+    tokens: TokenUsage = field(default_factory=TokenUsage)
+    duration_ms: Optional[float] = None
+    stderr: Optional[str] = None
+    raw_payload: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class ModelDescriptor:
+    """
+    Describes a model supported by a provider.
+
+    Attributes:
+        id: Provider-specific model identifier (e.g., "pro", "flash")
+        display_name: Human-friendly name for UIs/logs
+        capabilities: Feature flags supported by this model
+        routing_hints: Optional metadata for routing (cost, latency, etc.)
+    """
+
+    id: str
+    display_name: Optional[str] = None
+    capabilities: Set[ProviderCapability] = field(default_factory=set)
+    routing_hints: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class ProviderMetadata:
+    """
+    Provider-level metadata shared with registries and consumers.
+
+    This dataclass describes a provider's capabilities, supported models,
+    and configuration, enabling informed routing decisions.
+
+    Attributes:
+        provider_id: Canonical provider identifier (e.g., "gemini", "codex")
+        display_name: Human-friendly provider name
+        models: Supported model descriptors
+        default_model: Model ID used when no override supplied
+        capabilities: Aggregate capabilities across all models
+        security_flags: Provider-specific sandbox/safety configuration
+        extra: Arbitrary metadata (version info, auth requirements, etc.)
+    """
+
+    provider_id: str
+    display_name: Optional[str] = None
+    models: Sequence[ModelDescriptor] = field(default_factory=list)
+    default_model: Optional[str] = None
+    capabilities: Set[ProviderCapability] = field(default_factory=set)
+    security_flags: Dict[str, Any] = field(default_factory=dict)
+    extra: Dict[str, Any] = field(default_factory=dict)
+
+
+# =============================================================================
+# Error Hierarchy
+# =============================================================================
+
+
+class ProviderError(RuntimeError):
+    """Base exception for provider orchestration errors."""
+
+    def __init__(self, message: str, *, provider: Optional[str] = None):
+        self.provider = provider
+        super().__init__(message)
+
+
+class ProviderUnavailableError(ProviderError):
+    """Raised when a provider cannot be instantiated (binary missing, auth issues)."""
+
+
+class ProviderExecutionError(ProviderError):
+    """Raised when a provider command returns a non-retryable error."""
+
+
+class ProviderTimeoutError(ProviderError):
+    """Raised when a provider exceeds its allotted execution time."""
+
+
+class ContextWindowError(ProviderExecutionError):
+    """Raised when prompt exceeds the model's context window limit.
+
+    This error indicates the prompt/context size exceeded what the model
+    can process. It includes token counts to help with debugging and
+    provides actionable guidance for resolution.
+
+    Attributes:
+        prompt_tokens: Estimated tokens in the prompt (if known)
+        max_tokens: Maximum context window size (if known)
+        provider: Provider that raised the error
+        truncation_needed: How many tokens need to be removed
+    """
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        provider: Optional[str] = None,
+        prompt_tokens: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+    ):
+        """Initialize context window error.
+
+        Args:
+            message: Error message describing the issue
+            provider: Provider ID that raised the error
+            prompt_tokens: Number of tokens in the prompt (if known)
+            max_tokens: Maximum tokens allowed (if known)
+        """
+        super().__init__(message, provider=provider)
+        self.prompt_tokens = prompt_tokens
+        self.max_tokens = max_tokens
+        self.truncation_needed = (
+            (prompt_tokens - max_tokens) if prompt_tokens and max_tokens else None
+        )
+
+
+# =============================================================================
+# Lifecycle Hooks
+# =============================================================================
+
+
+# Type aliases for hook callables
+StreamChunkCallback = Callable[["StreamChunk", ProviderMetadata], None]
+BeforeExecuteHook = Callable[[ProviderRequest, ProviderMetadata], None]
+AfterResultHook = Callable[[ProviderResult, ProviderMetadata], None]
+
+
+@dataclass(frozen=True)
+class StreamChunk:
+    """Represents a streamed fragment emitted by the provider."""
+
+    content: str
+    index: int
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class ProviderHooks:
+    """
+    Optional lifecycle hooks wired by the registry.
+
+    Hooks default to None (no-ops) so providers can invoke them unconditionally.
+    Registries can wire hooks for observability, logging, or streaming.
+    """
+
+    before_execute: Optional[BeforeExecuteHook] = None
+    on_stream_chunk: Optional[StreamChunkCallback] = None
+    after_result: Optional[AfterResultHook] = None
+
+    def emit_before(self, request: ProviderRequest, metadata: ProviderMetadata) -> None:
+        """Emit before-execution hook if registered."""
+        if self.before_execute:
+            self.before_execute(request, metadata)
+
+    def emit_stream(self, chunk: StreamChunk, metadata: ProviderMetadata) -> None:
+        """Emit stream chunk hook if registered."""
+        if self.on_stream_chunk:
+            self.on_stream_chunk(chunk, metadata)
+
+    def emit_after(self, result: ProviderResult, metadata: ProviderMetadata) -> None:
+        """Emit after-result hook if registered."""
+        if self.after_result:
+            self.after_result(result, metadata)
+
+
+# =============================================================================
+# Abstract Base Class
+# =============================================================================
+
+
+class ProviderContext(ABC):
+    """
+    Base class for provider implementations.
+
+    Subclasses should:
+        * Resolve CLI/environment dependencies during initialization
+        * Implement `_execute()` to run the underlying provider
+        * Return a populated `ProviderResult` from `_execute()`
+        * Emit streaming chunks via `self._emit_stream_chunk()` when
+          `request.stream` is True and the provider supports streaming
+
+    The `generate()` method is a template method that:
+        1. Calls `_prepare_request()` for any request modifications
+        2. Emits the `before_execute` hook
+        3. Calls the abstract `_execute()` method
+        4. Normalizes exceptions into typed ProviderErrors
+        5. Emits the `after_result` hook
+        6. Returns the result
+    """
+
+    def __init__(
+        self,
+        metadata: ProviderMetadata,
+        hooks: Optional[ProviderHooks] = None,
+    ):
+        self._metadata = metadata
+        self._hooks = hooks or ProviderHooks()
+
+    @property
+    def metadata(self) -> ProviderMetadata:
+        """Return provider metadata."""
+        return self._metadata
+
+    def supports(self, capability: ProviderCapability) -> bool:
+        """Return True if any registered model advertises the capability."""
+        # Check provider-level capabilities first
+        if capability in self._metadata.capabilities:
+            return True
+        # Then check model-level capabilities
+        return any(capability in model.capabilities for model in self._metadata.models)
+
+    def generate(self, request: ProviderRequest) -> ProviderResult:
+        """
+        Execute the provider with the supplied request (template method).
+
+        Applies lifecycle hooks, normalizes errors, and ensures ProviderStatus
+        is consistent across implementations.
+
+        Args:
+            request: The generation request
+
+        Returns:
+            ProviderResult with the generation output
+
+        Raises:
+            ProviderUnavailableError: If provider binary/auth unavailable
+            ProviderTimeoutError: If request exceeds timeout
+            ProviderExecutionError: For other execution errors
+        """
+        normalized_request = self._prepare_request(request)
+        self._hooks.emit_before(normalized_request, self._metadata)
+
+        try:
+            result = self._execute(normalized_request)
+        except ProviderTimeoutError:
+            raise
+        except ProviderUnavailableError:
+            raise
+        except ProviderError:
+            raise
+        except FileNotFoundError as exc:
+            raise ProviderUnavailableError(
+                str(exc), provider=self._metadata.provider_id
+            ) from exc
+        except TimeoutError as exc:
+            raise ProviderTimeoutError(
+                str(exc), provider=self._metadata.provider_id
+            ) from exc
+        except Exception as exc:  # noqa: BLE001 - intentionally wrap all provider exceptions
+            raise ProviderExecutionError(
+                str(exc), provider=self._metadata.provider_id
+            ) from exc
+
+        self._hooks.emit_after(result, self._metadata)
+        return result
+
+    def _prepare_request(self, request: ProviderRequest) -> ProviderRequest:
+        """
+        Allow subclasses to adjust request metadata before execution.
+
+        The default implementation simply returns the request unchanged.
+        Subclasses can override to inject defaults, normalize parameters, etc.
+        """
+        return request
+
+    def _emit_stream_chunk(self, chunk: StreamChunk) -> None:
+        """Helper for subclasses to publish streaming output through hooks."""
+        self._hooks.emit_stream(chunk, self._metadata)
+
+    @abstractmethod
+    def _execute(self, request: ProviderRequest) -> ProviderResult:
+        """
+        Subclasses must implement the actual provider invocation.
+
+        Args:
+            request: The (possibly modified) generation request
+
+        Returns:
+            ProviderResult with generated content and metadata
+        """
+        raise NotImplementedError
+
+
+# =============================================================================
+# Module Exports
+# =============================================================================
+
+
+__all__ = [
+    # Enums
+    "ProviderCapability",
+    "ProviderStatus",
+    # Request/Response dataclasses
+    "ProviderRequest",
+    "ProviderResult",
+    "TokenUsage",
+    "StreamChunk",
+    # Metadata dataclasses
+    "ModelDescriptor",
+    "ProviderMetadata",
+    # Hooks
+    "ProviderHooks",
+    "StreamChunkCallback",
+    "BeforeExecuteHook",
+    "AfterResultHook",
+    # Errors
+    "ProviderError",
+    "ProviderUnavailableError",
+    "ProviderExecutionError",
+    "ProviderTimeoutError",
+    "ContextWindowError",
+    # ABC
+    "ProviderContext",
+]