chore(typing): clear basedpyright Any errors across proxy, rag and mcp modules

Replace Any seams in eight production modules with real types: Protocols for
the untyped Prisma client actions, TypedDicts for JSON and SSE payloads, and
precise parameter and return annotations on the db, endpoint and CLI helpers.

Repo-wide reportAny drops 17311 -> 16858 and reportExplicitAny 5892 -> 5865,
with every basedpyright rule at or below its baseline and total errors down
577. Budgets ratcheted to match.
This commit is contained in:
mateo-berri 2026-08-09 20:54:27 +00:00
parent ecba48dd7c
commit 1ac3700098
No known key found for this signature in database
11 changed files with 686 additions and 297 deletions

View file

@ -1,6 +1,6 @@
{
"reportAny": {
"limit": 27731
"limit": 27278
},
"reportArgumentType": {
"limit": 2626
@ -24,7 +24,7 @@
"limit": 19
},
"reportExplicitAny": {
"limit": 8807
"limit": 8780
},
"reportFunctionMemberAccess": {
"limit": 7
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5835
"limit": 5822
},
"reportMissingTypeArgument": {
"limit": 15790
"limit": 15778
},
"reportMissingTypeStubs": {
"limit": 40
@ -99,19 +99,19 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 45063
"limit": 45055
},
"reportUnknownLambdaType": {
"limit": 113
},
"reportUnknownMemberType": {
"limit": 39773
"limit": 39756
},
"reportUnknownParameterType": {
"limit": 20207
"limit": 20189
},
"reportUnknownVariableType": {
"limit": 31281
"limit": 31252
},
"reportUnnecessaryCast": {
"limit": 122

View file

@ -248,7 +248,38 @@ class WebSearchInterceptionLogger(CustomLogger):
)
return response
async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None:
@staticmethod
def _str_field(source: Mapping[str, object], key: str) -> str:
value: Final = source.get(key)
return value if isinstance(value, str) else ""
@staticmethod
def _tools_field(kwargs: Mapping[str, object]) -> Sequence[dict[str, object]] | None:
tools: Final = kwargs.get("tools")
return tools if isinstance(tools, list) else None
@classmethod
def _resolve_deployment_provider(cls, kwargs: Mapping[str, object]) -> str:
"""Provider from top-level kwargs, then nested litellm_params, then the model name."""
top_level: Final = cls._str_field(kwargs, "custom_llm_provider")
if top_level:
return top_level
litellm_params: Final = kwargs.get("litellm_params")
if isinstance(litellm_params, dict):
nested: Final = cls._str_field(litellm_params, "custom_llm_provider")
if nested:
return nested
try:
_, derived_provider, _, _ = litellm.get_llm_provider(model=cls._str_field(kwargs, "model"))
except Exception:
return ""
return derived_provider
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: CallTypes | None
) -> dict | None:
"""
Pre-call hook to convert native Anthropic web_search tools to regular tools.
@ -256,21 +287,12 @@ class WebSearchInterceptionLogger(CustomLogger):
Instead, we convert it to a regular tool so the model returns tool_use blocks
that we can intercept and execute ourselves.
"""
# Check if this is for an enabled provider
# Try top-level kwargs first, then nested litellm_params, then derive from model name
custom_llm_provider = kwargs.get("custom_llm_provider", "") or kwargs.get("litellm_params", {}).get(
"custom_llm_provider", ""
)
if not custom_llm_provider:
try:
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs.get("model", ""))
except Exception:
custom_llm_provider = ""
custom_llm_provider: Final = self._resolve_deployment_provider(kwargs)
if custom_llm_provider not in self.enabled_providers:
return None
# Check if request has tools with native web_search
tools: Final[Sequence[dict[str, object]] | None] = kwargs.get("tools")
tools: Final = self._tools_field(kwargs)
if not tools:
return None
@ -394,7 +416,7 @@ class WebSearchInterceptionLogger(CustomLogger):
return tool.get("name")
@classmethod
def _sync_forced_tool_choice(cls, tool_choice: Any, converted_tools: list[dict[str, object]]) -> object:
def _sync_forced_tool_choice(cls, tool_choice: object, converted_tools: list[dict[str, object]]) -> object:
"""Repoint a forced ``tool_choice`` at ``litellm_web_search`` when it
names a web-search tool that was just converted away.
@ -462,7 +484,7 @@ class WebSearchInterceptionLogger(CustomLogger):
kwargs[WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY] = True
# Convert native web search tools to LiteLLM standard
converted_tools: Final = []
converted_tools: Final[list[dict[str, object]]] = []
for tool in tools:
if is_web_search_tool(tool):
standard_tool = get_litellm_web_search_tool()
@ -482,7 +504,8 @@ class WebSearchInterceptionLogger(CustomLogger):
)
if "tool_choice" in kwargs:
kwargs["tool_choice"] = self._sync_forced_tool_choice(kwargs.get("tool_choice"), converted_tools)
forced_tool_choice: Final[object] = kwargs.get("tool_choice")
kwargs["tool_choice"] = self._sync_forced_tool_choice(forced_tool_choice, converted_tools)
# Also convert here for direct callers that bypass the deployment hook.
if kwargs.get("stream"):
@ -836,7 +859,8 @@ class WebSearchInterceptionLogger(CustomLogger):
native_blocks: Final = plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY)
if not native_blocks:
return response
return self._inject_native_blocks(response, native_blocks)
self._inject_native_blocks(response, native_blocks)
return response
@staticmethod
def _build_native_result_blocks(
@ -883,17 +907,17 @@ class WebSearchInterceptionLogger(CustomLogger):
)
@staticmethod
def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any:
"""Prepend native blocks to response content, dict or object form."""
def _inject_native_blocks(response: object, native_blocks: Sequence[Mapping[str, object]]) -> None:
"""Prepend native blocks to response content in place, dict or object form."""
if not native_blocks:
return response
return
if isinstance(response, dict):
existing = response.get("content") or []
response["content"] = list(native_blocks) + list(existing)
return response
existing = getattr(response, "content", None) or []
existing_items: Final = response.get("content") or []
response["content"] = list(native_blocks) + list(existing_items)
return
existing_attr: Final = getattr(response, "content", None) or []
try:
response.content = list(native_blocks) + list(existing)
response.__setattr__("content", list(native_blocks) + list(existing_attr))
except (AttributeError, TypeError):
# Object refused write — fall through and leave the response
# untouched rather than crash the request.
@ -901,7 +925,6 @@ class WebSearchInterceptionLogger(CustomLogger):
"WebSearchInterception: could not inject native blocks into response of type %s",
type(response).__name__,
)
return response
async def async_run_chat_completion_agentic_loop(
self,
@ -1177,7 +1200,7 @@ class WebSearchInterceptionLogger(CustomLogger):
if max_tokens is None:
max_tokens = cast(int, kwargs.get("max_tokens", 1024))
response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate(
response: Final[AnthropicMessagesResponse | AsyncIterator[object]] = await anthropic_messages.acreate(
max_tokens=max_tokens,
messages=request_patch.messages,
model=request_patch.model or model,
@ -1193,7 +1216,7 @@ class WebSearchInterceptionLogger(CustomLogger):
tool_calls=tool_calls,
structured_results=structured_results,
)
response = self._inject_native_blocks(response, native_blocks)
self._inject_native_blocks(response, native_blocks)
return response
@ -1400,7 +1423,7 @@ class WebSearchInterceptionLogger(CustomLogger):
valid_token=user_api_key_auth,
)
team_id: Final = getattr(user_api_key_auth, "team_id", None)
team_id: Final = user_api_key_auth.team_id
if team_id:
from litellm.proxy.proxy_server import (
prisma_client,

View file

@ -3,8 +3,9 @@ import binascii
import hashlib
import json
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from contextlib import AbstractAsyncContextManager
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -45,7 +46,6 @@ from litellm.types.mcp import MCPCredentials
if TYPE_CHECKING:
from prisma import models as prisma_db_models
from prisma import types as prisma_db_types
from prisma.actions import LiteLLM_MCPUserCredentialsActions, LiteLLM_MCPUserEnvVarsActions
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@ -434,23 +434,150 @@ def _credentials_blob_to_mutable_dict(blob: str | Mapping[str, object]) -> dict[
return parsed_blob
def _json_loads_object(raw: str) -> dict[str, object] | None:
"""Parse ``raw`` as JSON and return it only when it decodes to an object, else ``None``.
Decode errors propagate to the caller, which decides how a malformed blob is handled.
"""
parsed: Final[dict[str, object] | None] = json.loads(raw)
return parsed if isinstance(parsed, dict) else None
class _MCPServerTableActions(Protocol):
async def find_many(
self,
*,
where: Mapping[str, object] | None = None,
order: Mapping[str, str] | None = None,
take: int | None = None,
) -> "Sequence[prisma_db_models.LiteLLM_MCPServerTable]": ...
async def find_unique(self, *, where: Mapping[str, object]) -> "prisma_db_models.LiteLLM_MCPServerTable | None": ...
async def update(
self, *, where: Mapping[str, object], data: Mapping[str, object]
) -> "prisma_db_models.LiteLLM_MCPServerTable": ...
class _MCPServerOAuthClientActions(Protocol):
async def find_many(
self, *, where: Mapping[str, object] | None = None
) -> "Sequence[prisma_db_models.LiteLLM_MCPServerOAuthClient]": ...
async def find_unique(
self, *, where: Mapping[str, object]
) -> "prisma_db_models.LiteLLM_MCPServerOAuthClient | None": ...
async def upsert(
self, *, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]]
) -> "prisma_db_models.LiteLLM_MCPServerOAuthClient": ...
async def update(
self, *, where: Mapping[str, object], data: Mapping[str, object]
) -> "prisma_db_models.LiteLLM_MCPServerOAuthClient | None": ...
async def delete_many(self, *, where: Mapping[str, object]) -> int: ...
class _VerificationTokenActions(Protocol):
async def find_many(
self, *, where: Mapping[str, object] | None = None
) -> "list[prisma_db_models.LiteLLM_VerificationToken] | None": ...
async def find_unique(
self, *, where: Mapping[str, object], include: Mapping[str, bool] | None = None
) -> "prisma_db_models.LiteLLM_VerificationToken | None": ...
class _TeamTableActions(Protocol):
async def find_unique(
self, *, where: Mapping[str, object], include: Mapping[str, bool] | None = None
) -> "prisma_db_models.LiteLLM_TeamTable | None": ...
class _MCPUserCredentialsActions(Protocol):
async def find_many(
self, *, where: Mapping[str, object] | None = None
) -> "Sequence[prisma_db_models.LiteLLM_MCPUserCredentials]": ...
async def find_unique(
self, *, where: Mapping[str, object]
) -> "prisma_db_models.LiteLLM_MCPUserCredentials | None": ...
async def upsert(
self, *, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]]
) -> "prisma_db_models.LiteLLM_MCPUserCredentials": ...
async def update(
self, *, where: Mapping[str, object], data: Mapping[str, object]
) -> "prisma_db_models.LiteLLM_MCPUserCredentials | None": ...
async def delete(self, *, where: Mapping[str, object]) -> "prisma_db_models.LiteLLM_MCPUserCredentials | None": ...
async def delete_many(self, *, where: Mapping[str, object]) -> int: ...
class _MCPUserEnvVarsActions(Protocol):
async def find_many(
self, *, where: Mapping[str, object] | None = None
) -> "Sequence[prisma_db_models.LiteLLM_MCPUserEnvVars]": ...
async def find_unique(self, *, where: Mapping[str, object]) -> "prisma_db_models.LiteLLM_MCPUserEnvVars | None": ...
async def upsert(
self, *, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]]
) -> "prisma_db_models.LiteLLM_MCPUserEnvVars": ...
async def update(
self, *, where: Mapping[str, object], data: Mapping[str, object]
) -> "prisma_db_models.LiteLLM_MCPUserEnvVars | None": ...
async def delete_many(self, *, where: Mapping[str, object]) -> int: ...
class _MCPUserEnvVarsTransaction(Protocol):
async def execute_raw(self, query: str, *args: object) -> int: ...
@property
def litellm_mcpuserenvvars(self) -> _MCPUserEnvVarsActions: ...
def _mcp_server_actions(prisma_client: PrismaClient) -> _MCPServerTableActions:
table: Final[_MCPServerTableActions] = MCPServerRepository(prisma_client).table
return table
def _oauth_client_actions(prisma_client: PrismaClient) -> _MCPServerOAuthClientActions:
table: Final[_MCPServerOAuthClientActions] = MCPServerOAuthClientRepository(prisma_client).table
return table
def _verification_token_actions(prisma_client: PrismaClient) -> _VerificationTokenActions:
table: Final[_VerificationTokenActions] = VerificationTokenRepository(prisma_client).table
return table
def _team_actions(prisma_client: PrismaClient) -> _TeamTableActions:
table: Final[_TeamTableActions] = TeamRepository(prisma_client).table
return table
def _user_env_var_tx(prisma_client: PrismaClient) -> AbstractAsyncContextManager[_MCPUserEnvVarsTransaction]:
tx: Final[AbstractAsyncContextManager[_MCPUserEnvVarsTransaction]] = prisma_client.db.tx()
return tx
async def _db_find_mcp_server_rows(
prisma_client: PrismaClient,
where: "prisma_db_types.LiteLLM_MCPServerTableWhereInput | None" = None,
) -> "list[prisma_db_models.LiteLLM_MCPServerTable]":
rows: list[prisma_db_models.LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.find_many(
where=where
)
return rows
) -> "Sequence[prisma_db_models.LiteLLM_MCPServerTable]":
return await _mcp_server_actions(prisma_client).find_many(where=where)
async def _db_find_mcp_server_row(
prisma_client: PrismaClient, server_id: str
) -> "prisma_db_models.LiteLLM_MCPServerTable | None":
row: prisma_db_models.LiteLLM_MCPServerTable | None = await MCPServerRepository(prisma_client).table.find_unique(
where={"server_id": server_id}
)
return row
return await _mcp_server_actions(prisma_client).find_unique(where={"server_id": server_id})
async def _db_update_mcp_server_row(
@ -458,28 +585,16 @@ async def _db_update_mcp_server_row(
server_id: str,
data: "prisma_db_types.LiteLLM_MCPServerTableUpdateInput",
) -> "prisma_db_models.LiteLLM_MCPServerTable":
row: Final[prisma_db_models.LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.update(
where={"server_id": server_id},
data=data,
)
return row
return await _mcp_server_actions(prisma_client).update(where={"server_id": server_id}, data=data)
def _user_credential_actions(
prisma_client: PrismaClient,
) -> "LiteLLM_MCPUserCredentialsActions[prisma_db_models.LiteLLM_MCPUserCredentials]":
table: Final[LiteLLM_MCPUserCredentialsActions[prisma_db_models.LiteLLM_MCPUserCredentials]] = (
MCPUserCredentialsRepository(prisma_client).table
)
def _user_credential_actions(prisma_client: PrismaClient) -> _MCPUserCredentialsActions:
table: Final[_MCPUserCredentialsActions] = MCPUserCredentialsRepository(prisma_client).table
return table
def _user_env_var_actions(
prisma_client: PrismaClient,
) -> "LiteLLM_MCPUserEnvVarsActions[prisma_db_models.LiteLLM_MCPUserEnvVars]":
table: Final[LiteLLM_MCPUserEnvVarsActions[prisma_db_models.LiteLLM_MCPUserEnvVars]] = (
prisma_client.db.litellm_mcpuserenvvars
)
def _user_env_var_actions(prisma_client: PrismaClient) -> _MCPUserEnvVarsActions:
table: Final[_MCPUserEnvVarsActions] = prisma_client.db.litellm_mcpuserenvvars
return table
@ -494,14 +609,14 @@ async def _db_find_user_credential_row(
async def _db_find_user_credential_rows(
prisma_client: PrismaClient,
where: "prisma_db_types.LiteLLM_MCPUserCredentialsWhereInput | None" = None,
) -> "list[prisma_db_models.LiteLLM_MCPUserCredentials]":
) -> "Sequence[prisma_db_models.LiteLLM_MCPUserCredentials]":
return await _user_credential_actions(prisma_client).find_many(where=where)
async def _db_upsert_user_credential_row(
prisma_client: PrismaClient, user_id: str, server_id: str, credential_b64: str
) -> None:
await MCPUserCredentialsRepository(prisma_client).table.upsert(
await _user_credential_actions(prisma_client).upsert(
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}},
data={
"create": {
@ -517,7 +632,7 @@ async def _db_upsert_user_credential_row(
async def _db_find_user_env_var_rows(
prisma_client: PrismaClient,
where: "prisma_db_types.LiteLLM_MCPUserEnvVarsWhereInput | None" = None,
) -> "list[prisma_db_models.LiteLLM_MCPUserEnvVars]":
) -> "Sequence[prisma_db_models.LiteLLM_MCPUserEnvVars]":
return await _user_env_var_actions(prisma_client).find_many(where=where)
@ -585,9 +700,7 @@ async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str]
"""
Returns the matching mcp servers from the db with the server_ids
"""
_mcp_servers: Final[list[prisma_db_models.LiteLLM_MCPServerTable]] = await MCPServerRepository(
prisma_client
).table.find_many(
_mcp_servers: Final = await _mcp_server_actions(prisma_client).find_many(
where={
"server_id": {"in": server_ids},
}
@ -605,9 +718,7 @@ async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, toke
"""
Returns the mcp servers from the db for the verification token
"""
verification_token_record: prisma_db_models.LiteLLM_VerificationToken | None = await VerificationTokenRepository(
prisma_client
).table.find_unique(
verification_token_record: Final = await _verification_token_actions(prisma_client).find_unique(
where={
"token": token,
},
@ -626,7 +737,7 @@ async def get_mcp_servers_by_team(prisma_client: PrismaClient, team_id: str) ->
"""
Returns the mcp servers from the db for the team id
"""
team_record: prisma_db_models.LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.find_unique(
team_record: Final = await _team_actions(prisma_client).find_unique(
where={
"team_id": team_id,
},
@ -697,9 +808,9 @@ async def get_virtualkeys_for_mcp_server(
"""
Get all the virtual keys that have access to the mcp server
"""
virtual_keys: Final[list[prisma_db_models.LiteLLM_VerificationToken] | None] = await VerificationTokenRepository(
virtual_keys: Final[list[prisma_db_models.LiteLLM_VerificationToken] | None] = await _verification_token_actions(
prisma_client
).table.find_many(
).find_many(
where={
"mcp_servers": {"has": server_id},
},
@ -753,9 +864,9 @@ async def delete_mcp_server(
if deleted_server is not None:
credential_user_ids: list[str] = []
try:
credential_rows: Sequence[
prisma_db_models.LiteLLM_MCPUserCredentials
] = await prisma_client.db.litellm_mcpusercredentials.find_many(where={"server_id": server_id})
credential_rows: Final[
Sequence[prisma_db_models.LiteLLM_MCPUserCredentials]
] = await _user_credential_actions(prisma_client).find_many(where={"server_id": server_id})
credential_user_ids = [row.user_id for row in credential_rows]
except Exception as e: # noqa: BLE001 - enumeration is best-effort; cached tokens expire by TTL
verbose_proxy_logger.warning(
@ -764,9 +875,9 @@ async def delete_mcp_server(
e,
)
for model, label in (
(prisma_client.db.litellm_mcpusercredentials, "credential"),
(prisma_client.db.litellm_mcpuserenvvars, "env var"),
(prisma_client.db.litellm_mcpserveroauthclient, "OAuth client"),
(_user_credential_actions(prisma_client), "credential"),
(_user_env_var_actions(prisma_client), "env var"),
(_oauth_client_actions(prisma_client), "OAuth client"),
):
try:
await model.delete_many(where={"server_id": server_id})
@ -945,9 +1056,7 @@ async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, s
LiteLLM_MCPServerTable row, so their dynamically registered client lives here keyed
by server_id. The returned value is the raw credentials blob for
``_get_persisted_dcr_credentials`` to parse."""
row: Final[prisma_db_models.LiteLLM_MCPServerOAuthClient | None] = await MCPServerOAuthClientRepository(
prisma_client
).table.find_unique(where={"server_id": server_id})
row: Final = await _oauth_client_actions(prisma_client).find_unique(where={"server_id": server_id})
if row is None:
return None
return row.credentials
@ -965,7 +1074,7 @@ async def upsert_mcp_server_oauth_client_credentials(
encrypted: Final = encrypt_credentials(credentials=MCPCredentials(**credentials), encryption_key=_get_salt_key())
blob: Final = safe_dumps(encrypted)
await MCPServerOAuthClientRepository(prisma_client).table.upsert(
await _oauth_client_actions(prisma_client).upsert(
where={"server_id": server_id},
data={
"create": {"server_id": server_id, "credentials": blob},
@ -1012,21 +1121,19 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient,
continue
update_data["updated_by"] = touched_by
await MCPServerRepository(prisma_client).table.update(
await _mcp_server_actions(prisma_client).update(
where={"server_id": mcp_server.server_id},
data=update_data,
)
updated += 1
oauth_clients: Final[list[prisma_db_models.LiteLLM_MCPServerOAuthClient]] = await MCPServerOAuthClientRepository(
prisma_client
).table.find_many()
oauth_clients: Final = await _oauth_client_actions(prisma_client).find_many()
oauth_updated = 0
for oauth_client in oauth_clients:
rotated_credentials = _reencrypt_mcp_credentials_blob(oauth_client.credentials, new_master_key)
if rotated_credentials is None:
continue
await MCPServerOAuthClientRepository(prisma_client).table.update(
await _oauth_client_actions(prisma_client).update(
where={"server_id": oauth_client.server_id},
data={"credentials": rotated_credentials},
)
@ -1347,7 +1454,7 @@ def mcp_oauth_token_identity(server: object) -> tuple[object, ...]:
creds: Final = getattr(server, "credentials", None)
if isinstance(creds, str):
try:
parsed: dict[str, object] | None = json.loads(creds)
parsed: dict[str, object] | None = _json_loads_object(creds)
except ValueError:
parsed = None
else:
@ -1716,7 +1823,7 @@ async def get_mcp_submissions(
along with a summary count breakdown by approval_status.
Mirrors get_guardrail_submissions() from guardrail_endpoints.py.
"""
rows: list[prisma_db_models.LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.find_many(
rows: Final = await _mcp_server_actions(prisma_client).find_many(
where={"submitted_at": {"not": None}},
order={"submitted_at": "desc"},
take=500, # safety cap; paginate if needed in a future iteration
@ -1757,12 +1864,11 @@ def _decode_user_env_vars(stored: str) -> dict[str, str]:
"re-enter them rather than silently forwarding ciphertext"
)
return {}
parsed: dict[str, object] | None
try:
parsed = json.loads(decrypted)
parsed: Final = _json_loads_object(decrypted)
except (ValueError, TypeError):
return {}
if not isinstance(parsed, dict):
if parsed is None:
return {}
return {str(k): str(v) for k, v in parsed.items()}
@ -1818,9 +1924,9 @@ async def merge_user_env_vars(
"big",
signed=True,
)
async with prisma_client.db.tx() as tx:
async with _user_env_var_tx(prisma_client) as tx:
await tx.execute_raw("SELECT pg_advisory_xact_lock($1::bigint)", lock_key)
row: Final[prisma_db_models.LiteLLM_MCPUserEnvVars | None] = await tx.litellm_mcpuserenvvars.find_unique(
row: Final = await tx.litellm_mcpuserenvvars.find_unique(
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}
)
existing: Final = _decode_user_env_vars(row.values_b64) if row is not None else {}

View file

@ -11,10 +11,13 @@ MCP Spec Reference:
"""
import typing
from collections.abc import Mapping, Sequence
from collections.abc import Mapping, Sequence, Sized
from typing import Any, Final, NamedTuple, Optional, Protocol, Union, runtime_checkable
if typing.TYPE_CHECKING:
from collections.abc import Awaitable, Callable
from typing import TypeAlias
from fastapi import Request
from mcp.client.session import ClientSession
from mcp.shared.context import RequestContext
@ -28,8 +31,12 @@ if typing.TYPE_CHECKING:
ToolUseContent,
)
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.types.utils import ModelResponse
_AcompletionCallable: TypeAlias = Callable[..., Awaitable[ModelResponse | CustomStreamWrapper]]
from fastapi import HTTPException
from pydantic import TypeAdapter
@ -92,7 +99,7 @@ def _resolve_model_from_preferences(
if not available_model_names and litellm.model_list:
for entry in litellm.model_list:
if isinstance(entry, dict):
name = entry.get("model_name")
name: str | None = entry.get("model_name")
if name:
available_model_names.append(name)
elif isinstance(entry, str):
@ -439,9 +446,9 @@ def _convert_mcp_messages_to_openai(
)
# Separate marker items from regular content parts
tool_call_markers = []
tool_result_markers = []
regular_parts = []
tool_call_markers: list[Mapping[str, object]] = []
tool_result_markers: list[Mapping[str, object]] = []
regular_parts: list[Mapping[str, object]] = []
for part in converted_parts:
marker = part.get("_marker_type") if isinstance(part, dict) else None
if marker == "tool_use":
@ -520,7 +527,7 @@ def _extract_text_parts(
) -> str | None:
"""Extract text parts from mixed content."""
items: Final = content if isinstance(content, list) else [content]
texts: Final = []
texts: Final[list[str]] = []
for item in items:
if getattr(item, "type", None) == "text":
texts.append(getattr(item, "text", ""))
@ -1141,7 +1148,7 @@ async def _build_completion_kwargs(
user_api_key_auth: "UserAPIKeyAuth",
raw_headers: dict[str, str] | None,
client_ip: str | None,
) -> dict[str, Any]:
) -> dict[str, object]:
openai_messages: Final = _convert_mcp_messages_to_openai(
messages=params.messages,
system_prompt=params.systemPrompt,
@ -1176,8 +1183,21 @@ async def _build_completion_kwargs(
)
async def _call_acompletion(
acompletion: "_AcompletionCallable",
completion_kwargs: dict[str, object],
) -> "ModelResponse | CustomStreamWrapper":
"""Invoke a completion entrypoint with the dynamically assembled proxy payload.
The payload's keys are only known at runtime (MCP params plus whatever
``add_litellm_data_to_request`` and the pre-call hooks inject), so it is
forwarded through a signature-erased callable.
"""
return await acompletion(**completion_kwargs)
async def _run_guardrails_and_call_llm(
completion_kwargs: dict[str, Any],
completion_kwargs: dict[str, object],
user_api_key_auth: "UserAPIKeyAuth",
) -> Any:
try:
@ -1204,10 +1224,10 @@ async def _run_guardrails_and_call_llm(
from litellm.proxy.proxy_server import llm_router
if llm_router is not None:
return await llm_router.acompletion(**completion_kwargs)
return await litellm.acompletion(**completion_kwargs)
return await _call_acompletion(llm_router.acompletion, completion_kwargs)
return await _call_acompletion(litellm.acompletion, completion_kwargs)
except ImportError:
return await litellm.acompletion(**completion_kwargs)
return await _call_acompletion(litellm.acompletion, completion_kwargs)
async def handle_sampling_create_message(
@ -1284,12 +1304,12 @@ async def handle_sampling_create_message(
client_ip=client_ip,
)
openai_messages: Final[Sequence[Mapping[str, object]]] = completion_kwargs["messages"]
openai_messages: Final = completion_kwargs["messages"]
openai_tools: Final = completion_kwargs.get("tools")
verbose_logger.debug(
"MCP sampling: calling litellm.acompletion with model=%s, num_messages=%d, has_tools=%s",
model,
len(openai_messages),
len(openai_messages) if isinstance(openai_messages, Sized) else 0,
bool(openai_tools),
)

View file

@ -3,8 +3,9 @@ import os
import sys
import time
import webbrowser
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Any, Final
from typing import Any, Final, TypedDict
from urllib.parse import urlencode
import click
@ -18,6 +19,42 @@ from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh
from .private_json import write_private_json
class CliContext(TypedDict):
base_url: str
api_key: str | None
base_url_explicit: bool
class TeamRecord(TypedDict, total=False):
team_id: str | None
team_alias: str | None
models: list[str]
max_budget: float | None
class CliSsoStartResponse(TypedDict):
login_id: str
poll_secret: str
user_code: str
class PollResponse(TypedDict, total=False):
status: str
key: str
user_id: str
team_id: str
teams: list[str]
team_details: list[dict[str, str]]
requires_team_selection: bool
class AuthResult(TypedDict):
api_key: str
user_id: str | None
teams: list[str]
team_id: str | None
# Token storage utilities
def get_token_file_path() -> str:
"""Get the path to store the authentication token"""
@ -27,7 +64,7 @@ def get_token_file_path() -> str:
return str(config_dir / "token.json")
def save_token(token_data: dict[str, Any]) -> None:
def save_token(token_data: Mapping[str, object]) -> None:
"""Save token data to file"""
write_private_json(get_token_file_path(), token_data)
@ -65,7 +102,7 @@ def get_stored_api_key(expected_base_url: str | None = None) -> str | None:
# Team selection utilities
def display_teams_table(teams: list[dict[str, Any]]) -> None:
def display_teams_table(teams: Sequence[TeamRecord]) -> None:
"""Display teams in a formatted table"""
console: Final = Console()
@ -249,10 +286,11 @@ def prompt_team_selection_fallback(
while True:
try:
choice = click.prompt(
raw_choice: str = click.prompt(
"\nSelect a team by entering the index number (or 'skip' to continue without a team)",
type=str,
).strip()
)
choice = raw_choice.strip()
if choice.lower() == "skip":
return None
@ -275,7 +313,7 @@ def prompt_team_selection_fallback(
def _response_error_detail(response: requests.Response) -> str | None:
try:
body: Final = response.json()
body: Final[Mapping[str, object]] = response.json()
except ValueError:
return None
detail: Final = body.get("detail") if isinstance(body, dict) else None
@ -309,15 +347,16 @@ def _poll_for_ready_data(
other_status_log_every: int = 10,
http_error_log_every: int = 10,
connection_error_log_every: int = 10,
) -> dict[str, Any] | None:
) -> PollResponse | None:
for attempt in range(total_timeout // poll_interval):
try:
request_kwargs: dict[str, Any] = {"timeout": request_timeout}
if headers is not None:
request_kwargs["headers"] = headers
response = requests.get(url, **request_kwargs)
response = (
requests.get(url, timeout=request_timeout)
if headers is None
else requests.get(url, timeout=request_timeout, headers=headers)
)
if response.status_code == 200:
data = response.json()
data: PollResponse = response.json()
status = data.get("status")
if status == "ready":
return data
@ -341,7 +380,7 @@ def _poll_for_ready_data(
return None
def _normalize_teams(teams, team_details):
def _normalize_teams(teams: Sequence[str], team_details: Sequence[Mapping[str, str]] | None) -> list[TeamRecord]:
"""If team_details are a
Args:
@ -365,7 +404,7 @@ def _normalize_teams(teams, team_details):
return []
def _start_cli_sso_flow(base_url: str) -> dict[str, Any]:
def _start_cli_sso_flow(base_url: str) -> CliSsoStartResponse:
start_url: Final = f"{base_url}/sso/cli/start"
try:
response: Final = requests.post(start_url, timeout=10)
@ -389,7 +428,7 @@ def _start_cli_sso_flow(base_url: str) -> dict[str, Any]:
)
try:
data: Final = response.json()
data: Final[CliSsoStartResponse] = response.json()
except ValueError:
content_type: Final = response.headers.get("content-type", "unknown")
raise ValueError(
@ -398,8 +437,10 @@ def _start_cli_sso_flow(base_url: str) -> dict[str, Any]:
f"Response starts with: {response.text[:200]!r}"
)
required_fields: Final = ("login_id", "poll_secret", "user_code")
missing_fields: Final = tuple(field for field in required_fields if not isinstance(data.get(field), str))
string_fields: Final = frozenset(name for name, value in data.items() if isinstance(value, str))
missing_fields: Final = tuple(
field for field in ("login_id", "poll_secret", "user_code") if field not in string_fields
)
if missing_fields:
raise ValueError(
f"The response from {start_url} is missing required field(s): {', '.join(missing_fields)}. "
@ -412,7 +453,7 @@ def _get_cli_sso_poll_headers(poll_secret: str) -> dict[str, str]:
return {"x-litellm-cli-poll-secret": poll_secret}
def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> dict | None:
def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> AuthResult | None:
"""
Poll the server for authentication completion and handle team selection.
@ -431,7 +472,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> di
teams = data.get("teams", [])
team_details: Final = data.get("team_details")
user_id = data.get("user_id")
normalized_teams: Final[list[dict[str, Any]]] = _normalize_teams(teams, team_details)
normalized_teams: Final = _normalize_teams(teams, team_details)
if not normalized_teams:
click.echo("Warning: No teams available for selection.")
return None
@ -478,7 +519,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> di
def _handle_team_selection_during_polling(
base_url: str, key_id: str, poll_secret: str, teams: list[dict[str, Any]]
base_url: str, key_id: str, poll_secret: str, teams: Sequence[TeamRecord]
) -> str | None:
"""
Handle team selection and re-poll with selected team_id.
@ -522,7 +563,7 @@ def _handle_team_selection_during_polling(
return None
def _render_and_prompt_for_team_selection(teams: list[dict[str, Any]]) -> str | None:
def _render_and_prompt_for_team_selection(teams: Sequence[TeamRecord]) -> str | None:
"""Render teams table and prompt user for a team selection.
Returns the selected team_id as a string, or None if selection was
@ -546,10 +587,11 @@ def _render_and_prompt_for_team_selection(teams: list[dict[str, Any]]) -> str |
# Simple selection
while True:
try:
choice = click.prompt(
raw_choice: str = click.prompt(
"\nSelect a team by entering the index number (or 'skip' to use first team)",
type=str,
).strip()
)
choice = raw_choice.strip()
if choice.lower() == "skip":
# Default to the first team's ID if the user skips an
@ -582,7 +624,8 @@ def login(ctx: click.Context):
from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER
from litellm.proxy.client.cli.interface import show_commands
base_url: Final = ctx.obj["base_url"]
cli_obj: Final[CliContext] = ctx.obj
base_url: Final = cli_obj["base_url"]
try:
cli_sso_flow: Final = _start_cli_sso_flow(base_url=base_url)
@ -666,6 +709,7 @@ def print_token(ctx: click.Context):
expires after `LITELLM_CLI_JWT_EXPIRATION_HOURS` (default 24h); once
expired, run `lite login` again.
"""
cli_obj: Final[CliContext] = ctx.obj
token_data: Final = load_token()
if not token_data:
click.echo("Not authenticated. Run 'lite login'.", err=True)
@ -675,8 +719,8 @@ def print_token(ctx: click.Context):
# explicitly pointed us at a server, trust whichever one `lite login`
# actually issued this token for -- that's the whole point of not
# needing a wrapper command.
if ctx.obj.get("base_url_explicit"):
base_url: Final = ctx.obj["base_url"]
if cli_obj.get("base_url_explicit"):
base_url: Final = cli_obj["base_url"]
if token_data.get("base_url") != base_url.rstrip("/"):
click.echo("Not authenticated for this server. Run 'lite login'.", err=True)
sys.exit(1)

View file

@ -10,9 +10,9 @@ POST /cache/settings - Save cache settings to database
import asyncio
import json
from collections.abc import Mapping
from collections.abc import Coroutine, Mapping
from datetime import datetime, timezone
from typing import Any, Final
from typing import Any, Final, Protocol
from fastapi import APIRouter, Depends, Header, HTTPException
from pydantic import BaseModel, Field
@ -63,6 +63,35 @@ _REDACTED_VALUE: Final = "***REDACTED***"
_URL_OVERRIDDEN_CONNECTION_FIELDS: Final[frozenset] = frozenset({"host", "port", "db", "password", "username"})
class _CacheConfigRow(Protocol):
"""The single LiteLLM_CacheConfig row, narrowed to the column this module reads."""
cache_settings: object
class _CacheConfigTable(Protocol):
"""The prisma actions this module calls on the LiteLLM_CacheConfig table."""
def find_unique(self, *, where: Mapping[str, str]) -> Coroutine[None, None, _CacheConfigRow | None]: ...
def upsert(
self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, str]]
) -> Coroutine[None, None, _CacheConfigRow]: ...
class _CacheConfigTableProvider(Protocol):
@property
def table(self) -> _CacheConfigTable: ...
def _repository_table(repository: _CacheConfigTableProvider) -> _CacheConfigTable:
return repository.table
def _cache_config_table(prisma_client: object) -> _CacheConfigTable:
return _repository_table(CacheConfigRepository(prisma_client))
def _resolve_cache_url_precedence(settings: Mapping[str, object]) -> dict[str, Any]:
"""Return cache settings with the url-vs-discrete-fields ambiguity resolved.
@ -197,7 +226,7 @@ def _saved_secret_is_reusable(incoming: Mapping[str, object], saved: Mapping[str
return True
def _merge_over_saved(incoming: Mapping[str, object], saved: Mapping[str, object]) -> dict[str, Any]:
def _merge_over_saved(incoming: Mapping[str, object], saved: Mapping[str, object]) -> dict[str, object]:
"""Keep the stored secret behind any credential the caller echoed back redacted or omitted.
GET returns credentials as the marker and the form never re-prefills a
@ -339,7 +368,7 @@ class CacheSettingsManager:
return normalized1 == normalized2
@staticmethod
async def init_cache_settings_in_db(prisma_client, proxy_config):
async def init_cache_settings_in_db(prisma_client: object, proxy_config):
"""
Initialize cache settings from database into the router on startup.
Only reinitializes if cache params have changed.
@ -347,18 +376,17 @@ class CacheSettingsManager:
import json
try:
cache_config: Final = await call_with_db_reconnect_retry(
cache_config: Final[_CacheConfigRow | None] = await call_with_db_reconnect_retry(
prisma_client,
lambda: CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"}),
lambda: _cache_config_table(prisma_client).find_unique(where={"id": "cache_config"}),
reason="init_cache_settings_in_db_lookup_failure",
)
if cache_config is not None and cache_config.cache_settings:
# Parse cache settings JSON
cache_settings_json: Final = cache_config.cache_settings
if isinstance(cache_settings_json, str):
cache_settings_dict = json.loads(cache_settings_json)
else:
cache_settings_dict = cache_settings_json
cache_settings_json: Final[object] = cache_config.cache_settings
cache_settings_dict: Final[object] = (
json.loads(cache_settings_json) if isinstance(cache_settings_json, str) else cache_settings_json
)
# Decrypt cache settings
decrypted_settings: Final = proxy_config._decrypt_db_variables(variables_dict=cache_settings_dict)
@ -444,7 +472,9 @@ async def get_cache_settings(
# Read the stored settings (decrypted); an env-only cache has none.
stored: dict[str, object] = {}
if prisma_client is not None:
cache_config = await CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"})
cache_config: Final[_CacheConfigRow | None] = await _cache_config_table(prisma_client).find_unique(
where={"id": "cache_config"}
)
if cache_config is not None and cache_config.cache_settings:
stored = proxy_config._decrypt_db_variables(
variables_dict=_parse_stored_settings(cache_config.cache_settings)
@ -511,7 +541,7 @@ async def test_cache_connection(
saved_settings: dict[str, object] = {}
if prisma_client is not None:
try:
existing_row: Final = await CacheConfigRepository(prisma_client).table.find_unique(
existing_row: Final[_CacheConfigRow | None] = await _cache_config_table(prisma_client).find_unique(
where={"id": "cache_config"}
)
if existing_row is not None and existing_row.cache_settings:
@ -590,7 +620,9 @@ async def update_cache_settings(
try:
# Read the stored row first: its decrypted values back any credential the
# caller echoed back redacted, and its key set drives the audit diff.
existing_row: Final = await CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"})
existing_row: Final[_CacheConfigRow | None] = await _cache_config_table(prisma_client).find_unique(
where={"id": "cache_config"}
)
before_settings: dict[str, object] | None = None
saved_settings: dict[str, object] = {}
if existing_row is not None and existing_row.cache_settings:
@ -606,7 +638,7 @@ async def update_cache_settings(
encrypted_settings: Final = proxy_config._encrypt_env_variables(environment_variables=cache_settings)
# Save to database
await CacheConfigRepository(prisma_client).table.upsert(
await _cache_config_table(prisma_client).upsert(
where={"id": "cache_config"},
data={
"create": {

View file

@ -18,11 +18,14 @@ Scoping:
"""
import json
from typing import Any, Final
from collections.abc import Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Final, Protocol
from fastapi import APIRouter, Depends, HTTPException, Query
from litellm._logging import verbose_proxy_logger
from litellm.models.team import LiteLLM_TeamTable
from litellm.proxy._types import (
CommonProxyErrors,
LitellmUserRoles,
@ -40,10 +43,57 @@ from litellm.types.memory_management import (
MemoryUpdateRequest,
)
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
router: Final = APIRouter()
def _serialize_metadata_for_prisma(metadata: Any) -> str:
class _MemoryRecord(Protocol):
memory_id: str
key: str
value: str
user_id: str | None
team_id: str | None
created_at: datetime | None
created_by: str | None
updated_at: datetime | None
updated_by: str | None
class _MemoryTable(Protocol):
async def create(self, data: Mapping[str, object]) -> _MemoryRecord: ...
async def count(self, where: Mapping[str, object]) -> int: ...
async def find_many(
self,
where: Mapping[str, object],
order: Mapping[str, str],
take: int,
skip: int = 0,
) -> list[_MemoryRecord]: ...
async def update(self, where: Mapping[str, str], data: Mapping[str, object]) -> _MemoryRecord: ...
async def delete(self, where: Mapping[str, str]) -> _MemoryRecord | None: ...
class _TeamTable(Protocol):
async def find_unique(self, where: Mapping[str, str]) -> LiteLLM_TeamTable | None: ...
def _memory_table(prisma_client: "PrismaClient") -> _MemoryTable:
table: Final[_MemoryTable] = MemoryRepository(prisma_client).table
return table
def _team_table(prisma_client: "PrismaClient") -> _TeamTable:
table: Final[_TeamTable] = TeamRepository(prisma_client).table
return table
def _serialize_metadata_for_prisma(metadata: object) -> str:
"""
Encode a `metadata` payload for the `Json?` column.
@ -62,25 +112,25 @@ def _is_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:
return user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> dict | None:
def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object] | None:
"""
Prisma `where` fragment restricting rows to those the caller can see.
Returns None for admins (no restriction).
"""
if user_api_key_has_admin_view(user_api_key_dict):
return None
ors: Final[list[dict]] = []
if user_api_key_dict.user_id:
ors.append({"user_id": user_api_key_dict.user_id})
if user_api_key_dict.team_id:
ors.append({"team_id": user_api_key_dict.team_id})
ors: Final[list[Mapping[str, str]]] = [
{field: value}
for field, value in (("user_id", user_api_key_dict.user_id), ("team_id", user_api_key_dict.team_id))
if value
]
if not ors:
# Caller has neither user_id nor team_id — match nothing.
return {"memory_id": "__no_match__"}
return {"OR": ors}
def _row_to_model(row: Any) -> LiteLLM_MemoryRow:
def _row_to_model(row: _MemoryRecord) -> LiteLLM_MemoryRow:
return LiteLLM_MemoryRow(
memory_id=row.memory_id,
key=row.key,
@ -95,7 +145,7 @@ def _row_to_model(row: Any) -> LiteLLM_MemoryRow:
)
def _require_prisma():
def _require_prisma() -> "PrismaClient":
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
@ -113,7 +163,9 @@ def _internal_error(log_message: str, exc: Exception, default_detail: str) -> HT
return HTTPException(status_code=500, detail=default_detail)
async def _assert_write_access(prisma_client: Any, row: Any, user_api_key_dict: UserAPIKeyAuth) -> None:
async def _assert_write_access(
prisma_client: "PrismaClient", row: _MemoryRecord, user_api_key_dict: UserAPIKeyAuth
) -> None:
"""
Enforce ownership for mutations (PUT/DELETE).
@ -135,8 +187,8 @@ async def _assert_write_access(prisma_client: Any, row: Any, user_api_key_dict:
"""
if _is_admin(user_api_key_dict):
return
row_user_id: Final = getattr(row, "user_id", None)
row_team_id: Final = getattr(row, "team_id", None)
row_user_id: Final[str | None] = getattr(row, "user_id", None)
row_team_id: Final[str | None] = getattr(row, "team_id", None)
# Personal ownership.
if row_user_id and row_user_id == user_api_key_dict.user_id:
@ -153,7 +205,7 @@ async def _assert_write_access(prisma_client: Any, row: Any, user_api_key_dict:
)
async def _is_team_admin_for(prisma_client: Any, user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool:
async def _is_team_admin_for(prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool:
"""
True if the caller is a team admin of `team_id`, or an org admin for the
team's organization. Mirrors the auth pattern used by team-management
@ -168,7 +220,7 @@ async def _is_team_admin_for(prisma_client: Any, user_api_key_dict: UserAPIKeyAu
)
try:
team_obj: Final = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id})
team_obj: Final = await _team_table(prisma_client).find_unique(where={"team_id": team_id})
except Exception as e:
verbose_proxy_logger.exception("Error loading team for write-auth check (team_id=%s): %s", team_id, e)
return False
@ -261,7 +313,7 @@ def _resolve_scope(
async def create_memory(
body: MemoryCreateRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
) -> LiteLLM_MemoryRow:
"""Create a new memory entry for the caller (or, for admins, any scope)."""
prisma_client: Final = _require_prisma()
user_id, team_id = _resolve_scope(user_api_key_dict, body.user_id, body.team_id)
@ -269,19 +321,19 @@ async def create_memory(
# `metadata` is a `Json?` column — prisma-client-python rejects raw
# Python values, so JSON-encode any non-null payload and omit the field
# entirely when None so the column defaults to SQL NULL.
create_data: Final[dict] = {
metadata: Final[object] = body.metadata
create_data: Final[Mapping[str, object]] = {
"key": body.key,
"value": body.value,
"user_id": user_id,
"team_id": team_id,
"created_by": user_api_key_dict.user_id,
"updated_by": user_api_key_dict.user_id,
**({"metadata": _serialize_metadata_for_prisma(metadata)} if metadata is not None else {}),
}
if body.metadata is not None:
create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata)
try:
row: Final = await MemoryRepository(prisma_client).table.create(data=create_data)
row: Final = await _memory_table(prisma_client).create(data=create_data)
except Exception as e:
# Key is globally unique. Any duplicate → 409.
if _is_unique_violation(e):
@ -316,7 +368,7 @@ async def list_memory(
page: int = Query(1, ge=1),
page_size: int = Query(50, ge=1, le=500),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
) -> MemoryListResponse:
"""List memory entries visible to the caller."""
prisma_client: Final = _require_prisma()
@ -325,24 +377,18 @@ async def list_memory(
# top-level "AND" — safer than `dict.update` since future visibility
# filters could grow an "OR" key that would clobber this one if merged
# by key.
key_filter: Final[dict] = {}
if key_prefix is not None:
key_filter["key"] = {"startsWith": key_prefix}
elif key is not None:
key_filter["key"] = key
key_filter: Final[Mapping[str, object]] = (
{"key": {"startsWith": key_prefix}} if key_prefix is not None else {"key": key} if key is not None else {}
)
vis: Final = _visibility_filter(user_api_key_dict)
where: dict
if vis is None:
where = key_filter
elif not key_filter:
where = vis
else:
where = {"AND": [key_filter, vis]}
where: Final[Mapping[str, object]] = (
key_filter if vis is None else vis if not key_filter else {"AND": [key_filter, vis]}
)
try:
total: Final = await MemoryRepository(prisma_client).table.count(where=where)
rows: Final = await MemoryRepository(prisma_client).table.find_many(
total: Final = await _memory_table(prisma_client).count(where=where)
rows: Final = await _memory_table(prisma_client).find_many(
where=where,
order={"updated_at": "desc"},
skip=(page - 1) * page_size,
@ -354,12 +400,14 @@ async def list_memory(
return MemoryListResponse(memories=[_row_to_model(r) for r in rows], total=total)
async def _find_memory_for_caller(prisma_client: Any, key: str, user_api_key_dict: UserAPIKeyAuth) -> Any:
async def _find_memory_for_caller(
prisma_client: "PrismaClient", key: str, user_api_key_dict: UserAPIKeyAuth
) -> _MemoryRecord:
"""Look up a memory row by key, scoped to the caller's visibility."""
key_filter: Final[dict] = {"key": key}
key_filter: Final[Mapping[str, object]] = {"key": key}
vis: Final = _visibility_filter(user_api_key_dict)
where: Final[dict] = key_filter if vis is None else {"AND": [key_filter, vis]}
rows = await MemoryRepository(prisma_client).table.find_many(where=where, take=1, order={"updated_at": "desc"})
where: Final[Mapping[str, object]] = key_filter if vis is None else {"AND": [key_filter, vis]}
rows = await _memory_table(prisma_client).find_many(where=where, take=1, order={"updated_at": "desc"})
if not rows:
raise HTTPException(status_code=404, detail=f"Memory with key '{key}' not found")
return rows[0]
@ -374,7 +422,7 @@ async def _find_memory_for_caller(prisma_client: Any, key: str, user_api_key_dic
async def get_memory(
key: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
) -> LiteLLM_MemoryRow:
"""Get a single memory entry by key, scoped to the caller."""
prisma_client: Final = _require_prisma()
row: Final = await _find_memory_for_caller(prisma_client, key, user_api_key_dict)
@ -391,7 +439,7 @@ async def upsert_memory(
key: str,
body: MemoryUpdateRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
) -> LiteLLM_MemoryRow:
"""
Upsert a memory entry by key within the caller's scope.
@ -414,20 +462,24 @@ async def upsert_memory(
# `model_fields_set`), the column is preserved as-is.
fields_sent: Final = body.model_fields_set
metadata_in_payload: Final = "metadata" in fields_sent
metadata: Final[object] = body.metadata
data: Final[dict] = {}
if body.value is not None:
data["value"] = body.value
if metadata_in_payload:
data["metadata"] = _serialize_metadata_for_prisma(body.metadata)
if not data:
value_field: Final[Mapping[str, object]] = {"value": body.value} if body.value is not None else {}
metadata_field: Final[Mapping[str, object]] = (
{"metadata": _serialize_metadata_for_prisma(metadata)} if metadata_in_payload else {}
)
if not value_field and not metadata_field:
raise HTTPException(
status_code=400,
detail="Request body must include at least one of: value, metadata.",
)
data["updated_by"] = user_api_key_dict.user_id
data: Final[Mapping[str, object]] = {
**value_field,
**metadata_field,
"updated_by": user_api_key_dict.user_id,
}
async def _find_existing() -> Any:
async def _find_existing() -> _MemoryRecord | None:
"""Return the caller-visible row for `key`, or None."""
try:
return await _find_memory_for_caller(prisma_client, key, user_api_key_dict)
@ -444,7 +496,7 @@ async def upsert_memory(
# their team) — otherwise a teammate could overwrite a personal
# entry through the OR-based visibility filter.
await _assert_write_access(prisma_client, existing, user_api_key_dict)
row = await MemoryRepository(prisma_client).table.update(
row = await _memory_table(prisma_client).update(
where={"memory_id": existing.memory_id},
data=data,
)
@ -459,18 +511,17 @@ async def upsert_memory(
# Omit `metadata` when None so the column defaults to SQL NULL;
# otherwise JSON-encode for Prisma — same pattern as
# `create_memory` above.
create_data: Final[dict] = {
create_data: Final[Mapping[str, object]] = {
"key": key,
"value": body.value,
"user_id": user_id,
"team_id": team_id,
"created_by": user_api_key_dict.user_id,
"updated_by": user_api_key_dict.user_id,
**({"metadata": _serialize_metadata_for_prisma(metadata)} if metadata is not None else {}),
}
if body.metadata is not None:
create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata)
try:
row = await MemoryRepository(prisma_client).table.create(data=create_data)
row = await _memory_table(prisma_client).create(data=create_data)
except Exception as e:
# Race: a concurrent PUT/POST created the row after our check.
# Re-read and fall back to an update so the PUT stays idempotent
@ -487,7 +538,7 @@ async def upsert_memory(
)
# Same write-authorization check as the non-race path.
await _assert_write_access(prisma_client, existing_after_race, user_api_key_dict)
row = await MemoryRepository(prisma_client).table.update(
row = await _memory_table(prisma_client).update(
where={"memory_id": existing_after_race.memory_id},
data=data,
)
@ -508,14 +559,14 @@ async def upsert_memory(
async def delete_memory(
key: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
) -> MemoryDeleteResponse:
"""Delete a memory entry by key, scoped to the caller."""
prisma_client: Final = _require_prisma()
row: Final = await _find_memory_for_caller(prisma_client, key, user_api_key_dict)
# Visibility != write authority — see the upsert handler for the rationale.
await _assert_write_access(prisma_client, row, user_api_key_dict)
try:
await MemoryRepository(prisma_client).table.delete(where={"memory_id": row.memory_id})
await _memory_table(prisma_client).delete(where={"memory_id": row.memory_id})
except Exception as e:
raise _internal_error("Error deleting memory: %s", e, "Internal error deleting memory entry.")

View file

@ -10,7 +10,8 @@ https://platform.openai.com/docs/api-reference/responses-streaming
import asyncio
import json
from typing import Any, Final, cast
from collections.abc import AsyncIterable, Callable, Mapping
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, TypedDict
from fastapi import Request, Response
@ -18,28 +19,77 @@ from litellm._logging import verbose_proxy_logger
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
from litellm.types.llms.openai import ResponsesAPIStatus
if TYPE_CHECKING:
from litellm.proxy.proxy_server import ProxyConfig
_ContentEntry: TypeAlias = dict[str, object] | list[object] | str | int | float | bool | None
class _StreamingBodyResponse(Protocol):
body_iterator: AsyncIterable[str | bytes | memoryview]
class _OutputItem(TypedDict, total=False):
id: str
content: list[_ContentEntry]
class _TerminalResponse(TypedDict, total=False):
status: ResponsesAPIStatus
error: dict[str, object]
usage: dict[str, object]
reasoning: dict[str, object]
tool_choice: str | dict[str, object]
tools: list[dict[str, object]]
model: str
instructions: str
temperature: float
top_p: float
max_output_tokens: int
previous_response_id: str
text: dict[str, object]
truncation: str
parallel_tool_calls: bool
user: str
store: bool
incomplete_details: dict[str, object]
output: list[_OutputItem]
class _StreamEvent(TypedDict, total=False):
type: str
item: _OutputItem
item_id: str
part: dict[str, object]
content_index: int
delta: str
response: _TerminalResponse
async def background_streaming_task(
polling_id: str,
data: dict,
data: dict[str, object],
polling_handler: ResponsePollingHandler,
request: Request,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth,
general_settings: dict,
llm_router,
proxy_config,
proxy_logging_obj,
select_data_generator,
user_model,
user_temperature,
user_request_timeout,
user_max_tokens,
user_api_base,
version,
):
general_settings: dict[str, object],
llm_router: Router | None,
proxy_config: "ProxyConfig",
proxy_logging_obj: ProxyLogging,
select_data_generator: Callable[..., object] | None,
user_model: str | None,
user_temperature: float | None,
user_request_timeout: float | None,
user_max_tokens: int | None,
user_api_base: str | None,
version: str | None,
) -> None:
"""
Background task to stream response and update cache
@ -69,7 +119,7 @@ async def background_streaming_task(
# Make streaming request.
# Pre-call checks (rate limits, guardrails, budget) were already run
# before polling ID creation, so skip them here to avoid double-counting.
response: Final = await processor.base_process_llm_request(
response: Final[_StreamingBodyResponse] = await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
@ -91,8 +141,9 @@ async def background_streaming_task(
# Process streaming response following OpenAI events format
# https://platform.openai.com/docs/api-reference/responses-streaming
output_items: Final[dict[str, dict[str, Any]]] = {} # Track output items by ID
accumulated_text: Final = {} # Track accumulated text deltas by (item_id, content_index)
output_items: Final[dict[str, _OutputItem]] = {} # Track output items by ID
# Track accumulated text deltas by (item_id, content_index)
accumulated_text: Final[dict[tuple[str, int], str]] = {}
# ResponsesAPIResponse fields to extract from response.completed
usage_data = None
@ -121,7 +172,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",
@ -162,7 +213,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
@ -181,9 +232,10 @@ async def background_streaming_task(
if item_id and item_id in output_items:
# Update the output item with new content
if "content" not in output_items[item_id]:
output_items[item_id]["content"] = []
output_items[item_id]["content"].append(content_part)
added_to_item = output_items[item_id]
if "content" not in added_to_item:
added_to_item["content"] = []
added_to_item["content"].append(content_part)
state_dirty = True
elif event_type == "response.output_text.delta":
@ -201,12 +253,14 @@ async def background_streaming_task(
accumulated_text[key] += delta
# Update the content in output_items
if "content" in output_items[item_id]:
content_list = output_items[item_id]["content"]
delta_item = output_items[item_id]
if "content" in delta_item:
content_list = delta_item["content"]
if content_index < len(content_list):
# Update existing content part with accumulated text
if isinstance(content_list[content_index], dict):
content_list[content_index]["text"] = accumulated_text[key]
existing_part = content_list[content_index]
if isinstance(existing_part, dict):
existing_part["text"] = accumulated_text[key]
state_dirty = True
elif event_type == "response.content_part.done":
@ -217,8 +271,9 @@ async def background_streaming_task(
if item_id and item_id in output_items:
# Update with final content from event
if "content" in output_items[item_id]:
content_list = output_items[item_id]["content"]
done_item = output_items[item_id]
if "content" in done_item:
content_list = done_item["content"]
if content_index < len(content_list):
content_list[content_index] = content_part
state_dirty = True
@ -248,12 +303,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

View file

@ -17,7 +17,12 @@ from __future__ import annotations
import hashlib
import uuid
from typing import TYPE_CHECKING, Any, Final
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
import httpx
from typing_extensions import NotRequired, TypedDict
import litellm
from litellm._logging import verbose_logger
@ -39,6 +44,44 @@ if TYPE_CHECKING:
from litellm.types.rag import RAGIngestOptions
_NO_FILENAME_METADATA: Final[Mapping[str, str]] = MappingProxyType({})
class S3VectorsVectorData(TypedDict):
float32: Sequence[float]
class S3VectorsRecord(TypedDict):
key: str
data: S3VectorsVectorData
metadata: Mapping[str, str]
class S3VectorsMetadataConfiguration(TypedDict):
nonFilterableMetadataKeys: Sequence[str]
class S3VectorsCreateIndexRequest(TypedDict):
vectorBucketName: str
indexName: str | None
dataType: str
dimension: int | None
distanceMetric: str
metadataConfiguration: NotRequired[S3VectorsMetadataConfiguration]
class S3VectorsQueryResultMetadata(TypedDict, total=False):
source_text: str
class S3VectorsQueryResult(TypedDict, total=False):
metadata: S3VectorsQueryResultMetadata
class S3VectorsQueryResponse(TypedDict, total=False):
vectors: Sequence[S3VectorsQueryResult]
class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
"""
S3 Vectors RAG ingestion using httpx + AWS SigV4 signing.
@ -57,6 +100,12 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
- non_filterable_metadata_keys: List of metadata keys to exclude from filtering
"""
vector_bucket_name: str
index_name: str | None
distance_metric: str
non_filterable_metadata_keys: Sequence[str]
dimension: int | None
def __init__(
self,
ingest_options: RAGIngestOptions,
@ -78,7 +127,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
self.dimension = self._get_dimension_from_config()
# Get AWS region using BaseAWSLLM method
_aws_region: Final = self.vector_store_config.get("aws_region_name")
_aws_region: Final[str | None] = self.vector_store_config.get("aws_region_name")
self.aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(
aws_region_name=str(_aws_region) if _aws_region else None
)
@ -166,7 +215,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
url: str,
data: str | None = None,
headers: dict[str, str] | None = None,
) -> Any:
) -> httpx.Response:
"""
Helper to sign and execute AWS API requests using httpx + SigV4.
@ -223,18 +272,24 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
signed_headers: Final = dict(aws_request.headers.items())
# Make the request using specific method (pattern from s3_v2.py)
method_upper: Final = method.upper()
if method_upper == "PUT":
response = await self.async_httpx_client.put(url, data=data, headers=signed_headers)
elif method_upper == "POST":
response = await self.async_httpx_client.post(url, data=data, headers=signed_headers)
elif method_upper == "GET":
response = await self.async_httpx_client.get(url, headers=signed_headers)
else:
raise ValueError(f"Unsupported HTTP method: {method}")
response: Final = await self._execute_signed_request(method, url, data, signed_headers)
if response is None:
raise ValueError(f"No response from S3 Vectors for {method} {url}")
return response
async def _execute_signed_request(
self, method: str, url: str, data: str | None, headers: dict[str, str]
) -> httpx.Response | None:
method_upper: Final = method.upper()
if method_upper == "PUT":
return await self.async_httpx_client.put(url, data=data, headers=headers)
if method_upper == "POST":
return await self.async_httpx_client.post(url, data=data, headers=headers)
if method_upper == "GET":
return await self.async_httpx_client.get(url, headers=headers)
raise ValueError(f"Unsupported HTTP method: {method}")
async def _ensure_vector_bucket_exists(self):
"""Create vector bucket if it doesn't exist using GetVectorBucket and CreateVectorBucket APIs."""
verbose_logger.debug("Ensuring S3 vector bucket exists: %s", self.vector_bucket_name)
@ -311,16 +366,24 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
)
# Prepare index configuration per AWS API docs
index_config: Final = {
"vectorBucketName": self.vector_bucket_name,
"indexName": self.index_name,
"dataType": "float32",
"dimension": self.dimension,
"distanceMetric": self.distance_metric,
}
base_index_config: Final = S3VectorsCreateIndexRequest(
vectorBucketName=self.vector_bucket_name,
indexName=self.index_name,
dataType="float32",
dimension=self.dimension,
distanceMetric=self.distance_metric,
)
if self.non_filterable_metadata_keys:
index_config["metadataConfiguration"] = {"nonFilterableMetadataKeys": self.non_filterable_metadata_keys}
index_config: Final[S3VectorsCreateIndexRequest] = (
{
**base_index_config,
"metadataConfiguration": S3VectorsMetadataConfiguration(
nonFilterableMetadataKeys=self.non_filterable_metadata_keys
),
}
if self.non_filterable_metadata_keys
else base_index_config
)
create_url: Final = f"https://s3vectors.{self.aws_region_name}.api.aws/CreateIndex"
response = await self._sign_and_execute_request("POST", create_url, data=safe_dumps(index_config))
@ -336,7 +399,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
verbose_logger.exception("Error creating vector index: %s", e)
raise
async def _put_vectors(self, vectors: list[dict[str, Any]]):
async def _put_vectors(self, vectors: Sequence[S3VectorsRecord]):
"""
Call PutVectors API to store vectors in S3 Vectors.
@ -442,24 +505,19 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
raise ValueError(error_msg)
# Prepare vectors for PutVectors API
vectors: Final = []
for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
# Build metadata dict
metadata: dict[str, str] = {
"source_text": chunk, # Non-filterable (for reference)
"chunk_index": str(i), # Filterable
}
if filename:
metadata["filename"] = filename # Filterable
vector_obj = {
"key": f"{filename}_{i}" if filename else f"chunk_{i}",
"data": {"float32": embedding},
"metadata": metadata,
}
vectors.append(vector_obj)
filename_metadata: Final[Mapping[str, str]] = {"filename": filename} if filename else _NO_FILENAME_METADATA
vectors: Final = tuple(
S3VectorsRecord(
key=f"{filename}_{i}" if filename else f"chunk_{i}",
data=S3VectorsVectorData(float32=embedding),
metadata={
"source_text": chunk, # Non-filterable (for reference)
"chunk_index": str(i), # Filterable
**filename_metadata, # Filterable
},
)
for i, (chunk, embedding) in enumerate(zip(chunks, embeddings))
)
# Call PutVectors API
await self._put_vectors(vectors)
@ -468,7 +526,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
vector_store_id: Final = f"{self.vector_bucket_name}:{self.index_name}"
return vector_store_id, filename
async def query_vector_store(self, vector_store_id: str, query: str, top_k: int = 5) -> dict[str, Any] | None:
async def query_vector_store(
self, vector_store_id: str, query: str, top_k: int = 5
) -> S3VectorsQueryResponse | None:
"""
Query S3 Vectors using QueryVectors API.
@ -507,12 +567,13 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
response = await self._sign_and_execute_request("POST", url, data=safe_dumps(request_body))
if response.status_code == 200:
results: Final = response.json()
results: Final[S3VectorsQueryResponse] = response.json()
verbose_logger.debug("Query returned %s results", len(results.get("vectors", [])))
# Check if query terms appear in results
if results.get("vectors"):
for result in results["vectors"]:
matched_vectors: Final = results.get("vectors")
if matched_vectors:
for result in matched_vectors:
metadata = result.get("metadata", {})
source_text = metadata.get("source_text", "")
if query.lower() in source_text.lower():

View file

@ -1,6 +1,6 @@
{
"ANN001": {
"limit": 3114
"limit": 3101
},
"ANN002": {
"limit": 71
@ -9,10 +9,10 @@
"limit": 834
},
"ANN201": {
"limit": 2031
"limit": 2025
},
"ANN202": {
"limit": 865
"limit": 863
},
"ANN204": {
"limit": 713
@ -24,7 +24,7 @@
"limit": 133
},
"ANN401": {
"limit": 1555
"limit": 1543
},
"ASYNC230": {
"limit": 11
@ -201,7 +201,7 @@
"limit": 58
},
"SIM102": {
"limit": 322
"limit": 321
},
"SIM103": {
"limit": 119
@ -234,7 +234,7 @@
"limit": 5
},
"TID251": {
"limit": 1238
"limit": 1234
},
"TRY002": {
"limit": 528

View file

@ -1,6 +1,6 @@
{
"LIT001": {
"limit": 23149
"limit": 23144
},
"LIT002": {
"limit": 27166
@ -15,7 +15,7 @@
"limit": 0
},
"LIT006": {
"limit": 1086
"limit": 1085
},
"LIT007": {
"limit": 0
@ -27,9 +27,9 @@
"limit": 0
},
"LIT010": {
"limit": 16758
"limit": 16735
},
"LIT011": {
"limit": 5598
"limit": 5597
}
}