fix(utils): register ollama deployment prices from the built-in catalog without calling the provider (#51561)

* test(providers): keep registration lookups in the model catalog

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(providers): format catalog registration lookup

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/utils.py

@@ -4105,7 +4105,7 @@ def lookup_catalog_model_info(model: str, custom_llm_provider: str | None) -> Mo
 
 def _get_builtin_model_info_for_registration(
     model: str, custom_llm_provider: str | None
-) -> ModelInfo | ModelInfoLookupMiss:
+) -> ModelInfoBase | ModelInfoLookupMiss:
     """Resolve ``model`` to its built-in cost-map entry for registration merging.
 
     Misses when the lookup cannot map the model or resolved it via a
@@ -4114,9 +4114,7 @@ def _get_builtin_model_info_for_registration(
     changing the map mid-walk, is retried, so a transient failure registers exactly
     where the untripped lookup would.
     """
-    lookups: Final = (
-        lookup_model_info_for_registration(model, custom_llm_provider) for _ in range(_REGISTRATION_ATTEMPTS)
-    )
+    lookups: Final = (lookup_catalog_model_info(model, custom_llm_provider) for _ in range(_REGISTRATION_ATTEMPTS))
     return next(
         (lookup for lookup in lookups if not (isinstance(lookup, ModelInfoLookupMiss) and lookup.tripped)),
         ModelInfoLookupMiss(tripped=True),

tests/integration/providers/test_api_base_trailing_slash_wire.py

@@ -1,9 +1,12 @@
+from pathlib import Path
 from typing import Final
 
 import anthropic
+import httpx
 import openai
 import pytest
 from integration._support.client import Gateway, Scenario, list_value, object_value
+from integration._support.process import owned_proxy_process
 from integration._support.wire import Reply, Request, wire_server
 from integration.providers._api_base_trailing_slash_support import (
     KEY,
@@ -311,6 +314,24 @@ def test_embedding_base_that_already_names_the_route_is_sent_as_is(gateway: Gate
         assert_one_request(wire.drain(), row.target, marker)
 
 
+def test_explicit_ollama_embedding_base_does_not_probe_the_default_server(gateway: Gateway, tmp_path: Path) -> None:
+    marker: Final = new_marker()
+    row: Final = Row(f"ollama/{marker}", "/ol/api/embed", "")
+    with (
+        wire_server(lambda _: Reply()) as default,
+        wire_server(provider_peer) as wire,
+        owned_proxy_process(gateway, tmp_path, {"OLLAMA_API_BASE": default.url}) as owned,
+        owned.gateway.scenario() as scenario,
+        httpx.Client(trust_env=False) as client,
+    ):
+        model: Final = deployment(scenario, row, f"{wire.url}{row.base}")
+        assert _embedding(owned.gateway, model, marker) == list(VECTOR)
+        assert_one_request(wire.drain(), row.target, marker)
+        control: Final = client.get(f"{default.url}/control")
+        assert control.status_code == 200, control.text
+        assert [(request.method, request.target) for request in default.drain()] == [("GET", "/control")]
+
+
 @pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS)
 def test_search_base_that_already_names_the_route_is_sent_as_is(gateway: Gateway, trailing: str) -> None:
     marker: Final = new_marker()

tests/unit/test_register_model_custom_pricing.py

@@ -123,7 +123,7 @@ def test_register_model_strips_none_litellm_provider_from_get_model_info(monkeyp
 
     The companion test above hits the ``except Exception`` branch where
     ``existing_model`` is an empty dict, so the ``pop`` is a no-op. This
-    test patches ``get_model_info`` to return the failure mode the strip
+    test patches the catalog lookup to return the failure mode the strip
     was added to handle, namely a populated dict whose ``litellm_provider``
     is ``None``. Without the strip, the merged entry in
     ``litellm.model_cost`` would carry ``litellm_provider: None`` and
@@ -146,7 +146,7 @@ def test_register_model_strips_none_litellm_provider_from_get_model_info(monkeyp
             "max_tokens": 4096,
         }
 
-    monkeypatch.setattr(litellm_utils, "get_model_info", _fake_get_model_info)
+    monkeypatch.setattr(litellm_utils, "get_model_info_helper", _fake_get_model_info)
 
     try:
         litellm.register_model(
@@ -167,7 +167,7 @@ def test_register_model_strips_none_litellm_provider_from_get_model_info(monkeyp
             "register_model failed to strip litellm_provider=None returned "
             f"by get_model_info, got {registered.get('litellm_provider')!r}"
         )
-        # Metadata from the patched ``get_model_info`` must still flow
+        # Metadata from the patched catalog lookup must still flow
         # through, so we know the strip did not nuke the rest of the entry.
         assert registered.get("mode") == "chat"
         assert registered.get("max_tokens") == 4096
@@ -656,11 +656,13 @@ def test_register_model_keeps_a_catalog_rows_own_prices_when_the_lookup_fails(mo
     )
     publish(litellm.model_cost)
 
-    def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn:
+    def lookup_tripped_by_a_concurrent_registration(
+        model: str, custom_llm_provider: str | None = None, *, catalog_only: bool = False
+    ) -> NoReturn:
         raise RuntimeError("dictionary changed size during iteration")
 
     with pytest.MonkeyPatch.context() as lookup_patch:
-        lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration)
+        lookup_patch.setattr(litellm.utils, "get_model_info_helper", lookup_tripped_by_a_concurrent_registration)
         litellm.register_model({geo_key: {"mode": "chat"}}, persist_across_reloads=False)
 
     try:
@@ -694,16 +696,18 @@ def test_register_model_leaves_no_shadow_key_when_the_lookup_trips_once(
         },
     )
     publish(litellm.model_cost)
-    real_lookup: Final = litellm.utils.get_model_info
+    real_lookup: Final = litellm.utils.get_model_info_helper
     lookups: Final = itertools.count()
 
-    def lookup_tripped_once_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> object:
+    def lookup_tripped_once_by_a_concurrent_registration(
+        model: str, custom_llm_provider: str | None = None, *, catalog_only: bool = False
+    ) -> object:
         if next(lookups) == 0:
             raise RuntimeError("dictionary changed size during iteration")
-        return real_lookup(model=model, custom_llm_provider=custom_llm_provider)
+        return real_lookup(model=model, custom_llm_provider=custom_llm_provider, catalog_only=catalog_only)
 
     with pytest.MonkeyPatch.context() as lookup_patch:
-        lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_once_by_a_concurrent_registration)
+        lookup_patch.setattr(litellm.utils, "get_model_info_helper", lookup_tripped_once_by_a_concurrent_registration)
         litellm.register_model({alias: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False)
     litellm.register_model({catalog_key: {"input_cost_per_token": repriced_input_cost}}, persist_across_reloads=False)
 
@@ -733,14 +737,16 @@ def test_register_model_bills_a_catalog_rows_tier_rates_the_same_when_the_lookup
     builtin_cost_map: Final = litellm.get_model_cost_map(url="")
     monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
 
-    def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn:
+    def lookup_tripped_by_a_concurrent_registration(
+        model: str, custom_llm_provider: str | None = None, *, catalog_only: bool = False
+    ) -> NoReturn:
         raise RuntimeError("dictionary changed size during iteration")
 
     def long_prompt_cost_after_registering_with(lookup: Callable[..., object]) -> float:
         monkeypatch.setattr(litellm, "model_cost", {**builtin_cost_map, catalog_key: dict(catalog_row)})
         publish(litellm.model_cost)
         with pytest.MonkeyPatch.context() as lookup_patch:
-            lookup_patch.setattr(litellm.utils, "get_model_info", lookup)
+            lookup_patch.setattr(litellm.utils, "get_model_info_helper", lookup)
             litellm.register_model(
                 {
                     registered_key: {
@@ -762,7 +768,7 @@ def test_register_model_bills_a_catalog_rows_tier_rates_the_same_when_the_lookup
 
     try:
         assert long_prompt_cost_after_registering_with(lookup_tripped_by_a_concurrent_registration) == pytest.approx(
-            long_prompt_cost_after_registering_with(litellm.utils.get_model_info)
+            long_prompt_cost_after_registering_with(litellm.utils.get_model_info_helper)
         )
     finally:
         publish(litellm.model_cost)
@@ -788,12 +794,14 @@ def test_register_model_prices_a_provider_less_azure_ai_key_without_repricing_th
     )
     publish(litellm.model_cost)
 
-    def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn:
+    def lookup_tripped_by_a_concurrent_registration(
+        model: str, custom_llm_provider: str | None = None, *, catalog_only: bool = False
+    ) -> NoReturn:
         raise RuntimeError("dictionary changed size during iteration")
 
     with pytest.MonkeyPatch.context() as lookup_patch:
         if lookup_trips:
-            lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration)
+            lookup_patch.setattr(litellm.utils, "get_model_info_helper", lookup_tripped_by_a_concurrent_registration)
         litellm.register_model(
             {azure_ai_key: {"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}},
             persist_across_reloads=False,
@@ -830,16 +838,18 @@ def test_register_model_keeps_the_prices_a_reload_adopts_while_it_registers(monk
         **builtin_cost_map,
         catalog_key: {**catalog_row, "input_cost_per_token": reloaded_input_cost},
     }
-    real_lookup: Final = litellm.utils.get_model_info
+    real_lookup: Final = litellm.utils.get_model_info_helper
 
-    def lookup_racing_a_price_data_reload(model: str, custom_llm_provider: str | None = None) -> object:
-        info: Final = real_lookup(model=model, custom_llm_provider=custom_llm_provider)
+    def lookup_racing_a_price_data_reload(
+        model: str, custom_llm_provider: str | None = None, *, catalog_only: bool = False
+    ) -> object:
+        info: Final = real_lookup(model=model, custom_llm_provider=custom_llm_provider, catalog_only=catalog_only)
         if litellm.model_cost is not reloaded_cost_map:
             _publish(reloaded_cost_map)
         return info
 
     with pytest.MonkeyPatch.context() as lookup_patch:
-        lookup_patch.setattr(litellm.utils, "get_model_info", lookup_racing_a_price_data_reload)
+        lookup_patch.setattr(litellm.utils, "get_model_info_helper", lookup_racing_a_price_data_reload)
         litellm.register_model({deployment_key: {"litellm_provider": "bedrock"}}, persist_across_reloads=False)
 
     try: