diff --git a/cookbook/misc/test_responses_api.py b/cookbook/misc/test_responses_api.py new file mode 100644 index 00000000000..5fd19c6f66f --- /dev/null +++ b/cookbook/misc/test_responses_api.py @@ -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) + + diff --git a/deploy/charts/litellm-helm/templates/migrations-job.yaml b/deploy/charts/litellm-helm/templates/migrations-job.yaml index 143e62fceb3..cf10be0a76b 100644 --- a/deploy/charts/litellm-helm/templates/migrations-job.yaml +++ b/deploy/charts/litellm-helm/templates/migrations-job.yaml @@ -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 }} diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index cdf4b89bff4..388dc6d0766 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -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"] \ No newline at end of file +CMD ["--port", "4000"] diff --git a/docs/my-website/docs/observability/braintrust.md b/docs/my-website/docs/observability/braintrust.md index 79f3cf13be2..eb26680b18a 100644 --- a/docs/my-website/docs/observability/braintrust.md +++ b/docs/my-website/docs/observability/braintrust.md @@ -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) +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 diff --git a/docs/my-website/docs/providers/oci.md b/docs/my-website/docs/providers/oci.md index 36971376866..28beb71094a 100644 --- a/docs/my-website/docs/providers/oci.md +++ b/docs/my-website/docs/providers/oci.md @@ -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` diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index c8aee990415..3b903935a04 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -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 diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/session_handler.py b/enterprise/litellm_enterprise/enterprise_callbacks/session_handler.py deleted file mode 100644 index 1a08a8f9101..00000000000 --- a/enterprise/litellm_enterprise/enterprise_callbacks/session_handler.py +++ /dev/null @@ -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 diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.16-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.16-py3-none-any.whl new file mode 100644 index 00000000000..ce275d59451 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.16-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.16.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.16.tar.gz new file mode 100644 index 00000000000..16e8acf09ae Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.16.tar.gz differ diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index ceffd9adefd..86307e08238 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -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==", diff --git a/litellm/__init__.py b/litellm/__init__.py index 17fc6c00e12..1330584cdf1 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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 diff --git a/litellm/constants.py b/litellm/constants.py index 27cea0eb040..c7404f10a78 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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 ########################### ######################################################################################## diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index c68674f77ba..8149a6131e8 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -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 = { diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index b5c7101dde9..ee7e771faa6 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -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 diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 7df3e58b2da..efe18cb68ad 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -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 \ No newline at end of file diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 9606b47b9b8..fd82ecdf2b2 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -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] diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index dfa941d8301..5462257c9b7 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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 diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 44cb146f91a..9ec346c20a1 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -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] diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 0441e75beec..b38a4982471 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -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: """ diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9120f1f079f..0d4f1d1dfba 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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" } } diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b69d27775c1..bb59f2e94b6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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): diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 384d958946b..15953f4229b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -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 diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index ee38eb6515d..451f110a894 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -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 = """ -
- ⚠️ Invalid username or password. Please try again. -
- """ - - 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""" -
-

or

- - 🔐 Login with SSO - -
- """ - - 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""" - - - - - LiteLLM Login - - - - -
-
- -
-

Login

-

Access your LiteLLM Admin UI.

- - {error_message} - -
-
- - - - - - Default Credentials -
-

By default, Username is admin and Password is your set LiteLLM Proxy MASTER_KEY.

-

Need to set UI credentials or SSO? Check the documentation.

-
- - - - - - -
- - -
- - - {sso_button} -
- - - - """ - - 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) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index c8d148f7086..aa63369ef2b 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -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 \ No newline at end of file diff --git a/litellm/proxy/spend_tracking/cold_storage_handler.py b/litellm/proxy/spend_tracking/cold_storage_handler.py new file mode 100644 index 00000000000..21e785425ad --- /dev/null +++ b/litellm/proxy/spend_tracking/cold_storage_handler.py @@ -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) \ No newline at end of file diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index ad5cad29e60..653426e8bd2 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -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 diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py new file mode 100644 index 00000000000..28e76ca893f --- /dev/null +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -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 diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 13791666044..3706b2f7fc9 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -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): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6606b0e0b89..75c7d28460b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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): diff --git a/litellm/utils.py b/litellm/utils.py index ffd8bee382e..64d5f04a971 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9120f1f079f..0d4f1d1dfba 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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" } } diff --git a/poetry.lock b/poetry.lock index 2a1a664ee66..16b72b38887 100644 --- a/poetry.lock +++ b/poetry.lock @@ -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" diff --git a/pyproject.toml b/pyproject.toml index c4a0ed0477a..32baa74dc71 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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} diff --git a/requirements.txt b/requirements.txt index 84eb146aee2..6038bb5dab7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index ae8138f057c..158399305b5 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -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. ] diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index 46efd738b64..db403f81386 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -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, diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index d5a7271ae20..398a57e340f 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -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, ) diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index 184ede2222f..0f25ef4b6e3 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -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 diff --git a/tests/logging_callback_tests/test_amazing_s3_logs.py b/tests/logging_callback_tests/test_amazing_s3_logs.py index 37666d72b79..c9d0987e86f 100644 --- a/tests/logging_callback_tests/test_amazing_s3_logs.py +++ b/tests/logging_callback_tests/test_amazing_s3_logs.py @@ -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) + diff --git a/tests/logging_callback_tests/test_gcs_pub_sub.py b/tests/logging_callback_tests/test_gcs_pub_sub.py index f231d01d3fb..4172659e659 100644 --- a/tests/logging_callback_tests/test_gcs_pub_sub.py +++ b/tests/logging_callback_tests/test_gcs_pub_sub.py @@ -38,6 +38,7 @@ ignored_keys = [ "endTime", "metadata.model_map_information", "metadata.usage_object", + "metadata.cold_storage_object_key", ] diff --git a/tests/logging_callback_tests/test_otel_logging.py b/tests/logging_callback_tests/test_otel_logging.py index ff9d8300fef..aeb42bdaf79 100644 --- a/tests/logging_callback_tests/test_otel_logging.py +++ b/tests/logging_callback_tests/test_otel_logging.py @@ -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( diff --git a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py index 3c9d31890c7..a1760ca6371 100644 --- a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py +++ b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py @@ -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", diff --git a/tests/test_litellm/integrations/test_braintrust_logging.py b/tests/test_litellm/integrations/test_braintrust_logging.py new file mode 100644 index 00000000000..5ae40e82760 --- /dev/null +++ b/tests/test_litellm/integrations/test_braintrust_logging.py @@ -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)) \ No newline at end of file diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index e9b962f189f..317491d3172 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py new file mode 100644 index 00000000000..9e44a3eb419 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -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") diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 245f350be1b..53cd5a31aff 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -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" diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 93be9717e96..5b1e784caa3 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -29,6 +29,7 @@ ignored_keys = [ "endTime", "metadata.model_map_information", "metadata.usage_object", + "metadata.cold_storage_object_key", ] diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 26f41b76d3b..eac631bdb0d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -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)" ) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 00c55b9f60f..8e38011279d 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -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 diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_session_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py similarity index 58% rename from tests/enterprise/litellm_enterprise/enterprise_callbacks/test_session_handler.py rename to tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py index dbe163560b0..27edfef3eb5 100644 --- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_session_handler.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py @@ -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 diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py new file mode 100644 index 00000000000..976152db353 --- /dev/null +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py @@ -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") \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/entity_usage.tsx b/ui/litellm-dashboard/src/components/entity_usage.tsx index a481de0868d..7b910dab57b 100644 --- a/ui/litellm-dashboard/src/components/entity_usage.tsx +++ b/ui/litellm-dashboard/src/components/entity_usage.tsx @@ -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 = ({ {getProviderSpend().map((provider) => ( - {provider.provider} + +
+ {provider.provider && ( + {`${provider.provider} { + 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); + } + }} + /> + )} + {provider.provider} +
+ ${formatNumberWithCommas(provider.spend, 2)} diff --git a/ui/litellm-dashboard/src/components/new_usage.tsx b/ui/litellm-dashboard/src/components/new_usage.tsx index 3dda412e203..ec9d462fb92 100644 --- a/ui/litellm-dashboard/src/components/new_usage.tsx +++ b/ui/litellm-dashboard/src/components/new_usage.tsx @@ -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 = ({ accessToken, userRole, user {all_admin_roles.includes(userRole || "") ? Global Usage : Your Usage} Team Usage - {all_admin_roles.includes(userRole || "") ? ( - Tag Usage - ) : ( - <> - )} - {all_admin_roles.includes(userRole || "") ? ( - User Agent Activity - ) : ( - <> - )} + {all_admin_roles.includes(userRole || "") ? Tag Usage : <>} + {all_admin_roles.includes(userRole || "") ? User Agent Activity : <>} {/* Your Usage Panel */} @@ -658,7 +645,29 @@ const NewUsagePage: React.FC = ({ accessToken, userRole, user .filter((provider) => provider.spend > 0) .map((provider) => ( - {provider.provider} + +
+ {provider.provider && ( + {`${provider.provider} { + 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) + } + }} + /> + )} + {provider.provider} +
+ ${formatNumberWithCommas(provider.spend, 2)} {provider.successful_requests.toLocaleString()} @@ -725,10 +734,7 @@ const NewUsagePage: React.FC = ({ accessToken, userRole, user {/* User Agent Activity Panel */} - + diff --git a/ui/litellm-dashboard/src/components/view_logs/SessionView.tsx b/ui/litellm-dashboard/src/components/view_logs/SessionView.tsx index 1bd8a51192f..2d05680040b 100644 --- a/ui/litellm-dashboard/src/components/view_logs/SessionView.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/SessionView.tsx @@ -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 = ({ 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 = ({ sessionId, logs, onBac Total Cost ${formatNumberWithCommas(totalCost, 6)} - - Total Tokens - {totalTokens} - + +
Usage breakdown
+
+
+
Input usage:
+
+
+ input: + {formatNumberWithCommas(logs.reduce((sum, log) => sum + (log.prompt_tokens || 0), 0))} +
+ {totalCacheReadTokens > 0 && ( +
+ input_cached_tokens: + {formatNumberWithCommas(totalCacheReadTokens)} +
+ )} + {totalCacheCreationTokens > 0 && ( +
+ input_cache_creation_tokens: + {formatNumberWithCommas(totalCacheCreationTokens)} +
+ )} +
+
+
+
Output usage:
+
+
+ output: + {formatNumberWithCommas(logs.reduce((sum, log) => sum + (log.completion_tokens || 0), 0))} +
+
+
+
+
+ Total usage: + {formatNumberWithCommas(totalTokensWithCache)} +
+
+
+ + } + placement="top" + overlayStyle={{ minWidth: '300px' }} + > + +
+ Total Tokens + ⓘ +
+ {formatNumberWithCommas(totalTokensWithCache)} + + {/* Request Timeline */} Session Logs diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx index ea97d54a6f1..d2041b389fd 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx @@ -925,7 +925,15 @@ export function RequestViewer({ row }: { row: Row }) {
Tokens: - {row.original.total_tokens} ({row.original.prompt_tokens}+{row.original.completion_tokens}) + {row.original.total_tokens} ({row.original.prompt_tokens} prompt tokens + {row.original.completion_tokens} completion tokens) +
+
+ Cache Read Tokens: + {formatNumberWithCommas(row.original.metadata?.additional_usage_values?.cache_read_input_tokens || 0)} +
+
+ Cache Creation Tokens: + {formatNumberWithCommas(row.original.metadata?.additional_usage_values.cache_creation_input_tokens)}
Cost: