diff --git a/litellm/constants.py b/litellm/constants.py index e7c9c5dcbfb..d7606f948f6 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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 = [ diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index ee102e82ca4..551f082c2af 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -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: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4dd33fa5565..7b7697cd3d3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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( diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 404dfab0a6d..da9af1ec2ff 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -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) diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index e8e4b77e757..610bf5b7cdc 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -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) diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index f313c62f026..bba8aaa11b0 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -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)] diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index 826de8cfdc9..3fd3b912faf 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -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) diff --git a/litellm/router.py b/litellm/router.py index 20af2e5649e..eb7fa39bf4c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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() diff --git a/litellm/utils.py b/litellm/utils.py index e85e089289a..a61e8eb054f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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 diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 5af1532e9ed..70aedad5717 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -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, diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 0c357837d31..3259e5cc89f 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -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( diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index b3ae3bf6240..35a8cbb028c 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -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 diff --git a/tests/unit/llms/ollama/test_ollama_model_info.py b/tests/unit/llms/ollama/test_ollama_model_info.py index 053d4da035f..93fb29e5ebc 100644 --- a/tests/unit/llms/ollama/test_ollama_model_info.py +++ b/tests/unit/llms/ollama/test_ollama_model_info.py @@ -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.""" diff --git a/tests/unit/llms/reducto/test_cost.py b/tests/unit/llms/reducto/test_cost.py index 73340dc8729..79a9575ab00 100644 --- a/tests/unit/llms/reducto/test_cost.py +++ b/tests/unit/llms/reducto/test_cost.py @@ -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( diff --git a/tests/unit/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py index 35fbd28eae7..9b8791988ea 100644 --- a/tests/unit/test_register_model_custom_pricing.py +++ b/tests/unit/test_register_model_custom_pricing.py @@ -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