mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-07-07 07:59:29 +00:00
feat(cli): pull latest LLM models dynamically in the crew wizard (#6462)
* feat(cli): pull latest LLM models dynamically in the crew wizard
The JSON-crew creation wizard hardcoded a short model list per provider,
which goes stale as vendors ship new models every few weeks. Add a
three-tier resolver that prefers live data and falls back to a curated list.
- New `model_catalog.get_provider_models(provider, fallback)`:
1. Vendor API (openai/anthropic/gemini/groq/cerebras/ollama) when the
provider key is already in the environment — the only reliably-fresh
source (real release dates / display names).
2. Curated hardcoded fallback — hand-verified, used when no key is set.
3. LiteLLM feed — only for providers with no curated list; it lags real
releases, so it must never preempt the curated fallback.
- Rank by date/version parsed from model ids, humanize labels, 6h cache,
short timeouts, silent fallback on any error.
- Wire it into `create_json_crew._select_model()` (picker only).
- Refresh the curated fallback against each vendor's official model docs
(Anthropic Fable 5 / Opus 4.8 / Sonnet 5; OpenAI GPT-5.5(+pro); Gemini
3.5 Flash / 3.1 Pro preview / 3 Flash preview; Groq Llama 4 / GPT-OSS).
- Tests for ranking, chat filtering, caching, and the tier order (17 tests).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
* fix(cli): address model_catalog review findings
- Key the Ollama catalog cache by its base URL so a changed OLLAMA_API_BASE /
API_BASE no longer serves the previous host's models for up to the TTL.
- Negatively cache the curated fallback after a failed/empty fetch (short
_NEGATIVE_TTL) so the picker doesn't repeat a timeout-prone vendor/LiteLLM
request on every call — most impactful for a down local Ollama server.
- Guard _read_catalog_cache / _write_catalog_cache against a non-dict cache
root (corrupt JSON array no longer raises AttributeError).
- Replace the two empty `except OSError: pass` blocks with
contextlib.suppress(OSError) plus an explanatory comment (CodeQL empty-except).
- Tests: negative cache, base-keyed Ollama cache, corrupt-cache no-crash (20 total).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
* fix(cli): guard null litellm_provider and paginate Gemini models
- _from_litellm: coerce a present-but-null `litellm_provider` before string
ops so it's skipped instead of raising AttributeError (keeps the documented
"never raises" contract).
- _fetch_gemini: walk models.list pages via nextPageToken (bounded to 10) —
the API is paginated and not guaranteed newest-first, so a single page could
drop models the ranking should consider.
- Tests for both (22 total).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
* ci: ignore nltk PYSEC-2026-597 in pip-audit (no fix, not reachable)
pip-audit newly flags nltk 3.9.4 for PYSEC-2026-597 (CVE-2026-12243), a path
traversal via percent-encoded `..%2f` in nltk.data.load()/find(). It affects
all nltk versions <=3.9.4 with no patched release, so it can't be resolved by a
version bump — same situation as the already-ignored PYSEC-2026-97.
nltk is a transitive dependency (unstructured[local-inference, all-docs] in
crewai-tools) used for text tokenization; we never pass untrusted resource
URLs/paths to nltk.data, so the traversal is not reachable. Add it to the
curated --ignore-vuln list with a justification, matching the existing pattern.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
* fix(cli): address second Cursor review round on model_catalog
- Cache ignores new API keys: include API-key presence in the cache key
(`<provider>#key|#nokey`), so a key added after a no-key/negative-cached
lookup triggers a fresh live fetch instead of serving the stale fallback.
- Bad LiteLLM cache crashes picker: `_from_litellm` now requires a dict from
`_load_litellm_data` (a non-mapping JSON root is skipped, not `.items()`'d).
- Stale LiteLLM refetch loop: memoize the feed load once per process
(`_litellm_memo` + `_reset_litellm_memo` test hook) so repeated uncurated-
provider lookups don't each re-attempt a timed download when offline.
- Tests: new-key bypass, corrupt-litellm-cache no-crash, one-fetch-per-process
(25 total).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
* fix(cli): keep partial Gemini results on later-page fetch error
_fetch_gemini paginates; a network/HTTP error on page 2+ previously raised out
through _from_vendor, discarding models already parsed from earlier pages and
forcing the curated fallback. Catch per-page fetch errors and return the
partial set instead (a first-page failure still yields an empty list -> fallback).
Test: test_vendor_gemini_keeps_partial_on_later_page_error (26 total).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
* fix(cli): don't let an invalid fresh LiteLLM cache block download
_fetch_litellm_data treated any truthy JSON root in a fresh provider_cache.json
as the feed and returned it, so a non-mapping root (e.g. a JSON array) was
memoized and the tier never re-downloaded until the file aged out — leaving
uncurated providers with an empty picker despite a recoverable cache. Only
short-circuit on a usable dict; otherwise fall through to the download.
Test renamed to test_invalid_litellm_cache_falls_through_to_download (asserts
recovery via refetch).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
* fix(cli): honor OLLAMA_HOST and treat empty vendor list as authoritative
- Ollama host mismatch: _ollama_base now also reads OLLAMA_HOST (the Ollama
runtime convention) after OLLAMA_API_BASE/API_BASE, normalizing a scheme-less
value (e.g. "127.0.0.1:11434" -> "http://127.0.0.1:11434"), so users who set
only OLLAMA_HOST see models from the server the crew will actually use.
- Empty vendor list: a successful vendor fetch returning no models is now
authoritative instead of collapsing to the curated fallback. A reachable
Ollama with nothing installed yields an empty list (the picker prompts for
manual entry) rather than offering hardcoded models that aren't installed; a
failed fetch still falls back. _from_vendor now returns [] on success-empty
and None only when the tier is unavailable.
- Tests: ollama empty->manual, ollama down->fallback, OLLAMA_HOST resolution
(29 total).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
* fix(cli): Gemini first-page failure falls back instead of showing empty
Interaction from the prior two fixes: _fetch_gemini swallowed a first-page
error and returned [], which _from_vendor reported as a successful-empty result
and get_provider_models treated as authoritative — skipping the curated Gemini
fallback and jumping to manual entry. Now a first-page failure (nothing gathered
yet) re-raises so _from_vendor returns None and the curated list is used; a
later-page failure still keeps the partial results.
Test: test_vendor_gemini_first_page_error_uses_fallback (30 total).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
* fix(cli): Gemini GOOGLE_API_KEY + Ollama recovery not blocked by cache
- Gemini ignores GOOGLE_API_KEY: _PROVIDER_KEY_ENV now maps each provider to a
tuple of accepted env vars; Gemini accepts GEMINI_API_KEY or GOOGLE_API_KEY
(matching crewai's own Gemini provider). A new _provider_api_key() resolver
is used by both _from_vendor and the cache key, so a GOOGLE_API_KEY user gets
the live models API instead of the stale curated fallback.
- Ollama recovery blocked by cache: skip the negative (fallback) cache for
Ollama. It's a local, fast-failing server, so re-probing each call is cheap
and lets the picker pick up real installed models as soon as the server comes
up, instead of serving suggestions for the negative-cache TTL.
- Tests: GOOGLE_API_KEY live fetch, Ollama down->recover (32 total).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
* style(cli): ruff format model_catalog.py
Add the blank line ruff format expects after _provider_api_key; no behavior
change. Fixes the lint-run `ruff format --check lib/` step.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
* fix(cli): don't treat 'search' substring as a non-chat model marker
The 'search' entry in _NON_CHAT_MARKERS matched anywhere in a model id, dropping
legitimate completion models like gpt-4o-search-preview and anything containing
'research' (e.g. o3-deep-research, since 'search' is a substring). Remove it;
the remaining markers (embedding/audio/image/moderation/etc.) still filter
genuine non-chat models. Test: test_search_substring_not_treated_as_non_chat (33).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
* Revert "ci: ignore nltk PYSEC-2026-597 in pip-audit"
Do not suppress an unpatched security advisory to make CI green. Remove
PYSEC-2026-597 from the pip-audit ignore list; leave the scan failing so it
keeps surfacing the nltk path traversal (CVE-2026-12243). This PR should not be
merged until nltk ships a fix (or the vulnerable transitive dep is otherwise
resolved).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
* fix(cli): exclude fine-tuned models and checkpoints from the picker
With a live OPENAI_API_KEY, /v1/models returns the user's fine-tunes and
training checkpoints (ft:..., ...:ckpt-step-N). Their recent `created`
timestamps ranked them above the base models and filled every slot, so the
picker showed a wall of `ft:gpt-4o-mini-...:crewai::...` with mangled labels and
no foundation models at all. Skip fine-tunes/checkpoints in the OpenAI-shaped
fetcher so clean base models surface; a user who wants a fine-tune can still
enter it via the picker's "Other" option. Test:
test_openai_excludes_fine_tunes_and_checkpoints (34 total).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
* fix(cli): cleaner model labels + filter Ollama non-chat models
Anthropic/Gemini already use vendor display names; OpenAI/Groq/Cerebras/Ollama
fall to _humanize for anything outside the curated map, which produced mediocre
labels ("GPT Oss 120b", "qwen3 32b", "Deepseek r1", "llama3.3:70b").
Improve _humanize:
- split on ':' too (Ollama tags: llama3.3:70b -> "Llama3.3 70B")
- uppercase size suffixes (70b -> 70B), acronyms OSS/IT, brand casing
(DeepSeek, ChatGPT, QwQ)
- capitalize the leading letter of fused family+version tokens (qwen3 -> Qwen3)
while preserving OpenAI o-series lowercase (o3, o1-mini)
Also fix _fetch_ollama: /api/tags lists everything installed, so filter
non-chat (embedding) and fine-tune entries the same way the other tiers do.
Tests: expanded test_humanize + test_ollama_excludes_embedding_models (35 total).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RBYGqJHC2TMC6fonFziuuh
---------
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Lorenze Jay <63378463+lorenzejay@users.noreply.github.com>
This commit is contained in:
João Moura
@@ -14,6 +14,7 @@ from rich.text import Text
|
||||
|
||||
from crewai_cli.constants import ENV_VARS
|
||||
from crewai_cli.git import initialize_if_git_available
|
||||
from crewai_cli.model_catalog import get_provider_models
|
||||
from crewai_cli.tui_picker import pick_many, pick_one
|
||||
from crewai_cli.utils import (
|
||||
enable_prompt_line_editing,
|
||||
@@ -42,41 +43,50 @@ _PROVIDERS: list[tuple[str, str]] = [
|
||||
("watson", "IBM watsonx"),
|
||||
]
|
||||
|
||||
# Curated offline fallback / label source. The picker prefers models pulled
|
||||
# live from the vendor's own API via ``model_catalog.get_provider_models``;
|
||||
# this list is the hand-verified backstop used when no API key is available.
|
||||
# Keep entries to real, current model ids — last verified against each vendor's
|
||||
# official model docs on 2026-07-05.
|
||||
_PROVIDER_MODELS: dict[str, list[tuple[str, str]]] = {
|
||||
"openai": [
|
||||
("gpt-5.5", "GPT-5.5"),
|
||||
("gpt-5.5-pro", "GPT-5.5 Pro"),
|
||||
("gpt-5.4", "GPT-5.4"),
|
||||
("o4-mini", "o4-mini"),
|
||||
("gpt-5.4-mini", "GPT-5.4 Mini"),
|
||||
("gpt-5.2", "GPT-5.2"),
|
||||
("gpt-4.1", "GPT-4.1"),
|
||||
("gpt-4.1-mini", "GPT-4.1 Mini"),
|
||||
],
|
||||
"anthropic": [
|
||||
("claude-opus-4-6", "Claude Opus 4.6"),
|
||||
("claude-fable-5", "Claude Fable 5"),
|
||||
("claude-opus-4-8", "Claude Opus 4.8"),
|
||||
("claude-sonnet-5", "Claude Sonnet 5"),
|
||||
("claude-opus-4-7", "Claude Opus 4.7"),
|
||||
("claude-haiku-4-5", "Claude Haiku 4.5"),
|
||||
("claude-sonnet-4-6", "Claude Sonnet 4.6"),
|
||||
("claude-haiku-4-5-20251001", "Claude Haiku 4.5"),
|
||||
("claude-3-7-sonnet-20250219", "Claude 3.7 Sonnet"),
|
||||
("claude-3-5-sonnet-20241022", "Claude 3.5 Sonnet"),
|
||||
],
|
||||
"gemini": [
|
||||
("gemini-3-pro-preview", "Gemini 3 Pro (preview)"),
|
||||
("gemini-2.5-pro-exp-03-25", "Gemini 2.5 Pro"),
|
||||
("gemini-2.5-flash-preview-04-17", "Gemini 2.5 Flash"),
|
||||
("gemini-2.0-flash-001", "Gemini 2.0 Flash"),
|
||||
("gemini-1.5-pro", "Gemini 1.5 Pro"),
|
||||
("gemini-3.5-flash", "Gemini 3.5 Flash"),
|
||||
("gemini-3.1-pro-preview", "Gemini 3.1 Pro (preview)"),
|
||||
("gemini-3-flash-preview", "Gemini 3 Flash (preview)"),
|
||||
("gemini-2.5-pro", "Gemini 2.5 Pro"),
|
||||
("gemini-2.5-flash", "Gemini 2.5 Flash"),
|
||||
("gemini-2.5-flash-lite", "Gemini 2.5 Flash Lite"),
|
||||
],
|
||||
"groq": [
|
||||
("meta-llama/llama-4-maverick-17b-128e-instruct", "Llama 4 Maverick"),
|
||||
("meta-llama/llama-4-scout-17b-16e-instruct", "Llama 4 Scout"),
|
||||
("openai/gpt-oss-120b", "GPT-OSS 120B"),
|
||||
("qwen/qwen3-32b", "Qwen3 32B"),
|
||||
("moonshotai/kimi-k2-instruct-0905", "Kimi K2"),
|
||||
("llama-3.3-70b-versatile", "Llama 3.3 70B"),
|
||||
("llama-3.1-70b-versatile", "Llama 3.1 70B"),
|
||||
("llama-3.1-8b-instant", "Llama 3.1 8B"),
|
||||
("deepseek-r1-distill-llama-70b", "DeepSeek R1 70B"),
|
||||
("mixtral-8x7b-32768", "Mixtral 8x7B"),
|
||||
],
|
||||
"ollama": [
|
||||
("llama3.3", "Llama 3.3"),
|
||||
("llama3.1", "Llama 3.1"),
|
||||
("qwen3", "Qwen 3"),
|
||||
("deepseek-r1", "DeepSeek R1"),
|
||||
("qwen2.5", "Qwen 2.5"),
|
||||
("gpt-oss", "GPT-OSS"),
|
||||
("gemma3", "Gemma 3"),
|
||||
("mistral", "Mistral"),
|
||||
],
|
||||
}
|
||||
@@ -758,7 +768,9 @@ def _select_model() -> str:
|
||||
provider_key, provider_name = _PROVIDERS[p_idx]
|
||||
click.secho(f" → {provider_name}", fg="green")
|
||||
|
||||
models = _PROVIDER_MODELS.get(provider_key, [])
|
||||
# Prefer the latest models pulled live from the vendor / LiteLLM; the
|
||||
# curated ``_PROVIDER_MODELS`` entry is the offline fallback and label source.
|
||||
models = get_provider_models(provider_key, _PROVIDER_MODELS.get(provider_key, []))
|
||||
if not models:
|
||||
custom = click.prompt(
|
||||
click.style(f" Enter model name for {provider_key}/", fg="cyan"),
|
||||
|
||||
657
lib/cli/src/crewai_cli/model_catalog.py
Normal file
657
lib/cli/src/crewai_cli/model_catalog.py
Normal file
@@ -0,0 +1,657 @@
|
||||
"""Dynamic model catalog for the crew-creation wizard.
|
||||
|
||||
Resolves the models to offer for a given provider using a three-tier strategy:
|
||||
|
||||
1. **Vendor API** - when the provider's API key is already present in the
|
||||
environment, query the vendor's own model-listing endpoint. This is the only
|
||||
source that reliably reflects the *latest* models (real release dates /
|
||||
display names, straight from the vendor).
|
||||
2. **Curated hardcoded fallback** - the hand-verified list baked into the
|
||||
wizard, used when no API key is available. Authoritative but frozen, so it is
|
||||
refreshed periodically.
|
||||
3. **LiteLLM feed** - the community ``model_prices_and_context_window.json`` the
|
||||
CLI already caches. Only used for providers with *no* curated list: the feed
|
||||
lags real releases badly (it can miss a vendor's newest models entirely), so
|
||||
it must never preempt the curated fallback.
|
||||
|
||||
Every tier is best-effort: any network error, timeout, missing key, or empty
|
||||
result quietly falls through to the next tier, and the caller's hardcoded list
|
||||
is always the final backstop. The picker never blocks for long — network calls
|
||||
use a short timeout and successful results are cached.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
import contextlib
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import certifi
|
||||
import httpx
|
||||
|
||||
from crewai_cli.constants import JSON_URL
|
||||
|
||||
|
||||
# ── Tunables ─────────────────────────────────────────────────────
|
||||
|
||||
#: How many models to surface per provider.
|
||||
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 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
|
||||
#: repeated timeout-prone network attempt on every call within one session.
|
||||
_NEGATIVE_TTL = 300
|
||||
|
||||
#: How long the shared LiteLLM feed cache stays fresh.
|
||||
_LITELLM_TTL = 24 * 3600
|
||||
|
||||
#: Env vars that may hold each provider's API key, in priority order. A
|
||||
#: provider with an empty tuple (e.g. local Ollama) needs no key. Gemini accepts
|
||||
#: either name, matching crewai's own Gemini provider.
|
||||
_PROVIDER_KEY_ENV: dict[str, tuple[str, ...]] = {
|
||||
"openai": ("OPENAI_API_KEY",),
|
||||
"anthropic": ("ANTHROPIC_API_KEY",),
|
||||
"gemini": ("GEMINI_API_KEY", "GOOGLE_API_KEY"),
|
||||
"groq": ("GROQ_API_KEY",),
|
||||
"cerebras": ("CEREBRAS_API_KEY",),
|
||||
"ollama": (),
|
||||
}
|
||||
|
||||
|
||||
def _provider_api_key(provider_key: str) -> str | None:
|
||||
"""First non-empty API key found among the provider's env vars."""
|
||||
for env in _PROVIDER_KEY_ENV.get(provider_key, ()):
|
||||
value = os.environ.get(env)
|
||||
if value:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
# Substrings that mark a model id as *not* a chat/completion model. Used to
|
||||
# filter noisy OpenAI-compatible ``/models`` listings.
|
||||
_NON_CHAT_MARKERS = (
|
||||
"embedding",
|
||||
"embed",
|
||||
"whisper",
|
||||
"tts",
|
||||
"audio",
|
||||
"transcribe",
|
||||
"realtime",
|
||||
"dall-e",
|
||||
"dalle",
|
||||
"image",
|
||||
"moderation",
|
||||
"similarity",
|
||||
"-edit",
|
||||
"davinci-002",
|
||||
"babbage-002",
|
||||
"computer-use",
|
||||
"guard",
|
||||
)
|
||||
|
||||
_ACRONYMS = {
|
||||
"gpt": "GPT",
|
||||
"ai": "AI",
|
||||
"nim": "NIM",
|
||||
"llm": "LLM",
|
||||
"hd": "HD",
|
||||
"us": "US",
|
||||
"eu": "EU",
|
||||
"oss": "OSS",
|
||||
"it": "IT",
|
||||
}
|
||||
|
||||
# Tokens with non-title-case brand capitalization.
|
||||
_BRAND_TOKENS = {
|
||||
"deepseek": "DeepSeek",
|
||||
"chatgpt": "ChatGPT",
|
||||
"qwq": "QwQ",
|
||||
}
|
||||
|
||||
|
||||
# ── Public API ───────────────────────────────────────────────────
|
||||
|
||||
|
||||
def get_provider_models(
|
||||
provider_key: str, fallback: list[tuple[str, str]]
|
||||
) -> list[tuple[str, str]]:
|
||||
"""Return ``(model_id, label)`` pairs for ``provider_key``, newest first.
|
||||
|
||||
Tries the vendor API (if a key is in the environment) first, since it is the
|
||||
only reliably-fresh source. When no key is available it returns the curated
|
||||
``fallback`` verbatim — the LiteLLM feed is consulted **only** for providers
|
||||
with no curated list, because the feed lags real releases and would
|
||||
otherwise surface a staler list than the hand-verified fallback. Never
|
||||
raises: any failure degrades to the next tier.
|
||||
|
||||
Args:
|
||||
provider_key: Short provider identifier, e.g. ``"anthropic"``.
|
||||
fallback: Curated ``(model_id, label)`` pairs to use as the backstop and
|
||||
to source friendly labels for known models.
|
||||
|
||||
Returns:
|
||||
Up to :data:`MAX_MODELS` ``(model_id, label)`` pairs. Falls back to
|
||||
``fallback`` verbatim when no fresher list can be resolved.
|
||||
"""
|
||||
cached = _read_catalog_cache(provider_key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
label_map = {model_id: label for model_id, label in fallback}
|
||||
|
||||
# A non-None vendor result is authoritative — even when empty (e.g. a
|
||||
# reachable Ollama with no models installed): show that rather than
|
||||
# hardcoded suggestions the crew can't actually run. The picker handles an
|
||||
# empty list by prompting for manual entry.
|
||||
vendor = _from_vendor(provider_key)
|
||||
if vendor is not None:
|
||||
result = _finalize(vendor, label_map)
|
||||
if result:
|
||||
_write_catalog_cache(provider_key, result, source="dynamic")
|
||||
return result
|
||||
|
||||
# Vendor tier unavailable. The LiteLLM feed lags real releases, so only
|
||||
# reach for it when we have no curated fallback — never override the fallback.
|
||||
entries = _from_litellm(provider_key) if not fallback else None
|
||||
result = _finalize(entries, label_map) if entries else []
|
||||
if result:
|
||||
_write_catalog_cache(provider_key, result, source="dynamic")
|
||||
return result
|
||||
|
||||
# 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(provider_key, fallback, source="fallback")
|
||||
return fallback
|
||||
|
||||
|
||||
# ── Tier 1: vendor APIs ──────────────────────────────────────────
|
||||
|
||||
|
||||
def _from_vendor(provider_key: str) -> list[dict[str, Any]] | None:
|
||||
"""Fetch models from the vendor.
|
||||
|
||||
Returns the model list on a successful fetch — **including an empty list**,
|
||||
which is meaningful (e.g. a reachable Ollama server with nothing installed).
|
||||
Returns ``None`` only when the vendor tier is unavailable: no fetcher, no
|
||||
API key, or the request failed.
|
||||
"""
|
||||
fetcher = _VENDOR_FETCHERS.get(provider_key)
|
||||
if fetcher is None:
|
||||
return None
|
||||
|
||||
api_key = _provider_api_key(provider_key)
|
||||
if _PROVIDER_KEY_ENV.get(provider_key) and not api_key:
|
||||
# Provider needs a key and none is set — skip to the next tier.
|
||||
return None
|
||||
|
||||
try:
|
||||
return fetcher(api_key)
|
||||
except Exception:
|
||||
# Network error, auth failure, unexpected payload — degrade quietly.
|
||||
return None
|
||||
|
||||
|
||||
def _fetch_openai(api_key: str | None) -> list[dict[str, Any]]:
|
||||
return _fetch_openai_compatible("https://api.openai.com/v1", api_key)
|
||||
|
||||
|
||||
def _fetch_groq(api_key: str | None) -> list[dict[str, Any]]:
|
||||
return _fetch_openai_compatible("https://api.groq.com/openai/v1", api_key)
|
||||
|
||||
|
||||
def _fetch_cerebras(api_key: str | None) -> list[dict[str, Any]]:
|
||||
return _fetch_openai_compatible("https://api.cerebras.ai/v1", api_key)
|
||||
|
||||
|
||||
def _fetch_openai_compatible(
|
||||
base_url: str, api_key: str | None
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Parse an OpenAI-shaped ``GET /models`` response."""
|
||||
data = _http_get_json(
|
||||
f"{base_url}/models",
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
)
|
||||
entries: list[dict[str, Any]] = []
|
||||
for item in data.get("data", []):
|
||||
model_id = item.get("id")
|
||||
if not model_id or not _is_chat_model(model_id) or _is_fine_tune(model_id):
|
||||
continue
|
||||
created = _as_float(item.get("created"))
|
||||
entries.append(_entry(model_id, _humanize(model_id), created=created))
|
||||
return entries
|
||||
|
||||
|
||||
def _fetch_anthropic(api_key: str | None) -> list[dict[str, Any]]:
|
||||
data = _http_get_json(
|
||||
"https://api.anthropic.com/v1/models",
|
||||
headers={"x-api-key": api_key or "", "anthropic-version": "2023-06-01"},
|
||||
)
|
||||
entries: list[dict[str, Any]] = []
|
||||
for item in data.get("data", []):
|
||||
model_id = item.get("id")
|
||||
if not model_id:
|
||||
continue
|
||||
label = item.get("display_name") or _humanize(model_id)
|
||||
created = _parse_iso(item.get("created_at"))
|
||||
entries.append(_entry(model_id, label, created=created))
|
||||
return entries
|
||||
|
||||
|
||||
def _fetch_gemini(api_key: str | None) -> list[dict[str, Any]]:
|
||||
entries: list[dict[str, Any]] = []
|
||||
params: dict[str, Any] = {"key": api_key or "", "pageSize": 200}
|
||||
# models.list is paginated and not guaranteed newest-first, so walk pages
|
||||
# (bounded) to see the full set — _finalize does the sort + truncation.
|
||||
for _ in range(10):
|
||||
try:
|
||||
data = _http_get_json(
|
||||
"https://generativelanguage.googleapis.com/v1beta/models",
|
||||
params=params,
|
||||
)
|
||||
except Exception:
|
||||
# Later-page failure: keep the models already gathered. First-page
|
||||
# failure (nothing gathered yet) is a real outage — re-raise so the
|
||||
# caller falls back to the curated list rather than mistaking it for
|
||||
# a successful empty result.
|
||||
if entries:
|
||||
break
|
||||
raise
|
||||
for item in data.get("models", []):
|
||||
methods = item.get("supportedGenerationMethods") or []
|
||||
if "generateContent" not in methods:
|
||||
continue
|
||||
name = (item.get("name") or "").removeprefix("models/")
|
||||
if not name or not _is_chat_model(name) or "aqa" in name:
|
||||
continue
|
||||
label = item.get("displayName") or _humanize(name)
|
||||
# Gemini has no timestamp; rank by the version in name/version.
|
||||
version_hint = f"{name} {item.get('version') or ''}"
|
||||
entries.append(_entry(name, label, version_hint=version_hint))
|
||||
token = data.get("nextPageToken")
|
||||
if not token:
|
||||
break
|
||||
params = {"key": api_key or "", "pageSize": 200, "pageToken": token}
|
||||
return entries
|
||||
|
||||
|
||||
def _ollama_base() -> str:
|
||||
"""Resolve the Ollama server base URL from the environment.
|
||||
|
||||
Checks ``OLLAMA_API_BASE`` / ``API_BASE`` (what LiteLLM and the generated
|
||||
crew use) first, then ``OLLAMA_HOST`` (the Ollama runtime convention), so a
|
||||
user who only set ``OLLAMA_HOST`` sees models from the right server.
|
||||
"""
|
||||
base = (
|
||||
os.environ.get("OLLAMA_API_BASE")
|
||||
or os.environ.get("API_BASE")
|
||||
or os.environ.get("OLLAMA_HOST")
|
||||
or "http://localhost:11434"
|
||||
).strip()
|
||||
# OLLAMA_HOST is often scheme-less (e.g. "127.0.0.1:11434").
|
||||
if "://" not in base:
|
||||
base = f"http://{base}"
|
||||
return base.rstrip("/")
|
||||
|
||||
|
||||
def _fetch_ollama(_api_key: str | None) -> list[dict[str, Any]]:
|
||||
"""List models installed on the local Ollama server (no API key)."""
|
||||
data = _http_get_json(f"{_ollama_base()}/api/tags")
|
||||
entries: list[dict[str, Any]] = []
|
||||
for item in data.get("models", []):
|
||||
model_id = item.get("model") or item.get("name")
|
||||
if not model_id or not _is_chat_model(model_id) or _is_fine_tune(model_id):
|
||||
# /api/tags lists everything installed, including embedding models.
|
||||
continue
|
||||
# Ollama returns an ISO 8601 modified_at we can rank by.
|
||||
created = _parse_iso(item.get("modified_at"))
|
||||
entries.append(_entry(model_id, _humanize(model_id), created=created))
|
||||
return entries
|
||||
|
||||
|
||||
_VENDOR_FETCHERS: dict[str, Callable[[str | None], list[dict[str, Any]]]] = {
|
||||
"openai": _fetch_openai,
|
||||
"anthropic": _fetch_anthropic,
|
||||
"gemini": _fetch_gemini,
|
||||
"groq": _fetch_groq,
|
||||
"cerebras": _fetch_cerebras,
|
||||
"ollama": _fetch_ollama,
|
||||
}
|
||||
|
||||
|
||||
# ── Tier 2: LiteLLM feed ─────────────────────────────────────────
|
||||
|
||||
# Process-level memo so a single CLI run attempts the LiteLLM download at most
|
||||
# once — repeated picker calls otherwise each incur a multi-second timeout when
|
||||
# the feed is stale/unreachable. Reset via _reset_litellm_memo() in tests.
|
||||
_UNSET: Any = object()
|
||||
_litellm_memo: Any = _UNSET
|
||||
|
||||
|
||||
def _reset_litellm_memo() -> None:
|
||||
"""Clear the process-level LiteLLM memo (test hook)."""
|
||||
global _litellm_memo
|
||||
_litellm_memo = _UNSET
|
||||
|
||||
|
||||
def _from_litellm(provider_key: str) -> list[dict[str, Any]] | None:
|
||||
"""Build chat-model entries for ``provider_key`` from the LiteLLM feed."""
|
||||
data = _load_litellm_data()
|
||||
# A corrupt feed (non-mapping JSON root) must not crash the picker.
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
|
||||
entries: list[dict[str, Any]] = []
|
||||
for model_name, props in data.items():
|
||||
if not isinstance(props, dict):
|
||||
continue
|
||||
# `litellm_provider` can be present-but-null in the feed; coerce before
|
||||
# string ops so a null value is skipped rather than raising.
|
||||
if (props.get("litellm_provider") or "").strip().lower() != provider_key:
|
||||
continue
|
||||
if props.get("mode") != "chat":
|
||||
continue
|
||||
# LiteLLM keys are sometimes prefixed with the provider; the picker
|
||||
# re-adds ``provider/`` itself, so strip a leading one to avoid dupes.
|
||||
model_id = model_name
|
||||
if model_id.startswith(f"{provider_key}/"):
|
||||
model_id = model_id[len(provider_key) + 1 :]
|
||||
if not model_id:
|
||||
continue
|
||||
entries.append(_entry(model_id, _humanize(model_id), version_hint=model_id))
|
||||
return entries or None
|
||||
|
||||
|
||||
def _load_litellm_data() -> dict[str, Any] | None:
|
||||
"""Return the LiteLLM feed, memoized once per process (see _litellm_memo)."""
|
||||
global _litellm_memo
|
||||
if _litellm_memo is _UNSET:
|
||||
_litellm_memo = _fetch_litellm_data()
|
||||
memoized: dict[str, Any] | None = _litellm_memo
|
||||
return memoized
|
||||
|
||||
|
||||
def _fetch_litellm_data() -> dict[str, Any] | None:
|
||||
"""Read the cached LiteLLM feed, fetching it once if the cache is cold."""
|
||||
cache_file = _litellm_cache_file()
|
||||
fresh = (
|
||||
cache_file.exists()
|
||||
and (time.time() - cache_file.stat().st_mtime) < _LITELLM_TTL
|
||||
)
|
||||
if fresh:
|
||||
data = _read_json(cache_file)
|
||||
# A corrupt/non-mapping fresh cache must not block a recoverable
|
||||
# download — only short-circuit on a usable mapping.
|
||||
if isinstance(data, dict) and data:
|
||||
return data
|
||||
|
||||
try:
|
||||
data = _http_get_json(JSON_URL)
|
||||
except Exception:
|
||||
# Fall back to a stale cache if we have one, else give up on this tier.
|
||||
return _read_json(cache_file)
|
||||
|
||||
# Best-effort cache write; a failure (e.g. read-only home) is non-fatal
|
||||
# since we already hold the freshly-fetched data.
|
||||
with contextlib.suppress(OSError):
|
||||
cache_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
cache_file.write_text(json.dumps(data), encoding="utf-8")
|
||||
return data
|
||||
|
||||
|
||||
# ── Ranking + labelling ──────────────────────────────────────────
|
||||
|
||||
|
||||
def _finalize(
|
||||
entries: list[dict[str, Any]], label_map: dict[str, str]
|
||||
) -> list[tuple[str, str]]:
|
||||
"""Sort newest-first, dedupe, relabel with curated names, and truncate."""
|
||||
entries.sort(key=lambda e: e["sort"], reverse=True)
|
||||
seen: set[str] = set()
|
||||
out: list[tuple[str, str]] = []
|
||||
for entry in entries:
|
||||
model_id = entry["id"]
|
||||
if model_id in seen:
|
||||
continue
|
||||
seen.add(model_id)
|
||||
label = label_map.get(model_id) or entry["label"]
|
||||
out.append((model_id, label))
|
||||
if len(out) >= MAX_MODELS:
|
||||
break
|
||||
return out
|
||||
|
||||
|
||||
def _entry(
|
||||
model_id: str,
|
||||
label: str,
|
||||
*,
|
||||
created: float = 0.0,
|
||||
version_hint: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build a rankable catalog entry.
|
||||
|
||||
``sort`` is a comparable tuple ``(created, date_int, version_tuple)`` so a
|
||||
real vendor timestamp wins, then a date embedded in the id, then the numeric
|
||||
version. Types line up positionally, so entries compare cleanly.
|
||||
"""
|
||||
date_int, version = _version_key(version_hint or model_id)
|
||||
return {
|
||||
"id": model_id,
|
||||
"label": label,
|
||||
"sort": (created, date_int, version),
|
||||
}
|
||||
|
||||
|
||||
_DATE_RE = re.compile(r"(20\d{2})[-_]?(0[1-9]|1[0-2])[-_]?(0[1-9]|[12]\d|3[01])")
|
||||
_NUM_RE = re.compile(r"\d+")
|
||||
|
||||
|
||||
def _version_key(text: str) -> tuple[int, tuple[int, ...]]:
|
||||
"""Extract a ``(date_int, version_tuple)`` sort key from a model id.
|
||||
|
||||
A trailing/embedded ``YYYYMMDD`` (or ``YYYY-MM-DD``) becomes ``date_int``;
|
||||
remaining numbers become the version tuple. ``claude-opus-4-6`` → version
|
||||
``(4, 6)``; ``claude-3-5-sonnet-20241022`` → date ``20241022`` version
|
||||
``(3, 5)``.
|
||||
"""
|
||||
text = text or ""
|
||||
date_int = 0
|
||||
match = _DATE_RE.search(text)
|
||||
if match:
|
||||
date_int = int(match.group(1) + match.group(2) + match.group(3))
|
||||
text = _DATE_RE.sub(" ", text)
|
||||
version = tuple(int(n) for n in _NUM_RE.findall(text)[:4])
|
||||
return date_int, version
|
||||
|
||||
|
||||
def _is_chat_model(model_id: str) -> bool:
|
||||
"""Heuristically reject embedding/audio/image/etc. models by their id."""
|
||||
lowered = model_id.lower()
|
||||
return not any(marker in lowered for marker in _NON_CHAT_MARKERS)
|
||||
|
||||
|
||||
def _is_fine_tune(model_id: str) -> bool:
|
||||
"""A user fine-tune or training checkpoint (``ft:...`` / ``...:ckpt-step-N``).
|
||||
|
||||
These are account-specific artifacts: they clutter the picker, crowd out the
|
||||
foundation models (their recent ``created`` timestamps rank them first), and
|
||||
humanize into unreadable labels. Excluded from the auto-list; a user who
|
||||
wants one can still enter it via the picker's "Other" option.
|
||||
"""
|
||||
lowered = model_id.lower()
|
||||
return lowered.startswith("ft:") or ":ckpt" in lowered
|
||||
|
||||
|
||||
_SIZE_RE = re.compile(r"^\d+(?:\.\d+)?[bmk]$") # 8b, 70b, 1.5b, 120m, 32k
|
||||
_OSERIES_RE = re.compile(r"^o\d+$") # o1, o3, o4 — kept lowercase (OpenAI brand)
|
||||
|
||||
|
||||
def _humanize(model_id: str) -> str:
|
||||
"""Derive a readable label from a raw model id.
|
||||
|
||||
Best-effort only — vendor display names and the curated label map take
|
||||
precedence. Drops embedded dates and applies light casing so raw ids read
|
||||
cleanly: ``gpt-oss-120b`` → ``GPT OSS 120B``, ``qwen3-32b`` → ``Qwen3 32B``,
|
||||
``deepseek-r1:671b`` → ``DeepSeek R1 671B``, ``o3-mini`` → ``o3 Mini``.
|
||||
"""
|
||||
base = model_id.split("/")[-1]
|
||||
# Drop embedded release dates — they're noise in a label, and the picker
|
||||
# already shows the full model id alongside it.
|
||||
base = _DATE_RE.sub(" ", base)
|
||||
words: list[str] = []
|
||||
# Split on separators including ``:`` so Ollama tags (llama3.3:70b) read well.
|
||||
for part in re.split(r"[-_\s:]+", base):
|
||||
if not part:
|
||||
continue
|
||||
low = part.lower()
|
||||
if low in _ACRONYMS:
|
||||
words.append(_ACRONYMS[low])
|
||||
elif low in _BRAND_TOKENS:
|
||||
words.append(_BRAND_TOKENS[low])
|
||||
elif _SIZE_RE.match(low):
|
||||
words.append(low[:-1] + low[-1].upper()) # 70b -> 70B
|
||||
elif _OSERIES_RE.match(low):
|
||||
words.append(low) # o3 stays lowercase
|
||||
elif part[0].isalpha():
|
||||
# Capitalize the leading letter, preserve the rest (so a fused
|
||||
# family+version keeps its digits): qwen3 -> Qwen3, mini -> Mini.
|
||||
words.append(part[0].upper() + part[1:])
|
||||
else:
|
||||
words.append(part) # starts with a digit (4o, 4.1, 0905) — leave as-is
|
||||
return " ".join(words) or base
|
||||
|
||||
|
||||
# ── HTTP + parsing helpers ───────────────────────────────────────
|
||||
|
||||
|
||||
def _http_get_json(
|
||||
url: str,
|
||||
*,
|
||||
headers: dict[str, str] | None = None,
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""GET ``url`` and return parsed JSON, with a short timeout and TLS verify."""
|
||||
ssl_config = os.environ.get("SSL_CERT_FILE") or certifi.where()
|
||||
response = httpx.get(
|
||||
url,
|
||||
headers=headers,
|
||||
params=params,
|
||||
timeout=_TIMEOUT,
|
||||
verify=ssl_config,
|
||||
follow_redirects=True,
|
||||
)
|
||||
response.raise_for_status()
|
||||
result: dict[str, Any] = response.json()
|
||||
return result
|
||||
|
||||
|
||||
def _parse_iso(value: Any) -> float:
|
||||
"""Parse an ISO 8601 timestamp to an epoch float; ``0.0`` on failure."""
|
||||
if not value or not isinstance(value, str):
|
||||
return 0.0
|
||||
from datetime import datetime
|
||||
|
||||
try:
|
||||
return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp()
|
||||
except ValueError:
|
||||
return 0.0
|
||||
|
||||
|
||||
def _as_float(value: Any) -> float:
|
||||
try:
|
||||
return float(value)
|
||||
except (TypeError, ValueError):
|
||||
return 0.0
|
||||
|
||||
|
||||
def _read_json(path: Path) -> dict[str, Any] | None:
|
||||
try:
|
||||
data: dict[str, Any] = json.loads(path.read_text(encoding="utf-8"))
|
||||
return data
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
|
||||
|
||||
# ── Caching ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _cache_dir() -> Path:
|
||||
return Path.home() / ".crewai"
|
||||
|
||||
|
||||
def _catalog_cache_file() -> Path:
|
||||
return _cache_dir() / "model_catalog_cache.json"
|
||||
|
||||
|
||||
def _litellm_cache_file() -> Path:
|
||||
# Shared with crewai_cli.provider so both flows warm the same cache.
|
||||
return _cache_dir() / "provider_cache.json"
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
if provider_key == "ollama":
|
||||
return f"ollama@{_ollama_base()}"
|
||||
suffix = "key" if _provider_api_key(provider_key) else "nokey"
|
||||
return f"{provider_key}#{suffix}"
|
||||
|
||||
|
||||
def _read_catalog_cache(provider_key: str) -> list[tuple[str, str]] | None:
|
||||
"""Return a fresh cached catalog for ``provider_key``, or ``None``."""
|
||||
payload = _read_json(_catalog_cache_file())
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
entry = payload.get(_cache_key(provider_key))
|
||||
if not isinstance(entry, dict):
|
||||
return None
|
||||
# Fallback (negative) entries expire fast; dynamic ones live the full TTL.
|
||||
ttl = _NEGATIVE_TTL if entry.get("source") == "fallback" else _CATALOG_TTL
|
||||
if (time.time() - _as_float(entry.get("ts"))) >= ttl:
|
||||
return None
|
||||
models = entry.get("models")
|
||||
if not isinstance(models, list) or not models:
|
||||
return None
|
||||
try:
|
||||
return [(str(m[0]), str(m[1])) for m in models]
|
||||
except (IndexError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _write_catalog_cache(
|
||||
provider_key: str, models: list[tuple[str, str]], *, source: str
|
||||
) -> None:
|
||||
cache_file = _catalog_cache_file()
|
||||
payload = _read_json(cache_file)
|
||||
if not isinstance(payload, dict):
|
||||
payload = {}
|
||||
payload[_cache_key(provider_key)] = {
|
||||
"ts": time.time(),
|
||||
"source": source,
|
||||
"models": [[model_id, label] for model_id, label in models],
|
||||
}
|
||||
# Best-effort cache write; a failure (e.g. read-only home) is non-fatal.
|
||||
with contextlib.suppress(OSError):
|
||||
cache_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
cache_file.write_text(json.dumps(payload), encoding="utf-8")
|
||||
551
lib/cli/tests/test_model_catalog.py
Normal file
551
lib/cli/tests/test_model_catalog.py
Normal file
@@ -0,0 +1,551 @@
|
||||
"""Tests for the dynamic model catalog used by the crew-creation wizard."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
import crewai_cli.model_catalog as mc
|
||||
|
||||
_ALL_KEY_ENVS = [
|
||||
"OPENAI_API_KEY",
|
||||
"ANTHROPIC_API_KEY",
|
||||
"GEMINI_API_KEY",
|
||||
"GOOGLE_API_KEY",
|
||||
"GROQ_API_KEY",
|
||||
"CEREBRAS_API_KEY",
|
||||
"OLLAMA_API_BASE",
|
||||
"API_BASE",
|
||||
"OLLAMA_HOST",
|
||||
]
|
||||
|
||||
FALLBACK_ANTHROPIC = [
|
||||
("claude-opus-4-6", "Claude Opus 4.6"),
|
||||
("claude-sonnet-4-6", "Claude Sonnet 4.6"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolated_env(monkeypatch, tmp_path):
|
||||
"""Point the cache at a temp dir and clear provider keys for every test."""
|
||||
monkeypatch.setattr(mc, "_cache_dir", lambda: tmp_path)
|
||||
mc._reset_litellm_memo() # clear the process-level LiteLLM memo per test
|
||||
for key in _ALL_KEY_ENVS:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
|
||||
# ── version / label helpers ──────────────────────────────────────
|
||||
|
||||
|
||||
def test_version_key_parses_embedded_date():
|
||||
date_int, version = mc._version_key("claude-3-5-sonnet-20241022")
|
||||
assert date_int == 20241022
|
||||
assert version == (3, 5)
|
||||
|
||||
|
||||
def test_version_key_parses_dashed_date():
|
||||
date_int, _ = mc._version_key("gpt-4o-2024-08-06")
|
||||
assert date_int == 20240806
|
||||
|
||||
|
||||
def test_version_key_version_only():
|
||||
date_int, version = mc._version_key("claude-opus-4-6")
|
||||
assert date_int == 0
|
||||
assert version == (4, 6)
|
||||
|
||||
|
||||
def test_version_key_ranks_newer_higher():
|
||||
older = mc._version_key("claude-sonnet-4-5")
|
||||
newer = mc._version_key("claude-sonnet-4-6")
|
||||
assert newer > older
|
||||
|
||||
|
||||
def test_is_chat_model_rejects_non_chat():
|
||||
assert mc._is_chat_model("gpt-4.1-mini")
|
||||
assert not mc._is_chat_model("text-embedding-3-large")
|
||||
assert not mc._is_chat_model("whisper-1")
|
||||
assert not mc._is_chat_model("dall-e-3")
|
||||
|
||||
|
||||
def test_search_substring_not_treated_as_non_chat():
|
||||
# 'search' must not drop legitimate completion models: a token like
|
||||
# *-search-preview, or 'research' (which contains 'search' as a substring).
|
||||
assert mc._is_chat_model("gpt-4o-search-preview")
|
||||
assert mc._is_chat_model("o3-deep-research")
|
||||
# genuine non-chat markers still filter
|
||||
assert not mc._is_chat_model("text-embedding-3-large")
|
||||
|
||||
|
||||
def test_humanize():
|
||||
assert mc._humanize("gpt-4.1-mini") == "GPT 4.1 Mini"
|
||||
assert mc._humanize("anthropic/claude-opus-4-6") == "Claude Opus 4 6"
|
||||
# size suffixes uppercased, acronyms/brands cased, o-series preserved, ':' split
|
||||
assert mc._humanize("openai/gpt-oss-120b") == "GPT OSS 120B"
|
||||
assert mc._humanize("qwen/qwen3-32b") == "Qwen3 32B"
|
||||
assert mc._humanize("deepseek-r1-distill-llama-70b") == "DeepSeek R1 Distill Llama 70B"
|
||||
assert mc._humanize("o3-mini") == "o3 Mini"
|
||||
assert mc._humanize("chatgpt-4o-latest") == "ChatGPT 4o Latest"
|
||||
assert mc._humanize("llama3.3:70b") == "Llama3.3 70B"
|
||||
assert mc._humanize("gemma2-9b-it") == "Gemma2 9B IT"
|
||||
|
||||
|
||||
# ── vendor tier ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_vendor_anthropic_ranks_by_date_and_uses_display_name(monkeypatch):
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test")
|
||||
payload = {
|
||||
"data": [
|
||||
{
|
||||
"id": "claude-3-5-sonnet-20240620",
|
||||
"display_name": "Claude 3.5 Sonnet (old)",
|
||||
"created_at": "2024-06-20T00:00:00Z",
|
||||
},
|
||||
{
|
||||
"id": "claude-opus-4-6",
|
||||
"display_name": "Claude Opus 4.6",
|
||||
"created_at": "2026-02-01T00:00:00Z",
|
||||
},
|
||||
{
|
||||
"id": "claude-haiku-4-5-20251001",
|
||||
"display_name": "Claude Haiku 4.5",
|
||||
"created_at": "2025-10-01T00:00:00Z",
|
||||
},
|
||||
]
|
||||
}
|
||||
monkeypatch.setattr(mc, "_http_get_json", lambda *a, **k: payload)
|
||||
|
||||
models = mc.get_provider_models("anthropic", FALLBACK_ANTHROPIC)
|
||||
|
||||
# Newest first by created_at, display names preserved.
|
||||
assert models[0] == ("claude-opus-4-6", "Claude Opus 4.6")
|
||||
assert models[1] == ("claude-haiku-4-5-20251001", "Claude Haiku 4.5")
|
||||
assert models[2] == ("claude-3-5-sonnet-20240620", "Claude 3.5 Sonnet (old)")
|
||||
|
||||
|
||||
def test_vendor_openai_filters_non_chat_models(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
|
||||
payload = {
|
||||
"data": [
|
||||
{"id": "gpt-4.1", "created": 1_700_000_000},
|
||||
{"id": "text-embedding-3-large", "created": 1_800_000_000},
|
||||
{"id": "whisper-1", "created": 1_800_000_000},
|
||||
{"id": "gpt-5.5", "created": 1_750_000_000},
|
||||
]
|
||||
}
|
||||
monkeypatch.setattr(mc, "_http_get_json", lambda *a, **k: payload)
|
||||
|
||||
models = mc.get_provider_models("openai", [])
|
||||
ids = [m for m, _ in models]
|
||||
|
||||
assert ids == ["gpt-5.5", "gpt-4.1"] # embeddings/whisper dropped, newest first
|
||||
|
||||
|
||||
def test_vendor_gemini_requires_generate_content(monkeypatch):
|
||||
monkeypatch.setenv("GEMINI_API_KEY", "key")
|
||||
payload = {
|
||||
"models": [
|
||||
{
|
||||
"name": "models/gemini-2.5-pro",
|
||||
"displayName": "Gemini 2.5 Pro",
|
||||
"supportedGenerationMethods": ["generateContent"],
|
||||
},
|
||||
{
|
||||
"name": "models/text-embedding-004",
|
||||
"displayName": "Embedding",
|
||||
"supportedGenerationMethods": ["embedContent"],
|
||||
},
|
||||
{
|
||||
"name": "models/gemini-1.5-pro",
|
||||
"displayName": "Gemini 1.5 Pro",
|
||||
"supportedGenerationMethods": ["generateContent"],
|
||||
},
|
||||
]
|
||||
}
|
||||
monkeypatch.setattr(mc, "_http_get_json", lambda *a, **k: payload)
|
||||
|
||||
models = mc.get_provider_models("gemini", [])
|
||||
ids = [m for m, _ in models]
|
||||
|
||||
# "models/" prefix stripped, embedding excluded, newer version first.
|
||||
assert ids == ["gemini-2.5-pro", "gemini-1.5-pro"]
|
||||
|
||||
|
||||
def test_openai_excludes_fine_tunes_and_checkpoints(monkeypatch):
|
||||
# Fine-tunes/checkpoints have recent `created` timestamps and would otherwise
|
||||
# crowd out (and rank above) the base models — they must be excluded so the
|
||||
# picker shows clean foundation models.
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
|
||||
payload = {
|
||||
"data": [
|
||||
{"id": "ft:gpt-4o-mini-2024-07-18:crewai::DyJG86uF", "created": 1_900_000_000},
|
||||
{
|
||||
"id": "ft:gpt-4o-mini-2024-07-18:crewai::DyJG7Q9N:ckpt-step-84",
|
||||
"created": 1_900_000_001,
|
||||
},
|
||||
{"id": "gpt-5.5", "created": 1_800_000_000},
|
||||
{"id": "gpt-4.1", "created": 1_700_000_000},
|
||||
]
|
||||
}
|
||||
monkeypatch.setattr(mc, "_http_get_json", lambda *a, **k: payload)
|
||||
|
||||
ids = [m for m, _ in mc.get_provider_models("openai", [])]
|
||||
assert ids == ["gpt-5.5", "gpt-4.1"] # fine-tunes + checkpoints dropped
|
||||
|
||||
|
||||
def test_vendor_gemini_paginates(monkeypatch):
|
||||
monkeypatch.setenv("GEMINI_API_KEY", "key")
|
||||
pages = {
|
||||
None: {
|
||||
"models": [
|
||||
{
|
||||
"name": "models/gemini-3.5-flash",
|
||||
"displayName": "Gemini 3.5 Flash",
|
||||
"supportedGenerationMethods": ["generateContent"],
|
||||
}
|
||||
],
|
||||
"nextPageToken": "p2",
|
||||
},
|
||||
"p2": {
|
||||
"models": [
|
||||
{
|
||||
"name": "models/gemini-2.5-pro",
|
||||
"displayName": "Gemini 2.5 Pro",
|
||||
"supportedGenerationMethods": ["generateContent"],
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
|
||||
def fetch(url, headers=None, params=None):
|
||||
return pages[(params or {}).get("pageToken")]
|
||||
|
||||
monkeypatch.setattr(mc, "_http_get_json", fetch)
|
||||
|
||||
ids = sorted(m for m, _ in mc.get_provider_models("gemini", []))
|
||||
# Both pages contributed (newest-first ranking is _finalize's job).
|
||||
assert ids == ["gemini-2.5-pro", "gemini-3.5-flash"]
|
||||
|
||||
|
||||
def test_vendor_gemini_first_page_error_uses_fallback(monkeypatch):
|
||||
# A total (first-page) Gemini failure with a key set must fall back to the
|
||||
# curated list, not be mistaken for a successful empty result.
|
||||
monkeypatch.setenv("GEMINI_API_KEY", "key")
|
||||
|
||||
def boom(*a, **k):
|
||||
raise RuntimeError("gemini down")
|
||||
|
||||
monkeypatch.setattr(mc, "_http_get_json", boom)
|
||||
models = mc.get_provider_models("gemini", [("gemini-x", "Gemini X")])
|
||||
assert models == [("gemini-x", "Gemini X")]
|
||||
|
||||
|
||||
def test_vendor_gemini_keeps_partial_on_later_page_error(monkeypatch):
|
||||
monkeypatch.setenv("GEMINI_API_KEY", "key")
|
||||
|
||||
def fetch(url, headers=None, params=None):
|
||||
if (params or {}).get("pageToken"):
|
||||
raise RuntimeError("page 2 down")
|
||||
return {
|
||||
"models": [
|
||||
{
|
||||
"name": "models/gemini-3.5-flash",
|
||||
"displayName": "Gemini 3.5 Flash",
|
||||
"supportedGenerationMethods": ["generateContent"],
|
||||
}
|
||||
],
|
||||
"nextPageToken": "p2",
|
||||
}
|
||||
|
||||
monkeypatch.setattr(mc, "_http_get_json", fetch)
|
||||
|
||||
# Page-1 models are kept; the later-page error doesn't force the fallback.
|
||||
models = mc.get_provider_models("gemini", [("fallback-x", "Fallback X")])
|
||||
assert [m for m, _ in models] == ["gemini-3.5-flash"]
|
||||
|
||||
|
||||
def test_ollama_empty_response_not_filled_with_fallback(monkeypatch):
|
||||
# A reachable Ollama with nothing installed -> empty (manual entry), not the
|
||||
# curated suggestions the crew can't actually run.
|
||||
monkeypatch.setattr(mc, "_http_get_json", lambda *a, **k: {"models": []})
|
||||
assert mc.get_provider_models("ollama", [("llama3.3", "Llama 3.3")]) == []
|
||||
|
||||
|
||||
def test_ollama_unreachable_uses_fallback(monkeypatch):
|
||||
# Server down (fetch raises) is different from empty -> fall back to suggestions.
|
||||
def boom(*a, **k):
|
||||
raise RuntimeError("connection refused")
|
||||
|
||||
monkeypatch.setattr(mc, "_http_get_json", boom)
|
||||
models = mc.get_provider_models("ollama", [("llama3.3", "Llama 3.3")])
|
||||
assert models == [("llama3.3", "Llama 3.3")]
|
||||
|
||||
|
||||
def test_ollama_excludes_embedding_models(monkeypatch):
|
||||
# /api/tags lists everything installed, including embeddings — filter them.
|
||||
monkeypatch.setattr(
|
||||
mc,
|
||||
"_http_get_json",
|
||||
lambda *a, **k: {
|
||||
"models": [
|
||||
{"model": "llama3.3:70b"},
|
||||
{"model": "nomic-embed-text"},
|
||||
{"model": "mxbai-embed-large"},
|
||||
]
|
||||
},
|
||||
)
|
||||
ids = [m for m, _ in mc.get_provider_models("ollama", [])]
|
||||
assert ids == ["llama3.3:70b"]
|
||||
|
||||
|
||||
def test_ollama_base_honors_ollama_host(monkeypatch):
|
||||
# OLLAMA_HOST (scheme-less runtime convention) is resolved with a scheme.
|
||||
monkeypatch.setenv("OLLAMA_HOST", "10.0.0.5:11434")
|
||||
assert mc._ollama_base() == "http://10.0.0.5:11434"
|
||||
|
||||
|
||||
def test_ollama_recovery_not_blocked_by_negative_cache(monkeypatch):
|
||||
# Ollama down -> fallback, but not negatively cached; once the server is up
|
||||
# the next call fetches live models rather than serving suggestions.
|
||||
calls = {"n": 0}
|
||||
|
||||
def flaky(*a, **k):
|
||||
calls["n"] += 1
|
||||
if calls["n"] == 1:
|
||||
raise RuntimeError("connection refused")
|
||||
return {"models": [{"model": "llama-installed"}]}
|
||||
|
||||
monkeypatch.setattr(mc, "_http_get_json", flaky)
|
||||
first = mc.get_provider_models("ollama", [("llama3.3", "Llama 3.3")])
|
||||
assert first == [("llama3.3", "Llama 3.3")] # down -> fallback (not cached)
|
||||
second = mc.get_provider_models("ollama", [("llama3.3", "Llama 3.3")])
|
||||
assert [m for m, _ in second] == ["llama-installed"] # recovered live
|
||||
|
||||
|
||||
def test_gemini_honors_google_api_key(monkeypatch):
|
||||
# GOOGLE_API_KEY (equivalent to GEMINI_API_KEY in crewai) enables the live tier.
|
||||
monkeypatch.setenv("GOOGLE_API_KEY", "key")
|
||||
monkeypatch.setattr(
|
||||
mc,
|
||||
"_http_get_json",
|
||||
lambda *a, **k: {
|
||||
"models": [
|
||||
{
|
||||
"name": "models/gemini-3.5-flash",
|
||||
"displayName": "Gemini 3.5 Flash",
|
||||
"supportedGenerationMethods": ["generateContent"],
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
models = mc.get_provider_models("gemini", [("gemini-x", "Gemini X")])
|
||||
assert [m for m, _ in models] == ["gemini-3.5-flash"] # live, not fallback
|
||||
|
||||
|
||||
def test_curated_label_overrides_raw_vendor_label(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
|
||||
payload = {"data": [{"id": "gpt-5.5", "created": 1}]}
|
||||
monkeypatch.setattr(mc, "_http_get_json", lambda *a, **k: payload)
|
||||
|
||||
models = mc.get_provider_models("openai", [("gpt-5.5", "GPT-5.5 (curated)")])
|
||||
assert models == [("gpt-5.5", "GPT-5.5 (curated)")]
|
||||
|
||||
|
||||
def test_truncates_to_max_models(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
|
||||
payload = {
|
||||
"data": [{"id": f"gpt-test-{i}", "created": i} for i in range(20)]
|
||||
}
|
||||
monkeypatch.setattr(mc, "_http_get_json", lambda *a, **k: payload)
|
||||
|
||||
models = mc.get_provider_models("openai", [])
|
||||
assert len(models) == mc.MAX_MODELS
|
||||
|
||||
|
||||
# ── litellm tier ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_litellm_tier_for_uncurated_provider(monkeypatch):
|
||||
# A provider with no curated fallback ([]) -> the LiteLLM feed is consulted.
|
||||
litellm_data = {
|
||||
"claude-opus-4-6": {"litellm_provider": "anthropic", "mode": "chat"},
|
||||
"claude-sonnet-4-5": {"litellm_provider": "anthropic", "mode": "chat"},
|
||||
"voyage-embed": {"litellm_provider": "anthropic", "mode": "embedding"},
|
||||
"gpt-4.1": {"litellm_provider": "openai", "mode": "chat"},
|
||||
}
|
||||
mc._litellm_cache_file().write_text(json.dumps(litellm_data), encoding="utf-8")
|
||||
|
||||
models = mc.get_provider_models("anthropic", []) # empty == uncurated
|
||||
ids = [m for m, _ in models]
|
||||
|
||||
# Only anthropic chat models, embedding + other providers excluded.
|
||||
assert ids == ["claude-opus-4-6", "claude-sonnet-4-5"]
|
||||
|
||||
|
||||
def test_null_litellm_provider_does_not_crash(monkeypatch):
|
||||
# A present-but-null litellm_provider must be skipped, not raise.
|
||||
litellm_data = {
|
||||
"weird-model": {"litellm_provider": None, "mode": "chat"},
|
||||
"anthropic.claude-v2": {"litellm_provider": "bedrock", "mode": "chat"},
|
||||
}
|
||||
mc._litellm_cache_file().write_text(json.dumps(litellm_data), encoding="utf-8")
|
||||
|
||||
models = mc.get_provider_models("bedrock", [])
|
||||
assert [m for m, _ in models] == ["anthropic.claude-v2"]
|
||||
|
||||
|
||||
def test_litellm_strips_provider_prefix(monkeypatch):
|
||||
litellm_data = {
|
||||
"gemini/gemini-1.5-pro": {"litellm_provider": "gemini", "mode": "chat"},
|
||||
}
|
||||
mc._litellm_cache_file().write_text(json.dumps(litellm_data), encoding="utf-8")
|
||||
|
||||
models = mc.get_provider_models("gemini", [])
|
||||
assert models == [("gemini-1.5-pro", "Gemini 1.5 Pro")]
|
||||
|
||||
|
||||
# ── fallback + caching ───────────────────────────────────────────
|
||||
|
||||
|
||||
def test_falls_back_when_everything_fails(monkeypatch):
|
||||
# No key, no litellm cache, network raises -> curated fallback verbatim.
|
||||
def boom(*a, **k):
|
||||
raise RuntimeError("network down")
|
||||
|
||||
monkeypatch.setattr(mc, "_http_get_json", boom)
|
||||
models = mc.get_provider_models("anthropic", FALLBACK_ANTHROPIC)
|
||||
assert models == FALLBACK_ANTHROPIC
|
||||
|
||||
|
||||
def test_result_is_cached(monkeypatch):
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test")
|
||||
calls = {"n": 0}
|
||||
|
||||
def fetch(*a, **k):
|
||||
calls["n"] += 1
|
||||
return {"data": [{"id": "claude-opus-4-6", "created_at": "2026-01-01T00:00:00Z"}]}
|
||||
|
||||
monkeypatch.setattr(mc, "_http_get_json", fetch)
|
||||
|
||||
first = mc.get_provider_models("anthropic", FALLBACK_ANTHROPIC)
|
||||
# Second call must hit the cache and not touch the network again.
|
||||
monkeypatch.setattr(mc, "_http_get_json", lambda *a, **k: pytest.fail("refetched"))
|
||||
second = mc.get_provider_models("anthropic", FALLBACK_ANTHROPIC)
|
||||
|
||||
assert first == second
|
||||
assert calls["n"] == 1
|
||||
|
||||
|
||||
def test_curated_fallback_preferred_over_litellm(monkeypatch):
|
||||
# The feed lags real releases, so a non-empty curated fallback must win even
|
||||
# when a fresh LiteLLM cache is present (regression: Anthropic's feed lacked
|
||||
# Fable 5 / Opus 4.8 / Sonnet 5).
|
||||
monkeypatch.setattr(mc, "_http_get_json", lambda *a, **k: pytest.fail("no net"))
|
||||
litellm_data = {
|
||||
"claude-opus-4-6": {"litellm_provider": "anthropic", "mode": "chat"},
|
||||
}
|
||||
mc._litellm_cache_file().write_text(json.dumps(litellm_data), encoding="utf-8")
|
||||
|
||||
models = mc.get_provider_models("anthropic", FALLBACK_ANTHROPIC)
|
||||
assert models == FALLBACK_ANTHROPIC
|
||||
|
||||
|
||||
def test_added_key_bypasses_negative_cache(monkeypatch):
|
||||
# A no-key call negatively-caches the fallback; adding a key afterwards must
|
||||
# fetch live models rather than serve the cached fallback (distinct cache key).
|
||||
first = mc.get_provider_models("openai", [("gpt-x", "GPT X")])
|
||||
assert first == [("gpt-x", "GPT X")] # no key -> fallback
|
||||
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
|
||||
monkeypatch.setattr(
|
||||
mc, "_http_get_json", lambda *a, **k: {"data": [{"id": "gpt-5.5", "created": 1}]}
|
||||
)
|
||||
second = mc.get_provider_models("openai", [("gpt-x", "GPT X")])
|
||||
assert [m for m, _ in second] == ["gpt-5.5"] # live fetch, not cached fallback
|
||||
|
||||
|
||||
def test_invalid_litellm_cache_falls_through_to_download(monkeypatch):
|
||||
# A corrupt-but-fresh cache must neither crash the picker nor block a
|
||||
# recoverable download — it falls through and refetches.
|
||||
mc._litellm_cache_file().write_text("[1, 2, 3]", encoding="utf-8")
|
||||
monkeypatch.setattr(
|
||||
mc,
|
||||
"_http_get_json",
|
||||
lambda *a, **k: {
|
||||
"anthropic.claude-v2": {"litellm_provider": "bedrock", "mode": "chat"}
|
||||
},
|
||||
)
|
||||
models = mc.get_provider_models("bedrock", [])
|
||||
assert [m for m, _ in models] == ["anthropic.claude-v2"] # recovered via download
|
||||
|
||||
|
||||
def test_litellm_fetch_attempted_once_per_process(monkeypatch):
|
||||
# With no cache and a failing download, the feed is fetched at most once per
|
||||
# process — repeated lookups (across providers) must not re-hit the network.
|
||||
calls = {"n": 0}
|
||||
|
||||
def boom(*a, **k):
|
||||
calls["n"] += 1
|
||||
raise RuntimeError("offline")
|
||||
|
||||
monkeypatch.setattr(mc, "_http_get_json", boom)
|
||||
mc.get_provider_models("bedrock", [])
|
||||
mc.get_provider_models("azure", [])
|
||||
assert calls["n"] == 1 # memoized after the first failed attempt
|
||||
|
||||
|
||||
def test_litellm_fills_uncurated_bedrock(monkeypatch):
|
||||
# No vendor fetcher and no curated fallback -> LiteLLM feed fills the gap.
|
||||
monkeypatch.setattr(mc, "_http_get_json", lambda *a, **k: pytest.fail("no net"))
|
||||
litellm_data = {
|
||||
"anthropic.claude-v2": {"litellm_provider": "bedrock", "mode": "chat"},
|
||||
}
|
||||
mc._litellm_cache_file().write_text(json.dumps(litellm_data), encoding="utf-8")
|
||||
|
||||
models = mc.get_provider_models("bedrock", [])
|
||||
assert models == [("anthropic.claude-v2", "Anthropic.claude V2")]
|
||||
|
||||
|
||||
def test_failed_fetch_is_negatively_cached(monkeypatch):
|
||||
# A failed vendor fetch must not be retried on every call — the fallback is
|
||||
# cached briefly so the picker doesn't re-hit the timeout-prone endpoint.
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test")
|
||||
calls = {"n": 0}
|
||||
|
||||
def boom(*a, **k):
|
||||
calls["n"] += 1
|
||||
raise RuntimeError("down")
|
||||
|
||||
monkeypatch.setattr(mc, "_http_get_json", boom)
|
||||
first = mc.get_provider_models("anthropic", FALLBACK_ANTHROPIC)
|
||||
second = mc.get_provider_models("anthropic", FALLBACK_ANTHROPIC)
|
||||
|
||||
assert first == second == FALLBACK_ANTHROPIC
|
||||
assert calls["n"] == 1 # second call served from the negative cache
|
||||
|
||||
|
||||
def test_bad_cache_json_does_not_crash(monkeypatch):
|
||||
# A corrupt cache whose root is not a mapping must not raise (get_provider_models
|
||||
# is documented to never raise).
|
||||
mc._catalog_cache_file().write_text("[1, 2, 3]", encoding="utf-8")
|
||||
|
||||
models = mc.get_provider_models("anthropic", FALLBACK_ANTHROPIC)
|
||||
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"}]}
|
||||
)
|
||||
first = mc.get_provider_models("ollama", [])
|
||||
assert [m for m, _ in first] == ["llama-a"]
|
||||
|
||||
monkeypatch.setenv("OLLAMA_API_BASE", "http://host-b:11434")
|
||||
monkeypatch.setattr(
|
||||
mc, "_http_get_json", lambda *a, **k: {"models": [{"model": "llama-b"}]}
|
||||
)
|
||||
second = mc.get_provider_models("ollama", [])
|
||||
assert [m for m, _ in second] == ["llama-b"] # not the host-a cache
|
||||
Reference in New Issue
Block a user