mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
[Feature]: Add header support for spend_logs_metadata (#14186)
* fix: allow settings spend_logs_metadata * fix add_litellm_data_for_backend_llm_call * fix: add add_litellm_metadata_from_request_headers * fix add_litellm_metadata_from_request_headers * test_add_litellm_metadata_from_request_headers * add_litellm_metadata_from_request_headers * docs Tracking Spend with custom metadata * add_litellm_metadata_from_request_headers * add_litellm_metadata_from_request_headers
This commit is contained in:
parent
c821f1ddf1
commit
128d9a3488
5 changed files with 266 additions and 5 deletions
|
|
@ -439,6 +439,33 @@ response = client.chat.completions.create(
|
|||
|
||||
print(response)
|
||||
```
|
||||
|
||||
**Using Headers:**
|
||||
|
||||
```python
|
||||
import openai
|
||||
client = openai.OpenAI(
|
||||
api_key="sk-1234",
|
||||
base_url="http://0.0.0.0:4000"
|
||||
)
|
||||
|
||||
# Pass spend logs metadata via headers
|
||||
response = client.chat.completions.create(
|
||||
model="gpt-3.5-turbo",
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "this is a test request, write a short poem"
|
||||
}
|
||||
],
|
||||
extra_headers={
|
||||
"x-litellm-spend-logs-metadata": '{"user_id": "12345", "project_id": "proj_abc", "request_type": "chat_completion"}'
|
||||
}
|
||||
)
|
||||
|
||||
print(response)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
|
||||
|
|
@ -478,6 +505,43 @@ async function runOpenAI() {
|
|||
// Call the asynchronous function
|
||||
runOpenAI();
|
||||
```
|
||||
|
||||
**Using Headers:**
|
||||
|
||||
```js
|
||||
const openai = require('openai');
|
||||
|
||||
async function runOpenAI() {
|
||||
const client = new openai.OpenAI({
|
||||
apiKey: 'sk-1234',
|
||||
baseURL: 'http://0.0.0.0:4000'
|
||||
});
|
||||
|
||||
try {
|
||||
const response = await client.chat.completions.create({
|
||||
model: 'gpt-3.5-turbo',
|
||||
messages: [
|
||||
{
|
||||
role: 'user',
|
||||
content: "this is a test request, write a short poem"
|
||||
},
|
||||
]
|
||||
}, {
|
||||
headers: {
|
||||
'x-litellm-spend-logs-metadata': '{"user_id": "12345", "project_id": "proj_abc", "request_type": "chat_completion"}'
|
||||
}
|
||||
});
|
||||
console.log(response);
|
||||
} catch (error) {
|
||||
console.log("got this exception from server");
|
||||
console.error(error);
|
||||
}
|
||||
}
|
||||
|
||||
// Call the asynchronous function
|
||||
runOpenAI();
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="Curl" label="Curl Request">
|
||||
|
|
@ -502,6 +566,29 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
|
|||
}
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="headers" label="Using Headers">
|
||||
|
||||
Pass `x-litellm-spend-logs-metadata` as a request header with JSON string
|
||||
|
||||
```shell
|
||||
curl --location 'http://0.0.0.0:4000/chat/completions' \
|
||||
--header 'Content-Type: application/json' \
|
||||
--header 'Authorization: Bearer sk-1234' \
|
||||
--header 'x-litellm-spend-logs-metadata: {"user_id": "12345", "project_id": "proj_abc", "request_type": "chat_completion"}' \
|
||||
--data '{
|
||||
"model": "gpt-3.5-turbo",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "what llm are you"
|
||||
}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="langchain" label="Langchain">
|
||||
|
||||
|
|
|
|||
|
|
@ -14,6 +14,8 @@ Special headers that are supported by LiteLLM.
|
|||
|
||||
`x-litellm-num-retries`: Optional[int]: The number of retries for the request.
|
||||
|
||||
`x-litellm-spend-logs-metadata`: Optional[str]: JSON string containing custom metadata to include in spend logs. Example: `{"user_id": "12345", "project_id": "proj_abc", "request_type": "chat_completion"}`. [Learn More](./logging#tracking-spend-with-custom-metadata)
|
||||
|
||||
## Anthropic Headers
|
||||
|
||||
`anthropic-version` Optional[str]: The version of the Anthropic API to use.
|
||||
|
|
|
|||
|
|
@ -2908,6 +2908,12 @@ class LitellmDataForBackendLLMCall(TypedDict, total=False):
|
|||
user: Optional[str]
|
||||
num_retries: Optional[int]
|
||||
|
||||
class LitellmMetadataFromRequestHeaders(TypedDict, total=False):
|
||||
"""
|
||||
Headers a user can pass that will get added to litellm metadata for the request
|
||||
"""
|
||||
spend_logs_metadata: Optional[dict]
|
||||
|
||||
|
||||
class JWTKeyItem(TypedDict, total=False):
|
||||
kid: str
|
||||
|
|
|
|||
|
|
@ -291,6 +291,17 @@ class LiteLLMProxyRequestSetup:
|
|||
if num_retries_header is not None:
|
||||
return int(num_retries_header)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_spend_logs_metadata_from_request_headers(headers: dict) -> Optional[dict]:
|
||||
"""
|
||||
Get the `spend_logs_metadata` from the request headers.
|
||||
"""
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
spend_logs_metadata_header = headers.get("x-litellm-spend-logs-metadata", None)
|
||||
if spend_logs_metadata_header is not None:
|
||||
return safe_json_loads(spend_logs_metadata_header)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_forwardable_headers(
|
||||
|
|
@ -459,6 +470,30 @@ class LiteLLMProxyRequestSetup:
|
|||
data["num_retries"] = num_retries
|
||||
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def add_litellm_metadata_from_request_headers(
|
||||
headers: dict,
|
||||
data: dict,
|
||||
_metadata_variable_name: str,
|
||||
) -> dict:
|
||||
"""
|
||||
Add litellm metadata from request headers
|
||||
|
||||
Relevant issue: https://github.com/BerriAI/litellm/issues/14008
|
||||
"""
|
||||
from litellm.proxy._types import LitellmMetadataFromRequestHeaders
|
||||
metadata_from_headers = LitellmMetadataFromRequestHeaders()
|
||||
spend_logs_metadata = LiteLLMProxyRequestSetup._get_spend_logs_metadata_from_request_headers(headers)
|
||||
if spend_logs_metadata is not None:
|
||||
metadata_from_headers["spend_logs_metadata"] = spend_logs_metadata
|
||||
|
||||
#########################################################################################
|
||||
# Finally update the requests metadata with the `metadata_from_headers`
|
||||
#########################################################################################
|
||||
if isinstance(data[_metadata_variable_name], dict):
|
||||
data[_metadata_variable_name].update(metadata_from_headers)
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def get_sanitized_user_information_from_key(
|
||||
|
|
@ -643,6 +678,10 @@ async def add_litellm_data_to_request( # noqa: PLR0915
|
|||
from litellm.types.proxy.litellm_pre_call_utils import SecretFields
|
||||
|
||||
safe_add_api_version_from_query_params(data, request)
|
||||
_metadata_variable_name = _get_metadata_variable_name(request)
|
||||
if data.get(_metadata_variable_name, None) is None:
|
||||
data[_metadata_variable_name] = {}
|
||||
|
||||
|
||||
_headers = clean_headers(
|
||||
request.headers,
|
||||
|
|
@ -661,6 +700,14 @@ async def add_litellm_data_to_request( # noqa: PLR0915
|
|||
)
|
||||
)
|
||||
|
||||
data.update(
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=_headers,
|
||||
data=data,
|
||||
_metadata_variable_name=_metadata_variable_name,
|
||||
)
|
||||
)
|
||||
|
||||
# check for forwardable headers
|
||||
data = LiteLLMProxyRequestSetup.add_headers_to_llm_call_by_model_group(
|
||||
data=data, headers=_headers, user_api_key_dict=user_api_key_dict
|
||||
|
|
@ -711,11 +758,6 @@ async def add_litellm_data_to_request( # noqa: PLR0915
|
|||
|
||||
verbose_proxy_logger.debug("receiving data: %s", data)
|
||||
|
||||
_metadata_variable_name = _get_metadata_variable_name(request)
|
||||
|
||||
if data.get(_metadata_variable_name, None) is None:
|
||||
data[_metadata_variable_name] = {}
|
||||
|
||||
# Parse metadata if it's a string (e.g., from multipart/form-data)
|
||||
if "metadata" in data and data["metadata"] is not None:
|
||||
if isinstance(data["metadata"], str):
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from unittest.mock import MagicMock, patch
|
|||
import pytest
|
||||
from fastapi import Request
|
||||
|
||||
import litellm
|
||||
from litellm.proxy._types import TeamCallbackMetadata, UserAPIKeyAuth
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
KeyAndTeamLoggingSettings,
|
||||
|
|
@ -935,3 +936,126 @@ def test_add_headers_to_llm_call_by_model_group_existing_headers_in_data():
|
|||
finally:
|
||||
# Restore original model_group_settings
|
||||
litellm.model_group_settings = original_model_group_settings
|
||||
|
||||
import json
|
||||
import time
|
||||
from typing import Optional
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from fastapi.responses import Response
|
||||
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
|
||||
class TestCustomLogger(CustomLogger):
|
||||
def __init__(self):
|
||||
self.standard_logging_object: Optional[StandardLoggingPayload] = None
|
||||
super().__init__()
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
print(f"SUCCESS CALLBACK CALLED! kwargs keys: {list(kwargs.keys())}")
|
||||
self.standard_logging_object = kwargs.get("standard_logging_object")
|
||||
print(f"Captured standard_logging_object: {self.standard_logging_object}")
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
print(f"FAILURE CALLBACK CALLED! kwargs keys: {list(kwargs.keys())}")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_litellm_metadata_from_request_headers():
|
||||
"""
|
||||
Test that add_litellm_metadata_from_request_headers properly adds litellm metadata from request headers,
|
||||
makes an LLM request using base_process_llm_request, sleeps for 3 seconds, and checks standard_logging_payload has spend_logs_metadata from headers
|
||||
|
||||
Relevant issue: https://github.com/BerriAI/litellm/issues/14008
|
||||
"""
|
||||
# Set up test logger
|
||||
litellm._turn_on_debug()
|
||||
test_logger = TestCustomLogger()
|
||||
litellm.callbacks = [test_logger]
|
||||
|
||||
# Prepare test data (ensure no streaming, add mock_response and api_key to route to litellm.acompletion)
|
||||
headers = {"x-litellm-spend-logs-metadata": '{"user_id": "12345", "project_id": "proj_abc", "request_type": "chat_completion", "timestamp": "2025-09-02T10:30:00Z"}'}
|
||||
data = {"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}], "stream": False, "mock_response": "Hi", "api_key": "fake-key"}
|
||||
|
||||
# Create mock request with headers
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.headers = headers
|
||||
mock_request.url.path = "/chat/completions"
|
||||
|
||||
# Create mock response
|
||||
mock_fastapi_response = MagicMock(spec=Response)
|
||||
|
||||
# Create mock user API key dict
|
||||
mock_user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_id="test-user",
|
||||
org_id="test-org"
|
||||
)
|
||||
|
||||
# Create mock proxy logging object
|
||||
mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
|
||||
# Create async functions for the hooks
|
||||
async def mock_during_call_hook(*args, **kwargs):
|
||||
return None
|
||||
|
||||
async def mock_pre_call_hook(*args, **kwargs):
|
||||
return data
|
||||
|
||||
async def mock_post_call_success_hook(*args, **kwargs):
|
||||
# Return the response unchanged
|
||||
return kwargs.get('response', args[2] if len(args) > 2 else None)
|
||||
|
||||
mock_proxy_logging_obj.during_call_hook = mock_during_call_hook
|
||||
mock_proxy_logging_obj.pre_call_hook = mock_pre_call_hook
|
||||
mock_proxy_logging_obj.post_call_success_hook = mock_post_call_success_hook
|
||||
|
||||
# Create mock proxy config
|
||||
mock_proxy_config = MagicMock()
|
||||
|
||||
# Create mock general settings
|
||||
general_settings = {}
|
||||
|
||||
# Create mock select_data_generator with correct signature
|
||||
def mock_select_data_generator(response=None, user_api_key_dict=None, request_data=None):
|
||||
async def mock_generator():
|
||||
yield "data: " + json.dumps({"choices": [{"delta": {"content": "Hello"}}]}) + "\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
return mock_generator()
|
||||
|
||||
# Create the processor
|
||||
processor = ProxyBaseLLMRequestProcessing(data=data)
|
||||
|
||||
# Call base_process_llm_request (it will use the mock_response="Hi" parameter)
|
||||
result = await processor.base_process_llm_request(
|
||||
request=mock_request,
|
||||
fastapi_response=mock_fastapi_response,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
route_type="acompletion",
|
||||
proxy_logging_obj=mock_proxy_logging_obj,
|
||||
general_settings=general_settings,
|
||||
proxy_config=mock_proxy_config,
|
||||
select_data_generator=mock_select_data_generator,
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
is_streaming_request=False
|
||||
)
|
||||
|
||||
# Sleep for 3 seconds to allow logging to complete
|
||||
await asyncio.sleep(3)
|
||||
|
||||
# Check if standard_logging_object was set
|
||||
assert test_logger.standard_logging_object is not None, "standard_logging_object should be populated after LLM request"
|
||||
|
||||
# Verify the logging object contains expected metadata
|
||||
standard_logging_obj = test_logger.standard_logging_object
|
||||
|
||||
print(f"Standard logging object captured: {json.dumps(standard_logging_obj, indent=4, default=str)}")
|
||||
|
||||
SPEND_LOGS_METADATA = standard_logging_obj["metadata"]["spend_logs_metadata"]
|
||||
assert SPEND_LOGS_METADATA == dict(json.loads(headers["x-litellm-spend-logs-metadata"])), "spend_logs_metadata should be the same as the headers"
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue