From 128d9a348816ed7a2bc0444016f613340bfaeded Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 2 Sep 2025 15:13:15 -0700 Subject: [PATCH] [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 --- docs/my-website/docs/proxy/enterprise.md | 87 ++++++++++++ docs/my-website/docs/proxy/request_headers.md | 2 + litellm/proxy/_types.py | 6 + litellm/proxy/litellm_pre_call_utils.py | 52 +++++++- .../proxy/test_litellm_pre_call_utils.py | 124 ++++++++++++++++++ 5 files changed, 266 insertions(+), 5 deletions(-) diff --git a/docs/my-website/docs/proxy/enterprise.md b/docs/my-website/docs/proxy/enterprise.md index 468bcad2cf8..7d50aedb424 100644 --- a/docs/my-website/docs/proxy/enterprise.md +++ b/docs/my-website/docs/proxy/enterprise.md @@ -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) +``` + @@ -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(); +``` + @@ -502,6 +566,29 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ } }' ``` + + + + + +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" + } + ] +}' +``` + diff --git a/docs/my-website/docs/proxy/request_headers.md b/docs/my-website/docs/proxy/request_headers.md index c250d42f7bb..347dbe6a1bc 100644 --- a/docs/my-website/docs/proxy/request_headers.md +++ b/docs/my-website/docs/proxy/request_headers.md @@ -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. diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 87c5bd7a7c5..7206d3dcb0e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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 diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 32a283066ed..557aa3ceb9a 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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): diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 5104ffd80de..817f19d8d7d 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -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" + +