litellm/tests/unit/router_utils/test_batch_utils.py

import asyncio
import importlib
import json
from io import BytesIO
from typing import Dict, Final, List

import pytest

import litellm
from litellm.router_utils.batch_utils import (
    InMemoryFile,
    get_router_metadata_variable_name,
    parse_jsonl_with_embedded_newlines,
    replace_model_in_jsonl,
)
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome


# Fixtures
@pytest.fixture
def sample_jsonl_data() -> List[Dict]:
    """Fixture providing sample JSONL data"""
    return [
        {
            "body": {
                "model": "gpt-5-mini",
                "messages": [{"role": "user", "content": "Hello"}],
            }
        },
        {"body": {"model": "gpt-5.5", "messages": [{"role": "user", "content": "Hi"}]}},
    ]


@pytest.fixture
def sample_jsonl_bytes(sample_jsonl_data) -> bytes:
    """Fixture providing sample JSONL as bytes"""
    jsonl_str = "\n".join(json.dumps(line) for line in sample_jsonl_data)
    return jsonl_str.encode("utf-8")


@pytest.fixture
def sample_file_like(sample_jsonl_bytes):
    """Fixture providing a file-like object"""
    return BytesIO(sample_jsonl_bytes)


# Test cases
def test_bytes_input(sample_jsonl_bytes):
    """Test with bytes input"""
    new_model = "claude-3"
    result = replace_model_in_jsonl(sample_jsonl_bytes, new_model)

    assert result is not None
    assert isinstance(result, InMemoryFile)
    assert result.name == "modified_file.jsonl"
    assert result.content_type == "application/jsonl"


def test_tuple_input(sample_jsonl_bytes):
    """Test with tuple input"""
    new_model = "claude-3"
    test_tuple = ("test.jsonl", sample_jsonl_bytes, "application/json")
    result = replace_model_in_jsonl(test_tuple, new_model)

    assert result is not None
    assert isinstance(result, InMemoryFile)
    assert result.name == "modified_file.jsonl"
    assert result.content_type == "application/jsonl"


def test_tuple_with_file_handle_rewrites_model(sample_jsonl_bytes):
    """Security regression: when the tuple's content element is a file handle
    (batch uploads stream from the spooled upload handle), the model must still
    be rewritten. Otherwise a restricted body.model survives unmodified and
    bypasses the batch model allowlist, which only checks the upload target."""
    new_model = "approved-target-model"
    handle = BytesIO(sample_jsonl_bytes)
    test_tuple = ("test.jsonl", handle, "application/json")

    result = replace_model_in_jsonl(test_tuple, new_model)

    assert isinstance(result, InMemoryFile)
    rows = [json.loads(line) for line in result.getvalue().decode("utf-8").splitlines() if line.strip()]
    assert rows, "rewrite must produce rows"
    # every row now carries the rewritten target, not the original (restricted) model
    assert all(row["body"]["model"] == new_model for row in rows)
    assert all(row["body"]["model"] != "gpt-5.5" for row in rows)


def test_file_like_object(sample_file_like):
    """Test with file-like object input"""
    new_model = "claude-3"
    result = replace_model_in_jsonl(sample_file_like, new_model)

    assert result is not None
    assert isinstance(result, InMemoryFile)
    assert result.name == "modified_file.jsonl"
    assert result.content_type == "application/jsonl"


def test_router_metadata_variable_name():
    """Test that the variable name is correct"""
    assert get_router_metadata_variable_name(function_name="completion") == "metadata"
    assert (
        get_router_metadata_variable_name(function_name="batch") == "litellm_metadata"
    )
    assert (
        get_router_metadata_variable_name(function_name="acreate_file")
        == "litellm_metadata"
    )
    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") is True
    assert should_replace_model_in_jsonl(purpose="batch", passthrough=True) is False
    assert should_replace_model_in_jsonl(purpose="test") is False
    assert should_replace_model_in_jsonl(purpose="user_data") is False


def test_parse_jsonl_with_embedded_newlines_simple():
    """Test parsing simple JSONL without embedded newlines"""
    content = '{"id": 1, "name": "test"}\n{"id": 2, "name": "test2"}'
    result = parse_jsonl_with_embedded_newlines(content)

    assert len(result) == 2
    assert result[0] == {"id": 1, "name": "test"}
    assert result[1] == {"id": 2, "name": "test2"}


def test_parse_jsonl_with_embedded_newlines_in_strings():
    """Test parsing JSONL with newlines embedded in string values"""
    content = '{"id": 1, "message": "Line 1\\nLine 2\\nLine 3"}\n{"id": 2, "message": "test"}'
    result = parse_jsonl_with_embedded_newlines(content)

    assert len(result) == 2
    assert result[0] == {"id": 1, "message": "Line 1\nLine 2\nLine 3"}
    assert result[1] == {"id": 2, "message": "test"}


def test_parse_jsonl_with_embedded_newlines_real_world_example():
    """Test with the real-world example from the Cooler Master Shark X case"""
    # This simulates the actual problem case from the user's log
    content = """{"custom_id":"16546277850245725","method":"POST","url":"/v1/chat/completions","body":{"model":"openai-gpt-4o-mini-dp-items-translation-dag","messages":[{"role":"system","content":"Translate the product title and description for an e-commerce marketplace in Saudi Arabia and the UAE. Text may be in English or Arabic.\\n"},{"role":"user","content":"\\nOriginal Title: ```Cooler Master Shark X PC Case```\\nOriginal Description: ```UNIQUE MASTERPIECEShark X is a system that provides an impressive  unique alternative to traditional PC systems.  Shark X will stand out and can be the ultimate  trophy or conversation piece for people looking  for a unique setup that stands head and fins  above the res.```\\nStore Name: ```geekay```\\n"}]}}"""

    result = parse_jsonl_with_embedded_newlines(content)

    assert len(result) == 1
    assert result[0]["custom_id"] == "16546277850245725"
    assert result[0]["method"] == "POST"
    assert result[0]["body"]["model"] == "openai-gpt-4o-mini-dp-items-translation-dag"
    assert len(result[0]["body"]["messages"]) == 2
    assert "Translate the product title" in result[0]["body"]["messages"][0]["content"]
    assert "Cooler Master Shark X PC Case" in result[0]["body"]["messages"][1]["content"]
    assert "UNIQUE MASTERPIECEShark X" in result[0]["body"]["messages"][1]["content"]


def test_parse_jsonl_with_embedded_newlines_multiple_complex_objects():
    """Test parsing multiple complex JSON objects with embedded newlines"""
    content = """{"id":1,"text":"Line 1\\nLine 2"}
{"id":2,"nested":{"field":"Value\\nWith\\nNewlines"}}
{"id":3,"simple":"test"}"""

    result = parse_jsonl_with_embedded_newlines(content)

    assert len(result) == 3
    assert result[0]["id"] == 1
    assert result[0]["text"] == "Line 1\nLine 2"
    assert result[1]["id"] == 2
    assert result[1]["nested"]["field"] == "Value\nWith\nNewlines"
    assert result[2]["id"] == 3
    assert result[2]["simple"] == "test"


def test_parse_jsonl_with_embedded_newlines_no_trailing_newline():
    """Test parsing JSONL without trailing newline"""
    content = '{"id": 1, "name": "test"}'
    result = parse_jsonl_with_embedded_newlines(content)

    assert len(result) == 1
    assert result[0] == {"id": 1, "name": "test"}


def test_parse_jsonl_with_embedded_newlines_empty_string():
    """Test parsing empty string"""
    content = ""
    result = parse_jsonl_with_embedded_newlines(content)

    assert len(result) == 0


def test_parse_jsonl_with_embedded_newlines_whitespace_only():
    """Test parsing whitespace-only content"""
    content = "   \n  \n  "
    result = parse_jsonl_with_embedded_newlines(content)

    assert len(result) == 0


def test_replace_model_in_jsonl_malformed_middle_row_returns_original():
    """Regression: a malformed/truncated middle row must not silently drop the
    rows that follow it. The streaming rewrite accumulates physical lines into a
    buffer; a row that never parses poisons the buffer so every later valid row
    is concatenated into it and dropped. Returning that partial rewrite would
    ship a truncated batch with no error to the caller. Instead the original
    content is returned unchanged so the provider rejects the bad batch loudly."""
    content = (
        b'{"custom_id":"a","body":{"model":"x"}}\n'
        b'{"custom_id":"b","body":{"model":\n'  # truncated, never completes
        b'{"custom_id":"c","body":{"model":"x"}}\n'
    )

    result = replace_model_in_jsonl(content, "new-model")

    assert result == content, "must return the original unchanged, not a partial rewrite"


def test_replace_model_in_jsonl_malformed_row_seekable_handle_rewound():
    """When the source is a seekable handle that gets consumed during the failed
    rewrite, it must be rewound to 0 so the caller can re-read the full original."""
    content = (
        b'{"custom_id":"a","body":{"model":"x"}}\n'
        b'{"custom_id":"b","body":{"model":\n'
        b'{"custom_id":"c","body":{"model":"x"}}\n'
    )
    handle = BytesIO(content)

    result = replace_model_in_jsonl(handle, "new-model")

    assert result is handle
    assert handle.read() == content, "handle must be rewound for the caller to re-read"


def test_replace_model_in_jsonl_multi_row_rewrites_every_model():
    """Happy path: a well-formed multi-row file gets every row's model rewritten
    and no row is dropped."""
    content = (
        b'{"custom_id":"a","body":{"model":"old1"}}\n'
        b'{"custom_id":"b","body":{"model":"old2"}}\n'
        b'{"custom_id":"c","body":{"model":"old3"}}\n'
    )

    result = replace_model_in_jsonl(content, "new-model")

    assert isinstance(result, InMemoryFile)
    rows = [json.loads(line) for line in result.getvalue().decode("utf-8").splitlines() if line.strip()]
    assert [row["custom_id"] for row in rows] == ["a", "b", "c"]
    assert all(row["body"]["model"] == "new-model" for row in rows)


def test_replace_model_in_jsonl_keeps_every_row_on_its_original_line():
    content = (
        b'{"custom_id":"a","body":{"model":"old"}}\n'
        b"\n"
        b'{"custom_id":"c","body":{"model":"old"}}\n'
        b'{"custom_id":"d",\n'
        b'"body":{"model":"old"}}\n'
        b'{"custom_id":"f","body":{"model":"old"}}\n'
    )

    result = replace_model_in_jsonl(content, "new-model", keep_line_numbers=True)

    assert isinstance(result, InMemoryFile)
    lines = result.getvalue().decode("utf-8").split("\n")
    assert [json.loads(line)["custom_id"] if line else None for line in lines] == ["a", None, "c", None, "d", "f"]
    assert all(json.loads(line)["body"]["model"] == "new-model" for line in lines if line)


@pytest.mark.parametrize(
    "content",
    [
        b'\xef\xbb\xbf{"custom_id":"a","body":{"model":"old"}}\n\xef\xbb\xbf{"custom_id":"b","body":{"model":"old"}}\n',
        '\ufeff{"custom_id":"a","body":{"model":"old"}}\n\ufeff{"custom_id":"b","body":{"model":"old"}}\n',
    ],
    ids=["bytes", "str"],
)
def test_replace_model_in_jsonl_rewrites_rows_that_start_with_a_bom(content):
    result = replace_model_in_jsonl(content, "new-model")

    assert isinstance(result, InMemoryFile)
    rows = [json.loads(line) for line in result.getvalue().decode("utf-8").split("\n")]
    assert [(row["custom_id"], row["body"]["model"]) for row in rows] == [("a", "new-model"), ("b", "new-model")]


def test_replace_model_in_jsonl_with_embedded_newlines():
    """Test that replace_model_in_jsonl works correctly with embedded newlines in content"""
    # Create a JSONL with embedded newlines in the message content
    jsonl_data = {
        "custom_id": "test123",
        "body": {
            "model": "old-model",
            "messages": [{"role": "user", "content": "This is a message\nwith multiple\nlines"}],
        },
    }

    jsonl_bytes = json.dumps(jsonl_data).encode("utf-8")
    new_model = "new-model"

    result = replace_model_in_jsonl(jsonl_bytes, new_model)

    assert isinstance(result, InMemoryFile)

    # Read and parse the result
    result_content = result.read().decode("utf-8")
    result_json = json.loads(result_content)

    # Verify the model was replaced
    assert result_json["body"]["model"] == "new-model"
    # Verify the content with newlines is preserved
    assert result_json["body"]["messages"][0]["content"] == "This is a message\nwith multiple\nlines"
    assert result_json["custom_id"] == "test123"


def test_is_batch_retrieve_call_type_matches_only_batch_retrieves():
    from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
    from litellm.types.utils import CallTypes

    assert is_batch_retrieve_call_type(CallTypes.aretrieve_batch.value) is True
    assert is_batch_retrieve_call_type(CallTypes.retrieve_batch.value) is True

    for call_type in CallTypes:
        if call_type in (CallTypes.aretrieve_batch, CallTypes.retrieve_batch):
            continue
        assert is_batch_retrieve_call_type(call_type.value) is False

    assert is_batch_retrieve_call_type(None) is False


@pytest.fixture(autouse=True)
def _vcr_outcome_gate(request, vcr):
    install_live_call_probe(request, vcr)
    yield
    record_vcr_outcome(request, vcr)


@pytest.fixture(scope="function", autouse=True)
def setup_and_teardown():
    """
    This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained.
    """
    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER

    asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue())
    importlib.reload(litellm)
    try:
        if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"):
            importlib.reload(litellm.proxy.proxy_server)
    except Exception:
        pass
    loop = asyncio.get_event_loop_policy().new_event_loop()
    asyncio.set_event_loop(loop)
    yield
    loop.close()
    asyncio.set_event_loop(None)


@pytest.mark.parametrize(
    ("call_type", "cost", "eligible"),
    (("aretrieve_batch", 0.25, True), ("retrieve_batch", 0.25, True), ("acompletion", 0.25, False),
     ("aretrieve_batch", 0, False), ("aretrieve_batch", -0.25, False)),
)
def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, cost: float, eligible: bool) -> None:
    from openai.types import Batch
    from litellm.router_utils.batch_utils import batch_id_to_charge_once

    batch: Final = Batch(
        id="batch-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions",
        input_file_id="file-budget", object="batch", status="completed",
    )
    assert batch_id_to_charge_once(call_type, batch, cost) == (batch.id if eligible else None)
    assert batch_id_to_charge_once(call_type, {"id": batch.id}, cost) is None


def test_batch_charge_markers_keep_counters_and_batches_independent() -> None:
    from litellm.caching.in_memory_cache import InMemoryCache
    from litellm.router_utils.batch_utils import batch_charged_once_marker_key

    cache: Final = InMemoryCache()
    cache.set_cache(batch_charged_once_marker_key("provider-spend", "first-batch"), 1)
    assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "first-batch")) == 1
    assert cache.get_cache(batch_charged_once_marker_key("deployment-spend", "first-batch")) is None
    assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "second-batch")) is None