scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

Merge branch 'BerriAI:main' into main

Browse source
This commit is contained in:
Tasmay Pankaj Tibrewal 2025-08-15 04:36:07 +05:30 • committed by GitHub
commit 61ef0d42f9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 260 additions and 165 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.github/workflows/test-litellm.yml vendored
View file

@ -7,7 +7,7 @@ on:
jobs:
test:
runs-on: ubuntu-latest
timeout-minutes: 20
timeout-minutes: 25
steps:
- uses: actions/checkout@v4

6
litellm/integrations/mlflow.py
View file

@ -189,9 +189,9 @@ class MlflowLogger(CustomLogger):
{
"api_base": standard_obj.get("api_base"),
"cache_hit": standard_obj.get("cache_hit"),
"usage": {
"completion_tokens": standard_obj.get("completion_tokens"),
"prompt_tokens": standard_obj.get("prompt_tokens"),
"mlflow.chat.tokenUsage": {
"input_tokens": standard_obj.get("prompt_tokens"),
"output_tokens": standard_obj.get("completion_tokens"),
"total_tokens": standard_obj.get("total_tokens"),
},
"raw_llm_response": standard_obj.get("response"),

4
pyproject.toml
View file

@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm"
version = "1.75.5"
version = "1.75.6"
description = "Library to easily interface with LLM API providers"
authors = ["BerriAI"]
license = "MIT"
@ -155,7 +155,7 @@ requires = ["poetry-core", "wheel"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "1.75.5"
version = "1.75.6"
version_files = [
"pyproject.toml:^version"
]

41
tests/test_litellm/integrations/test_mlflow.py
View file

@ -71,5 +71,42 @@ async def test_mlflow_request_tags_functionality():
tags_param = call_args.kwargs.get('tags', {})
expected_tags = {"tag1": "", "tag2": "", "production": ""}
assert tags_param == expected_tags, f"Expected tags {expected_tags}, got {tags_param}"
print("✅ Request tags properly transformed and passed to MLflow trace")
def test_mlflow_token_usage_attribute_structure():
"""Ensure token usage attributes are formatted with mlflow.chat.tokenUsage."""
mock_mlflow_tracking = MagicMock()
mock_mlflow_tracking.MlflowClient = MagicMock()
with patch.dict(
"sys.modules",
{
"mlflow": MagicMock(),
"mlflow.tracking": mock_mlflow_tracking,
"mlflow.tracing.utils": MagicMock(),
},
):
from litellm.integrations.mlflow import MlflowLogger
mlflow_logger = MlflowLogger()
attrs = mlflow_logger._extract_attributes( # type: ignore
{
"litellm_call_id": "123",
"call_type": "completion",
"model": "gpt-3.5-turbo",
"standard_logging_object": {
"prompt_tokens": 5,
"completion_tokens": 7,
"total_tokens": 12,
},
}
)
assert attrs["mlflow.chat.tokenUsage"] == {
"input_tokens": 5,
"output_tokens": 7,
"total_tokens": 12,
}

228
tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py
View file

@ -14,6 +14,7 @@ from unittest.mock import patch
import litellm
from litellm.proxy.proxy_server import app
from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles
from litellm.types.tag_management import TagDeleteRequest, TagInfoRequest, TagNewRequest
client = TestClient(app)
@ -24,58 +25,71 @@ async def test_create_and_get_tag():
"""
Test creation of a new tag and retrieving its information
"""
# Mock the prisma client and _get_tags_config and _save_tags_config
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
"litellm.proxy.proxy_server.llm_router"
) as mock_router, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config"
) as mock_get_tags, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config"
) as mock_save_tags, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._add_tag_to_deployment"
) as mock_add_tag, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names"
) as mock_get_models:
# Setup mocks
mock_get_tags.return_value = {}
mock_get_models.return_value = {"model-1": "gpt-3.5-turbo"}
# Mock the user authentication
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
mock_user_auth = UserAPIKeyAuth(
user_id="test-user-123",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
try:
# Mock the prisma client and _get_tags_config and _save_tags_config
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
"litellm.proxy.proxy_server.llm_router"
) as mock_router, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config"
) as mock_get_tags, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config"
) as mock_save_tags, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._add_tag_to_deployment"
) as mock_add_tag, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names"
) as mock_get_models:
# Setup mocks
mock_get_tags.return_value = {}
mock_get_models.return_value = {"model-1": "gpt-3.5-turbo"}
# Create a new tag
tag_data = {
"name": "test-tag",
"description": "Test tag for unit testing",
"models": ["model-1"],
}
# Set admin access for the test
headers = {"Authorization": f"Bearer sk-1234"}
# Test tag creation
response = client.post("/tag/new", json=tag_data, headers=headers)
print(f"response: {response.text}")
assert response.status_code == 200
result = response.json()
assert result["message"] == "Tag test-tag created successfully"
assert result["tag"]["name"] == "test-tag"
assert result["tag"]["description"] == "Test tag for unit testing"
# Mock updated tag config for the get request
mock_get_tags.return_value = {
"test-tag": {
# Create a new tag
tag_data = {
"name": "test-tag",
"description": "Test tag for unit testing",
"models": ["model-1"],
"model_info": {"model-1": "gpt-3.5-turbo"},
}
}
# Test retrieving tag info
info_data = {"names": ["test-tag"]}
response = client.post("/tag/info", json=info_data, headers=headers)
assert response.status_code == 200
result = response.json()
assert "test-tag" in result
assert result["test-tag"]["description"] == "Test tag for unit testing"
# Set admin access for the test
headers = {"Authorization": f"Bearer sk-1234"}
# Test tag creation
response = client.post("/tag/new", json=tag_data, headers=headers)
print(f"response: {response.text}")
assert response.status_code == 200
result = response.json()
assert result["message"] == "Tag test-tag created successfully"
assert result["tag"]["name"] == "test-tag"
assert result["tag"]["description"] == "Test tag for unit testing"
# Mock updated tag config for the get request
mock_get_tags.return_value = {
"test-tag": {
"name": "test-tag",
"description": "Test tag for unit testing",
"models": ["model-1"],
"model_info": {"model-1": "gpt-3.5-turbo"},
}
}
# Test retrieving tag info
info_data = {"names": ["test-tag"]}
response = client.post("/tag/info", json=info_data, headers=headers)
assert response.status_code == 200
result = response.json()
assert "test-tag" in result
assert result["test-tag"]["description"] == "Test tag for unit testing"
finally:
# Clean up dependency overrides
app.dependency_overrides.clear()
@pytest.mark.asyncio
@ -83,16 +97,26 @@ async def test_update_tag():
"""
Test updating an existing tag
"""
# Mock the prisma client and _get_tags_config and _save_tags_config
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config"
) as mock_get_tags, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config"
) as mock_save_tags, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names"
) as mock_get_models:
# Setup mocks for existing tag
mock_get_tags.return_value = {
# Mock the user authentication
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
mock_user_auth = UserAPIKeyAuth(
user_id="test-user-123",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
try:
# Mock the prisma client and _get_tags_config and _save_tags_config
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config"
) as mock_get_tags, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config"
) as mock_save_tags, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names"
) as mock_get_models:
# Setup mocks for existing tag
mock_get_tags.return_value = {
"test-tag": {
"name": "test-tag",
"description": "Original description",
@ -101,27 +125,30 @@ async def test_update_tag():
"updated_at": "2023-01-01T00:00:00",
"created_by": "user-123",
}
}
mock_get_models.return_value = {"model-1": "gpt-3.5-turbo", "model-2": "gpt-4"}
}
mock_get_models.return_value = {"model-1": "gpt-3.5-turbo", "model-2": "gpt-4"}
# Update tag data
update_data = {
"name": "test-tag",
"description": "Updated description",
"models": ["model-1", "model-2"],
}
# Update tag data
update_data = {
"name": "test-tag",
"description": "Updated description",
"models": ["model-1", "model-2"],
}
# Set admin access for the test
headers = {"Authorization": f"Bearer sk-1234"}
# Set admin access for the test
headers = {"Authorization": f"Bearer sk-1234"}
# Test tag update
response = client.post("/tag/update", json=update_data, headers=headers)
assert response.status_code == 200
result = response.json()
assert result["message"] == "Tag test-tag updated successfully"
assert result["tag"]["description"] == "Updated description"
assert len(result["tag"]["models"]) == 2
assert "model-2" in result["tag"]["models"]
# Test tag update
response = client.post("/tag/update", json=update_data, headers=headers)
assert response.status_code == 200
result = response.json()
assert result["message"] == "Tag test-tag updated successfully"
assert result["tag"]["description"] == "Updated description"
assert len(result["tag"]["models"]) == 2
assert "model-2" in result["tag"]["models"]
finally:
# Clean up dependency overrides
app.dependency_overrides.clear()
@pytest.mark.asyncio
@ -129,14 +156,24 @@ async def test_delete_tag():
"""
Test deleting a tag
"""
# Mock the prisma client and _get_tags_config and _save_tags_config
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config"
) as mock_get_tags, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config"
) as mock_save_tags:
# Setup mocks for existing tag
mock_get_tags.return_value = {
# Mock the user authentication
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
mock_user_auth = UserAPIKeyAuth(
user_id="test-user-123",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
try:
# Mock the prisma client and _get_tags_config and _save_tags_config
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config"
) as mock_get_tags, patch(
"litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config"
) as mock_save_tags:
# Setup mocks for existing tag
mock_get_tags.return_value = {
"test-tag": {
"name": "test-tag",
"description": "Test tag for deletion",
@ -145,22 +182,25 @@ async def test_delete_tag():
"updated_at": "2023-01-01T00:00:00",
"created_by": "user-123",
}
}
}
# Delete tag data
delete_data = {"name": "test-tag"}
# Delete tag data
delete_data = {"name": "test-tag"}
# Set admin access for the test
headers = {"Authorization": f"Bearer sk-1234"}
# Set admin access for the test
headers = {"Authorization": f"Bearer sk-1234"}
# Test tag deletion
response = client.post("/tag/delete", json=delete_data, headers=headers)
assert response.status_code == 200
result = response.json()
assert result["message"] == "Tag test-tag deleted successfully"
# Test tag deletion
response = client.post("/tag/delete", json=delete_data, headers=headers)
assert response.status_code == 200
result = response.json()
assert result["message"] == "Tag test-tag deleted successfully"
# Verify _save_tags_config was called without the deleted tag
mock_save_tags.assert_called_once()
# Verify _save_tags_config was called without the deleted tag
mock_save_tags.assert_called_once()
finally:
# Clean up dependency overrides
app.dependency_overrides.clear()
@pytest.mark.asyncio

144
tests/test_litellm/test_cost_calculation_log_level.py
View file

@ -17,46 +17,55 @@ def test_cost_calculation_uses_debug_level(caplog):
This ensures cost calculation details don't appear in production logs.
Part of fix for issue #9815.
"""
# Create a mock completion response
mock_response = {
"id": "test",
"object": "chat.completion",
"created": 1234567890,
"model": "gpt-3.5-turbo",
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": "Test response"},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 20,
"total_tokens": 30
# Ensure verbose_logger is set to DEBUG level to capture the debug logs
from litellm._logging import verbose_logger
original_level = verbose_logger.level
verbose_logger.setLevel(logging.DEBUG)
try:
# Create a mock completion response
mock_response = {
"id": "test",
"object": "chat.completion",
"created": 1234567890,
"model": "gpt-3.5-turbo",
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": "Test response"},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 20,
"total_tokens": 30
}
}
}
# Test that cost calculation logs are at DEBUG level
with caplog.at_level(logging.DEBUG):
try:
cost = completion_cost(
completion_response=mock_response,
model="gpt-3.5-turbo"
)
except Exception:
pass # Cost calculation may fail, but we're checking log levels
# Find the cost calculation log records
cost_calc_records = [
record for record in caplog.records
if "selected model name for cost calculation" in record.message
]
# Verify that cost calculation logs are at DEBUG level
assert len(cost_calc_records) > 0, "No cost calculation logs found"
for record in cost_calc_records:
assert record.levelno == logging.DEBUG, \
f"Cost calculation log should be DEBUG level, but was {record.levelname}"
# Test that cost calculation logs are at DEBUG level
with caplog.at_level(logging.DEBUG, logger="LiteLLM"):
try:
cost = completion_cost(
completion_response=mock_response,
model="gpt-3.5-turbo"
)
except Exception:
pass # Cost calculation may fail, but we're checking log levels
# Find the cost calculation log records
cost_calc_records = [
record for record in caplog.records
if "selected model name for cost calculation" in record.message
]
# Verify that cost calculation logs are at DEBUG level
assert len(cost_calc_records) > 0, "No cost calculation logs found"
for record in cost_calc_records:
assert record.levelno == logging.DEBUG, \
f"Cost calculation log should be DEBUG level, but was {record.levelname}"
finally:
# Restore original logger level
verbose_logger.setLevel(original_level)
def test_batch_cost_calculation_uses_debug_level(caplog):
@ -65,29 +74,38 @@ def test_batch_cost_calculation_uses_debug_level(caplog):
"""
from litellm.cost_calculator import batch_cost_calculator
from litellm.types.utils import Usage
from litellm._logging import verbose_logger
# Create a mock usage object
usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300)
# Ensure verbose_logger is set to DEBUG level to capture the debug logs
original_level = verbose_logger.level
verbose_logger.setLevel(logging.DEBUG)
# Test that batch cost calculation logs are at DEBUG level
with caplog.at_level(logging.DEBUG):
try:
batch_cost_calculator(
usage=usage,
model="gpt-3.5-turbo",
custom_llm_provider="openai"
)
except Exception:
pass # May fail, but we're checking log levels
# Find batch cost calculation log records
batch_cost_records = [
record for record in caplog.records
if "Calculating batch cost per token" in record.message
]
# Verify logs exist and are at DEBUG level
if batch_cost_records: # May not always log depending on the code path
for record in batch_cost_records:
assert record.levelno == logging.DEBUG, \
f"Batch cost calculation log should be DEBUG level, but was {record.levelname}"
try:
# Create a mock usage object
usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300)
# Test that batch cost calculation logs are at DEBUG level
with caplog.at_level(logging.DEBUG, logger="LiteLLM"):
try:
batch_cost_calculator(
usage=usage,
model="gpt-3.5-turbo",
custom_llm_provider="openai"
)
except Exception:
pass # May fail, but we're checking log levels
# Find batch cost calculation log records
batch_cost_records = [
record for record in caplog.records
if "Calculating batch cost per token" in record.message
]
# Verify logs exist and are at DEBUG level
if batch_cost_records: # May not always log depending on the code path
for record in batch_cost_records:
assert record.levelno == logging.DEBUG, \
f"Batch cost calculation log should be DEBUG level, but was {record.levelname}"
finally:
# Restore original logger level
verbose_logger.setLevel(original_level)