Gemini MCP Server

"""Base interfaces and common behaviour for model providers.""" import logging import time from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, Callable, Optional if TYPE_CHECKING: from tools.models import ToolModelCategory from .shared import ModelCapabilities, ModelResponse, ProviderType logger = logging.getLogger(__name__) class ModelProvider(ABC): """Abstract base class for all model backends in the MCP server. Role Defines the interface every provider must implement so the registry, restriction service, and tools have a uniform surface for listing models, resolving aliases, and executing requests. Responsibilities * expose static capability metadata for each supported model via :class:`ModelCapabilities` * accept user prompts, forward them to the underlying SDK, and wrap responses in :class:`ModelResponse` * report tokenizer counts for budgeting and validation logic * advertise provider identity (``ProviderType``) so restriction policies can map environment configuration onto providers * validate whether a model name or alias is recognised by the provider Shared helpers like temperature validation, alias resolution, and restriction-aware ``list_models`` live here so concrete subclasses only need to supply their catalogue and wire up SDK-specific behaviour. """ # All concrete providers must define their supported models MODEL_CAPABILITIES: dict[str, Any] = {} def __init__(self, api_key: str, **kwargs): """Initialize the provider with API key and optional configuration.""" self.api_key = api_key self.config = kwargs self._sorted_capabilities_cache: Optional[list[tuple[str, ModelCapabilities]]] = None # ------------------------------------------------------------------ # Provider identity & capability surface # ------------------------------------------------------------------ @abstractmethod def get_provider_type(self) -> ProviderType: """Return the concrete provider identity.""" def get_capabilities(self, model_name: str) -> ModelCapabilities: """Resolve capability metadata for a model name. This centralises the alias resolution → lookup → restriction check pipeline so providers only override the pieces they genuinely need to customise. Subclasses usually only override ``_lookup_capabilities`` to integrate a registry or dynamic source, or ``_finalise_capabilities`` to tweak the returned object. Args: model_name: Canonical model name or its alias """ resolved_model_name = self._resolve_model_name(model_name) capabilities = self._lookup_capabilities(resolved_model_name, model_name) if capabilities is None: self._raise_unsupported_model(model_name) self._ensure_model_allowed(capabilities, resolved_model_name, model_name) return self._finalise_capabilities(capabilities, resolved_model_name, model_name) def get_all_model_capabilities(self) -> dict[str, ModelCapabilities]: """Return statically declared capabilities when available.""" model_map = getattr(self, "MODEL_CAPABILITIES", None) if isinstance(model_map, dict) and model_map: return {k: v for k, v in model_map.items() if isinstance(v, ModelCapabilities)} return {} def get_capabilities_by_rank(self) -> list[tuple[str, ModelCapabilities]]: """Return model capabilities sorted by effective capability rank.""" if self._sorted_capabilities_cache is not None: return list(self._sorted_capabilities_cache) model_configs = self.get_all_model_capabilities() if not model_configs: self._sorted_capabilities_cache = [] return [] items = list(model_configs.items()) items.sort(key=lambda item: (-item[1].get_effective_capability_rank(), item[0])) self._sorted_capabilities_cache = items return list(items) def _invalidate_capability_cache(self) -> None: """Clear cached sorted capability data (call after dynamic updates).""" self._sorted_capabilities_cache = None def list_models( self, *, respect_restrictions: bool = True, include_aliases: bool = True, lowercase: bool = False, unique: bool = False, ) -> list[str]: """Return formatted model names supported by this provider.""" model_configs = self.get_all_model_capabilities() if not model_configs: return [] restriction_service = None if respect_restrictions: from utils.model_restrictions import get_restriction_service restriction_service = get_restriction_service() if restriction_service: allowed_configs = {} for model_name, config in model_configs.items(): if restriction_service.is_allowed(self.get_provider_type(), model_name): allowed_configs[model_name] = config model_configs = allowed_configs if not model_configs: return [] return ModelCapabilities.collect_model_names( model_configs, include_aliases=include_aliases, lowercase=lowercase, unique=unique, ) # ------------------------------------------------------------------ # Request execution # ------------------------------------------------------------------ @abstractmethod def generate_content( self, prompt: str, model_name: str, system_prompt: Optional[str] = None, temperature: float = 0.3, max_output_tokens: Optional[int] = None, **kwargs, ) -> ModelResponse: """Generate content using the model. This is the core method that all providers must implement to generate responses from their models. Providers should handle model-specific capabilities and constraints appropriately. Args: prompt: The main user prompt/query to send to the model model_name: Canonical model name or its alias that the provider supports system_prompt: Optional system instructions to prepend to the prompt for establishing context, behavior, or role temperature: Controls randomness in generation (0.0=deterministic, 1.0=creative), default 0.3. Some models may not support temperature control max_output_tokens: Optional maximum number of tokens to generate in the response. If not specified, uses the model's default limit **kwargs: Additional provider-specific parameters that vary by implementation (e.g., thinking_mode for Gemini, top_p for OpenAI, images for vision models) Returns: ModelResponse: Standardized response object containing: - content: The generated text response - usage: Token usage statistics (input/output/total) - model_name: The model that was actually used - friendly_name: Human-readable provider/model identifier - provider: The ProviderType enum value - metadata: Provider-specific metadata (finish_reason, safety info, etc.) Raises: ValueError: If the model is not supported, parameters are invalid, or the model is restricted by policy RuntimeError: If the API call fails after retries """ def count_tokens(self, text: str, model_name: str) -> int: """Estimate token usage for a piece of text.""" resolved_model = self._resolve_model_name(model_name) if not text: return 0 estimated = max(1, len(text) // 4) logger.debug("Estimating %s tokens for model %s via character heuristic", estimated, resolved_model) return estimated def close(self) -> None: """Clean up any resources held by the provider.""" return # ------------------------------------------------------------------ # Retry helpers # ------------------------------------------------------------------ def _is_error_retryable(self, error: Exception) -> bool: """Return True when an error warrants another attempt. Subclasses with structured provider errors should override this hook. The default implementation only retries obvious transient failures such as timeouts or 5xx responses detected via string inspection. """ error_str = str(error).lower() if "429" in error_str or "rate limit" in error_str: return False retryable_indicators = [ "timeout", "connection", "temporary", "unavailable", "retry", "reset", "refused", "broken pipe", "tls", "handshake", "network", "500", "502", "503", "504", ] return any(indicator in error_str for indicator in retryable_indicators) def _run_with_retries( self, operation: Callable[[], Any], *, max_attempts: int, delays: Optional[list[float]] = None, log_prefix: str = "", ): """Execute ``operation`` with retry semantics. Args: operation: Callable returning the provider result. max_attempts: Maximum number of attempts (>=1). delays: Optional list of sleep durations between attempts. log_prefix: Optional identifier for log clarity. Returns: Whatever ``operation`` returns. Raises: The last exception when all retries fail or the error is not retryable. """ if max_attempts < 1: raise ValueError("max_attempts must be >= 1") attempts = max_attempts delays = delays or [] last_exc: Optional[Exception] = None for attempt_index in range(attempts): try: return operation() except Exception as exc: # noqa: BLE001 - bubble exact provider errors last_exc = exc attempt_number = attempt_index + 1 # Decide whether to retry based on subclass hook retryable = self._is_error_retryable(exc) if not retryable or attempt_number >= attempts: raise delay_idx = min(attempt_index, len(delays) - 1) if delays else -1 delay = delays[delay_idx] if delay_idx >= 0 else 0.0 if delay > 0: logger.warning( "%s retryable error (attempt %s/%s): %s. Retrying in %ss...", log_prefix or self.__class__.__name__, attempt_number, attempts, exc, delay, ) time.sleep(delay) else: logger.warning( "%s retryable error (attempt %s/%s): %s. Retrying...", log_prefix or self.__class__.__name__, attempt_number, attempts, exc, ) # Should never reach here because loop either returns or raises raise last_exc if last_exc else RuntimeError("Retry loop exited without result") # ------------------------------------------------------------------ # Validation hooks # ------------------------------------------------------------------ def validate_model_name(self, model_name: str) -> bool: """ Return ``True`` when the model resolves to an allowed capability. Args: model_name: Canonical model name or its alias """ try: self.get_capabilities(model_name) except ValueError: return False return True def validate_parameters(self, model_name: str, temperature: float, **kwargs) -> None: """ Validate model parameters against capabilities. Args: model_name: Canonical model name or its alias """ capabilities = self.get_capabilities(model_name) if not capabilities.temperature_constraint.validate(temperature): constraint_desc = capabilities.temperature_constraint.get_description() raise ValueError(f"Temperature {temperature} is invalid for model {model_name}. {constraint_desc}") # ------------------------------------------------------------------ # Preference / registry hooks # ------------------------------------------------------------------ def get_preferred_model(self, category: "ToolModelCategory", allowed_models: list[str]) -> Optional[str]: """Get the preferred model from this provider for a given category.""" return None def get_model_registry(self) -> Optional[dict[str, Any]]: """Return the model registry backing this provider, if any.""" return None # ------------------------------------------------------------------ # Capability lookup pipeline # ------------------------------------------------------------------ def _lookup_capabilities( self, canonical_name: str, requested_name: Optional[str] = None, ) -> Optional[ModelCapabilities]: """Return ``ModelCapabilities`` for the canonical model name.""" return self.get_all_model_capabilities().get(canonical_name) def _ensure_model_allowed( self, capabilities: ModelCapabilities, canonical_name: str, requested_name: str, ) -> None: """Raise ``ValueError`` if the model violates restriction policy.""" try: from utils.model_restrictions import get_restriction_service except Exception: # pragma: no cover - only triggered if service import breaks return restriction_service = get_restriction_service() if not restriction_service: return if restriction_service.is_allowed(self.get_provider_type(), canonical_name, requested_name): return raise ValueError( f"{self.get_provider_type().value} model '{canonical_name}' is not allowed by restriction policy." ) def _finalise_capabilities( self, capabilities: ModelCapabilities, canonical_name: str, requested_name: str, ) -> ModelCapabilities: """Allow subclasses to adjust capability metadata before returning.""" return capabilities def _raise_unsupported_model(self, model_name: str) -> None: """Raise the canonical unsupported-model error.""" raise ValueError(f"Unsupported model '{model_name}' for provider {self.get_provider_type().value}.") def _resolve_model_name(self, model_name: str) -> str: """Resolve model shorthand to full name. This implementation uses the hook methods to support different model configuration sources. Args: model_name: Canonical model name or its alias Returns: Resolved model name """ # Get model configurations from the hook method model_configs = self.get_all_model_capabilities() # First check if it's already a base model name (case-sensitive exact match) if model_name in model_configs: return model_name # Check case-insensitively for both base models and aliases model_name_lower = model_name.lower() # Check base model names case-insensitively for base_model in model_configs: if base_model.lower() == model_name_lower: return base_model # Check aliases from the model configurations alias_map = ModelCapabilities.collect_aliases(model_configs) for base_model, aliases in alias_map.items(): if any(alias.lower() == model_name_lower for alias in aliases): return base_model # If not found, return as-is return model_name