scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-05 07:37:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 360

fix(ollama): send model-info lookups to the call's api_base with its key

Browse source
This commit is contained in:
mateo-berri 2026-10-30 06:20:37 -07:00
parent 73aaf9e3da
commit 8127f296bf
15 changed files with 413 additions and 74 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/constants.py
View file

@ -897,6 +897,8 @@ PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO: Final = frozenset(
}
)
PROVIDERS_THAT_SEND_KEYS_ON_MODEL_INFO_LOOKUPS: Final = frozenset({"ollama", "ollama_chat", "lemonade"})
A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS: Final = frozenset({"databricks_agent"})
LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS: Final = [

48
litellm/cost_calculator.py
View file

@ -377,9 +377,13 @@ def _per_second_pricing_cost(
custom_llm_provider: str | None,
response_time_ms: float | None,
audio_seconds: float = 0.0,
api_base: str | None = None,
api_key: str | None = None,
) -> tuple[float, float] | None:
try:
model_info: Final = cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)
model_info: Final = cached_get_model_info_helper(
model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key
)
except Exception: # noqa: BLE001 # the lookup raises plain Exception for an unmapped model
return None
if _has_token_or_tiered_pricing(model_info) or not _bills_wall_clock_seconds(model_info):
@ -447,6 +451,8 @@ def cost_per_token(
request_model: str | None = None, # original request model for router detection
custom_model_info: OCRPricing | None = None,
bedrock_model_info: ModelInfo | None = None,
api_base: str | None = None,
api_key: str | None = None,
) -> tuple[float, float]:
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@ -600,7 +606,9 @@ def cost_per_token(
model = model_without_prefix
if is_decisions_call_type(call_type):
decisions_model_info: Final = cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)
decisions_model_info: Final = cached_get_model_info_helper(
model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key
)
decisions_overrides: Final = decisions_pricing_fields(
litellm.model_cost.get(decisions_model_info["key"], {}) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # model cost map is dynamically loaded
)
@ -719,6 +727,8 @@ def cost_per_token(
custom_llm_provider=custom_llm_provider,
response_time_ms=response_time_ms,
audio_seconds=audio_transcription_file_duration,
api_base=api_base,
api_key=api_key,
)
) is not None:
return per_second_cost
@ -804,7 +814,9 @@ def cost_per_token(
service_tier=service_tier,
)
else:
model_info: Final = cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)
model_info: Final = cached_get_model_info_helper(
model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key
)
if _has_token_or_tiered_pricing(model_info):
return generic_cost_per_token(
model=model,
@ -2146,6 +2158,8 @@ def completion_cost(
response=completion_response,
custom_model_info=_ocr_model_info(litellm_logging_obj, custom_pricing, router_model_id),
bedrock_model_info=bedrock_model_info,
api_base=_request_api_base(litellm_logging_obj),
api_key=_request_api_key(litellm_logging_obj),
)
# Get additional costs from provider (e.g., routing fees, infrastructure costs)
@ -2238,6 +2252,8 @@ def completion_cost(
vertex_location=vertex_location,
custom_cost_per_token=custom_cost_per_token,
model_info=bedrock_model_info,
api_base=_request_api_base(litellm_logging_obj),
api_key=_request_api_key(litellm_logging_obj),
)
_reasoning_cost = _token_type_breakdown.reasoning_cost
_cache_read_cost = _token_type_breakdown.cache_read_cost
@ -2456,6 +2472,19 @@ def _configured_search_cost_per_query(litellm_logging_obj: LitellmLoggingObject
return float(cost_per_query) if isinstance(cost_per_query, (int, float)) else None
def _request_api_base(litellm_logging_obj: LitellmLoggingObject | None) -> str | None:
logged_params: Final = getattr(litellm_logging_obj, "litellm_params", None)
if not isinstance(logged_params, Mapping):
return None
api_base: Final = _LOGGED_LITELLM_PARAMS.validate_python(logged_params).get("api_base")
return api_base if isinstance(api_base, str) and api_base else None
def _request_api_key(litellm_logging_obj: LitellmLoggingObject | None) -> str | None:
api_key: Final = getattr(litellm_logging_obj, "model_info_api_key", None)
return api_key if isinstance(api_key, str) else None
def _deployment_model_info(
litellm_logging_obj: LitellmLoggingObject | None,
custom_pricing: bool | None,
@ -2510,9 +2539,13 @@ def _layered_ocr_pricing(*sources: Mapping[str, object] | None) -> OCRPricing:
)
def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelInfo | None:
def _cost_map_model_info(
model: str, custom_llm_provider: str | None, api_base: str | None = None, api_key: str | None = None
) -> ModelInfo | None:
try:
return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
return litellm.get_model_info(
model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key
)
except Exception:
return None
@ -2574,8 +2607,11 @@ def pricing_entry_for_cost_calc(
if bedrock_model_info is not None:
return bedrock_model_info["key"], bedrock_model_info
candidates: Final = (selected_model, _get_response_model(completion_response), model)
api_base: Final = _request_api_base(litellm_logging_obj)
api_key: Final = _request_api_key(litellm_logging_obj)
lookups: Final = (_cost_map_model_info(name, custom_llm_provider, api_base, api_key) for name in candidates if name)
resolved: Final = next(
(info for info in (_cost_map_model_info(name, custom_llm_provider) for name in candidates if name) if info),
(info for info in lookups if info),
None,
)
if resolved is None:

13
litellm/litellm_core_utils/litellm_logging.py
View file

@ -205,7 +205,7 @@ from litellm.types.utils import (
batch_chat_usage,
)
from litellm.types.videos.main import VideoObject
from litellm.utils import get_base_model_from_metadata, print_verbose
from litellm.utils import get_base_model_from_metadata, model_info_lookup_api_key, print_verbose
from ..integrations.argilla import ArgillaLogger
from ..integrations.arize.arize_phoenix import ArizePhoenixLogger
@ -834,6 +834,8 @@ class Logging(LiteLLMLoggingBaseClass):
# Passthrough endpoint guardrails config for field targeting
self.passthrough_guardrails_config: dict[str, object] | None = None
self.model_info_api_key: str | None = None
self.model_call_details: dict[str, Any] = {
"litellm_trace_id": self.litellm_trace_id,
"litellm_call_id": litellm_call_id,
@ -1557,6 +1559,12 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["input"] = input
self.model_call_details["api_key"] = api_key
provider: Final = self.model_call_details.get("custom_llm_provider")
self.model_info_api_key = model_info_lookup_api_key(
self.model,
provider if isinstance(provider, str) else None,
api_key if isinstance(api_key, str) and api_key else None,
)
self.model_call_details["additional_args"] = additional_args
self.model_call_details["log_event_type"] = "pre_api_call"
if is_classifier_call(self.call_type, self.model_call_details.get("litellm_params") or EMPTY_MAPPING):
@ -6838,6 +6846,7 @@ class StandardLoggingPayloadSetup:
init_response_obj: object,
api_base: str | None = None,
router_model_id: str | None = None,
api_key: str | None = None,
) -> StandardLoggingModelInformation:
model_cost_name: Final = select_model_name_for_cost_calc(
model=base_model if custom_pricing else None,
@ -6864,6 +6873,7 @@ class StandardLoggingPayloadSetup:
model=pricing_model_name or model_cost_name,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_key=api_key,
)
model_cost_information = StandardLoggingModelInformation(
model_map_key=model_cost_name,
@ -7584,6 +7594,7 @@ def get_standard_logging_object_payload(
init_response_obj=init_response_obj,
api_base=litellm_params.get("api_base"),
router_model_id=logging_obj.get_router_model_id(),
api_key=logging_obj.model_info_api_key,
)
error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload(

10
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -1814,6 +1814,8 @@ def get_billed_token_rates(
current_time: datetime | None = None,
custom_cost_per_token: CostPerToken | None = None,
model_info: ModelInfo | None = None,
api_base: str | None = None,
api_key: str | None = None,
) -> BilledTokenRates | None:
"""Rates the cost calculator bills ``usage`` at, resolved exactly as the totals and the token-type
breakdown resolve them. None when the model's pricing cannot be resolved."""
@ -1823,7 +1825,9 @@ def get_billed_token_rates(
resolved_model_info: Final = apply_provider_cache_read_default(
model_info
if model_info is not None
else get_model_info(model=model, custom_llm_provider=custom_llm_provider),
else get_model_info(
model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key
),
custom_llm_provider,
)
except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model: no rates
@ -1849,6 +1853,8 @@ def get_token_type_cost_breakdown(
current_time: datetime | None = None,
custom_cost_per_token: CostPerToken | None = None,
model_info: ModelInfo | None = None,
api_base: str | None = None,
api_key: str | None = None,
) -> TokenTypeCostBreakdown:
"""
Provider-agnostic cost of reasoning and cache tokens, derived from the usage
@ -1874,6 +1880,8 @@ def get_token_type_cost_breakdown(
current_time=current_time,
custom_cost_per_token=custom_cost_per_token,
model_info=model_info,
api_base=api_base,
api_key=api_key,
)
if rates is None:
return TokenTypeCostBreakdown(0.0, 0.0, 0.0)

7
litellm/llms/base_llm/decisions/transformation.py
View file

@ -59,6 +59,7 @@ from litellm.types.decisions import (
systemone_choice_key,
)
from litellm.types.utils import is_decisions_model_mode
from litellm.utils import get_model_info_helper
_PAYLOAD_ADAPTER: Final[TypeAdapter[object]] = TypeAdapter(object)
_SYSTEMONE_RESPONSE_ADAPTER: Final[TypeAdapter[DecisionsResponse]] = TypeAdapter(DecisionsResponse)
@ -505,8 +506,10 @@ def ir_to_systemone_response(response: DecisionsIRResponse, request: DecisionsIR
def _is_decisions_model(model: str, custom_llm_provider: str) -> bool:
try:
model_mode: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider).get("mode")
except Exception: # noqa: BLE001 # get_model_info raises bare Exception for an unmapped model, which has no mode
model_mode: Final = get_model_info_helper(
model=model, custom_llm_provider=custom_llm_provider, catalog_only=True
).get("mode")
except Exception: # noqa: BLE001 # the lookup raises bare Exception for an unmapped model, which has no mode
return False
return is_decisions_model_mode(model_mode)

1
litellm/llms/ollama/common_utils.py
View file

@ -134,6 +134,7 @@ class OllamaModelInfo(BaseLLMModelInfo):
"/api/embeddings",
"/api/show",
"/api/tags",
"/v1/systemone",
):
if api_base.endswith(suffix):
return api_base[: -len(suffix)]

2
litellm/responses/bridge_check.py
View file

@ -167,7 +167,7 @@ def responses_api_bridge_check(
try:
model_info = cast(
dict[str, object],
get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider),
get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider, catalog_only=True),
)
except Exception as e:
verbose_logger.debug("Error getting model info: %s", e)

16
litellm/router.py
View file

@ -479,6 +479,7 @@ from litellm.types.services import ServiceTypes
from litellm.types.utils import (
AUTOROUTER_CLASSIFIER_CALL_ORIGIN,
PROMPT_QUOTING_ROUTING_DECISION_FIELDS,
AdminApiKeySecretReference,
CustomPricingLiteLLMParams,
GenericBudgetConfigType,
LiteLLMBatch,
@ -11868,6 +11869,11 @@ class Router(DeprecatedRouterMembers):
return (backend_key, backend_key.replace("responses/", ""))
return (backend_key,)
@staticmethod
def _model_info_lookup_api_key(deployment: Deployment) -> str | None:
api_key: Final = deployment.litellm_params.api_key
return None if isinstance(api_key, AdminApiKeySecretReference) else api_key
@staticmethod
def _deployment_model_cost_payload(deployment: Deployment) -> dict[str, object]: # mutable-ok: cost-map entry
"""The ``model_info`` a deployment contributes to ``litellm.model_cost``.
@ -11941,6 +11947,8 @@ class Router(DeprecatedRouterMembers):
model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
pricing_provider=deployment.model_info.pricing_provider,
api_base=deployment.litellm_params.api_base,
api_key=Router._model_info_lookup_api_key(deployment),
)
@staticmethod
@ -11951,6 +11959,8 @@ class Router(DeprecatedRouterMembers):
model: str,
custom_llm_provider: str | None,
pricing_provider: str | None = None,
api_base: str | None = None,
api_key: str | None = None,
) -> None:
"""Write a deployment's metadata into ``litellm.model_cost``.
@ -11976,6 +11986,8 @@ class Router(DeprecatedRouterMembers):
warning_display_name=model,
custom_llm_provider=pricing_provider or custom_llm_provider,
add_to_provider_model_sets=False,
api_base=api_base,
api_key=api_key,
)
## OLD MODEL REGISTRATION ## Kept to prevent breaking changes
@ -11991,6 +12003,8 @@ class Router(DeprecatedRouterMembers):
litellm.register_model(
model_cost={_key: keep_registered_mode(_key, shared_model_info) for _key in backend_keys},
persist_across_reloads=False,
api_base=api_base,
api_key=api_key,
)
def replay_model_cost_registrations(self) -> None:
@ -12017,6 +12031,8 @@ class Router(DeprecatedRouterMembers):
model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
pricing_provider=deployment.model_info.pricing_provider,
api_base=deployment.litellm_params.api_base,
api_key=Router._model_info_lookup_api_key(deployment),
)
self._invalidate_model_group_info_cache()

83
litellm/utils.py
View file

@ -81,6 +81,7 @@ from litellm.constants import (
NON_INFERENCE_CALL_TYPES,
OPENAI_EMBEDDING_PARAMS,
PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO,
PROVIDERS_THAT_SEND_KEYS_ON_MODEL_INFO_LOOKUPS,
TOOL_CHOICE_OBJECT_TOKEN_COUNT,
)
from litellm.litellm_core_utils.core_helpers import (
@ -3987,8 +3988,12 @@ def _lookup_mapped_model_info(lookup: Callable[[], _ModelInfoT]) -> _ModelInfoT
return ModelInfoLookupMiss(tripped=False) if is_generalized_model_info(info) else info
def lookup_model_info_for_registration(model: str, custom_llm_provider: str | None) -> ModelInfo | ModelInfoLookupMiss:
return _lookup_mapped_model_info(lambda: get_model_info(model=model, custom_llm_provider=custom_llm_provider))
def lookup_model_info_for_registration(
model: str, custom_llm_provider: str | None, api_base: str | None = None, api_key: str | None = None
) -> ModelInfo | ModelInfoLookupMiss:
return _lookup_mapped_model_info(
lambda: get_model_info(model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key)
)
def lookup_catalog_model_info(model: str, custom_llm_provider: str | None) -> ModelInfoBase | ModelInfoLookupMiss:
@ -3998,7 +4003,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
model: str, custom_llm_provider: str | None, api_base: str | None, api_key: str | None
) -> ModelInfo | ModelInfoLookupMiss:
"""Resolve ``model`` to its built-in cost-map entry for registration merging.
@ -4009,7 +4014,8 @@ def _get_builtin_model_info_for_registration(
where the untripped lookup would.
"""
lookups: Final = (
lookup_model_info_for_registration(model, custom_llm_provider) for _ in range(_REGISTRATION_ATTEMPTS)
lookup_model_info_for_registration(model, custom_llm_provider, api_base, api_key)
for _ in range(_REGISTRATION_ATTEMPTS)
)
return next(
(lookup for lookup in lookups if not (isinstance(lookup, ModelInfoLookupMiss) and lookup.tripped)),
@ -4034,19 +4040,29 @@ def model_cost_registration_row(key: str, provider: object) -> Mapping[str, obje
``litellm.model_cost`` row it resolves to (``openai/o3-pro`` lands on ``o3-pro``), else the row it seeds ``key``
with from the built-in entry it falls back to (an OpenAI fine-tune takes its base model's row)."""
base: Final = _registration_base(
key, {"litellm_provider": provider}, custom_llm_provider=None, warning_display_name=None
key,
{"litellm_provider": provider},
custom_llm_provider=None,
warning_display_name=None,
api_base=None,
api_key=None,
)
return _builtin_cost_map().get(base.key) or base.entry
def _registration_base(
key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None
key: str,
value: Mapping[str, object],
custom_llm_provider: str | None,
warning_display_name: str | None,
api_base: str | None,
api_key: str | None,
) -> _RegistrationWrite:
provider: Final = str(value.get("litellm_provider") or "")
if _skips_builtin_registration_lookup(key, provider):
return _RegistrationWrite(key=key, entry=dict(_builtin_cost_map().get(key, {})))
builtin_model_info: Final = _get_builtin_model_info_for_registration(
model=key, custom_llm_provider=custom_llm_provider
model=key, custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key
)
if not isinstance(builtin_model_info, ModelInfoLookupMiss):
return _RegistrationWrite(
@ -4077,9 +4093,14 @@ def _base_fields_replaced_by(value: Mapping[str, object]) -> frozenset[str]:
def _registration_write(
key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None
key: str,
value: Mapping[str, object],
custom_llm_provider: str | None,
warning_display_name: str | None,
api_base: str | None,
api_key: str | None,
) -> _RegistrationWrite:
base: Final = _registration_base(key, value, custom_llm_provider, warning_display_name)
base: Final = _registration_base(key, value, custom_llm_provider, warning_display_name, api_base, api_key)
# ``get_model_info`` returns ``litellm_provider: None`` when the
# provider is unknown (e.g. custom deployments registered via
# ``Router.add_deployment``). Persisting that None into
@ -4101,16 +4122,23 @@ def _registration_attempt(
value: Mapping[str, object],
custom_llm_provider: str | None,
warning_display_name: str | None,
api_base: str | None,
api_key: str | None,
final_attempt: bool,
) -> _RegistrationWrite | None:
source: Final = _builtin_cost_map()
write: Final = _registration_write(key, value, custom_llm_provider, warning_display_name)
write: Final = _registration_write(key, value, custom_llm_provider, warning_display_name, api_base, api_key)
_replace_cost_map_entry(_builtin_cost_map() if final_attempt else source, write.key, write.entry, value)
return write if final_attempt or litellm.model_cost is source else None
def _register_on_current_map(
key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None
key: str,
value: Mapping[str, object],
custom_llm_provider: str | None,
warning_display_name: str | None,
api_base: str | None,
api_key: str | None,
) -> _RegistrationWrite:
"""A price data reload that swaps the map while the entry is being built or written would otherwise get the old
map's prices written over its own, or drop the entry with the old map, so each attempt writes into the map it was
@ -4118,7 +4146,13 @@ def _register_on_current_map(
writes into the current map, so the entry is never lost."""
attempts: Final = (
_registration_attempt(
key, value, custom_llm_provider, warning_display_name, final_attempt=attempt == _REGISTRATION_ATTEMPTS - 1
key,
value,
custom_llm_provider,
warning_display_name,
api_base,
api_key,
final_attempt=attempt == _REGISTRATION_ATTEMPTS - 1,
)
for attempt in range(_REGISTRATION_ATTEMPTS)
)
@ -4217,6 +4251,8 @@ def register_model(
warning_display_name: str | None = None,
custom_llm_provider: str | None = None,
add_to_provider_model_sets: bool = True,
api_base: str | None = None,
api_key: str | None = None,
):
"""
Register new / Override existing models (and their pricing) to specific providers.
@ -4249,6 +4285,9 @@ def register_model(
``add_to_provider_model_sets`` adds the keys to the per-provider model sets
behind provider resolution and wildcard model listing. Pass False when the
keys are opaque ids rather than model names (e.g. router deployment ids).
``api_base`` and ``api_key`` send a provider's runtime model-info lookup (an Ollama ``/api/show``)
to the server the registered deployment calls, with its credential, instead of the provider's defaults.
"""
loaded_model_cost = {}
@ -4267,7 +4306,7 @@ def register_model(
)
for key, value in loaded_model_cost.items():
write = _register_on_current_map(str(key), value, custom_llm_provider, warning_display_name)
write = _register_on_current_map(str(key), value, custom_llm_provider, warning_display_name, api_base, api_key)
# Invalidate case-insensitive lookup map since model_cost was modified
_invalidate_model_cost_lowercase_map()
@ -6888,6 +6927,7 @@ def cached_get_model_info_helper(
model: str,
custom_llm_provider: str | None,
api_base: str | None = None,
api_key: str | None = None,
) -> ModelInfoBase:
"""
_get_model_info_helper wrapped with lru_cache
@ -6898,6 +6938,7 @@ def cached_get_model_info_helper(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_key=api_key,
)
@ -7430,13 +7471,19 @@ def _build_model_info(
return ModelInfo(**_model_info, supported_openai_params=supported_openai_params)
def model_info_lookup_api_key(model: str | None, custom_llm_provider: str | None, api_key: str | None) -> str | None:
provider: Final = custom_llm_provider or (model or "").partition("/")[0]
return api_key if provider in PROVIDERS_THAT_SEND_KEYS_ON_MODEL_INFO_LOOKUPS else None
@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
def _cached_get_model_info(
model: str,
custom_llm_provider: str | None = None,
api_base: str | None = None,
api_key: str | None = None,
) -> ModelInfo:
return _build_model_info(model=model, custom_llm_provider=custom_llm_provider, api_base=api_base)
return _build_model_info(model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key)
def get_model_info(
@ -7518,11 +7565,9 @@ def get_model_info(
"supported_openai_params": ["temperature", "max_tokens", "top_p", "frequency_penalty", "presence_penalty"]
}
"""
# api_key is a per-caller credential, not part of the model identity, so it is
# kept out of the cache key; explicit keys are resolved without the cache.
if api_key is not None:
return _build_model_info(model, custom_llm_provider, api_base, api_key)
return _cached_get_model_info(model, custom_llm_provider, api_base)
return _cached_get_model_info(
model, custom_llm_provider, api_base, model_info_lookup_api_key(model, custom_llm_provider, api_key)
)
get_model_info.cache_clear = _cached_get_model_info.cache_clear

30
tests/integration/providers/test_decisions_wire.py
View file

@ -262,23 +262,19 @@ def _upstream_calls(gateway: Gateway, handle: ScenarioHandle) -> list[dict[str,
return _calls_to(_observed_requests(gateway), handle)
def _sdk_inference_call(gateway: Gateway, handle: ScenarioHandle, provider: _Provider) -> dict[str, JsonValue]:
def _inference_call(gateway: Gateway, handle: ScenarioHandle, provider: _Provider) -> dict[str, JsonValue]:
observed: Final = _upstream_calls(gateway, handle)
inference_path: Final = f"/{handle.scenario_id}/v1/systemone"
(call,) = tuple(request for request in observed if request["path"] == inference_path)
metadata_calls: Final = tuple(request for request in observed if request["path"] != inference_path)
assert metadata_calls in (
(),
(
{
"path": f"{inference_path}/api/show",
"authorization": "",
"body": {"name": provider.body_model},
"method": "POST",
"api_key": "",
},
),
), observed
model_info_lookup: Final = {
"path": f"/{handle.scenario_id}/api/show",
"authorization": call["authorization"],
"body": {"name": provider.body_model},
"method": "POST",
"api_key": "",
}
assert all(request == model_info_lookup for request in metadata_calls), observed
return call
@ -442,7 +438,7 @@ def test_local_decisions_routes_preserve_wire_auth_usage_and_spend(
assert response.json() == expected_reply
assert response.headers["x-litellm-model-group"] == model
assert float(response.headers.get("x-litellm-response-cost", "0")) == 0.0
(call,) = _upstream_calls(gateway, handle)
call: Final = _inference_call(gateway, handle, provider)
assert (call["path"], call["authorization"], call["body"]) == (
f"/{handle.scenario_id}/v1/systemone",
f"Bearer {api_key}" if api_key else "",
@ -478,7 +474,7 @@ async def test_local_systemone_sdk_preserves_native_wire(
)
)
assert response.model_dump(mode="json") == _answer_body(provider)
call: Final = _sdk_inference_call(gateway, handle, provider)
call: Final = _inference_call(gateway, handle, provider)
assert (call["path"], call["authorization"], call["body"]) == (
f"/{handle.scenario_id}/v1/systemone",
"",
@ -513,7 +509,7 @@ async def test_local_decisions_sdk_translates_openai_format(
"answers": list(_LOCAL_OPENAI_ANSWERS),
"usage": dict(_LOCAL_OPENAI_USAGE),
}
call: Final = _sdk_inference_call(gateway, handle, provider)
call: Final = _inference_call(gateway, handle, provider)
assert (call["path"], call["authorization"], call["body"]) == (
f"/{handle.scenario_id}/v1/systemone",
"",
@ -556,7 +552,7 @@ def test_ollama_decisions_sdk_respects_environment_and_explicit_auth(
api_key=api_key,
)
assert response.model_dump(mode="json") == _answer_body(provider)
call: Final = _sdk_inference_call(gateway, handle, provider)
call: Final = _inference_call(gateway, handle, provider)
assert (call["path"], call["authorization"], call["body"]) == (
f"/{handle.scenario_id}/v1/systemone",
authorization,

43
tests/unit/decisions/test_main.py
View file

@ -4038,6 +4038,49 @@ async def test_ollama_decisions_resolves_environment_key_and_explicit_connection
assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9)
@pytest.mark.asyncio
@pytest.mark.respx(assert_all_called=False)
@pytest.mark.parametrize("provider", ("ollama", "ollama_chat"))
@pytest.mark.parametrize(
("api_base", "api_key", "authorization"),
(
(None, None, "Bearer environment-key"),
("http://ollama.local:11434/v1", "call-key", "Bearer call-key"),
("http://ollama.local:11434/v1", "fake-api-key", None),
),
ids=("environment", "explicit", "placeholder"),
)
async def test_ollama_decisions_model_info_lookups_go_to_the_call_server_with_the_call_key(
provider: str,
api_base: str | None,
api_key: str | None,
authorization: str | None,
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
from litellm.utils import cached_get_model_info_helper
monkeypatch.setenv("OLLAMA_API_BASE", "http://ignored.local:11434" if api_base else "http://ollama.local:11434")
monkeypatch.setenv("OLLAMA_API_KEY", "environment-key")
route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(json=_RESPONSE)
lookup: Final = respx_mock.post("http://ollama.local:11434/api/show").respond(json={"model_info": {}})
stray: Final = respx_mock.route().respond(404)
litellm.get_model_info.cache_clear()
cached_get_model_info_helper.cache_clear()
try:
await litellm.asystemone(
model=f"{provider}/d3:4b", state="review", questions=_QUESTIONS, api_base=api_base, api_key=api_key
)
finally:
litellm.get_model_info.cache_clear()
cached_get_model_info_helper.cache_clear()
assert route.calls.last.request.headers.get("authorization") == authorization
assert lookup.call_count > 0
assert {call.request.headers.get("authorization") for call in lookup.calls} == {authorization}
assert [str(call.request.url) for call in stray.calls] == []
@pytest.mark.asyncio
@pytest.mark.parametrize("provider", ("ollama", "ollama_chat"))
async def test_ollama_decisions_maps_upstream_bad_requests(

19
tests/unit/litellm_core_utils/test_litellm_logging.py
View file

@ -14,7 +14,7 @@ from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
from datetime import datetime as datetime_standard_logging, datetime as datetime_unit_test, datetime as dt_object
from importlib.machinery import ModuleSpec
from types import MappingProxyType, ModuleType
from typing import Final, Literal
from typing import Final, Literal, cast
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@ -465,6 +465,23 @@ def test_pre_call_tolerates_missing_api_base(logging_obj):
assert logging_obj.model_call_details["litellm_params"]["api_base"] == ""
def test_pre_call_without_a_model_still_logs_the_call():
files_api_logging_obj = LitellmLogging(
model=cast(str, None),
messages=[],
stream=False,
call_type="adelete_file",
start_time=time.time(),
litellm_call_id="12345",
function_id="1245",
)
files_api_logging_obj.update_environment_variables(litellm_params={}, optional_params={})
files_api_logging_obj.pre_call(input="", api_key="call-key", additional_args={"api_base": None, "headers": {}})
assert files_api_logging_obj.model_call_details["log_event_type"] == "pre_api_call"
def test_post_call_serializes_dict_with_datetime(logging_obj):
import datetime

173
tests/unit/llms/ollama/test_ollama_model_info.py
View file

@ -21,6 +21,7 @@ import httpx
import litellm
from litellm.llms.ollama.common_utils import OllamaModelInfo
from litellm.types.utils import AdminApiKeySecretReference
class DummyResponse:
@ -388,16 +389,7 @@ class TestOllamaGetModelInfo:
assert model_info["max_input_tokens"] == 32768
assert captured_headers[0] == {"Authorization": "Bearer explicit-api-key"}
def test_litellm_get_model_info_does_not_cache_on_api_key(self, monkeypatch):
"""Regression: api_key must not be part of the get_model_info cache key.
Distinct api_keys for the same (model, api_base) must not each create their
own cache entry (which would churn the shared LRU cache), and every explicit
key must still reach the backend rather than be served from a result cached
with a different key.
"""
from litellm.utils import _cached_get_model_info
def test_litellm_get_model_info_caches_ollama_lookups_per_api_key(self, monkeypatch):
captured_headers = []
def mock_post(url, json, headers=None):
@ -413,22 +405,35 @@ class TestOllamaGetModelInfo:
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
litellm.get_model_info.cache_clear()
try:
for api_key in ("key-one", "key-two", "key-three"):
for api_key in ("key-one", "key-two", "key-one", "key-two"):
litellm.get_model_info(
"ollama/unknown-model",
api_base="https://ollama.example",
api_key=api_key,
)
assert _cached_get_model_info.cache_info().currsize <= 1
assert captured_headers == [
{"Authorization": "Bearer key-one"},
{"Authorization": "Bearer key-two"},
{"Authorization": "Bearer key-three"},
]
finally:
litellm.get_model_info.cache_clear()
def test_litellm_get_model_info_shares_one_cache_entry_across_keys_when_the_lookup_sends_none(self):
openai_model = next(
name
for name, entry in litellm.model_cost.items()
if entry.get("litellm_provider") == "openai" and entry.get("mode") == "chat"
)
litellm.get_model_info.cache_clear()
try:
for api_key in ("key-one", "key-two", "key-three"):
litellm.get_model_info(openai_model, custom_llm_provider="openai", api_key=api_key)
assert litellm.get_model_info.cache_info().currsize == 1
finally:
litellm.get_model_info.cache_clear()
def test_get_model_info_normalizes_generate_api_base(self, monkeypatch):
"""When completion passes the final generate URL, model info should use the server base."""
from litellm.llms.ollama.completion.transformation import OllamaConfig
@ -576,6 +581,148 @@ class TestOllamaGetModelInfo:
assert model_info["litellm_provider"] == "ollama"
@pytest.mark.parametrize(
("api_key", "environment_key", "authorization"),
[
(None, None, None),
("call-key", "environment-key", "Bearer call-key"),
(None, "environment-key", "Bearer environment-key"),
],
ids=["no-key", "explicit-key", "environment-key"],
)
@pytest.mark.parametrize(
("model", "generation_path", "generation_body"),
[
("ollama/api-base-probe", "/api/generate", {"response": "hi"}),
("ollama_chat/api-base-probe", "/api/chat", {"message": {"role": "assistant", "content": "hi"}}),
],
)
def test_completion_sends_every_model_info_lookup_to_the_call_api_base_with_its_key(
monkeypatch, model, generation_path, generation_body, api_key, environment_key, authorization
):
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.utils import cached_get_model_info_helper
lookups = []
generation_authorizations = []
def respond(request):
if request.url.path == "/api/show":
lookups.append((request.url.host, request.headers.get("authorization")))
return httpx.Response(200, json={"template": "", "model_info": {}})
assert request.url.path == generation_path
generation_authorizations.append(request.headers.get("authorization"))
return httpx.Response(
200,
json={"model": "api-base-probe", "done": True, "prompt_eval_count": 3, "eval_count": 2, **generation_body},
)
transport = httpx.MockTransport(respond)
monkeypatch.setattr(litellm, "module_level_client", HTTPHandler(client=httpx.Client(transport=transport)))
monkeypatch.setattr(litellm, "api_key", None)
monkeypatch.setenv("OLLAMA_API_BASE", "http://ollama-default.test:11434")
if environment_key is None:
monkeypatch.delenv("OLLAMA_API_KEY", raising=False)
else:
monkeypatch.setenv("OLLAMA_API_KEY", environment_key)
client = HTTPHandler(client=httpx.Client(transport=transport))
def complete():
litellm.completion(
model=model,
messages=[{"role": "user", "content": "hi"}],
api_base="http://ollama-call.test:11434",
api_key=api_key,
client=client,
)
litellm.get_model_info.cache_clear()
cached_get_model_info_helper.cache_clear()
try:
complete()
lookups_after_first_call = len(lookups)
complete()
finally:
litellm.get_model_info.cache_clear()
cached_get_model_info_helper.cache_clear()
assert set(generation_authorizations) == {authorization}
assert set(lookups) == {("ollama-call.test", authorization)}
assert len(lookups) == lookups_after_first_call
@pytest.mark.parametrize("reload_catalog", [False, True])
@pytest.mark.parametrize(
"provider_params",
[
{"model": "ollama_chat/registration-probe"},
{"model": "registration-probe", "custom_llm_provider": "ollama_chat"},
],
)
@pytest.mark.parametrize(
("api_key", "authorization"),
[
(None, None),
("deployment-key", "Bearer deployment-key"),
(
AdminApiKeySecretReference(
"oidc/google/https://ollama-deployment.test", frozenset({"http://ollama-deployment.test:11434"})
),
None,
),
],
ids=["no-key", "deployment-key", "secret-reference"],
)
def test_router_registration_sends_model_info_lookup_to_the_deployment_api_base_with_its_key(
monkeypatch, reload_catalog, provider_params, api_key, authorization
):
from litellm import Router
from litellm.llms.custom_httpx.http_handler import HTTPHandler
lookups = []
def respond(request):
lookups.append((request.url.host, request.headers.get("authorization")))
return httpx.Response(200, json={"template": "", "model_info": {}})
transport = httpx.MockTransport(respond)
monkeypatch.setattr(litellm, "module_level_client", HTTPHandler(client=httpx.Client(transport=transport)))
monkeypatch.setenv("OLLAMA_API_BASE", "http://ollama-default.test:11434")
monkeypatch.setenv("OLLAMA_API_KEY", "environment-key")
catalog = dict(litellm.model_cost)
monkeypatch.setattr(litellm, "model_cost", dict(catalog))
litellm.get_model_info.cache_clear()
try:
router = Router(
model_list=[
{
"model_name": "registration-probe",
"litellm_params": {
**provider_params,
"api_base": "http://ollama-deployment.test:11434",
"api_key": api_key,
},
}
]
)
if reload_catalog:
lookups.clear()
monkeypatch.setattr(litellm, "model_cost", dict(catalog))
litellm.get_model_info.cache_clear()
router.replay_model_cost_registrations()
finally:
litellm.get_model_info.cache_clear()
assert set(lookups) == {("ollama-deployment.test", authorization)}
def test_model_info_lookup_for_a_decisions_endpoint_goes_to_the_server_root():
assert (
OllamaModelInfo.get_server_api_base("http://ollama.test:11434/prefix/v1/systemone")
== "http://ollama.test:11434/prefix"
)
class TestOllamaAuthHeaders:
"""Tests for Ollama authentication header handling in completion calls."""

10
tests/unit/llms/reducto/test_cost.py
View file

@ -9,7 +9,7 @@ def test_ocr_cost_prefers_credit_pricing_when_pages_processed_is_none(monkeypatc
monkeypatch.setattr(
litellm,
"get_model_info",
lambda model, custom_llm_provider=None: {"ocr_cost_per_credit": 0.003},
lambda model, custom_llm_provider=None, api_base=None, api_key=None: {"ocr_cost_per_credit": 0.003},
)
response = OCRResponse(
@ -32,7 +32,7 @@ def test_ocr_cost_prefers_zero_credit_pricing_over_page_pricing(monkeypatch):
monkeypatch.setattr(
litellm,
"get_model_info",
lambda model, custom_llm_provider=None: {
lambda model, custom_llm_provider=None, api_base=None, api_key=None: {
"ocr_cost_per_credit": 0.0,
"ocr_cost_per_page": 0.5,
},
@ -58,7 +58,7 @@ def test_ocr_cost_falls_back_to_page_pricing(monkeypatch):
monkeypatch.setattr(
litellm,
"get_model_info",
lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.5},
lambda model, custom_llm_provider=None, api_base=None, api_key=None: {"ocr_cost_per_page": 0.5},
)
response = OCRResponse(
@ -81,7 +81,7 @@ def test_ocr_cost_returns_zero_when_no_pricing_and_no_pages(monkeypatch):
monkeypatch.setattr(
litellm,
"get_model_info",
lambda model, custom_llm_provider=None: {},
lambda model, custom_llm_provider=None, api_base=None, api_key=None: {},
)
response = OCRResponse(
@ -104,7 +104,7 @@ def test_ocr_cost_raises_when_pages_processed_missing_for_page_pricing(monkeypat
monkeypatch.setattr(
litellm,
"get_model_info",
lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.5},
lambda model, custom_llm_provider=None, api_base=None, api_key=None: {"ocr_cost_per_page": 0.5},
)
response = OCRResponse(

30
tests/unit/test_register_model_custom_pricing.py
View file

@ -812,7 +812,9 @@ def test_register_model_prices_a_geo_alias_from_its_builtin_when_the_lookup_fail
alias: Final = f"bedrock/{route}{canonical_key}"
canonical: Final = litellm.model_cost[canonical_key]
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, api_base: str | None = None, api_key: str | None = None
) -> NoReturn:
raise RuntimeError("dictionary changed size during iteration")
with pytest.MonkeyPatch.context() as lookup_patch:
@ -867,7 +869,9 @@ def test_register_model_keeps_a_catalog_rows_own_prices_when_the_lookup_fails(mo
)
_invalidate_model_cost_lowercase_map()
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, api_base: str | None = None, api_key: str | None = None
) -> NoReturn:
raise RuntimeError("dictionary changed size during iteration")
with pytest.MonkeyPatch.context() as lookup_patch:
@ -908,10 +912,12 @@ def test_register_model_leaves_no_shadow_key_when_the_lookup_trips_once(
real_lookup: Final = litellm.utils.get_model_info
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, api_base: str | None = None, api_key: str | None = None
) -> 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, api_base=api_base, api_key=api_key)
with pytest.MonkeyPatch.context() as lookup_patch:
lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_once_by_a_concurrent_registration)
@ -944,7 +950,9 @@ 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, api_base: str | None = None, api_key: str | None = None
) -> NoReturn:
raise RuntimeError("dictionary changed size during iteration")
def long_prompt_cost_after_registering_with(lookup: Callable[..., object]) -> float:
@ -999,7 +1007,9 @@ def test_register_model_prices_a_provider_less_azure_ai_key_without_repricing_th
)
_invalidate_model_cost_lowercase_map()
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, api_base: str | None = None, api_key: str | None = None
) -> NoReturn:
raise RuntimeError("dictionary changed size during iteration")
with pytest.MonkeyPatch.context() as lookup_patch:
@ -1039,8 +1049,12 @@ def test_register_model_keeps_the_prices_a_reload_adopts_while_it_registers(monk
}
real_lookup: Final = litellm.utils.get_model_info
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, api_base: str | None = None, api_key: str | None = None
) -> object:
info: Final = real_lookup(
model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key
)
if litellm.model_cost is not reloaded_cost_map:
adopt_model_cost_map(reloaded_cost_map)
return info