refactor(model-connections): consolidate provider capability handling

This commit is contained in:
Anish Sarkar 2026-06-11 18:21:07 +05:30
parent c6a25cc1fe
commit 8f20a32571
11 changed files with 64 additions and 311 deletions

View file

@ -19,7 +19,7 @@ def _connection_key(conn: dict[str, Any]) -> tuple[Any, ...]:
# the same provider/base can have different quota/rate limits upstream.
return (
conn.get("protocol"),
conn.get("native_provider"),
conn.get("litellm_provider"),
conn.get("base_url"),
conn.get("api_key"),
_freeze(conn.get("extra") or {}),

View file

@ -323,10 +323,10 @@ def _generate_configs(
"seo_enabled": seo_enabled,
"seo_slug": None,
"quota_reserve_tokens": quota_reserve_tokens,
"provider": "OPENROUTER",
"litellm_provider": "openrouter",
"model_name": model_id,
"api_key": api_key,
"api_base": "",
"api_base": "https://openrouter.ai/api/v1",
"rpm": free_rpm if tier == "free" else rpm,
"tpm": free_tpm if tier == "free" else tpm,
"litellm_params": dict(litellm_params),
@ -420,14 +420,9 @@ def _generate_image_gen_configs(
"id": _stable_config_id(model_id, id_offset, taken),
"name": name,
"description": f"{name} via OpenRouter (image generation)",
"provider": "OPENROUTER",
"litellm_provider": "openrouter",
"model_name": model_id,
"api_key": api_key,
# Pin to OpenRouter's public base URL so a downstream call site
# that forgets ``resolve_api_base`` still doesn't inherit
# ``AZURE_OPENAI_ENDPOINT`` and 404 on
# ``image_generation/transformation`` (defense-in-depth, see
# ``provider_api_base`` docstring).
"api_base": "https://openrouter.ai/api/v1",
"api_version": None,
"rpm": free_rpm if tier == "free" else rpm,
@ -504,13 +499,9 @@ def _generate_vision_llm_configs(
"id": _stable_config_id(model_id, id_offset, taken),
"name": name,
"description": f"{name} via OpenRouter (vision)",
"provider": "OPENROUTER",
"litellm_provider": "openrouter",
"model_name": model_id,
"api_key": api_key,
# Pin to OpenRouter's public base URL so a downstream call site
# that forgets ``resolve_api_base`` still doesn't inherit
# ``AZURE_OPENAI_ENDPOINT`` (defense-in-depth, see
# ``provider_api_base`` docstring).
"api_base": "https://openrouter.ai/api/v1",
"api_version": None,
"rpm": free_rpm if tier == "free" else rpm,
@ -710,7 +701,7 @@ class OpenRouterIntegrationService:
)
# Re-blend health scores against the freshly fetched catalogue. Also
# re-stamps health for any YAML-curated cfg with provider==OPENROUTER
# re-stamps health for any YAML-curated cfg with litellm_provider=openrouter
# so a hand-picked dead OR model is gated like a dynamic one.
await self._enrich_health_safely(static_configs + new_configs, log_summary=True)
@ -787,7 +778,7 @@ class OpenRouterIntegrationService:
the entire previous cycle's cache for this run.
"""
or_cfgs = [
c for c in configs if str(c.get("provider", "")).upper() == "OPENROUTER"
c for c in configs if str(c.get("litellm_provider", "")).lower() == "openrouter"
]
if not or_cfgs:
return

View file

@ -143,12 +143,12 @@ def _register_chat_shape_configs(
sample_keys: list[str] = []
for cfg in configs:
provider = str(cfg.get("provider") or "").upper()
provider = str(cfg.get("litellm_provider") or "").lower()
model_name = str(cfg.get("model_name") or "").strip()
litellm_params = cfg.get("litellm_params") or {}
base_model = str(litellm_params.get("base_model") or model_name).strip()
if provider == "OPENROUTER":
if provider == "openrouter":
entry = or_pricing.get(model_name)
if entry:
input_cost = _safe_float(entry.get("prompt"))
@ -189,12 +189,11 @@ def _register_chat_shape_configs(
skipped_no_pricing += 1
continue
aliases = _alias_set_for_yaml(provider, model_name, base_model)
provider_slug = "azure" if provider == "AZURE_OPENAI" else provider.lower()
count = _register(
aliases,
input_cost=input_cost,
output_cost=output_cost,
provider=provider_slug,
provider=provider,
)
if count > 0:
registered_models += 1

View file

@ -1,106 +0,0 @@
"""Provider-aware ``api_base`` resolution shared by chat / image-gen / vision.
LiteLLM falls back to the module-global ``litellm.api_base`` when an
individual call doesn't pass one, which silently inherits provider-agnostic
env vars like ``AZURE_OPENAI_ENDPOINT`` / ``OPENAI_API_BASE``. Without an
explicit ``api_base``, an ``openrouter/<model>`` request can end up at an
Azure endpoint and 404 with ``Resource not found`` (real reproducer:
[litellm/llms/openrouter/image_generation/transformation.py:242-263] appends
``/chat/completions`` to whatever inherited base it gets, regardless of
provider).
The chat router has had this defense for a while
(``llm_router_service.py:466-478``). This module hoists the maps + cascade
into a tiny standalone helper so vision and image-gen can share the same
source of truth without an inter-service circular import.
"""
from __future__ import annotations
PROVIDER_DEFAULT_API_BASE: dict[str, str] = {
"openrouter": "https://openrouter.ai/api/v1",
"groq": "https://api.groq.com/openai/v1",
"mistral": "https://api.mistral.ai/v1",
"perplexity": "https://api.perplexity.ai",
"xai": "https://api.x.ai/v1",
"cerebras": "https://api.cerebras.ai/v1",
"deepinfra": "https://api.deepinfra.com/v1/openai",
"fireworks_ai": "https://api.fireworks.ai/inference/v1",
"together_ai": "https://api.together.xyz/v1",
"anyscale": "https://api.endpoints.anyscale.com/v1",
"cometapi": "https://api.cometapi.com/v1",
"sambanova": "https://api.sambanova.ai/v1",
}
"""Default ``api_base`` per LiteLLM provider prefix (lowercase).
Only providers with a well-known, stable public base URL are listed
self-hosted / BYO-endpoint providers (ollama, custom, bedrock, vertex_ai,
huggingface, databricks, cloudflare, replicate) are intentionally omitted
so their existing config-driven behaviour is preserved."""
PROVIDER_KEY_DEFAULT_API_BASE: dict[str, str] = {
"DEEPSEEK": "https://api.deepseek.com/v1",
"ALIBABA_QWEN": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
"MOONSHOT": "https://api.moonshot.ai/v1",
"ZHIPU": "https://open.bigmodel.cn/api/paas/v4",
"MINIMAX": "https://api.minimax.io/v1",
}
"""Canonical provider key (uppercase) → base URL.
Used when the LiteLLM provider prefix is the generic ``openai`` shim but the
config's ``provider`` field tells us which API it actually is (DeepSeek,
Alibaba, Moonshot, Zhipu, MiniMax all use the ``openai`` prefix but each
has its own base URL)."""
def resolve_api_base(
*,
provider: str | None,
provider_prefix: str | None,
config_api_base: str | None,
) -> str | None:
"""Resolve a non-Azure-leaking ``api_base`` for a deployment.
Cascade (first non-empty wins):
1. The config's own ``api_base`` (whitespace-only treated as missing).
2. ``PROVIDER_KEY_DEFAULT_API_BASE[provider.upper()]``.
3. ``PROVIDER_DEFAULT_API_BASE[provider_prefix.lower()]``.
4. ``None`` caller should NOT set ``api_base`` and let the LiteLLM
provider integration apply its own default (e.g. AzureOpenAI's
deployment-derived URL, custom provider's per-deployment URL).
Args:
provider: The config's ``provider`` field (e.g. ``"OPENROUTER"``,
``"DEEPSEEK"``). Case-insensitive.
provider_prefix: The LiteLLM model-string prefix the same call
site builds for the model id (e.g. ``"openrouter"``,
``"groq"``). Case-insensitive.
config_api_base: ``api_base`` from the global YAML / DB row /
OpenRouter dynamic config. Empty / whitespace-only means
"missing" the resolver still applies the cascade.
Returns:
A URL string, or ``None`` if no default applies for this provider.
"""
if config_api_base and config_api_base.strip():
return config_api_base
if provider:
key_default = PROVIDER_KEY_DEFAULT_API_BASE.get(provider.upper())
if key_default:
return key_default
if provider_prefix:
prefix_default = PROVIDER_DEFAULT_API_BASE.get(provider_prefix.lower())
if prefix_default:
return prefix_default
return None
__all__ = [
"PROVIDER_DEFAULT_API_BASE",
"PROVIDER_KEY_DEFAULT_API_BASE",
"resolve_api_base",
]

View file

@ -46,26 +46,12 @@ from collections.abc import Iterable
import litellm
from app.services.model_resolver import NATIVE_PROVIDER_PREFIX
logger = logging.getLogger(__name__)
# Provider-name → LiteLLM model-prefix map.
#
# Owned here because ``app.services.provider_capabilities`` is the
# only edge that's safe to call from ``app.config``'s YAML loader at
# class-body init time. ``app.agents.chat.runtime.llm_config`` re-exports
# this constant under the historical ``PROVIDER_MAP`` name; placing the
# map there directly would re-introduce the
# ``app.config -> ... -> deliverables/tools/generate_image ->
# app.config`` cycle that prompted the move.
_PROVIDER_PREFIX_MAP = NATIVE_PROVIDER_PREFIX
def _candidate_model_strings(
*,
provider: str | None,
litellm_provider: str | None,
model_name: str | None,
base_model: str | None,
custom_provider: str | None,
@ -92,12 +78,7 @@ def _candidate_model_strings(
seen.add(key)
candidates.append(key)
provider_prefix: str | None = None
if provider:
provider_prefix = _PROVIDER_PREFIX_MAP.get(provider.upper(), provider.lower())
if custom_provider:
# ``custom_provider`` overrides everything for CUSTOM/proxy setups.
provider_prefix = custom_provider
provider_prefix = custom_provider or litellm_provider
primary_model = base_model or model_name
bare_model = model_name
@ -132,7 +113,7 @@ def _candidate_model_strings(
def derive_supports_image_input(
*,
provider: str | None = None,
litellm_provider: str | None = None,
model_name: str | None = None,
base_model: str | None = None,
custom_provider: str | None = None,
@ -166,7 +147,7 @@ def derive_supports_image_input(
return False
for model_string, custom_llm_provider in _candidate_model_strings(
provider=provider,
litellm_provider=litellm_provider,
model_name=model_name,
base_model=base_model,
custom_provider=custom_provider,
@ -191,7 +172,7 @@ def derive_supports_image_input(
def is_known_text_only_chat_model(
*,
provider: str | None = None,
litellm_provider: str | None = None,
model_name: str | None = None,
base_model: str | None = None,
custom_provider: str | None = None,
@ -212,7 +193,7 @@ def is_known_text_only_chat_model(
leads to the regression we're fixing here.
"""
for model_string, custom_llm_provider in _candidate_model_strings(
provider=provider,
litellm_provider=litellm_provider,
model_name=model_name,
base_model=base_model,
custom_provider=custom_provider,

View file

@ -108,25 +108,23 @@ PROVIDER_PRESTIGE_OR: dict[str, int] = {
# YAML provider field (the upstream API shape the operator selected).
PROVIDER_PRESTIGE_YAML: dict[str, int] = {
"AZURE_OPENAI": 50,
"OPENAI": 50,
"ANTHROPIC": 50,
"GOOGLE": 50,
"VERTEX_AI": 50,
"GEMINI": 50,
"XAI": 50,
"MISTRAL": 38,
"DEEPSEEK": 38,
"COHERE": 38,
"GROQ": 30,
"TOGETHER_AI": 28,
"FIREWORKS_AI": 28,
"PERPLEXITY": 28,
"MINIMAX": 28,
"BEDROCK": 28,
"OPENROUTER": 25,
"OLLAMA": 12,
"CUSTOM": 12,
"azure": 50,
"openai": 50,
"anthropic": 50,
"gemini": 50,
"vertex_ai": 50,
"xai": 50,
"mistral": 38,
"deepseek": 38,
"cohere": 38,
"groq": 30,
"together_ai": 28,
"fireworks_ai": 28,
"perplexity": 28,
"bedrock": 28,
"openrouter": 25,
"ollama_chat": 12,
"custom": 12,
}
@ -275,7 +273,7 @@ def static_score_yaml(cfg: dict) -> int:
listed this model. Pricing / context fall through to lazy ``litellm``
lookups; failures are silent (we just lose those sub-points).
"""
provider = str(cfg.get("provider", "")).upper()
provider = str(cfg.get("litellm_provider", "")).lower()
base = PROVIDER_PRESTIGE_YAML.get(provider, 15)
model_name = cfg.get("model_name") or ""