Skip to content

fix: route OpenAI models to compatible API keys#9325

Open
zzz27578 wants to merge 2 commits into
AstrBotDevs:masterfrom
zzz27578:fix/openai-multi-key-model-routing
Open

fix: route OpenAI models to compatible API keys#9325
zzz27578 wants to merge 2 commits into
AstrBotDevs:masterfrom
zzz27578:fix/openai-multi-key-model-routing

Conversation

@zzz27578

@zzz27578 zzz27578 commented Jul 19, 2026

Copy link
Copy Markdown
Contributor

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

Key numbering

Per-model key ownership

Per-model key ownership

Verification performed:

  • uv sync --group dev
  • uv run ruff format --check .
  • uv run ruff check .
  • uv run pytest tests/test_openai_source.py::test_get_models_merges_models_from_all_api_keys tests/test_openai_source.py::test_get_models_keeps_successful_keys_when_one_key_fails tests/test_openai_source.py::test_text_chat_retries_with_another_key_for_model_access_error -q (3 passed)
  • pnpm typecheck
  • pnpm build

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

  • Feature discussion is not required because this PR fixes existing multi-key provider behavior.
  • The changes have been tested, with verification commands and screenshots provided above.
  • No new dependencies are introduced.
  • The changes do not introduce malicious code.

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:

  • Track per-model API key ownership for OpenAI-compatible providers and surface it via the provider-source models API and dashboard badges.
  • Support displaying ordinal labels for multiple OpenAI API keys in provider configuration lists.

Bug Fixes:

  • Ensure model discovery aggregates models from all configured OpenAI-compatible API keys so models are not hidden when keys expose different model sets.
  • Prefer API keys known to expose a requested model for chat and streaming requests, and retry with another key on authentication or model-access failures to avoid request errors.
  • Avoid fetching provider-source models against stale configuration by waiting for the source refresh to complete after saving.

Enhancements:

  • Refactor OpenAI client construction into a reusable helper and introduce a secret-safe cache for model-to-key mappings.

Tests:

  • Add targeted tests covering multi-key model discovery, partial key failures during discovery, and key retry behavior on model access errors.

@dosubot dosubot Bot added size:L This PR changes 100-499 lines, ignoring generated files. area:provider The bug / feature is about AI Provider, Models, LLM Agent, LLM Agent Runner. area:webui The bug / feature is about webui(dashboard) of astrbot. labels Jul 19, 2026
Comment thread astrbot/core/provider/sources/openai_source.py Fixed

@sourcery-ai sourcery-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey - I've found 3 issues, and left some high level feedback:

  • 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.
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>

Sourcery is free for open source - if you like our reviews please consider sharing them ✨
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.

Comment thread astrbot/core/provider/sources/openai_source.py Outdated
Comment thread tests/test_openai_source.py
Comment thread astrbot/core/provider/sources/openai_source.py Outdated

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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.

Suggested change
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
)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in 7dee31d: single-key discovery now filters missing, empty, and non-string model IDs using the same guarded extraction as multi-key discovery.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:provider The bug / feature is about AI Provider, Models, LLM Agent, LLM Agent Runner. area:webui The bug / feature is about webui(dashboard) of astrbot. size:L This PR changes 100-499 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants