From 8f20a3257189807136ddf56fa942683bf4b366ce Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Thu, 11 Jun 2026 18:21:07 +0530 Subject: [PATCH] refactor(model-connections): consolidate provider capability handling --- .../app/services/global_model_catalog.py | 2 +- .../openrouter_integration_service.py | 21 +--- .../app/services/pricing_registration.py | 7 +- .../app/services/provider_api_base.py | 106 ----------------- .../app/services/provider_capabilities.py | 31 +---- .../app/services/quality_score.py | 38 +++---- .../test_openrouter_integration_service.py | 12 +- .../services/test_pricing_registration.py | 17 ++- .../unit/services/test_provider_api_base.py | 107 ------------------ .../services/test_provider_capabilities.py | 28 ++--- .../tests/unit/services/test_quality_score.py | 6 +- 11 files changed, 64 insertions(+), 311 deletions(-) delete mode 100644 surfsense_backend/app/services/provider_api_base.py delete mode 100644 surfsense_backend/tests/unit/services/test_provider_api_base.py diff --git a/surfsense_backend/app/services/global_model_catalog.py b/surfsense_backend/app/services/global_model_catalog.py index a43f58b9e..e40b3a942 100644 --- a/surfsense_backend/app/services/global_model_catalog.py +++ b/surfsense_backend/app/services/global_model_catalog.py @@ -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 {}), diff --git a/surfsense_backend/app/services/openrouter_integration_service.py b/surfsense_backend/app/services/openrouter_integration_service.py index 6454e2d58..6996f0fde 100644 --- a/surfsense_backend/app/services/openrouter_integration_service.py +++ b/surfsense_backend/app/services/openrouter_integration_service.py @@ -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 diff --git a/surfsense_backend/app/services/pricing_registration.py b/surfsense_backend/app/services/pricing_registration.py index de98e50c2..6b99fe723 100644 --- a/surfsense_backend/app/services/pricing_registration.py +++ b/surfsense_backend/app/services/pricing_registration.py @@ -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 diff --git a/surfsense_backend/app/services/provider_api_base.py b/surfsense_backend/app/services/provider_api_base.py deleted file mode 100644 index dca1f9462..000000000 --- a/surfsense_backend/app/services/provider_api_base.py +++ /dev/null @@ -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/`` 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", -] diff --git a/surfsense_backend/app/services/provider_capabilities.py b/surfsense_backend/app/services/provider_capabilities.py index 9e1433214..9521ef7a4 100644 --- a/surfsense_backend/app/services/provider_capabilities.py +++ b/surfsense_backend/app/services/provider_capabilities.py @@ -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, diff --git a/surfsense_backend/app/services/quality_score.py b/surfsense_backend/app/services/quality_score.py index 2fb37de21..95484439b 100644 --- a/surfsense_backend/app/services/quality_score.py +++ b/surfsense_backend/app/services/quality_score.py @@ -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 "" diff --git a/surfsense_backend/tests/unit/services/test_openrouter_integration_service.py b/surfsense_backend/tests/unit/services/test_openrouter_integration_service.py index 88fcf2db3..9d4c1a04b 100644 --- a/surfsense_backend/tests/unit/services/test_openrouter_integration_service.py +++ b/surfsense_backend/tests/unit/services/test_openrouter_integration_service.py @@ -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. for c in cfgs: assert c["billing_tier"] in {"free", "premium"} - assert c["provider"] == "OPENROUTER" + assert c["litellm_provider"] == "openrouter" assert c[_OPENROUTER_DYNAMIC_MARKER] is True - # Defense-in-depth: emit the OpenRouter base URL at source so a - # downstream call site that forgets ``resolve_api_base`` still - # doesn't 404 against an inherited Azure endpoint. + # Emit the OpenRouter base URL at source so every call path passes an + # explicit api_base and cannot inherit a process-global endpoint. 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["output_cost_per_token"] == pytest.approx(15e-6) assert cfg[_OPENROUTER_DYNAMIC_MARKER] is True - # Defense-in-depth: emit the OpenRouter base URL at source so a - # downstream call site that forgets ``resolve_api_base`` still - # doesn't inherit an Azure endpoint. + # Emit the OpenRouter base URL at source so every call path passes an + # explicit api_base and cannot inherit a process-global endpoint. assert cfg["api_base"] == "https://openrouter.ai/api/v1" diff --git a/surfsense_backend/tests/unit/services/test_pricing_registration.py b/surfsense_backend/tests/unit/services/test_pricing_registration.py index e97250ff2..c9adc6aac 100644 --- a/surfsense_backend/tests/unit/services/test_pricing_registration.py +++ b/surfsense_backend/tests/unit/services/test_pricing_registration.py @@ -186,7 +186,7 @@ def test_openrouter_models_register_under_aliases(monkeypatch): [ { "id": 1, - "provider": "OPENROUTER", + "litellm_provider": "openrouter", "model_name": "anthropic/claude-3-5-sonnet", } ], @@ -228,7 +228,7 @@ def test_yaml_override_registers_under_alias_set(monkeypatch): [ { "id": 1, - "provider": "AZURE_OPENAI", + "litellm_provider": "azure", "model_name": "gpt-5.4", "litellm_params": { "base_model": "gpt-5.4", @@ -243,7 +243,6 @@ def test_yaml_override_registers_under_alias_set(monkeypatch): keys = spy.all_keys assert "gpt-5.4" in keys - assert "azure_openai/gpt-5.4" in keys assert "azure/gpt-5.4" in keys payload = spy.calls[0] @@ -271,7 +270,7 @@ def test_no_override_means_no_registration(monkeypatch): [ { "id": 1, - "provider": "OPENAI", + "litellm_provider": "openai", "model_name": "gpt-4o", "litellm_params": {"base_model": "gpt-4o"}, } @@ -302,7 +301,7 @@ def test_openrouter_skipped_when_pricing_missing(monkeypatch): [ { "id": 1, - "provider": "OPENROUTER", + "litellm_provider": "openrouter", "model_name": "anthropic/claude-3-5-sonnet", } ], @@ -349,12 +348,12 @@ def test_register_continues_after_individual_failure(monkeypatch, caplog): [ { "id": 1, - "provider": "OPENROUTER", + "litellm_provider": "openrouter", "model_name": "anthropic/claude-3-5-sonnet", }, { "id": 2, - "provider": "OPENAI", + "litellm_provider": "openai", "model_name": "custom-deployment", "litellm_params": { "base_model": "custom-deployment", @@ -396,7 +395,7 @@ def test_vision_configs_registered_with_chat_shape(monkeypatch): [ { "id": -1, - "provider": "OPENROUTER", + "litellm_provider": "openrouter", "model_name": "openai/gpt-4o", "billing_tier": "premium", "input_cost_per_token": 5e-6, @@ -433,7 +432,7 @@ def test_vision_with_inline_pricing_when_or_cache_missing(monkeypatch): [ { "id": -1, - "provider": "OPENROUTER", + "litellm_provider": "openrouter", "model_name": "google/gemini-2.5-flash", "billing_tier": "premium", "input_cost_per_token": 1e-6, diff --git a/surfsense_backend/tests/unit/services/test_provider_api_base.py b/surfsense_backend/tests/unit/services/test_provider_api_base.py deleted file mode 100644 index 12cd0a3d5..000000000 --- a/surfsense_backend/tests/unit/services/test_provider_api_base.py +++ /dev/null @@ -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 - ) diff --git a/surfsense_backend/tests/unit/services/test_provider_capabilities.py b/surfsense_backend/tests/unit/services/test_provider_capabilities.py index aac88977f..c327c2c87 100644 --- a/surfsense_backend/tests/unit/services/test_provider_capabilities.py +++ b/surfsense_backend/tests/unit/services/test_provider_capabilities.py @@ -32,7 +32,7 @@ pytestmark = pytest.mark.unit def test_or_modalities_with_image_returns_true(): assert ( derive_supports_image_input( - provider="OPENROUTER", + litellm_provider="openrouter", model_name="openai/gpt-4o", 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(): assert ( derive_supports_image_input( - provider="OPENROUTER", + litellm_provider="openrouter", model_name="deepseek/deepseek-v3.2-exp", openrouter_input_modalities=["text"], ) @@ -57,7 +57,7 @@ def test_or_modalities_empty_list_returns_false(): to LiteLLM.""" assert ( derive_supports_image_input( - provider="OPENROUTER", + litellm_provider="openrouter", model_name="weird/empty-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.""" assert ( derive_supports_image_input( - provider="OPENAI", + litellm_provider="openai", model_name="gpt-4o", 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(): assert ( derive_supports_image_input( - provider="OPENAI", + litellm_provider="openai", model_name="gpt-4o", ) is True @@ -100,7 +100,7 @@ def test_litellm_base_model_wins_over_model_name(): doesn't know) would shadow the real capability.""" assert ( derive_supports_image_input( - provider="AZURE_OPENAI", + litellm_provider="azure", model_name="my-azure-deployment-id", 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.""" assert ( derive_supports_image_input( - provider="CUSTOM", + litellm_provider="custom", model_name="brand-new-model-x9-unmapped", 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 # known not to support vision per the map. result = derive_supports_image_input( - provider="DEEPSEEK", + litellm_provider="openai", model_name="deepseek-chat", ) # 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(): assert ( is_known_text_only_chat_model( - provider="OPENAI", + litellm_provider="openai", model_name="gpt-4o", ) is False @@ -160,7 +160,7 @@ def test_is_known_text_only_returns_false_for_unknown_model(): fixing.""" assert ( is_known_text_only_chat_model( - provider="CUSTOM", + litellm_provider="custom", model_name="brand-new-model-x9-unmapped", custom_provider="brand_new_proxy", ) @@ -181,7 +181,7 @@ def test_is_known_text_only_returns_false_when_lookup_raises(monkeypatch): assert ( is_known_text_only_chat_model( - provider="OPENAI", + litellm_provider="openai", model_name="gpt-4o", ) is False @@ -201,7 +201,7 @@ def test_is_known_text_only_returns_true_on_explicit_false(monkeypatch): assert ( is_known_text_only_chat_model( - provider="OPENAI", + litellm_provider="openai", model_name="any-model", ) is True @@ -218,7 +218,7 @@ def test_is_known_text_only_returns_false_on_supports_vision_true(monkeypatch): assert ( is_known_text_only_chat_model( - provider="OPENAI", + litellm_provider="openai", model_name="any-model", ) is False @@ -237,7 +237,7 @@ def test_is_known_text_only_returns_false_on_missing_key(monkeypatch): assert ( is_known_text_only_chat_model( - provider="OPENAI", + litellm_provider="openai", model_name="any-model", ) is False diff --git a/surfsense_backend/tests/unit/services/test_quality_score.py b/surfsense_backend/tests/unit/services/test_quality_score.py index 6fbc8fd62..369c8b8f3 100644 --- a/surfsense_backend/tests/unit/services/test_quality_score.py +++ b/surfsense_backend/tests/unit/services/test_quality_score.py @@ -228,7 +228,7 @@ def test_static_score_or_recent_release_beats_year_old_same_provider(): def test_static_score_yaml_includes_operator_bonus(): cfg = { - "provider": "AZURE_OPENAI", + "litellm_provider": "azure", "model_name": "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(): cfg = { - "provider": "SOME_NEW_PROVIDER", + "litellm_provider": "some_new_provider", "model_name": "weird-model", } 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(): cfg = { - "provider": "AZURE_OPENAI", + "litellm_provider": "azure", "model_name": "gpt-5", "litellm_params": {"base_model": "azure/gpt-5"}, }