From 5ed564aeca25739389ae83d7d4bf0c945f3832e0 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 26 Feb 2026 10:42:01 +0530 Subject: [PATCH] Fix mypy issues --- .../litellm_core_utils/realtime_streaming.py | 8 +- .../amazon_qwen2_transformation.py | 15 +- .../amazon_qwen3_transformation.py | 15 +- .../mcp_server/auth/user_api_key_auth_mcp.py | 4 +- litellm/proxy/db/db_spend_update_writer.py | 284 +++++--- .../usage_endpoints/ai_usage_chat.py | 8 +- litellm/router.py | 641 +++++++++++++++--- 7 files changed, 766 insertions(+), 209 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 6ba1b48c647..898acadda60 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,7 +1,7 @@ import asyncio import concurrent.futures import json -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast import litellm from litellm._logging import verbose_logger @@ -92,7 +92,7 @@ class RealTimeStreaming: message_obj = message else: message_obj = json.loads(message) - self._collect_tool_calls_from_response_done(message_obj) + self._collect_tool_calls_from_response_done(cast(dict, message_obj)) try: if ( not isinstance(message, dict) @@ -428,11 +428,11 @@ class RealTimeStreaming: == "conversation.item.input_audio_transcription.completed" ): transcript = event.get("transcript", "") - self._collect_user_input_from_backend_event(event) + self._collect_user_input_from_backend_event(cast(dict, event)) self.store_message(event_str) await self.websocket.send_text(event_str) blocked = await self.run_realtime_guardrails( - transcript, item_id=event.get("item_id") + cast(str, transcript), item_id=cast(Optional[str], event.get("item_id")) ) if not blocked: await self._send_to_backend( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py index 2abcc679eef..0260eeafe63 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py @@ -11,7 +11,6 @@ from typing import Any, List, Optional import httpx -from litellm.types.utils import Usage from litellm.llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import ( AmazonQwen3Config, ) @@ -19,7 +18,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation LiteLLMLoggingObj, ) from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import ModelResponse +from litellm.types.utils import ModelResponse, Usage class AmazonQwen2Config(AmazonQwen3Config): @@ -80,10 +79,14 @@ class AmazonQwen2Config(AmazonQwen3Config): # Set usage information if available in response if "usage" in response_data: usage_data = response_data["usage"] - model_response.usage = Usage( - prompt_tokens=usage_data.get("prompt_tokens", 0), - completion_tokens=usage_data.get("completion_tokens", 0), - total_tokens=usage_data.get("total_tokens", 0), + setattr( + model_response, + "usage", + Usage( + prompt_tokens=usage_data.get("prompt_tokens", 0), + completion_tokens=usage_data.get("completion_tokens", 0), + total_tokens=usage_data.get("total_tokens", 0), + ), ) return model_response diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py index 12333623f51..6eddcccd631 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py @@ -10,14 +10,13 @@ from typing import Any, List, Optional import httpx -from litellm.types.utils import Usage from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( AmazonInvokeConfig, LiteLLMLoggingObj, ) from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import ModelResponse +from litellm.types.utils import ModelResponse, Usage class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): @@ -202,10 +201,14 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): # Set usage information if available in response if "usage" in response_data: usage_data = response_data["usage"] - model_response.usage = Usage( - prompt_tokens=usage_data.get("prompt_tokens", 0), - completion_tokens=usage_data.get("completion_tokens", 0), - total_tokens=usage_data.get("total_tokens", 0), + setattr( + model_response, + "usage", + Usage( + prompt_tokens=usage_data.get("prompt_tokens", 0), + completion_tokens=usage_data.get("completion_tokens", 0), + total_tokens=usage_data.get("total_tokens", 0), + ), ) return model_response diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 860569d24cb..6e78458cc0e 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Set, Tuple +from typing import Dict, List, Optional, Set, Tuple, cast from fastapi import HTTPException from starlette.datastructures import Headers @@ -539,7 +539,7 @@ class MCPRequestHandler: allowed_tools = team_tools else: # No team restrictions → use key restrictions - allowed_tools = key_tools + allowed_tools = cast(List[str], key_tools) # Intersect with agent's tool permissions if agent_id is set if user_api_key_auth.agent_id: diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index c543bc70340..0c25424ceaa 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -139,53 +139,20 @@ class DBSpendUpdateWriter: payload["startTime"] = payload["startTime"].isoformat() if isinstance(payload["endTime"], datetime): payload["endTime"] = payload["endTime"].isoformat() - + if org_id is not None and org_id != "": payload["organization_id"] = org_id if team_id is not None and team_id != "": payload["team_id"] = team_id - asyncio.create_task( - self._update_user_db( - response_cost=response_cost, - user_id=user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - litellm_proxy_budget_name=litellm_proxy_budget_name, - end_user_id=end_user_id, - ) - ) - asyncio.create_task( - self._update_key_db( - response_cost=response_cost, - hashed_token=hashed_token, - prisma_client=prisma_client, - ) - ) - asyncio.create_task( - self._update_team_db( - response_cost=response_cost, - team_id=team_id, - user_id=user_id, - prisma_client=prisma_client, - ) - ) - asyncio.create_task( - self._update_org_db( - response_cost=response_cost, - org_id=org_id, - prisma_client=prisma_client, - ) - ) - asyncio.create_task( - self._update_tag_db( - response_cost=response_cost, - request_tags=copy.deepcopy(payload.get("request_tags")), - prisma_client=prisma_client, - ) - ) + # One deepcopy shared by all 6 daily spend helpers (was 5, fixes agent bug) + payload_copy = copy.deepcopy(payload) + # Deepcopy request_tags for _update_tag_db + request_tags = copy.deepcopy(payload.get("request_tags")) + + # Keep _insert_spend_log_to_db awaited inline (not a task, preserve current behavior) if disable_spend_logs is False: await self._insert_spend_log_to_db( payload=copy.deepcopy(payload), @@ -196,44 +163,20 @@ class DBSpendUpdateWriter: "disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur." ) + # Single task replaces 11 create_task() calls asyncio.create_task( - self.add_spend_log_transaction_to_daily_user_transaction( - payload=copy.deepcopy(payload), - prisma_client=prisma_client, - ) - ) - - asyncio.create_task( - self.add_spend_log_transaction_to_daily_end_user_transaction( - payload=copy.deepcopy(payload), - prisma_client=prisma_client, - ) - ) - - asyncio.create_task( - self.add_spend_log_transaction_to_daily_agent_transaction( - payload=payload, - prisma_client=prisma_client, - ) - ) - - asyncio.create_task( - self.add_spend_log_transaction_to_daily_team_transaction( - payload=copy.deepcopy(payload), - prisma_client=prisma_client, - ) - ) - asyncio.create_task( - self.add_spend_log_transaction_to_daily_org_transaction( - payload=copy.deepcopy(payload), + self._batch_database_updates( + response_cost=response_cost, + user_id=user_id, + hashed_token=hashed_token, + team_id=team_id, org_id=org_id, + end_user_id=end_user_id, prisma_client=prisma_client, - ) - ) - asyncio.create_task( - self.add_spend_log_transaction_to_daily_tag_transaction( - payload=copy.deepcopy(payload), - prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + litellm_proxy_budget_name=litellm_proxy_budget_name, + payload_copy=payload_copy, + request_tags=request_tags, ) ) @@ -357,6 +300,157 @@ class DBSpendUpdateWriter: "_enqueue_tool_registry_upsert error (non-blocking): %s", e ) + async def _batch_database_updates( + self, + *, + response_cost: Optional[float], + user_id: Optional[str], + hashed_token: Optional[str], + team_id: Optional[str], + org_id: Optional[str], + end_user_id: Optional[str], + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + litellm_proxy_budget_name: Optional[str], + payload_copy: SpendLogsPayload, + request_tags: Optional[Any], + ): + """ + Runs all 11 spend-update helpers sequentially inside a single asyncio task. + + Each helper is wrapped in try/except so one failure doesn't prevent the others. + """ + try: + await self._update_user_db( + response_cost=response_cost, + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + litellm_proxy_budget_name=litellm_proxy_budget_name, + end_user_id=end_user_id, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_user_db failed: %s", + traceback.format_exc(), + ) + + try: + await self._update_key_db( + response_cost=response_cost, + hashed_token=hashed_token, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_key_db failed: %s", + traceback.format_exc(), + ) + + try: + await self._update_team_db( + response_cost=response_cost, + team_id=team_id, + user_id=user_id, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_team_db failed: %s", + traceback.format_exc(), + ) + + try: + await self._update_org_db( + response_cost=response_cost, + org_id=org_id, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_org_db failed: %s", + traceback.format_exc(), + ) + + try: + await self._update_tag_db( + response_cost=response_cost, + request_tags=request_tags, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_tag_db failed: %s", + traceback.format_exc(), + ) + + try: + await self.add_spend_log_transaction_to_daily_user_transaction( + payload=payload_copy, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: add_spend_log_transaction_to_daily_user_transaction failed: %s", + traceback.format_exc(), + ) + + try: + await self.add_spend_log_transaction_to_daily_end_user_transaction( + payload=payload_copy, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: add_spend_log_transaction_to_daily_end_user_transaction failed: %s", + traceback.format_exc(), + ) + + try: + await self.add_spend_log_transaction_to_daily_agent_transaction( + payload=payload_copy, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: add_spend_log_transaction_to_daily_agent_transaction failed: %s", + traceback.format_exc(), + ) + + try: + await self.add_spend_log_transaction_to_daily_team_transaction( + payload=payload_copy, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: add_spend_log_transaction_to_daily_team_transaction failed: %s", + traceback.format_exc(), + ) + + try: + await self.add_spend_log_transaction_to_daily_org_transaction( + payload=payload_copy, + org_id=org_id, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: add_spend_log_transaction_to_daily_org_transaction failed: %s", + traceback.format_exc(), + ) + + try: + await self.add_spend_log_transaction_to_daily_tag_transaction( + payload=payload_copy, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: add_spend_log_transaction_to_daily_tag_transaction failed: %s", + traceback.format_exc(), + ) + async def _update_key_db( self, response_cost: Optional[float], @@ -1061,7 +1155,7 @@ class DBSpendUpdateWriter: team_id = key.split("::")[1] user_id = key.split("::")[3] team_memberships_to_invalidate.append((user_id, team_id)) - + for i in range(n_retry_times + 1): start_time = time.time() try: @@ -1098,11 +1192,13 @@ class DBSpendUpdateWriter: _raise_failed_update_spend_exception( e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj ) - + # Invalidate cache for updated team memberships # This ensures budget checks read fresh spend data from the database if team_memberships_to_invalidate and proxy_logging_obj is not None: - user_api_key_cache = proxy_logging_obj.call_details.get("user_api_key_cache") + user_api_key_cache = proxy_logging_obj.call_details.get( + "user_api_key_cache" + ) if user_api_key_cache is not None: for user_id, team_id in team_memberships_to_invalidate: cache_key = "team_membership:{}:{}".format(user_id, team_id) @@ -1414,7 +1510,9 @@ class DBSpendUpdateWriter: ), "endpoint": transaction.get("endpoint") or "", "prompt_tokens": transaction["prompt_tokens"], - "completion_tokens": transaction["completion_tokens"], + "completion_tokens": transaction[ + "completion_tokens" + ], "spend": transaction["spend"], "api_requests": transaction["api_requests"], "successful_requests": transaction[ @@ -1425,12 +1523,14 @@ class DBSpendUpdateWriter: # Add cache-related fields if they exist if "cache_read_input_tokens" in transaction: - common_data["cache_read_input_tokens"] = ( - transaction.get("cache_read_input_tokens", 0) - ) + common_data[ + "cache_read_input_tokens" + ] = transaction.get("cache_read_input_tokens", 0) if "cache_creation_input_tokens" in transaction: - common_data["cache_creation_input_tokens"] = ( - transaction.get("cache_creation_input_tokens", 0) + common_data[ + "cache_creation_input_tokens" + ] = transaction.get( + "cache_creation_input_tokens", 0 ) if entity_type == "tag" and "request_id" in transaction: @@ -1473,10 +1573,14 @@ class DBSpendUpdateWriter: } if entity_type == "tag" and "request_id" in transaction: - update_data["request_id"] = transaction.get("request_id") + update_data["request_id"] = transaction.get( + "request_id" + ) # Add endpoint to update_data so existing rows get their endpoint field updated - update_data["endpoint"] = transaction.get("endpoint") or "" + update_data["endpoint"] = ( + transaction.get("endpoint") or "" + ) table.upsert( where=where_clause, @@ -1660,7 +1764,9 @@ class DBSpendUpdateWriter: self, payload: Union[dict, SpendLogsPayload], prisma_client: PrismaClient, - type: Literal["user", "team", "org", "request_tags", "end_user", "agent"] = "user", + type: Literal[ + "user", "team", "org", "request_tags", "end_user", "agent" + ] = "user", ) -> Optional[BaseDailySpendTransaction]: common_expected_keys = ["startTime", "api_key"] if type == "user": @@ -1719,7 +1825,7 @@ class DBSpendUpdateWriter: endpoint = None if call_type: endpoint = ROUTE_ENDPOINT_MAPPING.get(call_type, None) - + daily_transaction = BaseDailySpendTransaction( date=date, api_key=payload["api_key"], @@ -1931,7 +2037,7 @@ class DBSpendUpdateWriter: endpoint_str = base_daily_transaction.get("endpoint") or "" daily_transaction_key = f"{payload['agent_id']}_{base_daily_transaction['date']}_{payload_with_agent_id['api_key']}_{payload_with_agent_id['model']}_{payload_with_agent_id['custom_llm_provider']}_{endpoint_str}" daily_transaction = DailyAgentSpendTransaction( - agent_id=payload['agent_id'], **base_daily_transaction + agent_id=payload["agent_id"], **base_daily_transaction ) await self.daily_agent_spend_update_queue.add_update( update={daily_transaction_key: daily_transaction} diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py index f156be7d2cc..56c3c5f0477 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -5,7 +5,7 @@ usage/spend data by querying the aggregated daily activity endpoints. import json from datetime import date -from typing import Any, AsyncIterator, Callable, Dict, List, Literal, Optional +from typing import Any, AsyncIterator, Callable, Dict, List, Literal, Optional, cast import litellm from litellm._logging import verbose_proxy_logger @@ -492,17 +492,17 @@ async def _process_tool_call( "tool_label": handler["label"], "arguments": fn_args, } - yield _sse({**tool_event_base, "status": "running"}) + yield _sse(cast(SSEToolCallEvent, {**tool_event_base, "status": "running"})) try: tool_result = await _execute_tool_call( handler, fn_name, fn_args, user_id, is_admin ) - yield _sse({**tool_event_base, "status": "complete"}) + yield _sse(cast(SSEToolCallEvent, {**tool_event_base, "status": "complete"})) except Exception as e: verbose_proxy_logger.error("Tool %s failed: %s", fn_name, e) tool_result = f"Error fetching {handler['label']}. Please try again." - yield _sse({**tool_event_base, "status": "error"}) + yield _sse(cast(SSEToolCallEvent, {**tool_event_base, "status": "error"})) chat_messages.append( {"role": "tool", "tool_call_id": tc.id, "content": tool_result} diff --git a/litellm/router.py b/litellm/router.py index 5b72c3fb669..cbe5b414040 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -33,6 +33,7 @@ from typing import ( cast, ) +import anyio import httpx import openai from openai import AsyncOpenAI @@ -58,7 +59,6 @@ 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,15 +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.prompt_caching_deployment_check import ( - PromptCachingDeploymentCheck, -) -from litellm.router_utils.pre_call_checks.responses_api_deployment_check import ( - ResponsesApiDeploymentCheck, +from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, ) 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, @@ -190,11 +190,15 @@ 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 @@ -294,6 +298,7 @@ class Router: router_general_settings: Optional[ RouterGeneralSettings ] = RouterGeneralSettings(), + deployment_affinity_ttl_seconds: int = 3600, ignore_invalid_deployments: bool = False, ) -> None: """ @@ -327,6 +332,7 @@ 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. @@ -446,6 +452,7 @@ 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]] = ( @@ -470,6 +477,8 @@ 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: @@ -605,6 +614,7 @@ 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 @@ -619,11 +629,12 @@ class Router: self.retry_policy = RetryPolicy(**retry_policy) elif isinstance(retry_policy, RetryPolicy): self.retry_policy = retry_policy - verbose_router_logger.info( - "\033[32mRouter Custom Retry Policy Set:\n{}\033[0m".format( - self.retry_policy.model_dump(exclude_none=True) + 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) + ) ) - ) self.model_group_retry_policy: Optional[ Dict[str, RetryPolicy] @@ -636,11 +647,12 @@ class Router: elif isinstance(allowed_fails_policy, AllowedFailsPolicy): self.allowed_fails_policy = allowed_fails_policy - verbose_router_logger.info( - "\033[32mRouter Custom Allowed Fails Policy Set:\n{}\033[0m".format( - self.allowed_fails_policy.model_dump(exclude_none=True) + 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) + ) ) - ) self.alerting_config: Optional[AlertingConfig] = alerting_config @@ -895,14 +907,11 @@ class Router: def _initialize_vector_store_endpoints(self): """Initialize vector store endpoints.""" - from litellm.vector_stores.main import acreate, asearch, create, search + from litellm.vector_stores.main import 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" ) @@ -1158,6 +1167,8 @@ 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() @@ -1184,26 +1195,87 @@ class Router: def add_optional_pre_call_checks( self, optional_pre_call_checks: Optional[OptionalPreCallChecks] ): - 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) + 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) + + # --------------------------------------------------------------------- + # 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", + ): + 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) def print_deployment(self, deployment: dict): """ @@ -1269,19 +1341,20 @@ class Router: if silent_model is not None: # Mirroring traffic to a secondary model - # Use shared thread pool for background calls - executor.submit( - self._silent_experiment_completion, - silent_model, - messages, - **kwargs, + # 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, ) + thread.start() 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 = data["model"] + model_name = litellm_params["model"] potential_model_client = self._get_client( deployment=deployment, kwargs=kwargs ) @@ -1302,7 +1375,7 @@ class Router: self.routing_strategy_pre_call_checks(deployment=deployment) input_kwargs = { - **data, + **litellm_params, "messages": messages, "caching": self.cache_responses, "client": model_client, @@ -1463,7 +1536,7 @@ class Router: ) raise e - async def _acompletion_streaming_iterator( + async def _acompletion_streaming_iterator( # noqa: PLR0915 self, model_response: CustomStreamWrapper, messages: List[Dict[str, str]], @@ -1486,6 +1559,9 @@ 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 @@ -1494,6 +1570,7 @@ 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 @@ -1592,6 +1669,30 @@ 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()) @@ -1690,10 +1791,8 @@ 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 = data["model"] + model_name = litellm_params["model"] model_client = self._get_async_openai_model_client( deployment=deployment, @@ -1702,7 +1801,7 @@ class Router: self.total_calls[model_name] += 1 input_kwargs = { - **data, + **litellm_params, "messages": messages, "caching": self.cache_responses, "client": model_client, @@ -1893,6 +1992,26 @@ 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, @@ -1900,10 +2019,13 @@ class Router: function_name: Optional[str] = None, ) -> None: """ - 2 jobs: + 3 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") @@ -1928,6 +2050,28 @@ 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( @@ -2272,7 +2416,7 @@ class Router: item = FlowItem( priority=priority, # 👈 SET PRIORITY FOR REQUEST request_id=_request_id, # 👈 SET REQUEST ID - model_name="gpt-3.5-turbo", # 👈 SAME as 'Router' + model_name=model, # 👈 SAME as 'Router' ) ### [fin] ### @@ -2314,6 +2458,10 @@ 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, @@ -2375,6 +2523,10 @@ 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, @@ -2423,6 +2575,12 @@ 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) @@ -2435,7 +2593,9 @@ class Router: "litellm_params" ].get("prompt_label", None) - if prompt_id is None or not isinstance(prompt_id, str): + if not is_litellm_agent_model and ( + 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)}" ) @@ -3669,7 +3829,7 @@ class Router: ) raise e - async def _acreate_file( + async def _acreate_file( # noqa: PLR0915 self, model: str, **kwargs, @@ -3730,7 +3890,12 @@ 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, @@ -3801,6 +3966,114 @@ 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, @@ -4507,9 +4780,21 @@ 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( @@ -4815,7 +5100,8 @@ class Router: ) else: response = await self.async_function_with_retries(*args, **kwargs) - verbose_router_logger.debug(f"Async Response: {response}") + if verbose_router_logger.isEnabledFor(logging.DEBUG): + verbose_router_logger.debug(f"Async Response: {response}") response = add_fallback_headers_to_response( response=response, attempted_fallbacks=0, @@ -4892,6 +5178,10 @@ 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") @@ -4905,6 +5195,11 @@ 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 @@ -4938,19 +5233,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 self.model_group_retry_policy is not None - ): + if self.retry_policy is not None or 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") ) - _retry_policy_retries = self.get_num_retries_from_retry_policy( + # Use per-request model_group_retry_policy if provided, otherwise use self + _retry_policy_retries = _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 @@ -4967,6 +5262,9 @@ 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) @@ -4989,6 +5287,9 @@ 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( @@ -5019,7 +5320,7 @@ class Router: else: _healthy_deployments = [] _timeout = self._time_to_sleep_before_retry( - e=original_exception, + e=e, remaining_retries=remaining_retries, num_retries=num_retries, healthy_deployments=_healthy_deployments, @@ -5304,7 +5605,7 @@ class Router: return else: deployment_model_info = self.get_router_model_info( - deployment=deployment_info.model_dump(), + deployment=deployment_info, received_model_name=model_group, ) # get tpm/rpm from deployment info @@ -5920,9 +6221,18 @@ 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: _model_info, + _model_name: _shared_model_info, } ) @@ -5955,10 +6265,13 @@ class Router: def _is_auto_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: """ - Check if the deployment is an auto-router deployment. + Check if the deployment is an auto-router deployment (semantic router). 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 @@ -6010,6 +6323,61 @@ 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 @@ -6056,6 +6424,7 @@ class Router: self.model_id_to_deployment_index_map = {} # Reset the index self.model_name_to_deployment_indices = {} # Reset the model_name index self._invalidate_model_group_info_cache() + self._invalidate_access_groups_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: @@ -6220,6 +6589,12 @@ 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( @@ -6360,6 +6735,7 @@ class Router: idx = len(self.model_list) self.model_list.append(model) self._invalidate_model_group_info_cache() + self._invalidate_access_groups_cache() # Update model_id index for O(1) lookup if model_id is not None: @@ -6408,6 +6784,7 @@ class Router: if removal_idx is not None: self.model_list.pop(removal_idx) self._invalidate_model_group_info_cache() + self._invalidate_access_groups_cache() self._update_deployment_indices_after_removal( model_id=deployment_id, removal_idx=removal_idx ) @@ -6442,6 +6819,7 @@ class Router: # Pop the item from the list first item = self.model_list.pop(deployment_idx) self._invalidate_model_group_info_cache() + self._invalidate_access_groups_cache() self._update_deployment_indices_after_removal( model_id=id, removal_idx=deployment_idx ) @@ -6543,6 +6921,19 @@ 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[ @@ -6560,7 +6951,7 @@ class Router: @overload def get_router_model_info( - self, deployment: dict, received_model_name: str, id: None = None + self, deployment: Union[dict, "Deployment"], received_model_name: str, id: None = None ) -> ModelMapInfo: pass @@ -6572,7 +6963,7 @@ class Router: def get_router_model_info( self, - deployment: Optional[dict], + deployment: Optional[Union[dict, "Deployment"]], received_model_name: str, id: Optional[str] = None, ) -> ModelMapInfo: @@ -6592,22 +6983,34 @@ class Router: if id is not None: _deployment = self.get_deployment(model_id=id) if _deployment is not None: - deployment = _deployment.model_dump(exclude_none=True) + deployment = _deployment if deployment is None: raise ValueError("Deployment not found") ## GET BASE MODEL - base_model = deployment.get("model_info", {}).get("base_model", None) + base_model = (deployment.get("model_info") or {}).get("base_model", None) if base_model is None: - base_model = deployment.get("litellm_params", {}).get("base_model", None) + base_model = (deployment.get("litellm_params") or {}).get("base_model", None) model = base_model - ## GET PROVIDER + ## 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')}" + ) _model, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=deployment.get("litellm_params", {}).get("model", ""), - litellm_params=LiteLLM_Params(**deployment.get("litellm_params", {})), + model=litellm_params.model, + litellm_params=litellm_params, ) ## SET MODEL TO 'model=' - if base_model is None + not azure @@ -6623,10 +7026,10 @@ class Router: if potential_models is not None: for potential_model in potential_models: try: - if potential_model.get("model_info", {}).get( + if (potential_model.get("model_info") or {}).get( "id" - ) == deployment.get("model_info", {}).get("id"): - model = potential_model.get("litellm_params", {}).get( + ) == (deployment.get("model_info") or {}).get("id"): + model = (potential_model.get("litellm_params") or {}).get( "model" ) break @@ -6647,9 +7050,10 @@ class Router: model_info = litellm.get_model_info(model=model_info_name) ## CHECK USER SET MODEL INFO - user_model_info = deployment.get("model_info", {}) + user_model_info = deployment.get("model_info") or {} - model_info.update(user_model_info) + if model_info is not None: + model_info.update(cast(ModelInfo, user_model_info)) return model_info @@ -7177,6 +7581,7 @@ class Router: # First populate the model_list self.model_list = [] self._invalidate_model_group_info_cache() + self._invalidate_access_groups_cache() for _, model in enumerate(model_list): # Extract model_info from the model dict model_info = model.get("model_info", {}) @@ -7449,9 +7854,16 @@ class Router: Used by `.get_model_list` to get model list from model alias. """ returned_models: List[DeploymentTypedDict] = [] - for model_alias, model_value in self.model_group_alias.items(): - if model_name is not None and model_alias != model_name: - continue + + 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: if isinstance(model_value, str): _router_model_name: str = model_value elif isinstance(model_value, dict): @@ -7520,6 +7932,13 @@ class Router: """ self._cached_get_model_group_info.cache_clear() + def _invalidate_access_groups_cache(self) -> None: + """Invalidate the cached access groups. + + Call this whenever self.model_list is modified to ensure the cache is rebuilt. + """ + self._access_groups_cache = None + def get_model_access_groups( self, model_name: Optional[str] = None, @@ -7534,6 +7953,13 @@ 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) @@ -7552,6 +7978,11 @@ 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( @@ -8017,9 +8448,10 @@ class Router: # check if the user sent in a deployment name instead healthy_deployments = self._get_deployment_by_litellm_model(model=model) - verbose_router_logger.debug( - f"initial list of deployments: {healthy_deployments}" - ) + if verbose_router_logger.isEnabledFor(logging.DEBUG): + 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 @@ -8090,18 +8522,20 @@ class Router: request_kwargs=request_kwargs, ) - verbose_router_logger.debug( - f"healthy_deployments after team filter: {healthy_deployments}" - ) + if verbose_router_logger.isEnabledFor(logging.DEBUG): + 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, ) - verbose_router_logger.debug( - f"healthy_deployments after web search filter: {healthy_deployments}" - ) + if verbose_router_logger.isEnabledFor(logging.DEBUG): + verbose_router_logger.debug( + f"healthy_deployments after web search filter: {healthy_deployments}" + ) if isinstance(healthy_deployments, dict): return healthy_deployments @@ -8109,10 +8543,8 @@ class Router: cooldown_deployments = await _async_get_cooldown_deployments( litellm_router_instance=self, parent_otel_span=parent_otel_span ) - verbose_router_logger.debug( - f"async cooldown deployments: {cooldown_deployments}" - ) - verbose_router_logger.debug(f"cooldown_deployments: {cooldown_deployments}") + if verbose_router_logger.isEnabledFor(logging.DEBUG): + verbose_router_logger.debug(f"cooldown deployments: {cooldown_deployments}") healthy_deployments = self._filter_cooldown_deployments( healthy_deployments=healthy_deployments, cooldown_deployments=cooldown_deployments, @@ -8485,6 +8917,18 @@ 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( @@ -8788,7 +9232,8 @@ class Router: Returns: List of healthy deployments """ - verbose_router_logger.debug(f"cooldown deployments: {cooldown_deployments}") + if verbose_router_logger.isEnabledFor(logging.DEBUG): + 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 [