litellm/tests/batches_tests/test_batches_logging_unit_tests.py

import asyncio
import json
import os
import sys
import traceback
from unittest.mock import AsyncMock, MagicMock, patch
from dotenv import load_dotenv

load_dotenv()
sys.path.insert(
    0, os.path.abspath("../..")
)  # Adds the parent directory to the system-path
import logging
import time

import pytest
from typing import Optional
import litellm
from litellm import create_batch, create_file
from litellm._logging import verbose_logger
from litellm.batches.batch_utils import (
    _batch_cost_calculator,
    _get_file_content_as_dictionary,
    _get_batch_job_cost_from_file_content,
    _get_batch_job_total_usage_from_file_content,
    _get_batch_job_usage_from_response_body,
    _get_response_from_batch_job_output_file,
    _batch_response_was_successful,
)


@pytest.fixture
def sample_file_content():
    return b"""
{"id": "batch_req_6769ca596b38819093d7ae9f522de924", "custom_id": "request-1", "response": {"status_code": 200, "request_id": "07bc45ab4e7e26ac23a0c949973327e7", "body": {"id": "chatcmpl-AhjSMl7oZ79yIPHLRYgmgXSixTJr7", "object": "chat.completion", "created": 1734986202, "model": "gpt-4o-mini-2024-07-18", "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello! How can I assist you today?", "refusal": null}, "logprobs": null, "finish_reason": "stop"}], "usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30, "prompt_tokens_details": {"cached_tokens": 0, "audio_tokens": 0}, "completion_tokens_details": {"reasoning_tokens": 0, "audio_tokens": 0, "accepted_prediction_tokens": 0, "rejected_prediction_tokens": 0}}, "system_fingerprint": "fp_0aa8d3e20b"}}, "error": null}
{"id": "batch_req_6769ca597e588190920666612634e2b4", "custom_id": "request-2", "response": {"status_code": 200, "request_id": "82e04f4c001fe2c127cbad199f5fd31b", "body": {"id": "chatcmpl-AhjSNgVB4Oa4Hq0NruTRsBaEbRWUP", "object": "chat.completion", "created": 1734986203, "model": "gpt-4o-mini-2024-07-18", "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello! What can I do for you today?", "refusal": null}, "logprobs": null, "finish_reason": "length"}], "usage": {"prompt_tokens": 22, "completion_tokens": 10, "total_tokens": 32, "prompt_tokens_details": {"cached_tokens": 0, "audio_tokens": 0}, "completion_tokens_details": {"reasoning_tokens": 0, "audio_tokens": 0, "accepted_prediction_tokens": 0, "rejected_prediction_tokens": 0}}, "system_fingerprint": "fp_0aa8d3e20b"}}, "error": null}
"""


@pytest.fixture
def sample_file_content_dict():
    return [
        {
            "id": "batch_req_6769ca596b38819093d7ae9f522de924",
            "custom_id": "request-1",
            "response": {
                "status_code": 200,
                "request_id": "07bc45ab4e7e26ac23a0c949973327e7",
                "body": {
                    "id": "chatcmpl-AhjSMl7oZ79yIPHLRYgmgXSixTJr7",
                    "object": "chat.completion",
                    "created": 1734986202,
                    "model": "gpt-4o-mini-2024-07-18",
                    "choices": [
                        {
                            "index": 0,
                            "message": {
                                "role": "assistant",
                                "content": "Hello! How can I assist you today?",
                                "refusal": None,
                            },
                            "logprobs": None,
                            "finish_reason": "stop",
                        }
                    ],
                    "usage": {
                        "prompt_tokens": 20,
                        "completion_tokens": 10,
                        "total_tokens": 30,
                        "prompt_tokens_details": {
                            "cached_tokens": 0,
                            "audio_tokens": 0,
                        },
                        "completion_tokens_details": {
                            "reasoning_tokens": 0,
                            "audio_tokens": 0,
                            "accepted_prediction_tokens": 0,
                            "rejected_prediction_tokens": 0,
                        },
                    },
                    "system_fingerprint": "fp_0aa8d3e20b",
                },
            },
            "error": None,
        },
        {
            "id": "batch_req_6769ca597e588190920666612634e2b4",
            "custom_id": "request-2",
            "response": {
                "status_code": 200,
                "request_id": "82e04f4c001fe2c127cbad199f5fd31b",
                "body": {
                    "id": "chatcmpl-AhjSNgVB4Oa4Hq0NruTRsBaEbRWUP",
                    "object": "chat.completion",
                    "created": 1734986203,
                    "model": "gpt-4o-mini-2024-07-18",
                    "choices": [
                        {
                            "index": 0,
                            "message": {
                                "role": "assistant",
                                "content": "Hello! What can I do for you today?",
                                "refusal": None,
                            },
                            "logprobs": None,
                            "finish_reason": "length",
                        }
                    ],
                    "usage": {
                        "prompt_tokens": 22,
                        "completion_tokens": 10,
                        "total_tokens": 32,
                        "prompt_tokens_details": {
                            "cached_tokens": 0,
                            "audio_tokens": 0,
                        },
                        "completion_tokens_details": {
                            "reasoning_tokens": 0,
                            "audio_tokens": 0,
                            "accepted_prediction_tokens": 0,
                            "rejected_prediction_tokens": 0,
                        },
                    },
                    "system_fingerprint": "fp_0aa8d3e20b",
                },
            },
            "error": None,
        },
    ]


def test_get_file_content_as_dictionary(sample_file_content):
    result = _get_file_content_as_dictionary(sample_file_content)
    assert len(result) == 2
    assert result[0]["id"] == "batch_req_6769ca596b38819093d7ae9f522de924"
    assert result[0]["custom_id"] == "request-1"
    assert result[0]["response"]["status_code"] == 200
    assert result[0]["response"]["body"]["usage"]["total_tokens"] == 30


def test_get_batch_job_total_usage_from_file_content(sample_file_content_dict):
    usage = _get_batch_job_total_usage_from_file_content(
        sample_file_content_dict, custom_llm_provider="openai"
    )
    assert usage.total_tokens == 62  # 30 + 32
    assert usage.prompt_tokens == 42  # 20 + 22
    assert usage.completion_tokens == 20  # 10 + 10


@pytest.mark.asyncio
async def test_batch_cost_calculator(sample_file_content_dict):
    """
    mock litellm.completion_cost to return 0.5

    we know sample_file_content_dict has 2 successful responses

    so we expect the cost to be 0.5 * 2 = 1.0
    """
    with patch("litellm.completion_cost", return_value=0.5):
        cost = _batch_cost_calculator(
            file_content_dictionary=sample_file_content_dict,
            custom_llm_provider="openai",
        )
        assert cost == 1.0  # 0.5 * 2 successful responses


def test_get_response_from_batch_job_output_file(sample_file_content_dict):
    result = _get_response_from_batch_job_output_file(sample_file_content_dict[0])
    assert result["id"] == "chatcmpl-AhjSMl7oZ79yIPHLRYgmgXSixTJr7"
    assert result["object"] == "chat.completion"
    assert result["usage"]["total_tokens"] == 30


@pytest.mark.asyncio
async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cost():
    """
    Test that cost is calculated for completed batches when no explicit cost data is provided.

    Regression test for: When batch status is "completed" and explicit batch_cost/batch_usage/batch_models
    are not provided, the system should compute batch data by calling _handle_completed_batch.
    """
    from litellm.litellm_core_utils.litellm_logging import Logging
    from litellm.types.utils import CallTypes
    from litellm.types.utils import LiteLLMBatch
    from unittest.mock import AsyncMock, patch

    # Mock batch result with completed status
    mock_batch = LiteLLMBatch(
        id="batch-test-123",
        object="batch",
        endpoint="/v1/chat/completions",
        errors=None,
        input_file_id="file-input-123",
        completion_window="24h",
        status="completed",
        output_file_id="file-output-123",
        error_file_id=None,
        created_at=1234567890,
        in_progress_at=1234567900,
        expires_at=1234654290,
        finalizing_at=1234568000,
        completed_at=1234568100,
        failed_at=None,
        expired_at=None,
        cancelling_at=None,
        cancelled_at=None,
        request_counts={
            "total": 10,
            "completed": 10,
            "failed": 0,
        },
        metadata=None,
    )
    mock_batch._hidden_params = {}

    # Create logging object
    logging_obj = Logging(
        model="gpt-5-mini",
        messages=[{"role": "user", "content": "test"}],
        stream=False,
        call_type=CallTypes.aretrieve_batch.value,
        litellm_call_id="test-call-123",
        function_id="test-function",
        start_time=time.time(),
        dynamic_success_callbacks=[],
    )
    logging_obj.custom_llm_provider = "openai"

    # Mock _handle_completed_batch to return cost data
    expected_cost = 0.05
    expected_usage = litellm.Usage(
        prompt_tokens=100,
        completion_tokens=50,
        total_tokens=150,
    )
    expected_models = ["gpt-5-mini"]

    with patch(
        "litellm.litellm_core_utils.litellm_logging._handle_completed_batch",
        new=AsyncMock(return_value=(expected_cost, expected_usage, expected_models)),
    ) as mock_handle_batch:
        # Call async_success_handler
        await logging_obj.async_success_handler(
            result=mock_batch,
            start_time=time.time(),
            end_time=time.time() + 1,
        )

        # Verify _handle_completed_batch was called
        mock_handle_batch.assert_called_once()

        # Verify cost and usage were set on the batch result
        assert mock_batch._hidden_params["response_cost"] == expected_cost
        assert mock_batch._hidden_params["batch_models"] == expected_models
        assert mock_batch.usage == expected_usage


@pytest.mark.asyncio
async def test_handle_completed_batch_computes_real_cost_from_output_file(
    sample_file_content_dict,
):
    """Integration: a completed batch's cost and usage are computed from its output
    file via the real cost-calc chain (only the file download is stubbed). This is
    the function the retrieve handler invokes on completion; a dropped output line, a
    wrong token sum, or mispriced model fails this test.
    """
    from litellm.batches.batch_utils import _handle_completed_batch
    from litellm.types.utils import LiteLLMBatch

    batch = LiteLLMBatch(
        id="batch-real-cost-123",
        object="batch",
        endpoint="/v1/chat/completions",
        input_file_id="file-input-123",
        completion_window="24h",
        status="completed",
        output_file_id="file-output-123",
        created_at=1234567890,
    )

    with patch(
        "litellm.batches.batch_utils._get_batch_output_file_content_as_dictionary",
        new=AsyncMock(return_value=sample_file_content_dict),
    ):
        cost, usage, models = await _handle_completed_batch(
            batch=batch, custom_llm_provider="openai"
        )

    pricing = litellm.model_cost["gpt-4o-mini-2024-07-18"]
    expected_cost = (
        42 * pricing["input_cost_per_token_batches"]
        + 20 * pricing["output_cost_per_token_batches"]
    )

    assert cost == pytest.approx(expected_cost)
    assert cost > 0
    assert (
        cost
        < 42 * pricing["input_cost_per_token"] + 20 * pricing["output_cost_per_token"]
    )
    assert usage.prompt_tokens == 42
    assert usage.completion_tokens == 20
    assert usage.total_tokens == 62
    assert models == ["gpt-4o-mini-2024-07-18", "gpt-4o-mini-2024-07-18"]


@pytest.mark.asyncio
async def test_batch_retrieve_cost_tracking_with_explicit_cost_data():
    """
    Test that explicit cost data is used when provided, skipping computation.

    Regression test for: When batch_cost, batch_usage, and batch_models are explicitly
    provided in kwargs, they should be used directly without calling _handle_completed_batch.
    """
    from litellm.litellm_core_utils.litellm_logging import Logging
    from litellm.types.utils import CallTypes
    from litellm.types.utils import LiteLLMBatch
    from unittest.mock import AsyncMock, patch

    # Mock batch result with completed status
    mock_batch = LiteLLMBatch(
        id="batch-test-456",
        object="batch",
        endpoint="/v1/chat/completions",
        errors=None,
        input_file_id="file-input-456",
        completion_window="24h",
        status="completed",
        output_file_id="file-output-456",
        error_file_id=None,
        created_at=1234567890,
        in_progress_at=1234567900,
        expires_at=1234654290,
        finalizing_at=1234568000,
        completed_at=1234568100,
        failed_at=None,
        expired_at=None,
        cancelling_at=None,
        cancelled_at=None,
        request_counts={
            "total": 5,
            "completed": 5,
            "failed": 0,
        },
        metadata=None,
    )
    mock_batch._hidden_params = {}

    # Create logging object
    logging_obj = Logging(
        model="gpt-5-mini",
        messages=[{"role": "user", "content": "test"}],
        stream=False,
        call_type=CallTypes.aretrieve_batch.value,
        litellm_call_id="test-call-456",
        function_id="test-function",
        start_time=time.time(),
        dynamic_success_callbacks=[],
    )
    logging_obj.custom_llm_provider = "openai"

    # Explicit cost data to pass in kwargs
    explicit_cost = 0.10
    explicit_usage = litellm.Usage(
        prompt_tokens=200,
        completion_tokens=100,
        total_tokens=300,
    )
    explicit_models = ["gpt-5-mini", "gpt-5.5"]

    with patch(
        "litellm.litellm_core_utils.litellm_logging._handle_completed_batch",
        new=AsyncMock(),
    ) as mock_handle_batch:
        # Call async_success_handler with explicit cost data
        await logging_obj.async_success_handler(
            result=mock_batch,
            start_time=time.time(),
            end_time=time.time() + 1,
            batch_cost=explicit_cost,
            batch_usage=explicit_usage,
            batch_models=explicit_models,
        )

        # Verify _handle_completed_batch was NOT called (since explicit data provided)
        mock_handle_batch.assert_not_called()

        # Verify explicit cost data was used
        assert mock_batch._hidden_params["response_cost"] == explicit_cost
        assert mock_batch._hidden_params["batch_models"] == explicit_models
        assert mock_batch.usage == explicit_usage


@pytest.mark.asyncio
async def test_batch_retrieve_cost_tracking_with_unified_file_id_incomplete_batch():
    """
    Test that cost computation is skipped for unified file IDs with non-completed batches.

    Regression test for: For unified file IDs (base64 encoded), cost should only be computed
    when batch status is "completed" and explicit data is not provided.
    """
    import base64
    from litellm.litellm_core_utils.litellm_logging import Logging
    from litellm.types.utils import CallTypes, SpecialEnums
    from litellm.types.utils import LiteLLMBatch
    from unittest.mock import AsyncMock, patch

    # Create a proper unified file ID by encoding the correct prefix
    unified_id_str = f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}:test_file_789;unified_id:batch-789"
    encoded_unified_id = (
        base64.urlsafe_b64encode(unified_id_str.encode()).decode().rstrip("=")
    )

    # Mock batch result with in_progress status and unified file ID
    mock_batch = LiteLLMBatch(
        id=encoded_unified_id,  # Properly encoded unified ID
        object="batch",
        endpoint="/v1/chat/completions",
        errors=None,
        input_file_id="file-input-789",
        completion_window="24h",
        status="in_progress",  # Not completed
        output_file_id=None,
        error_file_id=None,
        created_at=1234567890,
        in_progress_at=1234567900,
        expires_at=1234654290,
        finalizing_at=None,
        completed_at=None,
        failed_at=None,
        expired_at=None,
        cancelling_at=None,
        cancelled_at=None,
        request_counts={
            "total": 10,
            "completed": 3,
            "failed": 0,
        },
        metadata=None,
    )
    mock_batch._hidden_params = {}

    # Create logging object
    logging_obj = Logging(
        model="gpt-5-mini",
        messages=[{"role": "user", "content": "test"}],
        stream=False,
        call_type=CallTypes.aretrieve_batch.value,
        litellm_call_id="test-call-789",
        function_id="test-function",
        start_time=time.time(),
        dynamic_success_callbacks=[],
    )
    logging_obj.custom_llm_provider = "openai"

    with patch(
        "litellm.litellm_core_utils.litellm_logging._handle_completed_batch",
        new=AsyncMock(),
    ) as mock_handle_batch:
        # Call async_success_handler with in_progress batch (unified file ID)
        await logging_obj.async_success_handler(
            result=mock_batch,
            start_time=time.time(),
            end_time=time.time() + 1,
        )

        # Verify _handle_completed_batch was NOT called (batch not completed and is unified file ID)
        mock_handle_batch.assert_not_called()

        # Verify cost data was not set
        assert "response_cost" not in mock_batch._hidden_params
        assert "batch_models" not in mock_batch._hidden_params
        assert not hasattr(mock_batch, "usage") or mock_batch.usage is None


@pytest.mark.asyncio
async def test_batch_retrieve_cost_tracking_with_partial_explicit_data():
    """
    Test that cost is computed when only partial explicit data is provided.

    Regression test for: If batch_cost, batch_usage, or batch_models is missing
    (not all three provided), and batch is completed, system should compute the data.
    """
    from litellm.litellm_core_utils.litellm_logging import Logging
    from litellm.types.utils import CallTypes
    from litellm.types.utils import LiteLLMBatch
    from unittest.mock import AsyncMock, patch

    # Mock batch result with completed status
    mock_batch = LiteLLMBatch(
        id="batch-test-partial",
        object="batch",
        endpoint="/v1/chat/completions",
        errors=None,
        input_file_id="file-input-partial",
        completion_window="24h",
        status="completed",
        output_file_id="file-output-partial",
        error_file_id=None,
        created_at=1234567890,
        in_progress_at=1234567900,
        expires_at=1234654290,
        finalizing_at=1234568000,
        completed_at=1234568100,
        failed_at=None,
        expired_at=None,
        cancelling_at=None,
        cancelled_at=None,
        request_counts={
            "total": 8,
            "completed": 8,
            "failed": 0,
        },
        metadata=None,
    )
    mock_batch._hidden_params = {}

    # Create logging object
    logging_obj = Logging(
        model="gpt-5-mini",
        messages=[{"role": "user", "content": "test"}],
        stream=False,
        call_type=CallTypes.aretrieve_batch.value,
        litellm_call_id="test-call-partial",
        function_id="test-function",
        start_time=time.time(),
        dynamic_success_callbacks=[],
    )

    logging_obj.custom_llm_provider = "openai"

    # Only provide batch_cost, missing batch_usage and batch_models
    partial_cost = 0.08

    expected_cost = 0.06
    expected_usage = litellm.Usage(
        prompt_tokens=150,
        completion_tokens=75,
        total_tokens=225,
    )
    expected_models = ["gpt-5-mini"]

    with patch(
        "litellm.litellm_core_utils.litellm_logging._handle_completed_batch",
        new=AsyncMock(return_value=(expected_cost, expected_usage, expected_models)),
    ) as mock_handle_batch:
        # Call async_success_handler with partial explicit data
        await logging_obj.async_success_handler(
            result=mock_batch,
            start_time=time.time(),
            end_time=time.time() + 1,
            batch_cost=partial_cost,  # Only cost provided, not usage or models
        )

        # Verify _handle_completed_batch WAS called (since not all data provided)
        mock_handle_batch.assert_called_once()

        # Verify computed cost data was used (not partial explicit data)
        assert mock_batch._hidden_params["response_cost"] == expected_cost
        assert mock_batch._hidden_params["batch_models"] == expected_models
        assert mock_batch.usage == expected_usage