tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-12 23:05:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

(feat) Audio transcription - cost tracking + (feat) image generation - accurate cost tracking based on output_format/quality/size

Browse Source 
* feat(audio_transcriptions/): calculate duration of audio file for cost calculation

Fixes https://github.com/BerriAI/litellm/issues/11846

Closes https://github.com/BerriAI/litellm/issues/14605

* fix(cost_calculator.py): correctly use base model, when set

Fixes issue where azure base model was being ignored

* feat(cost_calculator.py): fix default cost tracking quality param for image generation

* feat(image_generations/): return output_format, quality, size

aligns response to openai spec and improves cost tracking accuracy

* fix(cost_calculator.py): refactor cost calculation for image generation to use image response instead of hidden params

* build: update build

* fix: fix cost calculation

* build: update poetry lock

* fix: fix ruff checks

* fix: fix aembedding

* fix: fix ruff errors

* fix: modify to catch errors

* fix: test

* fix: loosen test to handle openai lib out of sync
This commit is contained in:
Krish Dholakia
2025-11-08 15:30:46 -08:00
committed by GitHub
parent 0a1fc0eeb2
commit c96da44265
47 changed files with 1483 additions and 1866 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.circleci/config.yml
+3 -3
View File
@@ -532,7 +532,7 @@ jobs:
command: |
pwd
ls
python -m pytest -vv tests/router_unit_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5
python -m pytest -vv tests/router_unit_tests --cov=litellm --cov-report=xml -x -s --junitxml=test-results/junit.xml --durations=5
no_output_timeout: 120m
- run:
name: Rename the coverage files
@@ -1164,7 +1164,7 @@ jobs:
command: |
pwd
ls
python -m pytest -vv tests/test_litellm --cov=litellm --cov-report=xml -s -v --junitxml=test-results/junit-litellm.xml --durations=10 -n 8
python -m pytest -vv tests/test_litellm --cov=litellm --cov-report=xml -v --junitxml=test-results/junit-litellm.xml --durations=10 -n 8
no_output_timeout: 120m
- run:
name: Rename the coverage files
@@ -1396,7 +1396,7 @@ jobs:
command: |
pwd
ls
python -m pytest -vv tests/image_gen_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5
python -m pytest -vv tests/image_gen_tests --cov=litellm --cov-report=xml -x -v --junitxml=test-results/junit.xml --durations=5
no_output_timeout: 120m
- run:
name: Rename the coverage files
enterprise/enterprise_hooks/aporia_ai.py
+3 -10
View File
@@ -8,6 +8,8 @@
import os
import sys
from litellm.types.utils import CallTypesLiteral
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
@@ -166,16 +168,7 @@ class AporiaGuardrail(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
enterprise/enterprise_hooks/google_text_moderation.py
+2 -12
View File
@@ -6,14 +6,13 @@
# +-----------------------------------------------+
# Thank you users! We ❤️ you! - Krrish & Ishaan
from typing import Literal
from fastapi import HTTPException
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import CallTypesLiteral
class _ENTERPRISE_GoogleTextModeration(CustomLogger):
@@ -89,16 +88,7 @@ class _ENTERPRISE_GoogleTextModeration(CustomLogger):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
"""
- Calls Google's Text Moderation API
enterprise/enterprise_hooks/openai_moderation.py
+2 -11
View File
@@ -12,7 +12,6 @@ sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import sys
from typing import Literal
from fastapi import HTTPException
@@ -20,6 +19,7 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import CallTypesLiteral
class _ENTERPRISE_OpenAI_Moderation(CustomLogger):
@@ -35,16 +35,7 @@ class _ENTERPRISE_OpenAI_Moderation(CustomLogger):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
text = ""
if "messages" in data and isinstance(data["messages"], list):
enterprise/litellm_enterprise/enterprise_callbacks/llama_guard.py
+2 -11
View File
@@ -23,7 +23,7 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import Choices, ModelResponse
from litellm.types.utils import CallTypesLiteral, Choices, ModelResponse
class _ENTERPRISE_LlamaGuard(CustomLogger):
@@ -98,16 +98,7 @@ class _ENTERPRISE_LlamaGuard(CustomLogger):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
"""
- Calls the Llama Guard Endpoint
enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py
+2 -10
View File
@@ -17,6 +17,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.secret_managers.main import get_secret_str
from litellm.types.utils import CallTypesLiteral
from litellm.utils import get_formatted_prompt
@@ -120,16 +121,7 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
"""
- Calls the LLM Guard Endpoint
enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py
+2 -12
View File
@@ -31,6 +31,7 @@ from litellm.types.integrations.pagerduty import (
PagerDutyRequestBody,
)
from litellm.types.utils import (
CallTypesLiteral,
StandardLoggingPayload,
StandardLoggingPayloadErrorInformation,
)
@@ -142,18 +143,7 @@ class PagerDutyAlerting(SlackAlerting):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Optional[Union[Exception, str, dict]]:
"""
Example of detecting hanging requests by waiting a given threshold.
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
+2 -22
View File
@@ -36,6 +36,7 @@ from litellm.types.llms.openai import (
OpenAIFilesPurpose,
)
from litellm.types.utils import (
CallTypesLiteral,
LiteLLMBatch,
LiteLLMFineTuningJob,
LLMResponseTypes,
@@ -272,28 +273,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: Dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"acreate_batch",
"aretrieve_batch",
"acreate_file",
"afile_list",
"afile_delete",
"afile_content",
"acreate_fine_tuning_job",
"aretrieve_fine_tuning_job",
"alist_fine_tuning_jobs",
"acancel_fine_tuning_job",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Union[Exception, str, Dict, None]:
"""
- Detect litellm_proxy/ file_id
litellm/cost_calculator.py
+34 -30
View File
@@ -327,12 +327,16 @@ def cost_per_token( # noqa: PLR0915
elif call_type == "search" or call_type == "asearch":
# Search providers use per-query pricing
from litellm.search import search_provider_cost_per_query
return search_provider_cost_per_query(
model=model,
custom_llm_provider=custom_llm_provider,
number_of_queries=number_of_queries or 1,
optional_params=response._hidden_params if response and hasattr(response, "_hidden_params") else None
optional_params=(
response._hidden_params
if response and hasattr(response, "_hidden_params")
else None
),
)
elif custom_llm_provider == "vertex_ai":
cost_router = google_cost_router(
@@ -509,16 +513,18 @@ def _select_model_name_for_cost_calc(
else:
return_model = model
if base_model is not None:
elif base_model is not None:
return_model = base_model
if completion_response_model is None and hidden_params is not None:
elif completion_response_model is None and hidden_params is not None:
if (
hidden_params.get("model", None) is not None
and len(hidden_params["model"]) > 0
):
return_model = hidden_params.get("model", model)
if hidden_params is not None and hidden_params.get("region_name", None) is not None:
elif (
hidden_params is not None and hidden_params.get("region_name", None) is not None
):
region_name = hidden_params.get("region_name", None)
if return_model is None and completion_response_model is not None:
@@ -853,15 +859,6 @@ def completion_cost( # noqa: PLR0915
"custom_llm_provider", custom_llm_provider or None
)
region_name = hidden_params.get("region_name", region_name)
size = hidden_params.get("optional_params", {}).get(
"size", "1024-x-1024"
) # openai default
quality = hidden_params.get("optional_params", {}).get(
"quality", "standard"
) # openai default
n = hidden_params.get("optional_params", {}).get(
"n", 1
) # openai default
else:
if model is None:
raise ValueError(
@@ -888,7 +885,9 @@ def completion_cost( # noqa: PLR0915
str(e)
)
)
if CostCalculatorUtils._call_type_has_image_response(call_type):
if CostCalculatorUtils._call_type_has_image_response(
call_type
) and isinstance(completion_response, ImageResponse):
### IMAGE GENERATION COST CALCULATION ###
return CostCalculatorUtils.route_image_generation_cost_calculator(
model=model,
@@ -906,27 +905,32 @@ def completion_cost( # noqa: PLR0915
or call_type == CallTypes.avideo_remix.value
):
### VIDEO GENERATION COST CALCULATION ###
if completion_response is not None and hasattr(completion_response, 'usage'):
usage_obj = completion_response.usage
usage_obj = getattr(completion_response, "usage", None)
if completion_response is not None and usage_obj:
# Handle both dict and Pydantic Usage object
if isinstance(usage_obj, dict):
duration_seconds = usage_obj.get('duration_seconds', None)
duration_seconds = usage_obj.get("duration_seconds", None)
else:
duration_seconds = getattr(usage_obj, 'duration_seconds', None)
duration_seconds = getattr(
usage_obj, "duration_seconds", None
)
if duration_seconds is not None:
# Calculate cost based on video duration using video-specific cost calculation
from litellm.llms.openai.cost_calculation import video_generation_cost
from litellm.llms.openai.cost_calculation import (
video_generation_cost,
)
return video_generation_cost(
model=model,
duration_seconds=duration_seconds,
custom_llm_provider=custom_llm_provider
custom_llm_provider=custom_llm_provider,
)
# Fallback to default video cost calculation if no duration available
return default_video_cost_calculator(
model=model,
duration_seconds=0.0, # Default to 0 if no duration available
custom_llm_provider=custom_llm_provider
custom_llm_provider=custom_llm_provider,
)
elif (
call_type == CallTypes.speech.value
@@ -1460,13 +1464,13 @@ def default_video_cost_calculator(
model_name_without_custom_llm_provider = model.replace(
f"{custom_llm_provider}/", ""
)
base_model_name = f"{custom_llm_provider}/{model_name_without_custom_llm_provider}"
base_model_name = (
f"{custom_llm_provider}/{model_name_without_custom_llm_provider}"
)
verbose_logger.debug(
f"Looking up cost for video model: {base_model_name}"
)
verbose_logger.debug(f"Looking up cost for video model: {base_model_name}")
model_without_provider = model.split('/')[-1]
model_without_provider = model.split("/")[-1]
# Try model with provider first, fall back to base model name
cost_info: Optional[dict] = None
@@ -1480,7 +1484,7 @@ def default_video_cost_calculator(
if _model is not None and _model in litellm.model_cost:
cost_info = litellm.model_cost[_model]
break
# If still not found, try with custom_llm_provider prefix
if cost_info is None and custom_llm_provider:
prefixed_model = f"{custom_llm_provider}/{model}"
@@ -1495,12 +1499,12 @@ def default_video_cost_calculator(
video_cost_per_second = cost_info.get("output_cost_per_video_per_second")
if video_cost_per_second is not None:
return video_cost_per_second * duration_seconds
# Fallback to general output cost per second
output_cost_per_second = cost_info.get("output_cost_per_second")
if output_cost_per_second is not None:
return output_cost_per_second * duration_seconds
# If no cost information found, return 0
verbose_logger.info(
f"No cost information found for video model {model}. Please add pricing to model_prices_and_context_window.json"
litellm/integrations/custom_logger.py
+43 -45
View File
@@ -8,7 +8,6 @@ from typing import (
AsyncGenerator,
Dict,
List,
Literal,
Optional,
Tuple,
Union,
@@ -24,6 +23,7 @@ from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest
from litellm.types.utils import (
AdapterCompletionStreamWrapper,
CallTypes,
CallTypesLiteral,
LLMResponseTypes,
ModelResponse,
ModelResponseStream,
@@ -65,12 +65,11 @@ _BASE64_INLINE_PATTERN = re.compile(
class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callback#callback-class
# Class variables or attributes
def __init__(
self,
self,
turn_off_message_logging: bool = False,
# deprecated param, use `turn_off_message_logging` instead
message_logging: bool = True,
**kwargs
**kwargs,
) -> None:
"""
Args:
@@ -221,7 +220,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
) -> Optional[Any]:
"""
Allow modifying streaming chunks just before they're returned to the user.
This is called for each streaming chunk in the response.
"""
pass
@@ -292,18 +291,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Optional[
Union[Exception, str, dict]
]: # raise exception if invalid, return a str for the user to receive - if rejected, or return a modified dictionary for passing into litellm
@@ -342,16 +330,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Any:
pass
@@ -435,7 +414,6 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
# MCP TOOL CALL HOOKS
#########################################################
async def async_post_mcp_tool_call_hook(
self, kwargs, response_obj: MCPPostCallResponseObject, start_time, end_time
) -> Optional[MCPPostCallResponseObject]:
@@ -519,18 +497,18 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
if LITELLM_METADATA_FIELD in request_kwargs:
return LITELLM_METADATA_FIELD
return OLD_LITELLM_METADATA_FIELD
def redact_standard_logging_payload_from_model_call_details(
self, model_call_details: Dict
) -> Dict:
"""
Only redacts messages and responses when self.turn_off_message_logging is True
By default, self.turn_off_message_logging is False and this does nothing.
Return a redacted deepcopy of the provided logging payload.
This is useful for logging payloads that contain sensitive information.
"""
from copy import copy
@@ -540,7 +518,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
if turn_off_message_logging is False:
return model_call_details
# Only make a shallow copy of the top-level dict to avoid deepcopy issues
# with complex objects like AuthenticationError that may be present
model_call_details_copy = copy(model_call_details)
@@ -553,7 +531,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
standard_logging_object_copy = copy(standard_logging_object)
if standard_logging_object_copy.get("messages") is not None:
standard_logging_object_copy["messages"] = [Message(content=redacted_str).model_dump()]
standard_logging_object_copy["messages"] = [
Message(content=redacted_str).model_dump()
]
if standard_logging_object_copy.get("response") is not None:
response = standard_logging_object_copy["response"]
@@ -580,11 +560,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
model_response_dict = model_response.model_dump()
standard_logging_object_copy["response"] = model_response_dict
model_call_details_copy["standard_logging_object"] = standard_logging_object_copy
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,
@@ -622,7 +602,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
verbose_logger.debug(f"Error in handle_callback_failure for {callback_name}: {str(e)}")
async def _strip_base64_from_messages(
self, payload: "StandardLoggingPayload", max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER
self,
payload: "StandardLoggingPayload",
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER,
) -> "StandardLoggingPayload":
"""
Removes or redacts base64-encoded file data (e.g., PDFs, images, audio)
@@ -671,7 +653,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
verbose_logger.debug(f"[CustomLogger] Stripping base64 from {len(messages)} messages")
if messages:
payload["messages"] = self._process_messages(messages=messages, max_depth=max_depth)
payload["messages"] = self._process_messages(
messages=messages, max_depth=max_depth
)
total_items = 0
for m in payload.get("messages", []) or []:
@@ -684,9 +668,13 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
f"[CustomLogger] Completed base64 strip; retained {total_items} content items"
)
return payload
def _redact_base64(self, value: Any, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER) -> Any:
def _redact_base64(
self,
value: Any,
depth: int = 0,
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER,
) -> Any:
"""Recursively redact inline base64 from any nested structure with a max recursion depth limit."""
if depth > max_depth:
verbose_logger.warning(
@@ -703,10 +691,16 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
return value
if isinstance(value, list):
return [self._redact_base64(value=v, depth=depth + 1, max_depth=max_depth) for v in value]
return [
self._redact_base64(value=v, depth=depth + 1, max_depth=max_depth)
for v in value
]
if isinstance(value, dict):
return {k: self._redact_base64(value=v, depth=depth + 1, max_depth=max_depth) for k, v in value.items()}
return {
k: self._redact_base64(value=v, depth=depth + 1, max_depth=max_depth)
for k, v in value.items()
}
return value
@@ -729,10 +723,14 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
cleaned: List[Any] = []
for c in contents:
if self._should_keep_content(content=c):
cleaned.append(self._redact_base64(value=c, max_depth=max_depth))
cleaned.append(
self._redact_base64(value=c, max_depth=max_depth)
)
msg["content"] = cleaned
else:
msg["content"] = self._redact_base64(value=contents, max_depth=max_depth)
msg["content"] = self._redact_base64(
value=contents, max_depth=max_depth
)
for key, val in list(msg.items()):
if key != "content":
litellm/litellm_core_utils/audio_utils/utils.py
+96 -23
View File
@@ -4,6 +4,7 @@ Utils used for litellm.transcription() and litellm.atranscription()
import os
from dataclasses import dataclass
from typing import Optional
from litellm.types.files import get_file_mime_type_from_extension
from litellm.types.utils import FileTypes
@@ -13,12 +14,13 @@ from litellm.types.utils import FileTypes
class ProcessedAudioFile:
"""
Processed audio file data.
Attributes:
file_content: The binary content of the audio file
filename: The filename (extracted or generated)
content_type: The MIME type of the audio file
"""
file_content: bytes
filename: str
content_type: str
@@ -27,61 +29,63 @@ class ProcessedAudioFile:
def process_audio_file(audio_file: FileTypes) -> ProcessedAudioFile:
"""
Common utility function to process audio files for audio transcription APIs.
Handles various input types:
- File paths (str, os.PathLike)
- Raw bytes/bytearray
- Tuples (filename, content, optional content_type)
- File-like objects with read() method
Args:
audio_file: The audio file input in various formats
Returns:
ProcessedAudioFile: Structured data with file content, filename, and content type
Raises:
ValueError: If audio_file type is unsupported or content cannot be extracted
"""
file_content = None
filename = None
if isinstance(audio_file, (bytes, bytearray)):
# Raw bytes
filename = 'audio.wav'
filename = "audio.wav"
file_content = bytes(audio_file)
elif isinstance(audio_file, (str, os.PathLike)):
# File path or PathLike
file_path = str(audio_file)
with open(file_path, 'rb') as f:
with open(file_path, "rb") as f:
file_content = f.read()
filename = file_path.split('/')[-1]
filename = file_path.split("/")[-1]
elif isinstance(audio_file, tuple):
# Tuple format: (filename, content, content_type) or (filename, content)
if len(audio_file) >= 2:
filename = audio_file[0] or 'audio.wav'
filename = audio_file[0] or "audio.wav"
content = audio_file[1]
if isinstance(content, (bytes, bytearray)):
file_content = bytes(content)
elif isinstance(content, (str, os.PathLike)):
# File path or PathLike
with open(str(content), 'rb') as f:
with open(str(content), "rb") as f:
file_content = f.read()
elif hasattr(content, 'read'):
elif hasattr(content, "read"):
# File-like object
file_content = content.read()
if hasattr(content, 'seek'):
if hasattr(content, "seek"):
content.seek(0)
else:
raise ValueError(f"Unsupported content type in tuple: {type(content)}")
else:
raise ValueError("Tuple must have at least 2 elements: (filename, content)")
elif hasattr(audio_file, 'read') and not isinstance(audio_file, (str, bytes, bytearray, tuple, os.PathLike)):
elif hasattr(audio_file, "read") and not isinstance(
audio_file, (str, bytes, bytearray, tuple, os.PathLike)
):
# File-like object (IO) - check this after all other types
filename = getattr(audio_file, 'name', 'audio.wav')
filename = getattr(audio_file, "name", "audio.wav")
file_content = audio_file.read() # type: ignore
# Reset file pointer if possible
if hasattr(audio_file, 'seek'):
if hasattr(audio_file, "seek"):
audio_file.seek(0) # type: ignore
else:
raise ValueError(f"Unsupported audio_file type: {type(audio_file)}")
@@ -90,20 +94,18 @@ def process_audio_file(audio_file: FileTypes) -> ProcessedAudioFile:
raise ValueError("Could not extract file content from audio_file")
# Determine content type using LiteLLM's file type utilities
content_type = 'audio/wav' # Default fallback
content_type = "audio/wav" # Default fallback
if filename:
try:
# Extract extension from filename
extension = filename.split('.')[-1].lower() if '.' in filename else 'wav'
extension = filename.split(".")[-1].lower() if "." in filename else "wav"
content_type = get_file_mime_type_from_extension(extension)
except ValueError:
# If extension is not recognized, fallback to audio/wav
content_type = 'audio/wav'
content_type = "audio/wav"
return ProcessedAudioFile(
file_content=file_content,
filename=filename,
content_type=content_type
file_content=file_content, filename=filename, content_type=content_type
)
@@ -134,3 +136,74 @@ def get_audio_file_for_health_check() -> FileTypes:
pwd = os.path.dirname(os.path.realpath(__file__))
file_path = os.path.join(pwd, "audio_health_check.wav")
return open(file_path, "rb")
def calculate_request_duration(file: FileTypes) -> Optional[float]:
"""
Calculate audio duration from file content.
Args:
file: The audio file (can be file path, bytes, or file-like object)
Returns:
Duration in seconds, or None if extraction fails or soundfile is not available
"""
try:
import soundfile as sf
except ImportError:
# soundfile not available, cannot extract duration
return None
try:
import io
# Handle different file input types
file_content: Optional[bytes] = None
if isinstance(file, (bytes, bytearray)):
# Raw bytes
file_content = bytes(file)
elif isinstance(file, (str, os.PathLike)):
# File path
with open(str(file), "rb") as f:
file_content = f.read()
elif isinstance(file, tuple):
# Tuple format: (filename, content, optional content_type)
if len(file) >= 2:
content = file[1]
if isinstance(content, bytes):
file_content = content
elif hasattr(content, "read") and not isinstance(
content, (str, os.PathLike)
):
# File-like object in tuple
current_pos = getattr(content, "tell", lambda: None)()
# Seek to start to ensure we read the entire content
if hasattr(content, "seek"):
content.seek(0)
file_content = content.read()
if current_pos is not None and hasattr(content, "seek"):
content.seek(current_pos)
elif hasattr(file, "read") and not isinstance(file, tuple):
# File-like object (including BytesIO)
current_position = file.tell() if hasattr(file, "tell") else None
# Seek to start to ensure we read the entire content
if hasattr(file, "seek"):
file.seek(0)
file_content = file.read()
# Reset file position if possible
if current_position is not None and hasattr(file, "seek"):
file.seek(current_position)
if file_content is None or not isinstance(file_content, bytes):
return None
# Extract duration using soundfile
file_object = io.BytesIO(file_content)
with sf.SoundFile(file_object) as audio:
duration = len(audio) / audio.samplerate
return duration
except Exception:
# Silently fail if duration extraction fails
return None
litellm/litellm_core_utils/litellm_logging.py
+185 -166
View File
@@ -308,9 +308,9 @@ class Logging(LiteLLMLoggingBaseClass):
self.litellm_trace_id: str = litellm_trace_id or str(uuid.uuid4())
self.function_id = function_id
self.streaming_chunks: List[Any] = [] # for generating complete stream response
self.sync_streaming_chunks: List[
Any
] = [] # for generating complete stream response
self.sync_streaming_chunks: List[Any] = (
[]
) # for generating complete stream response
self.log_raw_request_response = log_raw_request_response
# Initialize dynamic callbacks
@@ -686,9 +686,9 @@ class Logging(LiteLLMLoggingBaseClass):
if anthropic_cache_control_logger := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook(
non_default_params
):
self.model_call_details[
"prompt_integration"
] = anthropic_cache_control_logger.__class__.__name__
self.model_call_details["prompt_integration"] = (
anthropic_cache_control_logger.__class__.__name__
)
return anthropic_cache_control_logger
#########################################################
@@ -700,9 +700,9 @@ class Logging(LiteLLMLoggingBaseClass):
internal_usage_cache=None,
llm_router=None,
)
self.model_call_details[
"prompt_integration"
] = vector_store_custom_logger.__class__.__name__
self.model_call_details["prompt_integration"] = (
vector_store_custom_logger.__class__.__name__
)
# Add to global callbacks so post-call hooks are invoked
if (
vector_store_custom_logger
@@ -762,9 +762,9 @@ class Logging(LiteLLMLoggingBaseClass):
model
): # if model name was changes pre-call, overwrite the initial model call name with the new one
self.model_call_details["model"] = model
self.model_call_details["litellm_params"][
"api_base"
] = self._get_masked_api_base(additional_args.get("api_base", ""))
self.model_call_details["litellm_params"]["api_base"] = (
self._get_masked_api_base(additional_args.get("api_base", ""))
)
def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR0915
# Log the exact input to the LLM API
@@ -793,10 +793,10 @@ class Logging(LiteLLMLoggingBaseClass):
try:
# [Non-blocking Extra Debug Information in metadata]
if turn_off_message_logging is True:
_metadata[
"raw_request"
] = "redacted by litellm. \
_metadata["raw_request"] = (
"redacted by litellm. \
'litellm.turn_off_message_logging=True'"
)
else:
curl_command = self._get_request_curl_command(
api_base=additional_args.get("api_base", ""),
@@ -807,32 +807,32 @@ class Logging(LiteLLMLoggingBaseClass):
_metadata["raw_request"] = str(curl_command)
# split up, so it's easier to parse in the UI
self.model_call_details[
"raw_request_typed_dict"
] = RawRequestTypedDict(
raw_request_api_base=str(
additional_args.get("api_base") or ""
),
raw_request_body=self._get_raw_request_body(
additional_args.get("complete_input_dict", {})
),
raw_request_headers=self._get_masked_headers(
additional_args.get("headers", {}) or {},
ignore_sensitive_headers=True,
),
error=None,
self.model_call_details["raw_request_typed_dict"] = (
RawRequestTypedDict(
raw_request_api_base=str(
additional_args.get("api_base") or ""
),
raw_request_body=self._get_raw_request_body(
additional_args.get("complete_input_dict", {})
),
raw_request_headers=self._get_masked_headers(
additional_args.get("headers", {}) or {},
ignore_sensitive_headers=True,
),
error=None,
)
)
except Exception as e:
self.model_call_details[
"raw_request_typed_dict"
] = RawRequestTypedDict(
error=str(e),
self.model_call_details["raw_request_typed_dict"] = (
RawRequestTypedDict(
error=str(e),
)
)
_metadata[
"raw_request"
] = "Unable to Log \
_metadata["raw_request"] = (
"Unable to Log \
raw request: {}".format(
str(e)
str(e)
)
)
if getattr(self, "logger_fn", None) and callable(self.logger_fn):
try:
@@ -1133,13 +1133,13 @@ class Logging(LiteLLMLoggingBaseClass):
for callback in callbacks:
try:
if isinstance(callback, CustomLogger):
response: Optional[
MCPPostCallResponseObject
] = await callback.async_post_mcp_tool_call_hook(
kwargs=kwargs,
response_obj=post_mcp_tool_call_response_obj,
start_time=start_time,
end_time=end_time,
response: Optional[MCPPostCallResponseObject] = (
await callback.async_post_mcp_tool_call_hook(
kwargs=kwargs,
response_obj=post_mcp_tool_call_response_obj,
start_time=start_time,
end_time=end_time,
)
)
######################################################################
# if any of the callbacks modify the response, use the modified response
@@ -1243,6 +1243,7 @@ class Logging(LiteLLMLoggingBaseClass):
used for consistent cost calculation across response headers + logging integrations.
"""
if isinstance(result, BaseModel) and hasattr(result, "_hidden_params"):
hidden_params = getattr(result, "_hidden_params", {})
if (
@@ -1302,9 +1303,9 @@ class Logging(LiteLLMLoggingBaseClass):
verbose_logger.debug(
f"response_cost_failure_debug_information: {debug_info}"
)
self.model_call_details[
"response_cost_failure_debug_information"
] = debug_info
self.model_call_details["response_cost_failure_debug_information"] = (
debug_info
)
return None
try:
@@ -1330,9 +1331,9 @@ class Logging(LiteLLMLoggingBaseClass):
verbose_logger.debug(
f"response_cost_failure_debug_information: {debug_info}"
)
self.model_call_details[
"response_cost_failure_debug_information"
] = debug_info
self.model_call_details["response_cost_failure_debug_information"] = (
debug_info
)
return None
@@ -1476,9 +1477,9 @@ class Logging(LiteLLMLoggingBaseClass):
end_time = datetime.datetime.now()
if self.completion_start_time is None:
self.completion_start_time = end_time
self.model_call_details[
"completion_start_time"
] = self.completion_start_time
self.model_call_details["completion_start_time"] = (
self.completion_start_time
)
self.model_call_details["log_event_type"] = "successful_api_call"
self.model_call_details["end_time"] = end_time
self.model_call_details["cache_hit"] = cache_hit
@@ -1531,39 +1532,39 @@ class Logging(LiteLLMLoggingBaseClass):
"response_cost"
]
else:
self.model_call_details[
"response_cost"
] = self._response_cost_calculator(result=logging_result)
self.model_call_details["response_cost"] = (
self._response_cost_calculator(result=logging_result)
)
## STANDARDIZED LOGGING PAYLOAD
self.model_call_details[
"standard_logging_object"
] = get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=logging_result,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
self.model_call_details["standard_logging_object"] = (
get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=logging_result,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
)
elif isinstance(result, dict) or isinstance(result, list):
## STANDARDIZED LOGGING PAYLOAD
self.model_call_details[
"standard_logging_object"
] = get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=result,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
self.model_call_details["standard_logging_object"] = (
get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=result,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
)
elif standard_logging_object is not None:
self.model_call_details[
"standard_logging_object"
] = standard_logging_object
self.model_call_details["standard_logging_object"] = (
standard_logging_object
)
else: # streaming chunks + image gen.
self.model_call_details["response_cost"] = None
@@ -1571,17 +1572,31 @@ class Logging(LiteLLMLoggingBaseClass):
# MAP RESPONSES API USAGE OBJECT TO LITELLM USAGE OBJECT
if isinstance(result, ResponsesAPIResponse):
result = result.model_copy()
transformed_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
result.usage
transformed_usage = (
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
result.usage
)
)
# Set as dict instead of Usage object so model_dump() serializes it correctly
setattr(
result,
"usage",
transformed_usage.model_dump() if hasattr(transformed_usage, 'model_dump') else dict(transformed_usage),
(
transformed_usage.model_dump()
if hasattr(transformed_usage, "model_dump")
else dict(transformed_usage)
),
)
if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None:
standard_logging_payload["response"] = result.model_dump() if hasattr(result, 'model_dump') else dict(result)
if (
standard_logging_payload := self.model_call_details.get(
"standard_logging_object"
)
) is not None:
standard_logging_payload["response"] = (
result.model_dump()
if hasattr(result, "model_dump")
else dict(result)
)
if (
litellm.max_budget
@@ -1629,7 +1644,7 @@ class Logging(LiteLLMLoggingBaseClass):
or isinstance(logging_result, OCRResponse) # OCR
or isinstance(logging_result, dict)
and logging_result.get("object") == "vector_store.search_results.page"
or isinstance(logging_result, VideoObject)
or isinstance(logging_result, VideoObject)
or isinstance(logging_result, ContainerObject)
or (self.call_type == CallTypes.call_mcp_tool.value)
):
@@ -1723,23 +1738,23 @@ class Logging(LiteLLMLoggingBaseClass):
verbose_logger.debug(
"Logging Details LiteLLM-Success Call streaming complete"
)
self.model_call_details[
"complete_streaming_response"
] = complete_streaming_response
self.model_call_details[
"response_cost"
] = self._response_cost_calculator(result=complete_streaming_response)
self.model_call_details["complete_streaming_response"] = (
complete_streaming_response
)
self.model_call_details["response_cost"] = (
self._response_cost_calculator(result=complete_streaming_response)
)
## STANDARDIZED LOGGING PAYLOAD
self.model_call_details[
"standard_logging_object"
] = get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=complete_streaming_response,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
self.model_call_details["standard_logging_object"] = (
get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=complete_streaming_response,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
)
callbacks = self.get_combined_callback_list(
dynamic_success_callbacks=self.dynamic_success_callbacks,
@@ -2067,10 +2082,10 @@ class Logging(LiteLLMLoggingBaseClass):
)
else:
if self.stream and complete_streaming_response:
self.model_call_details[
"complete_response"
] = self.model_call_details.get(
"complete_streaming_response", {}
self.model_call_details["complete_response"] = (
self.model_call_details.get(
"complete_streaming_response", {}
)
)
result = self.model_call_details["complete_response"]
openMeterLogger.log_success_event(
@@ -2109,10 +2124,10 @@ class Logging(LiteLLMLoggingBaseClass):
)
else:
if self.stream and complete_streaming_response:
self.model_call_details[
"complete_response"
] = self.model_call_details.get(
"complete_streaming_response", {}
self.model_call_details["complete_response"] = (
self.model_call_details.get(
"complete_streaming_response", {}
)
)
result = self.model_call_details["complete_response"]
@@ -2255,9 +2270,9 @@ class Logging(LiteLLMLoggingBaseClass):
if complete_streaming_response is not None:
print_verbose("Async success callbacks: Got a complete streaming response")
self.model_call_details[
"async_complete_streaming_response"
] = complete_streaming_response
self.model_call_details["async_complete_streaming_response"] = (
complete_streaming_response
)
try:
if self.model_call_details.get("cache_hit", False) is True:
@@ -2268,10 +2283,10 @@ class Logging(LiteLLMLoggingBaseClass):
model_call_details=self.model_call_details
)
# base_model defaults to None if not set on model_info
self.model_call_details[
"response_cost"
] = self._response_cost_calculator(
result=complete_streaming_response
self.model_call_details["response_cost"] = (
self._response_cost_calculator(
result=complete_streaming_response
)
)
verbose_logger.debug(
@@ -2284,16 +2299,16 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["response_cost"] = None
## STANDARDIZED LOGGING PAYLOAD
self.model_call_details[
"standard_logging_object"
] = get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=complete_streaming_response,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
self.model_call_details["standard_logging_object"] = (
get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=complete_streaming_response,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
)
callbacks = self.get_combined_callback_list(
dynamic_success_callbacks=self.dynamic_async_success_callbacks,
@@ -2482,26 +2497,24 @@ class Logging(LiteLLMLoggingBaseClass):
def _handle_callback_failure(self, callback: Any):
"""
Handle callback logging failures by incrementing Prometheus metrics.
Works for both sync and async contexts since Prometheus counter increment is synchronous.
Args:
callback: The callback that failed
"""
try:
callback_name = self._get_callback_name(callback)
all_callbacks = litellm.logging_callback_manager._get_all_callbacks()
for callback_obj in all_callbacks:
if hasattr(callback_obj, 'increment_callback_logging_failure'):
if hasattr(callback_obj, "increment_callback_logging_failure"):
callback_obj.increment_callback_logging_failure(callback_name=callback_name) # type: ignore
break # Only increment once
except Exception as e:
verbose_logger.debug(
f"Error in _handle_callback_failure: {str(e)}"
)
verbose_logger.debug(f"Error in _handle_callback_failure: {str(e)}")
def _failure_handler_helper_fn(
self, exception, traceback_exception, start_time=None, end_time=None
@@ -2531,18 +2544,18 @@ class Logging(LiteLLMLoggingBaseClass):
## STANDARDIZED LOGGING PAYLOAD
self.model_call_details[
"standard_logging_object"
] = get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj={},
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="failure",
error_str=str(exception),
original_exception=exception,
standard_built_in_tools_params=self.standard_built_in_tools_params,
self.model_call_details["standard_logging_object"] = (
get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj={},
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="failure",
error_str=str(exception),
original_exception=exception,
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
)
return start_time, end_time
@@ -3040,14 +3053,20 @@ class Logging(LiteLLMLoggingBaseClass):
elif isinstance(result, ResponseCompletedEvent):
## return unified Usage object
if isinstance(result.response.usage, ResponseAPIUsage):
transformed_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
result.response.usage
transformed_usage = (
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
result.response.usage
)
)
# Set as dict instead of Usage object so model_dump() serializes it correctly
setattr(
result.response,
"usage",
transformed_usage.model_dump() if hasattr(transformed_usage, 'model_dump') else dict(transformed_usage),
(
transformed_usage.model_dump()
if hasattr(transformed_usage, "model_dump")
else dict(transformed_usage)
),
)
return result.response
else:
@@ -3447,9 +3466,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
endpoint=arize_config.endpoint,
)
os.environ[
"OTEL_EXPORTER_OTLP_TRACES_HEADERS"
] = f"space_id={arize_config.space_key},api_key={arize_config.api_key}"
os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
f"space_id={arize_config.space_key},api_key={arize_config.api_key}"
)
for callback in _in_memory_loggers:
if (
isinstance(callback, ArizeLogger)
@@ -3473,9 +3492,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
# auth can be disabled on local deployments of arize phoenix
if arize_phoenix_config.otlp_auth_headers is not None:
os.environ[
"OTEL_EXPORTER_OTLP_TRACES_HEADERS"
] = arize_phoenix_config.otlp_auth_headers
os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
arize_phoenix_config.otlp_auth_headers
)
for callback in _in_memory_loggers:
if (
@@ -3607,9 +3626,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
exporter="otlp_http",
endpoint="https://langtrace.ai/api/trace",
)
os.environ[
"OTEL_EXPORTER_OTLP_TRACES_HEADERS"
] = f"api_key={os.getenv('LANGTRACE_API_KEY')}"
os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
f"api_key={os.getenv('LANGTRACE_API_KEY')}"
)
for callback in _in_memory_loggers:
if (
isinstance(callback, OpenTelemetry)
@@ -4309,10 +4328,10 @@ class StandardLoggingPayloadSetup:
for key in StandardLoggingHiddenParams.__annotations__.keys():
if key in hidden_params:
if key == "additional_headers":
clean_hidden_params[
"additional_headers"
] = StandardLoggingPayloadSetup.get_additional_headers(
hidden_params[key]
clean_hidden_params["additional_headers"] = (
StandardLoggingPayloadSetup.get_additional_headers(
hidden_params[key]
)
)
else:
clean_hidden_params[key] = hidden_params[key] # type: ignore
@@ -4875,9 +4894,9 @@ def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]):
):
for k, v in metadata["user_api_key_metadata"].items():
if k == "logging": # prevent logging user logging keys
cleaned_user_api_key_metadata[
k
] = "scrubbed_by_litellm_for_sensitive_keys"
cleaned_user_api_key_metadata[k] = (
"scrubbed_by_litellm_for_sensitive_keys"
)
else:
cleaned_user_api_key_metadata[k] = v
litellm/litellm_core_utils/llm_cost_calc/utils.py
+31 -21
View File
@@ -1,7 +1,7 @@
# What is this?
## Helper utilities for cost_per_token()
from typing import Any, Literal, Optional, Tuple, TypedDict, cast
from typing import Literal, Optional, Tuple, TypedDict, cast
import litellm
from litellm._logging import verbose_logger
@@ -118,21 +118,21 @@ def _generic_cost_per_character(
def _get_service_tier_cost_key(base_key: str, service_tier: Optional[str]) -> str:
"""
Get the appropriate cost key based on service tier.
Args:
base_key: The base cost key (e.g., "input_cost_per_token")
service_tier: The service tier ("flex", "priority", or None for standard)
Returns:
str: The cost key to use (e.g., "input_cost_per_token_flex" or "input_cost_per_token")
"""
if service_tier is None:
return base_key
# Only use service tier specific keys for "flex" and "priority"
if service_tier.lower() in [ServiceTier.FLEX.value, ServiceTier.PRIORITY.value]:
return f"{base_key}_{service_tier.lower()}"
# For any other service tier, use standard pricing
return base_key
@@ -152,15 +152,15 @@ def _get_token_base_cost(
# Get service tier aware cost keys
input_cost_key = _get_service_tier_cost_key("input_cost_per_token", service_tier)
output_cost_key = _get_service_tier_cost_key("output_cost_per_token", service_tier)
cache_creation_cost_key = _get_service_tier_cost_key("cache_creation_input_token_cost", service_tier)
cache_read_cost_key = _get_service_tier_cost_key("cache_read_input_token_cost", service_tier)
prompt_base_cost = cast(
float, _get_cost_per_unit(model_info, input_cost_key)
cache_creation_cost_key = _get_service_tier_cost_key(
"cache_creation_input_token_cost", service_tier
)
completion_base_cost = cast(
float, _get_cost_per_unit(model_info, output_cost_key)
cache_read_cost_key = _get_service_tier_cost_key(
"cache_read_input_token_cost", service_tier
)
prompt_base_cost = cast(float, _get_cost_per_unit(model_info, input_cost_key))
completion_base_cost = cast(float, _get_cost_per_unit(model_info, output_cost_key))
cache_creation_cost = cast(
float, _get_cost_per_unit(model_info, cache_creation_cost_key)
)
@@ -168,9 +168,7 @@ def _get_token_base_cost(
float,
_get_cost_per_unit(model_info, "cache_creation_input_token_cost_above_1hr"),
)
cache_read_cost = cast(
float, _get_cost_per_unit(model_info, cache_read_cost_key)
)
cache_read_cost = cast(float, _get_cost_per_unit(model_info, cache_read_cost_key))
## CHECK IF ABOVE THRESHOLD
threshold: Optional[float] = None
@@ -278,7 +276,7 @@ def _get_cost_per_unit(
verbose_logger.exception(
f"litellm.litellm_core_utils.llm_cost_calc.utils.py::calculate_cost_per_component(): Exception occured - {cost_per_unit}\nDefaulting to 0.0"
)
# If the service tier key doesn't exist or is None, try to fall back to the standard key
if cost_per_unit is None:
# Check if any service tier suffix exists in the cost key using ServiceTier enum
@@ -286,7 +284,7 @@ def _get_cost_per_unit(
suffix = f"_{service_tier.value}"
if suffix in cost_key:
# Extract the base key by removing the matched suffix
base_key = cost_key.replace(suffix, '')
base_key = cost_key.replace(suffix, "")
fallback_cost = model_info.get(base_key)
if isinstance(fallback_cost, float):
return fallback_cost
@@ -300,7 +298,7 @@ def _get_cost_per_unit(
f"litellm.litellm_core_utils.llm_cost_calc.utils.py::_get_cost_per_unit(): Exception occured - {fallback_cost}\nDefaulting to 0.0"
)
break # Only try the first matching suffix
return default_value
@@ -495,7 +493,10 @@ def _calculate_input_cost(
def generic_cost_per_token(
model: str, usage: Usage, custom_llm_provider: str, service_tier: Optional[str] = None
model: str,
usage: Usage,
custom_llm_provider: str,
service_tier: Optional[str] = None,
) -> Tuple[float, float]:
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@@ -547,7 +548,9 @@ def generic_cost_per_token(
cache_creation_cost,
cache_creation_cost_above_1hr,
cache_read_cost,
) = _get_token_base_cost(model_info=model_info, usage=usage, service_tier=service_tier)
) = _get_token_base_cost(
model_info=model_info, usage=usage, service_tier=service_tier
)
prompt_cost = _calculate_input_cost(
prompt_tokens_details=prompt_tokens_details,
@@ -631,7 +634,7 @@ class CostCalculatorUtils:
@staticmethod
def route_image_generation_cost_calculator(
model: str,
completion_response: Any,
completion_response: ImageResponse,
custom_llm_provider: Optional[str] = None,
quality: Optional[str] = None,
n: Optional[int] = None,
@@ -658,6 +661,13 @@ class CostCalculatorUtils:
cost_calculator as vertex_ai_image_cost_calculator,
)
if size is None:
size = completion_response.size or "1024-x-1024"
if quality is None:
quality = completion_response.quality or "standard"
if n is None:
n = len(completion_response.data) if completion_response.data else 0
if custom_llm_provider == litellm.LlmProviders.VERTEX_AI.value:
if isinstance(completion_response, ImageResponse):
return vertex_ai_image_cost_calculator(
litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py
+19
View File
@@ -37,6 +37,8 @@ from litellm.types.utils import (
TextChoices,
TextCompletionResponse,
TranscriptionResponse,
TranscriptionUsageDurationObject,
TranscriptionUsageTokensObject,
Usage,
)
@@ -684,6 +686,23 @@ def convert_to_model_response_object( # noqa: PLR0915
if key in response_object:
setattr(model_response_object, key, response_object[key])
if "usage" in response_object and response_object["usage"] is not None:
tr_usage_object: Optional[
Union[
TranscriptionUsageDurationObject, TranscriptionUsageTokensObject
]
] = None
if response_object["usage"].get("type", None) == "duration":
tr_usage_object = TranscriptionUsageDurationObject(
**response_object["usage"]
)
elif response_object["usage"].get("type", None) == "tokens":
tr_usage_object = TranscriptionUsageTokensObject(
**response_object["usage"]
)
if tr_usage_object is not None:
setattr(model_response_object, "usage", tr_usage_object)
if hidden_params is not None:
model_response_object._hidden_params = hidden_params
litellm/llms/azure/azure.py
+34 -17
View File
@@ -36,6 +36,7 @@ from .common_utils import (
process_azure_headers,
select_azure_base_url_or_endpoint,
)
from .image_generation import get_azure_image_generation_config
class AzureOpenAIAssistantsAPIConfig:
@@ -1011,7 +1012,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
async def aimage_generation(
self,
data: dict,
model_response: ModelResponse,
model_response: Optional[ImageResponse],
azure_client_params: dict,
api_key: str,
input: list,
@@ -1020,6 +1021,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
client=None,
timeout=None,
) -> litellm.ImageResponse:
response: Optional[dict] = None
try:
# response = await azure_client.images.generate(**data, timeout=timeout)
@@ -1052,21 +1054,38 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
data=data,
headers=headers,
)
response = httpx_response.json()
stringified_response = response
## LOGGING
logging_obj.post_call(
input=input,
api_key=api_key,
additional_args={"complete_input_dict": data},
original_response=stringified_response,
)
return convert_to_model_response_object( # type: ignore
response_object=stringified_response,
model_response_object=model_response,
response_type="image_generation",
provider_config = get_azure_image_generation_config(
data.get("model", "dall-e-2")
)
if provider_config is not None:
return provider_config.transform_image_generation_response(
model=data.get("model", "dall-e-2"),
raw_response=httpx_response,
model_response=model_response or ImageResponse(),
logging_obj=logging_obj,
request_data=data,
optional_params=data,
litellm_params=data,
encoding=litellm.encoding,
)
else:
response = httpx_response.json()
stringified_response = response
## LOGGING
logging_obj.post_call(
input=input,
api_key=api_key,
additional_args={"complete_input_dict": data},
original_response=stringified_response,
)
return convert_to_model_response_object( # type: ignore
response_object=stringified_response,
model_response_object=model_response,
response_type="image_generation",
)
except Exception as e:
## LOGGING
logging_obj.post_call(
@@ -1124,9 +1143,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
if api_key is None and azure_ad_token_provider is not None:
azure_ad_token = azure_ad_token_provider()
if azure_ad_token:
headers.pop(
"api-key", None
)
headers.pop("api-key", None)
headers["Authorization"] = f"Bearer {azure_ad_token}"
# init AzureOpenAI Client
litellm/llms/openai/cost_calculation.py
+14 -10
View File
@@ -18,7 +18,9 @@ def cost_router(call_type: CallTypes) -> Literal["cost_per_token", "cost_per_sec
return "cost_per_token"
def cost_per_token(model: str, usage: Usage, service_tier: Optional[str] = None) -> Tuple[float, float]:
def cost_per_token(
model: str, usage: Usage, service_tier: Optional[str] = None
) -> Tuple[float, float]:
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@@ -31,7 +33,10 @@ def cost_per_token(model: str, usage: Usage, service_tier: Optional[str] = None)
"""
## CALCULATE INPUT COST
return generic_cost_per_token(
model=model, usage=usage, custom_llm_provider="openai", service_tier=service_tier
model=model,
usage=usage,
custom_llm_provider="openai",
service_tier=service_tier,
)
# ### Non-cached text tokens
# non_cached_text_tokens = usage.prompt_tokens
@@ -92,6 +97,7 @@ def cost_per_second(
Returns:
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
"""
## GET MODEL INFO
model_info = get_model_info(
model=model, custom_llm_provider=custom_llm_provider or "openai"
@@ -123,18 +129,16 @@ def cost_per_second(
def video_generation_cost(
model: str,
duration_seconds: float,
custom_llm_provider: Optional[str] = None
model: str, duration_seconds: float, custom_llm_provider: Optional[str] = None
) -> float:
"""
Calculates the cost for video generation based on duration in seconds.
Input:
- model: str, the model name without provider prefix
- duration_seconds: float, the duration of the generated video in seconds
- custom_llm_provider: str, the custom llm provider
Returns:
float - total_cost_in_usd
"""
@@ -142,7 +146,7 @@ def video_generation_cost(
model_info = get_model_info(
model=model, custom_llm_provider=custom_llm_provider or "openai"
)
# Check for video-specific cost per second
video_cost_per_second = model_info.get("output_cost_per_video_per_second")
if video_cost_per_second is not None:
@@ -150,7 +154,7 @@ def video_generation_cost(
f"For model={model} - output_cost_per_video_per_second: {video_cost_per_second}; duration: {duration_seconds}"
)
return video_cost_per_second * duration_seconds
# Fallback to general output cost per second
output_cost_per_second = model_info.get("output_cost_per_second")
if output_cost_per_second is not None:
@@ -158,7 +162,7 @@ def video_generation_cost(
f"For model={model} - output_cost_per_second: {output_cost_per_second}; duration: {duration_seconds}"
)
return output_cost_per_second * duration_seconds
# If no cost information found, return 0
verbose_logger.warning(
f"No cost information found for video model {model}. Please add pricing to model_prices_and_context_window.json"
litellm/llms/openai/image_generation/dall_e_2_transformation.py
+50 -1
View File
@@ -1,9 +1,16 @@
from typing import List
from typing import TYPE_CHECKING, Any, List, Optional
import httpx
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
from litellm.types.utils import ImageResponse
from litellm.utils import convert_to_model_response_object
if TYPE_CHECKING:
from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj
class DallE2ImageGenerationConfig(BaseImageGenerationConfig):
@@ -36,3 +43,45 @@ class DallE2ImageGenerationConfig(BaseImageGenerationConfig):
)
return optional_params
def transform_image_generation_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: "LiteLLMLoggingObj",
request_data: dict,
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> ImageResponse:
response = raw_response.json()
stringified_response = response
## LOGGING
logging_obj.post_call(
input=request_data.get("prompt", ""),
api_key=api_key,
additional_args={"complete_input_dict": request_data},
original_response=stringified_response,
)
image_response: ImageResponse = convert_to_model_response_object( # type: ignore
response_object=stringified_response,
model_response_object=model_response,
response_type="image_generation",
)
# set optional params
image_response.size = optional_params.get(
"size", "1024x1024"
) # default is always 1024x1024
image_response.quality = optional_params.get(
"quality", "standard"
) # always standard for dall-e-2
image_response.output_format = optional_params.get(
"output_format", "png"
) # always png for dall-e-2
return image_response
litellm/llms/openai/image_generation/dall_e_3_transformation.py
+50 -1
View File
@@ -1,9 +1,16 @@
from typing import List
from typing import TYPE_CHECKING, Any, List, Optional
import httpx
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
from litellm.types.utils import ImageResponse
from litellm.utils import convert_to_model_response_object
if TYPE_CHECKING:
from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj
class DallE3ImageGenerationConfig(BaseImageGenerationConfig):
@@ -36,3 +43,45 @@ class DallE3ImageGenerationConfig(BaseImageGenerationConfig):
)
return optional_params
def transform_image_generation_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: "LiteLLMLoggingObj",
request_data: dict,
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> ImageResponse:
response = raw_response.json()
stringified_response = response
## LOGGING
logging_obj.post_call(
input=request_data.get("prompt", ""),
api_key=api_key,
additional_args={"complete_input_dict": request_data},
original_response=stringified_response,
)
image_response: ImageResponse = convert_to_model_response_object( # type: ignore
response_object=stringified_response,
model_response_object=model_response,
response_type="image_generation",
)
# set optional params
image_response.size = optional_params.get(
"size", "1024x1024"
) # default is always 1024x1024
image_response.quality = optional_params.get(
"quality", "hd"
) # always hd for dall-e-3
image_response.output_format = optional_params.get(
"output_format", "png"
) # always png for dall-e-3
return image_response
litellm/llms/openai/image_generation/gpt_transformation.py
+50 -1
View File
@@ -1,9 +1,16 @@
from typing import List
from typing import TYPE_CHECKING, Any, List, Optional
import httpx
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
from litellm.types.utils import ImageResponse
from litellm.utils import convert_to_model_response_object
if TYPE_CHECKING:
from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj
class GPTImageGenerationConfig(BaseImageGenerationConfig):
@@ -45,3 +52,45 @@ class GPTImageGenerationConfig(BaseImageGenerationConfig):
)
return optional_params
def transform_image_generation_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: "LiteLLMLoggingObj",
request_data: dict,
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> ImageResponse:
response = raw_response.json()
stringified_response = response
## LOGGING
logging_obj.post_call(
input=request_data.get("prompt", ""),
api_key=api_key,
additional_args={"complete_input_dict": request_data},
original_response=stringified_response,
)
image_response: ImageResponse = convert_to_model_response_object( # type: ignore
response_object=stringified_response,
model_response_object=model_response,
response_type="image_generation",
)
# set optional params
image_response.size = optional_params.get(
"size", "1024x1024"
) # default is always 1024x1024
image_response.quality = optional_params.get(
"quality", "high"
) # always hd for dall-e-3
image_response.output_format = optional_params.get(
"response_format", "png"
) # always png for dall-e-3
return image_response
litellm/llms/openai/transcriptions/handler.py
+1
View File
@@ -213,6 +213,7 @@ class OpenAIAudioTranscription(OpenAIChatCompletion):
# Extract the actual model from data instead of hardcoding "whisper-1"
actual_model = data.get("model", "whisper-1")
hidden_params = {"model": actual_model, "custom_llm_provider": "openai"}
return convert_to_model_response_object(response_object=stringified_response, model_response_object=model_response, hidden_params=hidden_params, response_type="audio_transcription") # type: ignore
except Exception as e:
## LOGGING
litellm/main.py
+29 -1
View File
@@ -65,7 +65,10 @@ from litellm.constants import (
)
from litellm.exceptions import LiteLLMUnknownProvider
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.audio_utils.utils import get_audio_file_for_health_check
from litellm.litellm_core_utils.audio_utils.utils import (
calculate_request_duration,
get_audio_file_for_health_check,
)
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.get_provider_specific_headers import (
ProviderSpecificHeaderUtils,
@@ -5406,6 +5409,7 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse:
model = args[0] if len(args) > 0 else kwargs["model"]
### PASS ARGS TO Image Generation ###
kwargs["atranscription"] = True
file = kwargs.get("file", None)
custom_llm_provider = None
try:
# Use a partial function to pass your keyword arguments
@@ -5434,6 +5438,20 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse:
raise ValueError(
f"Invalid response from transcription provider, expected TranscriptionResponse, but got {type(response)}"
)
# Calculate and add duration if response is missing it
if (
response is not None
and not isinstance(response, Coroutine)
and file is not None
):
# Check if response is missing duration
existing_duration = getattr(response, "duration", None)
if existing_duration is None:
calculated_duration = calculate_request_duration(file)
if calculated_duration is not None:
setattr(response, "duration", calculated_duration)
return response
except Exception as e:
custom_llm_provider = custom_llm_provider or "openai"
@@ -5644,6 +5662,16 @@ def transcription(
headers={},
provider_config=provider_config,
)
# Calculate and add duration if response is missing it
if response is not None and not isinstance(response, Coroutine):
# Check if response is missing duration
existing_duration = getattr(response, "duration", None)
if existing_duration is None:
calculated_duration = calculate_request_duration(file)
if calculated_duration is not None:
setattr(response, "duration", calculated_duration)
if response is None:
raise ValueError("Unmapped provider passed in. Unable to get the response.")
return response
litellm/proxy/_new_secret_config.yaml
+4 -12
View File
@@ -6,16 +6,8 @@ model_list:
litellm_params:
model: openai/text-embedding-3-large
vector_store_registry:
- vector_store_name: "vertex-ai-litellm-website-knowledgebase"
- model_name: gpt-4o-mini-transcribe
litellm_params:
vector_store_id: "litellm-docs_1761094140318"
custom_llm_provider: "vertex_ai/search_api"
vertex_project: "test-vector-store-db"
vertex_location: "global"
- vector_store_name: "milvus-litellm-website-knowledgebase"
litellm_params:
vector_store_id: "can-be-anything"
custom_llm_provider: "milvus"
api_base: os.environ/MILVUS_API_BASE
api_key: os.environ/MILVUS_API_KEY
model: openai/gpt-4o-mini-transcribe
api_key: os.environ/OPENAI_API_KEY
litellm/proxy/example_config_yaml/custom_callbacks1.py
+3 -22
View File
@@ -3,6 +3,7 @@ from typing import Literal, Optional
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy.proxy_server import DualCache, UserAPIKeyAuth
from litellm.types.utils import CallTypesLiteral
# This file includes the custom callbacks for LiteLLM Proxy
@@ -21,18 +22,7 @@ class MyCustomHandler(
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
return data
@@ -58,16 +48,7 @@ class MyCustomHandler(
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
pass
litellm/proxy/example_config_yaml/custom_guardrail.py
+3 -22
View File
@@ -6,6 +6,7 @@ from litellm.caching.caching import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata
from litellm.types.utils import CallTypesLiteral
class myCustomGuardrail(CustomGuardrail):
@@ -23,18 +24,7 @@ class myCustomGuardrail(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Optional[Union[Exception, str, dict]]:
"""
Runs before the LLM API call
@@ -62,16 +52,7 @@ class myCustomGuardrail(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
"""
Runs in parallel to LLM API call
litellm/proxy/guardrails/guardrail_hooks/aim/aim.py
+4 -23
View File
@@ -7,7 +7,7 @@
import asyncio
import json
import os
from typing import TYPE_CHECKING, Any, AsyncGenerator, Literal, Optional, Type, Union
from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, Type, Union
from fastapi import HTTPException
from pydantic import BaseModel
@@ -23,6 +23,7 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import (
CallTypesLiteral,
Choices,
EmbeddingResponse,
ImageResponse,
@@ -67,18 +68,7 @@ class AimGuardrail(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Union[Exception, str, dict, None]:
verbose_proxy_logger.debug("Inside AIM Pre-Call Hook")
return await self.call_aim_guardrail(
@@ -89,16 +79,7 @@ class AimGuardrail(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Union[Exception, str, dict, None]:
verbose_proxy_logger.debug("Inside AIM Moderation Hook")
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
+6 -25
View File
@@ -42,6 +42,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
)
from litellm.types.utils import (
CallTypes,
CallTypesLiteral,
Choices,
GuardrailStatus,
ModelResponse,
@@ -608,18 +609,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Union[Exception, str, dict, None]:
verbose_proxy_logger.debug(
"Inside Bedrock Pre-Call Hook for call_type: %s", call_type
@@ -685,16 +675,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
@@ -1163,7 +1144,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
This method allows users to test Bedrock guardrails without making actual LLM calls.
It creates a mock request and response to test the guardrail functionality.
Args:
text: The text to analyze
language: Optional language parameter (not used by Bedrock)
@@ -1175,11 +1156,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
mock_messages: List[AllMessageValues] = [
ChatCompletionUserMessage(role="user", content=text)
]
# Use provided request_data or create a mock one for testing
if request_data is None:
request_data = {"messages": mock_messages}
bedrock_response = await self.make_bedrock_api_request(
source="INPUT",
messages=mock_messages,
litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py
+82 -97
View File
@@ -7,16 +7,7 @@
import os
from datetime import datetime
from typing import (
Any,
AsyncGenerator,
Dict,
List,
Literal,
Optional,
Type,
Union,
)
from typing import Any, AsyncGenerator, Dict, List, Optional, Type, Union
import httpx
@@ -36,7 +27,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.dynamoai import (
DynamoAIRequest,
DynamoAIResponse,
)
from litellm.types.utils import GuardrailStatus, ModelResponseStream
from litellm.types.utils import CallTypesLiteral, GuardrailStatus, ModelResponseStream
GUARDRAIL_NAME = "dynamoai"
@@ -44,9 +35,10 @@ GUARDRAIL_NAME = "dynamoai"
class DynamoAIGuardrails(CustomGuardrail):
"""
DynamoAI Guardrails integration for LiteLLM.
Provides content moderation and policy enforcement using DynamoAI's guardrail API.
"""
def __init__(
self,
guardrail_name: str = "litellm_test",
@@ -59,28 +51,30 @@ class DynamoAIGuardrails(CustomGuardrail):
self.async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback
)
# Set API configuration
self.api_key = api_key or os.getenv("DYNAMOAI_API_KEY")
if not self.api_key:
raise ValueError(
"DynamoAI API key is required. Set DYNAMOAI_API_KEY environment variable or pass api_key parameter."
)
self.api_base = api_base or os.getenv(
"DYNAMOAI_API_BASE", "https://api.dynamo.ai"
)
self.api_url = f"{self.api_base}/v1/moderation/analyze/"
# Model ID for tracking/logging purposes
self.model_id = model_id or os.getenv("DYNAMOAI_MODEL_ID", "")
# Policy IDs - get from parameter, env var, or use empty list
env_policy_ids = os.getenv("DYNAMOAI_POLICY_IDS", "")
self.policy_ids = policy_ids or (env_policy_ids.split(",") if env_policy_ids else [])
self.policy_ids = policy_ids or (
env_policy_ids.split(",") if env_policy_ids else []
)
self.guardrail_name = guardrail_name
self.guardrail_provider = "dynamoai"
# store kwargs as optional_params
self.optional_params = kwargs
@@ -108,38 +102,38 @@ class DynamoAIGuardrails(CustomGuardrail):
) -> DynamoAIResponse:
"""
Call DynamoAI Guardrails API to analyze messages for policy violations.
Args:
messages: List of messages to analyze
text_type: Type of text being analyzed ("input" or "output")
request_data: Optional request data for logging purposes
Returns:
DynamoAIResponse: Response from the DynamoAI Guardrails API
"""
start_time = datetime.now()
payload: DynamoAIRequest = {
"messages": messages,
}
# Add optional fields if provided
if self.policy_ids:
payload["policyIds"] = self.policy_ids
if self.model_id:
payload["modelId"] = self.model_id
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {self.api_key}"
"Authorization": f"Bearer {self.api_key}",
}
verbose_proxy_logger.debug(
"DynamoAI request to %s with payload=%s",
self.api_url,
payload,
)
try:
response = await self.async_handler.post(
url=self.api_url,
@@ -148,10 +142,10 @@ class DynamoAIGuardrails(CustomGuardrail):
)
response.raise_for_status()
response_json = response.json()
end_time = datetime.now()
duration = (end_time - start_time).total_seconds()
# Add guardrail information to request trace
if request_data:
guardrail_status = self._determine_guardrail_status(response_json)
@@ -164,17 +158,15 @@ class DynamoAIGuardrails(CustomGuardrail):
end_time=end_time.timestamp(),
duration=duration,
)
return response_json
except httpx.HTTPError as e:
end_time = datetime.now()
duration = (end_time - start_time).total_seconds()
verbose_proxy_logger.error(
"DynamoAI API request failed: %s", str(e)
)
verbose_proxy_logger.error("DynamoAI API request failed: %s", str(e))
# Add guardrail information with failure status
if request_data:
self.add_standard_logging_guardrail_information_to_request_data(
@@ -186,7 +178,7 @@ class DynamoAIGuardrails(CustomGuardrail):
end_time=end_time.timestamp(),
duration=duration,
)
raise
def _process_dynamoai_guardrails_response(
@@ -194,35 +186,35 @@ class DynamoAIGuardrails(CustomGuardrail):
) -> DynamoAIProcessedResult:
"""
Process the response from the DynamoAI Guardrails API
Args:
response: The response from the API with 'finalAction' and 'appliedPolicies' keys
Returns:
DynamoAIProcessedResult: Processed response with detected violations
"""
final_action = response.get("finalAction", "NONE")
applied_policies = response.get("appliedPolicies", [])
violations_detected: List[str] = []
violation_details: Dict[str, Any] = {}
# For now, only handle BLOCK action
if final_action == "BLOCK":
for applied_policy in applied_policies:
policy_info = applied_policy.get("policy", {})
policy_outputs = applied_policy.get("outputs", {})
# Get policy name and action
policy_name = policy_info.get("name", "unknown")
# Check for action in multiple places
policy_action = (
applied_policy.get("action") or
(policy_outputs.get("action") if policy_outputs else None) or
"NONE"
applied_policy.get("action")
or (policy_outputs.get("action") if policy_outputs else None)
or "NONE"
)
# Only include policies with BLOCK action
if policy_action == "BLOCK":
violations_detected.append(policy_name)
@@ -231,12 +223,14 @@ class DynamoAIGuardrails(CustomGuardrail):
"action": policy_action,
"method": policy_info.get("method"),
"description": policy_info.get("description"),
"message": policy_outputs.get("message") if policy_outputs else None,
"message": (
policy_outputs.get("message") if policy_outputs else None
),
}
return {
"violations_detected": violations_detected,
"violation_details": violation_details
"violation_details": violation_details,
}
def _determine_guardrail_status(
@@ -244,7 +238,7 @@ class DynamoAIGuardrails(CustomGuardrail):
) -> GuardrailStatus:
"""
Determine the guardrail status based on DynamoAI API response.
Returns:
"success": Content allowed through with no violations (finalAction is NONE)
"guardrail_intervened": Content blocked (finalAction is BLOCK)
@@ -253,21 +247,21 @@ class DynamoAIGuardrails(CustomGuardrail):
try:
if not isinstance(response_json, dict):
return "guardrail_failed_to_respond"
# Check for error in response
if response_json.get("error"):
return "guardrail_failed_to_respond"
final_action = response_json.get("finalAction", "NONE")
if final_action == "NONE":
return "success"
elif final_action == "BLOCK":
return "guardrail_intervened"
# For now, treat other actions as success (WARN, REDACT, SANITIZE not implemented yet)
return "success"
except Exception as e:
verbose_proxy_logger.error(
"Error determining DynamoAI guardrail status: %s", str(e)
@@ -277,22 +271,24 @@ class DynamoAIGuardrails(CustomGuardrail):
def _create_error_message(self, processed_result: DynamoAIProcessedResult) -> str:
"""
Create a detailed error message from processed guardrail results.
Args:
processed_result: Processed response with detected violations
Returns:
Formatted error message string
"""
violations_detected = processed_result["violations_detected"]
violation_details = processed_result["violation_details"]
error_message = f"Guardrail failed: {len(violations_detected)} violation(s) detected\n\n"
error_message = (
f"Guardrail failed: {len(violations_detected)} violation(s) detected\n\n"
)
for policy_name in violations_detected:
error_message += f"- {policy_name.upper()}:\n"
details = violation_details.get(policy_name, {})
# Format violation details
if details.get("action"):
error_message += f" Action: {details['action']}\n"
@@ -305,7 +301,7 @@ class DynamoAIGuardrails(CustomGuardrail):
if details.get("policyId"):
error_message += f" Policy ID: {details['policyId']}\n"
error_message += "\n"
return error_message.strip()
async def async_pre_call_hook(
@@ -313,18 +309,7 @@ class DynamoAIGuardrails(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Union[Exception, str, dict, None]:
"""
Runs before the LLM API call
@@ -349,7 +334,9 @@ class DynamoAIGuardrails(CustomGuardrail):
request_data=data,
)
verbose_proxy_logger.debug("Guardrails async_pre_call_hook result=%s", result)
verbose_proxy_logger.debug(
"Guardrails async_pre_call_hook result=%s", result
)
# Process the guardrails response
processed_result = self._process_dynamoai_guardrails_response(result)
@@ -364,23 +351,14 @@ class DynamoAIGuardrails(CustomGuardrail):
add_guardrail_to_applied_guardrails_header(
request_data=data, guardrail_name=self.guardrail_name
)
return data
async def async_moderation_hook(
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
"""
Runs in parallel to LLM API call
@@ -404,7 +382,9 @@ class DynamoAIGuardrails(CustomGuardrail):
request_data=data,
)
verbose_proxy_logger.debug("Guardrails async_moderation_hook result=%s", result)
verbose_proxy_logger.debug(
"Guardrails async_moderation_hook result=%s", result
)
# Process the guardrails response
processed_result = self._process_dynamoai_guardrails_response(result)
@@ -449,22 +429,26 @@ class DynamoAIGuardrails(CustomGuardrail):
return
verbose_proxy_logger.debug("async_post_call_success_hook response=%s", response)
# Check if the ModelResponse has text content in its choices
# to avoid sending empty content to DynamoAI (e.g., during tool calls)
if isinstance(response, litellm.ModelResponse):
has_text_content = False
dynamoai_messages: List[Dict[str, Any]] = []
for choice in response.choices:
if isinstance(choice, litellm.Choices):
if choice.message.content and isinstance(choice.message.content, str):
if choice.message.content and isinstance(
choice.message.content, str
):
has_text_content = True
dynamoai_messages.append({
"role": choice.message.role or "assistant",
"content": choice.message.content
})
dynamoai_messages.append(
{
"role": choice.message.role or "assistant",
"content": choice.message.content,
}
)
if not has_text_content:
verbose_proxy_logger.warning(
"DynamoAI: not running guardrail. No output text in response"
@@ -478,7 +462,9 @@ class DynamoAIGuardrails(CustomGuardrail):
request_data=data,
)
verbose_proxy_logger.debug("Guardrails async_post_call_success_hook result=%s", result)
verbose_proxy_logger.debug(
"Guardrails async_post_call_success_hook result=%s", result
)
# Process the guardrails response
processed_result = self._process_dynamoai_guardrails_response(result)
@@ -517,4 +503,3 @@ class DynamoAIGuardrails(CustomGuardrail):
)
return DynamoAIGuardrailConfigModel
litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py
+4 -24
View File
@@ -7,7 +7,7 @@
import os
from datetime import datetime
from typing import Any, AsyncGenerator, Dict, List, Literal, Optional, Union
from typing import Any, AsyncGenerator, Dict, List, Optional, Union
import httpx
@@ -25,7 +25,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import (
EnkryptAIProcessedResult,
EnkryptAIResponse,
)
from litellm.types.utils import GuardrailStatus, ModelResponseStream
from litellm.types.utils import CallTypesLiteral, GuardrailStatus, ModelResponseStream
GUARDRAIL_NAME = "enkryptai"
@@ -284,18 +284,7 @@ class EnkryptAIGuardrails(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Union[Exception, str, dict, None]:
"""
Runs before the LLM API call
@@ -348,16 +337,7 @@ class EnkryptAIGuardrails(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
"""
Runs in parallel to LLM API call
litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py
+5 -25
View File
@@ -7,7 +7,7 @@
import os
from datetime import datetime
from typing import Any, AsyncGenerator, Dict, List, Literal, Optional, Union
from typing import Any, AsyncGenerator, Dict, List, Optional, Union
from urllib.parse import urlencode
import httpx
@@ -26,7 +26,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.ibm import (
IBMDetectorDetection,
IBMDetectorResponseOrchestrator,
)
from litellm.types.utils import GuardrailStatus, ModelResponseStream
from litellm.types.utils import CallTypesLiteral, GuardrailStatus, ModelResponseStream
GUARDRAIL_NAME = "ibm_guardrails"
@@ -47,7 +47,7 @@ class IBMGuardrailDetector(CustomGuardrail):
):
self.async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback,
params={"ssl_verify": verify_ssl}
params={"ssl_verify": verify_ssl},
)
# Set API configuration
@@ -436,18 +436,7 @@ class IBMGuardrailDetector(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Union[Exception, str, dict, None]:
"""
Runs before the LLM API call
@@ -533,16 +522,7 @@ class IBMGuardrailDetector(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
"""
Runs in parallel to LLM API call
litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py
+3 -14
View File
@@ -1,5 +1,5 @@
from datetime import datetime
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Type, Union
from typing import TYPE_CHECKING, Dict, List, Optional, Type, Union
from fastapi import HTTPException
@@ -18,7 +18,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.javelin import (
JavelinGuardRequest,
JavelinGuardResponse,
)
from litellm.types.utils import GuardrailStatus
from litellm.types.utils import CallTypesLiteral, GuardrailStatus
if TYPE_CHECKING:
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
@@ -165,18 +165,7 @@ class JavelinGuardrail(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: litellm.DualCache,
data: Dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Optional[Union[Exception, str, Dict]]:
"""
Pre-call hook for the Javelin guardrail.
litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py
+6 -26
View File
@@ -1,7 +1,7 @@
import copy
import os
from datetime import datetime
from typing import Dict, List, Literal, Optional, Tuple, Union
from typing import Dict, List, Optional, Tuple, Union
from fastapi import HTTPException
@@ -20,7 +20,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import (
LakeraAIRequest,
LakeraAIResponse,
)
from litellm.types.utils import GuardrailStatus
from litellm.types.utils import CallTypesLiteral, GuardrailStatus
class LakeraAIGuardrail(CustomGuardrail):
@@ -183,18 +183,7 @@ class LakeraAIGuardrail(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: litellm.DualCache,
data: Dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Optional[Union[Exception, str, Dict]]:
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
@@ -257,16 +246,7 @@ class LakeraAIGuardrail(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
@@ -333,7 +313,7 @@ class LakeraAIGuardrail(CustomGuardrail):
breakdown = lakera_response.get("breakdown", []) or []
if not breakdown:
return False
has_violations = False
for item in breakdown:
if item.get("detected", False):
@@ -341,7 +321,7 @@ class LakeraAIGuardrail(CustomGuardrail):
detector_type = item.get("detector_type", "") or ""
if not detector_type.startswith("pii/"):
return False
# Return True only if there are violations and they are all PII
return has_violations
litellm/proxy/guardrails/guardrail_hooks/noma/noma.py
+80 -95
View File
@@ -6,11 +6,22 @@
# +-------------------------------------------------------------+
import asyncio
import json
import os
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, Final, Literal, Optional, Type, Union
from typing import (
TYPE_CHECKING,
Any,
AsyncGenerator,
Dict,
Final,
List,
Literal,
Optional,
Type,
Union,
)
from urllib.parse import urljoin
import json
from fastapi import HTTPException
@@ -18,25 +29,22 @@ import litellm
from litellm import DualCache, ModelResponse
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.main import stream_chunk_builder
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import EmbeddingResponse, GuardrailStatus, ImageResponse
from litellm.types.utils import (
CallTypesLiteral,
EmbeddingResponse,
GuardrailStatus,
ImageResponse,
ModelResponseStream,
TextCompletionResponse,
)
from typing import (
List,
AsyncGenerator
)
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
from litellm.main import stream_chunk_builder
from litellm.types.utils import TextCompletionResponse
# Constants
USER_ROLE: Final[Literal["user"]] = "user"
@@ -48,9 +56,9 @@ MessageRole = Literal["user", "assistant"]
LLMResponse = Union[Any, ModelResponse, EmbeddingResponse, ImageResponse]
if TYPE_CHECKING:
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
class NomaBlockedMessage(HTTPException):
"""Exception raised when Noma guardrail blocks a message"""
@@ -160,13 +168,7 @@ class NomaGuardrail(CustomGuardrail):
return None
payload = {
"input": [
{
"type": "message",
"role": "user",
"content": user_message
}
]
"input": [{"type": "message", "role": "user", "content": user_message}]
}
response_json = await self._call_noma_api(
payload=payload,
@@ -175,13 +177,13 @@ class NomaGuardrail(CustomGuardrail):
user_auth=user_auth,
extra_data=extra_data,
)
end_time = datetime.now()
duration = (end_time - start_time).total_seconds()
# Determine guardrail status based on response
guardrail_status = self._determine_guardrail_status(response_json)
# Always log guardrail information for consistency
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider="noma",
@@ -222,7 +224,7 @@ class NomaGuardrail(CustomGuardrail):
user_auth: UserAPIKeyAuth,
) -> Optional[str]:
"""Shared logic for processing LLM response checks"""
start_time = datetime.now()
extra_data = self.get_guardrail_dynamic_request_body_params(request_data)
@@ -243,12 +245,7 @@ class NomaGuardrail(CustomGuardrail):
{
"type": "message",
"role": "assistant",
"content": [
{
"type": "input_text",
"text": content
}
]
"content": [{"type": "input_text", "text": content}],
}
]
}
@@ -260,13 +257,13 @@ class NomaGuardrail(CustomGuardrail):
user_auth=user_auth,
extra_data=extra_data,
)
end_time = datetime.now()
duration = (end_time - start_time).total_seconds()
# Determine guardrail status based on response
guardrail_status = self._determine_guardrail_status(response_json)
# Always log guardrail information for consistency
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider="noma",
@@ -303,10 +300,10 @@ class NomaGuardrail(CustomGuardrail):
def _determine_guardrail_status(self, response_json: dict) -> GuardrailStatus:
"""
Determine the guardrail status based on NOMA API response.
Args:
response_json: Response from NOMA API
Returns:
"success": Content allowed through with no violations
"guardrail_intervened": Content blocked due to policy violations
@@ -316,24 +313,26 @@ class NomaGuardrail(CustomGuardrail):
# Check if we got a valid response structure
if not isinstance(response_json, dict):
return "guardrail_failed_to_respond"
# Get the aggregatedScanResult from the response
# aggregatedScanResult=True means unsafe (block), False means safe (allow)
aggregated_scan_result = response_json.get("aggregatedScanResult", False)
# If aggregatedScanResult is False, content is safe/allowed
if aggregated_scan_result is False:
return "success"
# If aggregatedScanResult is True, content is blocked/flagged
if aggregated_scan_result is True:
return "guardrail_intervened"
# If aggregatedScanResult is missing or invalid, treat as failure
return "guardrail_failed_to_respond"
except Exception as e:
verbose_proxy_logger.error(f"Error determining NOMA guardrail status: {str(e)}")
verbose_proxy_logger.error(
f"Error determining NOMA guardrail status: {str(e)}"
)
return "guardrail_failed_to_respond"
def _should_only_sensitive_data_failed(self, classification_obj: dict) -> bool:
@@ -392,12 +391,16 @@ class NomaGuardrail(CustomGuardrail):
scan_result = response_json.get("scanResult", [])
if not scan_result:
return None
# Find the scan result matching the message type (role)
for result_item in scan_result:
if result_item.get("role") == message_type:
return result_item.get("results", {}).get("anonymizedContent", {}).get("anonymized", "")
return (
result_item.get("results", {})
.get("anonymizedContent", {})
.get("anonymized", "")
)
return None
def _should_anonymize(self, response_json: dict, message_type: MessageRole) -> bool:
@@ -423,7 +426,7 @@ class NomaGuardrail(CustomGuardrail):
# aggregatedScanResult=False means safe, True means unsafe
aggregated_scan_result = response_json.get("aggregatedScanResult", False)
# If aggregatedScanResult is False, content is safe - anonymize if available
if not aggregated_scan_result:
return True
@@ -432,13 +435,15 @@ class NomaGuardrail(CustomGuardrail):
scan_result = response_json.get("scanResult", [])
if not scan_result:
return False
if not isinstance(scan_result, list) or len(scan_result) == 0:
return False
for result_item in scan_result:
if result_item.get("role") == message_type:
return self._should_only_sensitive_data_failed(result_item.get("results", {}))
return self._should_only_sensitive_data_failed(
result_item.get("results", {})
)
return False
@@ -534,7 +539,7 @@ class NomaGuardrail(CustomGuardrail):
try:
# aggregatedScanResult=True means blocked, False means allowed
aggregated_scan_result = response_json.get("aggregatedScanResult", False)
if aggregated_scan_result: # True = unsafe
msg = f"Noma guardrail blocked {type} message: {message}"
verbose_proxy_logger.warning(msg)
@@ -551,20 +556,9 @@ class NomaGuardrail(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Optional[Union[Exception, str, dict]]:
verbose_proxy_logger.debug("Running Noma pre-call hook")
if (
@@ -595,6 +589,7 @@ class NomaGuardrail(CustomGuardrail):
except Exception as e:
# Log technical failures
from datetime import datetime
start_time = datetime.now()
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider="noma",
@@ -605,7 +600,7 @@ class NomaGuardrail(CustomGuardrail):
end_time=start_time.timestamp(),
duration=0.0,
)
verbose_proxy_logger.error(f"Noma pre-call hook failed: {str(e)}")
if self.block_failures:
@@ -616,16 +611,7 @@ class NomaGuardrail(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Union[Exception, str, dict, None]:
event_type: GuardrailEventHooks = GuardrailEventHooks.during_call
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
@@ -651,6 +637,7 @@ class NomaGuardrail(CustomGuardrail):
except Exception as e:
# Log technical failures
from datetime import datetime
start_time = datetime.now()
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider="noma",
@@ -661,7 +648,7 @@ class NomaGuardrail(CustomGuardrail):
end_time=start_time.timestamp(),
duration=0.0,
)
verbose_proxy_logger.error(f"Noma moderation hook failed: {str(e)}")
if self.block_failures:
@@ -700,6 +687,7 @@ class NomaGuardrail(CustomGuardrail):
except Exception as e:
# Log technical failures
from datetime import datetime
start_time = datetime.now()
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider="noma",
@@ -710,7 +698,7 @@ class NomaGuardrail(CustomGuardrail):
end_time=start_time.timestamp(),
duration=0.0,
)
verbose_proxy_logger.error(f"Noma post-call hook failed: {str(e)}")
if self.block_failures:
raise
@@ -756,37 +744,31 @@ class NomaGuardrail(CustomGuardrail):
last_user_message = user_messages[-1].get("content", "")
if isinstance(last_user_message, str):
return [{
"type": "input_text",
"text": last_user_message
}]
return [{"type": "input_text", "text": last_user_message}]
elif isinstance(last_user_message, list):
converted_messages = []
for message in last_user_message:
converted_message = self._convert_single_user_message_to_payload(message)
converted_message = self._convert_single_user_message_to_payload(
message
)
if converted_message is not None:
converted_messages.append(converted_message)
return converted_messages
else:
return None
def _convert_single_user_message_to_payload(self, user_message: Any) -> Optional[dict]:
def _convert_single_user_message_to_payload(
self, user_message: Any
) -> Optional[dict]:
if isinstance(user_message, str):
return {
"type": "input_text",
"text": user_message
}
return {"type": "input_text", "text": user_message}
elif user_message.get("type", "") == "image_url":
return {
"type": "input_image",
"image_url": user_message.get("image_url", {}).get("url", "")
"image_url": user_message.get("image_url", {}).get("url", ""),
}
elif user_message.get("type", "") == "text":
return {
"type": "input_text",
"text": user_message.get("text", "")
}
return {"type": "input_text", "text": user_message.get("text", "")}
else:
return None
@@ -816,13 +798,16 @@ class NomaGuardrail(CustomGuardrail):
"applicationId": extra_data.get("application_id")
or request_data.get("metadata", {})
.get("headers", {})
.get("x-noma-application-id") or self.application_id,
.get("x-noma-application-id")
or self.application_id,
"ipAddress": request_data.get("metadata", {}).get(
"requester_ip_address", None
),
"userId": user_auth.user_email
if user_auth.user_email
else user_auth.user_id,
"userId": (
user_auth.user_email
if user_auth.user_email
else user_auth.user_id
),
"sessionId": call_id,
"requestId": llm_request_id,
},
@@ -844,7 +829,7 @@ class NomaGuardrail(CustomGuardrail):
"""
# aggregatedScanResult=True means blocked, False means allowed
aggregated_scan_result = response_json.get("aggregatedScanResult", False)
if aggregated_scan_result: # True = unsafe, block it
msg = f"Noma guardrail blocked {type} message: {message}"
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
+3 -14
View File
@@ -6,7 +6,7 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint
3. Implements a way to call /applyGuardrail endpoint for `/chat/completions` + `/v1/messages` requests on async_post_call_streaming_iterator_hook
"""
from typing import Any, AsyncGenerator, Literal, Union
from typing import Any, AsyncGenerator, Union
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
@@ -16,7 +16,7 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.llms import load_guardrail_translation_mappings
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import CallTypes, ModelResponseStream
from litellm.types.utils import CallTypes, CallTypesLiteral, ModelResponseStream
GUARDRAIL_NAME = "unified_llm_guardrails"
endpoint_guardrail_translation_mappings = None
@@ -43,18 +43,7 @@ class UnifiedLLMGuardrails(CustomLogger):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Union[Exception, str, dict, None]:
"""
Runs before the LLM API call
litellm/proxy/hooks/dynamic_rate_limiter.py
+17 -27
View File
@@ -4,7 +4,7 @@
import asyncio
import os
from typing import List, Literal, Optional, Tuple, Union
from typing import List, Optional, Tuple, Union
from fastapi import HTTPException
@@ -15,6 +15,7 @@ from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import CallTypesLiteral
from litellm.utils import get_utc_datetime
from .rate_limiter_utils import convert_priority_to_percent
@@ -102,10 +103,10 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
"""
try:
# Get model info first for conversion
model_group_info: Optional[
ModelGroupInfo
] = self.llm_router.get_model_group_info(model_group=model)
model_group_info: Optional[ModelGroupInfo] = (
self.llm_router.get_model_group_info(model_group=model)
)
weight: float = 1
if (
litellm.priority_reservation is None
@@ -193,18 +194,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Optional[
Union[Exception, str, dict]
]: # raise exception if invalid, return a str for the user to receive - if rejected, or return a modified dictionary for passing into litellm
@@ -287,16 +277,16 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
) = await self.check_available_usage(
model=model_info["model_name"], priority=key_priority
)
response._hidden_params[
"additional_headers"
] = { # Add additional response headers - easier debugging
"x-litellm-model_group": model_info["model_name"],
"x-ratelimit-remaining-litellm-project-tokens": available_tpm,
"x-ratelimit-remaining-litellm-project-requests": available_rpm,
"x-ratelimit-remaining-model-tokens": model_tpm,
"x-ratelimit-remaining-model-requests": model_rpm,
"x-ratelimit-current-active-projects": active_projects,
}
response._hidden_params["additional_headers"] = (
{ # Add additional response headers - easier debugging
"x-litellm-model_group": model_info["model_name"],
"x-ratelimit-remaining-litellm-project-tokens": available_tpm,
"x-ratelimit-remaining-litellm-project-requests": available_rpm,
"x-ratelimit-remaining-model-tokens": model_tpm,
"x-ratelimit-remaining-model-requests": model_rpm,
"x-ratelimit-current-active-projects": active_projects,
}
)
return response
return await super().async_post_call_success_hook(
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
+119 -105
View File
@@ -4,7 +4,7 @@ Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting
import os
from datetime import datetime
from typing import Callable, Dict, List, Literal, Optional, Union
from typing import Callable, Dict, List, Optional, Union
from fastapi import HTTPException
@@ -22,19 +22,20 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
from litellm.proxy.hooks.rate_limiter_utils import convert_priority_to_percent
from litellm.proxy.utils import InternalUsageCache
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import CallTypesLiteral
class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
"""
Saturation-aware priority-based rate limiter using v3 infrastructure.
Key features:
1. Model capacity ALWAYS enforced at 100% (prevents over-allocation)
2. Priority usage tracked from first request (accurate accounting)
3. Priority limits only enforced when saturated >= threshold
4. Three-phase checking prevents partial counter increments
5. Reuses v3 limiter's Redis-based tracking (multi-instance safe)
How it works:
- Phase 1: Read-only check of ALL limits (no increments)
- Phase 2: Decide enforcement based on saturation
@@ -43,6 +44,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
- When saturated: strict priority-based limits enforced (fair)
- Uses v3 limiter's atomic Lua scripts for race-free increments
"""
def __init__(
self,
internal_usage_cache: DualCache,
@@ -56,7 +58,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
def update_variables(self, llm_router: Router):
self.llm_router = llm_router
def _get_priority_weight(self, priority: Optional[str], model_info: Optional[ModelGroupInfo] = None) -> float:
def _get_priority_weight(
self, priority: Optional[str], model_info: Optional[ModelGroupInfo] = None
) -> float:
"""Get the weight for a given priority from litellm.priority_reservation"""
weight: float = litellm.priority_reservation_settings.default_priority
if (
@@ -76,30 +80,32 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
weight = convert_priority_to_percent(value, model_info)
return weight
def _normalize_priority_weights(self, model_info: ModelGroupInfo) -> Dict[str, float]:
def _normalize_priority_weights(
self, model_info: ModelGroupInfo
) -> Dict[str, float]:
"""
Normalize priority weights if they sum to > 1.0
Handles over-allocation: {key_a: 0.60, key_b: 0.80} -> {key_a: 0.43, key_b: 0.57}
Converts absolute rpm/tpm values to percentages based on model capacity.
"""
if litellm.priority_reservation is None:
return {}
# Convert all values to percentages first
weights: Dict[str, float] = {}
for k, v in litellm.priority_reservation.items():
weights[k] = convert_priority_to_percent(v, model_info)
total_weight = sum(weights.values())
if total_weight > 1.0:
normalized = {k: v / total_weight for k, v in weights.items()}
verbose_proxy_logger.debug(
f"Normalized over-allocated priorities: {weights} -> {normalized}"
)
return normalized
return weights
def _get_priority_allocation(
@@ -111,29 +117,31 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
) -> tuple[float, str]:
"""
Get priority weight and pool key for a given priority.
For explicit priorities: returns specific allocation and unique pool key
For default priority: returns default allocation and shared pool key
Args:
model: Model name
priority: Priority level (None for default)
normalized_weights: Pre-computed normalized weights
model_info: Model configuration (optional, for fallback conversion)
Returns:
tuple: (priority_weight, priority_key)
"""
# Check if this key has an explicit priority in litellm.priority_reservation
has_explicit_priority = (
priority is not None
and litellm.priority_reservation is not None
priority is not None
and litellm.priority_reservation is not None
and priority in litellm.priority_reservation
)
if has_explicit_priority and priority is not None:
# Explicit priority: get its specific allocation
priority_weight = normalized_weights.get(priority, self._get_priority_weight(priority, model_info))
priority_weight = normalized_weights.get(
priority, self._get_priority_weight(priority, model_info)
)
# Use unique key per priority level
priority_key = f"{model}:{priority}"
else:
@@ -141,7 +149,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
priority_weight = litellm.priority_reservation_settings.default_priority
# Use shared key for all default-priority requests
priority_key = f"{model}:default_pool"
return priority_weight, priority_key
async def _check_model_saturation(
@@ -151,16 +159,16 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
) -> float:
"""
Check current saturation by directly querying v3 limiter's cache keys.
Reuses v3 limiter's Redis-based tracking (works across multiple instances).
Reads counters WITHOUT incrementing them.
Returns:
float: Saturation ratio (0.0 = empty, 1.0 = at capacity, >1.0 = over)
"""
try:
max_saturation = 0.0
# Query RPM saturation
if model_group_info.rpm is not None and model_group_info.rpm > 0:
# Use v3 limiter's key format: {key:value}:rate_limit_type
@@ -169,24 +177,24 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
value=model,
rate_limit_type="requests",
)
# Query cache for current counter value
counter_value = await self.internal_usage_cache.async_get_cache(
key=counter_key,
litellm_parent_otel_span=None,
local_only=False, # Check Redis too
)
if counter_value is not None:
current_requests = int(counter_value)
rpm_saturation = current_requests / model_group_info.rpm
max_saturation = max(max_saturation, rpm_saturation)
verbose_proxy_logger.debug(
f"Model {model} RPM: {current_requests}/{model_group_info.rpm} "
f"({rpm_saturation:.1%})"
)
# Query TPM saturation
if model_group_info.tpm is not None and model_group_info.tpm > 0:
counter_key = self.v3_limiter.create_rate_limit_keys(
@@ -194,29 +202,29 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
value=model,
rate_limit_type="tokens",
)
counter_value = await self.internal_usage_cache.async_get_cache(
key=counter_key,
litellm_parent_otel_span=None,
local_only=False,
)
if counter_value is not None:
current_tokens = float(counter_value)
tpm_saturation = current_tokens / model_group_info.tpm
max_saturation = max(max_saturation, tpm_saturation)
verbose_proxy_logger.debug(
f"Model {model} TPM: {current_tokens}/{model_group_info.tpm} "
f"({tpm_saturation:.1%})"
)
verbose_proxy_logger.debug(
f"Model {model} overall saturation: {max_saturation:.1%}"
)
return max_saturation
except Exception as e:
verbose_proxy_logger.error(
f"Error checking saturation for {model}: {str(e)}"
@@ -232,17 +240,17 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
) -> List[RateLimitDescriptor]:
"""
Create rate limit descriptors with normalized priority weights.
Uses normalized weights to handle over-allocation scenarios.
For explicit priorities: each priority gets its own pool (e.g., prod gets 75%)
For default priority: ALL keys without explicit priority share ONE pool (e.g., all share 25%)
"""
descriptors: List[RateLimitDescriptor] = []
# Get model group info
model_group_info: Optional[ModelGroupInfo] = self.llm_router.get_model_group_info(
model_group=model
model_group_info: Optional[ModelGroupInfo] = (
self.llm_router.get_model_group_info(model_group=model)
)
if model_group_info is None:
return descriptors
@@ -255,21 +263,21 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
normalized_weights=normalized_weights,
model_info=model_group_info,
)
rate_limit_config: RateLimitDescriptorRateLimitObject = {}
# Apply priority weight to model limits
if model_group_info.tpm is not None:
reserved_tpm = int(model_group_info.tpm * priority_weight)
rate_limit_config["tokens_per_unit"] = reserved_tpm
if model_group_info.rpm is not None:
reserved_rpm = int(model_group_info.rpm * priority_weight)
rate_limit_config["requests_per_unit"] = reserved_rpm
if rate_limit_config:
rate_limit_config["window_size"] = self.v3_limiter.window_size
descriptors.append(
RateLimitDescriptor(
key="priority_model",
@@ -288,12 +296,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
) -> RateLimitDescriptor:
"""
Create a descriptor for tracking model-wide usage.
Args:
model: Model name
model_group_info: Model configuration with RPM/TPM limits
high_limit_multiplier: Multiplier for limits (use >1 for tracking-only)
Returns:
Rate limit descriptor for model-wide tracking
"""
@@ -302,18 +310,19 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
value=model,
rate_limit={
"requests_per_unit": (
model_group_info.rpm * high_limit_multiplier
if model_group_info.rpm else None
model_group_info.rpm * high_limit_multiplier
if model_group_info.rpm
else None
),
"tokens_per_unit": (
model_group_info.tpm * high_limit_multiplier
if model_group_info.tpm else None
model_group_info.tpm * high_limit_multiplier
if model_group_info.tpm
else None
),
"window_size": self.v3_limiter.window_size,
},
)
async def _check_rate_limits(
self,
model: str,
@@ -325,23 +334,23 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
) -> None:
"""
Check rate limits using THREE-PHASE approach to prevent partial increments.
Phase 1: Read-only check of ALL limits (no increments)
Phase 2: Decide which limits to enforce based on saturation
Phase 3: Increment ALL counters atomically (model + priority)
This prevents the bug where:
- Model counter increments in stage 1
- Priority check fails in stage 2
- Request blocked but model counter already incremented
Key behaviors:
- All checks performed first (read-only)
- Only increment counters if request will be allowed
- Model capacity: Always enforced at 100%
- Priority limits: Only enforced when saturated >= threshold
- Both counters tracked from first request (accurate accounting)
Args:
model: Model name
model_group_info: Model configuration
@@ -349,17 +358,20 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
key_priority: User's priority level
saturation: Current saturation level
data: Request data dictionary
Raises:
HTTPException: If any limit is exceeded
"""
import json
saturation_threshold = litellm.priority_reservation_settings.saturation_threshold
saturation_threshold = (
litellm.priority_reservation_settings.saturation_threshold
)
should_enforce_priority = saturation >= saturation_threshold
# Build ALL descriptors upfront
descriptors_to_check: List[RateLimitDescriptor] = []
# Model-wide descriptor (always enforce)
model_wide_descriptor = self._create_model_tracking_descriptor(
model=model,
@@ -367,7 +379,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
high_limit_multiplier=1,
)
descriptors_to_check.append(model_wide_descriptor)
# Priority descriptors (always track, conditionally enforce)
priority_descriptors = self._create_priority_based_descriptors(
model=model,
@@ -376,31 +388,33 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
)
if priority_descriptors:
descriptors_to_check.extend(priority_descriptors)
# PHASE 1: Read-only check of ALL limits (no increments)
check_response = await self.v3_limiter.should_rate_limit(
descriptors=descriptors_to_check,
parent_otel_span=user_api_key_dict.parent_otel_span,
read_only=True, # CRITICAL: Don't increment counters yet
)
verbose_proxy_logger.debug(f"Read-only check: {json.dumps(check_response, indent=2)}")
verbose_proxy_logger.debug(
f"Read-only check: {json.dumps(check_response, indent=2)}"
)
# PHASE 2: Decide which limits to enforce
if check_response["overall_code"] == "OVER_LIMIT":
for status in check_response["statuses"]:
if status["code"] == "OVER_LIMIT":
descriptor_key = status["descriptor_key"]
# Model-wide limit exceeded (ALWAYS enforce)
if descriptor_key == "model_saturation_check":
raise HTTPException(
status_code=429,
detail={
"error": f"Model capacity reached for {model}. "
f"Priority: {key_priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Remaining: {status['limit_remaining']}"
f"Priority: {key_priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Remaining: {status['limit_remaining']}"
},
headers={
"retry-after": str(self.v3_limiter.window_size),
@@ -408,7 +422,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
"x-litellm-priority": key_priority or "default",
},
)
# Priority limit exceeded (ONLY enforce when saturated)
elif descriptor_key == "priority_model" and should_enforce_priority:
verbose_proxy_logger.debug(
@@ -419,10 +433,10 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
status_code=429,
detail={
"error": f"Priority-based rate limit exceeded. "
f"Priority: {key_priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Remaining: {status['limit_remaining']}, "
f"Model saturation: {saturation:.1%}"
f"Priority: {key_priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Remaining: {status['limit_remaining']}, "
f"Model saturation: {saturation:.1%}"
},
headers={
"retry-after": str(self.v3_limiter.window_size),
@@ -431,12 +445,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
"x-litellm-saturation": f"{saturation:.2%}",
},
)
# PHASE 3: Increment counters separately to avoid early-exit issues
# Model counter must ALWAYS increment, but priority counter might be over limit
# If we increment them together, v3_limiter's in-memory check will exit early
# and skip incrementing the model counter
# Step 3a: Increment model-wide counter (always)
model_increment_response = await self.v3_limiter.should_rate_limit(
descriptors=[model_wide_descriptor],
@@ -451,11 +465,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
parent_otel_span=user_api_key_dict.parent_otel_span,
read_only=False,
)
# Combine responses for post-call hook
combined_response = {
"overall_code": model_increment_response["overall_code"],
"statuses": model_increment_response["statuses"] + priority_increment_response["statuses"]
"statuses": model_increment_response["statuses"]
+ priority_increment_response["statuses"],
}
data["litellm_proxy_rate_limit_response"] = combined_response
else:
@@ -466,50 +481,39 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Optional[Union[Exception, str, dict]]:
"""
Saturation-aware pre-call hook for priority-based rate limiting.
Flow:
1. Check current saturation level
2. THREE-PHASE rate limit check:
- PHASE 1: Read-only check of ALL limits (no increments)
- PHASE 2: Decide which limits to enforce based on saturation
- PHASE 3: Increment ALL counters atomically if request allowed
This three-phase approach ensures:
- Model capacity is NEVER exceeded (always enforced at 100%)
- Priority usage tracked from first request (accurate metrics)
- Counters only increment when request will be allowed (prevents phantom usage)
- When under-saturated: priorities can borrow unused capacity (generous)
- When saturated: fair allocation based on normalized priority weights (strict)
Example with 100 RPM model, 60% priority allocation, 80% threshold:
- Saturation < 80%: Priority can use up to 100 RPM (model limit enforced only)
- Saturation >= 80%: Priority limited to 60 RPM (both limits enforced)
Prevents bugs where:
- Model counter increments but priority check fails → model over-capacity
- Priority counter increments but not enforced → inaccurate metrics
Args:
user_api_key_dict: User authentication and metadata
cache: Dual cache instance
data: Request data containing model name
call_type: Type of API call being made
Returns:
None if request is allowed, otherwise raises HTTPException
"""
@@ -518,26 +522,29 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
model = data["model"]
key_priority: Optional[str] = user_api_key_dict.metadata.get("priority", None)
# Get model configuration
model_group_info: Optional[ModelGroupInfo] = self.llm_router.get_model_group_info(
model_group=model
model_group_info: Optional[ModelGroupInfo] = (
self.llm_router.get_model_group_info(model_group=model)
)
if model_group_info is None:
verbose_proxy_logger.debug(f"No model group info for {model}, allowing request")
verbose_proxy_logger.debug(
f"No model group info for {model}, allowing request"
)
return None
try:
# STEP 1: Check current saturation level
saturation = await self._check_model_saturation(model, model_group_info)
saturation_threshold = litellm.priority_reservation_settings.saturation_threshold
saturation_threshold = (
litellm.priority_reservation_settings.saturation_threshold
)
verbose_proxy_logger.debug(
f"[Dynamic Rate Limiter] Model={model}, Saturation={saturation:.1%}, "
f"Threshold={saturation_threshold:.1%}, Priority={key_priority}"
)
# STEP 2: Check rate limits in THREE phases
# Phase 1: Read-only check of ALL limits (no increments)
@@ -552,7 +559,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
saturation=saturation,
data=data,
)
except HTTPException:
raise
except Exception as e:
@@ -579,15 +586,22 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
# Add additional priority-specific headers
if isinstance(response, ModelResponse):
key_priority: Optional[str] = user_api_key_dict.metadata.get("priority", None)
key_priority: Optional[str] = user_api_key_dict.metadata.get(
"priority", None
)
# Get existing additional headers
additional_headers = getattr(response, "_hidden_params", {}).get("additional_headers", {}) or {}
additional_headers = (
getattr(response, "_hidden_params", {}).get(
"additional_headers", {}
)
or {}
)
# Add priority information
additional_headers["x-litellm-priority"] = key_priority or "default"
additional_headers["x-litellm-rate-limiter-version"] = "v3"
# Update response
if not hasattr(response, "_hidden_params"):
response._hidden_params = {}
litellm/proxy/hooks/responses_id_security.py
+5 -25
View File
@@ -5,16 +5,7 @@ This hook uses the DBSpendUpdateWriter to batch-write response IDs to the databa
instead of writing immediately on each request.
"""
from typing import (
TYPE_CHECKING,
Any,
AsyncGenerator,
Literal,
Optional,
Tuple,
Union,
cast,
)
from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, Tuple, Union, cast
from fastapi import HTTPException
@@ -29,7 +20,7 @@ from litellm.types.llms.openai import (
BaseLiteLLMOpenAIResponseObject,
ResponsesAPIResponse,
)
from litellm.types.utils import LLMResponseTypes, SpecialEnums
from litellm.types.utils import CallTypesLiteral, LLMResponseTypes, SpecialEnums
if TYPE_CHECKING:
from litellm.caching.caching import DualCache
@@ -45,18 +36,7 @@ class ResponsesIDSecurity(CustomLogger):
user_api_key_dict: "UserAPIKeyAuth",
cache: "DualCache",
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Optional[Union[Exception, str, dict]]:
# MAP all the responses api response ids to the encrypted response ids
responses_api_call_types = {
@@ -136,7 +116,7 @@ class ResponsesIDSecurity(CustomLogger):
split_result = response_id.split("resp_")
if len(split_result) < 2:
return False
remaining_string = split_result[1]
decrypted_value = decrypt_value_helper(
value=remaining_string, key="response_id", return_original_value=True
@@ -161,7 +141,7 @@ class ResponsesIDSecurity(CustomLogger):
split_result = response_id.split("resp_")
if len(split_result) < 2:
return response_id, None, None
remaining_string = split_result[1]
decrypted_value = decrypt_value_helper(
value=remaining_string, key="response_id", return_original_value=True
litellm/proxy/proxy_server.py
+40 -25
View File
@@ -270,9 +270,7 @@ from litellm.proxy.management_endpoints.customer_endpoints import (
from litellm.proxy.management_endpoints.internal_user_endpoints import (
router as internal_user_router,
)
from litellm.proxy.management_endpoints.internal_user_endpoints import (
user_update,
)
from litellm.proxy.management_endpoints.internal_user_endpoints import user_update
from litellm.proxy.management_endpoints.key_management_endpoints import (
delete_verification_tokens,
duration_in_seconds,
@@ -323,9 +321,7 @@ from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router
from litellm.proxy.openai_files_endpoints.files_endpoints import (
router as openai_files_router,
)
from litellm.proxy.openai_files_endpoints.files_endpoints import (
set_files_config,
)
from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
passthrough_endpoint_router,
)
@@ -412,9 +408,7 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
LiteLLM_UpperboundKeyGenerateParams,
)
from litellm.types.realtime import RealtimeQueryParams
from litellm.types.router import (
DeploymentTypedDict,
)
from litellm.types.router import DeploymentTypedDict
from litellm.types.router import ModelInfo as RouterModelInfo
from litellm.types.router import (
RouterGeneralSettings,
@@ -907,15 +901,23 @@ try:
# Support both "true" and "True" for case-insensitive comparison
if os.getenv("LITELLM_NON_ROOT", "").lower() == "true":
non_root_ui_path = "/tmp/litellm_ui"
# Check if the UI was built and exists at the expected location
if os.path.exists(non_root_ui_path) and os.listdir(non_root_ui_path):
verbose_proxy_logger.info(f"Using pre-built UI for non-root Docker: {non_root_ui_path}")
verbose_proxy_logger.info(f"UI files found: {len(os.listdir(non_root_ui_path))} items")
verbose_proxy_logger.info(
f"Using pre-built UI for non-root Docker: {non_root_ui_path}"
)
verbose_proxy_logger.info(
f"UI files found: {len(os.listdir(non_root_ui_path))} items"
)
ui_path = non_root_ui_path
else:
verbose_proxy_logger.error(f"UI not found at {non_root_ui_path}. UI will not be available.")
verbose_proxy_logger.error(f"Path exists: {os.path.exists(non_root_ui_path)}, Has content: {os.path.exists(non_root_ui_path) and bool(os.listdir(non_root_ui_path))}")
verbose_proxy_logger.error(
f"UI not found at {non_root_ui_path}. UI will not be available."
)
verbose_proxy_logger.error(
f"Path exists: {os.path.exists(non_root_ui_path)}, Has content: {os.path.exists(non_root_ui_path) and bool(os.listdir(non_root_ui_path))}"
)
# Only modify files if a custom server root path is set
if server_root_path and server_root_path != "/":
@@ -991,7 +993,9 @@ try:
dst = os.path.join(folder_path, "index.html")
os.rename(src, dst)
else:
verbose_proxy_logger.info("Skipping runtime HTML restructuring for non-root Docker (already done at build time)")
verbose_proxy_logger.info(
"Skipping runtime HTML restructuring for non-root Docker (already done at build time)"
)
except Exception:
pass
@@ -1852,8 +1856,12 @@ class ProxyConfig:
"""
global llm_router
import litellm
if llm_router is not None and litellm.cache is not None and llm_router.cache_responses is not True:
if (
llm_router is not None
and litellm.cache is not None
and llm_router.cache_responses is not True
):
llm_router.cache_responses = True
verbose_proxy_logger.debug(
"Set router.cache_responses=True after initializing cache"
@@ -2308,12 +2316,12 @@ class ProxyConfig:
litellm._key_management_settings = KeyManagementSettings(
**key_management_settings
)
### LOAD SECRET MANAGER ###
key_management_system = general_settings.get("key_management_system", None)
self.initialize_secret_manager(
key_management_system=key_management_system,
config_file_path=config_file_path
config_file_path=config_file_path,
)
### [DEPRECATED] LOAD FROM GOOGLE KMS ### old way of loading from google kms
use_google_kms = general_settings.get("use_google_kms", False)
@@ -2647,7 +2655,9 @@ class ProxyConfig:
pass
def initialize_secret_manager(
self, key_management_system: Optional[str], config_file_path: Optional[str] = None
self,
key_management_system: Optional[str],
config_file_path: Optional[str] = None,
):
"""
Initialize the relevant secret manager if `key_management_system` is provided
@@ -3029,14 +3039,16 @@ class ProxyConfig:
"Error setting env variable: %s - %s", k, str(e)
)
return decrypted_env_vars
def _decrypt_db_variables(self, variables_dict: dict) -> dict:
"""
Decrypts a dictionary of variables and returns them.
"""
decrypted_variables = {}
for k, v in variables_dict.items():
decrypted_value = decrypt_value_helper(value=v, key=k, return_original_value=True)
decrypted_value = decrypt_value_helper(
value=v, key=k, return_original_value=True
)
decrypted_variables[k] = decrypted_value
return decrypted_variables
@@ -3434,6 +3446,7 @@ class ProxyConfig:
from litellm.proxy.management_endpoints.cache_settings_endpoints import (
CacheSettingsManager,
)
await CacheSettingsManager.init_cache_settings_in_db(
prisma_client=prisma_client, proxy_config=self
)
@@ -3689,7 +3702,7 @@ class ProxyConfig:
Initialize search tools from database into the router on startup.
"""
global llm_router
from litellm.proxy.search_endpoints.search_tool_registry import (
SearchToolRegistry,
)
@@ -5081,7 +5094,9 @@ async def embeddings( # noqa: PLR0915
### CALL HOOKS ### - modify incoming data / reject request before calling the model
data = await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_dict, data=data, call_type="aembedding"
user_api_key_dict=user_api_key_dict,
data=data,
call_type=CallTypes.aembedding.value,
)
tasks = []
@@ -5494,7 +5509,7 @@ async def audio_transcriptions(
data = await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_dict,
data=data,
call_type="audio_transcription",
call_type="transcription",
)
## ROUTE TO CORRECT ENDPOINT ##
litellm/proxy/utils.py
+6 -63
View File
@@ -33,7 +33,7 @@ from litellm.proxy._types import (
SpendLogsPayload,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import CallTypes
from litellm.types.utils import CallTypes, CallTypesLiteral
try:
import backoff
@@ -796,19 +796,7 @@ class ProxyLogging:
callback: CustomGuardrail,
data: dict,
user_api_key_dict: Optional[UserAPIKeyAuth],
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"aembedding",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Optional[dict]:
"""
Process a guardrail callback during pre-call hook.
@@ -866,19 +854,7 @@ class ProxyLogging:
self,
user_api_key_dict: UserAPIKeyAuth,
data: None,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"aembedding",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> None:
pass
@@ -887,19 +863,7 @@ class ProxyLogging:
self,
user_api_key_dict: UserAPIKeyAuth,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"aembedding",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> dict:
pass
@@ -907,19 +871,7 @@ class ProxyLogging:
self,
user_api_key_dict: UserAPIKeyAuth,
data: Optional[dict],
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"aembedding",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Optional[dict]:
"""
Allows users to modify/reject the incoming request to the proxy, without having to deal with parsing Request body.
@@ -1046,16 +998,7 @@ class ProxyLogging:
self,
data: dict,
user_api_key_dict: Optional[UserAPIKeyAuth],
call_type: Literal[
"completion",
"responses",
"embeddings",
"aembedding",
"image_generation",
"moderation",
"audio_transcription",
"mcp_call",
],
call_type: CallTypesLiteral,
):
"""
Runs the CustomGuardrail's async_moderation_hook() in parallel
litellm/types/utils.py
+29 -12
View File
@@ -1,16 +1,7 @@
import json
import time
from enum import Enum
from typing import (
TYPE_CHECKING,
Any,
Dict,
List,
Literal,
Mapping,
Optional,
Union,
)
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Union
from aiohttp import FormData
from openai._models import BaseModel as OpenAIObject
@@ -299,7 +290,7 @@ class CallTypes(str, Enum):
avideo_retrieve_job = "avideo_retrieve_job"
video_delete = "video_delete"
avideo_delete = "avideo_delete"
#########################################################
# Container Call Types
#########################################################
@@ -311,7 +302,7 @@ class CallTypes(str, Enum):
aretrieve_container = "aretrieve_container"
delete_container = "delete_container"
adelete_container = "adelete_container"
acancel_fine_tuning_job = "acancel_fine_tuning_job"
cancel_fine_tuning_job = "cancel_fine_tuning_job"
alist_fine_tuning_jobs = "alist_fine_tuning_jobs"
@@ -374,6 +365,7 @@ CallTypesLiteral = Literal[
"aocr",
"avector_store_search",
"vector_store_search",
"call_mcp_tool",
]
@@ -1775,6 +1767,10 @@ class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject):
total_tokens=0,
)
super().__init__(created=created, data=_data, usage=_usage) # type: ignore
self.quality = kwargs.get("quality", None)
self.output_format = kwargs.get("output_format", None)
self.size = kwargs.get("size", None)
self._hidden_params = hidden_params or {}
def __contains__(self, key):
@@ -1801,8 +1797,29 @@ class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject):
return self.dict()
class TranscriptionUsageDurationObject(Usage):
type: Literal["duration"]
seconds: int
class TranscriptionUsageInputTokenDetailsObject(Usage):
audio_tokens: int
text_tokens: int
class TranscriptionUsageTokensObject(Usage):
type: Literal["tokens"]
input_tokens: int
output_tokens: int
total_tokens: int
input_token_details: TranscriptionUsageInputTokenDetailsObject
class TranscriptionResponse(OpenAIObject):
text: Optional[str] = None
usage: Optional[
Union[TranscriptionUsageDurationObject, TranscriptionUsageTokensObject]
] = None
_hidden_params: dict = {}
_response_headers: Optional[dict] = None
poetry.lock
Generated
+235 -721
View File
File diff suppressed because it is too large Load Diff
pyproject.toml
+2
View File
@@ -66,6 +66,7 @@ diskcache = {version = "^5.6.1", optional = true}
polars = {version = "^1.31.0", optional = true, python = ">=3.10"}
semantic-router = {version = "*", optional = true, python = ">=3.9"}
mlflow = {version = ">3.1.4", optional = true, python = ">=3.10"}
soundfile = {version = "^0.12.1", optional = true}
[tool.poetry.extras]
proxy = [
@@ -92,6 +93,7 @@ proxy = [
"litellm-enterprise",
"rich",
"polars",
"soundfile",
]
extra_proxy = [
requirements.txt
+1
View File
@@ -58,6 +58,7 @@ tenacity==8.5.0 # for retrying requests, when litellm.num_retries set
pydantic==2.10.2 # proxy + openai req.
jsonschema==4.22.0 # validating json schema
websockets==13.1.0 # for realtime API
soundfile==0.12.1 # for audio file processing
########################
# LITELLM ENTERPRISE DEPENDENCIES
tests/image_gen_tests/base_image_generation_test.py
-1
View File
@@ -79,7 +79,6 @@ class BaseImageGenTest(ABC):
if "usage" in response_dict:
response_dict["usage"] = dict(response_dict["usage"])
print("response usage=", response_dict.get("usage"))
ImagesResponse.model_validate(response_dict)
for d in response.data:
assert isinstance(d, Image)
tests/image_gen_tests/test_image_generation.py
+12 -11
View File
@@ -118,6 +118,7 @@ class TestVertexImageGeneration(BaseImageGenTest):
"n": 1,
}
class TestBedrockNovaCanvasTextToImage(BaseImageGenTest):
def get_base_image_generation_call_args(self) -> dict:
litellm.in_memory_llm_clients_cache = InMemoryCache()
@@ -164,10 +165,12 @@ class TestAimlImageGeneration(BaseImageGenTest):
def get_base_image_generation_call_args(self) -> dict:
return {"model": "aiml/flux-pro/v1.1"}
class TestFAL_AI_ImageGeneration(BaseImageGenTest):
def get_base_image_generation_call_args(self) -> dict:
return {"model": "fal_ai/fal-ai/imagen4/preview"}
class TestGoogleImageGen(BaseImageGenTest):
def get_base_image_generation_call_args(self) -> dict:
return {"model": "gemini/imagen-4.0-generate-001"}
@@ -175,7 +178,6 @@ class TestGoogleImageGen(BaseImageGenTest):
class TestAzureOpenAIDalle3(BaseImageGenTest):
def get_base_image_generation_call_args(self) -> dict:
litellm.set_verbose = True
return {
"model": "azure/dall-e-3",
"api_version": "2024-02-01",
@@ -189,7 +191,6 @@ class TestAzureOpenAIDalle3(BaseImageGenTest):
}
@pytest.mark.skip(reason="model EOL")
@pytest.mark.asyncio
async def test_aimage_generation_bedrock_with_optional_params():
@@ -281,13 +282,13 @@ async def test_aiml_image_generation_with_dynamic_api_key():
assert captured_json_data["prompt"] == "A cute baby sea otter"
assert captured_json_data["model"] == "flux-pro/v1.1"
@pytest.mark.asyncio
async def test_azure_image_generation_request_body():
from litellm import aimage_generation
test_dir = os.path.dirname(__file__)
expected_path = os.path.join(
test_dir, "request_payloads", "azure_gpt_image_1.json"
)
expected_path = os.path.join(test_dir, "request_payloads", "azure_gpt_image_1.json")
with open(expected_path, "r") as f:
expected_body = json.load(f)
@@ -299,12 +300,12 @@ async def test_azure_image_generation_request_body():
with pytest.raises(Exception):
await aimage_generation(
model="azure/gpt-image-1",
prompt="test prompt",
api_base="https://example.azure.com",
api_key="test-key",
api_version="2025-04-01-preview",
)
model="azure/gpt-image-1",
prompt="test prompt",
api_base="https://example.azure.com",
api_key="test-key",
api_version="2025-04-01-preview",
)
mock_post.assert_called_once()
call_args = mock_post.call_args
tests/local_testing/test_completion_cost.py
+6 -19
View File
@@ -311,17 +311,6 @@ def test_cost_azure_embedding():
# test_cost_azure_embedding()
def test_cost_openai_image_gen():
cost = litellm.completion_cost(
model="dall-e-2",
size="1024-x-1024",
quality="standard",
n=1,
call_type="image_generation",
)
assert cost == 0.019922944
def test_cost_bedrock_pricing_actual_calls():
litellm.set_verbose = True
model = "anthropic.claude-3-5-sonnet-20240620-v1:0"
@@ -1159,10 +1148,10 @@ def test_completion_cost_databricks_embedding(model, monkeypatch):
api_key = "dapimykey"
monkeypatch.setenv("DATABRICKS_API_BASE", base_url)
monkeypatch.setenv("DATABRICKS_API_KEY", api_key)
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
mock_response_data = {
"object": "list",
"model": model.split("/")[1],
@@ -1187,18 +1176,16 @@ def test_completion_cost_databricks_embedding(model, monkeypatch):
"prompt_tokens_details": None,
},
}
mock_response = MagicMock(spec=httpx.Response)
mock_response.status_code = 200
mock_response.json.return_value = mock_response_data
sync_handler = HTTPHandler()
with patch.object(HTTPHandler, "post", return_value=mock_response):
resp = litellm.embedding(
model=model,
input=["hey, how's it going?"],
client=sync_handler
model=model, input=["hey, how's it going?"], client=sync_handler
)
print(resp)
tests/test_litellm/litellm_core_utils/test_audio_utils.py
+84 -27
View File
@@ -11,6 +11,7 @@ import pytest
from litellm.litellm_core_utils.audio_utils.utils import (
ProcessedAudioFile,
calculate_request_duration,
get_audio_file_for_health_check,
get_audio_file_name,
process_audio_file,
@@ -24,7 +25,7 @@ class TestProcessAudioFile:
"""Test processing raw bytes input"""
audio_data = b"fake audio data"
result = process_audio_file(audio_data)
assert isinstance(result, ProcessedAudioFile)
assert result.file_content == audio_data
assert result.filename == "audio.wav"
@@ -34,7 +35,7 @@ class TestProcessAudioFile:
"""Test processing bytearray input"""
audio_data = bytearray(b"fake audio data")
result = process_audio_file(audio_data)
assert isinstance(result, ProcessedAudioFile)
assert result.file_content == bytes(audio_data)
assert result.filename == "audio.wav"
@@ -43,14 +44,14 @@ class TestProcessAudioFile:
def test_process_file_path_input(self):
"""Test processing file path input"""
test_content = b"test audio content"
with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as temp_file:
temp_file.write(test_content)
temp_file_path = temp_file.name
try:
result = process_audio_file(temp_file_path)
assert isinstance(result, ProcessedAudioFile)
assert result.file_content == test_content
assert result.filename == os.path.basename(temp_file_path)
@@ -63,9 +64,9 @@ class TestProcessAudioFile:
filename = "test.wav"
audio_data = b"fake audio data"
audio_tuple = (filename, audio_data)
result = process_audio_file(audio_tuple)
assert isinstance(result, ProcessedAudioFile)
assert result.file_content == audio_data
assert result.filename == filename
@@ -74,17 +75,17 @@ class TestProcessAudioFile:
def test_process_tuple_input_with_file_path(self):
"""Test processing tuple input with file path content"""
test_content = b"test audio content"
with tempfile.NamedTemporaryFile(suffix=".flac", delete=False) as temp_file:
temp_file.write(test_content)
temp_file_path = temp_file.name
try:
filename = "custom_name.flac"
audio_tuple = (filename, temp_file_path)
result = process_audio_file(audio_tuple)
assert isinstance(result, ProcessedAudioFile)
assert result.file_content == test_content
assert result.filename == filename
@@ -97,14 +98,14 @@ class TestProcessAudioFile:
test_content = b"test audio content"
file_obj = io.BytesIO(test_content)
file_obj.name = "test_audio.ogg"
result = process_audio_file(file_obj)
assert isinstance(result, ProcessedAudioFile)
assert result.file_content == test_content
assert result.filename == "test_audio.ogg"
assert result.content_type == "audio/ogg"
# Verify file pointer was reset
assert file_obj.tell() == 0
@@ -112,9 +113,9 @@ class TestProcessAudioFile:
"""Test processing file-like object without name attribute"""
test_content = b"test audio content"
file_obj = io.BytesIO(test_content)
result = process_audio_file(file_obj)
assert isinstance(result, ProcessedAudioFile)
assert result.file_content == test_content
assert result.filename == "audio.wav"
@@ -124,17 +125,17 @@ class TestProcessAudioFile:
"""Test processing tuple with file-like object as content"""
test_content = b"test audio content"
file_obj = io.BytesIO(test_content)
filename = "custom.mp3"
audio_tuple = (filename, file_obj)
result = process_audio_file(audio_tuple)
assert isinstance(result, ProcessedAudioFile)
assert result.file_content == test_content
assert result.filename == filename
assert result.content_type == "audio/mpeg"
# Verify file pointer was reset
assert file_obj.tell() == 0
@@ -148,7 +149,7 @@ class TestProcessAudioFile:
("test.aac", "audio/aac"),
("test.m4a", "audio/x-m4a"),
]
for filename, expected_mime_type in test_cases:
audio_tuple = (filename, b"fake content")
result = process_audio_file(audio_tuple)
@@ -158,24 +159,25 @@ class TestProcessAudioFile:
"""Test MIME type fallback for unknown file extensions"""
audio_tuple = ("test.unknown", b"fake content")
result = process_audio_file(audio_tuple)
assert result.content_type == "audio/wav" # Should fallback to default
def test_process_pathlike_object(self):
"""Test processing os.PathLike object"""
test_content = b"test audio content"
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
temp_file.write(test_content)
temp_file_path = temp_file.name
try:
# Convert to pathlib.Path
from pathlib import Path
path_obj = Path(temp_file_path)
result = process_audio_file(path_obj)
assert isinstance(result, ProcessedAudioFile)
assert result.file_content == test_content
assert result.filename == os.path.basename(temp_file_path)
@@ -202,7 +204,62 @@ class TestProcessAudioFile:
"""Test tuple with None filename gets default name"""
audio_tuple = (None, b"fake content")
result = process_audio_file(audio_tuple)
assert result.filename == "audio.wav"
assert result.content_type == "audio/wav"
class TestCalculateRequestDuration:
"""Test the calculate_request_duration function"""
@pytest.mark.skipif(
os.environ.get("SKIP_AUDIO_TESTS") == "true",
reason="Skipping audio tests - soundfile may not be available",
)
def test_bytesio_at_end_position(self):
"""
Test that calculate_request_duration handles BytesIO with file pointer at end.
This reproduces and verifies the fix for the OGG file bug where BytesIO
position was at the end after a previous read(), causing "Format not recognised" error.
"""
# Create a simple WAV file in memory (44 bytes header + some data)
# This is a minimal valid WAV file
wav_header = (
b"RIFF"
+ (36 + 8).to_bytes(4, "little") # ChunkSize
+ b"WAVE"
+ b"fmt "
+ (16).to_bytes(4, "little") # Subchunk1Size
+ (1).to_bytes(2, "little") # AudioFormat (PCM)
+ (1).to_bytes(2, "little") # NumChannels
+ (16000).to_bytes(4, "little") # SampleRate
+ (32000).to_bytes(4, "little") # ByteRate
+ (2).to_bytes(2, "little") # BlockAlign
+ (16).to_bytes(2, "little") # BitsPerSample
+ b"data"
+ (8).to_bytes(4, "little") # Subchunk2Size
+ b"\x00\x00\x00\x00\x00\x00\x00\x00" # Sample data
)
# Create BytesIO object
file_obj = io.BytesIO(wav_header)
file_obj.name = "test_audio.wav"
# Simulate the bug: something reads from the file first, moving position to end
_ = file_obj.read()
assert file_obj.tell() == len(wav_header), "File position should be at end"
# Call calculate_request_duration - this would fail before the fix
duration = calculate_request_duration(file_obj)
# Verify it succeeded (returns a duration, not None)
assert (
duration is not None
), "Duration should be calculated even when BytesIO is at end"
assert isinstance(duration, float), "Duration should be a float"
assert duration > 0, "Duration should be positive"
# Verify the file position was restored
assert file_obj.tell() == len(
wav_header
), "File position should be restored to original position"
tests/test_litellm/test_cost_calculator.py
+70 -24
View File
@@ -697,7 +697,7 @@ def test_cost_discount_vertex_ai():
# Save original config
original_discount_config = litellm.cost_discount_config.copy()
# Create mock response
response = ModelResponse(
id="test-id",
@@ -705,13 +705,9 @@ def test_cost_discount_vertex_ai():
created=1234567890,
model="gemini-pro",
object="chat.completion",
usage=Usage(
prompt_tokens=100,
completion_tokens=50,
total_tokens=150
)
usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
)
# Calculate cost without discount
litellm.cost_discount_config = {}
cost_without_discount = completion_cost(
@@ -719,24 +715,24 @@ def test_cost_discount_vertex_ai():
model="vertex_ai/gemini-pro",
custom_llm_provider="vertex_ai",
)
# Set 5% discount for vertex_ai
litellm.cost_discount_config = {"vertex_ai": 0.05}
# Calculate cost with discount
cost_with_discount = completion_cost(
completion_response=response,
model="vertex_ai/gemini-pro",
custom_llm_provider="vertex_ai",
)
# Restore original config
litellm.cost_discount_config = original_discount_config
# Verify discount is applied (5% off means 95% of original cost)
expected_cost = cost_without_discount * 0.95
assert cost_with_discount == pytest.approx(expected_cost, rel=1e-9)
print(f"✓ Cost discount test passed:")
print(f" - Original cost: ${cost_without_discount:.6f}")
print(f" - Discounted cost (5% off): ${cost_with_discount:.6f}")
@@ -752,7 +748,7 @@ def test_cost_discount_not_applied_to_other_providers():
# Save original config
original_discount_config = litellm.cost_discount_config.copy()
# Create mock response for OpenAI
response = ModelResponse(
id="test-id",
@@ -760,23 +756,19 @@ def test_cost_discount_not_applied_to_other_providers():
created=1234567890,
model="gpt-4",
object="chat.completion",
usage=Usage(
prompt_tokens=100,
completion_tokens=50,
total_tokens=150
)
usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
)
# Set discount only for vertex_ai (not openai)
litellm.cost_discount_config = {"vertex_ai": 0.05}
# Calculate cost for OpenAI - should NOT have discount applied
cost_with_selective_discount = completion_cost(
completion_response=response,
model="gpt-4",
custom_llm_provider="openai",
)
# Clear discount config
litellm.cost_discount_config = {}
cost_without_discount = completion_cost(
@@ -784,13 +776,67 @@ def test_cost_discount_not_applied_to_other_providers():
model="gpt-4",
custom_llm_provider="openai",
)
# Restore original config
litellm.cost_discount_config = original_discount_config
# Costs should be the same (no discount applied to OpenAI)
assert cost_with_selective_discount == cost_without_discount
print(f"✓ Selective discount test passed:")
print(f" - OpenAI cost (no discount configured): ${cost_without_discount:.6f}")
print(f" - Cost remains unchanged: ${cost_with_selective_discount:.6f}")
def test_azure_image_generation_cost_calculator():
from unittest.mock import MagicMock
from litellm.types.utils import (
ImageObject,
ImageResponse,
ImageUsage,
ImageUsageInputTokensDetails,
)
response_cost_calculator_kwargs = {
"response_object": ImageResponse(
created=1761785270,
background=None,
data=[
ImageObject(
b64_json=None,
revised_prompt="A futuristic, techno-inspired green duck wearing cool modern sunglasses. The duck has a sleek, metallic appearance with glowing neon green accents, standing on a high-tech urban background with holographic billboards and illuminated city lights in the distance. The duck's feathers have a glossy, high-tech sheen, resembling a robotic design but still maintaining its avian features. The scene has a vibrant, cyberpunk aesthetic with a neon color palette.",
url="https://dalleprodsec.blob.core.windows.net/private/images/caa17dc4-357d-4257-8938-eeea9baa8d0a/generated_00.png?se=2025-10-31T00%3A47%3A59Z&sig=KHRjLz3vMahbw94JtxL02S6t2AueeRMaiqj4z35HKDM%3D&ske=2025-11-05T00%3A26%3A20Z&skoid=e52d5ed7-0657-4f62-bc12-7e5dbb260a96&sks=b&skt=2025-10-29T00%3A26%3A20Z&sktid=33e01921-4d64-4f8c-a055-5bdaffd5e33d&skv=2020-10-02&sp=r&spr=https&sr=b&sv=2020-10-02",
)
],
output_format=None,
quality="hd",
size=None,
usage=ImageUsage(
input_tokens=0,
input_tokens_details=ImageUsageInputTokensDetails(
image_tokens=0, text_tokens=0
),
output_tokens=0,
total_tokens=0,
),
),
"model": "azure/dall-e-3",
"cache_hit": False,
"custom_llm_provider": "azure",
"base_model": "azure/dall-e-3",
"call_type": "aimage_generation",
"optional_params": {},
"custom_pricing": False,
"prompt": "",
"standard_built_in_tools_params": {
"web_search_options": None,
"file_search": None,
},
"router_model_id": "6738c432ffc9b733597c6b86613ca20dc5f49bde591fd3d03e7cd6aa25bb241e",
"litellm_logging_obj": MagicMock(),
"service_tier": None,
}
cost = response_cost_calculator(**response_cost_calculator_kwargs)
assert cost > 0.079