fix: route OpenAI models to compatible API keys#9325
Conversation
There was a problem hiding this comment.
Hey - I've found 3 issues, and left some high level feedback:
- The
_model_key_cacheis defined as a class-level dict but keyed off instance-specific configuration; consider making this an instance property or explicitly documenting the cross-instance sharing to avoid subtle leakage or interference between multiple provider instances in the same process. - In
_model_key_cache_id, you hashstr(key)for each API key, which will use the string form of whatever object is passed (e.g. secret wrapper classes); if those__str__implementations ever log or embed additional detail, it may be safer to rely on a stable but minimal identifier (e.g. index only, or a dedicated opaque key-id) to reduce the chance of surprising cache fragmentation or indirect disclosure.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The `_model_key_cache` is defined as a class-level dict but keyed off instance-specific configuration; consider making this an instance property or explicitly documenting the cross-instance sharing to avoid subtle leakage or interference between multiple provider instances in the same process.
- In `_model_key_cache_id`, you hash `str(key)` for each API key, which will use the string form of whatever object is passed (e.g. secret wrapper classes); if those `__str__` implementations ever log or embed additional detail, it may be safer to rely on a stable but minimal identifier (e.g. index only, or a dedicated opaque key-id) to reduce the chance of surprising cache fragmentation or indirect disclosure.
## Individual Comments
### Comment 1
<location path="astrbot/core/provider/sources/openai_source.py" line_range="665-666" />
<code_context>
+ 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
</code_context>
<issue_to_address>
**suggestion:** Clarify intent of returning sorted(model_key_indexes) from get_models in the multi-key branch.
In the multi-key branch, `model_key_indexes` is a dict of model ID → key index list, so `sorted(model_key_indexes)` returns a sorted list of IDs via implicit key iteration. To make this clearer and consistent with the single-key branch, and more robust if `model_key_indexes`’s type changes, consider returning `sorted(model_key_indexes.keys())` instead.
Suggested implementation:
```python
self._store_model_key_indexes(model_key_indexes)
return sorted(model_key_indexes.keys())
```
None needed beyond this change, assuming the single-key branch already returns a sorted list of model IDs and callers expect a list of IDs in both branches.
</issue_to_address>
### Comment 2
<location path="tests/test_openai_source.py" line_range="194-219" />
<code_context>
+
+
+@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()
+
+
</code_context>
<issue_to_address>
**suggestion (testing):** Add a test for the case where all keys fail during model discovery to assert the error propagation behavior.
Currently only the mixed-success case is tested (one key failing, one succeeding). There is also a path in `get_models` where all requests fail (`successful_requests == 0 and last_error is not None`) and the last stored exception is raised. Please add a test where `FakeClient.list()` fails for both keys to verify that `get_models()` raises the last error and that `model_key_indexes` remains empty in this failure-only path.
```suggestion
+
+
+@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):
+ # Both keys will trigger a failure path
+ raise RuntimeError(f"access denied for key: {self.key}")
+
+ async def close(self):
+ return None
+
+ provider = _make_provider({"key": ["bad-key-1", "bad-key-2"]})
+ provider._create_sdk_client = FakeClient
+ try:
+ with pytest.raises(RuntimeError) as excinfo:
+ await provider.get_models()
+
+ # The last stored exception should be raised
+ assert "bad-key-2" in str(excinfo.value)
+
+ # No successful requests means no model-key index mappings
+ assert provider.get_model_key_indexes() == {}
+ finally:
+ await provider.terminate()
+
+
```
</issue_to_address>
### Comment 3
<location path="astrbot/core/provider/sources/openai_source.py" line_range="54" />
<code_context>
)
class ProviderOpenAIOfficial(Provider):
_ERROR_TEXT_CANDIDATE_MAX_CHARS = 4096
+ _model_key_cache: dict[str, dict[str, tuple[int, ...]]] = {}
@classmethod
</code_context>
<issue_to_address>
**issue (complexity):** Consider replacing the global hashed `_model_key_cache` with a simple instance-level `model -> key index` mapping to keep multi-key/model behavior while reducing state and helper-method complexity.
The new multi-key/model logic adds useful behavior but the global hashed cache and tuple/list copying notably increase complexity. You can preserve all functionality with simpler, instance-scoped data structures.
### 1. Replace global hashed cache with instance-level mapping
Instead of `_model_key_cache` + `_model_key_cache_id` + per-call hashing/JSON, keep a simple mapping on the instance:
```python
class ProviderOpenAIOfficial(Provider):
# remove _model_key_cache and _model_key_cache_id
def __init__(self, provider_config, provider_settings) -> None:
super().__init__(provider_config, provider_settings)
self.api_keys: list = super().get_keys()
self.chosen_api_key = self.api_keys[0] if self.api_keys else None
self.timeout = provider_config.get("timeout", 120)
self.custom_headers = provider_config.get("custom_headers", {})
# ...
self.client = self._create_sdk_client(self.chosen_api_key)
# simple instance-scoped mapping: model -> set of key indexes
self._model_key_indexes: dict[str, set[int]] = {}
```
Then change `_store_model_key_indexes` / `get_model_key_indexes` to write directly into `self._model_key_indexes`:
```python
def _store_model_key_indexes(
self,
model_key_indexes: dict[str, list[int]],
) -> None:
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 {
model: sorted(indexes)
for model, indexes in self._model_key_indexes.items()
}
```
This removes the need for:
- class-level `_model_key_cache`
- `_model_key_cache_id`
- double SHA-256 hashing + JSON serialization
while keeping the same semantics for “which keys can access which models” for a given provider instance.
### 2. Simplify candidate key selection and remember/forget
With `self._model_key_indexes` as `dict[str, set[int]]`, the helper methods become much more direct:
```python
def _candidate_api_keys_for_model(self, model: str) -> list[str]:
if not self.api_keys:
return [""]
indexes = self._model_key_indexes.get(model)
if not indexes:
# no knowledge yet: try all keys (deduplicated)
return list(dict.fromkeys(self.api_keys))
return [
self.api_keys[i]
for i in sorted(indexes)
if 0 <= i < len(self.api_keys)
] or list(dict.fromkeys(self.api_keys))
def _remember_model_key(self, model: str, api_key: str) -> None:
if api_key not in self.api_keys:
return
key_index = self.api_keys.index(api_key)
self._model_key_indexes.setdefault(model, set()).add(key_index)
def _forget_model_key(self, model: str, api_key: str) -> None:
if api_key not in self.api_keys:
return
key_index = self.api_keys.index(api_key)
indexes = self._model_key_indexes.get(model)
if not indexes:
return
indexes.discard(key_index)
if not indexes:
self._model_key_indexes.pop(model, None)
```
This eliminates:
- repeated `dict(...)` cloning of the whole cache
- list/tuple conversions (`list(...)` → `tuple(...)` → back)
- dependence on the global `_model_key_cache` for per-instance behavior
while keeping:
- per-model multi-key tracking
- dynamic updates after success/failure
- deduplicated key selection
### 3. Keep `get_models` aligned with the simpler mapping
The multi-key discovery in `get_models` can remain essentially as-is, just writing into the new instance mapping:
```python
async def get_models(self):
api_keys = getattr(self, "api_keys", None)
if not isinstance(api_keys, list) or len(api_keys) <= 1:
# single-key case unchanged, just store index 0
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
unique_keys = list(dict.fromkeys(api_keys))
model_key_indexes: dict[str, list[int]] = {}
# ... same loop as you have now, filling model_key_indexes ...
self._store_model_key_indexes(model_key_indexes)
return sorted(model_key_indexes)
```
Because `_store_model_key_indexes` is now just a thin wrapper around `self._model_key_indexes`, the rest of the logic (temporary clients, error handling, logging) stays intact while avoiding the global cache.
---
These small changes:
- keep all current multi-key and retry behavior
- remove the global hashed cache and double hashing
- reduce list/tuple conversion churn
- consolidate state into a single, testable instance-level mapping (`_model_key_indexes`), making the overall flow in `text_chat`/`text_chat_stream` easier to reason about.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
There was a problem hiding this comment.
Code Review
This pull request implements multi-key support and automatic model-to-key mapping discovery for the OpenAI provider. It allows the system to fetch available models across all configured API keys, cache which keys have access to which models, and automatically retry requests with alternative keys upon encountering model access or authentication errors. The frontend is updated to display which keys are associated with each model. The reviewer suggests adding safety checks in the single-key model discovery path to handle potentially malformed model IDs, matching the robustness of the multi-key path.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| "OpenAI", | ||
| lambda: self.client.models.list(), | ||
| ) | ||
| models_str = sorted(model.id for model in models.data) |
There was a problem hiding this comment.
To prevent potential AttributeError or TypeError if a model object is missing an id or has a non-string id, we should apply the same robust safety checks here as we do in the multi-key discovery block below.
| 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 | |
| ) |
There was a problem hiding this comment.
Addressed in 7dee31d: single-key discovery now filters missing, empty, and non-string model IDs using the same guarded extraction as multi-key discovery.
c99e8de to
7dee31d
Compare
OpenAI-compatible provider sources already accept multiple API keys, but model discovery only queried one client. When keys expose different model sets, models can be hidden from the dashboard and chat requests may select a key that cannot access the configured model.
Modifications / Changes
Query every unique OpenAI-compatible API key during model discovery and merge the resulting model IDs.
Keep model-to-key ownership on each provider instance so chat and streaming requests prefer keys already known to expose the selected model.
Retry another configured key for authentication or model-access failures, and update the cached ownership after success/failure.
Return model ownership from the provider-source models API.
Label API key list items as Key 1, Key 2, etc., and show the matching key badges beside discovered/configured models.
Wait for the provider source refresh to finish before fetching models, avoiding stale state after saving.
This is NOT a breaking change.
Screenshots or Test Results / Verification
Key numbering in provider configuration
Per-model key ownership
Verification performed:
A full Windows run of ests/test_openai_source.py completed with 62 passed and three existing path-separator assertion failures in ile_uri_to_path; the multi-key tests above pass independently.
Checklist
Summary by Sourcery
Improve multi-key handling for OpenAI-compatible provider sources by aggregating models across all keys, caching model-to-key ownership, and using that information to route and display models correctly in the dashboard.
New Features:
Bug Fixes:
Enhancements:
Tests: