diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 6b3ccfd3694..64a86387c39 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -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 diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 73c06b891fb..7e800139228 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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={}): """ diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index f9fb782b052..36177b44930 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -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