File size: 2,786 Bytes
adca48b 1438f9e adca48b 1438f9e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
"""
Model registry for PIPS - centralized model management.
This module provides a pluggable model registry that makes it easy to add
new models from different providers without modifying the core codebase.
"""
from typing import Dict, Any, Optional
# Internal registry storage
_registry: Dict[str, Dict[str, Any]] = {}
def register_model(name: str, provider: str, display: str = "", **config):
"""
Register a new model in the registry.
Args:
name: Unique model identifier
provider: Provider name (openai, google, anthropic)
display: Human-readable display name
**config: Additional configuration parameters
"""
_registry[name] = {
"provider": provider,
"display": display or name,
**config
}
def list_models() -> Dict[str, Dict[str, Any]]:
"""
Get all registered models.
Returns:
Dictionary mapping model names to their configuration
"""
return _registry.copy()
def get_model_config(name: str) -> Optional[Dict[str, Any]]:
"""
Get configuration for a specific model.
Args:
name: Model identifier
Returns:
Model configuration or None if not found
"""
return _registry.get(name)
def get_available_models() -> Dict[str, str]:
"""
Get available models in the format expected by the UI.
Returns:
Dictionary mapping model IDs to display names
"""
return {name: config["display"] for name, config in _registry.items()}
# Initialize with default models
def _initialize_default_models():
"""Initialize the registry with default models."""
# OpenAI Models
register_model("gpt-5", "openai", "OpenAI GPT-5")
register_model("gpt-5-mini", "openai", "OpenAI GPT-5-mini")
register_model("gpt-4.1-2025-04-14", "openai", "OpenAI GPT-4.1")
register_model("gpt-4o-2024-08-06", "openai", "OpenAI GPT-4o")
register_model("gpt-4.1-mini-2025-04-14", "openai", "OpenAI GPT-4.1 Mini")
register_model("gpt-4o-mini", "openai", "OpenAI GPT-4o Mini")
register_model("o4-mini-2025-04-16", "openai", "OpenAI o4 Mini")
register_model("o3-2025-04-16", "openai", "OpenAI o3")
# Google Models
register_model("gemini-2.0-flash", "google", "Google Gemini 2.0 Flash")
register_model("gemini-2.0-flash-codeinterpreter", "google", "Google Gemini 2.0 Flash (Code Interpreter)")
# Anthropic Models
register_model("claude-sonnet-4-20250514", "anthropic", "Anthropic Claude 4 Sonnet")
register_model("claude-opus-4-20250514", "anthropic", "Anthropic Claude 4 Opus")
register_model("claude-3-5-haiku-latest", "anthropic", "Anthropic Claude 3.5 Haiku")
# Initialize default models when module is imported
_initialize_default_models()
|