scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 48

refactor(repositories): type prisma table access with one generic protocol

Browse source 
Every repository handed its `.table` back untyped, so a dozen modules had
each grown a private `_PrismaTableActions` Protocol to paper over it. They
had drifted: some declared `update` as returning the row, others the row or
None, and none agreed on whether `find_many` was covariant

Replace all of them with a single `TableActions[RowT_co]` in
`litellm/repositories/prisma_protocols.py`, keyed to the prisma row each
repository is bound to. Query inputs stay `Mapping[str, object]` so callers
keep passing plain dicts, and `find_many` returns `Sequence` so the row type
stays covariant

Typing the nullable returns honestly surfaced paths that were already
crashing. A team admin could never edit or delete a memory entry owned by
their team: the write-auth check fed a raw prisma row to a helper that
expects the domain model, so `members_with_roles` arrived as plain dicts and
the request died as a 500 instead of applying the edit. Non-admin members hit
the same 500 in place of the 403 they were owed, so refusal and breakage were
indistinguishable. `/v2/model/info?user_models_only=true` dereferenced a
missing user row rather than returning the 400 the route already had, three
team routes dereferenced a team deleted between the read and the write, and
the agent registry dereferenced a missing agent instead of naming it

basedpyright drops 2,132 errors, 1,454 of them reportAny and 73
reportExplicitAny. The dashboard's generated types pick up `string[]` where
they had `unknown[]` for a team's members, admins and models
This commit is contained in:
mateo-berri 2026-08-25 12:14:17 +00:00
parent a9c7b848f2
commit 9dabd72f2d
102 changed files with 2320 additions and 1325 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
basedpyright-code-budget.json
View file

@ -1,18 +1,18 @@
{
"reportAny": {
"limit": 19955
"limit": 18501
},
"reportArgumentType": {
"limit": 2566
"limit": 2564
},
"reportAssignmentType": {
"limit": 320
},
"reportAttributeAccessIssue": {
"limit": 488
"limit": 483
},
"reportCallIssue": {
"limit": 114
"limit": 113
},
"reportConstantRedefinition": {
"limit": 40
@ -24,7 +24,7 @@
"limit": 19
},
"reportExplicitAny": {
"limit": 6049
"limit": 5976
},
"reportFunctionMemberAccess": {
"limit": 7
@ -45,7 +45,7 @@
"limit": 35
},
"reportInvalidTypeForm": {
"limit": 35
"limit": 34
},
"reportInvalidTypeVarUse": {
"limit": 2
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5663
"limit": 5661
},
"reportMissingTypeArgument": {
"limit": 15555
"limit": 15504
},
"reportMissingTypeStubs": {
"limit": 40
@ -72,7 +72,7 @@
"limit": 0
},
"reportOptionalMemberAccess": {
"limit": 1061
"limit": 1058
},
"reportOptionalOperand": {
"limit": 0
@ -90,7 +90,7 @@
"limit": 8
},
"reportReturnType": {
"limit": 213
"limit": 212
},
"reportTypedDictNotRequiredAccess": {
"limit": 26
@ -99,31 +99,31 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44655
"limit": 44527
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 39011
"limit": 38827
},
"reportUnknownParameterType": {
"limit": 19885
"limit": 19848
},
"reportUnknownVariableType": {
"limit": 30569
"limit": 30384
},
"reportUnnecessaryCast": {
"limit": 117
},
"reportUnnecessaryComparison": {
"limit": 699
"limit": 697
},
"reportUnnecessaryContains": {
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 836
"limit": 833
},
"reportUntypedBaseClass": {
"limit": 0

4
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -14,6 +14,8 @@ from litellm.constants import (
)
if TYPE_CHECKING:
from prisma import models as prisma_models
from litellm.integrations.prometheus import PrometheusLogger
from litellm.proxy._types import LiteLLM_ManagedObjectTable
from litellm.proxy.utils import PrismaClient, ProxyLogging
@ -351,7 +353,7 @@ class CheckBatchCost:
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
async def _finalize_unbilled_terminal_job(
self, job: "LiteLLM_ManagedObjectTable", response: "LiteLLMBatch"
self, job: "prisma_models.LiteLLM_ManagedObjectTable", response: "LiteLLMBatch"
) -> None:
"""Persist a terminal batch that has nothing billable, converting any raw
provider file ids to managed ids, and take it out of the poll page."""

5
litellm/integrations/prometheus.py
View file

@ -96,7 +96,10 @@ class _PaginatedPrismaTable(Protocol[_TableRowT]):
def _paginated_table(repository: BaseRepository[_TableRowT]) -> _PaginatedPrismaTable[_TableRowT]:
"""View a repository's prisma table through the pagination surface budget metrics need."""
return repository.table
return cast(
_PaginatedPrismaTable[_TableRowT],
repository.table, # cast-ok: prisma rows carry the budget columns the domain model declares
)
class _OrgBudgetRow(Protocol):

6
litellm/models/team.py
View file

@ -64,8 +64,8 @@ class TeamBase(LiteLLMPydanticObjectBase):
team_alias: str | None = None
team_id: str | None = None
organization_id: str | None = None
admins: list = []
members: list = []
admins: list[str] = []
members: list[str] = []
members_with_roles: list[Member] = []
team_member_permissions: list[str] | None = None
metadata: dict | None = None
@ -75,7 +75,7 @@ class TeamBase(LiteLLMPydanticObjectBase):
soft_budget: float | None = None
budget_duration: str | None = None
budget_limits: list[BudgetLimitEntry] | None = None
models: list = []
models: list[str] = []
blocked: bool = False
router_settings: dict | None = None
access_group_ids: list[str] | None = None

99
litellm/proxy/_experimental/mcp_server/db.py
View file

@ -4,7 +4,7 @@ import hashlib
import json
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, TypeVar, cast
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -13,7 +13,6 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.proxy._experimental.mcp_server.oauth_utils import build_upstream_oauth2_token_request
from litellm.proxy._types import (
LiteLLM_MCPServerTable,
LiteLLM_ObjectPermissionTable,
MCPApprovalStatus,
MCPEnvVar,
MCPEnvVarScope,
@ -30,6 +29,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
)
from litellm.proxy.utils import PrismaClient
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
MCPServerOAuthClientRepository,
MCPServerRepository,
@ -48,34 +48,9 @@ if TYPE_CHECKING:
from litellm.types.mcp_server.mcp_server_manager import MCPServer
_RowT = TypeVar("_RowT")
class _TableActions(Protocol[_RowT]):
async def find_unique(
self, where: Mapping[str, object], include: Mapping[str, object] | None = None
) -> _RowT | None: ...
async def find_many(
self,
take: int | None = None,
where: Mapping[str, object] | None = None,
order: Mapping[str, object] | None = None,
) -> list[_RowT]: ...
async def create(self, data: Mapping[str, object]) -> _RowT: ...
async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> _RowT: ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _RowT | None: ...
async def delete(self, where: Mapping[str, object]) -> _RowT | None: ...
async def delete_many(self, where: Mapping[str, object] | None = None) -> int: ...
class _UserEnvVarsTransactionClient(Protocol):
litellm_mcpuserenvvars: "_TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]"
litellm_mcpuserenvvars: "TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]"
async def execute_raw(self, query: str, *args: object) -> int: ...
@ -473,15 +448,15 @@ def _credentials_blob_to_mutable_dict(blob: str | Mapping[str, object]) -> dict[
def _mcp_server_table_actions(
prisma_client: PrismaClient,
) -> "_TableActions[prisma_db_models.LiteLLM_MCPServerTable]":
table: Final[_TableActions[prisma_db_models.LiteLLM_MCPServerTable]] = MCPServerRepository(prisma_client).table
) -> "TableActions[prisma_db_models.LiteLLM_MCPServerTable]":
table: Final[TableActions[prisma_db_models.LiteLLM_MCPServerTable]] = MCPServerRepository(prisma_client).table
return table
def _verification_token_table_actions(
prisma_client: PrismaClient,
) -> "_TableActions[prisma_db_models.LiteLLM_VerificationToken]":
table: Final[_TableActions[prisma_db_models.LiteLLM_VerificationToken]] = VerificationTokenRepository(
) -> "TableActions[prisma_db_models.LiteLLM_VerificationToken]":
table: Final[TableActions[prisma_db_models.LiteLLM_VerificationToken]] = VerificationTokenRepository(
prisma_client
).table
return table
@ -489,15 +464,15 @@ def _verification_token_table_actions(
def _team_table_actions(
prisma_client: PrismaClient,
) -> "_TableActions[prisma_db_models.LiteLLM_TeamTable]":
table: Final[_TableActions[prisma_db_models.LiteLLM_TeamTable]] = TeamRepository(prisma_client).table
) -> "TableActions[prisma_db_models.LiteLLM_TeamTable]":
table: Final[TableActions[prisma_db_models.LiteLLM_TeamTable]] = TeamRepository(prisma_client).table
return table
def _oauth_client_table_actions(
prisma_client: PrismaClient,
) -> "_TableActions[prisma_db_models.LiteLLM_MCPServerOAuthClient]":
table: Final[_TableActions[prisma_db_models.LiteLLM_MCPServerOAuthClient]] = MCPServerOAuthClientRepository(
) -> "TableActions[prisma_db_models.LiteLLM_MCPServerOAuthClient]":
table: Final[TableActions[prisma_db_models.LiteLLM_MCPServerOAuthClient]] = MCPServerOAuthClientRepository(
prisma_client
).table
return table
@ -511,7 +486,7 @@ def _db_transaction_manager(prisma_client: PrismaClient) -> _UserEnvVarsTransact
async def _db_find_mcp_server_rows(
prisma_client: PrismaClient,
where: "prisma_db_types.LiteLLM_MCPServerTableWhereInput | None" = None,
) -> "list[prisma_db_models.LiteLLM_MCPServerTable]":
) -> "Sequence[prisma_db_models.LiteLLM_MCPServerTable]":
return await _mcp_server_table_actions(prisma_client).find_many(where=where)
@ -526,17 +501,19 @@ 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(
row: Final[prisma_db_models.LiteLLM_MCPServerTable | None] = await _mcp_server_table_actions(prisma_client).update(
where={"server_id": server_id},
data=data,
)
if row is None:
raise ValueError(f"MCP server not found, passed server_id={server_id}")
return row
def _user_credential_actions(
prisma_client: PrismaClient,
) -> "_TableActions[prisma_db_models.LiteLLM_MCPUserCredentials]":
table: Final[_TableActions[prisma_db_models.LiteLLM_MCPUserCredentials]] = MCPUserCredentialsRepository(
) -> "TableActions[prisma_db_models.LiteLLM_MCPUserCredentials]":
table: Final[TableActions[prisma_db_models.LiteLLM_MCPUserCredentials]] = MCPUserCredentialsRepository(
prisma_client
).table
return table
@ -544,8 +521,8 @@ def _user_credential_actions(
def _user_env_var_actions(
prisma_client: PrismaClient,
) -> "_TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]":
table: Final[_TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]] = prisma_client.db.litellm_mcpuserenvvars
) -> "TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]":
table: Final[TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]] = prisma_client.db.litellm_mcpuserenvvars
return table
@ -560,7 +537,7 @@ 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)
@ -583,7 +560,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)
@ -658,7 +635,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 _mcp_server_table_actions(
_mcp_servers: Final[Sequence[prisma_db_models.LiteLLM_MCPServerTable]] = await _mcp_server_table_actions(
prisma_client
).find_many(
where={
@ -745,13 +722,13 @@ async def get_all_mcp_servers_for_user(
async def get_objectpermissions_for_mcp_server(
prisma_client: PrismaClient, mcp_server_id: str
) -> list[LiteLLM_ObjectPermissionTable]:
) -> "Sequence[prisma_db_models.LiteLLM_ObjectPermissionTable]":
"""
Get all the object permissions records and the associated team and verficiationtoken records that have access to the mcp server
"""
object_permission_records: Final[list[LiteLLM_ObjectPermissionTable]] = await ObjectPermissionRepository(
prisma_client
).table.find_many(
object_permission_records: Final[
Sequence[prisma_db_models.LiteLLM_ObjectPermissionTable]
] = await ObjectPermissionRepository(prisma_client).table.find_many(
where={
"mcp_servers": {"has": mcp_server_id},
},
@ -766,19 +743,19 @@ async def get_objectpermissions_for_mcp_server(
async def get_virtualkeys_for_mcp_server(
prisma_client: PrismaClient, server_id: str
) -> "list[prisma_db_models.LiteLLM_VerificationToken]":
) -> "Sequence[prisma_db_models.LiteLLM_VerificationToken]":
"""
Get all the virtual keys that have access to the mcp server
"""
virtual_keys: Final[list[prisma_db_models.LiteLLM_VerificationToken] | None] = await VerificationTokenRepository(
prisma_client
).table.find_many(
virtual_keys: Final[
Sequence[prisma_db_models.LiteLLM_VerificationToken] | None
] = await VerificationTokenRepository(prisma_client).table.find_many(
where={
"mcp_servers": {"has": server_id},
},
)
if virtual_keys is None:
if virtual_keys is None: # pyright: ignore[reportUnnecessaryComparison] # unreachable per seam types; kept as-is
return []
return virtual_keys
@ -860,7 +837,7 @@ async def delete_mcp_server(
invalidate_token_cache = global_mcp_server_manager.invalidate_user_oauth_token_cache
for user_id in credential_user_ids:
await invalidate_token_cache(user_id, server_id)
return deleted_server
return deleted_server # pyright: ignore[reportReturnType] # prisma row, not domain LiteLLM_MCPServerTable
async def create_mcp_server(
@ -880,7 +857,7 @@ async def create_mcp_server(
data_dict["updated_by"] = touched_by
new_mcp_server: Final[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.create(
data=data_dict
data=data_dict, # pyright: ignore[reportAssignmentType] # prisma row, not domain LiteLLM_MCPServerTable
)
_decrypt_env_vars_on_returned_row(new_mcp_server)
@ -982,7 +959,7 @@ async def update_mcp_server(
data: UpdateMCPServerRequest,
touched_by: str,
fields_set: set[str] | None = None,
) -> LiteLLM_MCPServerTable:
) -> LiteLLM_MCPServerTable | None:
"""
Update a new mcp server record in the db
"""
@ -1093,9 +1070,9 @@ async def update_mcp_server(
data_dict["credentials"] = Json(None)
updated_mcp_server: Final[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.update(
updated_mcp_server: Final[LiteLLM_MCPServerTable | None] = await MCPServerRepository(prisma_client).table.update(
where={"server_id": data.server_id},
data=data_dict,
data=data_dict, # pyright: ignore[reportAssignmentType] # prisma row, not domain LiteLLM_MCPServerTable
)
_decrypt_env_vars_on_returned_row(updated_mcp_server)
@ -1181,7 +1158,7 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient,
)
updated += 1
oauth_clients: Final[list[prisma_db_models.LiteLLM_MCPServerOAuthClient]] = await _oauth_client_table_actions(
oauth_clients: Final[Sequence[prisma_db_models.LiteLLM_MCPServerOAuthClient]] = await _oauth_client_table_actions(
prisma_client
).find_many()
oauth_updated = 0
@ -1914,7 +1891,7 @@ async def get_mcp_submissions(
along with a summary count breakdown by approval_status.
Mirrors get_guardrail_submissions() from guardrail_endpoints.py.
"""
rows: Final[list[prisma_db_models.LiteLLM_MCPServerTable]] = await _mcp_server_table_actions(
rows: Final[Sequence[prisma_db_models.LiteLLM_MCPServerTable]] = await _mcp_server_table_actions(
prisma_client
).find_many(
where={"submitted_at": {"not": None}},

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

@ -4,7 +4,7 @@ import json
from collections.abc import Iterator, Mapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Any, Final, NamedTuple, Protocol, TypedDict
from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypedDict
import litellm
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
@ -12,9 +12,13 @@ from litellm.proxy.management_helpers.object_permission_utils import (
handle_update_object_permission_common,
)
from litellm.proxy.utils import PrismaClient
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import AgentsRepository, ObjectPermissionRepository
from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest
if TYPE_CHECKING:
from prisma import models as prisma_models
class AgentObjectPermissionRecord(Protocol):
def model_dump(self) -> dict[str, object]: ...
@ -42,11 +46,20 @@ class AgentRecordDump(TypedDict):
class AgentRecord(Protocol):
agent_id: str
agent_name: str
object_permission_id: str | None
object_permission: AgentObjectPermissionRecord | None
spend: float
@property
def agent_id(self) -> str: ...
@property
def agent_name(self) -> str: ...
@property
def object_permission_id(self) -> str | None: ...
@property
def object_permission(self) -> AgentObjectPermissionRecord | None: ...
@property
def spend(self) -> float: ...
def model_dump(self) -> AgentRecordDump: ...
@ -57,50 +70,47 @@ class AgentTableClient(Protocol):
async def create(
self,
data: Mapping[str, object],
include: Mapping[str, bool] | None = None,
include: Mapping[str, object] | None = None,
) -> AgentRecord: ...
async def find_unique(
self,
where: Mapping[str, object],
include: Mapping[str, bool] | None = None,
include: Mapping[str, object] | None = None,
) -> AgentRecord | None: ...
async def find_many(
self,
where: Mapping[str, object] | None = None,
order: Mapping[str, str] | None = None,
include: Mapping[str, bool] | None = None,
include: Mapping[str, object] | None = None,
) -> Sequence[AgentRecord]: ...
async def update(
self,
where: Mapping[str, object],
data: Mapping[str, object],
include: Mapping[str, bool] | None = None,
) -> AgentRecord: ...
where: Mapping[str, object],
include: Mapping[str, object] | None = None,
) -> AgentRecord | None: ...
async def delete(self, where: Mapping[str, object]) -> AgentRecord: ...
async def delete(
self,
where: Mapping[str, object],
include: Mapping[str, object] | None = None,
) -> AgentRecord | None: ...
def agents_table(prisma_client: PrismaClient) -> AgentTableClient:
table: Final[AgentTableClient] = AgentsRepository(prisma_client).table
table: Final[AgentTableClient] = AgentsRepository(prisma_client).table # pyright: ignore[reportAssignmentType] # prisma rows type model_dump() as dict[str, Any]
return table
class ObjectPermissionGrantRecord(Protocol):
object_permission_id: str
agents: list[str] | None
class ObjectPermissionTableClient(Protocol):
async def find_many(self, where: Mapping[str, object]) -> Sequence[ObjectPermissionGrantRecord]: ...
async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
def object_permission_table(prisma_client: PrismaClient) -> ObjectPermissionTableClient:
table: Final[ObjectPermissionTableClient] = ObjectPermissionRepository(prisma_client).table
def object_permission_table(
prisma_client: PrismaClient,
) -> "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]":
table: Final[TableActions[prisma_models.LiteLLM_ObjectPermissionTable]] = ObjectPermissionRepository(
prisma_client
).table
return table
@ -222,7 +232,9 @@ class AgentRegistry:
self.load_agents_from_config(agent_config if agent_config is not None else self.config_agents)
return self.agent_list
async def migrate_legacy_grant_ids(self, table: ObjectPermissionTableClient) -> GrantMigrationResult:
async def migrate_legacy_grant_ids(
self, table: "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]"
) -> GrantMigrationResult:
"""
Rewrite object_permission.agents rows holding a legacy full-entry hash to the
stable name-derived id.
@ -360,6 +372,8 @@ class AgentRegistry:
"""
try:
deleted_agent: Final = await agents_table(prisma_client).delete(where={"agent_id": agent_id})
if deleted_agent is None:
raise ValueError(f"Agent not found, passed agent_id={agent_id}")
return dict(deleted_agent)
except Exception as e:
raise Exception(f"Error deleting agent from DB: {e}")
@ -386,12 +400,12 @@ 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_row: Final = await AgentsRepository(prisma_client).table.find_unique(
where={"agent_id": agent_id} # mutable-ok: prisma filters are plain dicts
)
if existing_row is None:
raise Exception(f"Agent with ID {agent_id} not found")
existing_agent: Final = dict(existing_row)
augment_agent: Final = {**existing_agent, **agent}
update_data: Final[dict[str, Any]] = {}
@ -436,6 +450,8 @@ class AgentRegistry:
},
include={"object_permission": True},
)
if patched_agent is None:
raise ValueError(f"Agent not found, passed agent_id={agent_id}")
patched_agent_dict: Final = patched_agent.model_dump()
if patched_agent.object_permission is not None:
try:
@ -523,6 +539,8 @@ class AgentRegistry:
include={"object_permission": True},
)
if updated_agent is None:
raise ValueError(f"Agent not found, passed agent_id={agent_id}")
updated_agent_dict: Final = updated_agent.model_dump()
if updated_agent.object_permission is not None:
try:

4
litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py
View file

@ -543,7 +543,7 @@ async def update_plugin(
manifest: Final[Mapping[str, object]] = _build_plugin_manifest(plugin_name, request)
plugin: Final[_PluginRecord] = await ClaudeCodePluginRepository(prisma_client).table.update(
plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.update(
where={"name": plugin_name}, # mutable-ok: prisma query arguments must be plain dicts
data={ # mutable-ok: prisma query arguments must be plain dicts
"version": request.version,
@ -553,6 +553,8 @@ async def update_plugin(
"updated_at": datetime.now(timezone.utc),
},
)
if plugin is None:
raise _error_response(404, f"Plugin '{plugin_name}' not found")
verbose_proxy_logger.info("Plugin %s updated successfully", plugin_name)

18
litellm/proxy/auth/auth_checks.py
View file

@ -156,7 +156,12 @@ class _PrismaVectorStoreRow(Protocol):
class _PrismaUserRow(Protocol):
user_id: str
organization_memberships: Sequence[LiteLLM_OrganizationMembershipTable | None] | None
@property
def organization_memberships(self) -> Sequence[_PrismaModelDumpRow | None] | None: ...
@organization_memberships.setter
def organization_memberships(self, value: Sequence[_PrismaModelDumpRow] | None) -> None: ...
def __iter__(self) -> Iterator[tuple[str, object]]: ...
@ -214,9 +219,14 @@ def _user_table(repo: _PrismaTableHolder[_PrismaUserRow]) -> _PrismaAuthTable[_P
return repo.table
class _VectorStorePermissionsRow(Protocol):
@property
def vector_stores(self) -> Sequence[str] | None: ...
def _object_permission_table(
repo: _PrismaTableHolder[LiteLLM_ObjectPermissionTable],
) -> _PrismaAuthTable[LiteLLM_ObjectPermissionTable]:
repo: _PrismaTableHolder[_VectorStorePermissionsRow],
) -> _PrismaAuthTable[_VectorStorePermissionsRow]:
return repo.table
@ -5277,7 +5287,7 @@ async def vector_store_access_check(
def _can_object_call_vector_stores(
object_type: Literal["key", "team", "org"],
vector_store_ids_to_run: list[str],
object_permissions: LiteLLM_ObjectPermissionTable | None,
object_permissions: _VectorStorePermissionsRow | None,
):
"""
Raises ProxyException if the object (key, team, org) cannot access the specific vector store.

17
litellm/proxy/auth/user_api_key_auth.py
View file

@ -196,6 +196,15 @@ class _UserModelBudgetLimiter(Protocol):
) -> bool: ...
class _TokenTeamModels(Protocol):
@property
def team_models(self) -> list[str]: ...
def _token_team_models(valid_token: _TokenTeamModels) -> list[str]:
return valid_token.team_models
async def _read_user_model_max_budget(
user_id: str | None,
prisma_client: PrismaClient | None,
@ -1991,7 +2000,7 @@ async def _user_api_key_auth_builder(
include={"litellm_budget_table": True},
)
if _db_member is not None:
team_member_info = LiteLLM_TeamMembership(**_db_member.dict())
team_member_info = LiteLLM_TeamMembership(**_db_member.model_dump())
await user_api_key_cache.async_set_cache(
key=_cache_key,
value=team_member_info,
@ -2143,6 +2152,7 @@ async def _user_api_key_auth_builder(
proxy_logging_obj=proxy_logging_obj,
)
except HTTPException:
token_team_models: Final = _token_team_models(valid_token)
_team_obj = LiteLLM_TeamTableCachedObj(
team_id=valid_token.team_id,
max_budget=valid_token.team_max_budget,
@ -2151,7 +2161,7 @@ async def _user_api_key_auth_builder(
tpm_limit=valid_token.team_tpm_limit,
rpm_limit=valid_token.team_rpm_limit,
blocked=valid_token.team_blocked,
models=valid_token.team_models,
models=token_team_models,
metadata=valid_token.team_metadata,
object_permission_id=valid_token.team_object_permission_id,
object_permission=await _resolve_object_permission_for_unresolvable_team(
@ -2295,6 +2305,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached
UserAPIKeyAuth. Only called when valid_token.team_id is known to be
non-None (the caller gates on it)."""
assert valid_token.team_id is not None
token_team_models: Final = _token_team_models(valid_token)
return LiteLLM_TeamTableCachedObj(
team_id=valid_token.team_id,
max_budget=valid_token.team_max_budget,
@ -2303,7 +2314,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached
tpm_limit=valid_token.team_tpm_limit,
rpm_limit=valid_token.team_rpm_limit,
blocked=valid_token.team_blocked,
models=valid_token.team_models,
models=token_team_models,
metadata=valid_token.team_metadata,
object_permission_id=valid_token.team_object_permission_id,
)

8
litellm/proxy/common_utils/config_sync_pubsub.py
View file

@ -7,6 +7,7 @@ from dataclasses import asdict, dataclass
from typing import TYPE_CHECKING, Final, Protocol, cast # noqa: TID251 # untyped prisma/redis boundary needs cast
from litellm._logging import verbose_proxy_logger
from litellm.repositories.prisma_protocols import RowT_co, TableActions
if TYPE_CHECKING:
from litellm.caching.redis_cache import RedisCache
@ -163,13 +164,14 @@ class _PublishOnWriteActions:
def wrap_table_actions_for_config_sync(
actions: object,
actions: "TableActions[RowT_co]",
table_name: str,
publish: Callable[[str], Awaitable[None]] = publish_config_change_for_object_type,
) -> object:
) -> "TableActions[RowT_co]":
if table_name not in _CONFIG_SYNCED_TABLE_NAMES:
return actions
return _PublishOnWriteActions(actions=actions, object_type=table_name, publish=publish)
wrapped: Final = _PublishOnWriteActions(actions=actions, object_type=table_name, publish=publish)
return cast("TableActions[RowT_co]", wrapped) # cast-ok: dynamic write-through proxy keeps the wrapped row type
class ConfigSyncSubscriber:

12
litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py
View file

@ -4,8 +4,9 @@ Expired UI session key cleanup manager.
Deletes expired virtual keys created for LiteLLM dashboard sessions.
"""
from collections.abc import Sequence
from datetime import datetime, timezone
from typing import Any, Final
from typing import Any, Final, Protocol
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
@ -14,7 +15,7 @@ from litellm.constants import (
LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
UI_SESSION_TOKEN_TEAM_ID,
)
from litellm.proxy._types import KeyRequest, LiteLLM_VerificationToken, UserAPIKeyAuth
from litellm.proxy._types import KeyRequest, UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
from litellm.proxy.management_endpoints.key_management_endpoints import (
@ -26,6 +27,11 @@ from litellm.repositories.verification_token_repository import (
)
class _ExpiredSessionKeyRow(Protocol):
@property
def token(self) -> str | None: ...
class ExpiredUISessionKeyCleanupManager:
"""
Cleans up expired UI session keys.
@ -138,7 +144,7 @@ class ExpiredUISessionKeyCleanupManager:
return len(tokens)
async def _find_expired_ui_session_keys(self) -> list[LiteLLM_VerificationToken]:
async def _find_expired_ui_session_keys(self) -> Sequence[_ExpiredSessionKeyRow]:
"""
Find expired LiteLLM dashboard session keys.
"""

12
litellm/proxy/common_utils/key_rotation_manager.py
View file

@ -4,8 +4,9 @@ Key Rotation Manager - Automated key rotation based on rotation schedules
Handles finding keys that need rotation based on their individual schedules.
"""
from collections.abc import Sequence
from datetime import datetime, timezone
from typing import Final
from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
@ -31,6 +32,9 @@ from litellm.repositories.verification_token_repository import (
VerificationTokenRepository,
)
if TYPE_CHECKING:
from prisma import models as prisma_models
class KeyRotationManager:
"""
@ -106,7 +110,7 @@ class KeyRotationManager:
cronjob_id=KEY_ROTATION_JOB_NAME,
)
async def _find_keys_needing_rotation(self) -> list[LiteLLM_VerificationToken]:
async def _find_keys_needing_rotation(self) -> "Sequence[prisma_models.LiteLLM_VerificationToken]":
"""
Find keys that are due for rotation based on their key_rotation_at timestamp.
@ -156,7 +160,7 @@ class KeyRotationManager:
# Check if the rotation time has passed
return now >= key.key_rotation_at
async def _rotate_key(self, key: LiteLLM_VerificationToken):
async def _rotate_key(self, key: "prisma_models.LiteLLM_VerificationToken"):
"""
Rotate a single key using existing regenerate_key_fn and call the rotation hook
"""
@ -197,7 +201,7 @@ class KeyRotationManager:
if isinstance(response, GenerateKeyResponse):
await KeyManagementEventHooks.async_key_rotated_hook(
data=regenerate_request,
existing_key_row=key,
existing_key_row=key, # pyright: ignore[reportArgumentType] # prisma row, hook wants the domain model
response=response,
user_api_key_dict=system_user,
litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,

6
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -37,7 +37,7 @@ from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManage
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import ReadOnlyTable, SpendLinkedTable
from litellm.repositories.prisma_protocols import SpendLinkedTable
from litellm.repositories.table_repositories import (
EndUserRepository,
TagRepository,
@ -675,7 +675,7 @@ class ResetBudgetJob:
rely on the default budget (litellm.max_end_user_budget_id) applied
in-memory during auth checks.
"""
table: Final[ReadOnlyTable] = EndUserRepository(self.prisma_client).table
table: Final = EndUserRepository(self.prisma_client).table
rows: Final = await self._with_db_retry(
lambda: table.find_many(
where={
@ -685,7 +685,7 @@ class ResetBudgetJob:
),
reason="reset_budget_read_endusers_without_budget_id_failure",
)
return [LiteLLM_EndUserTable.model_validate(row.dict()) for row in rows]
return [LiteLLM_EndUserTable.model_validate(row.model_dump()) for row in rows]
async def _write_key_reset_updates(self, updated_keys: list[LiteLLM_VerificationToken]) -> None:
"""

37
litellm/proxy/container_endpoints/ownership.py
View file

@ -1,7 +1,7 @@
import json
from collections.abc import Mapping, Sequence
from collections.abc import Set as AbstractSet
from typing import TYPE_CHECKING, Any, Final, Protocol
from typing import TYPE_CHECKING, Any, Final
from fastapi import HTTPException
@ -18,28 +18,11 @@ from litellm.repositories.table_repositories import ManagedObjectRepository
from litellm.responses.utils import ResponsesAPIRequestUtils
if TYPE_CHECKING:
from prisma import models as prisma_models
from litellm.proxy.utils import PrismaClient
class _ManagedObjectRow(Protocol):
model_object_id: str
unified_object_id: str | None
file_purpose: str | None
created_by: str | None
class _ManagedObjectTable(Protocol):
async def find_unique(self, *, where: Mapping[str, str]) -> _ManagedObjectRow | None: ...
async def find_first(self, *, where: Mapping[str, str]) -> _ManagedObjectRow | None: ...
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_ManagedObjectRow]: ...
async def create(self, *, data: Mapping[str, str]) -> _ManagedObjectRow: ...
async def update(self, *, where: Mapping[str, str], data: Mapping[str, str]) -> _ManagedObjectRow | None: ...
CONTAINER_OBJECT_PURPOSE: Final = "container"
# 60s LRU/TTL cache absorbs every container access check before it reaches
@ -220,7 +203,7 @@ async def record_container_owner(
verbose_proxy_logger.warning("Skipping container ownership tracking because prisma_client is None")
return response
table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table
table: Final = ManagedObjectRepository(prisma_client).table
existing: Final = await table.find_unique(where={"model_object_id": model_object_id})
if existing is not None:
if getattr(existing, "file_purpose", None) != CONTAINER_OBJECT_PURPOSE:
@ -272,8 +255,8 @@ async def _get_container_owner(original_container_id: str, custom_llm_provider:
if prisma_client is None:
return None
table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table
row: Final[_ManagedObjectRow | None] = await table.find_first(
table: Final = ManagedObjectRepository(prisma_client).table
row: Final[prisma_models.LiteLLM_ManagedObjectTable | None] = await table.find_first(
where={
"model_object_id": model_object_id,
"file_purpose": CONTAINER_OBJECT_PURPOSE,
@ -309,8 +292,8 @@ async def _get_stored_container_id(original_container_id: str, custom_llm_provid
if prisma_client is None:
return None
table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table
row: Final[_ManagedObjectRow | None] = await table.find_first(
table: Final = ManagedObjectRepository(prisma_client).table
row: Final[prisma_models.LiteLLM_ManagedObjectTable | None] = await table.find_first(
where={
"model_object_id": model_object_id,
"file_purpose": CONTAINER_OBJECT_PURPOSE,
@ -394,8 +377,8 @@ async def _get_allowed_container_ids(
if prisma_client is None:
return set()
table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table
rows: Final[Sequence[_ManagedObjectRow]] = await table.find_many(
table: Final = ManagedObjectRepository(prisma_client).table
rows: Final[Sequence[prisma_models.LiteLLM_ManagedObjectTable]] = await table.find_many(
where={
"file_purpose": CONTAINER_OBJECT_PURPOSE,
"created_by": {"in": owner_scopes},

13
litellm/proxy/credential_endpoints/endpoints.py
View file

@ -2,7 +2,10 @@
CRUD endpoints for storing reusable credentials.
"""
from typing import Final
from typing import (
Final,
cast, # noqa: TID251 # jsonify_object in proxy/utils.py is annotated with a bare dict
)
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
@ -88,7 +91,9 @@ async def create_credential(
)
encrypted_credential: Final = CredentialHelperUtils.encrypt_credential_values(processed_credential)
credentials_dict: Final = encrypted_credential.model_dump()
credentials_dict_jsonified: Final = jsonify_object(credentials_dict)
credentials_dict_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str
"dict[str, object]", jsonify_object(credentials_dict)
)
await CredentialsRepository(prisma_client).create(
data={
**credentials_dict_jsonified,
@ -310,7 +315,9 @@ async def update_credential(
if db_credential is None:
raise HTTPException(status_code=404, detail="Credential not found in DB.")
merged_credential: Final = update_db_credential(db_credential, credential)
credential_object_jsonified: Final = jsonify_object(merged_credential.model_dump())
credential_object_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str
"dict[str, object]", jsonify_object(merged_credential.model_dump())
)
await credentials_repository.update_by_name(
credential_name,
data={

30
litellm/proxy/db/tool_registry_writer.py
View file

@ -6,14 +6,15 @@ Admins use the management endpoints to read and update input_policy / output_pol
"""
import uuid
from collections.abc import Mapping, Sequence
from collections.abc import Mapping
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar
from typing import TYPE_CHECKING, Any, Final
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import ToolDiscoveryQueueItem
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import ToolRepository
from litellm.types.tool_management import (
LiteLLM_ToolTableRow,
@ -25,33 +26,16 @@ if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
_RowT_co: Final = TypeVar("_RowT_co", covariant=True)
class _TableActions(Protocol[_RowT_co]):
async def find_unique(self, where: Mapping[str, object]) -> _RowT_co | None: ...
async def find_many(
self,
where: Mapping[str, object] | None = None,
order: Mapping[str, object] | None = None,
include: Mapping[str, object] | None = None,
) -> Sequence[_RowT_co]: ...
async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> _RowT_co: ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _RowT_co | None: ...
def _tool_table_actions(prisma_client: "PrismaClient") -> "_TableActions[prisma_db_models.LiteLLM_ToolTable]":
table: Final[_TableActions[prisma_db_models.LiteLLM_ToolTable]] = ToolRepository(prisma_client).table
def _tool_table_actions(prisma_client: "PrismaClient") -> "TableActions[prisma_db_models.LiteLLM_ToolTable]":
table: Final[TableActions[prisma_db_models.LiteLLM_ToolTable]] = ToolRepository(prisma_client).table
return table
def _object_permission_table_actions(
prisma_client: "PrismaClient",
) -> "_TableActions[prisma_db_models.LiteLLM_ObjectPermissionTable]":
table: Final[_TableActions[prisma_db_models.LiteLLM_ObjectPermissionTable]] = ObjectPermissionRepository(
) -> "TableActions[prisma_db_models.LiteLLM_ObjectPermissionTable]":
table: Final[TableActions[prisma_db_models.LiteLLM_ObjectPermissionTable]] = ObjectPermissionRepository(
prisma_client
).table
return table

22
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -29,6 +29,7 @@ from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import (
from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry
from litellm.proxy.guardrails.usage_endpoints import router as guardrails_usage_router
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import GuardrailsRepository
from litellm.types.guardrails import (
PII_ENTITY_CATEGORIES_MAP,
@ -65,29 +66,12 @@ router: Final = APIRouter()
GUARDRAIL_REGISTRY: Final = GuardrailRegistry()
class _GuardrailsTableActions(Protocol):
async def create(self, data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable": ...
async def delete(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ...
async def find_unique(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ...
async def find_many(
self, where: Mapping[str, object], order: Mapping[str, str]
) -> "Sequence[LiteLLM_GuardrailsTable]": ...
async def update(
self, where: Mapping[str, object], data: Mapping[str, object]
) -> "LiteLLM_GuardrailsTable | None": ...
def _as_str_object_mapping(mapping: Mapping[str, object]) -> Mapping[str, object]:
return mapping
def _guardrails_table(prisma_client: "PrismaClient") -> _GuardrailsTableActions:
table: Final[_GuardrailsTableActions] = GuardrailsRepository(prisma_client).table
return table
def _guardrails_table(prisma_client: "PrismaClient") -> "TableActions[LiteLLM_GuardrailsTable]":
return GuardrailsRepository(prisma_client).table
async def _create_guardrail_row(prisma_client: "PrismaClient", data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable":

24
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -3,12 +3,12 @@
import asyncio
import importlib
import os
from collections.abc import Callable, Iterator, Mapping, Sequence
from collections.abc import Callable, Iterator, Mapping
from datetime import datetime, timezone
from itertools import chain, count
from typing import Final, Literal, Optional, Protocol, cast
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, cast
from pydantic import BaseModel, ValidationError
from pydantic import ValidationError
import litellm
from litellm import Router
@ -39,6 +39,7 @@ from litellm.proxy.guardrails.guardrail_hooks.tool_permission import (
)
from litellm.proxy.types_utils.utils import get_instance_fn
from litellm.proxy.utils import PrismaClient
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import GuardrailsRepository
from litellm.secret_managers.main import get_secret
from litellm.types.guardrails import (
@ -61,6 +62,9 @@ from .guardrail_initializers import (
initialize_tool_permission,
)
if TYPE_CHECKING:
from prisma import models as prisma_models
class _GuardrailRowLike(Protocol):
@property
@ -68,15 +72,7 @@ class _GuardrailRowLike(Protocol):
def __iter__(self) -> Iterator[tuple[str, object]]: ...
class _GuardrailTableActions(Protocol):
async def create(self, *, data: Mapping[str, object]) -> _GuardrailRowLike: ...
async def delete(self, *, where: Mapping[str, str]) -> object: ...
async def update(self, *, where: Mapping[str, str], data: Mapping[str, object]) -> _GuardrailRowLike: ...
async def find_many(self, *, where: Mapping[str, str], order: Mapping[str, str]) -> Sequence[BaseModel]: ...
async def find_unique(self, *, where: Mapping[str, str]) -> BaseModel | None: ...
def _guardrail_table(prisma_client: PrismaClient) -> _GuardrailTableActions:
def _guardrail_table(prisma_client: PrismaClient) -> "TableActions[prisma_models.LiteLLM_GuardrailsTable]":
"""Typed view of the guardrails table actions exposed by the Prisma repository."""
return GuardrailsRepository(prisma_client).table
@ -347,7 +343,7 @@ class GuardrailRegistry:
guardrail_info: Final[str] = safe_dumps(guardrail.get("guardrail_info", {}))
# Update in DB
updated_guardrail: Final[_GuardrailRowLike] = await _guardrail_table(prisma_client).update(
updated_guardrail: Final[_GuardrailRowLike | None] = await _guardrail_table(prisma_client).update(
where={"guardrail_id": guardrail_id},
data={
"guardrail_name": guardrail_name,
@ -356,6 +352,8 @@ class GuardrailRegistry:
"updated_at": datetime.now(timezone.utc),
},
)
if updated_guardrail is None:
raise ValueError(f"Guardrail not found, passed guardrail_id={guardrail_id}")
# Convert to dict and return
return dict(updated_guardrail)

38
litellm/proxy/guardrails/usage_endpoints.py
View file

@ -17,6 +17,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
DailyGuardrailMetricsRepository,
DailyGuardrailUsageUnitsRepository,
@ -30,13 +31,6 @@ from litellm.repositories.table_repositories import (
if TYPE_CHECKING:
from prisma import models as prisma_models
from prisma import types as prisma_types
from prisma.actions import (
LiteLLM_DailyGuardrailMetricsActions,
LiteLLM_DailyGuardrailUsageUnitsActions,
LiteLLM_DailyPolicyMetricsActions,
LiteLLM_GuardrailsTableActions,
LiteLLM_PolicyTableActions,
)
from litellm.proxy.utils import PrismaClient
from litellm.types.guardrails import Guardrail
@ -85,8 +79,8 @@ def _resolve_usage_window(start_date: str | None, end_date: str | None) -> tuple
def _guardrails_table(
prisma_client: "PrismaClient",
) -> "LiteLLM_GuardrailsTableActions[prisma_models.LiteLLM_GuardrailsTable]":
guardrails_table: LiteLLM_GuardrailsTableActions[prisma_models.LiteLLM_GuardrailsTable] = GuardrailsRepository(
) -> "TableActions[prisma_models.LiteLLM_GuardrailsTable]":
guardrails_table: Final[TableActions[prisma_models.LiteLLM_GuardrailsTable]] = GuardrailsRepository(
prisma_client
).table
return guardrails_table
@ -94,28 +88,26 @@ def _guardrails_table(
def _policies_table(
prisma_client: "PrismaClient",
) -> "LiteLLM_PolicyTableActions[prisma_models.LiteLLM_PolicyTable]":
policies_table: Final[LiteLLM_PolicyTableActions[prisma_models.LiteLLM_PolicyTable]] = PolicyRepository(
prisma_client
).table
) -> "TableActions[prisma_models.LiteLLM_PolicyTable]":
policies_table: Final[TableActions[prisma_models.LiteLLM_PolicyTable]] = PolicyRepository(prisma_client).table
return policies_table
def _daily_guardrail_metrics_table(
prisma_client: "PrismaClient",
) -> "LiteLLM_DailyGuardrailMetricsActions[prisma_models.LiteLLM_DailyGuardrailMetrics]":
metrics_table: Final[LiteLLM_DailyGuardrailMetricsActions[prisma_models.LiteLLM_DailyGuardrailMetrics]] = (
DailyGuardrailMetricsRepository(prisma_client).table
)
) -> "TableActions[prisma_models.LiteLLM_DailyGuardrailMetrics]":
metrics_table: Final[TableActions[prisma_models.LiteLLM_DailyGuardrailMetrics]] = DailyGuardrailMetricsRepository(
prisma_client
).table
return metrics_table
def _daily_policy_metrics_table(
prisma_client: "PrismaClient",
) -> "LiteLLM_DailyPolicyMetricsActions[prisma_models.LiteLLM_DailyPolicyMetrics]":
metrics_table: Final[LiteLLM_DailyPolicyMetricsActions[prisma_models.LiteLLM_DailyPolicyMetrics]] = (
DailyPolicyMetricsRepository(prisma_client).table
)
) -> "TableActions[prisma_models.LiteLLM_DailyPolicyMetrics]":
metrics_table: Final[TableActions[prisma_models.LiteLLM_DailyPolicyMetrics]] = DailyPolicyMetricsRepository(
prisma_client
).table
return metrics_table
@ -135,8 +127,8 @@ async def _find_daily_policy_metrics(
def _daily_guardrail_usage_units_table(
prisma_client: "PrismaClient",
) -> "LiteLLM_DailyGuardrailUsageUnitsActions[prisma_models.LiteLLM_DailyGuardrailUsageUnits]":
units_table: Final[LiteLLM_DailyGuardrailUsageUnitsActions[prisma_models.LiteLLM_DailyGuardrailUsageUnits]] = (
) -> "TableActions[prisma_models.LiteLLM_DailyGuardrailUsageUnits]":
units_table: Final[TableActions[prisma_models.LiteLLM_DailyGuardrailUsageUnits]] = (
DailyGuardrailUsageUnitsRepository(prisma_client).table
)
return units_table

16
litellm/proxy/guardrails/usage_tracking.py
View file

@ -14,6 +14,8 @@ from operator import itemgetter
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeVar
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES
from litellm.proxy.utils import PrismaClient
@ -47,6 +49,18 @@ class _MetricsKey(NamedTuple):
date: str
class _UsageUnitCompoundKey(TypedDict):
guardrail_id: ReadOnly[str]
date: ReadOnly[str]
team_id: ReadOnly[str]
api_key: ReadOnly[str]
usage_unit: ReadOnly[str]
class _UsageUnitWhereUnique(TypedDict):
guardrail_id_date_team_id_api_key_usage_unit: ReadOnly[_UsageUnitCompoundKey]
class PendingRollups:
"""Rollup rows whose connection-error retries exhausted, held for the next flush."""
@ -229,7 +243,7 @@ async def _upsert_usage_unit_row(prisma_client: PrismaClient, key: _UsageUnitKey
"usage_unit": key.usage_unit,
"units": units,
}
where: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsWhereUniqueInput] = {
where: Final[_UsageUnitWhereUnique] = {
"guardrail_id_date_team_id_api_key_usage_unit": {
"guardrail_id": key.guardrail_id,
"date": key.date,

4
litellm/proxy/management_endpoints/access_group_endpoints.py
View file

@ -390,7 +390,7 @@ async def list_access_groups(
_require_admin_view(user_api_key_dict)
prisma_client: Final = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value)
table: Final[_AccessGroupTable] = AccessGroupRepository(prisma_client).table
table: Final = AccessGroupRepository(prisma_client).table
records: Final = await table.find_many(order={"created_at": "desc"})
return [_record_to_response(r) for r in records]
@ -406,7 +406,7 @@ async def get_access_group(
_require_admin_view(user_api_key_dict)
prisma_client: Final = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value)
table: Final[_AccessGroupTable] = AccessGroupRepository(prisma_client).table
table: Final = AccessGroupRepository(prisma_client).table
record: Final = await table.find_unique(where={"access_group_id": access_group_id})
if record is None:
raise HTTPException(

2
litellm/proxy/management_endpoints/budget_management_endpoints.py
View file

@ -93,7 +93,7 @@ async def new_budget(
budget_obj.budget_reset_at = get_budget_reset_time(budget_duration=budget_obj.budget_duration)
budget_obj_json: Final = budget_obj.model_dump(exclude_none=True)
budget_obj_jsonified: Final = jsonify_object(budget_obj_json) # json dump any dictionaries
budget_obj_jsonified: Final[dict[str, object]] = jsonify_object(budget_obj_json) # mutable-ok: prisma create input
try:
response: Final = await BudgetRepository(prisma_client).table.create(
data={

3
litellm/proxy/management_endpoints/cache_settings_endpoints.py
View file

@ -43,7 +43,8 @@ router: Final = APIRouter()
class _CacheConfigRow(Protocol):
cache_settings: str | Mapping[str, object] | None
@property
def cache_settings(self) -> str | Mapping[str, object] | None: ...
class _CacheConfigTable(Protocol):

8
litellm/proxy/management_endpoints/common_daily_activity.py
View file

@ -441,7 +441,7 @@ async def get_api_key_metadata(
This ensures that key_alias and team_id are preserved in historical activity logs
even after a key is deleted or regenerated.
"""
key_records: list[PrismaVerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many(
key_records: Sequence[PrismaVerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many(
where={"token": {"in": list(api_keys)}}
)
result: Final[dict[str, _KeyMetadataDict]] = {
@ -452,9 +452,9 @@ async def get_api_key_metadata(
missing_keys: Final = api_keys - set(result.keys())
if missing_keys:
try:
deleted_key_records: Final[list[PrismaDeletedVerificationToken]] = await DeletedVerificationTokenRepository(
prisma_client
).table.find_many(
deleted_key_records: Final[
Sequence[PrismaDeletedVerificationToken]
] = await DeletedVerificationTokenRepository(prisma_client).table.find_many(
where={"token": {"in": list(missing_keys)}},
order={"deleted_at": "desc"},
)

3
litellm/proxy/management_endpoints/config_override_endpoints.py
View file

@ -46,7 +46,8 @@ router: Final = APIRouter()
class _ConfigOverrideRow(Protocol):
config_value: str | Mapping[str, object] | None
@property
def config_value(self) -> str | Mapping[str, object] | None: ...
class _ConfigOverridesTableClient(Protocol):

121
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -15,9 +15,9 @@ These are members of a Team on LiteLLM
import asyncio
import json
import traceback
from collections.abc import Mapping, Sequence
from collections.abc import Awaitable, Mapping, Sequence
from datetime import datetime, timezone
from typing import Any, Final, Literal, Protocol, cast
from typing import Any, Final, Literal, cast
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
@ -58,6 +58,7 @@ from litellm.proxy.management_helpers.object_permission_utils import (
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
from litellm.proxy.utils import handle_exception_on_proxy, hash_password
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
InvitationLinkRepository,
OrganizationMembershipRepository,
@ -86,15 +87,6 @@ from litellm.types.proxy.management_endpoints.scim_v2 import (
if TYPE_CHECKING:
from prisma import models as prisma_models
from prisma import types as prisma_types
from prisma.actions import (
LiteLLM_InvitationLinkActions,
LiteLLM_OrganizationMembershipActions,
LiteLLM_OrganizationTableActions,
LiteLLM_TeamMembershipActions,
LiteLLM_TeamTableActions,
LiteLLM_UserTableActions,
LiteLLM_VerificationTokenActions,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.proxy_server import PrismaClient
@ -105,31 +97,31 @@ router: Final = APIRouter()
def _user_table(
prisma_client: "PrismaClient | None",
) -> "LiteLLM_UserTableActions[prisma_models.LiteLLM_UserTable]":
user_table: Final[LiteLLM_UserTableActions[prisma_models.LiteLLM_UserTable]] = UserRepository(prisma_client).table
) -> "TableActions[prisma_models.LiteLLM_UserTable]":
user_table: Final[TableActions[prisma_models.LiteLLM_UserTable]] = UserRepository(prisma_client).table
return user_table
def _team_table(
prisma_client: "PrismaClient | None",
) -> "LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]":
team_table: Final[LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]] = TeamRepository(prisma_client).table
) -> "TableActions[prisma_models.LiteLLM_TeamTable]":
team_table: Final[TableActions[prisma_models.LiteLLM_TeamTable]] = TeamRepository(prisma_client).table
return team_table
def _verification_token_table(
prisma_client: "PrismaClient | None",
) -> "LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]":
token_table: Final[LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]] = (
VerificationTokenRepository(prisma_client).table
)
) -> "TableActions[prisma_models.LiteLLM_VerificationToken]":
token_table: Final[TableActions[prisma_models.LiteLLM_VerificationToken]] = VerificationTokenRepository(
prisma_client
).table
return token_table
def _organization_membership_table(
prisma_client: "PrismaClient | None",
) -> "LiteLLM_OrganizationMembershipActions[prisma_models.LiteLLM_OrganizationMembership]":
membership_table: Final[LiteLLM_OrganizationMembershipActions[prisma_models.LiteLLM_OrganizationMembership]] = (
) -> "TableActions[prisma_models.LiteLLM_OrganizationMembership]":
membership_table: Final[TableActions[prisma_models.LiteLLM_OrganizationMembership]] = (
OrganizationMembershipRepository(prisma_client).table
)
return membership_table
@ -137,8 +129,8 @@ def _organization_membership_table(
def _invitation_link_table(
prisma_client: "PrismaClient | None",
) -> "LiteLLM_InvitationLinkActions[prisma_models.LiteLLM_InvitationLink]":
invitation_table: LiteLLM_InvitationLinkActions[prisma_models.LiteLLM_InvitationLink] = InvitationLinkRepository(
) -> "TableActions[prisma_models.LiteLLM_InvitationLink]":
invitation_table: Final[TableActions[prisma_models.LiteLLM_InvitationLink]] = InvitationLinkRepository(
prisma_client
).table
return invitation_table
@ -146,19 +138,19 @@ def _invitation_link_table(
def _organization_table(
prisma_client: "PrismaClient | None",
) -> "LiteLLM_OrganizationTableActions[prisma_models.LiteLLM_OrganizationTable]":
organization_table: Final[LiteLLM_OrganizationTableActions[prisma_models.LiteLLM_OrganizationTable]] = (
OrganizationRepository(prisma_client).table
)
) -> "TableActions[prisma_models.LiteLLM_OrganizationTable]":
organization_table: Final[TableActions[prisma_models.LiteLLM_OrganizationTable]] = OrganizationRepository(
prisma_client
).table
return organization_table
def _team_membership_table(
prisma_client: "PrismaClient | None",
) -> "LiteLLM_TeamMembershipActions[prisma_models.LiteLLM_TeamMembership]":
team_membership_table: Final[LiteLLM_TeamMembershipActions[prisma_models.LiteLLM_TeamMembership]] = (
TeamMembershipRepository(prisma_client).table
)
) -> "TableActions[prisma_models.LiteLLM_TeamMembership]":
team_membership_table: Final[TableActions[prisma_models.LiteLLM_TeamMembership]] = TeamMembershipRepository(
prisma_client
).table
return team_membership_table
@ -294,7 +286,7 @@ async def _add_user_to_organizations(
organization_member_add,
)
tasks: Final = []
tasks: Final[list[Awaitable[object]]] = []
for organization_id in organizations:
tasks.append(
organization_member_add(
@ -406,7 +398,7 @@ async def add_new_user_to_default_team(
teams: list[str] | list[NewUserRequestTeam],
prisma_client: "PrismaClient",
):
tasks: Final = []
tasks: Final[list[Awaitable[object]]] = []
for team in teams:
user_role: Literal["user", "admin"] = "user"
max_budget_in_team: float | None = None
@ -1479,7 +1471,8 @@ async def _update_single_user_helper(
# Create new user if not found
non_default_values["user_id"] = str(uuid.uuid4())
non_default_values["user_email"] = user_request.user_email
response = await prisma_client.insert_data(data=non_default_values, table_name="user")
inserted_user_row: Final = await prisma_client.insert_data(data=non_default_values, table_name="user")
response = inserted_user_row # pyright: ignore[reportAssignmentType] # insert_data returns a prisma row
if response is not None:
await _schedule_user_update_audit_log(
@ -1795,7 +1788,9 @@ async def bulk_user_update(
# Apply update transformations (reuse existing logic)
data_json: Final[dict] = data.user_updates.model_dump(exclude_unset=True)
non_default_values: Final = _update_internal_user_params(data_json=data_json, data=data.user_updates)
non_default_values: Final[dict[str, object]] = _update_internal_user_params(
data_json=data_json, data=data.user_updates
)
# Remove user identification fields since we're updating by user_id
non_default_values.pop("user_id", None)
@ -2149,7 +2144,7 @@ async def get_users(
_validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None
)
users: Sequence[prisma_models.LiteLLM_UserTable] | None = await UserRepository(prisma_client).table.find_many(
users: Final[Sequence[prisma_models.LiteLLM_UserTable]] = await UserRepository(prisma_client).table.find_many(
where=where_conditions,
skip=skip,
take=page_size,
@ -2160,10 +2155,7 @@ async def get_users(
total_count: Final[int] = await UserRepository(prisma_client).table.count(where=where_conditions)
# Get key count for each user
if users is not None:
user_key_counts = await get_user_key_counts(prisma_client, [user.user_id for user in users])
else:
user_key_counts = {}
user_key_counts: Final = await get_user_key_counts(prisma_client, [user.user_id for user in users])
verbose_proxy_logger.debug("Total count of users: %s", total_count)
@ -2172,17 +2164,14 @@ async def get_users(
# Prepare response
user_list: list[LiteLLM_UserTableWithKeyCount] = []
if users is not None:
for user in users:
user_dump = user.model_dump()
user_dump["metadata"] = _redact_scim_enterprise_metadata(user_dump.get("metadata"))
user_list.append(
LiteLLM_UserTableWithKeyCount.model_validate(
{**user_dump, "key_count": user_key_counts.get(user.user_id, 0)}
)
for user in users:
user_dump = user.model_dump()
user_dump["metadata"] = _redact_scim_enterprise_metadata(user_dump.get("metadata"))
user_list.append(
LiteLLM_UserTableWithKeyCount.model_validate(
{**user_dump, "key_count": user_key_counts.get(user.user_id, 0)}
)
else:
user_list = []
)
return {
"users": user_list,
@ -2193,13 +2182,6 @@ async def get_users(
}
class _DeleteTeamRow(Protocol):
team_id: str
members_with_roles: object
def model_dump(self) -> Mapping[str, object]: ...
@router.post(
"/user/delete",
tags=["Internal User management"],
@ -2258,9 +2240,9 @@ async def delete_user(
# loop an org-admin of org-A could delete users in org-B by supplying
# {"user_ids": [victim_in_org_B], "organization_id": "org-A"}.
caller_is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
caller_admin_org_ids: set = set()
caller_admin_org_ids: set[str] = set()
if not caller_is_proxy_admin:
caller_memberships: Final = (
caller_memberships: Final[Sequence[prisma_models.LiteLLM_OrganizationMembership]] = (
await _organization_membership_table(prisma_client).find_many(
where={
"user_id": user_api_key_dict.user_id,
@ -2279,7 +2261,7 @@ async def delete_user(
# Batch-fetch target memberships once before the per-user loop. Avoids
# an N+1 DB call when delete_user is called with a large user_ids list.
target_org_ids_by_user: Final[dict[str, set]] = {}
target_org_ids_by_user: Final[dict[str, set[str]]] = {}
if not caller_is_proxy_admin:
all_target_memberships: Final = await _organization_membership_table(prisma_client).find_many(
where={"user_id": {"in": data.user_ids}}
@ -2319,7 +2301,7 @@ async def delete_user(
# we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes
if is_audit_logging_enabled():
# make an audit log for each team deleted
_user_row = user_row.json(exclude_none=True)
_user_row = user_row.model_dump_json(exclude_none=True)
asyncio.create_task(
create_audit_log_for_update(
@ -2342,10 +2324,10 @@ async def delete_user(
)
## CLEANUP MEMBERS_WITH_ROLES
fetch_all_teams: Sequence[_DeleteTeamRow] = await TeamRepository(prisma_client).table.find_many(
where={"team_id": {"in": user_row.teams}}
)
teams_to_update = []
fetch_all_teams: Sequence[prisma_models.LiteLLM_TeamTable] = await TeamRepository(
prisma_client
).table.find_many(where={"team_id": {"in": user_row.teams}})
teams_to_update: list[tuple[str, str]] = []
for team in fetch_all_teams:
removed_team_members, new_team_members = _cleanup_members_with_roles(
existing_team_row=LiteLLM_TeamTable.model_validate(team.model_dump()),
@ -2357,15 +2339,14 @@ async def delete_user(
)
if removed_team_members:
_db_new_team_members: list[dict] = [m.model_dump() for m in new_team_members]
team.members_with_roles = json.dumps(_db_new_team_members)
teams_to_update.append(team)
teams_to_update.append((team.team_id, json.dumps(_db_new_team_members)))
## update teams
for team in teams_to_update:
for team_id, members_with_roles in teams_to_update:
await TeamRepository(prisma_client).table.update(
where={"team_id": team.team_id},
data={"members_with_roles": team.members_with_roles},
where={"team_id": team_id},
data={"members_with_roles": members_with_roles},
)
# End of Audit logging

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

@ -122,6 +122,9 @@ async def update_jwt_key_mapping(
where={"id": data.id}, data=update_data
)
if updated_mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
# Invalidate new cache key if claim fields changed
cache_key = f"jwt_key_mapping:{updated_mapping.jwt_claim_name}:{updated_mapping.jwt_claim_value}"
await user_api_key_cache.async_delete_cache(cache_key)

199
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -123,6 +123,7 @@ from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.config_repository import ConfigParam, ConfigRepository
from litellm.repositories.credentials_repository import CredentialsRepository
from litellm.repositories.model_repository import ModelRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
DeletedVerificationTokenRepository,
DeprecatedVerificationTokenRepository,
@ -151,65 +152,22 @@ from litellm.types.utils import (
if TYPE_CHECKING:
from prisma import Prisma
from prisma import models as prisma_models
_PrismaRowT = TypeVar("_PrismaRowT")
_RepositoryModelT = TypeVar("_RepositoryModelT", bound=BaseModel)
class _PrismaTableActions(Protocol[_PrismaRowT]):
"""Typed view of the Prisma table actions a repository exposes through its untyped ``table``."""
async def find_unique(
self,
*,
where: Mapping[str, object],
include: Mapping[str, object] | None = None,
) -> _PrismaRowT | None: ...
async def find_first(
self,
*,
where: Mapping[str, object],
include: Mapping[str, object] | None = None,
) -> _PrismaRowT | None: ...
async def find_many(
self,
*,
where: Mapping[str, object] | None = None,
include: Mapping[str, object] | None = None,
order: Mapping[str, object] | None = None,
skip: int | None = None,
take: int | None = None,
) -> list[_PrismaRowT]: ...
async def count(self, *, where: Mapping[str, object] | None = None) -> int: ...
async def create(self, *, data: Mapping[str, object]) -> _PrismaRowT: ...
async def create_many(self, *, data: Sequence[Mapping[str, object]]) -> int: ...
async def delete_many(self, *, where: Mapping[str, object] | None = None) -> int: ...
async def update(
self,
*,
where: Mapping[str, object],
data: Mapping[str, object],
) -> _PrismaRowT | None: ...
async def upsert(
self,
*,
where: Mapping[str, object],
data: Mapping[str, object],
) -> _PrismaRowT: ...
class _UserRowLike(Protocol):
user_id: str | None
user_email: str | None
user_alias: str | None
"""Read-only view of the user columns ``/key/list`` expands keys with."""
@property
def user_id(self) -> str | None: ...
@property
def user_email(self) -> str | None: ...
@property
def user_alias(self) -> str | None: ...
def model_dump(self) -> Mapping[str, object]: ...
@ -217,46 +175,56 @@ class _UserRowLike(Protocol):
class _TxTables(Protocol):
litellm_proxymodeltable: _PrismaTableActions[object]
litellm_proxymodeltable: TableActions[object]
class _TableSource(Protocol[_PrismaRowT]):
"""Repository view that exposes its untyped Prisma ``table`` with a concrete row type."""
class _ConfigTableActions(Protocol):
"""Config table surface this module needs; the shared repository seam exposes no ``update``."""
@property
def table(self) -> _PrismaTableActions[_PrismaRowT]: ...
async def find_many(self) -> Sequence[ConfigParam]: ...
def _table_of(source: _TableSource[_PrismaRowT]) -> _PrismaTableActions[_PrismaRowT]:
return source.table
async def update(
self,
*,
where: Mapping[str, object],
data: Mapping[str, object],
) -> ConfigParam | None: ...
def _prisma_table(
repository: BaseRepository[_RepositoryModelT],
) -> _PrismaTableActions[_RepositoryModelT]:
return _table_of(repository)
) -> TableActions[_RepositoryModelT]:
return cast( # cast-ok: callers read only the field names the prisma row and repository model share
"TableActions[_RepositoryModelT]", repository.table
)
def _deleted_verification_token_table(
prisma_client: PrismaClient,
) -> _PrismaTableActions[LiteLLM_DeletedVerificationToken]:
return _table_of(DeletedVerificationTokenRepository(prisma_client))
) -> "TableActions[prisma_models.LiteLLM_DeletedVerificationToken]":
return DeletedVerificationTokenRepository(prisma_client).table
def _deprecated_verification_token_table(prisma_client: PrismaClient) -> _PrismaTableActions[object]:
return _table_of(DeprecatedVerificationTokenRepository(prisma_client))
def _deprecated_verification_token_table(
prisma_client: PrismaClient,
) -> "TableActions[prisma_models.LiteLLM_DeprecatedVerificationToken]":
return DeprecatedVerificationTokenRepository(prisma_client).table
def _user_table(prisma_client: PrismaClient) -> _PrismaTableActions[_UserRowLike]:
return _table_of(UserRepository(prisma_client))
def _user_table(prisma_client: PrismaClient) -> TableActions[_UserRowLike]:
return UserRepository(prisma_client).table
def _credentials_table(prisma_client: PrismaClient) -> _PrismaTableActions[CredentialItem]:
return _table_of(CredentialsRepository(prisma_client))
def _credentials_table(prisma_client: PrismaClient) -> TableActions[CredentialItem]:
return cast( # cast-ok: the rotation loop reads and rewrites these rows through CredentialItem names only
"TableActions[CredentialItem]", CredentialsRepository(prisma_client).table
)
def _config_table(prisma_client: PrismaClient) -> _PrismaTableActions[ConfigParam]:
return _table_of(ConfigRepository(prisma_client))
def _config_table(prisma_client: PrismaClient) -> _ConfigTableActions:
return cast( # cast-ok: ConfigRepository.table hides the write actions this module needs on that same object
"_ConfigTableActions", ConfigRepository(prisma_client).table
)
async def _check_custom_key_allowed(custom_key_value: str | None) -> None:
@ -1046,7 +1014,7 @@ async def _common_key_generation_helper(
)
new_budget: Final = prisma_client.jsonify_object(budget_row.json(exclude_none=True))
_budget: Final[LiteLLM_BudgetTable] = await BudgetRepository(prisma_client).table.create(
_budget: Final[prisma_models.LiteLLM_BudgetTable] = await BudgetRepository(prisma_client).table.create(
data={
**new_budget,
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
@ -1252,7 +1220,7 @@ async def _common_key_generation_helper(
def _check_key_model_specific_limits(
keys: list[LiteLLM_VerificationToken],
keys: Sequence[LiteLLM_VerificationToken],
data: GenerateKeyRequest | UpdateKeyRequest,
entity_rpm_limit: int | None,
entity_tpm_limit: int | None,
@ -1323,7 +1291,7 @@ def _check_key_model_specific_limits(
def _check_key_rpm_tpm_limits(
keys: list[LiteLLM_VerificationToken],
keys: Sequence[LiteLLM_VerificationToken],
data: GenerateKeyRequest | UpdateKeyRequest,
entity_rpm_limit: int | None,
entity_tpm_limit: int | None,
@ -1361,7 +1329,7 @@ def _check_key_rpm_tpm_limits(
def check_team_key_model_specific_limits(
keys: list[LiteLLM_VerificationToken],
keys: Sequence[LiteLLM_VerificationToken],
team_table: LiteLLM_TeamTableCachedObj,
data: GenerateKeyRequest | UpdateKeyRequest,
) -> None:
@ -1386,7 +1354,7 @@ def check_team_key_model_specific_limits(
def check_team_key_rpm_tpm_limits(
keys: list[LiteLLM_VerificationToken],
keys: Sequence[LiteLLM_VerificationToken],
team_table: LiteLLM_TeamTableCachedObj,
data: GenerateKeyRequest | UpdateKeyRequest,
) -> None:
@ -1494,7 +1462,7 @@ async def _check_project_key_limits(
def check_org_key_model_specific_limits(
keys: list[LiteLLM_VerificationToken],
keys: Sequence[LiteLLM_VerificationToken],
org_table: LiteLLM_OrganizationTable,
data: GenerateKeyRequest | UpdateKeyRequest,
) -> None:
@ -1527,7 +1495,7 @@ def check_org_key_model_specific_limits(
def check_org_key_rpm_tpm_limits(
keys: list[LiteLLM_VerificationToken],
keys: Sequence[LiteLLM_VerificationToken],
org_table: LiteLLM_OrganizationTable,
data: GenerateKeyRequest | UpdateKeyRequest,
) -> None:
@ -2242,9 +2210,9 @@ async def _get_and_validate_existing_key(
code=status.HTTP_400_BAD_REQUEST,
)
rows: list[LiteLLM_VerificationToken] = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
where={"key_alias": key_alias}, take=2
)
rows: Sequence[LiteLLM_VerificationToken] = await _prisma_table(
VerificationTokenRepository(prisma_client)
).find_many(where={"key_alias": key_alias}, take=2)
if len(rows) == 0:
raise ProxyException(
@ -2407,7 +2375,10 @@ async def _process_single_key_update(
)
_data: Final = {**non_default_values, "token": update_key_request.key}
response: Final = await prisma_client.update_data(token=update_key_request.key, data=_data)
response: Final[Mapping[str, object] | None] = cast( # cast-ok: every update_data branch returns a str-keyed dict
"Mapping[str, object] | None",
await prisma_client.update_data(token=update_key_request.key, data=_data),
)
# Delete cache
await _delete_cache_key_object(
@ -3225,7 +3196,7 @@ async def bulk_update_team_keys(
# `blocked` is Boolean? with no default; `/key/generate` writes NULL. Prisma's `NOT`
# excludes NULLs, so explicitly OR `false` with `null` to include them.
now: Final = datetime.now(timezone.utc)
existing_keys = await VerificationTokenRepository(prisma_client).table.find_many(
existing_keys = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
where={
"team_id": data.team_id,
"AND": [
@ -3243,7 +3214,9 @@ async def bulk_update_team_keys(
"error": f"Team {data.team_id} has more than {MAX_BATCH_SIZE} keys. Use `key_ids` to update in batches of {MAX_BATCH_SIZE}."
},
)
requested_tokens = [row.token for row in existing_keys]
requested_tokens = cast( # cast-ok: token is the table's primary key, so a row read back always carries one
"list[str]", [row.token for row in existing_keys]
)
else:
if data.key_ids is None or len(data.key_ids) == 0:
raise HTTPException(
@ -3261,7 +3234,7 @@ async def bulk_update_team_keys(
seen_hashes.add(h)
requested_tokens.append(k)
hashed_key_ids.append(h)
existing_keys = await VerificationTokenRepository(prisma_client).table.find_many(
existing_keys = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
where={"team_id": data.team_id, "token": {"in": hashed_key_ids}}
)
@ -3698,7 +3671,7 @@ async def info_key_fn(
hashed_key: str | None = key
if key is not None:
hashed_key = _hash_token_if_needed(token=key)
key_info = await VerificationTokenRepository(prisma_client).table.find_unique(
key_info = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique(
where={"token": hashed_key},
include={"litellm_budget_table": True},
)
@ -3727,7 +3700,7 @@ async def info_key_fn(
key_info = key_info.model_dump()
except Exception:
# if using pydantic v1
key_info = key_info.dict()
key_info = key_info.dict() # pyright: ignore[reportDeprecated] # deliberate pydantic v1 fallback
key_token_hash: Final = key_info.pop("token")
model_max_budget = key_info.get("model_max_budget") or {}
@ -4012,7 +3985,10 @@ async def generate_key_helper_fn(
if table_name is None or table_name == "user": # do not auto-create users for `/key/generate`
## CREATE USER (If necessary)
if query_type == "insert_data":
user_row = await prisma_client.insert_data(data=user_data, table_name="user")
user_row = cast( # cast-ok: table_name="user" is the insert_data branch returning the user row
"prisma_models.LiteLLM_UserTable | None",
await prisma_client.insert_data(data=user_data, table_name="user"),
)
if user_row is None:
raise Exception("Failed to create user")
@ -4219,9 +4195,12 @@ async def delete_verification_tokens(
if prisma_client:
hashed_tokens: Final[list[str]] = [_hash_token_if_needed(token=key) for key in tokens]
tokens = hashed_tokens
_keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = await _prisma_table(
VerificationTokenRepository(prisma_client)
).find_many(where={"token": {"in": hashed_tokens}})
_keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = cast( # cast-ok: find_many returns a list
"list[LiteLLM_VerificationToken]",
await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
where={"token": {"in": hashed_tokens}}
),
)
if len(_keys_being_deleted) == 0:
raise HTTPException(
@ -4297,7 +4276,7 @@ async def delete_verification_tokens(
def _transform_verification_tokens_to_deleted_records(
keys: list[LiteLLM_VerificationToken],
keys: Sequence[LiteLLM_VerificationToken],
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None = None,
) -> list[dict[str, object]]:
@ -4372,7 +4351,7 @@ async def _save_deleted_verification_token_records(
async def _persist_deleted_verification_tokens(
keys: list[LiteLLM_VerificationToken],
keys: Sequence[LiteLLM_VerificationToken],
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None = None,
@ -4435,7 +4414,9 @@ async def _rotate_master_key(
from litellm.proxy.proxy_server import proxy_config
try:
models: list | None = await _prisma_table(ModelRepository(prisma_client)).find_many()
models: list | None = cast( # cast-ok: find_many returns a real list, which TableActions widens to Sequence
"list[object]", await _prisma_table(ModelRepository(prisma_client)).find_many()
)
except Exception:
models = None
# 2. process model table
@ -5361,9 +5342,9 @@ async def validate_key_list_check(
if key_hash:
try:
key_info: Final[LiteLLM_VerificationToken] = await VerificationTokenRepository(
prisma_client
).table.find_unique(
key_info: Final[LiteLLM_VerificationToken | None] = await _prisma_table(
VerificationTokenRepository(prisma_client)
).find_unique(
where={"token": key_hash},
)
except Exception:
@ -5373,6 +5354,13 @@ async def validate_key_list_check(
param="key_hash",
code=status.HTTP_403_FORBIDDEN,
)
if key_info is None:
raise ProxyException(
message="Key Hash not found.",
type=ProxyErrorTypes.bad_request_error,
param="key_hash",
code=status.HTTP_403_FORBIDDEN,
)
can_user_query_key_info: Final = await _can_user_query_key_info(
user_api_key_dict=user_api_key_dict,
key=key_hash,
@ -5394,8 +5382,9 @@ async def _fetch_user_team_objects(
if complete_user_info is None or not complete_user_info.teams:
return []
teams: Final[list[BaseModel] | None] = await TeamRepository(prisma_client).table.find_many(
where={"team_id": {"in": complete_user_info.teams}}
teams: Final[Sequence[BaseModel] | None] = cast( # cast-ok: the None guard below predates the non-optional seam
"Sequence[BaseModel] | None",
await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": complete_user_info.teams}}),
)
if teams is None:
return []
@ -6130,7 +6119,7 @@ async def _list_key_helper(
key_dict = key.model_dump()
except Exception:
# Fallback for Pydantic v1 compatibility
key_dict = key.dict()
key_dict = key.dict() # pyright: ignore[reportDeprecated] # deliberate pydantic v1 fallback
# Attach object_permission if object_permission_id is set (only for non-deleted keys)
if not use_deleted_table:
key_dict = await attach_object_permission_to_dict(key_dict, prisma_client)
@ -6155,7 +6144,9 @@ async def _list_key_helper(
# Use deleted key type to preserve deleted_at, deleted_by, etc.
key_list.append(LiteLLM_DeletedVerificationToken.model_validate(key_dict))
else:
key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object
key_list.append(
UserAPIKeyAuth(**key_dict) # pyright: ignore[reportAny] # model_dump() is dict[str, Any]
)
else:
_token = key_dict.get("token")
key_list.append(cast(str, _token)) # Return only the token

21
litellm/proxy/management_endpoints/model_access_group_management_endpoints.py
View file

@ -40,17 +40,22 @@ router: Final = APIRouter()
class _DeploymentRow(Protocol):
model_id: str
model_name: str
model_info: object
@property
def model_id(self) -> str: ...
@property
def model_name(self) -> str: ...
@property
def model_info(self) -> object: ...
class _ModelTableClient(Protocol):
async def find_many(self, where: Mapping[str, object] | None = None) -> Sequence[_DeploymentRow]: ...
async def find_many(self, *, where: Mapping[str, object] | None = None) -> Sequence[_DeploymentRow]: ...
async def find_unique(self, where: Mapping[str, object]) -> _DeploymentRow | None: ...
async def find_unique(self, *, where: Mapping[str, object]) -> _DeploymentRow | None: ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> object: ...
async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ...
def _model_table(prisma_client: PrismaClient) -> _ModelTableClient:
@ -322,7 +327,9 @@ async def get_all_access_groups_from_db(
for deployment in deployments:
model_info = deployment.model_info or {}
access_groups = model_info.get("access_groups", [])
access_groups = model_info.get( # pyright: ignore[reportAttributeAccessIssue] # Json reads back as a dict
"access_groups", []
)
model_name = deployment.model_name
for access_group in access_groups:

88
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -16,7 +16,7 @@ import json
from collections.abc import Awaitable, Mapping, Sequence
from json import JSONDecodeError
from types import MappingProxyType
from typing import Final, Literal, Protocol, cast
from typing import TYPE_CHECKING, Final, Literal, Protocol, cast
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from pydantic import BaseModel, ConfigDict, Field, ValidationError
@ -72,6 +72,7 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import (
)
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.model_repository import ModelRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import ModelTableRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.router import Router
@ -100,6 +101,9 @@ from litellm.types.router import (
)
from litellm.utils import get_utc_datetime
if TYPE_CHECKING:
from prisma import models as prisma_models
router: Final = APIRouter()
@ -120,10 +124,14 @@ class UpdatePublicModelGroupsRequest(BaseModel):
class _ProxyModelRow(Protocol):
model_id: str
model_name: str
litellm_params: Mapping[str, object]
model_info: Mapping[str, object] | None
@property
def model_id(self) -> str: ...
@property
def model_name(self) -> str: ...
@property
def model_info(self) -> object: ...
def model_dump_json(self, *, exclude_none: bool = False) -> str: ...
@ -133,7 +141,9 @@ class _ProxyModelTable(Protocol):
def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ...
def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[_ProxyModelRow]: ...
def update(
self, *, where: Mapping[str, object], data: Mapping[str, object]
) -> Awaitable[_ProxyModelRow | None]: ...
def delete(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ...
@ -144,41 +154,35 @@ class _TxModelTables(Protocol):
litellm_proxymodeltable: _ProxyModelTable
class _ExistingModelRow(Protocol):
@property
def litellm_params(self) -> Mapping[str, object]: ...
def model_dump_json(self, *, exclude_none: bool = False) -> str: ...
class _TeamRow(Protocol):
models: Sequence[str]
@property
def models(self) -> Sequence[str]: ...
def model_dump(self) -> Mapping[str, object]: ...
class _TeamTable(Protocol):
class _TeamLookupTable(Protocol):
def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_TeamRow | None]: ...
class _TeamTable(_TeamLookupTable, Protocol):
def update(
self, *, where: Mapping[str, object], data: Mapping[str, object], include: Mapping[str, bool]
) -> Awaitable[LiteLLM_TeamTable]: ...
class _TeamIdRef(Protocol):
team_id: str
class _ModelAliasRow(Protocol):
id: int
model_aliases: dict[str, str]
team: _TeamIdRef | None
class _ModelAliasTable(Protocol):
def find_many(self, *, include: Mapping[str, bool]) -> Awaitable[Sequence[_ModelAliasRow]]: ...
def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[object]: ...
def _proxy_model_table(prisma_client: PrismaClient) -> _ProxyModelTable:
return ModelRepository(prisma_client).table
def _repo_team_table(prisma_client: PrismaClient) -> _TeamTable:
def _repo_team_table(prisma_client: PrismaClient) -> _TeamLookupTable:
return TeamRepository(prisma_client).table
@ -186,7 +190,7 @@ def _db_team_table(prisma_client: PrismaClient) -> _TeamTable:
return prisma_client.db.litellm_teamtable
def _model_alias_table(prisma_client: PrismaClient) -> _ModelAliasTable:
def _model_alias_table(prisma_client: PrismaClient) -> "TableActions[prisma_models.LiteLLM_ModelTable]":
return ModelTableRepository(prisma_client).table
@ -677,6 +681,14 @@ async def patch_model(
data=update_data,
)
if updated_model is None:
raise ProxyException(
message=f"Model {model_id} not found on proxy.",
type=ProxyErrorTypes.not_found_error,
code=status.HTTP_404_NOT_FOUND,
param=None,
)
# Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates)
live_before_reload: Final = live_model_ids_snapshot()
reload_outcome: Final = await clear_cache()
@ -811,7 +823,7 @@ async def _set_model_blocked_status(
live_after=reload_outcome.live_after,
)
return updated_model
return updated_model # pyright: ignore[reportReturnType] # prisma row, coerced by this route's response_model
except Exception as e:
verbose_proxy_logger.exception("Error in model %s: %s", action, e)
@ -897,7 +909,7 @@ async def _add_model_to_db(
prisma_client: PrismaClient,
new_encryption_key: str | None = None,
should_create_model_in_db: bool = True,
) -> LiteLLM_ProxyModelTable | None:
) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None":
# encrypt litellm params #
_litellm_params_dict: Final = model_params.litellm_params.dict(exclude_none=True)
_original_litellm_model_name: Final = model_params.litellm_params.model
@ -914,8 +926,9 @@ async def _add_model_to_db(
}
if model_params.model_info.id is not None:
_data["model_id"] = model_params.model_info.id
_create_data: Final = cast("Mapping[str, object]", _data) # cast-ok: str-keyed json payload built just above
if should_create_model_in_db:
model_response = await ModelRepository(prisma_client).table.create(data=_data)
model_response = await ModelRepository(prisma_client).table.create(data=_create_data)
else:
model_response = LiteLLM_ProxyModelTable(**_data)
return model_response
@ -925,7 +938,7 @@ async def _add_team_model_to_db(
model_params: Deployment,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient,
) -> LiteLLM_ProxyModelTable | None:
) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None":
"""
If 'team_id' is provided,
@ -1638,7 +1651,9 @@ async def delete_team_model_alias(
tasks: Final = []
removed_model_aliases: Final[list[tuple[str, str]]] = []
for team_model_alias in team_model_aliases:
model_aliases = team_model_alias.model_aliases # {"alias": "public model name"}
model_aliases = cast( # cast-ok: prisma types Json columns as `str`; the driver hands back the parsed dict
"dict[str, str]", team_model_alias.model_aliases
)
id = team_model_alias.id
if public_model_name in model_aliases.values():
@ -1733,7 +1748,7 @@ async def add_new_model(
existing_params=None,
)
model_response: LiteLLM_ProxyModelTable | None = None
model_response: prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None = None
# update DB
incoming_model_info: Final = model_params.model_info.model_dump(exclude_none=True)
_raise_if_ptu_cost_attribution_disabled(incoming_model_info)
@ -1902,7 +1917,10 @@ async def update_model(
# update DB
if store_model_in_db is True:
_existing_litellm_params_dict: Final = dict(_existing_litellm_params.litellm_params)
existing_model_row: Final = cast( # cast-ok: prisma types Json columns as `str`; the driver parses them
"_ExistingModelRow", _existing_litellm_params
)
_existing_litellm_params_dict: Final = dict(existing_model_row.litellm_params)
if model_params.litellm_params is None:
raise Exception("litellm_params not provided")
@ -1946,8 +1964,8 @@ async def update_model(
user_api_key_dict=user_api_key_dict,
table_name=LitellmTableNames.PROXY_MODEL_TABLE_NAME,
before_value=(
_existing_litellm_params.model_dump_json(exclude_none=True)
if isinstance(_existing_litellm_params, BaseModel)
existing_model_row.model_dump_json(exclude_none=True)
if isinstance(existing_model_row, BaseModel)
else None
),
after_value=(

43
litellm/proxy/management_endpoints/organization_endpoints.py
View file

@ -14,7 +14,14 @@ Endpoints for /organization operations
#### ORGANIZATION MANAGEMENT ####
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Annotated, Final, Protocol, overload
from typing import (
TYPE_CHECKING,
Annotated,
Final,
Protocol,
cast, # noqa: TID251 # prisma types Json columns as fields.Json but reads back plain python values
overload,
)
import fastapi
from fastapi import APIRouter, Depends, HTTPException, Request, status
@ -74,6 +81,11 @@ if TYPE_CHECKING:
router: Final = APIRouter()
class _ObjectPermissionRow(Protocol):
@property
def object_permission_id(self) -> str | None: ...
class _UserTableClient(Protocol):
async def find_unique(self, where: Mapping[str, object]) -> "PrismaUserTable | None": ...
@ -681,7 +693,10 @@ async def update_organization(
existing_metadata: Final = existing_organization_row.metadata or {}
updated_metadata: Final = updated_organization_row_json.get("metadata", {})
merged_metadata: Final[Mapping[str, object]] = _update_dictionary(
existing_dict=existing_metadata.copy(), new_dict=updated_metadata
existing_dict=cast( # cast-ok: prisma de-serializes a Json column to the plain python dict it stores
"dict[str, object]", existing_metadata
).copy(),
new_dict=updated_metadata,
)
updated_organization_row_json["metadata"] = merged_metadata
@ -720,7 +735,7 @@ async def update_organization(
async def handle_update_object_permission(
data_json: dict[str, object],
existing_organization_row: LiteLLM_OrganizationTable,
existing_organization_row: _ObjectPermissionRow,
) -> dict[str, object]:
"""
Handle the update of object permission for an organization.
@ -1276,17 +1291,20 @@ async def find_member_if_email(user_email: str, prisma_client: PrismaClient) ->
Find a member if the user_email is in LiteLLM_UserTable
"""
not_unique_user_email_error: Final = HTTPException(
status_code=400,
detail={
"error": f"Unique user not found for user_email={user_email}. Potential duplicate OR non-existent user_email in LiteLLM_UserTable. Use 'user_id' instead."
},
)
try:
existing_user_email_row: Final[BaseModel] = await UserRepository(prisma_client).table.find_unique(
existing_user_email_row: Final = await UserRepository(prisma_client).table.find_unique(
where={"user_email": user_email}
)
except Exception:
raise HTTPException(
status_code=400,
detail={
"error": f"Unique user not found for user_email={user_email}. Potential duplicate OR non-existent user_email in LiteLLM_UserTable. Use 'user_id' instead."
},
)
raise not_unique_user_email_error
if existing_user_email_row is None:
raise not_unique_user_email_error
existing_user_email_row_pydantic: Final = LiteLLM_UserTable.model_validate(existing_user_email_row.model_dump())
return existing_user_email_row_pydantic
@ -1537,7 +1555,10 @@ async def add_member_to_organization(
_returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user")
if _returned_user is not None:
user_object = LiteLLM_UserTable.model_validate(_returned_user.model_dump())
elif existing_user_email_row is not None and len(existing_user_email_row) > 1:
elif existing_user_email_row is not None and (
len(existing_user_email_row) # pyright: ignore[reportArgumentType] # find_unique yields a row, not a list
> 1
):
raise HTTPException(
status_code=400,
detail={"error": "Multiple users with this email found in db. Please use 'user_id' instead."},

3
litellm/proxy/management_endpoints/scim/scim_transformations.py
View file

@ -33,7 +33,8 @@ class ScimTransformations:
# Get user's teams/groups
groups: Final = []
for team_id in user.teams or []:
team_ids: Final[list[str]] = user.teams or [] # mutable-ok: scim reads the user row's team ids
for team_id in team_ids:
team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id})
if team:
team_alias = getattr(team, "team_alias", team.team_id)

6
litellm/proxy/management_endpoints/scim/scim_v2.py
View file

@ -2761,6 +2761,12 @@ async def patch_group(
if final_team:
updated_team = final_team
if updated_team is None:
raise HTTPException(
status_code=404,
detail={"error": f"Group not found with ID: {group_id}"}, # mutable-ok: FastAPI detail contract
)
# Convert to SCIM format and return
scim_group: Final = await ScimTransformations.transform_litellm_team_to_scim_group(
LiteLLM_TeamTable.model_validate(updated_team.model_dump())

4
litellm/proxy/management_endpoints/tag_management_endpoints.py
View file

@ -369,10 +369,10 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str):
# Prisma returns litellm_params as dict (already parsed from JSON)
existing_params = db_model.litellm_params
if isinstance(existing_params, str):
if isinstance(existing_params, str): # pyright: ignore[reportUnnecessaryIsInstance] # prisma Json stub is str
# If it's a string, parse it
existing_params = json.loads(existing_params)
elif not isinstance(existing_params, dict):
elif not isinstance(existing_params, dict): # pyright: ignore[reportUnnecessaryIsInstance] # prisma Json stub
raise Exception(f"Unexpected litellm_params type: {type(existing_params)}")
# Add tag to tags array (preserve encryption of other fields)

3
litellm/proxy/management_endpoints/team_callback_endpoints.py
View file

@ -352,6 +352,9 @@ async def add_team_callbacks(
include={"object_permission": True}, # mutable-ok: prisma include takes a dict literal
)
if new_team_row is None:
raise _callback_error(400, f"Team id = {team_id} does not exist. Please use a different team id.")
# Without this a newly registered callback stays dormant for existing keys.
await _refresh_cached_team(
team_row=new_team_row,

337
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -16,7 +16,16 @@ import traceback
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Annotated, Final, NamedTuple, Protocol, TypedDict, TypeVar, cast
from typing import (
TYPE_CHECKING,
Annotated,
Final,
NamedTuple,
Protocol,
TypedDict,
TypeVar,
cast,
)
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
@ -33,21 +42,17 @@ from litellm.proxy._types import (
BudgetNewRequest,
CommonProxyErrors,
DeleteTeamRequest,
LiteLLM_AccessGroupTable,
LiteLLM_AuditLogs,
LiteLLM_BudgetTableFull,
LiteLLM_DeletedTeamTable,
LiteLLM_ManagementEndpoint_MetadataFields,
LiteLLM_ManagementEndpoint_MetadataFields_Premium,
LiteLLM_ModelTable,
LiteLLM_OrganizationMembershipTable,
LiteLLM_OrganizationTable,
LiteLLM_OrganizationTableWithMembers,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
LiteLLM_VerificationToken,
LitellmTableNames,
LitellmUserRoles,
Member,
@ -138,6 +143,7 @@ from litellm.proxy.management_helpers.utils import (
from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
AccessGroupRepository,
DeletedTeamRepository,
@ -169,6 +175,10 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
UpdateTeamMemberPermissionsRequest,
)
if TYPE_CHECKING:
from prisma import Prisma
from prisma import models as prisma_models
router: Final = APIRouter()
_DbRecordT = TypeVar("_DbRecordT")
@ -183,95 +193,14 @@ class _TeamIdGroupRow(TypedDict):
_count: _TeamIdKeyCount
class _PrismaTableActions(Protocol[_DbRecordT]):
async def find_unique(
self,
where: Mapping[str, object],
include: Mapping[str, bool] | None = None,
) -> _DbRecordT | None: ...
async def find_first(
self,
where: Mapping[str, object] | None = None,
order: Mapping[str, str] | None = None,
) -> _DbRecordT | None: ...
async def find_many(
self,
where: Mapping[str, object] | None = None,
include: Mapping[str, bool] | None = None,
order: Mapping[str, str] | None = None,
skip: int | None = None,
take: int | None = None,
cursor: Mapping[str, object] | None = None,
) -> list[_DbRecordT]: ...
async def create(
self,
data: Mapping[str, object],
include: Mapping[str, bool] | None = None,
) -> _DbRecordT: ...
async def create_many(
self,
data: Sequence[Mapping[str, object]],
skip_duplicates: bool | None = None,
) -> int: ...
async def update(
self,
where: Mapping[str, object],
data: Mapping[str, object],
include: Mapping[str, bool] | None = None,
) -> _DbRecordT: ...
async def update_many(
self,
where: Mapping[str, object],
data: Mapping[str, object],
) -> int: ...
async def upsert(
self,
where: Mapping[str, object],
data: Mapping[str, Mapping[str, object]],
) -> _DbRecordT: ...
async def delete_many(
self,
where: Mapping[str, object] | None = None,
) -> int: ...
async def count(
self,
where: Mapping[str, object] | None = None,
) -> int: ...
async def group_by(
self,
by: Sequence[str],
where: Mapping[str, object] | None = None,
count: Mapping[str, bool] | None = None,
) -> Sequence[_TeamIdGroupRow]: ...
class _HasTableActions(Protocol[_DbRecordT]):
@property
def table(self) -> "_PrismaTableActions[_DbRecordT]": ...
def _typed_table(
repo: "_HasTableActions[_DbRecordT]", record_type: type[_DbRecordT]
) -> "_PrismaTableActions[_DbRecordT]":
return repo.table
def _as_object(value: object) -> object:
return value
def _nullable(value: _DbRecordT | None) -> _DbRecordT | None:
return value
def _as_list(rows: Sequence[_DbRecordT]) -> list[_DbRecordT]: # mutable-ok: pydantic list[...] fields reject Sequence
return cast( # cast-ok: prisma-client-py find_many returns a list; TableActions only widens it to Sequence
"list[_DbRecordT]", rows
)
class _UserIdRow(Protocol):
@ -279,33 +208,75 @@ class _UserIdRow(Protocol):
def user_id(self) -> str | None: ...
class _HasUserIdTable(Protocol):
@property
def table(self) -> "_PrismaTableActions[_UserIdRow]": ...
def _user_id_rows_db(repo: "_HasUserIdTable") -> "_PrismaTableActions[_UserIdRow]":
def _user_id_rows_db(repo: UserRepository) -> "TableActions[_UserIdRow]":
return repo.table
class _RawTeamRow(Protocol):
class _ModelDumpRow(Protocol):
def model_dump(self) -> Mapping[str, object]: ...
class _TeamIdRow(Protocol):
@property
def members_with_roles(self) -> Sequence[Mapping[str, object]] | None: ...
def team_id(self) -> str: ...
class _HasRawTeamTable(Protocol):
class _CacheableTeamRow(_TeamIdRow, _ModelDumpRow, Protocol): ...
class _ObjectPermissionRow(Protocol):
@property
def table(self) -> "_PrismaTableActions[_RawTeamRow]": ...
def object_permission_id(self) -> str | None: ...
def _raw_team_db(repo: "_HasRawTeamTable") -> "_PrismaTableActions[_RawTeamRow]":
return repo.table
class _TeamAliasBudgetRow(Protocol):
@property
def team_alias(self) -> str | None: ...
@property
def budget_duration(self) -> str | None: ...
class _TeamBudgetRow(_TeamAliasBudgetRow, Protocol):
metadata: Mapping[str, JsonValue] | None
class _AuditableTeamRow(Protocol):
def json(self, *, exclude_none: bool = False) -> str: ...
class _RawTeamRow(_TeamIdRow, _ModelDumpRow, _ObjectPermissionRow, _TeamBudgetRow, _AuditableTeamRow, Protocol):
@property
def members_with_roles(
self,
) -> Sequence[dict[str, object]] | None: ... # mutable-ok: prisma deserializes this JSON column into plain dicts
@property
def organization_id(self) -> str | None: ...
@property
def max_budget(self) -> float | None: ...
@property
def soft_budget(self) -> float | None: ...
@property
def model_id(self) -> int | None: ...
def _raw_team_db(repo: TeamRepository) -> "TableActions[_RawTeamRow]":
return cast( # cast-ok: prisma types Json columns as str; the client hands back the deserialized value
"TableActions[_RawTeamRow]", repo.table
)
class _BudgetIdRow(Protocol):
@property
def budget_id(self) -> str: ...
class _BudgetWriteCall(Protocol):
async def __call__(
self, budget_obj: BudgetNewRequest, user_api_key_dict: UserAPIKeyAuth
) -> LiteLLM_BudgetTableFull: ...
async def __call__(self, budget_obj: BudgetNewRequest, user_api_key_dict: UserAPIKeyAuth) -> _BudgetIdRow: ...
def _as_budget_write(fn: "_BudgetWriteCall") -> "_BudgetWriteCall":
@ -330,7 +301,7 @@ class _TeamIdInFilter(TypedDict, total=False):
class _TeamCreateTx(AccessGroupSyncTx, Protocol):
@property
def litellm_teamtable(self) -> "_PrismaTableActions[LiteLLM_TeamTable]": ...
def litellm_teamtable(self) -> "TableActions[prisma_models.LiteLLM_TeamTable]": ...
_STRIP_DELETED_TEAM_FROM_USERS_SQL: Final = """
@ -340,46 +311,52 @@ UPDATE "LiteLLM_UserTable" SET teams = array_remove(teams, $1) WHERE $1 = ANY(te
_INCLUDE_MODEL_TABLE: Final = MappingProxyType({"litellm_model_table": True})
def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]":
return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable)
def _team_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_TeamTable]":
return TeamRepository(prisma_client).table
def _team_membership_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamMembership]":
return _typed_table(TeamMembershipRepository(prisma_client), LiteLLM_TeamMembership)
def _team_tx_db(tx: "Prisma") -> "TableActions[prisma_models.LiteLLM_TeamTable]":
return cast( # cast-ok: generated actions type Json columns as str; TableActions widens inputs to Mapping
"TableActions[prisma_models.LiteLLM_TeamTable]", tx.litellm_teamtable
)
def _user_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_UserTable]":
return _typed_table(UserRepository(prisma_client), LiteLLM_UserTable)
def _team_membership_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_TeamMembership]":
return TeamMembershipRepository(prisma_client).table
def _model_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_ModelTable]":
return _typed_table(ModelTableRepository(prisma_client), LiteLLM_ModelTable)
def _user_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_UserTable]":
return UserRepository(prisma_client).table
def _org_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_OrganizationTable]":
return _typed_table(OrganizationRepository(prisma_client), LiteLLM_OrganizationTable)
def _model_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_ModelTable]":
return ModelTableRepository(prisma_client).table
def _org_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_OrganizationTable]":
return OrganizationRepository(prisma_client).table
def _org_membership_db(
prisma_client: PrismaClient | None,
) -> "_PrismaTableActions[LiteLLM_OrganizationMembershipTable]":
return _typed_table(OrganizationMembershipRepository(prisma_client), LiteLLM_OrganizationMembershipTable)
) -> "TableActions[prisma_models.LiteLLM_OrganizationMembership]":
return OrganizationMembershipRepository(prisma_client).table
def _budget_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_BudgetTableFull]":
return _typed_table(BudgetRepository(prisma_client), LiteLLM_BudgetTableFull)
def _budget_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_BudgetTable]":
return BudgetRepository(prisma_client).table
def _deleted_team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_DeletedTeamTable]":
return _typed_table(DeletedTeamRepository(prisma_client), LiteLLM_DeletedTeamTable)
def _deleted_team_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_DeletedTeamTable]":
return DeletedTeamRepository(prisma_client).table
def _access_group_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_AccessGroupTable]":
return _typed_table(AccessGroupRepository(prisma_client), LiteLLM_AccessGroupTable)
def _access_group_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_AccessGroupTable]":
return AccessGroupRepository(prisma_client).table
def _tokens_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_VerificationToken]":
return _typed_table(VerificationTokenRepository(prisma_client), LiteLLM_VerificationToken)
def _tokens_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_VerificationToken]":
return VerificationTokenRepository(prisma_client).table
def _sanitize_for_log(value: object) -> str:
@ -392,7 +369,7 @@ def _sanitize_for_log(value: object) -> str:
async def _refresh_cached_team(
team_row: LiteLLM_TeamTable,
team_row: _CacheableTeamRow,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
) -> None:
@ -481,7 +458,7 @@ class TeamMemberBudgetHandler:
@staticmethod
async def create_team_member_budget_table(
data: NewTeamRequest | LiteLLM_TeamTable,
data: NewTeamRequest | _TeamAliasBudgetRow,
new_team_data_json: dict,
user_api_key_dict: UserAPIKeyAuth,
team_member_budget: float | None = None,
@ -532,7 +509,7 @@ class TeamMemberBudgetHandler:
@staticmethod
async def upsert_team_member_budget_table(
team_table: LiteLLM_TeamTable,
team_table: _TeamBudgetRow,
user_api_key_dict: UserAPIKeyAuth,
updated_kv: dict,
team_member_budget: float | None = None,
@ -603,7 +580,7 @@ class TeamMemberBudgetHandler:
@staticmethod
async def clear_team_member_budget_fields(
team_table: LiteLLM_TeamTable,
team_table: _TeamBudgetRow,
user_api_key_dict: "UserAPIKeyAuth",
updated_kv: dict,
explicitly_set_fields: set,
@ -1540,7 +1517,7 @@ async def new_team(
tx: _TeamCreateTx
async with prisma_client.db.tx() as tx:
team_row: Final[LiteLLM_TeamTable] = await tx.litellm_teamtable.create(
team_row: Final[prisma_models.LiteLLM_TeamTable] = await tx.litellm_teamtable.create(
data=team_creation_data,
include=_INCLUDE_MODEL_TABLE,
)
@ -1595,7 +1572,7 @@ async def new_team(
async def _create_team_update_audit_log(
existing_team_row: LiteLLM_TeamTable,
existing_team_row: _AuditableTeamRow,
updated_kv: dict,
team_id: str,
litellm_changed_by: str | None,
@ -1718,11 +1695,11 @@ async def _auto_add_team_members_to_organization(
async def fetch_and_validate_organization(
organization_id: str,
existing_team_row: LiteLLM_TeamTable,
existing_team_row: _ModelDumpRow,
llm_router: Router | None,
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth | None = None,
) -> LiteLLM_OrganizationTable:
) -> "prisma_models.LiteLLM_OrganizationTable":
"""
Fetch and validate an organization for team update operations.
@ -1996,7 +1973,9 @@ async def update_team(
validate_budget_duration(data.budget_duration)
validate_budget_duration(data.team_member_budget_duration)
existing_team_row = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id})
existing_team_row = await _raw_team_db(TeamRepository(prisma_client)).find_unique(
where={"team_id": data.team_id}
)
if existing_team_row is None:
raise HTTPException(
@ -2234,18 +2213,16 @@ async def update_team(
updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv)
team_update_data: Final[Mapping[str, object]] = updated_kv
team_row: Final[LiteLLM_TeamTable | None] = _nullable(
await _team_db(prisma_client).update(
where={"team_id": data.team_id},
data=team_update_data,
# `object_permission` is included so `_refresh_cached_team`
# doesn't write a cached team with the relation nulled out —
# see team_model_add for the full rationale.
include={
"litellm_model_table": True,
"object_permission": True,
},
)
team_row: Final = await _team_db(prisma_client).update(
where={"team_id": data.team_id},
data=team_update_data,
# `object_permission` is included so `_refresh_cached_team`
# doesn't write a cached team with the relation nulled out.
# See team_model_add for the full rationale.
include={
"litellm_model_table": True,
"object_permission": True,
},
)
if team_row is None or team_row.team_id is None:
@ -2375,7 +2352,7 @@ def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None:
updated_kv["budget_limits"] = json.dumps(initialized_windows)
async def handle_update_object_permission(data_json: dict, existing_team_row: LiteLLM_TeamTable) -> dict:
async def handle_update_object_permission(data_json: dict, existing_team_row: _ObjectPermissionRow) -> dict:
"""
Handle the update of object permission for a team.
@ -2705,7 +2682,7 @@ async def _add_team_members_to_team(
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_proxy_admin_name: str,
) -> tuple[LiteLLM_TeamTable, list[LiteLLM_UserTable], list[LiteLLM_TeamMembership]]:
) -> tuple["prisma_models.LiteLLM_TeamTable", list[LiteLLM_UserTable], list[LiteLLM_TeamMembership]]:
"""Add team members to the team.
The members_with_roles reconciliation runs inside a transaction that locks
@ -2750,7 +2727,7 @@ async def _write_members_with_roles_locked(
complete_team_data: LiteLLM_TeamTable,
prisma_client: PrismaClient,
updated_users: list[LiteLLM_UserTable],
) -> LiteLLM_TeamTable | None:
) -> "prisma_models.LiteLLM_TeamTable | None":
"""Reconcile members_with_roles under the team row lock. None when the team row is gone.
That read is at least as recent as the user and membership writes the caller
@ -2772,7 +2749,7 @@ async def _write_members_with_roles_locked(
)
_db_team_members: Final = [m.model_dump() for m in complete_team_data.members_with_roles]
return await tx.litellm_teamtable.update(
return await _team_tx_db(tx).update(
where={"team_id": data.team_id},
data={"members_with_roles": json.dumps(_db_team_members)},
)
@ -3292,7 +3269,9 @@ async def team_member_delete(
key_val: Final[Mapping[str, object]] = (
{"user_id": {"in": sorted(removed_user_ids)}} if removed_user_ids else {"user_email": data.user_email}
)
existing_user_rows: Final[Sequence[LiteLLM_UserTable]] = await _user_db(prisma_client).find_many(where=key_val)
existing_user_rows: Final[Sequence[prisma_models.LiteLLM_UserTable]] = await _user_db(prisma_client).find_many(
where=key_val
)
# Also clean up any existing team membership rows for this user and team
user_ids_to_delete: Final = removed_user_ids.union(
@ -3303,7 +3282,7 @@ async def team_member_delete(
## DELETE KEYS CREATED BY USER FOR THIS TEAM
# Fetch keys before deletion so their audit records can be persisted alongside the delete.
# An empty user_ids_to_delete still resolves cleanly: prisma's "in": [] matches no rows.
keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many(
keys_to_delete: Final = await _tokens_db(prisma_client).find_many(
where={
"user_id": {"in": sorted(user_ids_to_delete)},
"team_id": data.team_id,
@ -3313,7 +3292,7 @@ async def team_member_delete(
# All four cleanups run on one connection so a failure between them leaves
# no partial removal: either every write below lands, or none of them do.
async with prisma_client.tx() as tx:
await tx.litellm_teamtable.update(
await _team_tx_db(tx).update(
where={"team_id": data.team_id},
data={"members_with_roles": json.dumps(_db_new_team_members)},
)
@ -3826,9 +3805,7 @@ async def delete_team(
_persist_deleted_verification_tokens,
)
keys_to_delete: list[LiteLLM_VerificationToken] = await _tokens_db(prisma_client).find_many(
where={"team_id": {"in": data.team_ids}}
)
keys_to_delete: Final = await _tokens_db(prisma_client).find_many(where={"team_id": {"in": data.team_ids}})
if keys_to_delete:
await _persist_deleted_verification_tokens(
@ -3930,7 +3907,7 @@ async def _sweep_deleted_team_references(team_ids: Sequence[str], prisma_client:
async def _invalidate_deleted_key_cache(
keys: Sequence[LiteLLM_VerificationToken],
keys: "Sequence[prisma_models.LiteLLM_VerificationToken]",
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
) -> None:
@ -4115,7 +4092,7 @@ async def _hydrate_member_emails(
if not missing_user_ids:
return tuple(members)
user_rows: Final[Sequence[LiteLLM_UserTable]] = await _user_db(prisma_client).find_many(
user_rows: Final[Sequence[prisma_models.LiteLLM_UserTable]] = await _user_db(prisma_client).find_many(
where={ # mutable-ok: Prisma query filters are dict-shaped
"user_id": { # mutable-ok: Prisma query filters are dict-shaped
"in": sorted(missing_user_ids)
@ -4126,7 +4103,7 @@ async def _hydrate_member_emails(
return tuple(
m.model_copy(update={"user_email": email_by_user_id[m.user_id]}) # mutable-ok: pydantic update payload
if not m.user_email and m.user_id in email_by_user_id
if not m.user_email and m.user_id is not None and m.user_id in email_by_user_id
else m
for m in members
)
@ -4711,7 +4688,7 @@ async def _build_team_list_where_conditions(
async def _batch_resolve_access_group_resources(
all_access_group_ids: list[str],
) -> dict[str, LiteLLM_AccessGroupTable]:
) -> "dict[str, prisma_models.LiteLLM_AccessGroupTable]":
"""
Batch-fetch access groups in a single DB query and return them keyed by
access_group_id. Missing/invalid groups are silently omitted.
@ -4729,7 +4706,7 @@ async def _batch_resolve_access_group_resources(
def _convert_teams_to_response_models(
teams: list,
teams: Sequence,
use_deleted_table: bool,
keys_count_by_team: dict[str, int] | None = None,
) -> list[TeamListItem | LiteLLM_TeamTable | LiteLLM_DeletedTeamTable]:
@ -4763,7 +4740,7 @@ def _convert_teams_to_response_models(
async def _get_keys_count_by_team(
prisma_client: PrismaClient,
teams: Sequence[LiteLLM_TeamTable],
teams: Sequence[_TeamIdRow],
) -> dict[str, int]:
"""Aggregate virtual-key counts per team for the given page of teams.
@ -4775,10 +4752,13 @@ async def _get_keys_count_by_team(
if not page_team_ids:
return {}
grouped: Final = await _tokens_db(prisma_client).group_by(
by=["team_id"],
where={"team_id": {"in": page_team_ids}},
count={"team_id": True},
grouped: Final = cast( # cast-ok: prisma group_by returns one row per `by` key with `count=` nested under "_count"
"Sequence[_TeamIdGroupRow]",
await _tokens_db(prisma_client).group_by(
by=["team_id"],
where={"team_id": {"in": page_team_ids}},
count={"team_id": True},
),
)
return {row["team_id"]: row.get("_count", {}).get("team_id", 0) for row in grouped if row.get("team_id")}
@ -5168,7 +5148,7 @@ async def list_team(
_team_memberships.append(tm)
# add all keys that belong to the team
keys = await _tokens_db(prisma_client).find_many(where={"team_id": team.team_id})
keys = _as_list(await _tokens_db(prisma_client).find_many(where={"team_id": team.team_id}))
try:
returned_responses.append(
@ -5403,6 +5383,11 @@ async def team_model_add(
data={"updated_at": datetime.now(timezone.utc)},
include={"object_permission": True},
)
if updated_team is None:
raise HTTPException(
status_code=404,
detail={"error": f"Team not found, passed team_id={data.team_id}"},
)
await _refresh_cached_team(
team_row=updated_team,
@ -5485,6 +5470,11 @@ async def team_model_delete(
data={"models": updated_models},
include={"object_permission": True},
)
if updated_team is None:
raise HTTPException(
status_code=404,
detail={"error": f"Team not found, passed team_id={data.team_id}"},
)
await _refresh_cached_team(
team_row=updated_team,
@ -5619,8 +5609,13 @@ async def update_team_member_permissions(
where={"team_id": data.team_id},
data={"team_member_permissions": data.team_member_permissions},
)
if updated_team is None:
raise HTTPException(
status_code=404,
detail={"error": f"Team not found, passed team_id={data.team_id}"},
)
return updated_team
return updated_team # pyright: ignore[reportReturnType] # prisma row, coerced by this route's response_model
@router.post(
@ -5685,7 +5680,9 @@ async def bulk_update_team_member_permissions(
}
async def _compute_and_batch_updates(prisma_client, teams: Sequence[LiteLLM_TeamTable], permissions_to_add: set) -> int:
async def _compute_and_batch_updates(
prisma_client, teams: "Sequence[prisma_models.LiteLLM_TeamTable]", permissions_to_add: set
) -> int:
"""Compute merged permissions and batch-write updates. Returns count of teams updated."""
updates: Final = []
for team in teams:

63
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -29,7 +29,6 @@ from typing import (
NoReturn,
Optional,
Protocol,
TypeVar,
Union,
cast,
overload,
@ -122,6 +121,7 @@ from litellm.proxy.utils import (
get_custom_url,
get_server_root_path,
)
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import SSOConfigRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
@ -171,51 +171,16 @@ _CLI_SSO_SECRET_KEY_FRAGMENTS: Final = frozenset(
}
)
_DbRecordT: Final = TypeVar("_DbRecordT", covariant=True)
class _PrismaTableActions(Protocol[_DbRecordT]):
async def find_unique(
self,
where: Mapping[str, object],
) -> _DbRecordT | None: ...
async def find_first(
self,
where: Mapping[str, object] | None = None,
) -> _DbRecordT | None: ...
async def find_many(
self,
where: Mapping[str, object] | None = None,
include: Mapping[str, bool] | None = None,
) -> Sequence[_DbRecordT]: ...
async def update(
self,
where: Mapping[str, object],
data: Mapping[str, object],
) -> _DbRecordT: ...
async def update_many(
self,
where: Mapping[str, object],
data: Mapping[str, object],
) -> int: ...
class _UserMetadataRow(Protocol):
@property
def metadata(self) -> Mapping[str, object] | None: ...
class _HasUserMetadataTable(Protocol):
@property
def table(self) -> "_PrismaTableActions[_UserMetadataRow]": ...
def _user_meta_db(repo: "_HasUserMetadataTable") -> "_PrismaTableActions[_UserMetadataRow]":
return repo.table
def _user_meta_db(repo: UserRepository) -> "TableActions[_UserMetadataRow]":
return cast( # cast-ok: prisma types Json columns as str; the client hands back the deserialized value
"TableActions[_UserMetadataRow]", repo.table
)
class _SsoConfigRow(Protocol):
@ -223,25 +188,17 @@ class _SsoConfigRow(Protocol):
def sso_settings(self) -> Mapping[str, object] | None: ...
class _HasSsoConfigTable(Protocol):
@property
def table(self) -> "_PrismaTableActions[_SsoConfigRow]": ...
def _sso_config_db(repo: "_HasSsoConfigTable") -> "_PrismaTableActions[_SsoConfigRow]":
return repo.table
def _sso_config_db(repo: SSOConfigRepository) -> "TableActions[_SsoConfigRow]":
return cast( # cast-ok: prisma types Json columns as str; the client hands back the deserialized value
"TableActions[_SsoConfigRow]", repo.table
)
class _TeamDetailRow(Protocol):
def model_dump(self) -> Mapping[str, object]: ...
class _HasTeamDetailTable(Protocol):
@property
def table(self) -> "_PrismaTableActions[_TeamDetailRow]": ...
def _team_detail_db(repo: "_HasTeamDetailTable") -> "_PrismaTableActions[_TeamDetailRow]":
def _team_detail_db(repo: TeamRepository) -> "TableActions[_TeamDetailRow]":
return repo.table

14
litellm/proxy/management_helpers/object_permission_utils.py
View file

@ -4,7 +4,7 @@ organizations, teams, and keys.
"""
import json
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final, Optional
@ -19,6 +19,8 @@ from litellm.repositories.object_permission_repository import ObjectPermissionRe
from litellm.repositories.table_repositories import MCPServerRepository
if TYPE_CHECKING:
from prisma import models as prisma_models
from litellm.proxy._types import (
LiteLLM_ObjectPermissionTable,
LiteLLM_TeamTableCachedObj,
@ -26,7 +28,7 @@ if TYPE_CHECKING:
async def attach_object_permission_to_dict(
data_dict: dict,
data_dict: dict[str, object],
prisma_client: PrismaClient,
) -> dict:
"""
@ -61,7 +63,7 @@ async def attach_object_permission_to_dict(
try:
object_permission = object_permission.model_dump()
except Exception:
object_permission = object_permission.dict()
object_permission = object_permission.dict() # pyright: ignore[reportDeprecated] # pydantic v1 fallback
data_dict["object_permission"] = object_permission
return data_dict
@ -188,7 +190,9 @@ async def _set_object_permission(
return data_json
# Clean data: exclude None values and object_permission_id
clean_data: Final = {k: v for k, v in permission_data.items() if v is not None and k != "object_permission_id"}
clean_data: Final[dict[str, object]] = {
k: v for k, v in permission_data.items() if v is not None and k != "object_permission_id"
}
# Serialize mcp_tool_permissions to JSON string for GraphQL compatibility
if "mcp_tool_permissions" in clean_data:
@ -224,7 +228,7 @@ def _mcp_server_identifier_matches(server: Any, identifier: str) -> bool:
async def _get_db_mcp_servers_by_identifiers(
identifiers: set[str],
prisma_client: PrismaClient | None,
) -> list[Any]:
) -> "Sequence[prisma_models.LiteLLM_MCPServerTable]":
if prisma_client is None or not identifiers:
return []

26
litellm/proxy/management_helpers/utils.py
View file

@ -86,6 +86,20 @@ class _PrismaTeamMembershipTable(Protocol):
async def create(self, *, data: Mapping[str, object], include: Mapping[str, bool]) -> _PrismaRecord: ...
def _user_table(prisma_client: PrismaClient) -> _PrismaUserTable:
table: Final[_PrismaUserTable] = UserRepository(prisma_client).table
return table
async def _find_users_by_email(prisma_client: PrismaClient, user_email: str) -> Sequence[_PrismaUserRecord] | None:
rows: Final[Sequence[_PrismaUserRecord] | None] = await prisma_client.get_data(
key_val={"user_email": user_email},
table_name="user",
query_type="find_all",
)
return rows
def get_new_internal_user_defaults(user_id: str, user_email: str | None = None) -> dict[str, object]:
user_info: Final = litellm.default_internal_user_params or {}
@ -309,8 +323,7 @@ async def _append_team_id_if_absent(prisma_client: PrismaClient, user_id: str, t
number of teams a user belongs to). Teams added concurrently for a different
team id are unaffected, since each update filters on its own team id.
"""
user_table: Final[_PrismaUserTable] = UserRepository(prisma_client).table
await user_table.update_many(
await _user_table(prisma_client).update_many(
where={"user_id": user_id, "NOT": {"teams": {"has": team_id}}},
data={"teams": {"push": [team_id]}},
)
@ -348,8 +361,7 @@ async def add_new_member(
# Prisma only compiles an upsert down to INSERT ... ON CONFLICT when it
# is non-empty, and falls back to a racy SELECT-then-INSERT when it is
# not, so this re-states user_id as a no-op rather than being empty.
user_table: Final[_PrismaUserTable] = UserRepository(prisma_client).table
_returned_user: _PrismaUserRecord | None = await user_table.upsert(
_returned_user: _PrismaRecord | None = await _user_table(prisma_client).upsert(
where={"user_id": new_member.user_id},
data={
"create": {"teams": [team_id], **new_user_defaults},
@ -363,11 +375,7 @@ async def add_new_member(
new_user_defaults = get_new_internal_user_defaults(user_id=str(uuid.uuid4()), user_email=new_member.user_email)
## user email is not unique acc. to prisma schema -> future improvement
### for now: check if it exists in db, if not - insert it
existing_user_row: Final[list[_PrismaUserRecord] | None] = await prisma_client.get_data(
key_val={"user_email": new_member.user_email},
table_name="user",
query_type="find_all",
)
existing_user_row: Final = await _find_users_by_email(prisma_client, new_member.user_email)
if existing_user_row is None or (isinstance(existing_user_row, list) and len(existing_user_row) == 0):
new_user_defaults["teams"] = [team_id]
_returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user")

85
litellm/proxy/memory/memory_endpoints.py
View file

@ -18,21 +18,20 @@ Scoping:
"""
import json
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Final, Protocol
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final
from fastapi import APIRouter, Depends, HTTPException, Query
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
CommonProxyErrors,
LiteLLM_TeamTable,
LitellmUserRoles,
UserAPIKeyAuth,
user_api_key_has_admin_view,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import MemoryRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.types.memory_management import (
@ -44,54 +43,17 @@ from litellm.types.memory_management import (
)
if TYPE_CHECKING:
from prisma import models as prisma_models
from litellm.proxy.utils import PrismaClient
router: Final = APIRouter()
class _MemoryRecord(Protocol):
memory_id: str
key: str
value: str
metadata: object
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 _MemoryTableActions(Protocol):
async def create(self, data: Mapping[str, object]) -> _MemoryRecord: ...
async def find_many(
self,
where: Mapping[str, object] | None = ...,
order: Mapping[str, str] | None = ...,
skip: int = ...,
take: int = ...,
) -> Sequence[_MemoryRecord]: ...
async def count(self, where: Mapping[str, object] | None = ...) -> int: ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _MemoryRecord: ...
async def delete(self, where: Mapping[str, object]) -> _MemoryRecord | None: ...
def _memory_table(prisma_client: "PrismaClient") -> _MemoryTableActions:
def _memory_table(prisma_client: "PrismaClient") -> TableActions["prisma_models.LiteLLM_MemoryTable"]:
return MemoryRepository(prisma_client).table
class _TeamTableActions(Protocol):
async def find_unique(self, where: Mapping[str, str]) -> LiteLLM_TeamTable | None: ...
def _team_table(prisma_client: "PrismaClient") -> _TeamTableActions:
return TeamRepository(prisma_client).table
def _serialize_metadata_for_prisma(metadata: object) -> str:
"""
Encode a `metadata` payload for the `Json?` column.
@ -129,7 +91,7 @@ def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object
return {"OR": ors}
def _row_to_model(row: _MemoryRecord) -> LiteLLM_MemoryRow:
def _row_to_model(row: "prisma_models.LiteLLM_MemoryTable") -> LiteLLM_MemoryRow:
return LiteLLM_MemoryRow(
memory_id=row.memory_id,
key=row.key,
@ -163,7 +125,7 @@ def _internal_error(log_message: str, exc: Exception, default_detail: str) -> HT
async def _assert_write_access(
prisma_client: "PrismaClient", row: _MemoryRecord, user_api_key_dict: UserAPIKeyAuth
prisma_client: "PrismaClient", row: "prisma_models.LiteLLM_MemoryTable", user_api_key_dict: UserAPIKeyAuth
) -> None:
"""
Enforce ownership for mutations (PUT/DELETE).
@ -219,7 +181,7 @@ async def _is_team_admin_for(prisma_client: "PrismaClient", user_api_key_dict: U
)
try:
team_obj: Final = await _team_table(prisma_client).find_unique(where={"team_id": team_id})
team_obj: Final = await TeamRepository(prisma_client).find_by_id(team_id, id_field="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
@ -407,7 +369,7 @@ async def list_memory(
async def _find_memory_for_caller(
prisma_client: "PrismaClient", key: str, user_api_key_dict: UserAPIKeyAuth
) -> _MemoryRecord:
) -> "prisma_models.LiteLLM_MemoryTable":
"""Look up a memory row by key, scoped to the caller's visibility."""
key_filter: Final[Mapping[str, object]] = {"key": key}
vis: Final = _visibility_filter(user_api_key_dict)
@ -418,6 +380,18 @@ async def _find_memory_for_caller(
return rows[0]
async def _find_visible_memory_or_none(
prisma_client: "PrismaClient", key: str, user_api_key_dict: UserAPIKeyAuth
) -> "prisma_models.LiteLLM_MemoryTable | None":
"""The caller-visible row for `key`, or None when nothing is visible to them."""
try:
return await _find_memory_for_caller(prisma_client, key, user_api_key_dict)
except HTTPException as e:
if e.status_code == 404:
return None
raise
@router.get(
"/v1/memory/{key:path}",
tags=["memory management"],
@ -480,17 +454,8 @@ async def upsert_memory(
)
data["updated_by"] = user_api_key_dict.user_id
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)
except HTTPException as e:
if e.status_code == 404:
return None
raise
try:
existing: Final = await _find_existing()
existing: Final = await _find_visible_memory_or_none(prisma_client, key, user_api_key_dict)
if existing is not None:
# Visibility != write authority. Make sure the caller actually
# owns this row (their user_id matches, or it's a pure team row in
@ -530,7 +495,7 @@ async def upsert_memory(
# instead of surfacing a 500 on a unique-violation.
if not _is_unique_violation(e):
raise
existing_after_race: Final = await _find_existing()
existing_after_race: Final = await _find_visible_memory_or_none(prisma_client, key, user_api_key_dict)
if existing_after_race is None:
# Row exists globally but isn't visible to this caller
# (owned by someone else). Treat as conflict.
@ -549,6 +514,8 @@ async def upsert_memory(
except Exception as e:
raise _internal_error("Error upserting memory: %s", e, "Internal error updating memory entry.")
if row is None:
raise HTTPException(status_code=404, detail=f"Memory with key '{key}' not found")
return _row_to_model(row)

24
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -4,7 +4,16 @@ import re
from collections.abc import Mapping
from dataclasses import dataclass, field
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, get_args, runtime_checkable
from typing import (
TYPE_CHECKING,
Final,
Literal,
Optional,
Protocol,
cast, # noqa: TID251 # prisma types Json columns as fields.Json but de-serializes them to plain python on read
get_args,
runtime_checkable,
)
from litellm.proxy._types import ProxyException
from litellm.repositories.table_repositories import (
@ -1183,7 +1192,7 @@ async def ensure_batch_response_managed_file_ids(
prisma_client,
verbose_proxy_logger,
user_api_key_dict=None,
db_batch_object=None,
db_batch_object: "LiteLLM_ManagedObjectTable | None" = None,
unified_batch_id: str | Literal[False] | None = None,
) -> None:
"""Normalize batch file IDs to managed unified IDs before DB persistence."""
@ -1270,11 +1279,10 @@ async def get_batch_from_database(
return None, None
# Parse the batch object from database
batch_data: Final = (
json.loads(db_batch_object.file_object)
if isinstance(db_batch_object.file_object, str)
else db_batch_object.file_object
file_object: Final = cast( # cast-ok: prisma types the Json column as str; reads return the decoded value
"Mapping[str, object] | str", db_batch_object.file_object
)
batch_data: Final = json.loads(file_object) if isinstance(file_object, str) else file_object
response: Final = LiteLLMBatch.model_validate(batch_data)
response.id = batch_id
@ -1360,7 +1368,7 @@ async def update_batch_in_database(
managed_files_obj,
prisma_client,
verbose_proxy_logger,
db_batch_object=None,
db_batch_object: "LiteLLM_ManagedObjectTable | None" = None,
operation: str = "update",
user_api_key_dict=None,
poller_owns_accounting: bool | None = None,
@ -1427,7 +1435,7 @@ async def update_batch_in_database(
# Normalize status for database storage
db_status: Final = response.status if response.status != "completed" else "complete"
update_data: Final[dict] = {
update_data: Final[dict[str, object]] = {
"status": db_status,
"file_object": response.model_dump_json(),
"updated_at": litellm.utils.get_utc_datetime(),

16
litellm/proxy/pass_through_endpoints/managed_id_rewriter.py
View file

@ -33,7 +33,13 @@ from __future__ import annotations
import json
import re
from collections.abc import Callable, Mapping, Sequence
from typing import TYPE_CHECKING, Final, TypeVar, overload
from typing import (
TYPE_CHECKING,
Final,
TypeVar,
cast, # noqa: TID251 # prisma stubs type Json columns as fields.Json but de-serialize them on read
overload,
)
from urllib.parse import quote, unquote
from fastapi import HTTPException
@ -286,11 +292,15 @@ def _canonical_path(route: str) -> str:
def _file_table(prisma_client: PrismaClient) -> ManagedFileTable:
return ManagedFileRepository(prisma_client).table
return cast( # cast-ok: stub-only mismatch, prisma returns real lists and de-serialized Json
ManagedFileTable, ManagedFileRepository(prisma_client).table
)
def _object_table(prisma_client: PrismaClient) -> ManagedObjectTable:
return ManagedObjectRepository(prisma_client).table
return cast( # cast-ok: stub-only mismatch, prisma returns real lists and de-serialized Json
ManagedObjectTable, ManagedObjectRepository(prisma_client).table
)
async def _resolve_one(

11
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -5,7 +5,7 @@ import json
import posixpath
import traceback
from base64 import b64encode
from collections.abc import AsyncGenerator, Callable, Iterable, Mapping
from collections.abc import AsyncGenerator, Callable, Iterable, Mapping, Sequence
from datetime import datetime
from itertools import groupby
from typing import Any, Final, TypedDict, cast
@ -3183,13 +3183,18 @@ async def _filter_endpoints_by_team_allowed_routes(
)
# retrieve team metadata
team_metadata: Final = team.metadata
team_metadata: Final = cast( # cast-ok: prisma types the Json column as str; reads hand back the decoded value
"Mapping[str, object] | None", team.metadata
)
if team_metadata is not None and team_metadata.get("allowed_passthrough_routes") is not None:
## FILTER pass_through_endpoints by allowed_passthrough_routes
pass_through_endpoints = [
endpoint
for endpoint in pass_through_endpoints
if endpoint.path in team_metadata.get("allowed_passthrough_routes")
if endpoint.path
in cast( # cast-ok: guarded above; team metadata stores this key as a list of route paths
"Sequence[str]", team_metadata.get("allowed_passthrough_routes")
)
]
return pass_through_endpoints

68
litellm/proxy/policy_engine/policy_registry.py
View file

@ -10,9 +10,20 @@ by policy_attachments (see AttachmentRegistry).
import json
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypedDict, Union
from typing import (
TYPE_CHECKING,
Any,
Final,
Literal,
Optional,
Protocol,
TypedDict,
Union,
cast, # noqa: TID251 # prisma types the condition/pipeline Json columns as str, but reads return decoded values
)
from litellm._logging import verbose_proxy_logger
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import PolicyRepository
from litellm.types.proxy.policy_engine import (
GuardrailPipeline,
@ -65,15 +76,32 @@ class _PolicyRow(Protocol):
class _PolicyVersionSourceRow(Protocol):
policy_id: str
policy_name: str
version_number: int
inherit: str | None
description: str | None
guardrails_add: Sequence[str] | None
guardrails_remove: Sequence[str] | None
condition: Mapping[str, object] | str | None
pipeline: Mapping[str, object] | str | None
@property
def policy_id(self) -> str: ...
@property
def policy_name(self) -> str: ...
@property
def version_number(self) -> int: ...
@property
def inherit(self) -> str | None: ...
@property
def description(self) -> str | None: ...
@property
def guardrails_add(self) -> Sequence[str] | None: ...
@property
def guardrails_remove(self) -> Sequence[str] | None: ...
@property
def condition(self) -> Mapping[str, object] | str | None: ...
@property
def pipeline(self) -> Mapping[str, object] | str | None: ...
class _PolicyTableClient(Protocol):
@ -96,23 +124,15 @@ class _PolicyTableClient(Protocol):
async def delete_many(self, where: Mapping[str, object]) -> int: ...
class _PolicyVersionSourceTableClient(Protocol):
async def find_unique(self, where: Mapping[str, object]) -> _PolicyVersionSourceRow | None: ...
async def find_first(
self,
where: Mapping[str, object],
order: Mapping[str, str] | None = None,
) -> _PolicyVersionSourceRow | None: ...
def _policy_table(prisma_client: "PrismaClient") -> _PolicyTableClient:
table: Final[_PolicyTableClient] = PolicyRepository(prisma_client).table
return table
table: Final = PolicyRepository(prisma_client).table
return cast( # cast-ok: prisma types Json columns as str; the client hands back the decoded condition/pipeline
"_PolicyTableClient", table
)
def _policy_version_source_table(prisma_client: "PrismaClient") -> _PolicyVersionSourceTableClient:
table: Final[_PolicyVersionSourceTableClient] = PolicyRepository(prisma_client).table
def _policy_version_source_table(prisma_client: "PrismaClient") -> "TableActions[_PolicyVersionSourceRow]":
table: Final[TableActions[_PolicyVersionSourceRow]] = PolicyRepository(prisma_client).table
return table

56
litellm/proxy/policy_engine/policy_resolve_endpoints.py
View file

@ -6,7 +6,8 @@ Policy resolve and attachment impact estimation endpoints.
"""
import json
from typing import Final
from collections.abc import Sequence
from typing import TYPE_CHECKING, Final
from fastapi import APIRouter, Depends, HTTPException, Query
@ -30,25 +31,28 @@ from litellm.types.proxy.policy_engine import (
PolicyResolveResponse,
)
if TYPE_CHECKING:
from prisma import models as prisma_models
router: Final = APIRouter()
def _build_alias_where(field: str, patterns: list) -> dict:
def _build_alias_where(field: str, patterns: Sequence[str]) -> dict[str, object]:
"""Build a Prisma ``where`` clause for alias patterns.
Supports exact matches and suffix wildcards (``prefix*``).
Returns something like:
{"OR": [{"field": {"in": ["a","b"]}}, {"field": {"startsWith": "dev-"}}]}
"""
exact: Final[list] = []
prefix_conditions: Final[list] = []
exact: Final[list[str]] = []
prefix_conditions: Final[list[dict[str, object]]] = []
for pat in patterns:
if pat.endswith("*"):
prefix_conditions.append({field: {"startsWith": pat[:-1]}})
else:
exact.append(pat)
conditions: Final[list] = []
conditions: Final[list[dict[str, object]]] = []
if exact:
conditions.append({field: {"in": exact}})
conditions.extend(prefix_conditions)
@ -79,7 +83,7 @@ def _get_tags_from_metadata(metadata: object, json_metadata: object = None) -> l
return parsed.get("tags", []) or []
async def _fetch_all_teams(prisma_client: object) -> list:
async def _fetch_all_teams(prisma_client: object) -> "Sequence[prisma_models.LiteLLM_TeamTable]":
"""Fetch teams from DB once. Reuse the result across tag and alias lookups."""
return await TeamRepository(prisma_client).table.find_many(
where={},
@ -88,13 +92,15 @@ async def _fetch_all_teams(prisma_client: object) -> list:
)
def _filter_keys_by_tags(keys: list, tag_patterns: list) -> tuple:
def _filter_keys_by_tags(
keys: "Sequence[prisma_models.LiteLLM_VerificationToken]", tag_patterns: Sequence[str]
) -> tuple[list[str], int]:
"""Filter key rows whose metadata.tags match any of the given patterns.
Returns (named_aliases, unnamed_count).
"""
affected: Final[list] = []
affected: Final[list[str]] = []
unnamed_count = 0
for key in keys:
key_alias = key.key_alias or ""
@ -111,13 +117,15 @@ def _filter_keys_by_tags(keys: list, tag_patterns: list) -> tuple:
return affected, unnamed_count
def _filter_teams_by_tags(teams: list, tag_patterns: list) -> tuple:
def _filter_teams_by_tags(
teams: "Sequence[prisma_models.LiteLLM_TeamTable]", tag_patterns: Sequence[str]
) -> tuple[list[str], int]:
"""Filter pre-fetched team rows whose metadata.tags match any patterns.
Returns (named_aliases, unnamed_count).
"""
affected: Final[list] = []
affected: Final[list[str]] = []
unnamed_count = 0
for team in teams:
team_alias = team.team_alias or ""
@ -136,18 +144,18 @@ def _filter_teams_by_tags(teams: list, tag_patterns: list) -> tuple:
async def _find_affected_by_team_patterns(
prisma_client: object,
all_teams: list,
team_patterns: list,
existing_teams: list,
existing_keys: list,
) -> tuple:
all_teams: "Sequence[prisma_models.LiteLLM_TeamTable]",
team_patterns: Sequence[str],
existing_teams: Sequence[str],
existing_keys: Sequence[str],
) -> tuple[list[str], list[str], int]:
"""Filter pre-fetched teams by alias patterns, then fetch their keys.
Returns (new_teams, new_keys, unnamed_keys_count).
"""
new_teams: Final[list] = []
matched_team_ids: Final[list] = []
new_teams: Final[list[str]] = []
matched_team_ids: Final[list[str]] = []
for team in all_teams:
team_alias = team.team_alias or ""
@ -158,7 +166,7 @@ async def _find_affected_by_team_patterns(
new_teams.append(team_alias)
matched_team_ids.append(str(team.team_id))
new_keys: Final[list] = []
new_keys: Final[list[str]] = []
unnamed_keys_count = 0
if matched_team_ids:
keys: Final = await VerificationTokenRepository(prisma_client).table.find_many(
@ -177,10 +185,12 @@ async def _find_affected_by_team_patterns(
return new_teams, new_keys, unnamed_keys_count
async def _find_affected_keys_by_alias(prisma_client: object, key_patterns: list, existing_keys: list) -> list:
async def _find_affected_keys_by_alias(
prisma_client: object, key_patterns: Sequence[str], existing_keys: Sequence[str]
) -> list[str]:
"""Find keys whose alias matches the given patterns."""
affected: Final[list] = []
affected: Final[list[str]] = []
keys: Final = await VerificationTokenRepository(prisma_client).table.find_many(
where=_build_alias_where("key_alias", key_patterns),
@ -349,8 +359,8 @@ async def estimate_attachment_impact(
sample_teams=["(global scope — affects all teams)"],
)
affected_keys: list = []
affected_teams: list = []
affected_keys: list[str] = []
affected_teams: list[str] = []
unnamed_keys = 0
unnamed_teams = 0
@ -358,7 +368,7 @@ async def estimate_attachment_impact(
team_patterns: Final = request.teams or []
# Fetch teams once — reused by both tag-based and alias-based lookups
all_teams: list = []
all_teams: Sequence[prisma_models.LiteLLM_TeamTable] = []
if tag_patterns or team_patterns:
all_teams = await _fetch_all_teams(prisma_client)

8
litellm/proxy/prompts/prompt_endpoints.py
View file

@ -93,7 +93,7 @@ class _PromptTableActions(Protocol):
def create(self, *, data: Mapping[str, str | int | None]) -> Awaitable[_PromptRow]: ...
def update(self, *, where: Mapping[str, str | int], data: Mapping[str, str]) -> Awaitable[_PromptRow]: ...
def update(self, *, where: Mapping[str, str | int], data: Mapping[str, str]) -> Awaitable[_PromptRow | None]: ...
def delete_many(self, *, where: Mapping[str, str]) -> Awaitable[int]: ...
@ -1157,6 +1157,12 @@ async def patch_prompt(
data=update_data,
)
if updated_prompt_db_entry is None:
raise HTTPException(
status_code=404,
detail=f"Prompt with ID {base_prompt_id} not found in environment {env}",
)
updated_prompt_spec: Final = create_versioned_prompt_spec(db_prompt=updated_prompt_db_entry)
return _reload_prompt_in_registry(IN_MEMORY_PROMPT_REGISTRY, versioned_id, updated_prompt_spec)

107
litellm/proxy/proxy_server.py
View file

@ -133,6 +133,7 @@ if TYPE_CHECKING:
from aiohttp import ClientSession
from fastapi.routing import APIRoute
from opentelemetry.trace import Span as _Span
from prisma import models as prisma_models
from litellm.integrations.opentelemetry import OpenTelemetry
@ -634,6 +635,7 @@ from litellm.proxy.utils import (
from litellm.proxy.video_endpoints.endpoints import router as video_router
from litellm.repositories.base_repository import SupportsModelDump
from litellm.repositories.credentials_repository import CredentialsRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.router import (
AssistantsTypedDict,
Deployment,
@ -1642,12 +1644,21 @@ class _InvitationLinkRow(Protocol):
class _UserTableRow(Protocol):
user_id: str
user_email: str | None
user_role: str
user_role: str | None
class _ModelTableRow(Protocol):
model_id: str | None
created_by: str | None
class _UserTeamsRow(Protocol):
@property
def teams(self) -> Sequence[str]: ...
_ProxyModelRow: TypeAlias = "prisma_models.LiteLLM_ProxyModelTable"
def _config_param_table(client: PrismaClient | None) -> TableActions[_ConfigParamRow]:
return cast( # cast-ok: this is prisma's LiteLLM_Config actions object, which parses its Json column to a mapping
"TableActions[_ConfigParamRow]", ConfigRepository(client).table
)
class _TTFTRow(TypedDict):
@ -4370,7 +4381,7 @@ class ProxyConfig:
if prisma_client is None or not (general_settings.get("store_model_in_db", False) is True or store_model_in_db):
return
row: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
row: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
where={"param_name": "environment_variables"}
)
existing: Final[dict] = dict(row.param_value) if row is not None and row.param_value is not None else {}
@ -6226,7 +6237,7 @@ class ProxyConfig:
4. Update router settings
"""
if llm_router is not None and prisma_client is not None:
db_router_settings: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
db_router_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
where={"param_name": "router_settings"}
)
@ -6654,7 +6665,7 @@ class ProxyConfig:
def _should_load_db_object(self, object_type: str | SupportedDBObjectType) -> bool:
return should_load_db_object(object_type=object_type)
async def _get_models_from_db(self, prisma_client: PrismaClient) -> list | None:
async def _get_models_from_db(self, prisma_client: PrismaClient) -> Sequence[_ProxyModelRow] | None:
"""
Fetch all model deployments from the DB.
@ -6664,7 +6675,7 @@ class ProxyConfig:
as "all models deleted" and must not evict existing router deployments.
"""
try:
new_models: Final[list[_ModelTableRow]] = await ModelRepository(prisma_client).table.find_many()
new_models: Final[Sequence[_ProxyModelRow]] = await ModelRepository(prisma_client).table.find_many()
return new_models
except Exception as e:
verbose_proxy_logger.exception(
@ -6950,10 +6961,13 @@ class ProxyConfig:
"""
try:
sso_settings: Final[_SSOConfigRow | None] = await call_with_db_reconnect_retry(
prisma_client,
lambda: SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}),
reason="init_sso_settings_in_db_lookup_failure",
sso_settings: Final[_SSOConfigRow | None] = cast( # cast-ok: prisma Json stub is `str`, runtime is a dict
"_SSOConfigRow | None",
await call_with_db_reconnect_retry(
prisma_client,
lambda: SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}),
reason="init_sso_settings_in_db_lookup_failure",
),
)
if sso_settings is not None:
sso_settings.sso_settings.pop("role_mappings", None)
@ -6981,12 +6995,15 @@ class ProxyConfig:
)
try:
db_record: Final[_ConfigOverridesRow | None] = await call_with_db_reconnect_retry(
prisma_client,
lambda: ConfigOverridesRepository(prisma_client).table.find_unique(
where={"config_type": "hashicorp_vault"}
db_record: Final[_ConfigOverridesRow | None] = cast( # cast-ok: prisma Json stub is `str`, runtime dict
"_ConfigOverridesRow | None",
await call_with_db_reconnect_retry(
prisma_client,
lambda: ConfigOverridesRepository(prisma_client).table.find_unique(
where={"config_type": "hashicorp_vault"}
),
reason="init_hashicorp_vault_config_override_lookup_failure",
),
reason="init_hashicorp_vault_config_override_lookup_failure",
)
if db_record is None or db_record.config_value is None:
@ -8834,8 +8851,9 @@ class ProxyStartupEvent:
if prisma_client is None:
return
db_record: Final[_UISettingsRow | None] = await UISettingsRepository(prisma_client).table.find_unique(
where={"id": "ui_settings"}
db_record: Final[_UISettingsRow | None] = cast( # cast-ok: prisma Json stub is `str`, runtime is a dict
"_UISettingsRow | None",
await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}),
)
if db_record and db_record.ui_settings:
raw: Final = db_record.ui_settings
@ -8998,7 +9016,7 @@ class ProxyStartupEvent:
# but YAML config has False.
if store_model_in_db is not True and prisma_client is not None:
try:
_db_gs_record: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
_db_gs_record: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
where={"param_name": "general_settings"}
)
if _db_gs_record is not None and isinstance(_db_gs_record.param_value, dict):
@ -12143,14 +12161,14 @@ async def _check_if_model_is_user_added(
id = model.get("model_info", {}).get("id", None)
if id is None:
continue
db_model: _ModelTableRow | None = await ModelRepository(prisma_client).table.find_unique(where={"model_id": id})
db_model: _ProxyModelRow | None = await ModelRepository(prisma_client).table.find_unique(where={"model_id": id})
if db_model is not None:
if db_model.created_by == user_api_key_dict.user_id:
filtered_models.append(model)
return filtered_models
def _check_if_model_is_team_model(models: list[DeploymentTypedDict], user_row: LiteLLM_UserTable) -> list[dict]:
def _check_if_model_is_team_model(models: list[DeploymentTypedDict], user_row: _UserTeamsRow) -> list[dict]:
"""
Check if model is a team model
@ -12202,6 +12220,9 @@ async def non_admin_all_models(
except Exception:
raise HTTPException(status_code=400, detail={"error": "User not found"})
if user_row is None:
raise HTTPException(status_code=400, detail={"error": "User not found"})
# Get all models that are team models, when model team_id == user_row.teams
all_models += _check_if_model_is_team_model(
models=llm_router.get_model_list() or [],
@ -12630,7 +12651,7 @@ async def _fetch_db_models_for_search(
db_models_total_count: Final = await ModelRepository(prisma_client).table.count(where=db_where_condition)
db_models_raw: list = []
db_models_raw: Sequence[_ProxyModelRow] = []
if take_limit > 0:
db_models_raw = await ModelRepository(prisma_client).table.find_many(
where=db_where_condition,
@ -13020,7 +13041,7 @@ async def _gather_team_accessible_model_ids(
try:
if team_object.models and SpecialModelNames.all_proxy_models.value not in team_object.models:
_resolved_names: Final = _team_models_resolve_to_names(team_object.models, access_groups)
db_models: Final[Sequence[_ModelTableRow]] = await ModelRepository(prisma_client).table.find_many(
db_models: Final[Sequence[_ProxyModelRow]] = await ModelRepository(prisma_client).table.find_many(
where={"model_name": {"in": _resolved_names}}
)
for db_model in db_models:
@ -14494,14 +14515,18 @@ async def alerting_settings(
)
## get general settings from db
db_general_settings: Final = await ConfigRepository(prisma_client).table.find_first(
db_general_settings: Final = await _config_param_table(prisma_client).find_first(
where={"param_name": "general_settings"}
)
if db_general_settings is not None and db_general_settings.param_value is not None:
db_general_settings_dict: Final = dict(db_general_settings.param_value)
alerting_args_dict: dict = db_general_settings_dict.get("alerting_args", {})
alerting_values: list | None = db_general_settings_dict.get("alerting")
alerting_args_dict: dict = cast( # cast-ok: ConfigGeneralSettings validates alerting_args as a dict on write
dict[str, JsonValue], db_general_settings_dict.get("alerting_args", {})
)
alerting_values: list | None = cast( # cast-ok: ConfigGeneralSettings validates alerting as a list on write
list[JsonValue] | None, db_general_settings_dict.get("alerting")
)
else:
alerting_args_dict = {}
alerting_values = None
@ -15052,7 +15077,7 @@ async def onboarding(invite_link: str, request: Request):
user_id=user_obj.user_id,
key=onboarding_token,
user_email=user_obj.user_email,
user_role=user_obj.user_role,
user_role=user_obj.user_role, # pyright: ignore[reportArgumentType] # nullable DB column, no unset contract
login_method="username_password",
premium_user=premium_user,
auth_header_name=general_settings.get("litellm_key_header_name", "Authorization"),
@ -15161,7 +15186,7 @@ async def _generate_onboarding_ui_session_token(user_obj: _UserTableRow) -> str:
user_id=user_obj.user_id,
key=key,
user_email=user_obj.user_email,
user_role=user_obj.user_role,
user_role=user_obj.user_role, # pyright: ignore[reportArgumentType] # nullable DB column, no unset contract
login_method="username_password",
premium_user=premium_user,
auth_header_name=general_settings.get("litellm_key_header_name", "Authorization"),
@ -15722,7 +15747,7 @@ async def update_config(
raise Exception("No DB Connected")
async def _read_section(param_name: str) -> dict:
row: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
row: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
where={"param_name": param_name}
)
if row is None or row.param_value is None:
@ -15979,7 +16004,7 @@ async def update_config_general_settings(
)
## get general settings from db
db_general_settings: Final = await ConfigRepository(prisma_client).table.find_first(
db_general_settings: Final = await _config_param_table(prisma_client).find_first(
where={"param_name": "general_settings"}
)
### update value
@ -15997,7 +16022,7 @@ async def update_config_general_settings(
if data.field_name == "plugins":
field_value = _preserve_redacted_plugin_keys(field_value, general_settings.get("plugins"))
general_settings[data.field_name] = field_value
general_settings[data.field_name] = cast(JsonValue, field_value) # cast-ok: ConfigGeneralSettings validated it
response: Final = await ConfigRepository(prisma_client).table.upsert(
where={"param_name": "general_settings"},
@ -16017,7 +16042,7 @@ async def update_config_general_settings(
)
if data.field_name == "plugins":
register_plugins_from_config(general_settings)
register_plugins_from_config(cast(dict[str, object], general_settings)) # cast-ok: the callee only reads it
_apply_ssrf_general_settings(general_settings)
return response
@ -16193,7 +16218,7 @@ async def get_config_general_settings(
)
## get general settings from db
db_general_settings: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
db_general_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
where={"param_name": "general_settings"}
)
### pop the value
@ -16382,7 +16407,7 @@ async def get_config_list(
is_full_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
## get general settings from db
db_general_settings: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
db_general_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
where={"param_name": "general_settings"}
)
@ -16478,7 +16503,7 @@ async def get_config_list(
)
return_val.append(_response_obj)
db_litellm_settings_row: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
db_litellm_settings_row: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
where={"param_name": "litellm_settings"}
)
db_litellm_settings: Final[dict] = (
@ -16555,7 +16580,7 @@ async def delete_config_general_settings(
)
## get general settings from db
db_general_settings: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
db_general_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
where={"param_name": "general_settings"}
)
### pop the value
@ -17122,7 +17147,7 @@ async def reload_anthropic_beta_headers(
last_anthropic_beta_headers_reload = current_time.isoformat()
# Set force reload flag in database for other pods, preserving existing interval_hours
existing_beta_config: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_unique(
existing_beta_config: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_unique(
where={"param_name": "anthropic_beta_headers_reload_config"}
)
existing_beta_interval = None
@ -17300,7 +17325,7 @@ async def get_anthropic_beta_headers_reload_status(
}
# Get reload configuration from database
config_record: Final = await ConfigRepository(prisma_client).table.find_unique(
config_record: Final = await _config_param_table(prisma_client).find_unique(
where={"param_name": "anthropic_beta_headers_reload_config"}
)
@ -17314,7 +17339,9 @@ async def get_anthropic_beta_headers_reload_status(
}
config: Final = config_record.param_value
interval_hours: Final = config.get("interval_hours")
interval_hours: Final = cast( # cast-ok: every writer of this key stores `hours: int` or an explicit None
int | None, config.get("interval_hours")
)
if interval_hours is None:
verbose_proxy_logger.info("No interval configured, returning not scheduled")

32
litellm/proxy/spend_tracking/cloudzero_endpoints.py
View file

@ -1,5 +1,11 @@
import json
from typing import Final
from collections.abc import Mapping
from typing import (
TYPE_CHECKING,
Final,
Protocol,
cast, # noqa: TID251 # the config repository's table protocol omits find_first
)
from fastapi import APIRouter, Depends, HTTPException
@ -13,6 +19,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
)
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
from litellm.repositories.config_repository import ConfigRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.types.proxy.cloudzero_endpoints import (
CloudZeroExportRequest,
CloudZeroExportResponse,
@ -22,6 +29,9 @@ from litellm.types.proxy.cloudzero_endpoints import (
CloudZeroSettingsView,
)
if TYPE_CHECKING:
from litellm.proxy.proxy_server import PrismaClient
router: Final = APIRouter()
@ -29,6 +39,18 @@ router: Final = APIRouter()
_sensitive_masker: Final = SensitiveDataMasker()
class _CloudZeroConfigRow(Protocol):
"""The ``LiteLLM_Config`` row holding ``cloudzero_settings``, as this module reads it."""
@property
def param_value(self) -> str | Mapping[str, str] | None: ...
def _config_table(prisma_client: "PrismaClient") -> TableActions[_CloudZeroConfigRow]:
repository_table: Final = ConfigRepository(prisma_client).table
return cast(TableActions[_CloudZeroConfigRow], repository_table) # cast-ok: repo protocol omits find_first
async def _set_cloudzero_settings(api_key: str, connection_id: str, timezone: str):
"""
Store CloudZero settings in the database with encrypted API key.
@ -82,9 +104,7 @@ async def _get_cloudzero_settings():
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
cloudzero_config: Final = await ConfigRepository(prisma_client).table.find_first(
where={"param_name": "cloudzero_settings"}
)
cloudzero_config: Final = await _config_table(prisma_client).find_first(where={"param_name": "cloudzero_settings"})
if cloudzero_config is None or cloudzero_config.param_value is None:
return {}
@ -268,7 +288,7 @@ async def is_cloudzero_setup_in_db() -> bool:
return False
# Check for CloudZero settings in database
cloudzero_config: Final = await ConfigRepository(prisma_client).table.find_first(
cloudzero_config: Final = await _config_table(prisma_client).find_first(
where={"param_name": "cloudzero_settings"}
)
@ -530,7 +550,7 @@ async def delete_cloudzero_settings(
)
# Check if CloudZero settings exist
cloudzero_config: Final = await ConfigRepository(prisma_client).table.find_first(
cloudzero_config: Final = await _config_table(prisma_client).find_first(
where={"param_name": "cloudzero_settings"}
)

55
litellm/proxy/spend_tracking/spend_management_endpoints.py
View file

@ -4,10 +4,22 @@ import json
import os
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple, Protocol, TypedDict, TypeVar
from typing import (
TYPE_CHECKING,
Annotated,
Any,
Final,
Literal,
NamedTuple,
Protocol,
TypedDict,
TypeVar,
cast, # noqa: TID251 # prisma group_by returns untyped aggregate mappings
)
import fastapi
from fastapi import APIRouter, Depends, HTTPException, Request, status
from typing_extensions import ReadOnly
import litellm
from litellm._logging import verbose_proxy_logger
@ -23,6 +35,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
get_spend_by_team_and_customer,
)
from litellm.proxy.utils import handle_exception_on_proxy
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import SpendLogsRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.verification_token_repository import (
@ -30,6 +43,8 @@ from litellm.repositories.verification_token_repository import (
)
if TYPE_CHECKING:
from prisma import models as prisma_models
from litellm.proxy.proxy_server import PrismaClient
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
else:
@ -139,6 +154,18 @@ class _SessionSpendRow(TypedDict):
mcp_tool_call_spend: float
class _SpendSumAggregate(TypedDict, total=False):
spend: ReadOnly[float]
class _SpendGroupByRow(TypedDict):
api_key: ReadOnly[str]
user: ReadOnly[str | None]
model: ReadOnly[str]
startTime: ReadOnly[object]
_sum: ReadOnly[_SpendSumAggregate]
async def _query_raw(prisma_client: PrismaClient, sql_query: str, *args: object) -> Sequence[_RowT]:
"""Run a raw read query and return its rows as the row type the caller declares."""
return await prisma_client.db.query_raw(sql_query, *args)
@ -149,24 +176,6 @@ async def _query_raw_or_none(prisma_client: PrismaClient, sql_query: str, *args:
return await _query_raw(prisma_client, sql_query, *args)
class _SpendLogsTable(Protocol):
"""The subset of the Prisma spend-logs table API this module uses."""
async def find_many(
self, *, where: Mapping[str, object], order: Mapping[str, str]
) -> Sequence[_SupportsModelDump]: ...
async def find_unique(
self, *, where: Mapping[str, object], include: None = None
) -> _SpendLogOwnershipRow | None: ...
async def count(self, *, where: Mapping[str, object]) -> int: ...
async def group_by(
self, *, by: Sequence[str], where: Mapping[str, object], count: Mapping[str, bool]
) -> Sequence[_SessionCountRow]: ...
class _TeamTable(Protocol):
"""The subset of the Prisma team table API this module uses."""
@ -183,7 +192,7 @@ class _VerificationTokenTable(Protocol):
async def update_many(self, *, data: Mapping[str, float], where: Mapping[str, object]) -> int: ...
def _spend_logs_table(prisma_client: PrismaClient) -> _SpendLogsTable:
def _spend_logs_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_SpendLogs"]:
return SpendLogsRepository(prisma_client).table
@ -221,11 +230,12 @@ async def _count_logs_per_session(
prisma_client: PrismaClient, session_ids: Sequence[str | None]
) -> Sequence[_SessionCountRow]:
"""Count spend log rows per session for the given session ids."""
return await _spend_logs_table(prisma_client).group_by(
rows: Final = await _spend_logs_table(prisma_client).group_by(
by=["session_id"],
where={"session_id": {"in": session_ids}},
count={"session_id": True},
)
return cast(Sequence[_SessionCountRow], rows) # cast-ok: group_by(count=) shape is fixed by the by/count args
async def _find_team_row(prisma_client: PrismaClient, team_id: str) -> _SupportsModelDump | None:
@ -2974,8 +2984,9 @@ async def view_spend_logs(
)
if isinstance(response, list) and len(response) > 0 and isinstance(response[0], dict):
spend_rows: Final = cast(Sequence[_SpendGroupByRow], response) # cast-ok: by/sum fix the shape
result: Final[dict] = {}
for record in response:
for record in spend_rows:
dt_object = datetime.strptime(str(record["startTime"]), "%Y-%m-%dT%H:%M:%S.%fZ")
date = dt_object.date()
if date not in result:

30
litellm/proxy/spend_tracking/vantage_endpoints.py
View file

@ -1,5 +1,11 @@
import json
from typing import Final
from collections.abc import Mapping
from typing import (
TYPE_CHECKING,
Final,
Protocol,
cast, # noqa: TID251 # the config repository's table protocol omits find_first
)
from fastapi import APIRouter, Depends, HTTPException
@ -14,6 +20,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
)
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
from litellm.repositories.config_repository import ConfigRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.types.proxy.vantage_endpoints import (
VantageDryRunRequest,
VantageExportRequest,
@ -24,6 +31,9 @@ from litellm.types.proxy.vantage_endpoints import (
VantageSettingsView,
)
if TYPE_CHECKING:
from litellm.proxy.proxy_server import PrismaClient
router: Final = APIRouter()
_sensitive_masker: Final = SensitiveDataMasker()
@ -31,6 +41,18 @@ _sensitive_masker: Final = SensitiveDataMasker()
VANTAGE_SETTINGS_PARAM_NAME: Final = "vantage_settings"
class _VantageConfigRow(Protocol):
"""The ``LiteLLM_Config`` row holding ``vantage_settings``, as this module reads it."""
@property
def param_value(self) -> str | Mapping[str, str] | None: ...
def _config_table(prisma_client: "PrismaClient") -> TableActions[_VantageConfigRow]:
repository_table: Final = ConfigRepository(prisma_client).table
return cast(TableActions[_VantageConfigRow], repository_table) # cast-ok: repo protocol omits find_first
def _get_registered_vantage_logger():
"""Return the VantageLogger already registered in litellm.callbacks, if any."""
from litellm.integrations.vantage.vantage_logger import VantageLogger
@ -82,7 +104,7 @@ async def _get_vantage_settings():
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
vantage_config: Final = await ConfigRepository(prisma_client).table.find_first(
vantage_config: Final = await _config_table(prisma_client).find_first(
where={"param_name": VANTAGE_SETTINGS_PARAM_NAME}
)
if vantage_config is None or vantage_config.param_value is None:
@ -251,7 +273,7 @@ async def is_vantage_setup_in_db() -> bool:
if prisma_client is None:
return False
vantage_config: Final = await ConfigRepository(prisma_client).table.find_first(
vantage_config: Final = await _config_table(prisma_client).find_first(
where={"param_name": VANTAGE_SETTINGS_PARAM_NAME}
)
@ -525,7 +547,7 @@ async def delete_vantage_settings(
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
vantage_config: Final = await ConfigRepository(prisma_client).table.find_first(
vantage_config: Final = await _config_table(prisma_client).find_first(
where={"param_name": VANTAGE_SETTINGS_PARAM_NAME}
)

58
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
View file

@ -4,7 +4,12 @@ import json
import os
from collections import Counter
from collections.abc import Mapping
from typing import Any, Final, Protocol, TypeVar
from typing import (
Any,
Final,
Protocol,
cast, # noqa: TID251 # prisma types Json columns as fields.Json but de-serializes them to plain python on read
)
from urllib.parse import urlparse
from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile
@ -25,6 +30,7 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attributio
from litellm.proxy.utils import invalidate_config_param
from litellm.repositories.config_repository import ConfigRepository
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
SSOConfigRepository,
UISettingsRepository,
@ -37,29 +43,16 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
router: Final = APIRouter()
_DbRecordT: Final = TypeVar("_DbRecordT", covariant=True)
class _PrismaTableActions(Protocol[_DbRecordT]):
async def find_unique(self, where: Mapping[str, object]) -> _DbRecordT | None: ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _DbRecordT: ...
async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> _DbRecordT: ...
class _SsoSettingsMappingRow(Protocol):
@property
def sso_settings(self) -> Mapping[str, object] | None: ...
class _HasSsoSettingsMappingTable(Protocol):
@property
def table(self) -> _PrismaTableActions[_SsoSettingsMappingRow]: ...
def _sso_settings_mapping_db(repo: _HasSsoSettingsMappingTable) -> _PrismaTableActions[_SsoSettingsMappingRow]:
return repo.table
def _sso_settings_mapping_db(repo: SSOConfigRepository) -> TableActions[_SsoSettingsMappingRow]:
return cast( # cast-ok: prisma types Json columns as str; the client hands back the deserialized value
"TableActions[_SsoSettingsMappingRow]", repo.table
)
class _StoredSsoSettingsRow(Protocol):
@ -67,12 +60,7 @@ class _StoredSsoSettingsRow(Protocol):
def sso_settings(self) -> object: ...
class _HasStoredSsoSettingsTable(Protocol):
@property
def table(self) -> _PrismaTableActions[_StoredSsoSettingsRow]: ...
def _stored_sso_settings_db(repo: _HasStoredSsoSettingsTable) -> _PrismaTableActions[_StoredSsoSettingsRow]:
def _stored_sso_settings_db(repo: SSOConfigRepository) -> TableActions[_StoredSsoSettingsRow]:
return repo.table
@ -81,13 +69,10 @@ class _UiSettingsRow(Protocol):
def ui_settings(self) -> str | Mapping[str, JsonValue] | None: ...
class _HasUiSettingsTable(Protocol):
@property
def table(self) -> _PrismaTableActions[_UiSettingsRow]: ...
def _ui_settings_db(repo: _HasUiSettingsTable) -> _PrismaTableActions[_UiSettingsRow]:
return repo.table
def _ui_settings_db(repo: UISettingsRepository) -> TableActions[_UiSettingsRow]:
return cast( # cast-ok: prisma types Json columns as str; the client hands back the deserialized value
"TableActions[_UiSettingsRow]", repo.table
)
class _ConfigParamRow(Protocol):
@ -95,13 +80,10 @@ class _ConfigParamRow(Protocol):
def param_value(self) -> str | Mapping[str, object] | None: ...
class _HasConfigParamTable(Protocol):
@property
def table(self) -> _PrismaTableActions[_ConfigParamRow]: ...
def _config_param_db(repo: _HasConfigParamTable) -> _PrismaTableActions[_ConfigParamRow]:
return repo.table
def _config_param_db(repo: ConfigRepository) -> TableActions[_ConfigParamRow]:
return cast( # cast-ok: prisma's LiteLLM_Config actions object, whose Json column parses to a mapping
"TableActions[_ConfigParamRow]", repo.table
)
# Maps each UIThemeConfig field to the env var the UI branding path reads it

70
litellm/proxy/utils.py
View file

@ -177,6 +177,7 @@ from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams
if TYPE_CHECKING:
from mcp.types import CallToolResult
from opentelemetry.trace import Span as _Span
from prisma import models as prisma_models
from prisma.actions import LiteLLM_DeprecatedVerificationTokenActions
from prisma.client import TransactionManager
from prisma.models import LiteLLM_DeprecatedVerificationToken
@ -186,6 +187,7 @@ if TYPE_CHECKING:
from litellm.models.team import LiteLLM_TeamTableCachedObj
from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction
from litellm.repositories.prisma_protocols import TableActions
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline
Span = _Span | object
@ -3266,7 +3268,10 @@ async def prefetch_config_params(prisma_client: "PrismaClient | None", param_nam
if not param_names:
return
try:
rows: Final = await ConfigRepository(prisma_client).table.find_many(where={"param_name": {"in": param_names}})
config_table: Final = cast( # cast-ok: ConfigRepository.table is prisma's litellm_config actions object
"TableActions[prisma_models.LiteLLM_Config]", ConfigRepository(prisma_client).table
)
rows: Final = await config_table.find_many(where={"param_name": {"in": param_names}})
except Exception as e:
verbose_proxy_logger.debug(
"prefetch_config_params failed, falling through to per-param queries: %s",
@ -3555,8 +3560,8 @@ class PrismaClient:
return hashed_token
def jsonify_object(self, data: dict) -> dict:
db_data: Final = copy.deepcopy(data)
def jsonify_object(self, data: Mapping[str, object]) -> dict[str, object]:
db_data: Final[dict[str, object]] = copy.deepcopy(dict(data))
for k, v in db_data.items():
if isinstance(v, dict):
@ -3690,7 +3695,10 @@ class PrismaClient:
elif table_name == "keys":
return await VerificationTokenRepository(self).table.find_first(where={key: value})
elif table_name == "config":
return await ConfigRepository(self).table.find_first(where={key: value})
config_table: Final = cast( # cast-ok: ConfigRepository.table is prisma's litellm_config actions object
"TableActions[prisma_models.LiteLLM_Config]", ConfigRepository(self).table
)
return await config_table.find_first(where={key: value})
elif table_name == "spend":
return await self.db.l.find_first(where={key: value})
return None
@ -3793,9 +3801,9 @@ class PrismaClient:
self,
token: str | list | None = None,
user_id: str | None = None,
user_id_list: list | None = None,
user_id_list: Sequence[str] | None = None,
team_id: str | None = None,
team_id_list: list | None = None,
team_id_list: Sequence[str] | None = None,
key_val: dict | None = None,
table_name: Literal[
"user", "key", "config", "spend", "enduser", "budget", "team", "user_notification", "combined_view"
@ -3878,14 +3886,14 @@ class PrismaClient:
if isinstance(r.expires, datetime):
r.expires = r.expires.isoformat()
elif query_type == "find_all":
where_filter: Final[dict] = {}
where_filter: Final[dict[str, dict[str, Sequence[str]]]] = {}
if token is not None:
where_filter["token"] = {}
if isinstance(token, str):
token = _hash_token_if_needed(token=token)
where_filter["token"]["in"] = [token]
elif isinstance(token, list):
hashed_tokens: Final = []
hashed_tokens: Final[list[str]] = []
for t in token:
assert isinstance(t, str)
if t.startswith("sk-"):
@ -4182,7 +4190,7 @@ class PrismaClient:
)
raise e
def jsonify_team_object(self, db_data: dict):
def jsonify_team_object(self, db_data: Mapping[str, object]) -> dict[str, object]:
db_data = self.jsonify_object(data=db_data)
if db_data.get("members_with_roles", None) is not None and isinstance(db_data["members_with_roles"], list):
db_data["members_with_roles"] = json.dumps(db_data["members_with_roles"])
@ -4200,7 +4208,7 @@ class PrismaClient:
)
async def insert_data(
self,
data: dict,
data: Mapping[str, object],
table_name: Literal["user", "key", "config", "spend", "team", "user_notification"],
):
"""
@ -4210,10 +4218,12 @@ class PrismaClient:
try:
verbose_proxy_logger.debug(
"PrismaClient: insert_data: %s",
{**data, "token": self.hash_token(token=data["token"])} if data.get("token") is not None else data,
{**data, "token": self.hash_token(token=cast("str", data["token"]))} # cast-ok: a key token is a str
if data.get("token") is not None
else data,
)
if table_name == "key":
token: Final = data["token"]
token: Final = cast("str", data["token"]) # cast-ok: the key table's token column is a str
hashed_token: Final = self.hash_token(token=token)
db_data = self.jsonify_object(data=data)
db_data["token"] = hashed_token
@ -4348,14 +4358,14 @@ class PrismaClient:
async def update_data(
self,
token: str | None = None,
data: dict = {},
data: Mapping[str, object] = {},
data_list: list | None = None,
user_id: str | None = None,
team_id: str | None = None,
query_type: Literal["update", "update_many"] = "update",
table_name: Literal["user", "key", "config", "spend", "team", "enduser", "budget"] | None = None,
update_key_values: dict | None = None,
update_key_values_custom_query: dict | None = None,
update_key_values: dict[str, object] | None = None,
update_key_values_custom_query: dict[str, object] | None = None,
):
"""
Update existing data
@ -4381,14 +4391,14 @@ class PrismaClient:
try:
_data = response.model_dump()
except Exception:
_data = response.dict()
_data = response.dict() # pyright: ignore[reportDeprecated] # pydantic-v1 row fallback
return {"token": token, "data": _data}
elif user_id is not None or (table_name is not None and table_name == "user") and query_type == "update":
"""
If data['spend'] + data['user'], update the user table with spend info as well
"""
if user_id is None:
user_id = db_data["user_id"]
user_id = cast("str", db_data["user_id"]) # cast-ok: the user table's user_id column is a str
if update_key_values is None:
if update_key_values_custom_query is not None:
update_key_values = update_key_values_custom_query
@ -4410,7 +4420,7 @@ class PrismaClient:
If data['spend'] + data['user'], update the user table with spend info as well
"""
if team_id is None:
team_id = db_data["team_id"]
team_id = cast("str | None", db_data["team_id"]) # cast-ok: team_id column is a nullable str
if update_key_values is None:
update_key_values = db_data
if "team_id" not in db_data and team_id is not None:
@ -4584,8 +4594,8 @@ class PrismaClient:
)
async def delete_data(
self,
tokens: list | None = None,
team_id_list: list | None = None,
tokens: Sequence[str | None] | None = None,
team_id_list: Sequence[str] | None = None,
table_name: Literal["user", "key", "config", "spend", "team"] | None = None,
user_id: str | None = None,
):
@ -4597,14 +4607,14 @@ class PrismaClient:
start_time: Final = time.time()
try:
if tokens is not None and isinstance(tokens, list):
hashed_tokens: Final = []
hashed_tokens: Final[list[str | None]] = []
for token in tokens:
if isinstance(token, str) and token.startswith("sk-"):
hashed_token = self.hash_token(token=token)
else:
hashed_token = token
hashed_tokens.append(hashed_token)
filter_query: dict = {}
filter_query: dict[str, object] = {}
if user_id is not None:
filter_query = {"AND": [{"token": {"in": hashed_tokens}}, {"user_id": user_id}]}
else:
@ -5749,12 +5759,12 @@ class PrismaClient:
limit: int = 100,
offset: int = 0,
status_filter: str | None = None,
):
) -> "Sequence[prisma_models.LiteLLM_HealthCheckTable]":
"""
Get health check history with optional filtering
"""
try:
where_clause: Final = {}
where_clause: Final[dict[str, str]] = {}
if model_name:
where_clause["model_name"] = model_name
if status_filter:
@ -5771,7 +5781,7 @@ class PrismaClient:
verbose_proxy_logger.error("Error getting health check history: %s", e)
return []
async def get_all_latest_health_checks(self):
async def get_all_latest_health_checks(self) -> "Sequence[prisma_models.LiteLLM_HealthCheckTable]":
"""
Get the latest health check for each model.
@ -5949,15 +5959,17 @@ async def migrate_passwords_to_scrypt_async(prisma_client) -> str:
return len(s) == 64 and all(c in "0123456789abcdef" for c in s)
plaintext_users: Final = [
u for u in all_with_pw if u.password and not u.password.startswith("scrypt:") and not _is_sha256_hex(u.password)
(u.user_id, u.password)
for u in all_with_pw
if u.password and not u.password.startswith("scrypt:") and not _is_sha256_hex(u.password)
]
if not plaintext_users:
return "No plaintext passwords found"
for user in plaintext_users:
for user_id, plaintext_password in plaintext_users:
await UserRepository(prisma_client).table.update(
where={"user_id": user.user_id},
data={"password": hash_password(user.password)},
where={"user_id": user_id},
data={"password": hash_password(plaintext_password)},
)
return f"Migrated {len(plaintext_users)} plaintext passwords to scrypt"

13
litellm/proxy/vector_store_endpoints/endpoints.py
View file

@ -1,4 +1,9 @@
from typing import Annotated, Any, Final
from typing import (
Annotated,
Any, # noqa: TID251 # jsonify_object in proxy/utils.py is annotated with a bare dict
Final,
cast, # noqa: TID251 # jsonify_object in proxy/utils.py is annotated with a bare dict
)
from fastapi import APIRouter, Depends, HTTPException, Request, Response
@ -591,7 +596,11 @@ async def index_create(
index_data: Final = index_create_request.model_dump(exclude_none=True)
index_data["created_by"] = user_api_key_dict.user_id
index_data["updated_by"] = user_api_key_dict.user_id
new_index = await ManagedVectorStoreIndexRepository(prisma_client).table.create(data=jsonify_object(index_data))
new_index = await ManagedVectorStoreIndexRepository(prisma_client).table.create(
data=cast( # cast-ok: jsonify_object deep-copies a model_dump, so keys are str and values plain objects
"dict[str, object]", jsonify_object(index_data)
)
)
return new_index.model_dump()

24
litellm/proxy/vector_store_endpoints/management_endpoints.py
View file

@ -10,8 +10,7 @@ All /vector_store management endpoints
import copy
import json
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Protocol
from typing import TYPE_CHECKING, Any, Final
from fastapi import APIRouter, Depends, HTTPException
@ -37,6 +36,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helpe
from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user
from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store
from litellm.repositories.model_repository import ModelRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import ManagedVectorStoresRepository
from litellm.secret_managers.main import get_secret
from litellm.types.vector_stores import (
@ -51,17 +51,7 @@ from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
router: Final = APIRouter()
class _VectorStoreTableActions(Protocol):
async def find_unique(self, where: Mapping[str, str]) -> "_VectorStoreRow | None": ...
async def create(self, data: Mapping[str, object]) -> "_VectorStoreRow": ...
async def update(self, where: Mapping[str, str], data: Mapping[str, object]) -> "_VectorStoreRow": ...
async def delete(self, where: Mapping[str, str]) -> "_VectorStoreRow | None": ...
def _vector_store_table(prisma_client: "PrismaClient") -> _VectorStoreTableActions:
def _vector_store_table(prisma_client: "PrismaClient") -> "TableActions[_VectorStoreRow]":
return ManagedVectorStoresRepository(prisma_client).table
@ -277,7 +267,7 @@ async def _resolve_embedding_config_from_db(
if db_model and db_model.litellm_params:
# Extract litellm_params (could be dict or JSON string)
model_params = db_model.litellm_params
if isinstance(model_params, str):
if isinstance(model_params, str): # pyright: ignore[reportUnnecessaryIsInstance] # prisma Json is str
model_params = json.loads(model_params)
# Decrypt values from database (similar to how proxy_server.py does it)
@ -888,6 +878,12 @@ async def update_vector_store(
data=update_data,
)
if updated is None:
raise HTTPException(
status_code=404,
detail=f"Vector store with ID {vector_store_id} not found",
)
updated_vs: Final = _row_to_vector_store(updated)
# Immediately update in-memory registry to keep it in sync

31
litellm/repositories/base_repository.py
View file

@ -8,6 +8,8 @@ from typing import Any, Final, Generic, Protocol, TypeVar, runtime_checkable
from pydantic import BaseModel
from litellm.repositories.prisma_protocols import TableActions
T = TypeVar("T", bound=BaseModel)
@ -49,7 +51,7 @@ class BaseRepository(ABC, Generic[T]):
@property
@abstractmethod
def table(self) -> Any: # any-ok: Prisma table actions are reached through the untyped client wrapper
def table(self) -> TableActions[DbRecord]:
"""Return the Prisma table for this repository."""
...
@ -76,33 +78,28 @@ class BaseRepository(ABC, Generic[T]):
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]] = {}
if where:
kwargs["where"] = where
if skip is not None:
kwargs["skip"] = skip
if take is not None:
kwargs["take"] = take
if order:
kwargs["order"] = order
records: Final = await self.table.find_many(**kwargs)
records: Final = await self.table.find_many(
take=take,
skip=skip,
where=where or None,
order=order or None,
)
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)
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)
return self._to_model(record)
@ -112,7 +109,7 @@ class BaseRepository(ABC, Generic[T]):
record: Final = await self.table.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)

8
litellm/repositories/budget_repository.py
View file

@ -2,17 +2,21 @@
Budget repository for database operations on LiteLLM_BudgetTable.
"""
from typing import Any, Final
from typing import TYPE_CHECKING, Any, Final
from litellm.models.budget import LiteLLM_BudgetTable
from litellm.repositories.base_repository import BaseRepository
from litellm.repositories.prisma_protocols import TableActions
if TYPE_CHECKING:
from prisma import models as prisma_models
class BudgetRepository(BaseRepository[LiteLLM_BudgetTable]):
"""Repository for budget database operations."""
@property
def table(self) -> Any:
def table(self) -> TableActions["prisma_models.LiteLLM_BudgetTable"]:
return self.prisma_client.db.litellm_budgettable
@property

2
litellm/repositories/config_repository.py
View file

@ -77,7 +77,7 @@ class ConfigRepository:
return self.prisma_client.db.litellm_config
@property
def table(self) -> Any:
def table(self) -> _ConfigTable:
return self._config_table
async def get_param(self, param_name: str) -> ConfigParam | None:

53
litellm/repositories/credentials_repository.py
View file

@ -6,54 +6,77 @@ 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, TypeAlias
from litellm.models.credentials import CredentialItem
from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync
from litellm.repositories.base_repository import DbRecord, record_to_dict
from litellm.repositories.prisma_protocols import TableActions
if TYPE_CHECKING:
from prisma import models as prisma_models
_CredentialsTable: TypeAlias = TableActions[prisma_models.LiteLLM_CredentialsTable]
class _PrismaCredentialsDb(Protocol):
@property
def litellm_credentialstable(self) -> "_CredentialsTable": ...
class _PrismaClientView(Protocol):
@property
def db(self) -> _PrismaCredentialsDb: ...
class CredentialsRepository:
"""Repository for credentials database operations, keyed by credential name."""
def __init__(self, prisma_client: Any):
def __init__(self, prisma_client: Any): # any-ok: PrismaClient is an untyped runtime wrapper
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
client: Final[_PrismaClientView] = self._prisma_client
return client
@property
def table(self) -> Any:
def table(self) -> "_CredentialsTable":
return wrap_table_actions_for_config_sync(
actions=self.prisma_client.db.litellm_credentialstable,
table_name="litellm_credentialstable",
)
@staticmethod
def _to_model(record: Any) -> CredentialItem | None:
def _to_model(record: DbRecord | None) -> CredentialItem | None:
if record is None:
return None
data: Final = record.dict() if hasattr(record, "dict") else dict(record)
return CredentialItem(
credential_name=data["credential_name"],
credential_values=data.get("credential_values") or {},
credential_info=data.get("credential_info") or {},
data: Final = record_to_dict(record)
return CredentialItem.model_validate(
{
"credential_name": data["credential_name"],
"credential_values": data.get("credential_values") or {},
"credential_info": data.get("credential_info") or {},
}
)
async def find_all(self) -> Any:
async def find_all(self) -> Sequence["prisma_models.LiteLLM_CredentialsTable"]:
return await self.table.find_many()
async def create(self, data: dict[str, Any]) -> Any:
async def create(self, data: Mapping[str, object]) -> "prisma_models.LiteLLM_CredentialsTable":
return await self.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})
return self._to_model(record)
async def update_by_name(self, credential_name: str, data: dict[str, Any]) -> Any:
async def update_by_name(
self, credential_name: str, data: Mapping[str, object]
) -> "prisma_models.LiteLLM_CredentialsTable | None":
return await self.table.update(where={"credential_name": credential_name}, data=data)
async def delete_by_name(self, credential_name: str) -> Any:
async def delete_by_name(self, credential_name: str) -> "prisma_models.LiteLLM_CredentialsTable | None":
return await self.table.delete(where={"credential_name": credential_name})

42
litellm/repositories/model_repository.py
View file

@ -3,8 +3,8 @@ Model repository for database operations on LiteLLM_ProxyModelTable.
"""
import json
from collections.abc import Awaitable, Mapping, Sequence
from typing import Any, Final, Protocol
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Protocol
from litellm.models.model import LiteLLM_ProxyModelTable
from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync
@ -12,25 +12,21 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
)
from litellm.repositories.base_repository import BaseRepository, DbRecord
from litellm.repositories.base_repository import BaseRepository
from litellm.repositories.prisma_protocols import TableActions
if TYPE_CHECKING:
from prisma import models as prisma_models
class _PrismaModelDb(Protocol):
litellm_proxymodeltable: object
@property
def litellm_proxymodeltable(self) -> TableActions["prisma_models.LiteLLM_ProxyModelTable"]: ...
class _PrismaClientView(Protocol):
db: _PrismaModelDb
class _ProxyModelActions(Protocol):
"""Prisma table actions used by :class:`ModelRepository`."""
def find_many(self, *, where: Mapping[str, object] | None = None) -> Awaitable[Sequence[DbRecord]]: ...
def create(self, *, data: Mapping[str, object]) -> Awaitable[DbRecord]: ...
def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[DbRecord | None]: ...
@property
def db(self) -> _PrismaModelDb: ...
class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
@ -41,17 +37,13 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
self._encryption_key = encryption_key
@property
def table(self) -> Any:
def table(self) -> TableActions["prisma_models.LiteLLM_ProxyModelTable"]:
client: Final[_PrismaClientView] = self.prisma_client
return wrap_table_actions_for_config_sync(
actions=client.db.litellm_proxymodeltable,
table_name="litellm_proxymodeltable",
)
@property
def _model_table(self) -> _ProxyModelActions:
return self.table
@property
def model_class(self) -> type[LiteLLM_ProxyModelTable]:
return LiteLLM_ProxyModelTable
@ -100,17 +92,17 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
async def find_by_name(self, model_name: str) -> list[LiteLLM_ProxyModelTable]:
"""Find models by name."""
records: Final = await self._model_table.find_many(where={"model_name": model_name})
records: Final = await self.table.find_many(where={"model_name": model_name})
return self._to_model_list(records)
async def find_all(self) -> list[LiteLLM_ProxyModelTable]:
"""Find all models."""
records: Final = await self._model_table.find_many()
records: Final = await self.table.find_many()
return self._to_model_list(records)
async def find_unblocked(self) -> list[LiteLLM_ProxyModelTable]:
"""Find all models that are not blocked."""
records: Final = await self._model_table.find_many(where={"blocked": False})
records: Final = await self.table.find_many(where={"blocked": False})
return self._to_model_list(records)
async def find_by_team_id(self, team_id: str) -> list[LiteLLM_ProxyModelTable]:
@ -147,7 +139,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
if model_info is not None:
data["model_info"] = json.dumps(model_info)
record: Final = await self._model_table.create(data=data)
record: Final = await self.table.create(data=data)
model: Final = self._to_model(record)
assert model is not None
return model
@ -173,7 +165,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
if blocked is not None:
data["blocked"] = blocked
record: Final = await self._model_table.update(where={"model_id": model_id}, data=data)
record: Final = await self.table.update(where={"model_id": model_id}, data=data)
return self._to_model(record)
async def delete_model(self, model_id: str) -> LiteLLM_ProxyModelTable | None:

8
litellm/repositories/object_permission_repository.py
View file

@ -2,17 +2,21 @@
ObjectPermission repository for database operations on LiteLLM_ObjectPermissionTable.
"""
from typing import Any, Final
from typing import TYPE_CHECKING, Any, Final
from litellm.models.object_permission import LiteLLM_ObjectPermissionTable
from litellm.repositories.base_repository import BaseRepository
from litellm.repositories.prisma_protocols import TableActions
if TYPE_CHECKING:
from prisma import models as prisma_models
class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]):
"""Repository for object permission database operations."""
@property
def table(self) -> Any:
def table(self) -> TableActions["prisma_models.LiteLLM_ObjectPermissionTable"]:
return self.prisma_client.db.litellm_objectpermissiontable
@property

8
litellm/repositories/organization_repository.py
View file

@ -2,17 +2,21 @@
Organization repository for database operations on LiteLLM_OrganizationTable.
"""
from typing import Any, Final
from typing import TYPE_CHECKING, Any, Final
from litellm.models.organization import LiteLLM_OrganizationTable
from litellm.repositories.base_repository import BaseRepository
from litellm.repositories.prisma_protocols import TableActions
if TYPE_CHECKING:
from prisma import models as prisma_models
class OrganizationRepository(BaseRepository[LiteLLM_OrganizationTable]):
"""Repository for organization database operations."""
@property
def table(self) -> Any:
def table(self) -> TableActions["prisma_models.LiteLLM_OrganizationTable"]:
return self.prisma_client.db.litellm_organizationtable
@property

87
litellm/repositories/prisma_protocols.py
View file

@ -12,6 +12,93 @@ from typing import Protocol, TypeVar
RowT_co = TypeVar("RowT_co", covariant=True)
class TableActions(Protocol[RowT_co]):
"""The prisma-client-py per-model action surface, keyed to the row it returns.
Query inputs stay `Mapping[str, object]` rather than the generated
`types.*` TypedDicts so callers can keep passing plain dicts, while every
result carries the row type the repository is bound to.
"""
async def find_unique(
self, where: Mapping[str, object], include: Mapping[str, object] | None = None
) -> RowT_co | None: ...
async def find_first(
self,
skip: int | None = None,
where: Mapping[str, object] | None = None,
cursor: Mapping[str, object] | None = None,
include: Mapping[str, object] | None = None,
order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None,
distinct: Sequence[str] | None = None,
) -> RowT_co | None: ...
async def find_many(
self,
take: int | None = None,
skip: int | None = None,
where: Mapping[str, object] | None = None,
cursor: Mapping[str, object] | None = None,
include: Mapping[str, object] | None = None,
order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None,
distinct: Sequence[str] | None = None,
) -> Sequence[RowT_co]: ...
async def create(self, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> RowT_co: ...
async def create_many(
self, data: Sequence[Mapping[str, object]], *, skip_duplicates: bool | None = None
) -> int: ...
async def upsert(
self,
where: Mapping[str, object],
data: Mapping[str, object],
include: Mapping[str, object] | None = None,
) -> RowT_co: ...
async def update(
self,
data: Mapping[str, object],
where: Mapping[str, object],
include: Mapping[str, object] | None = None,
) -> RowT_co | None: ...
async def update_many(self, data: Mapping[str, object], where: Mapping[str, object]) -> int: ...
async def delete(
self, where: Mapping[str, object], include: Mapping[str, object] | None = None
) -> RowT_co | None: ...
async def delete_many(self, where: Mapping[str, object] | None = None) -> int: ...
async def count(
self,
select: None = None,
take: int | None = None,
skip: int | None = None,
where: Mapping[str, object] | None = None,
cursor: Mapping[str, object] | None = None,
) -> int: ...
async def group_by(
self,
by: Sequence[str],
*,
where: Mapping[str, object] | None = None,
take: int | None = None,
skip: int | None = None,
order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None,
having: Mapping[str, object] | None = None,
count: bool | Mapping[str, object] | None = None,
sum: bool | Mapping[str, object] | None = None,
avg: bool | Mapping[str, object] | None = None,
min: bool | Mapping[str, object] | None = None,
max: bool | Mapping[str, object] | None = None,
) -> Sequence[Mapping[str, object]]: ...
class PrismaRecord(Protocol):
def dict(self) -> Mapping[str, object]: ...

8
litellm/repositories/project_repository.py
View file

@ -2,17 +2,21 @@
Project repository for database operations on LiteLLM_ProjectTable.
"""
from typing import Any, Final
from typing import TYPE_CHECKING, Any, Final
from litellm.models.project import LiteLLM_ProjectTable
from litellm.repositories.base_repository import BaseRepository
from litellm.repositories.prisma_protocols import TableActions
if TYPE_CHECKING:
from prisma import models as prisma_models
class ProjectRepository(BaseRepository[LiteLLM_ProjectTable]):
"""Repository for project database operations."""
@property
def table(self) -> Any:
def table(self) -> TableActions["prisma_models.LiteLLM_ProjectTable"]:
return self.prisma_client.db.litellm_projecttable
@property

116
litellm/repositories/table_repositories.py
View file

@ -7,12 +7,16 @@ These are thin wrappers for tables that do not (yet) need domain-specific query
methods; richer repositories live in their own modules.
"""
from typing import Any
from typing import TYPE_CHECKING, Any, Final, Generic
from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync
from litellm.repositories.prisma_protocols import RowT_co, TableActions
if TYPE_CHECKING:
from prisma import models as prisma_models # noqa: F401 # used by quoted base-class subscripts
class PrismaTableRepository:
class PrismaTableRepository(Generic[RowT_co]):
"""Base for repositories that expose a single Prisma table."""
table_name: str
@ -27,208 +31,206 @@ class PrismaTableRepository:
return self._prisma_client
@property
def table(self) -> Any: # any-ok: Prisma table actions are reached through the untyped client wrapper
return wrap_table_actions_for_config_sync(
actions=getattr(self.prisma_client.db, self.table_name),
table_name=self.table_name,
)
def table(self) -> TableActions[RowT_co]:
actions: Final[TableActions[RowT_co]] = getattr(self.prisma_client.db, self.table_name)
return wrap_table_actions_for_config_sync(actions=actions, table_name=self.table_name)
class PolicyRepository(PrismaTableRepository):
class PolicyRepository(PrismaTableRepository["prisma_models.LiteLLM_PolicyTable"]):
table_name = "litellm_policytable"
class AgentsRepository(PrismaTableRepository):
class AgentsRepository(PrismaTableRepository["prisma_models.LiteLLM_AgentsTable"]):
table_name = "litellm_agentstable"
class ObjectPermissionRepository(PrismaTableRepository):
class ObjectPermissionRepository(PrismaTableRepository["prisma_models.LiteLLM_ObjectPermissionTable"]):
table_name = "litellm_objectpermissiontable"
class GuardrailsRepository(PrismaTableRepository):
class GuardrailsRepository(PrismaTableRepository["prisma_models.LiteLLM_GuardrailsTable"]):
table_name = "litellm_guardrailstable"
class MCPServerRepository(PrismaTableRepository):
class MCPServerRepository(PrismaTableRepository["prisma_models.LiteLLM_MCPServerTable"]):
table_name = "litellm_mcpservertable"
class ManagedObjectRepository(PrismaTableRepository):
class ManagedObjectRepository(PrismaTableRepository["prisma_models.LiteLLM_ManagedObjectTable"]):
table_name = "litellm_managedobjecttable"
class OrganizationMembershipRepository(PrismaTableRepository):
class OrganizationMembershipRepository(PrismaTableRepository["prisma_models.LiteLLM_OrganizationMembership"]):
table_name = "litellm_organizationmembership"
class SpendLogsRepository(PrismaTableRepository):
class SpendLogsRepository(PrismaTableRepository["prisma_models.LiteLLM_SpendLogs"]):
table_name = "litellm_spendlogs"
class ClaudeCodePluginRepository(PrismaTableRepository):
class ClaudeCodePluginRepository(PrismaTableRepository["prisma_models.LiteLLM_ClaudeCodePluginTable"]):
table_name = "litellm_claudecodeplugintable"
class TeamMembershipRepository(PrismaTableRepository):
class TeamMembershipRepository(PrismaTableRepository["prisma_models.LiteLLM_TeamMembership"]):
table_name = "litellm_teammembership"
class EndUserRepository(PrismaTableRepository):
class EndUserRepository(PrismaTableRepository["prisma_models.LiteLLM_EndUserTable"]):
table_name = "litellm_endusertable"
class ManagedVectorStoresRepository(PrismaTableRepository):
class ManagedVectorStoresRepository(PrismaTableRepository["prisma_models.LiteLLM_ManagedVectorStoresTable"]):
table_name = "litellm_managedvectorstorestable"
class MCPUserCredentialsRepository(PrismaTableRepository):
class MCPUserCredentialsRepository(PrismaTableRepository["prisma_models.LiteLLM_MCPUserCredentials"]):
table_name = "litellm_mcpusercredentials"
class MCPServerOAuthClientRepository(PrismaTableRepository):
class MCPServerOAuthClientRepository(PrismaTableRepository["prisma_models.LiteLLM_MCPServerOAuthClient"]):
table_name = "litellm_mcpserveroauthclient"
class PromptRepository(PrismaTableRepository):
class PromptRepository(PrismaTableRepository["prisma_models.LiteLLM_PromptTable"]):
table_name = "litellm_prompttable"
class TagRepository(PrismaTableRepository):
class TagRepository(PrismaTableRepository["prisma_models.LiteLLM_TagTable"]):
table_name = "litellm_tagtable"
class InvitationLinkRepository(PrismaTableRepository):
class InvitationLinkRepository(PrismaTableRepository["prisma_models.LiteLLM_InvitationLink"]):
table_name = "litellm_invitationlink"
class JWTKeyMappingRepository(PrismaTableRepository):
class JWTKeyMappingRepository(PrismaTableRepository["prisma_models.LiteLLM_JWTKeyMapping"]):
table_name = "litellm_jwtkeymapping"
class ManagedFileRepository(PrismaTableRepository):
class ManagedFileRepository(PrismaTableRepository["prisma_models.LiteLLM_ManagedFileTable"]):
table_name = "litellm_managedfiletable"
class MemoryRepository(PrismaTableRepository):
class MemoryRepository(PrismaTableRepository["prisma_models.LiteLLM_MemoryTable"]):
table_name = "litellm_memorytable"
class SearchToolsRepository(PrismaTableRepository):
class SearchToolsRepository(PrismaTableRepository["prisma_models.LiteLLM_SearchToolsTable"]):
table_name = "litellm_searchtoolstable"
class ConfigOverridesRepository(PrismaTableRepository):
class ConfigOverridesRepository(PrismaTableRepository["prisma_models.LiteLLM_ConfigOverrides"]):
table_name = "litellm_configoverrides"
class MCPToolsetRepository(PrismaTableRepository):
class MCPToolsetRepository(PrismaTableRepository["prisma_models.LiteLLM_MCPToolsetTable"]):
table_name = "litellm_mcptoolsettable"
class ToolRepository(PrismaTableRepository):
class ToolRepository(PrismaTableRepository["prisma_models.LiteLLM_ToolTable"]):
table_name = "litellm_tooltable"
class DeletedVerificationTokenRepository(PrismaTableRepository):
class DeletedVerificationTokenRepository(PrismaTableRepository["prisma_models.LiteLLM_DeletedVerificationToken"]):
table_name = "litellm_deletedverificationtoken"
class WorkflowRunRepository(PrismaTableRepository):
class WorkflowRunRepository(PrismaTableRepository["prisma_models.LiteLLM_WorkflowRun"]):
table_name = "litellm_workflowrun"
class ModelTableRepository(PrismaTableRepository):
class ModelTableRepository(PrismaTableRepository["prisma_models.LiteLLM_ModelTable"]):
table_name = "litellm_modeltable"
class AccessGroupRepository(PrismaTableRepository):
class AccessGroupRepository(PrismaTableRepository["prisma_models.LiteLLM_AccessGroupTable"]):
table_name = "litellm_accessgrouptable"
class SSOConfigRepository(PrismaTableRepository):
class SSOConfigRepository(PrismaTableRepository["prisma_models.LiteLLM_SSOConfig"]):
table_name = "litellm_ssoconfig"
class UISettingsRepository(PrismaTableRepository):
class UISettingsRepository(PrismaTableRepository["prisma_models.LiteLLM_UISettings"]):
table_name = "litellm_uisettings"
class DailyGuardrailMetricsRepository(PrismaTableRepository):
class DailyGuardrailMetricsRepository(PrismaTableRepository["prisma_models.LiteLLM_DailyGuardrailMetrics"]):
table_name = "litellm_dailyguardrailmetrics"
class DailyGuardrailUsageUnitsRepository(PrismaTableRepository):
class DailyGuardrailUsageUnitsRepository(PrismaTableRepository["prisma_models.LiteLLM_DailyGuardrailUsageUnits"]):
table_name = "litellm_dailyguardrailusageunits"
class PolicyAttachmentRepository(PrismaTableRepository):
class PolicyAttachmentRepository(PrismaTableRepository["prisma_models.LiteLLM_PolicyAttachmentTable"]):
table_name = "litellm_policyattachmenttable"
class DeletedTeamRepository(PrismaTableRepository):
class DeletedTeamRepository(PrismaTableRepository["prisma_models.LiteLLM_DeletedTeamTable"]):
table_name = "litellm_deletedteamtable"
class SkillsRepository(PrismaTableRepository):
class SkillsRepository(PrismaTableRepository["prisma_models.LiteLLM_SkillsTable"]):
table_name = "litellm_skillstable"
class CacheConfigRepository(PrismaTableRepository):
class CacheConfigRepository(PrismaTableRepository["prisma_models.LiteLLM_CacheConfig"]):
table_name = "litellm_cacheconfig"
class ManagedVectorStoreIndexRepository(PrismaTableRepository):
class ManagedVectorStoreIndexRepository(PrismaTableRepository["prisma_models.LiteLLM_ManagedVectorStoreIndexTable"]):
table_name = "litellm_managedvectorstoreindextable"
class WorkflowMessageRepository(PrismaTableRepository):
class WorkflowMessageRepository(PrismaTableRepository["prisma_models.LiteLLM_WorkflowMessage"]):
table_name = "litellm_workflowmessage"
class DailyTagSpendRepository(PrismaTableRepository):
class DailyTagSpendRepository(PrismaTableRepository["prisma_models.LiteLLM_DailyTagSpend"]):
table_name = "litellm_dailytagspend"
class SpendLogToolIndexRepository(PrismaTableRepository):
class SpendLogToolIndexRepository(PrismaTableRepository["prisma_models.LiteLLM_SpendLogToolIndex"]):
table_name = "litellm_spendlogtoolindex"
class DailyToolSpendRepository(PrismaTableRepository):
class DailyToolSpendRepository(PrismaTableRepository["prisma_models.LiteLLM_DailyToolSpend"]):
table_name = "litellm_dailytoolspend"
class SpendLogGuardrailIndexRepository(PrismaTableRepository):
class SpendLogGuardrailIndexRepository(PrismaTableRepository["prisma_models.LiteLLM_SpendLogGuardrailIndex"]):
table_name = "litellm_spendlogguardrailindex"
class UserNotificationsRepository(PrismaTableRepository):
class UserNotificationsRepository(PrismaTableRepository["prisma_models.LiteLLM_UserNotifications"]):
table_name = "litellm_usernotifications"
class HealthCheckRepository(PrismaTableRepository):
class HealthCheckRepository(PrismaTableRepository["prisma_models.LiteLLM_HealthCheckTable"]):
table_name = "litellm_healthchecktable"
class DeprecatedVerificationTokenRepository(PrismaTableRepository):
class DeprecatedVerificationTokenRepository(PrismaTableRepository["prisma_models.LiteLLM_DeprecatedVerificationToken"]):
table_name = "litellm_deprecatedverificationtoken"
class WorkflowEventRepository(PrismaTableRepository):
class WorkflowEventRepository(PrismaTableRepository["prisma_models.LiteLLM_WorkflowEvent"]):
table_name = "litellm_workflowevent"
class DailyPolicyMetricsRepository(PrismaTableRepository):
class DailyPolicyMetricsRepository(PrismaTableRepository["prisma_models.LiteLLM_DailyPolicyMetrics"]):
table_name = "litellm_dailypolicymetrics"
class AdaptiveRouterStateRepository(PrismaTableRepository):
class AdaptiveRouterStateRepository(PrismaTableRepository["prisma_models.LiteLLM_AdaptiveRouterState"]):
table_name = "litellm_adaptiverouterstate"
class AuditLogRepository(PrismaTableRepository):
class AuditLogRepository(PrismaTableRepository["prisma_models.LiteLLM_AuditLog"]):
table_name = "litellm_auditlog"
class AdaptiveRouterSessionRepository(PrismaTableRepository):
class AdaptiveRouterSessionRepository(PrismaTableRepository["prisma_models.LiteLLM_AdaptiveRouterSession"]):
table_name = "litellm_adaptiveroutersession"

8
litellm/repositories/team_repository.py
View file

@ -5,7 +5,7 @@ Team repository for database operations on LiteLLM_TeamTable.
import json
from collections.abc import Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Final
from pydantic import TypeAdapter
@ -15,9 +15,11 @@ from litellm.repositories.base_repository import (
DbRecord,
record_to_dict,
)
from litellm.repositories.prisma_protocols import TableActions
if TYPE_CHECKING:
from prisma import Prisma
from prisma import models as prisma_models
_MEMBERS_WITH_ROLES_ADAPTER: Final = TypeAdapter(list[Member])
_JSON_ENCODED_TEAM_FIELDS: Final = (
@ -34,11 +36,11 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
"""Repository for team database operations."""
@property
def table(self) -> Any: # any-ok: PrismaClient.db is an untyped runtime wrapper
def table(self) -> TableActions["prisma_models.LiteLLM_TeamTable"]:
return self.prisma_client.db.litellm_teamtable
@property
def deleted_table(self) -> Any: # any-ok: PrismaClient.db is an untyped runtime wrapper
def deleted_table(self) -> TableActions["prisma_models.LiteLLM_DeletedTeamTable"]:
return self.prisma_client.db.litellm_deletedteamtable
@property

7
litellm/repositories/user_banner_repository.py
View file

@ -1,11 +1,14 @@
from typing import Final
from typing import TYPE_CHECKING, Final
from litellm.repositories.table_repositories import PrismaTableRepository
if TYPE_CHECKING:
from prisma import models as prisma_models # noqa: F401 # resolved only from the quoted base-class subscript below
USER_BANNER_ROW_ID: Final = "user_banner"
class UserBannerRepository(PrismaTableRepository):
class UserBannerRepository(PrismaTableRepository["prisma_models.LiteLLM_UISettings"]):
table_name = "litellm_uisettings"
async def get_raw_settings(self) -> object:

8
litellm/repositories/user_repository.py
View file

@ -4,10 +4,14 @@ User repository for database operations on LiteLLM_UserTable.
import json
from collections.abc import Mapping
from typing import Any, Final
from typing import TYPE_CHECKING, Final
from litellm.models.user import LiteLLM_UserTable
from litellm.repositories.base_repository import BaseRepository, DbRecord, record_to_dict
from litellm.repositories.prisma_protocols import TableActions
if TYPE_CHECKING:
from prisma import models as prisma_models
_JSON_ENCODED_COLUMNS: Final = frozenset({"metadata", "model_spend", "model_max_budget"})
@ -16,7 +20,7 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]):
"""Repository for user database operations."""
@property
def table(self) -> Any: # any-ok: Prisma table actions are reached through the untyped client wrapper
def table(self) -> TableActions["prisma_models.LiteLLM_UserTable"]:
return self.prisma_client.db.litellm_usertable
@property

22
litellm/repositories/verification_token_repository.py
View file

@ -3,9 +3,9 @@ VerificationToken repository for database operations on LiteLLM_VerificationToke
"""
import json
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Final
from litellm.models.verification_token import (
LiteLLM_VerificationToken,
@ -15,8 +15,12 @@ from litellm.repositories.base_repository import (
DbRecord,
record_to_dict,
)
from litellm.repositories.prisma_protocols import TableActions
if TYPE_CHECKING:
from prisma.models import (
LiteLLM_DeletedVerificationToken as PrismaDeletedVerificationToken,
)
from prisma.models import (
LiteLLM_VerificationToken as PrismaVerificationToken,
)
@ -45,11 +49,11 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]):
return prisma_client
@property
def table(self) -> Any:
def table(self) -> TableActions["PrismaVerificationToken"]:
return self.prisma_client.db.litellm_verificationtoken
@property
def deleted_table(self) -> Any:
def deleted_table(self) -> TableActions["PrismaDeletedVerificationToken"]:
return self.prisma_client.db.litellm_deletedverificationtoken
@property
@ -79,29 +83,29 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]):
async def find_by_alias(self, key_alias: str) -> LiteLLM_VerificationToken | None:
"""Find a token by key alias."""
records: Final[list[PrismaVerificationToken]] = await self.table.find_many(where={"key_alias": key_alias})
records: Final[Sequence[PrismaVerificationToken]] = await self.table.find_many(where={"key_alias": key_alias})
if records:
return self._to_model(records[0])
return None
async def find_by_user_id(self, user_id: str) -> list[LiteLLM_VerificationToken]:
"""Find all tokens belonging to a user."""
records: Final[list[PrismaVerificationToken]] = await self.table.find_many(where={"user_id": user_id})
records: Final[Sequence[PrismaVerificationToken]] = await self.table.find_many(where={"user_id": user_id})
return self._to_model_list(records)
async def find_by_team_id(self, team_id: str) -> list[LiteLLM_VerificationToken]:
"""Find all tokens belonging to a team."""
records: Final[list[PrismaVerificationToken]] = await self.table.find_many(where={"team_id": team_id})
records: Final[Sequence[PrismaVerificationToken]] = await self.table.find_many(where={"team_id": team_id})
return self._to_model_list(records)
async def find_by_project_id(self, project_id: str) -> list[LiteLLM_VerificationToken]:
"""Find all tokens belonging to a project."""
records: Final[list[PrismaVerificationToken]] = await self.table.find_many(where={"project_id": project_id})
records: Final[Sequence[PrismaVerificationToken]] = await self.table.find_many(where={"project_id": project_id})
return self._to_model_list(records)
async def find_active_tokens(self) -> list[LiteLLM_VerificationToken]:
"""Find all active (non-expired, non-blocked) tokens."""
records: Final[list[PrismaVerificationToken]] = await self.table.find_many(
records: Final[Sequence[PrismaVerificationToken]] = await self.table.find_many(
where={
"blocked": {"not": True},
"OR": [{"expires": None}, {"expires": {"gt": datetime.utcnow()}}],

56
litellm/responses/file_search/emulated_handler.py
View file

@ -15,7 +15,9 @@ import json
import time
import uuid
from collections.abc import Iterable, Sequence
from typing import TYPE_CHECKING, Any, Final, TypeAlias, cast
from typing import TYPE_CHECKING, Any, Final, TypeAlias, cast # noqa: TID251 # see kwargs-ok / cast-ok markers
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._internal_context import is_internal_call
from litellm._logging import verbose_logger
@ -31,6 +33,12 @@ ToolParam: TypeAlias = object
FILE_SEARCH_FUNCTION_NAME: Final = "litellm_file_search"
class FileSearchToolCallArgs(TypedDict):
queries: ReadOnly[NotRequired[object]]
query: ReadOnly[NotRequired[object]]
vector_store_id: ReadOnly[NotRequired[object]]
# ---------------------------------------------------------------------------
# Detection
# ---------------------------------------------------------------------------
@ -175,13 +183,20 @@ async def _run_vector_searches(
# ---------------------------------------------------------------------------
def _get_field(result: object, key: str, default: object = None) -> Any:
def _get_field(result: object, key: str, default: object = None) -> object:
"""Read a field from either a dict/TypedDict or an attribute-based object."""
if isinstance(result, dict):
return result.get(key, default)
return getattr(result, key, default)
def _joined_content_text(result: object) -> str:
"""Concatenate the text of every content chunk on a search result."""
content_items: Final = cast(Iterable[object], _get_field(result, "content") or []) # cast-ok: iterated as today
text_chunks: Final = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items]
return " ".join(t for t in text_chunks if t)
def _format_search_results_as_tool_output(
results: list[VectorStoreSearchResult],
) -> str:
@ -194,9 +209,7 @@ def _format_search_results_as_tool_output(
score = _get_field(result, "score")
file_id = _get_field(result, "file_id")
filename = _get_field(result, "filename")
content_items = _get_field(result, "content") or []
text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items]
text = " ".join(t for t in text_chunks if t)
text = _joined_content_text(result)
header = f"[Result {i}"
if filename:
@ -226,9 +239,7 @@ def _build_search_results_for_include(
formatted: Final[list[dict[str, object]]] = []
for result in results:
file_id = _get_field(result, "file_id") or ""
content_items = _get_field(result, "content") or []
text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items]
text = " ".join(t for t in text_chunks if t)
text = _joined_content_text(result)
formatted.append(
{
"file_id": file_id,
@ -353,14 +364,14 @@ def _synthesize_responses_api_response(
created_at=getattr(original_response, "created_at", int(time.time())),
status="completed",
model=getattr(original_response, "model", ""),
output=cast(list[ResponseOutputItem | dict[str, Any]], synthesized_output),
output=cast(list[ResponseOutputItem | dict[str, object]], synthesized_output), # cast-ok: list is invariant
usage=getattr(original_response, "usage", None),
error=None,
)
if hasattr(original_response, "_hidden_params"):
hidden: Final = dict(getattr(original_response, "_hidden_params") or {})
if first_response is not None and hasattr(first_response, "_hidden_params"):
first_hidden: Final = getattr(first_response, "_hidden_params") or {}
first_hidden: Final[object] = getattr(first_response, "_hidden_params") or {}
first_cost: Final = (
first_hidden.get("response_cost")
if isinstance(first_hidden, dict)
@ -385,9 +396,10 @@ async def _call_aresponses(input, model, tools, **kwargs): # pragma: no cover
def _prepare_emulated_file_search_call(
kwargs: dict[str, Any],
kwargs: dict[str, object],
) -> tuple[bool, dict[str, object]]:
include_items: Final[list[str]] = list(kwargs.get("include") or [])
raw_include: Final = kwargs.get("include") or []
include_items: Final[list[object]] = list(cast(Iterable[object], raw_include)) # cast-ok: iterated as today
include_search_results: Final = "file_search_call.results" in include_items
original_stream: Final = kwargs.get("stream")
@ -413,16 +425,16 @@ def _extract_tool_call_fields(tool_call: object, fallback_call_id: str) -> tuple
return call_id, raw_args
def _resolve_queries_from_args(args: dict[str, Any], input: object) -> list[str]:
def _resolve_queries_from_args(args: FileSearchToolCallArgs, input: object) -> list[str]:
"""Pull the queries list out of parsed tool-call arguments, with backward-compat fallbacks."""
queries_from_call: Final = args.get("queries")
if not queries_from_call:
# Fallback: check for single "query" field (backward compat)
single_query: Final = args.get("query")
return [single_query] if single_query else [str(input)]
return [cast(str, single_query)] if single_query else [str(input)] # cast-ok: model-supplied, as today
if not isinstance(queries_from_call, list):
return [str(queries_from_call)]
return queries_from_call
return cast(list[str], queries_from_call) # cast-ok: model-supplied elements, forwarded unchecked as today
async def _execute_file_search_tool_calls(
@ -440,14 +452,14 @@ async def _execute_file_search_tool_calls(
call_id, raw_args = _extract_tool_call_fields(tool_call, fallback_call_id=file_search_call_id)
try:
args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
args: FileSearchToolCallArgs = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
except json.JSONDecodeError:
args = {}
queries_from_call = _resolve_queries_from_args(args, input)
vs_id_arg = args.get("vector_store_id")
vs_ids_for_call = [vs_id_arg] if vs_id_arg else all_vs_ids
vs_ids_for_call = [cast(str, vs_id_arg)] if vs_id_arg else all_vs_ids # cast-ok: model-supplied, as today
queries, results = await _run_vector_searches(
queries=queries_from_call,
@ -481,7 +493,7 @@ def _build_follow_up_input(
original_input_items: Final[list[object]] = (
list(input) if isinstance(input, (list, tuple)) else [{"role": "user", "content": str(input)}]
)
first_response_output_items: Final[list[Any]] = []
first_response_output_items: Final[list[object]] = []
for _item in first_response.output:
if isinstance(_item, dict):
first_response_output_items.append(_item)
@ -498,7 +510,7 @@ async def aresponses_with_emulated_file_search(
model: str,
tools: Iterable[ToolParam] | None = None,
# Pass-through params — forwarded as-is to the underlying aresponses call
**kwargs: Any,
**kwargs: Any, # kwargs-ok: `object` would surface the caller's partially-unknown dict at its call site
) -> ResponsesAPIResponse:
"""
Emulated file_search for providers that don't support it natively.
@ -507,7 +519,7 @@ async def aresponses_with_emulated_file_search(
runs vector search, and synthesizes an OpenAI-format response.
"""
# Determine whether caller wants search_results populated in the output.
_include_search_results, kwargs = _prepare_emulated_file_search_call(kwargs=kwargs)
_include_search_results, call_kwargs = _prepare_emulated_file_search_call(kwargs=kwargs)
# 1. Replace file_search tools with function tool
transformed_tools, all_vs_ids = _replace_file_search_tools(tools)
@ -524,7 +536,7 @@ async def aresponses_with_emulated_file_search(
input=input,
model=model,
tools=transformed_tools or None,
**kwargs,
**call_kwargs,
),
)
finally:
@ -588,7 +600,7 @@ async def aresponses_with_emulated_file_search(
input=follow_up_input,
model=model,
tools=None, # no tools needed for the answer step
**kwargs,
**call_kwargs,
),
)
finally:

10
litellm/responses/litellm_completion_transformation/custom_tools.py
View file

@ -16,8 +16,8 @@ logic.
"""
import json
from collections.abc import Mapping
from typing import Any, Final
from collections.abc import Mapping, Sequence
from typing import Final
from pydantic import BaseModel, TypeAdapter, ValidationError
@ -29,7 +29,7 @@ from litellm.types.llms.openai import (
_MAX_ARGUMENTS_LEN: Final = 1_000_000
def extract_custom_tool_names(tools: list[Any] | None) -> set[str]:
def extract_custom_tool_names(tools: Sequence[object] | None) -> set[str]:
"""Extract names of tools originally defined as ``type: "custom"``."""
if not tools:
return set()
@ -73,7 +73,7 @@ def build_tool_call_item_kwargs(
arguments_or_input: str,
status: str,
custom_tool_names: set[str],
) -> dict[str, Any]:
) -> dict[str, str]:
"""Build kwargs for an output item dict that is either a ``function_call``
or a ``custom_tool_call`` depending on whether *name* is in
*custom_tool_names*.
@ -86,7 +86,7 @@ def build_tool_call_item_kwargs(
"""
custom: Final = is_custom_tool_call(name, custom_tool_names)
item_type: Final = "custom_tool_call" if custom else "function_call"
kwargs: Final[dict[str, Any]] = {
kwargs: Final[dict[str, str]] = {
"type": item_type,
"id": call_id,
"call_id": call_id,

8
litellm/responses/litellm_completion_transformation/handler.py
View file

@ -2,8 +2,8 @@
Handler for transforming responses api requests to litellm.completion requests
"""
from collections.abc import Coroutine
from typing import Any, Final
from collections.abc import Coroutine, Mapping
from typing import Final
import litellm
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
@ -30,12 +30,12 @@ class LiteLLMCompletionTransformationHandler:
custom_llm_provider: str | None = None,
_is_async: bool = False,
stream: bool | None = None,
extra_headers: dict[str, Any] | None = None,
extra_headers: Mapping[str, object] | None = None,
**kwargs,
) -> (
ResponsesAPIResponse
| BaseResponsesAPIStreamingIterator
| Coroutine[Any, Any, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]
| Coroutine[object, object, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]
):
litellm_completion_request: Final[dict] = (
LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(

6
litellm/responses/litellm_completion_transformation/session_handler.py
View file

@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Any, Final, cast
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import SpendLogsPayload
from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import (
@ -143,7 +143,7 @@ class ResponsesSessionHandler:
model_response: Final = ModelResponse(**_response_output)
for choice in model_response.choices:
if hasattr(choice, "message"):
chat_completion_message_history.append(getattr(choice, "message"))
chat_completion_message_history.append(choice.message)
return chat_completion_message_history
@staticmethod
@ -195,7 +195,7 @@ class ResponsesSessionHandler:
try:
metadata_str: Final = spend_log.get("metadata", "{}")
if isinstance(metadata_str, str):
metadata_dict: Final = json.loads(metadata_str)
metadata_dict: Final[SpendLogsMetadata] = json.loads(metadata_str)
return metadata_dict.get("cold_storage_object_key")
elif isinstance(metadata_str, dict):
return metadata_str.get("cold_storage_object_key")

23
litellm/responses/litellm_completion_transformation/streaming_iterator.py
View file

@ -1,5 +1,6 @@
import time
import uuid
from collections.abc import Sequence
from typing import Any, Final, cast
import litellm
@ -48,14 +49,18 @@ from litellm.types.utils import (
)
def _index_of_output_item_type(items: Sequence[object], item_type: str) -> int | None:
return next(
(index for index, item in enumerate(items) if getattr(item, "type", None) == item_type),
None,
)
def _output_items_with_id(items: tuple[Any, ...], item_type: str, item_id: str | None) -> tuple[Any, ...]:
if item_id is None:
return items
target_index: Final = next(
(index for index, item in enumerate(items) if getattr(item, "type", None) == item_type),
None,
)
target_index: Final = _index_of_output_item_type(items, item_type)
if target_index is None:
return items
@ -86,7 +91,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self.litellm_metadata: dict | None = litellm_metadata or {}
# Store lightweight dict snapshots for stream_chunk_builder to reduce
# repeated Pydantic attribute access in end-of-stream assembly.
self.collected_chat_completion_chunks: list[dict[str, Any]] = []
self.collected_chat_completion_chunks: list[dict[str, object]] = []
self.finished: bool = False
self.litellm_logging_obj = litellm_custom_stream_wrapper.logging_obj
self.sent_response_created_event: bool = False
@ -98,7 +103,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self.sent_output_item_done_event: bool = False
self.sent_annotation_events: bool = False
self.litellm_model_response: ModelResponse | TextCompletionResponse | None = None
self.completed_response: Any = None
self.completed_response = None
self.final_text: str = ""
self._cached_item_id: str | None = None
self._cached_response_id: str | None = None
@ -123,7 +128,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._reasoning_done_emitted = False
self._reasoning_item_id: str | None = None
self._accumulated_reasoning_content_parts: list[str] = []
self._accumulated_provider_specific_fields: dict[str, Any] = {}
self._accumulated_provider_specific_fields: dict[str, object] = {}
self._custom_tool_names: set[str] = extract_custom_tool_names(self.responses_api_request.get("tools"))
self._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map(
self.responses_api_request.get("tools")
@ -543,7 +548,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
@staticmethod
def _snapshot_chunk_for_stream_chunk_builder(
chunk: ModelResponseStream,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Convert a streaming chunk into a plain dict for end-of-stream assembly.
Keep _hidden_params so downstream usage/header behavior is preserved.
@ -1161,7 +1166,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
if litellm_model_response:
# Add cost to usage object if include_cost_in_streaming_usage is True
if litellm.include_cost_in_streaming_usage and self.litellm_logging_obj is not None:
usage: Final = getattr(litellm_model_response, "usage", None)
usage: Final[object] = getattr(litellm_model_response, "usage", None)
if usage is not None:
setattr(
usage,

82
litellm/responses/litellm_completion_transformation/transformation.py
View file

@ -5,7 +5,7 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion
import json
import re
import uuid
from collections.abc import Iterator, Mapping, Sequence
from collections.abc import Iterable, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
@ -28,7 +28,7 @@ from openai.types.responses import ResponseFunctionToolCall
from openai.types.responses.response_create_params import ResponseInputParam
from openai.types.responses.tool_param import FunctionToolParam
from pydantic import TypeAdapter
from typing_extensions import TypedDict
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.caching import InMemoryCache
@ -46,6 +46,7 @@ from litellm.types.llms.openai import (
ChatCompletionRedactedThinkingBlock,
ChatCompletionResponseMessage,
ChatCompletionSystemMessage,
ChatCompletionTextObject,
ChatCompletionThinkingBlock,
ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk,
@ -129,6 +130,30 @@ class _HasId(Protocol):
id: object
class _ResponsesToolCallItem(Protocol):
name: object
arguments: object
def get(self, key: str, /) -> object: ...
class _ToolFunctionDefinition(TypedDict, total=False):
name: ReadOnly[str]
description: ReadOnly[str]
parameters: ReadOnly[dict[str, object]]
strict: ReadOnly[bool | None]
def _attribute_fields(value: object) -> dict[str, object]:
if not hasattr(value, "__dict__"):
return {} # mutable-ok: provider_specific_fields payload
return dict(cast("Iterable[tuple[str, object]]", value)) # cast-ok: dict() raises on non-pair values, as before
def _input_item_role(input_item: Mapping[str, object]) -> str:
return cast(str, input_item.get("role") or "user") # cast-ok: client-supplied role forwarded verbatim, unvalidated
class ChatCompletionSession(TypedDict, total=False):
messages: list[
AllMessageValues
@ -677,7 +702,7 @@ class LiteLLMCompletionResponsesConfig:
existing_text: Final = _reasoning_text(msg)
combined: Final = "\n".join(pending_texts + ((existing_text,) if existing_text else ()))
if isinstance(msg, dict):
cast(dict[str, Any], msg)["reasoning_content"] = combined # cast-ok: mutable reasoning carrier
cast(dict[str, object], msg)["reasoning_content"] = combined # cast-ok: mutable reasoning carrier
else:
setattr(msg, "reasoning_content", combined) # noqa: B010 # attribute name is fixed, not dynamic
if pending_blocks:
@ -685,7 +710,7 @@ class LiteLLMCompletionResponsesConfig:
pending_blocks + (_thinking_blocks(msg) or ())
)
if isinstance(msg, dict):
cast(dict[str, Any], msg)["thinking_blocks"] = replayed # cast-ok: mutable reasoning carrier
cast(dict[str, object], msg)["thinking_blocks"] = replayed # cast-ok: mutable reasoning carrier
else:
setattr(msg, "thinking_blocks", replayed) # noqa: B010 # attribute name is fixed, not dynamic
@ -1034,7 +1059,7 @@ class LiteLLMCompletionResponsesConfig:
def _add_tool_call_to_assistant(assistant_message: object, tool_call_chunk: ChatCompletionToolCallChunk) -> None:
"""Add a tool_call to an assistant message."""
if isinstance(assistant_message, dict):
prev_assistant_dict: Final = cast(dict[str, Any], assistant_message)
prev_assistant_dict: Final = cast(dict[str, object], assistant_message)
if "tool_calls" not in prev_assistant_dict:
prev_assistant_dict["tool_calls"] = []
tool_calls_list: Final = prev_assistant_dict["tool_calls"]
@ -1119,7 +1144,7 @@ class LiteLLMCompletionResponsesConfig:
# Type-safe way to set tool_call_id on tool message
if isinstance(message, dict):
# Cast to dict to allow setting tool_call_id
message_dict = cast(dict[str, Any], message)
message_dict = cast(dict[str, object], message)
message_dict["tool_call_id"] = tool_call_id
elif hasattr(message, "tool_call_id"):
setattr(message, "tool_call_id", tool_call_id)
@ -1171,7 +1196,7 @@ class LiteLLMCompletionResponsesConfig:
@staticmethod
def _transform_responses_api_input_item_to_chat_completion_message(
input_item: Any,
input_item: Mapping[str, object],
replay_reasoning: bool = False,
) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]:
"""
@ -1199,7 +1224,9 @@ class LiteLLMCompletionResponsesConfig:
elif LiteLLMCompletionResponsesConfig._is_input_item_function_call(input_item):
# handle function call input items
return LiteLLMCompletionResponsesConfig._transform_responses_api_function_call_to_chat_completion_message(
function_call=input_item
function_call=cast( # cast-ok: callee coerces every field it reads with `or ""` / str()
Mapping[str, str], input_item
)
)
elif input_item.get("type") == "reasoning":
# A ResponseReasoningItemParam carries the prior-turn chain-of-thought.
@ -1224,7 +1251,7 @@ class LiteLLMCompletionResponsesConfig:
return [] # mutable-ok: empty drop result
return [ # mutable-ok: single message result
GenericChatCompletionMessage(
role=input_item.get("role") or "user",
role=_input_item_role(input_item),
content=LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(
inspectable
),
@ -1252,7 +1279,7 @@ class LiteLLMCompletionResponsesConfig:
return []
return [
GenericChatCompletionMessage(
role=input_item.get("role") or "user",
role=_input_item_role(input_item),
content=LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(
content
),
@ -1339,7 +1366,7 @@ class LiteLLMCompletionResponsesConfig:
if not isinstance(encrypted_content, str) or not encrypted_content.strip():
return None
try:
decoded: Final[object] = json.loads(encrypted_content)
decoded: Final[object] = cast(object, json.loads(encrypted_content)) # cast-ok: json.loads returns Any
except ValueError:
return None
if not isinstance(decoded, list):
@ -1406,7 +1433,7 @@ class LiteLLMCompletionResponsesConfig:
def _normalize_function_call_output_to_tool_content(
output: object,
) -> Any:
) -> str | list[ChatCompletionTextObject | ChatCompletionImageObject]:
"""
Normalize Responses API function_call_output.output into a shape that downstream
chat adapters (esp. Gemini) can reliably consume.
@ -1428,7 +1455,7 @@ class LiteLLMCompletionResponsesConfig:
# Some adapters represent tool output as a list of "input_*" parts
if isinstance(output, list):
normalized_blocks: Final[list[dict[str, object]]] = []
normalized_blocks: Final[list[ChatCompletionTextObject | ChatCompletionImageObject]] = []
text_acc: Final[list[str]] = []
for part in output:
if not isinstance(part, dict):
@ -1899,7 +1926,7 @@ class LiteLLMCompletionResponsesConfig:
result.append(tool)
continue
if tool.get("type") == "function":
fn = cast(dict[str, Any], tool.get("function") or {})
fn = cast(_ToolFunctionDefinition, tool.get("function") or {})
parameters = dict(fn.get("parameters", {}) or {})
if not parameters or "type" not in parameters:
parameters["type"] = "object"
@ -2095,7 +2122,7 @@ class LiteLLMCompletionResponsesConfig:
@staticmethod
def convert_response_function_tool_call_to_chat_completion_tool_call(
tool_call_item: Any,
tool_call_item: object,
index: int = 0,
) -> dict[str, object]:
"""
@ -2108,24 +2135,25 @@ class LiteLLMCompletionResponsesConfig:
Returns:
Dictionary in ChatCompletionToolCallChunk format
"""
item: Final = cast( # cast-ok: duck-typed tool call item, .get access guarded by hasattr below
_ResponsesToolCallItem, tool_call_item
)
# Extract provider_specific_fields if present
provider_specific_fields = getattr(tool_call_item, "provider_specific_fields", None)
provider_specific_fields: object = getattr(tool_call_item, "provider_specific_fields", None)
if provider_specific_fields and not isinstance(provider_specific_fields, dict):
provider_specific_fields = (
dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {}
)
elif hasattr(tool_call_item, "get") and callable(tool_call_item.get):
provider_fields: Final = tool_call_item.get("provider_specific_fields")
provider_specific_fields = _attribute_fields(provider_specific_fields)
elif hasattr(tool_call_item, "get") and callable(item.get):
provider_fields: Final = item.get("provider_specific_fields")
if provider_fields:
provider_specific_fields = (
provider_fields
cast("dict[str, object]", provider_fields) # cast-ok: passed through as-is, keys unvalidated
if isinstance(provider_fields, dict)
else (dict(provider_fields) if hasattr(provider_fields, "__dict__") else {})
else _attribute_fields(provider_fields)
)
function_dict: Final[dict[str, object]] = {
"name": tool_call_item.name,
"arguments": tool_call_item.arguments,
"name": item.name,
"arguments": item.arguments,
}
if provider_specific_fields:
@ -2306,7 +2334,7 @@ class LiteLLMCompletionResponsesConfig:
"""
output_items: Final[list] = []
for choice in chat_completion_response.choices or []:
message = getattr(choice, "message", None)
message: object = getattr(choice, "message", None)
if not message:
continue
psf = getattr(message, "provider_specific_fields", None)
@ -2338,7 +2366,7 @@ class LiteLLMCompletionResponsesConfig:
for choice in choices:
if hasattr(choice, "message") and choice.message:
message = choice.message
reasoning_content = getattr(message, "reasoning_content", None) or ""
reasoning_content: str = getattr(message, "reasoning_content", None) or ""
encrypted_content = LiteLLMCompletionResponsesConfig._encode_thinking_blocks(message)
if reasoning_content or encrypted_content:
# Only check the first choice for reasoning content

4
litellm/responses/main.py
View file

@ -697,7 +697,7 @@ def _apply_managed_file_id_mapping(
tools = cast(
Iterable[ToolParam] | None,
update_responses_tools_with_model_file_ids(
tools=cast(list[dict[str, Any]] | None, tools),
tools=cast(list[dict[str, object]] | None, tools),
model_id=model_info_id,
model_file_id_mapping=model_file_id_mapping,
),
@ -734,7 +734,7 @@ def _responses_try_dispatch_mcp_gateway(
extra_body: dict[str, object] | None,
timeout: float | httpx.Timeout | None,
custom_llm_provider: str | None,
kwargs: dict[str, Any],
kwargs: dict[str, object],
_is_async: bool,
) -> Any | None:
"""Return a response when MCP gateway handles the call; otherwise None."""

22
litellm/responses/mcp/chat_completions_handler.py
View file

@ -1,7 +1,9 @@
"""Helpers for handling MCP-aware `/chat/completions` requests."""
import logging
from typing import TYPE_CHECKING, Any, Final, cast
from typing import TYPE_CHECKING, Final, cast
from typing_extensions import TypedDict, Unpack
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
@ -14,6 +16,10 @@ if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
class _MCPCompletionKwargs(TypedDict, total=False, extra_items=object):
"""Extra keywords forwarded verbatim to ``litellm.acompletion``, which owns their contract."""
def _add_mcp_metadata_to_response(
response: ModelResponse | CustomStreamWrapper,
openai_tools: list | None,
@ -79,7 +85,7 @@ async def acompletion_with_mcp(
model: str,
messages: list,
tools: list | None = None,
**kwargs: Any,
**kwargs: Unpack[_MCPCompletionKwargs], # kwargs-ok: forwarded verbatim to litellm.acompletion, which owns them
) -> ModelResponse | CustomStreamWrapper:
"""
Async completion with MCP integration.
@ -126,7 +132,7 @@ async def acompletion_with_mcp(
) = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_without_openai_transform(
user_api_key_auth=user_api_key_auth,
mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy,
litellm_trace_id=kwargs.get("litellm_trace_id"),
litellm_trace_id=context.litellm_trace_id,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
request_tags=request_tags,
@ -168,7 +174,7 @@ async def acompletion_with_mcp(
return response
# For auto-execute: handle streaming vs non-streaming differently
stream: Final[bool] = kwargs.get("stream", False)
stream: Final[object] = kwargs.get("stream", False)
mock_tool_calls: Final = base_call_args.pop("mock_tool_calls", None)
if stream:
@ -490,8 +496,8 @@ async def acompletion_with_mcp(
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
litellm_call_id=kwargs.get("litellm_call_id"),
litellm_trace_id=kwargs.get("litellm_trace_id"),
litellm_call_id=context.litellm_call_id,
litellm_trace_id=context.litellm_trace_id,
openai_tools=openai_tools,
base_call_args=base_call_args,
request_tags=request_tags,
@ -604,8 +610,8 @@ async def acompletion_with_mcp(
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
litellm_call_id=kwargs.get("litellm_call_id"),
litellm_trace_id=kwargs.get("litellm_trace_id"),
litellm_call_id=context.litellm_call_id,
litellm_trace_id=context.litellm_trace_id,
request_tags=request_tags,
)

14
litellm/responses/mcp/mcp_streaming_iterator.py
View file

@ -22,6 +22,8 @@ from litellm.types.llms.openai import (
)
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Iterator
from mcp.types import Tool as MCPTool
from litellm.proxy._types import UserAPIKeyAuth
@ -511,7 +513,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
if self.base_iterator:
if hasattr(self.base_iterator, "__anext__"):
try:
chunk: Final[ResponsesAPIStreamingResponse] = await cast(Any, self.base_iterator).__anext__()
chunk: Final[ResponsesAPIStreamingResponse] = await cast( # cast-ok: hasattr __anext__ checked
"AsyncIterator[ResponsesAPIStreamingResponse]", self.base_iterator
).__anext__()
# Capture the response ID from the first event to ensure consistency
if self._cached_response_id is None and hasattr(chunk, "response"):
@ -569,7 +573,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
if not self.base_iterator or not hasattr(self.base_iterator, "__anext__"):
raise StopAsyncIteration
chunk: Final[ResponsesAPIStreamingResponse] = await cast(Any, self.base_iterator).__anext__()
chunk: Final[ResponsesAPIStreamingResponse] = await cast( # cast-ok: hasattr __anext__ checked above
"AsyncIterator[ResponsesAPIStreamingResponse]", self.base_iterator
).__anext__()
if self._cached_response_id is None and hasattr(chunk, "response"):
new_response: Final[ResponsesAPIResponse | None] = getattr(chunk, "response", None)
@ -834,7 +840,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
if not self.is_async:
try:
if self.base_iterator and hasattr(self.base_iterator, "__next__"):
return next(cast(Any, self.base_iterator))
return next(
cast("Iterator[ResponsesAPIStreamingResponse]", self.base_iterator) # cast-ok: hasattr-checked
)
else:
raise StopIteration
except StopIteration:

23
litellm/responses/mcp/request_context.py
View file

@ -10,14 +10,25 @@ still executes the tool, just with no credentials.
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Final
from typing import TYPE_CHECKING, Any, Final
from typing_extensions import NotRequired, ReadOnly, TypedDict
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
class _AuthCarryingMetadata(TypedDict):
"""The one key this module reads out of a request's ``metadata`` / ``litellm_metadata``."""
user_api_key_auth: ReadOnly[NotRequired["UserAPIKeyAuth | None"]]
@dataclass(frozen=True, slots=True)
class MCPRequestContext:
"""Everything a gateway handler must forward to MCP tool listing and execution."""
user_api_key_auth: Any # any-ok: UserAPIKeyAuth is proxy-only; importing it here would create a cycle
user_api_key_auth: "UserAPIKeyAuth | None"
mcp_auth_header: str | None = None
mcp_server_auth_headers: Mapping[str, Mapping[str, str]] | None = None
oauth2_headers: Mapping[str, str] | None = None
@ -30,7 +41,7 @@ class MCPRequestContext:
def resolve(
cls,
kwargs: Mapping[str, Any],
tools: Iterable[Any] | None,
tools: Iterable[object] | None,
) -> "MCPRequestContext":
"""
Build the context from a gateway handler's kwargs.
@ -44,9 +55,9 @@ class MCPRequestContext:
)
from litellm.responses.utils import ResponsesAPIRequestUtils
litellm_metadata: Final = kwargs.get("litellm_metadata") or {}
metadata: Final = kwargs.get("metadata") or {}
user_api_key_auth: Final = (
litellm_metadata: Final[_AuthCarryingMetadata] = kwargs.get("litellm_metadata") or {}
metadata: Final[_AuthCarryingMetadata] = kwargs.get("metadata") or {}
user_api_key_auth: Final[UserAPIKeyAuth | None] = (
kwargs.get("user_api_key_auth")
or litellm_metadata.get("user_api_key_auth")
or metadata.get("user_api_key_auth")

53
litellm/responses/sse_output_recovery.py
View file

@ -8,14 +8,17 @@ caller automatically applies to all of them.
"""
import json
from typing import Any, Final
from collections.abc import Mapping
from typing import Final, SupportsInt, TypeAlias, cast # noqa: TID251 # int() re-checks the cast below at runtime
from litellm.constants import STREAM_SSE_DONE_STRING
_MAX_CONTENT_INDEX: Final = 1024
_ConvertibleToInt: TypeAlias = SupportsInt | str
def parse_sse_json_chunk(chunk: str) -> dict[str, Any] | None:
def parse_sse_json_chunk(chunk: str) -> dict[str, object] | None:
"""Parse a single raw SSE line into a JSON object dict.
Returns ``None`` for empty lines, ``event:`` lines, ``[DONE]`` markers,
@ -30,7 +33,7 @@ def parse_sse_json_chunk(chunk: str) -> dict[str, Any] | None:
if not stripped_chunk or stripped_chunk == STREAM_SSE_DONE_STRING or stripped_chunk.startswith("event:"):
return None
try:
parsed_chunk: Final = json.loads(stripped_chunk)
parsed_chunk: Final[object] = json.loads(stripped_chunk)
except json.JSONDecodeError:
return None
if not isinstance(parsed_chunk, dict):
@ -38,9 +41,19 @@ def parse_sse_json_chunk(chunk: str) -> dict[str, Any] | None:
return parsed_chunk
def _chunk_index(parsed_chunk: Mapping[str, object], key: str, fallback: int) -> int:
raw_index: Final = parsed_chunk.get(key)
if raw_index is None:
return fallback
try:
return int(cast(_ConvertibleToInt, raw_index)) # cast-ok: int() raises TypeError otherwise, caught below
except (TypeError, ValueError):
return fallback
def record_output_item_chunk(
parsed_chunk: dict[str, Any],
output_items: dict[int, dict[str, Any]],
parsed_chunk: Mapping[str, object],
output_items: dict[int, dict[str, object]],
) -> None:
"""Record an OUTPUT_ITEM_DONE chunk into ``output_items`` keyed by
``output_index`` (falling back to the next free slot when missing).
@ -48,20 +61,14 @@ def record_output_item_chunk(
item: Final = parsed_chunk.get("item")
if not isinstance(item, dict):
return
try:
output_index_raw: Final = parsed_chunk.get("output_index")
if output_index_raw is None:
raise ValueError("missing output_index")
output_index = int(output_index_raw)
except (TypeError, ValueError):
output_index = len(output_items)
output_index: Final = _chunk_index(parsed_chunk, "output_index", len(output_items))
output_items[output_index] = item
def record_output_text_chunk(
parsed_chunk: dict[str, Any],
output_items: dict[int, dict[str, Any]],
text_only_items: dict[int, dict[str, Any]],
parsed_chunk: Mapping[str, object],
output_items: Mapping[int, dict[str, object]],
text_only_items: dict[int, dict[str, object]],
) -> None:
"""Record an OUTPUT_TEXT_DONE chunk as a synthetic message item in
``text_only_items``. Real OUTPUT_ITEM_DONE events already captured in
@ -71,13 +78,7 @@ def record_output_text_chunk(
if not isinstance(text, str):
return
try:
output_index_raw: Final = parsed_chunk.get("output_index")
if output_index_raw is None:
raise ValueError("missing output_index")
output_index = int(output_index_raw)
except (TypeError, ValueError):
output_index = len(text_only_items)
output_index: Final = _chunk_index(parsed_chunk, "output_index", len(text_only_items))
if output_index in output_items:
return
@ -97,13 +98,7 @@ def record_output_text_chunk(
if not isinstance(content, list):
return
try:
content_index_raw: Final = parsed_chunk.get("content_index")
if content_index_raw is None:
raise ValueError("missing content_index")
content_index = int(content_index_raw)
except (TypeError, ValueError):
content_index = len(content)
content_index: Final = _chunk_index(parsed_chunk, "content_index", len(content))
if content_index < 0 or content_index > _MAX_CONTENT_INDEX:
return

42
litellm/responses/streaming_iterator.py
View file

@ -49,6 +49,21 @@ if TYPE_CHECKING:
ResponsesClientWebSocket,
)
class _StreamCachingHandler(Protocol):
"""The ``_llm_caching_handler`` attached to a logging object, as this module uses it."""
original_function: Callable[..., object]
def _should_store_result_in_cache(
self, original_function: Callable[..., object], kwargs: Mapping[str, object]
) -> bool: ...
class PiiUnmaskingGuardrailCallback(PresidioGuardrailCallback, Protocol):
"""Guardrail callback that can also reverse its own masking, selected by
``llm_http_handler`` on exactly this attribute."""
def _unmask_pii_text(self, text: str, pii_tokens: Mapping[str, str]) -> str: ...
class ProjectQuotaCallback(Protocol):
async def enforce_project_io_token_quota_for_frame(
@ -84,6 +99,11 @@ def _load_json_object(payload: str | bytes) -> dict[str, object]:
return json.loads(payload)
def _load_json_value(payload: str | bytes) -> object:
"""Parse a JSON payload whose top-level shape the caller narrows itself."""
return json.loads(payload)
def _model_id_from_metadata(litellm_metadata: dict[str, object] | None) -> str | None:
model_info: Final = litellm_metadata.get("model_info") if litellm_metadata else None
model_id: Final = model_info.get("id") if _is_json_object(model_info) else None
@ -243,10 +263,10 @@ class BaseResponsesAPIStreamingIterator:
try:
# Parse the JSON chunk
parsed_chunk: Final = json.loads(chunk)
parsed_chunk: Final = _load_json_value(chunk)
# Format as ResponsesAPIStreamingResponse
if isinstance(parsed_chunk, dict):
if _is_json_object(parsed_chunk):
if self.responses_api_provider_config is None:
raise ValueError("responses_api_provider_config is required to process live streaming chunks")
openai_responses_api_chunk: Final = self.responses_api_provider_config.transform_streaming_response(
@ -529,7 +549,7 @@ class BaseResponsesAPIStreamingIterator:
if response_obj is None:
return
caching_handler: Final = getattr(self.logging_obj, "_llm_caching_handler", None)
caching_handler: Final[_StreamCachingHandler | None] = getattr(self.logging_obj, "_llm_caching_handler", None)
if caching_handler is None:
return
@ -547,7 +567,7 @@ class BaseResponsesAPIStreamingIterator:
if preset_cache_key is not None:
request_kwargs["cache_key"] = preset_cache_key
if not caching_handler._should_store_result_in_cache(
if not caching_handler._should_store_result_in_cache( # pyright: ignore[reportPrivateUsage] # no public API
original_function=caching_handler.original_function,
kwargs=request_kwargs,
):
@ -1401,7 +1421,7 @@ async def _enforce_frame_project_quota(
if not quota_callbacks:
return
try:
msg_obj = json.loads(raw_message)
msg_obj: Final = _load_json_value(raw_message)
except (json.JSONDecodeError, TypeError):
return
if not _is_json_object(msg_obj) or msg_obj.get("type") != "response.create":
@ -1451,7 +1471,7 @@ class ResponsesWebSocketStreaming:
user_api_key_dict: UserAPIKeyAuth | None = None,
request_data: dict[str, object] | None = None,
first_message: str | None = None,
guardrail_callbacks: list[Any] | None = None,
guardrail_callbacks: list[PiiUnmaskingGuardrailCallback] | None = None,
output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None,
quota_callbacks: Sequence[ProjectQuotaCallback] | None = None,
authorized_model: str | None = None,
@ -1464,7 +1484,7 @@ class ResponsesWebSocketStreaming:
self.messages: list[dict[str, object]] = []
self.input_messages: list[dict[str, object]] = []
self.first_message = first_message
self.guardrail_callbacks: list[Any] = guardrail_callbacks or []
self.guardrail_callbacks: list[PiiUnmaskingGuardrailCallback] = guardrail_callbacks or []
self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or []
self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else ()
# Model name authorized at connection time; enforced on every
@ -1780,7 +1800,9 @@ class ResponsesWebSocketStreaming:
continue
text = content_block.get("text")
if isinstance(text, str):
unmasked = cb._unmask_pii_text(text, pii_tokens)
unmasked = cb._unmask_pii_text( # pyright: ignore[reportPrivateUsage] # no public unmasker
text, pii_tokens
)
if unmasked != text:
content_block["text"] = unmasked
modified = True
@ -1789,7 +1811,9 @@ class ResponsesWebSocketStreaming:
if event_type in self._DELTA_EVENT_TYPES:
delta: Final = evt_obj.get("delta")
if isinstance(delta, str):
unmasked = cb._unmask_pii_text(delta, pii_tokens)
unmasked = cb._unmask_pii_text( # pyright: ignore[reportPrivateUsage] # no public unmasker
delta, pii_tokens
)
if unmasked != delta:
evt_obj["delta"] = unmasked
return json.dumps(evt_obj)

56
litellm/responses/utils.py
View file

@ -1,15 +1,17 @@
import base64
import re
from collections.abc import Iterable, Mapping
from collections.abc import Iterable, Mapping, Sequence
from typing import Any, Final, Optional, Union, cast, get_type_hints, overload
from pydantic import BaseModel
from typing_extensions import TypeIs # noqa: TID251 # narrows untyped wire payloads without a runtime conversion
import litellm
from litellm._logging import verbose_logger
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.types.llms.openai import (
AllMessageValues,
OutputTokensDetails,
ResponseAPIUsage,
ResponseInputParam,
ResponsesAPIOptionalRequestParams,
@ -26,6 +28,16 @@ from litellm.types.utils import (
)
def _is_object_sequence(value: object) -> TypeIs[Sequence[object]]: # guard-ok: a list is a Sequence of anything
return isinstance(value, list)
def _is_object_dict(
value: object,
) -> TypeIs[dict[str, object]]: # guard-ok: wire dicts have str keys # mutable-ok: callers rewrite ids in place
return isinstance(value, dict)
def normalize_responses_api_stream_options(
stream_options: object,
) -> ResponsesAPIStreamOptions | None:
@ -703,12 +715,12 @@ class ResponsesAPIRequestUtils:
@staticmethod
def _encode_container_ids_in_annotations(
annotations: Any,
annotations: object,
custom_llm_provider: str | None,
model_id: str | None,
) -> None:
"""Encode ``container_id`` on each annotation (e.g. ``container_file_citation``)."""
if not annotations or not isinstance(annotations, list):
if not annotations or not _is_object_sequence(annotations):
return
for ann in annotations:
ResponsesAPIRequestUtils._encode_container_id_on_output_item(
@ -719,16 +731,16 @@ class ResponsesAPIRequestUtils:
@staticmethod
def _encode_container_ids_in_message_content(
content: Any,
content: object,
custom_llm_provider: str | None,
model_id: str | None,
) -> None:
"""Walk message ``content`` parts and encode citation ``container_id`` values."""
if not content:
return
if isinstance(content, list):
if _is_object_sequence(content):
for part in content:
if isinstance(part, dict):
if _is_object_dict(part):
ResponsesAPIRequestUtils._encode_container_ids_in_annotations(
part.get("annotations"),
custom_llm_provider,
@ -743,7 +755,7 @@ class ResponsesAPIRequestUtils:
@staticmethod
def _encode_container_id_on_output_item(
item: Any,
item: object,
custom_llm_provider: str | None,
model_id: str | None,
) -> None:
@ -770,14 +782,14 @@ class ResponsesAPIRequestUtils:
container_id=container_id,
)
if isinstance(item, dict):
if _is_object_dict(item):
cid: Final = item.get("container_id")
if isinstance(cid, str):
enc = _maybe_encode(cid)
if enc is not None:
item["container_id"] = enc
item["container_id"] = enc # rebind-ok: this helper's contract is to rewrite the item in place
nested: Final = item.get("code_interpreter_call")
if isinstance(nested, dict):
if _is_object_dict(nested):
nc: Final = nested.get("container_id")
if isinstance(nc, str):
enc = _maybe_encode(nc)
@ -803,7 +815,7 @@ class ResponsesAPIRequestUtils:
exc_info=True,
)
nested_obj: Final = getattr(item, "code_interpreter_call", None)
nested_obj: Final[object] = getattr(item, "code_interpreter_call", None)
if nested_obj is not None:
ResponsesAPIRequestUtils._encode_container_id_on_output_item(
nested_obj,
@ -820,24 +832,24 @@ class ResponsesAPIRequestUtils:
@staticmethod
def _collect_container_ids_from_annotations(
annotations: Any,
annotations: object,
collected: set[str],
) -> None:
if not annotations or not isinstance(annotations, list):
if not annotations or not _is_object_sequence(annotations):
return
for ann in annotations:
ResponsesAPIRequestUtils._collect_container_ids_from_output_item(ann, collected)
@staticmethod
def _collect_container_ids_from_message_content(
content: Any,
content: object,
collected: set[str],
) -> None:
if not content:
return
if isinstance(content, list):
if _is_object_sequence(content):
for part in content:
if isinstance(part, dict):
if _is_object_dict(part):
ResponsesAPIRequestUtils._collect_container_ids_from_annotations(
part.get("annotations"),
collected,
@ -850,19 +862,19 @@ class ResponsesAPIRequestUtils:
@staticmethod
def _collect_container_ids_from_output_item(
item: Any,
item: object,
collected: set[str],
) -> None:
"""Collect managed or raw ``container_id`` values from one output item."""
if item is None:
return
if isinstance(item, dict):
if _is_object_dict(item):
cid: Final = item.get("container_id")
if isinstance(cid, str) and cid:
collected.add(cid)
nested: Final = item.get("code_interpreter_call")
if isinstance(nested, dict):
if _is_object_dict(nested):
nc: Final = nested.get("container_id")
if isinstance(nc, str) and nc:
collected.add(nc)
@ -877,7 +889,7 @@ class ResponsesAPIRequestUtils:
if isinstance(cid_attr, str) and cid_attr:
collected.add(cid_attr)
nested_obj: Final = getattr(item, "code_interpreter_call", None)
nested_obj: Final[object] = getattr(item, "code_interpreter_call", None)
if nested_obj is not None:
ResponsesAPIRequestUtils._collect_container_ids_from_output_item(nested_obj, collected)
@ -1108,7 +1120,9 @@ class ResponseAPILoggingUtils:
cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None),
)
completion_tokens_details: CompletionTokensDetailsWrapper | None = None
output_tokens_details: Final = getattr(response_api_usage, "output_tokens_details", None)
output_tokens_details: Final[OutputTokensDetails | None] = getattr(
response_api_usage, "output_tokens_details", None
)
if output_tokens_details:
completion_tokens_details = CompletionTokensDetailsWrapper(
reasoning_tokens=getattr(output_tokens_details, "reasoning_tokens", None),

13
litellm/vector_stores/vector_store_registry.py
View file

@ -1,7 +1,14 @@
# litellm/proxy/vector_stores/vector_store_registry.py
import json
from collections.abc import Mapping
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Final, get_args
from typing import (
TYPE_CHECKING,
Any, # noqa: TID251 # untyped non_default_params dict is the only source of the unknown key type
Final,
cast, # noqa: TID251 # untyped non_default_params dict is the only source of the unknown key type
get_args,
)
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.core_helpers import remove_items_at_indices
@ -336,7 +343,9 @@ class VectorStoreRegistry:
try:
# Check if it still exists in database
db_vector_store = await ManagedVectorStoresRepository(prisma_client).table.find_unique(
where={"vector_store_id": vector_store_id}
where=cast( # cast-ok: every value is already an object, only the popped id is stub-untyped
"Mapping[str, object]", {"vector_store_id": vector_store_id}
)
)
if db_vector_store is None:
# Vector store was deleted from database, remove from cache

12
ruff-strict-budget.json
View file

@ -1,6 +1,6 @@
{
"ANN001": {
"limit": 3018
"limit": 3016
},
"ANN002": {
"limit": 71
@ -9,7 +9,7 @@
"limit": 827
},
"ANN201": {
"limit": 2016
"limit": 2013
},
"ANN202": {
"limit": 852
@ -24,7 +24,7 @@
"limit": 133
},
"ANN401": {
"limit": 1188
"limit": 1157
},
"ASYNC230": {
"limit": 11
@ -33,13 +33,13 @@
"limit": 2
},
"B006": {
"limit": 177
"limit": 176
},
"B008": {
"limit": 503
},
"B009": {
"limit": 59
"limit": 58
},
"B010": {
"limit": 190
@ -231,7 +231,7 @@
"limit": 5
},
"TID251": {
"limit": 1212
"limit": 1201
},
"TRY002": {
"limit": 524

28
tests/proxy_unit_tests/test_jwt_key_mapping.py
View file

@ -416,6 +416,34 @@ async def test_update_returns_404_when_not_found():
assert exc_info.value.status_code == 404
@pytest.mark.asyncio
async def test_update_returns_404_when_row_deleted_before_write():
"""A mapping deleted between the read and the write must 404, not 500.
Prisma's update returns None when the row is gone, and the endpoint used to
dereference it for the cache key.
"""
from litellm.proxy._types import UpdateJWTKeyMappingRequest
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = _mock_mapping()
mock_prisma.db.litellm_jwtkeymapping.update.return_value = None
mock_cache = AsyncMock()
data = UpdateJWTKeyMappingRequest(id="mapping-1", description="test")
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy_server module global is the endpoint's only injection point
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), # test-quality-ok: proxy_server module global is the endpoint's only injection point
):
with pytest.raises(HTTPException) as exc_info:
await update_jwt_key_mapping(
data=data, user_api_key_dict=_make_admin_auth()
)
assert exc_info.value.status_code == 404
assert exc_info.value.detail == "Mapping not found"
@pytest.mark.asyncio
async def test_info_returns_404_when_not_found():
"""Getting info for non-existent mapping should return 404."""

57
tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py
View file

@ -428,3 +428,60 @@ async def test_migrate_legacy_grant_ids_no_ops_without_config_agents():
assert await registry.migrate_legacy_grant_ids(table=table) == GrantMigrationResult(rewritten=0, missed=0)
table.find_many.assert_not_awaited()
@pytest.mark.asyncio
async def test_update_agent_in_db_raises_when_row_deleted_mid_update():
"""Prisma's update returns None when the row vanished between read and write. Without a
guard the code dereferences None and reports an opaque AttributeError instead of the id."""
registry: Final = AgentRegistry()
mock_prisma: Final = MagicMock()
mock_prisma.db.litellm_agentstable.update = AsyncMock(return_value=None)
with pytest.raises(Exception, match="Error updating agent in DB") as exc_info:
await registry.update_agent_in_db(
agent_id="agent-123",
agent={
"agent_name": "Updated Agent",
"agent_card_params": _sample_agent_card_params(),
"litellm_params": {},
},
prisma_client=mock_prisma,
updated_by="test-user",
)
assert str(exc_info.value) == "Error updating agent in DB: Agent not found, passed agent_id=agent-123"
@pytest.mark.asyncio
async def test_patch_agent_in_db_raises_when_row_deleted_mid_update():
"""Same race on PATCH: the existing row is read, then deleted before the update lands."""
registry: Final = AgentRegistry()
mock_prisma: Final = MagicMock()
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
return_value={"agent_id": "agent-123", "agent_name": "Old Agent", "object_permission_id": None}
)
mock_prisma.db.litellm_agentstable.update = AsyncMock(return_value=None)
with pytest.raises(Exception, match="Error patching agent in DB") as exc_info:
await registry.patch_agent_in_db(
agent_id="agent-123",
agent={"agent_name": "Patched Agent"},
prisma_client=mock_prisma,
updated_by="test-user",
)
assert str(exc_info.value) == "Error patching agent in DB: Agent not found, passed agent_id=agent-123"
@pytest.mark.asyncio
async def test_delete_agent_from_db_raises_when_row_already_gone():
"""Prisma's delete returns None for a missing row, which dict() cannot consume."""
registry: Final = AgentRegistry()
mock_prisma: Final = MagicMock()
mock_prisma.db.litellm_agentstable.delete = AsyncMock(return_value=None)
with pytest.raises(Exception, match="Error deleting agent from DB") as exc_info:
await registry.delete_agent_from_db(agent_id="agent-123", prisma_client=mock_prisma)
assert str(exc_info.value) == "Error deleting agent from DB: Agent not found, passed agent_id=agent-123"

24
tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py
View file

@ -231,6 +231,30 @@ async def test_update_plugin_db_error_maps_to_structured_500():
assert "connection lost" in exc_info.value.detail["error"]
@pytest.mark.asyncio
async def test_update_plugin_deleted_mid_update_returns_404():
"""A concurrent delete between the find_unique pre-check and the update makes prisma's
update return None; that must surface the same 404 as a plain miss, not an AttributeError."""
name = "my-monorepo-plugin"
await register_plugin(
request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"),
user_api_key_dict=_USER,
)
table = litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable
table.update = AsyncMock(return_value=None)
with pytest.raises(HTTPException) as exc_info:
await update_plugin(
plugin_name=name,
request=UpdatePluginRequest(source={"source": "github", "repo": "org/replacement"}),
user_api_key_dict=_USER,
)
assert exc_info.value.status_code == 404
assert exc_info.value.detail == {"error": f"Plugin '{name}' not found"}
@pytest.mark.asyncio
async def test_get_marketplace_skips_plugin_with_null_manifest():
await register_plugin(

2
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
View file

@ -695,7 +695,7 @@ def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock
"object_permission_id": None,
"object_permission": None,
"litellm_budget_table": None,
"dict": lambda self=None: {
"model_dump": lambda self=None: {
"spend": 25.0,
"user_id": "enduser-implicit",
"blocked": False,

33
tests/test_litellm/proxy/db/mcp_server/test_db.py
View file

@ -4,7 +4,11 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy._experimental.mcp_server.db import get_mcp_servers_by_team
from litellm.proxy._experimental.mcp_server.db import (
approve_mcp_server,
get_mcp_servers_by_team,
reject_mcp_server,
)
def _prisma_client_returning(team_record: object) -> MagicMock:
@ -38,3 +42,30 @@ async def test_fetch_mcp_servers_by_team(team_record, expected):
where={"team_id": "team-123"},
include={"object_permission": True},
)
def _prisma_client_with_missing_mcp_server_row() -> MagicMock:
prisma_client = MagicMock()
prisma_client.db.litellm_mcpservertable.update = AsyncMock(return_value=None)
return prisma_client
@pytest.mark.asyncio
async def test_approve_mcp_server_raises_value_error_when_row_missing():
prisma_client = _prisma_client_with_missing_mcp_server_row()
with pytest.raises(ValueError, match=r"^MCP server not found, passed server_id=server-gone$"):
await approve_mcp_server(prisma_client, "server-gone", touched_by="admin")
@pytest.mark.asyncio
async def test_reject_mcp_server_raises_value_error_when_row_missing():
prisma_client = _prisma_client_with_missing_mcp_server_row()
with pytest.raises(ValueError, match=r"^MCP server not found, passed server_id=server-gone$"):
await reject_mcp_server(
prisma_client,
"server-gone",
touched_by="admin",
review_notes="spam",
)

22
tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
View file

@ -1,8 +1,11 @@
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy.guardrails.guardrail_registry import (
get_guardrail_initializer_from_hooks,
GuardrailRegistry,
InMemoryGuardrailHandler,
)
from litellm.types.guardrails import GuardrailEventHooks, Guardrail, LitellmParams
@ -657,3 +660,22 @@ class TestScanOnlyToolResultsInitRefusal:
"scan_only_tool_results": True,
},
)
@pytest.mark.asyncio
async def test_update_guardrail_in_db_raises_when_row_missing():
prisma_client = MagicMock()
prisma_client.db.litellm_guardrailstable.update = AsyncMock(return_value=None)
with pytest.raises(
Exception,
match=r"^Error updating guardrail in DB: Guardrail not found, passed guardrail_id=missing-guardrail$",
):
await GuardrailRegistry().update_guardrail_in_db(
guardrail_id="missing-guardrail",
guardrail=Guardrail(
guardrail_name="missing-guardrail",
litellm_params=LitellmParams(guardrail="bedrock", mode="pre_call"),
),
prisma_client=prisma_client,
)

43
tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
View file

@ -5446,3 +5446,46 @@ async def test_handle_group_membership_changes_already_in_team_is_noop(mocker):
)
assert mock_team_member_add.await_count == 2
@pytest.mark.asyncio
async def test_patch_group_404s_when_team_deleted_mid_request(mocker):
"""A group deleted between the existence check and the write must 404.
Prisma returns None from both the update and the refresh reads once the row is
gone, and patch_group used to dereference that None while building the response.
"""
group_id = "team-gone"
snapshot_team = LiteLLM_TeamTable(
team_id=group_id,
team_alias="Group",
members_with_roles=[Member(user_id="zed", role="user")],
metadata={"externalId": "grp-ext"},
)
patch_ops = SCIMPatchOp(
schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"],
Operations=[SCIMPatchOperation(op="replace", path="displayName", value="Renamed")],
)
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=[snapshot_team, None, None])
mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None)
mocker.patch( # test-quality-ok: stubs the collaborator so the test pins the endpoint's own error contract
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
AsyncMock(return_value=mock_prisma_client),
)
mocker.patch( # test-quality-ok: stubs the collaborator so the test pins the endpoint's own error contract
"litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles",
AsyncMock(),
)
with pytest.raises(ProxyException) as exc_info:
await patch_group(group_id=group_id, patch_ops=patch_ops)
assert exc_info.value.code == "404"
assert f"Group not found with ID: {group_id}" in exc_info.value.message

2
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -2398,7 +2398,7 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker):
mock_user_row.user_id = "admin-creator"
mock_user_row.user_email = "admin@example.com"
mock_user_row.teams = []
mock_user_row.json.return_value = "{}"
mock_user_row.model_dump_json.return_value = "{}"
mock_user_row.model_dump.return_value = {
"user_id": "admin-creator",
"user_email": "admin@example.com",

39
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -7766,6 +7766,45 @@ async def test_validate_key_list_check_key_hash_not_found():
assert "Key Hash not found" in exc_info.value.message
@pytest.mark.asyncio
async def test_validate_key_list_check_key_hash_row_missing():
"""A key_hash with no row reaches the same 'Key Hash not found' 403 as a failed
lookup, instead of blowing up inside the ownership check on a None row."""
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
return_value=LiteLLM_UserTable(
user_id="test-user",
user_email="test@example.com",
teams=[],
organization_memberships=[],
)
)
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=None
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="test-user",
api_key="sk-caller",
)
with pytest.raises(ProxyException) as exc_info:
await validate_key_list_check(
user_api_key_dict=user_api_key_dict,
user_id=None,
team_id=None,
organization_id=None,
key_alias=None,
key_hash="hash-of-a-deleted-key",
prisma_client=mock_prisma_client,
)
assert exc_info.value.code == "403" or exc_info.value.code == 403
assert exc_info.value.param == "key_hash"
assert "Key Hash not found" in exc_info.value.message
@pytest.mark.asyncio
async def test_validate_key_list_check_proxy_admin_viewer_skips_db_lookup():
"""proxy_admin_viewer takes the same unscoped read fast-path as proxy_admin, so no

55
tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -3312,6 +3312,61 @@ class TestPatchModelBlockedAuthGate:
mock_prisma.db.litellm_proxymodeltable.update.assert_awaited_once()
class TestPatchModelRowDeletedBeforeWrite:
"""A row deleted between the read and the update makes prisma's `update`
return None. That must surface patch_model's own 404 not-found contract,
not a 500 from dereferencing the missing row."""
@pytest.mark.asyncio
async def test_patch_model_404s_when_update_returns_none(self):
from litellm.proxy.management_endpoints.model_management_endpoints import (
patch_model,
)
from litellm.proxy.proxy_server import ProxyException
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
existing_row = MagicMock()
existing_row.litellm_params = {"model": "openai/gpt-4o-mini"}
existing_row.model_dump.return_value = {
"model_name": "gpt-4o-mini",
"litellm_params": existing_row.litellm_params,
"model_info": {"id": "m1"},
}
existing_row.model_dump_json.return_value = "{}"
mock_prisma = MagicMock()
mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(
return_value=existing_row
)
mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=None)
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy_server module global is the endpoint's only injection point
patch("litellm.proxy.proxy_server.llm_router", MagicMock(**{"get_model_ids.return_value": ["m1"]})), # test-quality-ok: proxy_server module global is the endpoint's only injection point
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy_server module global is the endpoint's only injection point
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy_server module global is the endpoint's only injection point
patch( # test-quality-ok: stubs the auth gate so the test exercises the not-found branch under test
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
new=AsyncMock(return_value=None),
),
patch( # test-quality-ok: stubs the cache write so the test observes only the DB result handling
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
new=AsyncMock(
return_value=ReconcileOutcome(still_desired=None, live_after=None)
),
),
):
with pytest.raises(ProxyException) as exc_info:
await patch_model(
model_id="m1",
patch_data=updateDeployment(blocked=True),
user_api_key_dict=admin,
)
assert exc_info.value.code == "404"
assert exc_info.value.message == "Model m1 not found on proxy."
class TestWriteSurfacesReloadDrop:
"""A model-write endpoint may report success only if every row it wrote is, after the
reload it triggered, live in this pod's router or deliberately environment-inactive."""

26
tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py
View file

@ -1037,3 +1037,29 @@ async def test_get_organization_daily_activity_non_admin_without_org_admin_role_
assert get_daily_activity_mock.call_args.kwargs["entity_id"] == []
assert org_table_find_many.call_args.kwargs["where"] == {"organization_id": {"in": []}}
@pytest.mark.asyncio
async def test_find_member_if_email_missing_row_raises_documented_400():
"""A user_email lookup that matches nothing returns None instead of raising, so the
only failure the surrounding try/except models is never entered. Without an explicit
None guard the next line dereferences None and /organization/member_add answers with
an AttributeError-driven 500 rather than the documented 400.
"""
from litellm.proxy.management_endpoints.organization_endpoints import (
find_member_if_email,
)
prisma_client = AsyncMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
with pytest.raises(HTTPException) as exc_info:
await find_member_if_email("missing@example.com", prisma_client)
assert exc_info.value.status_code == 400
assert exc_info.value.detail == {
"error": (
"Unique user not found for user_email=missing@example.com. Potential duplicate OR "
"non-existent user_email in LiteLLM_UserTable. Use 'user_id' instead."
)
}

38
tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py
View file

@ -1169,6 +1169,44 @@ async def test_delete_team_callback_404s_for_unknown_team():
mock_prisma.db.litellm_teamtable.update.assert_not_called()
@pytest.mark.asyncio
async def test_add_team_callbacks_rejects_team_deleted_before_write():
"""A team deleted between the existence check and the write must be rejected.
Prisma's update returns None for a row that is gone, and add_team_callbacks
used to hand that None to the cache refresh and report success with a null
body. The rejection reuses this endpoint's own missing-team contract, so a
caller sees the same 400 whether the team vanished before or after the read.
"""
mock_prisma = _patch_prisma(_team_row(team_id="team-1", metadata={}))
mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=None)
data = AddTeamCallback(
callback_name="langfuse",
callback_type="success",
callback_vars={
"langfuse_public_key": "pk-demo",
"langfuse_secret_key": "sk-demo",
},
)
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy_server module global is the endpoint's only injection point
patch("litellm.proxy.proxy_server.master_key", None), # test-quality-ok: proxy_server module global is the endpoint's only injection point
):
with pytest.raises(HTTPException) as exc:
await add_team_callbacks(
data=data,
http_request=MagicMock(spec=Request),
team_id="team-1",
user_api_key_dict=_admin_auth(),
)
mock_prisma.db.litellm_teamtable.update.assert_called_once()
assert exc.value.status_code == 400
assert exc.value.detail == {"error": "Team id = team-1 does not exist. Please use a different team id."}
@pytest.mark.asyncio
async def test_delete_team_callback_keeps_last_removal_from_reviving_legacy_shape():
"""Removing the last entry must leave metadata["logging"] present and empty.

96
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -3,7 +3,7 @@ import json
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Optional, cast
from typing import Final, Optional, cast
from unittest.mock import AsyncMock, MagicMock, call, patch
import pytest
@ -2078,6 +2078,100 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name):
assert update_call_kwargs.get("include", {}).get("object_permission") is True
@pytest.mark.asyncio
@pytest.mark.parametrize(
"endpoint_name",
["team_model_add", "team_model_delete", "update_team_member_permissions"],
)
async def test_team_write_404s_when_row_vanishes_before_update(endpoint_name):
"""A team deleted between the read and the write must 404.
Prisma's `update` returns None when no row matches `where`, and the team
row can be deleted between the read these endpoints do first and the
update that follows it. Without the guard, `team_model_add` /
`team_model_delete` hand that None to `_refresh_cached_team` (which
reads `team_row.team_id`) and `/team/permissions_update` returns None
out of a route declared to return a team, so a plain race turns into a
500 instead of the 404 every other not-found path in this file raises.
"""
from unittest.mock import AsyncMock, MagicMock, Mock, patch
from fastapi import Request
from litellm.proxy._types import (
LitellmUserRoles,
TeamModelAddRequest,
TeamModelDeleteRequest,
UserAPIKeyAuth,
)
from litellm.proxy.management_endpoints.team_endpoints import (
team_model_add,
team_model_delete,
update_team_member_permissions,
)
mock_request = Mock(spec=Request)
mock_user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id"
)
existing_team = MagicMock()
existing_team.team_id = "team-1234"
existing_team.model_dump.return_value = {
"team_id": "team-1234",
"models": ["bedrock-claude-sonnet-4", "openai/*"],
"team_member_permissions": [],
"spend": 0.0,
}
call_endpoint_under_test: Final = {
"team_model_add": lambda: team_model_add(
data=TeamModelAddRequest(team_id="team-1234", models=["team-byok-1"]),
http_request=mock_request,
user_api_key_dict=mock_user_api_key_dict,
),
"team_model_delete": lambda: team_model_delete(
data=TeamModelDeleteRequest(team_id="team-1234", models=["openai/*"]),
http_request=mock_request,
user_api_key_dict=mock_user_api_key_dict,
),
"update_team_member_permissions": lambda: update_team_member_permissions(
data=UpdateTeamMemberPermissionsRequest(
team_id="team-1234",
team_member_permissions=["/key/generate"],
),
http_request=mock_request,
user_api_key_dict=mock_user_api_key_dict,
),
}[endpoint_name]
with (
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, # test-quality-ok: proxy_server module global is the endpoint's only injection point
patch("litellm.proxy.proxy_server.user_api_key_cache"), # test-quality-ok: proxy_server module global is the endpoint's only injection point
patch("litellm.proxy.proxy_server.proxy_logging_obj"), # test-quality-ok: proxy_server module global is the endpoint's only injection point
patch( # test-quality-ok: stubs the cache write so the test observes only the DB result handling
"litellm.proxy.management_endpoints.team_endpoints._cache_team_object",
new_callable=AsyncMock,
),
patch( # test-quality-ok: stubs the collaborator so the test pins the endpoint's own error contract
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
new_callable=AsyncMock,
return_value=existing_team,
),
):
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
return_value=existing_team
)
mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None)
mock_prisma_client.db.execute_raw = AsyncMock(return_value=None)
with pytest.raises(HTTPException) as exc_info:
await call_endpoint_under_test()
assert exc_info.value.status_code == 404
assert exc_info.value.detail == {"error": "Team not found, passed team_id=team-1234"}
@pytest.mark.asyncio
async def test_update_team_team_member_budget_not_passed_to_db(
disable_audit_logging_for_mocked_team,

72
tests/test_litellm/proxy/memory/test_memory_endpoints.py
View file

@ -7,6 +7,7 @@ We patch the endpoint module's `_require_prisma` helper so we never need the
real proxy_server import chain (which pulls heavy optional deps).
"""
import json
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
from unittest.mock import MagicMock, patch
@ -176,14 +177,36 @@ class _InMemoryTeamTable:
return None
def _make_team(team_id: str, *, admin_user_ids: List[str]) -> MagicMock:
"""Build a team-row stub with `members_with_roles` shaped like Prisma."""
members = [MagicMock(user_id=uid, role="admin") for uid in admin_user_ids]
team = MagicMock()
team.team_id = team_id
team.organization_id = None # skip org-admin path in tests
team.members_with_roles = members
return team
def _make_team(team_id: str, *, admin_user_ids: List[str]) -> Any:
"""Build a real Prisma team row.
`members_with_roles` is a JSON column, so Prisma deserializes it into plain
dicts, not `Member` objects. A stub that hands back attribute-style members
would let the router read `member.role` off something Prisma never returns.
"""
from prisma import models as prisma_models
now = datetime.now(timezone.utc)
return prisma_models.LiteLLM_TeamTable(
team_id=team_id,
organization_id=None,
members_with_roles=json.dumps([{"user_id": uid, "role": "admin"} for uid in admin_user_ids]),
metadata="{}",
models=[],
blocked=False,
created_at=now,
updated_at=now,
spend=0.0,
model_spend="{}",
model_max_budget="{}",
admins=[],
members=[],
team_member_permissions=[],
access_group_ids=[],
policies=[],
default_team_member_models=[],
allow_team_guardrail_config=False,
)
def _make_prisma() -> MagicMock:
@ -653,6 +676,39 @@ class TestMemoryEndpoints:
assert resp.json()["value"] == "new"
assert len(table.rows) == 1
def test_put_memory_row_deleted_mid_update_returns_404(self):
"""
A concurrent DELETE landing between the visibility read and the write
makes Prisma's `update` return None. That must surface the same 404 the
read path uses, not an AttributeError bubbling out as an unhandled 500.
"""
table = self.prisma.db.litellm_memorytable
table.rows.append(
_make_row(
memory_id="m1",
key="notes",
value="old",
user_id="user-a",
team_id="team-a",
)
)
async def vanished(*_args, **_kwargs):
return None
original_update = table.update
table.update = vanished
client = _make_client(_user_auth("user-a", "team-a"))
try:
with _patch_prisma(self.prisma):
resp = client.put("/v1/memory/notes", json={"value": "new"})
finally:
table.update = original_update
assert resp.status_code == 404, resp.text
assert resp.json()["detail"] == "Memory with key 'notes' not found"
def test_put_memory_explicit_null_metadata_clears_field(self):
"""
prisma-client-python can't write a true SQL NULL to a `Json?` column

55
tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py
View file

@ -191,3 +191,58 @@ async def test_get_prompt_info_by_base_id():
response.prompt_spec.prompt_id == "test_prompt"
) # Should return base ID in spec response
assert response.prompt_spec.version == 3 # Should identify it as version 3
@pytest.mark.asyncio
async def test_patch_prompt_row_deleted_mid_update_returns_404():
"""
A concurrent delete between the version lookup and the write makes Prisma's
`update` return None. That must reuse the endpoint's existing not-found 404
contract rather than blowing up into an opaque 500.
"""
from fastapi import HTTPException
from litellm.proxy.prompts.prompt_endpoints import PatchPromptRequest, patch_prompt
mock_user_auth = UserAPIKeyAuth(
api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN
)
target_row = MagicMock()
target_row.id = "row-1"
target_row.version = 1
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_prompttable.find_many = AsyncMock(
return_value=[target_row]
)
mock_prisma_client.db.litellm_prompttable.update = AsyncMock(return_value=None)
existing_prompt = PromptSpec(
prompt_id="test_prompt.v1",
litellm_params=PromptLiteLLMParams(
prompt_id="test_prompt", prompt_integration="dotprompt"
),
prompt_info=PromptInfo(prompt_type="db"),
)
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), # test-quality-ok: proxy_server module global is the endpoint's only injection point
patch( # test-quality-ok: stubs the collaborator so the test pins the endpoint's own error contract
"litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY"
) as mock_registry,
):
mock_registry.get_prompt_by_id.return_value = existing_prompt
with pytest.raises(HTTPException) as exc_info:
await patch_prompt(
prompt_id="test_prompt",
request=PatchPromptRequest(prompt_info=PromptInfo(prompt_type="db")),
user_api_key_dict=mock_user_auth,
)
assert exc_info.value.status_code == 404
assert (
exc_info.value.detail
== "Prompt with ID test_prompt not found in environment development"
)

33
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -1850,6 +1850,39 @@ def test_add_team_models_to_all_models_excludes_other_teams_byok_with_shared_nam
assert result == {"model-a-id": {"team-a"}}
@pytest.mark.asyncio
async def test_non_admin_all_models_raises_400_when_user_row_missing():
"""
Regression test: a key whose user row no longer exists made find_unique return
None, and _check_if_model_is_team_model then dereferenced it
(`model_team_id in user_row.teams`) and raised AttributeError, surfacing as a
500. The miss must reuse the 400 "User not found" contract the neighbouring
except-branch already raises.
"""
from fastapi import HTTPException
from litellm.proxy.proxy_server import non_admin_all_models
prisma_client = MagicMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
llm_router = MagicMock()
llm_router.get_model_list.return_value = [
{"model_info": {"id": "gpt-4-model-1", "team_id": "team-a"}},
]
with pytest.raises(HTTPException) as exc_info:
await non_admin_all_models(
all_models=[],
llm_router=llm_router,
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="deleted-user"),
prisma_client=prisma_client,
)
assert exc_info.value.status_code == 400
assert exc_info.value.detail == {"error": "User not found"}
@pytest.mark.asyncio
async def test_apply_search_filter_matches_team_public_model_name():
"""

51
tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py
View file

@ -2924,6 +2924,57 @@ class TestUpdateVectorStoreAccessControlAndRedaction:
assert params["api_key"] == REDACTED_BY_LITELM_STRING
assert params["api_base"] == "https://api.openai.com/v1"
@pytest.mark.asyncio
async def test_update_row_deleted_mid_update_returns_404(self):
"""A concurrent delete between the authorization read and the write makes Prisma's
``update`` return None. That must reuse the not-found 404 contract instead of
turning an AttributeError into an opaque 500."""
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.vector_store_endpoints.management_endpoints import (
update_vector_store,
)
from litellm.types.vector_stores import VectorStoreUpdateRequest
existing_row = MagicMock()
existing_row.model_dump = MagicMock(
return_value={"vector_store_id": "vs_owned", "team_id": "team-A"}
)
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(
return_value=existing_row
)
mock_prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock(
return_value=None
)
with (
patch( # test-quality-ok: stubs the auth gate so the test exercises the not-found branch under test
"litellm.proxy.vector_store_endpoints.management_endpoints.check_feature_access_for_user",
new_callable=AsyncMock,
),
patch( # test-quality-ok: stubs the auth gate so the test exercises the not-found branch under test
"litellm.proxy.vector_store_endpoints.management_endpoints._check_vector_store_access",
new_callable=AsyncMock,
return_value=True,
),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), # test-quality-ok: proxy_server module global is the endpoint's only injection point
patch("litellm.vector_store_registry", None), # test-quality-ok: litellm module global is the only injection point for the registry
):
with pytest.raises(HTTPException) as exc_info:
await update_vector_store(
data=VectorStoreUpdateRequest(
vector_store_id="vs_owned",
vector_store_description="new desc",
),
user_api_key_dict=UserAPIKeyAuth(user_id="owner", team_id="team-A"),
)
assert exc_info.value.status_code == 404
assert exc_info.value.detail == "Vector store with ID vs_owned not found"
class TestAzureAIDocumentWritePassthroughPermission:
"""Regression tests for the Azure AI Search passthrough write mapping.

Some files were not shown because too many files have changed in this diff Show more