mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-07-07 07:59:29 +00:00
fix(cli): key model-catalog cache by exact API key, shorten TTL, skip Ollama
Follow-ups to #6462's caching: 1. Key the catalog cache by the exact API key (via a short, non-reversible sha256 digest — never the key itself), not just key-present vs absent. Switching to a different key for the same provider now misses the previous account's entry and refetches, instead of showing the old account's models. 2. Never cache local providers (Ollama). /api/tags is fast and installed models change out-of-band, so caching could keep offering a model the user just deleted until the entry expired. _is_cacheable() gates both cache read and write; the picker now re-probes every call and reflects what's installed. 3. Shorten the dynamic catalog TTL from 6h to 5m — a stale list (new/removed models, account changes) is worse than a ~1s refetch, and the cache only needs to spare repeated fetches within a wizard session. Tests: distinct-key cache entries, digest never stores the raw key, Ollama not cached (reflects deletions / never written), and dynamic TTL expiry. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
This commit is contained in:
@@ -24,6 +24,7 @@ from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
import contextlib
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
@@ -45,8 +46,11 @@ MAX_MODELS = 8
|
||||
#: Timeout (seconds) for any network call made while resolving models.
|
||||
_TIMEOUT = 6.0
|
||||
|
||||
#: How long a resolved (dynamic) catalog stays fresh before we refetch.
|
||||
_CATALOG_TTL = 6 * 3600
|
||||
#: How long a resolved (dynamic) catalog stays fresh before we refetch. Kept
|
||||
#: short: it only spares the picker repeated fetches within a wizard session,
|
||||
#: and a stale list (new/removed models, account changes) is worse than a ~1s
|
||||
#: refetch. Local providers (Ollama) are not cached at all — see _is_cacheable.
|
||||
_CATALOG_TTL = 300
|
||||
|
||||
#: How long a fallback result is cached after a failed/empty fetch. Short, so a
|
||||
#: newly-added API key takes effect soon, but long enough to spare the picker a
|
||||
@@ -171,9 +175,8 @@ def get_provider_models(
|
||||
|
||||
# Nothing fresher than the curated list. Cache it briefly (negative cache)
|
||||
# so a failed vendor/LiteLLM fetch isn't retried on every subsequent call.
|
||||
# Skip Ollama: it's a local, fast-failing server, so re-probing is cheap and
|
||||
# avoids serving suggestions after the server comes up within the TTL.
|
||||
if fallback and provider_key != "ollama":
|
||||
# (_write_catalog_cache skips non-cacheable providers like Ollama.)
|
||||
if fallback:
|
||||
_write_catalog_cache(provider_key, fallback, source="fallback")
|
||||
return fallback
|
||||
|
||||
@@ -601,25 +604,36 @@ def _litellm_cache_file() -> Path:
|
||||
return _cache_dir() / "provider_cache.json"
|
||||
|
||||
|
||||
def _is_cacheable(provider_key: str) -> bool:
|
||||
"""Whether a provider's resolved catalog may be cached.
|
||||
|
||||
Ollama is a local server (``/api/tags`` is fast), and its installed models
|
||||
change out-of-band, so it is never cached — the picker re-probes every call
|
||||
and always reflects what is currently installed.
|
||||
"""
|
||||
return provider_key != "ollama"
|
||||
|
||||
|
||||
def _cache_key(provider_key: str) -> str:
|
||||
"""Cache key for a provider's resolved model list.
|
||||
|
||||
Includes the inputs that change what a fetch would return, so a cached
|
||||
entry is only reused when those inputs still match:
|
||||
|
||||
- Ollama lists models from a base URL that can change between runs.
|
||||
- Whether the vendor's API key is present flips between a live fetch and
|
||||
the negatively-cached fallback — so a key added after a no-key call is
|
||||
not shadowed by the cached fallback.
|
||||
Keyed by the exact API key (via a short, non-reversible digest — never the
|
||||
key itself), so switching to a different key for the same provider misses
|
||||
the previous account's cached entry and refetches. Absent key -> ``#nokey``,
|
||||
which also keeps a negatively-cached no-key fallback from shadowing a run
|
||||
after a key is added.
|
||||
"""
|
||||
if provider_key == "ollama":
|
||||
return f"ollama@{_ollama_base()}"
|
||||
suffix = "key" if _provider_api_key(provider_key) else "nokey"
|
||||
return f"{provider_key}#{suffix}"
|
||||
api_key = _provider_api_key(provider_key)
|
||||
if not api_key:
|
||||
return f"{provider_key}#nokey"
|
||||
digest = hashlib.sha256(api_key.encode("utf-8")).hexdigest()[:12]
|
||||
return f"{provider_key}#{digest}"
|
||||
|
||||
|
||||
def _read_catalog_cache(provider_key: str) -> list[tuple[str, str]] | None:
|
||||
"""Return a fresh cached catalog for ``provider_key``, or ``None``."""
|
||||
if not _is_cacheable(provider_key):
|
||||
return None
|
||||
payload = _read_json(_catalog_cache_file())
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
@@ -642,6 +656,8 @@ def _read_catalog_cache(provider_key: str) -> list[tuple[str, str]] | None:
|
||||
def _write_catalog_cache(
|
||||
provider_key: str, models: list[tuple[str, str]], *, source: str
|
||||
) -> None:
|
||||
if not _is_cacheable(provider_key):
|
||||
return
|
||||
cache_file = _catalog_cache_file()
|
||||
payload = _read_json(cache_file)
|
||||
if not isinstance(payload, dict):
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -534,18 +535,72 @@ def test_bad_cache_json_does_not_crash(monkeypatch):
|
||||
assert models == FALLBACK_ANTHROPIC
|
||||
|
||||
|
||||
def test_ollama_cache_keyed_by_base(monkeypatch):
|
||||
# Changing OLLAMA_API_BASE must not serve the previous host's cached models.
|
||||
monkeypatch.setenv("OLLAMA_API_BASE", "http://host-a:11434")
|
||||
monkeypatch.setattr(
|
||||
mc, "_http_get_json", lambda *a, **k: {"models": [{"model": "llama-a"}]}
|
||||
def test_ollama_is_not_cached_reflects_installed_changes(monkeypatch):
|
||||
# Ollama is local and never cached: the picker re-probes /api/tags on every
|
||||
# call, so a model deleted locally drops out immediately (no stale entry).
|
||||
responses = iter(
|
||||
[
|
||||
{"models": [{"model": "llama3.3"}, {"model": "qwen3"}]},
|
||||
{"models": [{"model": "llama3.3"}]}, # qwen3 deleted between calls
|
||||
]
|
||||
)
|
||||
first = mc.get_provider_models("ollama", [])
|
||||
assert [m for m, _ in first] == ["llama-a"]
|
||||
monkeypatch.setattr(mc, "_http_get_json", lambda *a, **k: next(responses))
|
||||
|
||||
monkeypatch.setenv("OLLAMA_API_BASE", "http://host-b:11434")
|
||||
first = {m for m, _ in mc.get_provider_models("ollama", [])}
|
||||
second = {m for m, _ in mc.get_provider_models("ollama", [])}
|
||||
|
||||
assert first == {"llama3.3", "qwen3"}
|
||||
assert second == {"llama3.3"} # re-probed, not served from a stale cache
|
||||
|
||||
|
||||
def test_ollama_never_written_to_catalog_cache(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
mc, "_http_get_json", lambda *a, **k: {"models": [{"model": "llama-b"}]}
|
||||
mc, "_http_get_json", lambda *a, **k: {"models": [{"model": "llama3.3"}]}
|
||||
)
|
||||
second = mc.get_provider_models("ollama", [])
|
||||
assert [m for m, _ in second] == ["llama-b"] # not the host-a cache
|
||||
mc.get_provider_models("ollama", [])
|
||||
assert not mc._catalog_cache_file().exists()
|
||||
|
||||
|
||||
def test_different_api_key_uses_separate_cache_entry(monkeypatch):
|
||||
# Cache is keyed by the exact key: switching keys must refetch, not serve
|
||||
# the previous account's cached list.
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-account-A")
|
||||
monkeypatch.setattr(
|
||||
mc, "_http_get_json", lambda *a, **k: {"data": [{"id": "gpt-5.5", "created": 1}]}
|
||||
)
|
||||
assert [m for m, _ in mc.get_provider_models("openai", [])] == ["gpt-5.5"]
|
||||
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-account-B")
|
||||
monkeypatch.setattr(
|
||||
mc, "_http_get_json", lambda *a, **k: {"data": [{"id": "gpt-4.1", "created": 1}]}
|
||||
)
|
||||
# New key -> distinct cache entry -> refetch, not the account-A cache.
|
||||
assert [m for m, _ in mc.get_provider_models("openai", [])] == ["gpt-4.1"]
|
||||
|
||||
|
||||
def test_cache_key_hashes_key_and_never_stores_it(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-super-secret")
|
||||
key = mc._cache_key("openai")
|
||||
assert key.startswith("openai#") and key != "openai#nokey"
|
||||
assert "sk-super-secret" not in key # only a digest, never the raw key
|
||||
|
||||
|
||||
def test_dynamic_cache_expires_after_catalog_ttl(monkeypatch):
|
||||
# A dynamic entry older than the (now short) catalog TTL is not served.
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test")
|
||||
entry_key = mc._cache_key("anthropic")
|
||||
stale_ts = time.time() - (mc._CATALOG_TTL + 5)
|
||||
mc._catalog_cache_file().write_text(
|
||||
json.dumps(
|
||||
{
|
||||
entry_key: {
|
||||
"ts": stale_ts,
|
||||
"source": "dynamic",
|
||||
"models": [["stale-model", "Stale"]],
|
||||
}
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
assert mc._read_catalog_cache("anthropic") is None
|
||||
|
||||
Reference in New Issue
Block a user