From f763876b07bbc669415f2e1ddb0320f441e03641 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 3 Feb 2026 16:51:45 -0800 Subject: [PATCH] fix: add cache invalidation for _cached_get_model_group_info on deployment changes _cached_get_model_group_info uses @lru_cache but had no invalidation, causing stale model group info (TPM/RPM limits) after dynamic deployment changes. Add cache_clear() at all 5 model_list mutation sites. --- litellm/router.py | 670 +++++------------------------- tests/test_litellm/test_router.py | 620 ++------------------------- 2 files changed, 137 insertions(+), 1153 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 0fc9ecc9bb5..5b72c3fb669 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -33,7 +33,6 @@ from typing import ( cast, ) -import anyio import httpx import openai from openai import AsyncOpenAI @@ -59,6 +58,7 @@ from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, get_metadata_variable_name_from_kwargs, ) +from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.dd_tracing import tracer @@ -111,18 +111,15 @@ from litellm.router_utils.handle_error import ( async_raise_no_deployment_exception, send_llm_exception_alert, ) -from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( - DeploymentAffinityCheck, +from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( + PromptCachingDeploymentCheck, ) -from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( - EncryptedContentAffinityCheck, +from litellm.router_utils.pre_call_checks.responses_api_deployment_check import ( + ResponsesApiDeploymentCheck, ) from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) -from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( - PromptCachingDeploymentCheck, -) from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_failures_for_current_minute, increment_deployment_successes_for_current_minute, @@ -193,15 +190,11 @@ if TYPE_CHECKING: AutoRouter, PreRoutingHookResponse, ) - from litellm.router_strategy.complexity_router.complexity_router import ( - ComplexityRouter, - ) Span = Union[_Span, Any] else: Span = Any AutoRouter = Any - ComplexityRouter = Any PreRoutingHookResponse = Any @@ -301,7 +294,6 @@ class Router: router_general_settings: Optional[ RouterGeneralSettings ] = RouterGeneralSettings(), - deployment_affinity_ttl_seconds: int = 3600, ignore_invalid_deployments: bool = False, ) -> None: """ @@ -335,7 +327,6 @@ class Router: routing_strategy_args (dict): Additional args for latency-based routing. Defaults to {}. alerting_config (AlertingConfig): Slack alerting configuration. Defaults to None. provider_budget_config (ProviderBudgetConfig): Provider budget configuration. Use this to set llm_provider budget limits. example $100/day to OpenAI, $100/day to Azure, etc. Defaults to None. - deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600. ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error. Returns: Router: An instance of the litellm.Router class. @@ -455,7 +446,6 @@ class Router: str, PatternMatchRouter ] = {} # {"TEAM_ID": PatternMatchRouter} self.auto_routers: Dict[str, "AutoRouter"] = {} - self.complexity_routers: Dict[str, "ComplexityRouter"] = {} # Initialize model_group_alias early since it's used in set_model_list self.model_group_alias: Dict[str, Union[str, RouterModelGroupAliasItem]] = ( @@ -480,8 +470,6 @@ class Router: [] ) # initialize an empty list - to allow _add_deployment and delete_deployment to work - self._access_groups_cache: Optional[Dict[str, List[str]]] = None - if allowed_fails is not None: self.allowed_fails = allowed_fails else: @@ -617,7 +605,6 @@ class Router: litellm.failure_callback = [self.deployment_callback_on_failure] self.routing_strategy_args = routing_strategy_args self.provider_budget_config = provider_budget_config - self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds self.router_budget_logger: Optional[RouterBudgetLimiting] = None if RouterBudgetLimiting.should_init_router_budget_limiter( model_list=model_list, provider_budget_config=self.provider_budget_config @@ -632,12 +619,11 @@ class Router: self.retry_policy = RetryPolicy(**retry_policy) elif isinstance(retry_policy, RetryPolicy): self.retry_policy = retry_policy - if self.retry_policy is not None: - verbose_router_logger.info( - "\033[32mRouter Custom Retry Policy Set:\n{}\033[0m".format( - self.retry_policy.model_dump(exclude_none=True) - ) + verbose_router_logger.info( + "\033[32mRouter Custom Retry Policy Set:\n{}\033[0m".format( + self.retry_policy.model_dump(exclude_none=True) ) + ) self.model_group_retry_policy: Optional[ Dict[str, RetryPolicy] @@ -650,12 +636,11 @@ class Router: elif isinstance(allowed_fails_policy, AllowedFailsPolicy): self.allowed_fails_policy = allowed_fails_policy - if self.allowed_fails_policy is not None: - verbose_router_logger.info( - "\033[32mRouter Custom Allowed Fails Policy Set:\n{}\033[0m".format( - self.allowed_fails_policy.model_dump(exclude_none=True) - ) + verbose_router_logger.info( + "\033[32mRouter Custom Allowed Fails Policy Set:\n{}\033[0m".format( + self.allowed_fails_policy.model_dump(exclude_none=True) ) + ) self.alerting_config: Optional[AlertingConfig] = alerting_config @@ -910,11 +895,14 @@ class Router: def _initialize_vector_store_endpoints(self): """Initialize vector store endpoints.""" - from litellm.vector_stores.main import asearch, create, search + from litellm.vector_stores.main import acreate, asearch, create, search self.avector_store_search = self.factory_function( asearch, call_type="avector_store_search" ) + self.avector_store_create = self.factory_function( + acreate, call_type="avector_store_create" + ) self.vector_store_search = self.factory_function( search, call_type="vector_store_search" ) @@ -1170,8 +1158,6 @@ class Router: self._initialize_vector_store_file_endpoints() self._initialize_google_genai_endpoints() self._initialize_ocr_search_endpoints() - # Override vector store methods with router-aware implementations - self._override_vector_store_methods_for_router() self._initialize_video_endpoints() self._initialize_container_endpoints() self._initialize_skills_endpoints() @@ -1198,104 +1184,26 @@ class Router: def add_optional_pre_call_checks( self, optional_pre_call_checks: Optional[OptionalPreCallChecks] ): - if optional_pre_call_checks is None: - return - - # --------------------------------------------------------------------- - # Unified deployment affinity (session stickiness) - # --------------------------------------------------------------------- - enable_user_key_affinity = "deployment_affinity" in optional_pre_call_checks - enable_responses_api_affinity = ( - "responses_api_deployment_check" in optional_pre_call_checks - ) - enable_session_id_affinity = "session_affinity" in optional_pre_call_checks - if ( - enable_user_key_affinity - or enable_responses_api_affinity - or enable_session_id_affinity - ): - if self.optional_callbacks is None: - self.optional_callbacks = [] - - existing_affinity_callback: Optional[DeploymentAffinityCheck] = None - for cb in self.optional_callbacks: - if isinstance(cb, DeploymentAffinityCheck): - existing_affinity_callback = cb - break - - if existing_affinity_callback is not None: - existing_affinity_callback.enable_user_key_affinity = ( - existing_affinity_callback.enable_user_key_affinity - or enable_user_key_affinity - ) - existing_affinity_callback.enable_responses_api_affinity = ( - existing_affinity_callback.enable_responses_api_affinity - or enable_responses_api_affinity - ) - existing_affinity_callback.enable_session_id_affinity = ( - existing_affinity_callback.enable_session_id_affinity - or enable_session_id_affinity - ) - existing_affinity_callback.ttl_seconds = ( - self.deployment_affinity_ttl_seconds - ) - else: - affinity_callback = DeploymentAffinityCheck( - cache=self.cache, - ttl_seconds=self.deployment_affinity_ttl_seconds, - enable_user_key_affinity=enable_user_key_affinity, - enable_responses_api_affinity=enable_responses_api_affinity, - enable_session_id_affinity=enable_session_id_affinity, - ) - self.optional_callbacks.append(affinity_callback) - litellm.logging_callback_manager.add_litellm_callback(affinity_callback) - - # --------------------------------------------------------------------- - # Encrypted content affinity - # --------------------------------------------------------------------- - if "encrypted_content_affinity" in optional_pre_call_checks: - if self.optional_callbacks is None: - self.optional_callbacks = [] - - already_registered = any( - isinstance(cb, EncryptedContentAffinityCheck) - for cb in self.optional_callbacks - ) - if not already_registered: - ec_callback = EncryptedContentAffinityCheck() - self.optional_callbacks.append(ec_callback) - litellm.logging_callback_manager.add_litellm_callback(ec_callback) - - # --------------------------------------------------------------------- - # Remaining optional pre-call checks - # --------------------------------------------------------------------- - for pre_call_check in optional_pre_call_checks: - _callback: Optional[CustomLogger] = None - if pre_call_check in ( - "deployment_affinity", - "responses_api_deployment_check", - "session_affinity", - "encrypted_content_affinity", - ): - continue - if pre_call_check == "prompt_caching": - _callback = PromptCachingDeploymentCheck(cache=self.cache) - elif pre_call_check == "router_budget_limiting": - _callback = RouterBudgetLimiting( - dual_cache=self.cache, - provider_budget_config=self.provider_budget_config, - model_list=self.model_list, - ) - elif pre_call_check == "enforce_model_rate_limits": - _callback = ModelRateLimitingCheck(dual_cache=self.cache) - - if _callback is None: - continue - - if self.optional_callbacks is None: - self.optional_callbacks = [] - self.optional_callbacks.append(_callback) - litellm.logging_callback_manager.add_litellm_callback(_callback) + if optional_pre_call_checks is not None: + for pre_call_check in optional_pre_call_checks: + _callback: Optional[CustomLogger] = None + if pre_call_check == "prompt_caching": + _callback = PromptCachingDeploymentCheck(cache=self.cache) + elif pre_call_check == "router_budget_limiting": + _callback = RouterBudgetLimiting( + dual_cache=self.cache, + provider_budget_config=self.provider_budget_config, + model_list=self.model_list, + ) + elif pre_call_check == "responses_api_deployment_check": + _callback = ResponsesApiDeploymentCheck() + elif pre_call_check == "enforce_model_rate_limits": + _callback = ModelRateLimitingCheck(dual_cache=self.cache) + if _callback is not None: + if self.optional_callbacks is None: + self.optional_callbacks = [] + self.optional_callbacks.append(_callback) + litellm.logging_callback_manager.add_litellm_callback(_callback) def print_deployment(self, deployment: dict): """ @@ -1361,20 +1269,19 @@ class Router: if silent_model is not None: # Mirroring traffic to a secondary model - # Use threading.Thread (not ThreadPoolExecutor) - executor.submit() - # requires pickling args, which fails when kwargs contain unpicklable - # objects (e.g. _thread.RLock from OTEL spans, loggers) in deployment. - thread = threading.Thread( - target=self._silent_experiment_completion, - args=(silent_model, messages), - kwargs=kwargs, - daemon=True, + # Use shared thread pool for background calls + executor.submit( + self._silent_experiment_completion, + silent_model, + messages, + **kwargs, ) - thread.start() self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) kwargs.pop("silent_model", None) # Ensure it's not in kwargs either - model_name = litellm_params["model"] + # No copy needed - data is only read and spread into new dict below + data = litellm_params.copy() # Use the local copy of litellm_params + model_name = data["model"] potential_model_client = self._get_client( deployment=deployment, kwargs=kwargs ) @@ -1395,7 +1302,7 @@ class Router: self.routing_strategy_pre_call_checks(deployment=deployment) input_kwargs = { - **litellm_params, + **data, "messages": messages, "caching": self.cache_responses, "client": model_client, @@ -1556,7 +1463,7 @@ class Router: ) raise e - async def _acompletion_streaming_iterator( # noqa: PLR0915 + async def _acompletion_streaming_iterator( self, model_response: CustomStreamWrapper, messages: List[Dict[str, str]], @@ -1579,9 +1486,6 @@ class Router: logging_obj=model_response.logging_obj, ) self._async_generator = async_generator - # Preserve hidden params (including litellm_overhead_time_ms) from original response - if hasattr(model_response, "_hidden_params"): - self._hidden_params = model_response._hidden_params.copy() def __aiter__(self): return self @@ -1590,7 +1494,6 @@ class Router: return await self._async_generator.__anext__() async def stream_with_fallbacks(): - fallback_response = None # Track for cleanup in finally try: async for item in model_response: yield item @@ -1689,30 +1592,6 @@ class Router: f"Fallback also failed: {fallback_error}" ) raise fallback_error - finally: - # Close the underlying streams to release HTTP connections - # back to the connection pool when the generator is closed - # (e.g. on client disconnect). - # Shield from anyio cancellation so the awaits can complete. - with anyio.CancelScope(shield=True): - if hasattr(model_response, "aclose"): - try: - await model_response.aclose() - except BaseException as e: - verbose_router_logger.debug( - "stream_with_fallbacks: error closing model_response: %s", - e, - ) - if fallback_response is not None and hasattr( - fallback_response, "aclose" - ): - try: - await fallback_response.aclose() - except BaseException as e: - verbose_router_logger.debug( - "stream_with_fallbacks: error closing fallback_response: %s", - e, - ) return FallbackStreamWrapper(stream_with_fallbacks()) @@ -1811,8 +1690,10 @@ class Router: self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) kwargs.pop("silent_model", None) # Ensure it's not in kwargs either + # No copy needed - data is only read and spread into new dict below + data = litellm_params.copy() # Use the local copy of litellm_params - model_name = litellm_params["model"] + model_name = data["model"] model_client = self._get_async_openai_model_client( deployment=deployment, @@ -1821,7 +1702,7 @@ class Router: self.total_calls[model_name] += 1 input_kwargs = { - **litellm_params, + **data, "messages": messages, "caching": self.cache_responses, "client": model_client, @@ -2012,26 +1893,6 @@ class Router: ) # add new deployment to router return deployment_pydantic_obj - @staticmethod - def _merge_tools_from_deployment(deployment: dict, kwargs: dict) -> None: - """ - Merge tools from deployment litellm_params with request kwargs. - When both have tools, concatenate them (deployment tools first, then request tools). - tool_choice: use request value if provided, else deployment's. - """ - dep_params_raw = deployment.get("litellm_params", {}) or {} - if isinstance(dep_params_raw, dict): - dep_params = dep_params_raw - else: - dep_params = dep_params_raw.model_dump(exclude_none=True) - dep_tools = dep_params.get("tools") or [] - req_tools = kwargs.get("tools") or [] - if dep_tools or req_tools: - merged = list(dep_tools) + list(req_tools) - kwargs["tools"] = merged - if "tool_choice" not in kwargs and dep_params.get("tool_choice") is not None: - kwargs["tool_choice"] = dep_params["tool_choice"] - def _update_kwargs_with_deployment( self, deployment: dict, @@ -2039,13 +1900,10 @@ class Router: function_name: Optional[str] = None, ) -> None: """ - 3 jobs: + 2 jobs: - Adds selected deployment, model_info and api_base to kwargs["metadata"] (used for logging) - Adds default litellm params to kwargs, if set. - - Merges tools from deployment with request (proxy-configured tools + request tools). """ - self._merge_tools_from_deployment(deployment=deployment, kwargs=kwargs) - model_info = deployment.get("model_info", {}).copy() deployment_litellm_model_name = deployment["litellm_params"]["model"] deployment_api_base = deployment["litellm_params"].get("api_base") @@ -2070,28 +1928,6 @@ class Router: "deployment_model_name": deployment_model_name, } ) - - ## DEPLOYMENT-LEVEL TAGS - deployment_tags = deployment.get("litellm_params", {}).get("tags") - if deployment_tags: - existing_tags = kwargs[metadata_variable_name].get("tags") or [] - merged_tags = list(existing_tags) - for tag in deployment_tags: - if tag not in merged_tags: - merged_tags.append(tag) - kwargs[metadata_variable_name]["tags"] = merged_tags - - ## CREDENTIAL NAME AS TAG - credential_name = deployment.get("litellm_params", {}).get( - "litellm_credential_name" - ) - if credential_name: - credential_tag = f"Credential: {credential_name}" - existing_tags = kwargs[metadata_variable_name].get("tags") or [] - if credential_tag not in existing_tags: - existing_tags.append(credential_tag) - kwargs[metadata_variable_name]["tags"] = existing_tags - kwargs["model_info"] = model_info kwargs["timeout"] = self._get_timeout( @@ -2436,7 +2272,7 @@ class Router: item = FlowItem( priority=priority, # 👈 SET PRIORITY FOR REQUEST request_id=_request_id, # 👈 SET REQUEST ID - model_name=model, # 👈 SAME as 'Router' + model_name="gpt-3.5-turbo", # 👈 SAME as 'Router' ) ### [fin] ### @@ -2478,10 +2314,6 @@ class Router: setattr(e, "priority", priority) raise e else: - # Clean up the request from the scheduler queue also before raising the timeout exception - await self.scheduler.remove_request( - request_id=item.request_id, model_name=item.model_name - ) raise litellm.Timeout( message="Request timed out while polling queue", model=model, @@ -2543,10 +2375,6 @@ class Router: setattr(e, "priority", priority) raise e else: - # Clean up the request from the scheduler queue also before raising the timeout exception - await self.scheduler.remove_request( - request_id=item.request_id, model_name=item.model_name - ) raise litellm.Timeout( message="Request timed out while polling queue", model=model, @@ -2595,12 +2423,6 @@ class Router: litellm_model = data.get("model", None) - # litellm_agent/ prefix only strips the model name, no prompt_id needed - is_litellm_agent_model = ( - isinstance(litellm_model, str) - and litellm_model.startswith("litellm_agent/") - ) - prompt_id = kwargs.get("prompt_id") or prompt_management_deployment[ "litellm_params" ].get("prompt_id", None) @@ -2613,9 +2435,7 @@ class Router: "litellm_params" ].get("prompt_label", None) - if not is_litellm_agent_model and ( - prompt_id is None or not isinstance(prompt_id, str) - ): + if prompt_id is None or not isinstance(prompt_id, str): raise ValueError( f"Prompt ID is not set or not a string. Got={prompt_id}, type={type(prompt_id)}" ) @@ -3849,7 +3669,7 @@ class Router: ) raise e - async def _acreate_file( # noqa: PLR0915 + async def _acreate_file( self, model: str, **kwargs, @@ -3910,12 +3730,7 @@ class Router: ) kwargs_copy["file"] = file - if ( - "gcs_bucket_name" in data - ): # TODO: Remove this once we have a better way to handle GCS bucket name: Problem is that we need to pass the gcs_bucket_name to the router for the create_file call but it doesn't show up there - kwargs_copy.setdefault("litellm_metadata", {})[ - "gcs_bucket_name" - ] = data["gcs_bucket_name"] + response = litellm.acreate_file( **{ **data, @@ -3986,114 +3801,6 @@ class Router: self.fail_calls[model] += 1 raise e - #### VECTOR STORES API #### - async def avector_store_create( - self, - model: Union[str, None], - **kwargs, - ): - """ - Create a vector store for a specific model. - - Args: - model: Model name from router config - **kwargs: Vector store creation parameters - - Returns: - VectorStoreCreateResponse - """ - try: - # If model is None, use the factory function approach (direct SDK call) - if model is None: - from litellm.vector_stores.main import acreate - - # Use the factory function to handle the call - factory_fn = self.factory_function( - acreate, call_type="avector_store_create" - ) - return await factory_fn(**kwargs) - - from litellm.vector_stores import acreate as avector_store_create_sdk - - parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) - deployment = await self.async_get_available_deployment( - model=model, - messages=[{"role": "user", "content": "vector-store-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - data = deployment["litellm_params"].copy() - model_name = data["model"] - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="avector_store_create", - ) - - model_client = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - # Get custom provider - _, custom_llm_provider, _, _ = get_llm_provider(model=data["model"]) - - response = avector_store_create_sdk( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self.cache_responses, - "client": model_client, - **kwargs, - } - ) - - rpm_semaphore = self._get_client( - deployment=deployment, - kwargs=kwargs, - client_type="max_parallel_requests", - ) - - if rpm_semaphore is not None and isinstance( - rpm_semaphore, asyncio.Semaphore - ): - async with rpm_semaphore: - await self.async_routing_strategy_pre_call_checks( - deployment=deployment, parent_otel_span=parent_otel_span - ) - response = await response - else: - await self.async_routing_strategy_pre_call_checks( - deployment=deployment, parent_otel_span=parent_otel_span - ) - response = await response - - self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.avector_store_create(model={model_name})\033[32m 200 OK\033[0m" - ) - - return response - except Exception as e: - verbose_router_logger.exception( - f"litellm.avector_store_create(model={model})\033[31m Exception {str(e)}\033[0m" - ) - if model is not None: - self.fail_calls[model] += 1 - raise e - - def _override_vector_store_methods_for_router(self): - """ - Override factory-generated vector store methods with router-aware implementations. - This is called after _initialize_vector_store_endpoints() to ensure our custom - methods that handle deployment selection and credential injection are used instead - of the generic factory-generated ones. - """ - # Store references to the custom methods defined above - # These methods handle proper routing through deployments - pass # The methods are already defined as instance methods above - async def acreate_batch( self, model: str, @@ -4800,21 +4507,9 @@ class Router: ): """ Initialize the Vector Store API endpoints on the router. - - If a model is provided in kwargs, use model-based routing to get - the deployment credentials. Otherwise, call the original function directly. """ if custom_llm_provider and "custom_llm_provider" not in kwargs: kwargs["custom_llm_provider"] = custom_llm_provider - - # If model is provided, use generic API call with fallbacks for proper routing - if kwargs.get("model"): - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - **kwargs, - ) - - # Otherwise, call the original function directly return await original_function(**kwargs) async def _init_containers_api_endpoints( @@ -5120,8 +4815,7 @@ class Router: ) else: response = await self.async_function_with_retries(*args, **kwargs) - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug(f"Async Response: {response}") + verbose_router_logger.debug(f"Async Response: {response}") response = add_fallback_headers_to_response( response=response, attempted_fallbacks=0, @@ -5198,10 +4892,6 @@ class Router: content_policy_fallbacks = kwargs.pop( "content_policy_fallbacks", self.content_policy_fallbacks ) - # Support per-request model_group_retry_policy override (from key/team settings) - model_group_retry_policy = kwargs.pop( - "model_group_retry_policy", self.model_group_retry_policy - ) model_group: Optional[str] = kwargs.get("model") num_retries = kwargs.pop("num_retries") @@ -5215,11 +4905,6 @@ class Router: verbose_router_logger.debug( f"async function w/ retries: original_function - {original_function}, num_retries - {num_retries}" ) - ## ADD RETRY TRACKING TO METADATA - used for spend logs retry tracking - _metadata["attempted_retries"] = 0 - _metadata[ - "max_retries" - ] = num_retries # Updated after overrides in exception handler try: self._handle_mock_testing_rate_limit_error( model_group=model_group, kwargs=kwargs @@ -5253,19 +4938,19 @@ class Router: # Check retry policy FIRST, before should_retry_this_error # This allows retry policies to override the healthy deployments check _retry_policy_applies = False - if self.retry_policy is not None or model_group_retry_policy is not None: + if ( + self.retry_policy is not None + or self.model_group_retry_policy is not None + ): # get num_retries from retry policy # Use the model_group captured at the start of the function, or get it from metadata # kwargs.get("model") at this point is the deployment model, not the model_group _model_group_for_retry_policy = ( model_group or _metadata.get("model_group") or kwargs.get("model") ) - # Use per-request model_group_retry_policy if provided, otherwise use self - _retry_policy_retries = _get_num_retries_from_retry_policy( + _retry_policy_retries = self.get_num_retries_from_retry_policy( exception=original_exception, model_group=_model_group_for_retry_policy, - model_group_retry_policy=model_group_retry_policy, - retry_policy=self.retry_policy, ) if _retry_policy_retries is not None: num_retries = _retry_policy_retries @@ -5282,9 +4967,6 @@ class Router: regular_fallbacks=fallbacks, content_policy_fallbacks=content_policy_fallbacks, ) - # Update max_retries after overrides (deployment_num_retries / retry_policy) - _metadata["max_retries"] = num_retries - ## LOGGING if num_retries > 0: kwargs = self.log_retry(kwargs=kwargs, e=original_exception) @@ -5307,9 +4989,6 @@ class Router: for current_attempt in range(num_retries): try: - # Update retry tracking metadata before each retry attempt - _metadata["attempted_retries"] = current_attempt + 1 - _metadata["max_retries"] = num_retries # if the function call is successful, no exception will be raised and we'll break out of the loop response = await self.make_call(original_function, *args, **kwargs) if coroutine_checker.is_async_callable( @@ -5340,7 +5019,7 @@ class Router: else: _healthy_deployments = [] _timeout = self._time_to_sleep_before_retry( - e=e, + e=original_exception, remaining_retries=remaining_retries, num_retries=num_retries, healthy_deployments=_healthy_deployments, @@ -5625,7 +5304,7 @@ class Router: return else: deployment_model_info = self.get_router_model_info( - deployment=deployment_info, + deployment=deployment_info.model_dump(), received_model_name=model_group, ) # get tpm/rpm from deployment info @@ -6241,18 +5920,9 @@ class Router: deployment.litellm_params.custom_llm_provider + "/" + _model_name ) - # For the shared backend key, strip custom pricing fields so that - # one deployment's pricing overrides don't pollute another - # deployment sharing the same backend model name. - # Each deployment's full pricing is already stored under its - # unique model_id above. - _custom_pricing_fields = CustomPricingLiteLLMParams.model_fields.keys() - _shared_model_info = { - k: v for k, v in _model_info.items() if k not in _custom_pricing_fields - } litellm.register_model( model_cost={ - _model_name: _shared_model_info, + _model_name: _model_info, } ) @@ -6285,13 +5955,10 @@ class Router: def _is_auto_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: """ - Check if the deployment is an auto-router deployment (semantic router). + Check if the deployment is an auto-router deployment. Returns True if the litellm_params model starts with "auto_router/" - but NOT "auto_router/complexity_router" (which uses complexity routing). """ - if litellm_params.model.startswith("auto_router/complexity_router"): - return False # This is handled by complexity_router if litellm_params.model.startswith("auto_router/"): return True return False @@ -6343,61 +6010,6 @@ class Router: ) self.auto_routers[deployment.model_name] = autor_router - def _is_complexity_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: - """ - Check if the deployment is a complexity-router deployment. - - Returns True if the litellm_params model starts with "auto_router/complexity_router" - """ - if litellm_params.model.startswith("auto_router/complexity_router"): - return True - return False - - def init_complexity_router_deployment(self, deployment: Deployment): - """ - Initialize the complexity-router deployment. - - This will initialize the complexity-router and add it to the complexity-routers dictionary. - """ - # Import here to avoid circular imports — ComplexityRouter is a CustomLogger - # subclass that imports litellm internals which depend on router.py. - # This matches the AutoRouter pattern in init_auto_router_deployment above. - from litellm.router_strategy.complexity_router.complexity_router import ( - ComplexityRouter, - ) - - complexity_router_config: Optional[ - dict - ] = deployment.litellm_params.complexity_router_config - - default_model: Optional[ - str - ] = deployment.litellm_params.complexity_router_default_model - - # If no default model specified, try to get from config tiers - if default_model is None and complexity_router_config: - tiers = complexity_router_config.get("tiers", {}) - # Use MEDIUM tier as fallback default - default_model = tiers.get("MEDIUM") or tiers.get("SIMPLE") - - if default_model is None: - raise ValueError( - "complexity_router_default_model is required for complexity-router deployments, " - "or configure tiers in complexity_router_config. Please set it in the litellm_params" - ) - - complexity_router: ComplexityRouter = ComplexityRouter( - model_name=deployment.model_name, - default_model=default_model, - litellm_router_instance=self, - complexity_router_config=complexity_router_config, - ) - if deployment.model_name in self.complexity_routers: - raise ValueError( - f"Complexity-router deployment {deployment.model_name} already exists. Please use a different model name." - ) - self.complexity_routers[deployment.model_name] = complexity_router - def deployment_is_active_for_environment(self, deployment: Deployment) -> bool: """ Function to check if a llm deployment is active for a given environment. Allows using the same config.yaml across multople environments @@ -6443,7 +6055,7 @@ class Router: self.model_list = [] self.model_id_to_deployment_index_map = {} # Reset the index self.model_name_to_deployment_indices = {} # Reset the model_name index - self._invalidate_access_groups_cache() + self._invalidate_model_group_info_cache() # we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works for model in original_model_list: @@ -6608,12 +6220,6 @@ class Router: if self._is_auto_router_deployment(litellm_params=deployment.litellm_params): self.init_auto_router_deployment(deployment=deployment) - ######################################################### - # Check if this is a complexity-router deployment - ######################################################### - if self._is_complexity_router_deployment(litellm_params=deployment.litellm_params): - self.init_complexity_router_deployment(deployment=deployment) - return deployment def _initialize_deployment_for_pass_through( @@ -6753,7 +6359,7 @@ class Router: """ idx = len(self.model_list) self.model_list.append(model) - self._invalidate_access_groups_cache() + self._invalidate_model_group_info_cache() # Update model_id index for O(1) lookup if model_id is not None: @@ -6801,7 +6407,7 @@ class Router: if removal_idx is not None: self.model_list.pop(removal_idx) - self._invalidate_access_groups_cache() + self._invalidate_model_group_info_cache() self._update_deployment_indices_after_removal( model_id=deployment_id, removal_idx=removal_idx ) @@ -6835,7 +6441,7 @@ class Router: if deployment_idx is not None: # Pop the item from the list first item = self.model_list.pop(deployment_idx) - self._invalidate_access_groups_cache() + self._invalidate_model_group_info_cache() self._update_deployment_indices_after_removal( model_id=id, removal_idx=deployment_idx ) @@ -6937,19 +6543,6 @@ class Router: **deployment.litellm_params.model_dump(exclude_none=True) ).model_dump(exclude_none=True) - # Resolve litellm_credential_name to actual credentials - if deployment.litellm_params.litellm_credential_name is not None: - credential_values = CredentialAccessor.get_credential_values( - deployment.litellm_params.litellm_credential_name - ) - if not credential_values: - verbose_router_logger.warning( - f"Credential '{deployment.litellm_params.litellm_credential_name}' not found in credential_list" - ) - credentials.update(credential_values) - # Remove the credential name since we've resolved it - credentials.pop("litellm_credential_name", None) - # Add custom_llm_provider if deployment.litellm_params.custom_llm_provider: credentials[ @@ -6967,7 +6560,7 @@ class Router: @overload def get_router_model_info( - self, deployment: Union[dict, "Deployment"], received_model_name: str, id: None = None + self, deployment: dict, received_model_name: str, id: None = None ) -> ModelMapInfo: pass @@ -6979,7 +6572,7 @@ class Router: def get_router_model_info( self, - deployment: Optional[Union[dict, "Deployment"]], + deployment: Optional[dict], received_model_name: str, id: Optional[str] = None, ) -> ModelMapInfo: @@ -6999,34 +6592,22 @@ class Router: if id is not None: _deployment = self.get_deployment(model_id=id) if _deployment is not None: - deployment = _deployment + deployment = _deployment.model_dump(exclude_none=True) if deployment is None: raise ValueError("Deployment not found") ## GET BASE MODEL - base_model = (deployment.get("model_info") or {}).get("base_model", None) + base_model = deployment.get("model_info", {}).get("base_model", None) if base_model is None: - base_model = (deployment.get("litellm_params") or {}).get("base_model", None) + base_model = deployment.get("litellm_params", {}).get("base_model", None) model = base_model - ## GET PROVIDER - reuse LiteLLM_Params if already constructed - litellm_params_data = deployment.get("litellm_params") - litellm_params: LiteLLM_Params - if isinstance(litellm_params_data, LiteLLM_Params): - litellm_params = litellm_params_data - elif isinstance(litellm_params_data, dict) and "model" in litellm_params_data: - litellm_params = LiteLLM_Params(**litellm_params_data) - else: - raise ValueError( - f"Deployment missing valid litellm_params. " - f"Got: {type(litellm_params_data).__name__}, " - f"deployment_id: {(deployment.get('model_info') or {}).get('id', 'unknown')}" - ) + ## GET PROVIDER _model, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=litellm_params.model, - litellm_params=litellm_params, + model=deployment.get("litellm_params", {}).get("model", ""), + litellm_params=LiteLLM_Params(**deployment.get("litellm_params", {})), ) ## SET MODEL TO 'model=' - if base_model is None + not azure @@ -7042,10 +6623,10 @@ class Router: if potential_models is not None: for potential_model in potential_models: try: - if (potential_model.get("model_info") or {}).get( + if potential_model.get("model_info", {}).get( "id" - ) == (deployment.get("model_info") or {}).get("id"): - model = (potential_model.get("litellm_params") or {}).get( + ) == deployment.get("model_info", {}).get("id"): + model = potential_model.get("litellm_params", {}).get( "model" ) break @@ -7066,10 +6647,9 @@ class Router: model_info = litellm.get_model_info(model=model_info_name) ## CHECK USER SET MODEL INFO - user_model_info = deployment.get("model_info") or {} + user_model_info = deployment.get("model_info", {}) - if model_info is not None: - model_info.update(user_model_info) + model_info.update(user_model_info) return model_info @@ -7596,7 +7176,7 @@ class Router: """ # First populate the model_list self.model_list = [] - self._invalidate_access_groups_cache() + self._invalidate_model_group_info_cache() for _, model in enumerate(model_list): # Extract model_info from the model dict model_info = model.get("model_info", {}) @@ -7869,16 +7449,9 @@ class Router: Used by `.get_model_list` to get model list from model alias. """ returned_models: List[DeploymentTypedDict] = [] - - if model_name is not None: - # Fast path: direct dict lookup avoids scanning all aliases for non-alias model names. - if model_name not in self.model_group_alias: - return returned_models - alias_items = [(model_name, self.model_group_alias[model_name])] - else: - alias_items = list(self.model_group_alias.items()) - - for model_alias, model_value in alias_items: + for model_alias, model_value in self.model_group_alias.items(): + if model_name is not None and model_alias != model_name: + continue if isinstance(model_value, str): _router_model_name: str = model_value elif isinstance(model_value, dict): @@ -7940,12 +7513,12 @@ class Router: return returned_models - def _invalidate_access_groups_cache(self) -> None: - """Invalidate the cached access groups. + def _invalidate_model_group_info_cache(self) -> None: + """Invalidate the cached model group info. Call this whenever self.model_list is modified to ensure the cache is rebuilt. """ - self._access_groups_cache = None + self._cached_get_model_group_info.cache_clear() def get_model_access_groups( self, @@ -7961,13 +7534,6 @@ class Router: - model_access_group: Optional[str] - the received model access group from the user. If set, will only return models for that access group. - team_id: Optional[str] - the team id, to resolve team-specific models """ - # Check if this is the no-args hot path (cacheable) - _use_cache = model_name is None and model_access_group is None and team_id is None - - # Return cached result for the no-args hot path - if _use_cache and self._access_groups_cache is not None: - return self._access_groups_cache - from collections import defaultdict access_groups = defaultdict(list) @@ -7986,11 +7552,6 @@ class Router: model_name = m["model_name"] access_groups[group].append(model_name) - # Cache the result for the no-args hot path - if _use_cache: - self._access_groups_cache = dict(access_groups) - return self._access_groups_cache - return access_groups def _is_model_access_group_for_wildcard_route( @@ -8456,10 +8017,9 @@ class Router: # check if the user sent in a deployment name instead healthy_deployments = self._get_deployment_by_litellm_model(model=model) - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug( - f"initial list of deployments: {healthy_deployments}" - ) + verbose_router_logger.debug( + f"initial list of deployments: {healthy_deployments}" + ) if len(healthy_deployments) == 0: # Check for default fallbacks if no deployments are found for the requested model @@ -8530,20 +8090,18 @@ class Router: request_kwargs=request_kwargs, ) - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug( - f"healthy_deployments after team filter: {healthy_deployments}" - ) + verbose_router_logger.debug( + f"healthy_deployments after team filter: {healthy_deployments}" + ) healthy_deployments = filter_web_search_deployments( healthy_deployments=healthy_deployments, request_kwargs=request_kwargs, ) - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug( - f"healthy_deployments after web search filter: {healthy_deployments}" - ) + verbose_router_logger.debug( + f"healthy_deployments after web search filter: {healthy_deployments}" + ) if isinstance(healthy_deployments, dict): return healthy_deployments @@ -8551,8 +8109,10 @@ class Router: cooldown_deployments = await _async_get_cooldown_deployments( litellm_router_instance=self, parent_otel_span=parent_otel_span ) - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug(f"cooldown deployments: {cooldown_deployments}") + verbose_router_logger.debug( + f"async cooldown deployments: {cooldown_deployments}" + ) + verbose_router_logger.debug(f"cooldown_deployments: {cooldown_deployments}") healthy_deployments = self._filter_cooldown_deployments( healthy_deployments=healthy_deployments, cooldown_deployments=cooldown_deployments, @@ -8653,13 +8213,6 @@ class Router: if isinstance(healthy_deployments, dict): return healthy_deployments - # When encrypted content affinity pins to a specific deployment, - if ( - request_kwargs.get("_encrypted_content_affinity_pinned") - and len(healthy_deployments) == 1 - ): - return healthy_deployments[0] - start_time = time.time() if ( self.routing_strategy == "usage-based-routing-v2" @@ -8932,18 +8485,6 @@ class Router: specific_deployment=specific_deployment, ) - ######################################################### - # Check if any complexity-router should be used - ######################################################### - if model in self.complexity_routers: - return await self.complexity_routers[model].async_pre_routing_hook( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) - return None def get_available_deployment( @@ -9247,8 +8788,7 @@ class Router: Returns: List of healthy deployments """ - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug(f"cooldown deployments: {cooldown_deployments}") + verbose_router_logger.debug(f"cooldown deployments: {cooldown_deployments}") # Convert to set for O(1) lookup and use list comprehension for O(n) filtering cooldown_set = set(cooldown_deployments) return [ diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 5732deda6fb..48aa435c3a9 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -925,10 +925,10 @@ def test_router_get_model_access_groups_team_only_models(): assert list(access_groups.keys()) == ["default-models"] -def test_get_model_access_groups_caching(): +def test_cached_get_model_group_info(): """ - Test that get_model_access_groups caches the no-args result - and invalidates on deployment changes. + Test that _cached_get_model_group_info caches results and + invalidates on deployment changes. """ from litellm.types.router import Deployment, LiteLLM_Params @@ -936,119 +936,60 @@ def test_get_model_access_groups_caching(): model_list=[ { "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4"}, - "model_info": {"access_groups": ["premium"]}, + "litellm_params": {"model": "gpt-4", "api_key": "fake"}, + "model_info": {"tpm": 1000, "rpm": 100}, }, ] ) - # First call computes and populates cache - result1 = router.get_model_access_groups() - assert "premium" in result1 + # First call should compute and cache + result1 = router._cached_get_model_group_info("gpt-4") + assert result1 is not None + assert result1.tpm == 1000 - # All subsequent calls should return the same cached object (including first) - result2 = router.get_model_access_groups() + # Second call should hit cache (same object) + result2 = router._cached_get_model_group_info("gpt-4") assert result1 is result2 - # Calls with args should bypass cache - result_with_args = router.get_model_access_groups(model_name="gpt-4") - assert result_with_args is not result2 - # Add a deployment — cache should be invalidated router.add_deployment( Deployment( - model_name="gpt-3.5", - litellm_params=LiteLLM_Params(model="gpt-3.5-turbo"), - model_info={"access_groups": ["default"]}, + model_name="gpt-4", + litellm_params=LiteLLM_Params(model="gpt-4", api_key="fake2"), + model_info={"tpm": 2000, "rpm": 200}, ) ) - result3 = router.get_model_access_groups() + result3 = router._cached_get_model_group_info("gpt-4") assert result3 is not result2 - assert "premium" in result3 - assert "default" in result3 + assert result3 is not None + assert result3.tpm == 3000 # 1000 + 2000 - # Delete the deployment — cache should be invalidated again - deployment_id = None - for m in router.model_list: - if m.get("model_name") == "gpt-3.5": - deployment_id = m.get("model_info", {}).get("id") - break - assert deployment_id is not None + # Delete a deployment — cache should be invalidated + deployment_id = router.model_list[-1]["model_info"]["id"] router.delete_deployment(id=deployment_id) - result4 = router.get_model_access_groups() + result4 = router._cached_get_model_group_info("gpt-4") assert result4 is not result3 - assert "default" not in result4 - assert "premium" in result4 + assert result4 is not None + assert result4.tpm == 1000 - -def test_get_model_access_groups_cache_invalidation_set_model_list(): - """ - Test that set_model_list invalidates the access groups cache. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4"}, - "model_info": {"access_groups": ["premium"]}, - }, - ] - ) - - # Populate cache - result1 = router.get_model_access_groups() - assert "premium" in result1 - - # set_model_list should invalidate cache + # set_model_list — cache should be invalidated router.set_model_list( [ - { - "model_name": "claude-3", - "litellm_params": {"model": "anthropic/claude-3-opus-20240229"}, - "model_info": {"access_groups": ["research"]}, - }, - ] - ) - result2 = router.get_model_access_groups() - assert result2 is not result1 - assert "research" in result2 - assert "premium" not in result2 - - -def test_get_model_access_groups_cache_invalidation_upsert_deployment(): - """ - Test that upsert_deployment invalidates the access groups cache. - """ - from litellm.types.router import Deployment, LiteLLM_Params - - router = litellm.Router( - model_list=[ { "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4"}, - "model_info": {"access_groups": ["premium"]}, + "litellm_params": {"model": "gpt-4", "api_key": "fake"}, + "model_info": {"tpm": 5000}, }, ] ) + result5 = router._cached_get_model_group_info("gpt-4") + assert result5 is not result4 + assert result5 is not None + assert result5.tpm == 5000 - # Populate cache - result1 = router.get_model_access_groups() - assert "premium" in result1 - - # Get the existing deployment's ID - existing_id = router.model_list[0]["model_info"]["id"] - - # Upsert with the same ID but different params — triggers pop + re-add - router.upsert_deployment( - Deployment( - model_name="gpt-4-updated", - litellm_params=LiteLLM_Params(model="gpt-4-turbo"), - model_info={"id": existing_id, "access_groups": ["updated-group"]}, - ) - ) - result2 = router.get_model_access_groups() - assert result2 is not result1 - assert "updated-group" in result2 + # Verify cache still works after invalidation + result6 = router._cached_get_model_group_info("gpt-4") + assert result5 is result6 @pytest.mark.asyncio @@ -1297,61 +1238,6 @@ async def test_acompletion_streaming_iterator_edge_cases(): print("✓ Edge case tests passed!") -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_preserves_hidden_params(): - """ - Regression test: FallbackStreamWrapper must copy _hidden_params from the - original CustomStreamWrapper so that x-litellm-overhead-duration-ms (and - other hidden params) are present in the proxy response headers for streaming. - """ - from unittest.mock import MagicMock - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - # Simulate a CustomStreamWrapper that already has timing metadata set by - # update_response_metadata (litellm_overhead_time_ms, _response_ms, etc.) - mock_response = MagicMock() - mock_response.model = "gpt-4" - mock_response.custom_llm_provider = "openai" - mock_response.logging_obj = MagicMock() - mock_response._hidden_params = { - "litellm_overhead_time_ms": 12.34, - "_response_ms": 500.0, - "litellm_call_id": "test-call-id", - "api_base": "https://api.openai.com", - "additional_headers": {}, - } - - # Make the mock iterable (yields nothing — we only care about hidden_params) - async def _empty(): - return - yield # make it an async generator - - mock_response.__aiter__ = lambda self: _empty().__aiter__() - - result = await router._acompletion_streaming_iterator( - model_response=mock_response, - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - - # The returned FallbackStreamWrapper must carry the original _hidden_params - assert hasattr(result, "_hidden_params"), "result must have _hidden_params" - assert result._hidden_params.get("litellm_overhead_time_ms") == 12.34, ( - "litellm_overhead_time_ms must be preserved — " - "this is what drives x-litellm-overhead-duration-ms in streaming responses" - ) - assert result._hidden_params.get("litellm_call_id") == "test-call-id" - assert result._hidden_params.get("_response_ms") == 500.0 - - @pytest.mark.asyncio async def test_async_function_with_fallbacks_common_utils(): """Test the async_function_with_fallbacks_common_utils method""" @@ -1907,54 +1793,6 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() assert credentials["custom_llm_provider"] == "bedrock" -def test_get_deployment_credentials_with_provider_resolves_credential_name(): - """ - Test that get_deployment_credentials_with_provider correctly resolves - litellm_credential_name to actual credential values (for UI-created models). - """ - from litellm.types.utils import CredentialItem - - # Setup credential list with a test credential - litellm.credential_list = [ - CredentialItem( - credential_name="test-azure-cred", - credential_info={"custom_llm_provider": "azure"}, - credential_values={ - "api_key": "resolved-api-key", - "api_base": "https://resolved.openai.azure.com", - "api_version": "2024-02-01" - } - ) - ] - - router = litellm.Router( - model_list=[ - { - "model_name": "azure-gpt-4", - "litellm_params": { - "model": "azure/gpt-4", - "litellm_credential_name": "test-azure-cred", - }, - } - ], - ) - - credentials = router.get_deployment_credentials_with_provider( - model_id="azure-gpt-4" - ) - - assert credentials is not None - assert credentials["api_key"] == "resolved-api-key" - assert credentials["api_base"] == "https://resolved.openai.azure.com" - assert credentials["api_version"] == "2024-02-01" - assert credentials["custom_llm_provider"] == "azure" - # Ensure credential name is removed after resolution - assert "litellm_credential_name" not in credentials - - # Cleanup - litellm.credential_list = [] - - def test_get_available_guardrail_single_deployment(): """ Test get_available_guardrail returns the single guardrail when only one exists. @@ -2098,397 +1936,3 @@ async def test_aguardrail(): assert result["result"] == "success" assert result["selected_guardrail"]["id"] == "guardrail-1" - -@pytest.mark.asyncio -async def test_anthropic_messages_call_type_is_cached(): - """ - Regression test: Verify that anthropic_messages call type is allowed - in PromptCachingDeploymentCheck.async_log_success_event. - """ - import asyncio - - from litellm.caching.dual_cache import DualCache - from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( - PromptCachingDeploymentCheck, - ) - from litellm.router_utils.prompt_caching_cache import PromptCachingCache - from litellm.types.utils import ( - CallTypes, - StandardLoggingHiddenParams, - StandardLoggingMetadata, - StandardLoggingModelInformation, - StandardLoggingPayload, - ) - - # Create mock standard logging payload inline - def create_standard_logging_payload() -> StandardLoggingPayload: - return StandardLoggingPayload( - id="test_id", - call_type="completion", - response_cost=0.1, - response_cost_failure_debug_info=None, - status="success", - total_tokens=30, - prompt_tokens=20, - completion_tokens=10, - startTime=1234567890.0, - endTime=1234567891.0, - completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), - model="gpt-3.5-turbo", - model_id="model-123", - model_group="openai-gpt", - api_base="https://api.openai.com", - metadata=StandardLoggingMetadata( - user_api_key_hash="test_hash", - user_api_key_org_id=None, - user_api_key_alias="test_alias", - user_api_key_team_id="test_team", - user_api_key_user_id="test_user", - user_api_key_team_alias="test_team_alias", - spend_logs_metadata=None, - requester_ip_address="127.0.0.1", - requester_metadata=None, - ), - cache_hit=False, - cache_key=None, - saved_cache_cost=0.0, - request_tags=[], - end_user=None, - requester_ip_address="127.0.0.1", - messages=[{"role": "user", "content": "Hello, world!"}], - response={"choices": [{"message": {"content": "Hi there!"}}]}, - error_str=None, - model_parameters={"stream": True}, - hidden_params=StandardLoggingHiddenParams( - model_id="model-123", - cache_key=None, - api_base="https://api.openai.com", - response_cost="0.1", - additional_headers=None, - ), - ) - - cache = DualCache() - deployment_check = PromptCachingDeploymentCheck(cache=cache) - prompt_cache = PromptCachingCache(cache=cache) - - # Create messages with enough tokens to pass the caching threshold - test_messages = [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "test long message here" * 1024, - "cache_control": { - "type": "ephemeral", - "ttl": "5m" - } - } - ] - } - ] - test_model_id = "test-model-id-123" - - # Create a payload with anthropic_messages call type - payload = create_standard_logging_payload() - payload["call_type"] = CallTypes.anthropic_messages.value - payload["messages"] = test_messages - payload["model"] = "anthropic/claude-3-5-sonnet-20240620" - payload["model_id"] = test_model_id - - # Log the success event (should cache the model_id) - await deployment_check.async_log_success_event( - kwargs={"standard_logging_object": payload}, - response_obj={}, - start_time=1234567890.0, - end_time=1234567891.0, - ) - - # Small delay to ensure cache write completes - await asyncio.sleep(0.1) - - # Verify that the model_id was actually cached - cached_result = await prompt_cache.async_get_model_id( - messages=test_messages, - tools=None, - ) - - # This assertion will FAIL if anthropic_messages is filtered out - assert cached_result is not None, "Model ID should be cached for anthropic_messages call type" - assert cached_result["model_id"] == test_model_id, f"Expected {test_model_id}, got {cached_result['model_id']}" - - -def test_update_kwargs_with_deployment_propagates_model_tags(): - """ - Test that deployment-level tags from litellm_params are merged into - kwargs metadata when _update_kwargs_with_deployment is called. - - This ensures model-level tags defined in config.yaml appear in SpendLogs. - See: https://github.com/BerriAI/litellm/issues/XXXX - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4o-mini", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "fake-key", - "tags": ["openai-account", "production"], - }, - }, - ], - ) - - kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - # Deployment tags should be propagated to kwargs metadata - assert "tags" in kwargs["metadata"] - assert "openai-account" in kwargs["metadata"]["tags"] - assert "production" in kwargs["metadata"]["tags"] - - -def test_update_kwargs_with_deployment_merges_tags_without_duplicates(): - """ - Test that when both request-level and deployment-level tags exist, - they are merged without duplicates. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4o-mini", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "fake-key", - "tags": ["openai-account", "shared-tag"], - }, - }, - ], - ) - - # Simulate request that already has tags (from request body or key/team level) - kwargs: dict = {"metadata": {"tags": ["user-tag", "shared-tag"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - # Both sources should be merged, no duplicates - assert "user-tag" in kwargs["metadata"]["tags"] - assert "openai-account" in kwargs["metadata"]["tags"] - assert "shared-tag" in kwargs["metadata"]["tags"] - assert kwargs["metadata"]["tags"].count("shared-tag") == 1 - - -def test_update_kwargs_with_deployment_no_tags(): - """ - Test that when deployment has no tags, kwargs metadata is not affected. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4o-mini", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "fake-key", - }, - }, - ], - ) - - kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - # No tags key should be added if deployment has no tags - assert "tags" not in kwargs["metadata"] - - -def test_update_kwargs_with_deployment_merges_tools(): - """ - Test that when both deployment litellm_params and request have tools, - they are merged (deployment tools first, then request tools). - - Supports proxy-configured tools (e.g. for o3 deep research) merged with - client-provided tools. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "o3-deep-research", - "litellm_params": { - "model": "openai/o3-deep-research", - "api_key": "fake-key", - "tools": [{"type": "web_search"}], - "tool_choice": "auto", - }, - }, - ], - ) - - kwargs: dict = { - "metadata": {}, - "tools": [ - { - "type": "function", - "function": {"name": "get_weather", "description": "Get weather"}, - }, - ], - } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - # Tools should be merged: deployment first, then request - assert "tools" in kwargs - assert len(kwargs["tools"]) == 2 - assert kwargs["tools"][0] == {"type": "web_search"} - assert kwargs["tools"][1]["function"]["name"] == "get_weather" - # tool_choice from request (none) - deployment's should be used - assert kwargs["tool_choice"] == "auto" - - -def test_update_kwargs_with_deployment_merge_tools_deployment_only(): - """ - Test that when only deployment has tools, they are applied to kwargs. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "o3-deep-research", - "litellm_params": { - "model": "openai/o3-deep-research", - "api_key": "fake-key", - "tools": [{"type": "web_search"}], - "tool_choice": "required", - }, - }, - ], - ) - - kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - assert kwargs["tools"] == [{"type": "web_search"}] - assert kwargs["tool_choice"] == "required" - - -def test_update_kwargs_with_deployment_merge_tools_request_overrides_tool_choice(): - """ - Test that when request has tool_choice, it overrides deployment's. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "o3-deep-research", - "litellm_params": { - "model": "openai/o3-deep-research", - "api_key": "fake-key", - "tools": [{"type": "web_search"}], - "tool_choice": "auto", - }, - }, - ], - ) - - kwargs: dict = { - "metadata": {}, - "tool_choice": "none", - } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - # Request tool_choice should be preserved (merged tools still applied) - assert kwargs["tool_choice"] == "none" - - -def test_credential_name_injected_as_tag(): - """ - Test that litellm_credential_name from deployment litellm_params - is injected as a tag into metadata during _update_kwargs_with_deployment. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "xai-model", - "litellm_params": { - "model": "xai/grok-4-1-fast", - "litellm_credential_name": "xAI", - }, - } - ], - ) - - kwargs: dict = {"metadata": {"tags": ["A.101"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="xai-model" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - assert "Credential: xAI" in kwargs["metadata"]["tags"] - assert "A.101" in kwargs["metadata"]["tags"] - - -def test_credential_name_not_duplicated_in_tags(): - """ - Test that if the credential tag already exists in the tags list, - it is not duplicated. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "xai-model", - "litellm_params": { - "model": "xai/grok-4-1-fast", - "litellm_credential_name": "xAI", - }, - } - ], - ) - - kwargs: dict = {"metadata": {"tags": ["Credential: xAI", "A.101"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="xai-model" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - assert kwargs["metadata"]["tags"].count("Credential: xAI") == 1 - - -def test_credential_name_not_injected_when_absent(): - """ - Test that when no litellm_credential_name is set, tags are unchanged. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-model", - "litellm_params": { - "model": "gpt-4o", - }, - } - ], - ) - - kwargs: dict = {"metadata": {"tags": ["A.101"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-model" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - assert kwargs["metadata"]["tags"] == ["A.101"]