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

fix(batches): mask api base credentials on batch cost rows

Browse source
This commit is contained in:
mateo-berri 2026-09-05 01:35:09 -07:00
parent ed2408f28a
commit 814c151b02
3 changed files with 94 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -685,7 +685,7 @@ class CheckBatchCost:
from litellm.files.main import afile_content
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
)
@ -858,7 +858,7 @@ class CheckBatchCost:
"user-agent": CHECK_BATCH_COST_USER_AGENT,
}
},
**({"api_base": deployment_api_base} if deployment_api_base else {}),
**({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}),
"metadata": {
**(await self._build_creator_attribution_metadata(job, batch_id)),
# spend logs read the deployment identity off these metadata keys, so

16
litellm/litellm_core_utils/litellm_logging.py
View file

@ -419,6 +419,13 @@ def _provider_response_id(source: object) -> str | None:
return candidate if isinstance(candidate, str) and candidate else None
def mask_api_base_credentials(api_base: str) -> str:
if "key=" not in api_base:
return api_base
key_end: Final = api_base.find("key=") + 4
return api_base[:key_end] + "*" * 5 + api_base[-4:]
class Logging(LiteLLMLoggingBaseClass):
global \
supabaseClient, \
@ -1160,14 +1167,7 @@ class Logging(LiteLLMLoggingBaseClass):
return data
def _get_masked_api_base(self, api_base: str) -> str:
if "key=" in api_base:
# Find the position of "key=" in the string
key_index: Final = api_base.find("key=") + 4
# Mask the last 5 characters after "key="
masked_api_base = api_base[:key_index] + "*" * 5 + api_base[-4:]
else:
masked_api_base = api_base
return str(masked_api_base)
return str(mask_api_base_credentials(api_base))
def _pre_call(self, input, api_key, model=None, additional_args={}):
"""

84
tests/proxy_unit_tests/test_check_batch_cost.py
View file

@ -583,6 +583,90 @@ class TestCheckBatchCost:
assert passed_model_info["input_cost_per_token_batches"] == 2e-06
assert passed_model_info["output_cost_per_token_batches"] == 4e-06
@pytest.mark.asyncio
async def test_poller_masks_api_base_credentials_before_logging(
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
):
"""Request rows mask `key=` query credentials out of api_base before it is
logged, but the poller skips that pre-call step, so an unmasked deployment
api_base would land verbatim on the batch cost row: regression test for the
poller masking the same way.
"""
import base64
from unittest.mock import patch
import httpx
import respx
from litellm.litellm_core_utils.litellm_logging import Logging
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_job = MagicMock()
mock_job.id = "job-masked-api-base-1"
mock_job.unified_object_id = base64.urlsafe_b64encode(
b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456"
).decode()
mock_job.created_by = "user-1"
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job])
mock_response = MagicMock()
mock_response.status = "completed"
mock_response.output_file_id = "file-output-123"
mock_response.error_file_id = None
mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}'
mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"})
mock_deployment = MagicMock()
mock_deployment.litellm_params.custom_llm_provider = "openai"
mock_deployment.litellm_params.model = "gpt-5.4-mini"
mock_deployment.litellm_params.api_base = "https://gateway.example.com/v1?key=AIzaSyVERYSECRET7890"
mock_deployment.model_info.model_dump.return_value = {}
mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment)
output_line = json.dumps(
{
"custom_id": "req-1",
"response": {
"status_code": 200,
"body": {
"id": "chatcmpl-1",
"object": "chat.completion",
"model": "gpt-5.4-mini",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
},
},
"error": None,
}
)
with (
respx.mock(assert_all_called=True) as provider,
patch.object( # test-quality-ok: the poller builds Logging inline, the only seam to the row it logs
Logging, "async_success_handler", autospec=True
) as success_handler,
):
provider.get("https://api.openai.com/v1/files/file-output-123/content").mock(
return_value=httpx.Response(200, content=f"{output_line}\n".encode())
)
await check_batch_cost_instance.check_batch_cost()
cost_row_calls = [call for call in success_handler.await_args_list if "batch_cost" in call.kwargs]
assert len(cost_row_calls) == 1
logged_api_base = cost_row_calls[0].args[0].litellm_params["api_base"]
assert logged_api_base == "https://gateway.example.com/v1?key=*****7890"
assert "VERYSECRET" not in logged_api_base
@pytest.mark.asyncio
async def test_primary_path_completion_update_includes_batch_processed(
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router