452 lines
19 KiB
Python
452 lines
19 KiB
Python
"""Model provider registry for managing available providers."""
|
||
|
||
import logging
|
||
import os
|
||
from typing import TYPE_CHECKING, Optional
|
||
|
||
from .base import ModelProvider, ProviderType
|
||
|
||
if TYPE_CHECKING:
|
||
from tools.models import ToolModelCategory
|
||
|
||
|
||
class ModelProviderRegistry:
|
||
"""Registry for managing model providers."""
|
||
|
||
_instance = None
|
||
|
||
def __new__(cls):
|
||
"""Singleton pattern for registry."""
|
||
if cls._instance is None:
|
||
logging.debug("REGISTRY: Creating new registry instance")
|
||
cls._instance = super().__new__(cls)
|
||
# Initialize instance dictionaries on first creation
|
||
cls._instance._providers = {}
|
||
cls._instance._initialized_providers = {}
|
||
logging.debug(f"REGISTRY: Created instance {cls._instance}")
|
||
else:
|
||
logging.debug(f"REGISTRY: Returning existing instance {cls._instance}")
|
||
return cls._instance
|
||
|
||
@classmethod
|
||
def register_provider(cls, provider_type: ProviderType, provider_class: type[ModelProvider]) -> None:
|
||
"""Register a new provider class.
|
||
|
||
Args:
|
||
provider_type: Type of the provider (e.g., ProviderType.GOOGLE)
|
||
provider_class: Class that implements ModelProvider interface
|
||
"""
|
||
instance = cls()
|
||
instance._providers[provider_type] = provider_class
|
||
|
||
@classmethod
|
||
def get_provider(cls, provider_type: ProviderType, force_new: bool = False) -> Optional[ModelProvider]:
|
||
"""Get an initialized provider instance.
|
||
|
||
Args:
|
||
provider_type: Type of provider to get
|
||
force_new: Force creation of new instance instead of using cached
|
||
|
||
Returns:
|
||
Initialized ModelProvider instance or None if not available
|
||
"""
|
||
instance = cls()
|
||
|
||
# Return cached instance if available and not forcing new
|
||
if not force_new and provider_type in instance._initialized_providers:
|
||
return instance._initialized_providers[provider_type]
|
||
|
||
# Check if provider class is registered
|
||
if provider_type not in instance._providers:
|
||
return None
|
||
|
||
# Get API key from environment
|
||
api_key = cls._get_api_key_for_provider(provider_type)
|
||
|
||
# Get provider class or factory function
|
||
provider_class = instance._providers[provider_type]
|
||
|
||
# For custom providers, handle special initialization requirements
|
||
if provider_type == ProviderType.CUSTOM:
|
||
# Check if it's a factory function (callable but not a class)
|
||
if callable(provider_class) and not isinstance(provider_class, type):
|
||
# Factory function - call it with api_key parameter
|
||
provider = provider_class(api_key=api_key)
|
||
else:
|
||
# Regular class - need to handle URL requirement
|
||
custom_url = os.getenv("CUSTOM_API_URL", "")
|
||
if not custom_url:
|
||
if api_key: # Key is set but URL is missing
|
||
logging.warning("CUSTOM_API_KEY set but CUSTOM_API_URL missing – skipping Custom provider")
|
||
return None
|
||
# Use empty string as API key for custom providers that don't need auth (e.g., Ollama)
|
||
# This allows the provider to be created even without CUSTOM_API_KEY being set
|
||
api_key = api_key or ""
|
||
# Initialize custom provider with both API key and base URL
|
||
provider = provider_class(api_key=api_key, base_url=custom_url)
|
||
else:
|
||
if not api_key:
|
||
return None
|
||
# Initialize non-custom provider with just API key
|
||
provider = provider_class(api_key=api_key)
|
||
|
||
# Cache the instance
|
||
instance._initialized_providers[provider_type] = provider
|
||
|
||
return provider
|
||
|
||
@classmethod
|
||
def get_provider_for_model(cls, model_name: str) -> Optional[ModelProvider]:
|
||
"""Get provider instance for a specific model name.
|
||
|
||
Provider priority order:
|
||
1. Native APIs (GOOGLE, OPENAI) - Most direct and efficient
|
||
2. CUSTOM - For local/private models with specific endpoints
|
||
3. OPENROUTER - Catch-all for cloud models via unified API
|
||
|
||
Args:
|
||
model_name: Name of the model (e.g., "gemini-2.5-flash", "o3-mini")
|
||
|
||
Returns:
|
||
ModelProvider instance that supports this model
|
||
"""
|
||
logging.debug(f"get_provider_for_model called with model_name='{model_name}'")
|
||
|
||
# Define explicit provider priority order
|
||
# Native APIs first, then custom endpoints, then catch-all providers
|
||
PROVIDER_PRIORITY_ORDER = [
|
||
ProviderType.GOOGLE, # Direct Gemini access
|
||
ProviderType.OPENAI, # Direct OpenAI access
|
||
ProviderType.XAI, # Direct X.AI GROK access
|
||
ProviderType.DIAL, # DIAL unified API access
|
||
ProviderType.CUSTOM, # Local/self-hosted models
|
||
ProviderType.OPENROUTER, # Catch-all for cloud models
|
||
]
|
||
|
||
# Check providers in priority order
|
||
instance = cls()
|
||
logging.debug(f"Registry instance: {instance}")
|
||
logging.debug(f"Available providers in registry: {list(instance._providers.keys())}")
|
||
|
||
for provider_type in PROVIDER_PRIORITY_ORDER:
|
||
if provider_type in instance._providers:
|
||
logging.debug(f"Found {provider_type} in registry")
|
||
# Get or create provider instance
|
||
provider = cls.get_provider(provider_type)
|
||
if provider and provider.validate_model_name(model_name):
|
||
logging.debug(f"{provider_type} validates model {model_name}")
|
||
return provider
|
||
else:
|
||
logging.debug(f"{provider_type} does not validate model {model_name}")
|
||
else:
|
||
logging.debug(f"{provider_type} not found in registry")
|
||
|
||
logging.debug(f"No provider found for model {model_name}")
|
||
return None
|
||
|
||
@classmethod
|
||
def get_available_providers(cls) -> list[ProviderType]:
|
||
"""Get list of registered provider types."""
|
||
instance = cls()
|
||
return list(instance._providers.keys())
|
||
|
||
@classmethod
|
||
def get_available_models(cls, respect_restrictions: bool = True) -> dict[str, ProviderType]:
|
||
"""Get mapping of all available models to their providers.
|
||
|
||
Args:
|
||
respect_restrictions: If True, filter out models not allowed by restrictions
|
||
|
||
Returns:
|
||
Dict mapping model names to provider types
|
||
"""
|
||
# Import here to avoid circular imports
|
||
from utils.model_restrictions import get_restriction_service
|
||
|
||
restriction_service = get_restriction_service() if respect_restrictions else None
|
||
models: dict[str, ProviderType] = {}
|
||
instance = cls()
|
||
|
||
for provider_type in instance._providers:
|
||
provider = cls.get_provider(provider_type)
|
||
if not provider:
|
||
continue
|
||
|
||
try:
|
||
available = provider.list_models(respect_restrictions=respect_restrictions)
|
||
except NotImplementedError:
|
||
logging.warning("Provider %s does not implement list_models", provider_type)
|
||
continue
|
||
|
||
for model_name in available:
|
||
# =====================================================================================
|
||
# CRITICAL: Prevent double restriction filtering (Fixed Issue #98)
|
||
# =====================================================================================
|
||
# Previously, both the provider AND registry applied restrictions, causing
|
||
# double-filtering that resulted in "no models available" errors.
|
||
#
|
||
# Logic: If respect_restrictions=True, provider already filtered models,
|
||
# so registry should NOT filter them again.
|
||
# TEST COVERAGE: tests/test_provider_routing_bugs.py::TestOpenRouterAliasRestrictions
|
||
# =====================================================================================
|
||
if (
|
||
restriction_service
|
||
and not respect_restrictions # Only filter if provider didn't already filter
|
||
and not restriction_service.is_allowed(provider_type, model_name)
|
||
):
|
||
logging.debug("Model %s filtered by restrictions", model_name)
|
||
continue
|
||
models[model_name] = provider_type
|
||
|
||
return models
|
||
|
||
@classmethod
|
||
def get_available_model_names(cls, provider_type: Optional[ProviderType] = None) -> list[str]:
|
||
"""Get list of available model names, optionally filtered by provider.
|
||
|
||
This respects model restrictions automatically.
|
||
|
||
Args:
|
||
provider_type: Optional provider to filter by
|
||
|
||
Returns:
|
||
List of available model names
|
||
"""
|
||
available_models = cls.get_available_models(respect_restrictions=True)
|
||
|
||
if provider_type:
|
||
# Filter by specific provider
|
||
return [name for name, ptype in available_models.items() if ptype == provider_type]
|
||
else:
|
||
# Return all available models
|
||
return list(available_models.keys())
|
||
|
||
@classmethod
|
||
def _get_api_key_for_provider(cls, provider_type: ProviderType) -> Optional[str]:
|
||
"""Get API key for a provider from environment variables.
|
||
|
||
Args:
|
||
provider_type: Provider type to get API key for
|
||
|
||
Returns:
|
||
API key string or None if not found
|
||
"""
|
||
key_mapping = {
|
||
ProviderType.GOOGLE: "GEMINI_API_KEY",
|
||
ProviderType.OPENAI: "OPENAI_API_KEY",
|
||
ProviderType.XAI: "XAI_API_KEY",
|
||
ProviderType.OPENROUTER: "OPENROUTER_API_KEY",
|
||
ProviderType.CUSTOM: "CUSTOM_API_KEY", # Can be empty for providers that don't need auth
|
||
ProviderType.DIAL: "DIAL_API_KEY",
|
||
}
|
||
|
||
env_var = key_mapping.get(provider_type)
|
||
if not env_var:
|
||
return None
|
||
|
||
return os.getenv(env_var)
|
||
|
||
@classmethod
|
||
def get_preferred_fallback_model(cls, tool_category: Optional["ToolModelCategory"] = None) -> str:
|
||
"""Get the preferred fallback model based on available API keys and tool category.
|
||
|
||
This method checks which providers have valid API keys and returns
|
||
a sensible default model for auto mode fallback situations.
|
||
|
||
Takes into account model restrictions when selecting fallback models.
|
||
|
||
Args:
|
||
tool_category: Optional category to influence model selection
|
||
|
||
Returns:
|
||
Model name string for fallback use
|
||
"""
|
||
# Import here to avoid circular import
|
||
from tools.models import ToolModelCategory
|
||
|
||
# Get available models respecting restrictions
|
||
available_models = cls.get_available_models(respect_restrictions=True)
|
||
|
||
# Group by provider
|
||
openai_models = [m for m, p in available_models.items() if p == ProviderType.OPENAI]
|
||
gemini_models = [m for m, p in available_models.items() if p == ProviderType.GOOGLE]
|
||
xai_models = [m for m, p in available_models.items() if p == ProviderType.XAI]
|
||
openrouter_models = [m for m, p in available_models.items() if p == ProviderType.OPENROUTER]
|
||
custom_models = [m for m, p in available_models.items() if p == ProviderType.CUSTOM]
|
||
|
||
openai_available = bool(openai_models)
|
||
gemini_available = bool(gemini_models)
|
||
xai_available = bool(xai_models)
|
||
openrouter_available = bool(openrouter_models)
|
||
custom_available = bool(custom_models)
|
||
|
||
if tool_category == ToolModelCategory.EXTENDED_REASONING:
|
||
# Prefer thinking-capable models for deep reasoning tools
|
||
if openai_available and "o3" in openai_models:
|
||
return "o3" # O3 for deep reasoning
|
||
elif openai_available and openai_models:
|
||
# Fall back to any available OpenAI model
|
||
return openai_models[0]
|
||
elif xai_available and "grok-3" in xai_models:
|
||
return "grok-3" # GROK-3 for deep reasoning
|
||
elif xai_available and xai_models:
|
||
# Fall back to any available XAI model
|
||
return xai_models[0]
|
||
elif gemini_available and any("pro" in m for m in gemini_models):
|
||
# Find the pro model (handles full names)
|
||
return next(m for m in gemini_models if "pro" in m)
|
||
elif gemini_available and gemini_models:
|
||
# Fall back to any available Gemini model
|
||
return gemini_models[0]
|
||
elif openrouter_available:
|
||
# Try to find thinking-capable model from openrouter
|
||
thinking_model = cls._find_extended_thinking_model()
|
||
if thinking_model:
|
||
return thinking_model
|
||
# Fallback to first available OpenRouter model
|
||
return openrouter_models[0]
|
||
elif custom_available:
|
||
# Fallback to custom models when available
|
||
return custom_models[0]
|
||
else:
|
||
# Fallback to pro if nothing found
|
||
return "gemini-2.5-pro"
|
||
|
||
elif tool_category == ToolModelCategory.FAST_RESPONSE:
|
||
# Prefer fast, cost-efficient models
|
||
if openai_available and "o4-mini" in openai_models:
|
||
return "o4-mini" # Latest, fast and efficient
|
||
elif openai_available and "o3-mini" in openai_models:
|
||
return "o3-mini" # Second choice
|
||
elif openai_available and openai_models:
|
||
# Fall back to any available OpenAI model
|
||
return openai_models[0]
|
||
elif xai_available and "grok-3-fast" in xai_models:
|
||
return "grok-3-fast" # GROK-3 Fast for speed
|
||
elif xai_available and xai_models:
|
||
# Fall back to any available XAI model
|
||
return xai_models[0]
|
||
elif gemini_available and any("flash" in m for m in gemini_models):
|
||
# Find the flash model (handles full names)
|
||
# Prefer 2.5 over 2.0 for backward compatibility
|
||
flash_models = [m for m in gemini_models if "flash" in m]
|
||
# Sort to ensure 2.5 comes before 2.0
|
||
flash_models_sorted = sorted(flash_models, reverse=True)
|
||
return flash_models_sorted[0]
|
||
elif gemini_available and gemini_models:
|
||
# Fall back to any available Gemini model
|
||
return gemini_models[0]
|
||
elif openrouter_available:
|
||
# Fallback to first available OpenRouter model
|
||
return openrouter_models[0]
|
||
elif custom_available:
|
||
# Fallback to custom models when available
|
||
return custom_models[0]
|
||
else:
|
||
# Default to flash
|
||
return "gemini-2.5-flash"
|
||
|
||
# BALANCED or no category specified - use existing balanced logic
|
||
if openai_available and "o4-mini" in openai_models:
|
||
return "o4-mini" # Latest balanced performance/cost
|
||
elif openai_available and "o3-mini" in openai_models:
|
||
return "o3-mini" # Second choice
|
||
elif openai_available and openai_models:
|
||
return openai_models[0]
|
||
elif xai_available and "grok-3" in xai_models:
|
||
return "grok-3" # GROK-3 as balanced choice
|
||
elif xai_available and xai_models:
|
||
return xai_models[0]
|
||
elif gemini_available and any("flash" in m for m in gemini_models):
|
||
# Prefer 2.5 over 2.0 for backward compatibility
|
||
flash_models = [m for m in gemini_models if "flash" in m]
|
||
flash_models_sorted = sorted(flash_models, reverse=True)
|
||
return flash_models_sorted[0]
|
||
elif gemini_available and gemini_models:
|
||
return gemini_models[0]
|
||
elif openrouter_available:
|
||
return openrouter_models[0]
|
||
elif custom_available:
|
||
# Fallback to custom models when available
|
||
return custom_models[0]
|
||
else:
|
||
# No models available due to restrictions - check if any providers exist
|
||
if not available_models:
|
||
# This might happen if all models are restricted
|
||
logging.warning("No models available due to restrictions")
|
||
# Return a reasonable default for backward compatibility
|
||
return "gemini-2.5-flash"
|
||
|
||
@classmethod
|
||
def _find_extended_thinking_model(cls) -> Optional[str]:
|
||
"""Find a model suitable for extended reasoning from custom/openrouter providers.
|
||
|
||
Returns:
|
||
Model name if found, None otherwise
|
||
"""
|
||
# Check custom provider first
|
||
custom_provider = cls.get_provider(ProviderType.CUSTOM)
|
||
if custom_provider:
|
||
# Check if it's a CustomModelProvider and has thinking models
|
||
try:
|
||
from providers.custom import CustomProvider
|
||
|
||
if isinstance(custom_provider, CustomProvider) and hasattr(custom_provider, "model_registry"):
|
||
for model_name, config in custom_provider.model_registry.items():
|
||
if config.get("supports_extended_thinking", False):
|
||
return model_name
|
||
except ImportError:
|
||
pass
|
||
|
||
# Then check OpenRouter for high-context/powerful models
|
||
openrouter_provider = cls.get_provider(ProviderType.OPENROUTER)
|
||
if openrouter_provider:
|
||
# Prefer models known for deep reasoning
|
||
preferred_models = [
|
||
"anthropic/claude-sonnet-4",
|
||
"anthropic/claude-opus-4",
|
||
"google/gemini-2.5-pro",
|
||
"google/gemini-pro-1.5",
|
||
"meta-llama/llama-3.1-70b-instruct",
|
||
"mistralai/mixtral-8x7b-instruct",
|
||
]
|
||
for model in preferred_models:
|
||
try:
|
||
if openrouter_provider.validate_model_name(model):
|
||
return model
|
||
except Exception as e:
|
||
# Log the error for debugging purposes but continue searching
|
||
import logging
|
||
|
||
logging.warning(f"Model validation for '{model}' on OpenRouter failed: {e}")
|
||
continue
|
||
|
||
return None
|
||
|
||
@classmethod
|
||
def get_available_providers_with_keys(cls) -> list[ProviderType]:
|
||
"""Get list of provider types that have valid API keys.
|
||
|
||
Returns:
|
||
List of ProviderType values for providers with valid API keys
|
||
"""
|
||
available = []
|
||
instance = cls()
|
||
for provider_type in instance._providers:
|
||
if cls.get_provider(provider_type) is not None:
|
||
available.append(provider_type)
|
||
return available
|
||
|
||
@classmethod
|
||
def clear_cache(cls) -> None:
|
||
"""Clear cached provider instances."""
|
||
instance = cls()
|
||
instance._initialized_providers.clear()
|
||
|
||
@classmethod
|
||
def unregister_provider(cls, provider_type: ProviderType) -> None:
|
||
"""Unregister a provider (mainly for testing)."""
|
||
instance = cls()
|
||
instance._providers.pop(provider_type, None)
|
||
instance._initialized_providers.pop(provider_type, None)
|