mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
Merge branch 'main' into feat/improve-opik
This commit is contained in:
mrFranklin
commit
cdfb32fd16
39 changed files with 1229 additions and 74 deletions
|
|
@ -144,7 +144,7 @@ jobs:
|
|||
name: Linting Testing
|
||||
command: |
|
||||
cd litellm
|
||||
pip install "cryptography<40.0.0"
|
||||
pip install "cryptography>=43.0.1"
|
||||
python -m pip install types-requests types-setuptools types-redis types-PyYAML
|
||||
if ! python -m mypy . \
|
||||
--config-file mypy.ini \
|
||||
|
|
|
|||
33
README.md
33
README.md
|
|
@ -350,13 +350,21 @@ curl 'http://0.0.0.0:4000/key/generate' \
|
|||
|
||||
[**Read the Docs**](https://docs.litellm.ai/docs/)
|
||||
|
||||
## Contributing
|
||||
## Run in Developer mode
|
||||
### Services
|
||||
1. Setup .env file in root
|
||||
2. Run dependant services `docker-compose up db prometheus`
|
||||
|
||||
Interested in contributing? Contributions to LiteLLM Python SDK, Proxy Server, and LLM integrations are both accepted and highly encouraged!
|
||||
### Backend
|
||||
1. (In root) create virtual environment `python -m venv .venv`
|
||||
2. Activate virtual environment `source .venv/bin/activate`
|
||||
3. Install dependencies `pip install -e ".[all]"`
|
||||
4. Start proxy backend `python litellm/proxy_cli.py`
|
||||
|
||||
**Quick start:** `git clone` → `make install-dev` → `make format` → `make lint` → `make test-unit`
|
||||
|
||||
See our comprehensive [Contributing Guide (CONTRIBUTING.md)](CONTRIBUTING.md) for detailed instructions.
|
||||
### Frontend
|
||||
1. Navigate to `ui/litellm-dashboard`
|
||||
2. Install dependencies `npm install`
|
||||
3. Run `npm run dev` to start the dashboard
|
||||
|
||||
# Enterprise
|
||||
For companies that need better security, user management and professional support
|
||||
|
|
@ -434,18 +442,3 @@ All these checks must pass before your PR can be merged.
|
|||
</a>
|
||||
|
||||
|
||||
## Run in Developer mode
|
||||
### Services
|
||||
1. Setup .env file in root
|
||||
2. Run dependant services `docker-compose up db prometheus`
|
||||
|
||||
### Backend
|
||||
1. (In root) create virtual environment `python -m venv .venv`
|
||||
2. Activate virtual environment `source .venv/bin/activate`
|
||||
3. Install dependencies `pip install -e ".[all]"`
|
||||
4. Start proxy backend `python3 /path/to/litellm/proxy_cli.py`
|
||||
|
||||
### Frontend
|
||||
1. Navigate to `ui/litellm-dashboard`
|
||||
2. Install dependencies `npm install`
|
||||
3. Run `npm run dev` to start the dashboard
|
||||
|
|
|
|||
14
docs/my-website/package-lock.json
generated
14
docs/my-website/package-lock.json
generated
|
|
@ -17120,9 +17120,10 @@
|
|||
}
|
||||
},
|
||||
"node_modules/prebuild-install/node_modules/tar-fs": {
|
||||
"version": "2.1.3",
|
||||
"resolved": "https://registry.npmjs.org/tar-fs/-/tar-fs-2.1.3.tgz",
|
||||
"integrity": "sha512-090nwYJDmlhwFwEW3QQl+vaNnxsO2yVsd45eTKRBzSzu+hlb1w2K9inVq5b0ngXuLVqQ4ApvsUHHnu/zQNkWAg==",
|
||||
"version": "2.1.4",
|
||||
"resolved": "https://registry.npmjs.org/tar-fs/-/tar-fs-2.1.4.tgz",
|
||||
"integrity": "sha512-mDAjwmZdh7LTT6pNleZ05Yt65HC3E+NiQzl672vQG38jIrehtJk/J3mNwIg+vShQPcLF/LV7CMnDW6vjj6sfYQ==",
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"chownr": "^1.1.1",
|
||||
"mkdirp-classic": "^0.5.2",
|
||||
|
|
@ -19295,9 +19296,10 @@
|
|||
}
|
||||
},
|
||||
"node_modules/tar-fs": {
|
||||
"version": "3.0.10",
|
||||
"resolved": "https://registry.npmjs.org/tar-fs/-/tar-fs-3.0.10.tgz",
|
||||
"integrity": "sha512-C1SwlQGNLe/jPNqapK8epDsXME7CAJR5RL3GcE6KWx1d9OUByzoHVcbu1VPI8tevg9H8Alae0AApHHFGzrD5zA==",
|
||||
"version": "3.1.1",
|
||||
"resolved": "https://registry.npmjs.org/tar-fs/-/tar-fs-3.1.1.tgz",
|
||||
"integrity": "sha512-LZA0oaPOc2fVo82Txf3gw+AkEd38szODlptMYejQUhndHMLQ9M059uXR+AfS7DNo0NpINvSqDsvyaCrBVkptWg==",
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"pump": "^3.0.0",
|
||||
"tar-stream": "^3.1.5"
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@ class InMemoryCache(BaseCache):
|
|||
max_size_in_memory [int]: Maximum number of items in cache. done to prevent memory leaks. Use 200 items as a default
|
||||
"""
|
||||
self.max_size_in_memory = (
|
||||
max_size_in_memory or 200
|
||||
max_size_in_memory if max_size_in_memory is not None else 200
|
||||
) # set an upper bound of 200 items in-memory
|
||||
self.default_ttl = default_ttl or 600
|
||||
self.max_size_per_item = (
|
||||
|
|
@ -103,20 +103,32 @@ class InMemoryCache(BaseCache):
|
|||
def evict_cache(self):
|
||||
"""
|
||||
Eviction policy:
|
||||
- check if any items in ttl_dict are expired -> remove them from ttl_dict and cache_dict
|
||||
1. First, remove expired items from ttl_dict and cache_dict
|
||||
2. If cache is still at or above max_size_in_memory, evict items with earliest expiration times
|
||||
|
||||
|
||||
This guarantees the following:
|
||||
- 1. When item ttl not set: At minimumm each item will remain in memory for 5 minutes
|
||||
- 2. When ttl is set: the item will remain in memory for at least that amount of time
|
||||
- 1. When item ttl not set: At minimum each item will remain in memory for the default ttl
|
||||
- 2. When ttl is set: the item will remain in memory for at least that amount of time, unless cache size requires eviction
|
||||
- 3. the size of in-memory cache is bounded
|
||||
|
||||
"""
|
||||
current_time = time.time()
|
||||
|
||||
# Step 1: Remove expired items
|
||||
expired_keys = [key for key, ttl in self.ttl_dict.items() if current_time > ttl]
|
||||
for key in expired_keys:
|
||||
self._remove_key(key)
|
||||
|
||||
# Step 2: If cache is still full, evict items with earliest expiration times
|
||||
if len(self.cache_dict) >= self.max_size_in_memory:
|
||||
# Sort by expiration time (earliest first) and evict until we're under the limit
|
||||
items_by_expiration = sorted(self.ttl_dict.items(), key=lambda x: x[1])
|
||||
keys_to_evict = items_by_expiration[:len(self.cache_dict) - self.max_size_in_memory + 1]
|
||||
|
||||
for key, _ in keys_to_evict:
|
||||
self._remove_key(key)
|
||||
|
||||
# de-reference the removed item
|
||||
# https://www.geeksforgeeks.org/diagnosing-and-fixing-memory-leaks-in-python/
|
||||
# One of the most common causes of memory leaks in Python is the retention of objects that are no longer being used.
|
||||
|
|
@ -135,6 +147,10 @@ class InMemoryCache(BaseCache):
|
|||
return False
|
||||
|
||||
def set_cache(self, key, value, **kwargs):
|
||||
# Handle the edge case where max_size_in_memory is 0
|
||||
if self.max_size_in_memory == 0:
|
||||
return # Don't cache anything if max size is 0
|
||||
|
||||
if len(self.cache_dict) >= self.max_size_in_memory:
|
||||
# only evict when cache is full
|
||||
self.evict_cache()
|
||||
|
|
|
|||
|
|
@ -812,6 +812,11 @@ BEDROCK_EMBEDDING_PROVIDERS_LITERAL = Literal[
|
|||
]
|
||||
|
||||
BEDROCK_CONVERSE_MODELS = [
|
||||
"qwen.qwen3-coder-480b-a35b-v1:0",
|
||||
"qwen.qwen3-235b-a22b-2507-v1:0",
|
||||
"qwen.qwen3-coder-30b-a3b-v1:0",
|
||||
"qwen.qwen3-32b-v1:0",
|
||||
"deepseek.v3-v1:0",
|
||||
"openai.gpt-oss-20b-1:0",
|
||||
"openai.gpt-oss-120b-1:0",
|
||||
"anthropic.claude-opus-4-1-20250805-v1:0",
|
||||
|
|
|
|||
|
|
@ -671,6 +671,7 @@ class LangFuseLogger:
|
|||
|
||||
generation_id = None
|
||||
usage = None
|
||||
usage_details = None
|
||||
if response_obj is not None:
|
||||
if (
|
||||
hasattr(response_obj, "id")
|
||||
|
|
@ -687,6 +688,11 @@ class LangFuseLogger:
|
|||
"completion_tokens": _usage_obj.completion_tokens,
|
||||
"total_cost": cost if self._supports_costs() else None,
|
||||
}
|
||||
usage_details = LangfuseUsageDetails(input=_usage_obj.prompt_tokens,
|
||||
output=_usage_obj.completion_tokens,
|
||||
cache_creation_input_tokens=_usage_obj.get('cache_creation_input_tokens', 0),
|
||||
cache_read_input_tokens=_usage_obj.get('cache_read_input_tokens', 0))
|
||||
|
||||
generation_name = clean_metadata.pop("generation_name", None)
|
||||
if generation_name is None:
|
||||
# if `generation_name` is None, use sensible default values
|
||||
|
|
@ -719,6 +725,7 @@ class LangFuseLogger:
|
|||
"input": input if not mask_input else "redacted-by-litellm",
|
||||
"output": output if not mask_output else "redacted-by-litellm",
|
||||
"usage": usage,
|
||||
"usage_details": usage_details,
|
||||
"metadata": log_requester_metadata(clean_metadata),
|
||||
"level": level,
|
||||
"version": clean_metadata.pop("version", None),
|
||||
|
|
|
|||
|
|
@ -16,11 +16,35 @@ Memory Management Strategy:
|
|||
|
||||
from typing import Any, Callable, Optional, Type, TypeVar
|
||||
|
||||
from pond import Pond, PooledObject, PooledObjectFactory
|
||||
try:
|
||||
from pond import Pond, PooledObject, PooledObjectFactory # type: ignore
|
||||
POND_AVAILABLE = True
|
||||
except ImportError: # pragma: no cover
|
||||
POND_AVAILABLE = False
|
||||
class Pond: # type: ignore
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
def register(self, *args: Any, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
def borrow(self, *args: Any, **kwargs: Any) -> Any:
|
||||
pass
|
||||
|
||||
def recycle(self, *args: Any, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
class PooledObject: # type: ignore
|
||||
def __init__(self, keeped_object: Any = None) -> None:
|
||||
self.keeped_object = keeped_object
|
||||
|
||||
class PooledObjectFactory: # type: ignore
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
T = TypeVar('T')
|
||||
|
||||
class GenericPooledObjectFactory(PooledObjectFactory):
|
||||
class GenericPooledObjectFactory(PooledObjectFactory): # type: ignore[misc]
|
||||
"""Generic factory class for creating pooled objects of any type."""
|
||||
|
||||
def __init__(
|
||||
|
|
@ -79,7 +103,7 @@ def get_object_pool(
|
|||
time_between_eviction_runs: int = 300, # Less frequent eviction to maintain high reuse ratio
|
||||
eviction_weight: float = 0.3, # Less aggressive eviction for better reuse
|
||||
prewarm_count: int = 5 # Lower pre-warm count to reduce initial memory usage
|
||||
) -> Pond:
|
||||
) -> Pond | None:
|
||||
"""Get or create a global object pool instance with balanced eviction-based memory control.
|
||||
|
||||
Memory is controlled through moderate eviction to balance reuse ratio and memory usage:
|
||||
|
|
@ -98,9 +122,13 @@ def get_object_pool(
|
|||
prewarm_count: Number of objects to pre-warm the pool with (default: 5)
|
||||
|
||||
Returns:
|
||||
Pond instance for the specified object type
|
||||
Pond instance for the specified object type or None if pond is not available
|
||||
"""
|
||||
|
||||
# If pond is not available, disable pooling gracefully
|
||||
if not POND_AVAILABLE:
|
||||
return None
|
||||
|
||||
if pool_name in _pools:
|
||||
return _pools[pool_name]
|
||||
|
||||
|
|
@ -134,4 +162,4 @@ def _prewarm_pool(pond: Pond, pool_name: str, prewarm_count: int = 20) -> None:
|
|||
pond.recycle(pooled_obj, name=f"{pool_name}Factory")
|
||||
except Exception:
|
||||
# If pre-warming fails, just continue
|
||||
break
|
||||
break
|
||||
|
|
|
|||
|
|
@ -2032,6 +2032,36 @@
|
|||
"supports_tool_choice": false,
|
||||
"supports_vision": true
|
||||
},
|
||||
"azure/gpt-5-codex": {
|
||||
"cache_read_input_token_cost": 1.25e-07,
|
||||
"input_cost_per_token": 1.25e-06,
|
||||
"litellm_provider": "azure",
|
||||
"max_input_tokens": 272000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"output_cost_per_token": 1e-05,
|
||||
"supported_endpoints": [
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"azure/gpt-5-mini": {
|
||||
"cache_read_input_token_cost": 2.5e-08,
|
||||
"input_cost_per_token": 2.5e-07,
|
||||
|
|
@ -5282,6 +5312,49 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"deepseek-chat": {
|
||||
"cache_read_input_token_cost": 6e-08,
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "deepseek",
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.7e-06,
|
||||
"source": "https://api-docs.deepseek.com/quick_start/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"deepseek-reasoner": {
|
||||
"cache_read_input_token_cost": 6e-08,
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "deepseek",
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 65536,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.7e-06,
|
||||
"source": "https://api-docs.deepseek.com/quick_start/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supports_function_calling": false,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": false,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": false
|
||||
},
|
||||
"dashscope/qwen-coder": {
|
||||
"input_cost_per_token": 3e-07,
|
||||
"litellm_provider": "dashscope",
|
||||
|
|
@ -7203,6 +7276,18 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"deepseek.v3-v1:0": {
|
||||
"input_cost_per_token": 5.8e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 163840,
|
||||
"max_output_tokens": 81920,
|
||||
"max_tokens": 163840,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.68e-06,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"dolphin": {
|
||||
"input_cost_per_token": 5e-07,
|
||||
"litellm_provider": "nlp_cloud",
|
||||
|
|
@ -12415,6 +12500,36 @@
|
|||
"supports_tool_choice": false,
|
||||
"supports_vision": true
|
||||
},
|
||||
"gpt-5-codex": {
|
||||
"cache_read_input_token_cost": 1.25e-07,
|
||||
"input_cost_per_token": 1.25e-06,
|
||||
"litellm_provider": "openai",
|
||||
"max_input_tokens": 272000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"output_cost_per_token": 1e-05,
|
||||
"supported_endpoints": [
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": false,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": false,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"gpt-5-mini": {
|
||||
"cache_read_input_token_cost": 2.5e-08,
|
||||
"cache_read_input_token_cost_flex": 1.25e-08,
|
||||
|
|
@ -17520,6 +17635,54 @@
|
|||
"mode": "chat",
|
||||
"output_cost_per_token": 2.8e-07
|
||||
},
|
||||
"qwen.qwen3-coder-480b-a35b-v1:0": {
|
||||
"input_cost_per_token": 2.2e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 262000,
|
||||
"max_output_tokens": 65536,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.8e-06,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"qwen.qwen3-235b-a22b-2507-v1:0": {
|
||||
"input_cost_per_token": 2.2e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 8.8e-07,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"qwen.qwen3-coder-30b-a3b-v1:0": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6.0e-07,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"qwen.qwen3-32b-v1:0": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 16384,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6.0e-07,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"recraft/recraftv2": {
|
||||
"litellm_provider": "recraft",
|
||||
"mode": "image_generation",
|
||||
|
|
@ -17946,6 +18109,32 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"sambanova/DeepSeek-V3.1": {
|
||||
"max_tokens": 32768,
|
||||
"max_input_tokens": 32768,
|
||||
"max_output_tokens": 32768,
|
||||
"input_cost_per_token": 3e-06,
|
||||
"output_cost_per_token": 4.5e-06,
|
||||
"litellm_provider": "sambanova",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_reasoning": true,
|
||||
"source": "https://cloud.sambanova.ai/plans/pricing"
|
||||
},
|
||||
"sambanova/gpt-oss-120b": {
|
||||
"max_tokens": 131072,
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 131072,
|
||||
"input_cost_per_token": 3e-06,
|
||||
"output_cost_per_token": 4.5e-06,
|
||||
"litellm_provider": "sambanova",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_reasoning": true,
|
||||
"source": "https://cloud.sambanova.ai/plans/pricing"
|
||||
},
|
||||
"sample_spec": {
|
||||
"code_interpreter_cost_per_session": 0.0,
|
||||
"computer_use_input_cost_per_1k_tokens": 0.0,
|
||||
|
|
@ -20467,6 +20656,24 @@
|
|||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"vertex_ai/deepseek-ai/deepseek-v3.1-maas": {
|
||||
"input_cost_per_token": 1.35e-06,
|
||||
"litellm_provider": "vertex_ai-deepseek_models",
|
||||
"max_input_tokens": 163840,
|
||||
"max_output_tokens": 32768,
|
||||
"max_tokens": 163840,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5.4e-06,
|
||||
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models",
|
||||
"supported_regions": [
|
||||
"us-west2"
|
||||
],
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"vertex_ai/deepseek-ai/deepseek-r1-0528-maas": {
|
||||
"input_cost_per_token": 1.35e-06,
|
||||
"litellm_provider": "vertex_ai-deepseek_models",
|
||||
|
|
@ -20880,6 +21087,30 @@
|
|||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "vertex_ai-qwen_models",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.2e-06,
|
||||
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"vertex_ai/qwen/qwen3-next-80b-a3b-thinking-maas": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "vertex_ai-qwen_models",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.2e-06,
|
||||
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"vertex_ai/veo-2.0-generate-001": {
|
||||
"litellm_provider": "vertex_ai-video-models",
|
||||
"max_input_tokens": 1024,
|
||||
|
|
|
|||
|
|
@ -422,6 +422,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/user/update",
|
||||
"/user/delete",
|
||||
"/user/info",
|
||||
"/user/list",
|
||||
# team
|
||||
"/team/new",
|
||||
"/team/update",
|
||||
|
|
|
|||
|
|
@ -3,9 +3,10 @@ import os
|
|||
from datetime import datetime
|
||||
from typing import Dict, List, Literal, Optional, Tuple, Union
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
|
|
@ -237,9 +238,8 @@ class LakeraAIGuardrail(CustomGuardrail):
|
|||
)
|
||||
else:
|
||||
# If there are other violations or not set to mask PII, raise exception
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message="Lakera AI flagged this request. Please review the request and try again.",
|
||||
raise self._get_http_exception_for_blocked_guardrail(
|
||||
lakera_guardrail_response
|
||||
)
|
||||
|
||||
#########################################################
|
||||
|
|
@ -304,9 +304,8 @@ class LakeraAIGuardrail(CustomGuardrail):
|
|||
)
|
||||
else:
|
||||
# If there are other violations or not set to mask PII, raise exception
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message="Lakera AI flagged this request. Please review the request and try again.",
|
||||
raise self._get_http_exception_for_blocked_guardrail(
|
||||
lakera_guardrail_response
|
||||
)
|
||||
|
||||
#########################################################
|
||||
|
|
@ -327,8 +326,32 @@ class LakeraAIGuardrail(CustomGuardrail):
|
|||
if not lakera_response:
|
||||
return False
|
||||
|
||||
for item in lakera_response.get("payload", []) or []:
|
||||
detector_type = item.get("detector_type", "") or ""
|
||||
if not detector_type.startswith("pii/"):
|
||||
return False
|
||||
return True
|
||||
# Check breakdown field for detected violations
|
||||
breakdown = lakera_response.get("breakdown", []) or []
|
||||
if not breakdown:
|
||||
return False
|
||||
|
||||
has_violations = False
|
||||
for item in breakdown:
|
||||
if item.get("detected", False):
|
||||
has_violations = True
|
||||
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
|
||||
|
||||
def _get_http_exception_for_blocked_guardrail(
|
||||
self, lakera_response: Optional[LakeraAIResponse]
|
||||
) -> HTTPException:
|
||||
"""
|
||||
Get the HTTP exception for a blocked guardrail, similar to Bedrock's implementation.
|
||||
"""
|
||||
return HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": "Violated guardrail policy",
|
||||
"lakera_guardrail_response": lakera_response,
|
||||
},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -23,3 +23,16 @@ model_list:
|
|||
litellm_params:
|
||||
model: gemini/*
|
||||
api_key: os.environ/GEMINI_API_KEY
|
||||
|
||||
|
||||
guardrails:
|
||||
- guardrail_name: lakera
|
||||
litellm_params:
|
||||
guardrail: lakera_v2
|
||||
mode: pre_call
|
||||
api_key: os.environ/LAKERA_API_KEY
|
||||
default_on: false
|
||||
project_id: project-9770817088
|
||||
breakdown: true
|
||||
payload: true
|
||||
dev_info: true
|
||||
|
|
|
|||
|
|
@ -7,3 +7,10 @@ class LangfuseLoggingConfig(TypedDict):
|
|||
langfuse_secret: Optional[str]
|
||||
langfuse_public_key: Optional[str]
|
||||
langfuse_host: Optional[str]
|
||||
|
||||
|
||||
class LangfuseUsageDetails(TypedDict):
|
||||
input: Optional[int]
|
||||
output: Optional[int]
|
||||
cache_creation_input_tokens: Optional[int]
|
||||
cache_read_input_tokens: Optional[int]
|
||||
|
|
|
|||
|
|
@ -2032,6 +2032,36 @@
|
|||
"supports_tool_choice": false,
|
||||
"supports_vision": true
|
||||
},
|
||||
"azure/gpt-5-codex": {
|
||||
"cache_read_input_token_cost": 1.25e-07,
|
||||
"input_cost_per_token": 1.25e-06,
|
||||
"litellm_provider": "azure",
|
||||
"max_input_tokens": 272000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"output_cost_per_token": 1e-05,
|
||||
"supported_endpoints": [
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"azure/gpt-5-mini": {
|
||||
"cache_read_input_token_cost": 2.5e-08,
|
||||
"input_cost_per_token": 2.5e-07,
|
||||
|
|
@ -5282,6 +5312,49 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"deepseek-chat": {
|
||||
"cache_read_input_token_cost": 6e-08,
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "deepseek",
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.7e-06,
|
||||
"source": "https://api-docs.deepseek.com/quick_start/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"deepseek-reasoner": {
|
||||
"cache_read_input_token_cost": 6e-08,
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "deepseek",
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 65536,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.7e-06,
|
||||
"source": "https://api-docs.deepseek.com/quick_start/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supports_function_calling": false,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": false,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": false
|
||||
},
|
||||
"dashscope/qwen-coder": {
|
||||
"input_cost_per_token": 3e-07,
|
||||
"litellm_provider": "dashscope",
|
||||
|
|
@ -7203,6 +7276,18 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"deepseek.v3-v1:0": {
|
||||
"input_cost_per_token": 5.8e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 163840,
|
||||
"max_output_tokens": 81920,
|
||||
"max_tokens": 163840,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.68e-06,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"dolphin": {
|
||||
"input_cost_per_token": 5e-07,
|
||||
"litellm_provider": "nlp_cloud",
|
||||
|
|
@ -12415,6 +12500,36 @@
|
|||
"supports_tool_choice": false,
|
||||
"supports_vision": true
|
||||
},
|
||||
"gpt-5-codex": {
|
||||
"cache_read_input_token_cost": 1.25e-07,
|
||||
"input_cost_per_token": 1.25e-06,
|
||||
"litellm_provider": "openai",
|
||||
"max_input_tokens": 272000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"output_cost_per_token": 1e-05,
|
||||
"supported_endpoints": [
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": false,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": false,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"gpt-5-mini": {
|
||||
"cache_read_input_token_cost": 2.5e-08,
|
||||
"cache_read_input_token_cost_flex": 1.25e-08,
|
||||
|
|
@ -17520,6 +17635,54 @@
|
|||
"mode": "chat",
|
||||
"output_cost_per_token": 2.8e-07
|
||||
},
|
||||
"qwen.qwen3-coder-480b-a35b-v1:0": {
|
||||
"input_cost_per_token": 2.2e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 262000,
|
||||
"max_output_tokens": 65536,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.8e-06,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"qwen.qwen3-235b-a22b-2507-v1:0": {
|
||||
"input_cost_per_token": 2.2e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 8.8e-07,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"qwen.qwen3-coder-30b-a3b-v1:0": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6.0e-07,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"qwen.qwen3-32b-v1:0": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 16384,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6.0e-07,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"recraft/recraftv2": {
|
||||
"litellm_provider": "recraft",
|
||||
"mode": "image_generation",
|
||||
|
|
@ -17946,6 +18109,32 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"sambanova/DeepSeek-V3.1": {
|
||||
"max_tokens": 32768,
|
||||
"max_input_tokens": 32768,
|
||||
"max_output_tokens": 32768,
|
||||
"input_cost_per_token": 3e-06,
|
||||
"output_cost_per_token": 4.5e-06,
|
||||
"litellm_provider": "sambanova",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_reasoning": true,
|
||||
"source": "https://cloud.sambanova.ai/plans/pricing"
|
||||
},
|
||||
"sambanova/gpt-oss-120b": {
|
||||
"max_tokens": 131072,
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 131072,
|
||||
"input_cost_per_token": 3e-06,
|
||||
"output_cost_per_token": 4.5e-06,
|
||||
"litellm_provider": "sambanova",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_reasoning": true,
|
||||
"source": "https://cloud.sambanova.ai/plans/pricing"
|
||||
},
|
||||
"sample_spec": {
|
||||
"code_interpreter_cost_per_session": 0.0,
|
||||
"computer_use_input_cost_per_1k_tokens": 0.0,
|
||||
|
|
@ -20467,6 +20656,24 @@
|
|||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"vertex_ai/deepseek-ai/deepseek-v3.1-maas": {
|
||||
"input_cost_per_token": 1.35e-06,
|
||||
"litellm_provider": "vertex_ai-deepseek_models",
|
||||
"max_input_tokens": 163840,
|
||||
"max_output_tokens": 32768,
|
||||
"max_tokens": 163840,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5.4e-06,
|
||||
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models",
|
||||
"supported_regions": [
|
||||
"us-west2"
|
||||
],
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"vertex_ai/deepseek-ai/deepseek-r1-0528-maas": {
|
||||
"input_cost_per_token": 1.35e-06,
|
||||
"litellm_provider": "vertex_ai-deepseek_models",
|
||||
|
|
@ -20880,6 +21087,30 @@
|
|||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "vertex_ai-qwen_models",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.2e-06,
|
||||
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"vertex_ai/qwen/qwen3-next-80b-a3b-thinking-maas": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "vertex_ai-qwen_models",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.2e-06,
|
||||
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"vertex_ai/veo-2.0-generate-001": {
|
||||
"litellm_provider": "vertex_ai-video-models",
|
||||
"max_input_tokens": 1024,
|
||||
|
|
|
|||
10
poetry.lock
generated
10
poetry.lock
generated
|
|
@ -2730,9 +2730,10 @@ files = [
|
|||
name = "madoka"
|
||||
version = "0.7.1"
|
||||
description = "Memory-efficient CountMin Sketch key-value structure (based on Madoka C++ library)"
|
||||
optional = false
|
||||
optional = true
|
||||
python-versions = "*"
|
||||
groups = ["main"]
|
||||
markers = "extra == \"proxy\""
|
||||
files = [
|
||||
{file = "madoka-0.7.1-cp37-cp37m-macosx_10_14_x86_64.whl", hash = "sha256:7521eee9ace30b376bb54fdcb2cb42bf6b7a0346b0d0b612f25f3299aa4a95af"},
|
||||
{file = "madoka-0.7.1.tar.gz", hash = "sha256:e258baa84fc0a3764365993b8bf5e1b065383a6ca8c9f862fb3e3e709843fae7"},
|
||||
|
|
@ -4044,9 +4045,10 @@ xlsxwriter = ["xlsxwriter"]
|
|||
name = "pondpond"
|
||||
version = "1.4.1"
|
||||
description = "Pond is a high performance object-pooling library for Python."
|
||||
optional = false
|
||||
optional = true
|
||||
python-versions = ">=3.8"
|
||||
groups = ["main"]
|
||||
markers = "extra == \"proxy\""
|
||||
files = [
|
||||
{file = "pondpond-1.4.1-py3-none-any.whl", hash = "sha256:641028ead4e8018ca6de1220c660ddd6d6fbf62a60e72f410655dd0451d82880"},
|
||||
{file = "pondpond-1.4.1.tar.gz", hash = "sha256:8afa34b869d1434d21dd2ec12644abc3b1733fcda8fcf355300338a13a79bb7b"},
|
||||
|
|
@ -6747,11 +6749,11 @@ type = ["pytest-mypy"]
|
|||
caching = ["diskcache"]
|
||||
extra-proxy = ["azure-identity", "azure-keyvault-secrets", "google-cloud-iam", "google-cloud-kms", "prisma", "redisvl", "resend"]
|
||||
mlflow = ["mlflow"]
|
||||
proxy = ["PyJWT", "apscheduler", "azure-identity", "azure-storage-blob", "backoff", "boto3", "cryptography", "fastapi", "fastapi-sso", "fastuuid", "gunicorn", "litellm-enterprise", "litellm-proxy-extras", "mcp", "orjson", "polars", "pynacl", "python-multipart", "pyyaml", "rich", "rq", "uvicorn", "uvloop", "websockets"]
|
||||
proxy = ["PyJWT", "apscheduler", "azure-identity", "azure-storage-blob", "backoff", "boto3", "cryptography", "fastapi", "fastapi-sso", "fastuuid", "gunicorn", "litellm-enterprise", "litellm-proxy-extras", "mcp", "orjson", "polars", "pondpond", "pynacl", "python-multipart", "pyyaml", "rich", "rq", "uvicorn", "uvloop", "websockets"]
|
||||
semantic-router = ["semantic-router"]
|
||||
utils = ["numpydoc"]
|
||||
|
||||
[metadata]
|
||||
lock-version = "2.1"
|
||||
python-versions = ">=3.8.1,<4.0, !=3.9.7"
|
||||
content-hash = "75004c6a23b70be86622fa417fd0d62fa3843e6e61c8dff8507ae5c967b7205d"
|
||||
content-hash = "16fdc1044b4bb316803cbf1825bc970895526949a8bef93eab741777b7210ca8"
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ jinja2 = "^3.1.2"
|
|||
aiohttp = ">=3.10"
|
||||
pydantic = "^2.5.0"
|
||||
jsonschema = "^4.22.0"
|
||||
pondpond = "^1.4.1"
|
||||
pondpond = {version = "^1.4.1", optional = true}
|
||||
numpydoc = {version = "*", optional = true} # used in utils.py
|
||||
fastuuid = {version = ">=0.12.0", optional = true}
|
||||
|
||||
|
|
@ -94,6 +94,7 @@ proxy = [
|
|||
"rich",
|
||||
"polars",
|
||||
"fastuuid",
|
||||
"pondpond",
|
||||
]
|
||||
|
||||
extra_proxy = [
|
||||
|
|
|
|||
|
|
@ -11,7 +11,8 @@ from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import LakeraAIGuardr
|
|||
from litellm.types.guardrails import PiiEntityType, PiiAction
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.exceptions import BlockedPiiEntityError
|
||||
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
|
||||
from fastapi import HTTPException
|
||||
from litellm.types.utils import CallTypes as LitellmCallTypes
|
||||
|
||||
|
||||
|
|
@ -54,3 +55,179 @@ async def test_lakera_pre_call_hook_for_pii_masking():
|
|||
assert "4111-1111-1111-1111" not in user_message
|
||||
assert "test@example.com" not in user_message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lakera_blocks_non_pii_violations():
|
||||
"""Test that Lakera guardrail blocks requests with non-PII violations like hate speech, violence, etc."""
|
||||
|
||||
lakera_guardrail = LakeraAIGuardrail(
|
||||
api_key="test_key",
|
||||
)
|
||||
|
||||
# Mock the call_v2_guard method to return a response similar to the user's example
|
||||
mock_response = {
|
||||
'payload': [],
|
||||
'flagged': True,
|
||||
'dev_info': {'git_revision': 'f0bc093a', 'git_timestamp': '2025-09-23T15:28:06+00:00', 'model_version': 'lakera-guard-1', 'version': '2.0.281'},
|
||||
'metadata': {'request_uuid': 'b7cd4c8a-28aa-4285-a245-2befee514dbf'},
|
||||
'breakdown': [
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/crime', 'detected': True, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/hate', 'detected': True, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/violence', 'detected': True, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-prompt-attack', 'detector_type': 'prompt_attack', 'detected': True, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/email', 'detected': False, 'message_id': 0},
|
||||
]
|
||||
}
|
||||
|
||||
with patch.object(lakera_guardrail, 'call_v2_guard', new_callable=AsyncMock) as mock_call:
|
||||
mock_call.return_value = (mock_response, {})
|
||||
|
||||
# Create a sample request that would trigger violations
|
||||
data = {
|
||||
"messages": [
|
||||
{"role": "user", "content": "Some harmful content that triggers violations"}
|
||||
],
|
||||
"model": "gpt-3.5-turbo",
|
||||
"metadata": {}
|
||||
}
|
||||
|
||||
# Mock objects needed for the pre-call hook
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
|
||||
cache = DualCache()
|
||||
|
||||
# The guardrail should raise an HTTPException for non-PII violations
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await lakera_guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=cache,
|
||||
data=data,
|
||||
call_type="completion"
|
||||
)
|
||||
|
||||
# Verify the exception details include the Lakera response
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "Violated guardrail policy" in str(exc_info.value.detail)
|
||||
assert "lakera_guardrail_response" in exc_info.value.detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lakera_only_pii_violations_are_masked():
|
||||
"""Test that Lakera guardrail only masks PII violations and doesn't block the request."""
|
||||
|
||||
lakera_guardrail = LakeraAIGuardrail(
|
||||
api_key="test_key",
|
||||
)
|
||||
|
||||
# Mock response with only PII violations
|
||||
mock_response = {
|
||||
'payload': [
|
||||
{'detector_type': 'pii/email', 'start': 10, 'end': 25, 'message_id': 0}
|
||||
],
|
||||
'flagged': True,
|
||||
'breakdown': [
|
||||
{'project_id': 'project-9770817088', 'detector_type': 'pii/email', 'detected': True, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'detector_type': 'moderated_content/hate', 'detected': False, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'detector_type': 'prompt_attack', 'detected': False, 'message_id': 0},
|
||||
]
|
||||
}
|
||||
|
||||
with patch.object(lakera_guardrail, 'call_v2_guard', new_callable=AsyncMock) as mock_call:
|
||||
mock_call.return_value = (mock_response, {})
|
||||
|
||||
data = {
|
||||
"messages": [
|
||||
{"role": "user", "content": "My email test@example.com here"}
|
||||
],
|
||||
"model": "gpt-3.5-turbo",
|
||||
"metadata": {}
|
||||
}
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
|
||||
cache = DualCache()
|
||||
|
||||
# Should not raise an exception, just mask the PII
|
||||
result = await lakera_guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=cache,
|
||||
data=data,
|
||||
call_type="completion"
|
||||
)
|
||||
|
||||
# Verify the request was not blocked
|
||||
assert result is not None
|
||||
assert "messages" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lakera_blocks_flagged_content_with_user_scenario():
|
||||
"""
|
||||
Test the exact user scenario where Lakera flagged content but request went through.
|
||||
This should now be blocked with the fix to check breakdown field instead of payload.
|
||||
"""
|
||||
|
||||
lakera_guardrail = LakeraAIGuardrail(
|
||||
api_key="test_key",
|
||||
)
|
||||
|
||||
# Mock response matching the exact user scenario
|
||||
mock_response = {
|
||||
'payload': [], # Empty payload like in user's case
|
||||
'flagged': True,
|
||||
'dev_info': {'git_revision': 'f0bc093a', 'git_timestamp': '2025-09-23T15:28:06+00:00', 'model_version': 'lakera-guard-1', 'version': '2.0.281'},
|
||||
'metadata': {'request_uuid': 'b7cd4c8a-28aa-4285-a245-2befee514dbf'},
|
||||
'breakdown': [
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/crime', 'detected': True, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/hate', 'detected': True, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/profanity', 'detected': False, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/sexual', 'detected': False, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/violence', 'detected': True, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-moderated-content', 'detector_type': 'moderated_content/weapons', 'detected': True, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/address', 'detected': False, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/credit_card', 'detected': False, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/email', 'detected': False, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/iban_code', 'detected': False, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/ip_address', 'detected': False, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/name', 'detected': False, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/phone_number', 'detected': False, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-pii', 'detector_type': 'pii/us_social_security_number', 'detected': False, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-prompt-attack', 'detector_type': 'prompt_attack', 'detected': True, 'message_id': 0},
|
||||
{'project_id': 'project-9770817088', 'policy_id': 'policy-lakera-default', 'detector_id': 'detector-lakera-default-unknown-links', 'detector_type': 'unknown_links', 'detected': False, 'message_id': 0}
|
||||
]
|
||||
}
|
||||
|
||||
with patch.object(lakera_guardrail, 'call_v2_guard', new_callable=AsyncMock) as mock_call:
|
||||
mock_call.return_value = (mock_response, {})
|
||||
|
||||
# Create a sample request that would trigger violations
|
||||
data = {
|
||||
"messages": [
|
||||
{"role": "user", "content": "Some harmful content that should be blocked"}
|
||||
],
|
||||
"model": "gpt-3.5-turbo",
|
||||
"metadata": {}
|
||||
}
|
||||
|
||||
# Mock objects needed for the pre-call hook
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
|
||||
cache = DualCache()
|
||||
|
||||
# With the fix, this should now raise an HTTPException instead of letting the request through
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await lakera_guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=cache,
|
||||
data=data,
|
||||
call_type="completion"
|
||||
)
|
||||
|
||||
# Verify the exception details
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "Violated guardrail policy" in str(exc_info.value.detail)
|
||||
assert "lakera_guardrail_response" in exc_info.value.detail
|
||||
|
||||
# Verify the full response is included in the exception
|
||||
lakera_response = exc_info.value.detail["lakera_guardrail_response"]
|
||||
assert lakera_response["flagged"] is True
|
||||
assert lakera_response["metadata"]["request_uuid"] == "b7cd4c8a-28aa-4285-a245-2befee514dbf"
|
||||
assert len(lakera_response["breakdown"]) == 16 # All the breakdown items from the user's scenario
|
||||
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ Simplified tests for object pooling utilities in litellm.
|
|||
"""
|
||||
|
||||
import pytest
|
||||
pytest.importorskip("pond")
|
||||
|
||||
from litellm.litellm_core_utils.object_pooling import (
|
||||
get_object_pool,
|
||||
|
|
@ -103,4 +104,4 @@ class TestObjectPooling:
|
|||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
pytest.main([__file__])
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ import litellm
|
|||
from litellm import (
|
||||
ModelResponse,
|
||||
RateLimitError,
|
||||
ServiceUnavailableError,
|
||||
Timeout,
|
||||
completion,
|
||||
completion_cost,
|
||||
|
|
@ -2020,10 +2021,16 @@ def test_bedrock_context_window_error():
|
|||
|
||||
def test_bedrock_converse_route():
|
||||
litellm.set_verbose = True
|
||||
litellm.completion(
|
||||
model="bedrock/converse/us.amazon.nova-pro-v1:0",
|
||||
messages=[{"role": "user", "content": "Hello, world!"}],
|
||||
)
|
||||
try:
|
||||
litellm.completion(
|
||||
model="bedrock/converse/us.amazon.nova-pro-v1:0",
|
||||
messages=[{"role": "user", "content": "Hello, world!"}],
|
||||
)
|
||||
except ServiceUnavailableError as e:
|
||||
if "Too many requests" in str(e):
|
||||
pytest.skip("Skipping test due to AWS Bedrock rate limiting")
|
||||
else:
|
||||
raise
|
||||
|
||||
|
||||
def test_bedrock_mapped_converse_models():
|
||||
|
|
|
|||
|
|
@ -68,7 +68,14 @@
|
|||
"output": 20,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 5.4999999999999995e-05
|
||||
}
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 20
|
||||
},
|
||||
"traceId": "litellm-test-6a51ae70-a4e7-499e-afcd-dce2a3b31850"
|
||||
},
|
||||
"timestamp": "2025-01-16T19:28:55.125258Z"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -63,7 +63,13 @@
|
|||
"output": 10,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 0.00018
|
||||
}
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 10
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-05-26T21:13:16.797156Z"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -109,6 +109,12 @@
|
|||
"output": 20,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 5.4999999999999995e-05
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 20
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-01-22T17:27:51.703046Z"
|
||||
|
|
|
|||
|
|
@ -87,6 +87,12 @@
|
|||
"output": 20,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 5.4999999999999995e-05
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 20
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-01-22T17:19:11.235541Z"
|
||||
|
|
|
|||
|
|
@ -58,6 +58,12 @@
|
|||
"output": 10,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 3.5e-05
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 10
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-02-07T00:23:27.670175Z"
|
||||
|
|
|
|||
|
|
@ -73,6 +73,12 @@
|
|||
"output": 10,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 3.5e-05
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 10
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-05-24T17:01:19.408586Z"
|
||||
|
|
|
|||
|
|
@ -77,6 +77,12 @@
|
|||
"output": 20,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 5.4999999999999995e-05
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 20
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-01-22T15:31:28.964179Z"
|
||||
|
|
|
|||
|
|
@ -77,6 +77,12 @@
|
|||
"output": 20,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 5.4999999999999995e-05
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 20
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-01-22T16:38:26.017252Z"
|
||||
|
|
|
|||
|
|
@ -61,6 +61,12 @@
|
|||
"output": 10,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 7.5e-06
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 10
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-05-26T21:15:40.610953Z"
|
||||
|
|
|
|||
|
|
@ -84,6 +84,12 @@
|
|||
"output": 20,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 5.4999999999999995e-05
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 20
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-01-22T17:59:39.368310Z"
|
||||
|
|
|
|||
|
|
@ -76,6 +76,12 @@
|
|||
"output": 20,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 5.4999999999999995e-05
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 20
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-01-22T18:06:50.959850Z"
|
||||
|
|
|
|||
|
|
@ -70,6 +70,12 @@
|
|||
"output": 20,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 5.4999999999999995e-05
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 20
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-01-22T17:59:32.889548Z"
|
||||
|
|
|
|||
|
|
@ -70,6 +70,12 @@
|
|||
"output": 20,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 5.4999999999999995e-05
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 20
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-01-22T17:59:36.162997Z"
|
||||
|
|
|
|||
|
|
@ -70,6 +70,12 @@
|
|||
"output": 20,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 5.4999999999999995e-05
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 20
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-01-22T17:59:32.889548Z"
|
||||
|
|
|
|||
|
|
@ -76,6 +76,12 @@
|
|||
"output": 20,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 5.4999999999999995e-05
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 20
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-01-22T17:55:28.855732Z"
|
||||
|
|
|
|||
|
|
@ -76,6 +76,12 @@
|
|||
"output": 20,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 5.4999999999999995e-05
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 20
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-01-22T17:53:53.754511Z"
|
||||
|
|
|
|||
|
|
@ -80,6 +80,12 @@
|
|||
"output": 20,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 5.4999999999999995e-05
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 20
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-01-22T17:56:35.478171Z"
|
||||
|
|
|
|||
|
|
@ -84,6 +84,12 @@
|
|||
"output": 20,
|
||||
"unit": "TOKENS",
|
||||
"totalCost": 5.4999999999999995e-05
|
||||
},
|
||||
"usageDetails": {
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0,
|
||||
"input": 10,
|
||||
"output": 20
|
||||
}
|
||||
},
|
||||
"timestamp": "2025-01-22T17:56:38.787196Z"
|
||||
|
|
|
|||
|
|
@ -43,8 +43,9 @@ async def test_anthropic_basic_completion_with_headers():
|
|||
json.dumps(response_json, indent=4, default=str),
|
||||
)
|
||||
reported_usage = response_json.get("usage", None)
|
||||
anthropic_api_input_tokens = reported_usage.get("input_tokens", None)
|
||||
anthropic_api_output_tokens = reported_usage.get("output_tokens", None)
|
||||
# fix null checks for reported_usage
|
||||
anthropic_api_input_tokens = reported_usage.get("input_tokens", None) if reported_usage else None
|
||||
anthropic_api_output_tokens = reported_usage.get("output_tokens", None) if reported_usage else None
|
||||
litellm_call_id = response_headers.get("x-litellm-call-id")
|
||||
|
||||
print(f"LiteLLM Call ID: {litellm_call_id}")
|
||||
|
|
|
|||
|
|
@ -88,3 +88,101 @@ def test_in_memory_cache_ttl_allow_override():
|
|||
new_ttl_time = in_memory_cache.ttl_dict["new-fake-key"]
|
||||
assert new_ttl_time is not None
|
||||
assert new_ttl_time != initial_ttl_time
|
||||
|
||||
|
||||
def test_in_memory_cache_max_size_with_ttl():
|
||||
"""
|
||||
Test that max_size_in_memory is respected even when all items have long TTLs.
|
||||
This tests the fix for the unbounded growth issue.
|
||||
"""
|
||||
in_memory_cache = InMemoryCache(max_size_in_memory=3)
|
||||
long_ttl = 86400 # 1 day
|
||||
|
||||
# Fill the cache to max capacity
|
||||
for i in range(3):
|
||||
in_memory_cache.set_cache(key=f"key_{i}", value=f"value_{i}", ttl=long_ttl)
|
||||
time.sleep(0.01) # Small delay to ensure different timestamps
|
||||
|
||||
assert len(in_memory_cache.cache_dict) == 3
|
||||
assert len(in_memory_cache.ttl_dict) == 3
|
||||
|
||||
# Add another item - should evict the earliest item
|
||||
in_memory_cache.set_cache(key="key_3", value="value_3", ttl=long_ttl)
|
||||
|
||||
# Cache should still be at max size, not larger
|
||||
assert len(in_memory_cache.cache_dict) == 3
|
||||
assert len(in_memory_cache.ttl_dict) == 3
|
||||
|
||||
# key_0 should have been evicted (it was added first)
|
||||
assert "key_0" not in in_memory_cache.cache_dict
|
||||
assert "key_0" not in in_memory_cache.ttl_dict
|
||||
|
||||
# Other keys should still be present
|
||||
assert "key_1" in in_memory_cache.cache_dict
|
||||
assert "key_2" in in_memory_cache.cache_dict
|
||||
assert "key_3" in in_memory_cache.cache_dict
|
||||
|
||||
|
||||
def test_in_memory_cache_expired_items_evicted_first():
|
||||
"""
|
||||
Test that expired items are evicted before non-expired items when cache is full.
|
||||
"""
|
||||
in_memory_cache = InMemoryCache(max_size_in_memory=3)
|
||||
|
||||
# Add items with short TTL that will expire
|
||||
in_memory_cache.set_cache(key="expired_1", value="value_1", ttl=1)
|
||||
in_memory_cache.set_cache(key="expired_2", value="value_2", ttl=1)
|
||||
|
||||
# Add item with long TTL
|
||||
in_memory_cache.set_cache(key="long_lived", value="value_long", ttl=86400)
|
||||
|
||||
assert len(in_memory_cache.cache_dict) == 3
|
||||
|
||||
# Wait for short TTL items to expire
|
||||
time.sleep(2)
|
||||
|
||||
# Add new item - should evict expired items first, not the long-lived one
|
||||
in_memory_cache.set_cache(key="new_item", value="new_value", ttl=86400)
|
||||
|
||||
# Long-lived item should still be present
|
||||
assert "long_lived" in in_memory_cache.cache_dict
|
||||
assert "new_item" in in_memory_cache.cache_dict
|
||||
|
||||
# Expired items should be gone
|
||||
assert "expired_1" not in in_memory_cache.cache_dict
|
||||
assert "expired_2" not in in_memory_cache.cache_dict
|
||||
assert "expired_1" not in in_memory_cache.ttl_dict
|
||||
assert "expired_2" not in in_memory_cache.ttl_dict
|
||||
|
||||
|
||||
def test_in_memory_cache_eviction_order():
|
||||
"""
|
||||
Test that when non-expired items need to be evicted, those with earliest expiration times are evicted first.
|
||||
"""
|
||||
in_memory_cache = InMemoryCache(max_size_in_memory=2)
|
||||
|
||||
# Add items with different TTLs
|
||||
now = time.time()
|
||||
in_memory_cache.set_cache(key="early_expire", value="value_1", ttl=100) # expires in 100 seconds
|
||||
time.sleep(0.01)
|
||||
in_memory_cache.set_cache(key="late_expire", value="value_2", ttl=200) # expires in 200 seconds
|
||||
|
||||
# Verify TTL order
|
||||
early_ttl = in_memory_cache.ttl_dict["early_expire"]
|
||||
late_ttl = in_memory_cache.ttl_dict["late_expire"]
|
||||
assert early_ttl < late_ttl, "early_expire should have earlier expiration time"
|
||||
|
||||
assert len(in_memory_cache.cache_dict) == 2
|
||||
|
||||
# Add third item - should evict the one with earliest expiration time
|
||||
in_memory_cache.set_cache(key="new_item", value="value_3", ttl=300)
|
||||
|
||||
assert len(in_memory_cache.cache_dict) == 2
|
||||
|
||||
# Item with earliest expiration should be evicted
|
||||
assert "early_expire" not in in_memory_cache.cache_dict
|
||||
assert "early_expire" not in in_memory_cache.ttl_dict
|
||||
|
||||
# Items with later expiration should remain
|
||||
assert "late_expire" in in_memory_cache.cache_dict
|
||||
assert "new_item" in in_memory_cache.cache_dict
|
||||
|
|
|
|||
|
|
@ -1,20 +1,222 @@
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
from typing import Optional
|
||||
|
||||
# Adds the grandparent directory to sys.path to allow importing project modules
|
||||
sys.path.insert(0, os.path.abspath("../.."))
|
||||
|
||||
import unittest
|
||||
import asyncio
|
||||
from unittest.mock import patch
|
||||
|
||||
from unittest.mock import patch, MagicMock
|
||||
from typing import Optional
|
||||
import sys
|
||||
import os
|
||||
import datetime
|
||||
import json
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.langfuse import langfuse as langfuse_module
|
||||
from litellm.integrations.langfuse.langfuse import LangFuseLogger
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../.."))
|
||||
from litellm.integrations.langfuse.langfuse import LangFuseLogger
|
||||
# Import LangfuseUsageDetails directly from the module where it's defined
|
||||
from litellm.types.integrations.langfuse import *
|
||||
|
||||
class TestLangfuseUsageDetails(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
# Set up environment variables for testing
|
||||
self.env_patcher = patch.dict('os.environ', {
|
||||
'LANGFUSE_SECRET_KEY': 'test-secret-key',
|
||||
'LANGFUSE_PUBLIC_KEY': 'test-public-key',
|
||||
'LANGFUSE_HOST': 'https://test.langfuse.com'
|
||||
})
|
||||
self.env_patcher.start()
|
||||
|
||||
# Create mock objects
|
||||
self.mock_langfuse_client = MagicMock()
|
||||
self.mock_langfuse_trace = MagicMock()
|
||||
self.mock_langfuse_generation = MagicMock()
|
||||
|
||||
# Setup the trace and generation chain
|
||||
self.mock_langfuse_trace.generation.return_value = self.mock_langfuse_generation
|
||||
self.mock_langfuse_client.trace.return_value = self.mock_langfuse_trace
|
||||
|
||||
# Mock the langfuse module that's imported locally in methods
|
||||
self.langfuse_module_patcher = patch.dict('sys.modules', {'langfuse': MagicMock()})
|
||||
self.mock_langfuse_module = self.langfuse_module_patcher.start()
|
||||
|
||||
# Create a mock for the langfuse module with version
|
||||
self.mock_langfuse = MagicMock()
|
||||
self.mock_langfuse.version = MagicMock()
|
||||
self.mock_langfuse.version.__version__ = "3.0.0" # Set a version that supports all features
|
||||
|
||||
# Mock the Langfuse class
|
||||
self.mock_langfuse_class = MagicMock()
|
||||
self.mock_langfuse_class.return_value = self.mock_langfuse_client
|
||||
|
||||
# Set up the sys.modules['langfuse'] mock
|
||||
sys.modules['langfuse'] = self.mock_langfuse
|
||||
sys.modules['langfuse'].Langfuse = self.mock_langfuse_class
|
||||
|
||||
# Mock the Langfuse client
|
||||
self.mock_langfuse_client = MagicMock()
|
||||
self.mock_langfuse_trace = MagicMock()
|
||||
self.mock_langfuse_generation = MagicMock()
|
||||
|
||||
# Setup the trace and generation chain
|
||||
self.mock_langfuse_trace.generation.return_value = self.mock_langfuse_generation
|
||||
self.mock_langfuse_client.trace.return_value = self.mock_langfuse_trace
|
||||
|
||||
# Mock the Langfuse class
|
||||
self.mock_langfuse_class = MagicMock()
|
||||
self.mock_langfuse_class.return_value = self.mock_langfuse_client
|
||||
self.mock_langfuse.Langfuse = self.mock_langfuse_class
|
||||
|
||||
# Create the logger
|
||||
self.logger = LangFuseLogger()
|
||||
|
||||
# Add the log_event_on_langfuse method to the instance
|
||||
def log_event_on_langfuse(self, kwargs, response_obj, start_time=None, end_time=None, user_id=None, level="DEFAULT", status_message=None):
|
||||
# This implementation calls _log_langfuse_v2 directly
|
||||
return self._log_langfuse_v2(
|
||||
user_id=user_id,
|
||||
metadata=kwargs.get("litellm_params", {}).get("metadata", {}),
|
||||
litellm_params=kwargs.get("litellm_params", {}),
|
||||
output=None,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
kwargs=kwargs,
|
||||
optional_params=kwargs.get("optional_params", {}),
|
||||
input=None,
|
||||
response_obj=response_obj,
|
||||
level=level,
|
||||
litellm_call_id=kwargs.get("litellm_call_id", None),
|
||||
print_verbose=True # Add the missing parameter
|
||||
)
|
||||
|
||||
# Bind the method to the instance
|
||||
import types
|
||||
self.logger.log_event_on_langfuse = types.MethodType(log_event_on_langfuse, self.logger)
|
||||
|
||||
# Make sure _is_langfuse_v2 returns True
|
||||
def mock_is_langfuse_v2(self):
|
||||
return True
|
||||
|
||||
self.logger._is_langfuse_v2 = types.MethodType(mock_is_langfuse_v2, self.logger)
|
||||
|
||||
def tearDown(self):
|
||||
self.env_patcher.stop()
|
||||
self.langfuse_module_patcher.stop()
|
||||
|
||||
def test_langfuse_usage_details_type(self):
|
||||
"""Test that LangfuseUsageDetails TypedDict is properly defined with the correct fields"""
|
||||
# Create an instance of LangfuseUsageDetails
|
||||
usage_details: LangfuseUsageDetails = {
|
||||
"input": 10,
|
||||
"output": 20,
|
||||
"cache_creation_input_tokens": 5,
|
||||
"cache_read_input_tokens": 3
|
||||
}
|
||||
|
||||
# Verify all fields are present
|
||||
self.assertEqual(usage_details["input"], 10)
|
||||
self.assertEqual(usage_details["output"], 20)
|
||||
self.assertEqual(usage_details["cache_creation_input_tokens"], 5)
|
||||
self.assertEqual(usage_details["cache_read_input_tokens"], 3)
|
||||
|
||||
# Test with all fields (all fields are required in TypedDict by default)
|
||||
minimal_usage_details: LangfuseUsageDetails = {
|
||||
"input": 10,
|
||||
"output": 20,
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0
|
||||
}
|
||||
|
||||
self.assertEqual(minimal_usage_details["input"], 10)
|
||||
self.assertEqual(minimal_usage_details["output"], 20)
|
||||
|
||||
def test_log_langfuse_v2_usage_details(self):
|
||||
"""Test that usage_details in _log_langfuse_v2 is correctly typed and assigned"""
|
||||
# Create a mock response object with usage information
|
||||
response_obj = MagicMock()
|
||||
response_obj.usage = MagicMock()
|
||||
response_obj.usage.prompt_tokens = 15
|
||||
response_obj.usage.completion_tokens = 25
|
||||
|
||||
# Add the cache token attributes using get method
|
||||
def mock_get(key, default=None):
|
||||
if key == 'cache_creation_input_tokens':
|
||||
return 7
|
||||
elif key == 'cache_read_input_tokens':
|
||||
return 4
|
||||
return default
|
||||
|
||||
response_obj.usage.get = mock_get
|
||||
|
||||
# Create kwargs for the log_event method
|
||||
kwargs = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"litellm_params": {"metadata": {}}
|
||||
}
|
||||
|
||||
# Create start and end times
|
||||
start_time = datetime.datetime.now()
|
||||
end_time = start_time + datetime.timedelta(seconds=1)
|
||||
|
||||
# Call the log_event method
|
||||
with patch.object(self.logger, '_log_langfuse_v2') as mock_log_langfuse_v2:
|
||||
self.logger.log_event_on_langfuse(
|
||||
kwargs=kwargs,
|
||||
response_obj=response_obj,
|
||||
start_time=start_time,
|
||||
end_time=end_time
|
||||
)
|
||||
|
||||
# Check if _log_langfuse_v2 was called
|
||||
mock_log_langfuse_v2.assert_called_once()
|
||||
|
||||
# Get the arguments passed to _log_langfuse_v2
|
||||
call_args = mock_log_langfuse_v2.call_args[1]
|
||||
|
||||
# Verify response_obj was passed correctly
|
||||
self.assertEqual(call_args["response_obj"], response_obj)
|
||||
|
||||
def test_langfuse_usage_details_optional_fields(self):
|
||||
"""Test that LangfuseUsageDetails fields are properly defined as Optional"""
|
||||
# Create an instance with None values for optional fields
|
||||
usage_details: LangfuseUsageDetails = {
|
||||
"input": 10,
|
||||
"output": 20,
|
||||
"cache_creation_input_tokens": None,
|
||||
"cache_read_input_tokens": None
|
||||
}
|
||||
|
||||
# Verify fields can be None
|
||||
self.assertEqual(usage_details["input"], 10)
|
||||
self.assertEqual(usage_details["output"], 20)
|
||||
self.assertIsNone(usage_details["cache_creation_input_tokens"])
|
||||
self.assertIsNone(usage_details["cache_read_input_tokens"])
|
||||
|
||||
def test_langfuse_usage_details_structure(self):
|
||||
"""Test that LangfuseUsageDetails has the correct structure as defined in the commit"""
|
||||
# This test directly verifies the structure of the TypedDict
|
||||
# without relying on the LangFuseLogger class
|
||||
|
||||
# Create a dictionary that matches the LangfuseUsageDetails structure
|
||||
usage_details = {
|
||||
"input": 15,
|
||||
"output": 25,
|
||||
"cache_creation_input_tokens": 7,
|
||||
"cache_read_input_tokens": 4
|
||||
}
|
||||
|
||||
# Verify the structure matches what we expect
|
||||
self.assertIn("input", usage_details)
|
||||
self.assertIn("output", usage_details)
|
||||
self.assertIn("cache_creation_input_tokens", usage_details)
|
||||
self.assertIn("cache_read_input_tokens", usage_details)
|
||||
|
||||
# Verify the values
|
||||
self.assertEqual(usage_details["input"], 15)
|
||||
self.assertEqual(usage_details["output"], 25)
|
||||
self.assertEqual(usage_details["cache_creation_input_tokens"], 7)
|
||||
self.assertEqual(usage_details["cache_read_input_tokens"], 4)
|
||||
|
||||
def test_max_langfuse_clients_limit():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue