From 931b2e487521ca930bec2987ee6661c1b5340be5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 20 Jun 2025 14:51:50 -0700 Subject: [PATCH] [Bug Fix] Fix model_group tracked for /v1/messages and /moderations (#11933) * fixes _get_router_metadata_variable_name * fixes _update_kwargs_before_fallbacks * test_anthropic_messages_litellm_router_non_streaming_with_logging * test_moderations_api_logging * fix _pass_through_moderation_endpoint_factory --- litellm/router.py | 15 +++++++- litellm/router_utils/batch_utils.py | 2 +- .../test_moderations_api_logging.py | 36 ++++++++++++++++--- .../test_anthropic_messages_passthrough.py | 9 +++-- 4 files changed, 53 insertions(+), 9 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index e7297b0d413..cd4b94d4965 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2454,6 +2454,14 @@ class Router: The response from the handler function """ handler_name = original_function.__name__ + function_name = "_ageneric_api_call_with_fallbacks" + self._update_kwargs_before_fallbacks( + model=model, + kwargs=kwargs, + metadata_variable_name = _get_router_metadata_variable_name( + function_name=function_name + ) + ) try: verbose_router_logger.debug( f"Inside _ageneric_api_call() - handler: {handler_name}, model: {model}; kwargs: {kwargs}" @@ -2467,7 +2475,7 @@ class Router: ) self._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name="generic_api_call" + deployment=deployment, kwargs=kwargs, function_name=function_name ) data = deployment["litellm_params"].copy() @@ -3180,6 +3188,11 @@ class Router: original_function: Callable, **kwargs, ): + # update kwargs with model_group + self._update_kwargs_before_fallbacks( + model=kwargs.get("model", ""), + kwargs=kwargs, + ) if kwargs.get("model") and self.get_model_list(model_name=kwargs["model"]): deployment = await self.async_get_available_deployment( model=kwargs["model"], diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index 1b9621c1832..6617ad1f68e 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -79,7 +79,7 @@ def _get_router_metadata_variable_name(function_name: Optional[str]) -> str: For ALL other endpoints we call this "metadata """ ROUTER_METHODS_USING_LITELLM_METADATA = set( - ["batch", "generic_api_call", "_acreate_batch", "file"] + ["batch", "generic_api_call", "_acreate_batch", "file", "_ageneric_api_call_with_fallbacks"] ) if function_name and any( method in function_name for method in ROUTER_METHODS_USING_LITELLM_METADATA diff --git a/tests/logging_callback_tests/test_moderations_api_logging.py b/tests/logging_callback_tests/test_moderations_api_logging.py index 129e802763b..1d3b5356b0f 100644 --- a/tests/logging_callback_tests/test_moderations_api_logging.py +++ b/tests/logging_callback_tests/test_moderations_api_logging.py @@ -17,6 +17,7 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import litellm +from litellm.router import Router import asyncio from typing import Optional from litellm.types.utils import StandardLoggingPayload, Usage, ModelInfoBase @@ -41,7 +42,8 @@ class TestCustomLogger(CustomLogger): @pytest.mark.asyncio @pytest.mark.parametrize("model", [ None, - "omni-moderation-latest" + "omni-moderation-latest", + "router-internal-moderation-model" ]) async def test_moderations_api_logging(model): @@ -51,12 +53,31 @@ async def test_moderations_api_logging(model): custom_logger = TestCustomLogger() litellm.logging_callback_manager.add_litellm_callback(custom_logger) - input_content = "Hello, how are you?" - response = await litellm.amoderation( - input=input_content, - model=model, + + MODEL_GROUP = "internal-moderation-model" + router = Router( + model_list=[ + { + "model_name": MODEL_GROUP, + "litellm_params": { + "model": "openai/omni-moderation-latest", + }, + } + ] ) + input_content = "Hello, how are you?" + if model == "router-internal-moderation-model": + response = await router.amoderation( + input=input_content, + model=MODEL_GROUP, + ) + else: + response = await litellm.amoderation( + input=input_content, + model=model, + ) + print("response", json.dumps(response, indent=4, default=str)) await asyncio.sleep(2) @@ -76,3 +97,8 @@ async def test_moderations_api_logging(model): # assert the logged response == response user received client side assert dict(standard_logging_payload["response"]) == response.model_dump() + + # if router used, validate model_group is logged as expected + if model == "router-internal-moderation-model": + assert standard_logging_payload["model_group"] == MODEL_GROUP + diff --git a/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py b/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py index 9f032ffaadd..0fbc95e4480 100644 --- a/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py +++ b/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py @@ -222,10 +222,11 @@ async def test_anthropic_messages_litellm_router_non_streaming_with_logging(): test_custom_logger = TestCustomLogger() litellm.callbacks = [test_custom_logger] litellm._turn_on_debug() + MODEL_GROUP = "claude-special-alias" router = Router( model_list=[ { - "model_name": "claude-special-alias", + "model_name": MODEL_GROUP, "litellm_params": { "model": "claude-3-haiku-20240307", "api_key": os.getenv("ANTHROPIC_API_KEY"), @@ -240,7 +241,7 @@ async def test_anthropic_messages_litellm_router_non_streaming_with_logging(): # Call the handler response = await router.aanthropic_messages( messages=messages, - model="claude-special-alias", + model=MODEL_GROUP, max_tokens=100, ) @@ -252,6 +253,7 @@ async def test_anthropic_messages_litellm_router_non_streaming_with_logging(): await asyncio.sleep(1) assert test_custom_logger.logged_standard_logging_payload is not None, "Logging payload should not be None" + print("tracked standard logging payload", json.dumps(test_custom_logger.logged_standard_logging_payload, indent=4, default=str)) assert test_custom_logger.logged_standard_logging_payload["messages"] == messages assert test_custom_logger.logged_standard_logging_payload["response"] is not None assert ( @@ -270,6 +272,9 @@ async def test_anthropic_messages_litellm_router_non_streaming_with_logging(): == response["usage"]["output_tokens"] ) + # assert model_group + assert test_custom_logger.logged_standard_logging_payload["model_group"] == MODEL_GROUP + @pytest.mark.asyncio async def test_anthropic_messages_litellm_router_streaming_with_logging():