From b731e432f88ee026f38ef93a6974782c62f68d9f Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Tue, 10 Mar 2026 12:45:48 +0530 Subject: [PATCH] feat: record silent metrics --- litellm/router.py | 107 +++++++++++------- .../test_router_silent_experiment.py | 40 +++++-- 2 files changed, 98 insertions(+), 49 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 43b53d14d79..fac9409b359 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -727,15 +727,11 @@ class Router: startup_nodes = cache_config.get("startup_nodes") if not startup_nodes: _env_cluster_nodes = get_secret("REDIS_CLUSTER_NODES") - if _env_cluster_nodes is not None and isinstance( - _env_cluster_nodes, str - ): + if _env_cluster_nodes is not None and isinstance(_env_cluster_nodes, str): startup_nodes = json.loads(_env_cluster_nodes) if startup_nodes: - return RedisClusterCache( - **{**cache_config, "startup_nodes": startup_nodes} - ) + return RedisClusterCache(**{**cache_config, "startup_nodes": startup_nodes}) else: return RedisCache(**cache_config) @@ -1466,12 +1462,15 @@ class Router: silent_kwargs["metadata"]["is_silent_experiment"] = True + # Force stream=False so the response is fully consumed and callbacks fire + silent_kwargs["stream"] = False + # Pop logging objects and call IDs to ensure a fresh logging context # This prevents collisions in the Proxy's database (spend_logs) silent_kwargs.pop("litellm_call_id", None) silent_kwargs.pop("litellm_logging_obj", None) silent_kwargs.pop("standard_logging_object", None) - silent_kwargs.pop("proxy_server_request", None) + # DON'T pop proxy_server_request — it's needed for spend log metadata return silent_kwargs @@ -1494,12 +1493,30 @@ class Router: silent_kwargs = self._get_silent_experiment_kwargs(**kwargs) - # Trigger the silent request - self.completion( - model=silent_model, - messages=cast(List[Dict[str, str]], messages), - **silent_kwargs, - ) + # Override model_group to correctly attribute metrics to the silent model + silent_kwargs["metadata"]["model_group"] = silent_model + + # Create a new event loop for this thread so that async success + # callbacks (e.g. _ProxyDBLogger) can schedule and run DB writes. + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + loop.run_until_complete( + self.acompletion( + model=silent_model, + messages=cast(List[AllMessageValues], messages), + **silent_kwargs, + ) + ) + # Drain any remaining fire-and-forget tasks (e.g. alerting hooks) + # scheduled via asyncio.create_task during the acompletion call. + pending = asyncio.all_tasks(loop) + if pending: + loop.run_until_complete( + asyncio.gather(*pending, return_exceptions=True) + ) + finally: + loop.close() except Exception as e: verbose_router_logger.error( f"Silent experiment failed for model {silent_model}: {str(e)}" @@ -1705,7 +1722,9 @@ class Router: and isinstance(fallback_item, ModelResponseStream) and hasattr(fallback_item, "usage") ): - self._combine_fallback_usage(fallback_item, complete_response_object_usage) + self._combine_fallback_usage( + fallback_item, complete_response_object_usage + ) yield fallback_item else: # If fallback returns a non-streaming response, yield None @@ -1825,13 +1844,11 @@ class Router: router_self._update_kwargs_before_fallbacks( model=model_group, kwargs=initial_kwargs ) - fallback_response = ( - router_self.function_with_fallbacks( - **initial_kwargs, - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - ) + fallback_response = router_self.function_with_fallbacks( + **initial_kwargs, + fallbacks=fallbacks, + context_window_fallbacks=context_window_fallbacks, + content_policy_fallbacks=content_policy_fallbacks, ) if hasattr(fallback_response, "__iter__"): @@ -1841,7 +1858,9 @@ class Router: and isinstance(fallback_item, ModelResponseStream) and hasattr(fallback_item, "usage") ): - router_self._combine_fallback_usage(fallback_item, complete_response_object_usage) + router_self._combine_fallback_usage( + fallback_item, complete_response_object_usage + ) yield fallback_item else: yield None @@ -1891,6 +1910,8 @@ class Router: ) silent_kwargs = self._get_silent_experiment_kwargs(**kwargs) + # Override model_group to correctly attribute metrics to the silent model + silent_kwargs["metadata"]["model_group"] = silent_model # Trigger the silent request await self.acompletion( @@ -2753,10 +2774,9 @@ 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/") - ) + 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" @@ -6538,7 +6558,7 @@ class Router: 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, " @@ -6771,7 +6791,9 @@ class Router: ######################################################### # Check if this is a complexity-router deployment ######################################################### - if self._is_complexity_router_deployment(litellm_params=deployment.litellm_params): + if self._is_complexity_router_deployment( + litellm_params=deployment.litellm_params + ): self.init_complexity_router_deployment(deployment=deployment) return deployment @@ -6863,9 +6885,7 @@ class Router: # zero-cost models, causing budget checks to block free models. _model_id = deployment.model_info.id if _model_id is not None: - _model_info_dict: dict = deployment.model_info.model_dump( - exclude_none=True - ) + _model_info_dict: dict = deployment.model_info.model_dump(exclude_none=True) for field in CustomPricingLiteLLMParams.model_fields.keys(): field_value = deployment.litellm_params.get(field) if field_value is not None: @@ -7157,7 +7177,10 @@ class Router: @overload def get_router_model_info( - self, deployment: Union[dict, "Deployment"], received_model_name: str, id: None = None + self, + deployment: Union[dict, "Deployment"], + received_model_name: str, + id: None = None, ) -> ModelMapInfo: pass @@ -7197,7 +7220,9 @@ class Router: ## GET BASE MODEL base_model = (deployment.get("model_info") or {}).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") or {}).get( + "base_model", None + ) model = base_model @@ -7232,12 +7257,12 @@ class Router: if potential_models is not None: for potential_model in potential_models: try: - if (potential_model.get("model_info") or {}).get( - "id" - ) == (deployment.get("model_info") or {}).get("id"): - model = (potential_model.get("litellm_params") or {}).get( - "model" - ) + if (potential_model.get("model_info") or {}).get("id") == ( + deployment.get("model_info") or {} + ).get("id"): + model = ( + potential_model.get("litellm_params") or {} + ).get("model") break except Exception: pass @@ -8160,7 +8185,9 @@ class Router: - 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 + _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: diff --git a/tests/test_litellm/test_router_silent_experiment.py b/tests/test_litellm/test_router_silent_experiment.py index a23ea80f7ce..79056805d54 100644 --- a/tests/test_litellm/test_router_silent_experiment.py +++ b/tests/test_litellm/test_router_silent_experiment.py @@ -19,11 +19,20 @@ def test_get_silent_experiment_kwargs(): }, ] router = Router(model_list=model_list) - kwargs = {"metadata": {"foo": "bar"}, "litellm_call_id": "call-123"} + kwargs = { + "metadata": {"foo": "bar"}, + "litellm_call_id": "call-123", + "stream": True, + "proxy_server_request": {"body": {"model": "test"}}, + } result = router._get_silent_experiment_kwargs(**kwargs) assert result["metadata"]["is_silent_experiment"] is True assert result["metadata"]["foo"] == "bar" assert "litellm_call_id" not in result + # stream must be forced to False so callbacks fire in background + assert result["stream"] is False + # proxy_server_request must be preserved for spend log metadata + assert "proxy_server_request" in result def test_silent_experiment_completion_direct(): @@ -39,7 +48,7 @@ def test_silent_experiment_completion_direct(): ] router = Router(model_list=model_list) messages = [{"role": "user", "content": "hi"}] - with patch.object(router, "completion", return_value=None): + with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None): router._silent_experiment_completion( silent_model="gpt-3.5-turbo", messages=messages, @@ -173,12 +182,20 @@ def test_router_silent_experiment_completion(): router = Router(model_list=model_list) - # Mock litellm.completion + # Mock litellm.acompletion mock_response = litellm.ModelResponse(choices=[{"message": {"content": "hello"}}]) - mock_completion = MagicMock(return_value=mock_response) + + # We need an async mock for acompletion + async def mock_acompletion(*args, **kwargs): + return mock_response + + mock_acompletion_mock = AsyncMock(side_effect=mock_acompletion) + mock_completion_mock = MagicMock(return_value=mock_response) # Patch at the litellm module level - with patch.object(litellm, "completion", mock_completion): + with patch.object(litellm, "acompletion", mock_acompletion_mock), patch.object( + litellm, "completion", mock_completion_mock + ): response = router.completion( model="primary-model", messages=[{"role": "user", "content": "hi"}], @@ -189,12 +206,15 @@ def test_router_silent_experiment_completion(): # The sync background call uses a thread pool. We might need to wait a bit. import time - time.sleep(0.5) + time.sleep(2.0) - # Should have 2 calls - assert mock_completion.call_count == 2 + # Should have 1 acompletion call (the silent background call) + # The primary completion call still goes to the real litellm.completion (or we can mock it separately, but here it's testing the background one) + # Wait, the primary call in the test is router.completion. + # Actually, let's just mock both to avoid real network calls if it's hitting one. + assert mock_acompletion_mock.call_count == 1 - call_args_list = mock_completion.call_args_list + call_args_list = mock_acompletion_mock.call_args_list # Verify no silent_model in any call for call in call_args_list: @@ -212,3 +232,5 @@ def test_router_silent_experiment_completion(): ) assert silent_call is not None assert silent_call[1]["model"] == "openai/gpt-4" + # Verify model_group is set to the silent model name for correct metric attribution + assert silent_call[1]["metadata"]["model_group"] == "silent-model"