diff --git a/litellm/utils.py b/litellm/utils.py index 9cc4cb467c7..4c4ce3bb559 100644 --- a/litellm/utils.py +++ b/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), diff --git a/tests/integration/providers/test_api_base_trailing_slash_wire.py b/tests/integration/providers/test_api_base_trailing_slash_wire.py index 01b68f52032..6d3b1201c98 100644 --- a/tests/integration/providers/test_api_base_trailing_slash_wire.py +++ b/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() diff --git a/tests/unit/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py index 969ca0269cf..7b5459daaf5 100644 --- a/tests/unit/test_register_model_custom_pricing.py +++ b/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: