Merge branch 'BerriAI:main' into fix/ollama-gpt-oss-thinking-field

This commit is contained in:
Cole McIntosh 2025-08-07 15:14:38 -06:00 • committed by GitHub
commit b66f4a7d4b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
55 changed files with 3795 additions and 840 deletions

View file

@ -0,0 +1,53 @@
import base64
from openai import OpenAI
import time
client = OpenAI(
base_url="http://0.0.0.0:4001",
api_key="sk-1234"
)
# Function to encode the image
def encode_image(image_path):
with open(image_path, "rb") as image_file:
return base64.b64encode(image_file.read()).decode("utf-8")
# Path to your image
image_path = "litellm/proxy/logo.jpg"
# Getting the Base64 string
base64_image = encode_image(image_path)
response = client.responses.create(
model="bedrock/us.anthropic.claude-3-5-sonnet-20241022-v2:0",
input=[
{
"role": "user",
"content": [
{ "type": "input_text", "text": "what color is the image"},
{
"type": "input_image",
"image_url": f"data:image/jpeg;base64,{base64_image}",
},
],
}
],
)
print(response.output_text)
print("response1 id===", response.id)
print("sleeping for 20 seconds...")
time.sleep(20)
print("making follow up request for existing id")
response2 = client.responses.create(
model="bedrock/us.anthropic.claude-3-5-sonnet-20241022-v2:0",
previous_response_id=response.id,
input="ok, and what objects are in the image?"
)
print(response2.output_text)

View file

@ -1,9 +1,11 @@
{{- if .Values.migrationJob.enabled }}
# This job runs the prisma migrations for the LiteLLM DB.
# This job runs the Prisma migrations for the LiteLLM DB.
apiVersion: batch/v1
kind: Job
metadata:
name: {{ include "litellm.fullname" . }}-migrations
labels:
{{- include "litellm.labels" . | nindent 4 }}
annotations:
{{- if .Values.migrationJob.hooks.argocd.enabled }}
argocd.argoproj.io/hook: PreSync
@ -18,6 +20,8 @@ metadata:
spec:
template:
metadata:
labels:
{{- include "litellm.labels" . | nindent 8 }}
annotations:
{{- with .Values.migrationJob.annotations }}
{{- toYaml . | nindent 8 }}

View file

@ -11,7 +11,7 @@ WORKDIR /app
# Install build dependencies
USER root
RUN apk add --no-cache build-base bash \
&& pip install --no-cache-dir --upgrade pip build
&& pip install --no-cache-dir --upgrade pip build
# Copy project files
COPY . .
@ -21,8 +21,8 @@ RUN chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
# Build package and wheel dependencies
RUN rm -rf dist/* && python -m build && \
pip install dist/*.whl && \
pip wheel --no-cache-dir --wheel-dir=/wheels/ -r requirements.txt
pip install dist/*.whl && \
pip wheel --no-cache-dir --wheel-dir=/wheels/ -r requirements.txt
# -----------------
# Runtime Stage
@ -33,9 +33,10 @@ WORKDIR /app
# Install runtime dependencies
USER root
RUN apk upgrade --no-cache && \
apk add --no-cache bash libstdc++ ca-certificates openssl
apk add --no-cache bash libstdc++ ca-certificates openssl
# Copy only necessary artifacts from builder stage for runtime
COPY . .
COPY --from=builder /app/docker/entrypoint.sh /app/docker/prod_entrypoint.sh /app/docker/
COPY --from=builder /app/schema.prisma /app/schema.prisma
COPY --from=builder /app/dist/*.whl .
@ -43,16 +44,16 @@ COPY --from=builder /wheels/ /wheels/
# Install package from wheel and dependencies
RUN pip install *.whl /wheels/* --no-index --find-links=/wheels/ \
&& rm -f *.whl \
&& rm -rf /wheels
&& rm -f *.whl \
&& rm -rf /wheels
# Install semantic_router without dependencies
RUN pip install semantic_router --no-deps
# Ensure correct JWT library is used (pyjwt not jwt)
RUN pip uninstall jwt -y && \
pip uninstall PyJWT -y && \
pip install PyJWT==2.9.0 --no-cache-dir
pip uninstall PyJWT -y && \
pip install PyJWT==2.9.0 --no-cache-dir
# --- Prisma Handling for Non-Root User ---
# Set Prisma cache directories
@ -61,29 +62,29 @@ ENV NPM_CONFIG_CACHE=/.npm
# Install prisma and make entrypoints executable
RUN pip install --no-cache-dir prisma && \
chmod +x docker/entrypoint.sh && \
chmod +x docker/prod_entrypoint.sh
chmod +x docker/entrypoint.sh && \
chmod +x docker/prod_entrypoint.sh
# Create directories and set permissions for non-root user
RUN mkdir -p /nonexistent /.npm && \
chown -R nobody:nogroup /app && \
chown -R nobody:nogroup /nonexistent /.npm && \
PRISMA_PATH=$(python -c "import os, prisma; print(os.path.dirname(prisma.__file__))") && \
chown -R nobody:nogroup $PRISMA_PATH
chown -R nobody:nogroup /app && \
chown -R nobody:nogroup /nonexistent /.npm && \
PRISMA_PATH=$(python -c "import os, prisma; print(os.path.dirname(prisma.__file__))") && \
chown -R nobody:nogroup $PRISMA_PATH
# --- OpenShift Compatibility: Apply Red Hat recommended pattern ---
# Get paths for directories that need write access at runtime
RUN PRISMA_PATH=$(python -c "import os, prisma; print(os.path.dirname(prisma.__file__))") && \
LITELLM_PROXY_EXTRAS_PATH=$(python -c "import os, litellm_proxy_extras; print(os.path.dirname(litellm_proxy_extras.__file__))" 2>/dev/null || echo "") && \
# Set group ownership to 0 (root group) for OpenShift compatibility && \
chgrp -R 0 $PRISMA_PATH && \
[ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chgrp -R 0 $LITELLM_PROXY_EXTRAS_PATH || true && \
# Mirror owner permissions to group (g=u) as recommended by Red Hat && \
chmod -R g=u $PRISMA_PATH && \
[ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chmod -R g=u $LITELLM_PROXY_EXTRAS_PATH || true && \
# Ensure directories are writable by group && \
chmod -R g+w $PRISMA_PATH && \
[ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chmod -R g+w $LITELLM_PROXY_EXTRAS_PATH || true
LITELLM_PROXY_EXTRAS_PATH=$(python -c "import os, litellm_proxy_extras; print(os.path.dirname(litellm_proxy_extras.__file__))" 2>/dev/null || echo "") && \
# Set group ownership to 0 (root group) for OpenShift compatibility && \
chgrp -R 0 $PRISMA_PATH && \
[ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chgrp -R 0 $LITELLM_PROXY_EXTRAS_PATH || true && \
# Mirror owner permissions to group (g=u) as recommended by Red Hat && \
chmod -R g=u $PRISMA_PATH && \
[ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chmod -R g=u $LITELLM_PROXY_EXTRAS_PATH || true && \
# Ensure directories are writable by group && \
chmod -R g+w $PRISMA_PATH && \
[ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chmod -R g+w $LITELLM_PROXY_EXTRAS_PATH || true
# Switch to non-root user
USER nobody
@ -100,4 +101,4 @@ ENTRYPOINT ["/app/docker/prod_entrypoint.sh"]
# Append "--detailed_debug" to the end of CMD to view detailed debug logs
# CMD ["--port", "4000", "--detailed_debug"]
CMD ["--port", "4000"]
CMD ["--port", "4000"]

View file

@ -15,6 +15,7 @@ import os
# set env
os.environ["BRAINTRUST_API_KEY"] = ""
os.environ["BRAINTRUST_API_BASE"] = "https://api.braintrustdata.com/v1"
os.environ['OPENAI_API_KEY']=""
# set braintrust as a callback, litellm will send the data to braintrust
@ -35,6 +36,7 @@ response = litellm.completion(
```env
BRAINTRUST_API_KEY=""
BRAINTRUST_API_BASE="https://api.braintrustdata.com/v1"
```
2. Add braintrust to callbacks
@ -157,6 +159,8 @@ For more examples, [**Click Here**](../proxy/user_keys.md#chatcompletions)
</TabItem>
</Tabs>
You can use `BRAINTRUST_API_BASE` to point to your self-hosted Braintrust data plane. Read more about this [here](https://www.braintrust.dev/docs/guides/self-hosting).
## Full API Spec
Here's everything you can pass in metadata for a braintrust request

View file

@ -6,20 +6,11 @@ LiteLLM supports the following models for OCI on-demand GenAI API.
Check the [OCI Models List](https://docs.oracle.com/en-us/iaas/Content/generative-ai/pretrained-models.htm) to see if the model is available for your region.
- `cohere.command-a-03-2025`
- `cohere.command-r-08-2024`
- `cohere.command-plus-latest` (alias `cohere.command-r-plus-08-2024`)
- `cohere.command-r-16k` (deprecated)
- `cohere.command-r-plus` (deprecated)
- `meta.llama-4-maverick-17b-128e-instruct-fp8`
- `meta.llama-4-scout-17b-16e-instruct`
- `meta.llama-3.3-70b-instruct`
- `meta.llama-3.2-90b-vision-instruct`
- `meta.llama-3.2-11b-vision-instruct`
- `meta.llama-3.1-405b-instruct`
- `meta.llama-3.1-70b-instruct`
- `meta.llama-3-70b-instruct`
- `xai.grok-4`
- `xai.grok-3`

View file

@ -369,6 +369,7 @@ router_settings:
| BEDROCK_MAX_POLICY_SIZE | Maximum size for Bedrock policy. Default is 75
| BERRISPEND_ACCOUNT_ID | Account ID for BerriSpend service
| BRAINTRUST_API_KEY | API key for Braintrust integration
| BRAINTRUST_API_BASE | Base URL for Braintrust API. Default is https://api.braintrustdata.com/v1
| CACHED_STREAMING_CHUNK_DELAY | Delay in seconds for cached streaming chunks. Default is 0.02
| CIRCLE_OIDC_TOKEN | OpenID Connect token for CircleCI
| CIRCLE_OIDC_TOKEN_V2 | Version 2 of the OpenID Connect token for CircleCI

View file

@ -1,160 +0,0 @@
import json
from typing import TYPE_CHECKING, Any, List, Optional, Union, cast
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import SpendLogsPayload
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionResponseMessage,
GenericChatCompletionMessage,
ResponseInputParam,
)
from litellm.types.utils import ChatCompletionMessageToolCall, Message, ModelResponse
if TYPE_CHECKING:
from litellm.responses.litellm_completion_transformation.transformation import (
ChatCompletionSession,
)
else:
ChatCompletionSession = Any
class _ENTERPRISE_ResponsesSessionHandler:
@staticmethod
async def get_chat_completion_message_history_for_previous_response_id(
previous_response_id: str,
) -> ChatCompletionSession:
"""
Return the chat completion message history for a previous response id
"""
from litellm.responses.litellm_completion_transformation.transformation import (
ChatCompletionSession,
LiteLLMCompletionResponsesConfig,
)
verbose_proxy_logger.debug(
"inside get_chat_completion_message_history_for_previous_response_id"
)
all_spend_logs: List[
SpendLogsPayload
] = await _ENTERPRISE_ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(
previous_response_id
)
verbose_proxy_logger.debug(
"found %s spend logs for this response id", len(all_spend_logs)
)
litellm_session_id: Optional[str] = None
if len(all_spend_logs) > 0:
litellm_session_id = all_spend_logs[0].get("session_id")
chat_completion_message_history: List[
Union[
AllMessageValues,
GenericChatCompletionMessage,
ChatCompletionMessageToolCall,
ChatCompletionResponseMessage,
Message,
]
] = []
for spend_log in all_spend_logs:
proxy_server_request: Union[str, dict] = (
spend_log.get("proxy_server_request") or "{}"
)
proxy_server_request_dict: Optional[dict] = None
response_input_param: Optional[Union[str, ResponseInputParam]] = None
if isinstance(proxy_server_request, dict):
proxy_server_request_dict = proxy_server_request
else:
proxy_server_request_dict = json.loads(proxy_server_request)
############################################################
# Add Input messages for this Spend Log
############################################################
if proxy_server_request_dict:
_response_input_param = proxy_server_request_dict.get("input", None)
if isinstance(_response_input_param, str):
response_input_param = _response_input_param
elif isinstance(_response_input_param, dict):
response_input_param = cast(
ResponseInputParam, _response_input_param
)
if response_input_param:
chat_completion_messages = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
input=response_input_param,
responses_api_request=proxy_server_request_dict or {},
)
chat_completion_message_history.extend(chat_completion_messages)
############################################################
# Add Output messages for this Spend Log
############################################################
_response_output = spend_log.get("response", "{}")
if isinstance(_response_output, dict):
# transform `ChatCompletion Response` to `ResponsesAPIResponse`
model_response = ModelResponse(**_response_output)
for choice in model_response.choices:
if hasattr(choice, "message"):
chat_completion_message_history.append(
getattr(choice, "message")
)
verbose_proxy_logger.debug(
"chat_completion_message_history %s",
json.dumps(chat_completion_message_history, indent=4, default=str),
)
return ChatCompletionSession(
messages=chat_completion_message_history,
litellm_session_id=litellm_session_id,
)
@staticmethod
async def get_all_spend_logs_for_previous_response_id(
previous_response_id: str,
) -> List[SpendLogsPayload]:
"""
Get all spend logs for a previous response id
SQL query
SELECT session_id FROM spend_logs WHERE response_id = previous_response_id, SELECT * FROM spend_logs WHERE session_id = session_id
"""
from litellm.proxy.proxy_server import prisma_client
verbose_proxy_logger.debug("decoding response id=%s", previous_response_id)
decoded_response_id = (
ResponsesAPIRequestUtils._decode_responses_api_response_id(
previous_response_id
)
)
previous_response_id = decoded_response_id.get(
"response_id", previous_response_id
)
if prisma_client is None:
return []
query = """
WITH matching_session AS (
SELECT session_id
FROM "LiteLLM_SpendLogs"
WHERE request_id = $1
)
SELECT *
FROM "LiteLLM_SpendLogs"
WHERE session_id IN (SELECT session_id FROM matching_session)
ORDER BY "endTime" ASC;
"""
spend_logs = await prisma_client.db.query_raw(query, previous_response_id)
verbose_proxy_logger.debug(
"Found the following spend logs for previous response id %s: %s",
previous_response_id,
json.dumps(spend_logs, indent=4, default=str),
)
return spend_logs

Binary file not shown.

View file

@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm-proxy-extras"
version = "0.2.15"
version = "0.2.16"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
authors = ["BerriAI"]
readme = "README.md"
@ -22,7 +22,7 @@ requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "0.2.15"
version = "0.2.16"
version_files = [
"pyproject.toml:version",
"../requirements.txt:litellm-proxy-extras==",

View file

@ -533,6 +533,7 @@ morph_models: List = []
lambda_ai_models: List = []
hyperbolic_models: List = []
recraft_models: List = []
oci_models: List = []
def is_bedrock_pricing_only_model(key: str) -> bool:
@ -722,6 +723,8 @@ def add_known_models():
hyperbolic_models.append(key)
elif value.get("litellm_provider") == "recraft":
recraft_models.append(key)
elif value.get("litellm_provider") == "oci":
oci_models.append(key)
add_known_models()
@ -810,6 +813,7 @@ model_list = (
+ morph_models
+ lambda_ai_models
+ recraft_models
+ oci_models
)
model_list_set = set(model_list)
@ -883,6 +887,7 @@ models_by_provider: dict = {
"lambda_ai": lambda_ai_models,
"hyperbolic": hyperbolic_models,
"recraft": recraft_models,
"oci": oci_models,
}
# mapping for those models which have larger equivalents

View file

@ -766,6 +766,7 @@ MAXIMUM_TRACEBACK_LINES_TO_LOG = int(os.getenv("MAXIMUM_TRACEBACK_LINES_TO_LOG",
X_LITELLM_DISABLE_CALLBACKS = "x-litellm-disable-callbacks"
LITELLM_METADATA_FIELD = "litellm_metadata"
OLD_LITELLM_METADATA_FIELD = "metadata"
LITELLM_TRUNCATED_PAYLOAD_FIELD = "litellm_truncated"
########################### LiteLLM Proxy Specific Constants ###########################
########################################################################################

View file

@ -42,7 +42,7 @@ class BraintrustLogger(CustomLogger):
) -> None:
super().__init__()
self.validate_environment(api_key=api_key)
self.api_base = api_base or API_BASE
self.api_base = api_base or os.getenv("BRAINTRUST_API_BASE") or API_BASE
self.default_project_id = None
self.api_key: str = api_key or os.getenv("BRAINTRUST_API_KEY") # type: ignore
self.headers = {

View file

@ -541,3 +541,14 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
model_call_details_copy["standard_logging_object"] = standard_logging_object_copy
return model_call_details_copy
async def get_proxy_server_request_from_cold_storage_with_object_key(
self,
object_key: str,
) -> Optional[dict]:
"""
Get the proxy server request from cold storage using the object key directly.
"""
pass

View file

@ -304,7 +304,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
data=prepped.body,
headers=prepped.headers,
)
SigV4Auth(credentials, "s3", self.s3_region_name).add_auth(aws_request)
aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(
aws_region_name=self.s3_region_name
)
SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request)
# Prepare the signed headers
signed_headers = dict(aws_request.headers.items())
@ -444,7 +447,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
data=prepped.body,
headers=prepped.headers,
)
SigV4Auth(credentials, "s3", self.s3_region_name).add_auth(aws_request)
aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(
aws_region_name=self.s3_region_name
)
SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request)
# Prepare the signed headers
signed_headers = dict(aws_request.headers.items())
@ -455,3 +461,108 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
response.raise_for_status()
except Exception as e:
verbose_logger.exception(f"Error uploading to s3: {str(e)}")
async def _download_object_from_s3(self, s3_object_key: str) -> Optional[dict]:
"""
Download and parse JSON object from S3.
Args:
s3_object_key: The S3 object key to download
Returns:
Optional[dict]: The parsed JSON object or None if not found/error
"""
try:
import hashlib
import requests
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call S3. Run 'pip install boto3'.")
try:
from litellm.litellm_core_utils.asyncify import asyncify
# Get AWS credentials
asyncified_get_credentials = asyncify(self.get_credentials)
credentials = await asyncified_get_credentials(
aws_access_key_id=self.s3_aws_access_key_id,
aws_secret_access_key=self.s3_aws_secret_access_key,
aws_session_token=self.s3_aws_session_token,
aws_region_name=self.s3_region_name,
aws_session_name=self.s3_aws_session_name,
aws_profile_name=self.s3_aws_profile_name,
aws_role_name=self.s3_aws_role_name,
aws_web_identity_token=self.s3_aws_web_identity_token,
aws_sts_endpoint=self.s3_aws_sts_endpoint,
)
verbose_logger.debug(
f"s3_v2 logger - downloading data from s3 - {s3_object_key}"
)
# Prepare the URL
url = f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{s3_object_key}"
if self.s3_endpoint_url:
url = self.s3_endpoint_url + "/" + s3_object_key
# Prepare the request for GET operation
# For GET requests, we need x-amz-content-sha256 with hash of empty string
empty_string_hash = hashlib.sha256(b"").hexdigest()
headers = {
"x-amz-content-sha256": empty_string_hash,
}
req = requests.Request("GET", url, headers=headers)
prepped = req.prepare()
# Sign the request
aws_request = AWSRequest(
method=prepped.method,
url=prepped.url,
headers=prepped.headers,
)
SigV4Auth(credentials, "s3", self.s3_region_name).add_auth(aws_request)
# Prepare the signed headers
signed_headers = dict(aws_request.headers.items())
# Make the request
response = await self.async_httpx_client.get(url, headers=signed_headers)
if response.status_code != 200:
verbose_logger.exception("S3 object not found, saw response=", response.text)
return None
# Parse JSON response
return response.json()
except Exception as e:
verbose_logger.exception(f"Error downloading from S3: {str(e)}")
return None
async def get_proxy_server_request_from_cold_storage_with_object_key(
self,
object_key: str,
) -> Optional[dict]:
"""
Get the proxy server request from cold storage
Allows fetching a dict of the proxy server request from s3 or GCS bucket.
Args:
request_id: The unique request ID to search for
start_time: The start time of the request (datetime or ISO string)
Returns:
Optional[dict]: The request data dictionary or None if not found
"""
try:
# Download and return the object from S3
downloaded_object = await self._download_object_from_s3(object_key)
return downloaded_object
except Exception as e:
verbose_logger.exception(f"Error retrieving object {object_key} from cold storage: {str(e)}")
return None

View file

@ -10,6 +10,7 @@ Example:
from typing import Union
from litellm import _custom_logger_compatible_callbacks_literal
from litellm.integrations.agentops import AgentOps
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
from litellm.integrations.argilla import ArgillaLogger
@ -150,3 +151,14 @@ class CustomLoggerRegistry:
if callback_class == class_type:
callback_strs.append(callback_str)
return callback_strs
@classmethod
def get_class_type_for_custom_logger_name(
cls,
custom_logger_name: _custom_logger_compatible_callbacks_literal,
) -> type:
"""
Get the class type for a given custom logger name
"""
return cls.CALLBACK_CLASS_STR_TO_CLASS_TYPE[custom_logger_name]

View file

@ -3830,6 +3830,8 @@ class StandardLoggingPayloadSetup:
] = None,
usage_object: Optional[dict] = None,
proxy_server_request: Optional[dict] = None,
start_time: Optional[dt_object] = None,
response_id: Optional[str] = None,
) -> StandardLoggingMetadata:
"""
Clean and filter the metadata dictionary to include only the specified keys in StandardLoggingMetadata.
@ -3881,6 +3883,7 @@ class StandardLoggingPayloadSetup:
usage_object=usage_object,
requester_custom_headers=None,
user_api_key_request_route=None,
cold_storage_object_key=None,
)
if isinstance(metadata, dict):
# Filter the metadata dictionary to include only the specified keys
@ -3913,6 +3916,16 @@ class StandardLoggingPayloadSetup:
proxy_server_request=proxy_server_request,
)
# Generate cold storage object key if cold storage is configured
if start_time is not None and response_id is not None:
cold_storage_object_key = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
start_time=start_time,
response_id=response_id,
team_alias=clean_metadata.get("user_api_key_team_alias"),
)
if cold_storage_object_key:
clean_metadata["cold_storage_object_key"] = cold_storage_object_key
return clean_metadata
@staticmethod
@ -4071,6 +4084,49 @@ class StandardLoggingPayloadSetup:
return api_base.rstrip("/")
return api_base
@staticmethod
def _generate_cold_storage_object_key(
start_time: dt_object,
response_id: str,
team_alias: Optional[str] = None,
) -> Optional[str]:
"""
Generate cold storage object key in the same format as S3Logger.
Args:
start_time: The start time of the request
response_id: The response ID
team_alias: Optional team alias for team-based prefixing
Returns:
Optional[str]: The generated object key or None if cold storage not configured
"""
# Generate object key in same format as S3Logger
from litellm.integrations.s3 import get_s3_object_key
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
# Only generate object key if cold storage is configured
configured_cold_storage_logger = ColdStorageHandler._get_configured_cold_storage_custom_logger()
if configured_cold_storage_logger is None:
return None
try:
# Generate file name in same format as litellm.utils.get_logging_id
s3_file_name = f"time-{start_time.strftime('%H-%M-%S-%f')}_{response_id}"
s3_object_key = get_s3_object_key(
s3_path="", # Use empty path as default
team_alias_prefix="", # Don't split by team alias for cold storage
start_time=start_time,
s3_file_name=s3_file_name,
)
return s3_object_key
except Exception:
# If any error occurs in generating the key, return None
return None
@staticmethod
def get_error_information(
original_exception: Optional[Exception],
@ -4322,6 +4378,8 @@ def get_standard_logging_object_payload(
),
usage_object=usage.model_dump(),
proxy_server_request=proxy_server_request,
start_time=start_time,
response_id=id,
)
_request_body = proxy_server_request.get("body", {})
@ -4469,6 +4527,7 @@ def get_standard_logging_metadata(
usage_object=None,
requester_custom_headers=None,
user_api_key_request_route=None,
cold_storage_object_key=None,
)
if isinstance(metadata, dict):
# Filter the metadata dictionary to include only the specified keys

View file

@ -1,4 +1,4 @@
from typing import Callable, List, Set, Type, Union
from typing import TYPE_CHECKING, Callable, List, Optional, Set, Type, Union
import litellm
from litellm._logging import verbose_logger
@ -6,6 +6,11 @@ from litellm.integrations.additional_logging_utils import AdditionalLoggingUtils
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import CallbacksByType
if TYPE_CHECKING:
from litellm import _custom_logger_compatible_callbacks_literal
else:
_custom_logger_compatible_callbacks_literal = str
class LoggingCallbackManager:
"""
@ -343,3 +348,26 @@ class LoggingCallbackManager:
elif callable(callback):
return getattr(callback, "__name__", str(callback))
return str(callback)
def get_active_custom_logger_for_callback_name(
self,
callback_name: _custom_logger_compatible_callbacks_literal,
) -> Optional[CustomLogger]:
"""
Get the active custom logger for a given callback name
"""
from litellm.litellm_core_utils.custom_logger_registry import (
CustomLoggerRegistry,
)
# get the custom logger class type
custom_logger_class_type = CustomLoggerRegistry.get_class_type_for_custom_logger_name(callback_name)
# get the active custom logger
custom_logger = self.get_custom_loggers_for_type(custom_logger_class_type)
if len(custom_logger) == 0:
raise ValueError(f"No active custom logger found for callback name: {callback_name}")
return custom_logger[0]

View file

@ -9,6 +9,7 @@ Docs - https://docs.mistral.ai/api/
from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, cast, overload
import httpx
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.prompt_templates.common_utils import (
handle_messages_with_content_list_to_str_conversion,
@ -147,7 +148,8 @@ class MistralConfig(OpenAIGPTConfig):
if param == "max_completion_tokens": # max_completion_tokens should take priority
optional_params["max_tokens"] = value
if param == "tools":
optional_params["tools"] = value
# Clean tools to remove problematic schema fields for Mistral API
optional_params["tools"] = self._clean_tool_schema_for_mistral(value)
if param == "stream" and value is True:
optional_params["stream"] = value
if param == "temperature":
@ -195,7 +197,8 @@ class MistralConfig(OpenAIGPTConfig):
@overload
def _transform_messages(
self, messages: List[AllMessageValues], model: str, is_async: Literal[True]
) -> Coroutine[Any, Any, List[AllMessageValues]]: ...
) -> Coroutine[Any, Any, List[AllMessageValues]]:
...
@overload
def _transform_messages(
@ -203,7 +206,8 @@ class MistralConfig(OpenAIGPTConfig):
messages: List[AllMessageValues],
model: str,
is_async: Literal[False] = False,
) -> List[AllMessageValues]: ...
) -> List[AllMessageValues]:
...
def _transform_messages(
self, messages: List[AllMessageValues], model: str, is_async: bool = False
@ -286,6 +290,38 @@ class MistralConfig(OpenAIGPTConfig):
optional_params.pop("_add_reasoning_prompt", None)
return messages
@classmethod
def _clean_tool_schema_for_mistral(cls, tools: list) -> list:
"""
Clean tool schemas to remove fields that cause issues with Mistral API.
Removes:
- $id and $schema fields (cause grammar validation errors)
- additionalProperties=False (causes OpenAI API schema errors)
- strict field (not supported by Mistral)
Args:
tools: List of tool definitions
max_depth: Maximum recursion depth for schema cleaning (default: 10)
Returns:
Cleaned tools list
"""
if not tools:
return tools
import copy
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.utils import _remove_json_schema_refs
cleaned_tools = copy.deepcopy(tools)
# Apply all cleaning functions with max_depth protection
cleaned_tools = _remove_json_schema_refs(cleaned_tools, max_depth=DEFAULT_MAX_RECURSE_DEPTH)
return cleaned_tools
@classmethod
def _handle_name_in_message(cls, message: AllMessageValues) -> AllMessageValues:
"""

View file

@ -612,6 +612,263 @@
"search_context_size_high": 0.03
}
},
"gpt-5": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 1.25e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 1.25e-07,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"gpt-5-mini": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 2.5e-07,
"output_cost_per_token": 2e-06,
"cache_read_input_token_cost": 2.5e-08,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"gpt-5-nano": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 5e-08,
"output_cost_per_token": 4e-07,
"cache_read_input_token_cost": 5e-09,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"gpt-5-chat": {
"max_tokens": 32768,
"max_input_tokens": 1047576,
"max_output_tokens": 32768,
"input_cost_per_token": 5e-06,
"output_cost_per_token": 2e-05,
"input_cost_per_token_batches": 2.5e-06,
"output_cost_per_token_batches": 1e-05,
"cache_read_input_token_cost": 1.25e-06,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true
},
"gpt-5-chat-latest": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 1.25e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 1.25e-07,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"gpt-5-2025-08-07": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 1.25e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 1.25e-07,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"gpt-5-mini-2025-08-07": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 2.5e-07,
"output_cost_per_token": 2e-06,
"cache_read_input_token_cost": 2.5e-08,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"gpt-5-nano-2025-08-07": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 5e-08,
"output_cost_per_token": 4e-07,
"cache_read_input_token_cost": 5e-09,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"codex-mini-latest": {
"max_tokens": 100000,
"max_input_tokens": 200000,
@ -2007,6 +2264,263 @@
"/v1/audio/speech"
]
},
"azure/gpt-5": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 1.25e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 1.25e-07,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"azure/gpt-5-2025-08-07": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 1.25e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 1.25e-07,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"azure/gpt-5-mini": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 2.5e-07,
"output_cost_per_token": 2e-06,
"cache_read_input_token_cost": 2.5e-08,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"azure/gpt-5-mini-2025-08-07": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 2.5e-07,
"output_cost_per_token": 2e-06,
"cache_read_input_token_cost": 2.5e-08,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"azure/gpt-5-nano-2025-08-07": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 5e-08,
"output_cost_per_token": 4e-07,
"cache_read_input_token_cost": 5e-09,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"azure/gpt-5-nano": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 5e-08,
"output_cost_per_token": 4e-07,
"cache_read_input_token_cost": 5e-09,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"azure/gpt-5-chat": {
"max_tokens": 32768,
"max_input_tokens": 1047576,
"max_output_tokens": 32768,
"input_cost_per_token": 5e-06,
"output_cost_per_token": 2e-05,
"input_cost_per_token_batches": 2.5e-06,
"output_cost_per_token_batches": 1e-05,
"cache_read_input_token_cost": 1.25e-06,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true
},
"azure/gpt-5-chat-latest": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 1.25e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 1.25e-07,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"azure/gpt-4o-mini-tts": {
"mode": "audio_speech",
"input_cost_per_token": 2.5e-06,
@ -5486,6 +6000,36 @@
"litellm_provider": "groq",
"mode": "audio_transcription"
},
"groq/openai/gpt-oss-20b": {
"max_tokens": 32768,
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"input_cost_per_token": 1e-07,
"output_cost_per_token": 5e-07,
"litellm_provider": "groq",
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_web_search": true
},
"groq/openai/gpt-oss-120b": {
"max_tokens": 32766,
"max_input_tokens": 131072,
"max_output_tokens": 32766,
"input_cost_per_token": 1.5e-07,
"output_cost_per_token": 7.5e-07,
"litellm_provider": "groq",
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_web_search": true
},
"cerebras/llama3.1-8b": {
"max_tokens": 128000,
"max_input_tokens": 128000,
@ -5741,6 +6285,32 @@
"supports_reasoning": true,
"supports_computer_use": true
},
"claude-opus-4-1": {
"max_tokens": 32000,
"max_input_tokens": 200000,
"max_output_tokens": 32000,
"input_cost_per_token": 1.5e-05,
"output_cost_per_token": 7.5e-05,
"search_context_cost_per_query": {
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01,
"search_context_size_high": 0.01
},
"cache_creation_input_token_cost": 1.875e-05,
"cache_read_input_token_cost": 1.5e-06,
"litellm_provider": "anthropic",
"mode": "chat",
"supports_function_calling": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159,
"supports_assistant_prefill": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_reasoning": true,
"supports_computer_use": true
},
"claude-opus-4-1-20250805": {
"max_tokens": 32000,
"max_input_tokens": 200000,
@ -17663,5 +18233,136 @@
"supports_vision": false,
"supports_system_messages": true,
"supports_tool_choice": false
},
"oci/meta.llama-4-maverick-17b-128e-instruct-fp8": {
"max_tokens": 512000,
"max_input_tokens": 512000,
"max_output_tokens": 4000,
"input_cost_per_token": 7.2e-07,
"output_cost_per_token": 7.2e-07,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/meta.llama-4-scout-17b-16e-instruct": {
"max_tokens": 192000,
"max_input_tokens": 192000,
"max_output_tokens": 4000,
"input_cost_per_token": 7.2e-07,
"output_cost_per_token": 7.2e-07,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/meta.llama-3.3-70b-instruct": {
"max_tokens": 128000,
"max_input_tokens": 128000,
"max_output_tokens": 4000,
"input_cost_per_token": 7.2e-07,
"output_cost_per_token": 7.2e-07,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/meta.llama-3.2-90b-vision-instruct": {
"max_tokens": 128000,
"max_input_tokens": 128000,
"max_output_tokens": 4000,
"input_cost_per_token": 2.0e-06,
"output_cost_per_token": 2.0e-06,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/meta.llama-3.1-405b-instruct": {
"max_tokens": 128000,
"max_input_tokens": 128000,
"max_output_tokens": 4000,
"input_cost_per_token": 1.068e-05,
"output_cost_per_token": 1.068e-05,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/xai.grok-4": {
"max_tokens": 128000,
"max_input_tokens": 128000,
"max_output_tokens": 128000,
"input_cost_per_token": 3.0e-06,
"output_cost_per_token": 1.5e-07,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/xai.grok-3": {
"max_tokens": 131072,
"max_input_tokens": 131072,
"max_output_tokens": 131072,
"input_cost_per_token": 3.0e-06,
"output_cost_per_token": 1.5e-07,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/xai.grok-3-mini": {
"max_tokens": 131072,
"max_input_tokens": 131072,
"max_output_tokens": 131072,
"input_cost_per_token": 3.0e-07,
"output_cost_per_token": 5.0e-07,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/xai.grok-3-fast": {
"max_tokens": 131072,
"max_input_tokens": 131072,
"max_output_tokens": 131072,
"input_cost_per_token": 5.0e-06,
"output_cost_per_token": 2.5e-05,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/xai.grok-3-mini-fast": {
"max_tokens": 131072,
"max_input_tokens": 131072,
"max_output_tokens": 131072,
"input_cost_per_token": 6.0e-07,
"output_cost_per_token": 4.0e-06,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
}
}

View file

@ -2210,7 +2210,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
braintrust: CallbackOnUI = CallbackOnUI(
litellm_callback_name="braintrust",
litellm_callback_params=["BRAINTRUST_API_KEY"],
litellm_callback_params=["BRAINTRUST_API_KEY","BRAINTRUST_API_BASE"],
ui_callback_name="Braintrust",
)
@ -2264,6 +2264,7 @@ class SpendLogsMetadata(TypedDict):
error_information: Optional[StandardLoggingPayloadErrorInformation]
usage_object: Optional[dict]
model_map_information: Optional[StandardLoggingModelInformation]
cold_storage_object_key: Optional[str] # S3/GCS object key for cold storage retrieval
class SpendLogsPayload(TypedDict):

View file

@ -5,6 +5,7 @@
# +-------------------------------------------------------------+
# Thank you users! We ❤️ you! - Krrish & Ishaan
import copy
import os
import sys
@ -50,6 +51,51 @@ from litellm.types.utils import (
GUARDRAIL_NAME = "bedrock"
def _redact_pii_matches(response_json: dict) -> dict:
try:
# Create a deep copy to avoid modifying the original response
redacted_response = copy.deepcopy(response_json)
# Get assessments from the response
assessments = redacted_response.get("assessments", [])
if not assessments:
return redacted_response
for assessment in assessments:
# Redact PII entities in sensitive information policy
sensitive_info_policy = assessment.get("sensitiveInformationPolicy")
if sensitive_info_policy:
pii_entities = sensitive_info_policy.get("piiEntities", [])
for pii_entity in pii_entities:
if "match" in pii_entity:
pii_entity["match"] = "[REDACTED]"
# Redact regex matches
regexes = sensitive_info_policy.get("regexes", [])
for regex_match in regexes:
if "match" in regex_match:
regex_match["match"] = "[REDACTED]"
# Redact custom word matches in word policy
word_policy = assessment.get("wordPolicy")
if word_policy:
custom_words = word_policy.get("customWords", [])
for custom_word in custom_words:
if "match" in custom_word:
custom_word["match"] = "[REDACTED]"
managed_words = word_policy.get("managedWordLists", [])
for managed_word in managed_words:
if "match" in managed_word:
managed_word["match"] = "[REDACTED]"
return redacted_response
except Exception as e:
# We do not want to fail in any case so this is just a warning
verbose_proxy_logger.warning("Guardrail log redaction failed: %s", str(e))
return response_json
class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
def __init__(
self,
@ -271,10 +317,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
data=prepared_request.body, # type: ignore
headers=prepared_request.headers, # type: ignore
)
verbose_proxy_logger.debug("Bedrock AI response: %s", response.text)
if response.status_code == 200:
# check if the response was flagged
_json_response = response.json()
redacted_response = _redact_pii_matches(_json_response)
verbose_proxy_logger.debug("Bedrock AI response : %s", redacted_response)
bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response)
if self._should_raise_guardrail_blocked_exception(
bedrock_guardrail_response

View file

@ -51,6 +51,7 @@ from litellm.proxy.common_utils.admin_ui_utils import (
from litellm.proxy.common_utils.html_forms.jwt_display_template import (
jwt_display_template,
)
from litellm.proxy.common_utils.html_forms.ui_login import html_form
from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
from litellm.proxy.management_endpoints.sso_helper_utils import (
check_is_admin_only_access,
@ -76,20 +77,16 @@ router = APIRouter()
@router.get("/sso/key/generate", tags=["experimental"], include_in_schema=False)
async def serve_login_page(
request: Request,
source: Optional[str] = None,
key: Optional[str] = None,
error: Optional[str] = None,
):
async def google_login(request: Request, source: Optional[str] = None, key: Optional[str] = None): # noqa: PLR0915
"""
Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env
PROXY_BASE_URL should be the your deployed proxy endpoint, e.g. PROXY_BASE_URL="https://litellm-production-7002.up.railway.app/"
Example:
Serves a unified login page with options for both normal
username/password login and SSO.
"""
from litellm.proxy.proxy_server import premium_user
from litellm.proxy.proxy_server import (
premium_user,
user_custom_ui_sso_sign_in_handler,
)
microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None)
google_client_id = os.getenv("GOOGLE_CLIENT_ID", None)
@ -102,334 +99,6 @@ async def serve_login_page(
if is_disabled:
return admin_ui_disabled()
####### Check if user is a Enterprise / Premium User for SSO #######
sso_available = False
if (
microsoft_client_id is not None
or google_client_id is not None
or generic_client_id is not None
):
if premium_user is True:
sso_available = True
####### Detect DB + MASTER KEY in .env #######
missing_env_vars = show_missing_vars_in_env()
if missing_env_vars is not None:
return missing_env_vars
#########################################################
# Construct Redirect URL
base_url_to_redirect_to: Optional[str] = None
base_url_to_redirect_to = os.getenv("PROXY_BASE_URL", "")
server_root_path = os.getenv("SERVER_ROOT_PATH", "")
if server_root_path != "":
base_url_to_redirect_to += server_root_path
#########################################################
# Build the unified login page HTML
error_message = ""
if error == "1":
error_message = """
<div style="
background-color: #fef2f2;
border-left: 4px solid #dc2626;
border-radius: 6px;
padding: 16px;
margin-bottom: 20px;
color: #dc2626;
font-size: 14px;
font-weight: 500;
">
⚠️ Invalid username or password. Please try again.
</div>
"""
sso_button = ""
if sso_available:
sso_login_url = base_url_to_redirect_to
if sso_login_url.endswith("/"):
sso_login_url += "sso/login"
else:
sso_login_url += "/sso/login"
sso_button = f"""
<div style="
margin-top: 20px;
padding-top: 20px;
border-top: 1px solid #e2e8f0;
text-align: center;
">
<p style="
color: #64748b;
font-size: 14px;
margin-bottom: 16px;
">or</p>
<a href="{sso_login_url}" style="
display: inline-block;
background-color: #f8fafc;
border: 1px solid #e2e8f0;
color: #374151;
padding: 10px 20px;
border-radius: 6px;
text-decoration: none;
font-weight: 500;
transition: all 0.2s;
font-size: 14px;
" onmouseover="this.style.backgroundColor='#f1f5f9'; this.style.borderColor='#cbd5e1';"
onmouseout="this.style.backgroundColor='#f8fafc'; this.style.borderColor='#e2e8f0';">
🔐 Login with SSO
</a>
</div>
"""
if base_url_to_redirect_to.endswith("/"):
url_to_redirect_to = base_url_to_redirect_to + "login"
else:
url_to_redirect_to = base_url_to_redirect_to + "/login"
unified_login_html = f"""
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>LiteLLM Login</title>
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<style>
body {{
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, sans-serif;
background-color: #f8fafc;
margin: 0;
padding: 20px;
display: flex;
justify-content: center;
align-items: center;
min-height: 100vh;
color: #333;
}}
form {{
background-color: #fff;
padding: 40px;
border-radius: 8px;
box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
width: 450px;
max-width: 100%;
}}
.logo-container {{
text-align: center;
margin-bottom: 30px;
}}
.logo {{
font-size: 24px;
font-weight: 600;
color: #1e293b;
}}
h2 {{
margin: 0 0 10px;
color: #1e293b;
font-size: 28px;
font-weight: 600;
text-align: center;
}}
.subtitle {{
color: #64748b;
margin: 0 0 20px;
font-size: 16px;
text-align: center;
}}
.info-box {{
background-color: #f1f5f9;
border-radius: 6px;
padding: 20px;
margin-bottom: 30px;
border-left: 4px solid #2563eb;
}}
.info-header {{
display: flex;
align-items: center;
margin-bottom: 12px;
color: #1e40af;
font-weight: 600;
font-size: 16px;
}}
.info-header svg {{
margin-right: 8px;
}}
.info-box p {{
color: #475569;
margin: 8px 0;
line-height: 1.5;
font-size: 14px;
}}
label {{
display: block;
margin-bottom: 8px;
font-weight: 500;
color: #334155;
font-size: 14px;
}}
.required {{
color: #dc2626;
margin-left: 2px;
}}
input[type="text"],
input[type="password"] {{
width: 100%;
padding: 10px 14px;
margin-bottom: 20px;
box-sizing: border-box;
border: 1px solid #e2e8f0;
border-radius: 6px;
font-size: 15px;
color: #1e293b;
background-color: #fff;
transition: border-color 0.2s, box-shadow 0.2s;
}}
input[type="text"]:focus,
input[type="password"]:focus {{
outline: none;
border-color: #3b82f6;
box-shadow: 0 0 0 2px rgba(59, 130, 246, 0.2);
}}
.toggle-password {{
display: flex;
align-items: center;
margin-top: -15px;
margin-bottom: 20px;
}}
.toggle-password input[type="checkbox"] {{
margin-right: 8px;
vertical-align: middle;
width: 16px;
height: 16px;
}}
.toggle-password label {{
margin-bottom: 0;
font-size: 14px;
cursor: pointer;
line-height: 1;
}}
input[type="submit"] {{
background-color: #6466E9;
color: #fff;
cursor: pointer;
font-weight: 500;
border: none;
padding: 10px 16px;
transition: background-color 0.2s;
border-radius: 6px;
margin-top: 10px;
font-size: 14px;
width: 100%;
}}
input[type="submit"]:hover {{
background-color: #4138C2;
}}
a {{
color: #3b82f6;
text-decoration: none;
}}
a:hover {{
text-decoration: underline;
}}
code {{
background-color: #f1f5f9;
padding: 2px 4px;
border-radius: 4px;
font-family: monospace;
font-size: 13px;
color: #334155;
}}
</style>
</head>
<body>
<form action="{url_to_redirect_to}" method="post">
<div class="logo-container">
<div class="logo">
🚅 LiteLLM
</div>
</div>
<h2>Login</h2>
<p class="subtitle">Access your LiteLLM Admin UI.</p>
{error_message}
<div class="info-box">
<div class="info-header">
<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
<circle cx="12" cy="12" r="10"></circle>
<line x1="12" y1="16" x2="12" y2="12"></line>
<line x1="12" y1="8" x2="12.01" y2="8"></line>
</svg>
Default Credentials
</div>
<p>By default, Username is <code>admin</code> and Password is your set LiteLLM Proxy <code>MASTER_KEY</code>.</p>
<p>Need to set UI credentials or SSO? <a href="https://docs.litellm.ai/docs/proxy/ui" target="_blank">Check the documentation</a>.</p>
</div>
<label for="username">Username<span class="required">*</span></label>
<input type="text" id="username" name="username" required placeholder="Enter your username" autocomplete="username">
<label for="password">Password<span class="required">*</span></label>
<input type="password" id="password" name="password" required placeholder="Enter your password" autocomplete="current-password">
<div class="toggle-password">
<input type="checkbox" id="show-password" onclick="togglePasswordVisibility()">
<label for="show-password">Show password</label>
</div>
<input type="submit" value="Login">
{sso_button}
</form>
<script>
function togglePasswordVisibility() {{
var passwordField = document.getElementById("password");
passwordField.type = passwordField.type === "password" ? "text" : "password";
}}
</script>
</body>
</html>
"""
from fastapi.responses import HTMLResponse
return HTMLResponse(content=unified_login_html, status_code=200)
@router.get("/sso/login", tags=["experimental"], include_in_schema=False)
async def sso_login_redirect(
request: Request, source: Optional[str] = None, key: Optional[str] = None
):
"""
Handles SSO login redirect - this is what the "Login with SSO" button points to
"""
from litellm.proxy.proxy_server import (
premium_user,
user_custom_ui_sso_sign_in_handler,
)
microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None)
google_client_id = os.getenv("GOOGLE_CLIENT_ID", None)
generic_client_id = os.getenv("GENERIC_CLIENT_ID", None)
####### Check if user is a Enterprise / Premium User #######
if (
microsoft_client_id is not None
@ -444,12 +113,18 @@ async def sso_login_redirect(
code=status.HTTP_403_FORBIDDEN,
)
####### Detect DB + MASTER KEY in .env #######
missing_env_vars = show_missing_vars_in_env()
if missing_env_vars is not None:
return missing_env_vars
ui_username = os.getenv("UI_USERNAME")
# get url from request - always use regular callback, but set state for CLI
redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso(
request=request,
sso_callback_route="sso/callback",
)
# Store CLI key in state for OAuth flow
cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state(
source=source,
@ -462,14 +137,11 @@ async def sso_login_redirect(
from litellm_enterprise.proxy.auth.custom_sso_handler import (
EnterpriseCustomSSOHandler,
)
return await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in(
request=request,
)
except ImportError:
raise ValueError(
"Enterprise features are not available. Custom UI SSO sign-in requires LiteLLM Enterprise."
)
raise ValueError("Enterprise features are not available. Custom UI SSO sign-in requires LiteLLM Enterprise.")
# Check if we should use SSO handler
if (
@ -488,9 +160,16 @@ async def sso_login_redirect(
generic_client_id=generic_client_id,
state=cli_state,
)
elif ui_username is not None:
# No Google, Microsoft SSO
# Use UI Credentials set in .env
from fastapi.responses import HTMLResponse
return HTMLResponse(content=html_form, status_code=200)
else:
# No SSO configured, redirect back to login page
return RedirectResponse(url="/sso/key/generate", status_code=303)
from fastapi.responses import HTMLResponse
return HTMLResponse(content=html_form, status_code=200)
def generic_response_convertor(
@ -846,16 +525,15 @@ async def check_and_update_if_proxy_admin_id(
async def auth_callback(request: Request, state: Optional[str] = None): # noqa: PLR0915
"""Verify login"""
verbose_proxy_logger.info(f"Starting SSO callback with state: {state}")
# Check if this is a CLI login (state starts with our CLI prefix)
from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX
if state and state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:"):
# Extract the key ID from the state
key_id = state.split(":", 1)[1]
verbose_proxy_logger.info(f"CLI SSO callback detected for key: {key_id}")
return await cli_sso_callback(request, key=key_id)
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.proxy_server import (
@ -930,7 +608,7 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa:
status_code=401,
detail="Result not returned by SSO provider.",
)
return await SSOAuthenticationHandler.get_redirect_response_from_openid(
result=result,
request=request,
@ -940,26 +618,28 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa:
)
async def cli_sso_callback(request: Request, key: Optional[str] = None):
"""CLI SSO callback - generates the key with pre-specified ID"""
verbose_proxy_logger.info(f"CLI SSO callback for key: {key}")
from litellm.proxy.management_endpoints.key_management_endpoints import (
generate_key_helper_fn,
)
from litellm.proxy.proxy_server import prisma_client
if not key or not key.startswith("sk-"):
if not key or not key.startswith('sk-'):
raise HTTPException(
status_code=400,
detail="Invalid key parameter. Must be a valid key ID starting with 'sk-'",
detail="Invalid key parameter. Must be a valid key ID starting with 'sk-'"
)
if prisma_client is None:
raise HTTPException(
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
)
# Generate a simple key for CLI usage with the pre-specified key ID
try:
await generate_key_helper_fn(
@ -973,57 +653,63 @@ async def cli_sso_callback(request: Request, key: Optional[str] = None):
table_name="key",
token=key, # Use the pre-specified key ID
)
verbose_proxy_logger.info(f"Generated CLI key: {key}")
# Return success page
from fastapi.responses import HTMLResponse
from litellm.proxy.common_utils.html_forms.cli_sso_success import (
render_cli_sso_success_page,
)
html_content = render_cli_sso_success_page()
return HTMLResponse(content=html_content, status_code=200)
except Exception as e:
verbose_proxy_logger.error(f"Error generating CLI key: {e}")
raise HTTPException(status_code=500, detail=f"Failed to generate key: {str(e)}")
raise HTTPException(
status_code=500,
detail=f"Failed to generate key: {str(e)}"
)
@router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False)
async def cli_poll_key(key_id: str):
"""CLI polling endpoint - checks if key exists in DB"""
from litellm.proxy.proxy_server import prisma_client
if not key_id.startswith("sk-"):
raise HTTPException(status_code=400, detail="Invalid key ID format")
if not key_id.startswith('sk-'):
raise HTTPException(
status_code=400,
detail="Invalid key ID format"
)
if prisma_client is None:
raise HTTPException(
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
)
try:
# Check if key exists in database
from litellm.proxy.utils import hash_token
hashed_token = hash_token(key_id)
key_obj = await prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": hashed_token}
)
if key_obj:
verbose_proxy_logger.info(f"CLI key found: {key_id}")
return {"status": "ready", "key": key_id}
else:
return {"status": "pending"}
except Exception as e:
verbose_proxy_logger.error(f"Error polling for CLI key: {e}")
raise HTTPException(
status_code=500, detail=f"Error checking key status: {str(e)}"
status_code=500,
detail=f"Error checking key status: {str(e)}"
)
@ -1125,7 +811,6 @@ class SSOAuthenticationHandler:
"""
Handler for SSO Authentication across all SSO providers
"""
@staticmethod
async def get_sso_login_redirect(
redirect_url: str,
@ -1478,6 +1163,7 @@ class SSOAuthenticationHandler:
_new_team_request.update(_default_team_params)
team_request = NewTeamRequest(**_new_team_request)
return team_request
@staticmethod
def _get_cli_state(source: Optional[str], key: Optional[str]) -> Optional[str]:
@ -1490,15 +1176,13 @@ class SSOAuthenticationHandler:
LITELLM_CLI_SESSION_TOKEN_PREFIX,
LITELLM_CLI_SOURCE_IDENTIFIER,
)
return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}" if source == LITELLM_CLI_SOURCE_IDENTIFIER and key else None
return (
f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}"
if source == LITELLM_CLI_SOURCE_IDENTIFIER and key
else None
)
@staticmethod
async def get_redirect_response_from_openid( # noqa: PLR0915
async def get_redirect_response_from_openid( # noqa: PLR0915
result: Union[OpenID, dict, CustomOpenID],
request: Request,
received_response: Optional[dict] = None,
@ -1518,18 +1202,14 @@ class SSOAuthenticationHandler:
)
from litellm.proxy.utils import get_prisma_client_or_throw
from litellm.types.proxy.ui_sso import ReturnedUITokenObject
prisma_client = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy")
prisma_client = get_prisma_client_or_throw(
"Prisma client is None, connect a database to your proxy"
)
# User is Authe'd in - generate key for the UI to access Proxy
verbose_proxy_logger.info(f"SSO callback result: {result}")
user_email: Optional[str] = getattr(result, "email", None)
user_id: Optional[str] = (
getattr(result, "id", None) if result is not None else None
)
user_id: Optional[str] = getattr(result, "id", None) if result is not None else None
if user_email is not None and os.getenv("ALLOWED_EMAIL_DOMAINS") is not None:
email_domain = user_email.split("@")[1]
@ -1714,8 +1394,7 @@ class SSOAuthenticationHandler:
redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
redirect_response.set_cookie(key="token", value=jwt_token)
return redirect_response
class MicrosoftSSOHandler:
"""
Handles Microsoft SSO callback response and returns a CustomOpenID object
@ -2215,28 +1894,3 @@ async def debug_sso_callback(request: Request):
)
return HTMLResponse(content=html_content)
@router.post("/sso/key/generate", tags=["experimental"], include_in_schema=False)
async def process_login(request: Request):
"""
Process username/password login from the unified login page
"""
try:
# Get form data
form_data = await request.form()
username = form_data.get("username")
password = form_data.get("password")
if not username or not password:
return RedirectResponse(url="/sso/key/generate?error=1", status_code=303)
# Import the actual login function from proxy_server
from litellm.proxy.proxy_server import login
# Call the real login function that handles all the authentication properly
return await login(request)
except Exception as e:
verbose_proxy_logger.error(f"Error processing login: {e}")
return RedirectResponse(url="/sso/key/generate?error=1", status_code=303)

View file

@ -1,6 +1,15 @@
model_list:
- model_name: anthropic/*
- model_name: bedrock/*
litellm_params:
model: anthropic/*
model: bedrock/*
litellm_settings:
callbacks: ["s3_v2"]
s3_callback_params:
s3_bucket_name: litellm-logs # AWS Bucket Name for S3
s3_region_name: us-west-2
general_settings:
cold_storage_custom_logger: s3_v2
store_prompts_in_cold_storage: true

View file

@ -0,0 +1,73 @@
"""
This module is responsible for handling Getting/Setting the proxy server request from cold storage.
It allows fetching a dict of the proxy server request from s3 or GCS bucket.
"""
from typing import Optional, cast
import litellm
from litellm import _custom_logger_compatible_callbacks_literal
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
class ColdStorageHandler:
"""
This class is responsible for handling Getting/Setting the proxy server request from cold storage.
It allows fetching a dict of the proxy server request from s3 or GCS bucket.
"""
async def get_proxy_server_request_from_cold_storage_with_object_key(
self,
object_key: str,
) -> Optional[dict]:
"""
Get the proxy server request from cold storage using the object key directly.
Args:
object_key: The S3/GCS object key to retrieve
Returns:
Optional[dict]: The proxy server request dict or None if not found
"""
# select the custom logger to use for cold storage
custom_logger_name: Optional[_custom_logger_compatible_callbacks_literal] = self._select_custom_logger_for_cold_storage()
# if no custom logger name is configured, return None
if custom_logger_name is None:
return None
# get the active/initialized custom logger
custom_logger: Optional[CustomLogger] = litellm.logging_callback_manager.get_active_custom_logger_for_callback_name(custom_logger_name)
# if no custom logger is found, return None
if custom_logger is None:
return None
proxy_server_request = await custom_logger.get_proxy_server_request_from_cold_storage_with_object_key(
object_key=object_key,
)
return proxy_server_request
def _select_custom_logger_for_cold_storage(
self,
) -> Optional[_custom_logger_compatible_callbacks_literal]:
cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = ColdStorageHandler._get_configured_cold_storage_custom_logger()
return cold_storage_custom_logger
@staticmethod
def _get_configured_cold_storage_custom_logger() -> Optional[_custom_logger_compatible_callbacks_literal]:
from litellm.proxy.proxy_server import general_settings
cold_storage_custom_logger: Optional[str] = general_settings.get("cold_storage_custom_logger")
if not cold_storage_custom_logger:
verbose_proxy_logger.debug("No cold storage custom logger found in general settings")
return None
return cast(_custom_logger_compatible_callbacks_literal, cold_storage_custom_logger)

View file

@ -53,6 +53,7 @@ def _get_spend_logs_metadata(
guardrail_information: Optional[StandardLoggingGuardrailInformation] = None,
usage_object: Optional[dict] = None,
model_map_information: Optional[StandardLoggingModelInformation] = None,
cold_storage_object_key: Optional[str] = None
) -> SpendLogsMetadata:
if metadata is None:
return SpendLogsMetadata(
@ -75,6 +76,7 @@ def _get_spend_logs_metadata(
model_map_information=None,
usage_object=None,
guardrail_information=None,
cold_storage_object_key=cold_storage_object_key,
)
verbose_proxy_logger.debug(
"getting payload for SpendLogs, available keys in metadata: "
@ -98,6 +100,8 @@ def _get_spend_logs_metadata(
clean_metadata["guardrail_information"] = guardrail_information
clean_metadata["usage_object"] = usage_object
clean_metadata["model_map_information"] = model_map_information
clean_metadata["cold_storage_object_key"] = cold_storage_object_key
return clean_metadata
@ -267,6 +271,11 @@ def get_logging_payload( # noqa: PLR0915
if standard_logging_payload is not None
else None
),
cold_storage_object_key=(
standard_logging_payload["metadata"].get("cold_storage_object_key", None)
if standard_logging_payload is not None
else None
),
)
special_usage_fields = ["completion_tokens", "prompt_tokens", "total_tokens"]
@ -474,6 +483,7 @@ def _sanitize_request_body_for_spend_logs_payload(
Recursively sanitize request body to prevent logging large base64 strings or other large values.
Truncates strings longer than 1000 characters and handles nested dictionaries.
"""
from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD
MAX_STRING_LENGTH = 1000
if visited is None:
@ -492,7 +502,7 @@ def _sanitize_request_body_for_spend_logs_payload(
return [_sanitize_value(item) for item in value]
elif isinstance(value, str):
if len(value) > MAX_STRING_LENGTH:
return f"{value[:MAX_STRING_LENGTH]}... (truncated {len(value) - MAX_STRING_LENGTH} chars)"
return f"{value[:MAX_STRING_LENGTH]}... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} {len(value) - MAX_STRING_LENGTH} chars)"
return value
return value

View file

@ -0,0 +1,301 @@
import json
from typing import TYPE_CHECKING, Any, List, Optional, Union, cast
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import SpendLogsPayload
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionResponseMessage,
GenericChatCompletionMessage,
ResponseInputParam,
)
from litellm.types.utils import ChatCompletionMessageToolCall, Message, ModelResponse
if TYPE_CHECKING:
from litellm.responses.litellm_completion_transformation.transformation import (
ChatCompletionSession,
)
else:
ChatCompletionSession = Any
########################################################
# Cold Storage Handler
########################################################
COLD_STORAGE_HANDLER = ColdStorageHandler()
########################################################
class ResponsesSessionHandler:
@staticmethod
async def get_chat_completion_message_history_for_previous_response_id(
previous_response_id: str,
) -> ChatCompletionSession:
"""
Return the chat completion message history for a previous response id
"""
from litellm.responses.litellm_completion_transformation.transformation import (
ChatCompletionSession,
)
verbose_proxy_logger.debug(
"inside get_chat_completion_message_history_for_previous_response_id"
)
all_spend_logs: List[
SpendLogsPayload
] = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(
previous_response_id
)
verbose_proxy_logger.debug(
"found %s spend logs for this response id", len(all_spend_logs)
)
litellm_session_id: Optional[str] = None
if len(all_spend_logs) > 0:
litellm_session_id = all_spend_logs[0].get("session_id")
chat_completion_message_history: List[
Union[
AllMessageValues,
GenericChatCompletionMessage,
ChatCompletionMessageToolCall,
ChatCompletionResponseMessage,
Message,
]
] = []
for spend_log in all_spend_logs:
chat_completion_message_history = await ResponsesSessionHandler.extend_chat_completion_message_with_spend_log_payload(
spend_log=spend_log,
chat_completion_message_history=chat_completion_message_history,
)
verbose_proxy_logger.debug(
"chat_completion_message_history %s",
json.dumps(chat_completion_message_history, indent=4, default=str),
)
return ChatCompletionSession(
messages=chat_completion_message_history,
litellm_session_id=litellm_session_id,
)
@staticmethod
async def extend_chat_completion_message_with_spend_log_payload(
spend_log: SpendLogsPayload,
chat_completion_message_history: List[
Union[
AllMessageValues,
GenericChatCompletionMessage,
ChatCompletionMessageToolCall,
ChatCompletionResponseMessage,
Message,
]
]
):
"""
Extend the chat completion message history with the spend log payload
"""
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
proxy_server_request_dict = await ResponsesSessionHandler.get_proxy_server_request_from_spend_log(
spend_log=spend_log,
)
response_input_param: Optional[Union[str, ResponseInputParam]] = None
_messages: Optional[Union[str, ResponseInputParam]] = None
############################################################
# Add Input messages for this Spend Log
############################################################
if proxy_server_request_dict:
_response_input_param = proxy_server_request_dict.get("input", None)
_messages = proxy_server_request_dict.get("messages", None)
if isinstance(_response_input_param, str):
response_input_param = _response_input_param
elif isinstance(_response_input_param, dict):
response_input_param = cast(
ResponseInputParam, _response_input_param
)
if response_input_param:
chat_completion_messages = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
input=response_input_param,
responses_api_request=proxy_server_request_dict or {},
)
chat_completion_message_history.extend(chat_completion_messages)
############################################################
# Check if `messages` field is present in the proxy server request dict
############################################################
elif _messages:
# ensure all messages are /chat/completions/messages
# certain requests can be stored as Responses API format - this ensures they are transformed to /chat/completions/messages
chat_completion_messages = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
input=_messages,
responses_api_request=proxy_server_request_dict or {},
)
chat_completion_message_history.extend(chat_completion_messages)
############################################################
# Add Output messages for this Spend Log
############################################################
_response_output = spend_log.get("response", "{}")
if isinstance(_response_output, dict):
# transform `ChatCompletion Response` to `ResponsesAPIResponse`
model_response = ModelResponse(**_response_output)
for choice in model_response.choices:
if hasattr(choice, "message"):
chat_completion_message_history.append(
getattr(choice, "message")
)
return chat_completion_message_history
@staticmethod
async def get_proxy_server_request_from_spend_log(
spend_log: SpendLogsPayload,
) -> Optional[dict]:
"""
Get the parsed proxy server request from the spend log
"""
proxy_server_request: Union[str, dict] = (
spend_log.get("proxy_server_request") or "{}"
)
proxy_server_request_dict: Optional[dict] = None
if isinstance(proxy_server_request, dict):
proxy_server_request_dict = proxy_server_request
else:
proxy_server_request_dict = json.loads(proxy_server_request)
############################################################
# Check if user has setup cold storage for session handling
############################################################
if ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_server_request_dict):
# Try to get cold storage object key from spend log metadata
_proxy_server_request_dict: Optional[dict] = None
cold_storage_object_key = ResponsesSessionHandler._get_cold_storage_object_key_from_spend_log(spend_log)
if cold_storage_object_key:
# Use the object key directly from metadata
_proxy_server_request_dict = await ResponsesSessionHandler.get_proxy_server_request_from_cold_storage_with_object_key(
object_key=cold_storage_object_key,
)
if _proxy_server_request_dict:
proxy_server_request_dict = _proxy_server_request_dict
return proxy_server_request_dict
@staticmethod
def _get_cold_storage_object_key_from_spend_log(spend_log: SpendLogsPayload) -> Optional[str]:
"""
Extract the cold storage object key from spend log metadata.
Args:
spend_log: The spend log payload containing metadata
Returns:
Optional[str]: The cold storage object key if found, None otherwise
"""
try:
metadata_str = spend_log.get("metadata", "{}")
if isinstance(metadata_str, str):
metadata_dict = json.loads(metadata_str)
return metadata_dict.get("cold_storage_object_key")
elif isinstance(metadata_str, dict):
return metadata_str.get("cold_storage_object_key")
return None
except (json.JSONDecodeError, TypeError, AttributeError):
verbose_proxy_logger.debug("Failed to parse metadata from spend log to extract cold storage object key")
return None
@staticmethod
async def get_proxy_server_request_from_cold_storage_with_object_key(
object_key: str,
) -> Optional[dict]:
"""
Get the proxy server request from cold storage using the object key directly.
Args:
object_key: The S3/GCS object key to retrieve
Returns:
Optional[dict]: The proxy server request dict or None if not found
"""
verbose_proxy_logger.debug("inside get_proxy_server_request_from_cold_storage_with_object_key...")
proxy_server_request_dict = await COLD_STORAGE_HANDLER.get_proxy_server_request_from_cold_storage_with_object_key(
object_key=object_key,
)
return proxy_server_request_dict
@staticmethod
def _should_check_cold_storage_for_full_payload(
proxy_server_request_dict: Optional[dict],
) -> bool:
"""
Only check cold storage when both are true
1. `LITELLM_TRUNCATED_PAYLOAD_FIELD` is in the proxy server request dict
2. `ColdStorageHandler._get_configured_cold_storage_custom_logger()` is not None
"""
from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD
configured_cold_storage_custom_logger = ColdStorageHandler._get_configured_cold_storage_custom_logger()
if configured_cold_storage_custom_logger is None:
return False
if proxy_server_request_dict is None:
return True
if len(proxy_server_request_dict) == 0:
return True
if LITELLM_TRUNCATED_PAYLOAD_FIELD in proxy_server_request_dict:
return True
return False
@staticmethod
async def get_all_spend_logs_for_previous_response_id(
previous_response_id: str,
) -> List[SpendLogsPayload]:
"""
Get all spend logs for a previous response id
SQL query
SELECT session_id FROM spend_logs WHERE response_id = previous_response_id, SELECT * FROM spend_logs WHERE session_id = session_id
"""
from litellm.proxy.proxy_server import prisma_client
verbose_proxy_logger.debug("decoding response id=%s", previous_response_id)
decoded_response_id = (
ResponsesAPIRequestUtils._decode_responses_api_response_id(
previous_response_id
)
)
previous_response_id = decoded_response_id.get(
"response_id", previous_response_id
)
if prisma_client is None:
return []
query = """
WITH matching_session AS (
SELECT session_id
FROM "LiteLLM_SpendLogs"
WHERE request_id = $1
)
SELECT *
FROM "LiteLLM_SpendLogs"
WHERE session_id IN (SELECT session_id FROM matching_session)
ORDER BY "endTime" ASC;
"""
spend_logs = await prisma_client.db.query_raw(query, previous_response_id)
verbose_proxy_logger.debug(
"Found the following spend logs for previous response id %s: %s",
previous_response_id,
json.dumps(spend_logs, indent=4, default=str),
)
return spend_logs

View file

@ -7,21 +7,14 @@ from typing import Any, Dict, List, Literal, Optional, Tuple, Union, cast
from openai.types.responses.tool_param import FunctionToolParam
from typing_extensions import TypedDict
from litellm._logging import verbose_logger
try:
from litellm_enterprise.enterprise_callbacks.session_handler import (
_ENTERPRISE_ResponsesSessionHandler,
)
except Exception as e:
verbose_logger.debug(
f"[Non-Blocking] Unable to import _ENTERPRISE_ResponsesSessionHandler - LiteLLM Enterprise Feature - {str(e)}"
)
_ENTERPRISE_ResponsesSessionHandler = None
from litellm.caching import InMemoryCache
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.responses.litellm_completion_transformation.session_handler import (
ResponsesSessionHandler,
)
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionImageObject,
ChatCompletionImageUrlObject,
ChatCompletionResponseMessage,
ChatCompletionSystemMessage,
@ -75,13 +68,6 @@ class ChatCompletionSession(TypedDict, total=False):
litellm_session_id: Optional[str]
class ChatCompletionImageItem(TypedDict):
"""TypedDict for image items in chat completion content"""
type: Literal["image"]
image_url: ChatCompletionImageUrlObject
########### End of Initialize Classes used for Responses API ###########
@ -210,20 +196,19 @@ class LiteLLMCompletionResponsesConfig:
"""
Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages
"""
if _ENTERPRISE_ResponsesSessionHandler is not None:
chat_completion_session = ChatCompletionSession(
messages=[], litellm_session_id=None
chat_completion_session = ChatCompletionSession(
messages=[], litellm_session_id=None
)
if previous_response_id:
chat_completion_session = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
previous_response_id=previous_response_id
)
if previous_response_id:
chat_completion_session = await _ENTERPRISE_ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
previous_response_id=previous_response_id
)
_messages = litellm_completion_request.get("messages") or []
session_messages = chat_completion_session.get("messages") or []
litellm_completion_request["messages"] = session_messages + _messages
litellm_completion_request[
"litellm_trace_id"
] = chat_completion_session.get("litellm_session_id")
_messages = litellm_completion_request.get("messages") or []
session_messages = chat_completion_session.get("messages") or []
litellm_completion_request["messages"] = session_messages + _messages
litellm_completion_request[
"litellm_trace_id"
] = chat_completion_session.get("litellm_session_id")
return litellm_completion_request
@staticmethod
@ -485,7 +470,7 @@ class LiteLLMCompletionResponsesConfig:
return new_item
@staticmethod
def _transform_input_image_item_to_image_item(item: Dict[str, Any]) -> ChatCompletionImageItem:
def _transform_input_image_item_to_image_item(item: Dict[str, Any]) -> ChatCompletionImageObject:
"""
Transform a Responses API input_image item to a Chat Completion image item
"""
@ -494,8 +479,8 @@ class LiteLLMCompletionResponsesConfig:
detail=item.get("detail") or "auto"
)
return ChatCompletionImageItem(
type="image",
return ChatCompletionImageObject(
type="image_url",
image_url=image_url_obj
)
@ -506,7 +491,6 @@ class LiteLLMCompletionResponsesConfig:
"""
Transform a Responses API content into a Chat Completion content
"""
if isinstance(content, str):
return content
elif isinstance(content, list):

View file

@ -1910,6 +1910,7 @@ class StandardLoggingMetadata(StandardLoggingUserAPIKeyMetadata):
vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]]
applied_guardrails: Optional[List[str]]
usage_object: Optional[dict]
cold_storage_object_key: Optional[str] # S3/GCS object key for cold storage retrieval
class StandardLoggingAdditionalHeaders(TypedDict, total=False):

View file

@ -2912,6 +2912,39 @@ def _remove_strict_from_schema(schema):
return schema
def _remove_json_schema_refs(schema, max_depth=10):
"""
Remove JSON schema reference fields like '$id' and '$schema' that can cause issues with some providers.
These fields are used for schema validation but can cause problems when the schema references
are not accessible to the provider's validation system.
Args:
schema: The schema object to clean (dict, list, or other)
max_depth: Maximum recursion depth to prevent infinite loops (default: 10)
Relevant Issues: Mistral API grammar validation fails when schema contains $id and $schema references
"""
if max_depth <= 0:
return schema
if isinstance(schema, dict):
# Remove JSON schema reference fields
schema.pop("$id", None)
schema.pop("$schema", None)
# Recursively process all dictionary values
for key, value in schema.items():
_remove_json_schema_refs(value, max_depth - 1)
elif isinstance(schema, list):
# Recursively process all items in the list
for item in schema:
_remove_json_schema_refs(item, max_depth - 1)
return schema
def _remove_unsupported_params(
non_default_params: dict, supported_openai_params: Optional[List[str]]
) -> dict:

View file

@ -612,6 +612,263 @@
"search_context_size_high": 0.03
}
},
"gpt-5": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 1.25e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 1.25e-07,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"gpt-5-mini": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 2.5e-07,
"output_cost_per_token": 2e-06,
"cache_read_input_token_cost": 2.5e-08,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"gpt-5-nano": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 5e-08,
"output_cost_per_token": 4e-07,
"cache_read_input_token_cost": 5e-09,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"gpt-5-chat": {
"max_tokens": 32768,
"max_input_tokens": 1047576,
"max_output_tokens": 32768,
"input_cost_per_token": 5e-06,
"output_cost_per_token": 2e-05,
"input_cost_per_token_batches": 2.5e-06,
"output_cost_per_token_batches": 1e-05,
"cache_read_input_token_cost": 1.25e-06,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true
},
"gpt-5-chat-latest": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 1.25e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 1.25e-07,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"gpt-5-2025-08-07": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 1.25e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 1.25e-07,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"gpt-5-mini-2025-08-07": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 2.5e-07,
"output_cost_per_token": 2e-06,
"cache_read_input_token_cost": 2.5e-08,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"gpt-5-nano-2025-08-07": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 5e-08,
"output_cost_per_token": 4e-07,
"cache_read_input_token_cost": 5e-09,
"litellm_provider": "openai",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"codex-mini-latest": {
"max_tokens": 100000,
"max_input_tokens": 200000,
@ -2007,6 +2264,263 @@
"/v1/audio/speech"
]
},
"azure/gpt-5": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 1.25e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 1.25e-07,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"azure/gpt-5-2025-08-07": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 1.25e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 1.25e-07,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"azure/gpt-5-mini": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 2.5e-07,
"output_cost_per_token": 2e-06,
"cache_read_input_token_cost": 2.5e-08,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"azure/gpt-5-mini-2025-08-07": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 2.5e-07,
"output_cost_per_token": 2e-06,
"cache_read_input_token_cost": 2.5e-08,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"azure/gpt-5-nano-2025-08-07": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 5e-08,
"output_cost_per_token": 4e-07,
"cache_read_input_token_cost": 5e-09,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"azure/gpt-5-nano": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 5e-08,
"output_cost_per_token": 4e-07,
"cache_read_input_token_cost": 5e-09,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"azure/gpt-5-chat": {
"max_tokens": 32768,
"max_input_tokens": 1047576,
"max_output_tokens": 32768,
"input_cost_per_token": 5e-06,
"output_cost_per_token": 2e-05,
"input_cost_per_token_batches": 2.5e-06,
"output_cost_per_token_batches": 1e-05,
"cache_read_input_token_cost": 1.25e-06,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true
},
"azure/gpt-5-chat-latest": {
"max_tokens": 128000,
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"input_cost_per_token": 1.25e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 1.25e-07,
"litellm_provider": "azure",
"mode": "chat",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_pdf_input": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_prompt_caching": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_native_streaming": true,
"supports_reasoning": true
},
"azure/gpt-4o-mini-tts": {
"mode": "audio_speech",
"input_cost_per_token": 2.5e-06,
@ -5486,6 +6000,36 @@
"litellm_provider": "groq",
"mode": "audio_transcription"
},
"groq/openai/gpt-oss-20b": {
"max_tokens": 32768,
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"input_cost_per_token": 1e-07,
"output_cost_per_token": 5e-07,
"litellm_provider": "groq",
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_web_search": true
},
"groq/openai/gpt-oss-120b": {
"max_tokens": 32766,
"max_input_tokens": 131072,
"max_output_tokens": 32766,
"input_cost_per_token": 1.5e-07,
"output_cost_per_token": 7.5e-07,
"litellm_provider": "groq",
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_web_search": true
},
"cerebras/llama3.1-8b": {
"max_tokens": 128000,
"max_input_tokens": 128000,
@ -5741,6 +6285,32 @@
"supports_reasoning": true,
"supports_computer_use": true
},
"claude-opus-4-1": {
"max_tokens": 32000,
"max_input_tokens": 200000,
"max_output_tokens": 32000,
"input_cost_per_token": 1.5e-05,
"output_cost_per_token": 7.5e-05,
"search_context_cost_per_query": {
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01,
"search_context_size_high": 0.01
},
"cache_creation_input_token_cost": 1.875e-05,
"cache_read_input_token_cost": 1.5e-06,
"litellm_provider": "anthropic",
"mode": "chat",
"supports_function_calling": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159,
"supports_assistant_prefill": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_reasoning": true,
"supports_computer_use": true
},
"claude-opus-4-1-20250805": {
"max_tokens": 32000,
"max_input_tokens": 200000,
@ -17663,5 +18233,136 @@
"supports_vision": false,
"supports_system_messages": true,
"supports_tool_choice": false
},
"oci/meta.llama-4-maverick-17b-128e-instruct-fp8": {
"max_tokens": 512000,
"max_input_tokens": 512000,
"max_output_tokens": 4000,
"input_cost_per_token": 7.2e-07,
"output_cost_per_token": 7.2e-07,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/meta.llama-4-scout-17b-16e-instruct": {
"max_tokens": 192000,
"max_input_tokens": 192000,
"max_output_tokens": 4000,
"input_cost_per_token": 7.2e-07,
"output_cost_per_token": 7.2e-07,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/meta.llama-3.3-70b-instruct": {
"max_tokens": 128000,
"max_input_tokens": 128000,
"max_output_tokens": 4000,
"input_cost_per_token": 7.2e-07,
"output_cost_per_token": 7.2e-07,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/meta.llama-3.2-90b-vision-instruct": {
"max_tokens": 128000,
"max_input_tokens": 128000,
"max_output_tokens": 4000,
"input_cost_per_token": 2.0e-06,
"output_cost_per_token": 2.0e-06,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/meta.llama-3.1-405b-instruct": {
"max_tokens": 128000,
"max_input_tokens": 128000,
"max_output_tokens": 4000,
"input_cost_per_token": 1.068e-05,
"output_cost_per_token": 1.068e-05,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/xai.grok-4": {
"max_tokens": 128000,
"max_input_tokens": 128000,
"max_output_tokens": 128000,
"input_cost_per_token": 3.0e-06,
"output_cost_per_token": 1.5e-07,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/xai.grok-3": {
"max_tokens": 131072,
"max_input_tokens": 131072,
"max_output_tokens": 131072,
"input_cost_per_token": 3.0e-06,
"output_cost_per_token": 1.5e-07,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/xai.grok-3-mini": {
"max_tokens": 131072,
"max_input_tokens": 131072,
"max_output_tokens": 131072,
"input_cost_per_token": 3.0e-07,
"output_cost_per_token": 5.0e-07,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/xai.grok-3-fast": {
"max_tokens": 131072,
"max_input_tokens": 131072,
"max_output_tokens": 131072,
"input_cost_per_token": 5.0e-06,
"output_cost_per_token": 2.5e-05,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
},
"oci/xai.grok-3-mini-fast": {
"max_tokens": 131072,
"max_input_tokens": 131072,
"max_output_tokens": 131072,
"input_cost_per_token": 6.0e-07,
"output_cost_per_token": 4.0e-06,
"litellm_provider": "oci",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": false,
"supports_tool_choice": false,
"source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing"
}
}

7
poetry.lock generated
View file

@ -2538,15 +2538,14 @@ files = [
[[package]]
name = "litellm-proxy-extras"
version = "0.2.15"
version = "0.2.16"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
optional = true
python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8"
groups = ["main"]
markers = "extra == \"proxy\""
files = [
{file = "litellm_proxy_extras-0.2.15-py3-none-any.whl", hash = "sha256:25e7d7cabe3f10233e2802d4521576bef438eea8b0800b1dbf38ba83ba751bb8"},
{file = "litellm_proxy_extras-0.2.15.tar.gz", hash = "sha256:62a9fdcb77d25aa7bfdfa04ac878e452f185dfb0b9538d4d8988a1b512360649"},
{file = "litellm_proxy_extras-0.2.16.tar.gz", hash = "sha256:81a1e8a172feb7da86985f529e891ca7be66ba293ae3e716bf69b266fa776a04"},
]
[[package]]
@ -6524,4 +6523,4 @@ utils = ["numpydoc"]
[metadata]
lock-version = "2.1"
python-versions = ">=3.8.1,<4.0, !=3.9.7"
content-hash = "8caa7dd3ee7d56562ccb799b97de7f2d206010d8300139bdc3614853c9ea1d31"
content-hash = "91f8d8cba2aa02a3eb205e91118891612c510790e828b56d19057fc509e6d25d"

View file

@ -58,7 +58,7 @@ websockets = {version = "^13.1.0", optional = true}
boto3 = {version = "1.34.34", optional = true}
redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"}
mcp = {version = "^1.10.0", optional = true, python = ">=3.10"}
litellm-proxy-extras = {version = "0.2.15", optional = true}
litellm-proxy-extras = {version = "0.2.16", optional = true}
rich = {version = "13.7.1", optional = true}
litellm-enterprise = {version = "0.1.19", optional = true}
diskcache = {version = "^5.6.1", optional = true}

View file

@ -41,7 +41,7 @@ sentry_sdk==2.21.0 # for sentry error handling
detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests
cryptography==43.0.1
tzdata==2025.1 # IANA time zone database
litellm-proxy-extras==0.2.15 # for proxy extras - e.g. prisma migrations
litellm-proxy-extras==0.2.16 # for proxy extras - e.g. prisma migrations
### LITELLM PACKAGE DEPENDENCIES
python-dotenv==1.0.1 # for env
tiktoken==0.8.0 # for calculating usage

View file

@ -25,6 +25,7 @@ IGNORE_FUNCTIONS = [
"filter_value_from_dict", # max depth set.
"normalize_json_schema_types", # max depth set.
"_extract_fields_recursive", # max depth set.
"_remove_json_schema_refs", # max depth set.
]

View file

@ -430,7 +430,7 @@ def test_gemini_with_empty_function_call_arguments():
async def test_claude_tool_use_with_gemini():
response = await litellm.anthropic.messages.acreate(
messages=[
{"role": "user", "content": "Hello, can you tell me the weather in Boston?"}
{"role": "user", "content": "Hello, can you tell me the weather in Boston. Please respond with a tool call?"}
],
model="gemini/gemini-2.5-flash",
stream=True,

View file

@ -504,75 +504,6 @@ async def test_async_vertexai_streaming_response():
pytest.fail(f"An exception occurred: {e}")
# asyncio.run(test_async_vertexai_streaming_response())
@pytest.mark.parametrize("provider", ["vertex_ai"]) # "vertex_ai_beta"
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.flaky(retries=3, delay=1)
@pytest.mark.asyncio
async def test_gemini_pro_vision(provider, sync_mode):
try:
load_vertex_ai_credentials()
litellm.set_verbose = True
litellm.num_retries = 3
if sync_mode:
resp = litellm.completion(
model="{}/gemini-2.5-flash-lite-preview-0514".format(provider),
messages=[
{"role": "system", "content": "Be a good bot"},
{
"role": "user",
"content": [
{"type": "text", "text": "Whats in this image?"},
{
"type": "image_url",
"image_url": {
"url": "gs://cloud-samples-data/generative-ai/image/boats.jpeg"
},
},
],
},
],
)
else:
resp = await litellm.acompletion(
model="{}/gemini-2.5-flash-lite-preview-0514".format(provider),
messages=[
{"role": "system", "content": "Be a good bot"},
{
"role": "user",
"content": [
{"type": "text", "text": "Whats in this image?"},
{
"type": "image_url",
"image_url": {
"url": "gs://cloud-samples-data/generative-ai/image/boats.jpeg"
},
},
],
},
],
)
print(resp)
prompt_tokens = resp.usage.prompt_tokens
# DO Not DELETE this ASSERT
# Google counts the prompt tokens for us, we should ensure we use the tokens from the orignal response
assert prompt_tokens == 267 # the gemini api returns 267 to us
except litellm.RateLimitError as e:
pass
except Exception as e:
if "500 Internal error encountered.'" in str(e):
pass
else:
pytest.fail(f"An exception occurred - {str(e)}")
# test_gemini_pro_vision()
@pytest.mark.parametrize("load_pdf", [False]) # True,
@pytest.mark.flaky(retries=3, delay=1)
@ -605,7 +536,7 @@ def test_completion_function_plus_pdf(load_pdf):
image_message = {"role": "user", "content": image_content}
response = completion(
model="vertex_ai_beta/gemini-2.5-flash-lite-preview-0514",
model="vertex_ai_beta/gemini-2.5-flash-lite",
messages=[image_message],
stream=False,
)

View file

@ -701,7 +701,7 @@ async def test_completion_gemini_stream(sync_mode):
},
}
]
messages = [{"role": "user", "content": "What is the weather like in Boston?"}]
messages = [{"role": "user", "content": "What is the weather like in Boston, MA?. You must provide me with a tool call in your response."}]
print("testing gemini streaming")
complete_response = ""
# Add any assertions here to check the response
@ -817,7 +817,7 @@ async def test_completion_gemini_stream_accumulated_json(sync_mode):
},
}
]
messages = [{"role": "user", "content": "What is the weather like in Boston?"}]
messages = [{"role": "user", "content": "What is the weather like in Boston, MA?. You must provide me with a tool call in your response."}]
print("testing gemini streaming")
complete_response = ""
# Add any assertions here to check the response

View file

@ -476,3 +476,18 @@ def test_s3_logging_r2():
# post, close log file and verify
# Reset stdout to the original value
print("Passed! Testing async s3 logging")
from litellm.integrations.s3_v2 import S3Logger
class TestS3Logger(S3Logger):
def __init__(self, *args, **kwargs):
self.recorded_requests = {}
self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None
super().__init__(*args, **kwargs)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.recorded_requests[response_obj["id"]] = start_time
print("recorded request", self.recorded_requests)
self.logged_standard_logging_payload = kwargs["standard_logging_object"]
return await super().async_log_success_event(kwargs, response_obj, start_time, end_time)

View file

@ -38,6 +38,7 @@ ignored_keys = [
"endTime",
"metadata.model_map_information",
"metadata.usage_object",
"metadata.cold_storage_object_key",
]

View file

@ -279,6 +279,7 @@ def validate_redacted_message_span_attributes(span):
"metadata.mcp_tool_call_metadata",
"metadata.vector_store_request_metadata",
"metadata.requester_custom_headers",
"metadata.cold_storage_object_key",
]
_all_attributes = set(

View file

@ -33,6 +33,7 @@ expected_env_vars = {
"LAGO_API_BASE": "mock_base",
"LAGO_API_EVENT_CODE": "mock_event_code",
"OPENMETER_API_KEY": "openmeter_api_key",
"BRAINTRUST_API_BASE": "braintrust_api_base",
"BRAINTRUST_API_KEY": "braintrust_api_key",
"GALILEO_API_KEY": "galileo_api_key",
"LITERAL_API_KEY": "literal_api_key",

View file

@ -0,0 +1,43 @@
import os
import unittest
from unittest.mock import patch
from litellm.integrations.braintrust_logging import BraintrustLogger
class TestBraintrustLogger(unittest.TestCase):
@patch.dict(os.environ, {"BRAINTRUST_API_KEY": "test-env-api-key"})
@patch.dict(os.environ, {"BRAINTRUST_API_BASE": "https://test-env-api.com/v1"})
def test_init_with_env_var(self):
"""Test BraintrustLogger initialization with environment variable."""
logger = BraintrustLogger()
self.assertEqual(logger.api_key, "test-env-api-key")
self.assertEqual(logger.api_base, "https://test-env-api.com/v1")
self.assertEqual(logger.headers["Authorization"], "Bearer test-env-api-key")
self.assertEqual(logger.headers["Content-Type"], "application/json")
def test_init_with_explicit_params(self):
"""Test BraintrustLogger initialization with explicit parameters."""
logger = BraintrustLogger(api_key="explicit-key", api_base="https://custom-api.com/v1")
self.assertEqual(logger.api_key, "explicit-key")
self.assertEqual(logger.api_base, "https://custom-api.com/v1")
self.assertEqual(logger.headers["Authorization"], "Bearer explicit-key")
@patch.dict(os.environ, {}, clear=True)
def test_init_missing_api_key(self):
"""Test BraintrustLogger initialization fails without API key."""
with self.assertRaises(Exception) as context:
BraintrustLogger()
self.assertIn("Missing keys=['BRAINTRUST_API_KEY']", str(context.exception))
def test_validate_environment_with_api_key(self):
"""Test validate_environment method with valid API key."""
logger = BraintrustLogger(api_key="test-key")
# Should not raise an exception
logger.validate_environment(api_key="test-key")
def test_validate_environment_missing_api_key(self):
"""Test validate_environment method with missing API key."""
with patch.dict(os.environ, {}, clear=True):
with self.assertRaises(Exception) as context:
BraintrustLogger(api_key=None)
self.assertIn("Missing keys=['BRAINTRUST_API_KEY']", str(context.exception))

View file

@ -394,3 +394,73 @@ def test_get_masked_values():
sensitive_object, unmasked_length=4, number_of_asterisks=4
)
assert masked_values["presidio_anonymizer_api_base"] is None
@pytest.mark.asyncio
async def test_e2e_generate_cold_storage_object_key_successful():
"""
Test end-to-end generation of cold storage object key when cold storage is properly configured.
"""
from datetime import datetime, timezone
from unittest.mock import patch
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
# Create test data
start_time = datetime(2025, 1, 15, 10, 30, 45, 123456, timezone.utc)
response_id = "chatcmpl-test-12345"
team_alias = "test-team"
with patch("litellm.proxy.spend_tracking.cold_storage_handler.ColdStorageHandler._get_configured_cold_storage_custom_logger", return_value="s3"), \
patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key:
# Mock the S3 object key generation to return a predictable result
mock_get_s3_key.return_value = "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
# Call the function
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
start_time=start_time,
response_id=response_id,
team_alias=team_alias
)
# Verify the S3 function was called with correct parameters
mock_get_s3_key.assert_called_once_with(
s3_path="", # Empty path as default
team_alias_prefix="", # No team alias prefix for cold storage
start_time=start_time,
s3_file_name="time-10-30-45-123456_chatcmpl-test-12345"
)
# Verify the result
assert result == "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
assert result is not None
assert isinstance(result, str)
@pytest.mark.asyncio
async def test_e2e_generate_cold_storage_object_key_not_configured():
"""
Test end-to-end generation of cold storage object key when cold storage is not configured.
"""
from datetime import datetime, timezone
from unittest.mock import patch
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
# Create test data
start_time = datetime(2025, 1, 15, 10, 30, 45, 123456, timezone.utc)
response_id = "chatcmpl-test-67890"
team_alias = "another-team"
with patch("litellm.proxy.spend_tracking.cold_storage_handler.ColdStorageHandler._get_configured_cold_storage_custom_logger", return_value=None):
# Call the function
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
start_time=start_time,
response_id=response_id,
team_alias=team_alias
)
# Verify the result is None when cold storage is not configured
assert result is None

View file

@ -0,0 +1,861 @@
"""
Unit tests for Bedrock Guardrails
"""
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
sys.path.insert(0, os.path.abspath("../../../../../.."))
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
BedrockGuardrail,
_redact_pii_matches,
)
@pytest.mark.asyncio
async def test__redact_pii_matches_function():
"""Test the _redact_pii_matches function directly"""
# Test case 1: Response with PII entities
response_with_pii = {
"action": "GUARDRAIL_INTERVENED",
"assessments": [
{
"sensitiveInformationPolicy": {
"piiEntities": [
{"type": "NAME", "match": "John Smith", "action": "BLOCKED"},
{
"type": "US_SOCIAL_SECURITY_NUMBER",
"match": "324-12-3212",
"action": "BLOCKED",
},
{"type": "PHONE", "match": "607-456-7890", "action": "BLOCKED"},
]
}
}
],
"outputs": [{"text": "Input blocked by PII policy"}],
}
# Call the redaction function
redacted_response = _redact_pii_matches(response_with_pii)
# Verify that PII matches are redacted
pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][
"piiEntities"
]
assert pii_entities[0]["match"] == "[REDACTED]", "Name should be redacted"
assert pii_entities[1]["match"] == "[REDACTED]", "SSN should be redacted"
assert pii_entities[2]["match"] == "[REDACTED]", "Phone should be redacted"
# Verify other fields remain unchanged
assert pii_entities[0]["type"] == "NAME"
assert pii_entities[1]["type"] == "US_SOCIAL_SECURITY_NUMBER"
assert pii_entities[2]["type"] == "PHONE"
assert redacted_response["action"] == "GUARDRAIL_INTERVENED"
assert redacted_response["outputs"][0]["text"] == "Input blocked by PII policy"
print("PII redaction function test passed")
@pytest.mark.asyncio
async def test__redact_pii_matches_no_pii():
"""Test _redact_pii_matches with response that has no PII"""
response_no_pii = {"action": "NONE", "assessments": [], "outputs": []}
# Call the redaction function
redacted_response = _redact_pii_matches(response_no_pii)
# Should return the same response unchanged
assert redacted_response == response_no_pii
print("No PII redaction test passed")
@pytest.mark.asyncio
async def test__redact_pii_matches_empty_assessments():
"""Test _redact_pii_matches with empty assessments"""
response_empty_assessments = {
"action": "GUARDRAIL_INTERVENED",
"assessments": [{"sensitiveInformationPolicy": {"piiEntities": []}}],
"outputs": [{"text": "Some output"}],
}
# Call the redaction function
redacted_response = _redact_pii_matches(response_empty_assessments)
# Should return the same response unchanged
assert redacted_response == response_empty_assessments
print("Empty assessments redaction test passed")
@pytest.mark.asyncio
async def test__redact_pii_matches_malformed_response():
"""Test _redact_pii_matches with malformed response (should not crash)"""
# Test with completely malformed response
malformed_response = {
"action": "GUARDRAIL_INTERVENED",
"assessments": "not_a_list", # This should cause an exception
}
# Should not crash and return original response
redacted_response = _redact_pii_matches(malformed_response)
assert redacted_response == malformed_response
# Test with missing keys
missing_keys_response = {
"action": "GUARDRAIL_INTERVENED"
# Missing assessments key
}
redacted_response = _redact_pii_matches(missing_keys_response)
assert redacted_response == missing_keys_response
print("Malformed response redaction test passed")
@pytest.mark.asyncio
async def test__redact_pii_matches_multiple_assessments():
"""Test _redact_pii_matches with multiple assessments containing PII"""
response_multiple_assessments = {
"action": "GUARDRAIL_INTERVENED",
"assessments": [
{
"sensitiveInformationPolicy": {
"piiEntities": [
{
"type": "EMAIL",
"match": "john@example.com",
"action": "ANONYMIZED",
}
]
}
},
{
"sensitiveInformationPolicy": {
"piiEntities": [
{
"type": "CREDIT_DEBIT_CARD_NUMBER",
"match": "1234-5678-9012-3456",
"action": "BLOCKED",
},
{
"type": "ADDRESS",
"match": "123 Main St, Anytown USA",
"action": "ANONYMIZED",
},
]
}
},
],
"outputs": [{"text": "Multiple PII detected"}],
}
# Call the redaction function
redacted_response = _redact_pii_matches(response_multiple_assessments)
# Verify all PII in all assessments are redacted
assessment1_pii = redacted_response["assessments"][0]["sensitiveInformationPolicy"][
"piiEntities"
]
assessment2_pii = redacted_response["assessments"][1]["sensitiveInformationPolicy"][
"piiEntities"
]
assert assessment1_pii[0]["match"] == "[REDACTED]", "Email should be redacted"
assert assessment2_pii[0]["match"] == "[REDACTED]", "Credit card should be redacted"
assert assessment2_pii[1]["match"] == "[REDACTED]", "Address should be redacted"
# Verify types remain unchanged
assert assessment1_pii[0]["type"] == "EMAIL"
assert assessment2_pii[0]["type"] == "CREDIT_DEBIT_CARD_NUMBER"
assert assessment2_pii[1]["type"] == "ADDRESS"
print("Multiple assessments redaction test passed")
@pytest.mark.asyncio
async def test_bedrock_guardrail_logging_uses_redacted_response():
"""Test that the Bedrock guardrail uses redacted response for logging"""
# Create proper mock objects
mock_user_api_key_dict = UserAPIKeyAuth()
guardrail = BedrockGuardrail(
guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
)
# Mock the Bedrock API response with PII
mock_bedrock_response = MagicMock()
mock_bedrock_response.status_code = 200
mock_bedrock_response.json.return_value = {
"action": "GUARDRAIL_INTERVENED",
"outputs": [{"text": "Hello, my phone number is {PHONE}"}],
"assessments": [
{
"sensitiveInformationPolicy": {
"piiEntities": [
{
"type": "PHONE",
"match": "+1 412 555 1212", # This should be redacted in logs
"action": "ANONYMIZED",
}
]
}
}
],
}
request_data = {
"model": "gpt-4o",
"messages": [
{"role": "user", "content": "Hello, my phone number is +1 412 555 1212"},
],
}
# Mock AWS credentials to avoid credential loading issues in CI
mock_credentials = MagicMock()
mock_credentials.access_key = "test-access-key"
mock_credentials.secret_key = "test-secret-key"
mock_credentials.token = None
# Mock AWS-related methods to ensure test runs without external dependencies
with patch.object(
guardrail.async_handler, "post", new_callable=AsyncMock
) as mock_post, patch(
"litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug"
) as mock_debug, patch.object(
guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")
) as mock_load_creds, patch.object(
guardrail, "_prepare_request", return_value=MagicMock()
) as mock_prepare_request:
mock_post.return_value = mock_bedrock_response
# Call the method that should log the redacted response
await guardrail.make_bedrock_api_request(
source="INPUT",
messages=request_data.get("messages"),
request_data=request_data,
)
# Verify that debug logging was called
mock_debug.assert_called()
# Get the logged response (second argument to debug call)
logged_calls = mock_debug.call_args_list
bedrock_response_log_call = None
for call in logged_calls:
args, kwargs = call
if len(args) >= 2 and "Bedrock AI response" in str(args[0]):
bedrock_response_log_call = call
break
assert (
bedrock_response_log_call is not None
), "Should have logged Bedrock AI response"
# Extract the logged response data
logged_response = bedrock_response_log_call[0][
1
] # Second argument to debug call
# Verify that the logged response has redacted PII
assert (
logged_response["assessments"][0]["sensitiveInformationPolicy"][
"piiEntities"
][0]["match"]
== "[REDACTED]"
)
# Verify other fields are preserved
assert logged_response["action"] == "GUARDRAIL_INTERVENED"
assert (
logged_response["assessments"][0]["sensitiveInformationPolicy"][
"piiEntities"
][0]["type"]
== "PHONE"
)
print("Bedrock guardrail logging redaction test passed")
@pytest.mark.asyncio
async def test_bedrock_guardrail_original_response_not_modified():
"""Test that the original response is not modified by redaction, only the logged version"""
# Create proper mock objects
mock_user_api_key_dict = UserAPIKeyAuth()
guardrail = BedrockGuardrail(
guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
)
# Mock the Bedrock API response with PII
original_response_data = {
"action": "GUARDRAIL_INTERVENED",
"outputs": [{"text": "Hello, my phone number is {PHONE}"}],
"assessments": [
{
"sensitiveInformationPolicy": {
"piiEntities": [
{
"type": "PHONE",
"match": "+1 412 555 1212", # This should NOT be modified in original
"action": "ANONYMIZED",
}
]
}
}
],
}
mock_bedrock_response = MagicMock()
mock_bedrock_response.status_code = 200
mock_bedrock_response.json.return_value = original_response_data
request_data = {
"model": "gpt-4o",
"messages": [
{"role": "user", "content": "Hello, my phone number is +1 412 555 1212"},
],
}
# Mock AWS credentials to avoid credential loading issues in CI
mock_credentials = MagicMock()
mock_credentials.access_key = "test-access-key"
mock_credentials.secret_key = "test-secret-key"
mock_credentials.token = None
# Mock AWS-related methods to ensure test runs without external dependencies
with patch.object(
guardrail.async_handler, "post", new_callable=AsyncMock
) as mock_post, patch.object(
guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")
) as mock_load_creds, patch.object(
guardrail, "_prepare_request", return_value=MagicMock()
) as mock_prepare_request:
mock_post.return_value = mock_bedrock_response
# Call the method
result = await guardrail.make_bedrock_api_request(
source="INPUT",
messages=request_data.get("messages"),
request_data=request_data,
)
# Verify that the original response data was not modified
# (The json() method should return the original data)
original_data = mock_bedrock_response.json()
assert (
original_data["assessments"][0]["sensitiveInformationPolicy"][
"piiEntities"
][0]["match"]
== "+1 412 555 1212"
)
# Verify that the returned BedrockGuardrailResponse contains original data
assert (
result["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][
"match"
]
== "+1 412 555 1212"
)
print("Original response not modified test passed")
@pytest.mark.asyncio
async def test__redact_pii_matches_preserves_non_pii_entities():
"""Test that _redact_pii_matches only affects PII-related entities and preserves other assessment data"""
response_with_mixed_data = {
"action": "GUARDRAIL_INTERVENED",
"assessments": [
{
"sensitiveInformationPolicy": {
"piiEntities": [
{
"type": "EMAIL",
"match": "user@example.com",
"action": "ANONYMIZED",
"confidence": "HIGH",
}
],
"regexes": [
{
"name": "custom_pattern",
"match": "some_pattern_match",
"action": "BLOCKED",
}
],
},
"contentPolicy": {
"filters": [
{
"type": "VIOLENCE",
"confidence": "MEDIUM",
"action": "BLOCKED",
}
]
},
"topicPolicy": {
"topics": [
{
"name": "Restricted Topic",
"type": "DENY",
"action": "BLOCKED",
}
]
},
}
],
"outputs": [{"text": "Content blocked"}],
}
# Call the redaction function
redacted_response = _redact_pii_matches(response_with_mixed_data)
# Verify that PII entity matches are redacted
pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][
"piiEntities"
]
assert pii_entities[0]["match"] == "[REDACTED]", "PII match should be redacted"
assert pii_entities[0]["type"] == "EMAIL", "PII type should be preserved"
assert pii_entities[0]["action"] == "ANONYMIZED", "PII action should be preserved"
assert pii_entities[0]["confidence"] == "HIGH", "PII confidence should be preserved"
# Verify that regex matches are also redacted (updated behavior)
regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][
"regexes"
]
assert regexes[0]["match"] == "[REDACTED]", "Regex match should be redacted"
assert regexes[0]["name"] == "custom_pattern", "Regex name should be preserved"
assert regexes[0]["action"] == "BLOCKED", "Regex action should be preserved"
# Verify that other policies are completely unchanged
content_policy = redacted_response["assessments"][0]["contentPolicy"]
assert content_policy["filters"][0]["type"] == "VIOLENCE"
assert content_policy["filters"][0]["confidence"] == "MEDIUM"
topic_policy = redacted_response["assessments"][0]["topicPolicy"]
assert topic_policy["topics"][0]["name"] == "Restricted Topic"
# Verify top-level fields are unchanged
assert redacted_response["action"] == "GUARDRAIL_INTERVENED"
assert redacted_response["outputs"][0]["text"] == "Content blocked"
print("Preserves non-PII entities test passed")
@pytest.mark.asyncio
async def test_pii_redaction_matches_debug_output_format():
"""Test that demonstrates the exact behavior shown in your debug output"""
# This matches the structure from your debug output
original_response = {
"action": "GUARDRAIL_INTERVENED",
"actionReason": "Guardrail blocked.",
"assessments": [
{
"invocationMetrics": {
"guardrailCoverage": {
"textCharacters": {"guarded": 84, "total": 84}
},
"guardrailProcessingLatency": 322,
"usage": {
"contentPolicyImageUnits": 0,
"contentPolicyUnits": 0,
"contextualGroundingPolicyUnits": 0,
"sensitiveInformationPolicyFreeUnits": 0,
"sensitiveInformationPolicyUnits": 1,
"topicPolicyUnits": 0,
"wordPolicyUnits": 0,
},
},
"sensitiveInformationPolicy": {
"piiEntities": [
{
"action": "BLOCKED",
"detected": True,
"match": "John Smith",
"type": "NAME",
},
{
"action": "BLOCKED",
"detected": True,
"match": "324-12-3212",
"type": "US_SOCIAL_SECURITY_NUMBER",
},
{
"action": "BLOCKED",
"detected": True,
"match": "607-456-7890",
"type": "PHONE",
},
]
},
}
],
"blockedResponse": "Input blocked by PII policy",
"guardrailCoverage": {"textCharacters": {"guarded": 84, "total": 84}},
"output": [{"text": "Input blocked by PII policy"}],
"outputs": [{"text": "Input blocked by PII policy"}],
"usage": {
"contentPolicyImageUnits": 0,
"contentPolicyUnits": 0,
"contextualGroundingPolicyUnits": 0,
"sensitiveInformationPolicyFreeUnits": 0,
"sensitiveInformationPolicyUnits": 1,
"topicPolicyUnits": 0,
"wordPolicyUnits": 0,
},
}
# Apply redaction
redacted_response = _redact_pii_matches(original_response)
# Verify the redacted response matches your expected debug output
pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][
"piiEntities"
]
# All PII matches should be redacted
assert pii_entities[0]["match"] == "[REDACTED]", "NAME should be redacted"
assert pii_entities[1]["match"] == "[REDACTED]", "SSN should be redacted"
assert pii_entities[2]["match"] == "[REDACTED]", "PHONE should be redacted"
# But all other fields should be preserved
assert pii_entities[0]["type"] == "NAME"
assert pii_entities[1]["type"] == "US_SOCIAL_SECURITY_NUMBER"
assert pii_entities[2]["type"] == "PHONE"
assert pii_entities[0]["action"] == "BLOCKED"
assert pii_entities[0]["detected"] == True
# Verify that the original response is unchanged
original_pii_entities = original_response["assessments"][0][
"sensitiveInformationPolicy"
]["piiEntities"]
assert (
original_pii_entities[0]["match"] == "John Smith"
), "Original should be unchanged"
assert (
original_pii_entities[1]["match"] == "324-12-3212"
), "Original should be unchanged"
assert (
original_pii_entities[2]["match"] == "607-456-7890"
), "Original should be unchanged"
# Verify all other metadata is preserved in redacted response
assert redacted_response["action"] == "GUARDRAIL_INTERVENED"
assert redacted_response["actionReason"] == "Guardrail blocked."
assert redacted_response["blockedResponse"] == "Input blocked by PII policy"
assert (
redacted_response["assessments"][0]["invocationMetrics"][
"guardrailProcessingLatency"
]
== 322
)
print("PII redaction matches debug output format test passed")
print(
f"Original PII values preserved: {[e['match'] for e in original_pii_entities]}"
)
print(f"Redacted PII values: {[e['match'] for e in pii_entities]}")
@pytest.mark.asyncio
async def test__redact_pii_matches_with_regex_matches():
"""Test redaction of regex matches in sensitive information policy"""
response_with_regex = {
"action": "GUARDRAIL_INTERVENED",
"assessments": [
{
"sensitiveInformationPolicy": {
"regexes": [
{
"name": "SSN_PATTERN",
"match": "123-45-6789",
"action": "BLOCKED",
},
{
"name": "CREDIT_CARD_PATTERN",
"match": "4111-1111-1111-1111",
"action": "ANONYMIZED",
},
]
}
}
],
"outputs": [{"text": "Regex patterns detected"}],
}
# Call the redaction function
redacted_response = _redact_pii_matches(response_with_regex)
# Verify that regex matches are redacted
regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][
"regexes"
]
assert regexes[0]["match"] == "[REDACTED]", "SSN regex match should be redacted"
assert (
regexes[1]["match"] == "[REDACTED]"
), "Credit card regex match should be redacted"
# Verify other fields are preserved
assert regexes[0]["name"] == "SSN_PATTERN", "Regex name should be preserved"
assert regexes[0]["action"] == "BLOCKED", "Regex action should be preserved"
assert regexes[1]["name"] == "CREDIT_CARD_PATTERN", "Regex name should be preserved"
assert regexes[1]["action"] == "ANONYMIZED", "Regex action should be preserved"
# Verify original response is unchanged
original_regexes = response_with_regex["assessments"][0][
"sensitiveInformationPolicy"
]["regexes"]
assert original_regexes[0]["match"] == "123-45-6789", "Original should be unchanged"
assert (
original_regexes[1]["match"] == "4111-1111-1111-1111"
), "Original should be unchanged"
print("Regex matches redaction test passed")
@pytest.mark.asyncio
async def test__redact_pii_matches_with_custom_words():
"""Test redaction of custom word matches in word policy"""
response_with_custom_words = {
"action": "GUARDRAIL_INTERVENED",
"assessments": [
{
"wordPolicy": {
"customWords": [
{
"match": "confidential_data",
"action": "BLOCKED",
},
{
"match": "secret_information",
"action": "ANONYMIZED",
},
]
}
}
],
"outputs": [{"text": "Custom words detected"}],
}
# Call the redaction function
redacted_response = _redact_pii_matches(response_with_custom_words)
# Verify that custom word matches are redacted
custom_words = redacted_response["assessments"][0]["wordPolicy"]["customWords"]
assert (
custom_words[0]["match"] == "[REDACTED]"
), "First custom word match should be redacted"
assert (
custom_words[1]["match"] == "[REDACTED]"
), "Second custom word match should be redacted"
# Verify other fields are preserved
assert (
custom_words[0]["action"] == "BLOCKED"
), "Custom word action should be preserved"
assert (
custom_words[1]["action"] == "ANONYMIZED"
), "Custom word action should be preserved"
# Verify original response is unchanged
original_custom_words = response_with_custom_words["assessments"][0]["wordPolicy"][
"customWords"
]
assert (
original_custom_words[0]["match"] == "confidential_data"
), "Original should be unchanged"
assert (
original_custom_words[1]["match"] == "secret_information"
), "Original should be unchanged"
print("Custom words redaction test passed")
@pytest.mark.asyncio
async def test__redact_pii_matches_with_managed_words():
"""Test redaction of managed word matches in word policy"""
response_with_managed_words = {
"action": "GUARDRAIL_INTERVENED",
"assessments": [
{
"wordPolicy": {
"managedWordLists": [
{
"match": "inappropriate_word",
"action": "BLOCKED",
"type": "PROFANITY",
},
{
"match": "offensive_term",
"action": "ANONYMIZED",
"type": "HATE_SPEECH",
},
]
}
}
],
"outputs": [{"text": "Managed words detected"}],
}
# Call the redaction function
redacted_response = _redact_pii_matches(response_with_managed_words)
# Verify that managed word matches are redacted
managed_words = redacted_response["assessments"][0]["wordPolicy"][
"managedWordLists"
]
assert (
managed_words[0]["match"] == "[REDACTED]"
), "First managed word match should be redacted"
assert (
managed_words[1]["match"] == "[REDACTED]"
), "Second managed word match should be redacted"
# Verify other fields are preserved
assert (
managed_words[0]["action"] == "BLOCKED"
), "Managed word action should be preserved"
assert (
managed_words[0]["type"] == "PROFANITY"
), "Managed word type should be preserved"
assert (
managed_words[1]["action"] == "ANONYMIZED"
), "Managed word action should be preserved"
assert (
managed_words[1]["type"] == "HATE_SPEECH"
), "Managed word type should be preserved"
# Verify original response is unchanged
original_managed_words = response_with_managed_words["assessments"][0][
"wordPolicy"
]["managedWordLists"]
assert (
original_managed_words[0]["match"] == "inappropriate_word"
), "Original should be unchanged"
assert (
original_managed_words[1]["match"] == "offensive_term"
), "Original should be unchanged"
print("Managed words redaction test passed")
@pytest.mark.asyncio
async def test__redact_pii_matches_comprehensive_coverage():
"""Test redaction across all supported policy types in a single response"""
comprehensive_response = {
"action": "GUARDRAIL_INTERVENED",
"assessments": [
{
"sensitiveInformationPolicy": {
"piiEntities": [
{
"type": "EMAIL",
"match": "user@example.com",
"action": "ANONYMIZED",
}
],
"regexes": [
{
"name": "PHONE_PATTERN",
"match": "555-123-4567",
"action": "BLOCKED",
}
],
},
"wordPolicy": {
"customWords": [
{
"match": "confidential",
"action": "BLOCKED",
}
],
"managedWordLists": [
{
"match": "inappropriate",
"action": "ANONYMIZED",
"type": "PROFANITY",
}
],
},
}
],
"outputs": [{"text": "Multiple policy violations detected"}],
}
# Call the redaction function
redacted_response = _redact_pii_matches(comprehensive_response)
# Verify all match fields are redacted
assessment = redacted_response["assessments"][0]
# PII entities
pii_entities = assessment["sensitiveInformationPolicy"]["piiEntities"]
assert (
pii_entities[0]["match"] == "[REDACTED]"
), "PII entity match should be redacted"
# Regex matches
regexes = assessment["sensitiveInformationPolicy"]["regexes"]
assert regexes[0]["match"] == "[REDACTED]", "Regex match should be redacted"
# Custom words
custom_words = assessment["wordPolicy"]["customWords"]
assert (
custom_words[0]["match"] == "[REDACTED]"
), "Custom word match should be redacted"
# Managed words
managed_words = assessment["wordPolicy"]["managedWordLists"]
assert (
managed_words[0]["match"] == "[REDACTED]"
), "Managed word match should be redacted"
# Verify all other fields are preserved
assert pii_entities[0]["type"] == "EMAIL"
assert regexes[0]["name"] == "PHONE_PATTERN"
assert managed_words[0]["type"] == "PROFANITY"
# Verify original response is unchanged
original_assessment = comprehensive_response["assessments"][0]
assert (
original_assessment["sensitiveInformationPolicy"]["piiEntities"][0]["match"]
== "user@example.com"
)
assert (
original_assessment["sensitiveInformationPolicy"]["regexes"][0]["match"]
== "555-123-4567"
)
assert (
original_assessment["wordPolicy"]["customWords"][0]["match"] == "confidential"
)
assert (
original_assessment["wordPolicy"]["managedWordLists"][0]["match"]
== "inappropriate"
)
print("Comprehensive coverage redaction test passed")

View file

@ -938,10 +938,10 @@ class TestUISSO_FunctionsExistence:
from litellm.proxy.management_endpoints.ui_sso import auth_callback
assert callable(auth_callback)
def test_sso_login_redirect_exists(self):
"""Test that sso_login_redirect function exists"""
from litellm.proxy.management_endpoints.ui_sso import sso_login_redirect
assert callable(sso_login_redirect)
def test_google_login_exists(self):
"""Test that google_login function exists"""
from litellm.proxy.management_endpoints.ui_sso import google_login
assert callable(google_login)
def test_sso_authentication_handler_exists(self):
"""Test that SSOAuthenticationHandler class exists with new methods"""
@ -1054,7 +1054,7 @@ class TestCustomUISSO:
"""Test that proper error is raised when enterprise module is not available"""
from unittest.mock import MagicMock, patch
from litellm.proxy.management_endpoints.ui_sso import sso_login_redirect
from litellm.proxy.management_endpoints.ui_sso import google_login
# Mock request
mock_request = MagicMock()
@ -1246,46 +1246,3 @@ class TestCustomUISSO:
assert result == mock_redirect_response
assert result.status_code == 303
@pytest.mark.asyncio
async def test_serve_login_page_server_root_path():
"""
Test that serve_login_page includes SERVER_ROOT_PATH in the SSO login URL
when SERVER_ROOT_PATH is set.
"""
# Arrange
mock_request = MagicMock(spec=Request)
captured_html = ""
# Mock environment variables
env_vars = {
"PROXY_BASE_URL": "https://example.com",
"SERVER_ROOT_PATH": "/api/v1",
"GOOGLE_CLIENT_ID": "mock_google_client_id", # Enable SSO
"DATABASE_URL": "mock_db_url", # Satisfy show_missing_vars_in_env
"LITELLM_MASTER_KEY": "mock_master_key", # Satisfy show_missing_vars_in_env
}
# Patch HTMLResponse to capture the content
def mock_html_response(content, status_code=200):
nonlocal captured_html
captured_html = content
return MagicMock()
with patch.dict(os.environ, env_vars):
with patch("litellm.proxy.proxy_server.premium_user", True):
with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()):
with patch("litellm.proxy.proxy_server.master_key", "mock_master_key"):
with patch("fastapi.responses.HTMLResponse", side_effect=mock_html_response):
# Import the function to test
from litellm.proxy.management_endpoints.ui_sso import (
serve_login_page,
)
# Act
result = await serve_login_page(request=mock_request)
# Assert
assert result is not None
expected_url = "https://example.com/api/v1/sso/login"
assert expected_url in captured_html, f"Expected URL '{expected_url}' not found in HTML content"

View file

@ -29,6 +29,7 @@ ignored_keys = [
"endTime",
"metadata.model_map_information",
"metadata.usage_object",
"metadata.cold_storage_object_key",
]

View file

@ -16,7 +16,7 @@ sys.path.insert(
from unittest.mock import MagicMock, patch
import litellm
from litellm.constants import REDACTED_BY_LITELM_STRING
from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, REDACTED_BY_LITELM_STRING
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy.spend_tracking.spend_tracking_utils import (
_get_vector_store_request_for_spend_logs_payload,
@ -35,7 +35,7 @@ def test_sanitize_request_body_for_spend_logs_payload_long_string():
long_string = "a" * 2000 # Create a string longer than MAX_STRING_LENGTH
request_body = {"text": long_string, "normal_text": "short text"}
sanitized = _sanitize_request_body_for_spend_logs_payload(request_body)
assert len(sanitized["text"]) == 1000 + len("... (truncated 1000 chars)")
assert len(sanitized["text"]) == 1000 + len(f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} 1000 chars)")
assert sanitized["normal_text"] == "short text"
@ -43,7 +43,7 @@ def test_sanitize_request_body_for_spend_logs_payload_nested_dict():
request_body = {"outer": {"inner": {"text": "a" * 2000, "normal": "short"}}}
sanitized = _sanitize_request_body_for_spend_logs_payload(request_body)
assert len(sanitized["outer"]["inner"]["text"]) == 1000 + len(
"... (truncated 1000 chars)"
f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} 1000 chars)"
)
assert sanitized["outer"]["inner"]["normal"] == "short"
@ -54,11 +54,11 @@ def test_sanitize_request_body_for_spend_logs_payload_nested_list():
}
sanitized = _sanitize_request_body_for_spend_logs_payload(request_body)
assert len(sanitized["items"][0]["text"]) == 1000 + len(
"... (truncated 1000 chars)"
f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} 1000 chars)"
)
assert sanitized["items"][1]["text"] == "short"
assert len(sanitized["items"][2][0]["text"]) == 1000 + len(
"... (truncated 1000 chars)"
f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} 1000 chars)"
)
@ -81,14 +81,14 @@ def test_sanitize_request_body_for_spend_logs_payload_mixed_types():
"nested": {"list": ["short", "a" * 2000], "dict": {"key": "a" * 2000}},
}
sanitized = _sanitize_request_body_for_spend_logs_payload(request_body)
assert len(sanitized["text"]) == 1000 + len("... (truncated 1000 chars)")
assert len(sanitized["text"]) == 1000 + len(f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} 1000 chars)")
assert sanitized["number"] == 42
assert sanitized["nested"]["list"][0] == "short"
assert len(sanitized["nested"]["list"][1]) == 1000 + len(
"... (truncated 1000 chars)"
f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} 1000 chars)"
)
assert len(sanitized["nested"]["dict"]["key"]) == 1000 + len(
"... (truncated 1000 chars)"
f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} 1000 chars)"
)

View file

@ -134,9 +134,9 @@ class TestLiteLLMCompletionResponsesConfig:
)
# Assert
expected = {"type": "image", "image_url": {"url": image_url, "detail": "high"}}
expected = {"type": "image_url", "image_url": {"url": image_url, "detail": "high"}}
assert result == expected
assert result["type"] == "image"
assert result["type"] == "image_url"
assert result["image_url"]["url"] == image_url
assert result["image_url"]["detail"] == "high"
@ -154,9 +154,9 @@ class TestLiteLLMCompletionResponsesConfig:
)
# Assert
expected = {"type": "image", "image_url": {"url": image_url, "detail": "high"}}
expected = {"type": "image_url", "image_url": {"url": image_url, "detail": "high"}}
assert result == expected
assert result["type"] == "image"
assert result["type"] == "image_url"
assert result["image_url"]["url"] == image_url
assert result["image_url"]["detail"] == "high"
@ -174,9 +174,9 @@ class TestLiteLLMCompletionResponsesConfig:
)
# Assert
expected = {"type": "image", "image_url": {"url": image_url, "detail": "auto"}}
expected = {"type": "image_url", "image_url": {"url": image_url, "detail": "auto"}}
assert result == expected
assert result["type"] == "image"
assert result["type"] == "image_url"
assert result["image_url"]["url"] == image_url
assert result["image_url"]["detail"] == "auto"
@ -193,9 +193,9 @@ class TestLiteLLMCompletionResponsesConfig:
)
# Assert
expected = {"type": "image", "image_url": {"url": "", "detail": "auto"}}
expected = {"type": "image_url", "image_url": {"url": "", "detail": "auto"}}
assert result == expected
assert result["type"] == "image"
assert result["type"] == "image_url"
assert result["image_url"]["url"] == ""
assert result["image_url"]["detail"] == "auto"
@ -217,9 +217,9 @@ class TestLiteLLMCompletionResponsesConfig:
)
# Assert
expected = {"type": "image", "image_url": {"url": "https://example.com/image.png", "detail": "auto"}}
expected = {"type": "image_url", "image_url": {"url": "https://example.com/image.png", "detail": "auto"}}
assert result == expected
assert result["type"] == "image"
assert result["type"] == "image_url"
assert result["image_url"]["url"] == "https://example.com/image.png"
assert result["image_url"]["detail"] == "auto"
assert "extra_field" not in result

View file

@ -11,9 +11,9 @@ sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from enterprise.litellm_enterprise.enterprise_callbacks.session_handler import (
_ENTERPRISE_ResponsesSessionHandler,
from litellm.responses.litellm_completion_transformation import session_handler
from litellm.responses.litellm_completion_transformation.session_handler import (
ResponsesSessionHandler,
)
@ -111,7 +111,7 @@ async def test_get_chat_completion_message_history_for_previous_response_id():
# Mock the get_all_spend_logs_for_previous_response_id method
with patch.object(
_ENTERPRISE_ResponsesSessionHandler,
ResponsesSessionHandler,
"get_all_spend_logs_for_previous_response_id",
new_callable=AsyncMock,
) as mock_get_spend_logs:
@ -119,7 +119,7 @@ async def test_get_chat_completion_message_history_for_previous_response_id():
# Test the function
previous_response_id = "chatcmpl-935b8dad-fdc2-466e-a8ca-e26e5a8a21bb"
result = await _ENTERPRISE_ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
previous_response_id
)
@ -166,17 +166,130 @@ async def test_get_chat_completion_message_history_empty_spend_logs():
Test get_chat_completion_message_history_for_previous_response_id with empty spend logs
"""
with patch.object(
_ENTERPRISE_ResponsesSessionHandler,
ResponsesSessionHandler,
"get_all_spend_logs_for_previous_response_id",
new_callable=AsyncMock,
) as mock_get_spend_logs:
mock_get_spend_logs.return_value = []
previous_response_id = "non-existent-id"
result = await _ENTERPRISE_ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
previous_response_id
)
# Verify empty result structure
assert result.get("messages") == []
assert result.get("litellm_session_id") is None
@pytest.mark.asyncio
async def test_e2e_cold_storage_successful_retrieval():
"""
Test end-to-end cold storage functionality with successful retrieval of full proxy request from cold storage.
"""
# Mock spend logs with cold storage object key in metadata
mock_spend_logs = [
{
"request_id": "chatcmpl-test-123",
"session_id": "session-456",
"metadata": '{"cold_storage_object_key": "s3://test-bucket/requests/session_456_req1.json"}',
"proxy_server_request": '{"litellm_truncated": true}', # Truncated payload
"response": {
"id": "chatcmpl-test-123",
"object": "chat.completion",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "I am an AI assistant."
}
}
]
}
}
]
# Full proxy request data from cold storage
full_proxy_request = {
"input": "Hello, who are you?",
"model": "gpt-4",
"messages": [{"role": "user", "content": "Hello, who are you?"}]
}
with patch.object(
ResponsesSessionHandler,
"get_all_spend_logs_for_previous_response_id",
new_callable=AsyncMock,
) as mock_get_spend_logs, \
patch.object(session_handler, "COLD_STORAGE_HANDLER") as mock_cold_storage, \
patch("litellm.proxy.spend_tracking.cold_storage_handler.ColdStorageHandler._get_configured_cold_storage_custom_logger", return_value="s3"):
# Setup mocks
mock_get_spend_logs.return_value = mock_spend_logs
mock_cold_storage.get_proxy_server_request_from_cold_storage_with_object_key = AsyncMock(return_value=full_proxy_request)
# Call the main function
result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
"chatcmpl-test-123"
)
# Verify cold storage was called with correct object key
mock_cold_storage.get_proxy_server_request_from_cold_storage_with_object_key.assert_called_once_with(
object_key="s3://test-bucket/requests/session_456_req1.json"
)
# Verify result structure
assert result.get("litellm_session_id") == "session-456"
assert len(result.get("messages", [])) >= 1 # At least the assistant response
@pytest.mark.asyncio
async def test_e2e_cold_storage_fallback_to_truncated_payload():
"""
Test end-to-end cold storage functionality when object key is missing, falling back to truncated payload.
"""
# Mock spend logs without cold storage object key
mock_spend_logs = [
{
"request_id": "chatcmpl-test-789",
"session_id": "session-999",
"metadata": '{"user_api_key": "test-key"}', # No cold storage object key
"proxy_server_request": '{"input": "Truncated message", "model": "gpt-4"}', # Regular payload
"response": {
"id": "chatcmpl-test-789",
"object": "chat.completion",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "This is a response."
}
}
]
}
}
]
with patch.object(
ResponsesSessionHandler,
"get_all_spend_logs_for_previous_response_id",
new_callable=AsyncMock,
) as mock_get_spend_logs, \
patch.object(session_handler, "COLD_STORAGE_HANDLER") as mock_cold_storage:
# Setup mocks
mock_get_spend_logs.return_value = mock_spend_logs
# Call the main function
result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
"chatcmpl-test-789"
)
# Verify cold storage was NOT called since no object key in metadata
mock_cold_storage.get_proxy_server_request_from_cold_storage_with_object_key.assert_not_called()
# Verify result structure
assert result.get("litellm_session_id") == "session-999"
assert len(result.get("messages", [])) >= 1 # At least the assistant response

View file

@ -0,0 +1,186 @@
"""
Unit tests for cold storage object key integration.
Tests for the changes to integrate cold storage handling across different components:
1. Add cold_storage_object_key field to StandardLoggingMetadata and SpendLogsMetadata
2. S3Logger generates object key when cold storage is enabled
3. Store object key in SpendLogsMetadata via spend_tracking_utils
4. Session handler uses object key from spend logs metadata
5. S3Logger supports retrieval using provided object key
"""
import json
from datetime import datetime, timezone
from typing import Optional
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.integrations.s3_v2 import S3Logger
from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
from litellm.proxy.spend_tracking.spend_tracking_utils import _get_spend_logs_metadata
from litellm.responses.litellm_completion_transformation.session_handler import (
ResponsesSessionHandler,
)
from litellm.types.utils import StandardLoggingMetadata, StandardLoggingPayload
class TestColdStorageObjectKeyIntegration:
"""Test suite for cold storage object key integration."""
def test_standard_logging_metadata_has_cold_storage_object_key_field(self):
"""
Test: Add cold_storage_object_key field to StandardLoggingMetadata.
This test verifies that the StandardLoggingMetadata TypedDict has the
cold_storage_object_key field for storing S3/GCS object keys.
"""
from litellm.types.utils import StandardLoggingMetadata
# Create a StandardLoggingMetadata instance with cold_storage_object_key
metadata = StandardLoggingMetadata(
user_api_key_hash="test_hash",
cold_storage_object_key="test/path/to/object.json"
)
# Verify the field can be set and accessed
assert metadata.get("cold_storage_object_key") == "test/path/to/object.json"
assert "cold_storage_object_key" in StandardLoggingMetadata.__annotations__
def test_spend_logs_metadata_has_cold_storage_object_key_field(self):
"""
Test: Add cold_storage_object_key field to SpendLogsMetadata.
This test verifies that the SpendLogsMetadata TypedDict has the
cold_storage_object_key field for storing S3/GCS object keys.
"""
# Create a SpendLogsMetadata instance with cold_storage_object_key
metadata = SpendLogsMetadata(
user_api_key="test_key",
cold_storage_object_key="test/path/to/object.json"
)
# Verify the field can be set and accessed
assert metadata.get("cold_storage_object_key") == "test/path/to/object.json"
# Verify it's part of the SpendLogsMetadata annotations
assert "cold_storage_object_key" in SpendLogsMetadata.__annotations__
def test_spend_tracking_utils_stores_object_key_in_metadata(self):
"""
Test: Store object key in SpendLogsMetadata via spend_tracking_utils.
This test verifies that the _get_spend_logs_metadata function extracts
the cold_storage_object_key from StandardLoggingPayload and stores it
in SpendLogsMetadata.
"""
# Create test data
metadata = {
"user_api_key": "test_key",
"user_api_key_team_id": "test_team"
}
# Call the function
result = _get_spend_logs_metadata(
metadata=metadata,
cold_storage_object_key="test/path/to/object.json"
)
# Verify the object key is stored in the result
assert result.get("cold_storage_object_key") == "test/path/to/object.json"
def test_session_handler_extracts_object_key_from_spend_log(self):
"""
Test: Session handler extracts object key from spend logs metadata.
This test verifies that the ResponsesSessionHandler can extract the
cold_storage_object_key from spend log metadata.
"""
# Create test spend log
spend_log = {
"request_id": "test_request_id",
"metadata": json.dumps({
"cold_storage_object_key": "test/path/to/object.json",
"user_api_key": "test_key"
})
}
# Test the extraction method
object_key = ResponsesSessionHandler._get_cold_storage_object_key_from_spend_log(spend_log)
assert object_key == "test/path/to/object.json"
def test_session_handler_handles_dict_metadata_in_spend_log(self):
"""
Test: Session handler handles dict metadata in spend log.
This test verifies that the method works when metadata is already a dict.
"""
# Create test spend log with dict metadata
spend_log = {
"request_id": "test_request_id",
"metadata": {
"cold_storage_object_key": "test/path/to/object.json",
"user_api_key": "test_key"
}
}
# Test the extraction method
object_key = ResponsesSessionHandler._get_cold_storage_object_key_from_spend_log(spend_log)
assert object_key == "test/path/to/object.json"
@pytest.mark.asyncio
async def test_cold_storage_handler_supports_object_key_retrieval(self):
"""
Test: ColdStorageHandler supports object key retrieval.
This test verifies that the ColdStorageHandler has the new method
for retrieving objects using object keys directly.
"""
handler = ColdStorageHandler()
# Mock the custom logger
mock_logger = AsyncMock()
mock_logger.get_proxy_server_request_from_cold_storage_with_object_key = AsyncMock(
return_value={"test": "data"}
)
with patch.object(handler, '_select_custom_logger_for_cold_storage', return_value="s3_v2"), \
patch('litellm.logging_callback_manager.get_active_custom_logger_for_callback_name', return_value=mock_logger):
result = await handler.get_proxy_server_request_from_cold_storage_with_object_key(
object_key="test/path/to/object.json"
)
assert result == {"test": "data"}
mock_logger.get_proxy_server_request_from_cold_storage_with_object_key.assert_called_once_with(
object_key="test/path/to/object.json"
)
@pytest.mark.asyncio
@patch('asyncio.create_task') # Mock asyncio.create_task to avoid event loop issues
async def test_s3_logger_supports_object_key_retrieval(self, mock_create_task):
"""
Test: S3Logger supports retrieval using provided object key.
This test verifies that the S3Logger can retrieve objects using
the object key directly without generating it from request_id and start_time.
"""
# Create S3Logger instance
s3_logger = S3Logger(s3_bucket_name="test-bucket")
# Mock the _download_object_from_s3 method
with patch.object(s3_logger, '_download_object_from_s3', return_value={"test": "data"}) as mock_download:
result = await s3_logger.get_proxy_server_request_from_cold_storage_with_object_key(
object_key="test/path/to/object.json"
)
assert result == {"test": "data"}
mock_download.assert_called_once_with("test/path/to/object.json")

View file

@ -29,6 +29,7 @@ import { tagDailyActivityCall, teamDailyActivityCall } from './networking';
import TopKeyView from "./top_key_view";
import { formatNumberWithCommas } from "@/utils/dataUtils";
import { valueFormatterSpend } from "./usage/utils/value_formatters";
import { getProviderLogoAndName } from "./provider_info_helpers";
interface EntityMetrics {
metrics: {
@ -688,7 +689,28 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
<TableBody>
{getProviderSpend().map((provider) => (
<TableRow key={provider.provider}>
<TableCell>{provider.provider}</TableCell>
<TableCell>
<div className="flex items-center space-x-2">
{provider.provider && (
<img
src={getProviderLogoAndName(provider.provider).logo}
alt={`${provider.provider} logo`}
className="w-4 h-4"
onError={(e) => {
const target = e.target as HTMLImageElement;
const parent = target.parentElement;
if (parent) {
const fallbackDiv = document.createElement('div');
fallbackDiv.className = 'w-4 h-4 rounded-full bg-gray-200 flex items-center justify-center text-xs';
fallbackDiv.textContent = provider.provider?.charAt(0) || '-';
parent.replaceChild(fallbackDiv, target);
}
}}
/>
)}
<span>{provider.provider}</span>
</div>
</TableCell>
<TableCell>
${formatNumberWithCommas(provider.spend, 2)}
</TableCell>

View file

@ -34,20 +34,14 @@ import {
import AdvancedDatePicker from "./shared/advanced_date_picker"
import { AreaChart } from "@tremor/react"
import { userDailyActivityCall, tagListCall } from "./networking";
import { Tag } from "./tag_management/types";
import ViewUserSpend from "./view_user_spend";
import TopKeyView from "./top_key_view";
import { ActivityMetrics, processActivityData } from "./activity_metrics";
import UserAgentActivity from "./user_agent_activity";
import {
SpendMetrics,
DailyData,
ModelActivityData,
MetricWithMetadata,
KeyMetricWithMetadata,
} from "./usage/types";
import EntityUsage from "./entity_usage";
import { userDailyActivityCall, tagListCall } from "./networking"
import { Tag } from "./tag_management/types"
import ViewUserSpend from "./view_user_spend"
import TopKeyView from "./top_key_view"
import { ActivityMetrics, processActivityData } from "./activity_metrics"
import UserAgentActivity from "./user_agent_activity"
import { SpendMetrics, DailyData, ModelActivityData, MetricWithMetadata, KeyMetricWithMetadata } from "./usage/types"
import EntityUsage from "./entity_usage"
import {
old_admin_roles,
v2_admin_role_names,
@ -62,6 +56,7 @@ import { formatNumberWithCommas } from "@/utils/dataUtils"
import { valueFormatterSpend } from "./usage/utils/value_formatters"
import CloudZeroExportModal from "./cloudzero_export_modal"
import { ChartLoader } from "./shared/chart_loader"
import { getProviderLogoAndName } from "./provider_info_helpers"
interface NewUsagePageProps {
accessToken: string | null
@ -416,16 +411,8 @@ const NewUsagePage: React.FC<NewUsagePageProps> = ({ accessToken, userRole, user
<TabList variant="solid" className="mt-1">
{all_admin_roles.includes(userRole || "") ? <Tab>Global Usage</Tab> : <Tab>Your Usage</Tab>}
<Tab>Team Usage</Tab>
{all_admin_roles.includes(userRole || "") ? (
<Tab>Tag Usage</Tab>
) : (
<></>
)}
{all_admin_roles.includes(userRole || "") ? (
<Tab>User Agent Activity</Tab>
) : (
<></>
)}
{all_admin_roles.includes(userRole || "") ? <Tab>Tag Usage</Tab> : <></>}
{all_admin_roles.includes(userRole || "") ? <Tab>User Agent Activity</Tab> : <></>}
</TabList>
<TabPanels>
{/* Your Usage Panel */}
@ -658,7 +645,29 @@ const NewUsagePage: React.FC<NewUsagePageProps> = ({ accessToken, userRole, user
.filter((provider) => provider.spend > 0)
.map((provider) => (
<TableRow key={provider.provider}>
<TableCell>{provider.provider}</TableCell>
<TableCell>
<div className="flex items-center space-x-2">
{provider.provider && (
<img
src={getProviderLogoAndName(provider.provider).logo}
alt={`${provider.provider} logo`}
className="w-4 h-4"
onError={(e) => {
const target = e.target as HTMLImageElement
const parent = target.parentElement
if (parent) {
const fallbackDiv = document.createElement("div")
fallbackDiv.className =
"w-4 h-4 rounded-full bg-gray-200 flex items-center justify-center text-xs"
fallbackDiv.textContent = provider.provider?.charAt(0) || "-"
parent.replaceChild(fallbackDiv, target)
}
}}
/>
)}
<span>{provider.provider}</span>
</div>
</TableCell>
<TableCell>${formatNumberWithCommas(provider.spend, 2)}</TableCell>
<TableCell className="text-green-600">
{provider.successful_requests.toLocaleString()}
@ -725,10 +734,7 @@ const NewUsagePage: React.FC<NewUsagePageProps> = ({ accessToken, userRole, user
</TabPanel>
{/* User Agent Activity Panel */}
<TabPanel>
<UserAgentActivity
accessToken={accessToken}
userRole={userRole}
/>
<UserAgentActivity accessToken={accessToken} userRole={userRole} />
</TabPanel>
</TabPanels>
</TabGroup>

View file

@ -9,6 +9,7 @@ import { ArrowLeftIcon } from "@heroicons/react/outline"
import { Button } from "antd"
import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"
import { CheckIcon, CopyIcon } from "lucide-react"
import { Tooltip } from "antd"
interface SessionViewProps {
sessionId: string
@ -24,6 +25,21 @@ export const SessionView: React.FC<SessionViewProps> = ({ sessionId, logs, onBac
// Calculate session metrics
const totalCost = logs.reduce((sum, log) => sum + (log.spend || 0), 0)
const totalTokens = logs.reduce((sum, log) => sum + (log.total_tokens || 0), 0)
// Calculate cache token totals from metadata
const totalCacheReadTokens = logs.reduce((sum, log) => {
const cacheReadTokens = log.metadata?.additional_usage_values?.cache_read_input_tokens || 0
return sum + cacheReadTokens
}, 0)
const totalCacheCreationTokens = logs.reduce((sum, log) => {
const cacheCreationTokens = log.metadata?.additional_usage_values?.cache_creation_input_tokens || 0
return sum + cacheCreationTokens
}, 0)
// Calculate total tokens including cache tokens
const totalTokensWithCache = totalTokens + totalCacheReadTokens + totalCacheCreationTokens
const startTime = logs.length > 0 ? new Date(logs[0].startTime) : new Date()
const endTime = logs.length > 0 ? new Date(logs[logs.length - 1].endTime) : new Date()
const durationMs = endTime.getTime() - startTime.getTime()
@ -100,10 +116,61 @@ export const SessionView: React.FC<SessionViewProps> = ({ sessionId, logs, onBac
<Text>Total Cost</Text>
<Metric>${formatNumberWithCommas(totalCost, 6)}</Metric>
</Card>
<Card>
<Text>Total Tokens</Text>
<Metric>{totalTokens}</Metric>
</Card>
<Tooltip
title={
<div className="text-white min-w-[200px]">
<div className="text-lg font-medium mb-3">Usage breakdown</div>
<div className="space-y-4">
<div>
<div className="text-base font-medium mb-2">Input usage:</div>
<div className="space-y-2 text-sm text-gray-300">
<div className="flex justify-between">
<span>input:</span>
<span className="ml-8">{formatNumberWithCommas(logs.reduce((sum, log) => sum + (log.prompt_tokens || 0), 0))}</span>
</div>
{totalCacheReadTokens > 0 && (
<div className="flex justify-between">
<span>input_cached_tokens:</span>
<span className="ml-8">{formatNumberWithCommas(totalCacheReadTokens)}</span>
</div>
)}
{totalCacheCreationTokens > 0 && (
<div className="flex justify-between">
<span>input_cache_creation_tokens:</span>
<span className="ml-8">{formatNumberWithCommas(totalCacheCreationTokens)}</span>
</div>
)}
</div>
</div>
<div className="border-t border-gray-600 pt-3">
<div className="text-base font-medium mb-2">Output usage:</div>
<div className="space-y-2 text-sm text-gray-300">
<div className="flex justify-between">
<span>output:</span>
<span className="ml-8">{formatNumberWithCommas(logs.reduce((sum, log) => sum + (log.completion_tokens || 0), 0))}</span>
</div>
</div>
</div>
<div className="border-t border-gray-600 pt-3">
<div className="flex justify-between items-center">
<span className="text-base font-medium">Total usage:</span>
<span className="text-sm text-gray-300">{formatNumberWithCommas(totalTokensWithCache)}</span>
</div>
</div>
</div>
</div>
}
placement="top"
overlayStyle={{ minWidth: '300px' }}
>
<Card>
<div className="flex items-center justify-between">
<Text>Total Tokens</Text>
<span className="text-gray-400 text-sm">ⓘ</span>
</div>
<Metric>{formatNumberWithCommas(totalTokensWithCache)}</Metric>
</Card>
</Tooltip>
</div>
{/* Request Timeline */}
<Title>Session Logs</Title>

View file

@ -925,7 +925,15 @@ export function RequestViewer({ row }: { row: Row<LogEntry> }) {
<div className="space-y-2">
<div className="flex">
<span className="font-medium w-1/3">Tokens:</span>
<span>{row.original.total_tokens} ({row.original.prompt_tokens}+{row.original.completion_tokens})</span>
<span>{row.original.total_tokens} ({row.original.prompt_tokens} prompt tokens + {row.original.completion_tokens} completion tokens)</span>
</div>
<div className="flex">
<span className="font-medium w-1/3">Cache Read Tokens:</span>
<span>{formatNumberWithCommas(row.original.metadata?.additional_usage_values?.cache_read_input_tokens || 0)}</span>
</div>
<div className="flex">
<span className="font-medium w-1/3">Cache Creation Tokens:</span>
<span>{formatNumberWithCommas(row.original.metadata?.additional_usage_values.cache_creation_input_tokens)}</span>
</div>
<div className="flex">
<span className="font-medium w-1/3">Cost:</span>