diff --git a/conf/model_providers.yaml b/conf/model_providers.yaml new file mode 100644 index 000000000..60ce6fb8c --- /dev/null +++ b/conf/model_providers.yaml @@ -0,0 +1,90 @@ +# Supported model providers for Agent Zero +# --------------------------------------- +# +# Each provider type ("chat", "embedding") contains a mapping of provider IDs +# to their configurations. +# +# The provider ID (e.g., "anthropic") is used: +# - in the settings UI dropdowns. +# - to construct the environment variable for the API key (e.g., ANTHROPIC_API_KEY). +# +# Each provider configuration requires: +# name: Human-readable name for the UI. +# litellm_provider: The corresponding provider name in LiteLLM. +# +# Optional fields: +# kwargs: A dictionary of extra parameters to pass to LiteLLM. +# This is useful for `api_base`, `extra_headers`, etc. + +chat: + anthropic: + name: Anthropic + litellm_provider: anthropic + deepseek: + name: DeepSeek + litellm_provider: deepseek + gemini: + name: Google + litellm_provider: gemini + groq: + name: Groq + litellm_provider: groq + huggingface: + name: HuggingFace + litellm_provider: huggingface + lm_studio: + name: LM Studio + litellm_provider: lm_studio + mistral: + name: Mistral AI + litellm_provider: mistral + ollama: + name: Ollama + litellm_provider: ollama + openai: + name: OpenAI + litellm_provider: openai + azure: + name: OpenAI Azure + litellm_provider: azure + openrouter: + name: OpenRouter + litellm_provider: openrouter + kwargs: + extra_headers: + "HTTP-Referer": "https://agent-zero.ai/" + "X-Title": "Agent Zero" + sambanova: + name: Sambanova + litellm_provider: sambanova + venice: + name: Venice + litellm_provider: openai + kwargs: + api_base: https://api.venice.ai/api/v1 + other: + name: Other OpenAI compatible + litellm_provider: openai + +embedding: + huggingface: + name: HuggingFace + litellm_provider: huggingface + lm_studio: + name: LM Studio + litellm_provider: lm_studio + mistral: + name: Mistral AI + litellm_provider: mistral + ollama: + name: Ollama + litellm_provider: ollama + openai: + name: OpenAI + litellm_provider: openai + azure: + name: OpenAI Azure + litellm_provider: azure + other: + name: Other OpenAI compatible + litellm_provider: openai \ No newline at end of file diff --git a/initialize.py b/initialize.py index 3ac974e88..111763ebe 100644 --- a/initialize.py +++ b/initialize.py @@ -1,5 +1,5 @@ -import models from agent import AgentConfig +import models from python.helpers import runtime, settings, defer from python.helpers.print_style import PrintStyle @@ -28,7 +28,7 @@ def initialize_agent(): # chat model from user settings chat_llm = models.ModelConfig( type=models.ModelType.CHAT, - provider=models.ModelProvider[current_settings["chat_model_provider"]], + provider=current_settings["chat_model_provider"], name=current_settings["chat_model_name"], api_base=current_settings["chat_model_api_base"], ctx_length=current_settings["chat_model_ctx_length"], @@ -42,7 +42,7 @@ def initialize_agent(): # utility model from user settings utility_llm = models.ModelConfig( type=models.ModelType.CHAT, - provider=models.ModelProvider[current_settings["util_model_provider"]], + provider=current_settings["util_model_provider"], name=current_settings["util_model_name"], api_base=current_settings["util_model_api_base"], ctx_length=current_settings["util_model_ctx_length"], @@ -54,7 +54,7 @@ def initialize_agent(): # embedding model from user settings embedding_llm = models.ModelConfig( type=models.ModelType.EMBEDDING, - provider=models.ModelProvider[current_settings["embed_model_provider"]], + provider=current_settings["embed_model_provider"], name=current_settings["embed_model_name"], api_base=current_settings["embed_model_api_base"], limit_requests=current_settings["embed_model_rl_requests"], @@ -63,7 +63,7 @@ def initialize_agent(): # browser model from user settings browser_llm = models.ModelConfig( type=models.ModelType.CHAT, - provider=models.ModelProvider[current_settings["browser_model_provider"]], + provider=current_settings["browser_model_provider"], name=current_settings["browser_model_name"], api_base=current_settings["browser_model_api_base"], vision=current_settings["browser_model_vision"], @@ -77,7 +77,7 @@ def initialize_agent(): browser_model=browser_llm, prompts_subdir=current_settings["agent_prompts_subdir"], memory_subdir=current_settings["agent_memory_subdir"], - knowledge_subdirs=["default", current_settings["agent_knowledge_subdir"]], + knowledge_subdirs=[current_settings["agent_knowledge_subdir"], "default"], mcp_servers=current_settings["mcp_servers"], code_exec_docker_enabled=False, # code_exec_docker_name = "A0-dev", diff --git a/models.py b/models.py index 1182f5251..de2f9ccb1 100644 --- a/models.py +++ b/models.py @@ -19,6 +19,7 @@ import litellm from python.helpers import dotenv from python.helpers.dotenv import load_dotenv +from python.helpers.providers import get_provider_config from python.helpers.rate_limiter import RateLimiter from python.helpers.tokens import approximate_tokens @@ -58,26 +59,10 @@ class ModelType(Enum): EMBEDDING = "Embedding" -class ModelProvider(Enum): - ANTHROPIC = "Anthropic" - DEEPSEEK = "DeepSeek" - GEMINI = "Google" - GROQ = "Groq" - HUGGINGFACE = "HuggingFace" - LM_STUDIO = "LM Studio" - MISTRAL = "Mistral AI" - OLLAMA = "Ollama" - OPENAI = "OpenAI" - AZURE = "OpenAI Azure" - OPENROUTER = "OpenRouter" - SAMBANOVA = "Sambanova" - OTHER = "Other OpenAI compatible" - - @dataclass class ModelConfig: type: ModelType - provider: ModelProvider + provider: str name: str api_base: str = "" ctx_length: int = 0 @@ -114,9 +99,9 @@ def get_api_key(service: str) -> str: def get_rate_limiter( - provider: ModelProvider, name: str, requests: int, input: int, output: int + provider: str, name: str, requests: int, input: int, output: int ) -> RateLimiter: - key = f"{provider.name}\\{name}" + key = f"{provider}\\{name}" rate_limiters[key] = limiter = rate_limiters.get(key, RateLimiter(seconds=60)) limiter.limits["requests"] = requests or 0 limiter.limits["input"] = input or 0 @@ -366,6 +351,9 @@ class LocalSentenceTransformerWrapper(Embeddings): """Local wrapper for sentence-transformers models to avoid HuggingFace API calls""" def __init__(self, provider: str, model: str, **kwargs: Any): + # Clean common user-input mistakes + model = model.strip().strip('"').strip("'") + # Remove the "sentence-transformers/" prefix if present if model.startswith("sentence-transformers/"): model = model[len("sentence-transformers/") :] @@ -467,8 +455,36 @@ def _adjust_call_args(provider_name: str, model_name: str, kwargs: dict): return provider_name, model_name, kwargs -def get_model(type: ModelType, provider: ModelProvider, name: str, **kwargs: Any): - provider_name = provider.name.lower() +def _merge_provider_defaults( + provider_type: str, original_provider: str, kwargs: dict +) -> tuple[str, dict]: + provider_name = original_provider # default: unchanged + cfg = get_provider_config(provider_type, original_provider) + if cfg: + provider_name = cfg.get("litellm_provider", original_provider).lower() + + # Extra arguments nested under `kwargs` for readability + extra_kwargs = cfg.get("kwargs") if isinstance(cfg, dict) else None # type: ignore[arg-type] + if isinstance(extra_kwargs, dict): + for k, v in extra_kwargs.items(): + kwargs.setdefault(k, v) + + # Copy any additional top-level fields except metadata keys + for k, v in cfg.items(): + if k not in ("id", "name", "value", "litellm_provider", "kwargs"): + kwargs.setdefault(k, v) + + # Inject API key based on the *original* provider id if still missing + if "api_key" not in kwargs: + key = get_api_key(original_provider) + if key and key not in ("None", "NA"): + kwargs["api_key"] = key + + return provider_name, kwargs + + +def get_model(type: ModelType, provider: str, name: str, **kwargs: Any): + provider_name = provider.lower() if type == ModelType.CHAT: return _get_litellm_chat(LiteLLMChatWrapper, name, provider_name, **kwargs) elif type == ModelType.EMBEDDING: @@ -478,26 +494,24 @@ def get_model(type: ModelType, provider: ModelProvider, name: str, **kwargs: Any def get_chat_model( - provider: ModelProvider, name: str, **kwargs: Any + provider: str, name: str, **kwargs: Any ) -> LiteLLMChatWrapper: - provider_name = provider.name.lower() - model = _get_litellm_chat(LiteLLMChatWrapper, name, provider_name, **kwargs) - return model + orig = provider.lower() + provider_name, kwargs = _merge_provider_defaults("chat", orig, kwargs) + return _get_litellm_chat(LiteLLMChatWrapper, name, provider_name, **kwargs) def get_browser_model( - provider: ModelProvider, name: str, **kwargs: Any + provider: str, name: str, **kwargs: Any ) -> BrowserCompatibleChatWrapper: - provider_name = provider.name.lower() - model = _get_litellm_chat( - BrowserCompatibleChatWrapper, name, provider_name, **kwargs - ) - return model + orig = provider.lower() + provider_name, kwargs = _merge_provider_defaults("chat", orig, kwargs) + return _get_litellm_chat(BrowserCompatibleChatWrapper, name, provider_name, **kwargs) def get_embedding_model( - provider: ModelProvider, name: str, **kwargs: Any + provider: str, name: str, **kwargs: Any ) -> LiteLLMEmbeddingWrapper | LocalSentenceTransformerWrapper: - provider_name = provider.name.lower() - model = _get_litellm_embedding(name, provider_name, **kwargs) - return model + orig = provider.lower() + provider_name, kwargs = _merge_provider_defaults("embedding", orig, kwargs) + return _get_litellm_embedding(name, provider_name, **kwargs) diff --git a/preload.py b/preload.py index 99625d51b..21c656612 100644 --- a/preload.py +++ b/preload.py @@ -18,11 +18,11 @@ async def preload(): # preload embedding model async def preload_embedding(): - if set["embed_model_provider"] == models.ModelProvider.HUGGINGFACE.name: + if set["embed_model_provider"].lower() == "huggingface": try: # Use the new LiteLLM-based model system emb_mod = models.get_embedding_model( - models.ModelProvider.HUGGINGFACE, set["embed_model_name"] + "huggingface", set["embed_model_name"] ) emb_txt = await emb_mod.aembed_query("test") return emb_txt diff --git a/python/helpers/memory.py b/python/helpers/memory.py index 23e2f0d86..68ad9454f 100644 --- a/python/helpers/memory.py +++ b/python/helpers/memory.py @@ -129,7 +129,7 @@ class Memory: **model_config.build_kwargs(), ) embeddings_model_id = files.safe_file_name( - model_config.provider.name + "_" + model_config.name + model_config.provider + "_" + model_config.name ) # here we setup the embeddings model with the chosen cache storage @@ -160,7 +160,7 @@ class Memory: if files.exists(emb_set_file): embedding_set = json.loads(files.read_file(emb_set_file)) if ( - embedding_set["model_provider"] == model_config.provider.name + embedding_set["model_provider"] == model_config.provider and embedding_set["model_name"] == model_config.name ): # model matches @@ -200,7 +200,7 @@ class Memory: meta_file_path, json.dumps( { - "model_provider": model_config.provider.name, + "model_provider": model_config.provider, "model_name": model_config.name, } ), diff --git a/python/helpers/providers.py b/python/helpers/providers.py new file mode 100644 index 000000000..cd139e88a --- /dev/null +++ b/python/helpers/providers.py @@ -0,0 +1,101 @@ +import yaml +from python.helpers import files +from typing import List, Dict, Optional, TypedDict + + +# Type alias for UI option items +class FieldOption(TypedDict): + value: str + label: str + +class ProviderManager: + _instance = None + _raw: Optional[Dict[str, List[Dict[str, str]]]] = None # full provider data + _options: Optional[Dict[str, List[FieldOption]]] = None # UI-friendly list + + @classmethod + def get_instance(cls): + if cls._instance is None: + cls._instance = cls() + return cls._instance + + def __init__(self): + if self._raw is None or self._options is None: + self._load_providers() + + def _load_providers(self): + """Loads provider configurations from the YAML file and normalises them.""" + try: + config_path = files.get_abs_path("conf/model_providers.yaml") + with open(config_path, "r", encoding="utf-8") as f: + raw_yaml = yaml.safe_load(f) or {} + except (FileNotFoundError, yaml.YAMLError): + raw_yaml = {} + + # ------------------------------------------------------------ + # Normalise the YAML so that internally we always work with a + # list-of-dicts [{id, name, ...}] for each provider type. This + # keeps existing callers unchanged while allowing the new nested + # mapping format in the YAML (id -> { ... }). + # ------------------------------------------------------------ + normalised: Dict[str, List[Dict[str, str]]] = {} + + for p_type, providers in (raw_yaml or {}).items(): + items: List[Dict[str, str]] = [] + + if isinstance(providers, dict): + # New format: mapping of id -> config + for pid, cfg in providers.items(): + entry = {"id": pid, **(cfg or {})} + items.append(entry) + elif isinstance(providers, list): + # Legacy list format – use as-is + items.extend(providers or []) + + normalised[p_type] = items + + # Save raw + self._raw = normalised + + # Build UI-friendly option list (value / label) + self._options = {} + for p_type, providers in normalised.items(): + opts: List[FieldOption] = [] + for p in providers: + pid = (p.get("id") or p.get("value") or "").lower() + name = p.get("name") or p.get("label") or pid + if pid: + opts.append({"value": pid, "label": name}) + self._options[p_type] = opts + + def get_providers(self, provider_type: str) -> List[FieldOption]: + """Returns a list of providers for a given type (e.g., 'chat', 'embedding').""" + return self._options.get(provider_type, []) if self._options else [] + + + def get_raw_providers(self, provider_type: str) -> List[Dict[str, str]]: + """Return raw provider dictionaries for advanced use-cases.""" + return self._raw.get(provider_type, []) if self._raw else [] + + def get_provider_config(self, provider_type: str, provider_id: str) -> Optional[Dict[str, str]]: + """Return the metadata dict for a single provider id (case-insensitive).""" + provider_id_low = provider_id.lower() + for p in self.get_raw_providers(provider_type): + if (p.get("id") or p.get("value", "")).lower() == provider_id_low: + return p + return None + + +def get_providers(provider_type: str) -> List[FieldOption]: + """Convenience function to get providers of a specific type.""" + return ProviderManager.get_instance().get_providers(provider_type) + + +def get_raw_providers(provider_type: str) -> List[Dict[str, str]]: + """Return full metadata for providers of a given type.""" + return ProviderManager.get_instance().get_raw_providers(provider_type) + + +def get_provider_config(provider_type: str, provider_id: str) -> Optional[Dict[str, str]]: + """Return metadata for a single provider (None if not found).""" + return ProviderManager.get_instance().get_provider_config(provider_type, provider_id) \ No newline at end of file diff --git a/python/helpers/settings.py b/python/helpers/settings.py index 522e7268d..b37c284c9 100644 --- a/python/helpers/settings.py +++ b/python/helpers/settings.py @@ -4,12 +4,13 @@ import json import os import re import subprocess -from typing import Any, Literal, TypedDict +from typing import Any, Literal, TypedDict, cast import models from python.helpers import runtime, whisper, defer, git from . import files, dotenv from python.helpers.print_style import PrintStyle +from python.helpers.providers import get_providers class Settings(TypedDict): @@ -122,8 +123,8 @@ SETTINGS_FILE = files.get_abs_path("tmp/settings.json") _settings: Settings | None = None + def convert_out(settings: Settings) -> SettingsOutput: - from models import ModelProvider default_settings = get_default_settings() # main model section @@ -135,7 +136,7 @@ def convert_out(settings: Settings) -> SettingsOutput: "description": "Select provider for main chat model used by Agent Zero", "type": "select", "value": settings["chat_model_provider"], - "options": [{"value": p.name, "label": p.value} for p in ModelProvider], + "options": cast(list[FieldOption], get_providers("chat")), } ) chat_model_fields.append( @@ -248,7 +249,7 @@ def convert_out(settings: Settings) -> SettingsOutput: "description": "Select provider for utility model used by the framework", "type": "select", "value": settings["util_model_provider"], - "options": [{"value": p.name, "label": p.value} for p in ModelProvider], + "options": cast(list[FieldOption], get_providers("chat")), } ) util_model_fields.append( @@ -328,7 +329,7 @@ def convert_out(settings: Settings) -> SettingsOutput: "description": "Select provider for embedding model used by the framework", "type": "select", "value": settings["embed_model_provider"], - "options": [{"value": p.name, "label": p.value} for p in ModelProvider], + "options": cast(list[FieldOption], get_providers("embedding")), } ) embed_model_fields.append( @@ -398,7 +399,7 @@ def convert_out(settings: Settings) -> SettingsOutput: "description": "Select provider for web browser model used by browser-use framework", "type": "select", "value": settings["browser_model_provider"], - "options": [{"value": p.name, "label": p.value} for p in ModelProvider], + "options": cast(list[FieldOption], get_providers("chat")), } ) browser_model_fields.append( @@ -499,8 +500,17 @@ def convert_out(settings: Settings) -> SettingsOutput: # api keys model section api_keys_fields: list[SettingsField] = [] - for provider in ModelProvider: - api_keys_fields.append(_get_api_key_field(settings, provider.name.lower(), provider.value)) + # Collect unique providers from both chat and embedding sections + providers_seen: set[str] = set() + for p_type in ("chat", "embedding"): + for provider in get_providers(p_type): + pid_lower = provider["value"].lower() + if pid_lower in providers_seen: + continue + providers_seen.add(pid_lower) + api_keys_fields.append( + _get_api_key_field(settings, pid_lower, provider["label"]) + ) api_keys_section: SettingsSection = { "id": "api_keys", @@ -990,11 +1000,9 @@ def _write_sensitive_settings(settings: Settings): def get_default_settings() -> Settings: - from models import ModelProvider - return Settings( version=_get_version(), - chat_model_provider=ModelProvider.OPENROUTER.name, + chat_model_provider="openrouter", chat_model_name="openai/gpt-4.1", chat_model_api_base="", chat_model_kwargs={"temperature": "0"}, @@ -1004,7 +1012,7 @@ def get_default_settings() -> Settings: chat_model_rl_requests=0, chat_model_rl_input=0, chat_model_rl_output=0, - util_model_provider=ModelProvider.OPENROUTER.name, + util_model_provider="openrouter", util_model_name="openai/gpt-4.1-nano", util_model_api_base="", util_model_ctx_length=100000, @@ -1013,13 +1021,13 @@ def get_default_settings() -> Settings: util_model_rl_requests=0, util_model_rl_input=0, util_model_rl_output=0, - embed_model_provider=ModelProvider.HUGGINGFACE.name, + embed_model_provider="huggingface", embed_model_name="sentence-transformers/all-MiniLM-L6-v2", embed_model_api_base="", embed_model_kwargs={}, embed_model_rl_requests=0, embed_model_rl_input=0, - browser_model_provider=ModelProvider.OPENROUTER.name, + browser_model_provider="openrouter", browser_model_name="openai/gpt-4.1", browser_model_api_base="", browser_model_vision=True,