mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(proxy): merge org-level spend_logs_metadata into request metadata
This commit is contained in:
parent
4d33964898
commit
eb8d69291a
2 changed files with 87 additions and 22 deletions
|
|
@ -1028,6 +1028,19 @@ class LiteLLMProxyRequestSetup:
|
|||
data[_metadata_variable_name]["user_api_key_auth_metadata"].update(added_metadata)
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def _merge_spend_logs_metadata(source_metadata: dict, data: dict, _metadata_variable_name: str) -> None:
|
||||
spend_logs_metadata = source_metadata.get("spend_logs_metadata")
|
||||
if not isinstance(spend_logs_metadata, dict):
|
||||
return
|
||||
existing_spend_logs_metadata = data[_metadata_variable_name].get("spend_logs_metadata")
|
||||
if isinstance(existing_spend_logs_metadata, dict):
|
||||
for key, value in spend_logs_metadata.items():
|
||||
if key not in existing_spend_logs_metadata:
|
||||
existing_spend_logs_metadata[key] = value
|
||||
else:
|
||||
data[_metadata_variable_name]["spend_logs_metadata"] = spend_logs_metadata
|
||||
|
||||
@staticmethod
|
||||
def add_key_level_controls(key_metadata: Optional[dict], data: dict, _metadata_variable_name: str):
|
||||
if key_metadata is None:
|
||||
|
|
@ -1047,17 +1060,11 @@ class LiteLLMProxyRequestSetup:
|
|||
)
|
||||
if "disable_global_guardrails" in key_metadata and isinstance(key_metadata["disable_global_guardrails"], bool):
|
||||
data[_metadata_variable_name]["disable_global_guardrails"] = key_metadata["disable_global_guardrails"]
|
||||
if "spend_logs_metadata" in key_metadata and isinstance(key_metadata["spend_logs_metadata"], dict):
|
||||
if "spend_logs_metadata" in data[_metadata_variable_name] and isinstance(
|
||||
data[_metadata_variable_name]["spend_logs_metadata"], dict
|
||||
):
|
||||
for key, value in key_metadata["spend_logs_metadata"].items():
|
||||
if (
|
||||
key not in data[_metadata_variable_name]["spend_logs_metadata"]
|
||||
): # don't override k-v pair sent by request (user request)
|
||||
data[_metadata_variable_name]["spend_logs_metadata"][key] = value
|
||||
else:
|
||||
data[_metadata_variable_name]["spend_logs_metadata"] = key_metadata["spend_logs_metadata"]
|
||||
LiteLLMProxyRequestSetup._merge_spend_logs_metadata(
|
||||
source_metadata=key_metadata,
|
||||
data=data,
|
||||
_metadata_variable_name=_metadata_variable_name,
|
||||
)
|
||||
|
||||
## KEY-LEVEL DISABLE FALLBACKS
|
||||
if "disable_fallbacks" in key_metadata and isinstance(key_metadata["disable_fallbacks"], bool):
|
||||
|
|
@ -1554,17 +1561,19 @@ async def add_litellm_data_to_request(
|
|||
team_metadata["opted_out_global_guardrails"], list
|
||||
):
|
||||
data[_metadata_variable_name]["opted_out_global_guardrails"] = team_metadata["opted_out_global_guardrails"]
|
||||
if "spend_logs_metadata" in team_metadata and isinstance(team_metadata["spend_logs_metadata"], dict):
|
||||
if "spend_logs_metadata" in data[_metadata_variable_name] and isinstance(
|
||||
data[_metadata_variable_name]["spend_logs_metadata"], dict
|
||||
):
|
||||
for key, value in team_metadata["spend_logs_metadata"].items():
|
||||
if (
|
||||
key not in data[_metadata_variable_name]["spend_logs_metadata"]
|
||||
): # don't override k-v pair sent by request (user request)
|
||||
data[_metadata_variable_name]["spend_logs_metadata"][key] = value
|
||||
else:
|
||||
data[_metadata_variable_name]["spend_logs_metadata"] = team_metadata["spend_logs_metadata"]
|
||||
LiteLLMProxyRequestSetup._merge_spend_logs_metadata(
|
||||
source_metadata=team_metadata,
|
||||
data=data,
|
||||
_metadata_variable_name=_metadata_variable_name,
|
||||
)
|
||||
|
||||
## ORG-LEVEL SPEND LOGS
|
||||
organization_metadata = user_api_key_dict.organization_metadata or {}
|
||||
LiteLLMProxyRequestSetup._merge_spend_logs_metadata(
|
||||
source_metadata=organization_metadata,
|
||||
data=data,
|
||||
_metadata_variable_name=_metadata_variable_name,
|
||||
)
|
||||
|
||||
## PROJECT-LEVEL TAGS
|
||||
project_metadata = user_api_key_dict.project_metadata or {}
|
||||
|
|
|
|||
|
|
@ -225,6 +225,62 @@ async def test_add_litellm_data_to_request_parses_string_metadata():
|
|||
assert updated_data["metadata"]["generation_name"] == "gen123"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_litellm_data_to_request_merges_org_level_spend_logs_metadata():
|
||||
"""Org-level spend_logs_metadata must be merged into the request, with
|
||||
key > team > org precedence and without overriding request-supplied keys.
|
||||
|
||||
Regression test for https://github.com/BerriAI/litellm/issues/33663 where
|
||||
organization-level spend_logs_metadata was silently dropped.
|
||||
"""
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
|
||||
request_mock = MagicMock(spec=Request)
|
||||
request_mock.url.path = "/v1/completions"
|
||||
request_mock.url = MagicMock()
|
||||
request_mock.url.__str__.return_value = "http://localhost/v1/completions"
|
||||
request_mock.method = "POST"
|
||||
request_mock.query_params = {}
|
||||
request_mock.headers = {"Content-Type": "application/json"}
|
||||
request_mock.client = MagicMock()
|
||||
request_mock.client.host = "127.0.0.1"
|
||||
|
||||
data = {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"metadata": {"spend_logs_metadata": {"source": "request"}},
|
||||
}
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="hashed-key",
|
||||
metadata={"spend_logs_metadata": {"key_only": "k", "shared": "key"}},
|
||||
team_metadata={"spend_logs_metadata": {"team_only": "t", "shared": "team"}},
|
||||
organization_metadata={
|
||||
"spend_logs_metadata": {"org_only": "o", "shared": "org", "source": "org"}
|
||||
},
|
||||
spend=0.0,
|
||||
max_budget=100.0,
|
||||
model_max_budget={},
|
||||
team_spend=0.0,
|
||||
team_max_budget=200.0,
|
||||
)
|
||||
|
||||
updated_data = await add_litellm_data_to_request(
|
||||
data=data,
|
||||
request=request_mock,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={},
|
||||
version="test-version",
|
||||
)
|
||||
|
||||
spend_logs_metadata = updated_data["metadata"]["spend_logs_metadata"]
|
||||
assert spend_logs_metadata["org_only"] == "o"
|
||||
assert spend_logs_metadata["key_only"] == "k"
|
||||
assert spend_logs_metadata["team_only"] == "t"
|
||||
assert spend_logs_metadata["source"] == "request"
|
||||
assert spend_logs_metadata["shared"] == "key"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_litellm_data_to_request_strips_admin_injection_slots():
|
||||
"""User-supplied user_api_key_metadata / user_api_key_team_metadata /
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue