| 1 | import yaml |
| 2 | from helpers import files, cache |
| 3 | from typing import List, Dict, Optional, TypedDict, Literal |
| 4 | |
| 5 | ModelType = Literal["chat", "embedding"] |
| 6 | |
| 7 | PROVIDER_MANAGER_CACHE_AREA = "model_providers(plugins)" |
| 8 | PROVIDER_MANAGER_CACHE_KEY = "manager" |
| 9 | |
| 10 | # Type alias for UI option items |
| 11 | class FieldOption(TypedDict): |
| 12 | value: str |
| 13 | label: str |
| 14 | |
| 15 | class ProviderManager: |
| 16 | _raw: Optional[Dict[str, List[Dict[str, str]]]] = None # full provider data |
| 17 | _options: Optional[Dict[str, List[FieldOption]]] = None # UI-friendly list |
| 18 | |
| 19 | @classmethod |
| 20 | def get_instance(cls): |
| 21 | instance = cache.get(PROVIDER_MANAGER_CACHE_AREA, PROVIDER_MANAGER_CACHE_KEY) |
| 22 | if instance is None: |
| 23 | instance = cls() |
| 24 | cache.add(PROVIDER_MANAGER_CACHE_AREA, PROVIDER_MANAGER_CACHE_KEY, instance) |
| 25 | return instance |
| 26 | |
| 27 | @classmethod |
| 28 | def reload(cls): |
| 29 | """Force reload of all provider configs (call after plugin changes).""" |
| 30 | cache.remove(PROVIDER_MANAGER_CACHE_AREA, PROVIDER_MANAGER_CACHE_KEY) |
| 31 | inst = cls.get_instance() |
| 32 | inst._load_providers() |
| 33 | |
| 34 | def __init__(self): |
| 35 | if self._raw is None or self._options is None: |
| 36 | self._load_providers() |
| 37 | |
| 38 | @staticmethod |
| 39 | def _load_yaml(path: str) -> dict: |
| 40 | try: |
| 41 | with open(path, "r", encoding="utf-8") as f: |
| 42 | return yaml.safe_load(f) or {} |
| 43 | except (FileNotFoundError, yaml.YAMLError): |
| 44 | return {} |
| 45 | |
| 46 | @staticmethod |
| 47 | def _normalise_yaml(raw_yaml: dict) -> Dict[str, Dict[str, Dict[str, str]]]: |
| 48 | """Normalise YAML into {type: {id: config}} mapping format.""" |
| 49 | result: Dict[str, Dict[str, Dict[str, str]]] = {} |
| 50 | for p_type, providers in (raw_yaml or {}).items(): |
| 51 | entries: Dict[str, Dict[str, str]] = {} |
| 52 | if isinstance(providers, dict): |
| 53 | for pid, cfg in providers.items(): |
| 54 | entries[pid] = cfg or {} |
| 55 | elif isinstance(providers, list): |
| 56 | for p in (providers or []): |
| 57 | pid = (p.get("id") or p.get("value") or "").lower() |
| 58 | if pid: |
| 59 | entries[pid] = {k: v for k, v in p.items() if k not in ("id", "value")} |
| 60 | result[p_type] = entries |
| 61 | return result |
| 62 | |
| 63 | def _load_providers(self): |
| 64 | """Loads provider configs from main YAML and enabled plugins, then merges.""" |
| 65 | # Load base config |
| 66 | base_path = files.get_abs_path("conf/model_providers.yaml") |
| 67 | merged = self._normalise_yaml(self._load_yaml(base_path)) |
| 68 | |
| 69 | # Merge plugin provider configs (enabled plugins only) |
| 70 | from helpers.plugins import get_enabled_plugin_paths |
| 71 | plugin_yamls = get_enabled_plugin_paths(None, "conf", "model_providers.yaml") |
| 72 | for plugin_yaml_path in plugin_yamls: |
| 73 | plugin_data = self._normalise_yaml(self._load_yaml(plugin_yaml_path)) |
| 74 | for p_type, providers in plugin_data.items(): |
| 75 | if p_type not in merged: |
| 76 | merged[p_type] = {} |
| 77 | # Overwrite matching keys, append new ones |
| 78 | merged[p_type].update(providers) |
| 79 | |
| 80 | # Convert merged {type: {id: config}} to normalised list format, |
| 81 | # sorted by name with "other" always last. |
| 82 | normalised: Dict[str, List[Dict[str, str]]] = {} |
| 83 | for p_type, providers in merged.items(): |
| 84 | items: List[Dict[str, str]] = [] |
| 85 | for pid, cfg in providers.items(): |
| 86 | entry = {"id": pid, **cfg} |
| 87 | items.append(entry) |
| 88 | items.sort(key=lambda p: ( |
| 89 | p.get("id") == "other", # False (0) first, True (1) last |
| 90 | (p.get("name") or p.get("id") or "").lower(), |
| 91 | )) |
| 92 | normalised[p_type] = items |
| 93 | |
| 94 | # Save raw |
| 95 | self._raw = normalised |
| 96 | |
| 97 | # Build UI-friendly option list (value / label) |
| 98 | self._options = {} |
| 99 | for p_type, providers in normalised.items(): |
| 100 | opts: List[FieldOption] = [] |
| 101 | for p in providers: |
| 102 | pid = (p.get("id") or p.get("value") or "").lower() |
| 103 | name = p.get("name") or p.get("label") or pid |
| 104 | if pid: |
| 105 | opts.append({"value": pid, "label": name}) |
| 106 | self._options[p_type] = opts |
| 107 | |
| 108 | def get_providers(self, provider_type: ModelType) -> List[FieldOption]: |
| 109 | """Returns a list of providers for a given type (e.g., 'chat', 'embedding').""" |
| 110 | return self._options.get(provider_type, []) if self._options else [] |
| 111 | |
| 112 | def get_raw_providers(self, provider_type: ModelType) -> List[Dict[str, str]]: |
| 113 | """Return raw provider dictionaries for advanced use-cases.""" |
| 114 | return self._raw.get(provider_type, []) if self._raw else [] |
| 115 | |
| 116 | def get_provider_config(self, provider_type: ModelType, provider_id: str) -> Optional[Dict[str, str]]: |
| 117 | """Return the metadata dict for a single provider id (case-insensitive).""" |
| 118 | provider_id_low = provider_id.lower() |
| 119 | for p in self.get_raw_providers(provider_type): |
| 120 | if (p.get("id") or p.get("value", "")).lower() == provider_id_low: |
| 121 | return p |
| 122 | return None |
| 123 | |
| 124 | |
| 125 | def get_providers(provider_type: ModelType) -> List[FieldOption]: |
| 126 | """Convenience function to get providers of a specific type.""" |
| 127 | return ProviderManager.get_instance().get_providers(provider_type) |
| 128 | |
| 129 | |
| 130 | def get_raw_providers(provider_type: ModelType) -> List[Dict[str, str]]: |
| 131 | """Return full metadata for providers of a given type.""" |
| 132 | return ProviderManager.get_instance().get_raw_providers(provider_type) |
| 133 | |
| 134 | |
| 135 | def get_provider_config(provider_type: ModelType, provider_id: str) -> Optional[Dict[str, str]]: |
| 136 | """Return metadata for a single provider (None if not found).""" |
| 137 | return ProviderManager.get_instance().get_provider_config(provider_type, provider_id) |
| 138 | |
| 139 | |
| 140 | def reload_providers(): |
| 141 | """Re-merge base + plugin provider configs. Call after plugin changes.""" |
| 142 | ProviderManager.reload() |