scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

chore(typing): clear 1.5k basedpyright Any errors across 54 files

Browse source 
Retypes the 54 highest-density reportAny/reportExplicitAny sources with real
types instead of shuffling the ceilings around: typed prisma table Protocols so
the untyped client surface stops at the query, local TypedDicts for JSON and
dict payloads, concrete chunk and logging types on the streaming and callback
surfaces, and 3-argument getattr with a Callable annotation where an SDK object
is genuinely duck-typed

No cast(), no type: ignore, no noqa, no suppression comments, and no new Any
annotations. Whole-tree basedpyright drops 1,941 errors with no rule rising
anywhere, and all three budget files are ratcheted so the cleared headroom
cannot silently grow back

Adds a GDC regression test pinning the named AttributeError that the typed
credential accessor now raises when with_gdch_audience is missing
This commit is contained in:
mateo-berri 2026-08-21 05:25:22 +00:00
parent 5290150a05
commit 64e993773d
58 changed files with 2172 additions and 823 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
basedpyright-code-budget.json
View file

@ -1,6 +1,6 @@
{
"reportAny": {
"limit": 19955
"limit": 18773
},
"reportArgumentType": {
"limit": 2566
@ -24,7 +24,7 @@
"limit": 19
},
"reportExplicitAny": {
"limit": 6049
"limit": 5774
},
"reportFunctionMemberAccess": {
"limit": 7
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5663
"limit": 5640
},
"reportMissingTypeArgument": {
"limit": 15555
"limit": 15498
},
"reportMissingTypeStubs": {
"limit": 40
@ -99,31 +99,31 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44655
"limit": 44589
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 39017
"limit": 38882
},
"reportUnknownParameterType": {
"limit": 19885
"limit": 19806
},
"reportUnknownVariableType": {
"limit": 30572
"limit": 30456
},
"reportUnnecessaryCast": {
"limit": 117
"limit": 116
},
"reportUnnecessaryComparison": {
"limit": 699
"limit": 698
},
"reportUnnecessaryContains": {
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 836
"limit": 835
},
"reportUntypedBaseClass": {
"limit": 0
@ -138,9 +138,9 @@
"limit": 139
},
"reportUnusedImport": {
"limit": 545
"limit": 544
},
"reportUnusedVariable": {
"limit": 146
"limit": 142
}
}

108
enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py
View file

@ -7,7 +7,9 @@ GET - /audit/{id} - Get audit log by id
GET - /audit - Get all audit logs
"""
from typing import Any, Dict, List, Optional
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import Final, Protocol
#### AUDIT LOGGING ####
from fastapi import APIRouter, Depends, HTTPException, Query
@ -15,6 +17,7 @@ from litellm_enterprise.types.proxy.audit_logging_endpoints import (
AuditLogResponse,
PaginatedAuditLogResponse,
)
from typing_extensions import ReadOnly, TypedDict
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -22,7 +25,44 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
router = APIRouter()
def _build_json_field_or_condition(json_key: str, value: str) -> Dict[str, Any]:
class _AuditLogFields(TypedDict):
"""Columns of the `LiteLLM_AuditLog` table, as returned by `model_dump()`."""
id: ReadOnly[str]
updated_at: ReadOnly[datetime]
changed_by: ReadOnly[str]
changed_by_api_key: ReadOnly[str]
action: ReadOnly[str]
table_name: ReadOnly[str]
object_id: ReadOnly[str]
before_value: ReadOnly[dict[str, object] | None]
updated_values: ReadOnly[dict[str, object] | None]
class _AuditLogRecord(Protocol):
"""Row of the `LiteLLM_AuditLog` table as materialised by the Prisma client."""
def model_dump(self) -> _AuditLogFields: ...
class _AuditLogTable(Protocol):
"""The `litellm_auditlog` accessor of the Prisma client."""
async def find_many(
self,
*,
where: Mapping[str, object],
order: Mapping[str, str],
skip: int,
take: int,
) -> Sequence[_AuditLogRecord]: ...
async def count(self, *, where: Mapping[str, object]) -> int: ...
async def find_unique(self, *, where: Mapping[str, str]) -> _AuditLogRecord | None: ...
def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, object]:
"""
Build an OR condition that matches a value inside a JSON column at the
given key, checking both before_value and updated_values.
@ -53,33 +93,33 @@ async def get_audit_logs(
page: int = Query(1, ge=1),
page_size: int = Query(10, ge=1, le=100),
# Filter parameters
changed_by: Optional[str] = Query(
changed_by: str | None = Query(
None, description="Filter by user or system that performed the action"
),
changed_by_api_key: Optional[str] = Query(
changed_by_api_key: str | None = Query(
None, description="Filter by API key hash that performed the action"
),
action: Optional[str] = Query(
action: str | None = Query(
None, description="Filter by action type (create, update, delete)"
),
table_name: Optional[str] = Query(
table_name: str | None = Query(
None, description="Filter by table name that was modified"
),
object_id: Optional[str] = Query(
object_id: str | None = Query(
None, description="Filter by ID of the object that was modified"
),
start_date: Optional[str] = Query(None, description="Filter logs after this date"),
end_date: Optional[str] = Query(None, description="Filter logs before this date"),
object_team_id: Optional[str] = Query(
start_date: str | None = Query(None, description="Filter logs after this date"),
end_date: str | None = Query(None, description="Filter logs before this date"),
object_team_id: str | None = Query(
None,
description="Filter by team_id present in before_value or updated_values JSON (PostgreSQL only)",
),
object_key_hash: Optional[str] = Query(
object_key_hash: str | None = Query(
None,
description="Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only)",
),
# Sorting parameters
sort_by: Optional[str] = Query(
sort_by: str | None = Query(
None,
description="Column to sort by (e.g. 'updated_at', 'action', 'table_name')",
),
@ -102,7 +142,7 @@ async def get_audit_logs(
)
# Build filter conditions
where_conditions: Dict[str, Any] = {}
where_conditions: Final[dict[str, object]] = {}
if changed_by:
where_conditions["changed_by"] = changed_by
if changed_by_api_key:
@ -114,33 +154,31 @@ async def get_audit_logs(
if object_id:
where_conditions["object_id"] = object_id
if start_date or end_date:
date_filter: Dict[str, Any] = {}
if start_date:
date_filter["gte"] = start_date
if end_date:
date_filter["lte"] = end_date
date_filter: Final[Mapping[str, str]] = {
bound: bound_value for bound, bound_value in (("gte", start_date), ("lte", end_date)) if bound_value
}
where_conditions["updated_at"] = date_filter
# JSON field filters (PostgreSQL only) — each filter is AND'd with the
# others, but checks both before_value and updated_values internally (OR).
if object_team_id:
where_conditions["AND"] = where_conditions.get("AND", []) + [
_build_json_field_or_condition("team_id", object_team_id)
]
if object_key_hash:
where_conditions["AND"] = where_conditions.get("AND", []) + [
_build_json_field_or_condition("token", object_key_hash)
if object_team_id or object_key_hash:
where_conditions["AND"] = [
_build_json_field_or_condition(json_key, json_value)
for json_key, json_value in (
("team_id", object_team_id),
("token", object_key_hash),
)
if json_value
]
# Build sort conditions
order_by: Dict[str, Any] = {}
if sort_by and isinstance(sort_by, str):
order_by[sort_by] = sort_order
else:
order_by["updated_at"] = sort_order # Default sort by updated_at
sort_column: Final[str] = sort_by if sort_by and isinstance(sort_by, str) else "updated_at"
order_by: Final[Mapping[str, str]] = {sort_column: sort_order}
audit_log_table: Final[_AuditLogTable] = prisma_client.db.litellm_auditlog
# Get paginated results
audit_logs = await prisma_client.db.litellm_auditlog.find_many(
audit_logs: Final = await audit_log_table.find_many(
where=where_conditions,
order=order_by,
skip=(page - 1) * page_size,
@ -148,8 +186,8 @@ async def get_audit_logs(
)
# Get total count for pagination
total_count = await prisma_client.db.litellm_auditlog.count(where=where_conditions)
total_pages = -(-total_count // page_size) # Ceiling division
total_count: Final = await audit_log_table.count(where=where_conditions)
total_pages: Final = -(-total_count // page_size) # Ceiling division
# Return paginated response
return PaginatedAuditLogResponse(
@ -198,8 +236,10 @@ async def get_audit_log_by_id(
detail={"message": CommonProxyErrors.db_not_connected_error.value},
)
audit_log_table: Final[_AuditLogTable] = prisma_client.db.litellm_auditlog
# Get the audit log by ID
audit_log = await prisma_client.db.litellm_auditlog.find_unique(where={"id": id})
audit_log: Final = await audit_log_table.find_unique(where={"id": id})
if audit_log is None:
raise HTTPException(

79
litellm/_lazy_imports.py
View file

@ -17,8 +17,14 @@ until they're actually needed.
import importlib
import sys
from collections.abc import Callable
from typing import Any, Final, cast
from collections.abc import Callable, Mapping
from typing import TYPE_CHECKING, Any, Final, cast
if TYPE_CHECKING:
import httpx
import tiktoken
from .caching.llm_caching_handler import LLMClientCache as LLMClientCacheType
# Import all the data structures that define what can be lazy-loaded
# These are just lists of names and maps of where to find them
@ -54,7 +60,7 @@ from ._lazy_imports_registry import (
)
def get_litellm_globals() -> dict:
def get_litellm_globals() -> dict[str, object]:
"""
Get the globals dictionary of the litellm module.
@ -64,7 +70,7 @@ def get_litellm_globals() -> dict:
return sys.modules["litellm"].__dict__
def _get_utils_globals() -> dict:
def _get_utils_globals() -> dict[str, object]:
"""
Get the globals dictionary of the utils module.
@ -74,14 +80,19 @@ def _get_utils_globals() -> dict:
return sys.modules["litellm.utils"].__dict__
def _get_module_level_client_timeout(litellm_globals: Mapping[str, Any]) -> "float | httpx.Timeout | None":
"""Read the configured `litellm.request_timeout` used for the module level http clients."""
return litellm_globals.get("request_timeout")
# These are special lazy loaders for things that are used internally
# They're separate from the main lazy import system because they have specific use cases
# Lazy loader for default encoding - avoids importing heavy tiktoken library at startup
_default_encoding: Any | None = None
_default_encoding: "tiktoken.Encoding | None" = None
def _get_default_encoding() -> Any:
def _get_default_encoding() -> "tiktoken.Encoding":
"""
Lazily load and cache the default OpenAI encoding.
@ -100,10 +111,10 @@ def _get_default_encoding() -> Any:
# Lazy loader for get_modified_max_tokens to avoid importing token_counter at module import time
_get_modified_max_tokens_func: Any | None = None
_get_modified_max_tokens_func: Callable[..., int | None] | None = None
def _get_modified_max_tokens() -> Any:
def _get_modified_max_tokens() -> Callable[..., int | None]:
"""
Lazily load and cache the get_modified_max_tokens function.
@ -124,10 +135,10 @@ def _get_modified_max_tokens() -> Any:
# Lazy loader for token_counter to avoid importing token_counter module at module import time
_token_counter_new_func: Any | None = None
_token_counter_new_func: Callable[..., int] | None = None
def _get_token_counter_new() -> Any:
def _get_token_counter_new() -> Callable[..., int]:
"""
Lazily load and cache the token_counter function (aliased as token_counter_new).
@ -154,10 +165,10 @@ def _get_token_counter_new() -> Any:
# This registry maps attribute names (like "ModelResponse") to handler functions
# It's built once the first time someone accesses a lazy-loaded attribute
# Example: {"ModelResponse": _lazy_import_utils, "Cache": _lazy_import_caching, ...}
_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], Any]] | None = None
_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], object]] | None = None
def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]:
def _get_lazy_import_registry() -> dict[str, Callable[[str], object]]:
"""
Build the registry that maps attribute names to their handler functions.
@ -206,7 +217,7 @@ def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]:
return _LAZY_IMPORT_REGISTRY
def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> Any:
def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> object:
"""
Generic function that handles lazy importing for most attributes.
@ -255,7 +266,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
# Step 6: Get the actual attribute from the module
# Example: getattr(utils_module, "ModelResponse") returns the ModelResponse class
value: Final = getattr(module, attr_name)
value: Final[object] = getattr(module, attr_name)
# Step 7: Cache it so we don't have to import again next time
_globals[name] = value
@ -272,62 +283,62 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
# The registry (above) maps attribute names to these handler functions.
def _lazy_import_utils(name: str) -> Any:
def _lazy_import_utils(name: str) -> object:
"""Handler for utils module attributes (ModelResponse, token_counter, etc.)"""
return _generic_lazy_import(name, _UTILS_IMPORT_MAP, "Utils")
def _lazy_import_cost_calculator(name: str) -> Any:
def _lazy_import_cost_calculator(name: str) -> object:
"""Handler for cost calculator functions (completion_cost, cost_per_token, etc.)"""
return _generic_lazy_import(name, _COST_CALCULATOR_IMPORT_MAP, "Cost calculator")
def _lazy_import_token_counter(name: str) -> Any:
def _lazy_import_token_counter(name: str) -> object:
"""Handler for token counter utilities"""
return _generic_lazy_import(name, _TOKEN_COUNTER_IMPORT_MAP, "Token counter")
def _lazy_import_bedrock_types(name: str) -> Any:
def _lazy_import_bedrock_types(name: str) -> object:
"""Handler for Bedrock type aliases"""
return _generic_lazy_import(name, _BEDROCK_TYPES_IMPORT_MAP, "Bedrock types")
def _lazy_import_types_utils(name: str) -> Any:
def _lazy_import_types_utils(name: str) -> object:
"""Handler for types from litellm.types.utils (BudgetConfig, ImageObject, etc.)"""
return _generic_lazy_import(name, _TYPES_UTILS_IMPORT_MAP, "Types utils")
def _lazy_import_caching(name: str) -> Any:
def _lazy_import_caching(name: str) -> object:
"""Handler for caching classes (Cache, DualCache, RedisCache, etc.)"""
return _generic_lazy_import(name, _CACHING_IMPORT_MAP, "Caching")
def _lazy_import_dotprompt(name: str) -> Any:
def _lazy_import_dotprompt(name: str) -> object:
"""Handler for dotprompt integration globals"""
return _generic_lazy_import(name, _DOTPROMPT_IMPORT_MAP, "Dotprompt")
def _lazy_import_types(name: str) -> Any:
def _lazy_import_types(name: str) -> object:
"""Handler for type classes (GuardrailItem, etc.)"""
return _generic_lazy_import(name, _TYPES_IMPORT_MAP, "Types")
def _lazy_import_llm_configs(name: str) -> Any:
def _lazy_import_llm_configs(name: str) -> object:
"""Handler for LLM config classes (AnthropicConfig, OpenAILikeChatConfig, etc.)"""
return _generic_lazy_import(name, _LLM_CONFIGS_IMPORT_MAP, "LLM config")
def _lazy_import_litellm_logging(name: str) -> Any:
def _lazy_import_litellm_logging(name: str) -> object:
"""Handler for litellm_logging module (Logging, modify_integration)"""
return _generic_lazy_import(name, _LITELLM_LOGGING_IMPORT_MAP, "Litellm logging")
def _lazy_import_llm_provider_logic(name: str) -> Any:
def _lazy_import_llm_provider_logic(name: str) -> object:
"""Handler for LLM provider logic functions (get_llm_provider, etc.)"""
return _generic_lazy_import(name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic")
def _lazy_import_utils_module(name: str) -> Any:
def _lazy_import_utils_module(name: str) -> object:
"""
Handler for utils module lazy imports.
@ -355,7 +366,7 @@ def _lazy_import_utils_module(name: str) -> Any:
module = importlib.import_module(module_path)
# Get the actual attribute from the module
value: Final = getattr(module, attr_name)
value: Final[object] = getattr(module, attr_name)
# Cache it so we don't have to import again next time
_globals[name] = value
@ -370,7 +381,7 @@ def _lazy_import_utils_module(name: str) -> Any:
# These handlers have custom logic that doesn't fit the generic pattern
def _lazy_import_llm_client_cache(name: str) -> Any:
def _lazy_import_llm_client_cache(name: str) -> object:
"""
Handler for LLM client cache - has special logic for singleton instance.
@ -387,7 +398,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any:
# Import the class
module: Final = importlib.import_module("litellm.caching.llm_caching_handler")
LLMClientCache: Final = getattr(module, "LLMClientCache")
LLMClientCache: Final[type[LLMClientCacheType]] = getattr(module, "LLMClientCache")
# If they want the class itself, return it
if name == "LLMClientCache":
@ -403,7 +414,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any:
raise AttributeError(f"LLM client cache lazy import: unknown attribute {name!r}")
def _lazy_import_http_handlers(name: str) -> Any:
def _lazy_import_http_handlers(name: str) -> object:
"""
Handler for HTTP clients - has special logic for creating client instances.
@ -419,8 +430,8 @@ def _lazy_import_http_handlers(name: str) -> Any:
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
# Get timeout from module config (if set)
timeout = _globals.get("request_timeout")
params: Final = {"timeout": timeout, "client_alias": "module level aclient"}
async_timeout: Final = _get_module_level_client_timeout(_globals)
params: Final = {"timeout": async_timeout, "client_alias": "module level aclient"}
# Create the client instance
provider_id: Final = cast(Any, "litellm_module_level_client")
@ -437,8 +448,8 @@ def _lazy_import_http_handlers(name: str) -> Any:
# Create a sync HTTP client
from litellm.llms.custom_httpx.http_handler import HTTPHandler
timeout = _globals.get("request_timeout")
sync_client: Final = HTTPHandler(timeout=timeout)
sync_timeout: Final = _get_module_level_client_timeout(_globals)
sync_client: Final = HTTPHandler(timeout=sync_timeout)
# Cache it
_globals["module_level_client"] = sync_client

46
litellm/a2a_protocol/litellm_completion_bridge/transformation.py
View file

@ -17,12 +17,16 @@ A2A Streaming Events:
- Artifact update (kind: "artifact-update") - Content/artifact delivery
"""
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from typing import Any, Final
from typing import TYPE_CHECKING, Any, Final
from uuid import uuid4
from litellm._logging import verbose_logger
if TYPE_CHECKING:
from litellm.types.utils import Choices
class A2AStreamingContext:
"""
@ -30,7 +34,7 @@ class A2AStreamingContext:
Tracks task_id, context_id, and message accumulation.
"""
def __init__(self, request_id: str, input_message: dict[str, Any]):
def __init__(self, request_id: str, input_message: Mapping[str, object]):
self.request_id = request_id
self.task_id = str(uuid4())
self.context_id = str(uuid4())
@ -46,7 +50,7 @@ class A2ACompletionBridgeTransformation:
"""
@staticmethod
def _extract_text_from_a2a_parts(parts: list[dict[str, Any]]) -> str:
def _extract_text_from_a2a_parts(parts: Sequence[Mapping[str, object]]) -> str:
"""Extract text from A2A parts (with or without explicit ``kind``)."""
content_parts: Final[list[str]] = []
for part in parts:
@ -62,16 +66,16 @@ class A2ACompletionBridgeTransformation:
@staticmethod
def get_forward_metadata(
a2a_message: dict[str, Any],
a2a_message: Mapping[str, object],
params: dict[str, Any] | None = None,
) -> dict[str, Any] | None:
) -> dict[str, object] | None:
"""
Merge A2A metadata from MessageSendParams and the message for downstream providers.
Forwarded once on the LangGraph run payload (``metadata``), not duplicated on
each input message — see ``apply_forward_metadata_to_completion_params``.
"""
merged: Final[dict[str, Any]] = {}
merged: Final[dict[str, object]] = {}
if params and isinstance(params.get("metadata"), dict):
merged.update(params["metadata"])
message_metadata: Final = a2a_message.get("metadata")
@ -81,8 +85,8 @@ class A2ACompletionBridgeTransformation:
@staticmethod
def apply_forward_metadata_to_completion_params(
completion_params: dict[str, Any],
a2a_message: dict[str, Any],
completion_params: dict[str, object],
a2a_message: Mapping[str, object],
params: dict[str, Any] | None = None,
) -> None:
"""
@ -104,8 +108,8 @@ class A2ACompletionBridgeTransformation:
# ``extra_body.metadata`` so the configured keys remain authoritative
# and an A2A caller cannot overwrite server-set run metadata.
existing_metadata: Final = extra_body.get("metadata")
existing_dict: Final[dict[str, Any]] = existing_metadata if isinstance(existing_metadata, dict) else {}
merged_metadata: Final[dict[str, Any]] = {**forward_metadata, **existing_dict}
existing_dict: Final[dict[str, object]] = existing_metadata if isinstance(existing_metadata, dict) else {}
merged_metadata: Final[dict[str, object]] = {**forward_metadata, **existing_dict}
extra_body = {**extra_body, "metadata": merged_metadata}
completion_params["extra_body"] = extra_body
@ -114,7 +118,7 @@ class A2ACompletionBridgeTransformation:
@staticmethod
def a2a_message_to_openai_messages(
a2a_message: dict[str, Any],
) -> list[dict[str, Any]]:
) -> list[dict[str, object]]:
"""
Transform an A2A message to OpenAI message format.
@ -124,8 +128,8 @@ class A2ACompletionBridgeTransformation:
Returns:
List of OpenAI-format messages
"""
role: Final = a2a_message.get("role", "user")
parts = a2a_message.get("parts", [])
role: Final[object] = a2a_message.get("role", "user")
parts: Sequence[Mapping[str, object]] = a2a_message.get("parts", [])
# Map A2A roles to OpenAI roles
openai_role = role
@ -143,7 +147,7 @@ class A2ACompletionBridgeTransformation:
# Do not attach A2A message.metadata here — the completion bridge forwards it
# once at run level via extra_body.metadata (LangGraph POST /runs/wait shape).
openai_message: Final[dict[str, Any]] = {"role": openai_role, "content": content}
openai_message: Final[dict[str, object]] = {"role": openai_role, "content": content}
verbose_logger.debug(
"A2A -> OpenAI transform: role=%s -> %s, content_length=%s", role, openai_role, len(content)
@ -155,7 +159,7 @@ class A2ACompletionBridgeTransformation:
def openai_response_to_a2a_response(
response: Any,
request_id: str | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Transform a LiteLLM ModelResponse to A2A SendMessageResponse format.
@ -169,7 +173,7 @@ class A2ACompletionBridgeTransformation:
# Extract content from response
content = ""
if hasattr(response, "choices") and response.choices:
choice: Final = response.choices[0]
choice: Final[Choices] = response.choices[0]
if hasattr(choice, "message") and choice.message:
content = choice.message.content or ""
@ -182,7 +186,7 @@ class A2ACompletionBridgeTransformation:
}
# Build A2A response
a2a_response: Final = {
a2a_response: Final[dict[str, object]] = {
"jsonrpc": "2.0",
"id": request_id,
"result": a2a_message,
@ -200,7 +204,7 @@ class A2ACompletionBridgeTransformation:
@staticmethod
def create_task_event(
ctx: A2AStreamingContext,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Create the initial task event with status 'submitted'.
@ -235,7 +239,7 @@ class A2ACompletionBridgeTransformation:
state: str,
final: bool = False,
message_text: str | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Create a status update event.
@ -245,7 +249,7 @@ class A2ACompletionBridgeTransformation:
final: Whether this is the final event
message_text: Optional message text for 'working' status
"""
status: Final[dict[str, Any]] = {
status: Final[dict[str, object]] = {
"state": state,
"timestamp": A2ACompletionBridgeTransformation._get_timestamp(),
}
@ -277,7 +281,7 @@ class A2ACompletionBridgeTransformation:
def create_artifact_update_event(
ctx: A2AStreamingContext,
text: str,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Create an artifact update event with content.

71
litellm/caching/valkey_semantic_cache.py
View file

@ -17,6 +17,7 @@ RedisSemanticCache since those are backend agnostic.
import asyncio
import hashlib
import os
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Final
@ -62,7 +63,7 @@ class ValkeySemanticCache(RedisSemanticCache):
sync_client: Redis | None = None,
async_client: AsyncRedis | None = None,
embedding_max_input_tokens: int | None = None,
**kwargs: Any,
**kwargs: object,
):
if similarity_threshold is None:
raise ValueError("similarity_threshold must be provided, passed None")
@ -115,7 +116,7 @@ class ValkeySemanticCache(RedisSemanticCache):
return hashlib.sha256(str(key).encode("utf-8")).hexdigest()
@staticmethod
def _embedding_to_bytes(embedding: list[float]) -> bytes:
def _embedding_to_bytes(embedding: Sequence[float]) -> bytes:
return pack_vector(embedding)
def _index_schema(self, dim: int) -> tuple[TagField, VectorField]:
@ -189,7 +190,9 @@ class ValkeySemanticCache(RedisSemanticCache):
def _doc_key(self, key: str) -> str:
return f"{self.key_prefix}{self._scope_tag(key)}:{uuid.uuid4()}"
def _doc_mapping(self, key: str, prompt: str, value_str: str, embedding: list[float]) -> dict:
def _doc_mapping(
self, key: str, prompt: str, value_str: str, embedding: Sequence[float]
) -> Mapping[str | bytes, str | bytes]:
return {
self.CACHE_KEY_FIELD_NAME: self._scope_tag(key),
self.PROMPT_FIELD_NAME: prompt,
@ -205,30 +208,49 @@ class ValkeySemanticCache(RedisSemanticCache):
)
return Query(query_string).return_fields(self.RESPONSE_FIELD_NAME, self.DISTANCE_FIELD_NAME).dialect(2)
async def _async_search(self, key: str, embedding: Sequence[float]) -> object:
"""Run the KNN query on the async client, stopping the untyped search surface here."""
return await self.async_client.ft(self.index_name).search(
self._knn_query(key),
query_params={"vec": self._embedding_to_bytes(embedding)},
)
@classmethod
def _first_hit(cls, search_result: Any) -> _ValkeyCacheHit | None:
docs: Final = getattr(search_result, "docs", [])
def _first_hit(cls, search_result: object) -> _ValkeyCacheHit | None:
docs: Final[Sequence[object]] = getattr(search_result, "docs", [])
if not docs:
return None
doc: Final = docs[0]
response_field: Final[object] = getattr(doc, cls.RESPONSE_FIELD_NAME)
distance_field: Final[str | bytes | float] = getattr(doc, cls.DISTANCE_FIELD_NAME)
return _ValkeyCacheHit(
response=str(getattr(doc, cls.RESPONSE_FIELD_NAME)),
distance=float(getattr(doc, cls.DISTANCE_FIELD_NAME)),
response=str(response_field),
distance=float(distance_field),
)
def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: Any) -> Any:
@staticmethod
def _record_similarity(kwargs: dict[str, Any], similarity: float) -> None:
"""Stamp the semantic-similarity score onto the request metadata carried in ``kwargs``."""
kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity
@staticmethod
def _embedding_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None:
"""The request metadata forwarded to the embedding call."""
return kwargs.get("metadata")
def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: object) -> object:
if hit is None:
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
self._record_similarity(kwargs, 0.0)
return None
similarity: Final = 1 - hit.distance
kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity
self._record_similarity(kwargs, similarity)
if similarity < self.similarity_threshold:
return None
return self._get_cache_logic(cached_response=hit.response)
def set_cache(self, key: str, value: Any, **kwargs: Any) -> None:
def set_cache(self, key: str, value: object, **kwargs: object) -> None:
print_verbose(f"Valkey semantic-cache set_cache, kwargs: {kwargs}")
try:
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
@ -247,12 +269,12 @@ class ValkeySemanticCache(RedisSemanticCache):
except Exception as e:
print_verbose(f"Error in Valkey semantic-cache set_cache: {e}")
def get_cache(self, key: str, **kwargs: Any) -> Any:
def get_cache(self, key: str, **kwargs: object) -> object:
print_verbose(f"Valkey semantic-cache get_cache, kwargs: {kwargs}")
try:
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
if prompt is None:
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
self._record_similarity(kwargs, 0.0)
return None
embedding: Final = self._get_embedding(prompt)
@ -265,9 +287,9 @@ class ValkeySemanticCache(RedisSemanticCache):
return self._resolve_hit(self._first_hit(search_result), key, **kwargs)
except Exception as e:
print_verbose(f"Error in Valkey semantic-cache get_cache: {e}")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
self._record_similarity(kwargs, 0.0)
async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> None:
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
print_verbose(f"Async Valkey semantic-cache set_cache, kwargs: {kwargs}")
try:
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
@ -275,7 +297,7 @@ class ValkeySemanticCache(RedisSemanticCache):
print_verbose("No prompt provided for semantic caching")
return
embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))
embedding: Final = await self._get_async_embedding(prompt, metadata=self._embedding_metadata(kwargs))
await self._ensure_index_async(len(embedding))
doc_key: Final = self._doc_key(key)
@ -286,31 +308,28 @@ class ValkeySemanticCache(RedisSemanticCache):
except Exception as e:
print_verbose(f"Error in async Valkey semantic-cache set_cache: {e}")
async def async_get_cache(self, key: str, **kwargs: Any) -> Any:
async def async_get_cache(self, key: str, **kwargs: object) -> object:
print_verbose(f"Async Valkey semantic-cache get_cache, kwargs: {kwargs}")
try:
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
if prompt is None:
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
self._record_similarity(kwargs, 0.0)
return None
embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))
embedding: Final = await self._get_async_embedding(prompt, metadata=self._embedding_metadata(kwargs))
await self._ensure_index_async(len(embedding))
search_result: Final = await self.async_client.ft(self.index_name).search(
self._knn_query(key),
query_params={"vec": self._embedding_to_bytes(embedding)},
)
search_result: Final[object] = await self._async_search(key, embedding)
return self._resolve_hit(self._first_hit(search_result), key, **kwargs)
except Exception as e:
print_verbose(f"Error in async Valkey semantic-cache get_cache: {e}")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
self._record_similarity(kwargs, 0.0)
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: Any) -> None:
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, object]], **kwargs: object) -> None:
try:
await asyncio.gather(*[self.async_set_cache(key, value, **kwargs) for key, value in cache_list])
except Exception as e:
print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {e}")
async def _index_info(self) -> dict:
async def _index_info(self) -> Mapping[str, object]:
return await self.async_client.ft(self.index_name).info()

49
litellm/integrations/bitbucket/bitbucket_client.py
View file

@ -4,11 +4,33 @@ BitBucket API client for fetching .prompt files from BitBucket repositories.
import base64
import urllib.parse
from typing import Any, Final
from collections.abc import Mapping, Sequence
from typing import Final
from typing_extensions import ReadOnly, TypedDict
from litellm.llms.custom_httpx.http_handler import HTTPHandler
class BitBucketSrcEntry(TypedDict, total=False):
"""One entry of a BitBucket ``src`` directory listing."""
type: ReadOnly[str]
path: ReadOnly[str]
class BitBucketSrcListing(TypedDict, total=False):
"""A page of a BitBucket ``src`` directory listing."""
values: ReadOnly[Sequence[BitBucketSrcEntry]]
class BitBucketBranchListing(TypedDict, total=False):
"""A page of a BitBucket ``refs/branches`` listing."""
values: ReadOnly[Sequence[Mapping[str, object]]]
def _sanitize_file_path(file_path: str) -> str:
"""Reject path traversal and URL-encode each path segment."""
if "#" in file_path or "?" in file_path:
@ -31,7 +53,7 @@ class BitBucketClient:
- Branch-specific file fetching
"""
def __init__(self, config: dict[str, Any]):
def __init__(self, config: Mapping[str, object]):
"""
Initialize the BitBucket client.
@ -135,16 +157,13 @@ class BitBucketClient:
response: Final = self.http_handler.get(url, headers=self.headers)
response.raise_for_status()
data: Final = response.json()
files: Final = []
data: Final[BitBucketSrcListing] = response.json()
for item in data.get("values", []):
if item.get("type") == "commit_file":
file_path = item.get("path", "")
if file_path.endswith(file_extension):
files.append(file_path)
return files
return [
file_path
for item in data.get("values", [])
if item.get("type") == "commit_file" and (file_path := item.get("path", "")).endswith(file_extension)
]
except Exception as e:
# Check if it's an HTTP error
@ -162,7 +181,7 @@ class BitBucketClient:
else:
raise Exception(f"Error listing files in '{directory_path}': {e}")
def get_repository_info(self) -> dict[str, Any]:
def get_repository_info(self) -> Mapping[str, object]:
"""
Get information about the repository.
@ -191,7 +210,7 @@ class BitBucketClient:
except Exception:
return False
def get_branches(self) -> list[dict[str, Any]]:
def get_branches(self) -> Sequence[Mapping[str, object]]:
"""
Get list of branches in the repository.
@ -204,12 +223,12 @@ class BitBucketClient:
response: Final = self.http_handler.get(url, headers=self.headers)
response.raise_for_status()
data: Final = response.json()
data: Final[BitBucketBranchListing] = response.json()
return data.get("values", [])
except Exception as e:
raise Exception(f"Failed to get branches: {e}")
def get_file_metadata(self, file_path: str) -> dict[str, Any] | None:
def get_file_metadata(self, file_path: str) -> Mapping[str, object] | None:
"""
Get metadata about a file (size, last modified, etc.).

84
litellm/integrations/compression_interception/handler.py
View file

@ -7,7 +7,10 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan.
import time
import uuid
from typing import Any, Final, cast
from collections.abc import Mapping, Sequence
from typing import Any, Final, Protocol, cast
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.compression import compress
@ -26,6 +29,19 @@ LITELLM_CONTENT_RETRIEVE_TOOL_NAME: Final = "litellm_content_retrieve"
_CACHE_TTL_SECONDS: Final = 15 * 60
class _AgenticLoopParams(TypedDict, total=False):
"""The ``agentic_loop_params`` entry the agentic loop driver records on the logging object."""
model: ReadOnly[str]
class _AgenticLoopLoggingObj(Protocol):
"""Logging object view exposing the untyped call details this handler reads."""
@property
def model_call_details(self) -> Mapping[str, _AgenticLoopParams]: ...
def _compression_savings_from_counts(
original_tokens: object, compressed_tokens: object
) -> CompressionSavingsMetadata | None:
@ -78,7 +94,7 @@ class CompressionInterceptionLogger(CustomLogger):
compression_trigger: int = 200_000,
compression_target: int | None = None,
embedding_model: str | None = None,
embedding_model_params: dict[str, Any] | None = None,
embedding_model_params: dict[str, object] | None = None,
):
super().__init__()
self.enabled = enabled
@ -101,7 +117,7 @@ class CompressionInterceptionLogger(CustomLogger):
@staticmethod
def initialize_from_proxy_config(
litellm_settings: dict[str, Any],
callback_specific_params: dict[str, Any],
callback_specific_params: Mapping[str, object],
) -> "CompressionInterceptionLogger":
compression_params: CompressionInterceptionConfig = {}
if "compression_interception_params" in litellm_settings:
@ -115,7 +131,9 @@ class CompressionInterceptionLogger(CustomLogger):
)
return CompressionInterceptionLogger.from_config_yaml(compression_params)
async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None:
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, Any], call_type: CallTypes | None
) -> dict[str, object] | None:
if not self.enabled:
return None
if call_type is not None and call_type != CallTypes.anthropic_messages:
@ -145,7 +163,7 @@ class CompressionInterceptionLogger(CustomLogger):
cache: Final = cast(dict[str, str], compressed.get("cache", {}))
skip_reason: Final = cast(str | None, compressed.get("compression_skipped_reason"))
compressed_tools: Final = cast(list[dict[str, Any]], compressed.get("tools", []))
compressed_tools: Final = cast(list[dict[str, object]], compressed.get("tools", []))
# Only mutate kwargs when compression actually produced a result.
# If compression was a no-op (below trigger, invalid tool sequence, etc.),
@ -156,7 +174,7 @@ class CompressionInterceptionLogger(CustomLogger):
kwargs["messages"] = compressed["messages"]
if compressed_tools:
kwargs["tools"] = self._merge_tools(
existing_tools=cast(list[dict[str, Any]] | None, kwargs.get("tools")),
existing_tools=cast(list[dict[str, object]] | None, kwargs.get("tools")),
compressed_tools=compressed_tools,
)
call_id = cast(str | None, kwargs.get("litellm_call_id"))
@ -189,14 +207,14 @@ class CompressionInterceptionLogger(CustomLogger):
async def async_should_run_agentic_loop(
self,
response: Any,
response: object,
model: str,
messages: list[dict],
tools: list[dict] | None,
messages: Sequence[Mapping[str, object]],
tools: Sequence[Mapping[str, object]] | None,
stream: bool,
custom_llm_provider: str,
kwargs: dict,
) -> tuple[bool, dict]:
kwargs: Mapping[str, object],
) -> tuple[bool, dict[str, object]]:
if not self.enabled:
return False, {}
if not self._has_retrieval_tool(tools):
@ -214,19 +232,19 @@ class CompressionInterceptionLogger(CustomLogger):
async def async_build_agentic_loop_plan(
self,
tools: dict,
tools: Mapping[str, object],
model: str,
messages: list[dict],
response: Any,
anthropic_messages_provider_config: Any,
anthropic_messages_optional_request_params: dict,
logging_obj: Any,
messages: list[dict[str, object]],
response: object,
anthropic_messages_provider_config: object,
anthropic_messages_optional_request_params: Mapping[str, object],
logging_obj: _AgenticLoopLoggingObj | None,
stream: bool,
kwargs: dict,
kwargs: Mapping[str, object],
) -> AgenticLoopPlan:
self._prune_expired_cache()
tool_calls: Final = cast(list[dict[str, Any]], tools.get("tool_calls", []))
thinking_blocks: Final = cast(list[dict[str, Any]], tools.get("thinking_blocks", []))
tool_calls: Final = cast(list[dict[str, object]], tools.get("tool_calls", []))
thinking_blocks: Final = cast(list[dict[str, object]], tools.get("thinking_blocks", []))
call_id: Final = self._resolve_call_id(logging_obj=logging_obj, kwargs=kwargs)
cache: Final = self._get_cache(call_id=call_id)
@ -269,7 +287,7 @@ class CompressionInterceptionLogger(CustomLogger):
full_model_name = model
if logging_obj is not None:
agentic_params: Final = logging_obj.model_call_details.get("agentic_loop_params", {})
full_model_name = cast(str, agentic_params.get("model", model))
full_model_name = agentic_params.get("model", model)
request_patch: Final = AgenticLoopRequestPatch(
model=full_model_name,
@ -304,15 +322,15 @@ class CompressionInterceptionLogger(CustomLogger):
return {}
return cache_entry[0]
def _resolve_call_id(self, logging_obj: Any, kwargs: dict[str, Any]) -> str | None:
def _resolve_call_id(self, logging_obj: _AgenticLoopLoggingObj | None, kwargs: Mapping[str, object]) -> str | None:
if logging_obj is not None:
logging_call_id: Final = getattr(logging_obj, "litellm_call_id", None)
if isinstance(logging_call_id, str) and logging_call_id:
return logging_call_id
kwargs_call_id: Final = kwargs.get("litellm_call_id")
return cast(str | None, kwargs_call_id if isinstance(kwargs_call_id, str) else None)
return kwargs_call_id if isinstance(kwargs_call_id, str) else None
def _resolve_retrieval_content(self, tool_call: dict[str, Any], cache: dict[str, str]) -> str:
def _resolve_retrieval_content(self, tool_call: Mapping[str, object], cache: Mapping[str, str]) -> str:
raw_input: Final = tool_call.get("input", {})
key = ""
if isinstance(raw_input, dict):
@ -323,7 +341,9 @@ class CompressionInterceptionLogger(CustomLogger):
return cache[key]
return f"[compressed content key '{key}' not found]"
def _extract_retrieval_tool_calls(self, response: Any) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
def _extract_retrieval_tool_calls(
self, response: object
) -> tuple[list[dict[str, object]], list[dict[str, object]]]:
if isinstance(response, dict):
content = response.get("content", [])
else:
@ -332,8 +352,8 @@ class CompressionInterceptionLogger(CustomLogger):
if not isinstance(content, list):
return [], []
tool_calls: Final[list[dict[str, Any]]] = []
thinking_blocks: Final[list[dict[str, Any]]] = []
tool_calls: Final[list[dict[str, object]]] = []
thinking_blocks: Final[list[dict[str, object]]] = []
for block in content:
if isinstance(block, dict):
@ -380,13 +400,13 @@ class CompressionInterceptionLogger(CustomLogger):
return tool_calls, thinking_blocks
def _prepare_followup_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
def _prepare_followup_kwargs(self, kwargs: Mapping[str, object]) -> dict[str, object]:
internal_keys: Final = {"litellm_logging_obj"}
return {
k: v for k, v in kwargs.items() if not k.startswith("_compression_interception") and k not in internal_keys
}
def _has_retrieval_tool(self, tools: Any) -> bool:
def _has_retrieval_tool(self, tools: object) -> bool:
if not isinstance(tools, list):
return False
for tool in tools:
@ -402,9 +422,9 @@ class CompressionInterceptionLogger(CustomLogger):
def _merge_tools(
self,
existing_tools: list[dict[str, Any]] | None,
compressed_tools: list[dict[str, Any]],
) -> list[dict[str, Any]]:
existing_tools: Sequence[Mapping[str, object]] | None,
compressed_tools: Sequence[Mapping[str, object]],
) -> list[Mapping[str, object]]:
merged: Final = list(existing_tools or [])
if self._has_retrieval_tool(merged):
return merged

28
litellm/integrations/custom_logger.py
View file

@ -2,7 +2,7 @@
# On success, logs events to Promptlayer
import re
import traceback
from collections.abc import AsyncGenerator
from collections.abc import AsyncGenerator, Sequence
from typing import TYPE_CHECKING, Any, Final, Optional
from pydantic import BaseModel
@ -103,11 +103,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
return []
callbacks: Final = AllCallbacks()
callback_info: Final = getattr(callbacks, lookup_name, None)
callback_info: Final[object] = getattr(callbacks, lookup_name, None)
if callback_info is None:
return []
params: Final = getattr(callback_info, "litellm_callback_params", None)
params: Final[Sequence[str] | None] = getattr(callback_info, "litellm_callback_params", None)
if not params:
return []
@ -783,7 +783,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
- Converting to string and then truncating the logged content catches this
2. We want to avoid modifying the original `messages`, `response`, and `error_str` in the logging payload since these are in kwargs and could be returned to the user
"""
field_value: Final = standard_logging_object.get(field_name)
field_value: Final[object] = standard_logging_object.get(field_name)
if field_value:
str_value: Final = str(field_value)
if len(str_value) > max_length:
@ -937,8 +937,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
• Keep untyped or text content.
• Recursively redact inline base64 blobs in *any* string field, at any depth.
"""
raw_messages: Final[Any] = payload.get("messages", [])
messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else []
raw_messages: Final[object] = payload.get("messages", [])
messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else []
verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages))
if messages:
@ -969,8 +969,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
• Keep untyped or text content.
• Recursively redact inline base64 blobs in *any* string field, at any depth.
"""
raw_messages: Final[Any] = payload.get("messages", [])
messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else []
raw_messages: Final[object] = payload.get("messages", [])
messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else []
verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages))
if messages:
@ -991,7 +991,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
value: Any,
depth: int = 0,
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER,
) -> Any:
) -> object:
"""Recursively redact inline base64 from any nested structure with a max recursion depth limit."""
if depth > max_depth:
verbose_logger.warning("[CustomLogger] Max recursion depth %s reached while redacting base64", max_depth)
@ -1022,16 +1022,16 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
def _process_messages(
self,
messages: list[Any],
messages: list[object],
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER,
) -> list[dict[str, Any]]:
filtered_messages: Final[list[dict[str, Any]]] = []
) -> list[dict[str, object]]:
filtered_messages: Final[list[dict[str, object]]] = []
for msg in messages:
if not isinstance(msg, dict):
continue
contents: Any = msg.get("content")
contents: object = msg.get("content")
if isinstance(contents, list):
cleaned: list[Any] = []
cleaned: list[object] = []
for c in contents:
if self._should_keep_content(content=c):
cleaned.append(self._redact_base64(value=c, max_depth=max_depth))

12
litellm/integrations/opik/opik_payload_builder/payload_builders.py
View file

@ -17,12 +17,12 @@ def build_trace_payload(
end_time: datetime,
input_data: Any,
output_data: Any,
metadata: dict[str, Any],
metadata: dict[str, object],
tags: list[str],
thread_id: str | None,
) -> types.TracePayload:
"""Build a complete trace payload."""
trace_name: Final = response_obj.get("object", "unknown type")
trace_name: Final[str] = response_obj.get("object", "unknown type")
return types.TracePayload(
project_name=project_name,
@ -47,7 +47,7 @@ def build_span_payload(
end_time: datetime,
input_data: Any,
output_data: Any,
metadata: dict[str, Any],
metadata: dict[str, object],
tags: list[str],
usage: dict[str, int],
provider: str | None = None,
@ -56,9 +56,9 @@ def build_span_payload(
"""Build a complete span payload."""
span_id: Final = utils.create_uuid7()
model: Final = response_obj.get("model", "unknown-model")
obj_type: Final = response_obj.get("object", "unknown-object")
created: Final = response_obj.get("created", 0)
model: Final[str] = response_obj.get("model", "unknown-model")
obj_type: Final[str] = response_obj.get("object", "unknown-object")
created: Final[int] = response_obj.get("created", 0)
span_name: Final = f"{model}_{obj_type}_{created}"
_logging.verbose_logger.debug("OpikLogger creating span with id %s for trace %s", span_id, trace_id)

54
litellm/integrations/prometheus_helpers/prometheus_api.py
View file

@ -4,9 +4,13 @@ Helper functions to query prometheus API
import json
import time
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta
from typing import Final
from httpx import Response
from typing_extensions import ReadOnly, TypedDict
from litellm import get_secret
from litellm._logging import verbose_logger
from litellm.llms.custom_httpx.http_handler import (
@ -19,9 +23,45 @@ PROMETHEUS_SELECTED_INSTANCE: Final[str | None] = get_secret("PROMETHEUS_SELECTE
async_http_handler: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
class _PrometheusSeries(TypedDict):
"""One time series in a Prometheus query result, with its labels and its ``[timestamp, value]`` samples."""
metric: ReadOnly[Mapping[str, str]]
values: ReadOnly[Sequence[Sequence[float | str]]]
class _PrometheusResultData(TypedDict):
"""The ``data`` envelope of a Prometheus query response."""
result: ReadOnly[Sequence[_PrometheusSeries]]
class _PrometheusQueryResponse(TypedDict):
"""The JSON body returned by the Prometheus ``/api/v1/query`` and ``/api/v1/query_range`` endpoints."""
data: ReadOnly[_PrometheusResultData]
class _DailySpend(TypedDict):
"""One day of spend, in the shape ``get_daily_spend_from_prometheus`` returns."""
date: ReadOnly[str]
spend: ReadOnly[float]
def _query_body(response: Response) -> _PrometheusQueryResponse:
"""Read the untyped JSON body of a Prometheus query response."""
return response.json()
def _query_series(response: Response) -> Sequence[_PrometheusSeries]:
"""Read the time series list out of the untyped JSON body of a Prometheus query response."""
return response.json()["data"]["result"]
async def get_metric_from_prometheus(
metric_name: str,
):
) -> Sequence[_PrometheusSeries]:
# Get the start of the current day in Unix timestamp
if PROMETHEUS_URL is None:
raise ValueError("PROMETHEUS_URL not set please set 'PROMETHEUS_URL=<>' in .env")
@ -31,13 +71,13 @@ async def get_metric_from_prometheus(
response: Final = await async_http_handler.get(
f"{PROMETHEUS_URL}/api/v1/query", params={"query": query, "time": now}
) # End of the day
_json_response: Final = response.json()
_json_response: Final = _query_body(response)
verbose_logger.debug("json response from prometheus /query api %s", _json_response)
results: Final = response.json()["data"]["result"]
results: Final = _query_series(response)
return results
async def get_fallback_metric_from_prometheus():
async def get_fallback_metric_from_prometheus() -> str:
"""
Gets fallback metrics from prometheus for the last 24 hours
"""
@ -96,7 +136,7 @@ def _quote_promql_string_literal(value: str) -> str:
return json.dumps(value, ensure_ascii=False)
async def get_daily_spend_from_prometheus(api_key: str | None):
async def get_daily_spend_from_prometheus(api_key: str | None) -> Sequence[_DailySpend]:
"""
Expected Response Format:
[
@ -133,9 +173,9 @@ async def get_daily_spend_from_prometheus(api_key: str | None):
}
response: Final = await async_http_handler.get(url, params=params)
_json_response: Final = response.json()
_json_response: Final = _query_body(response)
verbose_logger.debug("json response from prometheus /query api %s", _json_response)
results: Final = response.json()["data"]["result"]
results: Final = _query_series(response)
formatted_results: Final = []
for result in results:

16
litellm/integrations/websearch_interception/handler.py
View file

@ -10,7 +10,7 @@ import asyncio
import math
import uuid
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
from typing import TYPE_CHECKING, Any, Final, TypedDict, TypeVar, cast
from typing_extensions import ReadOnly
@ -74,6 +74,10 @@ WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: Final = "_websearch_interception_emit_native_b
# ``web_search_tool_result`` blocks to inject into the final response.
WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks"
_RESPONSE_CONTENT_FIELD: Final = "content"
_ResponseT: Final = TypeVar("_ResponseT")
class _PlanMetadataView(TypedDict):
websearch_native_blocks: Sequence[Mapping[str, object]] | None
@ -926,17 +930,17 @@ class WebSearchInterceptionLogger(CustomLogger):
)
@staticmethod
def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any:
def _inject_native_blocks(response: _ResponseT, native_blocks: Sequence[Mapping[str, object]]) -> _ResponseT:
"""Prepend native blocks to response content, dict or object form."""
if not native_blocks:
return response
if isinstance(response, dict):
existing = response.get("content") or []
response["content"] = list(native_blocks) + list(existing)
existing = response.get(_RESPONSE_CONTENT_FIELD) or []
response[_RESPONSE_CONTENT_FIELD] = list(native_blocks) + list(existing)
return response
existing = getattr(response, "content", None) or []
existing = getattr(response, _RESPONSE_CONTENT_FIELD, None) or []
try:
response.content = list(native_blocks) + list(existing)
setattr(response, _RESPONSE_CONTENT_FIELD, list(native_blocks) + list(existing))
except (AttributeError, TypeError):
# Object refused write — fall through and leave the response
# untouched rather than crash the request.

31
litellm/litellm_core_utils/llm_response_utils/response_metadata.py
View file

@ -1,6 +1,9 @@
import datetime
from collections.abc import Mapping
from typing import Any, Final
import httpx
from litellm.constants import LITELLM_DETAILED_TIMING
from litellm.litellm_core_utils.core_helpers import process_response_headers
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
@ -25,11 +28,7 @@ class ResponseMetadata:
@property
def supports_response_time(self) -> bool:
"""Check if response type supports timing metrics"""
return (
isinstance(self.result, ModelResponse)
or isinstance(self.result, EmbeddingResponse)
or isinstance(self.result, TranscriptionResponse)
)
return isinstance(self.result, (ModelResponse, EmbeddingResponse, TranscriptionResponse))
def set_hidden_params(self, logging_obj: LiteLLMLoggingObject, model: str | None, kwargs: dict) -> None:
"""Set hidden parameters on the response"""
@ -45,14 +44,14 @@ class ResponseMetadata:
result=self.result, litellm_model_name=model, router_model_id=model_id
),
"additional_headers": process_response_headers(
self._get_value_from_hidden_params("additional_headers") or {},
self._get_additional_headers_from_hidden_params() or {},
preserve_litellm_internal_headers=True,
),
"litellm_model_name": model,
}
self._update_hidden_params(new_params)
def _update_hidden_params(self, new_params: dict) -> None:
def _update_hidden_params(self, new_params: Mapping[str, object]) -> None:
"""
Update hidden params - handles when self._hidden_params is a dict or HiddenParams object
"""
@ -64,12 +63,12 @@ class ResponseMetadata:
for key, value in new_params.items():
setattr(self._hidden_params, key, value)
def _get_value_from_hidden_params(self, key: str) -> Any | None:
"""Get value from hidden params - handles when self._hidden_params is a dict or HiddenParams object"""
def _get_additional_headers_from_hidden_params(self) -> httpx.Headers | dict[str, str] | None:
"""Get `additional_headers` from hidden params - handles when self._hidden_params is a dict or HiddenParams object"""
if isinstance(self._hidden_params, dict):
return self._hidden_params.get(key, None)
return self._hidden_params.get("additional_headers", None)
elif isinstance(self._hidden_params, HiddenParams):
return getattr(self._hidden_params, key, None)
return getattr(self._hidden_params, "additional_headers", None)
def set_timing_metrics(
self,
@ -96,7 +95,7 @@ class ResponseMetadata:
#########################################################
# 2. Add LiteLLM overhead duration
#########################################################
llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms")
llm_api_duration_ms: Final[float | None] = logging_obj.model_call_details.get("llm_api_duration_ms")
if llm_api_duration_ms is not None:
overhead_ms = round(total_response_time_ms - llm_api_duration_ms, 4)
self._update_hidden_params(
@ -108,7 +107,7 @@ class ResponseMetadata:
#########################################################
# 3. Add callback processing duration
#########################################################
callback_duration_ms: Final = getattr(logging_obj, "callback_duration_ms", None)
callback_duration_ms: Final[float | None] = getattr(logging_obj, "callback_duration_ms", None)
if callback_duration_ms is not None:
self._update_hidden_params(
{
@ -136,17 +135,17 @@ class ResponseMetadata:
# 5. Detailed per-phase timing (opt-in via env var)
#########################################################
if LITELLM_DETAILED_TIMING and llm_api_duration_ms is not None:
detailed: Final[dict] = {
detailed: Final[dict[str, float]] = {
"timing_llm_api_ms": round(llm_api_duration_ms, 4),
}
# message copy time from Logging.__init__()
msg_copy_ms: Final = getattr(logging_obj, "message_copy_duration_ms", None)
msg_copy_ms: Final[float | None] = getattr(logging_obj, "message_copy_duration_ms", None)
if msg_copy_ms is not None:
detailed["timing_message_copy_ms"] = round(msg_copy_ms, 4)
# pre-processing = time from request start to LLM API call start
api_call_start: Final = logging_obj.model_call_details.get("api_call_start_time")
api_call_start: Final[datetime.datetime | None] = logging_obj.model_call_details.get("api_call_start_time")
if api_call_start is not None and start_time is not None:
pre_ms: Final = (api_call_start - start_time).total_seconds() * 1000
detailed["timing_pre_processing_ms"] = round(pre_ms, 4)

66
litellm/litellm_core_utils/model_response_utils.py
View file

@ -2,10 +2,54 @@
Utility functions for ModelResponse and ModelResponseStream objects.
"""
from typing import Any, Final
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final, Protocol
from litellm.types.utils import Delta, ModelResponseBase, ModelResponseStream
_NO_EXTRA_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
class _HasModelExtra(Protocol):
"""A Pydantic model, seen through the extra fields it collected."""
@property
def model_extra(self) -> Mapping[str, object] | None: ...
class _StreamingChoice(_HasModelExtra, Protocol):
"""The streaming-choice fields this emptiness check reads."""
@property
def finish_reason(self) -> object: ...
@property
def logprobs(self) -> object: ...
@property
def enhancements(self) -> object: ...
@property
def delta(self) -> Delta | None: ...
def _extra_fields(model: _HasModelExtra) -> Mapping[str, object]:
"""The dynamically added fields Pydantic stored on ``model``."""
return model.model_extra or _NO_EXTRA_FIELDS
def _attribute(obj: object, name: str) -> object:
"""The named attribute of ``obj``, or ``None`` when it is absent."""
attribute: Final[object] = getattr(obj, name, None)
return attribute
def _has_callable_attribute(obj: object, name: str) -> bool:
"""Whether the named attribute of ``obj`` is callable."""
attribute: Final[object] = getattr(obj, name)
return callable(attribute)
def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool:
"""
@ -41,7 +85,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool:
# Check model_extra for dynamically added fields (this is where Pydantic stores them)
if hasattr(model_response, "model_extra") and model_response.model_extra:
for extra_field_name, extra_field_value in model_response.model_extra.items():
for extra_field_name, extra_field_value in _extra_fields(model_response).items():
if _has_meaningful_content(extra_field_value):
return False
@ -57,7 +101,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool:
continue
# Check if any other field has meaningful content
model_response_value = getattr(model_response, model_response_field, None)
model_response_value = _attribute(model_response, model_response_field)
if _has_meaningful_content(model_response_value):
return False
@ -71,7 +115,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool:
return True
def _has_meaningful_content(value: Any) -> bool:
def _has_meaningful_content(value: object) -> bool:
"""
Check if a value contains meaningful content.
@ -102,7 +146,7 @@ def _has_meaningful_content(value: Any) -> bool:
return True
def _is_choice_non_empty(choice: Any) -> bool:
def _is_choice_non_empty(choice: _StreamingChoice) -> bool:
"""
Deep check if a choice contains any meaningful content.
@ -131,7 +175,7 @@ def _is_choice_non_empty(choice: Any) -> bool:
# Check model_extra for dynamically added fields on the choice
if hasattr(choice, "model_extra") and choice.model_extra:
for extra_field_name, extra_field_value in choice.model_extra.items():
for extra_field_name, extra_field_value in _extra_fields(choice).items():
# Skip certain structural fields that are just default/None placeholders
if extra_field_name == "index" and extra_field_value == 0:
continue
@ -147,7 +191,7 @@ def _is_choice_non_empty(choice: Any) -> bool:
# Skip private attributes, methods, and known empty fields
if (
attr_name.startswith("_")
or callable(getattr(choice, attr_name))
or _has_callable_attribute(choice, attr_name)
or attr_name.startswith("model_")
or attr_name
in {
@ -160,7 +204,7 @@ def _is_choice_non_empty(choice: Any) -> bool:
):
continue
attr_value = getattr(choice, attr_name, None)
attr_value = _attribute(choice, attr_name)
if _has_meaningful_content(attr_value):
return True
@ -179,7 +223,7 @@ def _is_delta_non_empty(delta: Delta) -> bool:
"""
# Check model_extra for dynamically added fields (this is where Pydantic stores them)
if hasattr(delta, "model_extra") and delta.model_extra:
for extra_field_name, extra_field_value in delta.model_extra.items():
for extra_field_name, extra_field_value in _extra_fields(delta).items():
# Even structural fields are meaningful if they have actual content
if _has_meaningful_content(extra_field_value):
return True
@ -187,10 +231,10 @@ def _is_delta_non_empty(delta: Delta) -> bool:
# Check all regular attributes of the delta object
for attr_name in dir(delta):
# Skip private attributes, methods, and Pydantic-specific fields
if attr_name.startswith("_") or callable(getattr(delta, attr_name)) or attr_name.startswith("model_"):
if attr_name.startswith("_") or _has_callable_attribute(delta, attr_name) or attr_name.startswith("model_"):
continue
attr_value = getattr(delta, attr_name, None)
attr_value = _attribute(delta, attr_name)
if _has_meaningful_content(attr_value):
return True

8
litellm/litellm_core_utils/streaming_handler.py
View file

@ -92,7 +92,7 @@ def print_verbose(print_statement: object):
@dataclass(frozen=True, slots=True)
class _ProviderChunkParsed:
response_obj: dict[str, Any]
response_obj: dict[str, object]
@dataclass(frozen=True, slots=True)
@ -1242,7 +1242,7 @@ class CustomStreamWrapper:
for key, value in anthropic_response_obj["provider_specific_fields"].items():
setattr(model_response, key, value)
response_obj = cast(dict[str, Any], anthropic_response_obj)
response_obj = cast(dict[str, object], anthropic_response_obj)
elif self.model == "replicate" or self.custom_llm_provider == "replicate":
response_obj = self.handle_replicate_chunk(chunk)
completion_obj["content"] = response_obj["text"]
@ -1398,7 +1398,7 @@ class CustomStreamWrapper:
if not isinstance(chunk, str):
raise ValueError(f"chunk is not a string: {chunk}")
response_obj = cast(
dict[str, Any],
dict[str, object],
litellm.CodestralTextCompletionConfig()._chunk_parser(chunk),
)
completion_obj["content"] = response_obj["text"]
@ -2462,7 +2462,7 @@ def calculate_total_usage(chunks: list[ModelResponse]) -> Usage:
prompt_tokens: int = 0
completion_tokens: int = 0
latest_usage_chunk = None
latest_usage_chunk: Usage | Mapping[str, int] | None = None
prompt_tokens_details: PromptTokensDetailsWrapper | None = None
completion_tokens_details: CompletionTokensDetailsWrapper | None = None
cache_creation_token_details: CacheCreationTokenDetails | None = None

25
litellm/litellm_core_utils/url_utils.py
View file

@ -20,6 +20,7 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config):
"""
import socket
from collections.abc import Sequence
from ipaddress import ip_address, ip_network
from typing import Any, Final
from urllib.parse import quote, urlparse, urlunparse
@ -44,7 +45,7 @@ class SSRFError(ValueError):
"""Raised when a URL targets a blocked network."""
def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") -> str:
def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str:
"""Percent-encode one user-controlled URL path segment.
``urllib.parse.quote(..., safe="")`` intentionally leaves RFC 3986
@ -64,7 +65,7 @@ def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") -
return quote(value_str, safe="")
def encode_url_path_segments(value: Any, *, field_name: str = "path") -> str:
def encode_url_path_segments(value: object, *, field_name: str = "path") -> str:
"""Percent-encode a user-controlled URL path made of multiple segments.
Empty segments are rejected, so leading, trailing, or consecutive slashes
@ -77,9 +78,9 @@ def encode_url_path_segments(value: Any, *, field_name: str = "path") -> str:
if value_str == "":
raise ValueError(f"{field_name} is required")
encoded_segments: Final = []
for segment in value_str.split("/"):
encoded_segments.append(encode_url_path_segment(segment, field_name=field_name))
encoded_segments: Final = tuple(
encode_url_path_segment(segment, field_name=field_name) for segment in value_str.split("/")
)
return "/".join(encoded_segments)
@ -202,7 +203,7 @@ def _format_host_header(hostname: str, port: int, default_port: int) -> str:
return f"{bracketed}:{port}"
def _sockaddr_host(sockaddr: Any) -> str:
def _sockaddr_host(sockaddr: Sequence[object]) -> str:
"""Return the host element of a ``getaddrinfo`` sockaddr as ``str``.
``getaddrinfo`` with ``IPPROTO_TCP`` returns AF_INET / AF_INET6 sockaddrs
@ -285,7 +286,7 @@ def validate_url(url: str) -> tuple[str, str]:
raise SSRFError(f"No addresses found for '{hostname}'")
if not is_allowlisted:
for family, type_, proto, canonname, sockaddr in addrinfo:
for _family, _type, _proto, _canonname, sockaddr in addrinfo:
resolved_ip = _sockaddr_host(sockaddr)
if _is_blocked_ip(resolved_ip):
raise SSRFError(
@ -363,7 +364,7 @@ def assert_same_origin(candidate_url: str, expected_url: str) -> None:
_MAX_REDIRECTS: Final = 10
def _extract_redirect_url(response: Any, request_url: str) -> str:
def _extract_redirect_url(response: httpx.Response, request_url: str) -> str:
"""Extract and resolve the redirect target from a response's Location header."""
location: Final = response.headers.get("location")
if not isinstance(location, str) or not location:
@ -372,7 +373,7 @@ def _extract_redirect_url(response: Any, request_url: str) -> str:
return str(httpx.URL(request_url).join(location))
def safe_get(client: Any, url: str, **kwargs: Any) -> Any:
def safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response:
"""
Fetch a user-supplied URL with SSRF protection on every redirect hop.
@ -398,7 +399,7 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> Any:
caller_headers: Final = kwargs.pop("headers", {})
for _ in range(_MAX_REDIRECTS):
validated_url, original_host = validate_url(url)
response = client.get(
response: httpx.Response = client.get(
validated_url,
headers={**caller_headers, "Host": original_host},
follow_redirects=False,
@ -412,7 +413,7 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> Any:
raise SSRFError("Too many redirects")
async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any:
async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response:
"""Async version of safe_get."""
if not getattr(litellm, "user_url_validation", True):
kwargs.setdefault("follow_redirects", True)
@ -421,7 +422,7 @@ async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any:
caller_headers: Final = kwargs.pop("headers", {})
for _ in range(_MAX_REDIRECTS):
validated_url, original_host = validate_url(url)
response = await client.get(
response: httpx.Response = await client.get(
validated_url,
headers={**caller_headers, "Host": original_host},
follow_redirects=False,

59
litellm/llms/anthropic/batches/transformation.py
View file

@ -1,9 +1,11 @@
import json
import time
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Literal, cast
import httpx
from httpx import Headers, Response
from typing_extensions import ReadOnly, TypedDict
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
@ -19,6 +21,29 @@ else:
LoggingClass = Any
class AnthropicBatchRequestCounts(TypedDict, total=False):
"""The ``request_counts`` object of an Anthropic Message Batch."""
processing: ReadOnly[int]
succeeded: ReadOnly[int]
errored: ReadOnly[int]
canceled: ReadOnly[int]
expired: ReadOnly[int]
class AnthropicMessageBatch(TypedDict, total=False):
"""The fields of an Anthropic Message Batch that map onto an OpenAI Batch."""
id: ReadOnly[str]
processing_status: ReadOnly[str]
created_at: ReadOnly[str | None]
ended_at: ReadOnly[str | None]
expires_at: ReadOnly[str | None]
cancel_initiated_at: ReadOnly[str | None]
archived_at: ReadOnly[str | None]
request_counts: ReadOnly[AnthropicBatchRequestCounts]
class AnthropicBatchesConfig(BaseBatchesConfig):
def __init__(self):
from ..chat.transformation import AnthropicConfig
@ -83,7 +108,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
create_batch_data: CreateBatchRequest,
optional_params: dict,
litellm_params: dict,
) -> bytes | str | dict[str, Any]:
) -> bytes | str | dict[str, object]:
"""
Transform the batch creation request to Anthropic format.
@ -133,7 +158,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
batch_id: str,
optional_params: dict,
litellm_params: dict,
) -> bytes | str | dict[str, Any]:
) -> bytes | str | dict[str, object]:
"""
Transform batch retrieval request for Anthropic.
@ -152,7 +177,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
) -> LiteLLMBatch:
"""Transform Anthropic MessageBatch retrieval response to LiteLLM format."""
try:
response_data: Final = raw_response.json()
response_data: Final[AnthropicMessageBatch] = raw_response.json()
except Exception as e:
raise ValueError(f"Failed to parse Anthropic batch response: {e}")
@ -161,18 +186,20 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
processing_status: Final = response_data.get("processing_status", "in_progress")
# Map Anthropic processing_status to OpenAI status
status_mapping: dict[
str,
Literal[
"validating",
"failed",
"in_progress",
"finalizing",
"completed",
"expired",
"cancelling",
"cancelled",
],
status_mapping: Final[
Mapping[
str,
Literal[
"validating",
"failed",
"in_progress",
"finalizing",
"completed",
"expired",
"cancelling",
"cancelled",
],
]
] = {
"in_progress": "in_progress",
"canceling": "cancelling",
@ -279,7 +306,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
if not line:
continue
try:
response_json = json.loads(line)
response_json: Mapping[str, Mapping[str, dict[str, object]]] = json.loads(line)
# Update model_response with the parsed JSON
completion_response = response_json["result"]["message"]
transformed_response = self.anthropic_chat_config.transform_parsed_response(

81
litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py
View file

@ -1,8 +1,10 @@
"""Dispatch ``context_management`` edits to registered polyfill editors."""
import inspect
from collections.abc import Awaitable, Callable
from typing import Any, Final, cast
from collections.abc import Awaitable, Callable, Mapping
from typing import TYPE_CHECKING, Final, TypeAlias, TypedDict, cast
from typing_extensions import ReadOnly
from litellm._logging import verbose_logger
from litellm.types.llms.anthropic import AppliedEdit
@ -11,31 +13,54 @@ from .constants import CLEAR_TOOL_USES_EDIT_TYPE, COMPACT_EDIT_TYPE
from .editors import apply_clear_tool_uses_20250919, apply_compact_20260112
from .result import PolyfillResult
EditorFn = Callable[..., Any]
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.router import Router
_EDITOR_REGISTRY: Final[dict[str, EditorFn]] = {
AnthropicMessages: TypeAlias = list[dict[str, object]]
AnthropicSystem: TypeAlias = str | list[dict[str, object]] | None
AnthropicTools: TypeAlias = list[dict[str, object]] | None
EditSpec: TypeAlias = dict[str, object]
ContextManagementSpec: TypeAlias = EditSpec | list[EditSpec] | None
SyncEditorReturn: TypeAlias = tuple[AnthropicMessages, AppliedEdit | None]
EditorFn: TypeAlias = Callable[..., object]
class EditorKwargs(TypedDict):
"""The keyword payload every registered editor accepts."""
model: ReadOnly[str]
messages: ReadOnly[AnthropicMessages]
tools: ReadOnly[AnthropicTools]
system: ReadOnly[AnthropicSystem]
edit_spec: ReadOnly[EditSpec]
_EDITOR_REGISTRY: Final[Mapping[str, EditorFn]] = {
CLEAR_TOOL_USES_EDIT_TYPE: apply_clear_tool_uses_20250919,
COMPACT_EDIT_TYPE: apply_compact_20260112,
}
def _normalize_spec(
spec: dict[str, Any] | list[dict[str, Any]] | None,
) -> list[dict[str, Any]] | None:
def _map_openai_spec(spec: list[EditSpec]) -> EditSpec | None:
"""Translate the OpenAI list form into the Anthropic-native dict form."""
# Local import to avoid an import cycle at module load.
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
return AnthropicConfig.map_openai_context_management_to_anthropic(spec)
def _normalize_spec(spec: ContextManagementSpec) -> list[EditSpec] | None:
"""Accept Anthropic-native dict form or OpenAI list form; return edits list."""
if isinstance(spec, list):
# Local import to avoid an import cycle at module load.
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
normalized: Final = _map_openai_spec(spec) if isinstance(spec, list) else spec
spec = AnthropicConfig.map_openai_context_management_to_anthropic(spec)
edits: Final = spec.get("edits") if isinstance(spec, dict) else None
edits: Final = normalized.get("edits") if isinstance(normalized, dict) else None
if not edits or not isinstance(edits, list):
return None
return [edit for edit in edits if isinstance(edit, dict)]
def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult:
def _wrap_editor_return(raw: object, *, fallback_system: AnthropicSystem) -> PolyfillResult:
"""Coerce an editor's native return shape into a ``PolyfillResult``.
v0 sync editors (e.g. ``clear_tool_uses_20250919``) return a 2-tuple
@ -46,7 +71,7 @@ def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult:
return raw
# Legacy 2-tuple return — sync editors don't mutate ``system``, so
# carry the caller's value forward.
messages, applied = cast(tuple[list[dict[str, Any]], Any], raw)
messages, applied = cast(SyncEditorReturn, raw)
return PolyfillResult(
messages=messages,
system=fallback_system,
@ -57,13 +82,13 @@ def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult:
async def apply_context_management(
*,
model: str,
messages: list[dict[str, Any]],
tools: list[dict[str, Any]] | None,
system: Any,
context_management_spec: dict[str, Any] | list[dict[str, Any]] | None,
litellm_metadata: dict[str, Any] | None = None,
llm_router: Any = None,
user_api_key_auth: Any = None,
messages: AnthropicMessages,
tools: AnthropicTools,
system: AnthropicSystem,
context_management_spec: ContextManagementSpec,
litellm_metadata: Mapping[str, object] | None = None,
llm_router: "Router | None" = None,
user_api_key_auth: "UserAPIKeyAuth | None" = None,
) -> PolyfillResult:
"""Run edits in order; return a single ``PolyfillResult``.
@ -92,7 +117,7 @@ async def apply_context_management(
)
continue
kwargs: dict[str, Any] = {
kwargs: EditorKwargs = {
"model": model,
"messages": current_messages,
"tools": tools,
@ -102,10 +127,12 @@ async def apply_context_management(
# Only async editors accept these — passing them to sync v0 editors
# would break their signature.
if inspect.iscoroutinefunction(editor):
kwargs["litellm_metadata"] = litellm_metadata
kwargs["llm_router"] = llm_router
kwargs["user_api_key_auth"] = user_api_key_auth
raw_result = await cast(Callable[..., Awaitable[Any]], editor)(**kwargs)
raw_result = await cast(Callable[..., Awaitable[PolyfillResult]], editor)(
**kwargs,
litellm_metadata=litellm_metadata,
llm_router=llm_router,
user_api_key_auth=user_api_key_auth,
)
else:
raw_result = editor(**kwargs)

40
litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py
View file

@ -33,18 +33,18 @@ def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str,
def _build_responses_kwargs(
*,
max_tokens: int,
messages: list[dict],
messages: list[dict[str, object]],
model: str,
context_management: dict | None = None,
metadata: dict | None = None,
context_management: dict[str, object] | None = None,
metadata: dict[str, object] | None = None,
output_config: AnthropicOutputConfig | None = None,
stop_sequences: list[str] | None = None,
stream: bool | None = False,
system: str | None = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
tools: list[AllAnthropicToolsValues | dict] | None = None,
thinking: dict[str, object] | None = None,
tool_choice: dict[str, object] | None = None,
tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None,
top_k: int | None = None,
top_p: float | None = None,
output_format: AnthropicOutputSchema | None = None,
@ -134,22 +134,22 @@ class LiteLLMMessagesToResponsesAPIHandler:
@staticmethod
async def async_anthropic_messages_handler(
max_tokens: int,
messages: list[dict],
messages: list[dict[str, object]],
model: str,
context_management: dict | None = None,
metadata: dict | None = None,
context_management: dict[str, object] | None = None,
metadata: dict[str, object] | None = None,
output_config: AnthropicOutputConfig | None = None,
stop_sequences: list[str] | None = None,
stream: bool | None = False,
system: str | None = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
tools: list[AllAnthropicToolsValues | dict] | None = None,
thinking: dict[str, object] | None = None,
tool_choice: dict[str, object] | None = None,
tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None,
top_k: int | None = None,
top_p: float | None = None,
output_format: AnthropicOutputSchema | None = None,
**kwargs,
**kwargs: object,
) -> AnthropicMessagesResponse | AsyncIterator[bytes]:
responses_kwargs: Final = _build_responses_kwargs(
max_tokens=max_tokens,
@ -185,23 +185,23 @@ class LiteLLMMessagesToResponsesAPIHandler:
@staticmethod
def anthropic_messages_handler(
max_tokens: int,
messages: list[dict],
messages: list[dict[str, object]],
model: str,
context_management: dict | None = None,
metadata: dict | None = None,
context_management: dict[str, object] | None = None,
metadata: dict[str, object] | None = None,
output_config: AnthropicOutputConfig | None = None,
stop_sequences: list[str] | None = None,
stream: bool | None = False,
system: str | None = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
tools: list[AllAnthropicToolsValues | dict] | None = None,
thinking: dict[str, object] | None = None,
tool_choice: dict[str, object] | None = None,
tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None,
top_k: int | None = None,
top_p: float | None = None,
output_format: AnthropicOutputSchema | None = None,
_is_async: bool = False,
**kwargs,
**kwargs: object,
) -> (
AnthropicMessagesResponse
| AsyncIterator[bytes]

30
litellm/llms/anthropic/skills/transformation.py
View file

@ -2,9 +2,10 @@
Anthropic Skills API configuration and transformations
"""
from typing import Any, Final
from typing import Final
import httpx
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
@ -23,6 +24,25 @@ from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
class _SkillPayload(TypedDict):
"""The JSON body Anthropic returns for a single skill, before ``Skill`` validates it."""
id: ReadOnly[str]
created_at: ReadOnly[str]
source: ReadOnly[str]
updated_at: ReadOnly[str]
display_title: NotRequired[ReadOnly[str | None]]
latest_version: NotRequired[ReadOnly[str | None]]
type: NotRequired[ReadOnly[str]]
class _DeleteSkillPayload(TypedDict):
"""The JSON body Anthropic returns for a skill deletion, before ``DeleteSkillResponse`` validates it."""
id: ReadOnly[str]
type: NotRequired[ReadOnly[str]]
class AnthropicSkillsConfig(BaseSkillsAPIConfig):
"""Anthropic-specific Skills API configuration"""
@ -104,7 +124,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig):
logging_obj: LiteLLMLoggingObj,
) -> Skill:
"""Transform Anthropic response to Skill object"""
response_json: Final = raw_response.json()
response_json: Final[_SkillPayload] = raw_response.json()
verbose_logger.debug("Transforming create skill response: %s", response_json)
return Skill(**response_json)
@ -122,7 +142,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig):
url: Final = self.get_complete_url(api_base=api_base, endpoint="skills")
# Build query parameters
query_params: Final[dict[str, Any]] = {}
query_params: Final[dict[str, int | str]] = {}
if "limit" in list_params and list_params["limit"]:
query_params["limit"] = list_params["limit"]
if "page" in list_params and list_params["page"]:
@ -168,7 +188,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig):
logging_obj: LiteLLMLoggingObj,
) -> Skill:
"""Transform Anthropic response to Skill object"""
response_json: Final = raw_response.json()
response_json: Final[_SkillPayload] = raw_response.json()
verbose_logger.debug("Transforming get skill response: %s", response_json)
return Skill(**response_json)
@ -193,7 +213,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig):
logging_obj: LiteLLMLoggingObj,
) -> DeleteSkillResponse:
"""Transform Anthropic response to DeleteSkillResponse"""
response_json: Final = raw_response.json()
response_json: Final[_DeleteSkillPayload] = raw_response.json()
verbose_logger.debug("Transforming delete skill response: %s", response_json)
return DeleteSkillResponse(**response_json)

13
litellm/llms/azure/files/handler.py
View file

@ -40,6 +40,11 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
data.pop("expires_after", None)
return data
@staticmethod
def _to_openai_file_object(response: FileObject) -> OpenAIFileObject:
"""Re-wrap the SDK's file object as litellm's, carrying every field across."""
return OpenAIFileObject(**response.model_dump())
async def acreate_file(
self,
create_file_data: CreateFileRequest,
@ -48,7 +53,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
verbose_logger.debug("create_file_data=%s", create_file_data)
response = await openai_client.files.create(**self._prepare_create_file_data(create_file_data))
verbose_logger.debug("create_file_response=%s", response)
return OpenAIFileObject(**response.model_dump())
return self._to_openai_file_object(response)
def create_file(
self,
@ -61,7 +66,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
max_retries: int | None,
client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None,
litellm_params: dict | None = None,
) -> OpenAIFileObject | Coroutine[Any, Any, OpenAIFileObject]:
) -> OpenAIFileObject | Coroutine[None, None, OpenAIFileObject]:
openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client(
litellm_params=litellm_params or {},
api_key=api_key,
@ -84,7 +89,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
response: Final = cast(AzureOpenAI | OpenAI, openai_client).files.create(
**self._prepare_create_file_data(create_file_data)
)
return OpenAIFileObject(**response.model_dump())
return self._to_openai_file_object(response)
async def afile_content(
self,
@ -105,7 +110,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
api_version: str | None = None,
client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None,
litellm_params: dict | None = None,
) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]:
) -> HttpxBinaryResponseContent | Coroutine[None, None, HttpxBinaryResponseContent]:
openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client(
litellm_params=litellm_params or {},
api_key=api_key,

104
litellm/llms/bedrock/realtime/handler.py
View file

@ -7,9 +7,10 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic.
import asyncio
import contextlib
import json
from typing import Any, Final
from typing import TYPE_CHECKING, Final, Protocol
from pydantic import TypeAdapter
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import _redact_string, verbose_proxy_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
@ -18,9 +19,74 @@ from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockError
from .transformation import BedrockRealtimeConfig
if TYPE_CHECKING:
from litellm.types.realtime import RealtimeResponseTransformInput
_CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None)
class _ClientWebSocket(Protocol):
"""The client-facing websocket surface used by the Bedrock realtime bridge."""
async def send_text(self, data: str) -> None: ...
async def receive_text(self) -> str: ...
async def close(self, code: int = ..., reason: str | None = ...) -> None: ...
class _BedrockInputStream(Protocol):
"""The write half of a Bedrock bidirectional stream."""
async def send(self, chunk: object) -> None: ...
async def close(self) -> None: ...
class _BedrockPayloadPart(Protocol):
"""A single Bedrock bidirectional output payload."""
bytes_: bytes | None
class _BedrockOutputChunk(Protocol):
"""A chunk read off the read half of a Bedrock bidirectional stream."""
value: _BedrockPayloadPart | None
class _BedrockOutputStream(Protocol):
"""The read half of a Bedrock bidirectional stream."""
async def receive(self) -> _BedrockOutputChunk | None: ...
class _BedrockBidirectionalStream(Protocol):
"""The bidirectional stream returned by ``invoke_model_with_bidirectional_stream``."""
input_stream: _BedrockInputStream
async def await_output(self) -> tuple[object, _BedrockOutputStream]: ...
class _ClientSessionPayload(TypedDict, total=False):
"""The ``session`` body of a client ``session.update`` frame."""
modalities: ReadOnly[object]
class _ClientRealtimeFrame(TypedDict, total=False):
"""The fields read off a client realtime frame."""
type: ReadOnly[str]
session: ReadOnly[_ClientSessionPayload]
def _decode_client_frame(payload: str) -> _ClientRealtimeFrame:
"""Decode a client realtime frame into the fields this bridge reads."""
return json.loads(payload)
class BedrockRealtime(BaseAWSLLM):
"""Handler for Bedrock Nova Sonic realtime speech-to-speech API."""
@ -30,7 +96,7 @@ class BedrockRealtime(BaseAWSLLM):
async def async_realtime(
self,
model: str,
websocket: Any,
websocket: _ClientWebSocket,
logging_obj: LiteLLMLogging,
api_base: str | None = None,
api_key: str | None = None,
@ -46,7 +112,7 @@ class BedrockRealtime(BaseAWSLLM):
aws_sts_endpoint: str | None = None,
aws_bedrock_runtime_endpoint: str | None = None,
aws_external_id: str | None = None,
**kwargs,
**kwargs: object,
):
"""
Establish bidirectional streaming connection with Bedrock Nova Sonic.
@ -118,13 +184,16 @@ class BedrockRealtime(BaseAWSLLM):
)
bedrock_client: Final = BedrockRuntimeClient(config=config)
async def open_bidirectional_stream() -> _BedrockBidirectionalStream:
return await bedrock_client.invoke_model_with_bidirectional_stream(
InvokeModelWithBidirectionalStreamOperationInput(model_id=model)
)
transformation_config: Final = BedrockRealtimeConfig()
try:
# Initialize the bidirectional stream
bedrock_stream: Final = await bedrock_client.invoke_model_with_bidirectional_stream(
InvokeModelWithBidirectionalStreamOperationInput(model_id=model)
)
bedrock_stream: Final = await open_bidirectional_stream()
verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established")
@ -132,7 +201,7 @@ class BedrockRealtime(BaseAWSLLM):
verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect")
# Track state for transformation
session_state: Final = {
session_state: Final[RealtimeResponseTransformInput] = {
"current_output_item_id": None,
"current_response_id": None,
"current_conversation_id": None,
@ -182,11 +251,11 @@ class BedrockRealtime(BaseAWSLLM):
async def _forward_client_to_bedrock(
self,
client_ws: Any,
bedrock_stream: Any,
client_ws: _ClientWebSocket,
bedrock_stream: _BedrockBidirectionalStream,
transformation_config: BedrockRealtimeConfig,
model: str,
session_state: dict,
session_state: "RealtimeResponseTransformInput",
logging_obj: LiteLLMLogging | None = None,
):
"""Forward messages from client WebSocket to Bedrock stream."""
@ -195,10 +264,11 @@ class BedrockRealtime(BaseAWSLLM):
InvokeModelWithBidirectionalStreamInputChunk,
)
def build_input_chunk(payload: bytes) -> object:
return InvokeModelWithBidirectionalStreamInputChunk(value=BidirectionalInputPayloadPart(bytes_=payload))
async def send_to_bedrock(bedrock_message: str) -> None:
event: Final = InvokeModelWithBidirectionalStreamInputChunk(
value=BidirectionalInputPayloadPart(bytes_=bedrock_message.encode("utf-8"))
)
event: Final = build_input_chunk(bedrock_message.encode("utf-8"))
await bedrock_stream.input_stream.send(event)
verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200])
@ -223,7 +293,7 @@ class BedrockRealtime(BaseAWSLLM):
client_message_type: str | None = None
requested_modalities: list[str] | None = None
with contextlib.suppress(Exception):
parsed_client_message = json.loads(message)
parsed_client_message = _decode_client_frame(message)
client_message_type = parsed_client_message.get("type")
if client_message_type == "session.update":
requested_modalities = _CLIENT_MODALITIES_ADAPTER.validate_python(
@ -246,12 +316,12 @@ class BedrockRealtime(BaseAWSLLM):
async def _forward_bedrock_to_client(
self,
bedrock_stream: Any,
client_ws: Any,
bedrock_stream: _BedrockBidirectionalStream,
client_ws: _ClientWebSocket,
transformation_config: BedrockRealtimeConfig,
model: str,
logging_obj: LiteLLMLogging,
session_state: dict,
session_state: "RealtimeResponseTransformInput",
):
"""Forward messages from Bedrock stream to client WebSocket."""
try:

41
litellm/llms/chatgpt/chat/streaming_utils.py
View file

@ -4,7 +4,32 @@ Streaming utilities for ChatGPT provider.
Normalizes non-spec-compliant tool_call chunks from the ChatGPT backend API.
"""
from typing import Any, Final
from collections.abc import Sequence
from typing import Final, Protocol
from litellm.types.utils import Delta
class ChatGPTStreamChoice(Protocol):
"""Streaming choice as read by :class:`ChatGPTToolCallNormalizer`."""
@property
def delta(self) -> Delta | None: ...
class ChatGPTStreamChunk(Protocol):
"""Streaming chunk as read by :class:`ChatGPTToolCallNormalizer`."""
@property
def choices(self) -> Sequence[ChatGPTStreamChoice]: ...
class ChatGPTChunkStream(Protocol):
"""Sync/async chunk source wrapped by :class:`ChatGPTToolCallNormalizer`."""
def __next__(self) -> ChatGPTStreamChunk: ...
async def __anext__(self) -> ChatGPTStreamChunk: ...
class ChatGPTToolCallNormalizer:
@ -20,36 +45,36 @@ class ChatGPTToolCallNormalizer:
chunks to the consumer.
"""
def __init__(self, stream: Any):
def __init__(self, stream: ChatGPTChunkStream):
self._stream = stream
self._seen_ids: dict[str, int] = {} # tool_call_id -> assigned_index
self._next_index: int = 0
self._last_id: str | None = None # tracks which tool call the next delta belongs to
def __getattr__(self, name: str) -> Any:
def __getattr__(self, name: str) -> object:
return getattr(self._stream, name)
def __iter__(self):
def __iter__(self) -> "ChatGPTToolCallNormalizer":
return self
def __aiter__(self):
def __aiter__(self) -> "ChatGPTToolCallNormalizer":
return self
def __next__(self):
def __next__(self) -> ChatGPTStreamChunk:
while True:
chunk = next(self._stream)
result = self._normalize(chunk)
if result is not None:
return result
async def __anext__(self):
async def __anext__(self) -> ChatGPTStreamChunk:
while True:
chunk = await self._stream.__anext__()
result = self._normalize(chunk)
if result is not None:
return result
def _normalize(self, chunk: Any) -> Any:
def _normalize(self, chunk: ChatGPTStreamChunk) -> ChatGPTStreamChunk | None:
"""Fix tool_calls in the chunk. Returns None to skip duplicate chunks."""
if not chunk.choices:
return chunk

33
litellm/llms/compactifai/chat/transformation.py
View file

@ -2,13 +2,16 @@
CompactifAI chat completion transformation
"""
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
import httpx
from typing_extensions import ReadOnly, TypedDict
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.openai.common_utils import OpenAIError
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelResponse
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
@ -21,6 +24,18 @@ else:
LiteLLMLoggingObj = Any
class CompactifAIResponseFields(TypedDict, total=False):
"""The chat completion fields of a CompactifAI response body."""
id: ReadOnly[str]
choices: ReadOnly[Sequence[Mapping[str, object]]]
created: ReadOnly[int]
model: ReadOnly[str]
system_fingerprint: ReadOnly[str | None]
usage: ReadOnly[Mapping[str, object]]
object: ReadOnly[str]
class CompactifAIChatConfig(OpenAIGPTConfig):
"""
Configuration class for CompactifAI chat completions.
@ -45,11 +60,11 @@ class CompactifAIChatConfig(OpenAIGPTConfig):
raw_response: httpx.Response,
model_response: ModelResponse,
logging_obj: LiteLLMLoggingObj,
request_data: dict,
messages: list,
optional_params: dict,
litellm_params: dict,
encoding: Any,
request_data: Mapping[str, object],
messages: Sequence[AllMessageValues],
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
encoding: object,
api_key: str | None = None,
json_mode: bool | None = None,
) -> ModelResponse:
@ -79,14 +94,18 @@ class CompactifAIChatConfig(OpenAIGPTConfig):
message["content"] = tool_calls[0]["function"].get("arguments", "")
message["tool_calls"] = None
returned_response: Final = ModelResponse(**response_json)
response_fields: Final[CompactifAIResponseFields] = response_json
returned_response: Final = ModelResponse(**response_fields)
# Set model name with provider prefix
returned_response.model = f"compactifai/{model}"
return returned_response
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
def get_error_class(
self, error_message: str, status_code: int, headers: dict[str, str] | httpx.Headers
) -> BaseLLMException:
"""
Get the appropriate error class for CompactifAI errors.
Since CompactifAI is OpenAI-compatible, we use OpenAI error handling.

178
litellm/llms/custom_httpx/container_handler.py
View file

@ -6,11 +6,13 @@ endpoint defined in endpoints.json, eliminating the need for individual handler
"""
import json
from collections.abc import Coroutine
from collections.abc import Coroutine, Mapping, Sequence
from pathlib import Path
from typing import TYPE_CHECKING, Any, Final
import httpx
from pydantic import BaseModel
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
@ -32,22 +34,54 @@ if TYPE_CHECKING:
from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
class EndpointConfig(TypedDict):
"""One endpoint entry of ``litellm/containers/endpoints.json``."""
name: ReadOnly[str]
async_name: ReadOnly[str]
path: ReadOnly[str]
method: ReadOnly[str]
path_params: ReadOnly[Sequence[str]]
query_params: ReadOnly[Sequence[str]]
response_type: ReadOnly[str]
is_multipart: NotRequired[ReadOnly[bool]]
returns_binary: NotRequired[ReadOnly[bool]]
class EndpointsConfig(TypedDict):
"""The parsed ``litellm/containers/endpoints.json`` document."""
endpoints: ReadOnly[Sequence[EndpointConfig]]
class ContainerErrorDetail(TypedDict, total=False):
"""The ``error`` object of a container API error body."""
message: ReadOnly[str]
class ContainerResponseBody(TypedDict, total=False):
"""The fields this handler reads off a container API JSON body."""
error: ReadOnly[ContainerErrorDetail]
# Response type mapping
RESPONSE_TYPES: Final[dict[str, type]] = {
RESPONSE_TYPES: Final[Mapping[str, type[BaseModel]]] = {
"ContainerFileListResponse": ContainerFileListResponse,
"ContainerFileObject": ContainerFileObject,
"DeleteContainerFileResponse": DeleteContainerFileResponse,
}
def _load_endpoints_config() -> dict:
def _load_endpoints_config() -> EndpointsConfig:
"""Load the endpoints configuration from JSON file."""
config_path: Final = Path(__file__).parent.parent.parent / "containers" / "endpoints.json"
with open(config_path) as f:
return json.load(f)
def _get_endpoint_config(endpoint_name: str) -> dict | None:
def _get_endpoint_config(endpoint_name: str) -> EndpointConfig | None:
"""Get config for a specific endpoint by name."""
config: Final = _load_endpoints_config()
for endpoint in config["endpoints"]:
@ -56,10 +90,15 @@ def _get_endpoint_config(endpoint_name: str) -> dict | None:
return None
def _response_model(response_type_name: str) -> type[BaseModel] | None:
"""The pydantic model a container endpoint's ``response_type`` names."""
return RESPONSE_TYPES.get(response_type_name)
def _build_url(
api_base: str,
path_template: str,
path_params: dict[str, str],
path_params: Mapping[str, object],
) -> str:
"""Build the full URL by substituting path parameters.
@ -89,22 +128,18 @@ def _build_url(
def _build_query_params(
query_param_names: list,
kwargs: dict[str, Any],
) -> dict[str, str]:
query_param_names: Sequence[str],
kwargs: Mapping[str, object],
) -> dict[str, object]:
"""Build query parameters from kwargs."""
params: Final = {}
for param_name in query_param_names:
value = kwargs.get(param_name)
if value is not None:
params[param_name] = str(value) if not isinstance(value, str) else value
return params
supplied: Final = ((param_name, kwargs.get(param_name)) for param_name in query_param_names)
return {name: value if isinstance(value, str) else str(value) for name, value in supplied if value is not None}
def _prepare_multipart_file_upload(
file: Any,
headers: dict[str, Any],
) -> tuple:
headers: dict[str, object],
) -> tuple[dict[str, tuple[str, bytes, str]], dict[str, object]]:
"""
Prepare file and headers for multipart upload.
@ -129,6 +164,52 @@ def _prepare_multipart_file_upload(
return files, headers_copy
def _request_headers(
container_provider_config: "BaseContainerConfig",
extra_headers: dict[str, object] | None,
litellm_params: GenericLiteLLMParams,
) -> dict[str, object]:
"""The provider auth headers for a container request."""
return container_provider_config.validate_environment(
headers=extra_headers or {},
api_key=litellm_params.get("api_key", None),
)
def _request_api_base(
container_provider_config: "BaseContainerConfig",
litellm_params: GenericLiteLLMParams,
) -> str:
"""The provider base URL for a container request."""
return container_provider_config.get_complete_url(
api_base=litellm_params.get("api_base", None),
litellm_params=dict(litellm_params),
)
def _sync_http_client(
client: HTTPHandler | AsyncHTTPHandler | None,
litellm_params: GenericLiteLLMParams,
) -> HTTPHandler:
"""The sync HTTP client for a container request, reusing the caller's when usable."""
if client is None or not isinstance(client, HTTPHandler):
return _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)})
return client
def _async_http_client(
client: HTTPHandler | AsyncHTTPHandler | None,
litellm_params: GenericLiteLLMParams,
) -> AsyncHTTPHandler:
"""The async HTTP client for a container request, reusing the caller's when usable."""
if client is None or not isinstance(client, AsyncHTTPHandler):
return get_async_httpx_client(
llm_provider=litellm.LlmProviders.OPENAI,
params={"ssl_verify": litellm_params.get("ssl_verify", None)},
)
return client
class GenericContainerHandler:
"""
Generic handler for container file API endpoints.
@ -143,13 +224,13 @@ class GenericContainerHandler:
container_provider_config: "BaseContainerConfig",
litellm_params: GenericLiteLLMParams,
logging_obj: "LiteLLMLoggingObj",
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout = 600,
_is_async: bool = False,
client: HTTPHandler | AsyncHTTPHandler | None = None,
**kwargs,
) -> Any | Coroutine[Any, Any, Any]:
**kwargs: object,
) -> Any | Coroutine[object, object, Any]:
"""
Generic handler for any container file endpoint.
@ -196,11 +277,11 @@ class GenericContainerHandler:
container_provider_config: "BaseContainerConfig",
litellm_params: GenericLiteLLMParams,
logging_obj: "LiteLLMLoggingObj",
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout = 600,
client: HTTPHandler | AsyncHTTPHandler | None = None,
**kwargs,
**kwargs: object,
) -> Any:
"""Synchronous request handler."""
endpoint_config: Final = _get_endpoint_config(endpoint_name)
@ -208,23 +289,14 @@ class GenericContainerHandler:
raise ValueError(f"Unknown endpoint: {endpoint_name}")
# Get HTTP client
if client is None or not isinstance(client, HTTPHandler):
http_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)})
else:
http_client = client
http_client: Final = _sync_http_client(client, litellm_params)
# Build request
headers = container_provider_config.validate_environment(
headers=extra_headers or {},
api_key=litellm_params.get("api_key", None),
)
headers = _request_headers(container_provider_config, extra_headers, litellm_params)
if extra_headers:
headers.update(extra_headers)
api_base: Final = container_provider_config.get_complete_url(
api_base=litellm_params.get("api_base", None),
litellm_params=dict(litellm_params),
)
api_base: Final = _request_api_base(container_provider_config, litellm_params)
# Build URL with path params
path_params: Final = {p: kwargs.get(p, "") for p in endpoint_config.get("path_params", [])}
@ -275,11 +347,11 @@ class GenericContainerHandler:
return response.content
# Check for error response
response_json: Final = response.json()
response_json: Final[ContainerResponseBody] = response.json()
if "error" in response_json:
from litellm.llms.base_llm.chat.transformation import BaseLLMException
error_msg: Final = response_json.get("error", {}).get("message", str(response_json))
error_msg: Final = response_json["error"].get("message", str(response_json))
raise BaseLLMException(
status_code=response.status_code,
message=error_msg,
@ -287,7 +359,7 @@ class GenericContainerHandler:
)
# Parse response
response_type: Final = RESPONSE_TYPES.get(endpoint_config["response_type"])
response_type: Final = _response_model(endpoint_config["response_type"])
if response_type:
return response_type(**response_json)
return response_json
@ -301,11 +373,11 @@ class GenericContainerHandler:
container_provider_config: "BaseContainerConfig",
litellm_params: GenericLiteLLMParams,
logging_obj: "LiteLLMLoggingObj",
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout = 600,
client: HTTPHandler | AsyncHTTPHandler | None = None,
**kwargs,
**kwargs: object,
) -> Any:
"""Asynchronous request handler."""
endpoint_config: Final = _get_endpoint_config(endpoint_name)
@ -313,26 +385,14 @@ class GenericContainerHandler:
raise ValueError(f"Unknown endpoint: {endpoint_name}")
# Get HTTP client
if client is None or not isinstance(client, AsyncHTTPHandler):
http_client = get_async_httpx_client(
llm_provider=litellm.LlmProviders.OPENAI,
params={"ssl_verify": litellm_params.get("ssl_verify", None)},
)
else:
http_client = client
http_client: Final = _async_http_client(client, litellm_params)
# Build request
headers = container_provider_config.validate_environment(
headers=extra_headers or {},
api_key=litellm_params.get("api_key", None),
)
headers = _request_headers(container_provider_config, extra_headers, litellm_params)
if extra_headers:
headers.update(extra_headers)
api_base: Final = container_provider_config.get_complete_url(
api_base=litellm_params.get("api_base", None),
litellm_params=dict(litellm_params),
)
api_base: Final = _request_api_base(container_provider_config, litellm_params)
# Build URL with path params
path_params: Final = {p: kwargs.get(p, "") for p in endpoint_config.get("path_params", [])}
@ -383,11 +443,11 @@ class GenericContainerHandler:
return response.content
# Check for error response
response_json: Final = response.json()
response_json: Final[ContainerResponseBody] = response.json()
if "error" in response_json:
from litellm.llms.base_llm.chat.transformation import BaseLLMException
error_msg: Final = response_json.get("error", {}).get("message", str(response_json))
error_msg: Final = response_json["error"].get("message", str(response_json))
raise BaseLLMException(
status_code=response.status_code,
message=error_msg,
@ -395,7 +455,7 @@ class GenericContainerHandler:
)
# Parse response
response_type: Final = RESPONSE_TYPES.get(endpoint_config["response_type"])
response_type: Final = _response_model(endpoint_config["response_type"])
if response_type:
return response_type(**response_json)
return response_json

45
litellm/llms/gdc/chat/transformation.py
View file

@ -6,11 +6,25 @@ import json
import os
import re
import threading
from typing import Any, Final
from collections.abc import Callable
from typing import Any, Final, Protocol
from urllib.parse import urlsplit
import litellm
from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig
from litellm.types.llms.openai import AllMessageValues
class _GDCHAudienceCredentials(Protocol):
"""A GDCH service account credential already bound to an audience, ready to mint a bearer token."""
@property
def valid(self) -> bool: ...
@property
def token(self) -> str: ...
def refresh(self, request: object) -> None: ...
class GDCGeminiConfig(OpenAILikeChatConfig):
@ -21,7 +35,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
def __init__(self, **kwargs: Any) -> None:
super().__init__(**kwargs)
self._creds_lock = threading.Lock()
self._gdch_creds_cache: dict = {}
self._gdch_creds_cache: dict[tuple[str, str], _GDCHAudienceCredentials] = {}
def get_supported_openai_params(self, model: str) -> list:
return [
@ -110,7 +124,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
return f"{api_base}/v1/projects/{project}/locations/{location}/chat/completions"
def _read_env_bool(self, val: Any, env_var: str, default: bool = True) -> bool | str:
def _read_env_bool(self, val: bool | str | None, env_var: str, default: bool = True) -> bool | str:
def _parse(s: str) -> bool | str:
cleaned: Final = s.strip().lower()
if cleaned in ("false", "0", "no", "off"):
@ -129,7 +143,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
return default
return _parse(_env_val)
def _fetch_auth(self, gdch_creds: Any, ssl_verify: bool | str) -> None:
def _fetch_auth(self, gdch_creds: _GDCHAudienceCredentials, ssl_verify: bool | str) -> None:
import requests
from google.auth.transport import requests as auth_requests
@ -138,13 +152,24 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
auth_request: Final = auth_requests.Request(session=auth_session)
gdch_creds.refresh(auth_request)
def _cached_fetch_token(self, creds: Any, audience: str, ssl_verify: bool | str, api_key: str | None = None) -> str:
def _with_gdch_audience(self, creds: object, audience: str) -> _GDCHAudienceCredentials:
"""The credential rebound to ``audience``, which GDCH requires before a token refresh."""
bind_audience: Final[Callable[[str], _GDCHAudienceCredentials] | None] = getattr(
creds, "with_gdch_audience", None
)
if bind_audience is None:
raise AttributeError("GDC credentials must expose with_gdch_audience to be bound to a request audience")
return bind_audience(audience)
def _cached_fetch_token(
self, creds: object, audience: str, ssl_verify: bool | str, api_key: str | None = None
) -> str:
# Key cache by both audience and credential identity to prevent cross-caller contamination
cache_key: Final = (audience.rstrip("/"), api_key or str(id(creds)))
with self._creds_lock:
if cache_key not in self._gdch_creds_cache:
self._gdch_creds_cache[cache_key] = creds.with_gdch_audience(audience.rstrip("/"))
self._gdch_creds_cache[cache_key] = self._with_gdch_audience(creds, audience.rstrip("/"))
gdch_creds: Final = self._gdch_creds_cache[cache_key]
@ -155,7 +180,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
return token
def _load_creds_from_key(self, api_key: str) -> tuple[Any, bool]:
def _load_creds_from_key(self, api_key: str) -> tuple[object | None, bool]:
import google.auth
try:
@ -175,7 +200,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
self,
headers: dict,
model: str,
messages: list[Any],
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: str | None = None,
@ -230,7 +255,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
if self._read_env_bool(litellm_params.get("gdc_token_caching"), "GDC_TOKEN_CACHING", default=False):
token = self._cached_fetch_token(creds, audience, ssl_verify, api_key)
else:
gdch_creds: Final = creds.with_gdch_audience(audience)
gdch_creds: Final = self._with_gdch_audience(creds, audience)
self._fetch_auth(gdch_creds, ssl_verify)
token = gdch_creds.token
headers["Authorization"] = f"Bearer {token}"
@ -252,7 +277,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
def transform_request(
self,
model: str,
messages: list[Any],
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,

29
litellm/llms/infinity/rerank/transformation.py
View file

@ -4,9 +4,11 @@ Transformation logic from Cohere's /v1/rerank format to Infinity's `/v1/rerank`
Why separate file? Make it easy to see how transformation works
"""
from collections.abc import Sequence
from typing import Final
import httpx
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._uuid import uuid
@ -25,6 +27,31 @@ from litellm.types.rerank import (
from ..common_utils import InfinityError
class _InfinityRerankUsage(TypedDict, extra_items=ReadOnly[int]):
"""The token counters Infinity reports in the ``usage`` block of a rerank response."""
class _InfinityRerankResult(TypedDict):
"""One scored document in an Infinity ``/v1/rerank`` response."""
index: ReadOnly[int]
relevance_score: ReadOnly[float]
document: ReadOnly[str]
class _InfinityRerankResponse(TypedDict):
"""The JSON body returned by Infinity's ``/v1/rerank`` endpoint."""
id: ReadOnly[NotRequired[str]]
usage: ReadOnly[NotRequired[_InfinityRerankUsage]]
results: ReadOnly[Sequence[_InfinityRerankResult]]
def _parse_rerank_response(raw_response: httpx.Response) -> _InfinityRerankResponse:
"""Read the untyped JSON body of an Infinity rerank response."""
return raw_response.json()
class InfinityRerankConfig(CohereRerankConfig):
def get_complete_url(
self,
@ -80,7 +107,7 @@ class InfinityRerankConfig(CohereRerankConfig):
No transformation required, Infinity follows Cohere API response format
"""
try:
raw_response_json: Final = raw_response.json()
raw_response_json: Final = _parse_rerank_response(raw_response)
except Exception:
raise InfinityError(message=raw_response.text, status_code=raw_response.status_code)

76
litellm/llms/litellm_proxy/skills/code_execution.py
View file

@ -13,12 +13,57 @@ Generated files are returned directly in the response - no separate storage need
import base64
import json
from collections.abc import Sequence
from enum import Enum
from typing import Any, Final
from typing import Any, Final, Protocol
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_logger
class _ToolCallFunction(Protocol):
"""Function payload of an assistant tool call."""
name: str | None
arguments: str
class _ToolCall(Protocol):
"""Tool call requested by the assistant on a chat completion choice."""
id: str
function: _ToolCallFunction
class _AssistantMessage(Protocol):
"""Assistant message carried by a chat completion choice."""
content: str | None
tool_calls: Sequence[_ToolCall] | None
class _CompletionChoice(Protocol):
"""Single choice of a chat completion response."""
finish_reason: str
message: _AssistantMessage
class _SandboxFile(TypedDict):
"""File generated inside the sandbox during a code execution run."""
name: ReadOnly[str]
mime_type: ReadOnly[str]
content_base64: ReadOnly[str]
class _CodeExecutionArguments(TypedDict):
"""Arguments the model passes to the `litellm_code_execution` tool."""
code: NotRequired[ReadOnly[str]]
class LiteLLMInternalTools(str, Enum):
"""
Enum for internal LiteLLM tools that are injected into requests.
@ -30,7 +75,7 @@ class LiteLLMInternalTools(str, Enum):
CODE_EXECUTION = "litellm_code_execution"
def get_litellm_code_execution_tool() -> dict[str, Any]:
def get_litellm_code_execution_tool() -> dict[str, object]:
"""
Returns the litellm_code_execution tool definition in OpenAI format.
@ -51,7 +96,7 @@ def get_litellm_code_execution_tool() -> dict[str, Any]:
}
def get_litellm_code_execution_tool_anthropic() -> dict[str, Any]:
def get_litellm_code_execution_tool_anthropic() -> dict[str, object]:
"""
Returns the litellm_code_execution tool definition in Anthropic/messages API format.
@ -103,7 +148,7 @@ class CodeExecutionHandler:
skill_files: dict[str, bytes],
skill_id: str | None = None,
**kwargs,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Execute an LLM call with automatic code execution handling.
@ -134,8 +179,8 @@ class CodeExecutionHandler:
)
current_messages: Final = list(messages)
generated_files: Final[list[dict[str, Any]]] = [] # Files returned directly
execution_results: Final[list[dict]] = []
generated_files: Final[list[dict[str, object]]] = [] # Files returned directly
execution_results: Final[list[dict[str, object]]] = []
executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout)
response: Any = None # Initialize to avoid possibly unbound error
@ -151,11 +196,12 @@ class CodeExecutionHandler:
**kwargs,
)
assistant_message = response.choices[0].message
stop_reason = response.choices[0].finish_reason
choice: _CompletionChoice = response.choices[0]
assistant_message = choice.message
stop_reason: str = choice.finish_reason
# Build assistant message for conversation history
assistant_msg_dict: dict[str, Any] = {
assistant_msg_dict: dict[str, object] = {
"role": "assistant",
"content": assistant_message.content,
}
@ -190,8 +236,8 @@ class CodeExecutionHandler:
if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value:
# Execute code in sandbox
try:
args = json.loads(tool_call.function.arguments)
code = args.get("code", "")
args: _CodeExecutionArguments = json.loads(tool_call.function.arguments)
code: str = args.get("code", "")
verbose_logger.debug("CodeExecutionHandler: Executing code (%s chars)", len(code))
@ -202,13 +248,15 @@ class CodeExecutionHandler:
verbose_logger.debug("CodeExecutionHandler: Execution result: %s", exec_result)
sandbox_files: Sequence[_SandboxFile] = exec_result["files"]
execution_results.append(
{
"iteration": iteration,
"success": exec_result["success"],
"output": exec_result["output"],
"error": exec_result["error"],
"files": [f["name"] for f in exec_result["files"]],
"files": [f["name"] for f in sandbox_files],
}
)
@ -216,9 +264,9 @@ class CodeExecutionHandler:
tool_result = exec_result["output"] or ""
# Collect generated files (returned directly, no storage)
if exec_result["files"]:
if sandbox_files:
tool_result += "\n\nGenerated files:"
for f in exec_result["files"]:
for f in sandbox_files:
file_content = base64.b64decode(f["content_base64"])
# Add to generated files list (returned in response)
generated_files.append(

38
litellm/llms/oci/chat/cohere.py
View file

@ -8,10 +8,12 @@ response parsing, and streaming chunk parsing for models served with
import datetime
import json
from typing import Any, Final
from collections.abc import Sequence
from typing import Any, Final, TypedDict
import httpx
from pydantic import ValidationError
from typing_extensions import ReadOnly
from litellm.llms.oci.chat.generic import (
_normalize_oci_finish_reason,
@ -46,6 +48,20 @@ from litellm.types.utils import (
)
class _OpenAIToolCallFunction(TypedDict, total=False):
"""The ``function`` block of an OpenAI-format assistant tool call."""
name: ReadOnly[str | None]
arguments: ReadOnly[str | dict[str, object]]
class _OpenAIToolCall(TypedDict, total=False):
"""An entry of an OpenAI-format assistant message's ``tool_calls``."""
id: ReadOnly[str | None]
function: ReadOnly[_OpenAIToolCallFunction]
def _extract_text_content(content: Any) -> str:
"""Return the plain-text representation of a message content value."""
if content is None:
@ -78,10 +94,10 @@ def adapt_messages_to_cohere_standard(
"""
# First pass: build tool_call_id → CohereToolCall so tool-result messages can
# reference the originating call by name and parameters.
tool_call_lookup: Final[dict[str, CohereToolCall]] = {}
tool_call_lookup: Final[dict[str | None, CohereToolCall]] = {}
for msg in messages:
if msg.get("role") == "assistant":
tool_calls_raw: Any = msg.get("tool_calls") or []
tool_calls_raw: Sequence[_OpenAIToolCall] = msg.get("tool_calls") or []
for tc in tool_calls_raw:
tc_id = tc.get("id", "")
raw_args = tc.get("function", {}).get("arguments", "{}")
@ -150,8 +166,22 @@ def adapt_messages_to_cohere_standard(
return chat_history
class _OpenAIToolDefinitionFunction(TypedDict, total=False):
"""The ``function`` block of an OpenAI-format tool definition."""
name: ReadOnly[str]
description: ReadOnly[str]
parameters: ReadOnly[dict[str, object]]
class _OpenAIToolDefinition(TypedDict, total=False):
"""An entry of an OpenAI-format ``tools`` array."""
function: ReadOnly[_OpenAIToolDefinitionFunction]
def adapt_tool_definitions_to_cohere_standard(
tools: list[dict[str, Any]],
tools: Sequence[_OpenAIToolDefinition],
) -> list[CohereTool]:
"""Adapt OpenAI-format tool definitions to the OCI Cohere format.

36
litellm/llms/ollama/completion/handler.py
View file

@ -4,16 +4,32 @@ Ollama /chat/completion calls handled in llm_http_handler.py
[TODO]: migrate embeddings to a base handler as well.
"""
from typing import Any, Final
from collections.abc import Mapping, Sequence
from typing import Any, Final, Protocol, TypedDict
from typing_extensions import NotRequired, ReadOnly
import litellm
from litellm.types.utils import EmbeddingResponse
class TokenEncoder(Protocol):
"""The tokenizer surface used to estimate prompt tokens."""
def encode(self, text: str, /) -> Sequence[int]: ...
class OllamaEmbeddingResponse(TypedDict):
"""Body of an Ollama ``/api/embed`` response."""
embeddings: ReadOnly[list[list[float]]]
prompt_eval_count: ReadOnly[NotRequired[int]]
def _prepare_ollama_embedding_payload(
model: str, prompts: list[str], optional_params: dict[str, Any]
) -> dict[str, Any]:
data: Final[dict[str, Any]] = {"model": model, "input": prompts}
model: str, prompts: list[str], optional_params: Mapping[str, object]
) -> dict[str, object]:
data: Final[dict[str, object]] = {"model": model, "input": prompts}
special_optional_params: Final = ["truncate", "options", "keep_alive", "dimensions"]
for k, v in optional_params.items():
@ -27,12 +43,12 @@ def _prepare_ollama_embedding_payload(
def _process_ollama_embedding_response(
response_json: dict,
response_json: OllamaEmbeddingResponse,
prompts: list[str],
model: str,
model_response: EmbeddingResponse,
logging_obj: Any,
encoding: Any,
encoding: TokenEncoder | None,
) -> EmbeddingResponse:
output_data: Final = []
embeddings: Final[list[list[float]]] = response_json["embeddings"]
@ -72,7 +88,7 @@ async def ollama_aembeddings(
model_response: EmbeddingResponse,
optional_params: dict,
logging_obj: Any,
encoding: Any,
encoding: TokenEncoder | None,
):
if not api_base.endswith("/api/embed"):
api_base += "/api/embed"
@ -80,7 +96,7 @@ async def ollama_aembeddings(
data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params)
response: Final = await litellm.module_level_aclient.post(url=api_base, json=data)
response_json: Final = response.json()
response_json: Final[OllamaEmbeddingResponse] = response.json()
return _process_ollama_embedding_response(
response_json=response_json,
@ -99,7 +115,7 @@ def ollama_embeddings(
optional_params: dict,
model_response: EmbeddingResponse,
logging_obj: Any,
encoding: Any = None,
encoding: TokenEncoder | None = None,
):
if not api_base.endswith("/api/embed"):
api_base += "/api/embed"
@ -107,7 +123,7 @@ def ollama_embeddings(
data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params)
response: Final = litellm.module_level_client.post(url=api_base, json=data)
response_json: Final = response.json()
response_json: Final[OllamaEmbeddingResponse] = response.json()
return _process_ollama_embedding_response(
response_json=response_json,

72
litellm/llms/openai/containers/transformation.py
View file

@ -1,6 +1,8 @@
from typing import TYPE_CHECKING, Any, Final
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Literal
import httpx
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
@ -11,9 +13,11 @@ from litellm.secret_managers.main import get_secret_str
from litellm.types.containers.main import (
ContainerCreateOptionalRequestParams,
ContainerFileListResponse,
ContainerFileObject,
ContainerListResponse,
ContainerObject,
DeleteContainerResult,
ExpiresAfter,
)
from litellm.types.router import GenericLiteLLMParams
@ -32,6 +36,46 @@ else:
BaseLLMException = Any
class OpenAIContainerPayload(TypedDict):
"""The JSON body OpenAI returns for a single container."""
id: ReadOnly[str]
object: ReadOnly[Literal["container"]]
created_at: ReadOnly[int]
status: ReadOnly[str]
expires_after: ReadOnly[ExpiresAfter | None]
last_active_at: ReadOnly[int | None]
name: ReadOnly[str | None]
class OpenAIContainerListPayload(TypedDict):
"""The JSON body OpenAI returns for a page of containers."""
object: ReadOnly[Literal["list"]]
data: ReadOnly[list[ContainerObject]]
first_id: ReadOnly[str | None]
last_id: ReadOnly[str | None]
has_more: ReadOnly[bool]
class OpenAIContainerDeletedPayload(TypedDict):
"""The JSON body OpenAI returns for a deleted container."""
id: ReadOnly[str]
object: ReadOnly[Literal["container.deleted"]]
deleted: ReadOnly[bool]
class OpenAIContainerFileListPayload(TypedDict):
"""The JSON body OpenAI returns for a page of container files."""
object: ReadOnly[Literal["list"]]
data: ReadOnly[list[ContainerFileObject]]
first_id: ReadOnly[str | None]
last_id: ReadOnly[str | None]
has_more: ReadOnly[bool]
class OpenAIContainerConfig(BaseContainerConfig):
"""Configuration class for OpenAI container API."""
@ -87,7 +131,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
def transform_container_create_request(
self,
name: str,
container_create_optional_request_params: dict,
container_create_optional_request_params: Mapping[str, object],
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> dict:
@ -111,7 +155,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
logging_obj: LiteLLMLoggingObj,
) -> ContainerObject:
"""Transform the OpenAI container creation response."""
response_data: Final = raw_response.json()
response_data: Final[OpenAIContainerPayload] = raw_response.json()
# Transform the response data
container_obj: Final = ContainerObject(**response_data)
@ -140,7 +184,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
after: str | None = None,
limit: int | None = None,
order: str | None = None,
extra_query: dict[str, Any] | None = None,
extra_query: Mapping[str, object] | None = None,
) -> tuple[str, dict]:
"""Transform the container list request for OpenAI API.
@ -151,7 +195,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
url: Final = api_base
# Prepare query parameters
params: Final = {}
params: Final[dict[str, object]] = {}
if after is not None:
params["after"] = after
if limit is not None:
@ -171,7 +215,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
logging_obj: LiteLLMLoggingObj,
) -> ContainerListResponse:
"""Transform the OpenAI container list response."""
response_data: Final = raw_response.json()
response_data: Final[OpenAIContainerListPayload] = raw_response.json()
# Transform the response data
container_list: Final = ContainerListResponse(**response_data)
@ -191,7 +235,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}")
# No additional data needed for GET request
data: Final[dict[str, Any]] = {}
data: Final[dict[str, object]] = {}
return url, data
@ -201,7 +245,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
logging_obj: LiteLLMLoggingObj,
) -> ContainerObject:
"""Transform the OpenAI container retrieve response."""
response_data: Final = raw_response.json()
response_data: Final[OpenAIContainerPayload] = raw_response.json()
# Transform the response data
container_obj: Final = ContainerObject(**response_data)
@ -224,7 +268,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}")
# No data needed for DELETE request
data: Final[dict[str, Any]] = {}
data: Final[dict[str, object]] = {}
return url, data
@ -234,7 +278,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
logging_obj: LiteLLMLoggingObj,
) -> DeleteContainerResult:
"""Transform the OpenAI container delete response."""
response_data: Final = raw_response.json()
response_data: Final[OpenAIContainerDeletedPayload] = raw_response.json()
# Transform the response data
delete_result: Final = DeleteContainerResult(**response_data)
@ -250,7 +294,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
after: str | None = None,
limit: int | None = None,
order: str | None = None,
extra_query: dict[str, Any] | None = None,
extra_query: Mapping[str, object] | None = None,
) -> tuple[str, dict]:
"""Transform the container file list request for OpenAI API.
@ -262,7 +306,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}/files")
# Prepare query parameters
params: Final[dict[str, Any]] = {}
params: Final[dict[str, object]] = {}
if after is not None:
params["after"] = after
if limit is not None:
@ -282,7 +326,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
logging_obj: LiteLLMLoggingObj,
) -> ContainerFileListResponse:
"""Transform the OpenAI container file list response."""
response_data: Final = raw_response.json()
response_data: Final[OpenAIContainerFileListPayload] = raw_response.json()
# Transform the response data
file_list: Final = ContainerFileListResponse(**response_data)
@ -308,7 +352,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}/files/{encoded_file_id}/content")
# No query parameters needed
params: Final[dict[str, Any]] = {}
params: Final[dict[str, object]] = {}
return url, params

127
litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py
View file

@ -1,6 +1,8 @@
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
import httpx
from typing_extensions import ReadOnly, TypedDict
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
@ -25,6 +27,68 @@ else:
LiteLLMLoggingObj = Any
class _AuthHeadersView(TypedDict):
"""Auth headers read out of an untyped ``BaseVectorStoreAuthCredentials``."""
headers: ReadOnly[Mapping[str, str]]
class _RagContext(TypedDict, total=False):
"""One context entry of a Vertex RAG ``:retrieveContexts`` response."""
text: ReadOnly[str]
sourceUri: ReadOnly[str]
sourceDisplayName: ReadOnly[str]
pageSpan: ReadOnly[Mapping[str, object]]
score: ReadOnly[float]
class _RagContexts(TypedDict, total=False):
"""The ``contexts`` envelope wrapping the context list."""
contexts: ReadOnly[Sequence[_RagContext]]
class _RetrieveContextsBody(TypedDict, total=False):
"""Body of a Vertex RAG ``:retrieveContexts`` response."""
contexts: ReadOnly[_RagContexts]
class _RetrieveContextsView(TypedDict):
"""Typed view over the untyped ``:retrieveContexts`` JSON payload."""
body: ReadOnly[_RetrieveContextsBody]
class _RagCorpusBody(TypedDict, total=False):
"""Body of a Vertex RAG ``ragCorpora`` create response."""
name: ReadOnly[str]
display_name: ReadOnly[str]
createTime: ReadOnly[str | int | float]
labels: ReadOnly[Mapping[str, str]]
class _RagCorpusView(TypedDict):
"""Typed view over the untyped ``ragCorpora`` create JSON payload."""
body: ReadOnly[_RagCorpusBody]
class _RagSearchQuery(TypedDict, total=False):
"""The ``query`` block of a Vertex RAG ``:retrieveContexts`` request."""
text: ReadOnly[str]
rag_retrieval_config: ReadOnly[Mapping[str, object]]
class _LoggedQueryView(TypedDict):
"""The search query recovered from the logging object's call details."""
search_query: ReadOnly[str]
class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
"""
Configuration for Vertex AI Vector Store RAG API
@ -35,7 +99,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
def __init__(self):
super().__init__()
def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials:
def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> BaseVectorStoreAuthCredentials:
# Get credentials and project info
vertex_credentials: Final = self.get_vertex_ai_credentials(dict(litellm_params))
vertex_project: Final = self.get_vertex_ai_project(dict(litellm_params))
@ -60,20 +124,23 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
"write": [("POST", "/ragCorpora")],
}
def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict:
def validate_environment(
self, headers: dict[str, str], litellm_params: GenericLiteLLMParams | None
) -> dict[str, str]:
"""
Validate and set up authentication for Vertex AI RAG API
"""
litellm_params = litellm_params or GenericLiteLLMParams()
auth_headers: Final = self.get_auth_credentials(litellm_params.model_dump())
headers.update(auth_headers.get("headers", {}))
auth_view: Final[_AuthHeadersView] = {"headers": auth_headers.get("headers", {})}
headers.update(auth_view["headers"])
return headers
def get_complete_url(
self,
api_base: str | None,
litellm_params: dict,
litellm_params: dict[str, object],
) -> str:
"""
Get the Base endpoint for Vertex AI RAG API
@ -95,9 +162,9 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: dict[str, Any] | None = None,
) -> tuple[str, dict[str, Any]]:
litellm_params: dict[str, object],
extra_body: Mapping[str, object] | None = None,
) -> tuple[str, dict[str, object]]:
"""
Transform search request for Vertex AI RAG API
"""
@ -120,35 +187,34 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
# Just the corpus ID, construct full path
full_rag_corpus = f"projects/{vertex_project}/locations/{vertex_location}/ragCorpora/{vector_store_id}"
# Build the request body for Vertex AI RAG API
request_body: Final[dict[str, Any]] = {
"vertex_rag_store": {"rag_resources": [{"rag_corpus": full_rag_corpus}]},
"query": {"text": query},
}
#########################################################
# Update logging object with details of the request
#########################################################
litellm_logging_obj.model_call_details["query"] = query
# Add optional parameters
rag_retrieval_config: Final[dict[str, object]] = {}
max_num_results: Final = vector_store_search_optional_params.get("max_num_results")
if max_num_results is not None:
request_body["query"]["rag_retrieval_config"] = {"top_k": max_num_results}
rag_retrieval_config["top_k"] = max_num_results
# Add filters if provided
filters: Final = vector_store_search_optional_params.get("filters")
filters: Final[object] = vector_store_search_optional_params.get("filters")
if filters is not None:
if "rag_retrieval_config" not in request_body["query"]:
request_body["query"]["rag_retrieval_config"] = {}
request_body["query"]["rag_retrieval_config"]["filter"] = filters
rag_retrieval_config["filter"] = filters
# Add ranking options if provided
ranking_options: Final = vector_store_search_optional_params.get("ranking_options")
ranking_options: Final[object] = vector_store_search_optional_params.get("ranking_options")
if ranking_options is not None:
if "rag_retrieval_config" not in request_body["query"]:
request_body["query"]["rag_retrieval_config"] = {}
request_body["query"]["rag_retrieval_config"]["ranking"] = ranking_options
rag_retrieval_config["ranking"] = ranking_options
query_body: Final[_RagSearchQuery] = (
{"text": query, "rag_retrieval_config": rag_retrieval_config} if rag_retrieval_config else {"text": query}
)
request_body: Final[dict[str, object]] = {
"vertex_rag_store": {"rag_resources": [{"rag_corpus": full_rag_corpus}]},
"query": query_body,
}
return url, request_body
@ -159,12 +225,13 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
Transform Vertex AI RAG API response to standard vector store search response
"""
try:
response_json: Final = response.json()
response_view: Final[_RetrieveContextsView] = {"body": response.json()}
response_json: Final = response_view["body"]
# Extract contexts from Vertex AI response - handle nested structure
contexts: Final = response_json.get("contexts", {}).get("contexts", [])
# Transform contexts to standard format
search_results: Final = []
search_results: Final[list[VectorStoreSearchResult]] = []
for context in contexts:
content = [
VectorStoreResultContent(
@ -202,9 +269,12 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
)
search_results.append(result)
query_view: Final[_LoggedQueryView] = {
"search_query": litellm_logging_obj.model_call_details.get("query", "")
}
return VectorStoreSearchResponse(
object="vector_store.search_results.page",
search_query=litellm_logging_obj.model_call_details.get("query", ""),
search_query=query_view["search_query"],
data=search_results,
)
@ -219,14 +289,14 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
self,
vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams,
api_base: str,
) -> tuple[str, dict[str, Any]]:
) -> tuple[str, dict[str, object]]:
"""
Transform create request for Vertex AI RAG Corpus
"""
url: Final = f"{api_base}/ragCorpora" # Base URL for creating RAG corpus
# Build the request body for Vertex AI RAG Corpus creation
request_body: Final[dict[str, Any]] = {
request_body: Final[dict[str, object]] = {
"display_name": vector_store_create_optional_params.get("name", "litellm-vector-store"),
"description": "Vector store created via LiteLLM",
}
@ -243,7 +313,8 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
Transform Vertex AI RAG Corpus creation response to standard vector store response
"""
try:
response_json: Final = response.json()
response_view: Final[_RagCorpusView] = {"body": response.json()}
response_json: Final = response_view["body"]
# Extract the corpus ID from the response name
corpus_name: Final = response_json.get("name", "")

2
litellm/models/base.py
View file

@ -33,6 +33,6 @@ class DomainModel(BaseModel):
return cls(**record.dict())
return cls(**dict(record))
def to_db_dict(self, exclude_unset: bool = False) -> dict[str, Any]:
def to_db_dict(self, exclude_unset: bool = False) -> dict[str, object]:
"""Convert domain model to a dictionary for database operations."""
return self.model_dump(exclude_none=True, exclude_unset=exclude_unset)

63
litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py
View file

@ -36,7 +36,10 @@ a healed fleet has no null rows and the backfill exits after one query.
import json
from collections import Counter
from typing import Any, Final, Literal
from collections.abc import Mapping, Sequence
from typing import Final, Literal, Protocol
from pydantic import JsonValue
from litellm._logging import verbose_proxy_logger
from litellm.proxy._experimental.mcp_server.db import _decode_oauth_payload, decrypt_credentials
@ -55,9 +58,59 @@ BackfillRule = Literal[
_BACKFILL_AUDIT_ACTOR: Final = "oauth2_flow_backfill"
def _decrypted_credentials(raw_credentials: Any) -> MCPCredentials | None:
class _MCPServerRow(Protocol):
"""The ``LiteLLM_MCPServerTable`` columns this backfill reads."""
@property
def server_id(self) -> str: ...
@property
def authorization_url(self) -> str | None: ...
@property
def registration_url(self) -> str | None: ...
@property
def token_url(self) -> str | None: ...
@property
def credentials(self) -> str | Mapping[str, JsonValue] | None: ...
class _MCPUserCredentialRow(Protocol):
"""The ``LiteLLM_MCPUserCredentials`` columns this backfill reads."""
@property
def server_id(self) -> str: ...
@property
def credential_b64(self) -> str: ...
class _MCPServerTable(Protocol):
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_MCPServerRow]: ...
async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, str]) -> object: ...
class _MCPUserCredentialsTable(Protocol):
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_MCPUserCredentialRow]: ...
def _mcp_server_table(prisma_client: PrismaClient) -> _MCPServerTable:
"""The MCP server table, typed so the untyped prisma client surface stops here."""
return prisma_client.db.litellm_mcpservertable
def _mcp_user_credentials_table(prisma_client: PrismaClient) -> _MCPUserCredentialsTable:
"""The per-user MCP credential table, typed so the untyped prisma client surface stops here."""
return prisma_client.db.litellm_mcpusercredentials
def _decrypted_credentials(raw_credentials: str | Mapping[str, JsonValue] | None) -> MCPCredentials | None:
if raw_credentials is None:
return None
parsed: JsonValue | Mapping[str, JsonValue]
if isinstance(raw_credentials, str):
try:
parsed = json.loads(raw_credentials)
@ -92,14 +145,14 @@ def classify_null_flow_row(
async def backfill_null_oauth2_flows(prisma_client: PrismaClient) -> dict[BackfillRule, int]:
"""Classify every ``auth_type=oauth2`` row whose ``oauth2_flow`` is null; stamp the provable
ones, warn on the ambiguous ones, and return counts per rule."""
null_rows: Final[list[Any]] = await prisma_client.db.litellm_mcpservertable.find_many(
null_rows: Final[Sequence[_MCPServerRow]] = await _mcp_server_table(prisma_client).find_many(
where={"auth_type": "oauth2", "oauth2_flow": None},
)
if not null_rows:
return {}
server_ids: Final = [row.server_id for row in null_rows]
token_rows: Final[list[Any]] = await prisma_client.db.litellm_mcpusercredentials.find_many(
token_rows: Final[Sequence[_MCPUserCredentialRow]] = await _mcp_user_credentials_table(prisma_client).find_many(
where={"server_id": {"in": server_ids}},
)
server_ids_with_oauth_tokens: Final[set[str]] = {
@ -141,7 +194,7 @@ async def backfill_null_oauth2_flows(prisma_client: PrismaClient) -> dict[Backfi
stamped_flows: Final = {flow for _, (flow, _) in classified if flow is not None}
for stamped_flow in stamped_flows:
server_ids_for_flow = [row.server_id for row, (row_flow, _) in classified if row_flow == stamped_flow]
await prisma_client.db.litellm_mcpservertable.update_many(
await _mcp_server_table(prisma_client).update_many(
where={"server_id": {"in": server_ids_for_flow}, "oauth2_flow": None},
data={"oauth2_flow": stamped_flow, "updated_by": _BACKFILL_AUDIT_ACTOR},
)

93
litellm/proxy/_experimental/mcp_server/toolset_db.py
View file

@ -1,5 +1,9 @@
import json
from typing import Final
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import Final, Protocol
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -7,18 +11,73 @@ from litellm.proxy.utils import PrismaClient
from litellm.repositories.table_repositories import MCPToolsetRepository
from litellm.types.mcp_server.mcp_toolset import (
MCPToolset,
MCPToolsetTool,
NewMCPToolsetRequest,
UpdateMCPToolsetRequest,
)
def _toolset_from_row(row) -> MCPToolset:
class MCPToolsetFields(TypedDict):
"""The ``MCPToolset`` constructor keywords a toolset row expands into."""
toolset_id: ReadOnly[str]
toolset_name: ReadOnly[str]
description: NotRequired[ReadOnly[str | None]]
tools: NotRequired[ReadOnly[list[MCPToolsetTool]]]
created_at: NotRequired[ReadOnly[datetime | None]]
created_by: NotRequired[ReadOnly[str | None]]
updated_at: NotRequired[ReadOnly[datetime | None]]
updated_by: NotRequired[ReadOnly[str | None]]
class MCPToolsetRowData(TypedDict):
"""A toolset table row, whose ``tools`` column is stored as JSON."""
toolset_id: ReadOnly[str]
toolset_name: ReadOnly[str]
description: NotRequired[ReadOnly[str | None]]
tools: NotRequired[ReadOnly[str | list[MCPToolsetTool]]]
created_at: NotRequired[ReadOnly[datetime | None]]
created_by: NotRequired[ReadOnly[str | None]]
updated_at: NotRequired[ReadOnly[datetime | None]]
updated_by: NotRequired[ReadOnly[str | None]]
class MCPToolsetRow(Protocol):
"""A row of the toolset table, as the prisma client returns it."""
def model_dump(self) -> MCPToolsetRowData: ...
class MCPToolsetTable(Protocol):
"""The prisma table actions this module runs against the toolset table."""
async def create(self, data: Mapping[str, object]) -> MCPToolsetRow: ...
async def find_unique(self, where: Mapping[str, object]) -> MCPToolsetRow | None: ...
async def find_first(self, where: Mapping[str, object]) -> MCPToolsetRow | None: ...
async def find_many(self, where: Mapping[str, object]) -> Sequence[MCPToolsetRow]: ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> MCPToolsetRow: ...
async def delete(self, where: Mapping[str, object]) -> MCPToolsetRow: ...
def _toolset_table(prisma_client: PrismaClient) -> MCPToolsetTable:
"""The toolset table actions of the prisma client."""
return MCPToolsetRepository(prisma_client).table
def _toolset_from_row(row: MCPToolsetRow) -> MCPToolset:
data: Final = row.model_dump()
tools = data.get("tools") or []
if isinstance(tools, str):
tools = json.loads(tools)
data["tools"] = tools
return MCPToolset(**data)
tools: Final = data.get("tools") or []
resolved: Final[MCPToolsetFields] = {
**data,
"tools": json.loads(tools) if isinstance(tools, str) else tools,
}
return MCPToolset(**resolved)
async def create_mcp_toolset(
@ -31,7 +90,7 @@ async def create_mcp_toolset(
data_dict["tools"] = json.dumps(data_dict.get("tools", []))
data_dict["created_by"] = touched_by
data_dict["updated_by"] = touched_by
row: Final = await MCPToolsetRepository(prisma_client).table.create(data=data_dict)
row: Final = await _toolset_table(prisma_client).create(data=data_dict)
return _toolset_from_row(row)
@ -39,7 +98,7 @@ async def get_mcp_toolset(
prisma_client: PrismaClient,
toolset_id: str,
) -> MCPToolset | None:
row: Final = await MCPToolsetRepository(prisma_client).table.find_unique(where={"toolset_id": toolset_id})
row: Final = await _toolset_table(prisma_client).find_unique(where={"toolset_id": toolset_id})
if row is None:
return None
return _toolset_from_row(row)
@ -47,13 +106,11 @@ async def get_mcp_toolset(
async def list_mcp_toolsets(
prisma_client: PrismaClient,
toolset_ids: list[str] | None = None,
) -> list[MCPToolset]:
toolset_ids: Sequence[str] | None = None,
) -> Sequence[MCPToolset]:
try:
where = {}
if toolset_ids is not None:
where = {"toolset_id": {"in": toolset_ids}}
rows: Final = await MCPToolsetRepository(prisma_client).table.find_many(where=where)
where: Final[Mapping[str, object]] = {} if toolset_ids is None else {"toolset_id": {"in": toolset_ids}}
rows: Final = await _toolset_table(prisma_client).find_many(where=where)
return [_toolset_from_row(r) for r in rows]
except Exception as e:
verbose_proxy_logger.warning("litellm.proxy._experimental.mcp_server.toolset_db::list_mcp_toolsets - %s", e)
@ -64,7 +121,7 @@ async def get_mcp_toolset_by_name(
prisma_client: PrismaClient,
toolset_name: str,
) -> MCPToolset | None:
row: Final = await MCPToolsetRepository(prisma_client).table.find_first(where={"toolset_name": toolset_name})
row: Final = await _toolset_table(prisma_client).find_first(where={"toolset_name": toolset_name})
if row is None:
return None
return _toolset_from_row(row)
@ -80,7 +137,7 @@ async def update_mcp_toolset(
data_dict["tools"] = json.dumps(data_dict["tools"])
data_dict["updated_by"] = touched_by
try:
row: Final = await MCPToolsetRepository(prisma_client).table.update(
row: Final = await _toolset_table(prisma_client).update(
where={"toolset_id": data.toolset_id},
data=data_dict,
)
@ -98,7 +155,7 @@ async def delete_mcp_toolset(
toolset_id: str,
) -> MCPToolset | None:
try:
row: Final = await MCPToolsetRepository(prisma_client).table.delete(where={"toolset_id": toolset_id})
row: Final = await _toolset_table(prisma_client).delete(where={"toolset_id": toolset_id})
except Exception as e:
from prisma.errors import RecordNotFoundError

81
litellm/proxy/agent_endpoints/agent_registry.py
View file

@ -12,6 +12,7 @@ from litellm.proxy.management_helpers.object_permission_utils import (
handle_update_object_permission_common,
)
from litellm.proxy.utils import PrismaClient
from litellm.repositories.base_repository import SupportsModelDump
from litellm.repositories.table_repositories import AgentsRepository, ObjectPermissionRepository
from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest
@ -83,9 +84,17 @@ class AgentTableClient(Protocol):
async def delete(self, where: Mapping[str, object]) -> AgentRecord: ...
class _AgentsRepositoryView(Protocol):
@property
def table(self) -> AgentTableClient: ...
def _agents_table_of(repository: _AgentsRepositoryView) -> AgentTableClient:
return repository.table
def agents_table(prisma_client: PrismaClient) -> AgentTableClient:
table: Final[AgentTableClient] = AgentsRepository(prisma_client).table
return table
return _agents_table_of(AgentsRepository(prisma_client))
class ObjectPermissionGrantRecord(Protocol):
@ -99,9 +108,17 @@ class ObjectPermissionTableClient(Protocol):
async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
class _ObjectPermissionRepositoryView(Protocol):
@property
def table(self) -> ObjectPermissionTableClient: ...
def _object_permission_table_of(repository: _ObjectPermissionRepositoryView) -> ObjectPermissionTableClient:
return repository.table
def object_permission_table(prisma_client: PrismaClient) -> ObjectPermissionTableClient:
table: Final[ObjectPermissionTableClient] = ObjectPermissionRepository(prisma_client).table
return table
return _object_permission_table_of(ObjectPermissionRepository(prisma_client))
class GrantMigrationResult(NamedTuple):
@ -283,19 +300,21 @@ class AgentRegistry:
agent_name: Final = agent.get("agent_name")
# Serialize litellm_params
litellm_params_obj: Final[Any] = agent.get("litellm_params", {})
if hasattr(litellm_params_obj, "model_dump"):
litellm_params_dict = litellm_params_obj.model_dump()
else:
litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {}
litellm_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("litellm_params", {})
litellm_params_dict: Final[Mapping[str, object]] = (
litellm_params_obj.model_dump()
if isinstance(litellm_params_obj, SupportsModelDump)
else (dict(litellm_params_obj) if litellm_params_obj else {})
)
litellm_params: Final[str] = safe_dumps(litellm_params_dict)
# Serialize agent_card_params
agent_card_params_obj: Final[Any] = agent.get("agent_card_params", {})
if hasattr(agent_card_params_obj, "model_dump"):
agent_card_params_dict = agent_card_params_obj.model_dump()
else:
agent_card_params_dict = dict(agent_card_params_obj) if agent_card_params_obj else {}
agent_card_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("agent_card_params", {})
agent_card_params_dict: Final[Mapping[str, object]] = (
agent_card_params_obj.model_dump()
if isinstance(agent_card_params_obj, SupportsModelDump)
else (dict(agent_card_params_obj) if agent_card_params_obj else {})
)
agent_card_params: Final[str] = safe_dumps(agent_card_params_dict)
# Handle object_permission (MCP tool access for agent)
@ -386,15 +405,13 @@ class AgentRegistry:
The patched agent
"""
try:
existing_agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id})
if existing_agent is not None:
existing_agent = dict(existing_agent)
if existing_agent is None:
existing_record: Final = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
if existing_record is None:
raise Exception(f"Agent with ID {agent_id} not found")
existing_agent: Final[Mapping[str, object]] = dict(existing_record)
augment_agent: Final = {**existing_agent, **agent}
update_data: Final[dict[str, Any]] = {}
update_data: Final[dict[str, object]] = {}
if augment_agent.get("agent_name"):
update_data["agent_name"] = augment_agent.get("agent_name")
if augment_agent.get("litellm_params"):
@ -418,7 +435,7 @@ class AgentRegistry:
update_data["extra_headers"] = extra_headers_value if extra_headers_value is not None else []
if agent.get("object_permission") is not None:
agent_copy: Final = dict(augment_agent)
existing_object_permission_id: Final = existing_agent.get("object_permission_id")
existing_object_permission_id: Final = existing_record.object_permission_id
object_permission_id: Final = await handle_update_object_permission_common(
agent_copy,
existing_object_permission_id,
@ -460,19 +477,21 @@ class AgentRegistry:
agent_name: Final = agent.get("agent_name")
# Serialize litellm_params
litellm_params_obj: Final[Any] = agent.get("litellm_params", {})
if hasattr(litellm_params_obj, "model_dump"):
litellm_params_dict = litellm_params_obj.model_dump()
else:
litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {}
litellm_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("litellm_params", {})
litellm_params_dict: Final[Mapping[str, object]] = (
litellm_params_obj.model_dump()
if isinstance(litellm_params_obj, SupportsModelDump)
else (dict(litellm_params_obj) if litellm_params_obj else {})
)
litellm_params: Final[str] = safe_dumps(litellm_params_dict)
# Serialize agent_card_params
agent_card_params_obj: Final[Any] = agent.get("agent_card_params", {})
if hasattr(agent_card_params_obj, "model_dump"):
agent_card_params_dict = agent_card_params_obj.model_dump()
else:
agent_card_params_dict = dict(agent_card_params_obj) if agent_card_params_obj else {}
agent_card_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("agent_card_params", {})
agent_card_params_dict: Final[Mapping[str, object]] = (
agent_card_params_obj.model_dump()
if isinstance(agent_card_params_obj, SupportsModelDump)
else (dict(agent_card_params_obj) if agent_card_params_obj else {})
)
agent_card_params: Final[str] = safe_dumps(agent_card_params_dict)
# Serialize static_headers for update

51
litellm/proxy/client/cli/commands/credentials.py
View file

@ -1,14 +1,42 @@
import json
from collections.abc import Mapping, Sequence
from typing import Final, Literal
import click
import requests
import rich
from rich.table import Table
from typing_extensions import NotRequired, ReadOnly, TypedDict
from ...credentials import CredentialsManagementClient
class _CliContext(TypedDict):
"""Values the top-level CLI group stores on the click context."""
base_url: ReadOnly[str]
api_key: ReadOnly[str | None]
class _CliContextView(TypedDict):
obj: ReadOnly[_CliContext]
class _CredentialRow(TypedDict):
"""Single credential entry as returned by ``GET /credentials``."""
credential_name: ReadOnly[NotRequired[str]]
credential_info: ReadOnly[NotRequired[Mapping[str, object]]]
class _CredentialRowsView(TypedDict):
rows: ReadOnly[Sequence[_CredentialRow]]
class _JsonBodyView(TypedDict):
body: ReadOnly[object]
@click.group()
def credentials():
"""Manage credentials for the LiteLLM proxy server"""
@ -25,7 +53,8 @@ def credentials():
@click.pass_context
def list(ctx: click.Context, output_format: Literal["table", "json"]):
"""List all credentials"""
client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"])
context: Final[_CliContextView] = {"obj": ctx.obj}
client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
response: Final = client.list()
assert isinstance(response, dict)
@ -39,7 +68,8 @@ def list(ctx: click.Context, output_format: Literal["table", "json"]):
table.add_column("Custom LLM Provider", style="green")
# Add rows
for cred in response.get("credentials", []):
credential_rows: Final[_CredentialRowsView] = {"rows": response.get("credentials", [])}
for cred in credential_rows["rows"]:
info = cred.get("credential_info", {})
table.add_row(
str(cred.get("credential_name", "")),
@ -66,7 +96,8 @@ def list(ctx: click.Context, output_format: Literal["table", "json"]):
@click.pass_context
def create(ctx: click.Context, credential_name: str, info: str, values: str):
"""Create a new credential"""
client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"])
context: Final[_CliContextView] = {"obj": ctx.obj}
client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
try:
credential_info: Final = json.loads(info)
credential_values: Final = json.loads(values)
@ -79,8 +110,8 @@ def create(ctx: click.Context, credential_name: str, info: str, values: str):
except requests.exceptions.HTTPError as e:
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
try:
error_body: Final = e.response.json()
rich.print_json(data=error_body)
error_body: Final[_JsonBodyView] = {"body": e.response.json()}
rich.print_json(data=error_body["body"])
except json.JSONDecodeError:
click.echo(e.response.text, err=True)
raise click.Abort()
@ -91,15 +122,16 @@ def create(ctx: click.Context, credential_name: str, info: str, values: str):
@click.pass_context
def delete(ctx: click.Context, credential_name: str):
"""Delete a credential by name"""
client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"])
context: Final[_CliContextView] = {"obj": ctx.obj}
client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
try:
response: Final = client.delete(credential_name)
rich.print_json(data=response)
except requests.exceptions.HTTPError as e:
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
try:
error_body: Final = e.response.json()
rich.print_json(data=error_body)
error_body: Final[_JsonBodyView] = {"body": e.response.json()}
rich.print_json(data=error_body["body"])
except json.JSONDecodeError:
click.echo(e.response.text, err=True)
raise click.Abort()
@ -110,6 +142,7 @@ def delete(ctx: click.Context, credential_name: str):
@click.pass_context
def get(ctx: click.Context, credential_name: str):
"""Get a credential by name"""
client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"])
context: Final[_CliContextView] = {"obj": ctx.obj}
client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
response: Final = client.get(credential_name)
rich.print_json(data=response)

52
litellm/proxy/client/cli/commands/teams.py
View file

@ -1,21 +1,48 @@
"""Team management commands for LiteLLM CLI."""
from collections.abc import Mapping, Sequence
from typing import Any, Final
import click
import requests
from rich.console import Console
from rich.table import Table
from typing_extensions import ReadOnly, TypedDict
from litellm.proxy.client import Client
class _CliContext(TypedDict):
"""The proxy connection settings the CLI group stores on the click context."""
base_url: ReadOnly[str]
api_key: ReadOnly[str | None]
def _cli_context(ctx: click.Context) -> _CliContext:
"""The proxy connection settings the CLI group stored on the click context."""
ctx_obj: Final[_CliContext] = ctx.obj
return ctx_obj
def _proxy_client(ctx: click.Context) -> Client:
"""A proxy client for the base URL and API key on the click context."""
ctx_obj: Final = _cli_context(ctx)
return Client(ctx_obj["base_url"], ctx_obj["api_key"])
def _http_error_detail(error: requests.exceptions.HTTPError) -> object:
"""The ``detail`` the proxy reported for a failed request."""
error_body: Final[Mapping[str, object]] = error.response.json()
return error_body.get("detail", "Unknown error")
@click.group()
def teams():
"""Manage teams and team assignments"""
def display_teams_table(teams: list[dict[str, Any]]) -> None:
def display_teams_table(teams: Sequence[dict[str, Any]]) -> None:
"""Display teams in a formatted table"""
console: Final = Console()
@ -33,8 +60,8 @@ def display_teams_table(teams: list[dict[str, Any]]) -> None:
for i, team in enumerate(teams):
team_alias = team.get("team_alias") or "N/A"
team_id = team.get("team_id", "N/A")
models = team.get("models", [])
team_id: str = team.get("team_id", "N/A")
models: Sequence[str] = team.get("models", [])
max_budget = team.get("max_budget")
# Format models list
@ -64,7 +91,7 @@ def display_teams_table(teams: list[dict[str, Any]]) -> None:
@click.pass_context
def list(ctx: click.Context):
"""List teams that you belong to"""
client: Final = Client(ctx.obj["base_url"], ctx.obj["api_key"])
client: Final = _proxy_client(ctx)
try:
# Use list() for simpler response structure (returns array directly)
@ -72,8 +99,7 @@ def list(ctx: click.Context):
display_teams_table(teams)
except requests.exceptions.HTTPError as e:
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
error_body: Final = e.response.json()
click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True)
click.echo(f"Details: {_http_error_detail(e)}", err=True)
raise click.Abort()
except Exception as e:
click.echo(f"Error: {e}", err=True)
@ -84,7 +110,7 @@ def list(ctx: click.Context):
@click.pass_context
def available(ctx: click.Context):
"""List teams that are available to join"""
client: Final = Client(ctx.obj["base_url"], ctx.obj["api_key"])
client: Final = _proxy_client(ctx)
try:
teams: Final = client.teams.get_available()
@ -96,8 +122,7 @@ def available(ctx: click.Context):
click.echo("No available teams to join.")
except requests.exceptions.HTTPError as e:
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
error_body: Final = e.response.json()
click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True)
click.echo(f"Details: {_http_error_detail(e)}", err=True)
except Exception as e:
click.echo(f"Error: {e}", err=True)
raise click.Abort()
@ -108,8 +133,8 @@ def available(ctx: click.Context):
@click.pass_context
def assign_key(ctx: click.Context, team_id: str | None):
"""Assign your current CLI key to a team"""
client: Final = Client(ctx.obj["base_url"], ctx.obj["api_key"])
api_key: Final = ctx.obj["api_key"]
client: Final = _proxy_client(ctx)
api_key: Final = _cli_context(ctx)["api_key"]
if not api_key:
click.echo("No API key found. Please login first using 'litellm login'")
@ -145,7 +170,7 @@ def assign_key(ctx: click.Context, team_id: str | None):
teams = client.teams.list()
for team in teams:
if team.get("team_id") == team_id:
models = team.get("models", [])
models: Sequence[str] = team.get("models", [])
if models:
click.echo(f"You can now access models: {', '.join(models)}")
else:
@ -154,8 +179,7 @@ def assign_key(ctx: click.Context, team_id: str | None):
except requests.exceptions.HTTPError as e:
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
error_body: Final = e.response.json()
click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True)
click.echo(f"Details: {_http_error_detail(e)}", err=True)
raise click.Abort()
except Exception as e:
click.echo(f"Error: {e}", err=True)

14
litellm/proxy/common_utils/callback_utils.py
View file

@ -1,6 +1,6 @@
import copy
import os
from collections.abc import Callable, Iterable
from collections.abc import Callable, Iterable, Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias
@ -525,8 +525,8 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset(
def sanitize_openai_provider_metadata(
metadata: dict[str, Any] | None,
) -> dict[str, str] | None:
metadata: Mapping[str, object] | None,
) -> Mapping[str, object] | None:
"""
Keep only provider-safe OpenAI metadata entries (string keys -> string values).
@ -644,7 +644,7 @@ def process_callback(_callback: str, callback_type: str, environment_variables:
return {"name": _callback, "variables": env_vars_dict, "type": callback_type}
def normalize_callback_names(callbacks: Iterable[Any]) -> list[Any]:
def normalize_callback_names(callbacks: Iterable[object] | None) -> list[object]:
if callbacks is None:
return []
return [c.lower() if isinstance(c, str) else c for c in callbacks]
@ -674,7 +674,7 @@ def decrypt_callback_vars(metadata: Any) -> Any:
return _transform_callback_vars(metadata, _decrypt_or_passthrough)
def _transform_callback_vars(metadata: Any, transform: Callable[[str, Any], Any]) -> Any:
def _transform_callback_vars(metadata: object, transform: Callable[[str, Any], Any]) -> object:
if not isinstance(metadata, dict):
return metadata
out: Final = copy.deepcopy(metadata)
@ -704,7 +704,7 @@ def is_sensitive_callback_key(
return _CALLBACK_VAR_MASKER.is_sensitive_key(key)
def _encrypt_if_plaintext(key: str, value: Any) -> Any:
def _encrypt_if_plaintext(key: str, value: object) -> object:
if not isinstance(value, str) or not value:
return value
if not is_sensitive_callback_key(key):
@ -725,7 +725,7 @@ def _encrypt_if_plaintext(key: str, value: Any) -> Any:
return value
def _decrypt_or_passthrough(key: str, value: Any) -> Any:
def _decrypt_or_passthrough(key: str, value: object) -> object:
if not isinstance(value, str) or not value:
return value
if not value.startswith(_CALLBACK_VAR_ENCRYPTED_PREFIX):

65
litellm/proxy/common_utils/get_routes.py
View file

@ -2,63 +2,78 @@
Utility class for getting routes from a FastAPI app.
"""
from collections.abc import Sequence
from typing import Any, Final
from starlette.routing import BaseRoute
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
class RouteInfo(TypedDict, total=False):
"""One entry of the app's route listing."""
path: ReadOnly[object]
methods: ReadOnly[object]
name: ReadOnly[object]
endpoint: ReadOnly[str | None]
mounted_app: ReadOnly[bool]
class GetRoutes:
@staticmethod
def get_app_routes(
route: BaseRoute,
endpoint_route: Any,
) -> list[dict[str, Any]]:
) -> Sequence[RouteInfo]:
"""
Get routes for a regular route.
"""
routes: Final[list[dict[str, Any]]] = []
route_info: Final = {
route_info: Final[RouteInfo] = {
"path": getattr(route, "path", None),
"methods": getattr(route, "methods", None),
"name": getattr(route, "name", None),
"endpoint": (endpoint_route.__name__ if getattr(route, "endpoint", None) else None),
}
routes.append(route_info)
return routes
return [route_info]
@staticmethod
def get_routes_for_mounted_app(
route: BaseRoute,
) -> list[dict[str, Any]]:
) -> Sequence[RouteInfo]:
"""
Get routes for a mounted sub-application.
"""
routes: Final[list[dict[str, Any]]] = []
routes: Final[list[RouteInfo]] = []
mount_path: Final = getattr(route, "path", "")
sub_app: Final = getattr(route, "app", None)
if sub_app and hasattr(sub_app, "routes"):
for sub_route in sub_app.routes:
# Get endpoint - either from endpoint attribute or app attribute
endpoint_func = getattr(sub_route, "endpoint", None) or getattr(sub_route, "app", None)
for sub_route in GetRoutes._mounted_app_routes(route):
endpoint_func: object = getattr(sub_route, "endpoint", None) or getattr(sub_route, "app", None)
if endpoint_func is not None:
sub_route_path = getattr(sub_route, "path", "")
full_path = mount_path.rstrip("/") + sub_route_path
if endpoint_func is not None:
sub_route_path = getattr(sub_route, "path", "")
full_path = mount_path.rstrip("/") + sub_route_path
route_info = {
"path": full_path,
"methods": getattr(sub_route, "methods", ["GET", "POST"]),
"name": getattr(sub_route, "name", None),
"endpoint": GetRoutes._safe_get_endpoint_name(endpoint_func),
"mounted_app": True,
}
routes.append(route_info)
route_info: RouteInfo = {
"path": full_path,
"methods": getattr(sub_route, "methods", ["GET", "POST"]),
"name": getattr(sub_route, "name", None),
"endpoint": GetRoutes._safe_get_endpoint_name(endpoint_func),
"mounted_app": True,
}
routes.append(route_info)
return routes
@staticmethod
def _safe_get_endpoint_name(endpoint_function: Any) -> str | None:
def _mounted_app_routes(route: BaseRoute) -> Sequence[BaseRoute]:
"""The routes of the sub-application mounted at ``route``, if it mounts one."""
sub_app: Final[object] = getattr(route, "app", None)
if sub_app and hasattr(sub_app, "routes"):
return getattr(sub_app, "routes")
return ()
@staticmethod
def _safe_get_endpoint_name(endpoint_function: object) -> str | None:
"""
Safely get the name of the endpoint function.
"""
@ -66,7 +81,7 @@ class GetRoutes:
if hasattr(endpoint_function, "__name__"):
return getattr(endpoint_function, "__name__")
elif hasattr(endpoint_function, "__class__") and hasattr(endpoint_function.__class__, "__name__"):
return getattr(endpoint_function.__class__, "__name__")
return endpoint_function.__class__.__name__
else:
return None
except Exception:

48
litellm/proxy/common_utils/user_api_key_cache.py
View file

@ -40,30 +40,31 @@ class UserApiKeyCache(DualCache):
@overload
def get_cache(
self,
key: Any,
parent_otel_span: Any = None,
key: object,
parent_otel_span: object = None,
local_only: bool = False,
*,
model_type: type[T],
**kwargs: Any,
**kwargs: object,
) -> T | None: ...
@overload
def get_cache(
self,
key: Any,
parent_otel_span: Any = None,
key: object,
parent_otel_span: object = None,
local_only: bool = False,
**kwargs: Any,
model_type: None = None,
**kwargs: object,
) -> Any: ...
def get_cache(
self,
key,
parent_otel_span=None,
key: object,
parent_otel_span: object = None,
local_only: bool = False,
model_type: type[BaseModel] | None = None,
**kwargs,
**kwargs: object,
) -> Any | BaseModel | None:
if model_type is None and "model_type" in kwargs:
model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
@ -85,30 +86,31 @@ class UserApiKeyCache(DualCache):
@overload
async def async_get_cache(
self,
key: Any,
parent_otel_span: Any = None,
key: object,
parent_otel_span: object = None,
local_only: bool = False,
*,
model_type: type[T],
**kwargs: Any,
**kwargs: object,
) -> T | None: ...
@overload
async def async_get_cache(
self,
key: Any,
parent_otel_span: Any = None,
key: object,
parent_otel_span: object = None,
local_only: bool = False,
**kwargs: Any,
model_type: None = None,
**kwargs: object,
) -> Any: ...
async def async_get_cache(
self,
key,
parent_otel_span=None,
key: object,
parent_otel_span: object = None,
local_only: bool = False,
model_type: type[BaseModel] | None = None,
**kwargs,
**kwargs: object,
) -> Any | BaseModel | None:
if model_type is None and "model_type" in kwargs:
model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
@ -129,17 +131,17 @@ class UserApiKeyCache(DualCache):
return None
return decoded
def set_cache(self, key, value, local_only: bool = False, **kwargs):
def set_cache(self, key: object, value: object, local_only: bool = False, **kwargs: object):
model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
payload: Final = CacheCodec.serialize(value, model_type=model_type)
payload: Final[object] = 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):
async def async_set_cache(self, key: object, value: object, local_only: bool = False, **kwargs: object):
model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
payload: Final = CacheCodec.serialize(value, model_type=model_type)
payload: Final[object] = 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(self, cache_list: list, local_only: bool = False, **kwargs) -> None:
async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs: object) -> None:
"""
Batch writes with the same Codec boundary as ``async_set_cache`` without
``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged.

36
litellm/proxy/db/routing_prisma_wrapper.py
View file

@ -6,11 +6,14 @@ otherwise PrismaClient uses the writer-only PrismaWrapper directly.
import os
from collections.abc import Callable
from typing import Any, Final
from datetime import timedelta
from typing import Any, Final, TypeAlias
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.prisma_client import PrismaWrapper
ConnectTimeout: TypeAlias = "int | timedelta | None"
# Per-model action methods that read from the database. These are routed to
# the read replica when one is configured.
_MODEL_READ_METHODS: Final = frozenset(
@ -31,6 +34,11 @@ _MODEL_READ_METHODS: Final = frozenset(
_TOP_LEVEL_READ_METHODS: Final = frozenset({"query_first", "query_raw"})
def _dynamic_attr(target: object, name: str) -> object:
"""Fetch `name` off an unstubbed Prisma object as an opaque value."""
return getattr(target, name)
class _RoutedActions:
"""Per-model accessor that sends reads to the reader and writes to the writer.
@ -43,18 +51,18 @@ class _RoutedActions:
def __init__(
self,
writer_actions: Any,
reader_actions: Any,
writer_actions: object,
reader_actions: object,
should_use_reader: Callable[[], bool],
):
self._writer_actions = writer_actions
self._reader_actions = reader_actions
self._should_use_reader = should_use_reader
def __getattr__(self, name: str) -> Any:
def __getattr__(self, name: str) -> object:
if name in _MODEL_READ_METHODS and self._should_use_reader():
return getattr(self._reader_actions, name)
return getattr(self._writer_actions, name)
return _dynamic_attr(self._reader_actions, name)
return _dynamic_attr(self._writer_actions, name)
class RoutingPrismaWrapper:
@ -135,7 +143,7 @@ class RoutingPrismaWrapper:
return not self._reader_unavailable
@staticmethod
async def _try_connect(client: PrismaWrapper, *args: Any, **kwargs: Any) -> Exception | None:
async def _try_connect(client: PrismaWrapper, *args: ConnectTimeout, **kwargs: ConnectTimeout) -> Exception | None:
if client.is_connected() is True:
return None
try:
@ -144,7 +152,7 @@ class RoutingPrismaWrapper:
except Exception as e:
return e
async def connect(self, *args: Any, **kwargs: Any) -> None:
async def connect(self, *args: ConnectTimeout, **kwargs: ConnectTimeout) -> None:
writer_error: Final = await self._try_connect(self._writer, *args, **kwargs)
if writer_error is None:
self._writer_unavailable = False
@ -176,7 +184,7 @@ class RoutingPrismaWrapper:
writer_error,
)
async def disconnect(self, *args: Any, **kwargs: Any) -> None:
async def disconnect(self, *args: object, **kwargs: object) -> None:
first_error: BaseException | None = None
for client in (self._writer, self._reader):
try:
@ -206,7 +214,7 @@ class RoutingPrismaWrapper:
async def recreate_prisma_client(
self,
new_db_url: str,
http_client: Any | None = None,
http_client: object | None = None,
*,
expected_generation: int | None = None,
) -> bool:
@ -245,7 +253,7 @@ class RoutingPrismaWrapper:
)
return True
async def _recreate_reader(self, http_client: Any | None = None) -> None:
async def _recreate_reader(self, http_client: object | None = None) -> None:
"""Resolve the reader URL and recreate its Prisma client.
IAM-enabled readers regenerate their token (host/port/user came from
@ -265,14 +273,14 @@ class RoutingPrismaWrapper:
def __getattr__(self, name: str) -> Any:
if name in _TOP_LEVEL_READ_METHODS:
return getattr(self.read_target, name)
writer_attr: Final = getattr(self._writer, name)
return _dynamic_attr(self.read_target, name)
writer_attr: Final = _dynamic_attr(self._writer, name)
# Per-model action accessors are non-callable instances that expose
# both `find_many` and `create`. Methods like execute_raw / batch_ /
# tx are callables and stay on the writer untouched.
if not callable(writer_attr) and hasattr(writer_attr, "find_many") and hasattr(writer_attr, "create"):
try:
reader_attr: Final = getattr(self._reader, name)
reader_attr: Final = _dynamic_attr(self._reader, name)
except AttributeError:
return writer_attr
return _RoutedActions(writer_attr, reader_attr, self._should_use_reader)

14
litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py
View file

@ -15,6 +15,8 @@ restriction intact.
"""
import operator
from collections.abc import Callable
from types import CodeType
from typing import Any, Final
from RestrictedPython import (
@ -58,7 +60,7 @@ class AsyncAwareTransformer(RestrictingNodeTransformer):
return self.node_contents_visit(node)
_INPLACE_OPS: Final[dict[str, Any]] = {
_INPLACE_OPS: Final[dict[str, Callable[[object, object], object]]] = {
"+=": operator.iadd,
"-=": operator.isub,
"*=": operator.imul,
@ -75,7 +77,7 @@ _INPLACE_OPS: Final[dict[str, Any]] = {
}
def _inplacevar_(op: str, x: Any, y: Any) -> Any:
def _inplacevar_(op: str, x: object, y: object) -> object:
# RestrictedPython rewrites ``x += 1`` on a simple name into
# ``x = _inplacevar_("+=", x, 1)``. The package deliberately ships no
# default, so we dispatch through ``operator``'s in-place helpers, which
@ -86,7 +88,7 @@ def _inplacevar_(op: str, x: Any, y: Any) -> Any:
return fn(x, y)
def _build_sandbox_builtins() -> dict[str, Any]:
def _build_sandbox_builtins() -> dict[str, object]:
# ``limited_builtins`` overrides ``list``/``tuple``/``range`` from
# ``safe_builtins`` with bounds-checking variants (e.g. ``limited_range``
# rejects ``range(10**18)``). ``utility_builtins`` adds ``set``,
@ -98,14 +100,14 @@ def _build_sandbox_builtins() -> dict[str, Any]:
}
def build_sandbox_globals() -> dict[str, Any]:
def build_sandbox_globals() -> dict[str, object]:
"""Assemble the globals dict for executing guardrail code.
Includes the LiteLLM-provided primitives (``regex_match``, ``http_get``,
``allow``/``block``/``modify``, etc.) plus the RestrictedPython guards
that the compiled bytecode expects to find by name.
"""
sandbox: Final[dict[str, Any]] = get_custom_code_primitives().copy()
sandbox: Final[dict[str, object]] = get_custom_code_primitives().copy()
sandbox["__builtins__"] = _build_sandbox_builtins()
sandbox["_getattr_"] = safer_getattr
sandbox["_getitem_"] = default_guarded_getitem
@ -116,7 +118,7 @@ def build_sandbox_globals() -> dict[str, Any]:
return sandbox
def compile_sandboxed(source: str, filename: str = "<guardrail>") -> Any:
def compile_sandboxed(source: str, filename: str = "<guardrail>") -> CodeType:
"""Compile guardrail source with RestrictedPython's AST transformer.
Raises ``SyntaxError`` on either a Python syntax error or a restricted

33
litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py
View file

@ -2,7 +2,7 @@ from __future__ import annotations
import os
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
from urllib.parse import urlparse
from uuid import uuid4
@ -11,7 +11,7 @@ import requests
from fastapi import HTTPException
from httpx import HTTPStatusError
from requests.auth import HTTPBasicAuth
from typing_extensions import ReadOnly
from typing_extensions import ReadOnly, TypedDict, Unpack
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
@ -36,24 +36,31 @@ if TYPE_CHECKING:
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
"""Base-class constructor options carried by this guardrail's forwarded keyword arguments."""
guardrail_name: ReadOnly[str | None]
supported_event_hooks: list[GuardrailEventHooks] | None
class _HiddenlayerEvaluation(TypedDict, total=False):
action: str
threat_level: str
action: ReadOnly[str]
threat_level: ReadOnly[str]
class _HiddenlayerAnalysisEntry(TypedDict, total=False):
name: str
detected: bool
name: ReadOnly[str]
detected: ReadOnly[bool]
class _HiddenlayerModifiedSide(TypedDict):
messages: Any
messages: ReadOnly[Any]
class _HiddenlayerResponse(TypedDict, total=False):
evaluation: _HiddenlayerEvaluation
analysis: Sequence[_HiddenlayerAnalysisEntry]
modified_data: Mapping[str, _HiddenlayerModifiedSide]
evaluation: ReadOnly[_HiddenlayerEvaluation]
analysis: ReadOnly[Sequence[_HiddenlayerAnalysisEntry]]
modified_data: ReadOnly[Mapping[str, _HiddenlayerModifiedSide]]
class _LoggedCallMetadata(TypedDict, total=False):
@ -151,7 +158,7 @@ class HiddenlayerGuardrail(CustomGuardrail):
api_key: str | None = None,
api_base: str | None = None,
auth_url: str | None = None,
**kwargs: Any,
**kwargs: Unpack[_CustomGuardrailOptions],
) -> None:
kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID")
@ -356,7 +363,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
api_key: str | None = None,
api_base: str | None = None,
auth_url: str | None = None,
**kwargs: Any,
**kwargs: Unpack[_CustomGuardrailOptions],
) -> None:
self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID")
self.hiddenlayer_client_secret = api_key or os.getenv("HIDDENLAYER_CLIENT_SECRET")
@ -486,7 +493,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
self,
payload: Any,
input_type: Literal["request", "response"],
hl_headers: dict[str, str],
hl_headers: Mapping[str, str],
) -> httpx.Response:
if input_type == "request":
path = "detection/v2/request-evaluations"

130
litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py
View file

@ -1,10 +1,11 @@
"""LLM-as-a-Judge guardrail: uses an LLM to score responses against weighted criteria."""
from collections.abc import Callable
from collections.abc import Callable, Mapping, MutableMapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
from fastapi import HTTPException
from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack
import litellm
from litellm._logging import verbose_logger
@ -41,28 +42,87 @@ _parse_judge_verdict: Final = parse_json_verdict
_extract_text_from_content: Final = extract_text_from_content
class _JudgeMessage(TypedDict):
"""Chat message, as far as the judge prompt builder reads it."""
role: ReadOnly[NotRequired[str]]
content: ReadOnly[NotRequired[object]]
class _GuardrailOptions(TypedDict, total=False):
"""Base :class:`CustomGuardrail` options forwarded untouched."""
mask_request_content: ReadOnly[bool]
mask_response_content: ReadOnly[bool]
violation_message_template: ReadOnly[str | None]
end_session_after_n_fails: ReadOnly[int | None]
on_violation: ReadOnly[str | None]
realtime_violation_message: ReadOnly[str | None]
on_sensitive_data: ReadOnly[str | None]
sensitive_data_route_to_model: ReadOnly[str | None]
sticky_session_routing: ReadOnly[bool]
run_in_parallel: ReadOnly[bool]
only_scan_new_messages: ReadOnly[bool]
class _RequestMessagesView(TypedDict):
messages: ReadOnly[Sequence[_JudgeMessage]]
class _RequestMetadataView(TypedDict):
metadata: ReadOnly[MutableMapping[str, object]]
class _OverallScoreView(TypedDict):
overall_score: ReadOnly[str | float]
class _JudgeModelView(TypedDict):
judge_model: ReadOnly[str]
class _CriteriaView(TypedDict):
criteria: ReadOnly[Sequence[Mapping[str, str | float]]]
class _OnFailureView(TypedDict):
on_failure: ReadOnly[Literal["block", "log"]]
class _ThresholdView(TypedDict):
overall_threshold: ReadOnly[str | float]
class _ModeView(TypedDict):
mode: ReadOnly[object]
class _DefaultOnView(TypedDict):
default_on: ReadOnly[object]
def _get_litellm_param(
litellm_params: "LitellmParams",
guardrail: "Guardrail",
key: str,
default: Any = None,
default: str | float | bool | None = None,
) -> Any:
val: Final = getattr(litellm_params, key, None)
val: Final[object] = getattr(litellm_params, key, None)
if val is not None:
return val
raw: Final = guardrail.get("litellm_params")
if isinstance(raw, dict) and key in raw:
return raw[key]
if raw is not None and not isinstance(raw, dict):
attr: Final = getattr(raw, key, None)
attr: Final[object] = getattr(raw, key, None)
if attr is not None:
return attr
return default
def _build_judge_prompt(
criteria: list[dict[str, Any]],
messages: list[dict[str, Any]],
criteria: Sequence[Mapping[str, object]],
messages: Sequence[_JudgeMessage],
response_text: str,
) -> str:
criteria_block: Final = "\n".join(
@ -87,13 +147,13 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
self,
guardrail_name: str,
judge_model: str,
criteria: list[dict[str, Any]],
criteria: Sequence[Mapping[str, object]],
overall_threshold: float = 80.0,
on_failure: Literal["block", "log"] = "block",
event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None,
default_on: bool = False,
router_provider: "Callable[[], Router | None] | None" = None,
**kwargs: Any,
**kwargs: Unpack[_GuardrailOptions],
) -> None:
_event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None
if event_hook is not None:
@ -121,9 +181,9 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
async def _run_judge(
self,
messages: list[dict[str, Any]],
messages: Sequence[_JudgeMessage],
response_text: str,
) -> dict[str, Any]:
) -> dict[str, object]:
judge_messages: Final = [
{"role": "system", "content": JUDGE_SYSTEM_PROMPT},
{
@ -162,10 +222,10 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
judge_result: dict[str, Any] = {}
try:
messages: Final[list[dict[str, Any]]] = request_data.get("messages") or []
request_messages: Final[_RequestMessagesView] = {"messages": request_data.get("messages") or []}
try:
judge_result = await self._run_judge(messages, response_text)
judge_result = await self._run_judge(request_messages["messages"], response_text)
except Exception as judge_err:
verbose_logger.warning(
"llm_as_a_judge guardrail: judge call failed, failing open. Error: %s", judge_err
@ -174,7 +234,8 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
return inputs
try:
overall_score: Final = max(0.0, min(100.0, float(judge_result.get("overall_score", 100))))
raw_score: Final[_OverallScoreView] = {"overall_score": judge_result.get("overall_score", 100)}
overall_score: Final = max(0.0, min(100.0, float(raw_score["overall_score"])))
except (TypeError, ValueError):
verbose_logger.warning("llm_as_a_judge: invalid overall_score from judge, failing open")
return inputs
@ -189,7 +250,8 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
"threshold": self.overall_threshold,
"verdicts": judge_result.get("verdicts", []),
}
_metadata: Final = request_data.setdefault("metadata", {})
request_metadata: Final[_RequestMetadataView] = {"metadata": request_data.setdefault("metadata", {})}
_metadata: Final = request_metadata["metadata"]
existing: Final = _metadata.get("eval_information")
if isinstance(existing, list):
existing.append(eval_info)
@ -238,37 +300,45 @@ def initialize_guardrail(
if not guardrail_name:
raise ValueError("llm_as_a_judge guardrail requires a guardrail_name")
judge_model: Final = _get_litellm_param(litellm_params, guardrail, "judge_model")
if not judge_model:
judge_model: Final[_JudgeModelView] = {"judge_model": _get_litellm_param(litellm_params, guardrail, "judge_model")}
if not judge_model["judge_model"]:
raise ValueError("llm_as_a_judge guardrail requires judge_model in litellm_params")
criteria: Final = _get_litellm_param(litellm_params, guardrail, "criteria") or []
if not criteria:
criteria: Final[_CriteriaView] = {"criteria": _get_litellm_param(litellm_params, guardrail, "criteria") or []}
if not criteria["criteria"]:
raise ValueError("llm_as_a_judge guardrail requires at least one criterion")
weight_total: Final = sum(float(c.get("weight", 0)) for c in criteria)
weight_total: Final = sum(float(c.get("weight", 0)) for c in criteria["criteria"])
if abs(weight_total - 100) > 0.5:
raise ValueError(f"llm_as_a_judge criterion weights must sum to 100 (got {weight_total})")
on_failure: Final = _get_litellm_param(litellm_params, guardrail, "on_failure", "block")
if on_failure not in _VALID_ON_FAILURE:
raise ValueError(f"llm_as_a_judge on_failure must be 'block' or 'log', got '{on_failure}'")
on_failure: Final[_OnFailureView] = {
"on_failure": _get_litellm_param(litellm_params, guardrail, "on_failure", "block")
}
if on_failure["on_failure"] not in _VALID_ON_FAILURE:
raise ValueError(f"llm_as_a_judge on_failure must be 'block' or 'log', got '{on_failure['on_failure']}'")
overall_threshold: Final = float(_get_litellm_param(litellm_params, guardrail, "overall_threshold", 80.0))
threshold: Final[_ThresholdView] = {
"overall_threshold": _get_litellm_param(litellm_params, guardrail, "overall_threshold", 80.0)
}
overall_threshold: Final = float(threshold["overall_threshold"])
mode: Final = _get_litellm_param(litellm_params, guardrail, "mode")
mode: Final[_ModeView] = {"mode": _get_litellm_param(litellm_params, guardrail, "mode")}
event_hook: GuardrailEventHooks | None = None
if isinstance(mode, str) and mode in {e.value for e in GuardrailEventHooks}:
event_hook = GuardrailEventHooks(mode)
if isinstance(mode["mode"], str) and mode["mode"] in {e.value for e in GuardrailEventHooks}:
event_hook = GuardrailEventHooks(mode["mode"])
default_on: Final[_DefaultOnView] = {
"default_on": _get_litellm_param(litellm_params, guardrail, "default_on", False)
}
instance: Final = LLMAsAJudgeGuardrail(
guardrail_name=guardrail_name,
judge_model=judge_model,
criteria=criteria,
judge_model=judge_model["judge_model"],
criteria=criteria["criteria"],
overall_threshold=overall_threshold,
on_failure=on_failure,
on_failure=on_failure["on_failure"],
event_hook=event_hook,
default_on=bool(_get_litellm_param(litellm_params, guardrail, "default_on", False)),
default_on=bool(default_on["default_on"]),
)
litellm.logging_callback_manager.add_litellm_callback(instance)
return instance

26
litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py
View file

@ -7,8 +7,9 @@
import enum
import json
import os
from collections.abc import Callable, Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, cast
from urllib.parse import urlparse
from litellm._logging import verbose_proxy_logger
@ -36,6 +37,8 @@ _AIDR_SCAN_ENDPOINT: Final = "/litellm/guardrail"
_INTERVENED_INPUT_FIELDS: Final = ("texts", "images", "tools", "tool_calls")
_DEFAULT_API_BASE_HOSTNAME: Final = urlparse(_DEFAULT_API_BASE).hostname
_GuardrailJsonResponse: TypeAlias = Exception | str | dict[str, object]
_KEYS_DUPLICATING_SCAN_INPUTS: Final = ("messages", "input")
_LOGGING_KEYS_DUPLICATING_SCAN_INPUTS: Final = _KEYS_DUPLICATING_SCAN_INPUTS + (
"additional_args",
@ -119,7 +122,7 @@ class NomaV2Guardrail(CustomGuardrail):
def _resolve_action_from_response(
self,
response_json: dict,
response_json: Mapping[str, object],
) -> _Action:
action: Final = response_json.get("action")
if isinstance(action, str):
@ -165,10 +168,11 @@ class NomaV2Guardrail(CustomGuardrail):
@staticmethod
def _sanitize_payload_for_transport(payload: dict) -> dict:
def _default(obj: Any) -> Any:
if hasattr(obj, "model_dump"):
def _default(obj: object) -> object:
model_dump: Final[Callable[[], Mapping[str, object]] | None] = getattr(obj, "model_dump", None)
if model_dump is not None:
try:
return obj.model_dump()
return model_dump()
except Exception:
pass
return str(obj)
@ -178,7 +182,7 @@ class NomaV2Guardrail(CustomGuardrail):
except (ValueError, TypeError):
json_str = safe_dumps(payload)
safe_payload: Final = safe_json_loads(json_str, default={})
safe_payload: Final[object] = safe_json_loads(json_str, default={})
if safe_payload == {} and payload:
verbose_proxy_logger.warning(
"Noma v2 guardrail: payload serialization failed, falling back to empty payload"
@ -196,7 +200,7 @@ class NomaV2Guardrail(CustomGuardrail):
async def _call_noma_scan(
self,
payload: dict,
) -> dict:
) -> dict[str, object]:
headers: Final[dict[str, str]] = {"Content-Type": "application/json"}
authorization_header: Final = self._get_authorization_header()
if authorization_header:
@ -215,7 +219,7 @@ class NomaV2Guardrail(CustomGuardrail):
response.text,
)
response.raise_for_status()
response_json: Final = response.json()
response_json: Final[dict[str, object]] = response.json()
verbose_proxy_logger.debug(
"Noma v2 AIDR response parsed: %s",
json.dumps(response_json, default=str),
@ -227,7 +231,7 @@ class NomaV2Guardrail(CustomGuardrail):
request_data: dict,
start_time: datetime,
guardrail_status: GuardrailStatus,
guardrail_json_response: Any,
guardrail_json_response: _GuardrailJsonResponse,
) -> None:
end_time: Final = datetime.now()
duration: Final = (end_time - start_time).total_seconds()
@ -270,11 +274,11 @@ class NomaV2Guardrail(CustomGuardrail):
) -> GenericGuardrailAPIInputs:
start_time: Final = datetime.now()
guardrail_status: GuardrailStatus = "success"
guardrail_json_response: Any = {}
guardrail_json_response: _GuardrailJsonResponse = {}
dynamic_params = self.get_guardrail_dynamic_request_body_params(request_data)
if not isinstance(dynamic_params, dict):
dynamic_params = {}
response_json: dict | None = None
response_json: dict[str, object] | None = None
# Per-request dynamic params can override configured application context.
application_id = self._get_non_empty_str(dynamic_params.get("application_id"))

20
litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py
View file

@ -7,8 +7,11 @@ before and after LLM calls.
"""
import os
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.exceptions import GuardrailRaisedException
from litellm.integrations.custom_guardrail import (
@ -20,6 +23,7 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
@ -34,6 +38,16 @@ _DEFAULT_API_BASE: Final = "https://api.promptguard.co"
_GUARD_ENDPOINT: Final = "/api/v1/guard"
class PromptGuardResult(TypedDict, total=False):
"""The fields this guardrail reads off a PromptGuard Guard API response."""
decision: ReadOnly[str]
threat_type: ReadOnly[str]
event_id: ReadOnly[str]
confidence: ReadOnly[float]
redacted_messages: ReadOnly[list[AllMessageValues]]
class PromptGuardMissingCredentials(Exception):
pass
@ -96,7 +110,7 @@ class PromptGuardGuardrail(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
request_data: Mapping[str, object],
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> GenericGuardrailAPIInputs:
@ -114,7 +128,7 @@ class PromptGuardGuardrail(CustomGuardrail):
direction: Final = "input" if input_type == "request" else "output"
payload: Final[dict[str, Any]] = {
payload: Final[dict[str, object]] = {
"messages": messages,
"direction": direction,
}
@ -144,7 +158,7 @@ class PromptGuardGuardrail(CustomGuardrail):
timeout=10.0,
)
response.raise_for_status()
result: Final = response.json()
result: Final[PromptGuardResult] = response.json()
except Exception as exc:
verbose_proxy_logger.error("PromptGuard API error: %s", str(exc))
if self.block_on_error:

76
litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py
View file

@ -1,4 +1,6 @@
from typing import Final
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import Final, Protocol
from fastapi import APIRouter, Depends, HTTPException, Query
@ -18,7 +20,59 @@ from litellm.repositories.table_repositories import JWTKeyMappingRepository
router: Final = APIRouter()
def _to_response(mapping) -> JWTKeyMappingResponse:
class _JWTKeyMappingRecord(Protocol):
"""A ``LiteLLM_JWTKeyMapping`` row, viewed through the columns these endpoints read."""
@property
def id(self) -> str: ...
@property
def jwt_claim_name(self) -> str: ...
@property
def jwt_claim_value(self) -> str: ...
@property
def description(self) -> str | None: ...
@property
def is_active(self) -> bool: ...
@property
def created_at(self) -> datetime: ...
@property
def updated_at(self) -> datetime: ...
@property
def created_by(self) -> str | None: ...
@property
def updated_by(self) -> str | None: ...
class _JWTKeyMappingTable(Protocol):
"""The Prisma table actions these endpoints issue against the JWT key mapping table."""
async def create(self, *, data: Mapping[str, object]) -> _JWTKeyMappingRecord: ...
async def find_unique(self, *, where: Mapping[str, object]) -> _JWTKeyMappingRecord | None: ...
async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> _JWTKeyMappingRecord: ...
async def delete(self, *, where: Mapping[str, object]) -> _JWTKeyMappingRecord | None: ...
async def find_many(self, *, skip: int, take: int, order: Mapping[str, str]) -> Sequence[_JWTKeyMappingRecord]: ...
async def count(self) -> int: ...
def _mapping_table(prisma_client: object) -> _JWTKeyMappingTable:
"""View the JWT key mapping repository's untyped Prisma table through the actions used here."""
return JWTKeyMappingRepository(prisma_client).table
def _to_response(mapping: _JWTKeyMappingRecord) -> JWTKeyMappingResponse:
"""Convert a Prisma mapping object to a safe response (no hashed token)."""
return JWTKeyMappingResponse(
id=mapping.id,
@ -62,7 +116,7 @@ async def create_jwt_key_mapping(
if data.description is not None:
create_data["description"] = data.description
new_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.create(data=create_data)
new_mapping: Final = await _mapping_table(prisma_client).create(data=create_data)
# Invalidate cache
cache_key: Final = f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}"
@ -110,7 +164,7 @@ async def update_jwt_key_mapping(
try:
# Get old mapping for cache invalidation
old_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": data.id})
old_mapping: Final = await _mapping_table(prisma_client).find_unique(where={"id": data.id})
if old_mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
@ -118,9 +172,7 @@ async def update_jwt_key_mapping(
cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}"
await user_api_key_cache.async_delete_cache(cache_key)
updated_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.update(
where={"id": data.id}, data=update_data
)
updated_mapping: Final = await _mapping_table(prisma_client).update(where={"id": data.id}, data=update_data)
# Invalidate new cache key if claim fields changed
cache_key = f"jwt_key_mapping:{updated_mapping.jwt_claim_name}:{updated_mapping.jwt_claim_value}"
@ -159,7 +211,7 @@ async def delete_jwt_key_mapping(
try:
# Get old mapping for cache invalidation
old_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": data.id})
old_mapping: Final = await _mapping_table(prisma_client).find_unique(where={"id": data.id})
if old_mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
@ -167,7 +219,7 @@ async def delete_jwt_key_mapping(
cache_key: Final = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}"
await user_api_key_cache.async_delete_cache(cache_key)
await JWTKeyMappingRepository(prisma_client).table.delete(where={"id": data.id})
await _mapping_table(prisma_client).delete(where={"id": data.id})
return {"status": "success"}
except HTTPException:
raise
@ -195,12 +247,12 @@ async def list_jwt_key_mappings(
try:
skip: Final = (page - 1) * size
mappings: Final = await JWTKeyMappingRepository(prisma_client).table.find_many(
mappings: Final = await _mapping_table(prisma_client).find_many(
skip=skip,
take=size,
order={"created_at": "desc"},
)
total_count: Final = await JWTKeyMappingRepository(prisma_client).table.count()
total_count: Final = await _mapping_table(prisma_client).count()
return {
"mappings": [_to_response(m) for m in mappings],
"total_count": total_count,
@ -232,7 +284,7 @@ async def info_jwt_key_mapping(
raise HTTPException(status_code=500, detail="Database not connected")
try:
mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": id})
mapping: Final = await _mapping_table(prisma_client).find_unique(where={"id": id})
if mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
return _to_response(mapping)

47
litellm/proxy/response_polling/background_streaming.py
View file

@ -10,8 +10,8 @@ https://platform.openai.com/docs/api-reference/responses-streaming
import asyncio
import json
from collections.abc import Sequence
from typing import TYPE_CHECKING, Final, TypedDict, cast
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final, TypedDict
from fastapi import Request, Response
from fastapi.responses import StreamingResponse
@ -38,6 +38,36 @@ class _StreamOutputItem(TypedDict, total=False):
content: ReadOnly[Sequence[_StreamContentPart | None]]
class _StreamTerminalResponse(TypedDict, total=False):
"""Fields of the ``response`` payload carried by a terminal streaming event."""
status: ReadOnly[ResponsesAPIStatus]
tool_choice: ReadOnly[object]
model: ReadOnly[str]
instructions: ReadOnly[str]
temperature: ReadOnly[float]
top_p: ReadOnly[float]
max_output_tokens: ReadOnly[int]
previous_response_id: ReadOnly[str]
truncation: ReadOnly[str]
parallel_tool_calls: ReadOnly[bool]
user: ReadOnly[str]
store: ReadOnly[bool]
output: ReadOnly[Sequence[_StreamOutputItem]]
class _StreamEvent(TypedDict, total=False):
"""One decoded ``data:`` frame of an OpenAI Responses streaming body."""
type: ReadOnly[str]
item: ReadOnly[_StreamOutputItem]
item_id: ReadOnly[str]
part: ReadOnly[_StreamContentPart]
content_index: ReadOnly[int]
delta: ReadOnly[str]
response: ReadOnly[_StreamTerminalResponse]
async def background_streaming_task(
polling_id: str,
data,
@ -139,7 +169,7 @@ async def background_streaming_task(
None # Will be set by response.completed/failed/incomplete/cancelled
)
terminal_error = None
_event_to_status: Final = {
_event_to_status: Final[Mapping[str, ResponsesAPIStatus]] = {
"response.completed": "completed",
"response.failed": "failed",
"response.incomplete": "incomplete",
@ -180,7 +210,7 @@ async def background_streaming_task(
break
try:
event = json.loads(chunk_data)
event: _StreamEvent = json.loads(chunk_data)
event_type = event.get("type", "")
# Process different event types based on OpenAI streaming spec
@ -288,12 +318,9 @@ async def background_streaming_task(
# Terminal event - extract all ResponsesAPIResponse fields
# https://platform.openai.com/docs/api-reference/responses-streaming
response_data = event.get("response", {})
terminal_status = cast(
ResponsesAPIStatus,
response_data.get(
"status",
_event_to_status.get(event_type, "completed"),
),
terminal_status = response_data.get(
"status",
_event_to_status.get(event_type, "completed"),
)
# Extract error for failed and incomplete responses

57
litellm/proxy/search_endpoints/search_tool_registry.py
View file

@ -2,8 +2,9 @@
Search Tool Registry for managing search tool configurations.
"""
from collections.abc import Iterator, Mapping, Sequence
from datetime import datetime, timezone
from typing import Final
from typing import Final, Protocol
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
@ -13,6 +14,40 @@ from litellm.repositories.table_repositories import SearchToolsRepository
from litellm.types.search import SearchTool
class SearchToolRecord(Protocol):
search_tool_id: str
search_tool_name: str
created_at: datetime
updated_at: datetime
def __iter__(self) -> Iterator[tuple[str, object]]: ...
class SearchToolTableClient(Protocol):
async def create(self, data: Mapping[str, object]) -> SearchToolRecord: ...
async def find_unique(self, where: Mapping[str, object]) -> SearchToolRecord | None: ...
async def find_many(self, order: Mapping[str, str] | None = None) -> Sequence[SearchToolRecord]: ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> SearchToolRecord: ...
async def delete(self, where: Mapping[str, object]) -> SearchToolRecord: ...
class _SearchToolsRepositoryView(Protocol):
@property
def table(self) -> SearchToolTableClient: ...
def _search_tools_table_of(repository: _SearchToolsRepositoryView) -> SearchToolTableClient:
return repository.table
def _search_tools_table(prisma_client: PrismaClient) -> SearchToolTableClient:
return _search_tools_table_of(SearchToolsRepository(prisma_client))
class SearchToolRegistry:
"""
Handles adding, removing, and getting search tools in DB + in memory.
@ -22,7 +57,7 @@ class SearchToolRegistry:
pass
@staticmethod
def _convert_prisma_to_dict(prisma_obj) -> dict:
def _convert_prisma_to_dict(prisma_obj: SearchToolRecord) -> dict:
"""
Convert Prisma result to dict with datetime objects as ISO format strings.
@ -35,9 +70,9 @@ class SearchToolRegistry:
result: Final = dict(prisma_obj)
# Convert datetime objects to ISO format strings
if "created_at" in result and result["created_at"]:
result["created_at"] = result["created_at"].isoformat()
result["created_at"] = prisma_obj.created_at.isoformat()
if "updated_at" in result and result["updated_at"]:
result["updated_at"] = result["updated_at"].isoformat()
result["updated_at"] = prisma_obj.updated_at.isoformat()
return result
###########################################################
@ -61,7 +96,7 @@ class SearchToolRegistry:
search_tool_info: Final[str] = safe_dumps(search_tool.get("search_tool_info", {}))
# Create search tool in DB
created_search_tool: Final = await SearchToolsRepository(prisma_client).table.create(
created_search_tool: Final = await _search_tools_table(prisma_client).create(
data={
"search_tool_name": search_tool_name,
"litellm_params": litellm_params,
@ -95,7 +130,7 @@ class SearchToolRegistry:
"""
try:
# Get search tool before deletion for response
existing_tool: Final = await SearchToolsRepository(prisma_client).table.find_unique(
existing_tool: Final = await _search_tools_table(prisma_client).find_unique(
where={"search_tool_id": search_tool_id}
)
@ -103,7 +138,7 @@ class SearchToolRegistry:
raise Exception(f"Search tool with ID {search_tool_id} not found")
# Delete from DB
await SearchToolsRepository(prisma_client).table.delete(where={"search_tool_id": search_tool_id})
await _search_tools_table(prisma_client).delete(where={"search_tool_id": search_tool_id})
return {
"message": f"Search tool {search_tool_id} deleted successfully",
@ -131,7 +166,7 @@ class SearchToolRegistry:
search_tool_info: Final[str] = safe_dumps(search_tool.get("search_tool_info", {}))
# Update in DB
updated_search_tool: Final = await SearchToolsRepository(prisma_client).table.update(
updated_search_tool: Final = await _search_tools_table(prisma_client).update(
where={"search_tool_id": search_tool_id},
data={
"search_tool_name": search_tool_name,
@ -163,7 +198,7 @@ class SearchToolRegistry:
try:
search_tools_from_db: Final = await call_with_db_reconnect_retry(
prisma_client,
lambda: SearchToolsRepository(prisma_client).table.find_many(
lambda: _search_tools_table(prisma_client).find_many(
order={"created_at": "desc"},
),
reason="get_all_search_tools_from_db_lookup_failure",
@ -194,7 +229,7 @@ class SearchToolRegistry:
Search tool configuration or None if not found
"""
try:
search_tool: Final = await SearchToolsRepository(prisma_client).table.find_unique(
search_tool: Final = await _search_tools_table(prisma_client).find_unique(
where={"search_tool_id": search_tool_id}
)
@ -222,7 +257,7 @@ class SearchToolRegistry:
Search tool configuration or None if not found
"""
try:
search_tool: Final = await SearchToolsRepository(prisma_client).table.find_unique(
search_tool: Final = await _search_tools_table(prisma_client).find_unique(
where={"search_tool_name": search_tool_name}
)

65
litellm/rag/rag_query.py
View file

@ -1,11 +1,45 @@
from collections.abc import Sequence
from typing import Any, Final
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
from litellm.types.utils import ModelResponse
from litellm.types.vector_stores import (
VectorStoreResultContent,
VectorStoreSearchResponse,
)
from litellm.types.vector_stores import VectorStoreSearchResponse
class _ResultContentView(TypedDict):
"""Content entry carried by a vector store search result."""
type: ReadOnly[NotRequired[str]]
text: ReadOnly[str]
class _SearchResultView(TypedDict):
"""Vector store search result, as far as :class:`RAGQuery` reads it."""
content: ReadOnly[NotRequired[Sequence[_ResultContentView]]]
text: ReadOnly[NotRequired[str]]
class _SearchDataView(TypedDict):
results: ReadOnly[Sequence[_SearchResultView]]
class _ContextChunksView(TypedDict):
chunks: ReadOnly[Sequence[_SearchResultView | str | None]]
class _RerankResultView(TypedDict):
index: ReadOnly[NotRequired[int]]
class _RerankResultsView(TypedDict):
results: ReadOnly[Sequence[_RerankResultView]]
class _MessageView(TypedDict):
message: ReadOnly[object]
class RAGQuery:
@ -42,9 +76,10 @@ class RAGQuery:
"""
context_content = RAGQuery.CONTENT_PREFIX_STRING
for chunk in context_chunks:
chunks: Final[_ContextChunksView] = {"chunks": context_chunks}
for chunk in chunks["chunks"]:
if isinstance(chunk, dict):
result_content: list[VectorStoreResultContent] | None = chunk.get("content")
result_content: Sequence[_ResultContentView] | None = chunk.get("content")
if result_content:
for content_item in result_content:
content_text: str | None = content_item.get("text")
@ -64,14 +99,15 @@ class RAGQuery:
def add_search_results_to_response(
response: ModelResponse,
search_results: VectorStoreSearchResponse,
rerank_results: Any | None = None,
rerank_results: object = None,
) -> ModelResponse:
"""
Add search results to the response choices.
"""
if hasattr(response, "choices") and response.choices:
for choice in response.choices:
message = getattr(choice, "message", None)
message_view: _MessageView = {"message": getattr(choice, "message", None)}
message = message_view["message"]
if message is not None:
# Get existing provider_specific_fields or create new dict
provider_fields = getattr(message, "provider_specific_fields", None) or {}
@ -91,7 +127,8 @@ class RAGQuery:
) -> list[str | dict[str, Any]]:
"""Extract text documents from vector store search response."""
documents: Final[list[str | dict[str, Any]]] = []
for result in search_response.get("data", []):
search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])}
for result in search_data["results"]:
content_list = result.get("content", [])
for content in content_list:
if content.get("type") == "text" and content.get("text"):
@ -99,11 +136,13 @@ class RAGQuery:
return documents
@staticmethod
def get_top_chunks_from_rerank(search_response: Any, rerank_response: Any) -> list[Any]:
def get_top_chunks_from_rerank(search_response: Any, rerank_response: Any) -> list[_SearchResultView]:
"""Get the original search results corresponding to the top reranked results."""
top_chunks: Final = []
original_results: Final = search_response.get("data", [])
for result in rerank_response.get("results", []):
top_chunks: Final[list[_SearchResultView]] = []
search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])}
original_results: Final = search_data["results"]
reranked: Final[_RerankResultsView] = {"results": rerank_response.get("results", [])}
for result in reranked["results"]:
index = result.get("index")
if index is not None and index < len(original_results):
top_chunks.append(original_results[index])

51
litellm/repositories/base_repository.py
View file

@ -3,7 +3,7 @@ Base repository class with common functionality.
"""
from abc import ABC, abstractmethod
from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from typing import Any, Final, Generic, Protocol, TypeVar, runtime_checkable
from pydantic import BaseModel
@ -24,6 +24,22 @@ class SupportsDict(Protocol):
DbRecord = Mapping[str, object] | SupportsModelDump | SupportsDict | Sequence[tuple[str, object]]
class PrismaCrudActions(Protocol):
"""The Prisma table actions reached by the generic repository CRUD helpers."""
async def find_unique(self, *, where: Mapping[str, object]) -> DbRecord | None: ...
find_many: Callable[..., Awaitable[Sequence[DbRecord]]]
async def create(self, *, data: Mapping[str, object]) -> DbRecord: ...
async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> DbRecord | None: ...
async def delete(self, *, where: Mapping[str, object]) -> DbRecord | None: ...
async def count(self, *, where: Mapping[str, object] | None = None) -> int: ...
def record_to_dict(record: DbRecord) -> Mapping[str, object]:
"""Project a database record into a mapping of column name to value."""
if isinstance(record, SupportsModelDump):
@ -38,7 +54,7 @@ def record_to_dict(record: DbRecord) -> Mapping[str, object]:
class BaseRepository(ABC, Generic[T]):
"""Abstract base class for all repositories."""
def __init__(self, prisma_client: Any): # any-ok: PrismaClient is an untyped runtime wrapper
def __init__(self, prisma_client: object):
self._prisma_client = prisma_client
@property
@ -53,6 +69,11 @@ class BaseRepository(ABC, Generic[T]):
"""Return the Prisma table for this repository."""
...
@property
def _crud_actions(self) -> PrismaCrudActions:
"""View ``table`` through the action surface the CRUD helpers below use."""
return self.table
@property
@abstractmethod
def model_class(self) -> type[T]:
@ -71,18 +92,18 @@ class BaseRepository(ABC, Generic[T]):
async def find_by_id(self, id_value: str, id_field: str = "id") -> T | None:
"""Find a record by its primary key."""
record: Final = await self.table.find_unique(where={id_field: id_value})
record: Final = await self._crud_actions.find_unique(where={id_field: id_value})
return self._to_model(record)
async def find_many(
self,
where: dict[str, Any] | None = None,
where: Mapping[str, object] | None = None,
skip: int | None = None,
take: int | None = None,
order: dict[str, str] | None = None,
order: Mapping[str, str] | None = None,
) -> list[T]:
"""Find multiple records matching the criteria."""
kwargs: Final[dict[str, Any]] = {}
kwargs: Final[dict[str, object]] = {}
if where:
kwargs["where"] = where
if skip is not None:
@ -92,31 +113,31 @@ class BaseRepository(ABC, Generic[T]):
if order:
kwargs["order"] = order
records: Final = await self.table.find_many(**kwargs)
records: Final = await self._crud_actions.find_many(**kwargs)
return self._to_model_list(records)
async def create(self, data: dict[str, Any]) -> T:
async def create(self, data: Mapping[str, object]) -> T:
"""Create a new record."""
record: Final = await self.table.create(data=data)
record: Final = await self._crud_actions.create(data=data)
model: Final = self._to_model(record)
assert model is not None
return model
async def update(self, id_value: str, data: dict[str, Any], id_field: str = "id") -> T | None:
async def update(self, id_value: str, data: Mapping[str, object], id_field: str = "id") -> T | None:
"""Update an existing record."""
record: Final = await self.table.update(where={id_field: id_value}, data=data)
record: Final = await self._crud_actions.update(where={id_field: id_value}, data=data)
return self._to_model(record)
async def delete(self, id_value: str, id_field: str = "id") -> T | None:
"""Delete a record by its primary key."""
record: Final = await self.table.delete(where={id_field: id_value})
record: Final = await self._crud_actions.delete(where={id_field: id_value})
return self._to_model(record)
async def count(self, where: dict[str, Any] | None = None) -> int:
async def count(self, where: Mapping[str, object] | None = None) -> int:
"""Count records matching the criteria."""
return await self.table.count(where=where)
return await self._crud_actions.count(where=where)
async def exists(self, id_value: str, id_field: str = "id") -> bool:
"""Check if a record exists."""
record: Final = await self.table.find_unique(where={id_field: id_value})
record: Final = await self._crud_actions.find_unique(where={id_field: id_value})
return record is not None

58
litellm/repositories/credentials_repository.py
View file

@ -6,11 +6,41 @@ credential values is the caller's responsibility (see ``CredentialHelperUtils``)
so reads return the stored values verbatim.
"""
from typing import Any, Final
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Protocol
from litellm.models.credentials import CredentialItem
from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync
if TYPE_CHECKING:
from prisma.models import LiteLLM_CredentialsTable
class _CredentialsDb(Protocol):
@property
def litellm_credentialstable(self) -> object: ...
class _PrismaClientView(Protocol):
@property
def db(self) -> _CredentialsDb: ...
class _CredentialsActions(Protocol):
"""Prisma table actions used by :class:`CredentialsRepository`."""
async def find_many(self) -> "Sequence[LiteLLM_CredentialsTable]": ...
async def create(self, *, data: Mapping[str, object]) -> "LiteLLM_CredentialsTable": ...
async def find_unique(self, *, where: Mapping[str, object]) -> "LiteLLM_CredentialsTable | None": ...
async def update(
self, *, where: Mapping[str, object], data: Mapping[str, object]
) -> "LiteLLM_CredentialsTable | None": ...
async def delete(self, *, where: Mapping[str, object]) -> "LiteLLM_CredentialsTable | None": ...
class CredentialsRepository:
"""Repository for credentials database operations, keyed by credential name."""
@ -19,7 +49,7 @@ class CredentialsRepository:
self._prisma_client = prisma_client
@property
def prisma_client(self) -> Any:
def prisma_client(self) -> _PrismaClientView:
if self._prisma_client is None:
raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
return self._prisma_client
@ -31,6 +61,10 @@ class CredentialsRepository:
table_name="litellm_credentialstable",
)
@property
def _credentials_table(self) -> _CredentialsActions:
return self.table
@staticmethod
def _to_model(record: Any) -> CredentialItem | None:
if record is None:
@ -42,18 +76,20 @@ class CredentialsRepository:
credential_info=data.get("credential_info") or {},
)
async def find_all(self) -> Any:
return await self.table.find_many()
async def find_all(self) -> "Sequence[LiteLLM_CredentialsTable]":
return await self._credentials_table.find_many()
async def create(self, data: dict[str, Any]) -> Any:
return await self.table.create(data=data)
async def create(self, data: Mapping[str, object]) -> "LiteLLM_CredentialsTable":
return await self._credentials_table.create(data=data)
async def find_by_name(self, credential_name: str) -> CredentialItem | None:
record: Final = await self.table.find_unique(where={"credential_name": credential_name})
record: Final = await self._credentials_table.find_unique(where={"credential_name": credential_name})
return self._to_model(record)
async def update_by_name(self, credential_name: str, data: dict[str, Any]) -> Any:
return await self.table.update(where={"credential_name": credential_name}, data=data)
async def update_by_name(
self, credential_name: str, data: Mapping[str, object]
) -> "LiteLLM_CredentialsTable | None":
return await self._credentials_table.update(where={"credential_name": credential_name}, data=data)
async def delete_by_name(self, credential_name: str) -> Any:
return await self.table.delete(where={"credential_name": credential_name})
async def delete_by_name(self, credential_name: str) -> "LiteLLM_CredentialsTable | None":
return await self._credentials_table.delete(where={"credential_name": credential_name})

76
litellm/repositories/team_repository.py
View file

@ -3,9 +3,10 @@ Team repository for database operations on LiteLLM_TeamTable.
"""
import json
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from contextlib import AbstractAsyncContextManager
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Any, Final, Protocol
from pydantic import TypeAdapter
@ -13,12 +14,45 @@ from litellm.models.team import LiteLLM_TeamTable, Member
from litellm.repositories.base_repository import (
BaseRepository,
DbRecord,
PrismaCrudActions,
record_to_dict,
)
if TYPE_CHECKING:
from prisma import Prisma
class _TeamTables(Protocol):
"""The Prisma tables this repository reaches, on the client or inside a transaction."""
litellm_teamtable: PrismaCrudActions
litellm_deletedteamtable: PrismaCrudActions
class _TeamArrays(Protocol):
"""The string array columns of a team row, which the domain model leaves untyped."""
@property
def members(self) -> Sequence[str]: ...
@property
def admins(self) -> Sequence[str]: ...
@property
def models(self) -> Sequence[str]: ...
def _team_arrays(team: LiteLLM_TeamTable) -> _TeamArrays:
"""View a team's untyped list columns as sequences of ids."""
return team
class _TeamDatabase(_TeamTables, Protocol):
"""The Prisma client surface used for team reads, writes, and archival transactions."""
def tx(self) -> AbstractAsyncContextManager[_TeamTables]: ...
_MEMBERS_WITH_ROLES_ADAPTER: Final = TypeAdapter(list[Member])
_JSON_ENCODED_TEAM_FIELDS: Final = (
"metadata",
@ -34,12 +68,16 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
"""Repository for team database operations."""
@property
def table(self) -> Any: # any-ok: PrismaClient.db is an untyped runtime wrapper
return self.prisma_client.db.litellm_teamtable
def _db(self) -> _TeamDatabase:
return self.prisma_client.db
@property
def deleted_table(self) -> Any: # any-ok: PrismaClient.db is an untyped runtime wrapper
return self.prisma_client.db.litellm_deletedteamtable
def table(self) -> Any: # any-ok: callers reach model-specific actions this repository does not use
return self._db.litellm_teamtable
@property
def deleted_table(self) -> PrismaCrudActions:
return self._db.litellm_deletedteamtable
@property
def model_class(self) -> type[LiteLLM_TeamTable]:
@ -75,8 +113,8 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
)
if not rows:
return None
raw_value: Final = rows[0]["members_with_roles"]
parsed: Final = json.loads(raw_value) if isinstance(raw_value, str) else raw_value
raw_value: Final[object] = rows[0]["members_with_roles"]
parsed: Final[object] = json.loads(raw_value) if isinstance(raw_value, str) else raw_value
if not parsed:
return []
return _MEMBERS_WITH_ROLES_ADAPTER.validate_python(parsed)
@ -86,24 +124,24 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
async def find_by_alias(self, team_alias: str) -> LiteLLM_TeamTable | None:
"""Find a team by alias."""
records: Final = await self.table.find_many(where={"team_alias": team_alias})
records: Final = await self._crud_actions.find_many(where={"team_alias": team_alias})
if records:
return self._to_model(records[0])
return None
async def find_by_organization_id(self, organization_id: str) -> list[LiteLLM_TeamTable]:
"""Find all teams belonging to an organization."""
records: Final = await self.table.find_many(where={"organization_id": organization_id})
records: Final = await self._crud_actions.find_many(where={"organization_id": organization_id})
return self._to_model_list(records)
async def find_by_member(self, user_id: str) -> list[LiteLLM_TeamTable]:
"""Find all teams where user is a member."""
records: Final = await self.table.find_many(where={"members": {"has": user_id}})
records: Final = await self._crud_actions.find_many(where={"members": {"has": user_id}})
return self._to_model_list(records)
async def find_by_admin(self, user_id: str) -> list[LiteLLM_TeamTable]:
"""Find all teams where user is an admin."""
records: Final = await self.table.find_many(where={"admins": {"has": user_id}})
records: Final = await self._crud_actions.find_many(where={"admins": {"has": user_id}})
return self._to_model_list(records)
async def create_team(
@ -232,7 +270,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
archive_data["litellm_changed_by"] = litellm_changed_by
archive_data["deleted_at"] = datetime.utcnow()
async with self.prisma_client.db.tx() as tx:
async with self._db.tx() as tx:
await tx.litellm_deletedteamtable.create(data=archive_data)
await tx.litellm_teamtable.delete(where={"team_id": team_id})
@ -293,7 +331,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
if not await self.exists(team_id, id_field="team_id"):
return None
record: Final = await self.table.update(
record: Final = await self._crud_actions.update(
where={"team_id": team_id},
data={"members": {"push": user_id}},
)
@ -310,7 +348,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
if team is None:
return None
members: Final = [m for m in team.members if m != user_id]
members: Final = [m for m in _team_arrays(team).members if m != user_id]
return await self.update(team_id, {"members": members}, id_field="team_id")
async def add_admin(self, team_id: str, user_id: str) -> LiteLLM_TeamTable | None:
@ -318,7 +356,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
if not await self.exists(team_id, id_field="team_id"):
return None
record: Final = await self.table.update(
record: Final = await self._crud_actions.update(
where={"team_id": team_id},
data={"admins": {"push": user_id}},
)
@ -335,7 +373,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
if team is None:
return None
admins: Final = [a for a in team.admins if a != user_id]
admins: Final = [a for a in _team_arrays(team).admins if a != user_id]
return await self.update(team_id, {"admins": admins}, id_field="team_id")
async def add_models(self, team_id: str, models: list[str]) -> LiteLLM_TeamTable | None:
@ -343,7 +381,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
if not await self.exists(team_id, id_field="team_id"):
return None
record: Final = await self.table.update(
record: Final = await self._crud_actions.update(
where={"team_id": team_id},
data={"models": {"push": models}},
)
@ -360,5 +398,5 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
if team is None:
return None
current_models: Final = [m for m in team.models if m not in models]
current_models: Final = [m for m in _team_arrays(team).models if m not in models]
return await self.update(team_id, {"models": current_models}, id_field="team_id")

23
litellm/rust_bridge/responses_websocket.py
View file

@ -2,8 +2,9 @@
from __future__ import annotations
from collections.abc import Awaitable
from dataclasses import dataclass
from typing import Any, Final, Protocol
from typing import Final, Protocol
import httpx
from websockets.exceptions import ConnectionClosedOK
@ -12,6 +13,16 @@ from litellm.rust_bridge.loader import get_native_bridge
from litellm.rust_bridge.timeouts import timeout_to_seconds
class RustResponsesWebSocketSocket(Protocol):
"""Open socket handle handed back by the native bridge."""
def send_text(self, text: str) -> Awaitable[None]: ...
def recv_text(self) -> Awaitable[str | None]: ...
def close(self) -> Awaitable[None]: ...
class RustResponsesWebSocketConnection(Protocol):
@classmethod
def connect(
@ -19,7 +30,7 @@ class RustResponsesWebSocketConnection(Protocol):
url: str,
headers: dict[str, str],
timeout_seconds: float | None,
) -> Any:
) -> Awaitable[RustResponsesWebSocketSocket]:
raise NotImplementedError
@ -32,7 +43,7 @@ _UNSET: Final[_Unset] = _Unset()
@dataclass(slots=True)
class _RustResponsesWebSocketState:
connection: Any = None
connection: type[RustResponsesWebSocketConnection] | None = None
_STATE: Final[_RustResponsesWebSocketState] = _RustResponsesWebSocketState()
@ -40,13 +51,13 @@ _STATE: Final[_RustResponsesWebSocketState] = _RustResponsesWebSocketState()
def set_rust_responses_websocket(
*,
connection: Any = _UNSET,
connection: type[RustResponsesWebSocketConnection] | None | _Unset = _UNSET,
) -> None:
if not isinstance(connection, _Unset):
_STATE.connection = connection
def load_rust_responses_websocket() -> Any:
def load_rust_responses_websocket() -> type[RustResponsesWebSocketConnection] | None:
if _STATE.connection is not None:
return _STATE.connection
native_bridge: Final = get_native_bridge()
@ -59,7 +70,7 @@ def load_rust_responses_websocket() -> Any:
class _ConnectionAdapter:
def __init__(self, connection: Any):
def __init__(self, connection: RustResponsesWebSocketSocket):
self._connection = connection
async def send(self, text: str) -> None:

60
litellm/secret_managers/secret_manager_handler.py
View file

@ -6,14 +6,62 @@ Handles retrieving secrets from different secret management systems.
import base64
import os
from typing import Any, Final
from collections.abc import Mapping
from typing import Final, Protocol, overload
import litellm
from litellm._logging import print_verbose
from litellm.types.secret_managers.main import KeyManagementSystem
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
def _is_base64(s):
class _VaultSecret(Protocol):
"""The secret object returned by the Azure Key Vault and Infisical clients."""
@property
def value(self) -> str | None: ...
@property
def secret_value(self) -> str | None: ...
class _KmsPlaintext(Protocol):
"""The decrypted payload the AWS KMS client exposes under the ``Plaintext`` response key."""
def decode(self, encoding: str) -> str | None: ...
class _KmsDecryptResponse(Protocol):
"""The decrypt response returned by the Google KMS and AWS KMS clients."""
@property
def plaintext(self) -> bytes: ...
def __getitem__(self, key: str) -> _KmsPlaintext: ...
class _SecretManagerClient(Protocol):
"""The untyped secret manager client surface reached by ``get_secret_from_manager``."""
def get_secret(self, secret_name: str, /) -> _VaultSecret: ...
@overload
def decrypt(self, *, request: Mapping[str, object]) -> _KmsDecryptResponse: ...
@overload
def decrypt(self, *, CiphertextBlob: bytes) -> _KmsDecryptResponse: ...
def sync_read_secret(
self,
*,
secret_name: str,
primary_secret_name: str | None = None,
optional_params: Mapping[str, object] | None = None,
) -> str | None: ...
def get_secret_from_google_secret_manager(self, secret_name: str, /) -> str | None: ...
def _is_base64(s: str | bytes) -> bool:
"""Check if a string is valid base64."""
import binascii
@ -24,10 +72,10 @@ def _is_base64(s):
def get_secret_from_manager(
client: Any,
client: _SecretManagerClient,
key_manager: str,
secret_name: str,
key_management_settings: Any | None = None,
key_management_settings: KeyManagementSettings | None = None,
) -> str | None:
"""
Get a secret from the configured secret manager.
@ -56,7 +104,7 @@ def get_secret_from_manager(
elif (
key_manager == KeyManagementSystem.GOOGLE_KMS.value or client.__class__.__name__ == "KeyManagementServiceClient"
):
encrypted_secret: Any = os.getenv(secret_name)
encrypted_secret: str | bytes | None = os.getenv(secret_name)
if encrypted_secret is None:
raise ValueError("Google KMS requires the encrypted secret to be in the environment!")
b64_flag: Final = _is_base64(encrypted_secret)

18
ruff-strict-budget.json
View file

@ -1,21 +1,21 @@
{
"ANN001": {
"limit": 3020
"limit": 3012
},
"ANN002": {
"limit": 71
},
"ANN003": {
"limit": 827
"limit": 818
},
"ANN201": {
"limit": 2017
"limit": 2014
},
"ANN202": {
"limit": 852
"limit": 851
},
"ANN204": {
"limit": 711
"limit": 707
},
"ANN205": {
"limit": 112
@ -24,7 +24,7 @@
"limit": 133
},
"ANN401": {
"limit": 1188
"limit": 1061
},
"ASYNC230": {
"limit": 11
@ -198,7 +198,7 @@
"limit": 22
},
"SIM101": {
"limit": 58
"limit": 57
},
"SIM102": {
"limit": 317
@ -234,7 +234,7 @@
"limit": 5
},
"TID251": {
"limit": 1212
"limit": 1202
},
"TRY002": {
"limit": 524
@ -249,7 +249,7 @@
"limit": 113
},
"TRY300": {
"limit": 859
"limit": 858
},
"UP028": {
"limit": 2

19
tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py
View file

@ -313,6 +313,25 @@ class TestGDCGeminiConfig:
api_base=TEST_API_BASE,
)
def test_validate_environment_credentials_missing_audience_binding_are_named(self):
config = GDCGeminiConfig()
creds_without_audience_binding = MagicMock(spec=[])
with patch(
"google.auth.load_credentials_from_dict",
return_value=(creds_without_audience_binding, None),
):
with pytest.raises(AttributeError, match="must expose with_gdch_audience"):
config.validate_environment(
headers={},
model=TEST_MODEL,
messages=[],
optional_params={},
litellm_params={"vertex_project": TEST_PROJECT},
api_key=TEST_API_KEY,
api_base=TEST_API_BASE,
)
def test_validate_environment_string_false_disables_token_caching(self):
config = GDCGeminiConfig()
mock_creds = MagicMock()

12
type-discipline-budget.json
View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 22805
"limit": 22728
},
"LIT002": {
"limit": 26878
"limit": 26860
},
"LIT003": {
"limit": 269
@ -15,7 +15,7 @@
"limit": 0
},
"LIT006": {
"limit": 1069
"limit": 1066
},
"LIT007": {
"limit": 0
@ -27,12 +27,12 @@
"limit": 0
},
"LIT010": {
"limit": 16695
"limit": 16673
},
"LIT011": {
"limit": 5588
"limit": 5586
},
"LIT012": {
"limit": 4519
"limit": 4512
}
}