diff --git a/astrbot/core/provider/sources/openai_source.py b/astrbot/core/provider/sources/openai_source.py index b1594d608a..ce5986d117 100644 --- a/astrbot/core/provider/sources/openai_source.py +++ b/astrbot/core/provider/sources/openai_source.py @@ -353,10 +353,165 @@ 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 _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. + """ + self._model_key_indexes = { + model: set(indexes) for model, indexes in model_key_indexes.items() + } + + 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, "_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. + + Args: + model: Model ID used by the request. + + Returns: + Deduplicated API keys eligible for the request. + """ + 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 sorted(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. + """ + if api_key not in self.api_keys: + return + 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. + + Args: + model: Model ID rejected by the provider. + api_key: API key that could not access the model. + """ + if api_key not in self.api_keys: + return + mapping = getattr(self, "_model_key_indexes", {}) + indexes = mapping.get(model) + if not indexes: + return + indexes.discard(self.api_keys.index(api_key)) + if not indexes: + mapping.pop(model, None) + + 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 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", {}) @@ -369,25 +524,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 +572,68 @@ 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(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 + 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.keys()) @staticmethod def _sanitize_assistant_messages(payloads: dict) -> None: @@ -1167,6 +1354,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 +1414,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 +1430,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 +1492,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 +1509,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..223a758fbe 100644 --- a/tests/test_openai_source.py +++ b/tests/test_openai_source.py @@ -148,6 +148,134 @@ 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_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, +): + 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] = {}