mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-06-24 21:38:09 +02:00
refactor(model-connections): consolidate provider capability handling
This commit is contained in:
parent
c6a25cc1fe
commit
8f20a32571
11 changed files with 64 additions and 311 deletions
|
|
@ -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 {}),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 ""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue