mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
Merge branch 'BerriAI:main' into fix/ollama-gpt-oss-thinking-field
This commit is contained in:
commit
b66f4a7d4b
55 changed files with 3795 additions and 840 deletions
53
cookbook/misc/test_responses_api.py
Normal file
53
cookbook/misc/test_responses_api.py
Normal 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)
|
||||
|
||||
|
||||
|
|
@ -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 }}
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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`
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
BIN
litellm-proxy-extras/dist/litellm_proxy_extras-0.2.16-py3-none-any.whl
vendored
Normal file
BIN
litellm-proxy-extras/dist/litellm_proxy_extras-0.2.16-py3-none-any.whl
vendored
Normal file
Binary file not shown.
BIN
litellm-proxy-extras/dist/litellm_proxy_extras-0.2.16.tar.gz
vendored
Normal file
BIN
litellm-proxy-extras/dist/litellm_proxy_extras-0.2.16.tar.gz
vendored
Normal file
Binary file not shown.
|
|
@ -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==",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 ###########################
|
||||
########################################################################################
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
73
litellm/proxy/spend_tracking/cold_storage_handler.py
Normal file
73
litellm/proxy/spend_tracking/cold_storage_handler.py
Normal 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)
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
7
poetry.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ ignored_keys = [
|
|||
"endTime",
|
||||
"metadata.model_map_information",
|
||||
"metadata.usage_object",
|
||||
"metadata.cold_storage_object_key",
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
43
tests/test_litellm/integrations/test_braintrust_logging.py
Normal file
43
tests/test_litellm/integrations/test_braintrust_logging.py
Normal 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))
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ ignored_keys = [
|
|||
"endTime",
|
||||
"metadata.model_map_information",
|
||||
"metadata.usage_object",
|
||||
"metadata.cold_storage_object_key",
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)"
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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")
|
||||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue