mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
fix(proxy): keep tool payloads and logprobs unmasked in stored spend logs (#44075)
* fix(proxy): keep tool payloads and logprobs unmasked in stored spend logs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): audit cells for stored spend-log tool payloads and logprobs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): assert cache-hit spend-log rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): handle model-list requests in spend-log audit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): assert unauthenticated requests never reach upstream Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng <yucheng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
d131c43782
commit
e32b25817f
4 changed files with 2030 additions and 5 deletions
|
|
@ -1095,16 +1095,32 @@ def _get_messages_for_spend_logs_payload(
|
|||
|
||||
_SENSITIVE_REQUEST_BODY_KEYS: Final = frozenset({"secret_fields"})
|
||||
_REQUEST_BODY_CREDENTIAL_MASKER: Final = SensitiveDataMasker(extra_sensitive_patterns=frozenset({"apikey"}))
|
||||
_TOOL_INPUT_BLOCK_TYPES: Final = frozenset({"tool_use", "server_tool_use", "mcp_tool_use"})
|
||||
_TOOL_OUTPUT_BLOCK_TYPES: Final = frozenset({"tool_result", "mcp_tool_result", "function_call_output"})
|
||||
|
||||
|
||||
def _is_request_body_credential(key: str, value: object) -> bool:
|
||||
return isinstance(value, str) and _REQUEST_BODY_CREDENTIAL_MASKER.is_sensitive_key(key)
|
||||
|
||||
|
||||
def _is_spend_log_content(parent: Mapping[str, object], key: str) -> bool:
|
||||
block_type: Final = parent.get("type")
|
||||
block_type_name: Final = block_type if isinstance(block_type, str) else None
|
||||
return (
|
||||
key in ("arguments", "logprobs")
|
||||
or (key == "input" and block_type_name in _TOOL_INPUT_BLOCK_TYPES)
|
||||
or (
|
||||
key in ("content", "output")
|
||||
and (block_type_name in _TOOL_OUTPUT_BLOCK_TYPES or parent.get("role") == "tool")
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _sanitize_request_body_for_spend_logs_payload(
|
||||
request_body: Mapping[str, object],
|
||||
visited: set | None = None,
|
||||
max_string_length_prompt_in_db: int | None = None,
|
||||
mask_credentials: bool = True,
|
||||
) -> dict:
|
||||
"""
|
||||
Recursively sanitize request body to prevent logging large base64 strings or other large values.
|
||||
|
|
@ -1112,7 +1128,8 @@ def _sanitize_request_body_for_spend_logs_payload(
|
|||
|
||||
At every nesting level, also strips keys listed in _SENSITIVE_REQUEST_BODY_KEYS (e.g. secret_fields,
|
||||
which holds raw HTTP headers including Authorization tokens), and replaces string values under keys
|
||||
SensitiveDataMasker classifies as credentials with REDACTED_BY_LITELM_STRING.
|
||||
SensitiveDataMasker classifies as credentials with REDACTED_BY_LITELM_STRING, except inside tool payloads
|
||||
and logprobs.
|
||||
"""
|
||||
from litellm.constants import (
|
||||
LITELLM_TRUNCATED_PAYLOAD_FIELD,
|
||||
|
|
@ -1130,11 +1147,13 @@ def _sanitize_request_body_for_spend_logs_payload(
|
|||
return {}
|
||||
visited.add(obj_id)
|
||||
|
||||
def _sanitize_value(value: object) -> object:
|
||||
def _sanitize_value(value: object, mask_credentials: bool) -> object:
|
||||
if isinstance(value, Mapping):
|
||||
return _sanitize_request_body_for_spend_logs_payload(value, visited, max_string_length_prompt_in_db)
|
||||
return _sanitize_request_body_for_spend_logs_payload(
|
||||
value, visited, max_string_length_prompt_in_db, mask_credentials
|
||||
)
|
||||
elif isinstance(value, list):
|
||||
return [_sanitize_value(item) for item in value]
|
||||
return [_sanitize_value(item, mask_credentials) for item in value]
|
||||
elif isinstance(value, str):
|
||||
if len(value) > max_string_length_prompt_in_db:
|
||||
# Keep 35% from beginning and 65% from end (end is usually more important)
|
||||
|
|
@ -1170,7 +1189,9 @@ def _sanitize_request_body_for_spend_logs_payload(
|
|||
return value
|
||||
|
||||
return {
|
||||
k: REDACTED_BY_LITELM_STRING if _is_request_body_credential(k, v) else _sanitize_value(v)
|
||||
k: REDACTED_BY_LITELM_STRING
|
||||
if mask_credentials and _is_request_body_credential(k, v)
|
||||
else _sanitize_value(v, mask_credentials and not _is_spend_log_content(request_body, k))
|
||||
for k, v in request_body.items()
|
||||
if k not in _SENSITIVE_REQUEST_BODY_KEYS
|
||||
}
|
||||
|
|
|
|||
1774
tests/integration/spend/test_spend_log_tool_payload_content.py
Normal file
1774
tests/integration/spend/test_spend_log_tool_payload_content.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -2789,6 +2789,167 @@ def test_sanitize_response_redacts_credential_named_fields() -> None:
|
|||
}
|
||||
|
||||
|
||||
def test_sanitize_response_keeps_logprob_tokens() -> None:
|
||||
response: Final = {
|
||||
"system_fingerprint": "fp_x",
|
||||
"choices": [
|
||||
{
|
||||
"logprobs": {
|
||||
"content": [
|
||||
{
|
||||
"token": "sort",
|
||||
"logprob": -0.1,
|
||||
"bytes": [115],
|
||||
"top_logprobs": [{"token": "sort", "logprob": -0.1}],
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
assert _sanitize_request_body_for_spend_logs_payload({"response": response}) == {
|
||||
"response": {
|
||||
"system_fingerprint": REDACTED_BY_LITELM_STRING,
|
||||
"choices": [
|
||||
{
|
||||
"logprobs": {
|
||||
"content": [
|
||||
{
|
||||
"token": "sort",
|
||||
"logprob": -0.1,
|
||||
"bytes": [115],
|
||||
"top_logprobs": [{"token": "sort", "logprob": -0.1}],
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
],
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_sanitize_request_body_keeps_key_named_tool_payload_fields() -> None:
|
||||
request_body: Final = {
|
||||
"model": "anthropic/claude",
|
||||
"aws_secret_access_key": "AKIAEXAMPLESECRET",
|
||||
"prompt_cache_key": "tenant-42-cache",
|
||||
"metadata": {"user_api_key_alias": "tenant-user"},
|
||||
"secret_fields": {"raw_headers": {"authorization": "Bearer secret"}},
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "tool_use", "input": {"key": "order-123", "sort_key": "created_at"}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_order",
|
||||
"arguments": {"key": "order-123", "sort_key": "created_at"},
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "content": {"token_type": "bearer", "partition_key": "tenant_42"}},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"content": [{"token_type": "bearer", "partition_key": "tenant_42"}],
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
"input": [
|
||||
{"type": "function_call", "arguments": {"key": "tenant-42", "access_level": "admin"}},
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"output": {"token_type": "bearer", "partition_key": "tenant_42"},
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
assert _sanitize_request_body_for_spend_logs_payload(request_body) == {
|
||||
"model": "anthropic/claude",
|
||||
"aws_secret_access_key": REDACTED_BY_LITELM_STRING,
|
||||
"prompt_cache_key": REDACTED_BY_LITELM_STRING,
|
||||
"metadata": {"user_api_key_alias": REDACTED_BY_LITELM_STRING},
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "tool_use", "input": {"key": "order-123", "sort_key": "created_at"}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_order",
|
||||
"arguments": {"key": "order-123", "sort_key": "created_at"},
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "content": {"token_type": "bearer", "partition_key": "tenant_42"}},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"content": [{"token_type": "bearer", "partition_key": "tenant_42"}],
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
"input": [
|
||||
{"type": "function_call", "arguments": {"key": "tenant-42", "access_level": "admin"}},
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"output": {"token_type": "bearer", "partition_key": "tenant_42"},
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def test_sanitize_request_body_masks_credentials_beside_tool_blocks() -> None:
|
||||
request_body: Final = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "tool_use", "api_key": "sk-live", "input": {"key": "order-123"}},
|
||||
{"type": {"nested": 1}, "input": {"api_key": "x"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
assert _sanitize_request_body_for_spend_logs_payload(request_body) == {
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"api_key": REDACTED_BY_LITELM_STRING,
|
||||
"input": {"key": "order-123"},
|
||||
},
|
||||
{"type": {"nested": 1}, "input": {"api_key": REDACTED_BY_LITELM_STRING}},
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs")
|
||||
def test_proxy_server_request_payload_excludes_secret_fields(mock_should_store):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import json
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -6,6 +7,9 @@ from fastapi import HTTPException
|
|||
from fastapi.testclient import TestClient
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import (
|
||||
_get_proxy_server_request_for_spend_logs_payload,
|
||||
)
|
||||
from litellm.responses.litellm_completion_transformation import session_handler
|
||||
from litellm.responses.litellm_completion_transformation.session_handler import (
|
||||
ResponsesSessionHandler,
|
||||
|
|
@ -718,3 +722,68 @@ async def test_message_history_normalizes_redacted_tool_call_arguments():
|
|||
tool_call = assistant_message.tool_calls[0]
|
||||
assert tool_call.function.arguments == "{}"
|
||||
assert json.loads(tool_call.function.arguments) == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_message_history_replays_real_key_named_tool_payloads() -> None:
|
||||
request_id: Final = "chatcmpl-tool-payload"
|
||||
function_arguments: Final = {"sort_key": "created_at", "access_level": "admin"}
|
||||
function_output: Final = {
|
||||
"status": "active",
|
||||
"token_type": "bearer",
|
||||
"partition_key": "tenant_42",
|
||||
}
|
||||
responses_request_body: Final = {
|
||||
"model": "anthropic/claude-sonnet-4-5",
|
||||
"input": [
|
||||
{"role": "user", "content": "Fetch my account settings."},
|
||||
{
|
||||
"type": "function_call",
|
||||
"call_id": "call_1",
|
||||
"name": "get_settings",
|
||||
"arguments": function_arguments,
|
||||
},
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_1",
|
||||
"output": function_output,
|
||||
},
|
||||
{"role": "user", "content": "Acknowledge with OK"},
|
||||
],
|
||||
"aws_secret_access_key": "AKIAEXAMPLESECRET",
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs",
|
||||
return_value=True,
|
||||
):
|
||||
proxy_server_request: Final = json.loads(
|
||||
_get_proxy_server_request_for_spend_logs_payload(
|
||||
metadata={},
|
||||
litellm_params={"proxy_server_request": {"body": responses_request_body}},
|
||||
kwargs={},
|
||||
)
|
||||
)
|
||||
|
||||
spend_log: Final = {
|
||||
"request_id": request_id,
|
||||
"call_type": "aresponses",
|
||||
"session_id": "session-tool-payload",
|
||||
"proxy_server_request": proxy_server_request,
|
||||
"response": _chat_completion_response(request_id, "OK"),
|
||||
}
|
||||
|
||||
with patch.object(
|
||||
ResponsesSessionHandler,
|
||||
"get_all_spend_logs_for_previous_response_id",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_get_spend_logs:
|
||||
mock_get_spend_logs.return_value = [spend_log]
|
||||
result: Final = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
|
||||
request_id
|
||||
)
|
||||
|
||||
assistant_message: Final = result["messages"][1]
|
||||
tool_message: Final = result["messages"][2]
|
||||
assert json.loads(assistant_message["tool_calls"][0]["function"]["arguments"]) == function_arguments
|
||||
assert json.loads(tool_message["content"]) == function_output
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue