scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 133
litellm/tests/batches_tests/test_batches_logging_unit_tests.py
yuneng-jiang 6a0d03914c
test: drop the cwd-relative sys.path.insert calls from the test suite (#37802) 
* test: drop the cwd-relative sys.path.insert calls from the test suite

TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.

Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.

Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.

* test: drop the duplicate imports the sys.path sweep exposed to F811

* test(pre-call-utils): restore the os import the new bedrock tests need
2026-08-22 09:25:58 -07:00

558 lines
21 KiB
Python
Raw Permalink Blame History

import asyncio
import json
import traceback
from unittest.mock import AsyncMock, MagicMock, patch
from dotenv import load_dotenv
load_dotenv()
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 (
_aggregate_batch_cost_usage_models,
_get_file_content_as_dictionary,
_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):
with patch("litellm.completion_cost", return_value=0.0):
_, usage, _ = _aggregate_batch_cost_usage_models(
entries=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, _, _ = _aggregate_batch_cost_usage_models(
entries=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,
)
sample_file_content_bytes = "\n".join(
json.dumps(row) for row in sample_file_content_dict
).encode()
with patch(
"litellm.batches.batch_utils._fetch_batch_output_file_content",
new=AsyncMock(return_value=sample_file_content_bytes),
):
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
Reference in a new issue View git blame Copy permalink