From b8fe0e057f7cf42f0f88de67a3f7fc6b68fe23a7 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 10 May 2025 12:04:01 -0700 Subject: [PATCH] complete unified batch id support - replace model in jsonl to be deployment model name (#10719) * feat(router.py): translate the model in jsonl for create file deployment to use the deployment model name * test: add unit test for replace model in jsonl * test(test_router.py): add unit tests * test: add unit tests --- litellm/proxy/_new_secret_config.yaml | 2 +- litellm/router.py | 33 ++++++- litellm/router_utils/batch_utils.py | 34 +++++-- tests/litellm/test_router.py | 88 +++++++++++++++++++ .../test_router_batch_utils.py | 31 +++++++ 5 files changed, 179 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 65ef4b89296..b8ea7aa9a89 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -71,4 +71,4 @@ model_list: api_key: 04d22fb7e9ad4d9c8afe7c6abf97a6fc model_info: id: my-unique-azure-deployment - mode: batch \ No newline at end of file + mode: batch diff --git a/litellm/router.py b/litellm/router.py index e0068cacbc9..528e8859ed7 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -69,7 +69,11 @@ from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, add_retry_headers_to_response, ) -from litellm.router_utils.batch_utils import _get_router_metadata_variable_name +from litellm.router_utils.batch_utils import ( + _get_router_metadata_variable_name, + replace_model_in_jsonl, + should_replace_model_in_jsonl, +) from litellm.router_utils.client_initalization_utils import InitalizeCachedClient from litellm.router_utils.clientside_credential_handler import ( get_dynamic_litellm_params, @@ -106,7 +110,12 @@ from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_successes_for_current_minute, ) from litellm.scheduler import FlowItem, Scheduler -from litellm.types.llms.openai import AllMessageValues, FileTypes, OpenAIFileObject +from litellm.types.llms.openai import ( + AllMessageValues, + FileTypes, + OpenAIFileObject, + OpenAIFilesPurpose, +) from litellm.types.router import ( CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, VALID_LITELLM_ENVIRONMENTS, @@ -2766,6 +2775,26 @@ class Router: stripped_model, custom_llm_provider, _, _ = get_llm_provider( model=data["model"] ) + + ## REPLACE MODEL IN FILE WITH SELECTED DEPLOYMENT ## + purpose = cast(Optional[OpenAIFilesPurpose], kwargs.get("purpose")) + file = cast(Optional[FileTypes], kwargs.get("file")) + if not file or not purpose: + raise Exception( + "file and file_purpose are required for create_file" + ) + + replace_model_in_jsonl_bool = should_replace_model_in_jsonl( + purpose=purpose, + ) + if replace_model_in_jsonl_bool: + file = replace_model_in_jsonl( + file_content=file, + new_model_name=stripped_model, + ) + + kwargs_copy["file"] = file + response = litellm.acreate_file( **{ **data, diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index a55be22913e..1b9621c1832 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -1,6 +1,9 @@ import io import json -from typing import Optional, Tuple, Union +from os import PathLike +from typing import Optional + +from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose class InMemoryFile(io.BytesIO): @@ -9,10 +12,25 @@ class InMemoryFile(io.BytesIO): self.name = name -def replace_model_in_jsonl( - file_content: Union[bytes, Tuple[str, bytes, str]], new_model_name: str -) -> Optional[InMemoryFile]: +def should_replace_model_in_jsonl( + purpose: OpenAIFilesPurpose, +) -> bool: + """ + Check if the model name should be replaced in the JSONL file for the deployment model name. + + Azure raises an error on create batch if the model name for deployment is not in the .jsonl. + """ + if purpose == "batch": + return True + return False + + +def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> FileTypes: try: + ## if pathlike, return the original file content + if isinstance(file_content, PathLike): + return file_content + # Decode the bytes to a string and split into lines # If file_content is a file-like object, read the bytes if hasattr(file_content, "read"): @@ -25,8 +43,11 @@ def replace_model_in_jsonl( # Decode the bytes to a string and split into lines if isinstance(file_content_bytes, bytes): file_content_str = file_content_bytes.decode("utf-8") - else: + elif isinstance(file_content_bytes, str): file_content_str = file_content_bytes + else: + return file_content + lines = file_content_str.splitlines() modified_lines = [] for line in lines: @@ -45,7 +66,8 @@ def replace_model_in_jsonl( return InMemoryFile(modified_file_content, name="modified_file.jsonl") # type: ignore except (json.JSONDecodeError, UnicodeDecodeError, TypeError): - return None + # return the original file content if there is an error replacing the model name + return file_content def _get_router_metadata_variable_name(function_name: Optional[str]) -> str: diff --git a/tests/litellm/test_router.py b/tests/litellm/test_router.py index e645e7f2174..fb40fa0b20e 100644 --- a/tests/litellm/test_router.py +++ b/tests/litellm/test_router.py @@ -148,6 +148,94 @@ async def test_router_acreate_file(): # assert that the mock_acreate_file was called twice assert mock_acreate_file.call_count == 2 +@pytest.mark.asyncio +async def test_router_acreate_file_with_jsonl(): + """ + Test router.acreate_file with both JSONL and non-JSONL files + """ + import json + from io import BytesIO + from unittest.mock import MagicMock, patch + + # Create test JSONL content + jsonl_data = [ + { + "body": { + "model": "gpt-3.5-turbo-router", + "messages": [{"role": "user", "content": "test"}], + } + }, + { + "body": { + "model": "gpt-3.5-turbo-router", + "messages": [{"role": "user", "content": "test2"}], + } + }, + ] + jsonl_content = "\n".join(json.dumps(item) for item in jsonl_data) + jsonl_file = BytesIO(jsonl_content.encode("utf-8")) + jsonl_file.name = "test.jsonl" + + # Create test non-JSONL content + non_jsonl_content = "This is not a JSONL file" + non_jsonl_file = BytesIO(non_jsonl_content.encode("utf-8")) + non_jsonl_file.name = "test.txt" + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo-router", + "litellm_params": {"model": "gpt-3.5-turbo"}, + }, + { + "model_name": "gpt-3.5-turbo-router", + "litellm_params": {"model": "gpt-4o-mini"}, + }, + ], + ) + + with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file: + # Test with JSONL file + response = await router.acreate_file( + model="gpt-3.5-turbo-router", + purpose="batch", + file=jsonl_file, + ) + + # Verify mock was called twice (once for each deployment) + print(f"mock_acreate_file.call_count: {mock_acreate_file.call_count}") + print(f"mock_acreate_file.call_args_list: {mock_acreate_file.call_args_list}") + assert mock_acreate_file.call_count == 2 + + # Get the file content passed to the first call + first_call_file = mock_acreate_file.call_args_list[0][1]["file"] + first_call_content = first_call_file.read().decode("utf-8") + + # Verify the model name was replaced in the JSONL content + first_line = json.loads(first_call_content.split("\n")[0]) + assert first_line["body"]["model"] == "gpt-3.5-turbo" + + # Reset mock for next test + mock_acreate_file.reset_mock() + + # Test with non-JSONL file + response = await router.acreate_file( + model="gpt-3.5-turbo-router", + purpose="user_data", + file=non_jsonl_file, + ) + + # Verify mock was called twice + assert mock_acreate_file.call_count == 2 + + # Get the file content passed to the first call + first_call_file = mock_acreate_file.call_args_list[0][1]["file"] + first_call_content = first_call_file.read().decode("utf-8") + + # Verify the non-JSONL content was not modified + assert first_call_content == non_jsonl_content + + @pytest.mark.asyncio async def test_router_async_get_healthy_deployments(): """ diff --git a/tests/router_unit_tests/test_router_batch_utils.py b/tests/router_unit_tests/test_router_batch_utils.py index 40a491a84a5..94cd6e001e4 100644 --- a/tests/router_unit_tests/test_router_batch_utils.py +++ b/tests/router_unit_tests/test_router_batch_utils.py @@ -88,3 +88,34 @@ def test_router_metadata_variable_name(): assert ( _get_router_metadata_variable_name(function_name="aget_file") == "litellm_metadata" ) + + +def test_non_json_input(): + """Test that replace_model_in_jsonl returns original content for non-JSON input""" + from litellm.router_utils.batch_utils import replace_model_in_jsonl + + # Test with non-JSON string + non_json_str = "This is not a JSON string" + result = replace_model_in_jsonl(non_json_str, "gpt-4") + assert result == non_json_str + + # Test with non-JSON bytes + non_json_bytes = b"This is not JSON bytes" + result = replace_model_in_jsonl(non_json_bytes, "gpt-4") + assert result == non_json_bytes + + # Test with non-JSON file-like object + from io import BytesIO + non_json_file = BytesIO(b"This is not JSON in a file") + result = replace_model_in_jsonl(non_json_file, "gpt-4") + assert result == non_json_file + + +def test_should_replace_model_in_jsonl(): + """Test that should_replace_model_in_jsonl returns the correct value""" + from litellm.router_utils.batch_utils import should_replace_model_in_jsonl + assert should_replace_model_in_jsonl(purpose="batch") == True + assert should_replace_model_in_jsonl(purpose="test") == False + assert should_replace_model_in_jsonl(purpose="user_data") == False + + \ No newline at end of file