|
|
|
|
|
"""
|
|
|
Provider management for Theory2Manim Gradio app.
|
|
|
Allows user to select provider, enter API key, and select model.
|
|
|
"""
|
|
|
|
|
|
import os
|
|
|
from typing import Dict, List, Optional
|
|
|
|
|
|
class ProviderManager:
|
|
|
def __init__(self):
|
|
|
|
|
|
self.providers = {
|
|
|
'OpenAI': {
|
|
|
'api_key_env': 'OPENAI_API_KEY',
|
|
|
'models': [
|
|
|
'gpt-4.1', 'gpt-4o', 'gpt-3.5-turbo'
|
|
|
]
|
|
|
},
|
|
|
'Google Gemini': {
|
|
|
'api_key_env': 'GOOGLE_API_KEY',
|
|
|
'models': [
|
|
|
'gemini-1.5-pro-002', 'gemini-2.5-flash-preview-04-17'
|
|
|
]
|
|
|
},
|
|
|
'Anthropic': {
|
|
|
'api_key_env': 'ANTHROPIC_API_KEY',
|
|
|
'models': [
|
|
|
'claude-3-5-sonnet-20241022', 'claude-3-haiku'
|
|
|
]
|
|
|
},
|
|
|
'OpenRouter': {
|
|
|
'api_key_env': 'OPENROUTER_API_KEY',
|
|
|
'models': [
|
|
|
'openai/gpt-4o', 'openai/gpt-4o-mini', 'anthropic/claude-3.5-sonnet',
|
|
|
'anthropic/claude-3-haiku', 'google/gemini-pro-1.5', 'deepseek/deepseek-chat',
|
|
|
'qwen/qwen-2.5-72b-instruct', 'meta-llama/llama-3.1-8b-instruct:free',
|
|
|
'microsoft/phi-3-mini-128k-instruct:free'
|
|
|
]
|
|
|
}
|
|
|
}
|
|
|
self.selected_provider = None
|
|
|
self.api_keys = {}
|
|
|
|
|
|
def get_providers(self) -> List[str]:
|
|
|
return list(self.providers.keys())
|
|
|
|
|
|
def get_models(self, provider: str) -> List[str]:
|
|
|
return self.providers.get(provider, {}).get('models', [])
|
|
|
|
|
|
def set_api_key(self, provider: str, api_key: str):
|
|
|
env_var = self.providers[provider]['api_key_env']
|
|
|
os.environ[env_var] = api_key
|
|
|
self.api_keys[provider] = api_key
|
|
|
|
|
|
def get_api_key(self, provider: str) -> Optional[str]:
|
|
|
env_var = self.providers[provider]['api_key_env']
|
|
|
return os.environ.get(env_var)
|
|
|
|
|
|
def get_selected_provider(self) -> Optional[str]:
|
|
|
return self.selected_provider
|
|
|
|
|
|
def set_selected_provider(self, provider: str):
|
|
|
self.selected_provider = provider
|
|
|
|
|
|
def get_selected_model(self) -> Optional[str]:
|
|
|
if self.selected_provider:
|
|
|
return self.get_models(self.selected_provider)[0]
|
|
|
return None
|
|
|
|
|
|
provider_manager = ProviderManager()
|
|
|
|
