diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 7e6ec1dc15..d091e34802 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -614,3 +614,121 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): f"Error in dynamic rate limiter v3 post-call hook: {str(e)}" ) return response + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + """ + Update token usage for priority-based rate limiting after successful API calls. + + Increments token counters for: + - model_saturation_check: Model-wide token tracking + - priority_model: Priority-specific token tracking + """ + from litellm.litellm_core_utils.core_helpers import ( + _get_parent_otel_span_from_kwargs, + ) + from litellm.proxy.common_utils.callback_utils import ( + get_model_group_from_litellm_kwargs, + ) + from litellm.types.caching import RedisPipelineIncrementOperation + from litellm.types.utils import Usage + + try: + verbose_proxy_logger.debug( + "INSIDE dynamic rate limiter ASYNC SUCCESS LOGGING" + ) + + litellm_parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) + + # Get metadata from standard_logging_object + standard_logging_object = kwargs.get("standard_logging_object") or {} + standard_logging_metadata = standard_logging_object.get("metadata") or {} + + # Get model and priority + model_group = get_model_group_from_litellm_kwargs(kwargs) + if not model_group: + return + + # Get priority from user_api_key_auth_metadata in standard_logging_metadata + # This is where user_api_key_dict.metadata is stored during pre-call + user_api_key_auth_metadata = standard_logging_metadata.get("user_api_key_auth_metadata") or {} + key_priority: Optional[str] = user_api_key_auth_metadata.get("priority") + + # Get total tokens from response + total_tokens = 0 + rate_limit_type = self.v3_limiter.get_rate_limit_type() + + if isinstance(response_obj, ModelResponse): + _usage = getattr(response_obj, "usage", None) + if _usage and isinstance(_usage, Usage): + if rate_limit_type == "output": + total_tokens = _usage.completion_tokens + elif rate_limit_type == "input": + total_tokens = _usage.prompt_tokens + elif rate_limit_type == "total": + total_tokens = _usage.total_tokens + + if total_tokens == 0: + return + + # Create pipeline operations for token increments + pipeline_operations: List[RedisPipelineIncrementOperation] = [] + + # Model-wide token tracking (model_saturation_check) + model_token_key = self.v3_limiter.create_rate_limit_keys( + key="model_saturation_check", + value=model_group, + rate_limit_type="tokens", + ) + pipeline_operations.append( + RedisPipelineIncrementOperation( + key=model_token_key, + increment_value=total_tokens, + ttl=self.v3_limiter.window_size, + ) + ) + + # Priority-specific token tracking (priority_model) + # Determine priority key (same logic as _get_priority_allocation) + has_explicit_priority = ( + key_priority is not None + and litellm.priority_reservation is not None + and key_priority in litellm.priority_reservation + ) + + if has_explicit_priority and key_priority is not None: + priority_key = f"{model_group}:{key_priority}" + else: + priority_key = f"{model_group}:default_pool" + + priority_token_key = self.v3_limiter.create_rate_limit_keys( + key="priority_model", + value=priority_key, + rate_limit_type="tokens", + ) + pipeline_operations.append( + RedisPipelineIncrementOperation( + key=priority_token_key, + increment_value=total_tokens, + ttl=self.v3_limiter.window_size, + ) + ) + + # Execute token increments with TTL preservation + if pipeline_operations: + await self.v3_limiter.async_increment_tokens_with_ttl_preservation( + pipeline_operations=pipeline_operations, + parent_otel_span=litellm_parent_otel_span, + ) + + # Only log 'priority' if it's known safe; otherwise, redact. + SAFE_PRIORITIES = {"low", "medium", "high", "default"} + logged_priority = key_priority if key_priority in SAFE_PRIORITIES else "REDACTED" + verbose_proxy_logger.debug( + f"[Dynamic Rate Limiter] Incremented tokens by {total_tokens} for " + f"model={model_group}, priority={logged_priority}" + ) + + except Exception as e: + verbose_proxy_logger.exception( + f"Error in dynamic rate limiter success event: {str(e)}" + ) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index c462493de6..6ef281e5dc 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -65,6 +65,12 @@ for i = 1, #KEYS, 2 do table.insert(results, increment_value) -- counter else local counter = redis.call('INCR', counter_key) + -- This happens when window_key exists but counter_key doesn't (e.g., tokens key + -- created after requests key when both share the same window_key) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 then + redis.call('EXPIRE', counter_key, window_size) + end table.insert(results, window_start) -- window_start table.insert(results, counter) -- counter end diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 098cdb80e0..a33f56b032 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,67 +1,16 @@ model_list: - - model_name: qwen-25vl-72b + - model_name: openai/gpt-4o-mini litellm_params: - model: bedrock/openai/arn:aws:bedrock:us-east-1:046319184608:imported-model/0m2lasirsp6z + model: openai/gpt-4o-mini + tpm: 1000 -guardrails: - - guardrail_name: "bedrock-pre-guard" - litellm_params: - guardrail: bedrock - mode: "pre_call" - guardrailIdentifier: ff6ujrregl1q - guardrailVersion: "DRAFT" - - -# like MCPs/vector stores -search_tools: - - search_tool_name: litellm-search - litellm_params: - search_provider: perplexity - api_key: os.environ/PERPLEXITYAI_API_KEY - - search_tool_name: firecrawl-search - litellm_params: - search_provider: firecrawl - api_key: os.environ/FIRECRAWL_API_KEY - litellm_settings: - max_end_user_budget_id: "2f6634cd-c631-4d3b-96c7-ad510ea06eaf" - # Comprehensive logging settings - store_audit_logs: true - verbose: true - log_level: "DEBUG" # Options: DEBUG, INFO, WARNING, ERROR - callbacks: ["s3_v2", "smtp_email"] - s3_callback_params: - s3_endpoint_url: "https://localhost:443" # Replace with your Minio server URL and port - s3_aws_access_key_id: "minioadmin" - s3_aws_secret_access_key: "minioadmin" - s3_region_name: "minio" # This can be any value for Minio - s3_bucket_name: "litellm-test" # Replace with your bucket name - s3_use_ssl: False - s3_verify: False - cache: True - cache_params: - type: local - drop_params: True + callbacks: ["dynamic_rate_limiter_v3"] + priority_reservation: + "prod": 0.9 # 90% reserved for production + "dev": 0.1 # 10% reserved for development -general_settings: - store_prompts_in_spend_logs: True - pass_through_endpoints: - - path: "/special/rerank" - target: "https://api.cohere.com/v1/rerank" - headers: - Authorization: "Bearer os.environ/COHERE_API_KEY" - guardrails: - bedrock-pre-guard: - request_fields: ["documents[*].text"] -vector_store_registry: - - vector_store_name: "bedrock-litellm-website-knowledgebase" - litellm_params: - vector_store_id: "T37J8R4WTM" - custom_llm_provider: "bedrock" - vector_store_description: "Bedrock vector store for the Litellm website knowledgebase" - vector_store_metadata: - source: "https://www.litellm.com/docs" diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index 1f76013e23..d9e10e6f4b 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -1323,3 +1323,187 @@ async def test_default_priority_shared_pool(): print(f" - 3 keys without priority share ONE pool: {desc_a[0]['value']}") print(f" - Shared pool limit: {desc_a[0]['rate_limit']['requests_per_unit']} RPM") print(f" - Explicit priority 'prod' uses separate pool: {desc_prod[0]['value']}") + + +@pytest.mark.asyncio +async def test_async_log_success_event_increments_by_actual_tokens(): + """ + Test that async_log_success_event increments token counters by actual token usage. + + This validates the fix for Bug 1: Token count was incrementing by 1 instead of actual usage. + The async_log_success_event should increment both model_saturation_check and priority_model + counters by the actual completion_tokens (when rate_limit_type=output). + """ + from unittest.mock import MagicMock + + from litellm.types.utils import ModelResponse, Usage + + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"dev": 0.1, "prod": 0.9} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "test-token-increment" + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": 1000, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Track what gets incremented + increment_calls = [] + + async def mock_increment(pipeline_operations, parent_otel_span=None): + for op in pipeline_operations: + increment_calls.append({ + "key": op["key"], + "increment_value": op["increment_value"], + }) + + handler.v3_limiter.async_increment_tokens_with_ttl_preservation = mock_increment + + # Create mock response with 50 completion tokens + mock_response = MagicMock(spec=ModelResponse) + mock_response.usage = MagicMock(spec=Usage) + mock_response.usage.prompt_tokens = 10 + mock_response.usage.completion_tokens = 50 + mock_response.usage.total_tokens = 60 + + # Create kwargs with priority in user_api_key_auth_metadata + kwargs = { + "standard_logging_object": { + "metadata": { + "user_api_key_auth_metadata": {"priority": "dev"}, + }, + "model_group": model, + }, + "litellm_params": { + "metadata": {"model_group": model}, + }, + } + + with patch( + "litellm.proxy.common_utils.callback_utils.get_model_group_from_litellm_kwargs", + return_value=model, + ): + await handler.async_log_success_event( + kwargs=kwargs, + response_obj=mock_response, + start_time=None, + end_time=None, + ) + + # Verify increments happened with actual token count (50 completion tokens) + assert len(increment_calls) == 2, f"Expected 2 increment calls, got {len(increment_calls)}" + + # Both should increment by 50 (completion_tokens, since rate_limit_type defaults to 'output') + for call in increment_calls: + assert call["increment_value"] == 50, ( + f"Expected increment of 50 tokens, got {call['increment_value']} for key {call['key']}" + ) + + # Verify correct keys were used + keys = [call["key"] for call in increment_calls] + assert any("model_saturation_check" in k for k in keys), "Should increment model_saturation_check" + assert any("priority_model" in k and "dev" in k for k in keys), "Should increment priority_model with 'dev' priority" + + +@pytest.mark.asyncio +async def test_async_log_success_event_uses_team_priority_from_auth_metadata(): + """ + Test that async_log_success_event correctly retrieves priority from user_api_key_auth_metadata. + + This validates the fix where priority is retrieved from standard_logging_metadata.user_api_key_auth_metadata + instead of just standard_logging_metadata.priority. This is important for team-based priority inheritance. + """ + from unittest.mock import MagicMock + + from litellm.types.utils import ModelResponse, Usage + + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"team_priority": 0.8, "default": 0.2} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "test-team-priority" + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": 1000, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Track incremented keys to verify priority is used correctly + incremented_keys = [] + + async def mock_increment(pipeline_operations, parent_otel_span=None): + for op in pipeline_operations: + incremented_keys.append(op["key"]) + + handler.v3_limiter.async_increment_tokens_with_ttl_preservation = mock_increment + + # Create mock response + mock_response = MagicMock(spec=ModelResponse) + mock_response.usage = MagicMock(spec=Usage) + mock_response.usage.prompt_tokens = 10 + mock_response.usage.completion_tokens = 20 + mock_response.usage.total_tokens = 30 + + # Simulate team metadata inheritance: priority is in user_api_key_auth_metadata + # This is how the proxy passes team metadata to the callback + kwargs = { + "standard_logging_object": { + "metadata": { + # Priority NOT at top level (this would fail before the fix) + # Priority IS in user_api_key_auth_metadata (team inheritance) + "user_api_key_auth_metadata": {"priority": "team_priority"}, + }, + "model_group": model, + }, + "litellm_params": { + "metadata": {"model_group": model}, + }, + } + + with patch( + "litellm.proxy.common_utils.callback_utils.get_model_group_from_litellm_kwargs", + return_value=model, + ): + await handler.async_log_success_event( + kwargs=kwargs, + response_obj=mock_response, + start_time=None, + end_time=None, + ) + + # Verify the priority_model key uses 'team_priority' (not 'default_pool') + priority_keys = [k for k in incremented_keys if "priority_model" in k] + assert len(priority_keys) == 1, f"Expected 1 priority_model key, got {len(priority_keys)}" + + # The key should contain 'team_priority', not 'default_pool' + assert "team_priority" in priority_keys[0], ( + f"Expected priority key to use 'team_priority' from user_api_key_auth_metadata, " + f"got key: {priority_keys[0]}" + ) + assert "default_pool" not in priority_keys[0], ( + f"Priority key should NOT use 'default_pool', should use team's priority. Got: {priority_keys[0]}" + )