apunkt/SurfSense
1
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-06-24 21:38:09 +02:00
Code Issues Projects Releases Packages Wiki Activity

refactor(model-connections): consolidate provider capability handling

Browse source
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
Download patch file Download diff file Expand all files Collapse all files

2
surfsense_backend/app/services/global_model_catalog.py
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. # the same provider/base can have different quota/rate limits upstream.
return ( return (
conn.get("protocol"), conn.get("protocol"),
conn.get("native_provider"), conn.get("litellm_provider"),
conn.get("base_url"), conn.get("base_url"),
conn.get("api_key"), conn.get("api_key"),
_freeze(conn.get("extra") or {}), _freeze(conn.get("extra") or {}),

21
surfsense_backend/app/services/openrouter_integration_service.py
View file

@ -323,10 +323,10 @@ def _generate_configs(
"seo_enabled": seo_enabled, "seo_enabled": seo_enabled,
"seo_slug": None, "seo_slug": None,
"quota_reserve_tokens": quota_reserve_tokens, "quota_reserve_tokens": quota_reserve_tokens,
"provider": "OPENROUTER", "litellm_provider": "openrouter",
"model_name": model_id, "model_name": model_id,
"api_key": api_key, "api_key": api_key,
"api_base": "", "api_base": "https://openrouter.ai/api/v1",
"rpm": free_rpm if tier == "free" else rpm, "rpm": free_rpm if tier == "free" else rpm,
"tpm": free_tpm if tier == "free" else tpm, "tpm": free_tpm if tier == "free" else tpm,
"litellm_params": dict(litellm_params), "litellm_params": dict(litellm_params),
@ -420,14 +420,9 @@ def _generate_image_gen_configs(
"id": _stable_config_id(model_id, id_offset, taken), "id": _stable_config_id(model_id, id_offset, taken),
"name": name, "name": name,
"description": f"{name} via OpenRouter (image generation)", "description": f"{name} via OpenRouter (image generation)",
"provider": "OPENROUTER", "litellm_provider": "openrouter",
"model_name": model_id, "model_name": model_id,
"api_key": api_key, "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_base": "https://openrouter.ai/api/v1",
"api_version": None, "api_version": None,
"rpm": free_rpm if tier == "free" else rpm, "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), "id": _stable_config_id(model_id, id_offset, taken),
"name": name, "name": name,
"description": f"{name} via OpenRouter (vision)", "description": f"{name} via OpenRouter (vision)",
"provider": "OPENROUTER", "litellm_provider": "openrouter",
"model_name": model_id, "model_name": model_id,
"api_key": api_key, "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_base": "https://openrouter.ai/api/v1",
"api_version": None, "api_version": None,
"rpm": free_rpm if tier == "free" else rpm, "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-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. # 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) 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. the entire previous cycle's cache for this run.
""" """
or_cfgs = [ 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: if not or_cfgs:
return return

7
surfsense_backend/app/services/pricing_registration.py
View file

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

106
surfsense_backend/app/services/provider_api_base.py
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",
]

31
surfsense_backend/app/services/provider_capabilities.py
View file

@ -46,26 +46,12 @@ from collections.abc import Iterable
import litellm import litellm
from app.services.model_resolver import NATIVE_PROVIDER_PREFIX
logger = logging.getLogger(__name__) 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( def _candidate_model_strings(
*, *,
provider: str | None, litellm_provider: str | None,
model_name: str | None, model_name: str | None,
base_model: str | None, base_model: str | None,
custom_provider: str | None, custom_provider: str | None,
@ -92,12 +78,7 @@ def _candidate_model_strings(
seen.add(key) seen.add(key)
candidates.append(key) candidates.append(key)
provider_prefix: str | None = None provider_prefix = custom_provider or litellm_provider
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
primary_model = base_model or model_name primary_model = base_model or model_name
bare_model = model_name bare_model = model_name
@ -132,7 +113,7 @@ def _candidate_model_strings(
def derive_supports_image_input( def derive_supports_image_input(
*, *,
provider: str | None = None, litellm_provider: str | None = None,
model_name: str | None = None, model_name: str | None = None,
base_model: str | None = None, base_model: str | None = None,
custom_provider: str | None = None, custom_provider: str | None = None,
@ -166,7 +147,7 @@ def derive_supports_image_input(
return False return False
for model_string, custom_llm_provider in _candidate_model_strings( for model_string, custom_llm_provider in _candidate_model_strings(
provider=provider, litellm_provider=litellm_provider,
model_name=model_name, model_name=model_name,
base_model=base_model, base_model=base_model,
custom_provider=custom_provider, custom_provider=custom_provider,
@ -191,7 +172,7 @@ def derive_supports_image_input(
def is_known_text_only_chat_model( def is_known_text_only_chat_model(
*, *,
provider: str | None = None, litellm_provider: str | None = None,
model_name: str | None = None, model_name: str | None = None,
base_model: str | None = None, base_model: str | None = None,
custom_provider: 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. leads to the regression we're fixing here.
""" """
for model_string, custom_llm_provider in _candidate_model_strings( for model_string, custom_llm_provider in _candidate_model_strings(
provider=provider, litellm_provider=litellm_provider,
model_name=model_name, model_name=model_name,
base_model=base_model, base_model=base_model,
custom_provider=custom_provider, custom_provider=custom_provider,

38
surfsense_backend/app/services/quality_score.py
View file

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

12
surfsense_backend/tests/unit/services/test_openrouter_integration_service.py
View file

@ -263,11 +263,10 @@ def test_generate_image_gen_configs_filters_by_image_output():
# Each config must carry ``billing_tier`` for routing in image_generation_routes. # Each config must carry ``billing_tier`` for routing in image_generation_routes.
for c in cfgs: for c in cfgs:
assert c["billing_tier"] in {"free", "premium"} assert c["billing_tier"] in {"free", "premium"}
assert c["provider"] == "OPENROUTER" assert c["litellm_provider"] == "openrouter"
assert c[_OPENROUTER_DYNAMIC_MARKER] is True assert c[_OPENROUTER_DYNAMIC_MARKER] is True
# Defense-in-depth: emit the OpenRouter base URL at source so a # Emit the OpenRouter base URL at source so every call path passes an
# downstream call site that forgets ``resolve_api_base`` still # explicit api_base and cannot inherit a process-global endpoint.
# doesn't 404 against an inherited Azure endpoint.
assert c["api_base"] == "https://openrouter.ai/api/v1" assert c["api_base"] == "https://openrouter.ai/api/v1"
@ -346,9 +345,8 @@ def test_generate_vision_llm_configs_filters_by_image_input_text_output():
assert cfg["input_cost_per_token"] == pytest.approx(5e-6) assert cfg["input_cost_per_token"] == pytest.approx(5e-6)
assert cfg["output_cost_per_token"] == pytest.approx(15e-6) assert cfg["output_cost_per_token"] == pytest.approx(15e-6)
assert cfg[_OPENROUTER_DYNAMIC_MARKER] is True assert cfg[_OPENROUTER_DYNAMIC_MARKER] is True
# Defense-in-depth: emit the OpenRouter base URL at source so a # Emit the OpenRouter base URL at source so every call path passes an
# downstream call site that forgets ``resolve_api_base`` still # explicit api_base and cannot inherit a process-global endpoint.
# doesn't inherit an Azure endpoint.
assert cfg["api_base"] == "https://openrouter.ai/api/v1" assert cfg["api_base"] == "https://openrouter.ai/api/v1"

17
surfsense_backend/tests/unit/services/test_pricing_registration.py
View file

@ -186,7 +186,7 @@ def test_openrouter_models_register_under_aliases(monkeypatch):
[ [
{ {
"id": 1, "id": 1,
"provider": "OPENROUTER", "litellm_provider": "openrouter",
"model_name": "anthropic/claude-3-5-sonnet", "model_name": "anthropic/claude-3-5-sonnet",
} }
], ],
@ -228,7 +228,7 @@ def test_yaml_override_registers_under_alias_set(monkeypatch):
[ [
{ {
"id": 1, "id": 1,
"provider": "AZURE_OPENAI", "litellm_provider": "azure",
"model_name": "gpt-5.4", "model_name": "gpt-5.4",
"litellm_params": { "litellm_params": {
"base_model": "gpt-5.4", "base_model": "gpt-5.4",
@ -243,7 +243,6 @@ def test_yaml_override_registers_under_alias_set(monkeypatch):
keys = spy.all_keys keys = spy.all_keys
assert "gpt-5.4" in keys assert "gpt-5.4" in keys
assert "azure_openai/gpt-5.4" in keys
assert "azure/gpt-5.4" in keys assert "azure/gpt-5.4" in keys
payload = spy.calls[0] payload = spy.calls[0]
@ -271,7 +270,7 @@ def test_no_override_means_no_registration(monkeypatch):
[ [
{ {
"id": 1, "id": 1,
"provider": "OPENAI", "litellm_provider": "openai",
"model_name": "gpt-4o", "model_name": "gpt-4o",
"litellm_params": {"base_model": "gpt-4o"}, "litellm_params": {"base_model": "gpt-4o"},
} }
@ -302,7 +301,7 @@ def test_openrouter_skipped_when_pricing_missing(monkeypatch):
[ [
{ {
"id": 1, "id": 1,
"provider": "OPENROUTER", "litellm_provider": "openrouter",
"model_name": "anthropic/claude-3-5-sonnet", "model_name": "anthropic/claude-3-5-sonnet",
} }
], ],
@ -349,12 +348,12 @@ def test_register_continues_after_individual_failure(monkeypatch, caplog):
[ [
{ {
"id": 1, "id": 1,
"provider": "OPENROUTER", "litellm_provider": "openrouter",
"model_name": "anthropic/claude-3-5-sonnet", "model_name": "anthropic/claude-3-5-sonnet",
}, },
{ {
"id": 2, "id": 2,
"provider": "OPENAI", "litellm_provider": "openai",
"model_name": "custom-deployment", "model_name": "custom-deployment",
"litellm_params": { "litellm_params": {
"base_model": "custom-deployment", "base_model": "custom-deployment",
@ -396,7 +395,7 @@ def test_vision_configs_registered_with_chat_shape(monkeypatch):
[ [
{ {
"id": -1, "id": -1,
"provider": "OPENROUTER", "litellm_provider": "openrouter",
"model_name": "openai/gpt-4o", "model_name": "openai/gpt-4o",
"billing_tier": "premium", "billing_tier": "premium",
"input_cost_per_token": 5e-6, "input_cost_per_token": 5e-6,
@ -433,7 +432,7 @@ def test_vision_with_inline_pricing_when_or_cache_missing(monkeypatch):
[ [
{ {
"id": -1, "id": -1,
"provider": "OPENROUTER", "litellm_provider": "openrouter",
"model_name": "google/gemini-2.5-flash", "model_name": "google/gemini-2.5-flash",
"billing_tier": "premium", "billing_tier": "premium",
"input_cost_per_token": 1e-6, "input_cost_per_token": 1e-6,

107
surfsense_backend/tests/unit/services/test_provider_api_base.py
View file

@ -1,107 +0,0 @@
"""Unit tests for the shared ``api_base`` resolver.
The cascade exists so vision and image-gen call sites can't silently
inherit ``litellm.api_base`` (commonly set by ``AZURE_OPENAI_ENDPOINT``)
when an OpenRouter / Groq / etc. config ships an empty string. See
``provider_api_base`` module docstring for the original repro
(OpenRouter image-gen 404-ing against an Azure endpoint).
"""
from __future__ import annotations
import pytest
from app.services.provider_api_base import (
PROVIDER_DEFAULT_API_BASE,
PROVIDER_KEY_DEFAULT_API_BASE,
resolve_api_base,
)
pytestmark = pytest.mark.unit
def test_config_value_wins_over_defaults():
"""A non-empty config value is always returned verbatim, even when the
provider has a default — the operator gets the last word."""
result = resolve_api_base(
provider="OPENROUTER",
provider_prefix="openrouter",
config_api_base="https://my-openrouter-mirror.example.com/v1",
)
assert result == "https://my-openrouter-mirror.example.com/v1"
def test_provider_key_default_when_config_missing():
"""``DEEPSEEK`` shares the ``openai`` LiteLLM prefix but has its own
base URL — the provider-key map must take precedence over the prefix
map so DeepSeek requests don't go to OpenAI."""
result = resolve_api_base(
provider="DEEPSEEK",
provider_prefix="openai",
config_api_base=None,
)
assert result == PROVIDER_KEY_DEFAULT_API_BASE["DEEPSEEK"]
def test_provider_prefix_default_when_no_key_default():
result = resolve_api_base(
provider="OPENROUTER",
provider_prefix="openrouter",
config_api_base=None,
)
assert result == PROVIDER_DEFAULT_API_BASE["openrouter"]
def test_unknown_provider_returns_none():
"""When neither map matches we return ``None`` so the caller can let
LiteLLM apply its own provider-integration default (Azure deployment
URL, custom-provider URL, etc.)."""
result = resolve_api_base(
provider="SOMETHING_NEW",
provider_prefix="something_new",
config_api_base=None,
)
assert result is None
def test_empty_string_config_treated_as_missing():
"""The original bug: OpenRouter dynamic configs ship ``api_base=""``
and downstream call sites use ``if cfg.get("api_base"):`` — empty
strings are falsy in Python but the cascade has to step in anyway."""
result = resolve_api_base(
provider="OPENROUTER",
provider_prefix="openrouter",
config_api_base="",
)
assert result == PROVIDER_DEFAULT_API_BASE["openrouter"]
def test_whitespace_only_config_treated_as_missing():
"""A config value of ``" "`` is a configuration mistake — treat it
as missing instead of forwarding whitespace to LiteLLM (which would
almost certainly 404)."""
result = resolve_api_base(
provider="OPENROUTER",
provider_prefix="openrouter",
config_api_base=" ",
)
assert result == PROVIDER_DEFAULT_API_BASE["openrouter"]
def test_provider_case_insensitive():
"""Some call sites pass the provider lowercase (DB enum value), others
uppercase (YAML key). Both must resolve."""
upper = resolve_api_base(
provider="DEEPSEEK", provider_prefix="openai", config_api_base=None
)
lower = resolve_api_base(
provider="deepseek", provider_prefix="openai", config_api_base=None
)
assert upper == lower == PROVIDER_KEY_DEFAULT_API_BASE["DEEPSEEK"]
def test_all_inputs_none_returns_none():
assert (
resolve_api_base(provider=None, provider_prefix=None, config_api_base=None)
is None
)

28
surfsense_backend/tests/unit/services/test_provider_capabilities.py
View file

@ -32,7 +32,7 @@ pytestmark = pytest.mark.unit
def test_or_modalities_with_image_returns_true(): def test_or_modalities_with_image_returns_true():
assert ( assert (
derive_supports_image_input( derive_supports_image_input(
provider="OPENROUTER", litellm_provider="openrouter",
model_name="openai/gpt-4o", model_name="openai/gpt-4o",
openrouter_input_modalities=["text", "image"], openrouter_input_modalities=["text", "image"],
) )
@ -43,7 +43,7 @@ def test_or_modalities_with_image_returns_true():
def test_or_modalities_text_only_returns_false(): def test_or_modalities_text_only_returns_false():
assert ( assert (
derive_supports_image_input( derive_supports_image_input(
provider="OPENROUTER", litellm_provider="openrouter",
model_name="deepseek/deepseek-v3.2-exp", model_name="deepseek/deepseek-v3.2-exp",
openrouter_input_modalities=["text"], openrouter_input_modalities=["text"],
) )
@ -57,7 +57,7 @@ def test_or_modalities_empty_list_returns_false():
to LiteLLM.""" to LiteLLM."""
assert ( assert (
derive_supports_image_input( derive_supports_image_input(
provider="OPENROUTER", litellm_provider="openrouter",
model_name="weird/empty-modalities", model_name="weird/empty-modalities",
openrouter_input_modalities=[], openrouter_input_modalities=[],
) )
@ -70,7 +70,7 @@ def test_or_modalities_none_falls_through_to_litellm():
to LiteLLM. Using ``openai/gpt-4o`` which is in LiteLLM's map.""" to LiteLLM. Using ``openai/gpt-4o`` which is in LiteLLM's map."""
assert ( assert (
derive_supports_image_input( derive_supports_image_input(
provider="OPENAI", litellm_provider="openai",
model_name="gpt-4o", model_name="gpt-4o",
openrouter_input_modalities=None, openrouter_input_modalities=None,
) )
@ -86,7 +86,7 @@ def test_or_modalities_none_falls_through_to_litellm():
def test_litellm_known_vision_model_returns_true(): def test_litellm_known_vision_model_returns_true():
assert ( assert (
derive_supports_image_input( derive_supports_image_input(
provider="OPENAI", litellm_provider="openai",
model_name="gpt-4o", model_name="gpt-4o",
) )
is True is True
@ -100,7 +100,7 @@ def test_litellm_base_model_wins_over_model_name():
doesn't know) would shadow the real capability.""" doesn't know) would shadow the real capability."""
assert ( assert (
derive_supports_image_input( derive_supports_image_input(
provider="AZURE_OPENAI", litellm_provider="azure",
model_name="my-azure-deployment-id", model_name="my-azure-deployment-id",
base_model="gpt-4o", base_model="gpt-4o",
) )
@ -112,7 +112,7 @@ def test_litellm_unknown_model_default_allows():
"""Default-allow on unknown — the safety net is the actual block.""" """Default-allow on unknown — the safety net is the actual block."""
assert ( assert (
derive_supports_image_input( derive_supports_image_input(
provider="CUSTOM", litellm_provider="custom",
model_name="brand-new-model-x9-unmapped", model_name="brand-new-model-x9-unmapped",
custom_provider="brand_new_proxy", custom_provider="brand_new_proxy",
) )
@ -128,7 +128,7 @@ def test_litellm_known_text_only_returns_false():
# Sanity: confirm the helper's negative path. We use a small model # Sanity: confirm the helper's negative path. We use a small model
# known not to support vision per the map. # known not to support vision per the map.
result = derive_supports_image_input( result = derive_supports_image_input(
provider="DEEPSEEK", litellm_provider="openai",
model_name="deepseek-chat", model_name="deepseek-chat",
) )
# We accept either False (LiteLLM said explicit no) or True # We accept either False (LiteLLM said explicit no) or True
@ -147,7 +147,7 @@ def test_litellm_known_text_only_returns_false():
def test_is_known_text_only_returns_false_for_vision_model(): def test_is_known_text_only_returns_false_for_vision_model():
assert ( assert (
is_known_text_only_chat_model( is_known_text_only_chat_model(
provider="OPENAI", litellm_provider="openai",
model_name="gpt-4o", model_name="gpt-4o",
) )
is False is False
@ -160,7 +160,7 @@ def test_is_known_text_only_returns_false_for_unknown_model():
fixing.""" fixing."""
assert ( assert (
is_known_text_only_chat_model( is_known_text_only_chat_model(
provider="CUSTOM", litellm_provider="custom",
model_name="brand-new-model-x9-unmapped", model_name="brand-new-model-x9-unmapped",
custom_provider="brand_new_proxy", custom_provider="brand_new_proxy",
) )
@ -181,7 +181,7 @@ def test_is_known_text_only_returns_false_when_lookup_raises(monkeypatch):
assert ( assert (
is_known_text_only_chat_model( is_known_text_only_chat_model(
provider="OPENAI", litellm_provider="openai",
model_name="gpt-4o", model_name="gpt-4o",
) )
is False is False
@ -201,7 +201,7 @@ def test_is_known_text_only_returns_true_on_explicit_false(monkeypatch):
assert ( assert (
is_known_text_only_chat_model( is_known_text_only_chat_model(
provider="OPENAI", litellm_provider="openai",
model_name="any-model", model_name="any-model",
) )
is True is True
@ -218,7 +218,7 @@ def test_is_known_text_only_returns_false_on_supports_vision_true(monkeypatch):
assert ( assert (
is_known_text_only_chat_model( is_known_text_only_chat_model(
provider="OPENAI", litellm_provider="openai",
model_name="any-model", model_name="any-model",
) )
is False is False
@ -237,7 +237,7 @@ def test_is_known_text_only_returns_false_on_missing_key(monkeypatch):
assert ( assert (
is_known_text_only_chat_model( is_known_text_only_chat_model(
provider="OPENAI", litellm_provider="openai",
model_name="any-model", model_name="any-model",
) )
is False is False

6
surfsense_backend/tests/unit/services/test_quality_score.py
View file

@ -228,7 +228,7 @@ def test_static_score_or_recent_release_beats_year_old_same_provider():
def test_static_score_yaml_includes_operator_bonus(): def test_static_score_yaml_includes_operator_bonus():
cfg = { cfg = {
"provider": "AZURE_OPENAI", "litellm_provider": "azure",
"model_name": "gpt-5", "model_name": "gpt-5",
"litellm_params": {"base_model": "azure/gpt-5"}, "litellm_params": {"base_model": "azure/gpt-5"},
} }
@ -238,7 +238,7 @@ def test_static_score_yaml_includes_operator_bonus():
def test_static_score_yaml_unknown_provider_still_carries_bonus(): def test_static_score_yaml_unknown_provider_still_carries_bonus():
cfg = { cfg = {
"provider": "SOME_NEW_PROVIDER", "litellm_provider": "some_new_provider",
"model_name": "weird-model", "model_name": "weird-model",
} }
score = static_score_yaml(cfg) score = static_score_yaml(cfg)
@ -247,7 +247,7 @@ def test_static_score_yaml_unknown_provider_still_carries_bonus():
def test_static_score_yaml_clamped_0_to_100(): def test_static_score_yaml_clamped_0_to_100():
cfg = { cfg = {
"provider": "AZURE_OPENAI", "litellm_provider": "azure",
"model_name": "gpt-5", "model_name": "gpt-5",
"litellm_params": {"base_model": "azure/gpt-5"}, "litellm_params": {"base_model": "azure/gpt-5"},
} }