mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
Merge remote-tracking branch 'upstream/litellm_internal_staging' into add-gemma-4-26b-vertex-ai
# Conflicts: # litellm/llms/vertex_ai/common_utils.py
This commit is contained in:
commit
c91e2ded8e
86 changed files with 6708 additions and 1022 deletions
|
|
@ -1,75 +0,0 @@
|
|||
name: Check Lazy OpenAPI Snapshot
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
checks: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
verify:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cache/uv
|
||||
.venv
|
||||
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Install dependencies
|
||||
run: uv sync --frozen --all-groups --all-extras
|
||||
|
||||
- name: Regenerate snapshot to /tmp
|
||||
id: regen
|
||||
run: |
|
||||
cp litellm/proxy/_lazy_openapi_snapshot.json /tmp/snapshot.committed.json
|
||||
uv run --no-sync python -m litellm.proxy._lazy_openapi_snapshot
|
||||
mv litellm/proxy/_lazy_openapi_snapshot.json /tmp/snapshot.fresh.json
|
||||
mv /tmp/snapshot.committed.json litellm/proxy/_lazy_openapi_snapshot.json
|
||||
|
||||
- name: Compare
|
||||
id: diff
|
||||
continue-on-error: true
|
||||
run: |
|
||||
diff -q /tmp/snapshot.fresh.json litellm/proxy/_lazy_openapi_snapshot.json
|
||||
|
||||
- name: Mark neutral if drift
|
||||
if: steps.diff.outcome == 'failure'
|
||||
uses: LouisBrunner/checks-action@6b626ffbad7cc56fd58627f774b9067e6118af23 # v2.0.0
|
||||
with:
|
||||
token: ${{ secrets.GITHUB_TOKEN }}
|
||||
name: lazy-openapi-snapshot
|
||||
conclusion: neutral
|
||||
output: |
|
||||
{
|
||||
"title": "Lazy openapi snapshot is stale",
|
||||
"summary": "Run `python -m litellm.proxy._lazy_openapi_snapshot` and commit the regenerated `litellm/proxy/_lazy_openapi_snapshot.json`. Not blocking — the snapshot will regenerate at release if not committed."
|
||||
}
|
||||
4
.gitignore
vendored
4
.gitignore
vendored
|
|
@ -90,7 +90,6 @@ test.py
|
|||
litellm_config.yaml
|
||||
!.github/observatory/litellm_config.yaml
|
||||
.cursor
|
||||
.vscode/launch.json
|
||||
litellm/proxy/to_delete_loadtest_work/*
|
||||
update_model_cost_map.py
|
||||
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
|
||||
|
|
@ -100,4 +99,5 @@ STABILIZATION_TODO.md
|
|||
**/test-results
|
||||
**/playwright-report
|
||||
**/*.storageState.json
|
||||
**/coverage
|
||||
**/coverage
|
||||
test-config
|
||||
|
|
@ -68,7 +68,7 @@ Managing LLM calls across providers gets complicated fast — different SDKs, au
|
|||
<td><img height="60" alt="Stripe" src="https://github.com/user-attachments/assets/f7296d4f-9fbd-460d-9d05-e4df31697c4b" /></td>
|
||||
<td><img height="60" alt="image" src="https://github.com/user-attachments/assets/436fca71-988b-40bb-b5fe-8450c80fdbd0" /></td>
|
||||
<td><img height="60" alt="Google ADK" src="https://github.com/user-attachments/assets/caf270a2-5aee-45c4-8222-41a2070c4f19" /></td>
|
||||
<td><img height="60" alt="Greptile" src="https://github.com/user-attachments/assets/0be4bd8a-7cfa-48d3-9090-f415fe948280" /></td>
|
||||
<td><img height="60" alt="Greptile" src="https://github.com/user-attachments/assets/3db0ae72-0843-4005-a56d-bba1dde2193d" /></td>
|
||||
<td><img height="60" alt="OpenHands" src="https://github.com/user-attachments/assets/a6150c4c-149e-4cae-888b-8b92be6e003f" /></td>
|
||||
<td><h2>Netflix</h2></td>
|
||||
<td><img height="60" alt="OpenAI Agents SDK" src="https://github.com/user-attachments/assets/c02f7be0-8c2e-4d27-aea7-7c024bfaebc0" /></td>
|
||||
|
|
|
|||
|
|
@ -857,10 +857,16 @@ async def project_info(
|
|||
where={"team_id": project.team_id}
|
||||
)
|
||||
if team:
|
||||
is_team_member = (
|
||||
user_api_key_dict.user_id in team.admins
|
||||
or user_api_key_dict.user_id in team.members
|
||||
)
|
||||
caller_user_id = user_api_key_dict.user_id
|
||||
for m in team.members_with_roles or []:
|
||||
m_user_id = (
|
||||
m.get("user_id")
|
||||
if isinstance(m, dict)
|
||||
else getattr(m, "user_id", None)
|
||||
)
|
||||
if m_user_id == caller_user_id:
|
||||
is_team_member = True
|
||||
break
|
||||
|
||||
if not (is_admin or is_team_member):
|
||||
raise HTTPException(
|
||||
|
|
@ -911,20 +917,20 @@ async def list_projects(
|
|||
include={"litellm_budget_table": True, "object_permission": True}
|
||||
)
|
||||
else:
|
||||
# Get projects for teams the user belongs to
|
||||
user_teams = await prisma_client.db.litellm_teamtable.find_many(
|
||||
where={
|
||||
"OR": [
|
||||
{"members": {"has": user_api_key_dict.user_id}},
|
||||
{"admins": {"has": user_api_key_dict.user_id}},
|
||||
]
|
||||
}
|
||||
# Look up the user's team memberships via the reverse-index on
|
||||
# LiteLLM_UserTable.teams (maintained by team_member_add alongside
|
||||
# members_with_roles). This avoids a full scan of all team rows.
|
||||
user_record = await prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": user_api_key_dict.user_id},
|
||||
)
|
||||
user_team_ids = (
|
||||
user_record.teams
|
||||
if user_record is not None and user_record.teams
|
||||
else []
|
||||
)
|
||||
|
||||
team_ids = [team.team_id for team in user_teams]
|
||||
|
||||
projects = await prisma_client.db.litellm_projecttable.find_many(
|
||||
where={"team_id": {"in": team_ids}},
|
||||
where={"team_id": {"in": user_team_ids}},
|
||||
include={"litellm_budget_table": True, "object_permission": True},
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -432,9 +432,10 @@ class Cache:
|
|||
str: The final hashed cache key with the redis namespace.
|
||||
"""
|
||||
dynamic_cache_control: DynamicCacheControl = kwargs.get("cache", {})
|
||||
metadata = kwargs.get("metadata") or {}
|
||||
namespace = (
|
||||
dynamic_cache_control.get("namespace")
|
||||
or kwargs.get("metadata", {}).get("redis_namespace")
|
||||
or metadata.get("redis_namespace")
|
||||
or self.namespace
|
||||
)
|
||||
if namespace:
|
||||
|
|
|
|||
|
|
@ -87,6 +87,18 @@ class CachingHandlerResponse(BaseModel):
|
|||
in_memory_cache_obj = InMemoryCache()
|
||||
|
||||
|
||||
def _should_defer_streaming_cache_hit_callbacks(*, kwargs: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
When stream=True, do not run success callbacks at cache-hit time.
|
||||
|
||||
Cached chat/text completion replay uses CustomStreamWrapper; cached Responses
|
||||
replay uses CachedResponsesAPIStreamingIterator. Both invoke logging success
|
||||
handlers when the stream finishes; firing them here too would double-count
|
||||
spend and callback records.
|
||||
"""
|
||||
return kwargs.get("stream", False) is True
|
||||
|
||||
|
||||
class LLMCachingHandler:
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -99,6 +111,7 @@ class LLMCachingHandler:
|
|||
self.async_streaming_chunks: List[ModelResponse] = []
|
||||
self.sync_streaming_chunks: List[ModelResponse] = []
|
||||
self.request_kwargs = request_kwargs
|
||||
self.preset_cache_key: Optional[str] = None
|
||||
self.original_function = original_function
|
||||
self.start_time = start_time
|
||||
if litellm.cache is not None and isinstance(litellm.cache.cache, RedisCache):
|
||||
|
|
@ -206,7 +219,7 @@ class LLMCachingHandler:
|
|||
custom_llm_provider=kwargs.get("custom_llm_provider", None),
|
||||
args=args,
|
||||
)
|
||||
if kwargs.get("stream", False) is False:
|
||||
if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs):
|
||||
# LOG SUCCESS
|
||||
self._async_log_cache_hit_on_callbacks(
|
||||
logging_obj=logging_obj,
|
||||
|
|
@ -215,11 +228,12 @@ class LLMCachingHandler:
|
|||
end_time=end_time,
|
||||
cache_hit=cache_hit,
|
||||
)
|
||||
cache_key = litellm.cache.get_cache_key(**kwargs)
|
||||
if (
|
||||
isinstance(cached_result, BaseModel)
|
||||
or isinstance(cached_result, CustomStreamWrapper)
|
||||
) and hasattr(cached_result, "_hidden_params"):
|
||||
cache_key = (
|
||||
self.preset_cache_key
|
||||
or self.request_kwargs.get("cache_key")
|
||||
or litellm.cache.get_cache_key(**self.request_kwargs)
|
||||
)
|
||||
if hasattr(cached_result, "_hidden_params"):
|
||||
cached_result._hidden_params["cache_key"] = cache_key # type: ignore
|
||||
return CachingHandlerResponse(cached_result=cached_result)
|
||||
elif (
|
||||
|
|
@ -265,8 +279,6 @@ class LLMCachingHandler:
|
|||
kwargs: Dict[str, Any],
|
||||
args: Optional[Tuple[Any, ...]] = None,
|
||||
) -> CachingHandlerResponse:
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
cached_result: Optional[Any] = None
|
||||
|
||||
# Check if caching should be performed BEFORE doing expensive kwargs copy
|
||||
|
|
@ -282,6 +294,11 @@ class LLMCachingHandler:
|
|||
args,
|
||||
)
|
||||
)
|
||||
if new_kwargs.get("metadata") is None:
|
||||
new_kwargs.pop("metadata", None)
|
||||
if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs:
|
||||
new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
|
||||
self.request_kwargs = new_kwargs
|
||||
print_verbose("Checking Sync Cache")
|
||||
cached_result = litellm.cache.get_cache(**new_kwargs)
|
||||
if cached_result is not None:
|
||||
|
|
@ -322,17 +339,19 @@ class LLMCachingHandler:
|
|||
is_async=False,
|
||||
)
|
||||
|
||||
logging_obj.handle_sync_success_callbacks_for_async_calls(
|
||||
result=cached_result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
cache_hit=cache_hit,
|
||||
if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs):
|
||||
logging_obj.handle_sync_success_callbacks_for_async_calls(
|
||||
result=cached_result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
cache_hit=cache_hit,
|
||||
)
|
||||
cache_key = (
|
||||
self.preset_cache_key
|
||||
or self.request_kwargs.get("cache_key")
|
||||
or litellm.cache.get_cache_key(**self.request_kwargs)
|
||||
)
|
||||
cache_key = litellm.cache.get_cache_key(**kwargs)
|
||||
if (
|
||||
isinstance(cached_result, BaseModel)
|
||||
or isinstance(cached_result, CustomStreamWrapper)
|
||||
) and hasattr(cached_result, "_hidden_params"):
|
||||
if hasattr(cached_result, "_hidden_params"):
|
||||
cached_result._hidden_params["cache_key"] = cache_key # type: ignore
|
||||
return CachingHandlerResponse(cached_result=cached_result)
|
||||
return CachingHandlerResponse(cached_result=cached_result)
|
||||
|
|
@ -686,6 +705,11 @@ class LLMCachingHandler:
|
|||
args,
|
||||
)
|
||||
)
|
||||
if new_kwargs.get("metadata") is None:
|
||||
new_kwargs.pop("metadata", None)
|
||||
if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs:
|
||||
new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
|
||||
self.request_kwargs = new_kwargs
|
||||
cached_result: Optional[Any] = None
|
||||
if call_type == CallTypes.aembedding.value:
|
||||
if isinstance(new_kwargs["input"], str):
|
||||
|
|
@ -710,14 +734,26 @@ class LLMCachingHandler:
|
|||
if all(result is None for result in cached_result):
|
||||
cached_result = None
|
||||
else:
|
||||
request_kwargs = new_kwargs.copy()
|
||||
request_cache_key = request_kwargs.pop("cache_key", None)
|
||||
if litellm.cache._supports_async() is True:
|
||||
## check if dual cache is supported ##
|
||||
self.preset_cache_key = (
|
||||
request_cache_key or litellm.cache.get_cache_key(**request_kwargs)
|
||||
)
|
||||
cached_result = await litellm.cache.async_get_cache(
|
||||
dynamic_cache_object=self.dual_cache, **new_kwargs
|
||||
dynamic_cache_object=self.dual_cache,
|
||||
cache_key=self.preset_cache_key,
|
||||
**request_kwargs,
|
||||
)
|
||||
else: # fallback for caches that don't support async
|
||||
self.preset_cache_key = (
|
||||
request_cache_key or litellm.cache.get_cache_key(**request_kwargs)
|
||||
)
|
||||
cached_result = litellm.cache.get_cache(
|
||||
dynamic_cache_object=self.dual_cache, **new_kwargs
|
||||
dynamic_cache_object=self.dual_cache,
|
||||
cache_key=self.preset_cache_key,
|
||||
**request_kwargs,
|
||||
)
|
||||
return cached_result
|
||||
|
||||
|
|
@ -825,8 +861,27 @@ class LLMCachingHandler:
|
|||
elif (call_type == "aresponses" or call_type == "responses") and isinstance(
|
||||
cached_result, dict
|
||||
):
|
||||
# Convert cached dict back to ResponsesAPIResponse object
|
||||
cached_result = ResponsesAPIResponse(**cached_result)
|
||||
from litellm.responses.streaming_iterator import (
|
||||
CachedResponsesAPIStreamingIterator,
|
||||
)
|
||||
|
||||
response_obj = ResponsesAPIResponse(**cached_result)
|
||||
if (
|
||||
hasattr(response_obj, "_hidden_params")
|
||||
and response_obj._hidden_params is not None
|
||||
and isinstance(response_obj._hidden_params, dict)
|
||||
):
|
||||
response_obj._hidden_params["cache_hit"] = True
|
||||
|
||||
if kwargs.get("stream", False) is True:
|
||||
cached_result = CachedResponsesAPIStreamingIterator(
|
||||
response=response_obj,
|
||||
logging_obj=logging_obj,
|
||||
request_data=kwargs,
|
||||
call_type=call_type,
|
||||
)
|
||||
else:
|
||||
cached_result = response_obj
|
||||
|
||||
if (
|
||||
hasattr(cached_result, "_hidden_params")
|
||||
|
|
|
|||
|
|
@ -92,6 +92,25 @@ class DualCache(BaseCache):
|
|||
if default_redis_ttl is not None:
|
||||
self.default_redis_ttl = default_redis_ttl
|
||||
|
||||
def attach_redis_cache(
|
||||
self,
|
||||
redis_cache: Optional[RedisCache] = None,
|
||||
*,
|
||||
default_redis_ttl: Optional[float] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Attach a Redis backend if this DualCache does not already have one.
|
||||
|
||||
No-op when ``redis_cache`` is None or when Redis was already set (constructor
|
||||
or a prior attach). Use this for lazy wiring after a shared Redis client exists.
|
||||
Does not backfill in-memory-only keys to Redis.
|
||||
"""
|
||||
if redis_cache is None or self.redis_cache is not None:
|
||||
return
|
||||
self.redis_cache = redis_cache
|
||||
if default_redis_ttl is not None:
|
||||
self.default_redis_ttl = default_redis_ttl
|
||||
|
||||
def set_cache(self, key, value, local_only: bool = False, **kwargs):
|
||||
# Update both Redis and in-memory cache
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -551,6 +551,13 @@ class RedisCache(BaseCache):
|
|||
async def async_set_cache(self, key, value, **kwargs):
|
||||
from redis.asyncio import Redis
|
||||
|
||||
if key is None:
|
||||
verbose_logger.debug(
|
||||
"LiteLLM Redis Caching: async set() skipped — key is None, value=%r",
|
||||
value,
|
||||
)
|
||||
return None
|
||||
|
||||
start_time = time.time()
|
||||
try:
|
||||
_redis_client: Redis = self.init_async_client() # type: ignore
|
||||
|
|
@ -569,8 +576,9 @@ class RedisCache(BaseCache):
|
|||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async set() - Got exception from REDIS %s, Writing value=%s",
|
||||
"LiteLLM Redis Caching: async set() - Got exception from REDIS %s, key=%r, value=%r",
|
||||
str(e),
|
||||
key,
|
||||
value,
|
||||
)
|
||||
raise e
|
||||
|
|
|
|||
|
|
@ -2,11 +2,23 @@
|
|||
Arize Phoenix API client for fetching prompt versions from Arize Phoenix.
|
||||
"""
|
||||
|
||||
import urllib.parse
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
|
||||
def _sanitize_id(identifier: str) -> str:
|
||||
"""Reject path traversal characters and URL-encode the identifier."""
|
||||
if any(c in identifier for c in ("/", "\\", "#", "?")):
|
||||
raise ValueError(
|
||||
f"Invalid identifier {identifier!r}: contains disallowed characters"
|
||||
)
|
||||
if ".." in identifier:
|
||||
raise ValueError(f"Invalid identifier {identifier!r}: path traversal detected")
|
||||
return urllib.parse.quote(identifier, safe="")
|
||||
|
||||
|
||||
class ArizePhoenixClient:
|
||||
"""
|
||||
Client for interacting with Arize Phoenix API to fetch prompt versions.
|
||||
|
|
@ -53,7 +65,8 @@ class ArizePhoenixClient:
|
|||
Returns:
|
||||
Dictionary containing prompt version data, or None if not found
|
||||
"""
|
||||
url = f"{self.api_base}/v1/prompt_versions/{prompt_version_id}"
|
||||
safe_id = _sanitize_id(prompt_version_id)
|
||||
url = f"{self.api_base}/v1/prompt_versions/{safe_id}"
|
||||
|
||||
try:
|
||||
# Use the underlying httpx client directly to avoid query param extraction
|
||||
|
|
|
|||
|
|
@ -3,11 +3,27 @@ BitBucket API client for fetching .prompt files from BitBucket repositories.
|
|||
"""
|
||||
|
||||
import base64
|
||||
import urllib.parse
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
|
||||
def _sanitize_file_path(file_path: str) -> str:
|
||||
"""Reject path traversal and URL-encode each path segment."""
|
||||
if "#" in file_path or "?" in file_path:
|
||||
raise ValueError(
|
||||
f"Invalid file path {file_path!r}: contains URL special characters"
|
||||
)
|
||||
parts = file_path.split("/")
|
||||
for part in parts:
|
||||
if part == "..":
|
||||
raise ValueError(
|
||||
f"Invalid file path {file_path!r}: path traversal detected"
|
||||
)
|
||||
return "/".join(urllib.parse.quote(part, safe="") for part in parts)
|
||||
|
||||
|
||||
class BitBucketClient:
|
||||
"""
|
||||
Client for interacting with BitBucket API to fetch .prompt files.
|
||||
|
|
@ -72,7 +88,8 @@ class BitBucketClient:
|
|||
Returns:
|
||||
File content as string, or None if file not found
|
||||
"""
|
||||
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{file_path}"
|
||||
safe_path = _sanitize_file_path(file_path)
|
||||
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_path}"
|
||||
|
||||
try:
|
||||
response = self.http_handler.get(url, headers=self.headers)
|
||||
|
|
@ -119,7 +136,8 @@ class BitBucketClient:
|
|||
Returns:
|
||||
List of file paths
|
||||
"""
|
||||
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{directory_path}"
|
||||
safe_dir = _sanitize_file_path(directory_path) if directory_path else ""
|
||||
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_dir}"
|
||||
|
||||
try:
|
||||
response = self.http_handler.get(url, headers=self.headers)
|
||||
|
|
@ -211,7 +229,8 @@ class BitBucketClient:
|
|||
Returns:
|
||||
Dictionary containing file metadata, or None if file not found
|
||||
"""
|
||||
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{file_path}"
|
||||
safe_path = _sanitize_file_path(file_path)
|
||||
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_path}"
|
||||
|
||||
try:
|
||||
# Use GET with Range header to get just the headers (HEAD equivalent)
|
||||
|
|
|
|||
|
|
@ -265,6 +265,7 @@ class PrometheusLogger(CustomLogger):
|
|||
########################################
|
||||
# LiteLLM Virtual API KEY metrics
|
||||
########################################
|
||||
|
||||
# Remaining MODEL RPM limit for API Key
|
||||
self.litellm_remaining_api_key_requests_for_model = self._gauge_factory(
|
||||
"litellm_remaining_api_key_requests_for_model",
|
||||
|
|
|
|||
|
|
@ -31,15 +31,23 @@ def load_cli_token() -> Optional[dict]:
|
|||
return None
|
||||
|
||||
|
||||
def get_litellm_gateway_api_key() -> Optional[str]:
|
||||
def get_litellm_gateway_api_key(
|
||||
expected_base_url: Optional[str] = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Get the stored CLI API key for use with LiteLLM SDK.
|
||||
|
||||
This function reads the token file created by `litellm-proxy login`
|
||||
and returns the API key for use in Python scripts.
|
||||
|
||||
Args:
|
||||
expected_base_url: When provided, the key is only returned if it was
|
||||
originally issued for this URL. Pass the target server URL to
|
||||
prevent credential leakage when the client is pointed at a
|
||||
different (possibly malicious) server.
|
||||
|
||||
Returns:
|
||||
str: The API key if found, None otherwise
|
||||
str: The API key if found (and origin matches), None otherwise
|
||||
|
||||
Example:
|
||||
>>> import litellm
|
||||
|
|
@ -53,6 +61,10 @@ def get_litellm_gateway_api_key() -> Optional[str]:
|
|||
>>> )
|
||||
"""
|
||||
token_data = load_cli_token()
|
||||
if token_data and "key" in token_data:
|
||||
return token_data["key"]
|
||||
return None
|
||||
if not token_data or "key" not in token_data:
|
||||
return None
|
||||
if expected_base_url is not None:
|
||||
stored_url = token_data.get("base_url")
|
||||
if stored_url != expected_base_url.rstrip("/"):
|
||||
return None
|
||||
return token_data["key"]
|
||||
|
|
|
|||
|
|
@ -4582,6 +4582,11 @@ class BedrockConverseMessagesProcessor:
|
|||
message=cast(ChatCompletionFileObject, element)
|
||||
)
|
||||
_parts.append(_part)
|
||||
elif element["type"] == "document":
|
||||
_part = BedrockConverseMessagesProcessor._process_document_message(
|
||||
element
|
||||
)
|
||||
_parts.append(_part)
|
||||
_cache_point_block = (
|
||||
litellm.AmazonConverseConfig()._get_cache_point_block(
|
||||
message_block=cast(
|
||||
|
|
@ -4864,6 +4869,44 @@ class BedrockConverseMessagesProcessor:
|
|||
image_url=cast(str, file_id or file_data), format=format
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _process_document_message(element: dict) -> BedrockContentBlock:
|
||||
"""Convert a document content block to a Bedrock DocumentBlock.
|
||||
|
||||
Handles the Anthropic-style document format:
|
||||
{"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "..."}}
|
||||
"""
|
||||
source = element["source"]
|
||||
source_type = source.get("type")
|
||||
if source_type != "base64":
|
||||
raise ValueError(
|
||||
f"Bedrock Converse only supports base64-encoded document sources, got '{source_type}'. "
|
||||
"Please convert the document to base64 before sending to Bedrock."
|
||||
)
|
||||
media_type: str = source["media_type"]
|
||||
data: str = source["data"]
|
||||
doc_format = BedrockImageProcessor._validate_format(
|
||||
mime_type=media_type, image_format=media_type.split("/")[1]
|
||||
)
|
||||
|
||||
# Deterministic name using the same hashing pattern as _create_bedrock_block
|
||||
HASH_SAMPLE_BYTES = 64 * 1024
|
||||
normalized = "".join(data.split()).encode("utf-8")
|
||||
sample = normalized[:HASH_SAMPLE_BYTES]
|
||||
hasher = hashlib.sha256()
|
||||
hasher.update(sample)
|
||||
hasher.update(str(len(normalized)).encode("utf-8"))
|
||||
content_hash = hasher.hexdigest()[:16]
|
||||
document_name = f"Document_{content_hash}_{doc_format}"
|
||||
|
||||
return BedrockContentBlock(
|
||||
document=BedrockDocumentBlock(
|
||||
source=BedrockSourceBlock(bytes=data),
|
||||
format=doc_format,
|
||||
name=document_name,
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def add_thinking_blocks_to_assistant_content(
|
||||
thinking_blocks: List[BedrockContentBlock],
|
||||
|
|
@ -4961,6 +5004,11 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915
|
|||
)
|
||||
)
|
||||
_parts.append(_part)
|
||||
elif element["type"] == "document":
|
||||
_part = BedrockConverseMessagesProcessor._process_document_message(
|
||||
element
|
||||
)
|
||||
_parts.append(_part)
|
||||
_cache_point_block = (
|
||||
litellm.AmazonConverseConfig()._get_cache_point_block(
|
||||
message_block=cast(
|
||||
|
|
|
|||
|
|
@ -97,7 +97,7 @@ def get_vertex_ai_model_route(
|
|||
Determine which handler to use for a Vertex AI model based on the model name.
|
||||
|
||||
Args:
|
||||
model: The model name (e.g., "llama3-405b", "gemini-pro", "gemma/gemma-3-12b-it", "google/gemma-4-26b-a4b-it-maas", "openai/gpt-oss-120b")
|
||||
model: The model name (e.g., "llama3-405b", "gemini-pro", "gemma/gemma-3-12b-it", "google/gemma-4-26b-a4b-it-maas", "openai/gpt-oss-120b", "xai/grok-4.1-fast-non-reasoning")
|
||||
litellm_params: Optional litellm parameters dict that may contain base_model for routing
|
||||
|
||||
Returns:
|
||||
|
|
@ -119,6 +119,9 @@ def get_vertex_ai_model_route(
|
|||
>>> get_vertex_ai_model_route("openai/gpt-oss-120b")
|
||||
VertexAIModelRoute.MODEL_GARDEN
|
||||
|
||||
>>> get_vertex_ai_model_route("xai/grok-4.1-fast-non-reasoning")
|
||||
VertexAIModelRoute.MODEL_GARDEN
|
||||
|
||||
>>> get_vertex_ai_model_route("1234567890", {"api_base": "http://10.96.32.8"})
|
||||
VertexAIModelRoute.GEMINI # Numeric endpoints with api_base use HTTP path
|
||||
"""
|
||||
|
|
@ -152,8 +155,11 @@ def get_vertex_ai_model_route(
|
|||
if "gemma/" in model or model.startswith("google/gemma-"):
|
||||
return VertexAIModelRoute.GEMMA
|
||||
|
||||
# Check for model garden openai models
|
||||
if "openai" in model:
|
||||
# Check for model garden OpenAI-compatible publisher models.
|
||||
# Examples:
|
||||
# - openai/gpt-oss-120b-maas
|
||||
# - xai/grok-4.1-fast-non-reasoning
|
||||
if "openai" in model or model.startswith("xai/"):
|
||||
return VertexAIModelRoute.MODEL_GARDEN
|
||||
|
||||
# Check for gemini models
|
||||
|
|
@ -259,8 +265,8 @@ def get_vertex_base_model_name(model: str) -> str:
|
|||
>>> get_vertex_base_model_name("gemma/gemma-3-12b-it")
|
||||
"gemma-3-12b-it"
|
||||
|
||||
>>> get_vertex_base_model_name("openai/gpt-oss-120b")
|
||||
"gpt-oss-120b"
|
||||
>>> get_vertex_base_model_name("xai/grok-4.1-fast-non-reasoning")
|
||||
"grok-4.1-fast-non-reasoning"
|
||||
|
||||
>>> get_vertex_base_model_name("1234567890")
|
||||
"1234567890"
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ Google AI Studio /batchEmbedContents Embeddings Endpoint
|
|||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, Literal, Optional, Union
|
||||
from typing import Any, Dict, List, Literal, Optional, Tuple, Union
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -13,8 +13,8 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
HTTPHandler,
|
||||
get_async_httpx_client,
|
||||
)
|
||||
from litellm.types.llms.openai import EmbeddingInput
|
||||
from litellm.types.llms.vertex_ai import (
|
||||
GeminiEmbeddingInput,
|
||||
VertexAIBatchEmbeddingsRequestBody,
|
||||
VertexAIBatchEmbeddingsResponseObject,
|
||||
)
|
||||
|
|
@ -23,7 +23,6 @@ from litellm.types.utils import EmbeddingResponse
|
|||
from ..gemini.vertex_and_google_ai_studio_gemini import VertexLLM
|
||||
from .batch_embed_content_transformation import (
|
||||
_is_file_reference,
|
||||
_is_multimodal_input,
|
||||
process_embed_content_response,
|
||||
process_response,
|
||||
transform_openai_input_gemini_content,
|
||||
|
|
@ -32,9 +31,24 @@ from .batch_embed_content_transformation import (
|
|||
|
||||
|
||||
class GoogleBatchEmbeddings(VertexLLM):
|
||||
@staticmethod
|
||||
def _flatten_and_detect_file_refs(
|
||||
input: GeminiEmbeddingInput,
|
||||
) -> Tuple[List[str], bool]:
|
||||
"""Flatten nested input lists and detect file references."""
|
||||
input_list = [input] if isinstance(input, str) else input
|
||||
flat_elements = [
|
||||
e
|
||||
for item in input_list
|
||||
for e in (item if isinstance(item, list) else [item])
|
||||
if isinstance(e, str)
|
||||
]
|
||||
has_file_refs = any(_is_file_reference(e) for e in flat_elements)
|
||||
return flat_elements, has_file_refs
|
||||
|
||||
def _resolve_file_references(
|
||||
self,
|
||||
input: EmbeddingInput,
|
||||
input: GeminiEmbeddingInput,
|
||||
api_key: str,
|
||||
sync_handler: HTTPHandler,
|
||||
) -> Dict[str, Dict[str, str]]:
|
||||
|
|
@ -42,7 +56,7 @@ class GoogleBatchEmbeddings(VertexLLM):
|
|||
Resolve Gemini file references (files/...) to get mime_type and uri.
|
||||
|
||||
Args:
|
||||
input: EmbeddingInput that may contain file references
|
||||
input: GeminiEmbeddingInput that may contain file references
|
||||
api_key: Gemini API key
|
||||
sync_handler: HTTP client
|
||||
|
||||
|
|
@ -73,7 +87,7 @@ class GoogleBatchEmbeddings(VertexLLM):
|
|||
|
||||
async def _async_resolve_file_references(
|
||||
self,
|
||||
input: EmbeddingInput,
|
||||
input: GeminiEmbeddingInput,
|
||||
api_key: str,
|
||||
async_handler: AsyncHTTPHandler,
|
||||
) -> Dict[str, Dict[str, str]]:
|
||||
|
|
@ -81,7 +95,7 @@ class GoogleBatchEmbeddings(VertexLLM):
|
|||
Async version of _resolve_file_references.
|
||||
|
||||
Args:
|
||||
input: EmbeddingInput that may contain file references
|
||||
input: GeminiEmbeddingInput that may contain file references
|
||||
api_key: Gemini API key
|
||||
async_handler: Async HTTP client
|
||||
|
||||
|
|
@ -110,10 +124,10 @@ class GoogleBatchEmbeddings(VertexLLM):
|
|||
|
||||
return resolved_files
|
||||
|
||||
def batch_embeddings(
|
||||
def batch_embeddings( # noqa: PLR0915
|
||||
self,
|
||||
model: str,
|
||||
input: EmbeddingInput,
|
||||
input: GeminiEmbeddingInput,
|
||||
print_verbose,
|
||||
model_response: EmbeddingResponse,
|
||||
custom_llm_provider: Literal["gemini", "vertex_ai"],
|
||||
|
|
@ -151,8 +165,7 @@ class GoogleBatchEmbeddings(VertexLLM):
|
|||
|
||||
optional_params = optional_params or {}
|
||||
|
||||
is_multimodal = _is_multimodal_input(input)
|
||||
use_embed_content = is_multimodal or (custom_llm_provider == "vertex_ai")
|
||||
use_embed_content = custom_llm_provider == "vertex_ai"
|
||||
mode: Literal["embedding", "batch_embedding"]
|
||||
if use_embed_content:
|
||||
mode = "embedding"
|
||||
|
|
@ -215,8 +228,22 @@ class GoogleBatchEmbeddings(VertexLLM):
|
|||
resolved_files=resolved_files,
|
||||
)
|
||||
else:
|
||||
flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input)
|
||||
if has_file_refs and not api_key:
|
||||
raise ValueError(
|
||||
"An API key is required to resolve Gemini file references (files/...). "
|
||||
"Pass api_key= or set GEMINI_API_KEY."
|
||||
)
|
||||
resolved_files = {}
|
||||
if api_key and has_file_refs:
|
||||
resolved_files = self._resolve_file_references(
|
||||
input=flat_elements, api_key=api_key, sync_handler=sync_handler
|
||||
)
|
||||
request_data = transform_openai_input_gemini_content(
|
||||
input=input, model=model, optional_params=optional_params
|
||||
input=input,
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
resolved_files=resolved_files,
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
|
|
@ -264,7 +291,7 @@ class GoogleBatchEmbeddings(VertexLLM):
|
|||
url: str,
|
||||
data: Optional[Union[VertexAIBatchEmbeddingsRequestBody, dict]],
|
||||
model_response: EmbeddingResponse,
|
||||
input: EmbeddingInput,
|
||||
input: GeminiEmbeddingInput,
|
||||
timeout: Optional[Union[float, httpx.Timeout]],
|
||||
headers={},
|
||||
client: Optional[AsyncHTTPHandler] = None,
|
||||
|
|
@ -303,8 +330,22 @@ class GoogleBatchEmbeddings(VertexLLM):
|
|||
resolved_files=resolved_files,
|
||||
)
|
||||
else:
|
||||
flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input)
|
||||
if has_file_refs and not api_key:
|
||||
raise ValueError(
|
||||
"An API key is required to resolve Gemini file references (files/...). "
|
||||
"Pass api_key= or set GEMINI_API_KEY."
|
||||
)
|
||||
resolved_files = {}
|
||||
if api_key and has_file_refs:
|
||||
resolved_files = await self._async_resolve_file_references(
|
||||
input=flat_elements, api_key=api_key, async_handler=async_handler
|
||||
)
|
||||
data = transform_openai_input_gemini_content(
|
||||
input=input, model=model, optional_params=optional_params or {}
|
||||
input=input,
|
||||
model=model,
|
||||
optional_params=optional_params or {},
|
||||
resolved_files=resolved_files,
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
|
|
|
|||
|
|
@ -6,12 +6,12 @@ Why separate file? Make it easy to see how transformation works
|
|||
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from litellm.types.llms.openai import EmbeddingInput
|
||||
from litellm.types.llms.vertex_ai import (
|
||||
BlobType,
|
||||
ContentType,
|
||||
EmbedContentRequest,
|
||||
FileDataType,
|
||||
GeminiEmbeddingInput,
|
||||
PartType,
|
||||
VertexAIBatchEmbeddingsRequestBody,
|
||||
VertexAIBatchEmbeddingsResponseObject,
|
||||
|
|
@ -114,33 +114,77 @@ def _parse_data_url(data_url: str) -> Tuple[str, str]:
|
|||
return media_type, base64_data
|
||||
|
||||
|
||||
def _is_multimodal_input(input: EmbeddingInput) -> bool:
|
||||
def _is_multimodal_input(input: GeminiEmbeddingInput) -> bool:
|
||||
"""
|
||||
Check if the input contains multimodal data (data URIs, file references, or GCS URLs).
|
||||
Check if the input contains multimodal data (data URIs, file references,
|
||||
GCS URLs, or nested lists for combined embeddings).
|
||||
|
||||
Args:
|
||||
input: EmbeddingInput (str or List[str])
|
||||
input: GeminiEmbeddingInput — str, List[str], or List[List[str]] for combined embeddings
|
||||
|
||||
Returns:
|
||||
bool: True if any element is a data URI, file reference, or GCS URL
|
||||
bool: True if any element is multimodal or a nested list
|
||||
"""
|
||||
if isinstance(input, str):
|
||||
input_list = [input]
|
||||
else:
|
||||
input_list = input
|
||||
return _is_multimodal_element(input)
|
||||
|
||||
for element in input_list:
|
||||
if isinstance(element, str):
|
||||
if element.startswith("data:") and ";base64," in element:
|
||||
return True
|
||||
if _is_file_reference(element):
|
||||
return True
|
||||
if _is_gcs_url(element):
|
||||
for element in input:
|
||||
if isinstance(element, list):
|
||||
if any(
|
||||
_is_multimodal_element(sub) for sub in element if isinstance(sub, str)
|
||||
):
|
||||
return True
|
||||
elif isinstance(element, str) and _is_multimodal_element(element):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _is_multimodal_element(element: str) -> bool:
|
||||
"""Check if a single string element is multimodal."""
|
||||
if element.startswith("data:") and ";base64," in element:
|
||||
return True
|
||||
if _is_file_reference(element):
|
||||
return True
|
||||
if _is_gcs_url(element):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _build_part_for_input(
|
||||
element: str,
|
||||
resolved_files: Optional[Dict[str, Dict[str, str]]] = None,
|
||||
) -> PartType:
|
||||
"""
|
||||
Build a single PartType for an input element, handling text, data URIs,
|
||||
file references, and GCS URLs.
|
||||
"""
|
||||
resolved_files = resolved_files or {}
|
||||
|
||||
if element.startswith("data:") and ";base64," in element:
|
||||
mime_type, base64_data = _parse_data_url(element)
|
||||
blob: BlobType = {"mime_type": mime_type, "data": base64_data}
|
||||
return PartType(inline_data=blob)
|
||||
elif _is_gcs_url(element):
|
||||
mime_type = _infer_mime_type_from_gcs_url(element)
|
||||
file_data: FileDataType = {
|
||||
"mime_type": mime_type,
|
||||
"file_uri": element,
|
||||
}
|
||||
return PartType(file_data=file_data)
|
||||
elif _is_file_reference(element):
|
||||
if element not in resolved_files:
|
||||
raise ValueError(f"File reference {element} not resolved")
|
||||
file_info = resolved_files[element]
|
||||
file_data_ref: FileDataType = {
|
||||
"mime_type": file_info["mime_type"],
|
||||
"file_uri": file_info["uri"],
|
||||
}
|
||||
return PartType(file_data=file_data_ref)
|
||||
else:
|
||||
return PartType(text=element)
|
||||
|
||||
|
||||
_SUPPORTED_EMBED_PARAMS = {"outputDimensionality", "taskType", "title"}
|
||||
|
||||
|
||||
|
|
@ -155,37 +199,60 @@ def _filter_embed_params(optional_params: dict) -> dict:
|
|||
|
||||
|
||||
def transform_openai_input_gemini_content(
|
||||
input: EmbeddingInput, model: str, optional_params: dict
|
||||
input: GeminiEmbeddingInput,
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
resolved_files: Optional[Dict[str, Dict[str, str]]] = None,
|
||||
) -> VertexAIBatchEmbeddingsRequestBody:
|
||||
"""
|
||||
The content to embed. Only the parts.text fields will be counted.
|
||||
Transform OpenAI embedding input to Gemini batchEmbedContents format.
|
||||
|
||||
Each input element becomes a separate EmbedContentRequest, supporting
|
||||
text, data URIs, file references, and GCS URLs.
|
||||
|
||||
If an element is a list (nested input), all sub-elements are combined
|
||||
into a single content with multiple parts, producing one combined
|
||||
embedding for the group.
|
||||
|
||||
Examples:
|
||||
input=["text", "image"] → 2 separate embeddings
|
||||
input=[["text", "image"]] → 1 combined embedding
|
||||
input=[["text", "image"], "x"] → 2 embeddings (1 combined + 1 separate)
|
||||
"""
|
||||
gemini_model_name = "models/{}".format(model)
|
||||
|
||||
gemini_params = _filter_embed_params(optional_params)
|
||||
|
||||
input_list = [input] if isinstance(input, str) else input
|
||||
requests: List[EmbedContentRequest] = []
|
||||
if isinstance(input, str):
|
||||
|
||||
for element in input_list:
|
||||
if isinstance(element, list):
|
||||
if not element:
|
||||
raise ValueError("Nested input list must not be empty")
|
||||
for sub in element:
|
||||
if not isinstance(sub, str):
|
||||
raise ValueError(
|
||||
f"Elements inside a nested input list must be strings, got {type(sub)}"
|
||||
)
|
||||
parts = [
|
||||
_build_part_for_input(sub, resolved_files=resolved_files)
|
||||
for sub in element
|
||||
]
|
||||
else:
|
||||
parts = [_build_part_for_input(element, resolved_files=resolved_files)]
|
||||
request = EmbedContentRequest(
|
||||
model=gemini_model_name,
|
||||
content=ContentType(parts=[PartType(text=input)]),
|
||||
content=ContentType(parts=parts),
|
||||
**gemini_params,
|
||||
)
|
||||
requests.append(request)
|
||||
else:
|
||||
for i in input:
|
||||
request = EmbedContentRequest(
|
||||
model=gemini_model_name,
|
||||
content=ContentType(parts=[PartType(text=i)]),
|
||||
**gemini_params,
|
||||
)
|
||||
requests.append(request)
|
||||
|
||||
return VertexAIBatchEmbeddingsRequestBody(requests=requests)
|
||||
|
||||
|
||||
def transform_openai_input_gemini_embed_content(
|
||||
input: EmbeddingInput,
|
||||
input: GeminiEmbeddingInput,
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
resolved_files: Optional[Dict[str, Dict[str, str]]] = None,
|
||||
|
|
@ -194,7 +261,7 @@ def transform_openai_input_gemini_embed_content(
|
|||
Transform OpenAI embedding input to Gemini embedContent format (multimodal).
|
||||
|
||||
Args:
|
||||
input: EmbeddingInput (str or List[str]) with text, data URIs, or file references
|
||||
input: GeminiEmbeddingInput with text, data URIs, or file references
|
||||
model: Model name
|
||||
optional_params: Additional parameters (taskType, outputDimensionality, etc.)
|
||||
resolved_files: Dict mapping file names (files/abc) to {mime_type, uri}
|
||||
|
|
@ -210,31 +277,14 @@ def transform_openai_input_gemini_embed_content(
|
|||
parts: List[PartType] = []
|
||||
|
||||
for element in input_list:
|
||||
if isinstance(element, list):
|
||||
raise ValueError(
|
||||
"Nested (combined) embeddings are not supported on the embedContent path. "
|
||||
"Use the batchEmbedContents path or pass a flat list instead."
|
||||
)
|
||||
if not isinstance(element, str):
|
||||
raise ValueError(f"Unsupported input type: {type(element)}")
|
||||
|
||||
if element.startswith("data:") and ";base64," in element:
|
||||
mime_type, base64_data = _parse_data_url(element)
|
||||
blob: BlobType = {"mime_type": mime_type, "data": base64_data}
|
||||
parts.append(PartType(inline_data=blob))
|
||||
elif _is_gcs_url(element):
|
||||
mime_type = _infer_mime_type_from_gcs_url(element)
|
||||
file_data: FileDataType = {
|
||||
"mime_type": mime_type,
|
||||
"file_uri": element,
|
||||
}
|
||||
parts.append(PartType(file_data=file_data))
|
||||
elif _is_file_reference(element):
|
||||
if element not in resolved_files:
|
||||
raise ValueError(f"File reference {element} not resolved")
|
||||
file_info = resolved_files[element]
|
||||
file_data_ref: FileDataType = {
|
||||
"mime_type": file_info["mime_type"],
|
||||
"file_uri": file_info["uri"],
|
||||
}
|
||||
parts.append(PartType(file_data=file_data_ref))
|
||||
else:
|
||||
parts.append(PartType(text=element))
|
||||
parts.append(_build_part_for_input(element, resolved_files=resolved_files))
|
||||
|
||||
request_body: dict = {
|
||||
"content": ContentType(parts=parts),
|
||||
|
|
@ -245,7 +295,7 @@ def transform_openai_input_gemini_embed_content(
|
|||
|
||||
|
||||
def process_embed_content_response(
|
||||
input: EmbeddingInput,
|
||||
input: GeminiEmbeddingInput,
|
||||
model_response: EmbeddingResponse,
|
||||
model: str,
|
||||
response_json: dict,
|
||||
|
|
@ -291,7 +341,7 @@ def process_embed_content_response(
|
|||
|
||||
|
||||
def process_response(
|
||||
input: EmbeddingInput,
|
||||
input: GeminiEmbeddingInput,
|
||||
model_response: EmbeddingResponse,
|
||||
model: str,
|
||||
_predictions: VertexAIBatchEmbeddingsResponseObject,
|
||||
|
|
@ -308,8 +358,29 @@ def process_response(
|
|||
model_response.data = openai_embeddings
|
||||
model_response.model = model
|
||||
|
||||
input_text = get_formatted_prompt(data={"input": input}, call_type="embedding")
|
||||
prompt_tokens = token_counter(model=model, text=input_text)
|
||||
has_nested = isinstance(input, list) and any(isinstance(e, list) for e in input)
|
||||
if _is_multimodal_input(input) or has_nested:
|
||||
input_list = input if isinstance(input, list) else [input]
|
||||
text_elements: List[str] = []
|
||||
for e in input_list:
|
||||
if isinstance(e, list):
|
||||
text_elements.extend(
|
||||
sub
|
||||
for sub in e
|
||||
if isinstance(sub, str) and not _is_multimodal_element(sub)
|
||||
)
|
||||
elif isinstance(e, str) and not _is_multimodal_element(e):
|
||||
text_elements.append(e)
|
||||
if text_elements:
|
||||
input_text = get_formatted_prompt(
|
||||
data={"input": text_elements}, call_type="embedding"
|
||||
)
|
||||
prompt_tokens = token_counter(model=model, text=input_text)
|
||||
else:
|
||||
prompt_tokens = 0
|
||||
else:
|
||||
input_text = get_formatted_prompt(data={"input": input}, call_type="embedding")
|
||||
prompt_tokens = token_counter(model=model, text=input_text)
|
||||
model_response.usage = Usage(
|
||||
prompt_tokens=prompt_tokens, total_tokens=prompt_tokens
|
||||
)
|
||||
|
|
|
|||
|
|
@ -27,6 +27,17 @@ from ..common_utils import VertexAIError, get_vertex_base_model_name
|
|||
from ..vertex_llm_base import VertexBase
|
||||
|
||||
|
||||
def _vertex_model_garden_model_id_in_json_body(model: str) -> bool:
|
||||
"""
|
||||
Vertex catalog / publisher models are addressed as publisher/model (e.g.
|
||||
xai/grok-4.1-fast-reasoning) on the shared OpenAPI URL, with the id in the JSON body.
|
||||
|
||||
Deployed Model Garden endpoints are typically a single segment (often numeric)
|
||||
and use .../endpoints/{ENDPOINT_ID}/chat/completions with an empty model field.
|
||||
"""
|
||||
return "/" in model
|
||||
|
||||
|
||||
def create_vertex_url(
|
||||
vertex_location: str,
|
||||
vertex_project: str,
|
||||
|
|
@ -34,8 +45,13 @@ def create_vertex_url(
|
|||
model: str,
|
||||
api_base: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Return the base url for the vertex garden models"""
|
||||
"""Return the api base for vertex model garden (without /chat/completions)."""
|
||||
base_url = get_vertex_base_url(vertex_location)
|
||||
if _vertex_model_garden_model_id_in_json_body(model):
|
||||
return (
|
||||
f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}"
|
||||
"/endpoints/openapi"
|
||||
)
|
||||
return f"{base_url}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}"
|
||||
|
||||
|
||||
|
|
@ -129,7 +145,10 @@ class VertexAIModelGardenModels(VertexBase):
|
|||
vertex_location=vertex_location or "us-central1",
|
||||
vertex_api_version="v1beta1",
|
||||
)
|
||||
model = ""
|
||||
# Publisher/catalog models: model id must be sent in the JSON body (OpenAPI route).
|
||||
# Single-segment endpoint ids: model is encoded in the URL path; body model stays empty.
|
||||
if not _vertex_model_garden_model_id_in_json_body(model):
|
||||
model = ""
|
||||
return openai_like_chat_completions.completion(
|
||||
model=model,
|
||||
messages=messages,
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ class XAIChatConfig(OpenAIGPTConfig):
|
|||
"logprobs",
|
||||
"max_tokens",
|
||||
"n",
|
||||
"parallel_tool_calls",
|
||||
"presence_penalty",
|
||||
"response_format",
|
||||
"seed",
|
||||
|
|
|
|||
|
|
@ -33351,6 +33351,72 @@
|
|||
"source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas",
|
||||
"supports_reasoning": true
|
||||
},
|
||||
"vertex_ai/xai/grok-4.1-fast-non-reasoning": {
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"max_input_tokens": 2000000,
|
||||
"max_output_tokens": 2000000,
|
||||
"max_tokens": 2000000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-07,
|
||||
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
|
||||
"supports_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"vertex_ai/xai/grok-4.1-fast-reasoning": {
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"max_input_tokens": 2000000,
|
||||
"max_output_tokens": 2000000,
|
||||
"max_tokens": 2000000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-07,
|
||||
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"vertex_ai/xai/grok-4.20-non-reasoning": {
|
||||
"cache_read_input_token_cost": 2e-07,
|
||||
"input_cost_per_token": 2e-06,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"max_input_tokens": 2000000,
|
||||
"max_output_tokens": 2000000,
|
||||
"max_tokens": 2000000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-06,
|
||||
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
|
||||
"supports_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"vertex_ai/xai/grok-4.20-reasoning": {
|
||||
"cache_read_input_token_cost": 2e-07,
|
||||
"input_cost_per_token": 2e-06,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"max_input_tokens": 2000000,
|
||||
"max_output_tokens": 2000000,
|
||||
"max_tokens": 2000000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-06,
|
||||
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": {
|
||||
"input_cost_per_token": 2.5e-07,
|
||||
"litellm_provider": "vertex_ai-qwen_models",
|
||||
|
|
|
|||
|
|
@ -169,6 +169,37 @@ def _deserialize_json_dict(data: Any) -> Optional[Dict[str, str]]:
|
|||
class MCPServerManager:
|
||||
_STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$")
|
||||
|
||||
@staticmethod
|
||||
def _resolve_oauth2_flow(
|
||||
*,
|
||||
auth_type: Optional[MCPAuthType],
|
||||
oauth2_flow: Optional[str],
|
||||
token_url: Optional[str],
|
||||
authorization_url: Optional[str],
|
||||
client_id: Optional[str],
|
||||
client_secret: Optional[str],
|
||||
) -> Optional[Literal["client_credentials", "authorization_code"]]:
|
||||
"""Infer oauth2_flow for legacy records that omit the field.
|
||||
|
||||
DB rows created before oauth2_flow support may have OAuth2 client
|
||||
credentials + token_url but a null oauth2_flow. Treat these as M2M,
|
||||
unless authorization_url is present (interactive OAuth).
|
||||
"""
|
||||
if oauth2_flow in ("client_credentials", "authorization_code"):
|
||||
return cast(
|
||||
Literal["client_credentials", "authorization_code"], oauth2_flow
|
||||
)
|
||||
if oauth2_flow:
|
||||
# Ignore unknown/untyped values and continue legacy inference.
|
||||
return None
|
||||
if auth_type != MCPAuth.oauth2:
|
||||
return None
|
||||
if authorization_url:
|
||||
return None
|
||||
if token_url and client_id and client_secret:
|
||||
return "client_credentials"
|
||||
return None
|
||||
|
||||
def __init__(self):
|
||||
self.registry: Dict[str, MCPServer] = {}
|
||||
self.config_mcp_servers: Dict[str, MCPServer] = {}
|
||||
|
|
@ -342,7 +373,14 @@ class MCPServerManager:
|
|||
# oauth specific fields
|
||||
client_id=server_config.get("client_id", None),
|
||||
client_secret=server_config.get("client_secret", None),
|
||||
oauth2_flow=server_config.get("oauth2_flow", None),
|
||||
oauth2_flow=self._resolve_oauth2_flow(
|
||||
auth_type=auth_type,
|
||||
oauth2_flow=server_config.get("oauth2_flow", None),
|
||||
token_url=resolved_token_url,
|
||||
authorization_url=resolved_authorization_url,
|
||||
client_id=server_config.get("client_id", None),
|
||||
client_secret=server_config.get("client_secret", None),
|
||||
),
|
||||
scopes=resolved_scopes,
|
||||
authorization_url=resolved_authorization_url,
|
||||
token_url=resolved_token_url,
|
||||
|
|
@ -679,7 +717,17 @@ class MCPServerManager:
|
|||
client_id=client_id_value or getattr(mcp_server, "client_id", None),
|
||||
client_secret=client_secret_value
|
||||
or getattr(mcp_server, "client_secret", None),
|
||||
oauth2_flow=getattr(mcp_server, "oauth2_flow", None),
|
||||
oauth2_flow=self._resolve_oauth2_flow(
|
||||
auth_type=auth_type,
|
||||
oauth2_flow=getattr(mcp_server, "oauth2_flow", None),
|
||||
token_url=mcp_server.token_url
|
||||
or getattr(mcp_oauth_metadata, "token_url", None),
|
||||
authorization_url=mcp_server.authorization_url
|
||||
or getattr(mcp_oauth_metadata, "authorization_url", None),
|
||||
client_id=client_id_value or getattr(mcp_server, "client_id", None),
|
||||
client_secret=client_secret_value
|
||||
or getattr(mcp_server, "client_secret", None),
|
||||
),
|
||||
scopes=resolved_scopes,
|
||||
authorization_url=mcp_server.authorization_url
|
||||
or getattr(mcp_oauth_metadata, "authorization_url", None),
|
||||
|
|
@ -2426,7 +2474,7 @@ class MCPServerManager:
|
|||
)
|
||||
)
|
||||
|
||||
async def _call_regular_mcp_tool(
|
||||
async def _call_regular_mcp_tool( # noqa: PLR0915
|
||||
self,
|
||||
mcp_server: MCPServer,
|
||||
original_tool_name: str,
|
||||
|
|
@ -2489,7 +2537,11 @@ class MCPServerManager:
|
|||
# oauth2 headers
|
||||
extra_headers: Optional[Dict[str, str]] = None
|
||||
if mcp_server.auth_type == MCPAuth.oauth2:
|
||||
extra_headers = oauth2_headers
|
||||
if mcp_server.has_client_credentials:
|
||||
# For M2M OAuth servers, Authorization must come from token fetch.
|
||||
extra_headers = None
|
||||
else:
|
||||
extra_headers = oauth2_headers
|
||||
|
||||
if mcp_server.extra_headers and raw_headers:
|
||||
if extra_headers is None:
|
||||
|
|
@ -2501,6 +2553,11 @@ class MCPServerManager:
|
|||
for header in mcp_server.extra_headers:
|
||||
if not isinstance(header, str):
|
||||
continue
|
||||
if (
|
||||
mcp_server.has_client_credentials
|
||||
and header.lower() == "authorization"
|
||||
):
|
||||
continue
|
||||
header_value = normalized_raw_headers.get(header.lower())
|
||||
if header_value is None:
|
||||
continue
|
||||
|
|
@ -2536,6 +2593,10 @@ class MCPServerManager:
|
|||
)
|
||||
extra_headers.update(hook_extra_headers)
|
||||
|
||||
# Reset to None if no headers were actually added
|
||||
if extra_headers is not None and len(extra_headers) == 0:
|
||||
extra_headers = None
|
||||
|
||||
stdio_env = self._build_stdio_env(mcp_server, raw_headers)
|
||||
|
||||
client = await self._create_mcp_client(
|
||||
|
|
|
|||
|
|
@ -153,6 +153,7 @@ if MCP_AVAILABLE:
|
|||
MCPAuthenticatedUser,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
|
||||
|
|
@ -900,6 +901,20 @@ if MCP_AVAILABLE:
|
|||
allowed_mcp_server_id
|
||||
)
|
||||
if mcp_server is not None:
|
||||
# Apply oauth2_flow resolution for legacy DB rows where it may be NULL
|
||||
resolved_flow = MCPServerManager._resolve_oauth2_flow(
|
||||
auth_type=mcp_server.auth_type,
|
||||
oauth2_flow=mcp_server.oauth2_flow,
|
||||
token_url=mcp_server.token_url,
|
||||
authorization_url=mcp_server.authorization_url,
|
||||
client_id=mcp_server.client_id,
|
||||
client_secret=mcp_server.client_secret,
|
||||
)
|
||||
if resolved_flow and resolved_flow != mcp_server.oauth2_flow:
|
||||
# Create a new instance with the resolved flow for this request
|
||||
mcp_server = mcp_server.model_copy(
|
||||
update={"oauth2_flow": resolved_flow}
|
||||
)
|
||||
allowed_mcp_servers.append(mcp_server)
|
||||
|
||||
if mcp_servers is not None:
|
||||
|
|
@ -1100,8 +1115,13 @@ if MCP_AVAILABLE:
|
|||
|
||||
extra_headers: Optional[Dict[str, str]] = None
|
||||
if server.auth_type == MCPAuth.oauth2:
|
||||
# Copy to avoid mutating the original dict (important for parallel fetching)
|
||||
extra_headers = oauth2_headers.copy() if oauth2_headers else None
|
||||
# For OAuth2 M2M servers, upstream Authorization must come from
|
||||
# client_credentials token fetch, never from caller headers.
|
||||
if server.has_client_credentials:
|
||||
extra_headers = None
|
||||
else:
|
||||
# Copy to avoid mutating the original dict (important for parallel fetching)
|
||||
extra_headers = oauth2_headers.copy() if oauth2_headers else None
|
||||
|
||||
if server.extra_headers and raw_headers:
|
||||
if extra_headers is None:
|
||||
|
|
@ -1114,11 +1134,17 @@ if MCP_AVAILABLE:
|
|||
for header in server.extra_headers:
|
||||
if not isinstance(header, str):
|
||||
continue
|
||||
if server.has_client_credentials and header.lower() == "authorization":
|
||||
continue
|
||||
header_value = normalized_raw_headers.get(header.lower())
|
||||
if header_value is None:
|
||||
continue
|
||||
extra_headers[header] = header_value
|
||||
|
||||
# Reset to None if no headers were actually added
|
||||
if extra_headers is not None and len(extra_headers) == 0:
|
||||
extra_headers = None
|
||||
|
||||
if server_auth_header is None:
|
||||
server_auth_header = mcp_auth_header
|
||||
|
||||
|
|
@ -1377,11 +1403,19 @@ if MCP_AVAILABLE:
|
|||
spend_meta["per_server_tool_counts"] = per_server_tool_counts
|
||||
|
||||
end_time = datetime.now()
|
||||
await litellm_logging_obj.async_success_handler(
|
||||
result=all_tools,
|
||||
start_time=list_tools_start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
try:
|
||||
await litellm_logging_obj.async_success_handler(
|
||||
result=all_tools,
|
||||
start_time=list_tools_start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
except Exception as log_exc:
|
||||
# list_tools responses must not be dropped due to non-blocking
|
||||
# observability/serialization failures.
|
||||
verbose_logger.warning(
|
||||
"MCP list_tools success logging failed (continuing): %s",
|
||||
log_exc,
|
||||
)
|
||||
|
||||
verbose_logger.info(
|
||||
f"Successfully fetched {len(all_tools)} tools total from all MCP servers"
|
||||
|
|
|
|||
|
|
@ -668,6 +668,8 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/models/{model_id}",
|
||||
"/guardrails/list",
|
||||
"/v2/guardrails/list",
|
||||
"/project/list",
|
||||
"/project/info",
|
||||
]
|
||||
+ spend_tracking_routes
|
||||
+ key_management_routes
|
||||
|
|
@ -692,6 +694,9 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/model/{model_id}/update",
|
||||
"/prompt/list",
|
||||
"/prompt/info",
|
||||
# Project read routes - endpoint scopes results to caller's teams (non-admin)
|
||||
"/project/list",
|
||||
"/project/info",
|
||||
# Invitation routes - org/team admins checked in endpoint via _user_has_admin_privileges
|
||||
"/invitation/new",
|
||||
"/invitation/delete",
|
||||
|
|
|
|||
|
|
@ -12,14 +12,13 @@ Run checks for:
|
|||
import asyncio
|
||||
import re
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union, cast
|
||||
|
||||
from fastapi import HTTPException, Request, status
|
||||
from pydantic import BaseModel
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.caching.dual_cache import LimitedSizeOrderedDict
|
||||
from litellm.constants import (
|
||||
CLI_JWT_EXPIRATION_HOURS,
|
||||
|
|
@ -66,6 +65,8 @@ from litellm.proxy.guardrails.tool_name_extraction import (
|
|||
TOOL_CAPABLE_CALL_TYPES,
|
||||
extract_request_tool_names,
|
||||
)
|
||||
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.route_llm_request import route_request
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics
|
||||
from litellm.router import Router
|
||||
|
|
@ -852,7 +853,7 @@ def get_actual_routes(allowed_routes: list) -> list:
|
|||
|
||||
async def get_default_end_user_budget(
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
) -> Optional[LiteLLM_BudgetTable]:
|
||||
"""
|
||||
|
|
@ -875,9 +876,12 @@ async def get_default_end_user_budget(
|
|||
cache_key = f"default_end_user_budget:{litellm.max_end_user_budget_id}"
|
||||
|
||||
# Check cache first
|
||||
cached_budget = await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
cached_budget = await user_api_key_cache.async_get_cache(
|
||||
key=cache_key,
|
||||
model_type=LiteLLM_BudgetTable,
|
||||
)
|
||||
if cached_budget is not None:
|
||||
return LiteLLM_BudgetTable(**cached_budget)
|
||||
return cached_budget
|
||||
|
||||
# Fetch from database
|
||||
try:
|
||||
|
|
@ -891,14 +895,16 @@ async def get_default_end_user_budget(
|
|||
)
|
||||
return None
|
||||
|
||||
_budget_obj = LiteLLM_BudgetTable(**budget_record.dict())
|
||||
# Cache the budget for 60 seconds
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=budget_record.dict(),
|
||||
value=_budget_obj,
|
||||
model_type=LiteLLM_BudgetTable,
|
||||
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
|
||||
)
|
||||
|
||||
return LiteLLM_BudgetTable(**budget_record.dict())
|
||||
return _budget_obj
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"Error fetching default end user budget: {str(e)}")
|
||||
|
|
@ -909,7 +915,7 @@ async def get_default_end_user_budget(
|
|||
async def get_team_member_default_budget(
|
||||
budget_id: str,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
) -> Optional[LiteLLM_BudgetTable]:
|
||||
"""
|
||||
Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"].
|
||||
|
|
@ -966,7 +972,7 @@ async def get_team_member_default_budget(
|
|||
async def _apply_default_budget_to_end_user(
|
||||
end_user_obj: LiteLLM_EndUserTable,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
) -> LiteLLM_EndUserTable:
|
||||
"""
|
||||
|
|
@ -1039,7 +1045,7 @@ def _check_end_user_budget(
|
|||
async def get_end_user_object(
|
||||
end_user_id: Optional[str],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
route: str,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
|
|
@ -1070,10 +1076,12 @@ async def get_end_user_object(
|
|||
_key = "end_user_id:{}".format(end_user_id)
|
||||
|
||||
# Check cache first
|
||||
cached_user_obj = await user_api_key_cache.async_get_cache(key=_key)
|
||||
cached_user_obj = await user_api_key_cache.async_get_cache(
|
||||
key=_key,
|
||||
model_type=LiteLLM_EndUserTable,
|
||||
)
|
||||
if cached_user_obj is not None:
|
||||
return_obj = LiteLLM_EndUserTable(**cached_user_obj)
|
||||
|
||||
return_obj = cached_user_obj
|
||||
# Apply default budget if needed
|
||||
return_obj = await _apply_default_budget_to_end_user(
|
||||
end_user_obj=return_obj,
|
||||
|
|
@ -1108,9 +1116,11 @@ async def get_end_user_object(
|
|||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
|
||||
# Save to cache (always store as dict for consistency)
|
||||
# Save to cache
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key="end_user_id:{}".format(end_user_id), value=_response.dict()
|
||||
key="end_user_id:{}".format(end_user_id),
|
||||
value=_response,
|
||||
model_type=LiteLLM_EndUserTable,
|
||||
)
|
||||
|
||||
# Check budget limits
|
||||
|
|
@ -1128,7 +1138,7 @@ async def get_end_user_object(
|
|||
async def get_tag_objects_batch(
|
||||
tag_names: List[str],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
) -> Dict[str, LiteLLM_TagTable]:
|
||||
|
|
@ -1161,12 +1171,12 @@ async def get_tag_objects_batch(
|
|||
# Try to get all tags from cache first
|
||||
for tag_name in tag_names:
|
||||
cache_key = f"tag:{tag_name}"
|
||||
cached_tag = await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
cached_tag = await user_api_key_cache.async_get_cache(
|
||||
key=cache_key,
|
||||
model_type=LiteLLM_TagTable,
|
||||
)
|
||||
if cached_tag is not None:
|
||||
if isinstance(cached_tag, dict):
|
||||
tag_objects[tag_name] = LiteLLM_TagTable(**cached_tag)
|
||||
else:
|
||||
tag_objects[tag_name] = cached_tag
|
||||
tag_objects[tag_name] = cached_tag
|
||||
else:
|
||||
uncached_tags.append(tag_name)
|
||||
|
||||
|
|
@ -1182,11 +1192,13 @@ async def get_tag_objects_batch(
|
|||
for db_tag in db_tags:
|
||||
tag_name = db_tag.tag_name
|
||||
cache_key = f"tag:{tag_name}"
|
||||
# Cache with default TTL (same as end_user objects)
|
||||
_tag_obj = LiteLLM_TagTable(**db_tag.dict())
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key, value=db_tag.dict()
|
||||
key=cache_key,
|
||||
value=_tag_obj,
|
||||
model_type=LiteLLM_TagTable,
|
||||
)
|
||||
tag_objects[tag_name] = LiteLLM_TagTable(**db_tag.dict())
|
||||
tag_objects[tag_name] = _tag_obj
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(f"Error batch fetching tags from database: {e}")
|
||||
|
||||
|
|
@ -1197,7 +1209,7 @@ async def get_tag_objects_batch(
|
|||
async def get_tag_object(
|
||||
tag_name: Optional[str],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
) -> Optional[LiteLLM_TagTable]:
|
||||
|
|
@ -1236,7 +1248,7 @@ async def get_team_membership(
|
|||
user_id: str,
|
||||
team_id: str,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
) -> Optional["LiteLLM_TeamMembership"]:
|
||||
|
|
@ -1256,9 +1268,12 @@ async def get_team_membership(
|
|||
_key = "team_membership:{}:{}".format(user_id, team_id)
|
||||
|
||||
# check if in cache
|
||||
cached_membership_obj = await user_api_key_cache.async_get_cache(key=_key)
|
||||
cached_membership_obj = await user_api_key_cache.async_get_cache(
|
||||
key=_key,
|
||||
model_type=LiteLLM_TeamMembership,
|
||||
)
|
||||
if cached_membership_obj is not None:
|
||||
return LiteLLM_TeamMembership(**cached_membership_obj)
|
||||
return cached_membership_obj
|
||||
|
||||
# else, check db
|
||||
try:
|
||||
|
|
@ -1270,10 +1285,12 @@ async def get_team_membership(
|
|||
if response is None:
|
||||
return None
|
||||
|
||||
# save the team membership object to cache (store as dict)
|
||||
await user_api_key_cache.async_set_cache(key=_key, value=response.dict())
|
||||
|
||||
_response = LiteLLM_TeamMembership(**response.dict())
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=_key,
|
||||
value=_response,
|
||||
model_type=LiteLLM_TeamMembership,
|
||||
)
|
||||
|
||||
return _response
|
||||
except Exception:
|
||||
|
|
@ -1441,7 +1458,7 @@ async def _get_fuzzy_user_object(
|
|||
async def get_user_object(
|
||||
user_id: Optional[str],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
user_id_upsert: bool,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
|
|
@ -1460,12 +1477,12 @@ async def get_user_object(
|
|||
|
||||
# check if in cache
|
||||
if not check_db_only:
|
||||
cached_user_obj = await user_api_key_cache.async_get_cache(key=user_id)
|
||||
cached_user_obj = await user_api_key_cache.async_get_cache(
|
||||
key=user_id,
|
||||
model_type=LiteLLM_UserTable,
|
||||
)
|
||||
if cached_user_obj is not None:
|
||||
if isinstance(cached_user_obj, dict):
|
||||
return LiteLLM_UserTable(**cached_user_obj)
|
||||
elif isinstance(cached_user_obj, LiteLLM_UserTable):
|
||||
return cached_user_obj
|
||||
return cached_user_obj
|
||||
# else, check db
|
||||
if prisma_client is None:
|
||||
raise Exception("No db connected")
|
||||
|
|
@ -1527,7 +1544,8 @@ async def get_user_object(
|
|||
# save the user object to cache
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=user_id,
|
||||
value=response_dict,
|
||||
value=_response,
|
||||
model_type=LiteLLM_UserTable,
|
||||
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
|
||||
)
|
||||
|
||||
|
|
@ -1548,13 +1566,21 @@ async def get_user_object(
|
|||
|
||||
async def _cache_management_object(
|
||||
key: str,
|
||||
value: BaseModel,
|
||||
user_api_key_cache: DualCache,
|
||||
value: Union[BaseModel, Dict[str, Any]],
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: Optional[ProxyLogging],
|
||||
*,
|
||||
model_type: Type[BaseModel],
|
||||
):
|
||||
"""
|
||||
Persist management objects via ``UserApiKeyCache`` (in-memory + optional Redis).
|
||||
|
||||
``UserApiKeyCache`` serializes with ``model_type`` so Redis and in-memory stay aligned.
|
||||
"""
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=key,
|
||||
value=value,
|
||||
model_type=model_type,
|
||||
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
|
||||
)
|
||||
|
||||
|
|
@ -1562,7 +1588,7 @@ async def _cache_management_object(
|
|||
async def _cache_team_object(
|
||||
team_id: str,
|
||||
team_table: LiteLLM_TeamTableCachedObj,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: Optional[ProxyLogging],
|
||||
):
|
||||
key = "team_id:{}".format(team_id)
|
||||
|
|
@ -1575,13 +1601,14 @@ async def _cache_team_object(
|
|||
value=team_table,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
model_type=LiteLLM_TeamTableCachedObj,
|
||||
)
|
||||
|
||||
|
||||
async def _cache_key_object(
|
||||
hashed_token: str,
|
||||
user_api_key_obj: UserAPIKeyAuth,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: Optional[ProxyLogging],
|
||||
):
|
||||
key = hashed_token
|
||||
|
|
@ -1594,12 +1621,13 @@ async def _cache_key_object(
|
|||
value=user_api_key_obj,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
model_type=UserAPIKeyAuth,
|
||||
)
|
||||
|
||||
|
||||
async def _delete_cache_key_object(
|
||||
hashed_token: str,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: Optional[ProxyLogging],
|
||||
):
|
||||
key = hashed_token
|
||||
|
|
@ -1647,7 +1675,7 @@ async def _get_team_object_from_db(team_id: str, prisma_client: PrismaClient):
|
|||
async def _get_team_object_from_user_api_key_cache(
|
||||
team_id: str,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
last_db_access_time: LimitedSizeOrderedDict,
|
||||
db_cache_expiry: int,
|
||||
proxy_logging_obj: Optional[ProxyLogging],
|
||||
|
|
@ -1708,38 +1736,38 @@ async def _get_team_object_from_user_api_key_cache(
|
|||
async def _get_team_object_from_cache(
|
||||
key: str,
|
||||
proxy_logging_obj: Optional[ProxyLogging],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span],
|
||||
) -> Optional[LiteLLM_TeamTableCachedObj]:
|
||||
cached_team_obj: Optional[LiteLLM_TeamTableCachedObj] = None
|
||||
|
||||
## CHECK REDIS CACHE ##
|
||||
## INTERNAL USAGE CACHE (plain DualCache) — checked before UserApiKeyCache stores ##
|
||||
if (
|
||||
proxy_logging_obj is not None
|
||||
and proxy_logging_obj.internal_usage_cache.dual_cache
|
||||
):
|
||||
cached_team_obj = (
|
||||
cached_raw = (
|
||||
await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache(
|
||||
key=key, parent_otel_span=parent_otel_span
|
||||
)
|
||||
)
|
||||
if cached_raw is not None:
|
||||
from_internal = CacheCodec.deserialize(
|
||||
cached_raw, LiteLLM_TeamTableCachedObj
|
||||
)
|
||||
if from_internal is not None:
|
||||
return from_internal
|
||||
|
||||
if cached_team_obj is None:
|
||||
cached_team_obj = await user_api_key_cache.async_get_cache(key=key)
|
||||
|
||||
if cached_team_obj is not None:
|
||||
if isinstance(cached_team_obj, dict):
|
||||
return LiteLLM_TeamTableCachedObj(**cached_team_obj)
|
||||
elif isinstance(cached_team_obj, LiteLLM_TeamTableCachedObj):
|
||||
return cached_team_obj
|
||||
|
||||
return None
|
||||
decoded = await user_api_key_cache.async_get_cache(
|
||||
key=key,
|
||||
parent_otel_span=parent_otel_span,
|
||||
model_type=LiteLLM_TeamTableCachedObj,
|
||||
)
|
||||
return decoded
|
||||
|
||||
|
||||
async def get_team_object(
|
||||
team_id: str,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
check_cache_only: Optional[bool] = None,
|
||||
|
|
@ -1805,20 +1833,21 @@ async def get_team_object(
|
|||
async def _cache_access_object(
|
||||
access_group_id: str,
|
||||
access_group_table: LiteLLM_AccessGroupTable,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
):
|
||||
key = "access_group_id:{}".format(access_group_id)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=key,
|
||||
value=access_group_table,
|
||||
model_type=LiteLLM_AccessGroupTable,
|
||||
ttl=DEFAULT_ACCESS_GROUP_CACHE_TTL,
|
||||
)
|
||||
|
||||
|
||||
async def _delete_cache_access_object(
|
||||
access_group_id: str,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
):
|
||||
key = "access_group_id:{}".format(access_group_id)
|
||||
|
|
@ -1836,7 +1865,7 @@ async def _delete_cache_access_object(
|
|||
async def get_access_object(
|
||||
access_group_id: str,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
) -> LiteLLM_AccessGroupTable:
|
||||
"""
|
||||
|
|
@ -1858,13 +1887,12 @@ async def get_access_object(
|
|||
|
||||
key = "access_group_id:{}".format(access_group_id)
|
||||
|
||||
# Always check cache first
|
||||
cached_access_obj = await user_api_key_cache.async_get_cache(key=key)
|
||||
cached_access_obj = await user_api_key_cache.async_get_cache(
|
||||
key=key,
|
||||
model_type=LiteLLM_AccessGroupTable,
|
||||
)
|
||||
if cached_access_obj is not None:
|
||||
if isinstance(cached_access_obj, dict):
|
||||
return LiteLLM_AccessGroupTable(**cached_access_obj)
|
||||
elif isinstance(cached_access_obj, LiteLLM_AccessGroupTable):
|
||||
return cached_access_obj
|
||||
return cached_access_obj
|
||||
|
||||
# Not in cache - fetch from DB
|
||||
try:
|
||||
|
|
@ -1910,7 +1938,7 @@ async def get_access_object(
|
|||
async def get_team_object_by_alias(
|
||||
team_alias: str,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional["Span"] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
) -> LiteLLM_TeamTableCachedObj:
|
||||
|
|
@ -1992,6 +2020,7 @@ async def get_team_object_by_alias(
|
|||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=team_obj,
|
||||
model_type=LiteLLM_TeamTableCachedObj,
|
||||
ttl=DEFAULT_IN_MEMORY_TTL,
|
||||
)
|
||||
# Also cache by team_id for consistency
|
||||
|
|
@ -1999,6 +2028,7 @@ async def get_team_object_by_alias(
|
|||
await user_api_key_cache.async_set_cache(
|
||||
key=team_id_cache_key,
|
||||
value=team_obj,
|
||||
model_type=LiteLLM_TeamTableCachedObj,
|
||||
ttl=DEFAULT_IN_MEMORY_TTL,
|
||||
)
|
||||
|
||||
|
|
@ -2020,7 +2050,7 @@ async def get_team_object_by_alias(
|
|||
async def get_org_object_by_alias(
|
||||
org_alias: str,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional["Span"] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
) -> Optional[LiteLLM_OrganizationTable]:
|
||||
|
|
@ -2047,12 +2077,12 @@ async def get_org_object_by_alias(
|
|||
|
||||
# Check cache first (keyed by alias)
|
||||
cache_key = "org_alias:{}".format(org_alias)
|
||||
cached_org_obj = await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
cached_org_obj = await user_api_key_cache.async_get_cache(
|
||||
key=cache_key,
|
||||
model_type=LiteLLM_OrganizationTable,
|
||||
)
|
||||
if cached_org_obj is not None:
|
||||
if isinstance(cached_org_obj, dict):
|
||||
return LiteLLM_OrganizationTable(**cached_org_obj)
|
||||
elif isinstance(cached_org_obj, LiteLLM_OrganizationTable):
|
||||
return cached_org_obj
|
||||
return cached_org_obj
|
||||
|
||||
# Query database by organization_alias
|
||||
try:
|
||||
|
|
@ -2082,13 +2112,15 @@ async def get_org_object_by_alias(
|
|||
# Cache the result
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=org_obj.model_dump(),
|
||||
value=org_obj,
|
||||
model_type=LiteLLM_OrganizationTable,
|
||||
ttl=DEFAULT_IN_MEMORY_TTL,
|
||||
)
|
||||
# Also cache by org_id for consistency
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key="org_id:{}".format(org_obj.organization_id),
|
||||
value=org_obj.model_dump(),
|
||||
value=org_obj,
|
||||
model_type=LiteLLM_OrganizationTable,
|
||||
ttl=DEFAULT_IN_MEMORY_TTL,
|
||||
)
|
||||
|
||||
|
|
@ -2291,7 +2323,7 @@ async def get_jwt_key_mapping_object(
|
|||
async def get_key_object(
|
||||
hashed_token: str,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
check_cache_only: Optional[bool] = None,
|
||||
|
|
@ -2309,15 +2341,14 @@ async def get_key_object(
|
|||
# check if in cache
|
||||
key = hashed_token
|
||||
|
||||
cached_key_obj: Optional[UserAPIKeyAuth] = await user_api_key_cache.async_get_cache(
|
||||
key=key
|
||||
# Same flow as before: use cache only when we have a hit we can turn into UserAPIKeyAuth
|
||||
# (dict from Redis / model_dump, or UserAPIKeyAuth from in-memory). Otherwise fall through to DB.
|
||||
user_api_key_auth = await user_api_key_cache.async_get_cache(
|
||||
key=key,
|
||||
model_type=UserAPIKeyAuth,
|
||||
)
|
||||
|
||||
if cached_key_obj is not None:
|
||||
if isinstance(cached_key_obj, dict):
|
||||
return UserAPIKeyAuth(**cached_key_obj)
|
||||
elif isinstance(cached_key_obj, UserAPIKeyAuth):
|
||||
return cached_key_obj
|
||||
if user_api_key_auth is not None:
|
||||
return user_api_key_auth
|
||||
|
||||
if check_cache_only:
|
||||
raise Exception(
|
||||
|
|
@ -2374,7 +2405,7 @@ async def get_key_object(
|
|||
async def get_object_permission(
|
||||
object_permission_id: str,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
) -> Optional[LiteLLM_ObjectPermissionTable]:
|
||||
|
|
@ -2390,12 +2421,12 @@ async def get_object_permission(
|
|||
|
||||
# check if in cache
|
||||
key = "object_permission_id:{}".format(object_permission_id)
|
||||
cached_obj_permission = await user_api_key_cache.async_get_cache(key=key)
|
||||
if cached_obj_permission is not None:
|
||||
if isinstance(cached_obj_permission, dict):
|
||||
return LiteLLM_ObjectPermissionTable(**cached_obj_permission)
|
||||
elif isinstance(cached_obj_permission, LiteLLM_ObjectPermissionTable):
|
||||
return cached_obj_permission
|
||||
deserialized_perm = await user_api_key_cache.async_get_cache(
|
||||
key=key,
|
||||
model_type=LiteLLM_ObjectPermissionTable,
|
||||
)
|
||||
if deserialized_perm is not None:
|
||||
return deserialized_perm
|
||||
|
||||
# else, check db
|
||||
try:
|
||||
|
|
@ -2406,14 +2437,15 @@ async def get_object_permission(
|
|||
if response is None:
|
||||
return None
|
||||
|
||||
# save the object permission to cache
|
||||
_perm_obj = LiteLLM_ObjectPermissionTable(**response.dict())
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=key,
|
||||
value=response.model_dump(),
|
||||
value=_perm_obj,
|
||||
model_type=LiteLLM_ObjectPermissionTable,
|
||||
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
|
||||
)
|
||||
|
||||
return LiteLLM_ObjectPermissionTable(**response.dict())
|
||||
return _perm_obj
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
|
@ -2422,7 +2454,7 @@ async def get_object_permission(
|
|||
async def get_managed_vector_store_rows_by_uuids(
|
||||
uuids: List[str],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
) -> List[LiteLLM_ManagedVectorStoresTable]:
|
||||
|
|
@ -2442,14 +2474,12 @@ async def get_managed_vector_store_rows_by_uuids(
|
|||
|
||||
for uuid in uuids:
|
||||
key = "managed_vector_store_id:{}".format(uuid)
|
||||
cached = await user_api_key_cache.async_get_cache(key=key)
|
||||
if cached is not None:
|
||||
if isinstance(cached, dict):
|
||||
result.append(LiteLLM_ManagedVectorStoresTable(**cached))
|
||||
elif isinstance(cached, LiteLLM_ManagedVectorStoresTable):
|
||||
result.append(cached)
|
||||
else:
|
||||
cache_misses.append(uuid)
|
||||
deserialized_vs = await user_api_key_cache.async_get_cache(
|
||||
key=key,
|
||||
model_type=LiteLLM_ManagedVectorStoresTable,
|
||||
)
|
||||
if deserialized_vs is not None:
|
||||
result.append(deserialized_vs)
|
||||
else:
|
||||
cache_misses.append(uuid)
|
||||
|
||||
|
|
@ -2475,7 +2505,8 @@ async def get_managed_vector_store_rows_by_uuids(
|
|||
key = "managed_vector_store_id:{}".format(cached_obj.vector_store_id)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=key,
|
||||
value=row_dict,
|
||||
value=cached_obj,
|
||||
model_type=LiteLLM_ManagedVectorStoresTable,
|
||||
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
|
||||
)
|
||||
result.append(cached_obj)
|
||||
|
|
@ -2487,7 +2518,7 @@ async def get_managed_vector_store_rows_by_uuids(
|
|||
async def get_org_object(
|
||||
org_id: str,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
include_budget_table: bool = False,
|
||||
|
|
@ -2518,12 +2549,12 @@ async def get_org_object(
|
|||
cache_key = "org_id:{}:with_budget".format(org_id)
|
||||
|
||||
# check if in cache
|
||||
cached_org_obj = user_api_key_cache.async_get_cache(key=cache_key)
|
||||
if cached_org_obj is not None:
|
||||
if isinstance(cached_org_obj, dict):
|
||||
return LiteLLM_OrganizationTable(**cached_org_obj)
|
||||
elif isinstance(cached_org_obj, LiteLLM_OrganizationTable):
|
||||
return cached_org_obj
|
||||
deserialized_org = await user_api_key_cache.async_get_cache(
|
||||
key=cache_key,
|
||||
model_type=LiteLLM_OrganizationTable,
|
||||
)
|
||||
if deserialized_org is not None:
|
||||
return deserialized_org
|
||||
# else, check db
|
||||
try:
|
||||
query_kwargs: Dict[str, Any] = {"where": {"organization_id": org_id}}
|
||||
|
|
@ -2537,16 +2568,16 @@ async def get_org_object(
|
|||
if response is None:
|
||||
raise Exception
|
||||
|
||||
_org_obj = LiteLLM_OrganizationTable(**response.model_dump())
|
||||
# Cache the result
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=(
|
||||
response.model_dump() if hasattr(response, "model_dump") else response
|
||||
),
|
||||
value=_org_obj,
|
||||
model_type=LiteLLM_OrganizationTable,
|
||||
ttl=DEFAULT_IN_MEMORY_TTL,
|
||||
)
|
||||
|
||||
return response
|
||||
return _org_obj
|
||||
except Exception:
|
||||
raise Exception(
|
||||
f"Organization doesn't exist in db. Organization={org_id}. Create organization via `/organization/new` call."
|
||||
|
|
@ -2559,7 +2590,7 @@ async def _get_resources_from_access_groups(
|
|||
"access_model_names", "access_mcp_server_ids", "access_agent_ids"
|
||||
],
|
||||
prisma_client: Optional[PrismaClient] = None,
|
||||
user_api_key_cache: Optional[DualCache] = None,
|
||||
user_api_key_cache: Optional[UserApiKeyCache] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
) -> List[str]:
|
||||
"""
|
||||
|
|
@ -2617,7 +2648,7 @@ async def _get_resources_from_access_groups(
|
|||
async def _get_models_from_access_groups(
|
||||
access_group_ids: List[str],
|
||||
prisma_client: Optional[PrismaClient] = None,
|
||||
user_api_key_cache: Optional[DualCache] = None,
|
||||
user_api_key_cache: Optional[UserApiKeyCache] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
) -> List[str]:
|
||||
"""
|
||||
|
|
@ -2636,7 +2667,7 @@ async def _get_models_from_access_groups(
|
|||
async def _get_mcp_server_ids_from_access_groups(
|
||||
access_group_ids: List[str],
|
||||
prisma_client: Optional[PrismaClient] = None,
|
||||
user_api_key_cache: Optional[DualCache] = None,
|
||||
user_api_key_cache: Optional[UserApiKeyCache] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
) -> List[str]:
|
||||
"""
|
||||
|
|
@ -2655,7 +2686,7 @@ async def _get_mcp_server_ids_from_access_groups(
|
|||
async def _get_agent_ids_from_access_groups(
|
||||
access_group_ids: List[str],
|
||||
prisma_client: Optional[PrismaClient] = None,
|
||||
user_api_key_cache: Optional[DualCache] = None,
|
||||
user_api_key_cache: Optional[UserApiKeyCache] = None,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
) -> List[str]:
|
||||
"""
|
||||
|
|
@ -3379,7 +3410,7 @@ async def _check_team_member_budget(
|
|||
user_object: Optional[LiteLLM_UserTable],
|
||||
valid_token: Optional[UserAPIKeyAuth],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
):
|
||||
"""Check if team member is over their max budget within the team."""
|
||||
|
|
@ -3447,7 +3478,7 @@ async def _check_team_member_model_access(
|
|||
valid_token: UserAPIKeyAuth,
|
||||
llm_router: Optional[Router],
|
||||
prisma_client: Optional["PrismaClient"],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> None:
|
||||
"""
|
||||
|
|
@ -3754,7 +3785,7 @@ async def _project_soft_budget_check(
|
|||
async def get_project_object(
|
||||
project_id: str,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: Optional[ProxyLogging] = None,
|
||||
) -> Optional[LiteLLM_ProjectTableCachedObj]:
|
||||
"""
|
||||
|
|
@ -3769,12 +3800,12 @@ async def get_project_object(
|
|||
|
||||
# Check cache first
|
||||
cache_key = "project_id:{}".format(project_id)
|
||||
cached_obj = await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
if cached_obj is not None:
|
||||
if isinstance(cached_obj, dict):
|
||||
return LiteLLM_ProjectTableCachedObj(**cached_obj)
|
||||
elif isinstance(cached_obj, LiteLLM_ProjectTableCachedObj):
|
||||
return cached_obj
|
||||
deserialized_project = await user_api_key_cache.async_get_cache(
|
||||
key=cache_key,
|
||||
model_type=LiteLLM_ProjectTableCachedObj,
|
||||
)
|
||||
if deserialized_project is not None:
|
||||
return deserialized_project
|
||||
|
||||
# Fetch from DB
|
||||
project_row = await prisma_client.db.litellm_projecttable.find_unique(
|
||||
|
|
@ -3793,6 +3824,7 @@ async def get_project_object(
|
|||
value=project_obj,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
model_type=LiteLLM_ProjectTableCachedObj,
|
||||
)
|
||||
|
||||
return project_obj
|
||||
|
|
@ -3802,7 +3834,7 @@ async def _organization_max_budget_check(
|
|||
valid_token: Optional[UserAPIKeyAuth],
|
||||
team_object: Optional[LiteLLM_TeamTable],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
):
|
||||
"""
|
||||
|
|
@ -3896,7 +3928,7 @@ async def _organization_max_budget_check(
|
|||
async def _tag_max_budget_check(
|
||||
request_body: dict,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
valid_token: Optional[UserAPIKeyAuth],
|
||||
):
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ Currently only supports admin.
|
|||
JWT token must have 'litellm_proxy_admin' in scope.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import fnmatch
|
||||
import hashlib
|
||||
import os
|
||||
|
|
@ -20,7 +22,6 @@ import jwt
|
|||
from jwt.api_jwk import PyJWK
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
|
||||
from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
|
||||
from litellm.llms.custom_httpx.httpx_handler import HTTPHandler
|
||||
|
|
@ -46,6 +47,7 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.proxy.auth.auth_checks import can_team_access_model
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
|
||||
from .auth_checks import (
|
||||
|
|
@ -73,7 +75,7 @@ class JWTHandler:
|
|||
"""
|
||||
|
||||
prisma_client: Optional[PrismaClient]
|
||||
user_api_key_cache: DualCache
|
||||
user_api_key_cache: UserApiKeyCache
|
||||
# Supported algos: https://pyjwt.readthedocs.io/en/stable/algorithms.html
|
||||
# "Warning: Make sure not to mix symmetric and asymmetric algorithms that interpret
|
||||
# the key in different ways (e.g. HS* and RS*)."
|
||||
|
|
@ -99,7 +101,7 @@ class JWTHandler:
|
|||
def update_environment(
|
||||
self,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
litellm_jwtauth: LiteLLM_JWTAuth,
|
||||
leeway: int = 0,
|
||||
) -> None:
|
||||
|
|
@ -952,7 +954,7 @@ class JWTAuthManager:
|
|||
jwt_handler: JWTHandler,
|
||||
jwt_valid_token: dict,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span],
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]:
|
||||
|
|
@ -1045,7 +1047,7 @@ class JWTAuthManager:
|
|||
route: str,
|
||||
jwt_handler: JWTHandler,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span],
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]:
|
||||
|
|
@ -1133,7 +1135,7 @@ class JWTAuthManager:
|
|||
valid_user_email: Optional[bool],
|
||||
jwt_handler: JWTHandler,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span],
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
route: str,
|
||||
|
|
@ -1349,7 +1351,7 @@ class JWTAuthManager:
|
|||
jwt_valid_token: dict,
|
||||
user_object: Optional[LiteLLM_UserTable],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: Optional[DualCache] = None,
|
||||
user_api_key_cache: Optional[UserApiKeyCache] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Sync user role and team memberships with JWT claims
|
||||
|
|
@ -1377,7 +1379,8 @@ class JWTAuthManager:
|
|||
if user_api_key_cache is not None:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=user_object.user_id,
|
||||
value=user_object.model_dump(),
|
||||
value=user_object,
|
||||
model_type=LiteLLM_UserTable,
|
||||
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
|
||||
)
|
||||
|
||||
|
|
@ -1400,7 +1403,8 @@ class JWTAuthManager:
|
|||
if user_api_key_cache is not None:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=user_object.user_id,
|
||||
value=user_object.model_dump(),
|
||||
value=user_object,
|
||||
model_type=LiteLLM_UserTable,
|
||||
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
|
||||
)
|
||||
return None
|
||||
|
|
@ -1412,7 +1416,7 @@ class JWTAuthManager:
|
|||
request_headers: Optional[dict],
|
||||
jwt_handler: JWTHandler,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span],
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> None:
|
||||
|
|
@ -1456,7 +1460,7 @@ class JWTAuthManager:
|
|||
user_object: Optional[LiteLLM_UserTable],
|
||||
user_id: Optional[str],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span],
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
team_id_upsert: Optional[bool],
|
||||
|
|
@ -1514,7 +1518,7 @@ class JWTAuthManager:
|
|||
general_settings: dict,
|
||||
route: str,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span],
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
request_headers: Optional[dict] = None,
|
||||
|
|
|
|||
|
|
@ -20,7 +20,6 @@ from fastapi.security.api_key import APIKeyHeader
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger, verbose_proxy_logger
|
||||
from litellm._service_logger import ServiceLogging
|
||||
from litellm.caching import DualCache
|
||||
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
|
||||
from litellm.litellm_core_utils.dd_tracing import tracer
|
||||
from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
|
||||
|
|
@ -60,6 +59,7 @@ from litellm.proxy.auth.oauth2_check import Oauth2Handler
|
|||
from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.common_utils.http_parsing_utils import (
|
||||
_read_request_body,
|
||||
_safe_get_request_headers,
|
||||
|
|
@ -329,7 +329,7 @@ _global_spend_coordinator = EventDrivenCacheCoordinator(log_prefix="[GLOBAL SPEN
|
|||
|
||||
async def _fetch_global_spend_with_event_coordination(
|
||||
cache_key: str,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
prisma_client: PrismaClient,
|
||||
) -> Optional[float]:
|
||||
"""
|
||||
|
|
@ -345,14 +345,14 @@ async def _fetch_global_spend_with_event_coordination(
|
|||
|
||||
return await _global_spend_coordinator.get_or_load(
|
||||
cache_key=cache_key,
|
||||
cache=user_api_key_cache,
|
||||
cache=user_api_key_cache, # pyright: ignore[reportArgumentType]
|
||||
load_fn=_load_global_spend,
|
||||
)
|
||||
|
||||
|
||||
async def get_global_proxy_spend(
|
||||
litellm_proxy_admin_name: str,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
token: str,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
|
|
@ -510,7 +510,7 @@ async def _resolve_jwt_to_virtual_key(
|
|||
jwt_claims: dict,
|
||||
jwt_handler: JWTHandler,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Optional[Span],
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> Optional[UserAPIKeyAuth]:
|
||||
|
|
@ -1112,9 +1112,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
|
|||
is_master_key_valid = False
|
||||
|
||||
## VALIDATE MASTER KEY ##
|
||||
try:
|
||||
assert isinstance(master_key, str)
|
||||
except Exception:
|
||||
if not isinstance(master_key, str):
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={
|
||||
|
|
@ -1184,11 +1182,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
|
|||
if len(api_key) > 8
|
||||
else "****"
|
||||
)
|
||||
assert api_key.startswith(
|
||||
"sk-"
|
||||
), "LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format(
|
||||
_masked_key
|
||||
) # prevent token hashes from being used
|
||||
if not api_key.startswith("sk-"):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=(
|
||||
"LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format(
|
||||
_masked_key
|
||||
)
|
||||
),
|
||||
) # prevent token hashes from being used
|
||||
else:
|
||||
verbose_logger.warning(
|
||||
"litellm.proxy.proxy_server.user_api_key_auth(): Warning - Key is not a string. Got type={}".format(
|
||||
|
|
@ -1296,7 +1298,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
|
|||
_cache_key = f"{valid_token.team_id}_{valid_token.user_id}"
|
||||
|
||||
team_member_info = await user_api_key_cache.async_get_cache(
|
||||
key=_cache_key
|
||||
key=_cache_key,
|
||||
model_type=LiteLLM_TeamMembership,
|
||||
)
|
||||
if team_member_info is None:
|
||||
# read from DB
|
||||
|
|
@ -1304,18 +1307,23 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
|
|||
_team_id = valid_token.team_id
|
||||
|
||||
if _user_id is not None and _team_id is not None:
|
||||
team_member_info = await prisma_client.db.litellm_teammembership.find_first(
|
||||
_db_member = await prisma_client.db.litellm_teammembership.find_first(
|
||||
where={
|
||||
"user_id": _user_id,
|
||||
"team_id": _team_id,
|
||||
}, # type: ignore
|
||||
include={"litellm_budget_table": True},
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=_cache_key,
|
||||
value=team_member_info,
|
||||
ttl=5,
|
||||
)
|
||||
if _db_member is not None:
|
||||
team_member_info = LiteLLM_TeamMembership(
|
||||
**_db_member.dict()
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=_cache_key,
|
||||
value=team_member_info,
|
||||
model_type=LiteLLM_TeamMembership,
|
||||
ttl=5,
|
||||
)
|
||||
|
||||
if (
|
||||
team_member_info is not None
|
||||
|
|
@ -1462,9 +1470,13 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
|
|||
else:
|
||||
valid_token.team_object_permission = None
|
||||
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=valid_token.team_id, value=_team_obj
|
||||
) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py
|
||||
# Only cache when the key is a real team_id (non-team keys must not use key=None).
|
||||
if valid_token.team_id is not None and _team_obj is not None:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=valid_token.team_id,
|
||||
value=_team_obj,
|
||||
model_type=LiteLLM_TeamTableCachedObj,
|
||||
) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py
|
||||
|
||||
# Fetch project object if key belongs to a project
|
||||
_project_obj = None
|
||||
|
|
|
|||
|
|
@ -53,12 +53,16 @@ def clear_token() -> None:
|
|||
os.remove(token_file)
|
||||
|
||||
|
||||
def get_stored_api_key() -> Optional[str]:
|
||||
"""Get the stored API key from token file"""
|
||||
# Use the SDK-level utility
|
||||
def get_stored_api_key(expected_base_url: Optional[str] = None) -> Optional[str]:
|
||||
"""Get the stored API key from token file.
|
||||
|
||||
If expected_base_url is provided, the key is only returned when it was
|
||||
originally issued for that URL. This prevents credential leakage when the
|
||||
CLI is pointed at a different (possibly malicious) server.
|
||||
"""
|
||||
from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_key
|
||||
|
||||
return get_litellm_gateway_api_key()
|
||||
return get_litellm_gateway_api_key(expected_base_url=expected_base_url)
|
||||
|
||||
|
||||
# Team selection utilities
|
||||
|
|
@ -572,9 +576,11 @@ def login(ctx: click.Context):
|
|||
api_key = auth_result["api_key"]
|
||||
user_id = auth_result["user_id"]
|
||||
|
||||
# Save token data (simplified for CLI - we just need the key)
|
||||
# Save token data. base_url is stored so we can verify origin
|
||||
# before reusing the key on a subsequent CLI invocation.
|
||||
save_token(
|
||||
{
|
||||
"base_url": base_url.rstrip("/"),
|
||||
"key": api_key,
|
||||
"user_id": user_id or "cli-user",
|
||||
"user_email": "unknown",
|
||||
|
|
|
|||
|
|
@ -74,9 +74,10 @@ def cli(ctx: click.Context, base_url: str, api_key: Optional[str]) -> None:
|
|||
"""LiteLLM Proxy CLI - Manage your LiteLLM proxy server"""
|
||||
ctx.ensure_object(dict)
|
||||
|
||||
# If no API key provided via flag or environment variable, try to load from saved token
|
||||
# If no API key provided via flag or environment variable, try to load from saved token.
|
||||
# Pass base_url so we only use the stored key when it was issued for this server.
|
||||
if api_key is None:
|
||||
api_key = get_stored_api_key()
|
||||
api_key = get_stored_api_key(expected_base_url=base_url)
|
||||
|
||||
ctx.obj["base_url"] = base_url
|
||||
ctx.obj["api_key"] = api_key
|
||||
|
|
|
|||
|
|
@ -28,12 +28,17 @@ class Client:
|
|||
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
|
||||
timeout: Request timeout in seconds (default: 30)
|
||||
"""
|
||||
self._base_url = base_url.rstrip("/") # Remove trailing slash if present
|
||||
self._api_key = get_litellm_gateway_api_key() or api_key
|
||||
self._base_url = base_url.rstrip("/")
|
||||
# Only use the stored CLI key when it was issued for this server.
|
||||
self._api_key = api_key or get_litellm_gateway_api_key(
|
||||
expected_base_url=self._base_url
|
||||
)
|
||||
|
||||
# Initialize resource clients
|
||||
|
||||
self.http = HTTPClient(base_url=base_url, api_key=api_key, timeout=timeout)
|
||||
self.http = HTTPClient(
|
||||
base_url=base_url, api_key=self._api_key, timeout=timeout
|
||||
)
|
||||
self.models = ModelsManagementClient(
|
||||
base_url=self._base_url, api_key=self._api_key
|
||||
)
|
||||
|
|
|
|||
|
|
@ -744,6 +744,11 @@ class ProxyBaseLLMRequestProcessing:
|
|||
"aingest",
|
||||
"aretrieve_container",
|
||||
"adelete_container",
|
||||
"aupload_container_file",
|
||||
"alist_container_files",
|
||||
"aretrieve_container_file",
|
||||
"adelete_container_file",
|
||||
"aretrieve_container_file_content",
|
||||
"acreate_skill",
|
||||
"alist_skills",
|
||||
"aget_skill",
|
||||
|
|
@ -1001,6 +1006,11 @@ class ProxyBaseLLMRequestProcessing:
|
|||
"aingest",
|
||||
"aretrieve_container",
|
||||
"adelete_container",
|
||||
"aupload_container_file",
|
||||
"alist_container_files",
|
||||
"aretrieve_container_file",
|
||||
"adelete_container_file",
|
||||
"aretrieve_container_file_content",
|
||||
"acreate_skill",
|
||||
"alist_skills",
|
||||
"aget_skill",
|
||||
|
|
|
|||
|
|
@ -20,11 +20,27 @@ T = TypeVar("T")
|
|||
|
||||
|
||||
class AsyncCacheProtocol(Protocol):
|
||||
"""Protocol for cache backends used by EventDrivenCacheCoordinator."""
|
||||
"""Protocol for cache backends used by EventDrivenCacheCoordinator.
|
||||
|
||||
async def async_get_cache(self, key: str, **kwargs: Any) -> Any: ...
|
||||
Matches ``DualCache`` / ``UserApiKeyCache`` call shapes (explicit optional params
|
||||
before ``**kwargs``), not only ``(key, **kwargs)``, so overloads validate.
|
||||
"""
|
||||
|
||||
async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> Any: ...
|
||||
async def async_get_cache(
|
||||
self,
|
||||
key: str,
|
||||
parent_otel_span: Any = None,
|
||||
local_only: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> Any: ...
|
||||
|
||||
async def async_set_cache(
|
||||
self,
|
||||
key: str,
|
||||
value: Any,
|
||||
local_only: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> Any: ...
|
||||
|
||||
|
||||
class EventDrivenCacheCoordinator:
|
||||
|
|
@ -36,6 +52,9 @@ class EventDrivenCacheCoordinator:
|
|||
- Other requests: wait for the signal, then read from cache.
|
||||
|
||||
Create one instance per resource (e.g. one for global spend, one for feature flags).
|
||||
|
||||
Args:
|
||||
log_prefix: Prefix for debug log messages.
|
||||
"""
|
||||
|
||||
def __init__(self, log_prefix: str = "[CACHE]"):
|
||||
|
|
|
|||
93
litellm/proxy/common_utils/cache_pydantic_utils.py
Normal file
93
litellm/proxy/common_utils/cache_pydantic_utils.py
Normal file
|
|
@ -0,0 +1,93 @@
|
|||
"""
|
||||
DualCache presents a single API for reads and writes, but the two backends behave
|
||||
differently: the in-memory layer can store arbitrary Python objects (including live
|
||||
``BaseModel`` instances), while Redis persists strings and therefore needs JSON-safe
|
||||
payloads (``json.dumps`` on the Redis side).
|
||||
|
||||
Call sites therefore see cache ``value`` / ``cached`` as effectively ``Any``: the same
|
||||
key may deserialize to a model on one process (memory hit) or to a ``dict`` after a
|
||||
Redis round-trip. ``CacheCodec`` centralizes encode/decode at that boundary:
|
||||
``CacheCodec.serialize`` before ``set``, ``CacheCodec.deserialize`` after ``get``
|
||||
when you need a typed ``BaseModel``.
|
||||
|
||||
``dataclasses`` are not supported: only ``dict`` and Pydantic ``BaseModel`` inputs
|
||||
are encoded; pass a Pydantic model or convert with e.g. ``dataclasses.asdict`` first.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Optional, Type, TypeVar
|
||||
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
T = TypeVar("T", bound=BaseModel)
|
||||
|
||||
|
||||
class CacheCodec:
|
||||
"""
|
||||
Encode/decode Pydantic models for DualCache (memory vs Redis safe payloads).
|
||||
|
||||
Dataclasses are not supported yet (only ``dict`` and ``BaseModel``).
|
||||
|
||||
Use ``serialize`` with ``model_type`` when writing so the same schema is used
|
||||
as on read (``deserialize``). Pass ``model_type`` whenever you know it
|
||||
(validates ``dict`` payloads and normalizes ``BaseModel`` instances).
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def serialize(value: Any, model_type: Optional[Type[T]] = None) -> Any:
|
||||
"""
|
||||
Encode a value for DualCache / Redis (``json.dumps``-safe).
|
||||
|
||||
If ``model_type`` is set, the payload is validated with that model, then
|
||||
``model_dump(mode="json", exclude_none=True)`` — symmetric with ``deserialize``.
|
||||
|
||||
If the value is already an instance of ``model_type`` (or a subclass),
|
||||
``model_validate`` is skipped to avoid an unnecessary Pydantic copy — the
|
||||
value is dumped directly.
|
||||
|
||||
If ``model_type`` is omitted, any ``BaseModel`` is dumped as above; other
|
||||
values (e.g. plain ``dict``) are returned unchanged.
|
||||
"""
|
||||
if model_type is not None:
|
||||
if isinstance(value, model_type):
|
||||
# Already the right type: dump directly, skip re-validation copy.
|
||||
return value.model_dump(mode="json", exclude_none=True)
|
||||
if isinstance(value, (dict, BaseModel)):
|
||||
return model_type.model_validate(value).model_dump(
|
||||
mode="json", exclude_none=True
|
||||
)
|
||||
return value
|
||||
if isinstance(value, BaseModel):
|
||||
return value.model_dump(mode="json", exclude_none=True)
|
||||
return value
|
||||
|
||||
@staticmethod
|
||||
def deserialize(cached: Any, model_type: Type[T]) -> Optional[T]:
|
||||
"""
|
||||
Decode a cache entry to ``model_type``.
|
||||
|
||||
- ``None`` → ``None``
|
||||
- Already an instance of ``model_type`` (including subclasses) → returned as-is
|
||||
- ``dict`` → ``model_type.model_validate(...)``; on ``ValidationError``,
|
||||
logs a warning and returns ``None`` (treat as cache miss; avoids serving
|
||||
malformed or schema-drifted entries)
|
||||
- Any other type → ``None`` (caller should treat as cache miss or log)
|
||||
"""
|
||||
if cached is None:
|
||||
return None
|
||||
if isinstance(cached, model_type):
|
||||
return cached
|
||||
if isinstance(cached, dict):
|
||||
try:
|
||||
return model_type.model_validate(cached)
|
||||
except ValidationError as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"CacheCodec.deserialize: validation failed for %s (%s)",
|
||||
model_type.__name__,
|
||||
e,
|
||||
)
|
||||
return None
|
||||
return None
|
||||
|
|
@ -8,7 +8,7 @@ from datetime import datetime, timezone
|
|||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching import DualCache
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.constants import (
|
||||
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
|
||||
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE,
|
||||
|
|
@ -31,7 +31,7 @@ class ExpiredUISessionKeyCleanupManager:
|
|||
def __init__(
|
||||
self,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
pod_lock_manager=None,
|
||||
):
|
||||
self.prisma_client = prisma_client
|
||||
|
|
|
|||
162
litellm/proxy/common_utils/user_api_key_cache.py
Normal file
162
litellm/proxy/common_utils/user_api_key_cache.py
Normal file
|
|
@ -0,0 +1,162 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Optional, Type, TypeVar, Union, cast, overload
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
|
||||
|
||||
T = TypeVar("T", bound=BaseModel)
|
||||
|
||||
|
||||
class UserApiKeyCache(DualCache):
|
||||
"""
|
||||
DualCache wrapper for UserAPIKeyAuth-like payloads.
|
||||
|
||||
Stores a Redis-safe JSON payload in BOTH in-memory and Redis to avoid
|
||||
"memory returns BaseModel, Redis returns dict" format drift.
|
||||
|
||||
When ``model_type`` is provided:
|
||||
- writes are serialized via ``CacheCodec.serialize(..., model_type=...)``
|
||||
- reads are deserialized via ``CacheCodec.deserialize(..., model_type)``
|
||||
and return ``Optional[T]``: the model on success, ``None`` on cache miss
|
||||
**or** if the cached payload fails validation (schema drift). On
|
||||
validation failure after a cache hit, an error line is emitted via
|
||||
``verbose_proxy_logger``.
|
||||
|
||||
When ``model_type`` is omitted, the interface behaves like ``DualCache``:
|
||||
raw cached payload is returned (dict/str/etc.).
|
||||
|
||||
``async_set_cache_pipeline`` applies the same untyped Codec pass as omitting
|
||||
``model_type`` on ``async_set_cache`` (so ``BaseModel`` rows are dumped before Redis).
|
||||
|
||||
``get_cache`` / ``async_get_cache`` overloads and implementations must be contiguous
|
||||
(no other methods in between) so mypy resolves ``@overload`` + implementation correctly.
|
||||
"""
|
||||
|
||||
@overload
|
||||
def get_cache(
|
||||
self,
|
||||
key: Any,
|
||||
parent_otel_span: Any = None,
|
||||
local_only: bool = False,
|
||||
*,
|
||||
model_type: Type[T],
|
||||
**kwargs: Any,
|
||||
) -> Optional[T]: ...
|
||||
|
||||
@overload
|
||||
def get_cache(
|
||||
self,
|
||||
key: Any,
|
||||
parent_otel_span: Any = None,
|
||||
local_only: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> Any: ...
|
||||
|
||||
def get_cache( # type: ignore[override]
|
||||
self,
|
||||
key,
|
||||
parent_otel_span=None,
|
||||
local_only: bool = False,
|
||||
model_type: Optional[Type[BaseModel]] = None,
|
||||
**kwargs,
|
||||
) -> Union[Any, Optional[BaseModel]]:
|
||||
if model_type is None and "model_type" in kwargs:
|
||||
model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None))
|
||||
cached = super().get_cache(
|
||||
key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs
|
||||
)
|
||||
if model_type is None:
|
||||
return cached
|
||||
if cached is None:
|
||||
return None
|
||||
decoded = CacheCodec.deserialize(cached, model_type=model_type)
|
||||
if decoded is None:
|
||||
verbose_proxy_logger.error(
|
||||
"UserApiKeyCache.get_cache failed to deserialize cached value for "
|
||||
"key=%r model_type=%s",
|
||||
key,
|
||||
getattr(model_type, "__name__", str(model_type)),
|
||||
)
|
||||
return None
|
||||
return decoded
|
||||
|
||||
@overload
|
||||
async def async_get_cache(
|
||||
self,
|
||||
key: Any,
|
||||
parent_otel_span: Any = None,
|
||||
local_only: bool = False,
|
||||
*,
|
||||
model_type: Type[T],
|
||||
**kwargs: Any,
|
||||
) -> Optional[T]: ...
|
||||
|
||||
@overload
|
||||
async def async_get_cache(
|
||||
self,
|
||||
key: Any,
|
||||
parent_otel_span: Any = None,
|
||||
local_only: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> Any: ...
|
||||
|
||||
async def async_get_cache( # type: ignore[override]
|
||||
self,
|
||||
key,
|
||||
parent_otel_span=None,
|
||||
local_only: bool = False,
|
||||
model_type: Optional[Type[BaseModel]] = None,
|
||||
**kwargs,
|
||||
) -> Union[Any, Optional[BaseModel]]:
|
||||
if model_type is None and "model_type" in kwargs:
|
||||
model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None))
|
||||
cached = await super().async_get_cache(
|
||||
key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs
|
||||
)
|
||||
if model_type is None:
|
||||
return cached
|
||||
if cached is None:
|
||||
return None
|
||||
decoded = CacheCodec.deserialize(cached, model_type=model_type)
|
||||
if decoded is None:
|
||||
verbose_proxy_logger.error(
|
||||
"UserApiKeyCache.async_get_cache failed to deserialize cached value for "
|
||||
"key=%r model_type=%s",
|
||||
key,
|
||||
getattr(model_type, "__name__", str(model_type)),
|
||||
)
|
||||
return None
|
||||
return decoded
|
||||
|
||||
def set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override]
|
||||
model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None))
|
||||
payload = CacheCodec.serialize(value, model_type=model_type)
|
||||
return super().set_cache(
|
||||
key=key, value=payload, local_only=local_only, **kwargs
|
||||
)
|
||||
|
||||
async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override]
|
||||
model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None))
|
||||
payload = CacheCodec.serialize(value, model_type=model_type)
|
||||
return await super().async_set_cache(
|
||||
key=key, value=payload, local_only=local_only, **kwargs
|
||||
)
|
||||
|
||||
async def async_set_cache_pipeline( # type: ignore[override]
|
||||
self, cache_list: list, local_only: bool = False, **kwargs
|
||||
) -> None:
|
||||
"""
|
||||
Batch writes with the same Codec boundary as ``async_set_cache`` without
|
||||
``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged.
|
||||
"""
|
||||
normalized = [
|
||||
(key, CacheCodec.serialize(value, model_type=None))
|
||||
for key, value in cache_list
|
||||
]
|
||||
return await super().async_set_cache_pipeline(
|
||||
cache_list=normalized, local_only=local_only, **kwargs
|
||||
)
|
||||
|
|
@ -19,7 +19,6 @@ from litellm.proxy.common_utils.openai_endpoint_utils import (
|
|||
get_custom_llm_provider_from_request_headers,
|
||||
get_custom_llm_provider_from_request_query,
|
||||
)
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
|
||||
|
||||
def _load_endpoints_config() -> Dict:
|
||||
|
|
@ -64,10 +63,12 @@ def _create_handler_for_path_params(
|
|||
request: Request,
|
||||
container_id: str,
|
||||
file_id: str,
|
||||
fastapi_response: Response,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
return await _process_binary_request(
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
container_id=container_id,
|
||||
file_id=file_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -152,63 +153,61 @@ def _create_handler_for_path_params(
|
|||
|
||||
async def _process_binary_request(
|
||||
request: Request,
|
||||
fastapi_response: Response,
|
||||
container_id: str,
|
||||
file_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
):
|
||||
"""
|
||||
Process binary content requests using the proper transformation pattern.
|
||||
Process binary content requests through the standard proxy/router pipeline.
|
||||
|
||||
This uses the provider config transformations and llm_http_handler
|
||||
to maintain consistency with the established pattern.
|
||||
The router owns managed container ID decoding and deployment selection. This
|
||||
handler only adapts the byte response to FastAPI.
|
||||
"""
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.proxy.proxy_server import (
|
||||
general_settings,
|
||||
llm_router,
|
||||
proxy_config,
|
||||
proxy_logging_obj,
|
||||
select_data_generator,
|
||||
user_api_base,
|
||||
user_max_tokens,
|
||||
user_model,
|
||||
user_request_timeout,
|
||||
user_temperature,
|
||||
version,
|
||||
)
|
||||
|
||||
# Extract custom_llm_provider
|
||||
custom_llm_provider = (
|
||||
get_custom_llm_provider_from_request_headers(request=request)
|
||||
or get_custom_llm_provider_from_request_query(request=request)
|
||||
or "openai"
|
||||
)
|
||||
|
||||
# Build litellm_params - credentials are resolved by provider config from env
|
||||
litellm_params = GenericLiteLLMParams()
|
||||
|
||||
# Decode container ID and extract provider info
|
||||
decoded = ResponsesAPIRequestUtils._decode_container_id(container_id)
|
||||
original_container_id = decoded.get("response_id", container_id)
|
||||
|
||||
# If container ID has encoded provider info and user didn't explicitly set provider, use it
|
||||
decoded_provider = decoded.get("custom_llm_provider")
|
||||
if decoded_provider and custom_llm_provider == "openai":
|
||||
custom_llm_provider = decoded_provider
|
||||
|
||||
# Get the provider config
|
||||
container_provider_config = _get_container_provider_config(custom_llm_provider)
|
||||
|
||||
# Create logging object
|
||||
logging_obj = Logging(
|
||||
model="container-file-content",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="container_file_content",
|
||||
start_time=None,
|
||||
litellm_call_id="",
|
||||
function_id="",
|
||||
)
|
||||
|
||||
# Use the HTTP handler to make the request
|
||||
handler = BaseLLMHTTPHandler()
|
||||
data: Dict[str, Any] = {
|
||||
"container_id": container_id,
|
||||
"file_id": file_id,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
}
|
||||
processor = ProxyBaseLLMRequestProcessing(data=data)
|
||||
|
||||
try:
|
||||
content = await handler.async_container_file_content_handler(
|
||||
container_id=original_container_id, # Use decoded original ID
|
||||
file_id=file_id,
|
||||
container_provider_config=container_provider_config,
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=logging_obj,
|
||||
content = await processor.base_process_llm_request(
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
route_type="aretrieve_container_file_content",
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
llm_router=llm_router,
|
||||
general_settings=general_settings,
|
||||
proxy_config=proxy_config,
|
||||
select_data_generator=select_data_generator,
|
||||
model=None,
|
||||
user_model=user_model,
|
||||
user_temperature=user_temperature,
|
||||
user_request_timeout=user_request_timeout,
|
||||
user_max_tokens=user_max_tokens,
|
||||
user_api_base=user_api_base,
|
||||
version=version,
|
||||
)
|
||||
|
||||
# Determine content type based on common file extensions in the file_id
|
||||
|
|
@ -229,13 +228,25 @@ async def _process_binary_request(
|
|||
elif ".pdf" in file_id_lower:
|
||||
content_type = "application/pdf"
|
||||
|
||||
if not isinstance(content, bytes):
|
||||
raise TypeError(
|
||||
"aretrieve_container_file_content expected bytes, got "
|
||||
f"{type(content).__name__}"
|
||||
)
|
||||
|
||||
return Response(
|
||||
content=content,
|
||||
headers=dict(fastapi_response.headers),
|
||||
media_type=content_type,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise e
|
||||
raise await processor._handle_llm_api_exception(
|
||||
e=e,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
version=version,
|
||||
)
|
||||
|
||||
|
||||
async def _process_multipart_upload_request(
|
||||
|
|
@ -284,16 +295,7 @@ async def _process_multipart_upload_request(
|
|||
or "openai"
|
||||
)
|
||||
|
||||
# Decode container ID and extract provider info
|
||||
decoded = ResponsesAPIRequestUtils._decode_container_id(container_id)
|
||||
original_container_id = decoded.get("response_id", container_id)
|
||||
|
||||
# If container ID has encoded provider info and user didn't explicitly set provider, use it
|
||||
decoded_provider = decoded.get("custom_llm_provider")
|
||||
if decoded_provider and custom_llm_provider == "openai":
|
||||
custom_llm_provider = decoded_provider
|
||||
|
||||
data["container_id"] = original_container_id # Use decoded original ID
|
||||
data["container_id"] = container_id
|
||||
data["custom_llm_provider"] = custom_llm_provider
|
||||
|
||||
processor = ProxyBaseLLMRequestProcessing(data=data)
|
||||
|
|
@ -359,21 +361,6 @@ async def _process_request(
|
|||
or "openai"
|
||||
)
|
||||
|
||||
# Decode container_id if present in path_params
|
||||
if "container_id" in path_params:
|
||||
decoded = ResponsesAPIRequestUtils._decode_container_id(
|
||||
path_params["container_id"]
|
||||
)
|
||||
original_container_id = decoded.get("response_id", path_params["container_id"])
|
||||
|
||||
# If container ID has encoded provider info and user didn't explicitly set provider, use it
|
||||
decoded_provider = decoded.get("custom_llm_provider")
|
||||
if decoded_provider and custom_llm_provider == "openai":
|
||||
custom_llm_provider = decoded_provider
|
||||
|
||||
# Update path_params with decoded original ID
|
||||
data["container_id"] = original_container_id
|
||||
|
||||
data["custom_llm_provider"] = custom_llm_provider
|
||||
|
||||
processor = ProxyBaseLLMRequestProcessing(data=data)
|
||||
|
|
|
|||
|
|
@ -29,6 +29,10 @@ ILLEGAL_DISPLAY_PARAMS = [
|
|||
"exception", # internal; not JSON-serializable, never for display
|
||||
"litellm_metadata", # internal tracking metadata with auth objects; not for display
|
||||
]
|
||||
# Provider routing fields. Allowed for proxy admins so they can see which
|
||||
# region/version a deployment is checking; gated at the endpoint layer for
|
||||
# non-admin callers (see _strip_admin_only_fields_from_health_result).
|
||||
ADMIN_ONLY_HEALTH_DISPLAY_PARAMS = ("api_base", "api_version")
|
||||
|
||||
MINIMAL_DISPLAY_PARAMS = ["model", "mode_error"]
|
||||
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from litellm.proxy._types import (
|
|||
CallInfo,
|
||||
EnterpriseLicenseData,
|
||||
Litellm_EntityType,
|
||||
LitellmUserRoles,
|
||||
ProxyErrorTypes,
|
||||
ProxyException,
|
||||
UserAPIKeyAuth,
|
||||
|
|
@ -28,6 +29,7 @@ from litellm.proxy._types import (
|
|||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.health_check import (
|
||||
ADMIN_ONLY_HEALTH_DISPLAY_PARAMS,
|
||||
_clean_endpoint_data,
|
||||
_update_litellm_params_for_health_check,
|
||||
perform_health_check,
|
||||
|
|
@ -723,6 +725,90 @@ async def _save_background_health_checks_to_db(
|
|||
# Continue execution - don't let database save failure break health checks
|
||||
|
||||
|
||||
_PROXY_ADMIN_ROLES = frozenset(
|
||||
{
|
||||
LitellmUserRoles.PROXY_ADMIN.value,
|
||||
# View-only admins are operators (oncall, support); they need the
|
||||
# routing fields (api_base, api_version) to diagnose health and tell
|
||||
# which provider region a check is hitting. They cannot mutate config
|
||||
# so granting them the read-only view is safe.
|
||||
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:
|
||||
"""
|
||||
Return True if the caller has a proxy-admin role (full or view-only).
|
||||
|
||||
user_role on UserAPIKeyAuth can be either a LitellmUserRoles enum or its
|
||||
string value depending on how the auth path constructed the object, so we
|
||||
compare against the raw value rather than the enum identity.
|
||||
"""
|
||||
role = user_api_key_dict.user_role
|
||||
if role is None:
|
||||
return False
|
||||
role_value = role.value if hasattr(role, "value") else role
|
||||
return role_value in _PROXY_ADMIN_ROLES
|
||||
|
||||
|
||||
def _strip_admin_only_fields_from_health_result(result: dict) -> dict:
|
||||
"""
|
||||
Return a copy of the /health response with provider routing fields
|
||||
(``api_base``, ``api_version``) removed from each healthy/unhealthy
|
||||
endpoint entry. Used to hide those fields from non-admin callers while
|
||||
still showing them which deployments they own and whether each one is
|
||||
healthy. Proxy admins receive the unmodified result.
|
||||
"""
|
||||
out = dict(result)
|
||||
drop = set(ADMIN_ONLY_HEALTH_DISPLAY_PARAMS)
|
||||
for key in ("healthy_endpoints", "unhealthy_endpoints"):
|
||||
eps = out.get(key)
|
||||
if isinstance(eps, list):
|
||||
out[key] = [
|
||||
(
|
||||
{k: v for k, v in ep.items() if k not in drop}
|
||||
if isinstance(ep, dict)
|
||||
else ep
|
||||
)
|
||||
for ep in eps
|
||||
]
|
||||
return out
|
||||
|
||||
|
||||
def _filter_health_check_results_by_model_ids(
|
||||
results: dict, allowed_model_ids: set
|
||||
) -> dict:
|
||||
"""
|
||||
Restrict a cached background health-check result dict to endpoints whose
|
||||
model_id is in ``allowed_model_ids``.
|
||||
|
||||
Endpoints without a model_id (e.g. CLI-model entries that predate the
|
||||
model_id wiring) are dropped conservatively — we cannot prove they belong
|
||||
to the caller, so they are excluded rather than leaked.
|
||||
|
||||
Each retained endpoint is shallow-copied before being returned, so any
|
||||
downstream transform (e.g. _strip_admin_only_fields_from_health_result)
|
||||
cannot accidentally mutate the shared ``health_check_results`` cache.
|
||||
"""
|
||||
healthy = [
|
||||
dict(ep)
|
||||
for ep in (results.get("healthy_endpoints") or [])
|
||||
if ep.get("model_id") in allowed_model_ids
|
||||
]
|
||||
unhealthy = [
|
||||
dict(ep)
|
||||
for ep in (results.get("unhealthy_endpoints") or [])
|
||||
if ep.get("model_id") in allowed_model_ids
|
||||
]
|
||||
return {
|
||||
"healthy_endpoints": healthy,
|
||||
"unhealthy_endpoints": unhealthy,
|
||||
"healthy_count": len(healthy),
|
||||
"unhealthy_count": len(unhealthy),
|
||||
}
|
||||
|
||||
|
||||
async def _perform_health_check_and_save(
|
||||
model_list,
|
||||
target_model,
|
||||
|
|
@ -771,6 +857,7 @@ async def _perform_health_check_and_save(
|
|||
|
||||
@router.get("/health", tags=["health"], dependencies=[Depends(user_api_key_auth)])
|
||||
async def health_endpoint(
|
||||
response: Response,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
model: Optional[str] = fastapi.Query(
|
||||
None, description="Specify the model name (optional)"
|
||||
|
|
@ -838,11 +925,26 @@ async def health_endpoint(
|
|||
detail={"error": f"Model with ID {model_id} not found"},
|
||||
)
|
||||
|
||||
is_admin = _is_proxy_admin(user_api_key_dict)
|
||||
|
||||
def _post_process(result: dict) -> dict:
|
||||
# api_base / api_version reveal which provider/region/internal host the
|
||||
# deployment talks to; only proxy admins receive them. Non-admin keys
|
||||
# still see model/model_id and the healthy/unhealthy status. We also
|
||||
# set a header so non-admin clients that previously parsed those
|
||||
# fields can detect the change programmatically.
|
||||
if is_admin:
|
||||
return result
|
||||
response.headers["Litellm-Health-Field-Notice"] = (
|
||||
"api_base and api_version are admin-only on this endpoint"
|
||||
)
|
||||
return _strip_admin_only_fields_from_health_result(result)
|
||||
|
||||
try:
|
||||
if llm_model_list is None:
|
||||
# if no router set, check if user set a model using litellm --model ollama/llama2
|
||||
if user_model is not None:
|
||||
return await _perform_health_check_and_save(
|
||||
cli_result = await _perform_health_check_and_save(
|
||||
model_list=[],
|
||||
target_model=None,
|
||||
cli_model=user_model,
|
||||
|
|
@ -853,20 +955,59 @@ async def health_endpoint(
|
|||
model_id=None, # CLI model doesn't have model_id
|
||||
max_concurrency=health_check_concurrency,
|
||||
)
|
||||
return _post_process(cli_result)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail={"error": "Model list not initialized"},
|
||||
)
|
||||
_llm_model_list = copy.deepcopy(llm_model_list)
|
||||
### FILTER MODELS FOR ONLY THOSE USER HAS ACCESS TO ###
|
||||
# Live path: scope by model_name (every deployment has one).
|
||||
# Cache path: scope by model_id (the cache is keyed on model_id).
|
||||
# Consequence: a deployment whose model_name the caller can access
|
||||
# but which lacks model_info.id will appear in the live /health
|
||||
# response but NOT in the background-cache /health response. This is
|
||||
# surfaced via the "warnings" field below so operators can fix the
|
||||
# missing model_info.id rather than guess at the discrepancy.
|
||||
if len(user_api_key_dict.models) > 0:
|
||||
pass
|
||||
else:
|
||||
pass #
|
||||
allowed_models = set(user_api_key_dict.models)
|
||||
_llm_model_list = [
|
||||
m for m in _llm_model_list if m.get("model_name") in allowed_models
|
||||
]
|
||||
if use_background_health_checks:
|
||||
return health_check_results
|
||||
if len(user_api_key_dict.models) > 0:
|
||||
allowed_model_ids = {
|
||||
(m.get("model_info") or {}).get("id")
|
||||
for m in _llm_model_list
|
||||
if (m.get("model_info") or {}).get("id")
|
||||
}
|
||||
filtered = _filter_health_check_results_by_model_ids(
|
||||
health_check_results, allowed_model_ids
|
||||
)
|
||||
if not allowed_model_ids:
|
||||
# Caller has accessible model_names but none of the
|
||||
# matching deployments expose a model_info.id, so the
|
||||
# cache filter (which keys on model_id) drops every
|
||||
# entry. Surface this both as a warning log and a
|
||||
# structured "warnings" field on the response so the
|
||||
# caller can distinguish "no deployments found" from
|
||||
# "deployments excluded due to missing model_info.id".
|
||||
verbose_proxy_logger.warning(
|
||||
"health_endpoint: scoped key %s has accessible models %s "
|
||||
"but none of the matching deployments carry a model_info.id; "
|
||||
"background health-check cache will return an empty result.",
|
||||
user_api_key_dict.user_id,
|
||||
list(user_api_key_dict.models),
|
||||
)
|
||||
filtered["warnings"] = [
|
||||
"Some accessible deployments are missing model_info.id "
|
||||
"and were excluded from this response. Ask a proxy admin "
|
||||
"to populate model_info.id for these models."
|
||||
]
|
||||
return _post_process(filtered)
|
||||
return _post_process(health_check_results)
|
||||
else:
|
||||
return await _perform_health_check_and_save(
|
||||
router_result = await _perform_health_check_and_save(
|
||||
model_list=_llm_model_list,
|
||||
target_model=target_model,
|
||||
cli_model=None,
|
||||
|
|
@ -877,6 +1018,7 @@ async def health_endpoint(
|
|||
model_id=model_id,
|
||||
max_concurrency=health_check_concurrency,
|
||||
)
|
||||
return _post_process(router_result)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
"litellm.proxy.proxy_server.py::health_endpoint(): Exception occured - {}".format(
|
||||
|
|
|
|||
|
|
@ -236,13 +236,12 @@ async def _patch_key_caches_add_access_group(
|
|||
) -> None:
|
||||
"""Patch cached key objects to include access_group_id."""
|
||||
for token in key_tokens:
|
||||
cached_key = await user_api_key_cache.async_get_cache(key=token)
|
||||
cached_key = await user_api_key_cache.async_get_cache(
|
||||
key=token,
|
||||
model_type=UserAPIKeyAuth,
|
||||
)
|
||||
if cached_key is None:
|
||||
continue
|
||||
if isinstance(cached_key, dict):
|
||||
cached_key = UserAPIKeyAuth(**cached_key)
|
||||
if not isinstance(cached_key, UserAPIKeyAuth):
|
||||
continue
|
||||
if cached_key.access_group_ids is None:
|
||||
cached_key.access_group_ids = [access_group_id]
|
||||
elif access_group_id not in cached_key.access_group_ids:
|
||||
|
|
@ -267,12 +266,11 @@ async def _patch_key_caches_remove_access_group(
|
|||
) -> None:
|
||||
"""Patch cached key objects to remove access_group_id."""
|
||||
for token in key_tokens:
|
||||
cached_key = await user_api_key_cache.async_get_cache(key=token)
|
||||
if cached_key is None:
|
||||
continue
|
||||
if isinstance(cached_key, dict):
|
||||
cached_key = UserAPIKeyAuth(**cached_key)
|
||||
if isinstance(cached_key, UserAPIKeyAuth) and cached_key.access_group_ids:
|
||||
cached_key = await user_api_key_cache.async_get_cache(
|
||||
key=token,
|
||||
model_type=UserAPIKeyAuth,
|
||||
)
|
||||
if cached_key is not None and cached_key.access_group_ids:
|
||||
cached_key.access_group_ids = [
|
||||
ag for ag in cached_key.access_group_ids if ag != access_group_id
|
||||
]
|
||||
|
|
|
|||
|
|
@ -27,7 +27,7 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, s
|
|||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.caching import DualCache
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.constants import (
|
||||
LENGTH_OF_LITELLM_GENERATED_KEY,
|
||||
LITELLM_PROXY_ADMIN_NAME,
|
||||
|
|
@ -1059,7 +1059,7 @@ async def _check_project_key_limits(
|
|||
project_id: str,
|
||||
data: Union[GenerateKeyRequest, UpdateKeyRequest],
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
) -> None:
|
||||
"""
|
||||
Validate that key's models and budget respect its project's limits.
|
||||
|
|
@ -1834,7 +1834,7 @@ async def _process_single_key_update(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_changed_by: Optional[str],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: Any,
|
||||
llm_router: Optional[Router],
|
||||
user_custom_key_update: Optional[Callable] = None,
|
||||
|
|
@ -3298,7 +3298,7 @@ async def _team_key_deletion_check(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
key_info: LiteLLM_VerificationToken,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
):
|
||||
is_team_key = _is_team_key(data=key_info)
|
||||
|
||||
|
|
@ -3341,7 +3341,7 @@ async def _team_key_deletion_check(
|
|||
|
||||
async def can_modify_verification_token(
|
||||
key_info: LiteLLM_VerificationToken,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
) -> bool:
|
||||
|
|
@ -3415,7 +3415,7 @@ async def can_modify_verification_token(
|
|||
|
||||
async def delete_verification_tokens(
|
||||
tokens: List,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_changed_by: Optional[str] = None,
|
||||
) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]:
|
||||
|
|
@ -3605,7 +3605,7 @@ async def _persist_deleted_verification_tokens(
|
|||
|
||||
async def delete_key_aliases(
|
||||
key_aliases: List[str],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_changed_by: Optional[str] = None,
|
||||
|
|
@ -3862,7 +3862,7 @@ async def _execute_virtual_key_regeneration(
|
|||
data: Optional[RegenerateKeyRequest],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_changed_by: Optional[str],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> GenerateKeyResponse:
|
||||
"""Generate new token, update DB, invalidate cache, and return response."""
|
||||
|
|
@ -4152,7 +4152,7 @@ async def _check_proxy_or_team_admin_for_key(
|
|||
key_in_db: LiteLLM_VerificationToken,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
) -> None:
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||||
return
|
||||
|
|
@ -5173,7 +5173,7 @@ async def _check_key_admin_access(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
hashed_token: str,
|
||||
prisma_client: Any,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
route: str,
|
||||
) -> None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -39,9 +39,9 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
|
|||
from fastapi.responses import RedirectResponse
|
||||
|
||||
import litellm
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.caching import DualCache
|
||||
from litellm.constants import (
|
||||
CLI_SSO_SESSION_CACHE_KEY_PREFIX,
|
||||
CLI_SSO_SESSION_TTL_SECONDS,
|
||||
|
|
@ -75,7 +75,11 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object
|
||||
from litellm.proxy.auth.auth_utils import _get_request_ip_address, _has_user_setup_sso
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.auth.auth_utils import (
|
||||
_get_request_ip_address,
|
||||
_has_user_setup_sso,
|
||||
)
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.admin_ui_utils import (
|
||||
|
|
@ -1301,7 +1305,7 @@ async def get_existing_user_info_from_db(
|
|||
user_id: Optional[str],
|
||||
user_email: Optional[str],
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> Optional[LiteLLM_UserTable]:
|
||||
try:
|
||||
|
|
@ -1325,7 +1329,7 @@ async def get_existing_user_info_from_db(
|
|||
async def get_user_info_from_db(
|
||||
result: Union[CustomOpenID, OpenID, dict],
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
user_email: Optional[str],
|
||||
user_defined_values: Optional[SSOUserDefinedValues],
|
||||
|
|
@ -1445,7 +1449,7 @@ async def _sync_user_role_from_jwt_role_map(
|
|||
received_response: Optional[dict],
|
||||
user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]],
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
user_defined_values: Optional[SSOUserDefinedValues],
|
||||
) -> None:
|
||||
"""
|
||||
|
|
@ -1484,11 +1488,8 @@ async def _sync_user_role_from_jwt_role_map(
|
|||
user_info.user_role = mapped_role.value
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=user_info.user_id,
|
||||
value=(
|
||||
user_info.model_dump()
|
||||
if hasattr(user_info, "model_dump")
|
||||
else dict(user_info)
|
||||
),
|
||||
value=user_info,
|
||||
model_type=LiteLLM_UserTable,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
from typing import List, Optional
|
||||
|
||||
from litellm.caching import DualCache
|
||||
from litellm.proxy._types import (
|
||||
KeyManagementRoutes,
|
||||
LiteLLM_TeamTableCachedObj,
|
||||
|
|
@ -12,6 +11,7 @@ from litellm.proxy._types import (
|
|||
ProxyException,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.auth.auth_checks import get_team_object
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
|
@ -65,7 +65,7 @@ class TeamMemberPermissionChecks:
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
route: KeyManagementRoutes,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
existing_key_row: LiteLLM_VerificationToken,
|
||||
):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ Prometheus Auth Middleware - Pure ASGI implementation
|
|||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, List, MutableMapping
|
||||
|
||||
from fastapi import Request
|
||||
from starlette.types import ASGIApp, Receive, Scope, Send
|
||||
|
|
@ -40,8 +41,17 @@ class PrometheusAuthMiddleware:
|
|||
|
||||
# Only run auth if configured to do so
|
||||
if litellm.require_auth_for_metrics_endpoint is True:
|
||||
# Construct Request only when auth is actually needed
|
||||
request = Request(scope, receive)
|
||||
# user_api_key_auth reads the request body, which consumes ASGI `receive`.
|
||||
# Buffer those messages and replay them for the inner app; otherwise a
|
||||
# successful auth would forward an exhausted receive and /metrics hangs.
|
||||
buffered_messages: List[MutableMapping[str, Any]] = []
|
||||
|
||||
async def receive_for_auth() -> MutableMapping[str, Any]:
|
||||
message = await receive()
|
||||
buffered_messages.append(message)
|
||||
return message
|
||||
|
||||
request = Request(scope, receive_for_auth)
|
||||
api_key = request.headers.get(_AUTHORIZATION_HEADER) or ""
|
||||
|
||||
try:
|
||||
|
|
@ -70,5 +80,18 @@ class PrometheusAuthMiddleware:
|
|||
)
|
||||
return
|
||||
|
||||
replay_idx = 0
|
||||
|
||||
async def receive_replay() -> MutableMapping[str, Any]:
|
||||
nonlocal replay_idx
|
||||
if replay_idx < len(buffered_messages):
|
||||
msg = buffered_messages[replay_idx]
|
||||
replay_idx += 1
|
||||
return msg
|
||||
return await receive()
|
||||
|
||||
await self.app(scope, receive_replay, send)
|
||||
return
|
||||
|
||||
# Pass through to the inner application
|
||||
await self.app(scope, receive, send)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import json
|
||||
import time
|
||||
import urllib.parse
|
||||
from datetime import datetime
|
||||
from typing import Literal, Optional
|
||||
from urllib.parse import urlparse
|
||||
|
|
@ -203,8 +204,16 @@ class AssemblyAIPassthroughLoggingHandler:
|
|||
)
|
||||
if _api_key is None:
|
||||
raise ValueError("AssemblyAI API key not found")
|
||||
if (
|
||||
any(c in transcript_id for c in ("/", "\\", "#", "?"))
|
||||
or ".." in transcript_id
|
||||
):
|
||||
raise ValueError(
|
||||
f"Invalid transcript_id {transcript_id!r}: contains disallowed characters"
|
||||
)
|
||||
safe_transcript_id = urllib.parse.quote(transcript_id, safe="")
|
||||
try:
|
||||
url = f"{_base_url}/v2/transcript/{transcript_id}"
|
||||
url = f"{_base_url}/v2/transcript/{safe_transcript_id}"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {_api_key}",
|
||||
"Content-Type": "application/json",
|
||||
|
|
|
|||
|
|
@ -78,8 +78,11 @@ from litellm.proxy._types import (
|
|||
InvitationNew,
|
||||
InvitationUpdate,
|
||||
Litellm_EntityType,
|
||||
LiteLLM_EndUserTable,
|
||||
LiteLLM_JWTAuth,
|
||||
LiteLLM_TagTable,
|
||||
LiteLLM_TeamTable,
|
||||
LiteLLM_TeamTableCachedObj,
|
||||
LiteLLM_UserTable,
|
||||
LitellmUserRoles,
|
||||
PassThroughGenericEndpoint,
|
||||
|
|
@ -94,6 +97,7 @@ from litellm.proxy._types import (
|
|||
UI_TEAM_ID,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
normalize_callback_names,
|
||||
process_callback,
|
||||
|
|
@ -206,6 +210,7 @@ from litellm import Router
|
|||
from litellm._logging import verbose_proxy_logger, verbose_router_logger
|
||||
from litellm.caching.caching import DualCache, RedisCache
|
||||
from litellm.caching.redis_cluster_cache import RedisClusterCache
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.constants import (
|
||||
_REALTIME_BODY_CACHE_SIZE,
|
||||
APSCHEDULER_COALESCE,
|
||||
|
|
@ -1612,7 +1617,7 @@ prisma_client: Optional[PrismaClient] = None
|
|||
shared_aiohttp_session: Optional["ClientSession"] = (
|
||||
None # Global shared session for connection reuse
|
||||
)
|
||||
user_api_key_cache = DualCache(
|
||||
user_api_key_cache: UserApiKeyCache = UserApiKeyCache(
|
||||
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
|
||||
)
|
||||
spend_counter_cache = DualCache(
|
||||
|
|
@ -2014,14 +2019,16 @@ async def update_cache( # noqa: PLR0915
|
|||
else:
|
||||
hashed_token = token
|
||||
verbose_proxy_logger.debug("_update_key_cache: hashed_token=%s", hashed_token)
|
||||
existing_spend_obj: LiteLLM_VerificationTokenView = await user_api_key_cache.async_get_cache(key=hashed_token) # type: ignore
|
||||
existing_spend_obj = await user_api_key_cache.async_get_cache(
|
||||
key=hashed_token, model_type=UserAPIKeyAuth
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
f"_update_key_cache: existing_spend_obj={existing_spend_obj}"
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
return
|
||||
else:
|
||||
existing_spend = existing_spend_obj.spend
|
||||
|
||||
existing_spend = existing_spend_obj.spend or 0.0
|
||||
# Calculate the new cost by adding the existing cost and response_cost
|
||||
new_spend = existing_spend + response_cost
|
||||
|
||||
|
|
@ -2079,41 +2086,48 @@ async def update_cache( # noqa: PLR0915
|
|||
existing_team_member_spend + response_cost
|
||||
)
|
||||
|
||||
# Update the cost column for the given token
|
||||
# Existing spend_obj is mutated; UserApiKeyCache.async_set_cache_pipeline turns
|
||||
# BaseModel values into dicts for Redis (same Codec path as async_set_cache).
|
||||
existing_spend_obj.spend = new_spend
|
||||
values_to_update_in_cache.append((hashed_token, existing_spend_obj))
|
||||
|
||||
### UPDATE USER SPEND ###
|
||||
async def _update_user_cache():
|
||||
## UPDATE CACHE FOR USER ID + GLOBAL PROXY
|
||||
if response_cost is None:
|
||||
return
|
||||
user_ids = [user_id]
|
||||
try:
|
||||
for _id in user_ids:
|
||||
# Fetch the existing cost for the given user
|
||||
if _id is None:
|
||||
continue
|
||||
existing_spend_obj = await user_api_key_cache.async_get_cache(key=_id)
|
||||
if existing_spend_obj is None:
|
||||
cached_user = await user_api_key_cache.async_get_cache(key=_id)
|
||||
if cached_user is None:
|
||||
# do nothing if there is no cache value
|
||||
return
|
||||
existing_spend_obj = CacheCodec.deserialize(
|
||||
cached_user, LiteLLM_UserTable
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
return
|
||||
verbose_proxy_logger.debug(
|
||||
f"_update_user_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}"
|
||||
)
|
||||
|
||||
if isinstance(existing_spend_obj, dict):
|
||||
existing_spend = existing_spend_obj["spend"]
|
||||
else:
|
||||
existing_spend = existing_spend_obj.spend
|
||||
existing_spend = existing_spend_obj.spend or 0.0
|
||||
# Calculate the new cost by adding the existing cost and response_cost
|
||||
new_spend = existing_spend + response_cost
|
||||
|
||||
# Update the cost column for the given user
|
||||
if isinstance(existing_spend_obj, dict):
|
||||
existing_spend_obj["spend"] = new_spend
|
||||
values_to_update_in_cache.append((_id, existing_spend_obj))
|
||||
else:
|
||||
existing_spend_obj.spend = new_spend
|
||||
values_to_update_in_cache.append((_id, existing_spend_obj.json()))
|
||||
existing_spend_obj.spend = new_spend
|
||||
values_to_update_in_cache.append(
|
||||
(
|
||||
_id,
|
||||
CacheCodec.serialize(
|
||||
existing_spend_obj, model_type=LiteLLM_UserTable
|
||||
),
|
||||
)
|
||||
)
|
||||
## UPDATE GLOBAL PROXY ##
|
||||
global_proxy_spend = await user_api_key_cache.async_get_cache(
|
||||
key="{}:spend".format(litellm_proxy_admin_name)
|
||||
|
|
@ -2145,31 +2159,33 @@ async def update_cache( # noqa: PLR0915
|
|||
_id = "end_user_id:{}".format(end_user_id)
|
||||
try:
|
||||
# Fetch the existing cost for the given user
|
||||
existing_spend_obj = await user_api_key_cache.async_get_cache(key=_id)
|
||||
if existing_spend_obj is None:
|
||||
cached_end_user = await user_api_key_cache.async_get_cache(key=_id)
|
||||
if cached_end_user is None:
|
||||
# if user does not exist in LiteLLM_UserTable, create a new user
|
||||
# do nothing if end-user not in api key cache
|
||||
return
|
||||
existing_spend_obj = CacheCodec.deserialize(
|
||||
cached_end_user, LiteLLM_EndUserTable
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
return
|
||||
verbose_proxy_logger.debug(
|
||||
f"_update_end_user_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}"
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
existing_spend = 0
|
||||
else:
|
||||
if isinstance(existing_spend_obj, dict):
|
||||
existing_spend = existing_spend_obj["spend"]
|
||||
else:
|
||||
existing_spend = existing_spend_obj.spend
|
||||
|
||||
existing_spend = existing_spend_obj.spend or 0.0
|
||||
# Calculate the new cost by adding the existing cost and response_cost
|
||||
new_spend = existing_spend + response_cost
|
||||
|
||||
# Update the cost column for the given user
|
||||
if isinstance(existing_spend_obj, dict):
|
||||
existing_spend_obj["spend"] = new_spend
|
||||
values_to_update_in_cache.append((_id, existing_spend_obj))
|
||||
else:
|
||||
existing_spend_obj.spend = new_spend
|
||||
values_to_update_in_cache.append((_id, existing_spend_obj.json()))
|
||||
existing_spend_obj.spend = new_spend
|
||||
values_to_update_in_cache.append(
|
||||
(
|
||||
_id,
|
||||
CacheCodec.serialize(
|
||||
existing_spend_obj, model_type=LiteLLM_EndUserTable
|
||||
),
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - failed to update end user spend in cache. "
|
||||
|
|
@ -2188,36 +2204,32 @@ async def update_cache( # noqa: PLR0915
|
|||
|
||||
_id = "team_id:{}".format(team_id)
|
||||
try:
|
||||
# Fetch the existing cost for the given user
|
||||
existing_spend_obj: Optional[LiteLLM_TeamTable] = (
|
||||
await user_api_key_cache.async_get_cache(key=_id)
|
||||
cached_team = await user_api_key_cache.async_get_cache(key=_id)
|
||||
if cached_team is None:
|
||||
# do nothing if team not in api key cache
|
||||
return
|
||||
existing_spend_obj: Optional[LiteLLM_TeamTableCachedObj] = (
|
||||
CacheCodec.deserialize(cached_team, LiteLLM_TeamTableCachedObj)
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
# do nothing if team not in api key cache
|
||||
return
|
||||
verbose_proxy_logger.debug(
|
||||
f"_update_team_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}"
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
existing_spend: Optional[float] = 0.0
|
||||
else:
|
||||
if isinstance(existing_spend_obj, dict):
|
||||
existing_spend = existing_spend_obj["spend"]
|
||||
else:
|
||||
existing_spend = existing_spend_obj.spend
|
||||
|
||||
if existing_spend is None:
|
||||
existing_spend = 0.0
|
||||
existing_spend: float = existing_spend_obj.spend or 0.0
|
||||
# Calculate the new cost by adding the existing cost and response_cost
|
||||
new_spend = existing_spend + response_cost
|
||||
|
||||
# Update the cost column for the given user
|
||||
if isinstance(existing_spend_obj, dict):
|
||||
existing_spend_obj["spend"] = new_spend
|
||||
values_to_update_in_cache.append((_id, existing_spend_obj))
|
||||
else:
|
||||
existing_spend_obj.spend = new_spend
|
||||
values_to_update_in_cache.append((_id, existing_spend_obj))
|
||||
existing_spend_obj.spend = new_spend
|
||||
values_to_update_in_cache.append(
|
||||
(
|
||||
_id,
|
||||
CacheCodec.serialize(
|
||||
existing_spend_obj, model_type=LiteLLM_TeamTableCachedObj
|
||||
),
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - failed to update team spend in cache. "
|
||||
|
|
@ -2244,32 +2256,32 @@ async def update_cache( # noqa: PLR0915
|
|||
|
||||
cache_key = f"tag:{tag_name}"
|
||||
# Fetch the existing tag object from cache
|
||||
existing_tag_obj = await user_api_key_cache.async_get_cache(
|
||||
key=cache_key
|
||||
)
|
||||
if existing_tag_obj is None:
|
||||
cached_tag = await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
if cached_tag is None:
|
||||
# do nothing if tag not in api key cache
|
||||
continue
|
||||
|
||||
existing_tag_obj = CacheCodec.deserialize(cached_tag, LiteLLM_TagTable)
|
||||
if existing_tag_obj is None:
|
||||
continue
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"_update_tag_cache: existing spend for tag={tag_name}: {existing_tag_obj}; response_cost: {response_cost}"
|
||||
)
|
||||
|
||||
if isinstance(existing_tag_obj, dict):
|
||||
existing_spend = existing_tag_obj.get("spend", 0) or 0
|
||||
else:
|
||||
existing_spend = getattr(existing_tag_obj, "spend", 0) or 0
|
||||
|
||||
existing_spend = existing_tag_obj.spend or 0.0
|
||||
# Calculate the new cost by adding the existing cost and response_cost
|
||||
new_spend = existing_spend + response_cost
|
||||
|
||||
# Update the spend column for the given tag
|
||||
if isinstance(existing_tag_obj, dict):
|
||||
existing_tag_obj["spend"] = new_spend
|
||||
values_to_update_in_cache.append((cache_key, existing_tag_obj))
|
||||
else:
|
||||
existing_tag_obj.spend = new_spend
|
||||
values_to_update_in_cache.append((cache_key, existing_tag_obj))
|
||||
existing_tag_obj.spend = new_spend
|
||||
values_to_update_in_cache.append(
|
||||
(
|
||||
cache_key,
|
||||
CacheCodec.serialize(
|
||||
existing_tag_obj, model_type=LiteLLM_TagTable
|
||||
),
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - failed to update tag spend in cache. "
|
||||
|
|
@ -2937,8 +2949,9 @@ class ProxyConfig:
|
|||
def _init_cache(
|
||||
self,
|
||||
cache_params: dict,
|
||||
enable_redis_auth_cache: bool = False,
|
||||
):
|
||||
global redis_usage_cache, llm_router
|
||||
global redis_usage_cache, llm_router, general_settings
|
||||
from litellm import Cache
|
||||
|
||||
if "default_in_memory_ttl" in cache_params:
|
||||
|
|
@ -2954,7 +2967,29 @@ class ProxyConfig:
|
|||
):
|
||||
## INIT PROXY REDIS USAGE CLIENT ##
|
||||
redis_usage_cache = litellm.cache.cache
|
||||
spend_counter_cache.redis_cache = redis_usage_cache
|
||||
spend_counter_cache.attach_redis_cache(
|
||||
redis_usage_cache,
|
||||
default_redis_ttl=litellm.default_redis_ttl,
|
||||
)
|
||||
# Note: PKCE verifier storage uses redis_usage_cache directly (not
|
||||
# user_api_key_cache) to avoid routing all API-key lookups through Redis.
|
||||
if enable_redis_auth_cache is True:
|
||||
user_api_key_cache.attach_redis_cache(
|
||||
redis_usage_cache,
|
||||
default_redis_ttl=litellm.default_redis_ttl,
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
"enable_redis_auth_cache=True: attached Redis to "
|
||||
"user_api_key_cache — virtual-key lookups are now "
|
||||
"shared across all proxy workers."
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.info(
|
||||
"enable_redis_auth_cache is not set: user_api_key_cache "
|
||||
"remains in-memory only (per-worker). Set "
|
||||
"litellm_settings.enable_redis_auth_cache: true to share "
|
||||
"the auth cache across workers and reduce DB load."
|
||||
)
|
||||
litellm_config_cache.redis_cache = redis_usage_cache
|
||||
# Note: PKCE verifier storage uses redis_usage_cache directly (not
|
||||
# user_api_key_cache) to avoid routing all API-key lookups through Redis.
|
||||
|
|
@ -3280,7 +3315,13 @@ class ProxyConfig:
|
|||
cache_params[key] = get_secret(value)
|
||||
|
||||
## to pass a complete url, or set ssl=True, etc. just set it as `os.environ[REDIS_URL] = <your-redis-url>`, _redis.py checks for REDIS specific environment variables
|
||||
self._init_cache(cache_params=cache_params)
|
||||
self._init_cache(
|
||||
cache_params=cache_params,
|
||||
enable_redis_auth_cache=litellm_settings.get(
|
||||
"enable_redis_auth_cache", False
|
||||
)
|
||||
is True,
|
||||
)
|
||||
if litellm.cache is not None:
|
||||
verbose_proxy_logger.debug(
|
||||
f"{blue_color_code}Set Cache on LiteLLM Proxy{reset_color_code}"
|
||||
|
|
@ -3551,21 +3592,23 @@ class ProxyConfig:
|
|||
verbose_proxy_logger.critical(
|
||||
"LITELLM_MASTER_KEY is not set! All requests will be treated as INTERNAL_USER with no admin access. Set LITELLM_MASTER_KEY for production use."
|
||||
)
|
||||
### USER API KEY CACHE IN-MEMORY TTL ###
|
||||
### USER API KEY CACHE TTL (in-memory + Redis when Redis auth sharing is enabled) ###
|
||||
user_api_key_cache_ttl = general_settings.get(
|
||||
"user_api_key_cache_ttl", None
|
||||
)
|
||||
if user_api_key_cache_ttl is not None:
|
||||
ttl = float(user_api_key_cache_ttl)
|
||||
# Mirror TTL on Redis as well when ``litellm_settings.enable_redis_auth_cache``
|
||||
# attaches Redis to ``user_api_key_cache``; otherwise DualCache misses in
|
||||
# memory fall back to a key that outlasts ``user_api_key_cache_ttl``.
|
||||
user_api_key_cache.update_cache_ttl(
|
||||
default_in_memory_ttl=float(user_api_key_cache_ttl),
|
||||
default_redis_ttl=None, # user_api_key_cache uses in-memory TTL only; Redis not configured for key lookups
|
||||
default_in_memory_ttl=ttl,
|
||||
default_redis_ttl=ttl,
|
||||
)
|
||||
|
||||
### PKCE MULTI-INSTANCE PREREQUISITE CHECK ###
|
||||
# PKCE verifiers are stored in redis_usage_cache when available so they can
|
||||
# be read back by any instance (not just the one that started the auth flow).
|
||||
# user_api_key_cache is intentionally left in-memory-only to avoid routing
|
||||
# all API-key lookups through Redis.
|
||||
use_pkce = os.getenv("GENERIC_CLIENT_USE_PKCE", "false").lower() == "true"
|
||||
if use_pkce and redis_usage_cache is None:
|
||||
global _pkce_no_redis_warning_emitted
|
||||
|
|
@ -6294,7 +6337,7 @@ class ProxyStartupEvent:
|
|||
cls,
|
||||
general_settings: dict,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
):
|
||||
"""Initialize JWT auth on startup"""
|
||||
if general_settings.get("litellm_jwtauth", None) is not None:
|
||||
|
|
@ -6343,7 +6386,7 @@ class ProxyStartupEvent:
|
|||
async def _warm_global_spend_cache(
|
||||
cls,
|
||||
litellm_proxy_admin_name: str,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
"""Warm global spend cache once at startup to reduce impact of first wave of requests."""
|
||||
|
|
@ -6983,7 +7026,7 @@ class ProxyStartupEvent:
|
|||
cls,
|
||||
database_url: Optional[str],
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
) -> Optional[PrismaClient]:
|
||||
"""
|
||||
- Sets up prisma client
|
||||
|
|
@ -10285,6 +10328,101 @@ def _paginate_models_response(
|
|||
}
|
||||
|
||||
|
||||
def _team_models_resolve_to_names(
|
||||
team_models: List[str], access_groups: Dict[str, Any]
|
||||
) -> List[str]:
|
||||
"""Expand team model entries (including access group names) to concrete model names."""
|
||||
resolved: List[str] = []
|
||||
for name in team_models:
|
||||
if name in access_groups:
|
||||
resolved.extend(access_groups[name])
|
||||
else:
|
||||
resolved.append(name)
|
||||
return resolved
|
||||
|
||||
|
||||
async def _load_team_object_for_model_filter(
|
||||
team_id: str, prisma_client: PrismaClient
|
||||
) -> Optional[LiteLLM_TeamTable]:
|
||||
"""Load team row from DB; returns None if missing or on error."""
|
||||
try:
|
||||
team_db_object = await prisma_client.db.litellm_teamtable.find_unique(
|
||||
where={"team_id": team_id}
|
||||
)
|
||||
if team_db_object is None:
|
||||
verbose_proxy_logger.warning(f"Team {team_id} not found in database")
|
||||
return None
|
||||
return LiteLLM_TeamTable(**team_db_object.model_dump())
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(f"Error fetching team {team_id}: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
async def _gather_team_accessible_model_ids(
|
||||
team_object: LiteLLM_TeamTable,
|
||||
team_id: str,
|
||||
prisma_client: PrismaClient,
|
||||
llm_router: Router,
|
||||
) -> Set[str]:
|
||||
"""Collect model IDs the team can use from router config and DB."""
|
||||
team_accessible_model_ids: Set[str] = set()
|
||||
access_groups = llm_router.get_model_access_groups() if llm_router else {}
|
||||
|
||||
if (
|
||||
not team_object.models
|
||||
or SpecialModelNames.all_proxy_models.value in team_object.models
|
||||
):
|
||||
model_list = llm_router.get_model_list() if llm_router else []
|
||||
if model_list is not None:
|
||||
for model in model_list:
|
||||
model_id = model.get("model_info", {}).get("id", None)
|
||||
if model_id is None:
|
||||
continue
|
||||
team_model_id = model.get("model_info", {}).get("team_id", None)
|
||||
if team_model_id is None or team_model_id == team_id:
|
||||
team_accessible_model_ids.add(model_id)
|
||||
else:
|
||||
resolved_model_names: Set[str] = set()
|
||||
for model_name in team_object.models:
|
||||
if model_name in access_groups:
|
||||
resolved_model_names.update(access_groups[model_name])
|
||||
else:
|
||||
resolved_model_names.add(model_name)
|
||||
|
||||
for model_name in resolved_model_names:
|
||||
_models = (
|
||||
llm_router.get_model_list(model_name=model_name, team_id=team_id)
|
||||
if llm_router
|
||||
else []
|
||||
)
|
||||
if _models is not None:
|
||||
for model in _models:
|
||||
model_id = model.get("model_info", {}).get("id", None)
|
||||
if model_id is not None:
|
||||
team_accessible_model_ids.add(model_id)
|
||||
|
||||
try:
|
||||
if (
|
||||
team_object.models
|
||||
and SpecialModelNames.all_proxy_models.value not in team_object.models
|
||||
):
|
||||
_resolved_names = _team_models_resolve_to_names(
|
||||
team_object.models, access_groups
|
||||
)
|
||||
db_models = await prisma_client.db.litellm_proxymodeltable.find_many(
|
||||
where={"model_name": {"in": _resolved_names}}
|
||||
)
|
||||
for db_model in db_models:
|
||||
if db_model.model_id:
|
||||
team_accessible_model_ids.add(db_model.model_id)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Error querying database models for team {team_id}: {str(e)}"
|
||||
)
|
||||
|
||||
return team_accessible_model_ids
|
||||
|
||||
|
||||
async def _filter_models_by_team_id(
|
||||
all_models: List[Dict[str, Any]],
|
||||
team_id: str,
|
||||
|
|
@ -10307,78 +10445,13 @@ async def _filter_models_by_team_id(
|
|||
Returns:
|
||||
Filtered list of models
|
||||
"""
|
||||
# Get team from database
|
||||
try:
|
||||
team_db_object = await prisma_client.db.litellm_teamtable.find_unique(
|
||||
where={"team_id": team_id}
|
||||
)
|
||||
if team_db_object is None:
|
||||
verbose_proxy_logger.warning(f"Team {team_id} not found in database")
|
||||
# If team doesn't exist, return empty list
|
||||
return []
|
||||
|
||||
team_object = LiteLLM_TeamTable(**team_db_object.model_dump())
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(f"Error fetching team {team_id}: {str(e)}")
|
||||
team_object = await _load_team_object_for_model_filter(team_id, prisma_client)
|
||||
if team_object is None:
|
||||
return []
|
||||
|
||||
# Get models accessible to this team (similar to _add_team_models_to_all_models)
|
||||
team_accessible_model_ids: Set[str] = set()
|
||||
|
||||
if (
|
||||
not team_object.models # empty list = all model access
|
||||
or SpecialModelNames.all_proxy_models.value in team_object.models
|
||||
):
|
||||
# Team has access to all models
|
||||
model_list = llm_router.get_model_list() if llm_router else []
|
||||
if model_list is not None:
|
||||
for model in model_list:
|
||||
model_id = model.get("model_info", {}).get("id", None)
|
||||
if model_id is None:
|
||||
continue
|
||||
# if team model id set, check if team id matches
|
||||
team_model_id = model.get("model_info", {}).get("team_id", None)
|
||||
can_add_model = False
|
||||
if team_model_id is None:
|
||||
can_add_model = True
|
||||
elif team_model_id == team_id:
|
||||
can_add_model = True
|
||||
|
||||
if can_add_model:
|
||||
team_accessible_model_ids.add(model_id)
|
||||
else:
|
||||
# Team has access to specific models
|
||||
for model_name in team_object.models:
|
||||
_models = (
|
||||
llm_router.get_model_list(model_name=model_name, team_id=team_id)
|
||||
if llm_router
|
||||
else []
|
||||
)
|
||||
if _models is not None:
|
||||
for model in _models:
|
||||
model_id = model.get("model_info", {}).get("id", None)
|
||||
if model_id is not None:
|
||||
team_accessible_model_ids.add(model_id)
|
||||
|
||||
# Also search database for models accessible to this team
|
||||
# This complements the config search done above
|
||||
try:
|
||||
if (
|
||||
team_object.models
|
||||
and SpecialModelNames.all_proxy_models.value not in team_object.models
|
||||
):
|
||||
# Team has specific models - check database for those model names
|
||||
db_models = await prisma_client.db.litellm_proxymodeltable.find_many(
|
||||
where={"model_name": {"in": team_object.models}}
|
||||
)
|
||||
for db_model in db_models:
|
||||
model_id = db_model.model_id
|
||||
if model_id:
|
||||
team_accessible_model_ids.add(model_id)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Error querying database models for team {team_id}: {str(e)}"
|
||||
)
|
||||
team_accessible_model_ids = await _gather_team_accessible_model_ids(
|
||||
team_object, team_id, prisma_client, llm_router
|
||||
)
|
||||
|
||||
# Filter models based on direct_access or access_via_team_ids
|
||||
# Models are already enriched with these fields before this function is called
|
||||
|
|
|
|||
|
|
@ -101,6 +101,7 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.db.create_views import (
|
||||
create_missing_views,
|
||||
should_create_missing_views,
|
||||
|
|
@ -340,7 +341,7 @@ class ProxyLogging:
|
|||
|
||||
def __init__(
|
||||
self,
|
||||
user_api_key_cache: DualCache,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
premium_user: bool = False,
|
||||
):
|
||||
## INITIALIZE LITELLM CALLBACKS ##
|
||||
|
|
@ -5715,7 +5716,7 @@ async def get_available_models_for_user(
|
|||
include_model_access_groups: bool = False,
|
||||
only_model_access_groups: bool = False,
|
||||
return_wildcard_routes: bool = False,
|
||||
user_api_key_cache: Optional["DualCache"] = None,
|
||||
user_api_key_cache: Optional["UserApiKeyCache"] = None,
|
||||
) -> List[str]:
|
||||
"""
|
||||
Get the list of models available to a user based on their API key and team permissions.
|
||||
|
|
|
|||
|
|
@ -1,9 +1,12 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional
|
||||
from functools import lru_cache
|
||||
from typing import Any, Dict, List, Literal, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -22,19 +25,26 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
|
|||
from litellm.litellm_core_utils.thread_pool_executor import executor
|
||||
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.openai import (
|
||||
OutputTextDeltaEvent,
|
||||
ResponseAPIUsage,
|
||||
ResponseCompletedEvent,
|
||||
ResponsesAPIRequestParams,
|
||||
ResponsesAPIResponse,
|
||||
ResponsesAPIStreamEvents,
|
||||
ResponsesAPIStreamingResponse,
|
||||
)
|
||||
from litellm.types.llms.openai import ResponsesAPIStreamEvents
|
||||
from litellm.types.utils import CallTypes
|
||||
from litellm.utils import CustomStreamWrapper, async_post_call_success_deployment_hook
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _get_openai_response_types():
|
||||
from litellm.types.llms import openai as openai_types
|
||||
|
||||
return openai_types
|
||||
|
||||
|
||||
def _log_background_task_failure(task: "asyncio.Task[Any]", *, task_name: str) -> None:
|
||||
if task.cancelled():
|
||||
return
|
||||
exception = task.exception()
|
||||
if exception is not None:
|
||||
verbose_logger.error("%s failed: %s", task_name, exception)
|
||||
|
||||
|
||||
class BaseResponsesAPIStreamingIterator:
|
||||
"""
|
||||
Base class for streaming iterators that process responses from the Responses API.
|
||||
|
|
@ -46,7 +56,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self,
|
||||
response: httpx.Response,
|
||||
model: str,
|
||||
responses_api_provider_config: BaseResponsesAPIConfig,
|
||||
responses_api_provider_config: Optional[BaseResponsesAPIConfig],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
litellm_metadata: Optional[Dict[str, Any]] = None,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
|
|
@ -58,9 +68,13 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self.logging_obj = logging_obj
|
||||
self.finished = False
|
||||
self.responses_api_provider_config = responses_api_provider_config
|
||||
self.completed_response: Optional[ResponsesAPIStreamingResponse] = None
|
||||
self.completed_response: Optional[Any] = None
|
||||
self.start_time = getattr(logging_obj, "start_time", datetime.now())
|
||||
self._failure_handled = False # Track if failure handler has been called
|
||||
self._completed_response_cached = False
|
||||
self._completed_response_logged = False
|
||||
self._completed_response_cache_hit: Optional[bool] = None
|
||||
self._persist_completed_response_before_logging = True
|
||||
self._stream_created_time: float = time.time()
|
||||
|
||||
# track request context for hooks
|
||||
|
|
@ -101,7 +115,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
llm_provider=self.custom_llm_provider or "",
|
||||
)
|
||||
|
||||
def _process_chunk(self, chunk) -> Optional[ResponsesAPIStreamingResponse]:
|
||||
def _process_chunk(self, chunk) -> Optional[Any]:
|
||||
"""Process a single chunk of data from the stream"""
|
||||
if not chunk:
|
||||
return None
|
||||
|
|
@ -122,6 +136,10 @@ class BaseResponsesAPIStreamingIterator:
|
|||
|
||||
# Format as ResponsesAPIStreamingResponse
|
||||
if isinstance(parsed_chunk, dict):
|
||||
if self.responses_api_provider_config is None:
|
||||
raise ValueError(
|
||||
"responses_api_provider_config is required to process live streaming chunks"
|
||||
)
|
||||
openai_responses_api_chunk = (
|
||||
self.responses_api_provider_config.transform_streaming_response(
|
||||
model=self.model,
|
||||
|
|
@ -195,10 +213,11 @@ class BaseResponsesAPIStreamingIterator:
|
|||
if self.litellm_metadata and self.litellm_metadata.get(
|
||||
"encrypted_content_affinity_enabled"
|
||||
):
|
||||
openai_types = _get_openai_response_types()
|
||||
event_type = getattr(openai_responses_api_chunk, "type", None)
|
||||
if event_type in (
|
||||
ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
|
||||
ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
|
||||
openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
|
||||
openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
|
||||
):
|
||||
item = getattr(openai_responses_api_chunk, "item", None)
|
||||
if item:
|
||||
|
|
@ -219,10 +238,11 @@ class BaseResponsesAPIStreamingIterator:
|
|||
|
||||
# Store the completed response (also for incomplete/failed so logging still fires)
|
||||
_chunk_type = getattr(openai_responses_api_chunk, "type", None)
|
||||
openai_types = _get_openai_response_types()
|
||||
if openai_responses_api_chunk and _chunk_type in (
|
||||
ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
|
||||
ResponsesAPIStreamEvents.RESPONSE_FAILED,
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED,
|
||||
):
|
||||
self.completed_response = openai_responses_api_chunk
|
||||
# Add cost to usage object if include_cost_in_streaming_usage is True
|
||||
|
|
@ -230,11 +250,11 @@ class BaseResponsesAPIStreamingIterator:
|
|||
litellm.include_cost_in_streaming_usage
|
||||
and self.logging_obj is not None
|
||||
):
|
||||
response_obj: Optional[ResponsesAPIResponse] = getattr(
|
||||
response_obj: Optional[Any] = getattr(
|
||||
openai_responses_api_chunk, "response", None
|
||||
)
|
||||
if response_obj:
|
||||
usage_obj: Optional[ResponseAPIUsage] = getattr(
|
||||
usage_obj: Optional[Any] = getattr(
|
||||
response_obj, "usage", None
|
||||
)
|
||||
if usage_obj is not None:
|
||||
|
|
@ -247,9 +267,13 @@ class BaseResponsesAPIStreamingIterator:
|
|||
if cost is not None:
|
||||
setattr(usage_obj, "cost", cost)
|
||||
except Exception:
|
||||
# Best-effort usage cost annotation should not break stream replay.
|
||||
pass
|
||||
|
||||
if _chunk_type == ResponsesAPIStreamEvents.RESPONSE_FAILED:
|
||||
if (
|
||||
_chunk_type
|
||||
== openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED
|
||||
):
|
||||
self._handle_logging_failed_response()
|
||||
else:
|
||||
self._handle_logging_completed_response()
|
||||
|
|
@ -266,6 +290,59 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self._handle_failure(e)
|
||||
raise
|
||||
|
||||
def _log_completed_response(self, *, is_async: bool) -> None:
|
||||
if self._completed_response_logged:
|
||||
return
|
||||
self._completed_response_logged = True
|
||||
|
||||
if self._persist_completed_response_before_logging:
|
||||
self._persist_completed_response_to_cache(is_async=is_async)
|
||||
|
||||
# Create a copy for logging to avoid modifying the response object that will be returned to the user
|
||||
# The logging handlers may transform usage from Responses API format (input_tokens/output_tokens)
|
||||
# to chat completion format (prompt_tokens/completion_tokens) for internal logging
|
||||
# Use model_dump + model_validate instead of deepcopy to avoid pickle errors with
|
||||
# Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192)
|
||||
logging_response = self.completed_response
|
||||
if self.completed_response is not None and hasattr(
|
||||
self.completed_response, "model_dump"
|
||||
):
|
||||
try:
|
||||
logging_response = type(self.completed_response).model_validate(
|
||||
self.completed_response.model_dump()
|
||||
)
|
||||
except Exception:
|
||||
# Fallback to original if serialization fails
|
||||
pass
|
||||
|
||||
end_time = datetime.now()
|
||||
if is_async:
|
||||
asyncio.create_task(
|
||||
self.logging_obj.async_success_handler(
|
||||
result=logging_response,
|
||||
start_time=self.start_time,
|
||||
end_time=end_time,
|
||||
cache_hit=self._completed_response_cache_hit,
|
||||
)
|
||||
)
|
||||
else:
|
||||
run_async_function(
|
||||
async_function=self.logging_obj.async_success_handler,
|
||||
result=logging_response,
|
||||
start_time=self.start_time,
|
||||
end_time=end_time,
|
||||
cache_hit=self._completed_response_cache_hit,
|
||||
)
|
||||
|
||||
executor.submit(
|
||||
self.logging_obj.success_handler,
|
||||
result=logging_response,
|
||||
cache_hit=self._completed_response_cache_hit,
|
||||
start_time=self.start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
self._run_post_success_hooks(end_time=end_time)
|
||||
|
||||
def _handle_logging_completed_response(self):
|
||||
"""Base implementation - should be overridden by subclasses"""
|
||||
pass
|
||||
|
|
@ -296,6 +373,88 @@ class BaseResponsesAPIStreamingIterator:
|
|||
)
|
||||
self._handle_failure(exception)
|
||||
|
||||
def _get_completed_response_object(self) -> Optional[Any]:
|
||||
openai_types = _get_openai_response_types()
|
||||
completed_response = self.completed_response
|
||||
if isinstance(completed_response, openai_types.ResponsesAPIResponse):
|
||||
return completed_response
|
||||
|
||||
response_obj = getattr(completed_response, "response", None)
|
||||
if isinstance(response_obj, openai_types.ResponsesAPIResponse):
|
||||
return response_obj
|
||||
|
||||
return None
|
||||
|
||||
def _persist_completed_response_to_cache(self, *, is_async: bool) -> None:
|
||||
if self._completed_response_cached:
|
||||
return
|
||||
|
||||
completed_response = self.completed_response
|
||||
openai_types = _get_openai_response_types()
|
||||
if (
|
||||
getattr(completed_response, "type", None)
|
||||
!= openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED
|
||||
):
|
||||
return
|
||||
|
||||
response_obj = self._get_completed_response_object()
|
||||
if response_obj is None:
|
||||
return
|
||||
|
||||
caching_handler = getattr(self.logging_obj, "_llm_caching_handler", None)
|
||||
if caching_handler is None:
|
||||
return
|
||||
|
||||
request_kwargs = getattr(caching_handler, "request_kwargs", None)
|
||||
if (
|
||||
not isinstance(request_kwargs, dict)
|
||||
or request_kwargs.get("stream") is not True
|
||||
):
|
||||
return
|
||||
request_kwargs = request_kwargs.copy()
|
||||
preset_cache_key = getattr(caching_handler, "preset_cache_key", None)
|
||||
request_cache_key = request_kwargs.pop("cache_key", None)
|
||||
if preset_cache_key is None:
|
||||
preset_cache_key = request_cache_key
|
||||
if request_kwargs.get("metadata") is None:
|
||||
request_kwargs.pop("metadata", None)
|
||||
request_kwargs.pop("custom_llm_provider", None)
|
||||
if preset_cache_key is not None:
|
||||
request_kwargs["cache_key"] = preset_cache_key
|
||||
|
||||
if not caching_handler._should_store_result_in_cache(
|
||||
original_function=caching_handler.original_function,
|
||||
kwargs=request_kwargs,
|
||||
):
|
||||
return
|
||||
|
||||
if litellm.cache is None:
|
||||
return
|
||||
|
||||
cached_response = response_obj.model_dump_json()
|
||||
if is_async:
|
||||
cache_write_task = asyncio.create_task(
|
||||
litellm.cache.async_add_cache(
|
||||
cached_response,
|
||||
dynamic_cache_object=getattr(caching_handler, "dual_cache", None),
|
||||
**request_kwargs,
|
||||
)
|
||||
)
|
||||
cache_write_task.add_done_callback(
|
||||
lambda task: _log_background_task_failure(
|
||||
task,
|
||||
task_name="Responses stream cache write",
|
||||
)
|
||||
)
|
||||
else:
|
||||
litellm.cache.add_cache(
|
||||
cached_response,
|
||||
dynamic_cache_object=getattr(caching_handler, "dual_cache", None),
|
||||
**request_kwargs,
|
||||
)
|
||||
|
||||
self._completed_response_cached = True
|
||||
|
||||
async def _call_post_streaming_deployment_hook(self, chunk):
|
||||
"""
|
||||
Allow callbacks to modify streaming chunks before returning (parity with chat).
|
||||
|
|
@ -480,7 +639,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> ResponsesAPIStreamingResponse:
|
||||
async def __anext__(self) -> Any:
|
||||
try:
|
||||
self._check_max_streaming_duration()
|
||||
while True:
|
||||
|
|
@ -520,40 +679,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
|
||||
def _handle_logging_completed_response(self):
|
||||
"""Handle logging for completed responses in async context"""
|
||||
# Create a copy for logging to avoid modifying the response object that will be returned to the user
|
||||
# The logging handlers may transform usage from Responses API format (input_tokens/output_tokens)
|
||||
# to chat completion format (prompt_tokens/completion_tokens) for internal logging
|
||||
# Use model_dump + model_validate instead of deepcopy to avoid pickle errors with
|
||||
# Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192)
|
||||
logging_response = self.completed_response
|
||||
if self.completed_response is not None and hasattr(
|
||||
self.completed_response, "model_dump"
|
||||
):
|
||||
try:
|
||||
logging_response = type(self.completed_response).model_validate(
|
||||
self.completed_response.model_dump()
|
||||
)
|
||||
except Exception:
|
||||
# Fallback to original if serialization fails
|
||||
pass
|
||||
|
||||
asyncio.create_task(
|
||||
self.logging_obj.async_success_handler(
|
||||
result=logging_response,
|
||||
start_time=self.start_time,
|
||||
end_time=datetime.now(),
|
||||
cache_hit=None,
|
||||
)
|
||||
)
|
||||
|
||||
executor.submit(
|
||||
self.logging_obj.success_handler,
|
||||
result=logging_response,
|
||||
cache_hit=None,
|
||||
start_time=self.start_time,
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
self._run_post_success_hooks(end_time=datetime.now())
|
||||
self._log_completed_response(is_async=True)
|
||||
|
||||
|
||||
class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
||||
|
|
@ -627,39 +753,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
|
||||
def _handle_logging_completed_response(self):
|
||||
"""Handle logging for completed responses in sync context"""
|
||||
# Create a copy for logging to avoid modifying the response object that will be returned to the user
|
||||
# The logging handlers may transform usage from Responses API format (input_tokens/output_tokens)
|
||||
# to chat completion format (prompt_tokens/completion_tokens) for internal logging
|
||||
# Use model_dump + model_validate instead of deepcopy to avoid pickle errors with
|
||||
# Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192)
|
||||
logging_response = self.completed_response
|
||||
if self.completed_response is not None and hasattr(
|
||||
self.completed_response, "model_dump"
|
||||
):
|
||||
try:
|
||||
logging_response = type(self.completed_response).model_validate(
|
||||
self.completed_response.model_dump()
|
||||
)
|
||||
except Exception:
|
||||
# Fallback to original if serialization fails
|
||||
pass
|
||||
|
||||
run_async_function(
|
||||
async_function=self.logging_obj.async_success_handler,
|
||||
result=logging_response,
|
||||
start_time=self.start_time,
|
||||
end_time=datetime.now(),
|
||||
cache_hit=None,
|
||||
)
|
||||
|
||||
executor.submit(
|
||||
self.logging_obj.success_handler,
|
||||
result=logging_response,
|
||||
cache_hit=None,
|
||||
start_time=self.start_time,
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
self._run_post_success_hooks(end_time=datetime.now())
|
||||
self._log_completed_response(is_async=False)
|
||||
|
||||
|
||||
class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
||||
|
|
@ -683,90 +777,441 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
request_data: Optional[Dict[str, Any]] = None,
|
||||
call_type: Optional[str] = None,
|
||||
):
|
||||
super().__init__(
|
||||
response=response,
|
||||
transformed = responses_api_provider_config.transform_response_api_response(
|
||||
model=model,
|
||||
responses_api_provider_config=responses_api_provider_config,
|
||||
raw_response=response,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
super().__init__(
|
||||
response=httpx.Response(200),
|
||||
model=model,
|
||||
responses_api_provider_config=None,
|
||||
logging_obj=logging_obj,
|
||||
litellm_metadata=litellm_metadata,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
request_data=request_data,
|
||||
call_type=call_type,
|
||||
)
|
||||
self._set_events_from_response(transformed=transformed, logging_obj=logging_obj)
|
||||
|
||||
# one-time transform
|
||||
transformed = (
|
||||
self.responses_api_provider_config.transform_response_api_response(
|
||||
model=self.model,
|
||||
raw_response=response,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
def _set_events_from_response(
|
||||
self,
|
||||
transformed: Any,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> None:
|
||||
self._events = _build_synthetic_response_events(
|
||||
transformed=transformed,
|
||||
logging_obj=logging_obj,
|
||||
chunk_size=self.CHUNK_SIZE,
|
||||
)
|
||||
full_text = self._collect_text(transformed)
|
||||
|
||||
# build a list of 5‑char delta events
|
||||
deltas = [
|
||||
OutputTextDeltaEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
|
||||
delta=full_text[i : i + self.CHUNK_SIZE],
|
||||
item_id=transformed.id,
|
||||
output_index=0,
|
||||
content_index=0,
|
||||
)
|
||||
for i in range(0, len(full_text), self.CHUNK_SIZE)
|
||||
]
|
||||
|
||||
# Add cost to usage object if include_cost_in_streaming_usage is True
|
||||
if litellm.include_cost_in_streaming_usage and logging_obj is not None:
|
||||
usage_obj: Optional[ResponseAPIUsage] = getattr(transformed, "usage", None)
|
||||
if usage_obj is not None:
|
||||
try:
|
||||
cost: Optional[float] = logging_obj._response_cost_calculator(
|
||||
result=transformed
|
||||
)
|
||||
if cost is not None:
|
||||
setattr(usage_obj, "cost", cost)
|
||||
except Exception:
|
||||
# If cost calculation fails, continue without cost
|
||||
pass
|
||||
|
||||
# append the completed event
|
||||
self._events = deltas + [
|
||||
ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=transformed,
|
||||
)
|
||||
]
|
||||
self._idx = 0
|
||||
self.completed_response = self._events[-1]
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> ResponsesAPIStreamingResponse:
|
||||
async def __anext__(self) -> Any:
|
||||
if self._idx >= len(self._events):
|
||||
raise StopAsyncIteration
|
||||
evt = self._events[self._idx]
|
||||
self._idx += 1
|
||||
openai_types = _get_openai_response_types()
|
||||
if (
|
||||
getattr(evt, "type", None)
|
||||
== openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED
|
||||
):
|
||||
self.completed_response = evt
|
||||
self._log_completed_response(is_async=True)
|
||||
return evt
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self) -> ResponsesAPIStreamingResponse:
|
||||
def __next__(self) -> Any:
|
||||
if self._idx >= len(self._events):
|
||||
raise StopIteration
|
||||
evt = self._events[self._idx]
|
||||
self._idx += 1
|
||||
openai_types = _get_openai_response_types()
|
||||
if (
|
||||
getattr(evt, "type", None)
|
||||
== openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED
|
||||
):
|
||||
self.completed_response = evt
|
||||
self._log_completed_response(is_async=False)
|
||||
return evt
|
||||
|
||||
def _collect_text(self, resp: ResponsesAPIResponse) -> str:
|
||||
out = ""
|
||||
for out_item in resp.output:
|
||||
item_type = getattr(out_item, "type", None)
|
||||
if item_type == "message":
|
||||
for c in getattr(out_item, "content", []):
|
||||
out += c.text
|
||||
return out
|
||||
|
||||
class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
||||
def __init__(
|
||||
self,
|
||||
response: Any,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
request_data: Optional[Dict[str, Any]] = None,
|
||||
call_type: Optional[str] = None,
|
||||
):
|
||||
BaseResponsesAPIStreamingIterator.__init__(
|
||||
self,
|
||||
response=httpx.Response(200),
|
||||
model=getattr(response, "model", ""),
|
||||
responses_api_provider_config=None,
|
||||
logging_obj=logging_obj,
|
||||
litellm_metadata=None,
|
||||
custom_llm_provider="cached_response",
|
||||
request_data=request_data,
|
||||
call_type=call_type,
|
||||
)
|
||||
self._completed_response_cache_hit = True
|
||||
self._persist_completed_response_before_logging = False
|
||||
self._events: List[Any] = []
|
||||
self._idx = 0
|
||||
self._set_events_from_response(transformed=response, logging_obj=logging_obj)
|
||||
|
||||
def _set_events_from_response(
|
||||
self,
|
||||
transformed: Any,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> None:
|
||||
self._events = _build_synthetic_response_events(
|
||||
transformed=transformed,
|
||||
logging_obj=logging_obj,
|
||||
chunk_size=MockResponsesAPIStreamingIterator.CHUNK_SIZE,
|
||||
)
|
||||
self._idx = 0
|
||||
self.completed_response = self._events[-1]
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> Any:
|
||||
if self._idx >= len(self._events):
|
||||
raise StopAsyncIteration
|
||||
evt = self._events[self._idx]
|
||||
self._idx += 1
|
||||
openai_types = _get_openai_response_types()
|
||||
if (
|
||||
getattr(evt, "type", None)
|
||||
== openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED
|
||||
):
|
||||
self.completed_response = evt
|
||||
self._log_completed_response(is_async=True)
|
||||
return evt
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self) -> Any:
|
||||
if self._idx >= len(self._events):
|
||||
raise StopIteration
|
||||
evt = self._events[self._idx]
|
||||
self._idx += 1
|
||||
openai_types = _get_openai_response_types()
|
||||
if (
|
||||
getattr(evt, "type", None)
|
||||
== openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED
|
||||
):
|
||||
self.completed_response = evt
|
||||
self._log_completed_response(is_async=False)
|
||||
return evt
|
||||
|
||||
|
||||
def _dump_response_object(obj: Any) -> Dict[str, Any]:
|
||||
if hasattr(obj, "model_dump"):
|
||||
return obj.model_dump()
|
||||
if isinstance(obj, dict):
|
||||
return obj
|
||||
return {}
|
||||
|
||||
|
||||
def _build_response_status_event(
|
||||
event_type: Literal[
|
||||
"response.created",
|
||||
"response.in_progress",
|
||||
],
|
||||
transformed: Any,
|
||||
) -> Any:
|
||||
openai_types = _get_openai_response_types()
|
||||
in_progress_response = transformed.model_copy(
|
||||
deep=True,
|
||||
update={"status": "in_progress", "output": []},
|
||||
)
|
||||
if event_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED:
|
||||
return openai_types.ResponseCreatedEvent(
|
||||
type=event_type, response=in_progress_response
|
||||
)
|
||||
return openai_types.ResponseInProgressEvent(
|
||||
type=event_type, response=in_progress_response
|
||||
)
|
||||
|
||||
|
||||
def _build_content_part_done_event(
|
||||
*,
|
||||
item_id: str,
|
||||
output_index: int,
|
||||
content_index: int,
|
||||
part_payload: Dict[str, Any],
|
||||
) -> Optional[Any]:
|
||||
openai_types = _get_openai_response_types()
|
||||
part_type = part_payload.get("type")
|
||||
part: Any
|
||||
if part_type == "output_text":
|
||||
annotations = [
|
||||
openai_types.BaseLiteLLMOpenAIResponseObject(**annotation)
|
||||
for annotation in part_payload.get("annotations", []) or []
|
||||
]
|
||||
part = openai_types.ContentPartDonePartOutputText(
|
||||
type="output_text",
|
||||
text=str(part_payload.get("text") or ""),
|
||||
annotations=annotations,
|
||||
logprobs=part_payload.get("logprobs"),
|
||||
)
|
||||
elif part_type == "refusal":
|
||||
part = openai_types.ContentPartDonePartRefusal(
|
||||
type="refusal",
|
||||
refusal=str(part_payload.get("refusal") or ""),
|
||||
)
|
||||
elif part_type == "reasoning_text":
|
||||
part = openai_types.ContentPartDonePartReasoningText(
|
||||
type="reasoning_text",
|
||||
reasoning=str(part_payload.get("reasoning") or ""),
|
||||
)
|
||||
else:
|
||||
return None
|
||||
|
||||
return openai_types.ContentPartDoneEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.CONTENT_PART_DONE,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
content_index=content_index,
|
||||
part=part,
|
||||
)
|
||||
|
||||
|
||||
def _add_text_like_part_events(
|
||||
*,
|
||||
events: List[Any],
|
||||
item_id: str,
|
||||
output_index: int,
|
||||
content_index: int,
|
||||
part_payload: Dict[str, Any],
|
||||
chunk_size: int,
|
||||
) -> None:
|
||||
openai_types = _get_openai_response_types()
|
||||
part_type = part_payload.get("type")
|
||||
if part_type == "output_text":
|
||||
text = str(part_payload.get("text") or "")
|
||||
for i in range(0, len(text), chunk_size):
|
||||
events.append(
|
||||
openai_types.OutputTextDeltaEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
content_index=content_index,
|
||||
delta=text[i : i + chunk_size],
|
||||
)
|
||||
)
|
||||
for annotation_index, annotation in enumerate(
|
||||
part_payload.get("annotations", []) or []
|
||||
):
|
||||
events.append(
|
||||
openai_types.OutputTextAnnotationAddedEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
content_index=content_index,
|
||||
annotation_index=annotation_index,
|
||||
annotation=annotation,
|
||||
)
|
||||
)
|
||||
events.append(
|
||||
openai_types.OutputTextDoneEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
content_index=content_index,
|
||||
text=text,
|
||||
)
|
||||
)
|
||||
elif part_type == "refusal":
|
||||
refusal = str(part_payload.get("refusal") or "")
|
||||
for i in range(0, len(refusal), chunk_size):
|
||||
events.append(
|
||||
openai_types.RefusalDeltaEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.REFUSAL_DELTA,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
content_index=content_index,
|
||||
delta=refusal[i : i + chunk_size],
|
||||
)
|
||||
)
|
||||
events.append(
|
||||
openai_types.RefusalDoneEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.REFUSAL_DONE,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
content_index=content_index,
|
||||
refusal=refusal,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _build_synthetic_response_events(
|
||||
*,
|
||||
transformed: Any,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
chunk_size: int,
|
||||
) -> List[Any]:
|
||||
openai_types = _get_openai_response_types()
|
||||
if litellm.include_cost_in_streaming_usage and logging_obj is not None:
|
||||
usage_obj: Optional[Any] = getattr(transformed, "usage", None)
|
||||
if usage_obj is not None:
|
||||
try:
|
||||
cost: Optional[float] = logging_obj._response_cost_calculator(
|
||||
result=transformed
|
||||
)
|
||||
if cost is not None:
|
||||
setattr(usage_obj, "cost", cost)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
events: List[Any] = [
|
||||
_build_response_status_event(
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed
|
||||
),
|
||||
_build_response_status_event(
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed
|
||||
),
|
||||
]
|
||||
|
||||
sequence_number = 0
|
||||
for output_index, output_item in enumerate(
|
||||
getattr(transformed, "output", []) or []
|
||||
):
|
||||
output_item_payload = _dump_response_object(output_item)
|
||||
item_id = str(output_item_payload.get("id") or transformed.id)
|
||||
item_type = output_item_payload.get("type")
|
||||
|
||||
events.append(
|
||||
openai_types.OutputItemAddedEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
|
||||
output_index=output_index,
|
||||
item=openai_types.BaseLiteLLMOpenAIResponseObject(
|
||||
**output_item_payload
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
if item_type == "message":
|
||||
for content_index, part in enumerate(
|
||||
output_item_payload.get("content", []) or []
|
||||
):
|
||||
part_payload = _dump_response_object(part)
|
||||
events.append(
|
||||
openai_types.ContentPartAddedEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.CONTENT_PART_ADDED,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
content_index=content_index,
|
||||
part=openai_types.BaseLiteLLMOpenAIResponseObject(
|
||||
**part_payload
|
||||
),
|
||||
)
|
||||
)
|
||||
_add_text_like_part_events(
|
||||
events=events,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
content_index=content_index,
|
||||
part_payload=part_payload,
|
||||
chunk_size=chunk_size,
|
||||
)
|
||||
done_event = _build_content_part_done_event(
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
content_index=content_index,
|
||||
part_payload=part_payload,
|
||||
)
|
||||
if done_event is not None:
|
||||
events.append(done_event)
|
||||
elif item_type == "function_call":
|
||||
arguments = str(output_item_payload.get("arguments") or "")
|
||||
for i in range(0, len(arguments), chunk_size):
|
||||
events.append(
|
||||
openai_types.FunctionCallArgumentsDeltaEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
delta=arguments[i : i + chunk_size],
|
||||
)
|
||||
)
|
||||
events.append(
|
||||
openai_types.FunctionCallArgumentsDoneEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
arguments=arguments,
|
||||
)
|
||||
)
|
||||
elif item_type == "reasoning":
|
||||
for summary_index, summary in enumerate(
|
||||
output_item_payload.get("summary", []) or []
|
||||
):
|
||||
summary_payload = _dump_response_object(summary)
|
||||
summary_text = str(summary_payload.get("text") or "")
|
||||
for i in range(0, len(summary_text), chunk_size):
|
||||
events.append(
|
||||
openai_types.ReasoningSummaryTextDeltaEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
summary_index=summary_index,
|
||||
delta=summary_text[i : i + chunk_size],
|
||||
)
|
||||
)
|
||||
sequence_number += 1
|
||||
events.append(
|
||||
openai_types.ReasoningSummaryTextDoneEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
sequence_number=sequence_number,
|
||||
summary_index=summary_index,
|
||||
text=summary_text,
|
||||
)
|
||||
)
|
||||
sequence_number += 1
|
||||
events.append(
|
||||
openai_types.ReasoningSummaryPartDoneEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
sequence_number=sequence_number,
|
||||
summary_index=summary_index,
|
||||
part=openai_types.BaseLiteLLMOpenAIResponseObject(
|
||||
**summary_payload
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
sequence_number += 1
|
||||
events.append(
|
||||
openai_types.OutputItemDoneEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
|
||||
output_index=output_index,
|
||||
sequence_number=sequence_number,
|
||||
item=openai_types.BaseLiteLLMOpenAIResponseObject(
|
||||
**output_item_payload
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
events.append(
|
||||
openai_types.ResponseCompletedEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=transformed,
|
||||
)
|
||||
)
|
||||
return events
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -951,8 +1396,8 @@ class ResponsesWebSocketStreaming:
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
_RESPONSE_CREATE_PARAMS: frozenset = (
|
||||
ResponsesAPIRequestParams.__required_keys__
|
||||
| ResponsesAPIRequestParams.__optional_keys__
|
||||
_get_openai_response_types().ResponsesAPIRequestParams.__required_keys__
|
||||
| _get_openai_response_types().ResponsesAPIRequestParams.__optional_keys__
|
||||
)
|
||||
|
||||
_MANAGED_WS_SKIP_KWARGS: frozenset = frozenset(
|
||||
|
|
@ -1085,7 +1530,7 @@ class ManagedResponsesWebSocketHandler:
|
|||
|
||||
@staticmethod
|
||||
def _extract_output_messages(
|
||||
completed_event: Dict[str, Any]
|
||||
completed_event: Dict[str, Any],
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Convert the output items in a ``response.completed`` event into
|
||||
|
|
|
|||
|
|
@ -5261,11 +5261,34 @@ class Router:
|
|||
"""
|
||||
Initialize the Containers API endpoints on the router.
|
||||
|
||||
Container operations don't need model-based routing, so we call the
|
||||
original function directly with the custom_llm_provider.
|
||||
LiteLLM-managed container IDs (``cntr_...``) encode ``model_id`` and provider
|
||||
metadata. When present, decode the ID, replace ``container_id`` with the
|
||||
upstream value, and route through ``_ageneric_api_call_with_fallbacks`` so
|
||||
deployment credentials (e.g. regional ``api_base`` for Azure) match
|
||||
:meth:`_init_responses_api_endpoints`. Otherwise call the handler directly.
|
||||
"""
|
||||
if custom_llm_provider and "custom_llm_provider" not in kwargs:
|
||||
kwargs["custom_llm_provider"] = custom_llm_provider
|
||||
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
|
||||
container_id = kwargs.get("container_id")
|
||||
if isinstance(container_id, str):
|
||||
decoded = ResponsesAPIRequestUtils._decode_container_id(container_id)
|
||||
original_id = decoded.get("response_id", container_id)
|
||||
if original_id != container_id:
|
||||
kwargs["container_id"] = original_id
|
||||
decoded_provider = decoded.get("custom_llm_provider")
|
||||
if decoded_provider and kwargs.get("custom_llm_provider") == "openai":
|
||||
kwargs["custom_llm_provider"] = decoded_provider
|
||||
model_id = decoded.get("model_id")
|
||||
if model_id:
|
||||
kwargs["model"] = model_id
|
||||
return await self._ageneric_api_call_with_fallbacks(
|
||||
original_function=original_function,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
return await original_function(**kwargs)
|
||||
|
||||
async def _init_responses_api_endpoints(
|
||||
|
|
|
|||
|
|
@ -106,7 +106,8 @@ def _match_deployment(
|
|||
# check either didn't run (no request tags) or failed (step 1 returned
|
||||
# None). Block the regex path so it cannot circumvent the operator's
|
||||
# strict-tag policy.
|
||||
strict_tag_check_failed = not match_any and bool(deployment_tags)
|
||||
deployment_has_plain_tags = deployment_tags is not None and len(deployment_tags) > 0
|
||||
strict_tag_check_failed = not match_any and deployment_has_plain_tags
|
||||
if deployment_tag_regex and header_strings and not strict_tag_check_failed:
|
||||
regex_match = _is_valid_deployment_tag_regex(
|
||||
deployment_tag_regex, header_strings
|
||||
|
|
|
|||
|
|
@ -23,21 +23,28 @@ def add_model_file_id_mappings(
|
|||
healthy_deployments: Union[List[Dict], Dict], responses: List["OpenAIFileObject"]
|
||||
) -> dict:
|
||||
"""
|
||||
Create a mapping of model name to file id
|
||||
Create a mapping of model id to file id
|
||||
{
|
||||
"model_id": "file_id",
|
||||
"model_id": "file_id",
|
||||
}
|
||||
|
||||
`healthy_deployments` may be either a list of deployment dicts (multiple
|
||||
matched deployments) or a single deployment dict (when the router resolved
|
||||
a specific deployment, e.g. because the requested model matched a
|
||||
`model_info.id`). Both shapes must be handled by extracting
|
||||
`model_info.id` from each deployment.
|
||||
"""
|
||||
model_file_id_mapping = {}
|
||||
if isinstance(healthy_deployments, list):
|
||||
for deployment, response in zip(healthy_deployments, responses):
|
||||
model_file_id_mapping[deployment.get("model_info", {}).get("id")] = (
|
||||
response.id
|
||||
)
|
||||
elif isinstance(healthy_deployments, dict):
|
||||
for model_id, file_id in healthy_deployments.items():
|
||||
model_file_id_mapping[model_id] = file_id
|
||||
model_file_id_mapping: Dict[str, str] = {}
|
||||
deployments_list: List[Dict] = (
|
||||
healthy_deployments
|
||||
if isinstance(healthy_deployments, list)
|
||||
else [healthy_deployments]
|
||||
)
|
||||
for deployment, response in zip(deployments_list, responses):
|
||||
model_id = deployment.get("model_info", {}).get("id")
|
||||
if model_id is not None:
|
||||
model_file_id_mapping[model_id] = response.id
|
||||
return model_file_id_mapping
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1482,6 +1482,7 @@ class ReasoningSummaryTextDeltaEvent(BaseLiteLLMOpenAIResponseObject):
|
|||
type: Literal[ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA]
|
||||
item_id: str
|
||||
output_index: int
|
||||
summary_index: int = 0
|
||||
delta: str
|
||||
|
||||
|
||||
|
|
@ -1490,7 +1491,7 @@ class ReasoningSummaryTextDoneEvent(BaseLiteLLMOpenAIResponseObject):
|
|||
item_id: str
|
||||
output_index: int
|
||||
sequence_number: int
|
||||
summary_index: int
|
||||
summary_index: int = 0
|
||||
text: str
|
||||
|
||||
|
||||
|
|
@ -1499,7 +1500,7 @@ class ReasoningSummaryPartDoneEvent(BaseLiteLLMOpenAIResponseObject):
|
|||
item_id: str
|
||||
output_index: int
|
||||
sequence_number: int
|
||||
summary_index: int
|
||||
summary_index: int = 0
|
||||
part: BaseLiteLLMOpenAIResponseObject
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -6,6 +6,13 @@ from typing_extensions import (
|
|||
TypedDict,
|
||||
)
|
||||
|
||||
from litellm.types.llms.openai import EmbeddingInput
|
||||
|
||||
# Gemini supports nested-list inputs (e.g. [["text", "image"]]) as an explicit
|
||||
# opt-in for combined embeddings — a provider-specific extension of the
|
||||
# OpenAI-faithful EmbeddingInput shape.
|
||||
GeminiEmbeddingInput = Union[EmbeddingInput, List[List[str]]]
|
||||
|
||||
|
||||
class FunctionResponse(TypedDict):
|
||||
name: str
|
||||
|
|
|
|||
|
|
@ -33405,6 +33405,72 @@
|
|||
"source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas",
|
||||
"supports_reasoning": true
|
||||
},
|
||||
"vertex_ai/xai/grok-4.1-fast-non-reasoning": {
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"max_input_tokens": 2000000,
|
||||
"max_output_tokens": 2000000,
|
||||
"max_tokens": 2000000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-07,
|
||||
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
|
||||
"supports_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"vertex_ai/xai/grok-4.1-fast-reasoning": {
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"max_input_tokens": 2000000,
|
||||
"max_output_tokens": 2000000,
|
||||
"max_tokens": 2000000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-07,
|
||||
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"vertex_ai/xai/grok-4.20-non-reasoning": {
|
||||
"cache_read_input_token_cost": 2e-07,
|
||||
"input_cost_per_token": 2e-06,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"max_input_tokens": 2000000,
|
||||
"max_output_tokens": 2000000,
|
||||
"max_tokens": 2000000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-06,
|
||||
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
|
||||
"supports_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"vertex_ai/xai/grok-4.20-reasoning": {
|
||||
"cache_read_input_token_cost": 2e-07,
|
||||
"input_cost_per_token": 2e-06,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"max_input_tokens": 2000000,
|
||||
"max_output_tokens": 2000000,
|
||||
"max_tokens": 2000000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-06,
|
||||
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": {
|
||||
"input_cost_per_token": 2.5e-07,
|
||||
"litellm_provider": "vertex_ai-qwen_models",
|
||||
|
|
@ -34842,6 +34908,20 @@
|
|||
"supports_tool_choice": true,
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/"
|
||||
},
|
||||
"zai.glm-5": {
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3.2e-06,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/"
|
||||
},
|
||||
"zai.glm-4.7-flash": {
|
||||
"input_cost_per_token": 7e-08,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
import asyncio
|
||||
from contextlib import suppress
|
||||
from datetime import datetime
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -8,8 +11,17 @@ import pytest
|
|||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.responses import streaming_iterator as streaming_module
|
||||
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
|
||||
from litellm.types.llms.openai import ResponsesAPIStreamEvents
|
||||
from litellm.responses.streaming_iterator import (
|
||||
CachedResponsesAPIStreamingIterator,
|
||||
MockResponsesAPIStreamingIterator,
|
||||
ResponsesAPIStreamingIterator,
|
||||
SyncResponsesAPIStreamingIterator,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
ResponseCompletedEvent,
|
||||
ResponsesAPIResponse,
|
||||
ResponsesAPIStreamEvents,
|
||||
)
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
|
||||
|
|
@ -19,15 +31,19 @@ class _FakeLoggingObj:
|
|||
self.async_success_calls = 0
|
||||
self.failure_calls = 0
|
||||
self.async_failure_calls = 0
|
||||
self.last_success_kwargs = None
|
||||
self.last_async_success_kwargs = None
|
||||
self.start_time = datetime.now()
|
||||
self.model_call_details = {"litellm_params": {}}
|
||||
|
||||
# Signature alignment with Logging handlers
|
||||
def success_handler(self, *args, **kwargs):
|
||||
self.success_calls += 1
|
||||
self.last_success_kwargs = kwargs
|
||||
|
||||
async def async_success_handler(self, *args, **kwargs):
|
||||
self.async_success_calls += 1
|
||||
self.last_async_success_kwargs = kwargs
|
||||
|
||||
def failure_handler(self, *args, **kwargs):
|
||||
self.failure_calls += 1
|
||||
|
|
@ -36,6 +52,115 @@ class _FakeLoggingObj:
|
|||
self.async_failure_calls += 1
|
||||
|
||||
|
||||
def _make_completed_response(response_id: str = "resp_test") -> ResponseCompletedEvent:
|
||||
return ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=ResponsesAPIResponse(
|
||||
id=response_id,
|
||||
created_at=int(datetime.now().timestamp()),
|
||||
status="completed",
|
||||
model="test-model",
|
||||
object="response",
|
||||
output=[
|
||||
{
|
||||
"type": "message",
|
||||
"id": f"msg_{response_id}",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "output_text",
|
||||
"text": "cached streamed response",
|
||||
"annotations": [],
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_background_task_failure_logs_task_exceptions(monkeypatch):
|
||||
error_logger = MagicMock()
|
||||
monkeypatch.setattr(streaming_module.verbose_logger, "error", error_logger)
|
||||
|
||||
async def _boom():
|
||||
raise RuntimeError("boom")
|
||||
|
||||
task = asyncio.create_task(_boom())
|
||||
with suppress(RuntimeError):
|
||||
await task
|
||||
|
||||
streaming_module._log_background_task_failure(task, task_name="cache write")
|
||||
|
||||
error_logger.assert_called_once()
|
||||
assert error_logger.call_args.args == (
|
||||
"%s failed: %s",
|
||||
"cache write",
|
||||
task.exception(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_background_task_failure_ignores_cancelled_tasks(monkeypatch):
|
||||
error_logger = MagicMock()
|
||||
monkeypatch.setattr(streaming_module.verbose_logger, "error", error_logger)
|
||||
|
||||
task = asyncio.create_task(asyncio.sleep(1))
|
||||
task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
streaming_module._log_background_task_failure(task, task_name="cache write")
|
||||
|
||||
error_logger.assert_not_called()
|
||||
|
||||
|
||||
def test_content_part_done_event_supports_refusal_and_reasoning_text():
|
||||
refusal_event = streaming_module._build_content_part_done_event(
|
||||
item_id="msg_1",
|
||||
output_index=0,
|
||||
content_index=0,
|
||||
part_payload={"type": "refusal", "refusal": "no"},
|
||||
)
|
||||
reasoning_event = streaming_module._build_content_part_done_event(
|
||||
item_id="msg_1",
|
||||
output_index=0,
|
||||
content_index=1,
|
||||
part_payload={"type": "reasoning_text", "reasoning": "because"},
|
||||
)
|
||||
unsupported_event = streaming_module._build_content_part_done_event(
|
||||
item_id="msg_1",
|
||||
output_index=0,
|
||||
content_index=2,
|
||||
part_payload={"type": "image"},
|
||||
)
|
||||
|
||||
assert refusal_event.part.type == "refusal"
|
||||
assert refusal_event.part.refusal == "no"
|
||||
assert reasoning_event.part.type == "reasoning_text"
|
||||
assert reasoning_event.part.reasoning == "because"
|
||||
assert unsupported_event is None
|
||||
|
||||
|
||||
def test_dump_response_object_handles_model_and_unknown_values():
|
||||
response = ResponsesAPIResponse(
|
||||
id="resp_dump",
|
||||
created_at=int(datetime.now().timestamp()),
|
||||
status="completed",
|
||||
model="gpt-4.1-mini",
|
||||
object="response",
|
||||
output=[],
|
||||
)
|
||||
|
||||
assert streaming_module._dump_response_object(response)["id"] == "resp_dump"
|
||||
assert streaming_module._dump_response_object({"type": "message"}) == {
|
||||
"type": "message"
|
||||
}
|
||||
assert streaming_module._dump_response_object(object()) == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_streaming_triggers_hooks(monkeypatch):
|
||||
"""
|
||||
|
|
@ -167,3 +292,768 @@ async def test_responses_streaming_failure_triggers_failure_handlers():
|
|||
await asyncio.sleep(0.2)
|
||||
assert logging_obj.failure_calls >= 1
|
||||
assert logging_obj.async_failure_calls >= 1
|
||||
|
||||
|
||||
def test_process_chunk_requires_provider_config():
|
||||
iterator = ResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="test-model",
|
||||
responses_api_provider_config=None,
|
||||
logging_obj=_FakeLoggingObj(),
|
||||
request_data={"foo": "bar"},
|
||||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="responses_api_provider_config is required"):
|
||||
iterator._process_chunk(json.dumps({"type": "response.completed"}))
|
||||
|
||||
|
||||
def test_process_chunk_wraps_encrypted_content_with_model_id():
|
||||
openai_types = streaming_module._get_openai_response_types()
|
||||
|
||||
class _EncryptedConfig:
|
||||
def transform_streaming_response(self, **kwargs):
|
||||
return openai_types.OutputItemAddedEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
|
||||
output_index=0,
|
||||
item=openai_types.BaseLiteLLMOpenAIResponseObject(
|
||||
id="rs_123",
|
||||
type="reasoning",
|
||||
encrypted_content="ciphertext",
|
||||
),
|
||||
)
|
||||
|
||||
iterator = ResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="test-model",
|
||||
responses_api_provider_config=_EncryptedConfig(),
|
||||
logging_obj=_FakeLoggingObj(),
|
||||
litellm_metadata={
|
||||
"encrypted_content_affinity_enabled": True,
|
||||
"model_info": {"id": "model-123"},
|
||||
},
|
||||
request_data={"foo": "bar"},
|
||||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
|
||||
event = iterator._process_chunk(json.dumps({"type": "response.output_item.added"}))
|
||||
|
||||
assert event.item.encrypted_content.startswith("litellm_enc:")
|
||||
assert event.item.encrypted_content.endswith(";ciphertext")
|
||||
|
||||
|
||||
def test_process_chunk_completed_response_updates_id_and_usage_cost(monkeypatch):
|
||||
original_include_cost = litellm.include_cost_in_streaming_usage
|
||||
litellm.include_cost_in_streaming_usage = True
|
||||
openai_types = streaming_module._get_openai_response_types()
|
||||
|
||||
class _CompletedConfig:
|
||||
def transform_streaming_response(self, **kwargs):
|
||||
return openai_types.ResponseCompletedEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=ResponsesAPIResponse(
|
||||
id="resp_live",
|
||||
created_at=int(datetime.now().timestamp()),
|
||||
status="completed",
|
||||
model="test-model",
|
||||
object="response",
|
||||
output=[],
|
||||
usage=openai_types.ResponseAPIUsage(
|
||||
input_tokens=1,
|
||||
output_tokens=2,
|
||||
total_tokens=3,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
logging_obj = _FakeLoggingObj()
|
||||
logging_obj._response_cost_calculator = MagicMock(return_value=1.23)
|
||||
iterator = ResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="test-model",
|
||||
responses_api_provider_config=_CompletedConfig(),
|
||||
logging_obj=logging_obj,
|
||||
litellm_metadata={"model_info": {"id": "model-123"}},
|
||||
custom_llm_provider="openai",
|
||||
request_data={"foo": "bar"},
|
||||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
completion_handler = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
iterator, "_handle_logging_completed_response", completion_handler
|
||||
)
|
||||
|
||||
try:
|
||||
# Chunk must include a top-level "response" key so BaseResponsesAPIStreamingIterator
|
||||
# runs _update_responses_api_response_id_with_model_id (see streaming_iterator.py).
|
||||
event = iterator._process_chunk(
|
||||
json.dumps(
|
||||
{"type": "response.completed", "response": {"id": "resp_live"}}
|
||||
)
|
||||
)
|
||||
finally:
|
||||
litellm.include_cost_in_streaming_usage = original_include_cost
|
||||
|
||||
assert iterator.completed_response is event
|
||||
assert event.response.id != "resp_live"
|
||||
assert event.response.id.startswith("resp_")
|
||||
assert event.response.usage.cost == 1.23
|
||||
completion_handler.assert_called_once()
|
||||
|
||||
|
||||
def test_process_chunk_failed_response_triggers_failure_logging(monkeypatch):
|
||||
openai_types = streaming_module._get_openai_response_types()
|
||||
|
||||
class _FailedConfig:
|
||||
def transform_streaming_response(self, **kwargs):
|
||||
return openai_types.ResponseFailedEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED,
|
||||
response=ResponsesAPIResponse(
|
||||
id="resp_failed",
|
||||
created_at=int(datetime.now().timestamp()),
|
||||
status="failed",
|
||||
model="test-model",
|
||||
object="response",
|
||||
output=[],
|
||||
error={"message": "provider failed"},
|
||||
),
|
||||
)
|
||||
|
||||
iterator = ResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="test-model",
|
||||
responses_api_provider_config=_FailedConfig(),
|
||||
logging_obj=_FakeLoggingObj(),
|
||||
request_data={"foo": "bar"},
|
||||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
failure_handler = MagicMock()
|
||||
monkeypatch.setattr(iterator, "_handle_logging_failed_response", failure_handler)
|
||||
|
||||
event = iterator._process_chunk(json.dumps({"type": "response.failed"}))
|
||||
|
||||
assert iterator.completed_response is event
|
||||
failure_handler.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_logging_failed_response_uses_response_error_message():
|
||||
openai_types = streaming_module._get_openai_response_types()
|
||||
logging_obj = _FakeLoggingObj()
|
||||
iterator = ResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="test-model",
|
||||
responses_api_provider_config=SimpleNamespace(),
|
||||
logging_obj=logging_obj,
|
||||
request_data={"foo": "bar"},
|
||||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
iterator.completed_response = openai_types.ResponseFailedEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED,
|
||||
response=ResponsesAPIResponse(
|
||||
id="resp_failed_real",
|
||||
created_at=int(datetime.now().timestamp()),
|
||||
status="failed",
|
||||
model="test-model",
|
||||
object="response",
|
||||
output=[],
|
||||
error={"message": "provider failed"},
|
||||
),
|
||||
)
|
||||
|
||||
iterator._handle_logging_failed_response()
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
assert logging_obj.failure_calls == 1
|
||||
assert logging_obj.async_failure_calls == 1
|
||||
|
||||
|
||||
def test_process_chunk_returns_none_for_invalid_json_and_non_dict_payload():
|
||||
class _NoopConfig:
|
||||
def transform_streaming_response(self, **kwargs):
|
||||
raise AssertionError("should not be called")
|
||||
|
||||
iterator = ResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="test-model",
|
||||
responses_api_provider_config=_NoopConfig(),
|
||||
logging_obj=_FakeLoggingObj(),
|
||||
request_data={"foo": "bar"},
|
||||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
|
||||
assert iterator._process_chunk("not-json") is None
|
||||
assert iterator._process_chunk(json.dumps(["not", "a", "dict"])) is None
|
||||
|
||||
|
||||
def test_process_chunk_cost_annotation_failure_is_nonfatal(monkeypatch):
|
||||
original_include_cost = litellm.include_cost_in_streaming_usage
|
||||
litellm.include_cost_in_streaming_usage = True
|
||||
openai_types = streaming_module._get_openai_response_types()
|
||||
|
||||
class _CompletedConfig:
|
||||
def transform_streaming_response(self, **kwargs):
|
||||
return openai_types.ResponseCompletedEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=ResponsesAPIResponse(
|
||||
id="resp_cost_failure",
|
||||
created_at=int(datetime.now().timestamp()),
|
||||
status="completed",
|
||||
model="test-model",
|
||||
object="response",
|
||||
output=[],
|
||||
usage=openai_types.ResponseAPIUsage(
|
||||
input_tokens=1,
|
||||
output_tokens=2,
|
||||
total_tokens=3,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
logging_obj = _FakeLoggingObj()
|
||||
logging_obj._response_cost_calculator = MagicMock(side_effect=RuntimeError("boom"))
|
||||
iterator = ResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="test-model",
|
||||
responses_api_provider_config=_CompletedConfig(),
|
||||
logging_obj=logging_obj,
|
||||
request_data={"foo": "bar"},
|
||||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
completion_handler = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
iterator, "_handle_logging_completed_response", completion_handler
|
||||
)
|
||||
|
||||
try:
|
||||
event = iterator._process_chunk(json.dumps({"type": "response.completed"}))
|
||||
finally:
|
||||
litellm.include_cost_in_streaming_usage = original_include_cost
|
||||
|
||||
assert iterator.completed_response is event
|
||||
assert event.response.usage.cost is None
|
||||
completion_handler.assert_called_once()
|
||||
|
||||
|
||||
def test_get_completed_response_object_accepts_direct_response():
|
||||
logging_obj = _FakeLoggingObj()
|
||||
iterator = SyncResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="test-model",
|
||||
responses_api_provider_config=SimpleNamespace(),
|
||||
logging_obj=logging_obj,
|
||||
request_data={"foo": "bar"},
|
||||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
direct_response = _make_completed_response("resp_direct").response
|
||||
iterator.completed_response = direct_response
|
||||
|
||||
assert iterator._get_completed_response_object() is direct_response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_streaming_completed_event_persists_async_cache():
|
||||
logging_obj = _FakeLoggingObj()
|
||||
original_cache = litellm.cache
|
||||
litellm.cache = SimpleNamespace(
|
||||
async_add_cache=AsyncMock(),
|
||||
add_cache=MagicMock(),
|
||||
)
|
||||
caching_handler = SimpleNamespace(
|
||||
request_kwargs={
|
||||
"model": "test-model",
|
||||
"input": "hello",
|
||||
"stream": True,
|
||||
"caching": True,
|
||||
"cache_key": "stale-request-cache-key",
|
||||
"metadata": None,
|
||||
"custom_llm_provider": "openai",
|
||||
},
|
||||
preset_cache_key="responses-stream-cache-key",
|
||||
original_function=litellm.aresponses,
|
||||
async_set_cache=AsyncMock(),
|
||||
_should_store_result_in_cache=lambda original_function, kwargs: True,
|
||||
)
|
||||
logging_obj._llm_caching_handler = caching_handler
|
||||
|
||||
iterator = ResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="test-model",
|
||||
responses_api_provider_config=SimpleNamespace(),
|
||||
logging_obj=logging_obj,
|
||||
request_data=caching_handler.request_kwargs,
|
||||
call_type=CallTypes.aresponses.value,
|
||||
)
|
||||
iterator.completed_response = _make_completed_response()
|
||||
|
||||
iterator._handle_logging_completed_response()
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
litellm.cache.async_add_cache.assert_called_once()
|
||||
assert litellm.cache.async_add_cache.call_args.kwargs["stream"] is True
|
||||
assert (
|
||||
litellm.cache.async_add_cache.call_args.kwargs["cache_key"]
|
||||
== "responses-stream-cache-key"
|
||||
)
|
||||
assert "metadata" not in litellm.cache.async_add_cache.call_args.kwargs
|
||||
assert "custom_llm_provider" not in litellm.cache.async_add_cache.call_args.kwargs
|
||||
assert (
|
||||
json.loads(litellm.cache.async_add_cache.call_args.args[0])["id"]
|
||||
== iterator.completed_response.response.id
|
||||
)
|
||||
litellm.cache = original_cache
|
||||
|
||||
|
||||
def test_responses_streaming_completed_event_persists_sync_cache():
|
||||
logging_obj = _FakeLoggingObj()
|
||||
original_cache = litellm.cache
|
||||
litellm.cache = SimpleNamespace(
|
||||
async_add_cache=AsyncMock(),
|
||||
add_cache=MagicMock(),
|
||||
)
|
||||
caching_handler = SimpleNamespace(
|
||||
request_kwargs={
|
||||
"model": "test-model",
|
||||
"input": "hello",
|
||||
"stream": True,
|
||||
"caching": True,
|
||||
"cache_key": "stale-request-cache-key",
|
||||
"metadata": None,
|
||||
"custom_llm_provider": "openai",
|
||||
},
|
||||
preset_cache_key="responses-stream-cache-key",
|
||||
original_function=litellm.responses,
|
||||
sync_set_cache=MagicMock(),
|
||||
_should_store_result_in_cache=lambda original_function, kwargs: True,
|
||||
)
|
||||
logging_obj._llm_caching_handler = caching_handler
|
||||
|
||||
iterator = SyncResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="test-model",
|
||||
responses_api_provider_config=SimpleNamespace(),
|
||||
logging_obj=logging_obj,
|
||||
request_data=caching_handler.request_kwargs,
|
||||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
iterator.completed_response = _make_completed_response("resp_sync")
|
||||
|
||||
iterator._handle_logging_completed_response()
|
||||
|
||||
litellm.cache.add_cache.assert_called_once()
|
||||
assert litellm.cache.add_cache.call_args.kwargs["stream"] is True
|
||||
assert (
|
||||
litellm.cache.add_cache.call_args.kwargs["cache_key"]
|
||||
== "responses-stream-cache-key"
|
||||
)
|
||||
assert "metadata" not in litellm.cache.add_cache.call_args.kwargs
|
||||
assert "custom_llm_provider" not in litellm.cache.add_cache.call_args.kwargs
|
||||
assert (
|
||||
json.loads(litellm.cache.add_cache.call_args.args[0])["id"]
|
||||
== iterator.completed_response.response.id
|
||||
)
|
||||
litellm.cache = original_cache
|
||||
|
||||
|
||||
def test_log_completed_response_sync_direct_path(monkeypatch):
|
||||
hook_calls = {"post_call": 0, "metadata": 0}
|
||||
|
||||
async def fake_post_call(request_data, response, call_type):
|
||||
hook_calls["post_call"] += 1
|
||||
|
||||
def fake_update_metadata(**kwargs):
|
||||
hook_calls["metadata"] += 1
|
||||
|
||||
monkeypatch.setattr(
|
||||
streaming_module,
|
||||
"async_post_call_success_deployment_hook",
|
||||
fake_post_call,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
streaming_module,
|
||||
"update_response_metadata",
|
||||
fake_update_metadata,
|
||||
)
|
||||
|
||||
logging_obj = _FakeLoggingObj()
|
||||
iterator = SyncResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="test-model",
|
||||
responses_api_provider_config=SimpleNamespace(),
|
||||
logging_obj=logging_obj,
|
||||
request_data={"foo": "bar"},
|
||||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
iterator._persist_completed_response_before_logging = False
|
||||
iterator.completed_response = _make_completed_response("resp_log_sync")
|
||||
|
||||
iterator._log_completed_response(is_async=False)
|
||||
asyncio.run(asyncio.sleep(0.2))
|
||||
|
||||
assert logging_obj.success_calls == 1
|
||||
assert logging_obj.async_success_calls == 1
|
||||
assert hook_calls["post_call"] == 1
|
||||
assert hook_calls["metadata"] == 1
|
||||
|
||||
|
||||
def test_log_completed_response_falls_back_when_model_validate_fails(monkeypatch):
|
||||
class _BadSerializableResponse:
|
||||
@classmethod
|
||||
def model_validate(cls, value):
|
||||
raise RuntimeError("nope")
|
||||
|
||||
def model_dump(self):
|
||||
return {"id": "bad"}
|
||||
|
||||
logging_obj = _FakeLoggingObj()
|
||||
iterator = SyncResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="test-model",
|
||||
responses_api_provider_config=SimpleNamespace(),
|
||||
logging_obj=logging_obj,
|
||||
request_data={"foo": "bar"},
|
||||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
iterator._persist_completed_response_before_logging = False
|
||||
iterator.completed_response = _BadSerializableResponse()
|
||||
monkeypatch.setattr(iterator, "_run_post_success_hooks", MagicMock())
|
||||
|
||||
iterator._log_completed_response(is_async=False)
|
||||
asyncio.run(asyncio.sleep(0.2))
|
||||
|
||||
assert logging_obj.success_calls == 1
|
||||
assert logging_obj.async_success_calls == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"scenario",
|
||||
[
|
||||
"already_cached",
|
||||
"not_completed",
|
||||
"missing_caching_handler",
|
||||
"not_streaming",
|
||||
"store_disabled",
|
||||
"missing_cache_backend",
|
||||
],
|
||||
)
|
||||
def test_persist_completed_response_to_cache_guard_branches(monkeypatch, scenario):
|
||||
logging_obj = _FakeLoggingObj()
|
||||
iterator = SyncResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="test-model",
|
||||
responses_api_provider_config=SimpleNamespace(),
|
||||
logging_obj=logging_obj,
|
||||
request_data={"foo": "bar"},
|
||||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
openai_types = streaming_module._get_openai_response_types()
|
||||
completed_event = _make_completed_response("resp_guard")
|
||||
iterator.completed_response = completed_event
|
||||
|
||||
if scenario == "already_cached":
|
||||
iterator._completed_response_cached = True
|
||||
elif scenario == "not_completed":
|
||||
iterator.completed_response = openai_types.ResponseIncompleteEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
|
||||
response=completed_event.response,
|
||||
)
|
||||
elif scenario == "missing_caching_handler":
|
||||
logging_obj._llm_caching_handler = None
|
||||
else:
|
||||
logging_obj._llm_caching_handler = SimpleNamespace(
|
||||
request_kwargs={
|
||||
"model": "test-model",
|
||||
"input": "hello",
|
||||
"stream": scenario != "not_streaming",
|
||||
"cache_key": "request-cache-key",
|
||||
"metadata": None,
|
||||
"custom_llm_provider": "openai",
|
||||
},
|
||||
preset_cache_key=None,
|
||||
original_function=litellm.responses,
|
||||
dual_cache=None,
|
||||
_should_store_result_in_cache=lambda original_function, kwargs: (
|
||||
scenario != "store_disabled"
|
||||
),
|
||||
)
|
||||
if scenario == "missing_cache_backend":
|
||||
monkeypatch.setattr(streaming_module.litellm, "cache", None)
|
||||
else:
|
||||
monkeypatch.setattr(
|
||||
streaming_module.litellm,
|
||||
"cache",
|
||||
SimpleNamespace(add_cache=MagicMock(), async_add_cache=AsyncMock()),
|
||||
)
|
||||
|
||||
iterator._persist_completed_response_to_cache(is_async=False)
|
||||
|
||||
expected_cached_flag = scenario == "already_cached"
|
||||
assert iterator._completed_response_cached is expected_cached_flag
|
||||
|
||||
|
||||
def test_build_synthetic_response_events_covers_annotations_function_calls_and_refusals():
|
||||
original_include_cost = litellm.include_cost_in_streaming_usage
|
||||
litellm.include_cost_in_streaming_usage = True
|
||||
logging_obj = _FakeLoggingObj()
|
||||
logging_obj._response_cost_calculator = MagicMock(side_effect=RuntimeError("boom"))
|
||||
transformed = ResponsesAPIResponse(
|
||||
id="resp_events",
|
||||
created_at=int(datetime.now().timestamp()),
|
||||
status="completed",
|
||||
model="gpt-4.1-mini",
|
||||
object="response",
|
||||
output=[
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_events",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "output_text",
|
||||
"text": "hello world",
|
||||
"annotations": [{"type": "file_citation", "file_id": "file_1"}],
|
||||
},
|
||||
{
|
||||
"type": "refusal",
|
||||
"refusal": "no thanks",
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"type": "function_call",
|
||||
"id": "fc_events",
|
||||
"call_id": "call_123",
|
||||
"name": "lookup",
|
||||
"arguments": '{"id":1}',
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
try:
|
||||
events = streaming_module._build_synthetic_response_events(
|
||||
transformed=transformed,
|
||||
logging_obj=logging_obj,
|
||||
chunk_size=5,
|
||||
)
|
||||
finally:
|
||||
litellm.include_cost_in_streaming_usage = original_include_cost
|
||||
|
||||
event_types = [
|
||||
event.type.value if hasattr(event.type, "value") else str(event.type)
|
||||
for event in events
|
||||
]
|
||||
|
||||
assert "response.output_text.annotation.added" in event_types
|
||||
assert "response.refusal.delta" in event_types
|
||||
assert "response.refusal.done" in event_types
|
||||
assert "response.function_call_arguments.delta" in event_types
|
||||
assert "response.function_call_arguments.done" in event_types
|
||||
assert event_types[-1] == "response.completed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mock_responses_streaming_iterator_async_iteration_logs_completion(
|
||||
monkeypatch,
|
||||
):
|
||||
hook_calls = {"post_call": 0, "metadata": 0}
|
||||
|
||||
async def fake_post_call(request_data, response, call_type):
|
||||
hook_calls["post_call"] += 1
|
||||
|
||||
def fake_update_metadata(**kwargs):
|
||||
hook_calls["metadata"] += 1
|
||||
|
||||
monkeypatch.setattr(
|
||||
streaming_module,
|
||||
"async_post_call_success_deployment_hook",
|
||||
fake_post_call,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
streaming_module,
|
||||
"update_response_metadata",
|
||||
fake_update_metadata,
|
||||
)
|
||||
|
||||
class _MockTransformConfig:
|
||||
def transform_response_api_response(self, **kwargs):
|
||||
return _make_completed_response("resp_mock").response
|
||||
|
||||
logging_obj = _FakeLoggingObj()
|
||||
|
||||
iterator = MockResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="test-model",
|
||||
responses_api_provider_config=_MockTransformConfig(),
|
||||
logging_obj=logging_obj,
|
||||
request_data={"model": "test-model", "stream": True},
|
||||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
|
||||
streamed_events = [event async for event in iterator]
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
assert streamed_events[0].type == ResponsesAPIStreamEvents.RESPONSE_CREATED
|
||||
assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
|
||||
assert logging_obj.success_calls == 1
|
||||
assert logging_obj.async_success_calls == 1
|
||||
assert hook_calls["post_call"] == 1
|
||||
assert hook_calls["metadata"] == 1
|
||||
|
||||
|
||||
def test_mock_responses_streaming_iterator_sync_iteration_logs_completion(monkeypatch):
|
||||
hook_calls = {"post_call": 0, "metadata": 0}
|
||||
|
||||
async def fake_post_call(request_data, response, call_type):
|
||||
hook_calls["post_call"] += 1
|
||||
|
||||
def fake_update_metadata(**kwargs):
|
||||
hook_calls["metadata"] += 1
|
||||
|
||||
monkeypatch.setattr(
|
||||
streaming_module,
|
||||
"async_post_call_success_deployment_hook",
|
||||
fake_post_call,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
streaming_module,
|
||||
"update_response_metadata",
|
||||
fake_update_metadata,
|
||||
)
|
||||
|
||||
class _MockTransformConfig:
|
||||
def transform_response_api_response(self, **kwargs):
|
||||
return _make_completed_response("resp_mock_sync").response
|
||||
|
||||
logging_obj = _FakeLoggingObj()
|
||||
iterator = MockResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="test-model",
|
||||
responses_api_provider_config=_MockTransformConfig(),
|
||||
logging_obj=logging_obj,
|
||||
request_data={"model": "test-model", "stream": True},
|
||||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
|
||||
streamed_events = list(iterator)
|
||||
asyncio.run(asyncio.sleep(0.2))
|
||||
|
||||
assert streamed_events[0].type == ResponsesAPIStreamEvents.RESPONSE_CREATED
|
||||
assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
|
||||
assert logging_obj.success_calls == 1
|
||||
assert logging_obj.async_success_calls == 1
|
||||
assert hook_calls["post_call"] == 1
|
||||
assert hook_calls["metadata"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cached_responses_stream_async_hit_triggers_success_callbacks(
|
||||
monkeypatch,
|
||||
):
|
||||
hook_calls = {"post_call": 0, "metadata": 0}
|
||||
|
||||
async def fake_post_call(request_data, response, call_type):
|
||||
hook_calls["post_call"] += 1
|
||||
|
||||
def fake_update_metadata(**kwargs):
|
||||
hook_calls["metadata"] += 1
|
||||
|
||||
monkeypatch.setattr(
|
||||
streaming_module,
|
||||
"async_post_call_success_deployment_hook",
|
||||
fake_post_call,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
streaming_module,
|
||||
"update_response_metadata",
|
||||
fake_update_metadata,
|
||||
)
|
||||
|
||||
logging_obj = _FakeLoggingObj()
|
||||
original_cache = litellm.cache
|
||||
litellm.cache = SimpleNamespace(
|
||||
async_add_cache=AsyncMock(),
|
||||
add_cache=MagicMock(),
|
||||
)
|
||||
logging_obj._llm_caching_handler = SimpleNamespace(
|
||||
request_kwargs={"model": "test-model", "input": "hello", "stream": True},
|
||||
preset_cache_key="responses-stream-cache-key",
|
||||
original_function=litellm.aresponses,
|
||||
_should_store_result_in_cache=lambda original_function, kwargs: True,
|
||||
)
|
||||
|
||||
iterator = CachedResponsesAPIStreamingIterator(
|
||||
response=_make_completed_response("resp_cached_async").response,
|
||||
logging_obj=logging_obj,
|
||||
request_data={"model": "test-model", "input": "hello", "stream": True},
|
||||
call_type=CallTypes.aresponses.value,
|
||||
)
|
||||
|
||||
streamed_events = [event async for event in iterator]
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
|
||||
assert logging_obj.success_calls == 1
|
||||
assert logging_obj.async_success_calls == 1
|
||||
assert logging_obj.last_success_kwargs["cache_hit"] is True
|
||||
assert logging_obj.last_async_success_kwargs["cache_hit"] is True
|
||||
assert hook_calls["post_call"] == 1
|
||||
assert hook_calls["metadata"] == 1
|
||||
litellm.cache.async_add_cache.assert_not_called()
|
||||
litellm.cache.add_cache.assert_not_called()
|
||||
litellm.cache = original_cache
|
||||
|
||||
|
||||
def test_cached_responses_stream_sync_hit_triggers_success_callbacks(monkeypatch):
|
||||
hook_calls = {"post_call": 0, "metadata": 0}
|
||||
|
||||
async def fake_post_call(request_data, response, call_type):
|
||||
hook_calls["post_call"] += 1
|
||||
|
||||
def fake_update_metadata(**kwargs):
|
||||
hook_calls["metadata"] += 1
|
||||
|
||||
monkeypatch.setattr(
|
||||
streaming_module,
|
||||
"async_post_call_success_deployment_hook",
|
||||
fake_post_call,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
streaming_module,
|
||||
"update_response_metadata",
|
||||
fake_update_metadata,
|
||||
)
|
||||
|
||||
logging_obj = _FakeLoggingObj()
|
||||
original_cache = litellm.cache
|
||||
litellm.cache = SimpleNamespace(
|
||||
async_add_cache=AsyncMock(),
|
||||
add_cache=MagicMock(),
|
||||
)
|
||||
logging_obj._llm_caching_handler = SimpleNamespace(
|
||||
request_kwargs={"model": "test-model", "input": "hello", "stream": True},
|
||||
preset_cache_key="responses-stream-cache-key",
|
||||
original_function=litellm.responses,
|
||||
_should_store_result_in_cache=lambda original_function, kwargs: True,
|
||||
)
|
||||
|
||||
iterator = CachedResponsesAPIStreamingIterator(
|
||||
response=_make_completed_response("resp_cached_sync").response,
|
||||
logging_obj=logging_obj,
|
||||
request_data={"model": "test-model", "input": "hello", "stream": True},
|
||||
call_type=CallTypes.responses.value,
|
||||
)
|
||||
|
||||
streamed_events = list(iterator)
|
||||
asyncio.run(asyncio.sleep(0.2))
|
||||
|
||||
assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
|
||||
assert logging_obj.success_calls == 1
|
||||
assert logging_obj.async_success_calls == 1
|
||||
assert logging_obj.last_success_kwargs["cache_hit"] is True
|
||||
assert logging_obj.last_async_success_kwargs["cache_hit"] is True
|
||||
assert hook_calls["post_call"] == 1
|
||||
assert hook_calls["metadata"] == 1
|
||||
litellm.cache.async_add_cache.assert_not_called()
|
||||
litellm.cache.add_cache.assert_not_called()
|
||||
litellm.cache = original_cache
|
||||
|
|
|
|||
|
|
@ -19,9 +19,14 @@ import pytest
|
|||
import litellm
|
||||
from litellm import aembedding, completion, embedding, aresponses, responses
|
||||
from litellm.caching.caching import Cache
|
||||
from litellm.responses.streaming_iterator import CachedResponsesAPIStreamingIterator
|
||||
|
||||
from unittest.mock import AsyncMock, patch, MagicMock
|
||||
from litellm.caching.caching_handler import LLMCachingHandler, CachingHandlerResponse
|
||||
from litellm.caching.caching_handler import (
|
||||
LLMCachingHandler,
|
||||
CachingHandlerResponse,
|
||||
_should_defer_streaming_cache_hit_callbacks,
|
||||
)
|
||||
from litellm.caching.caching import LiteLLMCacheType
|
||||
from litellm.types.utils import CallTypes
|
||||
from litellm.types.rerank import RerankResponse
|
||||
|
|
@ -627,6 +632,55 @@ async def test_async_responses_api_caching():
|
|||
assert cached_response.cached_result._hidden_params["cache_hit"] == True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get_cache_updates_request_kwargs_for_streaming_responses():
|
||||
"""
|
||||
Ensure streamed responses retain the normalized lookup kwargs so a later
|
||||
cache write can reuse the exact cache key from the read path.
|
||||
"""
|
||||
setup_cache()
|
||||
|
||||
caching_handler = LLMCachingHandler(
|
||||
original_function=aresponses,
|
||||
request_kwargs={"stale": True},
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
|
||||
logging_obj = LiteLLMLogging(
|
||||
litellm_call_id=str(datetime.now()),
|
||||
call_type=CallTypes.aresponses.value,
|
||||
model="gpt-4o",
|
||||
messages=[],
|
||||
function_id=str(uuid.uuid4()),
|
||||
stream=True,
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
|
||||
kwargs = {
|
||||
"model": "gpt-4o",
|
||||
"input": "hello",
|
||||
"stream": True,
|
||||
"caching": True,
|
||||
}
|
||||
|
||||
await caching_handler._async_get_cache(
|
||||
model="gpt-4o",
|
||||
original_function=aresponses,
|
||||
logging_obj=logging_obj,
|
||||
start_time=datetime.now(),
|
||||
call_type=CallTypes.aresponses.value,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
assert "stale" not in caching_handler.request_kwargs
|
||||
assert caching_handler.request_kwargs["model"] == "gpt-4o"
|
||||
assert caching_handler.request_kwargs["input"] == "hello"
|
||||
assert caching_handler.request_kwargs["stream"] is True
|
||||
assert caching_handler.request_kwargs["cache_key"] == litellm.cache.get_cache_key(
|
||||
**caching_handler.request_kwargs
|
||||
)
|
||||
|
||||
|
||||
def test_sync_responses_api_caching():
|
||||
"""
|
||||
Test that synchronous responses API calls are properly cached and retrieved.
|
||||
|
|
@ -769,6 +823,339 @@ def test_convert_cached_responses_api_result_to_model_response():
|
|||
assert len(result.output) == 1
|
||||
|
||||
|
||||
def test_sync_get_cache_does_not_eagerly_log_streaming_responses_hits():
|
||||
litellm.set_verbose = True
|
||||
setup_cache()
|
||||
caching_handler = LLMCachingHandler(
|
||||
original_function=responses, request_kwargs={}, start_time=datetime.now()
|
||||
)
|
||||
|
||||
original_model = "gpt-4o"
|
||||
responses_api_response = ResponsesAPIResponse(
|
||||
id="resp_stream_sync_hit",
|
||||
created_at=int(time.time()),
|
||||
status="completed",
|
||||
model=original_model,
|
||||
object="response",
|
||||
output=[
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_stream_sync_hit",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "output_text",
|
||||
"text": "Sync streamed cache hit response.",
|
||||
"annotations": [],
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
logging_obj = LiteLLMLogging(
|
||||
litellm_call_id=str(datetime.now()),
|
||||
call_type=CallTypes.responses.value,
|
||||
model=original_model,
|
||||
messages=[],
|
||||
function_id=str(uuid.uuid4()),
|
||||
stream=True,
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
|
||||
|
||||
kwargs = {
|
||||
"model": original_model,
|
||||
"input": "Tell me a cached story",
|
||||
"stream": True,
|
||||
"caching": True,
|
||||
}
|
||||
|
||||
caching_handler.sync_set_cache(result=responses_api_response, kwargs=kwargs)
|
||||
time.sleep(0.2)
|
||||
|
||||
cached_response = caching_handler._sync_get_cache(
|
||||
model=original_model,
|
||||
original_function=responses,
|
||||
logging_obj=logging_obj,
|
||||
start_time=datetime.now(),
|
||||
call_type=CallTypes.responses.value,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
assert cached_response.cached_result is not None
|
||||
assert isinstance(
|
||||
cached_response.cached_result, CachedResponsesAPIStreamingIterator
|
||||
)
|
||||
logging_obj.handle_sync_success_callbacks_for_async_calls.assert_not_called()
|
||||
|
||||
|
||||
def test_sync_get_cache_defers_streaming_completion_hit_callbacks():
|
||||
litellm.set_verbose = True
|
||||
setup_cache()
|
||||
caching_handler = LLMCachingHandler(
|
||||
original_function=completion, request_kwargs={}, start_time=datetime.now()
|
||||
)
|
||||
|
||||
original_model = "gpt-4o"
|
||||
logging_obj = LiteLLMLogging(
|
||||
litellm_call_id=str(datetime.now()),
|
||||
call_type=CallTypes.completion.value,
|
||||
model=original_model,
|
||||
messages=[],
|
||||
function_id=str(uuid.uuid4()),
|
||||
stream=True,
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
|
||||
|
||||
kwargs = {
|
||||
"model": original_model,
|
||||
"messages": [{"role": "user", "content": "Tell me a cached joke"}],
|
||||
"stream": True,
|
||||
"caching": True,
|
||||
}
|
||||
|
||||
caching_handler.sync_set_cache(result=chat_completion_response, kwargs=kwargs)
|
||||
time.sleep(0.2)
|
||||
|
||||
cached_response = caching_handler._sync_get_cache(
|
||||
model=original_model,
|
||||
original_function=completion,
|
||||
logging_obj=logging_obj,
|
||||
start_time=datetime.now(),
|
||||
call_type=CallTypes.completion.value,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
assert cached_response.cached_result is not None
|
||||
logging_obj.handle_sync_success_callbacks_for_async_calls.assert_not_called()
|
||||
|
||||
|
||||
def test_should_defer_streaming_cache_hit_callbacks_for_any_streaming_request():
|
||||
assert (
|
||||
_should_defer_streaming_cache_hit_callbacks(
|
||||
kwargs={"stream": True},
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
_should_defer_streaming_cache_hit_callbacks(
|
||||
kwargs={"stream": False},
|
||||
)
|
||||
is False
|
||||
)
|
||||
assert (
|
||||
_should_defer_streaming_cache_hit_callbacks(
|
||||
kwargs={},
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get_cache_defers_streaming_completion_hit_callbacks():
|
||||
litellm.set_verbose = True
|
||||
setup_cache()
|
||||
caching_handler = LLMCachingHandler(
|
||||
original_function=completion, request_kwargs={}, start_time=datetime.now()
|
||||
)
|
||||
|
||||
original_model = "gpt-4o"
|
||||
kwargs = {
|
||||
"model": original_model,
|
||||
"messages": [{"role": "user", "content": "Tell me a cached joke"}],
|
||||
"stream": True,
|
||||
"caching": True,
|
||||
}
|
||||
|
||||
await caching_handler.async_set_cache(
|
||||
result=chat_completion_response,
|
||||
original_function=litellm.acompletion,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
logging_obj = LiteLLMLogging(
|
||||
litellm_call_id=str(datetime.now()),
|
||||
call_type=CallTypes.acompletion.value,
|
||||
model=original_model,
|
||||
messages=[],
|
||||
function_id=str(uuid.uuid4()),
|
||||
stream=True,
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
caching_handler._async_log_cache_hit_on_callbacks = MagicMock()
|
||||
|
||||
cached_response = await caching_handler._async_get_cache(
|
||||
model=original_model,
|
||||
original_function=litellm.acompletion,
|
||||
logging_obj=logging_obj,
|
||||
start_time=datetime.now(),
|
||||
call_type=CallTypes.acompletion.value,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
assert cached_response is not None
|
||||
assert cached_response.cached_result is not None
|
||||
caching_handler._async_log_cache_hit_on_callbacks.assert_not_called()
|
||||
|
||||
|
||||
def test_convert_cached_streaming_responses_result_to_iterator():
|
||||
"""
|
||||
Test that cached streaming Responses results are replayed through a synthetic
|
||||
streaming iterator instead of being returned as a full response object.
|
||||
"""
|
||||
caching_handler = LLMCachingHandler(
|
||||
original_function=responses, request_kwargs={}, start_time=datetime.now()
|
||||
)
|
||||
|
||||
logging_obj = LiteLLMLogging(
|
||||
litellm_call_id=str(datetime.now()),
|
||||
call_type=CallTypes.responses.value,
|
||||
model="gpt-4o",
|
||||
messages=[],
|
||||
function_id=str(uuid.uuid4()),
|
||||
stream=True,
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
|
||||
cached_result = {
|
||||
"id": "resp_stream_cache_test",
|
||||
"created_at": int(time.time()),
|
||||
"status": "completed",
|
||||
"model": "gpt-4o",
|
||||
"object": "response",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_stream_cache_test",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "output_text",
|
||||
"text": "Streaming cache replay test.",
|
||||
"annotations": [],
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
result = caching_handler._convert_cached_result_to_model_response(
|
||||
cached_result=cached_result,
|
||||
call_type=CallTypes.responses.value,
|
||||
kwargs={"model": "gpt-4o", "input": "test", "stream": True},
|
||||
logging_obj=logging_obj,
|
||||
model="gpt-4o",
|
||||
args=(),
|
||||
)
|
||||
|
||||
assert isinstance(result, CachedResponsesAPIStreamingIterator)
|
||||
assert result.completed_response is not None
|
||||
assert result.completed_response.response.id == cached_result["id"]
|
||||
|
||||
streamed_events = list(result)
|
||||
assert streamed_events[0].type == "response.created"
|
||||
assert streamed_events[1].type == "response.in_progress"
|
||||
assert streamed_events[2].type == "response.output_item.added"
|
||||
assert streamed_events[3].type == "response.content_part.added"
|
||||
assert streamed_events[-4].type == "response.output_text.done"
|
||||
assert streamed_events[-3].type == "response.content_part.done"
|
||||
assert streamed_events[-2].type == "response.output_item.done"
|
||||
assert streamed_events[-1].type == "response.completed"
|
||||
assert streamed_events[-1].response.id == cached_result["id"]
|
||||
assert streamed_events[-1].response.output[0].content[0].text == (
|
||||
"Streaming cache replay test."
|
||||
)
|
||||
|
||||
|
||||
def test_convert_cached_streaming_reasoning_result_to_iterator():
|
||||
caching_handler = LLMCachingHandler(
|
||||
original_function=responses, request_kwargs={}, start_time=datetime.now()
|
||||
)
|
||||
|
||||
logging_obj = LiteLLMLogging(
|
||||
litellm_call_id=str(datetime.now()),
|
||||
call_type=CallTypes.responses.value,
|
||||
model="gpt-4o",
|
||||
messages=[],
|
||||
function_id=str(uuid.uuid4()),
|
||||
stream=True,
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
|
||||
cached_result = {
|
||||
"id": "resp_stream_reasoning_cache_test",
|
||||
"created_at": int(time.time()),
|
||||
"status": "completed",
|
||||
"model": "gpt-4o",
|
||||
"object": "response",
|
||||
"output": [
|
||||
{
|
||||
"type": "reasoning",
|
||||
"id": "rs_stream_cache_test",
|
||||
"summary": [
|
||||
{
|
||||
"type": "summary_text",
|
||||
"text": "Cached reasoning summary.",
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
result = caching_handler._convert_cached_result_to_model_response(
|
||||
cached_result=cached_result,
|
||||
call_type=CallTypes.responses.value,
|
||||
kwargs={"model": "gpt-4o", "input": "test", "stream": True},
|
||||
logging_obj=logging_obj,
|
||||
model="gpt-4o",
|
||||
args=(),
|
||||
)
|
||||
|
||||
assert isinstance(result, CachedResponsesAPIStreamingIterator)
|
||||
|
||||
streamed_events = list(result)
|
||||
streamed_event_types = [
|
||||
event.type.value if hasattr(event.type, "value") else str(event.type)
|
||||
for event in streamed_events
|
||||
]
|
||||
|
||||
assert streamed_event_types[:3] == [
|
||||
"response.created",
|
||||
"response.in_progress",
|
||||
"response.output_item.added",
|
||||
]
|
||||
assert streamed_event_types[-4:] == [
|
||||
"response.reasoning_summary_text.done",
|
||||
"response.reasoning_summary_part.done",
|
||||
"response.output_item.done",
|
||||
"response.completed",
|
||||
]
|
||||
assert streamed_event_types.count("response.reasoning_summary_text.delta") >= 1
|
||||
|
||||
delta_events = [
|
||||
event
|
||||
for event in streamed_events
|
||||
if (event.type.value if hasattr(event.type, "value") else str(event.type))
|
||||
== "response.reasoning_summary_text.delta"
|
||||
]
|
||||
text_done_event = streamed_events[-4]
|
||||
part_done_event = streamed_events[-3]
|
||||
output_item_done_event = streamed_events[-2]
|
||||
|
||||
assert all(delta_event.summary_index == 0 for delta_event in delta_events)
|
||||
assert text_done_event.text == "Cached reasoning summary."
|
||||
assert text_done_event.summary_index == 0
|
||||
assert part_done_event.part.type == "summary_text"
|
||||
assert part_done_event.part.text == "Cached reasoning summary."
|
||||
assert output_item_done_event.item.type == "reasoning"
|
||||
assert output_item_done_event.item.summary[0]["text"] == "Cached reasoning summary."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_api_cache_with_different_inputs():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -477,3 +477,4 @@ def test_get_llm_provider_use_proxy_arg_true_with_direct_args():
|
|||
assert provider == "litellm_proxy"
|
||||
assert key == arg_api_key # Should use the argument key
|
||||
assert base == arg_api_base # Should use the argument base
|
||||
|
||||
|
|
|
|||
141
tests/local_testing/test_responses_stream_cache_keys.py
Normal file
141
tests/local_testing/test_responses_stream_cache_keys.py
Normal file
|
|
@ -0,0 +1,141 @@
|
|||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import aresponses
|
||||
from litellm._uuid import uuid
|
||||
from litellm.caching.caching_handler import LLMCachingHandler
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.types.llms import openai as openai_types
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get_cache_reuses_preset_cache_key_for_responses():
|
||||
caching_handler = LLMCachingHandler(
|
||||
original_function=aresponses,
|
||||
request_kwargs={},
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
logging_obj = LiteLLMLogging(
|
||||
litellm_call_id=str(datetime.now()),
|
||||
call_type=CallTypes.aresponses.value,
|
||||
model="gpt-4.1-mini",
|
||||
messages=[],
|
||||
function_id=str(uuid.uuid4()),
|
||||
stream=True,
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
|
||||
original_cache = litellm.cache
|
||||
mock_cache = MagicMock()
|
||||
mock_cache.supported_call_types = [CallTypes.aresponses.value]
|
||||
mock_cache._supports_async.return_value = True
|
||||
mock_cache.get_cache_key.return_value = "responses-stream-cache-key"
|
||||
mock_cache.async_get_cache = AsyncMock(return_value=None)
|
||||
litellm.cache = mock_cache
|
||||
|
||||
kwargs = {
|
||||
"model": "gpt-4.1-mini",
|
||||
"input": "hello",
|
||||
"stream": True,
|
||||
"litellm_params": {},
|
||||
}
|
||||
await caching_handler._async_get_cache(
|
||||
model="gpt-4.1-mini",
|
||||
original_function=aresponses,
|
||||
logging_obj=logging_obj,
|
||||
start_time=datetime.now(),
|
||||
call_type=CallTypes.aresponses.value,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
assert caching_handler.preset_cache_key == "responses-stream-cache-key"
|
||||
mock_cache.async_get_cache.assert_awaited_once()
|
||||
assert (
|
||||
mock_cache.async_get_cache.call_args.kwargs["cache_key"]
|
||||
== "responses-stream-cache-key"
|
||||
)
|
||||
|
||||
litellm.cache = original_cache
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get_cache_falls_back_to_sync_cache_for_responses():
|
||||
caching_handler = LLMCachingHandler(
|
||||
original_function=aresponses,
|
||||
request_kwargs={},
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
logging_obj = LiteLLMLogging(
|
||||
litellm_call_id=str(datetime.now()),
|
||||
call_type=CallTypes.aresponses.value,
|
||||
model="gpt-4.1-mini",
|
||||
messages=[],
|
||||
function_id=str(uuid.uuid4()),
|
||||
stream=True,
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
|
||||
original_cache = litellm.cache
|
||||
mock_cache = MagicMock()
|
||||
mock_cache.supported_call_types = [CallTypes.aresponses.value]
|
||||
mock_cache._supports_async.return_value = False
|
||||
mock_cache.get_cache_key.return_value = "responses-stream-cache-key"
|
||||
mock_cache.get_cache.return_value = None
|
||||
litellm.cache = mock_cache
|
||||
|
||||
kwargs = {
|
||||
"model": "gpt-4.1-mini",
|
||||
"input": "hello",
|
||||
"stream": True,
|
||||
"litellm_params": {},
|
||||
}
|
||||
await caching_handler._async_get_cache(
|
||||
model="gpt-4.1-mini",
|
||||
original_function=aresponses,
|
||||
logging_obj=logging_obj,
|
||||
start_time=datetime.now(),
|
||||
call_type=CallTypes.aresponses.value,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
assert caching_handler.preset_cache_key == "responses-stream-cache-key"
|
||||
mock_cache.get_cache.assert_called_once()
|
||||
assert mock_cache.get_cache.call_args.kwargs["cache_key"] == (
|
||||
"responses-stream-cache-key"
|
||||
)
|
||||
|
||||
litellm.cache = original_cache
|
||||
|
||||
|
||||
def test_reasoning_summary_events_default_summary_index():
|
||||
delta_event = openai_types.ReasoningSummaryTextDeltaEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA,
|
||||
item_id="rs_1",
|
||||
output_index=0,
|
||||
delta="abc",
|
||||
)
|
||||
text_done_event = openai_types.ReasoningSummaryTextDoneEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE,
|
||||
item_id="rs_1",
|
||||
output_index=0,
|
||||
sequence_number=1,
|
||||
text="abc",
|
||||
)
|
||||
part_done_event = openai_types.ReasoningSummaryPartDoneEvent(
|
||||
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE,
|
||||
item_id="rs_1",
|
||||
output_index=0,
|
||||
sequence_number=2,
|
||||
part=openai_types.BaseLiteLLMOpenAIResponseObject(
|
||||
type="summary_text",
|
||||
text="abc",
|
||||
),
|
||||
)
|
||||
|
||||
assert delta_event.summary_index == 0
|
||||
assert text_done_event.summary_index == 0
|
||||
assert part_done_event.summary_index == 0
|
||||
|
|
@ -134,3 +134,62 @@ def test_is_assemblyai_route():
|
|||
== False
|
||||
)
|
||||
assert handler.is_assemblyai_route("") == False
|
||||
|
||||
|
||||
# --- Security: SSRF via transcript_id path traversal ---
|
||||
|
||||
|
||||
def test_get_assembly_transcript_rejects_slash_in_id(assembly_handler):
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials",
|
||||
return_value="test-key",
|
||||
):
|
||||
with pytest.raises(ValueError, match="disallowed characters"):
|
||||
assembly_handler._get_assembly_transcript("../../admin/credentials")
|
||||
|
||||
|
||||
def test_get_assembly_transcript_rejects_dotdot_in_id(assembly_handler):
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials",
|
||||
return_value="test-key",
|
||||
):
|
||||
with pytest.raises(ValueError, match="disallowed characters"):
|
||||
assembly_handler._get_assembly_transcript("..evil")
|
||||
|
||||
|
||||
def test_get_assembly_transcript_rejects_fragment_in_id(assembly_handler):
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials",
|
||||
return_value="test-key",
|
||||
):
|
||||
with pytest.raises(ValueError, match="disallowed characters"):
|
||||
assembly_handler._get_assembly_transcript("abc#suffix")
|
||||
|
||||
|
||||
def test_get_assembly_transcript_rejects_query_in_id(assembly_handler):
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials",
|
||||
return_value="test-key",
|
||||
):
|
||||
with pytest.raises(ValueError, match="disallowed characters"):
|
||||
assembly_handler._get_assembly_transcript("abc?x=1")
|
||||
|
||||
|
||||
def test_get_assembly_transcript_allows_valid_id(
|
||||
assembly_handler, mock_transcript_response
|
||||
):
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials",
|
||||
return_value="test-key",
|
||||
):
|
||||
with patch("httpx.get") as mock_get:
|
||||
mock_get.return_value.json.return_value = mock_transcript_response
|
||||
mock_get.return_value.raise_for_status.return_value = None
|
||||
|
||||
transcript = assembly_handler._get_assembly_transcript(
|
||||
"abc123-valid-id_xyz"
|
||||
)
|
||||
assert transcript == mock_transcript_response
|
||||
called_url = mock_get.call_args[0][0]
|
||||
assert "abc123-valid-id_xyz" in called_url
|
||||
assert ".." not in called_url
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ import httpx
|
|||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import get_end_user_object
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_EndUserTable,
|
||||
LiteLLM_BudgetTable,
|
||||
|
|
@ -48,9 +49,15 @@ async def test_get_end_user_object(customer_spend, customer_budget):
|
|||
litellm_budget_table=_budget,
|
||||
blocked=False,
|
||||
)
|
||||
_cache = DualCache()
|
||||
# UserApiKeyCache applies model_type on get/set; plain DualCache returns raw dicts
|
||||
# and breaks get_end_user_object's typed async_get_cache path.
|
||||
_cache = UserApiKeyCache()
|
||||
_key = "end_user_id:{}".format(end_user_id)
|
||||
_cache.set_cache(key=_key, value=end_user_obj.model_dump())
|
||||
await _cache.async_set_cache(
|
||||
key=_key,
|
||||
value=end_user_obj,
|
||||
model_type=LiteLLM_EndUserTable,
|
||||
)
|
||||
try:
|
||||
await get_end_user_object(
|
||||
end_user_id=end_user_id,
|
||||
|
|
|
|||
|
|
@ -268,7 +268,12 @@ async def test_aaauser_personal_budgets(key_ownership):
|
|||
|
||||
test_user_cache = getattr(litellm.proxy.proxy_server, "user_api_key_cache")
|
||||
|
||||
assert test_user_cache.get_cache(key=hash_token(user_key)) == valid_token
|
||||
assert (
|
||||
test_user_cache.get_cache(
|
||||
key=hash_token(user_key), model_type=UserAPIKeyAuth
|
||||
)
|
||||
== valid_token
|
||||
)
|
||||
|
||||
try:
|
||||
await user_api_key_auth(request=request, api_key="Bearer " + user_key)
|
||||
|
|
|
|||
|
|
@ -1110,7 +1110,7 @@ def test_initialize_skills_endpoints():
|
|||
async def test_init_containers_api_endpoints():
|
||||
"""
|
||||
Test that _init_containers_api_endpoints calls the original function
|
||||
directly without model-based routing.
|
||||
directly when there is no managed container ID (no embedded model_id).
|
||||
"""
|
||||
router = Router(model_list=[])
|
||||
|
||||
|
|
@ -1127,3 +1127,112 @@ async def test_init_containers_api_endpoints():
|
|||
custom_llm_provider="openai", name="Test Container"
|
||||
)
|
||||
assert result == mock_response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_containers_api_endpoints_managed_id_routes_via_generic_fallbacks():
|
||||
"""
|
||||
Managed ``cntr_`` IDs embed ``model_id``; router should decode and use
|
||||
``_ageneric_api_call_with_fallbacks`` so deployment credentials apply.
|
||||
"""
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "azure-router-model",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-4",
|
||||
"api_key": "fake-key",
|
||||
"api_base": "https://westus.api.cognitive.microsoft.com",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
router._ageneric_api_call_with_fallbacks = AsyncMock()
|
||||
|
||||
managed_id = ResponsesAPIRequestUtils._build_container_id(
|
||||
custom_llm_provider="azure",
|
||||
model_id="azure-router-model",
|
||||
container_id="cfile_upstream_abc",
|
||||
)
|
||||
|
||||
await router._init_containers_api_endpoints(
|
||||
original_function=AsyncMock(),
|
||||
custom_llm_provider="openai",
|
||||
container_id=managed_id,
|
||||
file_id="cfile_xyz",
|
||||
)
|
||||
|
||||
router._ageneric_api_call_with_fallbacks.assert_called_once()
|
||||
call_kw = router._ageneric_api_call_with_fallbacks.call_args.kwargs
|
||||
assert call_kw["model"] == "azure-router-model"
|
||||
assert call_kw["container_id"] == "cfile_upstream_abc"
|
||||
assert call_kw["file_id"] == "cfile_xyz"
|
||||
assert call_kw["custom_llm_provider"] == "azure"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_containers_api_endpoints_managed_id_without_model_id_unwraps():
|
||||
"""
|
||||
Managed ``cntr_`` IDs may be encoded with an empty ``model_id`` (e.g. when a
|
||||
streaming response had no router metadata). The router must still unwrap the
|
||||
managed ID before calling the upstream provider — otherwise the raw
|
||||
``cntr_...`` token leaks downstream and the provider rejects it.
|
||||
"""
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
|
||||
router = Router(model_list=[])
|
||||
mock_original_function = AsyncMock(return_value={"ok": True})
|
||||
|
||||
managed_id = ResponsesAPIRequestUtils._build_container_id(
|
||||
custom_llm_provider="openai",
|
||||
model_id=None,
|
||||
container_id="cfile_upstream_abc",
|
||||
)
|
||||
|
||||
await router._init_containers_api_endpoints(
|
||||
original_function=mock_original_function,
|
||||
custom_llm_provider="openai",
|
||||
container_id=managed_id,
|
||||
file_id="cfile_xyz",
|
||||
)
|
||||
|
||||
mock_original_function.assert_called_once()
|
||||
call_kw = mock_original_function.call_args.kwargs
|
||||
assert call_kw["container_id"] == "cfile_upstream_abc"
|
||||
assert call_kw["file_id"] == "cfile_xyz"
|
||||
assert call_kw["custom_llm_provider"] == "openai"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_containers_api_endpoints_managed_id_without_model_id_applies_decoded_provider():
|
||||
"""
|
||||
A managed ``cntr_`` ID can encode a non-OpenAI provider (e.g. ``azure``) with
|
||||
an empty ``model_id`` (streaming events without router ``model_info.id``).
|
||||
The router must still apply the decoded provider so the request routes to
|
||||
the correct upstream — not stay on the default ``openai``.
|
||||
"""
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
|
||||
router = Router(model_list=[])
|
||||
mock_original_function = AsyncMock(return_value={"ok": True})
|
||||
|
||||
managed_id = ResponsesAPIRequestUtils._build_container_id(
|
||||
custom_llm_provider="azure",
|
||||
model_id=None,
|
||||
container_id="cfile_upstream_abc",
|
||||
)
|
||||
|
||||
await router._init_containers_api_endpoints(
|
||||
original_function=mock_original_function,
|
||||
custom_llm_provider="openai",
|
||||
container_id=managed_id,
|
||||
file_id="cfile_xyz",
|
||||
)
|
||||
|
||||
mock_original_function.assert_called_once()
|
||||
call_kw = mock_original_function.call_args.kwargs
|
||||
assert call_kw["container_id"] == "cfile_upstream_abc"
|
||||
assert call_kw["file_id"] == "cfile_xyz"
|
||||
assert call_kw["custom_llm_provider"] == "azure"
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import time
|
||||
import uuid
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -260,3 +261,72 @@ async def test_async_increment_cache_returns_none_when_no_in_memory_cache_and_re
|
|||
f"Expected None when in_memory_cache is absent and Redis fails, got {result!r}. "
|
||||
"Returning the delta (1.0) would silently miscalculate rate-limit counters."
|
||||
)
|
||||
|
||||
|
||||
def test_dual_cache_late_attach_redis_wires_writes_and_ttl_sync():
|
||||
"""
|
||||
Typical lazy startup (sync): DualCache runs with in-memory only, then Redis
|
||||
becomes available and is attached. New writes must reach Redis; keys written
|
||||
before attach are not backfilled. Optional default_redis_ttl is applied on attach.
|
||||
"""
|
||||
in_memory = InMemoryCache()
|
||||
dual_cache = DualCache(in_memory_cache=in_memory, redis_cache=None)
|
||||
|
||||
mock_redis = MagicMock()
|
||||
mock_redis.set_cache = MagicMock()
|
||||
mock_redis.async_set_cache = AsyncMock()
|
||||
|
||||
key_before = f"before_attach_{uuid.uuid4()}"
|
||||
val_before = {"phase": "memory_only"}
|
||||
dual_cache.set_cache(key_before, val_before)
|
||||
|
||||
assert in_memory.get_cache(key_before) == val_before
|
||||
|
||||
dual_cache.attach_redis_cache(mock_redis, default_redis_ttl=99.0)
|
||||
assert dual_cache.redis_cache is mock_redis
|
||||
assert dual_cache.default_redis_ttl == 99.0
|
||||
|
||||
mock_redis.set_cache.assert_not_called()
|
||||
|
||||
key_after = f"after_attach_{uuid.uuid4()}"
|
||||
val_after = {"phase": "memory_and_redis"}
|
||||
dual_cache.set_cache(key_after, val_after)
|
||||
mock_redis.set_cache.assert_called_once()
|
||||
assert mock_redis.set_cache.call_args[0][:2] == (key_after, val_after)
|
||||
|
||||
assert in_memory.get_cache(key_after) == val_after
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dual_cache_late_attach_redis_wires_writes_and_ttl_async():
|
||||
"""
|
||||
Typical lazy startup (async): DualCache runs with in-memory only, then Redis
|
||||
becomes available and is attached. New writes must reach Redis; keys written
|
||||
before attach are not backfilled. Optional default_redis_ttl is applied on attach.
|
||||
"""
|
||||
in_memory = InMemoryCache()
|
||||
dual_cache = DualCache(in_memory_cache=in_memory, redis_cache=None)
|
||||
|
||||
mock_redis = MagicMock()
|
||||
mock_redis.set_cache = MagicMock()
|
||||
mock_redis.async_set_cache = AsyncMock()
|
||||
|
||||
key_before = f"before_attach_{uuid.uuid4()}"
|
||||
val_before = {"phase": "memory_only"}
|
||||
await dual_cache.async_set_cache(key_before, val_before)
|
||||
|
||||
assert in_memory.get_cache(key_before) == val_before
|
||||
|
||||
dual_cache.attach_redis_cache(mock_redis, default_redis_ttl=99.0)
|
||||
assert dual_cache.redis_cache is mock_redis
|
||||
assert dual_cache.default_redis_ttl == 99.0
|
||||
|
||||
mock_redis.async_set_cache.assert_not_called()
|
||||
|
||||
key_after = f"after_attach_{uuid.uuid4()}"
|
||||
val_after = {"phase": "memory_and_redis"}
|
||||
await dual_cache.async_set_cache(key_after, val_after)
|
||||
mock_redis.async_set_cache.assert_called_once()
|
||||
assert mock_redis.async_set_cache.call_args[0][:2] == (key_after, val_after)
|
||||
|
||||
assert in_memory.get_cache(key_after) == val_after
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ sys.path.insert(0, os.path.abspath("../../../"))
|
|||
import litellm
|
||||
from litellm.llms.azure.containers.transformation import AzureContainerConfig
|
||||
from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.containers.main import (
|
||||
ContainerFileListResponse,
|
||||
ContainerListResponse,
|
||||
|
|
@ -518,3 +519,206 @@ class TestAzureContainerKnownFailureRegressions:
|
|||
c2 = _get_container_provider_config("azure_text")
|
||||
assert type(c1) is type(c2)
|
||||
assert isinstance(c1, AzureContainerConfig)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_process_request_preserves_managed_container_id(
|
||||
self, monkeypatch
|
||||
):
|
||||
from starlette.requests import Request
|
||||
|
||||
from litellm.proxy.container_endpoints import handler_factory
|
||||
|
||||
encoded_id = ResponsesAPIRequestUtils._build_container_id(
|
||||
custom_llm_provider="azure",
|
||||
model_id="model_abc123",
|
||||
container_id="cntr_123",
|
||||
)
|
||||
captured = {}
|
||||
|
||||
async def _mock_base_process_llm_request(
|
||||
self,
|
||||
request,
|
||||
fastapi_response,
|
||||
user_api_key_dict,
|
||||
route_type,
|
||||
**kwargs,
|
||||
):
|
||||
captured["data"] = self.data
|
||||
captured["route_type"] = route_type
|
||||
return {"id": "cfile_abc"}
|
||||
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"base_process_llm_request",
|
||||
_mock_base_process_llm_request,
|
||||
)
|
||||
|
||||
request = Request(
|
||||
{
|
||||
"type": "http",
|
||||
"method": "GET",
|
||||
"path": "/v1/containers/id/files/id/content",
|
||||
"headers": [],
|
||||
"query_string": b"",
|
||||
}
|
||||
)
|
||||
fastapi_response = MagicMock()
|
||||
|
||||
await handler_factory._process_request(
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
user_api_key_dict=MagicMock(),
|
||||
route_type="alist_container_files",
|
||||
path_params={"container_id": encoded_id},
|
||||
)
|
||||
|
||||
assert captured["route_type"] == "alist_container_files"
|
||||
assert captured["data"]["container_id"] == encoded_id
|
||||
assert captured["data"]["custom_llm_provider"] == "openai"
|
||||
assert "model_id" not in captured["data"]
|
||||
assert "api_base" not in captured["data"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regression_binary_file_request_routes_through_proxy_processor(
|
||||
self, monkeypatch
|
||||
):
|
||||
from fastapi import Response
|
||||
from starlette.requests import Request
|
||||
|
||||
from litellm.proxy.container_endpoints import handler_factory
|
||||
|
||||
encoded_id = ResponsesAPIRequestUtils._build_container_id(
|
||||
custom_llm_provider="azure",
|
||||
model_id="model_abc123",
|
||||
container_id="cntr_123",
|
||||
)
|
||||
captured = {}
|
||||
|
||||
async def _mock_base_process_llm_request(
|
||||
self,
|
||||
request,
|
||||
fastapi_response,
|
||||
user_api_key_dict,
|
||||
route_type,
|
||||
**kwargs,
|
||||
):
|
||||
captured["data"] = self.data
|
||||
captured["route_type"] = route_type
|
||||
fastapi_response.headers["x-litellm-call-id"] = "call-123"
|
||||
return b"csv-bytes"
|
||||
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"base_process_llm_request",
|
||||
_mock_base_process_llm_request,
|
||||
)
|
||||
|
||||
request = Request(
|
||||
{
|
||||
"type": "http",
|
||||
"method": "GET",
|
||||
"path": "/v1/containers/id/files/id/content",
|
||||
"headers": [],
|
||||
"query_string": b"",
|
||||
}
|
||||
)
|
||||
fastapi_response = Response()
|
||||
|
||||
response = await handler_factory._process_binary_request(
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
container_id=encoded_id,
|
||||
file_id="cfile_abc",
|
||||
user_api_key_dict=MagicMock(),
|
||||
)
|
||||
|
||||
assert captured["route_type"] == "aretrieve_container_file_content"
|
||||
assert captured["data"]["container_id"] == encoded_id
|
||||
assert captured["data"]["file_id"] == "cfile_abc"
|
||||
assert captured["data"]["custom_llm_provider"] == "openai"
|
||||
assert response.status_code == 200
|
||||
assert response.body == b"csv-bytes"
|
||||
assert response.headers["x-litellm-call-id"] == "call-123"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regression_multipart_upload_request_uses_provider_from_managed_id(
|
||||
self, monkeypatch
|
||||
):
|
||||
from starlette.requests import Request
|
||||
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
from litellm.proxy.common_utils import http_parsing_utils
|
||||
from litellm.proxy.container_endpoints import handler_factory
|
||||
|
||||
encoded_id = ResponsesAPIRequestUtils._build_container_id(
|
||||
custom_llm_provider="azure",
|
||||
model_id="model_abc123",
|
||||
container_id="cntr_123",
|
||||
)
|
||||
captured = {}
|
||||
|
||||
async def _mock_get_form_data(request):
|
||||
return {"file": "ignored"}
|
||||
|
||||
async def _mock_convert_upload_files_to_file_data(form_data):
|
||||
return {"file": [("data.csv", b"csv-bytes", "text/csv")]}
|
||||
|
||||
async def _mock_base_process_llm_request(
|
||||
self,
|
||||
request,
|
||||
fastapi_response,
|
||||
user_api_key_dict,
|
||||
route_type,
|
||||
**kwargs,
|
||||
):
|
||||
captured["data"] = self.data
|
||||
captured["route_type"] = route_type
|
||||
return {"id": "cfile_abc"}
|
||||
|
||||
monkeypatch.setattr(
|
||||
http_parsing_utils,
|
||||
"get_form_data",
|
||||
_mock_get_form_data,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
http_parsing_utils,
|
||||
"convert_upload_files_to_file_data",
|
||||
_mock_convert_upload_files_to_file_data,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"base_process_llm_request",
|
||||
_mock_base_process_llm_request,
|
||||
)
|
||||
|
||||
request = Request(
|
||||
{
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/v1/containers/id/files",
|
||||
"headers": [],
|
||||
"query_string": b"",
|
||||
}
|
||||
)
|
||||
|
||||
await handler_factory._process_multipart_upload_request(
|
||||
request=request,
|
||||
fastapi_response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(),
|
||||
route_type="aupload_container_file",
|
||||
container_id=encoded_id,
|
||||
)
|
||||
|
||||
assert captured["route_type"] == "aupload_container_file"
|
||||
assert captured["data"]["container_id"] == encoded_id
|
||||
assert captured["data"]["custom_llm_provider"] == "openai"
|
||||
|
|
|
|||
|
|
@ -280,3 +280,55 @@ class TestDynamicProjectNameOnSpan:
|
|||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
|
||||
# --- Security: SSRF via prompt_version_id path traversal ---
|
||||
|
||||
|
||||
def test_arize_phoenix_client_sanitize_id_rejects_traversal():
|
||||
from litellm.integrations.arize.arize_phoenix_client import _sanitize_id
|
||||
|
||||
# dotdot without slashes
|
||||
with pytest.raises(ValueError, match="path traversal"):
|
||||
_sanitize_id("..something")
|
||||
# full traversal (slash caught first)
|
||||
with pytest.raises(ValueError, match="disallowed characters"):
|
||||
_sanitize_id("../../projects")
|
||||
|
||||
|
||||
def test_arize_phoenix_client_sanitize_id_rejects_slash():
|
||||
from litellm.integrations.arize.arize_phoenix_client import _sanitize_id
|
||||
|
||||
with pytest.raises(ValueError, match="disallowed characters"):
|
||||
_sanitize_id("valid/extra")
|
||||
|
||||
|
||||
def test_arize_phoenix_client_sanitize_id_rejects_fragment():
|
||||
from litellm.integrations.arize.arize_phoenix_client import _sanitize_id
|
||||
|
||||
with pytest.raises(ValueError, match="disallowed characters"):
|
||||
_sanitize_id("abc#suffix")
|
||||
|
||||
|
||||
def test_arize_phoenix_client_sanitize_id_rejects_query():
|
||||
from litellm.integrations.arize.arize_phoenix_client import _sanitize_id
|
||||
|
||||
with pytest.raises(ValueError, match="disallowed characters"):
|
||||
_sanitize_id("abc?x=1")
|
||||
|
||||
|
||||
def test_arize_phoenix_client_sanitize_id_allows_uuid():
|
||||
from litellm.integrations.arize.arize_phoenix_client import _sanitize_id
|
||||
|
||||
uid = "550e8400-e29b-41d4-a716-446655440000"
|
||||
assert _sanitize_id(uid) == uid
|
||||
|
||||
|
||||
def test_arize_phoenix_client_get_prompt_version_rejects_traversal():
|
||||
from litellm.integrations.arize.arize_phoenix_client import ArizePhoenixClient
|
||||
|
||||
client = ArizePhoenixClient(
|
||||
api_key="test-key", api_base="https://app.phoenix.arize.com"
|
||||
)
|
||||
with pytest.raises(ValueError, match="disallowed characters"):
|
||||
client.get_prompt_version("../../projects")
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ sys.path.insert(
|
|||
|
||||
import litellm
|
||||
from litellm.integrations.bitbucket import BitBucketPromptManager
|
||||
from litellm.integrations.bitbucket.bitbucket_client import _sanitize_file_path
|
||||
|
||||
|
||||
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
|
||||
|
|
@ -370,3 +371,45 @@ def test_bitbucket_prompt_manager_list_templates(mock_client_class):
|
|||
templates = manager.prompt_manager.list_templates()
|
||||
assert isinstance(templates, list)
|
||||
assert "test_prompt" in templates
|
||||
|
||||
|
||||
# --- Security: path traversal / SSRF ---
|
||||
|
||||
|
||||
def test_sanitize_file_path_rejects_traversal():
|
||||
with pytest.raises(ValueError, match="path traversal"):
|
||||
_sanitize_file_path("../../etc/passwd")
|
||||
|
||||
|
||||
def test_sanitize_file_path_rejects_fragment():
|
||||
with pytest.raises(ValueError, match="URL special characters"):
|
||||
_sanitize_file_path("secret#.prompt")
|
||||
|
||||
|
||||
def test_sanitize_file_path_rejects_query():
|
||||
with pytest.raises(ValueError, match="URL special characters"):
|
||||
_sanitize_file_path("secret?.prompt")
|
||||
|
||||
|
||||
def test_sanitize_file_path_encodes_special_chars():
|
||||
result = _sanitize_file_path("prompts/my prompt.prompt")
|
||||
assert result == "prompts/my%20prompt.prompt"
|
||||
|
||||
|
||||
def test_sanitize_file_path_allows_normal_paths():
|
||||
assert _sanitize_file_path("prompts/my-prompt") == "prompts/my-prompt"
|
||||
assert _sanitize_file_path("simple") == "simple"
|
||||
|
||||
|
||||
def test_bitbucket_client_rejects_traversal_in_get_file_content():
|
||||
from litellm.integrations.bitbucket.bitbucket_client import BitBucketClient
|
||||
|
||||
client = BitBucketClient(
|
||||
{
|
||||
"workspace": "ws",
|
||||
"repository": "repo",
|
||||
"access_token": "tok",
|
||||
}
|
||||
)
|
||||
with pytest.raises(ValueError, match="path traversal"):
|
||||
client.get_file_content("../../admin/credentials")
|
||||
|
|
|
|||
|
|
@ -9,8 +9,10 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
BAD_MESSAGE_ERROR_STR,
|
||||
BedrockConverseMessagesProcessor,
|
||||
BedrockImageProcessor,
|
||||
anthropic_messages_pt,
|
||||
_bedrock_converse_messages_pt,
|
||||
_convert_to_bedrock_tool_call_invoke,
|
||||
_convert_to_bedrock_tool_call_result,
|
||||
anthropic_messages_pt,
|
||||
convert_to_gemini_tool_call_result,
|
||||
ollama_pt,
|
||||
sanitize_messages_for_tool_calling,
|
||||
|
|
@ -2485,10 +2487,6 @@ def test_convert_to_anthropic_tool_result_openai_file_pdf_becomes_document():
|
|||
inside the tool_result content. Reuses anthropic_process_openai_file_message,
|
||||
which already handles this for user messages.
|
||||
"""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
convert_to_anthropic_tool_result,
|
||||
)
|
||||
|
||||
pdf_b64 = "JVBERi0xLjQKJeLjz9MK"
|
||||
message = {
|
||||
"tool_call_id": "toolu_pdf_1",
|
||||
|
|
@ -2505,157 +2503,105 @@ def test_convert_to_anthropic_tool_result_openai_file_pdf_becomes_document():
|
|||
],
|
||||
}
|
||||
|
||||
result = convert_to_anthropic_tool_result(message)
|
||||
result = _convert_to_bedrock_tool_call_result(message)
|
||||
|
||||
assert result["type"] == "tool_result"
|
||||
assert result["tool_use_id"] == "toolu_pdf_1"
|
||||
content = result["content"]
|
||||
assert isinstance(content, list) and len(content) == 1
|
||||
block = content[0]
|
||||
assert block["type"] == "document"
|
||||
assert block["source"]["type"] == "base64"
|
||||
assert block["source"]["media_type"] == "application/pdf"
|
||||
assert block["source"]["data"] == pdf_b64
|
||||
tool_result = result["toolResult"]
|
||||
assert len(tool_result["content"]) == 1
|
||||
assert "document" in tool_result["content"][0]
|
||||
assert tool_result["content"][0]["document"]["format"] == "pdf"
|
||||
assert tool_result["content"][0]["document"]["source"]["bytes"] == pdf_b64
|
||||
|
||||
|
||||
def test_convert_to_anthropic_tool_result_image_url_pdf_data_uri_becomes_document():
|
||||
"""
|
||||
Regression: a PDF sent as an `image_url` data URI on the tool-result path
|
||||
must translate to an Anthropic document block (not an image block — Anthropic
|
||||
rejects image blocks whose media_type is a non-image like application/pdf).
|
||||
"""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
convert_to_anthropic_tool_result,
|
||||
)
|
||||
def test_bedrock_converse_messages_pt_document_various_formats():
|
||||
"""Test that various document media types produce the correct format value."""
|
||||
test_cases = [
|
||||
("application/pdf", "pdf"),
|
||||
("text/csv", "csv"),
|
||||
("text/html", "html"),
|
||||
("text/plain", "txt"),
|
||||
("text/markdown", "md"),
|
||||
(
|
||||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
"docx",
|
||||
),
|
||||
]
|
||||
|
||||
pdf_b64 = "JVBERi0xLjQKJeLjz9MK"
|
||||
message = {
|
||||
"tool_call_id": "toolu_pdf_img_1",
|
||||
"role": "tool",
|
||||
"name": "fetch_document",
|
||||
"content": [
|
||||
for media_type, expected_format in test_cases:
|
||||
messages = [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:application/pdf;base64,{pdf_b64}",
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "document",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": media_type,
|
||||
"data": "dGVzdA==",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
result = _bedrock_converse_messages_pt(
|
||||
messages, "anthropic.claude-sonnet-4-6", "bedrock"
|
||||
)
|
||||
|
||||
doc_block = result[0]["content"][0]
|
||||
assert doc_block["document"]["format"] == expected_format, (
|
||||
f"Expected format '{expected_format}' for media_type '{media_type}', "
|
||||
f"got '{doc_block['document']['format']}'"
|
||||
)
|
||||
|
||||
|
||||
def test_bedrock_converse_messages_pt_document_deterministic_name():
|
||||
"""Test that the same document data always produces the same name."""
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "document",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": "application/pdf",
|
||||
"data": "dGVzdA==",
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
result = convert_to_anthropic_tool_result(message)
|
||||
|
||||
content = result["content"]
|
||||
assert isinstance(content, list) and len(content) == 1
|
||||
block = content[0]
|
||||
assert block["type"] == "document"
|
||||
assert block["source"]["media_type"] == "application/pdf"
|
||||
assert block["source"]["data"] == pdf_b64
|
||||
|
||||
|
||||
def test_convert_to_anthropic_tool_result_image_url_unsupported_mime_stays_image_path():
|
||||
"""
|
||||
An `image_url` data URI whose mime is neither application/pdf nor text/plain
|
||||
(e.g. application/json) must NOT be routed through the document path. Anthropic
|
||||
only accepts application/pdf and text/plain as base64 document media_types —
|
||||
anything else would produce a document block the API rejects. The old
|
||||
(pre-fix) behavior was to wrap such data as an image block, which also
|
||||
fails but stays on the image code path; preserve that failure mode rather
|
||||
than switching to a document path that is equally broken.
|
||||
"""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
convert_to_anthropic_tool_result,
|
||||
result1 = _bedrock_converse_messages_pt(
|
||||
messages, "anthropic.claude-sonnet-4-6", "bedrock"
|
||||
)
|
||||
result2 = _bedrock_converse_messages_pt(
|
||||
messages, "anthropic.claude-sonnet-4-6", "bedrock"
|
||||
)
|
||||
|
||||
message = {
|
||||
"tool_call_id": "toolu_json_1",
|
||||
"role": "tool",
|
||||
"name": "fetch_json",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "data:application/json;base64,eyJrIjoidiJ9",
|
||||
name1 = result1[0]["content"][0]["document"]["name"]
|
||||
name2 = result2[0]["content"][0]["document"]["name"]
|
||||
assert name1 == name2
|
||||
|
||||
|
||||
def test_bedrock_converse_messages_pt_document_rejects_url_source():
|
||||
"""Test that a URL-type document source raises a clear error instead of KeyError."""
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "document",
|
||||
"source": {
|
||||
"type": "url",
|
||||
"url": "https://example.com/doc.pdf",
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
result = convert_to_anthropic_tool_result(message)
|
||||
|
||||
content = result["content"]
|
||||
assert isinstance(content, list) and len(content) == 1
|
||||
block = content[0]
|
||||
assert block["type"] == "image", (
|
||||
f"unsupported mime {block.get('source', {}).get('media_type')!r} "
|
||||
f"should not be routed to document path; got {block}"
|
||||
)
|
||||
|
||||
|
||||
def test_convert_to_anthropic_tool_result_image_url_text_plain_data_uri_becomes_document():
|
||||
"""
|
||||
text/plain is one of the two mimes Anthropic accepts as a base64 document
|
||||
media_type. Confirm it routes through the document path so tightening the
|
||||
gate to {application/pdf, text/plain} (not "application/*") covers both.
|
||||
"""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
convert_to_anthropic_tool_result,
|
||||
)
|
||||
|
||||
txt_b64 = "aGVsbG8=" # "hello"
|
||||
message = {
|
||||
"tool_call_id": "toolu_txt_1",
|
||||
"role": "tool",
|
||||
"name": "fetch_text",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:text/plain;base64,{txt_b64}",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
result = convert_to_anthropic_tool_result(message)
|
||||
|
||||
content = result["content"]
|
||||
assert isinstance(content, list) and len(content) == 1
|
||||
block = content[0]
|
||||
assert block["type"] == "document"
|
||||
assert block["source"]["media_type"] == "text/plain"
|
||||
assert block["source"]["data"] == txt_b64
|
||||
|
||||
|
||||
def test_convert_to_anthropic_tool_result_image_url_png_still_becomes_image():
|
||||
"""
|
||||
Regression: image_url with a real image mime type must continue to translate
|
||||
to an Anthropic image block. Locks in existing behavior after the
|
||||
data-URI-mime-type branching for PDFs.
|
||||
"""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
convert_to_anthropic_tool_result,
|
||||
)
|
||||
|
||||
png_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGNgYGBgAAAABQABXvMqOgAAAABJRU5ErkJggg=="
|
||||
message = {
|
||||
"tool_call_id": "toolu_png_1",
|
||||
"role": "tool",
|
||||
"name": "fetch_image",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/png;base64,{png_b64}",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
result = convert_to_anthropic_tool_result(message)
|
||||
|
||||
content = result["content"]
|
||||
assert isinstance(content, list) and len(content) == 1
|
||||
block = content[0]
|
||||
assert block["type"] == "image"
|
||||
assert block["source"]["media_type"] == "image/png"
|
||||
with pytest.raises(ValueError, match="only supports base64-encoded"):
|
||||
_bedrock_converse_messages_pt(
|
||||
messages, "anthropic.claude-sonnet-4-6", "bedrock"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,290 @@
|
|||
"""
|
||||
Tests for Gemini batchEmbedContents transformation logic.
|
||||
|
||||
Covers:
|
||||
- Text-only inputs (single and batch)
|
||||
- Multimodal inputs (data URIs, GCS URLs, file references)
|
||||
- Mixed text + multimodal inputs
|
||||
- Response processing with correct indices
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import (
|
||||
_build_part_for_input,
|
||||
_is_multimodal_input,
|
||||
process_response,
|
||||
transform_openai_input_gemini_content,
|
||||
transform_openai_input_gemini_embed_content,
|
||||
)
|
||||
from litellm.types.llms.vertex_ai import VertexAIBatchEmbeddingsResponseObject
|
||||
from litellm.types.utils import EmbeddingResponse
|
||||
|
||||
|
||||
IMAGE_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII"
|
||||
GCS_URL = "gs://my-bucket/image.png"
|
||||
|
||||
|
||||
class TestIsMultimodalInput:
|
||||
def test_text_only_string(self):
|
||||
assert _is_multimodal_input("hello world") is False
|
||||
|
||||
def test_text_only_list(self):
|
||||
assert _is_multimodal_input(["hello", "world"]) is False
|
||||
|
||||
def test_data_uri(self):
|
||||
assert _is_multimodal_input([IMAGE_DATA_URI]) is True
|
||||
|
||||
def test_gcs_url(self):
|
||||
assert _is_multimodal_input([GCS_URL]) is True
|
||||
|
||||
def test_file_reference(self):
|
||||
assert _is_multimodal_input(["files/abc123"]) is True
|
||||
|
||||
def test_mixed_text_and_image(self):
|
||||
assert _is_multimodal_input(["hello", IMAGE_DATA_URI]) is True
|
||||
|
||||
def test_nested_text_is_not_multimodal(self):
|
||||
"""Nested list with text is not multimodal."""
|
||||
assert _is_multimodal_input([["text_a", "text_b"]]) is False
|
||||
|
||||
def test_nested_list_with_image_is_multimodal(self):
|
||||
assert _is_multimodal_input([["a red shoe", IMAGE_DATA_URI]]) is True
|
||||
|
||||
|
||||
class TestBuildPartForInput:
|
||||
def test_text_input(self):
|
||||
part = _build_part_for_input("hello")
|
||||
assert part["text"] == "hello"
|
||||
assert part.get("inline_data") is None
|
||||
|
||||
def test_data_uri_input(self):
|
||||
part = _build_part_for_input(IMAGE_DATA_URI)
|
||||
assert part.get("text") is None
|
||||
assert part["inline_data"] is not None
|
||||
assert part["inline_data"]["mime_type"] == "image/png"
|
||||
|
||||
def test_gcs_url_input(self):
|
||||
part = _build_part_for_input(GCS_URL)
|
||||
assert part.get("text") is None
|
||||
assert part["file_data"] is not None
|
||||
assert part["file_data"]["mime_type"] == "image/png"
|
||||
assert part["file_data"]["file_uri"] == GCS_URL
|
||||
|
||||
def test_file_reference_resolved(self):
|
||||
resolved = {"files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"}}
|
||||
part = _build_part_for_input("files/abc", resolved_files=resolved)
|
||||
assert part["file_data"] is not None
|
||||
assert part["file_data"]["mime_type"] == "image/jpeg"
|
||||
|
||||
def test_file_reference_unresolved_raises(self):
|
||||
with pytest.raises(ValueError, match="not resolved"):
|
||||
_build_part_for_input("files/abc")
|
||||
|
||||
|
||||
class TestTransformOpenaiInputGeminiContent:
|
||||
"""Test that transform_openai_input_gemini_content creates separate requests per input."""
|
||||
|
||||
def test_single_text(self):
|
||||
result = transform_openai_input_gemini_content(
|
||||
input="hello", model="gemini-embedding-2-preview", optional_params={}
|
||||
)
|
||||
assert len(result["requests"]) == 1
|
||||
assert result["requests"][0]["content"]["parts"][0]["text"] == "hello"
|
||||
|
||||
def test_multiple_texts(self):
|
||||
result = transform_openai_input_gemini_content(
|
||||
input=["hello", "world"], model="gemini-embedding-2-preview", optional_params={}
|
||||
)
|
||||
assert len(result["requests"]) == 2
|
||||
assert result["requests"][0]["content"]["parts"][0]["text"] == "hello"
|
||||
assert result["requests"][1]["content"]["parts"][0]["text"] == "world"
|
||||
|
||||
def test_multimodal_inputs_are_separate_requests(self):
|
||||
"""Key regression test for #24209: each input becomes its own request."""
|
||||
result = transform_openai_input_gemini_content(
|
||||
input=["The food was delicious", IMAGE_DATA_URI],
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={},
|
||||
)
|
||||
assert len(result["requests"]) == 2
|
||||
# First request is text
|
||||
assert result["requests"][0]["content"]["parts"][0]["text"] == "The food was delicious"
|
||||
# Second request is image
|
||||
assert result["requests"][1]["content"]["parts"][0]["inline_data"] is not None
|
||||
|
||||
def test_dimensions_mapped_to_output_dimensionality(self):
|
||||
result = transform_openai_input_gemini_content(
|
||||
input="hello",
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={"dimensions": 256},
|
||||
)
|
||||
assert result["requests"][0]["outputDimensionality"] == 256
|
||||
|
||||
def test_model_name_prefixed(self):
|
||||
result = transform_openai_input_gemini_content(
|
||||
input="hello", model="gemini-embedding-2-preview", optional_params={}
|
||||
)
|
||||
assert result["requests"][0]["model"] == "models/gemini-embedding-2-preview"
|
||||
|
||||
def test_gcs_url_input(self):
|
||||
result = transform_openai_input_gemini_content(
|
||||
input=[GCS_URL], model="gemini-embedding-2-preview", optional_params={}
|
||||
)
|
||||
assert len(result["requests"]) == 1
|
||||
assert result["requests"][0]["content"]["parts"][0]["file_data"] is not None
|
||||
|
||||
def test_mixed_text_image_gcs(self):
|
||||
result = transform_openai_input_gemini_content(
|
||||
input=["hello", IMAGE_DATA_URI, GCS_URL],
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={},
|
||||
)
|
||||
assert len(result["requests"]) == 3
|
||||
|
||||
def test_nested_input_combined_embedding(self):
|
||||
"""Nested list produces one request with multiple parts (combined embedding)."""
|
||||
result = transform_openai_input_gemini_content(
|
||||
input=[["a red shoe", IMAGE_DATA_URI]],
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={},
|
||||
)
|
||||
assert len(result["requests"]) == 1
|
||||
parts = result["requests"][0]["content"]["parts"]
|
||||
assert len(parts) == 2
|
||||
assert parts[0]["text"] == "a red shoe"
|
||||
assert parts[1]["inline_data"] is not None
|
||||
|
||||
def test_mixed_nested_and_flat(self):
|
||||
"""Mixed nested + flat produces correct number of requests."""
|
||||
result = transform_openai_input_gemini_content(
|
||||
input=[["text", IMAGE_DATA_URI], "standalone"],
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={},
|
||||
)
|
||||
assert len(result["requests"]) == 2
|
||||
# First: combined (2 parts)
|
||||
assert len(result["requests"][0]["content"]["parts"]) == 2
|
||||
# Second: standalone (1 part)
|
||||
assert len(result["requests"][1]["content"]["parts"]) == 1
|
||||
assert result["requests"][1]["content"]["parts"][0]["text"] == "standalone"
|
||||
|
||||
|
||||
class TestTransformOpenaiInputGeminiEmbedContent:
|
||||
"""Test transform_openai_input_gemini_embed_content (vertex_ai / embedContent path)."""
|
||||
|
||||
def test_text_and_image_combined(self):
|
||||
result = transform_openai_input_gemini_embed_content(
|
||||
input=["hello", IMAGE_DATA_URI],
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={},
|
||||
)
|
||||
assert "content" in result
|
||||
parts = result["content"]["parts"]
|
||||
assert len(parts) == 2
|
||||
assert parts[0]["text"] == "hello"
|
||||
assert parts[1]["inline_data"] is not None
|
||||
|
||||
def test_gcs_url(self):
|
||||
result = transform_openai_input_gemini_embed_content(
|
||||
input=[GCS_URL],
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={},
|
||||
)
|
||||
parts = result["content"]["parts"]
|
||||
assert len(parts) == 1
|
||||
assert parts[0]["file_data"]["file_uri"] == GCS_URL
|
||||
|
||||
def test_dimensions_mapped(self):
|
||||
result = transform_openai_input_gemini_embed_content(
|
||||
input="hello",
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={"dimensions": 256},
|
||||
)
|
||||
assert result["outputDimensionality"] == 256
|
||||
|
||||
|
||||
class TestProcessResponse:
|
||||
"""Test that process_response sets correct indices."""
|
||||
|
||||
def test_single_embedding_index(self):
|
||||
predictions: VertexAIBatchEmbeddingsResponseObject = {
|
||||
"embeddings": [{"values": [0.1, 0.2]}]
|
||||
}
|
||||
model_response = EmbeddingResponse()
|
||||
result = process_response(
|
||||
input="hello",
|
||||
model_response=model_response,
|
||||
model="gemini-embedding-2-preview",
|
||||
_predictions=predictions,
|
||||
)
|
||||
assert len(result.data) == 1
|
||||
assert result.data[0]["index"] == 0
|
||||
|
||||
def test_multiple_embeddings_have_correct_indices(self):
|
||||
"""Regression test: indices should be 0, 1, 2... not all 0."""
|
||||
predictions: VertexAIBatchEmbeddingsResponseObject = {
|
||||
"embeddings": [
|
||||
{"values": [0.1, 0.2]},
|
||||
{"values": [0.3, 0.4]},
|
||||
{"values": [0.5, 0.6]},
|
||||
]
|
||||
}
|
||||
model_response = EmbeddingResponse()
|
||||
result = process_response(
|
||||
input=["a", "b", "c"],
|
||||
model_response=model_response,
|
||||
model="gemini-embedding-2-preview",
|
||||
_predictions=predictions,
|
||||
)
|
||||
assert len(result.data) == 3
|
||||
assert result.data[0]["index"] == 0
|
||||
assert result.data[1]["index"] == 1
|
||||
assert result.data[2]["index"] == 2
|
||||
|
||||
def test_multimodal_mixed_input(self):
|
||||
"""process_response works with mixed text + multimodal inputs."""
|
||||
predictions: VertexAIBatchEmbeddingsResponseObject = {
|
||||
"embeddings": [{"values": [0.1, 0.2]}, {"values": [0.3, 0.4]}]
|
||||
}
|
||||
result = process_response(
|
||||
input=["hello", IMAGE_DATA_URI],
|
||||
model_response=EmbeddingResponse(),
|
||||
model="gemini-embedding-2-preview",
|
||||
_predictions=predictions,
|
||||
)
|
||||
assert len(result.data) == 2
|
||||
assert result.data[0]["index"] == 0
|
||||
assert result.data[1]["index"] == 1
|
||||
# Should count tokens only for the text element, not the image
|
||||
assert result.usage.prompt_tokens > 0
|
||||
|
||||
def test_nested_input_token_counting(self):
|
||||
"""Nested list: only plain-text sub-elements should be counted."""
|
||||
predictions: VertexAIBatchEmbeddingsResponseObject = {
|
||||
"embeddings": [{"values": [0.1, 0.2]}]
|
||||
}
|
||||
result = process_response(
|
||||
input=[["a red shoe", IMAGE_DATA_URI]],
|
||||
model_response=EmbeddingResponse(),
|
||||
model="gemini-embedding-2-preview",
|
||||
_predictions=predictions,
|
||||
)
|
||||
assert len(result.data) == 1
|
||||
assert result.usage.prompt_tokens > 0
|
||||
|
||||
def test_nested_empty_list_raises(self):
|
||||
with pytest.raises(ValueError, match="must not be empty"):
|
||||
transform_openai_input_gemini_content(
|
||||
input=[[]],
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={},
|
||||
)
|
||||
|
||||
def test_nested_non_string_element_raises(self):
|
||||
with pytest.raises(ValueError, match="must be strings"):
|
||||
transform_openai_input_gemini_content(
|
||||
input=[[["doubly", "nested"]]],
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={},
|
||||
)
|
||||
|
|
@ -0,0 +1,41 @@
|
|||
"""Vertex Model Garden: OpenAPI base URL for publisher/model ids vs per-endpoint path."""
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.llms.vertex_ai.vertex_model_garden.main import (
|
||||
_vertex_model_garden_model_id_in_json_body,
|
||||
create_vertex_url,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model,expect_openapi_base",
|
||||
[
|
||||
("xai/grok-4.1-fast-reasoning", True),
|
||||
("openai/foo/bar", True),
|
||||
("5464397967697903616", False),
|
||||
("gpt-oss-20b-maas", False),
|
||||
],
|
||||
)
|
||||
def test_create_vertex_url_openapi_vs_deployed_endpoint(
|
||||
model: str, expect_openapi_base: bool
|
||||
) -> None:
|
||||
url = create_vertex_url(
|
||||
vertex_location="us-central1",
|
||||
vertex_project="my-project",
|
||||
stream=False,
|
||||
model=model,
|
||||
)
|
||||
if expect_openapi_base:
|
||||
assert "/v1/projects/my-project/locations/us-central1/endpoints/openapi" in url
|
||||
else:
|
||||
assert (
|
||||
"/v1beta1/projects/my-project/locations/us-central1/endpoints/"
|
||||
f"{model}" in url
|
||||
)
|
||||
assert "openapi" not in url
|
||||
|
||||
|
||||
def test_model_id_in_json_body_heuristic() -> None:
|
||||
assert _vertex_model_garden_model_id_in_json_body("xai/grok-4.1-fast-reasoning") is True
|
||||
assert _vertex_model_garden_model_id_in_json_body("5464397967697903616") is False
|
||||
39
tests/test_litellm/llms/xai/test_xai_chat_transformation.py
Normal file
39
tests/test_litellm/llms/xai/test_xai_chat_transformation.py
Normal file
|
|
@ -0,0 +1,39 @@
|
|||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
from litellm.llms.xai.chat.transformation import XAIChatConfig
|
||||
|
||||
|
||||
class TestXAIParallelToolCalls:
|
||||
"""Test suite for XAI parallel tool calls functionality."""
|
||||
|
||||
def test_get_supported_openai_params_includes_parallel_tool_calls(self):
|
||||
"""Test that parallel_tool_calls is in supported parameters."""
|
||||
config = XAIChatConfig()
|
||||
supported_params = config.get_supported_openai_params(
|
||||
"xai/grok-4.20"
|
||||
)
|
||||
assert "parallel_tool_calls" in supported_params
|
||||
|
||||
def test_transform_request_preserves_parallel_tool_calls(self):
|
||||
"""Test that transform_request preserves parallel_tool_calls parameter."""
|
||||
config = XAIChatConfig()
|
||||
|
||||
messages = [{"role": "user", "content": "What's the weather like?"}]
|
||||
optional_params = {"parallel_tool_calls": True}
|
||||
|
||||
result = config.transform_request(
|
||||
model="xai/grok-4.20",
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert result.get("parallel_tool_calls") is True
|
||||
assert len(result["messages"]) == 1
|
||||
assert result["messages"][0]["role"] == "user"
|
||||
|
|
@ -673,6 +673,106 @@ class TestHookHeaderMergePriority:
|
|||
assert headers["X-OAuth"] == "yes"
|
||||
assert headers["X-Trace-Id"] == "trace-123"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_m2m_oauth2_does_not_forward_litellm_caller_authorization(self):
|
||||
"""M2M must not put caller Bearer (LiteLLM API key) into extra_headers (#23652)."""
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="test-id",
|
||||
name="Test Server",
|
||||
server_name="test_server",
|
||||
url="https://example.com",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
oauth2_flow="client_credentials",
|
||||
token_url="https://auth.example.com/token",
|
||||
)
|
||||
|
||||
captured_extra_headers: Dict[str, Any] = {}
|
||||
|
||||
async def fake_create_mcp_client(
|
||||
server, mcp_auth_header=None, extra_headers=None, stdio_env=None
|
||||
):
|
||||
captured_extra_headers["value"] = extra_headers
|
||||
mock_client = MagicMock()
|
||||
mock_client.call_tool = AsyncMock(return_value=MagicMock())
|
||||
return mock_client
|
||||
|
||||
with patch.object(
|
||||
manager, "_create_mcp_client", side_effect=fake_create_mcp_client
|
||||
):
|
||||
with patch.object(manager, "_build_stdio_env", return_value=None):
|
||||
try:
|
||||
await manager._call_regular_mcp_tool(
|
||||
mcp_server=server,
|
||||
original_tool_name="test_tool",
|
||||
arguments={"key": "val"},
|
||||
tasks=[],
|
||||
mcp_auth_header=None,
|
||||
mcp_server_auth_headers=None,
|
||||
oauth2_headers={"Authorization": "Bearer sk-1234"},
|
||||
raw_headers={"authorization": "Bearer sk-1234"},
|
||||
proxy_logging_obj=None,
|
||||
hook_extra_headers=None,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
assert captured_extra_headers.get("value") is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_m2m_oauth2_skips_authorization_in_configured_extra_headers(self):
|
||||
"""M2M must not take Authorization from raw_headers even if extra_headers lists it."""
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="test-id",
|
||||
name="Test Server",
|
||||
server_name="test_server",
|
||||
url="https://example.com",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
oauth2_flow="client_credentials",
|
||||
token_url="https://auth.example.com/token",
|
||||
extra_headers=["Authorization", "X-Custom"],
|
||||
)
|
||||
|
||||
captured_extra_headers: Dict[str, Any] = {}
|
||||
|
||||
async def fake_create_mcp_client(
|
||||
server, mcp_auth_header=None, extra_headers=None, stdio_env=None
|
||||
):
|
||||
captured_extra_headers["value"] = extra_headers
|
||||
mock_client = MagicMock()
|
||||
mock_client.call_tool = AsyncMock(return_value=MagicMock())
|
||||
return mock_client
|
||||
|
||||
with patch.object(
|
||||
manager, "_create_mcp_client", side_effect=fake_create_mcp_client
|
||||
):
|
||||
with patch.object(manager, "_build_stdio_env", return_value=None):
|
||||
try:
|
||||
await manager._call_regular_mcp_tool(
|
||||
mcp_server=server,
|
||||
original_tool_name="test_tool",
|
||||
arguments={"key": "val"},
|
||||
tasks=[],
|
||||
mcp_auth_header=None,
|
||||
mcp_server_auth_headers=None,
|
||||
oauth2_headers={"Authorization": "Bearer sk-1234"},
|
||||
raw_headers={
|
||||
"authorization": "Bearer sk-1234",
|
||||
"x-custom": "from-client",
|
||||
},
|
||||
proxy_logging_obj=None,
|
||||
hook_extra_headers=None,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
headers = captured_extra_headers.get("value") or {}
|
||||
assert "Authorization" not in headers
|
||||
assert headers.get("X-Custom") == "from-client"
|
||||
|
||||
|
||||
class TestUserAPIKeyAuthJwtClaims:
|
||||
"""Tests that UserAPIKeyAuth correctly carries jwt_claims."""
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from litellm.proxy._types import (
|
|||
MCPTransport,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.types.mcp import MCPAuth
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
|
||||
|
|
@ -135,6 +136,152 @@ def test_prepare_mcp_server_headers_case_insensitive_extra_headers():
|
|||
assert extra_headers == {"Authorization": "Bearer token"}
|
||||
|
||||
|
||||
def test_prepare_mcp_server_headers_oauth2_m2m_omits_litellm_caller_authorization():
|
||||
"""M2M OAuth must not put caller Bearer (LiteLLM API key) into extra_headers (#23652)."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_prepare_mcp_server_headers,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
server = MCPServer(
|
||||
server_id="m2m-server",
|
||||
name="m2m",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
oauth2_flow="client_credentials",
|
||||
token_url="https://auth.example.com/token",
|
||||
)
|
||||
caller_key = {"Authorization": "Bearer sk-litellm-caller"}
|
||||
|
||||
server_auth_header, extra_headers = _prepare_mcp_server_headers(
|
||||
server=server,
|
||||
mcp_server_auth_headers=None,
|
||||
mcp_auth_header=None,
|
||||
oauth2_headers=caller_key,
|
||||
raw_headers=None,
|
||||
)
|
||||
|
||||
assert server_auth_header is None
|
||||
assert extra_headers is None
|
||||
|
||||
|
||||
def test_prepare_mcp_server_headers_oauth2_interactive_copies_oauth2_headers():
|
||||
"""Interactive OAuth still forwards the user's OAuth token in extra_headers."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_prepare_mcp_server_headers,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
user_oauth = {"Authorization": "Bearer upstream-user-token"}
|
||||
|
||||
server = MCPServer(
|
||||
server_id="3lo-server",
|
||||
name="3lo",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
oauth2_flow=None,
|
||||
)
|
||||
|
||||
server_auth_header, extra_headers = _prepare_mcp_server_headers(
|
||||
server=server,
|
||||
mcp_server_auth_headers=None,
|
||||
mcp_auth_header=None,
|
||||
oauth2_headers=user_oauth,
|
||||
raw_headers=None,
|
||||
)
|
||||
|
||||
assert server_auth_header is None
|
||||
assert extra_headers == user_oauth
|
||||
|
||||
|
||||
def test_prepare_mcp_server_headers_m2m_skips_authorization_from_raw_extra_headers():
|
||||
"""M2M must not merge caller Authorization from raw_headers when extra_headers lists it."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_prepare_mcp_server_headers,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
server = MCPServer(
|
||||
server_id="m2m-raw",
|
||||
name="m2m",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
oauth2_flow="client_credentials",
|
||||
token_url="https://auth.example.com/token",
|
||||
extra_headers=["Authorization", "X-Custom"],
|
||||
)
|
||||
|
||||
server_auth_header, extra_headers = _prepare_mcp_server_headers(
|
||||
server=server,
|
||||
mcp_server_auth_headers=None,
|
||||
mcp_auth_header=None,
|
||||
oauth2_headers={"Authorization": "Bearer sk-1234"},
|
||||
raw_headers={
|
||||
"authorization": "Bearer sk-1234",
|
||||
"x-custom": "trace",
|
||||
},
|
||||
)
|
||||
|
||||
assert server_auth_header is None
|
||||
assert extra_headers is not None
|
||||
assert "Authorization" not in extra_headers
|
||||
assert extra_headers.get("X-Custom") == "trace"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_tool_m2m_skips_authorization_headers():
|
||||
"""M2M call_tool must not forward caller Authorization in oauth2/raw headers."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="m2m-call-tool",
|
||||
name="m2m-call-tool",
|
||||
server_name="m2m-call-tool",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
oauth2_flow="client_credentials",
|
||||
token_url="https://auth.example.com/token",
|
||||
client_id="cid",
|
||||
client_secret="csecret",
|
||||
extra_headers=["Authorization", "X-Custom"],
|
||||
)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.call_tool = AsyncMock(return_value=MagicMock())
|
||||
|
||||
with patch.object(
|
||||
manager, "_create_mcp_client", new=AsyncMock(return_value=mock_client)
|
||||
) as create_client_mock:
|
||||
await manager._call_regular_mcp_tool(
|
||||
mcp_server=server,
|
||||
original_tool_name="echo",
|
||||
arguments={"message": "hello"},
|
||||
tasks=[],
|
||||
mcp_auth_header=None,
|
||||
mcp_server_auth_headers=None,
|
||||
oauth2_headers={"Authorization": "Bearer sk-1234"},
|
||||
raw_headers={"authorization": "Bearer sk-1234", "x-custom": "trace"},
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
|
||||
create_kwargs = create_client_mock.await_args.kwargs
|
||||
extra_headers = create_kwargs["extra_headers"] or {}
|
||||
assert "Authorization" not in extra_headers
|
||||
assert extra_headers.get("X-Custom") == "trace"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_prompts_from_mcp_servers_success():
|
||||
try:
|
||||
|
|
@ -2288,6 +2435,79 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab
|
|||
assert spend_meta["per_server_tool_counts"]["server_a"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_tools_from_mcp_servers_returns_tools_when_success_logging_fails():
|
||||
"""
|
||||
Regression test: list_tools should still return fetched tools even if
|
||||
async_success_handler raises (e.g. serialization errors in logging path).
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_get_tools_from_mcp_servers,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user")
|
||||
|
||||
server_a = MagicMock(name="server_a_obj")
|
||||
server_a.name = "server_a"
|
||||
server_a.alias = "server_a"
|
||||
server_a.server_name = "server_a"
|
||||
server_a.server_id = "a"
|
||||
server_a.auth_type = None
|
||||
server_a.extra_headers = None
|
||||
|
||||
tool_1 = MagicMock()
|
||||
tool_1.name = "server_a-tool_1"
|
||||
|
||||
dummy_logging_obj = MagicMock()
|
||||
dummy_logging_obj.model_call_details = {"metadata": {"spend_logs_metadata": {}}}
|
||||
dummy_logging_obj.async_success_handler = AsyncMock(
|
||||
side_effect=TypeError("Object of type Tool is not JSON serializable")
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
|
||||
new=AsyncMock(return_value=[server_a]),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._prepare_mcp_server_headers",
|
||||
return_value=(None, None),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
|
||||
) as mock_manager,
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools",
|
||||
side_effect=lambda tools, _server: tools,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions",
|
||||
new=AsyncMock(side_effect=lambda tools, **_: tools),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.function_setup",
|
||||
return_value=(dummy_logging_obj, None),
|
||||
),
|
||||
):
|
||||
mock_manager._get_tools_from_server = AsyncMock(return_value=[tool_1])
|
||||
|
||||
tools = await _get_tools_from_mcp_servers(
|
||||
user_api_key_auth=user_auth,
|
||||
mcp_auth_header=None,
|
||||
mcp_servers=["server_a"],
|
||||
mcp_server_auth_headers=None,
|
||||
log_list_tools_to_spendlogs=True,
|
||||
list_tools_log_source="mcp_protocol",
|
||||
)
|
||||
|
||||
assert tools == [tool_1]
|
||||
dummy_logging_obj.async_success_handler.assert_awaited_once()
|
||||
|
||||
|
||||
def test_tool_name_matches_case_insensitive():
|
||||
"""Test that _tool_name_matches performs case-insensitive comparison.
|
||||
|
||||
|
|
@ -2719,3 +2939,177 @@ class TestGatewayCreateInitializationOptions:
|
|||
_mcp_gateway_initialize_instructions.reset(tok)
|
||||
opts = server.create_initialization_options()
|
||||
assert getattr(opts, "instructions", None) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_tools_with_legacy_db_m2m_server_resolves_oauth2_flow():
|
||||
"""
|
||||
P1 Regression: list_tools path must apply _resolve_oauth2_flow to legacy DB
|
||||
rows where oauth2_flow is NULL but M2M credentials are present.
|
||||
|
||||
Without this fix, has_client_credentials returns False and the caller's
|
||||
Authorization header is forwarded upstream instead of being blocked.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_get_tools_from_mcp_servers,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.mcp import MCPAuth
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
user_auth = UserAPIKeyAuth(api_key="sk-1234", user_id="test-user")
|
||||
|
||||
# Simulate a legacy DB row: OAuth2 with M2M credentials but oauth2_flow=None
|
||||
legacy_server = MagicMock(name="legacy_m2m_server")
|
||||
legacy_server.name = "legacy_m2m"
|
||||
legacy_server.alias = "legacy_m2m"
|
||||
legacy_server.server_name = "legacy_m2m"
|
||||
legacy_server.server_id = "legacy-m2m-id"
|
||||
legacy_server.auth_type = MCPAuth.oauth2
|
||||
legacy_server.oauth2_flow = None # Legacy: field not set in DB
|
||||
legacy_server.token_url = "https://oauth.example.com/token"
|
||||
legacy_server.authorization_url = None
|
||||
legacy_server.client_id = "client-id"
|
||||
legacy_server.client_secret = "client-secret"
|
||||
legacy_server.extra_headers = None
|
||||
legacy_server.has_client_credentials = False # This is the bug: should be True
|
||||
legacy_server.model_copy = MagicMock(
|
||||
side_effect=lambda update: MCPServer(
|
||||
server_id=legacy_server.server_id,
|
||||
name=legacy_server.name,
|
||||
transport=MCPTransport.http,
|
||||
auth_type=legacy_server.auth_type,
|
||||
oauth2_flow=update.get("oauth2_flow", legacy_server.oauth2_flow),
|
||||
token_url=legacy_server.token_url,
|
||||
authorization_url=legacy_server.authorization_url,
|
||||
client_id=legacy_server.client_id,
|
||||
client_secret=legacy_server.client_secret,
|
||||
)
|
||||
)
|
||||
|
||||
tool_1 = MagicMock()
|
||||
tool_1.name = "legacy_m2m-tool"
|
||||
|
||||
captured_extra_headers = None
|
||||
|
||||
async def capture_extra_headers(*args, **kwargs):
|
||||
nonlocal captured_extra_headers
|
||||
captured_extra_headers = kwargs.get("extra_headers")
|
||||
return [tool_1]
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
|
||||
) as mock_manager,
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools",
|
||||
side_effect=lambda tools, _server: tools,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions",
|
||||
new=AsyncMock(side_effect=lambda tools, **_: tools),
|
||||
),
|
||||
):
|
||||
mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=["legacy-m2m-id"])
|
||||
mock_manager.get_mcp_server_by_id = MagicMock(return_value=legacy_server)
|
||||
mock_manager.filter_server_ids_by_ip_with_info = MagicMock(
|
||||
return_value=(["legacy-m2m-id"], 0)
|
||||
)
|
||||
mock_manager._get_tools_from_server = AsyncMock(
|
||||
side_effect=capture_extra_headers
|
||||
)
|
||||
|
||||
tools = await _get_tools_from_mcp_servers(
|
||||
user_api_key_auth=user_auth,
|
||||
mcp_auth_header=None,
|
||||
mcp_servers=["legacy_m2m"],
|
||||
mcp_server_auth_headers=None,
|
||||
oauth2_headers={"Authorization": "Bearer sk-1234"}, # Caller's token
|
||||
)
|
||||
|
||||
# With P1 fix: _get_allowed_mcp_servers applies _resolve_oauth2_flow,
|
||||
# so has_client_credentials becomes True and extra_headers should be None
|
||||
# (caller's Authorization blocked)
|
||||
assert captured_extra_headers is None, (
|
||||
"P1 security issue: caller's Authorization header was forwarded to M2M server. "
|
||||
"Expected None, got: " + str(captured_extra_headers)
|
||||
)
|
||||
assert tools == [tool_1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_tool_empty_extra_headers_returns_none():
|
||||
"""
|
||||
P2 Regression: When all configured extra_headers are filtered out (e.g.
|
||||
Authorization for M2M), the resulting extra_headers should be None, not {}.
|
||||
|
||||
Downstream code that checks `if extra_headers is None` will behave
|
||||
differently if an empty dict is passed instead.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
)
|
||||
from litellm.types.mcp import MCPAuth
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
manager = MCPServerManager()
|
||||
|
||||
# M2M server with only Authorization in extra_headers
|
||||
m2m_server = MCPServer(
|
||||
server_id="m2m-srv",
|
||||
name="m2m_test",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
oauth2_flow="client_credentials",
|
||||
token_url="https://oauth.example.com/token",
|
||||
client_id="client-id",
|
||||
client_secret="client-secret",
|
||||
extra_headers=["Authorization"], # Will be filtered out for M2M
|
||||
)
|
||||
|
||||
raw_headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
||||
|
||||
captured_extra_headers = None
|
||||
|
||||
async def capture_create_mcp_client(*args, **kwargs):
|
||||
nonlocal captured_extra_headers
|
||||
captured_extra_headers = kwargs.get("extra_headers")
|
||||
# Return a mock client
|
||||
mock_client = AsyncMock()
|
||||
mock_client.call_tool = AsyncMock(return_value=MagicMock(content=[]))
|
||||
return mock_client
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
manager,
|
||||
"_create_mcp_client",
|
||||
side_effect=capture_create_mcp_client,
|
||||
),
|
||||
patch.object(
|
||||
manager,
|
||||
"get_mcp_server_by_id",
|
||||
return_value=m2m_server,
|
||||
),
|
||||
):
|
||||
try:
|
||||
await manager._call_regular_mcp_tool(
|
||||
mcp_server=m2m_server,
|
||||
original_tool_name="test_tool",
|
||||
arguments={},
|
||||
mcp_auth_header=None,
|
||||
oauth2_headers=None,
|
||||
raw_headers=raw_headers,
|
||||
)
|
||||
except Exception:
|
||||
pass # We only care about the captured headers
|
||||
|
||||
# With P2 fix: extra_headers should be None (not {}) when all headers filtered
|
||||
assert captured_extra_headers is None, (
|
||||
"P2 API consistency issue: expected None for empty extra_headers, got: "
|
||||
+ str(captured_extra_headers)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -922,19 +922,21 @@ async def test_get_tag_objects_batch():
|
|||
# Simulate 5 tags: 2 cached, 3 uncached
|
||||
tag_names = ["cached-1", "uncached-1", "cached-2", "uncached-2", "uncached-3"]
|
||||
|
||||
# Mock cached tags
|
||||
cached_tag_1 = {
|
||||
"tag_name": "cached-1",
|
||||
"spend": 10.0,
|
||||
"models": [],
|
||||
"litellm_budget_table": None,
|
||||
}
|
||||
cached_tag_2 = {
|
||||
"tag_name": "cached-2",
|
||||
"spend": 20.0,
|
||||
"models": [],
|
||||
"litellm_budget_table": None,
|
||||
}
|
||||
# Mock cached tags — must be LiteLLM_TagTable instances: the mocked async_get_cache
|
||||
# bypasses UserApiKeyCache deserialization, so returning plain dicts would flow through
|
||||
# as dict (production returns models after Codec.deserialize inside the cache).
|
||||
cached_tag_1 = LiteLLM_TagTable(
|
||||
tag_name="cached-1",
|
||||
spend=10.0,
|
||||
models=[],
|
||||
litellm_budget_table=None,
|
||||
)
|
||||
cached_tag_2 = LiteLLM_TagTable(
|
||||
tag_name="cached-2",
|
||||
spend=20.0,
|
||||
models=[],
|
||||
litellm_budget_table=None,
|
||||
)
|
||||
|
||||
# Mock DB response for uncached tags
|
||||
uncached_tag_1 = MagicMock()
|
||||
|
|
@ -980,13 +982,13 @@ async def test_get_tag_objects_batch():
|
|||
)
|
||||
|
||||
# Mock cache behavior - return cached tags, None for uncached
|
||||
async def mock_get_cache(key):
|
||||
async def mock_get_cache(*args, **kwargs):
|
||||
key = kwargs.get("key")
|
||||
if key == "tag:cached-1":
|
||||
return cached_tag_1
|
||||
elif key == "tag:cached-2":
|
||||
if key == "tag:cached-2":
|
||||
return cached_tag_2
|
||||
else:
|
||||
return None
|
||||
return None
|
||||
|
||||
mock_cache.async_get_cache = AsyncMock(side_effect=mock_get_cache)
|
||||
mock_cache.async_set_cache = AsyncMock()
|
||||
|
|
|
|||
|
|
@ -405,9 +405,9 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_role_change():
|
|||
mock_cache.async_set_cache.assert_called_once()
|
||||
call_kwargs = mock_cache.async_set_cache.call_args
|
||||
assert call_kwargs.kwargs["key"] == "u1"
|
||||
assert (
|
||||
call_kwargs.kwargs["value"]["user_role"] == LitellmUserRoles.PROXY_ADMIN.value
|
||||
)
|
||||
assert isinstance(call_kwargs.kwargs["value"], LiteLLM_UserTable)
|
||||
assert call_kwargs.kwargs["value"].user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
assert call_kwargs.kwargs["model_type"] == LiteLLM_UserTable
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -452,7 +452,9 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_team_change():
|
|||
mock_cache.async_set_cache.assert_called_once()
|
||||
call_kwargs = mock_cache.async_set_cache.call_args
|
||||
assert call_kwargs.kwargs["key"] == "u1"
|
||||
assert set(call_kwargs.kwargs["value"]["teams"]) == {"team1", "team2"}
|
||||
assert isinstance(call_kwargs.kwargs["value"], LiteLLM_UserTable)
|
||||
assert set(call_kwargs.kwargs["value"].teams) == {"team1", "team2"}
|
||||
assert call_kwargs.kwargs["model_type"] == LiteLLM_UserTable
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -231,6 +231,50 @@ class TestTokenUtilities:
|
|||
result = get_stored_api_key()
|
||||
assert result is None
|
||||
|
||||
def test_get_stored_api_key_base_url_match(self):
|
||||
"""Stored key is returned when expected_base_url matches stored origin"""
|
||||
token_data = {"key": "sk-prod", "base_url": "https://real-proxy.com"}
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.cli_token_utils.load_cli_token",
|
||||
return_value=token_data,
|
||||
):
|
||||
assert (
|
||||
get_stored_api_key(expected_base_url="https://real-proxy.com")
|
||||
== "sk-prod"
|
||||
)
|
||||
|
||||
def test_get_stored_api_key_base_url_match_trailing_slash(self):
|
||||
"""Trailing slash on expected_base_url is normalised before comparison"""
|
||||
token_data = {"key": "sk-prod", "base_url": "https://real-proxy.com"}
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.cli_token_utils.load_cli_token",
|
||||
return_value=token_data,
|
||||
):
|
||||
assert (
|
||||
get_stored_api_key(expected_base_url="https://real-proxy.com/")
|
||||
== "sk-prod"
|
||||
)
|
||||
|
||||
def test_get_stored_api_key_base_url_mismatch(self):
|
||||
"""Stored key is NOT returned when expected_base_url differs from stored origin"""
|
||||
token_data = {"key": "sk-prod", "base_url": "https://real-proxy.com"}
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.cli_token_utils.load_cli_token",
|
||||
return_value=token_data,
|
||||
):
|
||||
assert get_stored_api_key(expected_base_url="https://evil.com") is None
|
||||
|
||||
def test_get_stored_api_key_old_token_no_base_url(self):
|
||||
"""Old tokens without a base_url field are rejected when origin check is requested"""
|
||||
token_data = {"key": "sk-old-token"}
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.cli_token_utils.load_cli_token",
|
||||
return_value=token_data,
|
||||
):
|
||||
assert (
|
||||
get_stored_api_key(expected_base_url="https://real-proxy.com") is None
|
||||
)
|
||||
|
||||
|
||||
class TestLoginCommand:
|
||||
"""Test login CLI command"""
|
||||
|
|
|
|||
126
tests/test_litellm/proxy/common_utils/test_cache_codec.py
Normal file
126
tests/test_litellm/proxy/common_utils/test_cache_codec.py
Normal file
|
|
@ -0,0 +1,126 @@
|
|||
import logging
|
||||
from typing import Optional
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
|
||||
|
||||
|
||||
class _SampleModel(BaseModel):
|
||||
name: str
|
||||
count: Optional[int] = None
|
||||
|
||||
|
||||
class _SampleSubModel(_SampleModel):
|
||||
pass
|
||||
|
||||
|
||||
class TestCacheCodecSerialize:
|
||||
def test_without_model_type_base_model_dumped_json_safe(self):
|
||||
m = _SampleModel(name="a", count=1)
|
||||
out = CacheCodec.serialize(m)
|
||||
assert out == {"name": "a", "count": 1}
|
||||
|
||||
def test_without_model_type_dict_unchanged(self):
|
||||
d = {"name": "x"}
|
||||
assert CacheCodec.serialize(d) is d
|
||||
|
||||
def test_without_model_type_primitive_unchanged(self):
|
||||
assert CacheCodec.serialize(42) == 42
|
||||
|
||||
def test_with_model_type_dict_validated_and_dumped(self):
|
||||
out = CacheCodec.serialize({"name": "b", "count": 2}, model_type=_SampleModel)
|
||||
assert out == {"name": "b", "count": 2}
|
||||
|
||||
def test_with_model_type_base_model_validated_and_dumped(self):
|
||||
m = _SampleModel(name="c", count=None)
|
||||
out = CacheCodec.serialize(m, model_type=_SampleModel)
|
||||
assert out == {"name": "c"}
|
||||
|
||||
def test_with_model_type_exclude_none_on_dump(self):
|
||||
out = CacheCodec.serialize({"name": "d"}, model_type=_SampleModel)
|
||||
assert out == {"name": "d"}
|
||||
assert "count" not in out
|
||||
|
||||
def test_with_model_type_non_dict_non_model_passthrough(self):
|
||||
assert CacheCodec.serialize("raw", model_type=_SampleModel) == "raw"
|
||||
|
||||
def test_with_model_type_invalid_dict_raises(self):
|
||||
with pytest.raises(ValidationError):
|
||||
CacheCodec.serialize({"count": 1}, model_type=_SampleModel)
|
||||
|
||||
def test_with_model_type_already_correct_instance_skips_revalidation(self):
|
||||
"""Fast-path: value is already model_type — model_validate must NOT be called."""
|
||||
m = _SampleModel(name="fast", count=7)
|
||||
with patch.object(_SampleModel, "model_validate", wraps=_SampleModel.model_validate) as mock_validate:
|
||||
out = CacheCodec.serialize(m, model_type=_SampleModel)
|
||||
assert out == {"name": "fast", "count": 7}
|
||||
mock_validate.assert_not_called()
|
||||
|
||||
def test_with_model_type_subclass_instance_skips_revalidation(self):
|
||||
"""Subclass is isinstance of base → should also take the fast path."""
|
||||
sub = _SampleSubModel(name="sub", count=2)
|
||||
with patch.object(_SampleModel, "model_validate", wraps=_SampleModel.model_validate) as mock_validate:
|
||||
out = CacheCodec.serialize(sub, model_type=_SampleModel)
|
||||
assert out == {"name": "sub", "count": 2}
|
||||
mock_validate.assert_not_called()
|
||||
|
||||
def test_with_model_type_dict_input_goes_through_model_validate(self):
|
||||
"""A dict value (not yet an instance) must still go through model_validate."""
|
||||
raw = {"name": "via-dict", "count": 5}
|
||||
with patch.object(
|
||||
_SampleModel, "model_validate", wraps=_SampleModel.model_validate
|
||||
) as mock_validate:
|
||||
out = CacheCodec.serialize(raw, model_type=_SampleModel)
|
||||
assert out == {"name": "via-dict", "count": 5}
|
||||
mock_validate.assert_called_once()
|
||||
|
||||
def test_with_model_type_incompatible_model_raises_validation_error(self):
|
||||
"""Passing a BaseModel whose fields don't satisfy model_type's required fields raises.
|
||||
|
||||
_IncompatibleModel only has `foo: int`, so when Pydantic v2 extracts its
|
||||
data and validates it against _SampleModel (which requires `name: str`),
|
||||
a ValidationError is raised.
|
||||
"""
|
||||
|
||||
class _IncompatibleModel(BaseModel):
|
||||
foo: int # missing required 'name' field of _SampleModel
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
CacheCodec.serialize(_IncompatibleModel(foo=1), model_type=_SampleModel)
|
||||
|
||||
|
||||
class TestCacheCodecDeserialize:
|
||||
def test_none_returns_none(self):
|
||||
assert CacheCodec.deserialize(None, _SampleModel) is None
|
||||
|
||||
def test_dict_validates_to_model(self):
|
||||
m = CacheCodec.deserialize({"name": "e", "count": 3}, _SampleModel)
|
||||
assert isinstance(m, _SampleModel)
|
||||
assert m.name == "e"
|
||||
assert m.count == 3
|
||||
|
||||
def test_instance_same_type_returned_as_is(self):
|
||||
original = _SampleModel(name="f")
|
||||
m = CacheCodec.deserialize(original, _SampleModel)
|
||||
assert m is original
|
||||
|
||||
def test_subclass_instance_accepted(self):
|
||||
sub = _SampleSubModel(name="g")
|
||||
m = CacheCodec.deserialize(sub, _SampleModel)
|
||||
assert m is sub
|
||||
|
||||
def test_wrong_type_returns_none(self):
|
||||
assert CacheCodec.deserialize("not-a-dict", _SampleModel) is None
|
||||
|
||||
def test_invalid_dict_returns_none_and_logs_warning(self, caplog):
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
out = CacheCodec.deserialize({"count": 1}, _SampleModel)
|
||||
assert out is None
|
||||
assert any(
|
||||
"CacheCodec.deserialize" in r.message and "_SampleModel" in r.message
|
||||
for r in caplog.records
|
||||
if r.levelno >= logging.WARNING
|
||||
), f"Expected deserialize validation warning. Records: {[r.message for r in caplog.records]}"
|
||||
219
tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py
Normal file
219
tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py
Normal file
|
|
@ -0,0 +1,219 @@
|
|||
import json
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.proxy_server import UserAPIKeyCacheTTLEnum
|
||||
|
||||
|
||||
class CapturingInMemoryCache(InMemoryCache):
|
||||
"""Records ``ttl`` passed into ``set_cache`` (what DualCache injects)."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.last_ttl: Any = None
|
||||
|
||||
def set_cache(self, key, value, **kwargs): # type: ignore[override]
|
||||
self.last_ttl = kwargs.get("ttl")
|
||||
super().set_cache(key, value, **kwargs)
|
||||
|
||||
|
||||
class FakeRedisCache(RedisCache):
|
||||
"""
|
||||
In-memory fake that enforces the UserApiKeyCache Redis payload contract.
|
||||
|
||||
For user_api_key_cache entries we expect Redis to store a JSON object (dict)
|
||||
produced by `CacheCodec.serialize(..., model_type=...)`.
|
||||
|
||||
This fake:
|
||||
- raises TypeError if the value is not a dict
|
||||
- raises TypeError if the dict is not JSON-serializable
|
||||
|
||||
Records the ``ttl`` kwarg DualCache forwards on each Redis write for tests.
|
||||
"""
|
||||
|
||||
def __init__(self): # noqa: super().__init__ skipped intentionally
|
||||
self._store: dict[str, str] = {}
|
||||
self.last_ttl: Any = None
|
||||
|
||||
def set_cache(self, key: str, value: Any, **kwargs): # type: ignore[override]
|
||||
if not isinstance(value, dict):
|
||||
raise TypeError("FakeRedisCache only accepts dict payloads")
|
||||
self.last_ttl = kwargs.get("ttl")
|
||||
self._store[key] = json.dumps(value)
|
||||
return True
|
||||
|
||||
def get_cache(self, key: str, **kwargs): # type: ignore[override]
|
||||
raw = self._store.get(key)
|
||||
if raw is None:
|
||||
return None
|
||||
return json.loads(raw)
|
||||
|
||||
async def async_set_cache(self, key: str, value: Any, **kwargs): # type: ignore[override]
|
||||
if not isinstance(value, dict):
|
||||
raise TypeError("FakeRedisCache only accepts dict payloads")
|
||||
self.last_ttl = kwargs.get("ttl")
|
||||
self._store[key] = json.dumps(value)
|
||||
return True
|
||||
|
||||
async def async_get_cache(self, key: str, **kwargs): # type: ignore[override]
|
||||
raw = self._store.get(key)
|
||||
if raw is None:
|
||||
return None
|
||||
return json.loads(raw)
|
||||
|
||||
def delete_cache(self, key: str): # type: ignore[override]
|
||||
self._store.pop(key, None)
|
||||
|
||||
async def async_delete_cache(self, key: str): # type: ignore[override]
|
||||
self._store.pop(key, None)
|
||||
|
||||
|
||||
def _make_key_obj(token: str = "tok") -> UserAPIKeyAuth:
|
||||
# Minimal object (UserAPIKeyAuth inherits token from base view).
|
||||
return UserAPIKeyAuth(token=token)
|
||||
|
||||
|
||||
class TestUserApiKeyCache:
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_set_in_memory_gets_enum_default_when_user_api_key_cache_ttl_omitted(
|
||||
self,
|
||||
):
|
||||
"""
|
||||
If ``general_settings.user_api_key_cache_ttl`` is absent, the proxy never
|
||||
calls ``update_cache_ttl``; ``user_api_key_cache`` keeps
|
||||
``default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl``.
|
||||
DualCache must forward that as the in-memory ``ttl`` kwarg on each set.
|
||||
"""
|
||||
mem = CapturingInMemoryCache()
|
||||
cache = UserApiKeyCache(
|
||||
in_memory_cache=mem,
|
||||
redis_cache=FakeRedisCache(),
|
||||
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value,
|
||||
)
|
||||
await cache.async_set_cache(
|
||||
"k",
|
||||
_make_key_obj("t"),
|
||||
model_type=UserAPIKeyAuth,
|
||||
)
|
||||
expected = UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
|
||||
assert mem.last_ttl == expected
|
||||
|
||||
def test_sync_set_in_memory_gets_enum_default_when_user_api_key_cache_ttl_omitted(
|
||||
self,
|
||||
):
|
||||
mem = CapturingInMemoryCache()
|
||||
cache = UserApiKeyCache(
|
||||
in_memory_cache=mem,
|
||||
redis_cache=FakeRedisCache(),
|
||||
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value,
|
||||
)
|
||||
cache.set_cache("sk", _make_key_obj("s"), model_type=UserAPIKeyAuth)
|
||||
assert mem.last_ttl == UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_set_forwards_default_in_memory_ttl_to_redis_layer(self):
|
||||
"""
|
||||
DualCache injects missing ``ttl`` from ``default_in_memory_ttl`` into kwargs
|
||||
before calling ``redis_cache.async_set_cache`` — Redis should receive the same
|
||||
TTL as memory (matches proxy defaults: enum 60s).
|
||||
"""
|
||||
fake = FakeRedisCache()
|
||||
cache = UserApiKeyCache(
|
||||
redis_cache=fake,
|
||||
default_in_memory_ttl=60,
|
||||
)
|
||||
|
||||
await cache.async_set_cache(
|
||||
key="ttl-key",
|
||||
value=_make_key_obj("ttl-tok"),
|
||||
model_type=UserAPIKeyAuth,
|
||||
)
|
||||
|
||||
assert fake.last_ttl == 60
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_set_explicit_ttl_override_reaches_redis(self):
|
||||
fake = FakeRedisCache()
|
||||
cache = UserApiKeyCache(
|
||||
redis_cache=fake,
|
||||
default_in_memory_ttl=60,
|
||||
)
|
||||
|
||||
await cache.async_set_cache(
|
||||
key="k",
|
||||
value=_make_key_obj("x"),
|
||||
model_type=UserAPIKeyAuth,
|
||||
ttl=900,
|
||||
)
|
||||
|
||||
assert fake.last_ttl == 900
|
||||
|
||||
def test_sync_set_forwards_default_in_memory_ttl_to_redis_layer(self):
|
||||
fake = FakeRedisCache()
|
||||
cache = UserApiKeyCache(
|
||||
redis_cache=fake,
|
||||
default_in_memory_ttl=45,
|
||||
)
|
||||
cache.set_cache(
|
||||
"sk",
|
||||
_make_key_obj("sync"),
|
||||
model_type=UserAPIKeyAuth,
|
||||
)
|
||||
assert fake.last_ttl == 45
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_set_typed_stores_serialized_payload_in_memory_and_redis(self):
|
||||
cache = UserApiKeyCache(redis_cache=FakeRedisCache())
|
||||
obj = _make_key_obj("abc")
|
||||
|
||||
await cache.async_set_cache("k", obj, model_type=UserAPIKeyAuth)
|
||||
|
||||
# In-memory hit should still be raw dict (not BaseModel) because wrapper
|
||||
# stores the serialized payload into both layers.
|
||||
raw = await cache.in_memory_cache.async_get_cache("k") # type: ignore[union-attr]
|
||||
assert isinstance(raw, dict)
|
||||
assert raw["token"] == "abc"
|
||||
|
||||
# Redis should also hold the same serialized dict
|
||||
redis_raw = await cache.redis_cache.async_get_cache("k") # type: ignore[union-attr]
|
||||
assert redis_raw == raw
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get_typed_returns_model_on_valid_hit(self):
|
||||
cache = UserApiKeyCache(redis_cache=FakeRedisCache())
|
||||
await cache.async_set_cache("k", {"token": "abc"}, model_type=UserAPIKeyAuth)
|
||||
|
||||
value = await cache.async_get_cache("k", model_type=UserAPIKeyAuth)
|
||||
assert value is not None
|
||||
assert isinstance(value, UserAPIKeyAuth)
|
||||
assert value.token == "abc"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get_typed_returns_none_on_validation_failure_after_hit(self):
|
||||
cache = UserApiKeyCache(redis_cache=FakeRedisCache())
|
||||
|
||||
# Bypass UserApiKeyCache.serialize: CacheCodec rejects non-dict cached values
|
||||
# for dict-based models (deserialize returns None).
|
||||
await cache.in_memory_cache.async_set_cache(
|
||||
key="k", value="invalid-payload-not-a-dict"
|
||||
)
|
||||
|
||||
value = await cache.async_get_cache("k", model_type=UserAPIKeyAuth)
|
||||
assert value is None
|
||||
|
||||
def test_fake_redis_cache_rejects_non_json_serializable_values(self):
|
||||
fake = FakeRedisCache()
|
||||
|
||||
class NotSerializable:
|
||||
pass
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
fake.set_cache("k", NotSerializable())
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
fake.set_cache("k2", {"ok": NotSerializable()})
|
||||
|
|
@ -778,3 +778,373 @@ def test_get_callback_identifier_custom_logger_registry_and_fallback():
|
|||
result = get_callback_identifier(my_callback_function)
|
||||
# Should fall back to callback_name() which returns __name__
|
||||
assert result == "my_callback_function"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# /health response shape: model-access scoping and display-field allowlist
|
||||
# ---------------------------------------------------------------------------
|
||||
# These tests pin the contract that the /health response (a) only includes
|
||||
# deployments the calling key is allowed to see, and (b) does not return
|
||||
# provider routing fields like api_base / api_version. They guard against
|
||||
# regressions that would widen the response shape.
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_endpoint_filters_model_list_by_user_access():
|
||||
"""
|
||||
health_endpoint() should restrict _llm_model_list to deployments whose
|
||||
model_name appears in user_api_key_dict.models before running the health
|
||||
check. A key scoped to ["model-a"] should only see model-a in the result,
|
||||
not other deployments configured on the proxy.
|
||||
"""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
|
||||
|
||||
full_model_list = [
|
||||
{
|
||||
"model_name": "model-a",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o",
|
||||
"api_base": "https://example-a.test",
|
||||
},
|
||||
"model_info": {"id": "id-a"},
|
||||
},
|
||||
{
|
||||
"model_name": "model-b",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o",
|
||||
"api_base": "https://example-b.test",
|
||||
"api_version": "2024-10-21",
|
||||
},
|
||||
"model_info": {"id": "id-b"},
|
||||
},
|
||||
]
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="hashed-test-key",
|
||||
models=["model-a"],
|
||||
)
|
||||
|
||||
captured: dict = {}
|
||||
|
||||
async def fake_perform(**kwargs):
|
||||
captured["model_list"] = kwargs["model_list"]
|
||||
return {
|
||||
"healthy_endpoints": [],
|
||||
"unhealthy_endpoints": [],
|
||||
"healthy_count": 0,
|
||||
"unhealthy_count": 0,
|
||||
}
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_model_list", full_model_list),
|
||||
patch("litellm.proxy.proxy_server.llm_router", None),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", None),
|
||||
patch("litellm.proxy.proxy_server.use_background_health_checks", False),
|
||||
patch("litellm.proxy.proxy_server.user_model", None),
|
||||
patch("litellm.proxy.proxy_server.health_check_results", {}),
|
||||
patch("litellm.proxy.proxy_server.health_check_details", True),
|
||||
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
|
||||
patch(
|
||||
"litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save",
|
||||
side_effect=fake_perform,
|
||||
),
|
||||
):
|
||||
from fastapi import Response
|
||||
|
||||
await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict)
|
||||
|
||||
assert (
|
||||
"model_list" in captured
|
||||
), "health_endpoint did not call _perform_health_check_and_save"
|
||||
returned_names = {m["model_name"] for m in captured["model_list"]}
|
||||
assert returned_names == {
|
||||
"model-a"
|
||||
}, f"health_endpoint did not scope model_list to caller access: {returned_names}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_endpoint_filters_background_cache_by_user_access():
|
||||
"""
|
||||
When background_health_checks is enabled, health_endpoint() should also
|
||||
scope the cached result to the caller's allowed models rather than
|
||||
returning the cache verbatim.
|
||||
"""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
|
||||
|
||||
full_model_list = [
|
||||
{
|
||||
"model_name": "model-a",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o",
|
||||
"api_base": "https://example-a.test",
|
||||
},
|
||||
"model_info": {"id": "id-a"},
|
||||
},
|
||||
{
|
||||
"model_name": "model-b",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o",
|
||||
"api_base": "https://example-b.test",
|
||||
},
|
||||
"model_info": {"id": "id-b"},
|
||||
},
|
||||
]
|
||||
|
||||
cached_results = {
|
||||
"healthy_endpoints": [
|
||||
{
|
||||
"model": "openai/gpt-4o",
|
||||
"model_id": "id-a",
|
||||
"api_base": "https://example-a.test",
|
||||
},
|
||||
{
|
||||
"model": "openai/gpt-4o",
|
||||
"model_id": "id-b",
|
||||
"api_base": "https://example-b.test",
|
||||
},
|
||||
],
|
||||
"unhealthy_endpoints": [],
|
||||
"healthy_count": 2,
|
||||
"unhealthy_count": 0,
|
||||
}
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="hashed-test-key",
|
||||
models=["model-a"],
|
||||
)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_model_list", full_model_list),
|
||||
patch("litellm.proxy.proxy_server.llm_router", None),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", None),
|
||||
patch("litellm.proxy.proxy_server.use_background_health_checks", True),
|
||||
patch("litellm.proxy.proxy_server.user_model", None),
|
||||
patch("litellm.proxy.proxy_server.health_check_results", cached_results),
|
||||
patch("litellm.proxy.proxy_server.health_check_details", True),
|
||||
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
|
||||
):
|
||||
from fastapi import Response
|
||||
|
||||
result = await health_endpoint(
|
||||
response=Response(), user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
|
||||
# Sanity: the source cache had two entries before scoping; the scoping
|
||||
# step is what reduces it to one. (This guards against the test passing
|
||||
# vacuously when the cache filter drops everything because cached
|
||||
# entries lack the model_id key — both entries carry model_id above.)
|
||||
assert len(cached_results["healthy_endpoints"]) == 2
|
||||
assert all(
|
||||
ep.get("model_id") for ep in cached_results["healthy_endpoints"]
|
||||
), "test fixture invariant: every cached entry must carry a model_id"
|
||||
|
||||
# The non-admin caller must not see api_base on the returned cache entries.
|
||||
returned = result.get("healthy_endpoints", [])
|
||||
assert (
|
||||
len(returned) == 1
|
||||
), f"expected exactly one cached entry after scoping, got {len(returned)}"
|
||||
assert returned[0]["model_id"] == "id-a"
|
||||
assert "api_base" not in returned[0]
|
||||
assert result["healthy_count"] == 1
|
||||
assert result["unhealthy_count"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_endpoint_admin_sees_routing_fields_non_admin_does_not():
|
||||
"""
|
||||
A proxy admin should still see ``api_base`` and ``api_version`` in the
|
||||
/health response so they can tell which Vertex region / Azure resource
|
||||
+ API version is healthy. A non-admin caller must not — both fields
|
||||
should be stripped, and the response should carry a notice header so
|
||||
non-admin clients can detect the change programmatically.
|
||||
"""
|
||||
from fastapi import Response
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
|
||||
|
||||
full_model_list = [
|
||||
{
|
||||
"model_name": "model-a",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o",
|
||||
"api_base": "https://example-a.test",
|
||||
},
|
||||
"model_info": {"id": "id-a"},
|
||||
},
|
||||
]
|
||||
cached_results = {
|
||||
"healthy_endpoints": [
|
||||
{
|
||||
"model": "openai/gpt-4o",
|
||||
"model_id": "id-a",
|
||||
"api_base": "https://us-central1-aiplatform.googleapis.com/v1/projects/p",
|
||||
"api_version": "2024-10-21",
|
||||
},
|
||||
],
|
||||
"unhealthy_endpoints": [],
|
||||
"healthy_count": 1,
|
||||
"unhealthy_count": 0,
|
||||
}
|
||||
|
||||
admin_key = UserAPIKeyAuth(
|
||||
api_key="hashed-admin-key",
|
||||
models=["model-a"],
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
non_admin_key = UserAPIKeyAuth(
|
||||
api_key="hashed-user-key",
|
||||
models=["model-a"],
|
||||
)
|
||||
|
||||
common_patches = [
|
||||
patch("litellm.proxy.proxy_server.llm_model_list", full_model_list),
|
||||
patch("litellm.proxy.proxy_server.llm_router", None),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", None),
|
||||
patch("litellm.proxy.proxy_server.use_background_health_checks", True),
|
||||
patch("litellm.proxy.proxy_server.user_model", None),
|
||||
patch("litellm.proxy.proxy_server.health_check_results", cached_results),
|
||||
patch("litellm.proxy.proxy_server.health_check_details", True),
|
||||
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
|
||||
]
|
||||
|
||||
for p in common_patches:
|
||||
p.start()
|
||||
try:
|
||||
admin_response = Response()
|
||||
non_admin_response = Response()
|
||||
admin_result = await health_endpoint(
|
||||
response=admin_response, user_api_key_dict=admin_key
|
||||
)
|
||||
non_admin_result = await health_endpoint(
|
||||
response=non_admin_response, user_api_key_dict=non_admin_key
|
||||
)
|
||||
finally:
|
||||
for p in common_patches:
|
||||
p.stop()
|
||||
|
||||
admin_eps = admin_result.get("healthy_endpoints", [])
|
||||
non_admin_eps = non_admin_result.get("healthy_endpoints", [])
|
||||
|
||||
assert len(admin_eps) == 1
|
||||
assert (
|
||||
admin_eps[0]["api_base"]
|
||||
== "https://us-central1-aiplatform.googleapis.com/v1/projects/p"
|
||||
), "admin must see the full api_base so they can identify the region"
|
||||
assert (
|
||||
admin_eps[0]["api_version"] == "2024-10-21"
|
||||
), "admin must see api_version so they can distinguish provider deployments"
|
||||
|
||||
assert len(non_admin_eps) == 1
|
||||
assert "api_base" not in non_admin_eps[0]
|
||||
assert "api_version" not in non_admin_eps[0]
|
||||
|
||||
# Non-admin response must advertise that api_base/api_version were
|
||||
# withheld so clients that previously parsed them can detect the change.
|
||||
assert (
|
||||
non_admin_response.headers.get("Litellm-Health-Field-Notice")
|
||||
== "api_base and api_version are admin-only on this endpoint"
|
||||
)
|
||||
assert "Litellm-Health-Field-Notice" not in admin_response.headers
|
||||
|
||||
# Stripping must produce a copy — the shared cache must still carry the
|
||||
# routing fields so the next admin caller can read them.
|
||||
cached_first = cached_results["healthy_endpoints"][0]
|
||||
assert (
|
||||
cached_first["api_base"]
|
||||
== "https://us-central1-aiplatform.googleapis.com/v1/projects/p"
|
||||
)
|
||||
assert cached_first["api_version"] == "2024-10-21"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_endpoint_warns_when_scoped_models_lack_model_id():
|
||||
"""
|
||||
When a scoped key's accessible models exist on the proxy but none of the
|
||||
matching deployments expose a ``model_info.id``, the cache filter drops
|
||||
everything. The response should include a structured ``warnings`` field
|
||||
so the caller can distinguish "no deployments configured" from
|
||||
"deployments excluded due to missing model_info.id".
|
||||
"""
|
||||
from fastapi import Response
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
|
||||
|
||||
full_model_list = [
|
||||
{
|
||||
"model_name": "model-a",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o",
|
||||
"api_base": "https://example-a.test",
|
||||
},
|
||||
# Intentionally no model_info.id — this is the misconfiguration
|
||||
# the warnings field is meant to flag.
|
||||
"model_info": {},
|
||||
},
|
||||
]
|
||||
cached_results = {
|
||||
"healthy_endpoints": [
|
||||
{
|
||||
"model": "openai/gpt-4o",
|
||||
"model_id": "id-a",
|
||||
"api_base": "https://example-a.test",
|
||||
},
|
||||
],
|
||||
"unhealthy_endpoints": [],
|
||||
"healthy_count": 1,
|
||||
"unhealthy_count": 0,
|
||||
}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="hashed-user-key",
|
||||
models=["model-a"],
|
||||
)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_model_list", full_model_list),
|
||||
patch("litellm.proxy.proxy_server.llm_router", None),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", None),
|
||||
patch("litellm.proxy.proxy_server.use_background_health_checks", True),
|
||||
patch("litellm.proxy.proxy_server.user_model", None),
|
||||
patch("litellm.proxy.proxy_server.health_check_results", cached_results),
|
||||
patch("litellm.proxy.proxy_server.health_check_details", True),
|
||||
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
|
||||
):
|
||||
result = await health_endpoint(
|
||||
response=Response(), user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
|
||||
assert result["healthy_count"] == 0
|
||||
assert result["unhealthy_count"] == 0
|
||||
assert "warnings" in result, (
|
||||
"empty cache result must surface a warnings field so the caller "
|
||||
"can distinguish 'no deployments' from 'deployments excluded'"
|
||||
)
|
||||
assert any("model_info.id" in w for w in result["warnings"])
|
||||
|
||||
|
||||
def test_clean_endpoint_data_strips_credentials_keeps_routing_fields():
|
||||
"""
|
||||
_clean_endpoint_data() drops credentials but leaves api_base /
|
||||
api_version intact — the per-caller hide/show happens in the endpoint
|
||||
layer based on user role, not in the cleaning helper. This guarantees
|
||||
proxy admins continue to see those fields in the /health response.
|
||||
"""
|
||||
from litellm.proxy.health_check import _clean_endpoint_data
|
||||
|
||||
raw = {
|
||||
"model": "openai/gpt-4o",
|
||||
"api_key": "sk-test",
|
||||
"api_base": "https://example.test/v1",
|
||||
"api_version": "2024-10-21",
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
}
|
||||
|
||||
cleaned = _clean_endpoint_data(raw, details=True)
|
||||
|
||||
assert "api_key" not in cleaned
|
||||
assert "aws_access_key_id" not in cleaned
|
||||
assert cleaned.get("api_base") == "https://example.test/v1"
|
||||
assert cleaned.get("api_version") == "2024-10-21"
|
||||
|
|
|
|||
|
|
@ -738,15 +738,32 @@ def test_delete_access_group_patches_cached_team_and_key(
|
|||
return_value=None
|
||||
)
|
||||
|
||||
# Build cached key object (returned from user_api_key_cache)
|
||||
if key_cache_group_ids is not None:
|
||||
cached_key = UserAPIKeyAuth(
|
||||
token="hashed-key-1",
|
||||
access_group_ids=list(key_cache_group_ids),
|
||||
# user_api_key_cache is queried both for teams (fallback after dual_cache) and
|
||||
# hashed keys — return the right stub per ``key``. A single AsyncMock(return_value=key)
|
||||
# would wrongly serve the key blob for ``team_id:team-1`` and trigger team patching.
|
||||
# Use a synchronous side_effect (not async def): AsyncMock awaits coroutine side_effects
|
||||
# inconsistently across Python/unittest versions; sync returns are awaited as immediate results.
|
||||
def user_cache_get_side_effect(*args, **kwargs):
|
||||
cache_key = (
|
||||
kwargs.get("key") if "key" in kwargs else (args[0] if args else None)
|
||||
)
|
||||
mock_cache.async_get_cache = AsyncMock(return_value=cached_key)
|
||||
else:
|
||||
mock_cache.async_get_cache = AsyncMock(return_value=None)
|
||||
if cache_key == "team_id:team-1":
|
||||
if team_cache_group_ids is None:
|
||||
return None
|
||||
return LiteLLM_TeamTableCachedObj(
|
||||
team_id="team-1",
|
||||
access_group_ids=list(team_cache_group_ids),
|
||||
)
|
||||
if cache_key == "hashed-key-1":
|
||||
if key_cache_group_ids is None:
|
||||
return None
|
||||
return UserAPIKeyAuth(
|
||||
token="hashed-key-1",
|
||||
access_group_ids=list(key_cache_group_ids),
|
||||
)
|
||||
return None
|
||||
|
||||
mock_cache.async_get_cache = AsyncMock(side_effect=user_cache_get_side_effect)
|
||||
|
||||
resp = client.delete("/v1/access_group/ag-to-delete")
|
||||
assert resp.status_code == 204
|
||||
|
|
@ -803,7 +820,7 @@ def test_delete_access_group_patches_cached_team_and_key(
|
|||
|
||||
|
||||
def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks):
|
||||
"""Delete correctly patches a key cached as a raw dict (not UserAPIKeyAuth)."""
|
||||
"""Delete patches key cache — mock returns UserAPIKeyAuth (what UserApiKeyCache emits after deserialize)."""
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
|
||||
client_and_mocks
|
||||
)
|
||||
|
|
@ -826,12 +843,24 @@ def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks):
|
|||
return_value=None
|
||||
)
|
||||
|
||||
# Key cached as a plain dict (as can happen with Redis serialization)
|
||||
# Serialized shape from Redis dict; UserApiKeyCache.async_get_cache(model_type=...) yields a model — simulate that.
|
||||
cached_key_payload = {
|
||||
"token": "hashed-key-dict",
|
||||
"access_group_ids": ["ag-to-delete", "ag-other"],
|
||||
}
|
||||
|
||||
def user_cache_get_dict_when_key_matches(*args, **kwargs):
|
||||
cache_key = (
|
||||
kwargs.get("key") if "key" in kwargs else (args[0] if args else None)
|
||||
)
|
||||
if cache_key == "team_id:team-1":
|
||||
return None
|
||||
if cache_key == "hashed-key-dict":
|
||||
return UserAPIKeyAuth.model_validate(cached_key_payload)
|
||||
return None
|
||||
|
||||
mock_cache.async_get_cache = AsyncMock(
|
||||
return_value={
|
||||
"token": "hashed-key-dict",
|
||||
"access_group_ids": ["ag-to-delete", "ag-other"],
|
||||
}
|
||||
side_effect=user_cache_get_dict_when_key_matches
|
||||
)
|
||||
|
||||
resp = client.delete("/v1/access_group/ag-to-delete")
|
||||
|
|
|
|||
|
|
@ -5512,6 +5512,9 @@ async def test_update_team_guardrails_with_org_id():
|
|||
return_value=mock_updated_team
|
||||
)
|
||||
mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data)
|
||||
# async_get_cache must be an AsyncMock so `await` in get_org_object works
|
||||
mock_cache.async_get_cache = AsyncMock(return_value=None)
|
||||
mock_cache.async_set_cache = AsyncMock()
|
||||
|
||||
# Mock llm_router
|
||||
mock_router = MagicMock()
|
||||
|
|
|
|||
|
|
@ -26,6 +26,15 @@ async def fake_valid_auth(request, api_key):
|
|||
return
|
||||
|
||||
|
||||
async def fake_valid_auth_reads_body(request, api_key, **kwargs):
|
||||
"""
|
||||
Like real user_api_key_auth, consumes the ASGI body stream. Regression test
|
||||
for successful auth passing a drained receive to the inner app (hang).
|
||||
"""
|
||||
await request.body()
|
||||
return
|
||||
|
||||
|
||||
async def fake_invalid_auth(request, api_key):
|
||||
print("running fake invalid auth", request, api_key)
|
||||
# Simulate invalid auth by raising an exception.
|
||||
|
|
@ -62,6 +71,28 @@ def app_with_middleware():
|
|||
return app
|
||||
|
||||
|
||||
def test_valid_auth_metrics_after_body_consumed(app_with_middleware, monkeypatch):
|
||||
"""
|
||||
Auth that reads the request body must not cause /metrics to hang on success.
|
||||
"""
|
||||
litellm.require_auth_for_metrics_endpoint = True
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth",
|
||||
fake_valid_auth_reads_body,
|
||||
)
|
||||
|
||||
client = TestClient(app_with_middleware)
|
||||
headers = {SpecialHeaders.openai_authorization.value: "valid"}
|
||||
|
||||
response = client.get("/metrics", headers=headers)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == {"msg": "metrics OK"}
|
||||
|
||||
response = client.get("/metrics/", headers=headers)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == {"msg": "metrics OK"}
|
||||
|
||||
|
||||
def test_valid_auth_metrics(app_with_middleware, monkeypatch):
|
||||
"""
|
||||
Test that a request to /metrics (and /metrics/) with valid auth headers passes.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,236 @@
|
|||
"""
|
||||
Tests for _filter_models_by_team_id resolving access group names.
|
||||
|
||||
Verifies that when a team's `models` field contains an access group name
|
||||
(e.g., "Group-A"), the filter resolves it to the member model names before
|
||||
looking up deployments — matching the behavior of the auth path in
|
||||
auth_checks.py:model_in_access_group().
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
from litellm.proxy.proxy_server import _filter_models_by_team_id
|
||||
|
||||
|
||||
def _make_model(model_name: str, model_id: str, access_groups: list[str] = None):
|
||||
"""Helper to build a model dict matching the router's format."""
|
||||
return {
|
||||
"model_name": model_name,
|
||||
"litellm_params": {"model": model_name},
|
||||
"model_info": {
|
||||
"id": model_id,
|
||||
"access_groups": access_groups or [],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _make_team(models: list[str], team_id: str = "team_alpha"):
|
||||
"""Helper to build a mock team DB object."""
|
||||
mock = MagicMock()
|
||||
mock.model_dump.return_value = {
|
||||
"team_id": team_id,
|
||||
"team_alias": "Team Alpha",
|
||||
"models": models,
|
||||
"max_budget": None,
|
||||
"spend": 0.0,
|
||||
"blocked": False,
|
||||
"members_with_roles": [],
|
||||
"metadata": {},
|
||||
}
|
||||
return mock
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_resolves_access_group_names():
|
||||
"""
|
||||
When team.models contains an access group name, _filter_models_by_team_id
|
||||
should resolve it to the member models and return only those deployments.
|
||||
"""
|
||||
# Models on the proxy
|
||||
gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"])
|
||||
gpt5 = _make_model("gpt-5", "id-2", ["Group-A"])
|
||||
claude = _make_model("claude-3", "id-3", ["Group-B"])
|
||||
|
||||
all_models = [gpt4o, gpt5, claude]
|
||||
|
||||
# Router mock
|
||||
mock_router = MagicMock()
|
||||
# get_model_access_groups returns {group_name: [model_names]}
|
||||
mock_router.get_model_access_groups.return_value = {
|
||||
"Group-A": ["gpt-4o", "gpt-5"],
|
||||
"Group-B": ["claude-3"],
|
||||
}
|
||||
|
||||
# get_model_list returns deployments matching a model_name
|
||||
def fake_get_model_list(model_name=None, team_id=None):
|
||||
return [m for m in all_models if m["model_name"] == model_name]
|
||||
|
||||
mock_router.get_model_list = MagicMock(side_effect=fake_get_model_list)
|
||||
|
||||
# Team has models: ["Group-A"] — an access group name, not a literal model
|
||||
team_db = _make_team(models=["Group-A"])
|
||||
|
||||
# Prisma mock
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db)
|
||||
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
|
||||
|
||||
result = await _filter_models_by_team_id(
|
||||
all_models=all_models,
|
||||
team_id="team_alpha",
|
||||
prisma_client=mock_prisma,
|
||||
llm_router=mock_router,
|
||||
)
|
||||
|
||||
result_ids = {m["model_info"]["id"] for m in result}
|
||||
# Should include gpt-4o and gpt-5 (Group-A), but NOT claude-3 (Group-B)
|
||||
assert result_ids == {
|
||||
"id-1",
|
||||
"id-2",
|
||||
}, f"Expected Group-A models only, got {result_ids}"
|
||||
|
||||
# Verify DB fallback query received resolved model names, not access group name
|
||||
call_kwargs = mock_prisma.db.litellm_proxymodeltable.find_many.call_args[1]
|
||||
assert set(call_kwargs["where"]["model_name"]["in"]) == {
|
||||
"gpt-4o",
|
||||
"gpt-5",
|
||||
}, "find_many should receive resolved model names, not the access group name"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_resolves_mix_of_access_groups_and_literal_names():
|
||||
"""
|
||||
When team.models contains both an access group name and a literal model name,
|
||||
both should be resolved correctly.
|
||||
"""
|
||||
gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"])
|
||||
gpt5 = _make_model("gpt-5", "id-2", ["Group-A"])
|
||||
claude = _make_model("claude-3", "id-3", ["Group-B"])
|
||||
mistral = _make_model("mistral-large", "id-4", []) # no access group
|
||||
|
||||
all_models = [gpt4o, gpt5, claude, mistral]
|
||||
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_model_access_groups.return_value = {
|
||||
"Group-A": ["gpt-4o", "gpt-5"],
|
||||
"Group-B": ["claude-3"],
|
||||
}
|
||||
|
||||
def fake_get_model_list(model_name=None, team_id=None):
|
||||
return [m for m in all_models if m["model_name"] == model_name]
|
||||
|
||||
mock_router.get_model_list = MagicMock(side_effect=fake_get_model_list)
|
||||
|
||||
# Team has access to Group-A (access group) + mistral-large (literal name)
|
||||
team_db = _make_team(models=["Group-A", "mistral-large"])
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db)
|
||||
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
|
||||
|
||||
result = await _filter_models_by_team_id(
|
||||
all_models=all_models,
|
||||
team_id="team_alpha",
|
||||
prisma_client=mock_prisma,
|
||||
llm_router=mock_router,
|
||||
)
|
||||
|
||||
result_ids = {m["model_info"]["id"] for m in result}
|
||||
# Group-A models + mistral-large, but NOT claude-3
|
||||
assert result_ids == {
|
||||
"id-1",
|
||||
"id-2",
|
||||
"id-4",
|
||||
}, f"Expected Group-A + mistral-large, got {result_ids}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_excludes_models_from_other_access_group():
|
||||
"""
|
||||
Models belonging only to a different access group must not appear in results.
|
||||
"""
|
||||
gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"])
|
||||
claude = _make_model("claude-3", "id-3", ["Group-B"])
|
||||
llama = _make_model("llama-4", "id-4", ["Group-B"])
|
||||
|
||||
all_models = [gpt4o, claude, llama]
|
||||
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_model_access_groups.return_value = {
|
||||
"Group-A": ["gpt-4o"],
|
||||
"Group-B": ["claude-3", "llama-4"],
|
||||
}
|
||||
|
||||
def fake_get_model_list(model_name=None, team_id=None):
|
||||
return [m for m in all_models if m["model_name"] == model_name]
|
||||
|
||||
mock_router.get_model_list = MagicMock(side_effect=fake_get_model_list)
|
||||
|
||||
team_db = _make_team(models=["Group-A"])
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db)
|
||||
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
|
||||
|
||||
result = await _filter_models_by_team_id(
|
||||
all_models=all_models,
|
||||
team_id="team_alpha",
|
||||
prisma_client=mock_prisma,
|
||||
llm_router=mock_router,
|
||||
)
|
||||
|
||||
result_names = {m["model_name"] for m in result}
|
||||
assert "claude-3" not in result_names, "Group-B model should not be accessible"
|
||||
assert "llama-4" not in result_names, "Group-B model should not be accessible"
|
||||
assert "gpt-4o" in result_names, "Group-A model should be accessible"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_db_fallback_receives_resolved_model_names():
|
||||
"""
|
||||
When get_model_list returns no results (forcing the DB fallback path),
|
||||
the DB query should receive resolved model names, not the raw access group name.
|
||||
"""
|
||||
gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"])
|
||||
all_models = [gpt4o]
|
||||
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_model_access_groups.return_value = {
|
||||
"Group-A": ["gpt-4o", "gpt-5"],
|
||||
}
|
||||
# get_model_list returns nothing — forces reliance on the DB fallback
|
||||
mock_router.get_model_list = MagicMock(return_value=[])
|
||||
|
||||
team_db = _make_team(models=["Group-A"])
|
||||
|
||||
# DB returns a model that the router didn't find
|
||||
mock_db_model = MagicMock()
|
||||
mock_db_model.model_id = "id-db-1"
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db)
|
||||
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(
|
||||
return_value=[mock_db_model]
|
||||
)
|
||||
|
||||
result = await _filter_models_by_team_id(
|
||||
all_models=all_models,
|
||||
team_id="team_alpha",
|
||||
prisma_client=mock_prisma,
|
||||
llm_router=mock_router,
|
||||
)
|
||||
|
||||
# Verify DB query received resolved names, not "Group-A"
|
||||
call_kwargs = mock_prisma.db.litellm_proxymodeltable.find_many.call_args[1]
|
||||
queried_names = set(call_kwargs["where"]["model_name"]["in"])
|
||||
assert queried_names == {
|
||||
"gpt-4o",
|
||||
"gpt-5",
|
||||
}, f"DB query should receive resolved model names, got {queried_names}"
|
||||
assert "Group-A" not in queried_names, "Raw access group name should not be in DB query"
|
||||
145
tests/test_litellm/proxy/test_redis_auth_cache_flag.py
Normal file
145
tests/test_litellm/proxy/test_redis_auth_cache_flag.py
Normal file
|
|
@ -0,0 +1,145 @@
|
|||
"""
|
||||
Tests for the enable_redis_auth_cache litellm_settings flag.
|
||||
|
||||
Verifies that _init_cache attaches Redis to user_api_key_cache only when
|
||||
the flag is explicitly set to True, and leaves it in-memory-only otherwise.
|
||||
"""
|
||||
|
||||
from contextlib import contextmanager
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.caching.caching import RedisCache
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeRedisCache(RedisCache):
|
||||
"""
|
||||
Minimal RedisCache subclass that passes isinstance checks without
|
||||
requiring a real Redis connection. __init__ is bypassed so no
|
||||
network calls are made.
|
||||
"""
|
||||
|
||||
def __init__(self): # noqa: super().__init__ skipped intentionally
|
||||
self._store = {}
|
||||
|
||||
def set_cache(self, key, value, **kwargs): # type: ignore[override]
|
||||
# Enforce Redis JSON-serializable payload contract.
|
||||
self._store[key] = json.dumps(value)
|
||||
return True
|
||||
|
||||
def get_cache(self, key, **kwargs): # type: ignore[override]
|
||||
raw = self._store.get(key)
|
||||
if raw is None:
|
||||
return None
|
||||
return json.loads(raw)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patched_init_cache(litellm_settings: dict, cache_params: dict):
|
||||
"""
|
||||
Context manager that:
|
||||
1. Replaces the module-level globals with fresh DualCache instances.
|
||||
2. Patches ``litellm.Cache`` (locally imported inside _init_cache) so
|
||||
it returns a fake cache whose ``.cache`` attribute is a
|
||||
_FakeRedisCache (passes the isinstance guard in _init_cache).
|
||||
3. Extracts enable_redis_auth_cache from litellm_settings and passes it
|
||||
as the second argument to _init_cache (matching production behaviour).
|
||||
4. Yields (user_api_key_cache, spend_counter_cache) after calling
|
||||
_init_cache, then restores everything.
|
||||
"""
|
||||
fake_redis = _FakeRedisCache()
|
||||
|
||||
mock_litellm_cache = MagicMock()
|
||||
mock_litellm_cache.cache = fake_redis
|
||||
|
||||
fresh_user_cache = DualCache()
|
||||
fresh_spend_cache = DualCache()
|
||||
|
||||
enable_redis_auth_cache = litellm_settings.get("enable_redis_auth_cache", False)
|
||||
|
||||
with (
|
||||
patch.object(ps, "user_api_key_cache", fresh_user_cache),
|
||||
patch.object(ps, "spend_counter_cache", fresh_spend_cache),
|
||||
patch.object(ps, "llm_router", None),
|
||||
# Cache is locally imported inside _init_cache: patch it at source.
|
||||
patch("litellm.Cache", return_value=mock_litellm_cache),
|
||||
):
|
||||
litellm.cache = None
|
||||
ps.ProxyConfig()._init_cache(cache_params, enable_redis_auth_cache)
|
||||
yield fresh_user_cache, fresh_spend_cache
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRedisAuthCacheFlag:
|
||||
def test_flag_true_attaches_redis_to_user_api_key_cache(self):
|
||||
"""When enable_redis_auth_cache=True, user_api_key_cache.redis_cache must be set."""
|
||||
with _patched_init_cache(
|
||||
litellm_settings={"enable_redis_auth_cache": True},
|
||||
cache_params={"type": "redis", "host": "localhost", "port": 6379},
|
||||
) as (user_cache, _):
|
||||
assert user_cache.redis_cache is not None, (
|
||||
"Redis should be attached to user_api_key_cache when "
|
||||
"enable_redis_auth_cache=True"
|
||||
)
|
||||
|
||||
def test_flag_false_leaves_user_api_key_cache_in_memory_only(self):
|
||||
"""When enable_redis_auth_cache=False, user_api_key_cache must stay in-memory."""
|
||||
with _patched_init_cache(
|
||||
litellm_settings={"enable_redis_auth_cache": False},
|
||||
cache_params={"type": "redis", "host": "localhost", "port": 6379},
|
||||
) as (user_cache, _):
|
||||
assert user_cache.redis_cache is None, (
|
||||
"user_api_key_cache must remain in-memory-only when "
|
||||
"enable_redis_auth_cache=False"
|
||||
)
|
||||
|
||||
def test_flag_absent_leaves_user_api_key_cache_in_memory_only(self):
|
||||
"""When enable_redis_auth_cache is not set at all, default is in-memory-only."""
|
||||
with _patched_init_cache(
|
||||
litellm_settings={},
|
||||
cache_params={"type": "redis", "host": "localhost", "port": 6379},
|
||||
) as (user_cache, _):
|
||||
assert user_cache.redis_cache is None, (
|
||||
"user_api_key_cache must remain in-memory-only when "
|
||||
"enable_redis_auth_cache is absent from litellm_settings"
|
||||
)
|
||||
|
||||
def test_spend_counter_cache_always_gets_redis_regardless_of_flag(self):
|
||||
"""spend_counter_cache must receive Redis regardless of the auth-cache flag."""
|
||||
for flag_value in (True, False, None):
|
||||
ls = (
|
||||
{"enable_redis_auth_cache": flag_value}
|
||||
if flag_value is not None
|
||||
else {}
|
||||
)
|
||||
with _patched_init_cache(
|
||||
litellm_settings=ls,
|
||||
cache_params={"type": "redis", "host": "localhost", "port": 6379},
|
||||
) as (_, spend_cache):
|
||||
assert spend_cache.redis_cache is not None, (
|
||||
f"spend_counter_cache must always get Redis "
|
||||
f"(enable_redis_auth_cache={flag_value!r})"
|
||||
)
|
||||
|
||||
def test_flag_false_spend_gets_redis_but_user_cache_does_not(self):
|
||||
"""Explicit False: spend cache wired, auth cache left in-memory."""
|
||||
with _patched_init_cache(
|
||||
litellm_settings={"enable_redis_auth_cache": False},
|
||||
cache_params={"type": "redis", "host": "localhost", "port": 6379},
|
||||
) as (user_cache, spend_cache):
|
||||
assert spend_cache.redis_cache is not None
|
||||
assert user_cache.redis_cache is None
|
||||
|
|
@ -346,6 +346,34 @@ def test_tag_routing_with_list_of_tags_match_all():
|
|||
assert not is_valid_deployment_tag(["default"], ["teamA"], match_any=False)
|
||||
|
||||
|
||||
def test_strict_tag_routing_without_request_tags_blocks_header_regex_fallback():
|
||||
"""
|
||||
When tag_filtering_match_any=False, deployments with plain tags must require
|
||||
those request tags before header regex can match. A spoofed User-Agent must
|
||||
not route to a tagged deployment when the request has no tags.
|
||||
"""
|
||||
from litellm.router_strategy.tag_based_routing import _match_deployment
|
||||
|
||||
deployment = {
|
||||
"model_name": "restricted-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-4o",
|
||||
"tags": ["internal"],
|
||||
"tag_regex": ["^User-Agent: internal-tool"],
|
||||
},
|
||||
}
|
||||
|
||||
assert (
|
||||
_match_deployment(
|
||||
deployment=deployment,
|
||||
request_tags=None,
|
||||
header_strings=["User-Agent: internal-tool"],
|
||||
match_any=False,
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_router_free_paid_tier_with_responses_api():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import pytest
|
|||
from litellm import Router
|
||||
from litellm.router_utils.common_utils import (
|
||||
_deployment_supports_web_search,
|
||||
add_model_file_id_mappings,
|
||||
filter_team_based_models,
|
||||
filter_web_search_deployments,
|
||||
)
|
||||
|
|
@ -362,3 +363,112 @@ def test_invalidate_model_group_info_cache():
|
|||
# Invalidate and verify cache is cleared
|
||||
router._invalidate_model_group_info_cache()
|
||||
assert router._cached_get_model_group_info.cache_info().currsize == 0
|
||||
|
||||
|
||||
class TestAddModelFileIdMappings:
|
||||
"""Test cases for add_model_file_id_mappings.
|
||||
|
||||
The router may pass either a list of deployment dicts (multiple matched
|
||||
deployments) or a single deployment dict (when a specific deployment was
|
||||
resolved, e.g. because the requested model matched a `model_info.id`).
|
||||
Both shapes must produce a `{model_id: file_id}` mapping by extracting
|
||||
`model_info.id` from each deployment.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _make_response(file_id: str):
|
||||
response = Mock()
|
||||
response.id = file_id
|
||||
return response
|
||||
|
||||
def test_should_map_each_deployment_id_when_given_list(self):
|
||||
deployments = [
|
||||
{
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4"},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
{
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4"},
|
||||
"model_info": {"id": "deployment-2"},
|
||||
},
|
||||
]
|
||||
responses = [self._make_response("file-1"), self._make_response("file-2")]
|
||||
|
||||
result = add_model_file_id_mappings(deployments, responses)
|
||||
|
||||
assert result == {"deployment-1": "file-1", "deployment-2": "file-2"}
|
||||
|
||||
def test_should_extract_model_info_id_when_given_single_deployment_dict(self):
|
||||
"""Regression test: when `_common_checks_available_deployment` resolves
|
||||
a specific deployment (returned as a dict, not a list), the function
|
||||
must still extract `model_info.id` rather than iterate over the
|
||||
deployment's own keys (`model_name`, `litellm_params`, `model_info`).
|
||||
"""
|
||||
deployment = {
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4", "api_key": "sk-test"},
|
||||
"model_info": {"id": "deployment-1", "mode": "chat"},
|
||||
}
|
||||
responses = [self._make_response("file-1")]
|
||||
|
||||
result = add_model_file_id_mappings(deployment, responses)
|
||||
|
||||
assert result == {"deployment-1": "file-1"}
|
||||
assert all(isinstance(v, str) for v in result.values())
|
||||
|
||||
def test_should_handle_batch_model_when_id_matches_model_name(self):
|
||||
"""Regression test for the batch-model case: when `model_info.id` is
|
||||
intentionally set equal to `model_name`, the router resolves a single
|
||||
deployment via `has_model_id` and returns it as a dict. The mapping
|
||||
must contain only `{id: file_id}` with string values so the resulting
|
||||
`LiteLLM_ManagedFileTable` Pydantic validation passes.
|
||||
"""
|
||||
deployment = {
|
||||
"model_name": "openai/openai/gpt-5.5-batch",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-5.5",
|
||||
"api_key": "sk-test",
|
||||
"tpm": 40000000,
|
||||
"rpm": 15000,
|
||||
},
|
||||
"model_info": {
|
||||
"id": "openai/openai/gpt-5.5-batch",
|
||||
"mode": "batch",
|
||||
"base_model": "gpt-5.5",
|
||||
"access_groups": ["default-models"],
|
||||
},
|
||||
}
|
||||
responses = [self._make_response("file-batch-1")]
|
||||
|
||||
result = add_model_file_id_mappings(deployment, responses)
|
||||
|
||||
# Bug case would have produced keys ["model_name", "litellm_params",
|
||||
# "model_info"] with non-string values.
|
||||
assert result == {"openai/openai/gpt-5.5-batch": "file-batch-1"}
|
||||
assert "litellm_params" not in result
|
||||
assert "model_info" not in result
|
||||
|
||||
def test_should_skip_deployment_when_model_info_id_missing(self):
|
||||
deployments = [
|
||||
{
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4"},
|
||||
"model_info": {},
|
||||
},
|
||||
{
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4"},
|
||||
"model_info": {"id": "deployment-2"},
|
||||
},
|
||||
]
|
||||
responses = [self._make_response("file-1"), self._make_response("file-2")]
|
||||
|
||||
result = add_model_file_id_mappings(deployments, responses)
|
||||
|
||||
assert result == {"deployment-2": "file-2"}
|
||||
|
||||
def test_should_return_empty_mapping_when_given_empty_list(self):
|
||||
result = add_model_file_id_mappings([], [])
|
||||
assert result == {}
|
||||
|
|
|
|||
|
|
@ -6,8 +6,8 @@ import {
|
|||
deriveErrorMessage,
|
||||
handleError,
|
||||
} from "@/components/networking";
|
||||
import { all_admin_roles } from "@/utils/roles";
|
||||
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
|
||||
import { all_admin_roles } from "@/utils/roles";
|
||||
|
||||
// ── Types ────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
|
@ -81,7 +81,6 @@ export const useProjects = () => {
|
|||
return useQuery<ProjectResponse[]>({
|
||||
queryKey: projectKeys.list({}),
|
||||
queryFn: async () => fetchProjects(accessToken!),
|
||||
enabled:
|
||||
Boolean(accessToken) && all_admin_roles.includes(userRole || ""),
|
||||
enabled: Boolean(accessToken) && all_admin_roles.includes(userRole!),
|
||||
});
|
||||
};
|
||||
|
|
|
|||
|
|
@ -169,8 +169,8 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
|
|||
}
|
||||
}, [isAdmin, userID]);
|
||||
|
||||
// For non-admins, always pass their own user_id
|
||||
const effectiveUserId = isAdmin ? selectedUserId : userID || null;
|
||||
// For non-admins or "my-usage" view, always pass their own user_id
|
||||
const effectiveUserId = usageView === "my-usage" || !isAdmin ? userID || null : selectedUserId;
|
||||
|
||||
const startTime = useMemo(() => (dateValue.from ? new Date(dateValue.from) : null), [dateValue.from]);
|
||||
const endTime = useMemo(() => (dateValue.to ? new Date(dateValue.to) : null), [dateValue.to]);
|
||||
|
|
@ -477,10 +477,10 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
|
|||
}
|
||||
/>
|
||||
)}
|
||||
{/* Your Usage Panel */}
|
||||
{usageView === "global" && (
|
||||
{/* Your Usage / Global Usage Panel */}
|
||||
{(usageView === "global" || usageView === "my-usage") && (
|
||||
<>
|
||||
{isAdmin && (
|
||||
{isAdmin && usageView === "global" && (
|
||||
<div className="mb-4">
|
||||
<Text className="mb-2">Filter by user</Text>
|
||||
<Select
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ import {
|
|||
} from "@ant-design/icons";
|
||||
import { Badge, Select } from "antd";
|
||||
import React from "react";
|
||||
export type UsageOption = "global" | "organization" | "team" | "customer" | "tag" | "agent" | "user" | "user-agent-activity";
|
||||
export type UsageOption = "global" | "my-usage" | "organization" | "team" | "customer" | "tag" | "agent" | "user" | "user-agent-activity";
|
||||
export interface UsageViewSelectProps {
|
||||
value: UsageOption;
|
||||
onChange: (value: UsageOption) => void;
|
||||
|
|
@ -43,6 +43,13 @@ const OPTIONS: OptionConfig[] = [
|
|||
descriptionForNonAdmin: "View your usage",
|
||||
icon: <GlobalOutlined style={{ fontSize: "16px" }} />,
|
||||
},
|
||||
{
|
||||
value: "my-usage",
|
||||
label: "Your Usage",
|
||||
description: "View your own usage",
|
||||
icon: <UserOutlined style={{ fontSize: "16px" }} />,
|
||||
adminOnly: true,
|
||||
},
|
||||
{
|
||||
value: "organization",
|
||||
label: "Organization Usage",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue