mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into HEAD
# Conflicts: # litellm/proxy/middleware/prometheus_auth_middleware.py # tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py
This commit is contained in:
commit
0c864880a8
159 changed files with 44782 additions and 2548 deletions
75
.github/workflows/check-lazy-openapi-snapshot.yml
vendored
Normal file
75
.github/workflows/check-lazy-openapi-snapshot.yml
vendored
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
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
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.69"
|
||||
version = "0.4.70"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.69"
|
||||
version = "0.4.70"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
@ -392,6 +411,7 @@ class DualCache(BaseCache):
|
|||
value: float,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
local_only: bool = False,
|
||||
refresh_ttl: bool = False,
|
||||
**kwargs,
|
||||
) -> Optional[float]:
|
||||
"""
|
||||
|
|
@ -399,6 +419,9 @@ class DualCache(BaseCache):
|
|||
|
||||
Value - float - the value you want to increment by
|
||||
|
||||
Refresh_ttl - bool - if True, resets the Redis TTL on every write.
|
||||
Default False preserves window-style semantics.
|
||||
|
||||
Returns - the incremented value, or None if no cache backend is
|
||||
available (in_memory_cache is None and Redis failed/is absent).
|
||||
"""
|
||||
|
|
@ -415,6 +438,7 @@ class DualCache(BaseCache):
|
|||
value,
|
||||
parent_otel_span=parent_otel_span,
|
||||
ttl=kwargs.get("ttl", None),
|
||||
refresh_ttl=refresh_ttl,
|
||||
)
|
||||
|
||||
return result
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -824,6 +832,7 @@ class RedisCache(BaseCache):
|
|||
value: float,
|
||||
ttl: Optional[int] = None,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
refresh_ttl: bool = False,
|
||||
) -> float:
|
||||
from redis.asyncio import Redis
|
||||
|
||||
|
|
@ -834,11 +843,12 @@ class RedisCache(BaseCache):
|
|||
try:
|
||||
result = await _redis_client.incrbyfloat(name=key, amount=value)
|
||||
if _used_ttl is not None:
|
||||
# check if key already has ttl, if not -> set ttl
|
||||
current_ttl = await _redis_client.ttl(key)
|
||||
if current_ttl == -1:
|
||||
# Key has no expiration
|
||||
if refresh_ttl:
|
||||
await _redis_client.expire(key, _used_ttl)
|
||||
else:
|
||||
current_ttl = await _redis_client.ttl(key)
|
||||
if current_ttl == -1:
|
||||
await _redis_client.expire(key, _used_ttl)
|
||||
|
||||
## LOGGING ##
|
||||
end_time = time.time()
|
||||
|
|
|
|||
|
|
@ -1425,6 +1425,7 @@ LITELLM_PROXY_ADMIN_NAME = "default_user_id"
|
|||
LITELLM_CLI_SOURCE_IDENTIFIER = "litellm-cli"
|
||||
LITELLM_CLI_SESSION_TOKEN_PREFIX = "litellm-session-token"
|
||||
CLI_SSO_SESSION_CACHE_KEY_PREFIX = "cli_sso_session"
|
||||
CLI_SSO_SESSION_TTL_SECONDS = 600
|
||||
CLI_JWT_TOKEN_NAME = "cli-jwt-token"
|
||||
# Support both CLI_JWT_EXPIRATION_HOURS and LITELLM_CLI_JWT_EXPIRATION_HOURS for backwards compatibility
|
||||
CLI_JWT_EXPIRATION_HOURS = int(
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -23,6 +23,13 @@ def _raise_env_reference_error(param: str, *, source: str) -> None:
|
|||
)
|
||||
|
||||
|
||||
def validate_no_callback_env_reference(
|
||||
param: str, value: object, *, source: str
|
||||
) -> None:
|
||||
if _is_env_reference(value):
|
||||
_raise_env_reference_error(param, source=source)
|
||||
|
||||
|
||||
# Hardcoded list of supported callback params to avoid runtime inspection issues with TypedDict
|
||||
_supported_callback_params = [
|
||||
"langfuse_public_key",
|
||||
|
|
@ -66,8 +73,9 @@ def initialize_standard_callback_dynamic_params(
|
|||
for param in _supported_callback_params:
|
||||
if param in kwargs:
|
||||
_param_value = kwargs.get(param)
|
||||
if _is_env_reference(_param_value):
|
||||
_raise_env_reference_error(param, source="request body")
|
||||
validate_no_callback_env_reference(
|
||||
param, _param_value, source="request body"
|
||||
)
|
||||
standard_callback_dynamic_params[param] = _param_value # type: ignore
|
||||
|
||||
# 2. Fallback: check "metadata" or "litellm_params" -> "metadata"
|
||||
|
|
@ -80,8 +88,9 @@ def initialize_standard_callback_dynamic_params(
|
|||
for param in _supported_callback_params:
|
||||
if param not in standard_callback_dynamic_params and param in metadata:
|
||||
_param_value = metadata.get(param)
|
||||
if _is_env_reference(_param_value):
|
||||
_raise_env_reference_error(param, source="metadata")
|
||||
validate_no_callback_env_reference(
|
||||
param, _param_value, source="metadata"
|
||||
)
|
||||
standard_callback_dynamic_params[param] = _param_value # type: ignore
|
||||
|
||||
return standard_callback_dynamic_params
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -60,6 +60,9 @@ def _redact_choice_content(choice):
|
|||
def _redact_responses_api_output(output_items):
|
||||
"""Helper to redact ResponsesAPIResponse output items."""
|
||||
for output_item in output_items:
|
||||
if hasattr(output_item, "text"):
|
||||
output_item.text = "redacted-by-litellm"
|
||||
|
||||
if hasattr(output_item, "content") and isinstance(output_item.content, list):
|
||||
for content_part in output_item.content:
|
||||
if hasattr(content_part, "text"):
|
||||
|
|
@ -75,6 +78,28 @@ def _redact_responses_api_output(output_items):
|
|||
summary_item.text = "redacted-by-litellm"
|
||||
|
||||
|
||||
def _redact_responses_api_output_dict(output_items, redacted_str: str):
|
||||
"""Helper to redact ResponsesAPIResponse output items in dict form."""
|
||||
for output_item in output_items:
|
||||
if not isinstance(output_item, dict):
|
||||
continue
|
||||
|
||||
if "text" in output_item:
|
||||
output_item["text"] = redacted_str
|
||||
|
||||
if isinstance(output_item.get("content"), list):
|
||||
for content_item in output_item["content"]:
|
||||
if isinstance(content_item, dict) and "text" in content_item:
|
||||
content_item["text"] = redacted_str
|
||||
|
||||
if output_item.get("type") == "reasoning" and isinstance(
|
||||
output_item.get("summary"), list
|
||||
):
|
||||
for summary_item in output_item["summary"]:
|
||||
if isinstance(summary_item, dict) and "text" in summary_item:
|
||||
summary_item["text"] = redacted_str
|
||||
|
||||
|
||||
def _redact_standard_logging_object(model_call_details: dict):
|
||||
"""Redact messages and response inside standard_logging_object if present."""
|
||||
standard_logging_object = model_call_details.get("standard_logging_object")
|
||||
|
|
@ -93,28 +118,11 @@ def _redact_standard_logging_object(model_call_details: dict):
|
|||
if isinstance(response, dict) and "output" in response:
|
||||
# ResponsesAPIResponse format - redact content in output items
|
||||
if isinstance(response.get("output"), list):
|
||||
for output_item in response["output"]:
|
||||
if isinstance(output_item, dict) and "content" in output_item:
|
||||
if isinstance(output_item["content"], list):
|
||||
for content_item in output_item["content"]:
|
||||
if (
|
||||
isinstance(content_item, dict)
|
||||
and "text" in content_item
|
||||
):
|
||||
content_item["text"] = redacted_str
|
||||
_redact_responses_api_output_dict(response["output"], redacted_str)
|
||||
elif isinstance(response, dict) and "choices" in response:
|
||||
# ModelResponse dict format - redact content in choices
|
||||
if isinstance(response.get("choices"), list):
|
||||
for choice in response["choices"]:
|
||||
if isinstance(choice, dict):
|
||||
if "message" in choice and isinstance(choice["message"], dict):
|
||||
choice["message"]["content"] = redacted_str
|
||||
if "audio" in choice["message"]:
|
||||
choice["message"]["audio"] = None
|
||||
elif "delta" in choice and isinstance(choice["delta"], dict):
|
||||
choice["delta"]["content"] = redacted_str
|
||||
if "audio" in choice["delta"]:
|
||||
choice["delta"]["audio"] = None
|
||||
_redact_model_response_dict_choices(response["choices"], redacted_str)
|
||||
elif isinstance(response, str):
|
||||
standard_logging_object["response"] = redacted_str
|
||||
else:
|
||||
|
|
@ -122,6 +130,29 @@ def _redact_standard_logging_object(model_call_details: dict):
|
|||
standard_logging_object["response"] = {"text": redacted_str}
|
||||
|
||||
|
||||
def _redact_model_response_dict_choices(choices, redacted_str: str):
|
||||
for choice in choices:
|
||||
if isinstance(choice, dict):
|
||||
if "message" in choice and isinstance(choice["message"], dict):
|
||||
choice["message"]["content"] = redacted_str
|
||||
if "reasoning_content" in choice["message"]:
|
||||
choice["message"]["reasoning_content"] = redacted_str
|
||||
if "thinking_blocks" in choice["message"]:
|
||||
choice["message"]["thinking_blocks"] = None
|
||||
if "audio" in choice["message"]:
|
||||
choice["message"]["audio"] = None
|
||||
elif "delta" in choice and isinstance(choice["delta"], dict):
|
||||
choice["delta"]["content"] = redacted_str
|
||||
if "reasoning_content" in choice["delta"]:
|
||||
choice["delta"]["reasoning_content"] = redacted_str
|
||||
if "thinking_blocks" in choice["delta"]:
|
||||
choice["delta"]["thinking_blocks"] = None
|
||||
if "audio" in choice["delta"]:
|
||||
choice["delta"]["audio"] = None
|
||||
else:
|
||||
_redact_choice_content(choice)
|
||||
|
||||
|
||||
def perform_redaction(model_call_details: dict, result):
|
||||
"""
|
||||
Performs the actual redaction on the logging object and result.
|
||||
|
|
@ -132,6 +163,7 @@ def perform_redaction(model_call_details: dict, result):
|
|||
]
|
||||
model_call_details["prompt"] = ""
|
||||
model_call_details["input"] = ""
|
||||
_redact_standard_logging_object(model_call_details)
|
||||
|
||||
# Redact streaming response
|
||||
if (
|
||||
|
|
@ -171,30 +203,14 @@ def perform_redaction(model_call_details: dict, result):
|
|||
elif isinstance(_result, dict) and "choices" in _result:
|
||||
# Handle dict representation of ModelResponse (e.g., from model_dump())
|
||||
if _result.get("choices") is not None:
|
||||
for choice in _result["choices"]:
|
||||
if isinstance(choice, dict):
|
||||
if "message" in choice and isinstance(choice["message"], dict):
|
||||
choice["message"]["content"] = "redacted-by-litellm"
|
||||
if "reasoning_content" in choice["message"]:
|
||||
choice["message"][
|
||||
"reasoning_content"
|
||||
] = "redacted-by-litellm"
|
||||
if "thinking_blocks" in choice["message"]:
|
||||
choice["message"]["thinking_blocks"] = None
|
||||
if "audio" in choice["message"]:
|
||||
choice["message"]["audio"] = None
|
||||
elif "delta" in choice and isinstance(choice["delta"], dict):
|
||||
choice["delta"]["content"] = "redacted-by-litellm"
|
||||
if "reasoning_content" in choice["delta"]:
|
||||
choice["delta"][
|
||||
"reasoning_content"
|
||||
] = "redacted-by-litellm"
|
||||
if "thinking_blocks" in choice["delta"]:
|
||||
choice["delta"]["thinking_blocks"] = None
|
||||
if "audio" in choice["delta"]:
|
||||
choice["delta"]["audio"] = None
|
||||
else:
|
||||
_redact_choice_content(choice)
|
||||
_redact_model_response_dict_choices(
|
||||
_result["choices"], "redacted-by-litellm"
|
||||
)
|
||||
elif isinstance(_result, dict) and "output" in _result:
|
||||
if isinstance(_result.get("output"), list):
|
||||
_redact_responses_api_output_dict(
|
||||
_result["output"], "redacted-by-litellm"
|
||||
)
|
||||
elif isinstance(_result, litellm.ResponsesAPIResponse):
|
||||
if hasattr(_result, "output"):
|
||||
_redact_responses_api_output(_result.output)
|
||||
|
|
|
|||
|
|
@ -1553,25 +1553,43 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
)
|
||||
data["output_config"] = output_config
|
||||
|
||||
def _transform_response_for_json_mode(
|
||||
def _resolve_json_mode_non_streaming(
|
||||
self,
|
||||
json_mode: Optional[bool],
|
||||
tool_calls: List[ChatCompletionToolCallChunk],
|
||||
) -> Optional[LitellmMessage]:
|
||||
_message: Optional[LitellmMessage] = None
|
||||
if json_mode is True and len(tool_calls) == 1:
|
||||
# check if tool name is the default tool name
|
||||
json_mode_content_str: Optional[str] = None
|
||||
if (
|
||||
"name" in tool_calls[0]["function"]
|
||||
and tool_calls[0]["function"]["name"] == RESPONSE_FORMAT_TOOL_NAME
|
||||
):
|
||||
json_mode_content_str = tool_calls[0]["function"].get("arguments")
|
||||
if json_mode_content_str is not None:
|
||||
_message = AnthropicConfig._convert_tool_response_to_message(
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
return _message
|
||||
) -> Tuple[
|
||||
Optional[LitellmMessage],
|
||||
List[ChatCompletionToolCallChunk],
|
||||
Optional[str],
|
||||
]:
|
||||
"""Strip internal response_format tool calls; merge payload into content when mixed with user tools."""
|
||||
if json_mode is not True or not tool_calls:
|
||||
return None, tool_calls, None
|
||||
|
||||
json_indices = [
|
||||
i
|
||||
for i, t in enumerate(tool_calls)
|
||||
if t.get("function", {}).get("name") == RESPONSE_FORMAT_TOOL_NAME
|
||||
]
|
||||
if not json_indices:
|
||||
return None, tool_calls, None
|
||||
|
||||
if len(json_indices) == len(tool_calls):
|
||||
json_tool = tool_calls[json_indices[0]]
|
||||
if json_tool.get("function", {}).get("arguments") is None:
|
||||
return None, tool_calls, None
|
||||
_message = AnthropicConfig._convert_tool_response_to_message(
|
||||
tool_calls=[json_tool]
|
||||
)
|
||||
return _message, [], None
|
||||
|
||||
first_json = tool_calls[json_indices[0]]
|
||||
json_msg = AnthropicConfig._convert_tool_response_to_message([first_json])
|
||||
extra_content: Optional[str] = (
|
||||
json_msg.content if json_msg is not None else None
|
||||
)
|
||||
filtered_tools = [t for i, t in enumerate(tool_calls) if i not in json_indices]
|
||||
return None, filtered_tools, extra_content
|
||||
|
||||
def extract_response_content(self, completion_response: dict) -> Tuple[
|
||||
str,
|
||||
|
|
@ -1931,19 +1949,27 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
tool_calls,
|
||||
)
|
||||
|
||||
json_mode_message, tool_calls_for_message, json_extra_content = (
|
||||
self._resolve_json_mode_non_streaming(
|
||||
json_mode=json_mode,
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
)
|
||||
merged_text = text_content or ""
|
||||
if json_extra_content:
|
||||
merged_text = (
|
||||
merged_text + json_extra_content if merged_text else json_extra_content
|
||||
)
|
||||
|
||||
_message = litellm.Message(
|
||||
tool_calls=tool_calls,
|
||||
content=text_content or None,
|
||||
tool_calls=tool_calls_for_message,
|
||||
content=merged_text or None,
|
||||
provider_specific_fields=provider_specific_fields,
|
||||
thinking_blocks=thinking_blocks,
|
||||
reasoning_content=reasoning_content,
|
||||
)
|
||||
_message.provider_specific_fields = provider_specific_fields
|
||||
|
||||
json_mode_message = self._transform_response_for_json_mode(
|
||||
json_mode=json_mode,
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
if json_mode_message is not None:
|
||||
completion_response["stop_reason"] = "stop"
|
||||
_message = json_mode_message
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ class BaseRerankConfig(ABC):
|
|||
model: str,
|
||||
optional_rerank_params: Dict,
|
||||
headers: dict,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> dict:
|
||||
return {}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import urllib.parse
|
||||
from datetime import datetime
|
||||
from typing import (
|
||||
|
|
@ -37,6 +38,11 @@ else:
|
|||
AWSPreparedRequest = Any
|
||||
|
||||
|
||||
# Real AWS region names are lowercase letters, digits, and hyphens
|
||||
# (e.g. "us-east-1", "eu-west-2", "us-gov-west-1", "cn-north-1").
|
||||
_VALID_AWS_REGION_PATTERN = re.compile(r"\A[a-z0-9-]+\Z")
|
||||
|
||||
|
||||
class Boto3CredentialsInfo(BaseModel):
|
||||
credentials: Credentials
|
||||
aws_region_name: str
|
||||
|
|
@ -284,6 +290,9 @@ class BaseAWSLLM:
|
|||
if not region: # Check if region is empty
|
||||
return None
|
||||
|
||||
if not _VALID_AWS_REGION_PATTERN.match(region):
|
||||
return None
|
||||
|
||||
return region
|
||||
except Exception:
|
||||
# Catch any unexpected errors and return None
|
||||
|
|
@ -481,6 +490,7 @@ class BaseAWSLLM:
|
|||
str: The AWS region name
|
||||
"""
|
||||
aws_region_name = optional_params.get("aws_region_name", None)
|
||||
self._validate_aws_region_name(aws_region_name)
|
||||
### SET REGION NAME ###
|
||||
if aws_region_name is None:
|
||||
# check model arn #
|
||||
|
|
@ -519,8 +529,25 @@ class BaseAWSLLM:
|
|||
except Exception:
|
||||
aws_region_name = "us-west-2"
|
||||
|
||||
self._validate_aws_region_name(aws_region_name)
|
||||
return aws_region_name
|
||||
|
||||
@staticmethod
|
||||
def _validate_aws_region_name(aws_region_name: Optional[str]) -> None:
|
||||
"""
|
||||
Validate that an AWS region name conforms to the expected format
|
||||
(lowercase alphanumerics and hyphens). Raises ValueError otherwise.
|
||||
"""
|
||||
if aws_region_name is None:
|
||||
return
|
||||
if not isinstance(aws_region_name, str) or not _VALID_AWS_REGION_PATTERN.match(
|
||||
aws_region_name
|
||||
):
|
||||
raise ValueError(
|
||||
f"Invalid AWS region format: {aws_region_name!r}. "
|
||||
"Region names must contain only lowercase letters, digits, and hyphens."
|
||||
)
|
||||
|
||||
def get_aws_region_name_for_non_llm_api_calls(
|
||||
self,
|
||||
aws_region_name: Optional[str] = None,
|
||||
|
|
@ -532,6 +559,7 @@ class BaseAWSLLM:
|
|||
|
||||
For non-llm api calls eg. Guardrails, Vector Stores we just need to check the dynamic param or env vars.
|
||||
"""
|
||||
self._validate_aws_region_name(aws_region_name)
|
||||
if aws_region_name is None:
|
||||
# check env #
|
||||
litellm_aws_region_name = get_secret("AWS_REGION_NAME", None)
|
||||
|
|
@ -549,6 +577,8 @@ class BaseAWSLLM:
|
|||
|
||||
if aws_region_name is None:
|
||||
aws_region_name = "us-west-2"
|
||||
|
||||
self._validate_aws_region_name(aws_region_name)
|
||||
return aws_region_name
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -111,6 +111,7 @@ class CohereRerankConfig(BaseRerankConfig):
|
|||
model: str,
|
||||
optional_rerank_params: Dict,
|
||||
headers: dict,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> dict:
|
||||
if "query" not in optional_rerank_params:
|
||||
raise ValueError("query is required for Cohere rerank")
|
||||
|
|
|
|||
|
|
@ -71,6 +71,7 @@ class CohereRerankV2Config(CohereRerankConfig):
|
|||
model: str,
|
||||
optional_rerank_params: Dict,
|
||||
headers: dict,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> dict:
|
||||
if "query" not in optional_rerank_params:
|
||||
raise ValueError("query is required for Cohere rerank")
|
||||
|
|
|
|||
|
|
@ -1007,6 +1007,7 @@ class BaseLLMHTTPHandler:
|
|||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
|
||||
litellm_params: Optional[Dict[str, Any]] = None,
|
||||
) -> RerankResponse:
|
||||
# get config from model, custom llm provider
|
||||
headers = provider_config.validate_environment(
|
||||
|
|
@ -1026,6 +1027,7 @@ class BaseLLMHTTPHandler:
|
|||
model=model,
|
||||
optional_rerank_params=optional_rerank_params,
|
||||
headers=headers,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
|
|
@ -2535,10 +2537,16 @@ class BaseLLMHTTPHandler:
|
|||
},
|
||||
)
|
||||
|
||||
delete_kwargs: Dict[str, Any] = {
|
||||
"url": url,
|
||||
"headers": headers,
|
||||
"timeout": timeout,
|
||||
}
|
||||
if data:
|
||||
delete_kwargs["json"] = data
|
||||
|
||||
try:
|
||||
response = await async_httpx_client.delete(
|
||||
url=url, headers=headers, json=data, timeout=timeout
|
||||
)
|
||||
response = await async_httpx_client.delete(**delete_kwargs)
|
||||
|
||||
except Exception as e:
|
||||
raise self._handle_error(
|
||||
|
|
@ -2619,10 +2627,16 @@ class BaseLLMHTTPHandler:
|
|||
},
|
||||
)
|
||||
|
||||
delete_kwargs: Dict[str, Any] = {
|
||||
"url": url,
|
||||
"headers": headers,
|
||||
"timeout": timeout,
|
||||
}
|
||||
if data:
|
||||
delete_kwargs["json"] = data
|
||||
|
||||
try:
|
||||
response = sync_httpx_client.delete(
|
||||
url=url, headers=headers, json=data, timeout=timeout
|
||||
)
|
||||
response = sync_httpx_client.delete(**delete_kwargs)
|
||||
|
||||
except Exception as e:
|
||||
raise self._handle_error(
|
||||
|
|
|
|||
|
|
@ -132,6 +132,7 @@ class DeepinfraRerankConfig(BaseRerankConfig):
|
|||
model: str,
|
||||
optional_rerank_params: Dict,
|
||||
headers: dict,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> dict:
|
||||
# Convert OptionalRerankParams to dict as expected by parent class
|
||||
if optional_rerank_params is None:
|
||||
|
|
|
|||
|
|
@ -127,6 +127,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig):
|
|||
model: str,
|
||||
optional_rerank_params: Dict,
|
||||
headers: dict,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Transform request to Fireworks AI rerank format
|
||||
|
|
|
|||
|
|
@ -121,6 +121,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
|
|||
model: str,
|
||||
optional_rerank_params: Dict,
|
||||
headers: dict,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> dict:
|
||||
if "query" not in optional_rerank_params:
|
||||
raise ValueError("query is required for Hosted VLLM rerank")
|
||||
|
|
|
|||
|
|
@ -146,6 +146,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig):
|
|||
model: str,
|
||||
optional_rerank_params: Union[OptionalRerankParams, dict],
|
||||
headers: dict,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> dict:
|
||||
if "query" not in optional_rerank_params:
|
||||
raise ValueError("query is required for HuggingFace rerank")
|
||||
|
|
|
|||
|
|
@ -74,7 +74,11 @@ class JinaAIRerankConfig(BaseRerankConfig):
|
|||
return cleaned_base
|
||||
|
||||
def transform_rerank_request(
|
||||
self, model: str, optional_rerank_params: Dict, headers: Dict
|
||||
self,
|
||||
model: str,
|
||||
optional_rerank_params: Dict,
|
||||
headers: Dict,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> Dict:
|
||||
return {"model": model, **optional_rerank_params}
|
||||
|
||||
|
|
|
|||
|
|
@ -25,7 +25,6 @@ else:
|
|||
LiteLLMLoggingObj = Any
|
||||
|
||||
MILVUS_OPTIONAL_PARAMS = {
|
||||
"dbName",
|
||||
"annsField",
|
||||
"limit",
|
||||
"filter",
|
||||
|
|
@ -33,7 +32,6 @@ MILVUS_OPTIONAL_PARAMS = {
|
|||
"groupingField",
|
||||
"outputFields",
|
||||
"searchParams",
|
||||
"partitionNames",
|
||||
"consistencyLevel",
|
||||
}
|
||||
|
||||
|
|
@ -173,13 +171,21 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig):
|
|||
url = f"{api_base}/v2/vectordb/entities/search"
|
||||
|
||||
# Build the request body for Azure AI Search with vector search
|
||||
request_body = {
|
||||
request_body: Dict[str, Any] = {
|
||||
"collectionName": index_name,
|
||||
"data": [query_vector],
|
||||
"annsField": "book_intro_vector",
|
||||
**vector_store_search_optional_params,
|
||||
}
|
||||
|
||||
db_name = litellm_params.get("milvus_db_name")
|
||||
if db_name:
|
||||
request_body["dbName"] = db_name
|
||||
|
||||
partition_names = litellm_params.get("milvus_partition_names")
|
||||
if partition_names:
|
||||
request_body["partitionNames"] = partition_names
|
||||
|
||||
#########################################################
|
||||
# Update logging object with details of the request
|
||||
#########################################################
|
||||
|
|
|
|||
|
|
@ -66,6 +66,7 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig):
|
|||
model: str,
|
||||
optional_rerank_params: Dict,
|
||||
headers: dict,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Transform request, using clean model name without 'ranking/' prefix.
|
||||
|
|
@ -75,4 +76,5 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig):
|
|||
model=clean_model,
|
||||
optional_rerank_params=optional_rerank_params,
|
||||
headers=headers,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -177,6 +177,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig):
|
|||
model: str,
|
||||
optional_rerank_params: Dict,
|
||||
headers: dict,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Transform request to Nvidia NIM format.
|
||||
|
|
|
|||
|
|
@ -27,6 +27,53 @@ class VertexAIError(BaseLLMException):
|
|||
super().__init__(message=message, status_code=status_code, headers=headers)
|
||||
|
||||
|
||||
def vertex_request_labels_from_litellm_params(
|
||||
litellm_params: Optional[dict],
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""
|
||||
Build Vertex/GCP billing labels from LiteLLM user metadata on ``litellm_params``:
|
||||
``metadata`` (``completion(..., metadata=...)``) or ``litellm_metadata``,
|
||||
using ``requester_metadata`` string key-value pairs (same convention as Gemini).
|
||||
``metadata`` is tried first when both are present.
|
||||
"""
|
||||
if not litellm_params:
|
||||
return None
|
||||
for key in ("metadata", "litellm_metadata"):
|
||||
if key not in litellm_params:
|
||||
continue
|
||||
metadata = litellm_params[key]
|
||||
if metadata is None or not isinstance(metadata, dict):
|
||||
continue
|
||||
if "requester_metadata" not in metadata:
|
||||
continue
|
||||
rm = metadata["requester_metadata"]
|
||||
if not isinstance(rm, dict):
|
||||
continue
|
||||
labels = {k: v for k, v in rm.items() if isinstance(v, str)}
|
||||
if labels:
|
||||
return labels
|
||||
return None
|
||||
|
||||
|
||||
def pop_vertex_request_labels(
|
||||
optional_params: Optional[dict],
|
||||
litellm_params: Optional[dict],
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""
|
||||
Resolve labels from optional ``labels`` (Gemini-style) and/or
|
||||
``litellm_params["metadata"]`` / ``litellm_params["litellm_metadata"]``
|
||||
(``requester_metadata``). Pops ``labels`` from optional_params when present.
|
||||
"""
|
||||
labels: Optional[Dict[str, str]] = None
|
||||
if optional_params is not None and "labels" in optional_params:
|
||||
raw = optional_params.pop("labels")
|
||||
if isinstance(raw, dict):
|
||||
labels = {k: v for k, v in raw.items() if isinstance(v, str)}
|
||||
if not labels:
|
||||
labels = vertex_request_labels_from_litellm_params(litellm_params)
|
||||
return labels if labels else None
|
||||
|
||||
|
||||
class VertexAIModelRoute(str, Enum):
|
||||
"""Enum for Vertex AI model routing"""
|
||||
|
||||
|
|
@ -50,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", "openai/gpt-oss-120b")
|
||||
model: The model name (e.g., "llama3-405b", "gemini-pro", "gemma/gemma-3-12b-it", "xai/grok-4.1-fast-non-reasoning")
|
||||
litellm_params: Optional litellm parameters dict that may contain base_model for routing
|
||||
|
||||
Returns:
|
||||
|
|
@ -66,7 +113,7 @@ def get_vertex_ai_model_route(
|
|||
>>> get_vertex_ai_model_route("gemma/gemma-3-12b-it")
|
||||
VertexAIModelRoute.GEMMA
|
||||
|
||||
>>> get_vertex_ai_model_route("openai/gpt-oss-120b")
|
||||
>>> 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"})
|
||||
|
|
@ -102,8 +149,11 @@ def get_vertex_ai_model_route(
|
|||
if "gemma/" in model:
|
||||
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
|
||||
|
|
@ -209,8 +259,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"
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
response_schema_prompt,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
|
||||
from litellm.types.files import (
|
||||
get_file_mime_type_for_file_type,
|
||||
get_file_type_from_extension,
|
||||
|
|
@ -714,16 +715,8 @@ def _transform_request_body( # noqa: PLR0915
|
|||
optional_params.pop("output_config", None)
|
||||
config_fields = GenerationConfig.__annotations__.keys()
|
||||
|
||||
# If the LiteLLM client sends Gemini-supported parameter "labels", add it
|
||||
# as "labels" field to the request sent to the Gemini backend.
|
||||
labels: Optional[dict[str, str]] = optional_params.pop("labels", None)
|
||||
# If the LiteLLM client sends OpenAI-supported parameter "metadata", add it
|
||||
# as "labels" field to the request sent to the Gemini backend.
|
||||
if labels is None and "metadata" in litellm_params:
|
||||
metadata = litellm_params["metadata"]
|
||||
if metadata is not None and "requester_metadata" in metadata:
|
||||
rm = metadata["requester_metadata"]
|
||||
labels = {k: v for k, v in rm.items() if isinstance(v, str)}
|
||||
# labels: optional explicit param and/or metadata.requester_metadata (OpenAI metadata)
|
||||
labels = pop_vertex_request_labels(optional_params, litellm_params)
|
||||
|
||||
filtered_params = {
|
||||
k: v
|
||||
|
|
|
|||
|
|
@ -7,7 +7,10 @@ import litellm
|
|||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
)
|
||||
from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
get_vertex_base_url,
|
||||
pop_vertex_request_labels,
|
||||
)
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import (
|
||||
|
|
@ -203,13 +206,16 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
|
|||
"sampleCount": 1,
|
||||
}
|
||||
|
||||
# Merge with optional params
|
||||
labels = pop_vertex_request_labels(optional_params, litellm_params)
|
||||
# Merge with optional params (after popping labels so they are not sent as Imagen parameters)
|
||||
parameters = {**default_params, **optional_params}
|
||||
|
||||
request_body = {
|
||||
request_body: dict = {
|
||||
"instances": [{"prompt": prompt}],
|
||||
"parameters": parameters,
|
||||
}
|
||||
if labels:
|
||||
request_body["labels"] = labels
|
||||
|
||||
return request_body
|
||||
|
||||
|
|
|
|||
|
|
@ -11,12 +11,15 @@ import httpx
|
|||
import litellm
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
vertex_request_labels_from_litellm_params,
|
||||
)
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.rerank import (
|
||||
RerankBilledUnits,
|
||||
RerankResponse,
|
||||
RerankResponseMeta,
|
||||
RerankBilledUnits,
|
||||
RerankResponseResult,
|
||||
)
|
||||
|
||||
|
|
@ -109,6 +112,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
|
|||
model: str,
|
||||
optional_rerank_params: Dict,
|
||||
headers: dict,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Transform the request from Cohere format to Vertex AI Discovery Engine format
|
||||
|
|
@ -145,6 +149,10 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
|
|||
# When return_documents is False, we want to ignore record details (return only IDs)
|
||||
request_data["ignoreRecordDetailsInResponse"] = not return_documents
|
||||
|
||||
user_labels = vertex_request_labels_from_litellm_params(litellm_params)
|
||||
if user_labels:
|
||||
request_data["userLabels"] = user_labels
|
||||
|
||||
return request_data
|
||||
|
||||
def transform_rerank_response(
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from typing import Literal, Optional, Union
|
||||
from typing import Dict, Literal, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -44,6 +44,7 @@ class VertexEmbedding(VertexBase):
|
|||
vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES] = None,
|
||||
gemini_api_key: Optional[str] = None,
|
||||
extra_headers: Optional[dict] = None,
|
||||
litellm_params: Optional[Dict] = None,
|
||||
) -> EmbeddingResponse:
|
||||
if aembedding is True:
|
||||
return self.async_embedding( # type: ignore
|
||||
|
|
@ -61,6 +62,7 @@ class VertexEmbedding(VertexBase):
|
|||
vertex_credentials=vertex_credentials,
|
||||
gemini_api_key=gemini_api_key,
|
||||
extra_headers=extra_headers,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
should_use_v1beta1_features = self.is_using_v1beta1_features(
|
||||
|
|
@ -92,7 +94,10 @@ class VertexEmbedding(VertexBase):
|
|||
headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers)
|
||||
vertex_request: VertexEmbeddingRequest = (
|
||||
litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request(
|
||||
input=input, optional_params=optional_params, model=model
|
||||
input=input,
|
||||
optional_params=optional_params,
|
||||
model=model,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -156,6 +161,7 @@ class VertexEmbedding(VertexBase):
|
|||
gemini_api_key: Optional[str] = None,
|
||||
extra_headers: Optional[dict] = None,
|
||||
encoding=None,
|
||||
litellm_params: Optional[Dict] = None,
|
||||
) -> EmbeddingResponse:
|
||||
"""
|
||||
Async embedding implementation
|
||||
|
|
@ -188,7 +194,10 @@ class VertexEmbedding(VertexBase):
|
|||
headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers)
|
||||
vertex_request: VertexEmbeddingRequest = (
|
||||
litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request(
|
||||
input=input, optional_params=optional_params, model=model
|
||||
input=input,
|
||||
optional_params=optional_params,
|
||||
model=model,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from typing import List, Literal, Optional, Union
|
|||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
|
||||
from litellm.types.utils import EmbeddingResponse, Usage
|
||||
|
||||
from .types import *
|
||||
|
|
@ -100,7 +101,11 @@ class VertexAITextEmbeddingConfig(BaseModel):
|
|||
return optional_params
|
||||
|
||||
def transform_openai_request_to_vertex_embedding_request(
|
||||
self, input: Union[list, str], optional_params: dict, model: str
|
||||
self,
|
||||
input: Union[list, str],
|
||||
optional_params: dict,
|
||||
model: str,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> VertexEmbeddingRequest:
|
||||
"""
|
||||
Transforms an openai request to a vertex embedding request.
|
||||
|
|
@ -108,16 +113,26 @@ class VertexAITextEmbeddingConfig(BaseModel):
|
|||
# Import here to avoid circular import issues with litellm.__init__
|
||||
from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig
|
||||
|
||||
labels = pop_vertex_request_labels(optional_params, litellm_params)
|
||||
|
||||
if model.isdigit():
|
||||
return self._transform_openai_request_to_fine_tuned_embedding_request(
|
||||
input, optional_params, model
|
||||
vertex_request = (
|
||||
self._transform_openai_request_to_fine_tuned_embedding_request(
|
||||
input, optional_params, model
|
||||
)
|
||||
)
|
||||
if labels:
|
||||
vertex_request["labels"] = labels
|
||||
return vertex_request
|
||||
if VertexBGEConfig.is_bge_model(model):
|
||||
return VertexBGEConfig.transform_request(
|
||||
vertex_request = VertexBGEConfig.transform_request(
|
||||
input=input, optional_params=optional_params, model=model
|
||||
)
|
||||
if labels:
|
||||
vertex_request["labels"] = labels
|
||||
return vertex_request
|
||||
|
||||
vertex_request: VertexEmbeddingRequest = VertexEmbeddingRequest()
|
||||
vertex_request = VertexEmbeddingRequest()
|
||||
vertex_text_embedding_input_list: List[TextEmbeddingInput] = []
|
||||
task_type: Optional[TaskType] = optional_params.get("task_type")
|
||||
title = optional_params.get("title")
|
||||
|
|
@ -133,6 +148,8 @@ class VertexAITextEmbeddingConfig(BaseModel):
|
|||
|
||||
vertex_request["instances"] = vertex_text_embedding_input_list
|
||||
vertex_request["parameters"] = EmbeddingParameters(**optional_params)
|
||||
if labels:
|
||||
vertex_request["labels"] = labels
|
||||
|
||||
return vertex_request
|
||||
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ Types for Vertex Embeddings Requests
|
|||
"""
|
||||
|
||||
from enum import Enum
|
||||
from typing import List, Optional, Union
|
||||
from typing import Dict, List, Optional, Union
|
||||
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
|
|
@ -56,6 +56,7 @@ class VertexEmbeddingRequest(TypedDict, total=False):
|
|||
List[TextEmbeddingFineTunedInput],
|
||||
]
|
||||
parameters: Optional[Union[EmbeddingParameters, TextEmbeddingFineTunedParameters]]
|
||||
labels: Optional[Dict[str, str]]
|
||||
|
||||
|
||||
# Example usage:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -67,7 +67,11 @@ class VoyageRerankConfig(BaseRerankConfig):
|
|||
return api_base
|
||||
|
||||
def transform_rerank_request(
|
||||
self, model: str, optional_rerank_params: Dict, headers: Dict
|
||||
self,
|
||||
model: str,
|
||||
optional_rerank_params: Dict,
|
||||
headers: Dict,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> Dict:
|
||||
return {"model": model, **optional_rerank_params}
|
||||
|
||||
|
|
|
|||
|
|
@ -143,6 +143,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig):
|
|||
model: str,
|
||||
optional_rerank_params: Dict,
|
||||
headers: dict,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Transform request to IBM watsonx.ai rerank format
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ class XAIChatConfig(OpenAIGPTConfig):
|
|||
"logprobs",
|
||||
"max_tokens",
|
||||
"n",
|
||||
"parallel_tool_calls",
|
||||
"presence_penalty",
|
||||
"response_format",
|
||||
"seed",
|
||||
|
|
|
|||
|
|
@ -5311,6 +5311,7 @@ def embedding( # noqa: PLR0915
|
|||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
client=client,
|
||||
litellm_params=litellm_params_dict,
|
||||
)
|
||||
elif custom_llm_provider == "oobabooga":
|
||||
response = oobabooga.embedding(
|
||||
|
|
|
|||
|
|
@ -33337,6 +33337,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",
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import base64
|
||||
import binascii
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast
|
||||
|
|
@ -498,6 +499,82 @@ async def rotate_mcp_server_credentials_master_key(
|
|||
)
|
||||
|
||||
|
||||
def _decode_user_credential(stored: str) -> Optional[str]:
|
||||
"""Read back a value persisted in ``LiteLLM_MCPUserCredentials.credential_b64``.
|
||||
|
||||
Tries nacl decryption first (current write format). Falls back to a
|
||||
plain ``urlsafe_b64decode`` for rows persisted by older code that wrote
|
||||
the credential without encryption. Returns ``None`` when neither path
|
||||
yields a valid string.
|
||||
"""
|
||||
decrypted = decrypt_value_helper(
|
||||
value=stored,
|
||||
key="mcp_user_credential",
|
||||
exception_type="debug",
|
||||
return_original_value=False,
|
||||
)
|
||||
if decrypted is not None:
|
||||
return decrypted
|
||||
try:
|
||||
return base64.urlsafe_b64decode(stored).decode()
|
||||
except (binascii.Error, UnicodeDecodeError, ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _decode_oauth_payload(stored: str) -> Optional[Dict[str, Any]]:
|
||||
"""Return the OAuth2 payload dict if ``stored`` holds one, else ``None``.
|
||||
|
||||
A row is considered an OAuth2 credential iff its decoded value parses as
|
||||
a JSON object with ``"type": "oauth2"``. Plain BYOK credentials (which
|
||||
share the same column) decode to a non-JSON string and return ``None``.
|
||||
"""
|
||||
decoded = _decode_user_credential(stored)
|
||||
if decoded is None:
|
||||
return None
|
||||
try:
|
||||
parsed = json.loads(decoded)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
if isinstance(parsed, dict) and parsed.get("type") == "oauth2":
|
||||
return parsed
|
||||
return None
|
||||
|
||||
|
||||
async def rotate_mcp_user_credentials_master_key(
|
||||
prisma_client: PrismaClient, new_master_key: str
|
||||
):
|
||||
"""Re-encrypt every ``LiteLLM_MCPUserCredentials`` row with ``new_master_key``.
|
||||
|
||||
Reads each ``credential_b64`` with the current salt key (falling back to
|
||||
legacy plain base64 for unmigrated rows) and writes it back encrypted
|
||||
under the new master key. Rows that are unreadable under both paths
|
||||
are logged and skipped so one corrupt row does not abort the rotation.
|
||||
"""
|
||||
rows = await prisma_client.db.litellm_mcpusercredentials.find_many()
|
||||
for row in rows:
|
||||
plaintext = _decode_user_credential(row.credential_b64)
|
||||
if plaintext is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"rotate_mcp_user_credentials_master_key: could not decode "
|
||||
"credential for user_id=%s server_id=%s, skipping",
|
||||
row.user_id,
|
||||
row.server_id,
|
||||
)
|
||||
continue
|
||||
re_encrypted = encrypt_value_helper(
|
||||
plaintext, new_encryption_key=new_master_key
|
||||
)
|
||||
await prisma_client.db.litellm_mcpusercredentials.update(
|
||||
where={
|
||||
"user_id_server_id": {
|
||||
"user_id": row.user_id,
|
||||
"server_id": row.server_id,
|
||||
}
|
||||
},
|
||||
data={"credential_b64": re_encrypted},
|
||||
)
|
||||
|
||||
|
||||
async def store_user_credential(
|
||||
prisma_client: PrismaClient,
|
||||
user_id: str,
|
||||
|
|
@ -506,7 +583,7 @@ async def store_user_credential(
|
|||
) -> None:
|
||||
"""Store a user credential for a BYOK MCP server."""
|
||||
|
||||
encoded = base64.urlsafe_b64encode(credential.encode()).decode()
|
||||
encoded = encrypt_value_helper(credential)
|
||||
await prisma_client.db.litellm_mcpusercredentials.upsert(
|
||||
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}},
|
||||
data={
|
||||
|
|
@ -532,16 +609,7 @@ async def get_user_credential(
|
|||
)
|
||||
if row is None:
|
||||
return None
|
||||
try:
|
||||
return base64.urlsafe_b64decode(row.credential_b64).decode()
|
||||
except Exception:
|
||||
# Fall back to nacl decryption for credentials stored by older code
|
||||
return decrypt_value_helper(
|
||||
value=row.credential_b64,
|
||||
key="byok_credential",
|
||||
exception_type="debug",
|
||||
return_original_value=False,
|
||||
)
|
||||
return _decode_user_credential(row.credential_b64)
|
||||
|
||||
|
||||
async def has_user_credential(
|
||||
|
|
@ -582,7 +650,7 @@ async def store_user_oauth_credential(
|
|||
) -> None:
|
||||
"""Persist an OAuth2 access token for a user+server pair.
|
||||
|
||||
The payload is JSON-serialised and stored base64-encoded in the same
|
||||
The payload is JSON-serialised and stored encrypted in the same
|
||||
``credential_b64`` column used by BYOK. A ``"type": "oauth2"`` key
|
||||
differentiates it from plain BYOK API keys.
|
||||
"""
|
||||
|
|
@ -606,29 +674,27 @@ async def store_user_oauth_credential(
|
|||
payload["scopes"] = scopes
|
||||
|
||||
# Guard against silently overwriting a BYOK credential with an OAuth token.
|
||||
# BYOK credentials lack a "type" field (or use a non-"oauth2" type).
|
||||
# Skip the guard when the caller knows the row is already an OAuth2 credential
|
||||
# (e.g. during token refresh), saving an extra DB round-trip.
|
||||
if not skip_byok_guard:
|
||||
existing = await prisma_client.db.litellm_mcpusercredentials.find_unique(
|
||||
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}
|
||||
)
|
||||
if existing is not None:
|
||||
_byok_error = ValueError(
|
||||
f"A non-OAuth2 credential already exists for user {user_id} "
|
||||
f"and server {server_id}. Refusing to overwrite."
|
||||
if (
|
||||
existing is not None
|
||||
and _decode_oauth_payload(existing.credential_b64) is None
|
||||
):
|
||||
# Existing row is either a BYOK secret or an OAuth2 row that no
|
||||
# longer decrypts (e.g. after a salt-key rotation). In either
|
||||
# case, refuse to overwrite — the caller would clobber data
|
||||
# that may still be recoverable.
|
||||
raise ValueError(
|
||||
f"Existing credential for user {user_id} and server "
|
||||
f"{server_id} could not be verified as an OAuth2 token. "
|
||||
f"Refusing to overwrite."
|
||||
)
|
||||
try:
|
||||
raw = json.loads(
|
||||
base64.urlsafe_b64decode(existing.credential_b64).decode()
|
||||
)
|
||||
except Exception:
|
||||
# Credential is not base64+JSON — it's a plain-text BYOK key.
|
||||
raise _byok_error
|
||||
if raw.get("type") != "oauth2":
|
||||
raise _byok_error
|
||||
|
||||
encoded = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode()
|
||||
encoded = encrypt_value_helper(json.dumps(payload))
|
||||
await prisma_client.db.litellm_mcpusercredentials.upsert(
|
||||
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}},
|
||||
data={
|
||||
|
|
@ -672,15 +738,7 @@ async def get_user_oauth_credential(
|
|||
)
|
||||
if row is None:
|
||||
return None
|
||||
try:
|
||||
decoded = base64.urlsafe_b64decode(row.credential_b64).decode()
|
||||
parsed = json.loads(decoded)
|
||||
if isinstance(parsed, dict) and parsed.get("type") == "oauth2":
|
||||
return parsed
|
||||
# Row exists but is a BYOK (plain string), not an OAuth token
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
return _decode_oauth_payload(row.credential_b64)
|
||||
|
||||
|
||||
async def list_user_oauth_credentials(
|
||||
|
|
@ -694,14 +752,11 @@ async def list_user_oauth_credentials(
|
|||
)
|
||||
results: List[Dict[str, Any]] = []
|
||||
for row in rows:
|
||||
try:
|
||||
decoded = base64.urlsafe_b64decode(row.credential_b64).decode()
|
||||
parsed = json.loads(decoded)
|
||||
if isinstance(parsed, dict) and parsed.get("type") == "oauth2":
|
||||
parsed["server_id"] = row.server_id
|
||||
results.append(parsed)
|
||||
except Exception:
|
||||
pass # Skip non-OAuth rows (BYOK plain strings)
|
||||
payload = _decode_oauth_payload(row.credential_b64)
|
||||
if payload is None:
|
||||
continue
|
||||
payload["server_id"] = row.server_id
|
||||
results.append(payload)
|
||||
return results
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -131,6 +131,22 @@ def decode_state_hash(encrypted_state: str) -> dict:
|
|||
return state_data
|
||||
|
||||
|
||||
def _get_validated_client_redirect_uri(state_data: Dict[str, Any]) -> str:
|
||||
"""Return a loopback client redirect URI from OAuth state."""
|
||||
redirect_uri = state_data.get("client_redirect_uri") or state_data.get("base_url")
|
||||
if not redirect_uri or not isinstance(redirect_uri, str):
|
||||
raise HTTPException(status_code=400, detail="Invalid redirect URI")
|
||||
validate_loopback_redirect_uri(redirect_uri)
|
||||
return redirect_uri
|
||||
|
||||
|
||||
def _append_query_params(url: str, params: Dict[str, str]) -> str:
|
||||
parsed = urlparse(url)
|
||||
query_params = parse_qsl(parsed.query, keep_blank_values=True)
|
||||
query_params.extend(params.items())
|
||||
return urlunparse(parsed._replace(query=urlencode(query_params)))
|
||||
|
||||
|
||||
def _resolve_oauth2_server_for_root_endpoints(
|
||||
client_ip: Optional[str] = None,
|
||||
) -> Optional[MCPServer]:
|
||||
|
|
@ -568,7 +584,7 @@ async def authorize(
|
|||
else None
|
||||
)
|
||||
if mcp_server is None and mcp_server_name is None:
|
||||
mcp_server = _resolve_oauth2_server_for_root_endpoints()
|
||||
mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
|
||||
if mcp_server is None:
|
||||
raise HTTPException(status_code=404, detail="MCP server not found")
|
||||
# Use server's stored client_id when caller doesn't supply one.
|
||||
|
|
@ -630,7 +646,7 @@ async def token_endpoint(
|
|||
lookup_name, client_ip=client_ip
|
||||
)
|
||||
if mcp_server is None and mcp_server_name is None:
|
||||
mcp_server = _resolve_oauth2_server_for_root_endpoints()
|
||||
mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
|
||||
if mcp_server is None:
|
||||
raise HTTPException(status_code=404, detail="MCP server not found")
|
||||
return await exchange_token_with_server(
|
||||
|
|
@ -651,7 +667,6 @@ async def token_endpoint(
|
|||
async def callback(code: str, state: str):
|
||||
try:
|
||||
state_data = decode_state_hash(state)
|
||||
base_url = state_data["base_url"]
|
||||
original_state = state_data["original_state"]
|
||||
|
||||
# Re-validate loopback at the sink. /authorize rejects non-loopback
|
||||
|
|
@ -659,10 +674,10 @@ async def callback(code: str, state: str):
|
|||
# minted before that check was added have no expiry and remain
|
||||
# valid indefinitely. Validating here blocks the open-redirect +
|
||||
# code-theft primitive even for pre-fix states.
|
||||
validate_loopback_redirect_uri(base_url)
|
||||
redirect_uri = _get_validated_client_redirect_uri(state_data)
|
||||
|
||||
params = {"code": code, "state": original_state}
|
||||
complete_returned_url = f"{base_url}?{urlencode(params)}"
|
||||
complete_returned_url = _append_query_params(redirect_uri, params)
|
||||
return RedirectResponse(url=complete_returned_url, status_code=302)
|
||||
|
||||
except HTTPException:
|
||||
|
|
@ -719,16 +734,16 @@ def _build_oauth_protected_resource_response(
|
|||
)
|
||||
|
||||
request_base_url = get_request_base_url(request)
|
||||
client_ip = IPAddressUtils.get_mcp_client_ip(request)
|
||||
|
||||
# When no server name provided, try to resolve the single OAuth2 server
|
||||
if mcp_server_name is None:
|
||||
resolved = _resolve_oauth2_server_for_root_endpoints()
|
||||
resolved = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
|
||||
if resolved:
|
||||
mcp_server_name = resolved.server_name or resolved.name
|
||||
|
||||
mcp_server: Optional[MCPServer] = None
|
||||
if mcp_server_name:
|
||||
client_ip = IPAddressUtils.get_mcp_client_ip(request)
|
||||
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(
|
||||
mcp_server_name, client_ip=client_ip
|
||||
)
|
||||
|
|
@ -835,10 +850,11 @@ def _build_oauth_authorization_server_response(
|
|||
)
|
||||
|
||||
request_base_url = get_request_base_url(request)
|
||||
client_ip = IPAddressUtils.get_mcp_client_ip(request)
|
||||
|
||||
# When no server name provided, try to resolve the single OAuth2 server
|
||||
if mcp_server_name is None:
|
||||
resolved = _resolve_oauth2_server_for_root_endpoints()
|
||||
resolved = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
|
||||
if resolved:
|
||||
mcp_server_name = resolved.server_name or resolved.name
|
||||
|
||||
|
|
@ -855,7 +871,6 @@ def _build_oauth_authorization_server_response(
|
|||
|
||||
mcp_server: Optional[MCPServer] = None
|
||||
if mcp_server_name:
|
||||
client_ip = IPAddressUtils.get_mcp_client_ip(request)
|
||||
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(
|
||||
mcp_server_name, client_ip=client_ip
|
||||
)
|
||||
|
|
@ -1007,8 +1022,9 @@ async def register_client(request: Request, mcp_server_name: Optional[str] = Non
|
|||
"client_secret": "dummy",
|
||||
"redirect_uris": [f"{request_base_url}/callback"],
|
||||
}
|
||||
client_ip = IPAddressUtils.get_mcp_client_ip(request)
|
||||
if not mcp_server_name:
|
||||
resolved = _resolve_oauth2_server_for_root_endpoints()
|
||||
resolved = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
|
||||
if resolved:
|
||||
return await register_client_with_server(
|
||||
request=request,
|
||||
|
|
@ -1021,7 +1037,6 @@ async def register_client(request: Request, mcp_server_name: Optional[str] = Non
|
|||
)
|
||||
return dummy_return
|
||||
|
||||
client_ip = IPAddressUtils.get_mcp_client_ip(request)
|
||||
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(
|
||||
mcp_server_name, client_ip=client_ip
|
||||
)
|
||||
|
|
|
|||
|
|
@ -41,6 +41,7 @@ from litellm.constants import (
|
|||
MCP_TOOL_LISTING_TIMEOUT,
|
||||
)
|
||||
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get
|
||||
from litellm.experimental_mcp_client.client import MCPClient, MCPSigV4Auth
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
|
||||
|
|
@ -168,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] = {}
|
||||
|
|
@ -341,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,
|
||||
|
|
@ -678,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),
|
||||
|
|
@ -1499,6 +1548,47 @@ class MCPServerManager:
|
|||
)
|
||||
return await client.get_prompt(get_prompt_request_params)
|
||||
|
||||
@staticmethod
|
||||
def _is_same_authority_metadata_url(url: str, server_url: str) -> bool:
|
||||
"""
|
||||
Whether ``url`` shares scheme, host, and port with ``server_url``.
|
||||
|
||||
Same-authority metadata URLs are produced by our well-known discovery
|
||||
construction and by resource servers that publish protected-resource
|
||||
metadata on the resource origin. These must keep working for
|
||||
administrator-configured internal MCP servers, so they are fetched
|
||||
directly. Cross-origin URLs are fetched through ``async_safe_get``.
|
||||
"""
|
||||
try:
|
||||
target = urlparse(url)
|
||||
base = urlparse(server_url)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
if target.scheme not in ("http", "https") or not target.hostname:
|
||||
return False
|
||||
|
||||
target_port = target.port or (443 if target.scheme == "https" else 80)
|
||||
base_port = base.port or (443 if base.scheme == "https" else 80)
|
||||
return (
|
||||
base.scheme == target.scheme
|
||||
and (base.hostname or "").lower() == target.hostname.lower()
|
||||
and base_port == target_port
|
||||
)
|
||||
|
||||
async def _fetch_oauth_discovery_url(self, url: str, server_url: str) -> Any:
|
||||
client = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.MCP,
|
||||
params={"timeout": MCP_METADATA_TIMEOUT},
|
||||
)
|
||||
if self._is_same_authority_metadata_url(url, server_url):
|
||||
# Same-authority URLs may point at administrator-configured
|
||||
# internal MCP servers. Do not run them through user URL
|
||||
# validation, but also do not follow redirects because the
|
||||
# redirect target would not inherit the same-authority guarantee.
|
||||
return await client.get(url, follow_redirects=False)
|
||||
return await async_safe_get(client, url)
|
||||
|
||||
async def _descovery_metadata(
|
||||
self,
|
||||
server_url: str,
|
||||
|
|
@ -1514,7 +1604,7 @@ class MCPServerManager:
|
|||
resource_scopes,
|
||||
) = await self._attempt_well_known_discovery(server_url)
|
||||
metadata = await self._fetch_authorization_server_metadata(
|
||||
authorization_servers
|
||||
authorization_servers, server_url
|
||||
)
|
||||
if (
|
||||
metadata is None
|
||||
|
|
@ -1555,7 +1645,7 @@ class MCPServerManager:
|
|||
authorization_servers,
|
||||
resource_scopes,
|
||||
) = await self._fetch_oauth_metadata_from_resource(
|
||||
resource_metadata_url
|
||||
resource_metadata_url, server_url
|
||||
)
|
||||
else:
|
||||
(
|
||||
|
|
@ -1576,7 +1666,7 @@ class MCPServerManager:
|
|||
|
||||
if authorization_servers:
|
||||
metadata = await self._fetch_authorization_server_metadata(
|
||||
authorization_servers
|
||||
authorization_servers, server_url
|
||||
)
|
||||
|
||||
preferred_scopes = scopes or resource_scopes
|
||||
|
|
@ -1616,19 +1706,26 @@ class MCPServerManager:
|
|||
return resource_metadata_url, scopes
|
||||
|
||||
async def _fetch_oauth_metadata_from_resource(
|
||||
self, resource_metadata_url: str
|
||||
self, resource_metadata_url: str, server_url: str
|
||||
) -> Tuple[List[str], Optional[List[str]]]:
|
||||
if not resource_metadata_url:
|
||||
return [], None
|
||||
|
||||
try:
|
||||
client = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.MCP,
|
||||
params={"timeout": MCP_METADATA_TIMEOUT},
|
||||
response = await self._fetch_oauth_discovery_url(
|
||||
resource_metadata_url, server_url
|
||||
)
|
||||
response = await client.get(resource_metadata_url)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
except SSRFError as exc:
|
||||
verbose_logger.warning(
|
||||
"MCP OAuth discovery: refusing to fetch resource metadata from %s "
|
||||
"(rejected by SSRF guard for server %s): %s",
|
||||
resource_metadata_url,
|
||||
server_url,
|
||||
exc,
|
||||
)
|
||||
return [], None
|
||||
except Exception as exc: # pragma: no cover - network issues
|
||||
verbose_logger.debug(
|
||||
"Failed to fetch MCP OAuth metadata from %s: %s",
|
||||
|
|
@ -1677,23 +1774,25 @@ class MCPServerManager:
|
|||
(
|
||||
authorization_servers,
|
||||
scopes,
|
||||
) = await self._fetch_oauth_metadata_from_resource(url)
|
||||
) = await self._fetch_oauth_metadata_from_resource(url, server_url)
|
||||
if authorization_servers:
|
||||
return authorization_servers, scopes
|
||||
|
||||
return [], None
|
||||
|
||||
async def _fetch_authorization_server_metadata(
|
||||
self, authorization_servers: List[str]
|
||||
self, authorization_servers: List[str], server_url: str
|
||||
) -> Optional[MCPOAuthMetadata]:
|
||||
for issuer in authorization_servers:
|
||||
metadata = await self._fetch_single_authorization_server_metadata(issuer)
|
||||
metadata = await self._fetch_single_authorization_server_metadata(
|
||||
issuer, server_url
|
||||
)
|
||||
if metadata is not None:
|
||||
return metadata
|
||||
return None
|
||||
|
||||
async def _fetch_single_authorization_server_metadata(
|
||||
self, issuer_url: str
|
||||
self, issuer_url: str, server_url: str
|
||||
) -> Optional[MCPOAuthMetadata]:
|
||||
try:
|
||||
parsed = urlparse(issuer_url)
|
||||
|
|
@ -1721,13 +1820,18 @@ class MCPServerManager:
|
|||
|
||||
for url in candidate_urls:
|
||||
try:
|
||||
client = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.MCP,
|
||||
params={"timeout": MCP_METADATA_TIMEOUT},
|
||||
)
|
||||
response = await client.get(url)
|
||||
response = await self._fetch_oauth_discovery_url(url, server_url)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
except SSRFError as exc:
|
||||
verbose_logger.warning(
|
||||
"MCP OAuth discovery: refusing to fetch authorization-server "
|
||||
"metadata from %s (rejected by SSRF guard for server %s): %s",
|
||||
url,
|
||||
server_url,
|
||||
exc,
|
||||
)
|
||||
continue
|
||||
except Exception as exc: # pragma: no cover - network issues
|
||||
verbose_logger.debug(
|
||||
"Failed to fetch authorization metadata from %s: %s",
|
||||
|
|
@ -2370,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,
|
||||
|
|
@ -2433,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:
|
||||
|
|
@ -2445,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
|
||||
|
|
@ -2480,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"
|
||||
|
|
|
|||
432
litellm/proxy/_lazy_features.py
Normal file
432
litellm/proxy/_lazy_features.py
Normal file
|
|
@ -0,0 +1,432 @@
|
|||
"""
|
||||
Lazy registration for optional feature routers. Each LAZY_FEATURES entry
|
||||
imports its module only on the first request matching its path prefix,
|
||||
saving ~700 MB at idle for deployments that don't use these features.
|
||||
First hit pays the import cost (1-3 s for heavy modules); /openapi.json
|
||||
omits each feature's routes until the feature is warmed.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import importlib
|
||||
import sys
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Callable, Dict, Tuple
|
||||
|
||||
from starlette.types import Receive, Scope, Send
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi import APIRouter, FastAPI
|
||||
|
||||
|
||||
def _include_router(attr_name: str = "router") -> Callable[["FastAPI", object], None]:
|
||||
def _register(app: "FastAPI", module: object) -> None:
|
||||
app.include_router(getattr(module, attr_name))
|
||||
|
||||
return _register
|
||||
|
||||
|
||||
def _mount_app(
|
||||
prefix: str, attr_name: str = "app"
|
||||
) -> Callable[["FastAPI", object], None]:
|
||||
def _register(app: "FastAPI", module: object) -> None:
|
||||
app.mount(path=prefix, app=getattr(module, attr_name))
|
||||
|
||||
return _register
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LazyFeature:
|
||||
name: str
|
||||
module_path: str
|
||||
path_prefixes: Tuple[str, ...]
|
||||
register_fn: Callable[["FastAPI", object], None] = field(
|
||||
default_factory=lambda: _include_router("router")
|
||||
)
|
||||
# For routes whose path has a leading parameter (e.g. /{server}/authorize)
|
||||
# — startswith can't match those, so the matcher also checks endswith.
|
||||
path_suffixes: Tuple[str, ...] = ()
|
||||
# Keep the stub injected even after load — for mounted ASGI sub-apps
|
||||
# whose routes don't appear in the parent app's openapi spec.
|
||||
persistent_swagger_stub: bool = False
|
||||
|
||||
|
||||
LAZY_FEATURES: Tuple[LazyFeature, ...] = (
|
||||
LazyFeature(
|
||||
name="guardrails",
|
||||
module_path="litellm.proxy.guardrails.guardrail_endpoints",
|
||||
path_prefixes=(
|
||||
"/guardrails",
|
||||
"/v2/guardrails",
|
||||
"/apply_guardrail",
|
||||
"/policies/usage",
|
||||
),
|
||||
),
|
||||
LazyFeature(
|
||||
name="policies",
|
||||
module_path="litellm.proxy.management_endpoints.policy_endpoints",
|
||||
# Trailing slash to avoid matching /policies/... (policy_engine).
|
||||
path_prefixes=("/policy/", "/utils/test_policies_and_guardrails"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="policy_engine",
|
||||
module_path="litellm.proxy.policy_engine.policy_endpoints",
|
||||
path_prefixes=("/policies",),
|
||||
),
|
||||
LazyFeature(
|
||||
name="policy_resolve",
|
||||
module_path="litellm.proxy.policy_engine.policy_resolve_endpoints",
|
||||
path_prefixes=("/policies/resolve", "/policies/attachments/estimate-impact"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="agents",
|
||||
module_path="litellm.proxy.agent_endpoints.endpoints",
|
||||
path_prefixes=("/v1/agents", "/agents", "/agent/"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="a2a",
|
||||
module_path="litellm.proxy.agent_endpoints.a2a_endpoints",
|
||||
path_prefixes=("/a2a", "/v1/a2a"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="vector_stores",
|
||||
module_path="litellm.proxy.vector_store_endpoints.endpoints",
|
||||
path_prefixes=("/v1/vector_stores", "/vector_stores", "/v1/indexes"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="vector_store_management",
|
||||
module_path="litellm.proxy.vector_store_endpoints.management_endpoints",
|
||||
# Trailing slash to avoid matching /vector_stores/... (vector_stores).
|
||||
path_prefixes=("/vector_store/", "/v1/vector_store/"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="vector_store_files",
|
||||
# Routes appear under both /v1/vector_stores/{id}/files and the
|
||||
# un-versioned form, so both prefixes must trigger the load.
|
||||
module_path="litellm.proxy.vector_store_files_endpoints.endpoints",
|
||||
path_prefixes=("/v1/vector_stores", "/vector_stores"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="tools",
|
||||
module_path="litellm.proxy.management_endpoints.tool_management_endpoints",
|
||||
path_prefixes=("/v1/tool", "/tool"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="search_tools",
|
||||
module_path="litellm.proxy.search_endpoints.search_tool_management",
|
||||
path_prefixes=("/search_tools",),
|
||||
),
|
||||
# mcp_management owns most /v1/mcp/* admin routes; mcp_app is the mounted
|
||||
# streaming sub-app at /mcp.
|
||||
LazyFeature(
|
||||
name="mcp_management",
|
||||
module_path="litellm.proxy.management_endpoints.mcp_management_endpoints",
|
||||
path_prefixes=("/v1/mcp/",),
|
||||
),
|
||||
LazyFeature(
|
||||
# Also serves /.well-known/oauth-* (OAuth metadata discovery).
|
||||
# No /mcp/oauth prefix here: the mounted /mcp sub-app would
|
||||
# shadow it, and there are no actual routes there anyway.
|
||||
name="mcp_byok_oauth",
|
||||
module_path="litellm.proxy._experimental.mcp_server.byok_oauth_endpoints",
|
||||
path_prefixes=("/v1/mcp/oauth", "/.well-known/oauth-"),
|
||||
),
|
||||
LazyFeature(
|
||||
# Serves OAuth dance endpoints (/authorize, /token, /callback,
|
||||
# /register) plus several /.well-known/ discovery URLs at the proxy
|
||||
# root — needed for MCP-over-OAuth flows even before /mcp is hit.
|
||||
name="mcp_discoverable",
|
||||
module_path="litellm.proxy._experimental.mcp_server.discoverable_endpoints",
|
||||
path_prefixes=(
|
||||
"/.well-known/oauth-",
|
||||
"/.well-known/openid-configuration",
|
||||
"/.well-known/jwks.json",
|
||||
"/authorize",
|
||||
"/token",
|
||||
"/callback",
|
||||
"/register",
|
||||
),
|
||||
# Catches the /{mcp_server_name}/authorize|token|register variants.
|
||||
path_suffixes=("/authorize", "/token", "/register"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="mcp_rest",
|
||||
module_path="litellm.proxy._experimental.mcp_server.rest_endpoints",
|
||||
path_prefixes=("/mcp-rest",),
|
||||
),
|
||||
LazyFeature(
|
||||
# Hardcoded /mcp matches BASE_MCP_ROUTE; importing the constant
|
||||
# here would defeat lazy loading.
|
||||
name="mcp_app",
|
||||
module_path="litellm.proxy._experimental.mcp_server.server",
|
||||
path_prefixes=("/mcp",),
|
||||
register_fn=_mount_app("/mcp", attr_name="app"),
|
||||
persistent_swagger_stub=True,
|
||||
),
|
||||
LazyFeature(
|
||||
name="config_overrides",
|
||||
module_path="litellm.proxy.management_endpoints.config_override_endpoints",
|
||||
path_prefixes=("/config_overrides",),
|
||||
),
|
||||
LazyFeature(
|
||||
name="realtime",
|
||||
module_path="litellm.proxy.realtime_endpoints.endpoints",
|
||||
path_prefixes=("/openai/v1/realtime", "/v1/realtime", "/realtime"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="anthropic_passthrough",
|
||||
module_path="litellm.proxy.anthropic_endpoints.endpoints",
|
||||
path_prefixes=("/v1/messages", "/anthropic", "/api/event_logging"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="anthropic_skills",
|
||||
module_path="litellm.proxy.anthropic_endpoints.skills_endpoints",
|
||||
path_prefixes=("/v1/skills", "/skills"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="langfuse_passthrough",
|
||||
module_path="litellm.proxy.vertex_ai_endpoints.langfuse_endpoints",
|
||||
path_prefixes=("/langfuse",),
|
||||
),
|
||||
LazyFeature(
|
||||
name="evals",
|
||||
module_path="litellm.proxy.openai_evals_endpoints.endpoints",
|
||||
path_prefixes=("/v1/evals", "/evals"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="claude_code_marketplace",
|
||||
module_path="litellm.proxy.anthropic_endpoints.claude_code_endpoints",
|
||||
path_prefixes=("/claude-code",),
|
||||
register_fn=_include_router("claude_code_marketplace_router"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="scim",
|
||||
module_path="litellm.proxy.management_endpoints.scim.scim_v2",
|
||||
path_prefixes=("/scim",),
|
||||
register_fn=_include_router("scim_router"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="cloudzero",
|
||||
module_path="litellm.proxy.spend_tracking.cloudzero_endpoints",
|
||||
path_prefixes=("/cloudzero",),
|
||||
),
|
||||
LazyFeature(
|
||||
name="vantage",
|
||||
module_path="litellm.proxy.spend_tracking.vantage_endpoints",
|
||||
path_prefixes=("/vantage",),
|
||||
),
|
||||
LazyFeature(
|
||||
name="usage_ai",
|
||||
module_path="litellm.proxy.management_endpoints.usage_endpoints",
|
||||
path_prefixes=("/usage/ai",),
|
||||
),
|
||||
LazyFeature(
|
||||
name="prompts",
|
||||
module_path="litellm.proxy.prompts.prompt_endpoints",
|
||||
path_prefixes=("/prompts", "/utils/dotprompt_json_converter"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="jwt_mappings",
|
||||
module_path="litellm.proxy.management_endpoints.jwt_key_mapping_endpoints",
|
||||
path_prefixes=("/jwt/key/mapping",),
|
||||
),
|
||||
LazyFeature(
|
||||
name="compliance",
|
||||
module_path="litellm.proxy.management_endpoints.compliance_endpoints",
|
||||
path_prefixes=("/compliance",),
|
||||
),
|
||||
LazyFeature(
|
||||
name="access_groups",
|
||||
module_path="litellm.proxy.management_endpoints.access_group_endpoints",
|
||||
path_prefixes=("/access_group", "/v1/access_group", "/v1/unified_access_group"),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class LazyFeatureMiddleware:
|
||||
"""ASGI middleware that imports + registers a feature router on first
|
||||
matching request. Idempotent; once loaded, subsequent requests skip."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
app,
|
||||
fastapi_app: "FastAPI",
|
||||
features: Tuple[LazyFeature, ...] = LAZY_FEATURES,
|
||||
):
|
||||
self.app = app
|
||||
self._fastapi_app = fastapi_app
|
||||
self._features = features
|
||||
# Loaded set / per-feature locks live on app.state so the warm endpoint
|
||||
# and the middleware share them — preventing duplicate registrations
|
||||
# when both paths fire for the same feature.
|
||||
if not hasattr(fastapi_app.state, "lazy_loaded"):
|
||||
fastapi_app.state.lazy_loaded = set()
|
||||
fastapi_app.state.lazy_locks = {}
|
||||
|
||||
@property
|
||||
def _loaded(self) -> set:
|
||||
return self._fastapi_app.state.lazy_loaded
|
||||
|
||||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
# Short-circuit once every feature has loaded.
|
||||
if scope["type"] in ("http", "websocket") and len(self._loaded) < len(
|
||||
self._features
|
||||
):
|
||||
path = scope.get("path", "")
|
||||
for feat in self._features:
|
||||
if feat.module_path in self._loaded:
|
||||
continue
|
||||
if any(path.startswith(p) for p in feat.path_prefixes) or any(
|
||||
path.endswith(s) for s in feat.path_suffixes
|
||||
):
|
||||
await _force_load(self._fastapi_app, feat)
|
||||
await self.app(scope, receive, send)
|
||||
|
||||
|
||||
async def _force_load(app: "FastAPI", feat: LazyFeature) -> bool:
|
||||
"""Import + register a lazy feature exactly once per (app, module).
|
||||
Shared by the middleware and the /lazy/warm endpoint."""
|
||||
if not hasattr(app.state, "lazy_loaded"):
|
||||
app.state.lazy_loaded = set()
|
||||
app.state.lazy_locks = {}
|
||||
lock = app.state.lazy_locks.setdefault(feat.module_path, asyncio.Lock())
|
||||
async with lock:
|
||||
if feat.module_path in app.state.lazy_loaded:
|
||||
return False
|
||||
try:
|
||||
# Import on a thread (heavy modules take 1-3 s). register_fn
|
||||
# mutates app.router.routes, so it stays on the loop thread.
|
||||
loop = asyncio.get_running_loop()
|
||||
module = await loop.run_in_executor(
|
||||
None, importlib.import_module, feat.module_path
|
||||
)
|
||||
feat.register_fn(app, module)
|
||||
app.state.lazy_loaded.add(feat.module_path)
|
||||
app.openapi_schema = None
|
||||
verbose_proxy_logger.info(
|
||||
"Lazy-loaded optional feature %r (module: %s)",
|
||||
feat.name,
|
||||
feat.module_path,
|
||||
)
|
||||
return True
|
||||
except Exception as exc:
|
||||
# Mark loaded anyway so we don't retry on every request.
|
||||
app.state.lazy_loaded.add(feat.module_path)
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to lazy-load optional feature %r (module: %s): %s. "
|
||||
"This feature's endpoints will return 404 until restart.",
|
||||
feat.name,
|
||||
feat.module_path,
|
||||
exc,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def attach_lazy_features(app: "FastAPI") -> None:
|
||||
app.include_router(_make_warmup_router(app))
|
||||
app.add_middleware(LazyFeatureMiddleware, fastapi_app=app)
|
||||
|
||||
|
||||
def _make_warmup_router(app: "FastAPI") -> "APIRouter":
|
||||
"""POST /lazy/warm/{name}: load a feature and return its partial openapi
|
||||
so the Swagger plugin can merge in-place without a full /openapi.json refetch.
|
||||
Requires auth — anyone who can hit the proxy can already trigger the same
|
||||
imports by sending a real request to a feature's prefix, but gating this
|
||||
debug endpoint avoids unauthenticated callers forcing the import chain."""
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from fastapi.openapi.utils import get_openapi
|
||||
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
@router.post(
|
||||
"/lazy/warm/{name}",
|
||||
include_in_schema=False,
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def warm(name: str):
|
||||
feat = next((f for f in LAZY_FEATURES if f.name == name), None)
|
||||
if feat is None:
|
||||
raise HTTPException(404, f"unknown lazy feature: {name}")
|
||||
if feat.persistent_swagger_stub:
|
||||
return {"stub_path": None, "paths": {}, "components": {"schemas": {}}}
|
||||
|
||||
await _force_load(app, feat)
|
||||
|
||||
feat_routes = [
|
||||
r
|
||||
for r in app.routes
|
||||
if any(getattr(r, "path", "").startswith(p) for p in feat.path_prefixes)
|
||||
]
|
||||
full = get_openapi(title=app.title, version=app.version, routes=feat_routes)
|
||||
# Force all operations under one tag so they group under a single Swagger
|
||||
# section — many lazy modules tag routes inconsistently.
|
||||
for path_ops in full.get("paths", {}).values():
|
||||
for op in path_ops.values():
|
||||
if isinstance(op, dict):
|
||||
op["tags"] = [feat.name]
|
||||
return {
|
||||
"stub_path": feat.path_prefixes[0],
|
||||
"paths": full.get("paths", {}),
|
||||
"components": {"schemas": full.get("components", {}).get("schemas", {})},
|
||||
}
|
||||
|
||||
return router
|
||||
|
||||
|
||||
def inject_lazy_stubs(schema: Dict) -> Dict:
|
||||
"""Inject openapi entries for unloaded features. Uses the snapshot file
|
||||
when available (full route info), otherwise falls back to a single
|
||||
placeholder per feature. Any failure logs and returns the schema unchanged
|
||||
so /openapi.json never 500s on a cosmetic injection bug."""
|
||||
try:
|
||||
from litellm.proxy._lazy_openapi_snapshot import load_snapshot
|
||||
|
||||
snapshot = load_snapshot()
|
||||
paths = schema.setdefault("paths", {})
|
||||
schemas = schema.setdefault("components", {}).setdefault("schemas", {})
|
||||
|
||||
for feat in LAZY_FEATURES:
|
||||
if feat.module_path in sys.modules and not feat.persistent_swagger_stub:
|
||||
continue
|
||||
|
||||
fragment = (snapshot or {}).get(feat.name)
|
||||
if fragment:
|
||||
for p, ops in fragment.get("paths", {}).items():
|
||||
paths.setdefault(p, ops)
|
||||
for name, sch in (
|
||||
fragment.get("components", {}).get("schemas", {}).items()
|
||||
):
|
||||
schemas.setdefault(name, sch)
|
||||
continue
|
||||
|
||||
prefix = feat.path_prefixes[0]
|
||||
if prefix in paths:
|
||||
continue
|
||||
paths[prefix] = {
|
||||
"get": {
|
||||
"tags": [feat.name],
|
||||
"summary": feat.name,
|
||||
"responses": {"200": {"description": "OK"}},
|
||||
}
|
||||
}
|
||||
except Exception as exc:
|
||||
verbose_proxy_logger.warning("inject_lazy_stubs failed: %s", exc)
|
||||
return schema
|
||||
|
||||
|
||||
def lazy_tag_to_prefix() -> Dict[str, str]:
|
||||
"""feature.name -> first prefix, used by the Swagger warmup JS plugin.
|
||||
Returns empty when the snapshot is loaded — the plugin is unnecessary
|
||||
because /openapi.json already has full route info."""
|
||||
from litellm.proxy._lazy_openapi_snapshot import load_snapshot
|
||||
|
||||
if load_snapshot():
|
||||
return {}
|
||||
return {
|
||||
feat.name: feat.path_prefixes[0]
|
||||
for feat in LAZY_FEATURES
|
||||
if not feat.persistent_swagger_stub
|
||||
}
|
||||
31651
litellm/proxy/_lazy_openapi_snapshot.json
Normal file
31651
litellm/proxy/_lazy_openapi_snapshot.json
Normal file
File diff suppressed because it is too large
Load diff
70
litellm/proxy/_lazy_openapi_snapshot.py
Normal file
70
litellm/proxy/_lazy_openapi_snapshot.py
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
"""
|
||||
Per-feature OpenAPI snapshot for lazy-loaded routers.
|
||||
|
||||
The committed JSON is generated by `python -m litellm.proxy._lazy_openapi_snapshot`
|
||||
and consumed at runtime so /openapi.json can show full route info for unloaded
|
||||
features without importing them. CI verifies the file is current and surfaces
|
||||
any drift as a neutral check.
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional
|
||||
|
||||
SNAPSHOT_FILE = Path(__file__).parent / "_lazy_openapi_snapshot.json"
|
||||
|
||||
|
||||
def load_snapshot() -> Optional[Dict[str, Dict]]:
|
||||
if not SNAPSHOT_FILE.exists():
|
||||
return None
|
||||
try:
|
||||
with SNAPSHOT_FILE.open() as f:
|
||||
return json.load(f)
|
||||
except (json.JSONDecodeError, OSError):
|
||||
return None
|
||||
|
||||
|
||||
def generate_snapshot() -> Dict[str, Dict]:
|
||||
import importlib
|
||||
|
||||
from fastapi.openapi.utils import get_openapi
|
||||
|
||||
from litellm.proxy._lazy_features import LAZY_FEATURES
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
for feat in LAZY_FEATURES:
|
||||
if feat.module_path in sys.modules:
|
||||
continue
|
||||
try:
|
||||
module = importlib.import_module(feat.module_path)
|
||||
feat.register_fn(app, module)
|
||||
except Exception as exc:
|
||||
sys.stderr.write(f"warning: skip {feat.name}: {exc}\n")
|
||||
|
||||
fragments: Dict[str, Dict] = {}
|
||||
for feat in LAZY_FEATURES:
|
||||
feat_routes = [
|
||||
r
|
||||
for r in app.routes
|
||||
if any(getattr(r, "path", "").startswith(p) for p in feat.path_prefixes)
|
||||
]
|
||||
if not feat_routes:
|
||||
continue
|
||||
full = get_openapi(title=app.title, version=app.version, routes=feat_routes)
|
||||
# Group all of a feature's routes under one tag.
|
||||
for path_ops in full.get("paths", {}).values():
|
||||
for op in path_ops.values():
|
||||
if isinstance(op, dict):
|
||||
op["tags"] = [feat.name]
|
||||
fragments[feat.name] = {
|
||||
"paths": full.get("paths", {}),
|
||||
"components": {"schemas": full.get("components", {}).get("schemas", {})},
|
||||
}
|
||||
return fragments
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
fragments = generate_snapshot()
|
||||
SNAPSHOT_FILE.write_text(json.dumps(fragments, indent=2, sort_keys=True) + "\n")
|
||||
sys.stdout.write(f"wrote {len(fragments)} feature fragments to {SNAPSHOT_FILE}\n")
|
||||
|
|
@ -17,6 +17,9 @@ from typing_extensions import Required, TypedDict
|
|||
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import MCP_STDIO_ALLOWED_COMMANDS
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
|
||||
validate_no_callback_env_reference,
|
||||
)
|
||||
from litellm.types.integrations.slack_alerting import AlertType
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
|
|
@ -1863,8 +1866,10 @@ class AddTeamCallback(LiteLLMPydanticObjectBase):
|
|||
raise ValueError(
|
||||
f"Invalid callback variable: {key}. Must be one of {valid_keys}"
|
||||
)
|
||||
if not isinstance(value, str):
|
||||
callback_vars[key] = str(value)
|
||||
callback_vars[key] = str(value)
|
||||
validate_no_callback_env_reference(
|
||||
key, callback_vars[key], source="key/team callback metadata"
|
||||
)
|
||||
return values
|
||||
|
||||
|
||||
|
|
@ -2150,8 +2155,8 @@ class PassThroughGenericEndpoint(LiteLLMPydanticObjectBase):
|
|||
description="The USD cost per request to the target endpoint. This is used to calculate the cost of the request to the target endpoint.",
|
||||
)
|
||||
auth: bool = Field(
|
||||
default=False,
|
||||
description="Whether authentication is required for the pass-through endpoint. If True, requests to the endpoint will require a valid LiteLLM API key.",
|
||||
default=True,
|
||||
description="Whether authentication is required for the pass-through endpoint. Defaults to True so a pass-through silently created without an explicit value still requires a valid LiteLLM API key — set to False only if the endpoint is meant to be a public forwarder (e.g. an unauthenticated webhook target).",
|
||||
)
|
||||
guardrails: Optional[PassThroughGuardrailsConfig] = Field(
|
||||
default=None,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -361,7 +361,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]:
|
||||
"""
|
||||
|
|
@ -377,14 +377,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,
|
||||
|
|
@ -505,7 +505,12 @@ async def check_api_key_for_custom_headers_or_pass_through_endpoints(
|
|||
for endpoint in pass_through_endpoints:
|
||||
if isinstance(endpoint, dict) and endpoint.get("path", "") == route:
|
||||
## IF AUTH DISABLED
|
||||
if endpoint.get("auth") is not True:
|
||||
# Default to True: a config dict with no ``auth`` key
|
||||
# otherwise produced an unauthenticated forwarder. The
|
||||
# Pydantic ``PassThroughGenericEndpoint.auth`` default
|
||||
# is also True, but raw config dicts skip that path —
|
||||
# so this runtime check has to default to True too.
|
||||
if endpoint.get("auth", True) is not True:
|
||||
return UserAPIKeyAuth()
|
||||
## IF AUTH ENABLED
|
||||
### IF CUSTOM PARSER REQUIRED
|
||||
|
|
@ -537,7 +542,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]:
|
||||
|
|
@ -1141,9 +1146,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={
|
||||
|
|
@ -1213,11 +1216,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(
|
||||
|
|
@ -1325,7 +1332,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
|
||||
|
|
@ -1333,18 +1341,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
|
||||
|
|
@ -1491,9 +1504,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
|
||||
|
|
|
|||
|
|
@ -313,23 +313,24 @@ sequenceDiagram
|
|||
participant Proxy as LiteLLM Proxy
|
||||
participant SSO as SSO Provider
|
||||
|
||||
CLI->>CLI: Generate key ID (sk-uuid)
|
||||
CLI->>Browser: Open /sso/key/generate?source=litellm-cli&key=sk-uuid
|
||||
CLI->>Proxy: POST /sso/cli/start
|
||||
Proxy->>CLI: Return login_id, poll_secret, user_code
|
||||
CLI->>Browser: Open /sso/key/generate?source=litellm-cli&key=login_id
|
||||
|
||||
Browser->>Proxy: GET /sso/key/generate?source=litellm-cli&key=sk-uuid
|
||||
Proxy->>Proxy: Set cli_state = litellm-session-token:sk-uuid
|
||||
Proxy->>SSO: Redirect with state=litellm-session-token:sk-uuid
|
||||
Browser->>Proxy: GET /sso/key/generate?source=litellm-cli&key=login_id
|
||||
Proxy->>Proxy: Set cli_state = litellm-session-token:login_id
|
||||
Proxy->>SSO: Redirect with state=litellm-session-token:login_id
|
||||
|
||||
SSO->>Browser: Show login page
|
||||
Browser->>SSO: User authenticates
|
||||
SSO->>Proxy: Redirect to /sso/callback?state=litellm-session-token:sk-uuid
|
||||
SSO->>Proxy: Redirect to /sso/callback?state=litellm-session-token:login_id
|
||||
|
||||
Proxy->>Proxy: Check if state starts with "litellm-session-token:"
|
||||
Proxy->>Proxy: Generate API key with ID=sk-uuid
|
||||
Proxy->>Browser: Show success page
|
||||
Proxy->>Browser: Prompt for user_code
|
||||
Browser->>Proxy: POST /sso/cli/complete/login_id
|
||||
|
||||
CLI->>Proxy: Poll /sso/cli/poll/sk-uuid
|
||||
Proxy->>CLI: Return {"status": "ready", "key": "sk-uuid"}
|
||||
CLI->>Proxy: Poll /sso/cli/poll/login_id with poll_secret header
|
||||
Proxy->>CLI: Return {"status": "ready", "key": "jwt"}
|
||||
CLI->>CLI: Save key to ~/.litellm/token.json
|
||||
```
|
||||
|
||||
|
|
@ -343,13 +344,13 @@ The CLI provides three authentication commands:
|
|||
|
||||
### Authentication Flow Steps
|
||||
|
||||
1. **Generate Session ID**: CLI generates a unique key ID (`sk-{uuid}`)
|
||||
2. **Open Browser**: CLI opens browser to `/sso/key/generate` with CLI source and key parameters
|
||||
3. **SSO Redirect**: Proxy sets the formatted state (`litellm-session-token:sk-uuid`) as OAuth state parameter and redirects to SSO provider
|
||||
1. **Start Session**: CLI creates a short-lived login session with `/sso/cli/start`
|
||||
2. **Open Browser**: CLI opens browser to `/sso/key/generate` with CLI source and login ID parameters
|
||||
3. **SSO Redirect**: Proxy sets the formatted state (`litellm-session-token:{login_id}`) as OAuth state parameter and redirects to SSO provider
|
||||
4. **User Authentication**: User completes SSO authentication in browser
|
||||
5. **Callback Processing**: SSO provider redirects back to proxy with state parameter
|
||||
6. **Key Generation**: Proxy detects CLI login (state starts with "litellm-session-token:") and generates API key with pre-specified ID
|
||||
7. **Polling**: CLI polls `/sso/cli/poll/{key_id}` endpoint until key is ready
|
||||
6. **User Code Verification**: Browser confirms the verification code shown in the CLI
|
||||
7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready
|
||||
8. **Token Storage**: CLI saves the authentication token to `~/.litellm/token.json`
|
||||
|
||||
### Benefits of This Approach
|
||||
|
|
@ -357,7 +358,7 @@ The CLI provides three authentication commands:
|
|||
- **No Local Server**: No need to run a local callback server
|
||||
- **Standard OAuth**: Uses OAuth 2.0 state parameter correctly
|
||||
- **Remote Compatible**: Works with remote proxy servers
|
||||
- **Secure**: Uses UUID session identifiers
|
||||
- **Secure**: Keeps the polling secret out of the browser handoff
|
||||
- **Simple Setup**: No additional OAuth redirect URL configuration needed
|
||||
|
||||
### Token Storage
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ import time
|
|||
import webbrowser
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
from urllib.parse import urlencode
|
||||
|
||||
import click
|
||||
import requests
|
||||
|
|
@ -241,7 +242,7 @@ def prompt_team_selection(teams: List[Dict[str, Any]]) -> Optional[Dict[str, Any
|
|||
|
||||
|
||||
def prompt_team_selection_fallback(
|
||||
teams: List[Dict[str, Any]]
|
||||
teams: List[Dict[str, Any]],
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Fallback team selection for non-interactive environments"""
|
||||
if not teams:
|
||||
|
|
@ -279,6 +280,7 @@ def prompt_team_selection_fallback(
|
|||
def _poll_for_ready_data(
|
||||
url: str,
|
||||
*,
|
||||
headers: Optional[Dict[str, str]] = None,
|
||||
total_timeout: int = 300,
|
||||
poll_interval: int = 2,
|
||||
request_timeout: int = 10,
|
||||
|
|
@ -291,7 +293,10 @@ def _poll_for_ready_data(
|
|||
) -> Optional[Dict[str, Any]]:
|
||||
for attempt in range(total_timeout // poll_interval):
|
||||
try:
|
||||
response = requests.get(url, timeout=request_timeout)
|
||||
request_kwargs: Dict[str, Any] = {"timeout": request_timeout}
|
||||
if headers is not None:
|
||||
request_kwargs["headers"] = headers
|
||||
response = requests.get(url, **request_kwargs)
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
status = data.get("status")
|
||||
|
|
@ -346,7 +351,23 @@ def _normalize_teams(teams, team_details):
|
|||
return []
|
||||
|
||||
|
||||
def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
|
||||
def _start_cli_sso_flow(base_url: str) -> Dict[str, Any]:
|
||||
response = requests.post(f"{base_url}/sso/cli/start", timeout=10)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
required_fields = ("login_id", "poll_secret", "user_code")
|
||||
if not all(isinstance(data.get(field), str) for field in required_fields):
|
||||
raise ValueError("Invalid CLI SSO start response")
|
||||
return data
|
||||
|
||||
|
||||
def _get_cli_sso_poll_headers(poll_secret: str) -> Dict[str, str]:
|
||||
return {"x-litellm-cli-poll-secret": poll_secret}
|
||||
|
||||
|
||||
def _poll_for_authentication(
|
||||
base_url: str, key_id: str, poll_secret: str
|
||||
) -> Optional[dict]:
|
||||
"""
|
||||
Poll the server for authentication completion and handle team selection.
|
||||
|
||||
|
|
@ -356,6 +377,7 @@ def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
|
|||
poll_url = f"{base_url}/sso/cli/poll/{key_id}"
|
||||
data = _poll_for_ready_data(
|
||||
poll_url,
|
||||
headers=_get_cli_sso_poll_headers(poll_secret),
|
||||
pending_message="Still waiting for authentication...",
|
||||
)
|
||||
if not data:
|
||||
|
|
@ -373,6 +395,7 @@ def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
|
|||
jwt_with_team = _handle_team_selection_during_polling(
|
||||
base_url=base_url,
|
||||
key_id=key_id,
|
||||
poll_secret=poll_secret,
|
||||
teams=normalized_teams,
|
||||
)
|
||||
|
||||
|
|
@ -410,7 +433,7 @@ def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
|
|||
|
||||
|
||||
def _handle_team_selection_during_polling(
|
||||
base_url: str, key_id: str, teams: List[Dict[str, Any]]
|
||||
base_url: str, key_id: str, poll_secret: str, teams: List[Dict[str, Any]]
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Handle team selection and re-poll with selected team_id.
|
||||
|
|
@ -441,6 +464,7 @@ def _handle_team_selection_during_polling(
|
|||
poll_url = f"{base_url}/sso/cli/poll/{key_id}?team_id={team_id}"
|
||||
data = _poll_for_ready_data(
|
||||
poll_url,
|
||||
headers=_get_cli_sso_poll_headers(poll_secret),
|
||||
pending_message="Still waiting for team authentication...",
|
||||
other_status_message="Waiting for team authentication to complete...",
|
||||
http_error_log_every=10,
|
||||
|
|
@ -514,29 +538,24 @@ def _render_and_prompt_for_team_selection(teams: List[Dict[str, Any]]) -> Option
|
|||
@click.pass_context
|
||||
def login(ctx: click.Context):
|
||||
"""Login to LiteLLM proxy using SSO authentication"""
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER
|
||||
from litellm.proxy.client.cli.interface import show_commands
|
||||
|
||||
base_url = ctx.obj["base_url"]
|
||||
|
||||
# Check if we have an existing key to regenerate
|
||||
existing_key = get_stored_api_key()
|
||||
|
||||
# Generate unique key ID for this login session
|
||||
key_id = f"sk-{str(uuid.uuid4())}"
|
||||
|
||||
try:
|
||||
# Construct SSO login URL with CLI source and pre-generated key
|
||||
sso_url = f"{base_url}/sso/key/generate?source={LITELLM_CLI_SOURCE_IDENTIFIER}&key={key_id}"
|
||||
cli_sso_flow = _start_cli_sso_flow(base_url=base_url)
|
||||
key_id = cli_sso_flow["login_id"]
|
||||
poll_secret = cli_sso_flow["poll_secret"]
|
||||
user_code = cli_sso_flow["user_code"]
|
||||
|
||||
# If we have an existing key, include it as a parameter to the login endpoint
|
||||
# The server will encode it in the OAuth state parameter for the SSO flow
|
||||
if existing_key:
|
||||
sso_url += f"&existing_key={existing_key}"
|
||||
sso_url = f"{base_url}/sso/key/generate?" + urlencode(
|
||||
{"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": key_id}
|
||||
)
|
||||
|
||||
click.echo(f"Opening browser to: {sso_url}")
|
||||
click.echo("Please complete the SSO authentication in your browser...")
|
||||
click.echo(f"Verification code: {user_code}")
|
||||
click.echo(f"Session ID: {key_id}")
|
||||
|
||||
# Open browser
|
||||
|
|
@ -545,7 +564,9 @@ def login(ctx: click.Context):
|
|||
# Poll for authentication completion
|
||||
click.echo("Waiting for authentication...")
|
||||
|
||||
auth_result = _poll_for_authentication(base_url=base_url, key_id=key_id)
|
||||
auth_result = _poll_for_authentication(
|
||||
base_url=base_url, key_id=key_id, poll_secret=poll_secret
|
||||
)
|
||||
|
||||
if auth_result:
|
||||
api_key = auth_result["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
|
||||
|
|
@ -14,6 +14,8 @@ from litellm.types.utils import (
|
|||
blue_color_code = "\033[94m"
|
||||
reset_color_code = "\033[0m"
|
||||
|
||||
TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY = "_pillar_response_headers_trusted"
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
|
||||
|
|
@ -417,10 +419,19 @@ def get_logging_caching_headers(request_data: Dict) -> Optional[Dict]:
|
|||
if "semantic-similarity" in _metadata:
|
||||
headers["x-litellm-semantic-similarity"] = str(_metadata["semantic-similarity"])
|
||||
|
||||
is_trusted_pillar_metadata = (
|
||||
_metadata.get(TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY) is True
|
||||
)
|
||||
pillar_headers = _metadata.get("pillar_response_headers")
|
||||
if isinstance(pillar_headers, dict):
|
||||
headers.update(pillar_headers)
|
||||
elif "pillar_flagged" in _metadata:
|
||||
if is_trusted_pillar_metadata and isinstance(pillar_headers, dict):
|
||||
headers.update(
|
||||
{
|
||||
key: str(value)
|
||||
for key, value in pillar_headers.items()
|
||||
if isinstance(key, str) and key.lower().startswith("x-pillar-")
|
||||
}
|
||||
)
|
||||
elif is_trusted_pillar_metadata and "pillar_flagged" in _metadata:
|
||||
headers["x-pillar-flagged"] = str(_metadata["pillar_flagged"]).lower()
|
||||
|
||||
return headers
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -52,6 +52,37 @@ class ResetBudgetJob:
|
|||
### RESET MULTI-WINDOW BUDGETS ###
|
||||
await self.reset_budget_windows()
|
||||
|
||||
@staticmethod
|
||||
async def _invalidate_spend_counter(counter_key: str) -> None:
|
||||
"""Zero a spend counter so a DB-row reset takes effect immediately.
|
||||
|
||||
Call AFTER the DB write commits. Clearing Redis before the DB
|
||||
commit opens a window where get_current_spend reads 0 from Redis
|
||||
while the DB still holds the pre-reset value, allowing bypass.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import spend_counter_cache
|
||||
|
||||
spend_counter_cache.in_memory_cache.set_cache(
|
||||
key=counter_key, value=0.0, ttl=60
|
||||
)
|
||||
if spend_counter_cache.redis_cache is not None:
|
||||
try:
|
||||
await spend_counter_cache.redis_cache.async_set_cache(
|
||||
key=counter_key, value=0.0, ttl=60
|
||||
)
|
||||
except Exception as redis_err:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to reset spend counter %s in Redis: %s. "
|
||||
"Budget may be over-enforced until counter expires.",
|
||||
counter_key,
|
||||
redis_err,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to reset spend counter %s: %s", counter_key, e
|
||||
)
|
||||
|
||||
async def reset_budget_for_litellm_team_members(
|
||||
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
|
||||
):
|
||||
|
|
@ -64,46 +95,30 @@ class ResetBudgetJob:
|
|||
if budget.budget_id is not None
|
||||
]
|
||||
|
||||
# Reset spend counters for affected team members.
|
||||
# Reset Redis directly so a transient failure doesn't leave stale
|
||||
# counters that get_current_spend would read as authoritative.
|
||||
try:
|
||||
from litellm.proxy.proxy_server import spend_counter_cache
|
||||
|
||||
memberships = await self.prisma_client.db.litellm_teammembership.find_many(
|
||||
where={"budget_id": {"in": budget_ids}}
|
||||
)
|
||||
for m in memberships:
|
||||
counter_key = f"spend:team_member:{m.user_id}:{m.team_id}"
|
||||
# Always reset in-memory
|
||||
spend_counter_cache.in_memory_cache.set_cache(
|
||||
key=counter_key, value=0.0
|
||||
)
|
||||
# Explicitly reset Redis with warning on failure
|
||||
if spend_counter_cache.redis_cache is not None:
|
||||
try:
|
||||
await spend_counter_cache.redis_cache.async_set_cache(
|
||||
key=counter_key, value=0.0
|
||||
)
|
||||
except Exception as redis_err:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to reset team member spend counter in Redis %s: %s. "
|
||||
"Budget may be over-enforced until counter expires.",
|
||||
counter_key,
|
||||
redis_err,
|
||||
)
|
||||
except Exception as e:
|
||||
memberships = []
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to reset team member spend counters: %s", e
|
||||
"Failed to fetch team memberships for counter invalidation: %s", e
|
||||
)
|
||||
|
||||
return await self.prisma_client.db.litellm_teammembership.update_many(
|
||||
update_result = await self.prisma_client.db.litellm_teammembership.update_many(
|
||||
where={"budget_id": {"in": budget_ids}},
|
||||
data={
|
||||
"spend": 0,
|
||||
},
|
||||
)
|
||||
|
||||
for m in memberships:
|
||||
await self._invalidate_spend_counter(
|
||||
f"spend:team_member:{m.user_id}:{m.team_id}"
|
||||
)
|
||||
|
||||
return update_result
|
||||
|
||||
async def reset_budget_for_keys_linked_to_budgets(
|
||||
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
|
||||
):
|
||||
|
|
@ -126,17 +141,36 @@ class ResetBudgetJob:
|
|||
if not budget_ids:
|
||||
return
|
||||
|
||||
return await self.prisma_client.db.litellm_verificationtoken.update_many(
|
||||
where={
|
||||
"budget_id": {"in": budget_ids},
|
||||
"budget_duration": None, # only keys without their own reset schedule
|
||||
"spend": {"gt": 0}, # only reset keys that have accumulated spend
|
||||
},
|
||||
data={
|
||||
"spend": 0,
|
||||
},
|
||||
where_clause: dict = {
|
||||
"budget_id": {"in": budget_ids},
|
||||
"budget_duration": None, # only keys without their own reset schedule
|
||||
"spend": {"gt": 0}, # only reset keys that have accumulated spend
|
||||
}
|
||||
|
||||
try:
|
||||
keys = await self.prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where=where_clause
|
||||
)
|
||||
except Exception as e:
|
||||
keys = []
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to fetch keys for counter invalidation: %s", e
|
||||
)
|
||||
|
||||
update_result = (
|
||||
await self.prisma_client.db.litellm_verificationtoken.update_many(
|
||||
where=where_clause,
|
||||
data={
|
||||
"spend": 0,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
for k in keys:
|
||||
await self._invalidate_spend_counter(f"spend:key:{k.token}")
|
||||
|
||||
return update_result
|
||||
|
||||
async def reset_budget_for_litellm_budget_table(self):
|
||||
"""
|
||||
Resets the budget for all LiteLLM End-Users (Customers), and Team Members if their budget has expired
|
||||
|
|
@ -365,6 +399,10 @@ class ResetBudgetJob:
|
|||
data_list=updated_keys,
|
||||
table_name="key",
|
||||
)
|
||||
for k in updated_keys:
|
||||
token = getattr(k, "token", None)
|
||||
if token:
|
||||
await self._invalidate_spend_counter(f"spend:key:{token}")
|
||||
|
||||
end_time = time.time()
|
||||
if len(failed_keys) > 0: # If any keys failed to reset
|
||||
|
|
@ -450,6 +488,12 @@ class ResetBudgetJob:
|
|||
data_list=updated_users,
|
||||
table_name="user",
|
||||
)
|
||||
for u in updated_users:
|
||||
user_id = getattr(u, "user_id", None)
|
||||
if user_id:
|
||||
await self._invalidate_spend_counter(
|
||||
f"spend:user:{user_id}"
|
||||
)
|
||||
|
||||
end_time = time.time()
|
||||
if len(failed_users) > 0: # If any users failed to reset
|
||||
|
|
@ -541,6 +585,12 @@ class ResetBudgetJob:
|
|||
data_list=updated_teams,
|
||||
table_name="team",
|
||||
)
|
||||
for t in updated_teams:
|
||||
team_id = getattr(t, "team_id", None)
|
||||
if team_id:
|
||||
await self._invalidate_spend_counter(
|
||||
f"spend:team:{team_id}"
|
||||
)
|
||||
|
||||
end_time = time.time()
|
||||
if len(failed_teams) > 0: # If any teams failed to reset
|
||||
|
|
|
|||
52
litellm/proxy/common_utils/static_asset_utils.py
Normal file
52
litellm/proxy/common_utils/static_asset_utils.py
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
"""Helpers for unauthenticated logo / favicon endpoints."""
|
||||
|
||||
import os
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
LOCAL_IMAGE_HEADER_BYTES = 512
|
||||
|
||||
|
||||
def detect_local_image_media_type(header: bytes) -> Optional[str]:
|
||||
"""Return a browser image media type for supported local image signatures."""
|
||||
if header[0:8] == b"\x89PNG\r\n\x1a\n":
|
||||
return "image/png"
|
||||
if header[0:4] == b"GIF8" and header[5:6] == b"a":
|
||||
return "image/gif"
|
||||
if header[0:3] == b"\xff\xd8\xff":
|
||||
return "image/jpeg"
|
||||
if header[0:4] == b"RIFF" and header[8:12] == b"WEBP":
|
||||
return "image/webp"
|
||||
if header[0:4] in (b"\x00\x00\x01\x00", b"\x00\x00\x02\x00"):
|
||||
return "image/x-icon"
|
||||
return None
|
||||
|
||||
|
||||
def resolve_validated_local_image_path(candidate: str) -> Optional[Tuple[str, str]]:
|
||||
"""Resolve ``candidate`` only when it is an existing supported image file."""
|
||||
if not candidate:
|
||||
return None
|
||||
try:
|
||||
resolved = os.path.realpath(os.path.expanduser(candidate))
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
if not os.path.isfile(resolved):
|
||||
return None
|
||||
|
||||
try:
|
||||
with open(resolved, "rb") as f:
|
||||
header = f.read(LOCAL_IMAGE_HEADER_BYTES)
|
||||
except OSError as exc:
|
||||
verbose_proxy_logger.debug("Could not read local asset %r: %s", candidate, exc)
|
||||
return None
|
||||
|
||||
media_type = detect_local_image_media_type(header)
|
||||
if media_type is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"Local asset %r is not a supported image file; falling back to default.",
|
||||
candidate,
|
||||
)
|
||||
return None
|
||||
|
||||
return resolved, media_type
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -129,7 +129,9 @@ class SpendCounterReseed:
|
|||
"""
|
||||
lock = await SpendCounterReseed._get_lock(counter_key)
|
||||
async with lock:
|
||||
# Re-check after acquiring the lock - another waiter may have warmed it.
|
||||
# Re-check after acquiring the lock. Skip in-memory on a clean
|
||||
# Redis miss - in-memory is per-pod-stale.
|
||||
redis_clean_miss = False
|
||||
if spend_counter_cache.redis_cache is not None:
|
||||
try:
|
||||
val = await spend_counter_cache.redis_cache.async_get_cache(
|
||||
|
|
@ -137,11 +139,13 @@ class SpendCounterReseed:
|
|||
)
|
||||
if val is not None:
|
||||
return float(val)
|
||||
redis_clean_miss = True
|
||||
except Exception:
|
||||
pass
|
||||
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
|
||||
if val is not None:
|
||||
return float(val)
|
||||
if not redis_clean_miss:
|
||||
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
|
||||
if val is not None:
|
||||
return float(val)
|
||||
|
||||
db_spend = await SpendCounterReseed.from_db(prisma_client, counter_key)
|
||||
if db_spend is None:
|
||||
|
|
@ -149,7 +153,7 @@ class SpendCounterReseed:
|
|||
# Warm even when 0 so subsequent reads hit cache, not DB.
|
||||
try:
|
||||
await spend_counter_cache.async_increment_cache(
|
||||
key=counter_key, value=db_spend
|
||||
key=counter_key, value=db_spend, refresh_ttl=True
|
||||
)
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception(
|
||||
|
|
|
|||
|
|
@ -1,10 +1,6 @@
|
|||
from datetime import datetime
|
||||
from fastapi import APIRouter, Depends, Request, Response
|
||||
from fastapi.responses import ORJSONResponse
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response
|
||||
from fastapi.responses import ORJSONResponse, StreamingResponse
|
||||
|
||||
import litellm
|
||||
from litellm._uuid import uuid
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
|
@ -30,12 +26,17 @@ async def google_generate_content(
|
|||
fastapi_response: Response,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
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,
|
||||
)
|
||||
|
||||
|
|
@ -43,48 +44,33 @@ async def google_generate_content(
|
|||
if "model" not in data:
|
||||
data["model"] = model_name
|
||||
|
||||
# Extract generationConfig and pass it as config parameter
|
||||
generation_config = data.pop("generationConfig", None)
|
||||
if generation_config:
|
||||
data["config"] = generation_config
|
||||
|
||||
# Add user authentication metadata for cost tracking
|
||||
data = await add_litellm_data_to_request(
|
||||
data=data,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_config=proxy_config,
|
||||
general_settings=general_settings,
|
||||
version=version,
|
||||
)
|
||||
|
||||
# Create logging object with full request metadata so callbacks (e.g. S3) get user/trace_id
|
||||
data["litellm_call_id"] = request.headers.get(
|
||||
"x-litellm-call-id", str(uuid.uuid4())
|
||||
)
|
||||
logging_obj, data = litellm.utils.function_setup(
|
||||
original_function="agenerate_content",
|
||||
rules_obj=litellm.utils.Rules(),
|
||||
start_time=datetime.now(),
|
||||
**data,
|
||||
)
|
||||
data["litellm_logging_obj"] = logging_obj
|
||||
|
||||
# call router
|
||||
if llm_router is None:
|
||||
raise HTTPException(status_code=500, detail="Router not initialized")
|
||||
response = await llm_router.agenerate_content(**data)
|
||||
success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
|
||||
response=response,
|
||||
request_data=data,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
logging_obj=logging_obj,
|
||||
version=version,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
fastapi_response.headers.update(success_headers)
|
||||
return response
|
||||
processor = ProxyBaseLLMRequestProcessing(data=data)
|
||||
try:
|
||||
return await processor.base_process_llm_request(
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
route_type="agenerate_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=model_name,
|
||||
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,
|
||||
)
|
||||
except Exception as 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,
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
|
|
@ -101,73 +87,52 @@ async def google_stream_generate_content(
|
|||
fastapi_response: Response,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
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,
|
||||
)
|
||||
|
||||
data = await _read_request_body(request=request)
|
||||
|
||||
if "model" not in data:
|
||||
data["model"] = model_name
|
||||
data["stream"] = True
|
||||
|
||||
data["stream"] = True # enforce streaming for this endpoint
|
||||
|
||||
# Extract generationConfig and pass it as config parameter
|
||||
generation_config = data.pop("generationConfig", None)
|
||||
if generation_config:
|
||||
data["config"] = generation_config
|
||||
|
||||
# Add user authentication metadata for cost tracking
|
||||
data = await add_litellm_data_to_request(
|
||||
data=data,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_config=proxy_config,
|
||||
general_settings=general_settings,
|
||||
version=version,
|
||||
)
|
||||
|
||||
# Create logging object with full request metadata so streaming END callbacks (e.g. S3) get user/trace_id
|
||||
data["litellm_call_id"] = request.headers.get(
|
||||
"x-litellm-call-id", str(uuid.uuid4())
|
||||
)
|
||||
logging_obj, data = litellm.utils.function_setup(
|
||||
original_function="agenerate_content_stream",
|
||||
rules_obj=litellm.utils.Rules(),
|
||||
start_time=datetime.now(),
|
||||
**data,
|
||||
)
|
||||
data["litellm_logging_obj"] = logging_obj
|
||||
|
||||
# call router
|
||||
if llm_router is None:
|
||||
raise HTTPException(status_code=500, detail="Router not initialized")
|
||||
response = await llm_router.agenerate_content_stream(**data)
|
||||
|
||||
success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
|
||||
response=response,
|
||||
request_data=data,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
logging_obj=logging_obj,
|
||||
version=version,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
# Check if response is an async iterator (streaming response)
|
||||
if response is not None and hasattr(response, "__aiter__"):
|
||||
return StreamingResponse(
|
||||
content=response,
|
||||
media_type="text/event-stream",
|
||||
headers=success_headers,
|
||||
processor = ProxyBaseLLMRequestProcessing(data=data)
|
||||
try:
|
||||
return await processor.base_process_llm_request(
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
route_type="agenerate_content_stream",
|
||||
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=model_name,
|
||||
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,
|
||||
)
|
||||
except Exception as 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,
|
||||
)
|
||||
fastapi_response.headers.update(success_headers)
|
||||
return response
|
||||
|
||||
|
||||
@router.post(
|
||||
|
|
|
|||
|
|
@ -71,6 +71,7 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
GUARDRAIL_NAME = "bedrock"
|
||||
_BEDROCK_DYNAMIC_BODY_DENYLIST = frozenset({"content", "source"})
|
||||
|
||||
|
||||
class GuardrailMessageFilterResult(NamedTuple):
|
||||
|
|
@ -413,11 +414,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
)
|
||||
api_key: Optional[str] = None
|
||||
if request_data:
|
||||
bedrock_request_data.update(
|
||||
dynamic_request_body_params = (
|
||||
self.get_guardrail_dynamic_request_body_params(
|
||||
request_data=request_data
|
||||
)
|
||||
)
|
||||
bedrock_request_data.update(
|
||||
{
|
||||
key: value
|
||||
for key, value in dynamic_request_body_params.items()
|
||||
if key not in _BEDROCK_DYNAMIC_BODY_DENYLIST
|
||||
}
|
||||
)
|
||||
if request_data.get("api_key") is not None:
|
||||
api_key = request_data["api_key"]
|
||||
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY,
|
||||
add_guardrail_to_applied_guardrails_header,
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
)
|
||||
|
|
@ -144,6 +145,7 @@ def build_pillar_response_headers(metadata_store: Dict[str, Any]) -> Dict[str, s
|
|||
|
||||
if headers:
|
||||
metadata_store["pillar_response_headers"] = headers
|
||||
metadata_store[TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY] = True
|
||||
|
||||
return headers
|
||||
|
||||
|
|
|
|||
|
|
@ -41,6 +41,7 @@ class KeyManagementEventHooks:
|
|||
"""
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
get_audit_log_changed_by,
|
||||
)
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name
|
||||
|
||||
|
|
@ -61,9 +62,11 @@ class KeyManagementEventHooks:
|
|||
request_data=LiteLLM_AuditLogs(
|
||||
id=str(uuid.uuid4()),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
changed_by=litellm_changed_by
|
||||
or user_api_key_dict.user_id
|
||||
or litellm_proxy_admin_name,
|
||||
changed_by=get_audit_log_changed_by(
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
),
|
||||
changed_by_api_key=user_api_key_dict.api_key,
|
||||
table_name=LitellmTableNames.KEY_TABLE_NAME,
|
||||
object_id=response.token_id or "",
|
||||
|
|
@ -102,6 +105,7 @@ class KeyManagementEventHooks:
|
|||
"""
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
get_audit_log_changed_by,
|
||||
)
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name
|
||||
|
||||
|
|
@ -117,9 +121,11 @@ class KeyManagementEventHooks:
|
|||
request_data=LiteLLM_AuditLogs(
|
||||
id=str(uuid.uuid4()),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
changed_by=litellm_changed_by
|
||||
or user_api_key_dict.user_id
|
||||
or litellm_proxy_admin_name,
|
||||
changed_by=get_audit_log_changed_by(
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
),
|
||||
changed_by_api_key=user_api_key_dict.api_key,
|
||||
table_name=LitellmTableNames.KEY_TABLE_NAME,
|
||||
object_id=data.key,
|
||||
|
|
@ -140,6 +146,7 @@ class KeyManagementEventHooks:
|
|||
):
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
get_audit_log_changed_by,
|
||||
)
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name
|
||||
|
||||
|
|
@ -189,9 +196,11 @@ class KeyManagementEventHooks:
|
|||
request_data=LiteLLM_AuditLogs(
|
||||
id=str(uuid.uuid4()),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
changed_by=litellm_changed_by
|
||||
or user_api_key_dict.user_id
|
||||
or litellm_proxy_admin_name,
|
||||
changed_by=get_audit_log_changed_by(
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
),
|
||||
changed_by_api_key=user_api_key_dict.token,
|
||||
table_name=LitellmTableNames.KEY_TABLE_NAME,
|
||||
object_id=existing_key_row.token,
|
||||
|
|
@ -220,6 +229,7 @@ class KeyManagementEventHooks:
|
|||
"""
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
get_audit_log_changed_by,
|
||||
)
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name
|
||||
|
||||
|
|
@ -237,9 +247,11 @@ class KeyManagementEventHooks:
|
|||
request_data=LiteLLM_AuditLogs(
|
||||
id=str(uuid.uuid4()),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
changed_by=litellm_changed_by
|
||||
or user_api_key_dict.user_id
|
||||
or litellm_proxy_admin_name,
|
||||
changed_by=get_audit_log_changed_by(
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
),
|
||||
changed_by_api_key=user_api_key_dict.token,
|
||||
table_name=LitellmTableNames.KEY_TABLE_NAME,
|
||||
object_id=key.token,
|
||||
|
|
|
|||
|
|
@ -192,13 +192,19 @@ class UserManagementEventHooks:
|
|||
if not litellm.store_audit_logs:
|
||||
return
|
||||
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
get_audit_log_changed_by,
|
||||
)
|
||||
|
||||
await create_audit_log_for_update(
|
||||
request_data=LiteLLM_AuditLogs(
|
||||
id=str(uuid.uuid4()),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
changed_by=litellm_changed_by
|
||||
or user_api_key_dict.user_id
|
||||
or litellm_proxy_admin_name,
|
||||
changed_by=get_audit_log_changed_by(
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
),
|
||||
changed_by_api_key=user_api_key_dict.api_key,
|
||||
table_name=LitellmTableNames.USER_TABLE_NAME,
|
||||
object_id=user_id,
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from collections import OrderedDict
|
|||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
|
||||
|
||||
from fastapi import Request
|
||||
from pydantic import ValidationError as PydanticValidationError
|
||||
from starlette.datastructures import Headers
|
||||
|
||||
import litellm
|
||||
|
|
@ -104,6 +105,112 @@ LITELLM_METADATA_ROUTES = (
|
|||
"files",
|
||||
)
|
||||
|
||||
_UNTRUSTED_ROOT_CONTROL_FIELDS = (
|
||||
"proxy_server_request",
|
||||
"standard_logging_object",
|
||||
"secret_fields",
|
||||
"mock_response",
|
||||
"mock_tool_calls",
|
||||
"disable_global_guardrails",
|
||||
"disable_global_guardrail",
|
||||
"opted_out_global_guardrails",
|
||||
"applied_guardrails",
|
||||
"applied_policies",
|
||||
"policy_sources",
|
||||
"pillar_response_headers",
|
||||
"_guardrail_pipelines",
|
||||
"_pipeline_managed_guardrails",
|
||||
)
|
||||
|
||||
_UNTRUSTED_METADATA_CONTROL_FIELDS = (
|
||||
"disable_global_guardrails",
|
||||
"disable_global_guardrail",
|
||||
"opted_out_global_guardrails",
|
||||
"pillar_response_headers",
|
||||
"_pillar_response_headers_trusted",
|
||||
"pillar_flagged",
|
||||
"pillar_scanners",
|
||||
"pillar_evidence",
|
||||
"pillar_evidence_truncated",
|
||||
"pillar_session_id_response",
|
||||
"applied_guardrails",
|
||||
"applied_policies",
|
||||
"policy_sources",
|
||||
"standard_logging_object",
|
||||
"proxy_server_request",
|
||||
"secret_fields",
|
||||
"_guardrail_pipelines",
|
||||
"_pipeline_managed_guardrails",
|
||||
)
|
||||
|
||||
_UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS = frozenset(
|
||||
{
|
||||
"litellm-disable-message-redaction",
|
||||
}
|
||||
)
|
||||
_CLIENT_MOCK_CONTROL_FIELDS = frozenset({"mock_response", "mock_tool_calls"})
|
||||
_ALLOW_CLIENT_MOCK_RESPONSE_METADATA_KEY = "allow_client_mock_response"
|
||||
_ALLOW_CLIENT_MESSAGE_REDACTION_OPT_OUT_METADATA_KEY = (
|
||||
"allow_client_message_redaction_opt_out"
|
||||
)
|
||||
|
||||
|
||||
def _strip_untrusted_request_header_controls(
|
||||
headers: Any,
|
||||
*,
|
||||
allow_client_message_redaction_opt_out: bool = False,
|
||||
) -> None:
|
||||
if not isinstance(headers, dict):
|
||||
return
|
||||
|
||||
for header_name in list(headers.keys()):
|
||||
if (
|
||||
isinstance(header_name, str)
|
||||
and header_name.lower() in _UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS
|
||||
):
|
||||
if allow_client_message_redaction_opt_out:
|
||||
continue
|
||||
headers.pop(header_name, None)
|
||||
|
||||
|
||||
def _is_false_like(value: Any) -> bool:
|
||||
if isinstance(value, bool):
|
||||
return value is False
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() in {"false", "0", "no", "off"}
|
||||
return False
|
||||
|
||||
|
||||
def _key_or_team_metadata_flag_is_true(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
metadata_key: str,
|
||||
) -> bool:
|
||||
for admin_metadata in (user_api_key_dict.metadata, user_api_key_dict.team_metadata):
|
||||
if (
|
||||
isinstance(admin_metadata, dict)
|
||||
and admin_metadata.get(metadata_key) is True
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _key_or_team_allows_client_mock_response(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> bool:
|
||||
return _key_or_team_metadata_flag_is_true(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
metadata_key=_ALLOW_CLIENT_MOCK_RESPONSE_METADATA_KEY,
|
||||
)
|
||||
|
||||
|
||||
def _key_or_team_allows_client_message_redaction_opt_out(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> bool:
|
||||
return _key_or_team_metadata_flag_is_true(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
metadata_key=_ALLOW_CLIENT_MESSAGE_REDACTION_OPT_OUT_METADATA_KEY,
|
||||
)
|
||||
|
||||
|
||||
def _get_metadata_variable_name(request: Request) -> str:
|
||||
"""
|
||||
|
|
@ -228,13 +335,25 @@ def convert_key_logging_metadata_to_callback(
|
|||
for var, value in data.callback_vars.items():
|
||||
if team_callback_settings_obj.callback_vars is None:
|
||||
team_callback_settings_obj.callback_vars = {}
|
||||
team_callback_settings_obj.callback_vars[var] = str(
|
||||
litellm.utils.get_secret(value, default_value=value) or value
|
||||
)
|
||||
team_callback_settings_obj.callback_vars[var] = str(value)
|
||||
|
||||
return team_callback_settings_obj
|
||||
|
||||
|
||||
def _get_validated_callback_metadata(
|
||||
item: dict, *, source: str
|
||||
) -> Optional[AddTeamCallback]:
|
||||
try:
|
||||
return AddTeamCallback(**item)
|
||||
except (PydanticValidationError, ValueError) as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Ignoring invalid %s callback metadata: %s",
|
||||
source,
|
||||
_sanitize_for_log(str(e)),
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
class KeyAndTeamLoggingSettings:
|
||||
"""
|
||||
Helper class to get the dynamic logging settings for the key and team
|
||||
|
|
@ -274,8 +393,11 @@ def _get_dynamic_logging_metadata(
|
|||
#########################################################################################
|
||||
if key_dynamic_logging_settings is not None:
|
||||
for item in key_dynamic_logging_settings:
|
||||
callback = _get_validated_callback_metadata(item=item, source="key-level")
|
||||
if callback is None:
|
||||
continue
|
||||
callback_settings_obj = convert_key_logging_metadata_to_callback(
|
||||
data=AddTeamCallback(**item),
|
||||
data=callback,
|
||||
team_callback_settings_obj=callback_settings_obj,
|
||||
)
|
||||
#########################################################################################
|
||||
|
|
@ -283,8 +405,11 @@ def _get_dynamic_logging_metadata(
|
|||
#########################################################################################
|
||||
elif team_dynamic_logging_settings is not None:
|
||||
for item in team_dynamic_logging_settings:
|
||||
callback = _get_validated_callback_metadata(item=item, source="team-level")
|
||||
if callback is None:
|
||||
continue
|
||||
callback_settings_obj = convert_key_logging_metadata_to_callback(
|
||||
data=AddTeamCallback(**item),
|
||||
data=callback,
|
||||
team_callback_settings_obj=callback_settings_obj,
|
||||
)
|
||||
#########################################################################################
|
||||
|
|
@ -904,6 +1029,14 @@ class LiteLLMProxyRequestSetup:
|
|||
callback_vars_dict.pop("team_id", None)
|
||||
callback_vars_dict.pop("success_callback", None)
|
||||
callback_vars_dict.pop("failure_callback", None)
|
||||
callback_vars_dict = {
|
||||
key: (
|
||||
litellm.utils.get_secret(value, default_value=value) or value
|
||||
if isinstance(value, str)
|
||||
else value
|
||||
)
|
||||
for key, value in callback_vars_dict.items()
|
||||
}
|
||||
|
||||
return TeamCallbackMetadata(
|
||||
success_callback=team_config.get("success_callback", None),
|
||||
|
|
@ -962,11 +1095,15 @@ async def add_litellm_data_to_request( # noqa: PLR0915
|
|||
# Strip internal-only keys from user input before the proxy sets its own.
|
||||
# These keys are injected by the proxy itself below — user-supplied values
|
||||
# must not be trusted.
|
||||
for _internal_key in (
|
||||
"proxy_server_request",
|
||||
"standard_logging_object",
|
||||
"secret_fields",
|
||||
):
|
||||
_allow_client_mock_response = _key_or_team_allows_client_mock_response(
|
||||
user_api_key_dict
|
||||
)
|
||||
_allow_client_message_redaction_opt_out = (
|
||||
_key_or_team_allows_client_message_redaction_opt_out(user_api_key_dict)
|
||||
)
|
||||
for _internal_key in _UNTRUSTED_ROOT_CONTROL_FIELDS:
|
||||
if _allow_client_mock_response and _internal_key in _CLIENT_MOCK_CONTROL_FIELDS:
|
||||
continue
|
||||
data.pop(_internal_key, None)
|
||||
# Strip spoofable auth metadata from user-supplied metadata dict
|
||||
_user_metadata = data.get("metadata")
|
||||
|
|
@ -1007,6 +1144,17 @@ async def add_litellm_data_to_request( # noqa: PLR0915
|
|||
forward_llm_provider_auth_headers=forward_llm_auth,
|
||||
authenticated_with_header=authenticated_with_header,
|
||||
)
|
||||
_strip_untrusted_request_header_controls(
|
||||
_headers,
|
||||
allow_client_message_redaction_opt_out=_allow_client_message_redaction_opt_out,
|
||||
)
|
||||
if (
|
||||
not _allow_client_message_redaction_opt_out
|
||||
and litellm.turn_off_message_logging is True
|
||||
and "turn_off_message_logging" in data
|
||||
and _is_false_like(data["turn_off_message_logging"])
|
||||
):
|
||||
data.pop("turn_off_message_logging", None)
|
||||
verbose_proxy_logger.debug(f"Request Headers: {_headers}")
|
||||
verbose_proxy_logger.debug(f"Raw Headers: {_raw_headers}")
|
||||
|
||||
|
|
@ -1144,8 +1292,18 @@ async def add_litellm_data_to_request( # noqa: PLR0915
|
|||
for _meta_key in ("metadata", "litellm_metadata"):
|
||||
_user_meta = data.get(_meta_key)
|
||||
if isinstance(_user_meta, dict):
|
||||
_user_meta.pop("_pipeline_managed_guardrails", None)
|
||||
for _k in [k for k in _user_meta if k.startswith("user_api_key_")]:
|
||||
_strip_untrusted_request_header_controls(
|
||||
_user_meta.get("headers"),
|
||||
allow_client_message_redaction_opt_out=(
|
||||
_allow_client_message_redaction_opt_out
|
||||
),
|
||||
)
|
||||
for _k in [
|
||||
k
|
||||
for k in _user_meta
|
||||
if k.startswith("user_api_key_")
|
||||
or k in _UNTRUSTED_METADATA_CONTROL_FIELDS
|
||||
]:
|
||||
_user_meta.pop(_k, None)
|
||||
|
||||
# Strip caller-supplied routing/budget tags unless the admin has opted
|
||||
|
|
|
|||
|
|
@ -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
|
||||
]
|
||||
|
|
|
|||
|
|
@ -2069,6 +2069,9 @@ async def delete_user(
|
|||
litellm_proxy_admin_name,
|
||||
prisma_client,
|
||||
)
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
get_audit_log_changed_by,
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
||||
|
|
@ -2162,9 +2165,11 @@ async def delete_user(
|
|||
request_data=LiteLLM_AuditLogs(
|
||||
id=str(uuid.uuid4()),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
changed_by=litellm_changed_by
|
||||
or user_api_key_dict.user_id
|
||||
or litellm_proxy_admin_name,
|
||||
changed_by=get_audit_log_changed_by(
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
),
|
||||
changed_by_api_key=user_api_key_dict.api_key,
|
||||
table_name=LitellmTableNames.USER_TABLE_NAME,
|
||||
object_id=user_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,
|
||||
|
|
@ -37,6 +37,7 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
|||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.proxy._experimental.mcp_server.db import (
|
||||
rotate_mcp_server_credentials_master_key,
|
||||
rotate_mcp_user_credentials_master_key,
|
||||
)
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy._types import LiteLLM_VerificationToken
|
||||
|
|
@ -1058,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.
|
||||
|
|
@ -1833,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,
|
||||
|
|
@ -3297,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)
|
||||
|
||||
|
|
@ -3340,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:
|
||||
|
|
@ -3414,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]]:
|
||||
|
|
@ -3604,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,
|
||||
|
|
@ -3718,6 +3719,17 @@ async def _rotate_master_key( # noqa: PLR0915
|
|||
"Failed to rotate MCP server credentials: %s", str(e)
|
||||
)
|
||||
|
||||
# 4b. process MCP user-scoped credentials table (BYOK + OAuth2 tokens)
|
||||
try:
|
||||
await rotate_mcp_user_credentials_master_key(
|
||||
prisma_client=prisma_client,
|
||||
new_master_key=new_master_key,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to rotate MCP user credentials: %s", str(e)
|
||||
)
|
||||
|
||||
# 5. process credentials table
|
||||
try:
|
||||
credentials = await prisma_client.db.litellm_credentialstable.find_many()
|
||||
|
|
@ -3850,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."""
|
||||
|
|
@ -4140,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
|
||||
|
|
@ -5161,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:
|
||||
"""
|
||||
|
|
@ -5254,6 +5266,9 @@ async def block_key(
|
|||
proxy_logging_obj,
|
||||
user_api_key_cache,
|
||||
)
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
get_audit_log_changed_by,
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
raise Exception("{}".format(CommonProxyErrors.db_not_connected_error.value))
|
||||
|
|
@ -5297,9 +5312,11 @@ async def block_key(
|
|||
request_data=LiteLLM_AuditLogs(
|
||||
id=str(uuid.uuid4()),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
changed_by=litellm_changed_by
|
||||
or user_api_key_dict.user_id
|
||||
or litellm_proxy_admin_name,
|
||||
changed_by=get_audit_log_changed_by(
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
),
|
||||
changed_by_api_key=user_api_key_dict.api_key,
|
||||
table_name=LitellmTableNames.KEY_TABLE_NAME,
|
||||
object_id=hashed_token,
|
||||
|
|
@ -5363,6 +5380,9 @@ async def unblock_key(
|
|||
proxy_logging_obj,
|
||||
user_api_key_cache,
|
||||
)
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
get_audit_log_changed_by,
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
raise Exception("{}".format(CommonProxyErrors.db_not_connected_error.value))
|
||||
|
|
@ -5406,9 +5426,11 @@ async def unblock_key(
|
|||
request_data=LiteLLM_AuditLogs(
|
||||
id=str(uuid.uuid4()),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
changed_by=litellm_changed_by
|
||||
or user_api_key_dict.user_id
|
||||
or litellm_proxy_admin_name,
|
||||
changed_by=get_audit_log_changed_by(
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
),
|
||||
changed_by_api_key=user_api_key_dict.api_key,
|
||||
table_name=LitellmTableNames.KEY_TABLE_NAME,
|
||||
object_id=hashed_token,
|
||||
|
|
@ -5589,7 +5611,6 @@ async def test_key_logging(
|
|||
"content": "Hello, this is a test from litellm /key/health. No LLM API call was made for this",
|
||||
}
|
||||
],
|
||||
"mock_response": "test response",
|
||||
}
|
||||
data = await add_litellm_data_to_request(
|
||||
data=data,
|
||||
|
|
@ -5598,6 +5619,7 @@ async def test_key_logging(
|
|||
general_settings=general_settings,
|
||||
request=request,
|
||||
)
|
||||
data["mock_response"] = "test response"
|
||||
await litellm.acompletion(
|
||||
**data
|
||||
) # make mock completion call to trigger key based callbacks
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
|||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
)
|
||||
from litellm.proxy.management_helpers.audit_logs import get_audit_log_changed_by
|
||||
|
||||
router = APIRouter(prefix="/v1/mcp", tags=["mcp"])
|
||||
|
||||
|
|
@ -2230,7 +2231,12 @@ if MCP_AVAILABLE:
|
|||
detail={"error": "Only proxy admins can create MCP toolsets."},
|
||||
)
|
||||
touched_by = (
|
||||
litellm_changed_by or user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME
|
||||
get_audit_log_changed_by(
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=LITELLM_PROXY_ADMIN_NAME,
|
||||
)
|
||||
or LITELLM_PROXY_ADMIN_NAME
|
||||
)
|
||||
try:
|
||||
result = await create_mcp_toolset(prisma_client, payload, touched_by)
|
||||
|
|
@ -2321,7 +2327,12 @@ if MCP_AVAILABLE:
|
|||
detail={"error": "Only proxy admins can update MCP toolsets."},
|
||||
)
|
||||
touched_by = (
|
||||
litellm_changed_by or user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME
|
||||
get_audit_log_changed_by(
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=LITELLM_PROXY_ADMIN_NAME,
|
||||
)
|
||||
or LITELLM_PROXY_ADMIN_NAME
|
||||
)
|
||||
try:
|
||||
result = await update_mcp_toolset(prisma_client, payload, touched_by)
|
||||
|
|
|
|||
|
|
@ -4,15 +4,22 @@ Endpoints to control callbacks per team
|
|||
Use this when each team should control its own callbacks
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import json
|
||||
import traceback
|
||||
from typing import List, Optional
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.proxy._types import (
|
||||
AddTeamCallback,
|
||||
LiteLLM_AuditLogs,
|
||||
LitellmTableNames,
|
||||
ProxyErrorTypes,
|
||||
ProxyException,
|
||||
TeamCallbackMetadata,
|
||||
|
|
@ -24,6 +31,106 @@ from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
|
|||
router = APIRouter()
|
||||
|
||||
|
||||
_CALLBACK_VARS_REDACTED = "***REDACTED***"
|
||||
|
||||
|
||||
def _redact_callback_secrets(metadata: Any) -> Any:
|
||||
"""Strip secret values out of a team-metadata snapshot before audit logging.
|
||||
|
||||
Both ``team_metadata["logging"]`` (list of ``AddTeamCallback`` dicts) and
|
||||
``team_metadata["callback_settings"]["callback_vars"]`` carry provider
|
||||
credentials such as ``langfuse_secret_key``, ``langsmith_api_key``, and
|
||||
``gcs_path_service_account``. Persisting them verbatim into
|
||||
``LiteLLM_AuditLogs`` would let anyone with read access to the audit
|
||||
table harvest team callback credentials, so we replace each value with
|
||||
a fixed marker. The keys themselves are kept so the audit reader can
|
||||
still see *which* fields changed.
|
||||
"""
|
||||
if not isinstance(metadata, dict):
|
||||
return metadata
|
||||
redacted = copy.deepcopy(metadata)
|
||||
logging_entries = redacted.get("logging")
|
||||
if isinstance(logging_entries, list):
|
||||
for entry in logging_entries:
|
||||
if isinstance(entry, dict) and isinstance(entry.get("callback_vars"), dict):
|
||||
entry["callback_vars"] = {
|
||||
k: _CALLBACK_VARS_REDACTED for k in entry["callback_vars"]
|
||||
}
|
||||
callback_settings = redacted.get("callback_settings")
|
||||
if isinstance(callback_settings, dict) and isinstance(
|
||||
callback_settings.get("callback_vars"), dict
|
||||
):
|
||||
callback_settings["callback_vars"] = {
|
||||
k: _CALLBACK_VARS_REDACTED for k in callback_settings["callback_vars"]
|
||||
}
|
||||
return redacted
|
||||
|
||||
|
||||
def _log_audit_task_exception(task: "asyncio.Task[None]") -> None:
|
||||
"""Surface a fire-and-forget audit-log task failure.
|
||||
|
||||
``asyncio.create_task`` swallows exceptions silently — if the audit
|
||||
write fails (transient DB error etc.) we'd otherwise lose the row
|
||||
without any signal. Log at warning level so the operator sees there's
|
||||
a gap in the audit trail.
|
||||
"""
|
||||
if task.cancelled():
|
||||
return
|
||||
exc = task.exception()
|
||||
if exc is not None:
|
||||
verbose_proxy_logger.warning("Failed to write team-callback audit log: %s", exc)
|
||||
|
||||
|
||||
async def _emit_team_callback_audit_log(
|
||||
*,
|
||||
team_id: str,
|
||||
before_metadata: Any,
|
||||
after_metadata: Any,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_changed_by: Optional[str],
|
||||
) -> None:
|
||||
"""Emit an audit-log row for a team-callback mutation.
|
||||
|
||||
Mirrors the ``store_audit_logs``-gated pattern used in
|
||||
``team_endpoints.py``: the call is async-fire-and-forget and is a no-op
|
||||
when audit logging is not enabled on the proxy. Captured under
|
||||
``LitellmTableNames.TEAM_TABLE_NAME`` so the row co-locates with other
|
||||
team mutations in the audit table.
|
||||
|
||||
Callback secrets are redacted before serialization so the audit table
|
||||
cannot itself become a credential-harvest sink.
|
||||
"""
|
||||
if litellm.store_audit_logs is not True:
|
||||
return
|
||||
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
)
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name
|
||||
|
||||
redacted_before = _redact_callback_secrets(before_metadata)
|
||||
redacted_after = _redact_callback_secrets(after_metadata)
|
||||
|
||||
task = asyncio.create_task(
|
||||
create_audit_log_for_update(
|
||||
request_data=LiteLLM_AuditLogs(
|
||||
id=str(uuid.uuid4()),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
changed_by=litellm_changed_by
|
||||
or user_api_key_dict.user_id
|
||||
or litellm_proxy_admin_name,
|
||||
changed_by_api_key=user_api_key_dict.api_key,
|
||||
table_name=LitellmTableNames.TEAM_TABLE_NAME,
|
||||
object_id=team_id,
|
||||
action="updated",
|
||||
updated_values=json.dumps({"metadata": redacted_after}, default=str),
|
||||
before_value=json.dumps({"metadata": redacted_before}, default=str),
|
||||
)
|
||||
)
|
||||
)
|
||||
task.add_done_callback(_log_audit_task_exception)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/team/{team_id:path}/callback",
|
||||
tags=["team management"],
|
||||
|
|
@ -123,6 +230,7 @@ async def add_team_callbacks(
|
|||
param="callback_name",
|
||||
)
|
||||
|
||||
before_metadata = copy.deepcopy(team_metadata)
|
||||
team_callback_settings.append(data.model_dump())
|
||||
|
||||
team_metadata["logging"] = team_callback_settings
|
||||
|
|
@ -132,6 +240,14 @@ async def add_team_callbacks(
|
|||
where={"team_id": team_id}, data={"metadata": team_metadata_json} # type: ignore
|
||||
)
|
||||
|
||||
await _emit_team_callback_audit_log(
|
||||
team_id=team_id,
|
||||
before_metadata=before_metadata,
|
||||
after_metadata=team_metadata,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
)
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"data": new_team_row,
|
||||
|
|
@ -165,6 +281,10 @@ async def disable_team_logging(
|
|||
http_request: Request,
|
||||
team_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
litellm_changed_by: Optional[str] = Header(
|
||||
None,
|
||||
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
|
||||
),
|
||||
):
|
||||
"""
|
||||
Disable all logging callbacks for a team
|
||||
|
|
@ -198,6 +318,7 @@ async def disable_team_logging(
|
|||
|
||||
# Update team metadata to disable logging
|
||||
team_metadata = _existing_team.metadata
|
||||
before_metadata = copy.deepcopy(team_metadata)
|
||||
team_callback_settings = team_metadata.get("callback_settings", {})
|
||||
team_callback_settings_obj = TeamCallbackMetadata(**team_callback_settings)
|
||||
|
||||
|
|
@ -222,6 +343,17 @@ async def disable_team_logging(
|
|||
},
|
||||
)
|
||||
|
||||
# Disabling a team's logging callbacks is itself a logging-control
|
||||
# action — emit an audit-log row so the action remains traceable
|
||||
# even though the team's own observability is now off.
|
||||
await _emit_team_callback_audit_log(
|
||||
team_id=team_id,
|
||||
before_metadata=before_metadata,
|
||||
after_metadata=team_metadata,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
)
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"message": f"Logging disabled for team {team_id}",
|
||||
|
|
|
|||
|
|
@ -906,6 +906,9 @@ async def new_team( # noqa: PLR0915
|
|||
prisma_client,
|
||||
user_api_key_cache,
|
||||
)
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
get_audit_log_changed_by,
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
||||
|
|
@ -1174,9 +1177,11 @@ async def new_team( # noqa: PLR0915
|
|||
request_data=LiteLLM_AuditLogs(
|
||||
id=str(uuid.uuid4()),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
changed_by=litellm_changed_by
|
||||
or user_api_key_dict.user_id
|
||||
or litellm_proxy_admin_name,
|
||||
changed_by=get_audit_log_changed_by(
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
),
|
||||
changed_by_api_key=user_api_key_dict.api_key,
|
||||
table_name=LitellmTableNames.TEAM_TABLE_NAME,
|
||||
object_id=data.team_id,
|
||||
|
|
@ -1214,7 +1219,10 @@ async def _create_team_update_audit_log(
|
|||
user_api_key_dict: User API key authentication details
|
||||
litellm_proxy_admin_name: Name of the proxy admin
|
||||
"""
|
||||
from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
get_audit_log_changed_by,
|
||||
)
|
||||
|
||||
_before_value = existing_team_row.json(exclude_none=True)
|
||||
_before_value = json.dumps(_before_value, default=str)
|
||||
|
|
@ -1225,9 +1233,11 @@ async def _create_team_update_audit_log(
|
|||
request_data=LiteLLM_AuditLogs(
|
||||
id=str(uuid.uuid4()),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
changed_by=litellm_changed_by
|
||||
or user_api_key_dict.user_id
|
||||
or litellm_proxy_admin_name,
|
||||
changed_by=get_audit_log_changed_by(
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
),
|
||||
changed_by_api_key=user_api_key_dict.api_key,
|
||||
table_name=LitellmTableNames.TEAM_TABLE_NAME,
|
||||
object_id=team_id,
|
||||
|
|
@ -2003,21 +2013,34 @@ def team_member_add_duplication_check(
|
|||
async def _validate_team_member_add_permissions(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
complete_team_data: LiteLLM_TeamTable,
|
||||
data: TeamMemberAddRequest,
|
||||
) -> None:
|
||||
"""Validate if user has permission to add members to the team."""
|
||||
"""Validate if user has permission to add members to the team.
|
||||
|
||||
Standard users can self-join an *available team*, but the bypass
|
||||
must not be allowed to escalate them to ``role=admin`` or to add
|
||||
other users into the team. When access is granted via the
|
||||
available-team bypass we therefore enforce that every member in
|
||||
the request matches the caller's own ``user_id`` and is being
|
||||
added with ``role="user"``.
|
||||
"""
|
||||
if (
|
||||
hasattr(user_api_key_dict, "user_role")
|
||||
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
|
||||
and not _is_user_team_admin(
|
||||
user_api_key_dict=user_api_key_dict, team_obj=complete_team_data
|
||||
)
|
||||
and not await _is_user_org_admin_for_team(
|
||||
user_api_key_dict=user_api_key_dict, team_obj=complete_team_data
|
||||
)
|
||||
and not _is_available_team(
|
||||
team_id=complete_team_data.team_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
getattr(user_api_key_dict, "user_role", None)
|
||||
== LitellmUserRoles.PROXY_ADMIN.value
|
||||
):
|
||||
return
|
||||
if _is_user_team_admin(
|
||||
user_api_key_dict=user_api_key_dict, team_obj=complete_team_data
|
||||
):
|
||||
return
|
||||
if await _is_user_org_admin_for_team(
|
||||
user_api_key_dict=user_api_key_dict, team_obj=complete_team_data
|
||||
):
|
||||
return
|
||||
|
||||
if not _is_available_team(
|
||||
team_id=complete_team_data.team_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
|
|
@ -2029,6 +2052,34 @@ async def _validate_team_member_add_permissions(
|
|||
},
|
||||
)
|
||||
|
||||
# Available-team self-join: caller may add only themselves, only as a
|
||||
# standard user. Enforce that here so the bypass cannot be used as a
|
||||
# privilege-escalation or cross-user-injection primitive.
|
||||
members = data.member if isinstance(data.member, list) else [data.member]
|
||||
caller_user_id = getattr(user_api_key_dict, "user_id", None)
|
||||
for member in members:
|
||||
if getattr(member, "role", "user") != "user":
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": (
|
||||
"Available-team self-join cannot assign 'admin' role. "
|
||||
"Only proxy/team/org admins can add admins to a team."
|
||||
)
|
||||
},
|
||||
)
|
||||
member_user_id = getattr(member, "user_id", None)
|
||||
if not caller_user_id or not member_user_id or member_user_id != caller_user_id:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": (
|
||||
"Available-team self-join can only add the caller "
|
||||
"(user_id must match the authenticated user's user_id)."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def _process_team_members(
|
||||
data: TeamMemberAddRequest,
|
||||
|
|
@ -2384,6 +2435,7 @@ async def team_member_add(
|
|||
await _validate_team_member_add_permissions(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
complete_team_data=complete_team_data,
|
||||
data=data,
|
||||
)
|
||||
|
||||
# Validate and populate user_email/user_id for members before processing
|
||||
|
|
@ -2995,6 +3047,9 @@ async def delete_team(
|
|||
litellm_proxy_admin_name,
|
||||
prisma_client,
|
||||
)
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
get_audit_log_changed_by,
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
||||
|
|
@ -3054,9 +3109,11 @@ async def delete_team(
|
|||
request_data=LiteLLM_AuditLogs(
|
||||
id=str(uuid.uuid4()),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
changed_by=litellm_changed_by
|
||||
or user_api_key_dict.user_id
|
||||
or litellm_proxy_admin_name,
|
||||
changed_by=get_audit_log_changed_by(
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
),
|
||||
changed_by_api_key=user_api_key_dict.api_key,
|
||||
table_name=LitellmTableNames.TEAM_TABLE_NAME,
|
||||
object_id=team_id,
|
||||
|
|
@ -4697,6 +4754,8 @@ async def update_team_member_permissions(
|
|||
|
||||
complete_team_data = LiteLLM_TeamTable(**existing_team_row.model_dump())
|
||||
|
||||
# Available-team self-join must NOT grant write access to team-wide
|
||||
# permission policies; only proxy/team/org admins can update them.
|
||||
if (
|
||||
hasattr(user_api_key_dict, "user_role")
|
||||
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
|
||||
|
|
@ -4706,16 +4765,12 @@ async def update_team_member_permissions(
|
|||
and not await _is_user_org_admin_for_team(
|
||||
user_api_key_dict=user_api_key_dict, team_obj=complete_team_data
|
||||
)
|
||||
and not _is_available_team(
|
||||
team_id=complete_team_data.team_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "Call not allowed. User not proxy admin OR team admin. route={}, team_id={}".format(
|
||||
"/team/member_add",
|
||||
"/team/permissions_update",
|
||||
complete_team_data.team_id,
|
||||
)
|
||||
},
|
||||
|
|
|
|||
|
|
@ -13,7 +13,9 @@ import base64
|
|||
import hashlib
|
||||
import inspect
|
||||
import os
|
||||
import re
|
||||
import secrets
|
||||
from html import escape
|
||||
from copy import deepcopy
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
|
|
@ -27,20 +29,23 @@ from typing import (
|
|||
Union,
|
||||
cast,
|
||||
)
|
||||
from urllib.parse import urlencode, urlparse
|
||||
from urllib.parse import parse_qs, urlencode, urlparse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import httpx
|
||||
|
||||
import jwt
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
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,
|
||||
LITELLM_CLI_SOURCE_IDENTIFIER,
|
||||
LITELLM_UI_SESSION_DURATION,
|
||||
MAX_SPENDLOG_ROWS_TO_QUERY,
|
||||
MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE,
|
||||
|
|
@ -70,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 _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 (
|
||||
|
|
@ -123,6 +132,250 @@ router = APIRouter()
|
|||
# Metadata fields (token_type, expires_in, scope) are intentionally kept so
|
||||
# response convertors see the same fields in the PKCE path as in the non-PKCE path.
|
||||
_OAUTH_TOKEN_FIELDS = frozenset({"access_token", "id_token", "refresh_token"})
|
||||
_CLI_SSO_FLOW_CACHE_KEY_PREFIX = f"{CLI_SSO_SESSION_CACHE_KEY_PREFIX}:flow"
|
||||
_CLI_SSO_START_RATE_LIMIT_CACHE_KEY_PREFIX = (
|
||||
f"{_CLI_SSO_FLOW_CACHE_KEY_PREFIX}:start_rate_limit"
|
||||
)
|
||||
_CLI_SSO_START_RATE_LIMIT_WINDOW_SECONDS = 60
|
||||
_CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS = 30
|
||||
_CLI_SSO_USER_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"
|
||||
_CLI_SSO_LOGIN_ID_RE = re.compile(r"^cli-[A-Za-z0-9_-]{12,124}$")
|
||||
|
||||
|
||||
def _hash_cli_sso_secret(secret: str) -> str:
|
||||
return hashlib.sha256(secret.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _normalize_cli_sso_user_code(user_code: str) -> str:
|
||||
return "".join(ch for ch in user_code.upper() if ch.isalnum())
|
||||
|
||||
|
||||
def _generate_cli_sso_user_code() -> str:
|
||||
user_code = "".join(secrets.choice(_CLI_SSO_USER_CODE_ALPHABET) for _ in range(8))
|
||||
return f"{user_code[:4]}-{user_code[4:]}"
|
||||
|
||||
|
||||
def _get_cli_sso_flow_cache_key(login_id: str) -> str:
|
||||
return f"{_CLI_SSO_FLOW_CACHE_KEY_PREFIX}:{login_id}"
|
||||
|
||||
|
||||
def _is_valid_cli_sso_login_id(login_id: Optional[str]) -> bool:
|
||||
return isinstance(login_id, str) and bool(_CLI_SSO_LOGIN_ID_RE.fullmatch(login_id))
|
||||
|
||||
|
||||
def _get_cli_sso_start_rate_limit_cache_key(
|
||||
request: Request, use_x_forwarded_for: Optional[bool] = False
|
||||
) -> str:
|
||||
client_ip = (
|
||||
_get_request_ip_address(
|
||||
request=request, use_x_forwarded_for=use_x_forwarded_for
|
||||
)
|
||||
or "unknown"
|
||||
)
|
||||
client_ip_hash = _hash_cli_sso_secret(client_ip)
|
||||
return f"{_CLI_SSO_START_RATE_LIMIT_CACHE_KEY_PREFIX}:{client_ip_hash}"
|
||||
|
||||
|
||||
def _check_cli_sso_start_rate_limit(
|
||||
request: Request,
|
||||
cache: DualCache,
|
||||
use_x_forwarded_for: Optional[bool] = False,
|
||||
) -> None:
|
||||
rate_limit_cache_key = _get_cli_sso_start_rate_limit_cache_key(
|
||||
request=request, use_x_forwarded_for=use_x_forwarded_for
|
||||
)
|
||||
current_attempts = cache.increment_cache(
|
||||
key=rate_limit_cache_key,
|
||||
value=1,
|
||||
ttl=_CLI_SSO_START_RATE_LIMIT_WINDOW_SECONDS,
|
||||
)
|
||||
if current_attempts > _CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS:
|
||||
raise HTTPException(
|
||||
status_code=429,
|
||||
detail="Too many CLI login attempts. Try again later.",
|
||||
)
|
||||
|
||||
|
||||
def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dict:
|
||||
if not _is_valid_cli_sso_login_id(login_id):
|
||||
raise HTTPException(status_code=400, detail="Invalid CLI login session")
|
||||
|
||||
cache_key = _get_cli_sso_flow_cache_key(cast(str, login_id))
|
||||
flow = cache.get_cache(key=cache_key)
|
||||
if not isinstance(flow, dict) or "poll_secret_hash" not in flow:
|
||||
raise HTTPException(status_code=400, detail="Invalid CLI login session")
|
||||
return flow
|
||||
|
||||
|
||||
def _set_cli_sso_flow(login_id: str, cache: DualCache, flow: dict) -> None:
|
||||
cache.set_cache(
|
||||
key=_get_cli_sso_flow_cache_key(login_id),
|
||||
value=flow,
|
||||
ttl=CLI_SSO_SESSION_TTL_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
def _verify_cli_sso_poll_secret(flow: dict, poll_secret: Optional[str]) -> bool:
|
||||
expected_poll_secret_hash = flow.get("poll_secret_hash")
|
||||
if not isinstance(expected_poll_secret_hash, str) or not isinstance(
|
||||
poll_secret, str
|
||||
):
|
||||
return False
|
||||
supplied_poll_secret_hash = _hash_cli_sso_secret(poll_secret)
|
||||
return secrets.compare_digest(supplied_poll_secret_hash, expected_poll_secret_hash)
|
||||
|
||||
|
||||
def _render_cli_sso_verification_page(
|
||||
verify_url: str, browser_complete_token: str
|
||||
) -> str:
|
||||
escaped_verify_url = escape(verify_url, quote=True)
|
||||
escaped_browser_complete_token = escape(browser_complete_token, quote=True)
|
||||
return f"""
|
||||
<!doctype html>
|
||||
<html>
|
||||
<head>
|
||||
<title>LiteLLM CLI Login</title>
|
||||
<style>
|
||||
body {{
|
||||
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
|
||||
margin: 0;
|
||||
min-height: 100vh;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
background: #f8fafc;
|
||||
color: #0f172a;
|
||||
}}
|
||||
main {{
|
||||
width: min(420px, calc(100vw - 32px));
|
||||
background: #ffffff;
|
||||
border: 1px solid #e2e8f0;
|
||||
border-radius: 8px;
|
||||
padding: 28px;
|
||||
box-shadow: 0 12px 32px rgba(15, 23, 42, 0.08);
|
||||
}}
|
||||
h1 {{ font-size: 22px; margin: 0 0 12px; }}
|
||||
p {{ line-height: 1.5; margin: 0 0 18px; color: #334155; }}
|
||||
label {{ display: block; font-weight: 600; margin-bottom: 8px; }}
|
||||
input {{
|
||||
box-sizing: border-box;
|
||||
width: 100%;
|
||||
padding: 12px;
|
||||
border: 1px solid #cbd5e1;
|
||||
border-radius: 6px;
|
||||
font-size: 20px;
|
||||
letter-spacing: 0.08em;
|
||||
text-transform: uppercase;
|
||||
}}
|
||||
button {{
|
||||
margin-top: 16px;
|
||||
width: 100%;
|
||||
padding: 12px;
|
||||
border: 0;
|
||||
border-radius: 6px;
|
||||
background: #0f172a;
|
||||
color: #ffffff;
|
||||
font-weight: 600;
|
||||
cursor: pointer;
|
||||
}}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<main>
|
||||
<h1>Complete CLI Login</h1>
|
||||
<p>Enter the verification code shown in your terminal to finish this login.</p>
|
||||
<form method="post" action="{escaped_verify_url}">
|
||||
<input type="hidden" name="browser_complete_token" value="{escaped_browser_complete_token}" />
|
||||
<label for="user_code">Verification code</label>
|
||||
<input id="user_code" name="user_code" autocomplete="one-time-code" required autofocus />
|
||||
<button type="submit">Continue</button>
|
||||
</form>
|
||||
</main>
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
|
||||
|
||||
@router.post("/sso/cli/start", tags=["experimental"], include_in_schema=False)
|
||||
async def cli_sso_start(request: Request):
|
||||
from litellm.proxy.proxy_server import general_settings, user_api_key_cache
|
||||
|
||||
_check_cli_sso_start_rate_limit(
|
||||
request=request,
|
||||
cache=user_api_key_cache,
|
||||
use_x_forwarded_for=bool(
|
||||
(general_settings or {}).get("use_x_forwarded_for", False)
|
||||
),
|
||||
)
|
||||
|
||||
login_id = f"cli-{secrets.token_urlsafe(24)}"
|
||||
poll_secret = secrets.token_urlsafe(32)
|
||||
user_code = _generate_cli_sso_user_code()
|
||||
|
||||
flow = {
|
||||
"poll_secret_hash": _hash_cli_sso_secret(poll_secret),
|
||||
"user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code(user_code)),
|
||||
"sso_complete": False,
|
||||
"user_code_verified": False,
|
||||
"session_data": None,
|
||||
}
|
||||
_set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow)
|
||||
|
||||
return {
|
||||
"login_id": login_id,
|
||||
"poll_secret": poll_secret,
|
||||
"user_code": user_code,
|
||||
"expires_in": CLI_SSO_SESSION_TTL_SECONDS,
|
||||
}
|
||||
|
||||
|
||||
@router.post(
|
||||
"/sso/cli/complete/{login_id}", tags=["experimental"], include_in_schema=False
|
||||
)
|
||||
async def cli_sso_complete(request: Request, login_id: str):
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
from litellm.proxy.common_utils.html_forms.cli_sso_success import (
|
||||
render_cli_sso_success_page,
|
||||
)
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=user_api_key_cache)
|
||||
if not flow.get("sso_complete") or not flow.get("session_data"):
|
||||
raise HTTPException(status_code=400, detail="CLI login is not ready")
|
||||
|
||||
body = (await request.body()).decode("utf-8")
|
||||
form_values = parse_qs(body)
|
||||
supplied_user_code = (form_values.get("user_code") or [""])[0]
|
||||
supplied_browser_complete_token = (
|
||||
form_values.get("browser_complete_token") or [""]
|
||||
)[0]
|
||||
supplied_user_code_hash = _hash_cli_sso_secret(
|
||||
_normalize_cli_sso_user_code(supplied_user_code)
|
||||
)
|
||||
supplied_browser_complete_token_hash = _hash_cli_sso_secret(
|
||||
supplied_browser_complete_token
|
||||
)
|
||||
|
||||
expected_user_code_hash = flow.get("user_code_hash")
|
||||
if not isinstance(expected_user_code_hash, str) or not secrets.compare_digest(
|
||||
supplied_user_code_hash, expected_user_code_hash
|
||||
):
|
||||
raise HTTPException(status_code=400, detail="Invalid verification code")
|
||||
|
||||
expected_browser_complete_token_hash = flow.get("browser_complete_token_hash")
|
||||
if not isinstance(
|
||||
expected_browser_complete_token_hash, str
|
||||
) or not secrets.compare_digest(
|
||||
supplied_browser_complete_token_hash, expected_browser_complete_token_hash
|
||||
):
|
||||
raise HTTPException(status_code=400, detail="Invalid verification code")
|
||||
|
||||
flow["user_code_verified"] = True
|
||||
_set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow)
|
||||
|
||||
html_content = render_cli_sso_success_page()
|
||||
return HTMLResponse(content=html_content, status_code=200)
|
||||
|
||||
|
||||
def normalize_email(email: Optional[str]) -> Optional[str]:
|
||||
|
|
@ -333,6 +586,7 @@ async def google_login(
|
|||
from litellm.proxy.proxy_server import (
|
||||
premium_user,
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
user_custom_ui_sso_sign_in_handler,
|
||||
)
|
||||
|
||||
|
|
@ -382,14 +636,15 @@ async def google_login(
|
|||
redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso(
|
||||
request=request,
|
||||
sso_callback_route="sso/callback",
|
||||
existing_key=existing_key,
|
||||
)
|
||||
|
||||
# Store CLI key in state for OAuth flow
|
||||
if source == LITELLM_CLI_SOURCE_IDENTIFIER:
|
||||
_get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache)
|
||||
|
||||
# Store CLI login handle in state for OAuth flow
|
||||
cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state(
|
||||
source=source,
|
||||
key=key,
|
||||
existing_key=existing_key,
|
||||
)
|
||||
|
||||
# check if user defined a custom auth sso sign in handler, if yes, use it
|
||||
|
|
@ -1050,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:
|
||||
|
|
@ -1074,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],
|
||||
|
|
@ -1194,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:
|
||||
"""
|
||||
|
|
@ -1233,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,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1392,18 +1644,12 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa:
|
|||
)
|
||||
|
||||
if state and state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:"):
|
||||
# Extract the key ID and existing_key from the state
|
||||
# State format: {PREFIX}:{key}:{existing_key} or {PREFIX}:{key}
|
||||
state_parts = state.split(":", 2) # Split into max 3 parts
|
||||
# State format: {PREFIX}:{login_id}
|
||||
state_parts = state.split(":", 1)
|
||||
key_id = state_parts[1] if len(state_parts) > 1 else None
|
||||
existing_key = state_parts[2] if len(state_parts) > 2 else None
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
f"CLI SSO callback detected for key: {key_id}, existing_key: {existing_key}"
|
||||
)
|
||||
return await cli_sso_callback(
|
||||
request=request, key=key_id, existing_key=existing_key, result=result
|
||||
)
|
||||
verbose_proxy_logger.info("CLI SSO callback detected")
|
||||
return await cli_sso_callback(request=request, key=key_id, result=result)
|
||||
|
||||
# Control-plane cross-origin: read return_to from cookie.
|
||||
# Starlette's cookie_parser already handles RFC 2109 unquoting.
|
||||
|
|
@ -1424,13 +1670,10 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa:
|
|||
async def cli_sso_callback(
|
||||
request: Request,
|
||||
key: Optional[str] = None,
|
||||
existing_key: Optional[str] = None,
|
||||
result: Optional[Union[OpenID, dict]] = None,
|
||||
):
|
||||
"""CLI SSO callback - stores session info for JWT generation on polling"""
|
||||
verbose_proxy_logger.info(
|
||||
f"CLI SSO callback for key: {key}, existing_key: {existing_key}"
|
||||
)
|
||||
verbose_proxy_logger.info("CLI SSO callback")
|
||||
|
||||
from litellm.proxy.proxy_server import (
|
||||
prisma_client,
|
||||
|
|
@ -1438,11 +1681,7 @@ async def cli_sso_callback(
|
|||
user_api_key_cache,
|
||||
)
|
||||
|
||||
if not key or not key.startswith("sk-"):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Invalid key parameter. Must be a valid key ID starting with 'sk-'",
|
||||
)
|
||||
flow = _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache)
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
|
|
@ -1480,9 +1719,6 @@ async def cli_sso_callback(
|
|||
status_code=500, detail="Failed to retrieve user information from SSO"
|
||||
)
|
||||
|
||||
# Store session info in cache (10 min TTL)
|
||||
from litellm.constants import CLI_SSO_SESSION_CACHE_KEY_PREFIX
|
||||
|
||||
# Get all teams from user_info - CLI will let user select which one
|
||||
teams: List[str] = []
|
||||
if hasattr(user_info, "teams") and user_info.teams:
|
||||
|
|
@ -1523,21 +1759,25 @@ async def cli_sso_callback(
|
|||
"team_details": team_details,
|
||||
}
|
||||
|
||||
cache_key = f"{CLI_SSO_SESSION_CACHE_KEY_PREFIX}:{key}"
|
||||
user_api_key_cache.set_cache(key=cache_key, value=session_data, ttl=600)
|
||||
flow["session_data"] = session_data
|
||||
flow["sso_complete"] = True
|
||||
browser_complete_token = secrets.token_urlsafe(32)
|
||||
flow["browser_complete_token_hash"] = _hash_cli_sso_secret(
|
||||
browser_complete_token
|
||||
)
|
||||
_set_cli_sso_flow(login_id=cast(str, key), cache=user_api_key_cache, flow=flow)
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}"
|
||||
)
|
||||
|
||||
# Return success page
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
from litellm.proxy.common_utils.html_forms.cli_sso_success import (
|
||||
render_cli_sso_success_page,
|
||||
verify_url = str(request.url_for("cli_sso_complete", login_id=key))
|
||||
html_content = _render_cli_sso_verification_page(
|
||||
verify_url=verify_url,
|
||||
browser_complete_token=browser_complete_token,
|
||||
)
|
||||
|
||||
html_content = render_cli_sso_success_page()
|
||||
return HTMLResponse(content=html_content, status_code=200)
|
||||
|
||||
except Exception as e:
|
||||
|
|
@ -1548,7 +1788,11 @@ async def cli_sso_callback(
|
|||
|
||||
|
||||
@router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False)
|
||||
async def cli_poll_key(key_id: str, team_id: Optional[str] = None):
|
||||
async def cli_poll_key(
|
||||
key_id: str,
|
||||
team_id: Optional[str] = None,
|
||||
x_litellm_cli_poll_secret: Optional[str] = Header(default=None),
|
||||
):
|
||||
"""
|
||||
CLI polling endpoint - retrieves session from cache and generates JWT.
|
||||
|
||||
|
|
@ -1557,22 +1801,25 @@ async def cli_poll_key(key_id: str, team_id: Optional[str] = None):
|
|||
2. Second poll (with team_id): Generates JWT with selected team and deletes session
|
||||
|
||||
Args:
|
||||
key_id: The session key ID
|
||||
key_id: The CLI login session ID
|
||||
team_id: Optional team ID to assign to the JWT. If provided, must be one of user's teams.
|
||||
"""
|
||||
from litellm.constants import CLI_SSO_SESSION_CACHE_KEY_PREFIX
|
||||
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
if not key_id.startswith("sk-"):
|
||||
raise HTTPException(status_code=400, detail="Invalid key ID format")
|
||||
|
||||
try:
|
||||
# Look up session in cache
|
||||
cache_key = f"{CLI_SSO_SESSION_CACHE_KEY_PREFIX}:{key_id}"
|
||||
session_data = user_api_key_cache.get_cache(key=cache_key)
|
||||
flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=user_api_key_cache)
|
||||
if not _verify_cli_sso_poll_secret(
|
||||
flow=flow, poll_secret=x_litellm_cli_poll_secret
|
||||
):
|
||||
raise HTTPException(status_code=403, detail="Invalid CLI polling secret")
|
||||
|
||||
if session_data:
|
||||
if not flow.get("sso_complete") or not flow.get("user_code_verified"):
|
||||
return {"status": "pending"}
|
||||
|
||||
session_data = flow.get("session_data")
|
||||
|
||||
if isinstance(session_data, dict):
|
||||
user_teams = session_data.get("teams", [])
|
||||
user_team_details = session_data.get("team_details")
|
||||
user_id = session_data["user_id"]
|
||||
|
|
@ -1632,7 +1879,7 @@ async def cli_poll_key(key_id: str, team_id: Optional[str] = None):
|
|||
)
|
||||
|
||||
# Delete cache entry (single-use)
|
||||
user_api_key_cache.delete_cache(key=cache_key)
|
||||
user_api_key_cache.delete_cache(key=_get_cli_sso_flow_cache_key(key_id))
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
f"CLI JWT generated for user: {user_id}, team: {team_id}"
|
||||
|
|
@ -1650,6 +1897,8 @@ async def cli_poll_key(key_id: str, team_id: Optional[str] = None):
|
|||
else:
|
||||
return {"status": "pending"}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"Error polling for CLI JWT: {e}")
|
||||
raise HTTPException(
|
||||
|
|
@ -2393,20 +2642,15 @@ class SSOAuthenticationHandler:
|
|||
|
||||
This is used to authenticate through the CLI login flow.
|
||||
|
||||
The state parameter format is: {PREFIX}:{key}:{existing_key}
|
||||
- If existing_key is provided, it's included in the state
|
||||
The state parameter format is: {PREFIX}:{login_id}
|
||||
- The state parameter is used to pass data through the OAuth flow without changing the callback URL
|
||||
"""
|
||||
from litellm.constants import (
|
||||
LITELLM_CLI_SESSION_TOKEN_PREFIX,
|
||||
LITELLM_CLI_SOURCE_IDENTIFIER,
|
||||
)
|
||||
|
||||
if source == LITELLM_CLI_SOURCE_IDENTIFIER and key:
|
||||
if existing_key:
|
||||
return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}:{existing_key}"
|
||||
else:
|
||||
return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}"
|
||||
return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}"
|
||||
else:
|
||||
return None
|
||||
|
||||
|
|
|
|||
|
|
@ -21,6 +21,28 @@ from litellm.proxy._types import (
|
|||
from litellm.types.utils import StandardAuditLogPayload
|
||||
|
||||
_audit_log_callback_cache: Dict[str, CustomLogger] = {}
|
||||
ALLOW_LITELLM_CHANGED_BY_HEADER_METADATA_KEY = "allow_litellm_changed_by_header"
|
||||
|
||||
|
||||
def _allows_litellm_changed_by_header(user_api_key_dict: UserAPIKeyAuth) -> bool:
|
||||
for admin_metadata in (user_api_key_dict.metadata, user_api_key_dict.team_metadata):
|
||||
if (
|
||||
isinstance(admin_metadata, dict)
|
||||
and admin_metadata.get(ALLOW_LITELLM_CHANGED_BY_HEADER_METADATA_KEY) is True
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def get_audit_log_changed_by(
|
||||
*,
|
||||
litellm_changed_by: Optional[str],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_proxy_admin_name: Optional[str],
|
||||
) -> Optional[str]:
|
||||
if litellm_changed_by and _allows_litellm_changed_by_header(user_api_key_dict):
|
||||
return litellm_changed_by
|
||||
return user_api_key_dict.user_id or litellm_proxy_admin_name
|
||||
|
||||
|
||||
def _resolve_audit_log_callback(name: str) -> Optional[CustomLogger]:
|
||||
|
|
@ -143,8 +165,10 @@ async def create_object_audit_log(
|
|||
if _store_audit_logs is not True:
|
||||
return
|
||||
|
||||
_changed_by = (
|
||||
litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name
|
||||
_changed_by = get_audit_log_changed_by(
|
||||
litellm_changed_by=litellm_changed_by,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
)
|
||||
|
||||
await create_audit_log_for_update(
|
||||
|
|
|
|||
|
|
@ -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:
|
|||
|
||||
# Run auth by default; allow legacy public metrics only when explicitly disabled.
|
||||
if litellm.require_auth_for_metrics_endpoint is not False:
|
||||
# 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)
|
||||
|
||||
try:
|
||||
await user_api_key_auth(
|
||||
|
|
@ -89,5 +99,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)
|
||||
|
|
|
|||
|
|
@ -41,7 +41,6 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
|||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.passthrough import BasePassthroughUtils
|
||||
from litellm.proxy._types import (
|
||||
CommonProxyErrors,
|
||||
ConfigFieldInfo,
|
||||
ConfigFieldUpdate,
|
||||
LiteLLMRoutes,
|
||||
|
|
@ -2325,12 +2324,14 @@ async def _register_pass_through_endpoint(
|
|||
dependencies = None
|
||||
|
||||
if auth is not None and str(auth).lower() == "true":
|
||||
if premium_user is not True:
|
||||
raise ValueError(
|
||||
"Error Setting Authentication on Pass Through Endpoint: {}".format(
|
||||
CommonProxyErrors.not_premium_user.value
|
||||
)
|
||||
)
|
||||
# Authentication on a pass-through endpoint used to be enterprise-
|
||||
# only — which left the OSS tier with no safe configuration: the
|
||||
# default was ``auth=False`` (unauthenticated forwarder) and the
|
||||
# safe ``auth=True`` raised at startup unless the operator had a
|
||||
# license. The default is now ``True`` (safe-by-default), and
|
||||
# turning it on no longer requires a license: an unauthenticated
|
||||
# forwarder is a deployment choice the operator should be allowed
|
||||
# to make explicitly, but the safe option must always be free.
|
||||
dependencies = [Depends(user_api_key_auth)]
|
||||
if path not in LiteLLMRoutes.openai_routes.value:
|
||||
LiteLLMRoutes.openai_routes.value.append(path)
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -163,19 +163,21 @@ def rerank( # noqa: PLR0915
|
|||
|
||||
model_response = RerankResponse()
|
||||
|
||||
rerank_litellm_params = {
|
||||
"litellm_call_id": litellm_call_id,
|
||||
"proxy_server_request": proxy_server_request,
|
||||
"model_info": model_info,
|
||||
"preset_cache_key": None,
|
||||
"stream_response": {},
|
||||
**optional_params.model_dump(exclude_unset=True),
|
||||
}
|
||||
|
||||
litellm_logging_obj.update_from_kwargs(
|
||||
kwargs=kwargs,
|
||||
model=model,
|
||||
user=user,
|
||||
optional_params=dict(optional_rerank_params),
|
||||
litellm_params={
|
||||
"litellm_call_id": litellm_call_id,
|
||||
"proxy_server_request": proxy_server_request,
|
||||
"model_info": model_info,
|
||||
"preset_cache_key": None,
|
||||
"stream_response": {},
|
||||
**optional_params.model_dump(exclude_unset=True),
|
||||
},
|
||||
litellm_params=dict(rerank_litellm_params),
|
||||
custom_llm_provider=_custom_llm_provider,
|
||||
)
|
||||
|
||||
|
|
@ -214,6 +216,7 @@ def rerank( # noqa: PLR0915
|
|||
headers=headers or litellm.headers or {},
|
||||
client=client,
|
||||
model_response=model_response,
|
||||
litellm_params=rerank_litellm_params,
|
||||
)
|
||||
elif _custom_llm_provider == litellm.LlmProviders.AZURE_AI:
|
||||
api_base = (
|
||||
|
|
@ -235,6 +238,7 @@ def rerank( # noqa: PLR0915
|
|||
headers=headers or litellm.headers or {},
|
||||
client=client,
|
||||
model_response=model_response,
|
||||
litellm_params=rerank_litellm_params,
|
||||
)
|
||||
elif _custom_llm_provider == litellm.LlmProviders.INFINITY:
|
||||
# Implement Infinity rerank logic
|
||||
|
|
@ -265,6 +269,7 @@ def rerank( # noqa: PLR0915
|
|||
headers=headers or litellm.headers or {},
|
||||
client=client,
|
||||
model_response=model_response,
|
||||
litellm_params=rerank_litellm_params,
|
||||
)
|
||||
elif _custom_llm_provider == litellm.LlmProviders.TOGETHER_AI:
|
||||
# Implement Together AI rerank logic
|
||||
|
|
@ -318,6 +323,7 @@ def rerank( # noqa: PLR0915
|
|||
headers=headers or litellm.headers or {},
|
||||
client=client,
|
||||
model_response=model_response,
|
||||
litellm_params=rerank_litellm_params,
|
||||
)
|
||||
elif _custom_llm_provider == litellm.LlmProviders.NVIDIA_NIM:
|
||||
if dynamic_api_key is None:
|
||||
|
|
@ -346,6 +352,7 @@ def rerank( # noqa: PLR0915
|
|||
headers=headers or litellm.headers or {},
|
||||
client=client,
|
||||
model_response=model_response,
|
||||
litellm_params=rerank_litellm_params,
|
||||
)
|
||||
elif _custom_llm_provider == litellm.LlmProviders.BEDROCK:
|
||||
api_base = (
|
||||
|
|
@ -409,6 +416,7 @@ def rerank( # noqa: PLR0915
|
|||
headers=headers or litellm.headers or {},
|
||||
client=client,
|
||||
model_response=model_response,
|
||||
litellm_params=rerank_litellm_params,
|
||||
)
|
||||
|
||||
elif _custom_llm_provider == litellm.LlmProviders.DEEPINFRA:
|
||||
|
|
@ -442,6 +450,7 @@ def rerank( # noqa: PLR0915
|
|||
headers=headers or litellm.headers or {},
|
||||
client=client,
|
||||
model_response=model_response,
|
||||
litellm_params=rerank_litellm_params,
|
||||
)
|
||||
elif _custom_llm_provider == litellm.LlmProviders.FIREWORKS_AI:
|
||||
api_key = (
|
||||
|
|
@ -472,6 +481,7 @@ def rerank( # noqa: PLR0915
|
|||
headers=headers or litellm.headers or {},
|
||||
client=client,
|
||||
model_response=model_response,
|
||||
litellm_params=rerank_litellm_params,
|
||||
)
|
||||
elif _custom_llm_provider == litellm.LlmProviders.VOYAGE:
|
||||
api_key = (
|
||||
|
|
@ -500,6 +510,7 @@ def rerank( # noqa: PLR0915
|
|||
headers=headers or litellm.headers or {},
|
||||
client=client,
|
||||
model_response=model_response,
|
||||
litellm_params=rerank_litellm_params,
|
||||
)
|
||||
elif _custom_llm_provider == litellm.LlmProviders.WATSONX:
|
||||
credentials = IBMWatsonXMixin.get_watsonx_credentials(
|
||||
|
|
@ -527,6 +538,7 @@ def rerank( # noqa: PLR0915
|
|||
headers=headers or litellm.headers or {},
|
||||
client=client,
|
||||
model_response=model_response,
|
||||
litellm_params=rerank_litellm_params,
|
||||
)
|
||||
else:
|
||||
# Generic handler for all providers that use base_llm_http_handler
|
||||
|
|
@ -559,6 +571,7 @@ def rerank( # noqa: PLR0915
|
|||
headers=headers or litellm.headers or {},
|
||||
client=client,
|
||||
model_response=model_response,
|
||||
litellm_params=rerank_litellm_params,
|
||||
)
|
||||
|
||||
# Placeholder return
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -237,6 +237,8 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
|
|||
# Vector Store Params
|
||||
vector_store_id: Optional[str] = None
|
||||
milvus_text_field: Optional[str] = None
|
||||
milvus_db_name: Optional[str] = None
|
||||
milvus_partition_names: Optional[List[str]] = None
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
|
|
|
|||
|
|
@ -33391,6 +33391,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",
|
||||
|
|
|
|||
|
|
@ -52,7 +52,7 @@ proxy = [
|
|||
"azure-identity==1.25.2",
|
||||
"azure-storage-blob==12.28.0",
|
||||
"mcp==1.26.0",
|
||||
"litellm-proxy-extras==0.4.69",
|
||||
"litellm-proxy-extras==0.4.70",
|
||||
"litellm-enterprise==0.1.39",
|
||||
"RestrictedPython==8.1",
|
||||
"rich==13.9.4",
|
||||
|
|
|
|||
|
|
@ -812,7 +812,7 @@ def test_redact_msgs_from_logs_with_dynamic_params():
|
|||
# Assert redaction occurred
|
||||
assert _redacted_response_obj.choices[0].message.content == "redacted-by-litellm"
|
||||
|
||||
# Test Case 3: standard_callback_dynamic_params does not override litellm.turn_off_message_logging
|
||||
# Test Case 3: standard_callback_dynamic_params does not set turn_off_message_logging
|
||||
# since litellm.turn_off_message_logging is True redaction should occur
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
litellm_logging_obj.model_call_details["standard_callback_dynamic_params"] = (
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -870,7 +870,7 @@ from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
|
|||
def test_anthropic_json_mode_and_tool_call_response(
|
||||
json_mode, tool_calls, expect_null_response
|
||||
):
|
||||
result = litellm.AnthropicConfig()._transform_response_for_json_mode(
|
||||
result, _, _ = litellm.AnthropicConfig()._resolve_json_mode_non_streaming(
|
||||
json_mode=json_mode,
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
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
|
||||
|
|
@ -13,11 +13,13 @@ import logging
|
|||
import time
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.responses.main import mock_responses_api_response
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
|
||||
|
|
@ -126,17 +128,10 @@ async def test_redaction_responses_api():
|
|||
test_custom_logger = TestCustomLogger(turn_off_message_logging=True)
|
||||
litellm.callbacks = [test_custom_logger]
|
||||
|
||||
# Mock a ResponsesAPIResponse-style response
|
||||
mock_response = {
|
||||
"output": [{"text": "This is a test response"}],
|
||||
"model": "gpt-3.5-turbo",
|
||||
"usage": {"input_tokens": 5, "output_tokens": 5, "total_tokens": 10},
|
||||
}
|
||||
|
||||
response = await litellm.aresponses(
|
||||
model="gpt-3.5-turbo",
|
||||
input="hi",
|
||||
mock_response=mock_response,
|
||||
mock_response="This is a test response",
|
||||
)
|
||||
|
||||
await asyncio.sleep(1)
|
||||
|
|
@ -163,6 +158,7 @@ async def test_redaction_responses_api():
|
|||
assert (
|
||||
content_item["text"] == "redacted-by-litellm"
|
||||
), f"Expected redacted text but got: {content_item['text']}"
|
||||
assert "This is a test response" not in json.dumps(standard_logging_payload)
|
||||
print(
|
||||
"logged standard logging payload for ResponsesAPIResponse",
|
||||
json.dumps(standard_logging_payload, indent=2),
|
||||
|
|
@ -176,29 +172,36 @@ async def test_redaction_responses_api_stream():
|
|||
test_custom_logger = TestCustomLogger(turn_off_message_logging=True)
|
||||
litellm.callbacks = [test_custom_logger]
|
||||
|
||||
# Mock a ResponsesAPIResponse-style response with streaming chunks
|
||||
mock_response = [
|
||||
{
|
||||
"output": [{"text": "This"}],
|
||||
"model": "gpt-3.5-turbo",
|
||||
},
|
||||
{
|
||||
"output": [{"text": " is"}],
|
||||
"model": "gpt-3.5-turbo",
|
||||
},
|
||||
{
|
||||
"output": [{"text": " a test response"}],
|
||||
"model": "gpt-3.5-turbo",
|
||||
"usage": {"input_tokens": 5, "output_tokens": 5, "total_tokens": 10},
|
||||
},
|
||||
]
|
||||
mocked_response_payload = mock_responses_api_response(
|
||||
"This is a test response"
|
||||
).model_dump()
|
||||
|
||||
response = await litellm.aresponses(
|
||||
model="gpt-3.5-turbo",
|
||||
input="hi",
|
||||
mock_response=mock_response,
|
||||
stream=True,
|
||||
)
|
||||
async def mock_post(self, url, headers, timeout, stream=False, **kwargs):
|
||||
stream_content = (
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": mocked_response_payload,
|
||||
}
|
||||
)
|
||||
+ "\n\ndata: [DONE]\n\n"
|
||||
)
|
||||
return httpx.Response(
|
||||
status_code=200,
|
||||
content=stream_content,
|
||||
request=httpx.Request("POST", url),
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
new=mock_post,
|
||||
):
|
||||
response = await litellm.aresponses(
|
||||
model="gpt-3.5-turbo",
|
||||
input="hi",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Consume the stream
|
||||
chunks = []
|
||||
|
|
@ -445,18 +448,11 @@ async def test_disable_redaction_header_responses_api():
|
|||
test_custom_logger = TestCustomLogger()
|
||||
litellm.callbacks = [test_custom_logger]
|
||||
|
||||
# Mock a ResponsesAPIResponse-style response
|
||||
mock_response = {
|
||||
"output": [{"text": "This is a test response"}],
|
||||
"model": "gpt-3.5-turbo",
|
||||
"usage": {"input_tokens": 5, "output_tokens": 5, "total_tokens": 10},
|
||||
}
|
||||
|
||||
# Pass the header via litellm_metadata (as the proxy does for Responses API)
|
||||
response = await litellm.aresponses(
|
||||
model="gpt-3.5-turbo",
|
||||
input="hi",
|
||||
mock_response=mock_response,
|
||||
mock_response="This is a test response",
|
||||
litellm_metadata={"headers": {"litellm-disable-message-redaction": "true"}},
|
||||
)
|
||||
|
||||
|
|
@ -464,14 +460,14 @@ async def test_disable_redaction_header_responses_api():
|
|||
standard_logging_payload = test_custom_logger.logged_standard_logging_payload
|
||||
assert standard_logging_payload is not None
|
||||
|
||||
# Verify that messages are NOT redacted because the header was set
|
||||
# Verify that the direct SDK path still honors the explicit header.
|
||||
print(
|
||||
"logged standard logging payload for ResponsesAPI with disable header",
|
||||
json.dumps(standard_logging_payload, indent=2, default=str),
|
||||
)
|
||||
|
||||
# The content should NOT be redacted
|
||||
assert standard_logging_payload["response"] != {"text": "redacted-by-litellm"}
|
||||
response = standard_logging_payload["response"]
|
||||
assert response["output"][0]["content"][0]["text"] == "This is a test response"
|
||||
assert standard_logging_payload["messages"][0]["content"] == "hi"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -6,13 +6,19 @@ from httpx import AsyncClient
|
|||
from typing import Any, Optional, List, Literal
|
||||
|
||||
|
||||
# The proxy strips client-supplied `mock_response` unless the calling key or
|
||||
# team has this admin-metadata flag set. See `_UNTRUSTED_ROOT_CONTROL_FIELDS`
|
||||
# in litellm/proxy/litellm_pre_call_utils.py.
|
||||
_ALLOW_CLIENT_MOCK_METADATA = {"allow_client_mock_response": True}
|
||||
|
||||
|
||||
async def generate_key(
|
||||
session, models: Optional[List[str]] = None, team_id: Optional[str] = None
|
||||
):
|
||||
"""Helper function to generate a key with specific model access controls"""
|
||||
url = "http://0.0.0.0:4000/key/generate"
|
||||
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
||||
data = {}
|
||||
data: dict = {"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA)}
|
||||
if models is not None:
|
||||
data["models"] = models
|
||||
if team_id is not None:
|
||||
|
|
@ -25,7 +31,7 @@ async def generate_team(session, models: Optional[List[str]] = None):
|
|||
"""Helper function to generate a team with specific model access"""
|
||||
url = "http://0.0.0.0:4000/team/new"
|
||||
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
||||
data = {}
|
||||
data: dict = {"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA)}
|
||||
if models is not None:
|
||||
data["models"] = models
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
|
|
@ -111,7 +117,12 @@ async def test_model_access_update():
|
|||
|
||||
# Create initial key with restricted access
|
||||
response = await client.post(
|
||||
"/key/generate", json={"models": ["openai/gpt-4"]}, headers=headers
|
||||
"/key/generate",
|
||||
json={
|
||||
"models": ["openai/gpt-4"],
|
||||
"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA),
|
||||
},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200
|
||||
key_data = response.json()
|
||||
|
|
@ -214,7 +225,11 @@ async def test_team_model_access_update():
|
|||
# Create initial team with restricted access
|
||||
response = await client.post(
|
||||
"/team/new",
|
||||
json={"models": ["openai/gpt-4"], "name": "test-team"},
|
||||
json={
|
||||
"models": ["openai/gpt-4"],
|
||||
"name": "test-team",
|
||||
"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA),
|
||||
},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
|
@ -223,7 +238,12 @@ async def test_team_model_access_update():
|
|||
|
||||
# Generate a key for this team
|
||||
response = await client.post(
|
||||
"/key/generate", json={"team_id": team_id}, headers=headers
|
||||
"/key/generate",
|
||||
json={
|
||||
"team_id": team_id,
|
||||
"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA),
|
||||
},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200
|
||||
key = response.json()["key"]
|
||||
|
|
|
|||
|
|
@ -45,8 +45,11 @@ verbose_proxy_logger.setLevel(level=logging.DEBUG)
|
|||
|
||||
from starlette.datastructures import URL
|
||||
|
||||
from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update
|
||||
from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
get_audit_log_changed_by,
|
||||
)
|
||||
from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, UserAPIKeyAuth
|
||||
from litellm.caching.caching import DualCache
|
||||
from unittest.mock import patch, AsyncMock
|
||||
|
||||
|
|
@ -54,6 +57,119 @@ proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
|||
import json
|
||||
|
||||
|
||||
def test_get_audit_log_changed_by_prefers_authenticated_user():
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_id="authenticated-user",
|
||||
)
|
||||
|
||||
assert (
|
||||
get_audit_log_changed_by(
|
||||
litellm_changed_by="spoofed-user",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name="proxy-admin",
|
||||
)
|
||||
== "authenticated-user"
|
||||
)
|
||||
|
||||
|
||||
def test_get_audit_log_changed_by_honors_header_with_admin_opt_in():
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_id="service-account",
|
||||
metadata={"allow_litellm_changed_by_header": True},
|
||||
)
|
||||
|
||||
assert (
|
||||
get_audit_log_changed_by(
|
||||
litellm_changed_by="delegated-user",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name="proxy-admin",
|
||||
)
|
||||
== "delegated-user"
|
||||
)
|
||||
|
||||
|
||||
def test_get_audit_log_changed_by_honors_header_with_team_opt_in():
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_id="service-account",
|
||||
team_metadata={"allow_litellm_changed_by_header": True},
|
||||
)
|
||||
|
||||
assert (
|
||||
get_audit_log_changed_by(
|
||||
litellm_changed_by="delegated-user",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name="proxy-admin",
|
||||
)
|
||||
== "delegated-user"
|
||||
)
|
||||
|
||||
|
||||
def test_get_audit_log_changed_by_ignores_header_without_opt_in_when_user_id_missing():
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
||||
|
||||
assert (
|
||||
get_audit_log_changed_by(
|
||||
litellm_changed_by="spoofed-user",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name="proxy-admin",
|
||||
)
|
||||
== "proxy-admin"
|
||||
)
|
||||
|
||||
|
||||
def test_get_audit_log_changed_by_honors_header_with_opt_in_when_user_id_missing():
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
metadata={"allow_litellm_changed_by_header": True},
|
||||
)
|
||||
|
||||
assert (
|
||||
get_audit_log_changed_by(
|
||||
litellm_changed_by="delegated-user",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name="proxy-admin",
|
||||
)
|
||||
== "delegated-user"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_internal_user_audit_log_uses_changed_by_helper():
|
||||
from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_id="service-account",
|
||||
metadata={"allow_litellm_changed_by_header": True},
|
||||
)
|
||||
|
||||
with (
|
||||
patch("litellm.store_audit_logs", True),
|
||||
patch(
|
||||
"litellm.proxy.hooks.user_management_event_hooks.create_audit_log_for_update",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_create_audit_log_for_update,
|
||||
):
|
||||
await UserManagementEventHooks.create_internal_user_audit_log(
|
||||
user_id="target-user",
|
||||
action="updated",
|
||||
litellm_changed_by="delegated-user",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name="proxy-admin",
|
||||
before_value='{"before": true}',
|
||||
after_value='{"after": true}',
|
||||
)
|
||||
|
||||
request_data = mock_create_audit_log_for_update.await_args.kwargs["request_data"]
|
||||
assert request_data.changed_by == "delegated-user"
|
||||
assert request_data.changed_by_api_key == "test-key"
|
||||
assert request_data.object_id == "target-user"
|
||||
assert request_data.action == "updated"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_audit_log_for_update_premium_user():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
import os
|
||||
import sys
|
||||
from unittest import mock
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../.."))
|
||||
|
||||
|
|
@ -26,50 +25,30 @@ async def test_get_favicon_default():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_favicon_with_custom_url():
|
||||
"""Test that get_favicon fetches from a custom URL."""
|
||||
os.environ["LITELLM_FAVICON_URL"] = "https://example.com/favicon.ico"
|
||||
async def test_get_favicon_with_custom_url(monkeypatch):
|
||||
"""Test that get_favicon redirects browser-loaded custom URLs."""
|
||||
monkeypatch.setenv("LITELLM_FAVICON_URL", "https://example.com/favicon.ico")
|
||||
|
||||
mock_response = mock.Mock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.content = b"\x00\x00\x01\x00"
|
||||
mock_response.headers = {"content-type": "image/x-icon"}
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=app),
|
||||
base_url="http://testserver",
|
||||
) as ac:
|
||||
response = await ac.get("/get_favicon")
|
||||
|
||||
try:
|
||||
with mock.patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get"
|
||||
) as mock_get:
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=app),
|
||||
base_url="http://testserver",
|
||||
) as ac:
|
||||
response = await ac.get("/get_favicon")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "image/x-icon"
|
||||
finally:
|
||||
os.environ.pop("LITELLM_FAVICON_URL", None)
|
||||
assert response.status_code == 307
|
||||
assert response.headers["location"] == "https://example.com/favicon.ico"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_favicon_url_error_fallback():
|
||||
"""Test that get_favicon falls back to default on error."""
|
||||
os.environ["LITELLM_FAVICON_URL"] = "https://invalid.com/favicon.ico"
|
||||
async def test_get_favicon_remote_url_is_not_server_fetched(monkeypatch):
|
||||
"""Test that get_favicon does not validate remote URLs server-side."""
|
||||
monkeypatch.setenv("LITELLM_FAVICON_URL", "https://invalid.com/favicon.ico")
|
||||
|
||||
try:
|
||||
with mock.patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get"
|
||||
) as mock_get:
|
||||
mock_get.side_effect = httpx.ConnectError("unreachable")
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=app),
|
||||
base_url="http://testserver",
|
||||
) as ac:
|
||||
response = await ac.get("/get_favicon")
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=app),
|
||||
base_url="http://testserver",
|
||||
) as ac:
|
||||
response = await ac.get("/get_favicon")
|
||||
|
||||
assert response.status_code in [200, 404]
|
||||
finally:
|
||||
os.environ.pop("LITELLM_FAVICON_URL", None)
|
||||
assert response.status_code == 307
|
||||
assert response.headers["location"] == "https://invalid.com/favicon.ico"
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue