refactor: move LLM configs to llm_providers.py

Author: marginal23326

src/utils/llm_providers.py

@@ -0,0 +1,84 @@
+import os
+
+PROVIDER_DISPLAY_NAMES = {
+    "openai": "OpenAI",
+    "azure_openai": "Azure OpenAI",
+    "anthropic": "Anthropic",
+    "deepseek": "DeepSeek",
+    "google": "Google",
+    "mistral": "Mistral",
+    "alibaba": "Alibaba",
+    "moonshot": "MoonShot"
+}
+
+PROVIDER_CONFIGS = {
+    "openai": {
+        "api_key_env": "OPENAI_API_KEY", "base_url_env": "OPENAI_ENDPOINT",
+        "default_base_url": "https://api.openai.com/v1", "default_model": "gpt-4o"
+    },
+    "azure_openai": {
+        "api_key_env": "AZURE_OPENAI_API_KEY", "base_url_env": "AZURE_OPENAI_ENDPOINT",
+        "api_version_env": "AZURE_OPENAI_API_VERSION",
+        "default_api_version": "2025-01-01-preview", "default_model": "gpt-4o"
+    },
+    "anthropic": {
+        "api_key_env": "ANTHROPIC_API_KEY", "base_url_env": "ANTHROPIC_ENDPOINT",
+        "default_base_url": "https://api.anthropic.com", "default_model": "claude-3-5-sonnet-20241022"
+    },
+    "google": {
+        "api_key_env": "GOOGLE_API_KEY",
+        "default_model": "gemini-2.0-flash-exp"
+    },
+    "deepseek": {
+        "api_key_env": "DEEPSEEK_API_KEY", "base_url_env": "DEEPSEEK_ENDPOINT",
+        "default_base_url": "https://api.deepseek.com", "default_model": "deepseek-chat"
+    },
+    "mistral": {
+        "api_key_env": "MISTRAL_API_KEY", "base_url_env": "MISTRAL_ENDPOINT",
+        "default_base_url": "https://api.mistral.ai/v1", "default_model": "mistral-large-latest"
+    },
+    "alibaba": {
+        "api_key_env": "ALIBABA_API_KEY", "base_url_env": "ALIBABA_ENDPOINT",
+        "default_base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1", "default_model": "qwen-plus"
+    },
+    "moonshot": {
+        "api_key_env": "MOONSHOT_API_KEY", "base_url_env": "MOONSHOT_ENDPOINT",
+        "default_base_url": "https://api.moonshot.cn/v1", "default_model": "moonshot-v1-32k-vision-preview"
+    },
+    "ollama": {
+        "base_url_env": "OLLAMA_ENDPOINT",
+        "default_base_url": "http://localhost:11434", "default_model": "qwen2.5:7b"
+    }
+}
+
+# Predefined model names for common providers
+MODEL_NAMES = {
+    "anthropic": ["claude-3-5-sonnet-20241022", "claude-3-5-sonnet-20240620", "claude-3-opus-20240229"],
+    "openai": ["gpt-4o", "gpt-4", "gpt-3.5-turbo", "o3-mini"],
+    "deepseek": ["deepseek-chat", "deepseek-reasoner"],
+    "google": ["gemini-2.0-flash", "gemini-2.0-flash-thinking-exp", "gemini-1.5-flash-latest",
+               "gemini-1.5-flash-8b-latest", "gemini-2.0-flash-thinking-exp-01-21", "gemini-2.0-pro-exp-02-05"],
+    "ollama": ["qwen2.5:7b", "qwen2.5:14b", "qwen2.5:32b", "qwen2.5-coder:14b", "qwen2.5-coder:32b", "llama2:7b",
+               "deepseek-r1:14b", "deepseek-r1:32b"],
+    "azure_openai": ["gpt-4o", "gpt-4", "gpt-3.5-turbo"],
+    "mistral": ["pixtral-large-latest", "mistral-large-latest", "mistral-small-latest", "ministral-8b-latest"],
+    "alibaba": ["qwen-plus", "qwen-max", "qwen-turbo", "qwen-long"],
+    "moonshot": ["moonshot-v1-32k-vision-preview", "moonshot-v1-8k-vision-preview"],
+}
+
+def get_provider_config(provider: str):
+    return PROVIDER_CONFIGS.get(provider, {})
+
+def get_config_value(provider: str, key: str, **kwargs):
+    config = get_provider_config(provider)
+    
+    if key in kwargs and kwargs[key]:
+        return kwargs[key]
+    
+    env_key = config.get(f"{key}_env")
+    if env_key:
+        env_value = os.getenv(env_key)
+        if env_value:
+            return env_value
+    
+    return config.get(f"default_{key}")

src/utils/utils.py

@@ -14,15 +14,10 @@
 
 from .llm import DeepSeekR1ChatOpenAI, DeepSeekR1ChatOllama
 
-PROVIDER_DISPLAY_NAMES = {
-    "openai": "OpenAI",
-    "azure_openai": "Azure OpenAI",
-    "anthropic": "Anthropic",
-    "deepseek": "DeepSeek",
-    "google": "Google",
-    "alibaba": "Alibaba",
-    "moonshot": "MoonShot"
-}
+from .llm_providers import (
+    get_provider_config, get_config_value,
+    PROVIDER_DISPLAY_NAMES, MODEL_NAMES
+)
 
 
 def get_llm_model(provider: str, **kwargs):
@@ -32,176 +27,123 @@ def get_llm_model(provider: str, **kwargs):
     :param kwargs:
     :return:
     """
-    if provider not in ["ollama"]:
-        env_var = f"{provider.upper()}_API_KEY"
-        api_key = kwargs.get("api_key", "") or os.getenv(env_var, "")
+    if provider not in {"ollama"}:
+        api_key = get_config_value(provider, "api_key", **kwargs)
         if not api_key:
-            handle_api_key_error(provider, env_var)
-        kwargs["api_key"] = api_key
+            handle_api_key_error(provider)
 
-    if provider == "anthropic":
-        if not kwargs.get("base_url", ""):
-            base_url = "https://api.anthropic.com"
-        else:
-            base_url = kwargs.get("base_url")
+    base_url = get_config_value(provider, "base_url", **kwargs)
+    model_name = get_config_value(provider, "model", **kwargs)
+    temperature = kwargs.get("temperature", 0.0)
 
+    if provider == "anthropic":
         return ChatAnthropic(
-            model=kwargs.get("model_name", "claude-3-5-sonnet-20241022"),
-            temperature=kwargs.get("temperature", 0.0),
+            model=model_name,
+            temperature=temperature,
             base_url=base_url,
             api_key=api_key,
         )
-    elif provider == 'mistral':
-        if not kwargs.get("base_url", ""):
-            base_url = os.getenv("MISTRAL_ENDPOINT", "https://api.mistral.ai/v1")
-        else:
-            base_url = kwargs.get("base_url")
-        if not kwargs.get("api_key", ""):
-            api_key = os.getenv("MISTRAL_API_KEY", "")
-        else:
-            api_key = kwargs.get("api_key")
-
+    elif provider == "mistral":
         return ChatMistralAI(
-            model=kwargs.get("model_name", "mistral-large-latest"),
-            temperature=kwargs.get("temperature", 0.0),
+            model=model_name,
+            temperature=temperature,
             base_url=base_url,
             api_key=api_key,
         )
     elif provider == "openai":
-        if not kwargs.get("base_url", ""):
-            base_url = os.getenv("OPENAI_ENDPOINT", "https://api.openai.com/v1")
-        else:
-            base_url = kwargs.get("base_url")
-
         return ChatOpenAI(
-            model=kwargs.get("model_name", "gpt-4o"),
-            temperature=kwargs.get("temperature", 0.0),
+            model=model_name,
+            temperature=temperature,
             base_url=base_url,
             api_key=api_key,
         )
     elif provider == "deepseek":
-        if not kwargs.get("base_url", ""):
-            base_url = os.getenv("DEEPSEEK_ENDPOINT", "")
-        else:
-            base_url = kwargs.get("base_url")
-
-        if kwargs.get("model_name", "deepseek-chat") == "deepseek-reasoner":
+        if model_name == "deepseek-reasoner":
             return DeepSeekR1ChatOpenAI(
-                model=kwargs.get("model_name", "deepseek-reasoner"),
-                temperature=kwargs.get("temperature", 0.0),
+                model=model_name,
+                temperature=temperature,
                 base_url=base_url,
                 api_key=api_key,
             )
         else:
             return ChatOpenAI(
-                model=kwargs.get("model_name", "deepseek-chat"),
-                temperature=kwargs.get("temperature", 0.0),
+                model=model_name,
+                temperature=temperature,
                 base_url=base_url,
                 api_key=api_key,
             )
     elif provider == "google":
         return ChatGoogleGenerativeAI(
-            model=kwargs.get("model_name", "gemini-2.0-flash-exp"),
-            temperature=kwargs.get("temperature", 0.0),
+            model=model_name,
+            temperature=temperature,
             api_key=api_key,
         )
     elif provider == "ollama":
-        if not kwargs.get("base_url", ""):
-            base_url = os.getenv("OLLAMA_ENDPOINT", "http://localhost:11434")
-        else:
-            base_url = kwargs.get("base_url")
-
-        if "deepseek-r1" in kwargs.get("model_name", "qwen2.5:7b"):
+        num_ctx = kwargs.get("num_ctx", 32000)
+        if "deepseek-r1" in model_name:
             return DeepSeekR1ChatOllama(
                 model=kwargs.get("model_name", "deepseek-r1:14b"),
-                temperature=kwargs.get("temperature", 0.0),
-                num_ctx=kwargs.get("num_ctx", 32000),
+                temperature=temperature,
+                num_ctx=num_ctx,
                 base_url=base_url,
             )
         else:
             return ChatOllama(
-                model=kwargs.get("model_name", "qwen2.5:7b"),
-                temperature=kwargs.get("temperature", 0.0),
-                num_ctx=kwargs.get("num_ctx", 32000),
+                model=model_name,
+                temperature=temperature,
+                num_ctx=num_ctx,
                 num_predict=kwargs.get("num_predict", 1024),
                 base_url=base_url,
             )
     elif provider == "azure_openai":
-        if not kwargs.get("base_url", ""):
-            base_url = os.getenv("AZURE_OPENAI_ENDPOINT", "")
-        else:
-            base_url = kwargs.get("base_url")
-        api_version = kwargs.get("api_version", "") or os.getenv("AZURE_OPENAI_API_VERSION", "2025-01-01-preview")
+        api_version = get_config_value(provider, "api_version", **kwargs)
         return AzureChatOpenAI(
-            model=kwargs.get("model_name", "gpt-4o"),
-            temperature=kwargs.get("temperature", 0.0),
+            model=model_name,
+            temperature=temperature,
             api_version=api_version,
             azure_endpoint=base_url,
             api_key=api_key,
         )
     elif provider == "alibaba":
-        if not kwargs.get("base_url", ""):
-            base_url = os.getenv("ALIBABA_ENDPOINT", "https://dashscope.aliyuncs.com/compatible-mode/v1")
-        else:
-            base_url = kwargs.get("base_url")
-
         return ChatOpenAI(
-            model=kwargs.get("model_name", "qwen-plus"),
-            temperature=kwargs.get("temperature", 0.0),
+            model=model_name,
+            temperature=temperature,
             base_url=base_url,
             api_key=api_key,
         )
 
     elif provider == "moonshot":
         return ChatOpenAI(
-            model=kwargs.get("model_name", "moonshot-v1-32k-vision-preview"),
-            temperature=kwargs.get("temperature", 0.0),
-            base_url=os.getenv("MOONSHOT_ENDPOINT"),
-            api_key=os.getenv("MOONSHOT_API_KEY"),
+            model=model_name,
+            temperature=temperature,
+            base_url=base_url,
+            api_key=api_key,
         )
     else:
         raise ValueError(f"Unsupported provider: {provider}")
 
 
-# Predefined model names for common providers
-model_names = {
-    "anthropic": ["claude-3-5-sonnet-20241022", "claude-3-5-sonnet-20240620", "claude-3-opus-20240229"],
-    "openai": ["gpt-4o", "gpt-4", "gpt-3.5-turbo", "o3-mini"],
-    "deepseek": ["deepseek-chat", "deepseek-reasoner"],
-    "google": ["gemini-2.0-flash", "gemini-2.0-flash-thinking-exp", "gemini-1.5-flash-latest",
-               "gemini-1.5-flash-8b-latest", "gemini-2.0-flash-thinking-exp-01-21", "gemini-2.0-pro-exp-02-05"],
-    "ollama": ["qwen2.5:7b", "qwen2.5:14b", "qwen2.5:32b", "qwen2.5-coder:14b", "qwen2.5-coder:32b", "llama2:7b",
-               "deepseek-r1:14b", "deepseek-r1:32b"],
-    "azure_openai": ["gpt-4o", "gpt-4", "gpt-3.5-turbo"],
-    "mistral": ["mixtral-large-latest", "mistral-large-latest", "mistral-small-latest", "ministral-8b-latest"],
-    "alibaba": ["qwen-plus", "qwen-max", "qwen-turbo", "qwen-long"],
-    "moonshot": ["moonshot-v1-32k-vision-preview", "moonshot-v1-8k-vision-preview"],
-}
-
-
-# Callback to update the model name dropdown based on the selected provider
 def update_model_dropdown(llm_provider, api_key=None, base_url=None):
     """
     Update the model name dropdown with predefined models for the selected provider.
     """
     # Use API keys from .env if not provided
     if not api_key:
-        api_key = os.getenv(f"{llm_provider.upper()}_API_KEY", "")
+        api_key = get_config_value(llm_provider, "api_key")
     if not base_url:
-        base_url = os.getenv(f"{llm_provider.upper()}_BASE_URL", "")
+        base_url = get_config_value(llm_provider, "base_url")
 
     # Use predefined models for the selected provider
-    if llm_provider in model_names:
-        return gr.Dropdown(choices=model_names[llm_provider], value=model_names[llm_provider][0], interactive=True)
-    else:
-        return gr.Dropdown(choices=[], value="", interactive=True, allow_custom_value=True)
-
+    if llm_provider in MODEL_NAMES:
+        return gr.Dropdown(choices=MODEL_NAMES[llm_provider], value=MODEL_NAMES[llm_provider][0], interactive=True)
+    return gr.Dropdown(choices=[], value="", interactive=True, allow_custom_value=True)
 
-def handle_api_key_error(provider: str, env_var: str):
+def handle_api_key_error(provider: str):
     """
     Handles the missing API key error by raising a gr.Error with a clear message.
     """
     provider_display = PROVIDER_DISPLAY_NAMES.get(provider, provider.upper())
+    env_var = get_provider_config(provider).get("api_key_env")
     raise gr.Error(
         f"💥 {provider_display} API key not found! 🔑 Please set the "
         f"`{env_var}` environment variable or provide it in the UI."

tests/test_llm_api.py

@@ -32,21 +32,6 @@ def create_message_content(text, image_path=None):
         })
     return content
 
-def get_env_value(key, provider):
-    env_mappings = {
-        "openai": {"api_key": "OPENAI_API_KEY", "base_url": "OPENAI_ENDPOINT"},
-        "azure_openai": {"api_key": "AZURE_OPENAI_API_KEY", "base_url": "AZURE_OPENAI_ENDPOINT"},
-        "google": {"api_key": "GOOGLE_API_KEY"},
-        "deepseek": {"api_key": "DEEPSEEK_API_KEY", "base_url": "DEEPSEEK_ENDPOINT"},
-        "mistral": {"api_key": "MISTRAL_API_KEY", "base_url": "MISTRAL_ENDPOINT"},
-        "alibaba": {"api_key": "ALIBABA_API_KEY", "base_url": "ALIBABA_ENDPOINT"},
-        "moonshot":{"api_key": "MOONSHOT_API_KEY", "base_url": "MOONSHOT_ENDPOINT"},
-    }
-
-    if provider in env_mappings and key in env_mappings[provider]:
-        return os.getenv(env_mappings[provider][key], "")
-    return ""
-
 def test_llm(config, query, image_path=None, system_message=None):
     from src.utils import utils
 
@@ -69,8 +54,8 @@ def test_llm(config, query, image_path=None, system_message=None):
         provider=config.provider,
         model_name=config.model_name,
         temperature=config.temperature,
-        base_url=config.base_url or get_env_value("base_url", config.provider),
-        api_key=config.api_key or get_env_value("api_key", config.provider)
+        base_url=config.base_url,
+        api_key=config.api_key
     )
 
     # Prepare messages for non-Ollama models
@@ -130,7 +115,7 @@ def test_moonshot_model():
     # test_openai_model()
     # test_google_model()
     # test_azure_openai_model()
-    #test_deepseek_model()
+    # test_deepseek_model()
     # test_ollama_model()
     test_deepseek_r1_model()
     # test_deepseek_r1_ollama_model()

webui.py

@@ -24,6 +24,7 @@
 from langchain_ollama import ChatOllama
 from playwright.async_api import async_playwright
 from src.utils.agent_state import AgentState
+from src.utils.llm_providers import MODEL_NAMES
 
 from src.utils import utils
 from src.agent.custom_agent import CustomAgent
@@ -796,14 +797,14 @@ def create_ui(config, theme_name="Ocean"):
             with gr.TabItem("🔧 LLM Settings", id=2):
                 with gr.Group():
                     llm_provider = gr.Dropdown(
-                        choices=[provider for provider, model in utils.model_names.items()],
+                        choices=[provider for provider,model in MODEL_NAMES.items()],
                         label="LLM Provider",
                         value=config['llm_provider'],
                         info="Select your preferred language model provider"
                     )
                     llm_model_name = gr.Dropdown(
                         label="Model Name",
-                        choices=utils.model_names['openai'],
+                        choices=MODEL_NAMES['openai'],
                         value=config['llm_model_name'],
                         interactive=True,
                         allow_custom_value=True,  # Allow users to input custom model names