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

@@ -71,4 +71,4 @@ model_list:
       api_key: 04d22fb7e9ad4d9c8afe7c6abf97a6fc
     model_info: 
       id: my-unique-azure-deployment
-      mode: batch
+      mode: batch

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,

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:

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():
     """

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
+    
+