From fc28783fada978f298b8d76910bcc9ea3d0245d1 Mon Sep 17 00:00:00 2001 From: zzz27578 <2950506809@qq.com> Date: Mon, 20 Jul 2026 03:12:13 +0800 Subject: [PATCH 1/2] fix: route OpenAI models to compatible API keys --- .../core/provider/sources/openai_source.py | 314 ++++++++++++++++-- astrbot/dashboard/services/config_service.py | 7 +- dashboard/src/api/v1.ts | 1 + .../provider/ProviderModelsPanel.vue | 18 + .../components/shared/ConfigItemRenderer.vue | 1 + .../src/components/shared/ListConfigItem.vue | 18 + .../src/composables/useProviderSources.ts | 15 +- tests/test_openai_source.py | 105 ++++++ 8 files changed, 440 insertions(+), 39 deletions(-) diff --git a/astrbot/core/provider/sources/openai_source.py b/astrbot/core/provider/sources/openai_source.py index b1594d608a..37b1319e0a 100644 --- a/astrbot/core/provider/sources/openai_source.py +++ b/astrbot/core/provider/sources/openai_source.py @@ -1,5 +1,6 @@ import asyncio import copy +import hashlib import inspect import json import random @@ -50,6 +51,7 @@ ) class ProviderOpenAIOfficial(Provider): _ERROR_TEXT_CANDIDATE_MAX_CHARS = 4096 + _model_key_cache: dict[str, dict[str, tuple[int, ...]]] = {} @classmethod def _truncate_error_text_candidate(cls, text: str) -> str: @@ -353,6 +355,194 @@ def _create_http_client(self, provider_config: dict) -> httpx.AsyncClient: pass return create_proxy_client("OpenAI", proxy, httpx_module=httpx_module) + def _create_sdk_client(self, api_key: str | None): + """Create an OpenAI SDK client for one configured key. + + Args: + api_key: API key assigned to the client. + + Returns: + An Azure OpenAI or OpenAI async client matching the provider config. + """ + if "api_version" in self.provider_config: + return AsyncAzureOpenAI( + api_key=api_key, + api_version=self.provider_config.get("api_version", None), + default_headers=self.custom_headers, + base_url=self.provider_config.get("api_base", ""), + timeout=self.timeout, + http_client=self._create_http_client(self.provider_config), + ) + return AsyncOpenAI( + api_key=api_key, + base_url=self.provider_config.get("api_base", None), + default_headers=self.custom_headers, + timeout=self.timeout, + http_client=self._create_http_client(self.provider_config), + ) + + def _model_key_cache_id(self) -> str | None: + """Build a secret-safe identity for cached model-to-key mappings. + + Returns: + A stable cache ID, or ``None`` before provider initialization. + """ + provider_config = getattr(self, "provider_config", None) + api_keys = getattr(self, "api_keys", None) + if not isinstance(provider_config, dict) or not isinstance(api_keys, list): + return None + identity = { + "api_base": provider_config.get("api_base"), + "api_version": provider_config.get("api_version"), + "custom_headers": getattr(self, "custom_headers", None) or {}, + "keys": [ + hashlib.sha256(str(key).encode("utf-8")).hexdigest() for key in api_keys + ], + } + payload = json.dumps(identity, ensure_ascii=True, sort_keys=True, default=str) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + def _store_model_key_indexes( + self, + model_key_indexes: dict[str, list[int]], + ) -> None: + """Store model ownership discovered from all configured keys. + + Args: + model_key_indexes: Mapping from model ID to zero-based key indexes. + """ + normalized = { + model: tuple(dict.fromkeys(indexes)) + for model, indexes in model_key_indexes.items() + } + self._last_model_key_indexes = normalized + cache_id = self._model_key_cache_id() + if cache_id: + self._model_key_cache[cache_id] = normalized + + def get_model_key_indexes(self) -> dict[str, list[int]]: + """Return the model ownership found by the latest model discovery. + + Returns: + A copy of the model-to-key-index mapping. + """ + mapping = getattr(self, "_last_model_key_indexes", {}) + return {model: list(indexes) for model, indexes in mapping.items()} + + def _candidate_api_keys_for_model(self, model: str) -> list[str]: + """Select keys known to expose a model, falling back to all keys. + + Args: + model: Model ID used by the request. + + Returns: + Deduplicated API keys eligible for the request. + """ + api_keys = list(dict.fromkeys(getattr(self, "api_keys", []) or [""])) + cache_id = self._model_key_cache_id() + if not cache_id: + return api_keys + indexes = self._model_key_cache.get(cache_id, {}).get(model) + if not indexes: + return api_keys + matched = [ + self.api_keys[index] for index in indexes if 0 <= index < len(self.api_keys) + ] + return list(dict.fromkeys(matched)) or api_keys + + def _remember_model_key(self, model: str, api_key: str) -> None: + """Remember that a request succeeded for a model and key. + + Args: + model: Model ID used by the request. + api_key: API key that completed the request. + """ + cache_id = self._model_key_cache_id() + if not cache_id or api_key not in self.api_keys: + return + mapping = dict(self._model_key_cache.get(cache_id, {})) + indexes = list(mapping.get(model, ())) + key_index = self.api_keys.index(api_key) + if key_index not in indexes: + indexes.append(key_index) + mapping[model] = tuple(indexes) + self._model_key_cache[cache_id] = mapping + + def _forget_model_key(self, model: str, api_key: str) -> None: + """Remove a model/key association after an access failure. + + Args: + model: Model ID rejected by the provider. + api_key: API key that could not access the model. + """ + cache_id = self._model_key_cache_id() + if not cache_id or api_key not in self.api_keys: + return + mapping = dict(self._model_key_cache.get(cache_id, {})) + indexes = list(mapping.get(model, ())) + key_index = self.api_keys.index(api_key) + if key_index not in indexes: + return + indexes.remove(key_index) + if indexes: + mapping[model] = tuple(indexes) + else: + mapping.pop(model, None) + self._model_key_cache[cache_id] = mapping + + def _key_label(self, api_key: str) -> str: + """Return a log-safe ordinal label for an API key. + + Args: + api_key: Configured API key to identify. + + Returns: + A label such as ``Key #2`` without exposing key material. + """ + try: + return f"Key #{self.api_keys.index(api_key) + 1}" + except (AttributeError, ValueError): + return "configured key" + + def _is_key_or_model_access_error(self, error: Exception) -> bool: + """Check whether retrying the model with another key is appropriate. + + Args: + error: Provider exception raised by a chat request. + + Returns: + Whether the error indicates key authentication or model access failure. + """ + status_codes = ( + getattr(error, "status_code", None), + getattr(error, "status", None), + getattr(getattr(error, "response", None), "status_code", None), + ) + if any(code in {401, 403, 404} for code in status_codes): + return True + error_text = " ".join(self._extract_error_text_candidates(error)).lower() + if "model" not in error_text and "模型" not in error_text: + return False + return any( + marker in error_text + for marker in ( + "does not exist", + "not found", + "not available", + "not accessible", + "no access", + "access denied", + "insufficient permission", + "permission", + "not allowed", + "unauthorized", + "不存在", + "不可用", + "无权", + "权限", + ) + ) + def __init__(self, provider_config, provider_settings) -> None: super().__init__(provider_config, provider_settings) self.chosen_api_key = None @@ -369,25 +559,7 @@ def __init__(self, provider_config, provider_settings) -> None: for key in self.custom_headers: self.custom_headers[key] = str(self.custom_headers[key]) - if "api_version" in provider_config: - # Using Azure OpenAI API - self.client = AsyncAzureOpenAI( - api_key=self.chosen_api_key, - api_version=provider_config.get("api_version", None), - default_headers=self.custom_headers, - base_url=provider_config.get("api_base", ""), - timeout=self.timeout, - http_client=self._create_http_client(provider_config), - ) - else: - # Using OpenAI Official API - self.client = AsyncOpenAI( - api_key=self.chosen_api_key, - base_url=provider_config.get("api_base", None), - default_headers=self.custom_headers, - timeout=self.timeout, - http_client=self._create_http_client(provider_config), - ) + self.client = self._create_sdk_client(self.chosen_api_key) self.default_params = inspect.signature( self.client.chat.completions.create, @@ -435,18 +607,63 @@ def _apply_provider_specific_request_overrides( extra_body["reasoning_effort"] = "none" async def get_models(self): - try: - models_str = [] - models = await retry_provider_request( - "OpenAI", - lambda: self.client.models.list(), - ) - models = sorted(models.data, key=lambda x: x.id) - for model in models: - models_str.append(model.id) - return models_str - except NotFoundError as e: - raise Exception(f"获取模型列表失败：{e}") + api_keys = getattr(self, "api_keys", None) + if not isinstance(api_keys, list) or len(api_keys) <= 1: + try: + models = await retry_provider_request( + "OpenAI", + lambda: self.client.models.list(), + ) + models_str = sorted(model.id for model in models.data) + if isinstance(api_keys, list): + self._store_model_key_indexes({model: [0] for model in models_str}) + return models_str + except NotFoundError as e: + raise Exception(f"获取模型列表失败：{e}") + + unique_keys = list(dict.fromkeys(api_keys)) + model_key_indexes: dict[str, list[int]] = {} + last_error: Exception | None = None + successful_requests = 0 + + for api_key in unique_keys: + key_index = api_keys.index(api_key) + client = None + try: + client = self._create_sdk_client(api_key) + models = await retry_provider_request( + f"OpenAI {self._key_label(api_key)}", + lambda client=client: client.models.list(), + ) + successful_requests += 1 + for model in models.data: + model_id = getattr(model, "id", None) + if not isinstance(model_id, str) or not model_id: + continue + model_key_indexes.setdefault(model_id, []).append(key_index) + except Exception as exc: + last_error = exc + logger.warning( + "Failed to fetch the OpenAI model list with %s: %s", + self._key_label(api_key), + type(exc).__name__, + ) + finally: + if client is not None: + try: + await client.close() + except Exception as exc: + logger.debug( + "Failed to close the temporary OpenAI client for %s: %s", + self._key_label(api_key), + type(exc).__name__, + ) + + if successful_requests == 0 and last_error is not None: + raise last_error + + self._store_model_key_indexes(model_key_indexes) + return sorted(model_key_indexes) @staticmethod def _sanitize_assistant_messages(payloads: dict) -> None: @@ -1167,6 +1384,29 @@ async def _handle_api_error( ) # logger.error(f"发生了错误。Provider 配置如下: {self.provider_config}") + if self._is_key_or_model_access_error(e) and len(self.api_keys) > 1: + model_id = str(payloads.get("model") or self.get_model()) + self._forget_model_key(model_id, chosen_key) + if chosen_key in available_api_keys: + available_api_keys.remove(chosen_key) + if available_api_keys: + next_key = random.choice(available_api_keys) + logger.warning( + "%s cannot access model %s; retrying with %s.", + self._key_label(chosen_key), + model_id, + self._key_label(next_key), + ) + return ( + False, + next_key, + available_api_keys, + payloads, + context_query, + func_tool, + image_fallback_used, + ) + if is_connection_error(e): proxy = self.provider_config.get("proxy", "") log_connection_failure("OpenAI", e, proxy) @@ -1204,8 +1444,9 @@ async def text_chat( payloads["tool_choice"] = tool_choice llm_response = None - max_retries = 10 - available_api_keys = self.api_keys.copy() + model_id = str(payloads.get("model") or self.get_model()) + available_api_keys = self._candidate_api_keys_for_model(model_id) + max_retries = max(10, len(available_api_keys) + 2) chosen_key = random.choice(available_api_keys) image_fallback_used = False @@ -1219,6 +1460,7 @@ async def text_chat( func_tool, request_max_retries=request_max_retries, ) + self._remember_model_key(model_id, chosen_key) break except Exception as e: last_exception = e @@ -1280,8 +1522,9 @@ async def text_chat_stream( if func_tool and not func_tool.empty(): payloads["tool_choice"] = tool_choice - max_retries = 10 - available_api_keys = self.api_keys.copy() + model_id = str(payloads.get("model") or self.get_model()) + available_api_keys = self._candidate_api_keys_for_model(model_id) + max_retries = max(10, len(available_api_keys) + 2) chosen_key = random.choice(available_api_keys) image_fallback_used = False @@ -1296,6 +1539,7 @@ async def text_chat_stream( request_max_retries=request_max_retries, ): yield response + self._remember_model_key(model_id, chosen_key) break except Exception as e: last_exception = e diff --git a/astrbot/dashboard/services/config_service.py b/astrbot/dashboard/services/config_service.py index a35dda8610..ed4ba8a159 100644 --- a/astrbot/dashboard/services/config_service.py +++ b/astrbot/dashboard/services/config_service.py @@ -1428,7 +1428,7 @@ async def list_provider_source_models(self, source_id: str) -> dict: try: models = await inst.get_models() models = models or [] - return { + result = { "models": models, "provider_source_id": source_id, "model_metadata": { @@ -1437,6 +1437,11 @@ async def list_provider_source_models(self, source_id: str) -> dict: if model_id in LLM_METADATAS }, } + key_mapping_getter = getattr(inst, "get_model_key_indexes", None) + if callable(key_mapping_getter): + model_key_indexes = key_mapping_getter() + result["model_key_indexes"] = model_key_indexes + return result finally: terminate_fn = getattr(inst, "terminate", None) if callable(terminate_fn): diff --git a/dashboard/src/api/v1.ts b/dashboard/src/api/v1.ts index 5df8c58561..6bb05e4245 100644 --- a/dashboard/src/api/v1.ts +++ b/dashboard/src/api/v1.ts @@ -98,6 +98,7 @@ export interface ProviderByIdData { export interface ProviderSourceModelsData { models?: string[]; model_metadata?: Record; + model_key_indexes?: Record; } export interface ProviderTestData { diff --git a/dashboard/src/components/provider/ProviderModelsPanel.vue b/dashboard/src/components/provider/ProviderModelsPanel.vue index 0420e93bd3..fe6bb3e20a 100644 --- a/dashboard/src/components/provider/ProviderModelsPanel.vue +++ b/dashboard/src/components/provider/ProviderModelsPanel.vue @@ -66,6 +66,15 @@
{{ entry.provider.id }}
{{ entry.provider.model }}
+
{{ entry.model }}
+ diff --git a/dashboard/src/components/shared/ListConfigItem.vue b/dashboard/src/components/shared/ListConfigItem.vue index 5ea6be7b3f..8f6b76957d 100644 --- a/dashboard/src/components/shared/ListConfigItem.vue +++ b/dashboard/src/components/shared/ListConfigItem.vue @@ -75,6 +75,9 @@ rounded="md" class="ma-1 list-item-clickable" @click="startEdit(index, item)"> + {{ item }} @@ -188,6 +191,10 @@ const props = defineProps({ preferSingleItem: { type: Boolean, default: true + }, + showItemIndex: { + type: Boolean, + default: false } }) @@ -202,6 +209,7 @@ const editItem = ref('') const showBatchImport = ref(false) const batchImportText = ref('') const isSingleItemMode = computed(() => (props.modelValue?.length ?? 0) <= 1 && props.preferSingleItem) +const shouldShowItemIndex = computed(() => props.showItemIndex && localItems.value.length > 1) const singleItemValue = computed({ get: () => props.modelValue?.[0] ?? '', set: (value) => { @@ -342,6 +350,16 @@ function cancelBatchImport() { user-select: none; } +.item-index-label { + width: 44px; + flex-shrink: 0; + margin-right: 10px; + color: rgba(var(--v-theme-on-surface), 0.56); + font-size: 12px; + font-weight: 500; + line-height: 1.4; +} + .v-chip { margin: 2px; } diff --git a/dashboard/src/composables/useProviderSources.ts b/dashboard/src/composables/useProviderSources.ts index f2ff86ed61..52f6ecae3e 100644 --- a/dashboard/src/composables/useProviderSources.ts +++ b/dashboard/src/composables/useProviderSources.ts @@ -56,6 +56,7 @@ export function useProviderSources(options: UseProviderSourcesOptions) { const editableProviderSource = ref(null) const availableModels = ref([]) const modelMetadata = ref>({}) + const modelKeyIndexes = ref>({}) const loadingModels = ref(false) const savingSource = ref(false) const savingProviderToggles = ref([]) @@ -157,7 +158,8 @@ export function useProviderSources(options: UseProviderSourcesOptions) { type: 'configured', provider, metadata: metadata || buildMetadataFromProvider(provider), - hasModelMetadata: Boolean(metadata) + hasModelMetadata: Boolean(metadata), + keyIndexes: modelKeyIndexes.value?.[provider.model] || [] } }) @@ -172,7 +174,8 @@ export function useProviderSources(options: UseProviderSourcesOptions) { type: 'available', model: name, metadata: typeof item === 'object' ? item?.metadata : getModelMetadata(name), - hasModelMetadata: Boolean(typeof item === 'object' ? item?.metadata : getModelMetadata(name)) + hasModelMetadata: Boolean(typeof item === 'object' ? item?.metadata : getModelMetadata(name)), + keyIndexes: modelKeyIndexes.value?.[name] || [] } }) @@ -380,6 +383,7 @@ export function useProviderSources(options: UseProviderSourcesOptions) { suppressSourceWatch = false }) availableModels.value = [] + modelKeyIndexes.value = {} modelMetadata.value = {} isSourceModified.value = false } @@ -439,6 +443,7 @@ export function useProviderSources(options: UseProviderSourcesOptions) { selectedProviderSourceOriginalId.value = null editableProviderSource.value = null availableModels.value = [] + modelKeyIndexes.value = {} modelMetadata.value = {} isSourceModified.value = false } @@ -467,6 +472,7 @@ export function useProviderSources(options: UseProviderSourcesOptions) { selectedProviderSourceOriginalId.value = newId editableProviderSource.value = JSON.parse(JSON.stringify(newSource)) availableModels.value = [] + modelKeyIndexes.value = {} modelMetadata.value = {} isSourceModified.value = true } @@ -538,7 +544,7 @@ export function useProviderSources(options: UseProviderSourcesOptions) { return false } finally { savingSource.value = false - loadConfig() + await loadConfig() } } @@ -558,7 +564,9 @@ export function useProviderSources(options: UseProviderSourcesOptions) { const response = await providerApi.sourceModels(sourceId) if (response.data.status === 'ok') { const metadataMap = (response.data.data.model_metadata || {}) as Record + const keyIndexMap = (response.data.data.model_key_indexes || {}) as Record modelMetadata.value = metadataMap + modelKeyIndexes.value = keyIndexMap availableModels.value = (response.data.data.models || []).map((model: string) => ({ name: model, metadata: metadataMap?.[model] || null @@ -571,6 +579,7 @@ export function useProviderSources(options: UseProviderSourcesOptions) { } } catch (error: any) { modelMetadata.value = {} + modelKeyIndexes.value = {} showMessage(error.response?.data?.message || error.message || tm('models.fetchError'), 'error') } finally { loadingModels.value = false diff --git a/tests/test_openai_source.py b/tests/test_openai_source.py index a45a232938..ca864125ee 100644 --- a/tests/test_openai_source.py +++ b/tests/test_openai_source.py @@ -148,6 +148,111 @@ async def list(self): assert models.calls == 2 +@pytest.mark.asyncio +async def test_get_models_merges_models_from_all_api_keys(): + class FakeClient: + def __init__(self, key: str): + self.key = key + self.closed = False + self.models = self + + async def list(self): + models_by_key = { + "key-a": ["model-a", "shared-model"], + "key-b": ["model-b", "shared-model"], + } + return SimpleNamespace( + data=[SimpleNamespace(id=model) for model in models_by_key[self.key]] + ) + + async def close(self): + self.closed = True + + provider = _make_provider({"key": ["key-a", "key-b"]}) + created_clients: list[FakeClient] = [] + + def create_client(key): + client = FakeClient(key) + created_clients.append(client) + return client + + provider._create_sdk_client = create_client + try: + assert await provider.get_models() == [ + "model-a", + "model-b", + "shared-model", + ] + assert provider.get_model_key_indexes() == { + "model-a": [0], + "shared-model": [0, 1], + "model-b": [1], + } + assert all(client.closed for client in created_clients) + finally: + await provider.terminate() + + +@pytest.mark.asyncio +async def test_get_models_keeps_successful_keys_when_one_key_fails(): + class FakeClient: + def __init__(self, key: str): + self.key = key + self.models = self + + async def list(self): + if self.key == "bad-key": + raise RuntimeError("access denied") + return SimpleNamespace(data=[SimpleNamespace(id="model-b")]) + + async def close(self): + return None + + provider = _make_provider({"key": ["bad-key", "good-key"]}) + provider._create_sdk_client = FakeClient + try: + assert await provider.get_models() == ["model-b"] + assert provider.get_model_key_indexes() == {"model-b": [1]} + finally: + await provider.terminate() + + +@pytest.mark.asyncio +async def test_text_chat_retries_with_another_key_for_model_access_error( + monkeypatch, +): + class ModelAccessError(Exception): + status_code = 404 + + provider = _make_provider({"key": ["key-a", "key-b"], "model": "model-b"}) + attempted_keys: list[str] = [] + + async def fake_prepare_chat_payload(*args, **kwargs): + return {"messages": [], "model": "model-b"}, [] + + async def fake_query(payloads, func_tool, *, request_max_retries=None): + attempted_keys.append(provider.client.api_key) + if provider.client.api_key == "key-a": + raise ModelAccessError("model not found") + return LLMResponse(role="assistant", completion_text="ok") + + choices = iter(["key-a", "key-b"]) + monkeypatch.setattr( + openai_source_module.random, + "choice", + lambda _keys: next(choices), + ) + provider._prepare_chat_payload = fake_prepare_chat_payload + provider._query = fake_query + try: + response = await provider.text_chat(prompt="hello") + assert response.completion_text == "ok" + assert attempted_keys == ["key-a", "key-b"] + assert provider._candidate_api_keys_for_model("model-b") == ["key-b"] + finally: + await provider.terminate() + + @pytest.mark.asyncio async def test_text_chat_passes_request_max_retries_to_query(): captured: dict[str, object] = {} From 7dee31d8828146c49ac919f60d8fa6d2cf1846ec Mon Sep 17 00:00:00 2001 From: zzz27578 <2950506809@qq.com> Date: Mon, 20 Jul 2026 03:30:07 +0800 Subject: [PATCH 2/2] refactor: simplify OpenAI model key tracking --- .../core/provider/sources/openai_source.py | 90 +++++++------------ tests/test_openai_source.py | 23 +++++ 2 files changed, 53 insertions(+), 60 deletions(-) diff --git a/astrbot/core/provider/sources/openai_source.py b/astrbot/core/provider/sources/openai_source.py index 37b1319e0a..ce5986d117 100644 --- a/astrbot/core/provider/sources/openai_source.py +++ b/astrbot/core/provider/sources/openai_source.py @@ -1,6 +1,5 @@ import asyncio import copy -import hashlib import inspect import json import random @@ -51,7 +50,6 @@ ) class ProviderOpenAIOfficial(Provider): _ERROR_TEXT_CANDIDATE_MAX_CHARS = 4096 - _model_key_cache: dict[str, dict[str, tuple[int, ...]]] = {} @classmethod def _truncate_error_text_candidate(cls, text: str) -> str: @@ -381,27 +379,6 @@ def _create_sdk_client(self, api_key: str | None): http_client=self._create_http_client(self.provider_config), ) - def _model_key_cache_id(self) -> str | None: - """Build a secret-safe identity for cached model-to-key mappings. - - Returns: - A stable cache ID, or ``None`` before provider initialization. - """ - provider_config = getattr(self, "provider_config", None) - api_keys = getattr(self, "api_keys", None) - if not isinstance(provider_config, dict) or not isinstance(api_keys, list): - return None - identity = { - "api_base": provider_config.get("api_base"), - "api_version": provider_config.get("api_version"), - "custom_headers": getattr(self, "custom_headers", None) or {}, - "keys": [ - hashlib.sha256(str(key).encode("utf-8")).hexdigest() for key in api_keys - ], - } - payload = json.dumps(identity, ensure_ascii=True, sort_keys=True, default=str) - return hashlib.sha256(payload.encode("utf-8")).hexdigest() - def _store_model_key_indexes( self, model_key_indexes: dict[str, list[int]], @@ -411,14 +388,9 @@ def _store_model_key_indexes( Args: model_key_indexes: Mapping from model ID to zero-based key indexes. """ - normalized = { - model: tuple(dict.fromkeys(indexes)) - for model, indexes in model_key_indexes.items() + self._model_key_indexes = { + model: set(indexes) for model, indexes in model_key_indexes.items() } - self._last_model_key_indexes = normalized - cache_id = self._model_key_cache_id() - if cache_id: - self._model_key_cache[cache_id] = normalized def get_model_key_indexes(self) -> dict[str, list[int]]: """Return the model ownership found by the latest model discovery. @@ -426,8 +398,8 @@ def get_model_key_indexes(self) -> dict[str, list[int]]: Returns: A copy of the model-to-key-index mapping. """ - mapping = getattr(self, "_last_model_key_indexes", {}) - return {model: list(indexes) for model, indexes in mapping.items()} + mapping = getattr(self, "_model_key_indexes", {}) + return {model: sorted(indexes) for model, indexes in mapping.items()} def _candidate_api_keys_for_model(self, model: str) -> list[str]: """Select keys known to expose a model, falling back to all keys. @@ -438,15 +410,15 @@ def _candidate_api_keys_for_model(self, model: str) -> list[str]: Returns: Deduplicated API keys eligible for the request. """ - api_keys = list(dict.fromkeys(getattr(self, "api_keys", []) or [""])) - cache_id = self._model_key_cache_id() - if not cache_id: - return api_keys - indexes = self._model_key_cache.get(cache_id, {}).get(model) + api_keys = list(dict.fromkeys(self.api_keys or [""])) + mapping = getattr(self, "_model_key_indexes", {}) + indexes = mapping.get(model) if not indexes: return api_keys matched = [ - self.api_keys[index] for index in indexes if 0 <= index < len(self.api_keys) + self.api_keys[index] + for index in sorted(indexes) + if 0 <= index < len(self.api_keys) ] return list(dict.fromkeys(matched)) or api_keys @@ -457,16 +429,13 @@ def _remember_model_key(self, model: str, api_key: str) -> None: model: Model ID used by the request. api_key: API key that completed the request. """ - cache_id = self._model_key_cache_id() - if not cache_id or api_key not in self.api_keys: + if api_key not in self.api_keys: return - mapping = dict(self._model_key_cache.get(cache_id, {})) - indexes = list(mapping.get(model, ())) - key_index = self.api_keys.index(api_key) - if key_index not in indexes: - indexes.append(key_index) - mapping[model] = tuple(indexes) - self._model_key_cache[cache_id] = mapping + mapping = getattr(self, "_model_key_indexes", None) + if mapping is None: + mapping = {} + self._model_key_indexes = mapping + mapping.setdefault(model, set()).add(self.api_keys.index(api_key)) def _forget_model_key(self, model: str, api_key: str) -> None: """Remove a model/key association after an access failure. @@ -475,20 +444,15 @@ def _forget_model_key(self, model: str, api_key: str) -> None: model: Model ID rejected by the provider. api_key: API key that could not access the model. """ - cache_id = self._model_key_cache_id() - if not cache_id or api_key not in self.api_keys: + if api_key not in self.api_keys: return - mapping = dict(self._model_key_cache.get(cache_id, {})) - indexes = list(mapping.get(model, ())) - key_index = self.api_keys.index(api_key) - if key_index not in indexes: + mapping = getattr(self, "_model_key_indexes", {}) + indexes = mapping.get(model) + if not indexes: return - indexes.remove(key_index) - if indexes: - mapping[model] = tuple(indexes) - else: + indexes.discard(self.api_keys.index(api_key)) + if not indexes: mapping.pop(model, None) - self._model_key_cache[cache_id] = mapping def _key_label(self, api_key: str) -> str: """Return a log-safe ordinal label for an API key. @@ -547,6 +511,7 @@ def __init__(self, provider_config, provider_settings) -> None: super().__init__(provider_config, provider_settings) self.chosen_api_key = None self.api_keys: list = super().get_keys() + self._model_key_indexes: dict[str, set[int]] = {} self.chosen_api_key = self.api_keys[0] if len(self.api_keys) > 0 else None self.timeout = provider_config.get("timeout", 120) self.custom_headers = provider_config.get("custom_headers", {}) @@ -614,7 +579,12 @@ async def get_models(self): "OpenAI", lambda: self.client.models.list(), ) - models_str = sorted(model.id for model in models.data) + models_str = sorted( + model_id + for model in models.data + if isinstance(model_id := getattr(model, "id", None), str) + and model_id + ) if isinstance(api_keys, list): self._store_model_key_indexes({model: [0] for model in models_str}) return models_str @@ -663,7 +633,7 @@ async def get_models(self): raise last_error self._store_model_key_indexes(model_key_indexes) - return sorted(model_key_indexes) + return sorted(model_key_indexes.keys()) @staticmethod def _sanitize_assistant_messages(payloads: dict) -> None: diff --git a/tests/test_openai_source.py b/tests/test_openai_source.py index ca864125ee..223a758fbe 100644 --- a/tests/test_openai_source.py +++ b/tests/test_openai_source.py @@ -217,6 +217,29 @@ async def close(self): await provider.terminate() +@pytest.mark.asyncio +async def test_get_models_raises_when_all_keys_fail(): + class FakeClient: + def __init__(self, key: str): + self.key = key + self.models = self + + async def list(self): + raise RuntimeError(f"access denied for {self.key}") + + async def close(self): + return None + + provider = _make_provider({"key": ["bad-key-a", "bad-key-b"]}) + provider._create_sdk_client = FakeClient + try: + with pytest.raises(RuntimeError, match="bad-key-b"): + await provider.get_models() + assert provider.get_model_key_indexes() == {} + finally: + await provider.terminate() + + @pytest.mark.asyncio async def test_text_chat_retries_with_another_key_for_model_access_error( monkeypatch,