mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
chore(typing): clear basedpyright Any errors across proxy, rag and mcp modules
Replace Any seams in eight production modules with real types: Protocols for the untyped Prisma client actions, TypedDicts for JSON and SSE payloads, and precise parameter and return annotations on the db, endpoint and CLI helpers. Repo-wide reportAny drops 17311 -> 16858 and reportExplicitAny 5892 -> 5865, with every basedpyright rule at or below its baseline and total errors down 577. Budgets ratcheted to match.
This commit is contained in:
parent
ecba48dd7c
commit
1ac3700098
11 changed files with 686 additions and 297 deletions
|
|
@ -1,6 +1,6 @@
|
|||
{
|
||||
"reportAny": {
|
||||
"limit": 27731
|
||||
"limit": 27278
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"limit": 2626
|
||||
|
|
@ -24,7 +24,7 @@
|
|||
"limit": 19
|
||||
},
|
||||
"reportExplicitAny": {
|
||||
"limit": 8807
|
||||
"limit": 8780
|
||||
},
|
||||
"reportFunctionMemberAccess": {
|
||||
"limit": 7
|
||||
|
|
@ -54,10 +54,10 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportMissingParameterType": {
|
||||
"limit": 5835
|
||||
"limit": 5822
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15790
|
||||
"limit": 15778
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 40
|
||||
|
|
@ -99,19 +99,19 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 45063
|
||||
"limit": 45055
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 113
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 39773
|
||||
"limit": 39756
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 20207
|
||||
"limit": 20189
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 31281
|
||||
"limit": 31252
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 122
|
||||
|
|
|
|||
|
|
@ -248,7 +248,38 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
)
|
||||
return response
|
||||
|
||||
async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None:
|
||||
@staticmethod
|
||||
def _str_field(source: Mapping[str, object], key: str) -> str:
|
||||
value: Final = source.get(key)
|
||||
return value if isinstance(value, str) else ""
|
||||
|
||||
@staticmethod
|
||||
def _tools_field(kwargs: Mapping[str, object]) -> Sequence[dict[str, object]] | None:
|
||||
tools: Final = kwargs.get("tools")
|
||||
return tools if isinstance(tools, list) else None
|
||||
|
||||
@classmethod
|
||||
def _resolve_deployment_provider(cls, kwargs: Mapping[str, object]) -> str:
|
||||
"""Provider from top-level kwargs, then nested litellm_params, then the model name."""
|
||||
top_level: Final = cls._str_field(kwargs, "custom_llm_provider")
|
||||
if top_level:
|
||||
return top_level
|
||||
|
||||
litellm_params: Final = kwargs.get("litellm_params")
|
||||
if isinstance(litellm_params, dict):
|
||||
nested: Final = cls._str_field(litellm_params, "custom_llm_provider")
|
||||
if nested:
|
||||
return nested
|
||||
|
||||
try:
|
||||
_, derived_provider, _, _ = litellm.get_llm_provider(model=cls._str_field(kwargs, "model"))
|
||||
except Exception:
|
||||
return ""
|
||||
return derived_provider
|
||||
|
||||
async def async_pre_call_deployment_hook(
|
||||
self, kwargs: dict[str, object], call_type: CallTypes | None
|
||||
) -> dict | None:
|
||||
"""
|
||||
Pre-call hook to convert native Anthropic web_search tools to regular tools.
|
||||
|
||||
|
|
@ -256,21 +287,12 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
Instead, we convert it to a regular tool so the model returns tool_use blocks
|
||||
that we can intercept and execute ourselves.
|
||||
"""
|
||||
# Check if this is for an enabled provider
|
||||
# Try top-level kwargs first, then nested litellm_params, then derive from model name
|
||||
custom_llm_provider = kwargs.get("custom_llm_provider", "") or kwargs.get("litellm_params", {}).get(
|
||||
"custom_llm_provider", ""
|
||||
)
|
||||
if not custom_llm_provider:
|
||||
try:
|
||||
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs.get("model", ""))
|
||||
except Exception:
|
||||
custom_llm_provider = ""
|
||||
custom_llm_provider: Final = self._resolve_deployment_provider(kwargs)
|
||||
if custom_llm_provider not in self.enabled_providers:
|
||||
return None
|
||||
|
||||
# Check if request has tools with native web_search
|
||||
tools: Final[Sequence[dict[str, object]] | None] = kwargs.get("tools")
|
||||
tools: Final = self._tools_field(kwargs)
|
||||
if not tools:
|
||||
return None
|
||||
|
||||
|
|
@ -394,7 +416,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
return tool.get("name")
|
||||
|
||||
@classmethod
|
||||
def _sync_forced_tool_choice(cls, tool_choice: Any, converted_tools: list[dict[str, object]]) -> object:
|
||||
def _sync_forced_tool_choice(cls, tool_choice: object, converted_tools: list[dict[str, object]]) -> object:
|
||||
"""Repoint a forced ``tool_choice`` at ``litellm_web_search`` when it
|
||||
names a web-search tool that was just converted away.
|
||||
|
||||
|
|
@ -462,7 +484,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
kwargs[WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY] = True
|
||||
|
||||
# Convert native web search tools to LiteLLM standard
|
||||
converted_tools: Final = []
|
||||
converted_tools: Final[list[dict[str, object]]] = []
|
||||
for tool in tools:
|
||||
if is_web_search_tool(tool):
|
||||
standard_tool = get_litellm_web_search_tool()
|
||||
|
|
@ -482,7 +504,8 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
)
|
||||
|
||||
if "tool_choice" in kwargs:
|
||||
kwargs["tool_choice"] = self._sync_forced_tool_choice(kwargs.get("tool_choice"), converted_tools)
|
||||
forced_tool_choice: Final[object] = kwargs.get("tool_choice")
|
||||
kwargs["tool_choice"] = self._sync_forced_tool_choice(forced_tool_choice, converted_tools)
|
||||
|
||||
# Also convert here for direct callers that bypass the deployment hook.
|
||||
if kwargs.get("stream"):
|
||||
|
|
@ -836,7 +859,8 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
native_blocks: Final = plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY)
|
||||
if not native_blocks:
|
||||
return response
|
||||
return self._inject_native_blocks(response, native_blocks)
|
||||
self._inject_native_blocks(response, native_blocks)
|
||||
return response
|
||||
|
||||
@staticmethod
|
||||
def _build_native_result_blocks(
|
||||
|
|
@ -883,17 +907,17 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any:
|
||||
"""Prepend native blocks to response content, dict or object form."""
|
||||
def _inject_native_blocks(response: object, native_blocks: Sequence[Mapping[str, object]]) -> None:
|
||||
"""Prepend native blocks to response content in place, dict or object form."""
|
||||
if not native_blocks:
|
||||
return response
|
||||
return
|
||||
if isinstance(response, dict):
|
||||
existing = response.get("content") or []
|
||||
response["content"] = list(native_blocks) + list(existing)
|
||||
return response
|
||||
existing = getattr(response, "content", None) or []
|
||||
existing_items: Final = response.get("content") or []
|
||||
response["content"] = list(native_blocks) + list(existing_items)
|
||||
return
|
||||
existing_attr: Final = getattr(response, "content", None) or []
|
||||
try:
|
||||
response.content = list(native_blocks) + list(existing)
|
||||
response.__setattr__("content", list(native_blocks) + list(existing_attr))
|
||||
except (AttributeError, TypeError):
|
||||
# Object refused write — fall through and leave the response
|
||||
# untouched rather than crash the request.
|
||||
|
|
@ -901,7 +925,6 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
"WebSearchInterception: could not inject native blocks into response of type %s",
|
||||
type(response).__name__,
|
||||
)
|
||||
return response
|
||||
|
||||
async def async_run_chat_completion_agentic_loop(
|
||||
self,
|
||||
|
|
@ -1177,7 +1200,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
if max_tokens is None:
|
||||
max_tokens = cast(int, kwargs.get("max_tokens", 1024))
|
||||
|
||||
response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate(
|
||||
response: Final[AnthropicMessagesResponse | AsyncIterator[object]] = await anthropic_messages.acreate(
|
||||
max_tokens=max_tokens,
|
||||
messages=request_patch.messages,
|
||||
model=request_patch.model or model,
|
||||
|
|
@ -1193,7 +1216,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
tool_calls=tool_calls,
|
||||
structured_results=structured_results,
|
||||
)
|
||||
response = self._inject_native_blocks(response, native_blocks)
|
||||
self._inject_native_blocks(response, native_blocks)
|
||||
|
||||
return response
|
||||
|
||||
|
|
@ -1400,7 +1423,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
valid_token=user_api_key_auth,
|
||||
)
|
||||
|
||||
team_id: Final = getattr(user_api_key_auth, "team_id", None)
|
||||
team_id: Final = user_api_key_auth.team_id
|
||||
if team_id:
|
||||
from litellm.proxy.proxy_server import (
|
||||
prisma_client,
|
||||
|
|
|
|||
|
|
@ -3,8 +3,9 @@ import binascii
|
|||
import hashlib
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -45,7 +46,6 @@ from litellm.types.mcp import MCPCredentials
|
|||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_db_models
|
||||
from prisma import types as prisma_db_types
|
||||
from prisma.actions import LiteLLM_MCPUserCredentialsActions, LiteLLM_MCPUserEnvVarsActions
|
||||
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
|
|
@ -434,23 +434,150 @@ def _credentials_blob_to_mutable_dict(blob: str | Mapping[str, object]) -> dict[
|
|||
return parsed_blob
|
||||
|
||||
|
||||
def _json_loads_object(raw: str) -> dict[str, object] | None:
|
||||
"""Parse ``raw`` as JSON and return it only when it decodes to an object, else ``None``.
|
||||
|
||||
Decode errors propagate to the caller, which decides how a malformed blob is handled.
|
||||
"""
|
||||
parsed: Final[dict[str, object] | None] = json.loads(raw)
|
||||
return parsed if isinstance(parsed, dict) else None
|
||||
|
||||
|
||||
class _MCPServerTableActions(Protocol):
|
||||
async def find_many(
|
||||
self,
|
||||
*,
|
||||
where: Mapping[str, object] | None = None,
|
||||
order: Mapping[str, str] | None = None,
|
||||
take: int | None = None,
|
||||
) -> "Sequence[prisma_db_models.LiteLLM_MCPServerTable]": ...
|
||||
|
||||
async def find_unique(self, *, where: Mapping[str, object]) -> "prisma_db_models.LiteLLM_MCPServerTable | None": ...
|
||||
|
||||
async def update(
|
||||
self, *, where: Mapping[str, object], data: Mapping[str, object]
|
||||
) -> "prisma_db_models.LiteLLM_MCPServerTable": ...
|
||||
|
||||
|
||||
class _MCPServerOAuthClientActions(Protocol):
|
||||
async def find_many(
|
||||
self, *, where: Mapping[str, object] | None = None
|
||||
) -> "Sequence[prisma_db_models.LiteLLM_MCPServerOAuthClient]": ...
|
||||
|
||||
async def find_unique(
|
||||
self, *, where: Mapping[str, object]
|
||||
) -> "prisma_db_models.LiteLLM_MCPServerOAuthClient | None": ...
|
||||
|
||||
async def upsert(
|
||||
self, *, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]]
|
||||
) -> "prisma_db_models.LiteLLM_MCPServerOAuthClient": ...
|
||||
|
||||
async def update(
|
||||
self, *, where: Mapping[str, object], data: Mapping[str, object]
|
||||
) -> "prisma_db_models.LiteLLM_MCPServerOAuthClient | None": ...
|
||||
|
||||
async def delete_many(self, *, where: Mapping[str, object]) -> int: ...
|
||||
|
||||
|
||||
class _VerificationTokenActions(Protocol):
|
||||
async def find_many(
|
||||
self, *, where: Mapping[str, object] | None = None
|
||||
) -> "list[prisma_db_models.LiteLLM_VerificationToken] | None": ...
|
||||
|
||||
async def find_unique(
|
||||
self, *, where: Mapping[str, object], include: Mapping[str, bool] | None = None
|
||||
) -> "prisma_db_models.LiteLLM_VerificationToken | None": ...
|
||||
|
||||
|
||||
class _TeamTableActions(Protocol):
|
||||
async def find_unique(
|
||||
self, *, where: Mapping[str, object], include: Mapping[str, bool] | None = None
|
||||
) -> "prisma_db_models.LiteLLM_TeamTable | None": ...
|
||||
|
||||
|
||||
class _MCPUserCredentialsActions(Protocol):
|
||||
async def find_many(
|
||||
self, *, where: Mapping[str, object] | None = None
|
||||
) -> "Sequence[prisma_db_models.LiteLLM_MCPUserCredentials]": ...
|
||||
|
||||
async def find_unique(
|
||||
self, *, where: Mapping[str, object]
|
||||
) -> "prisma_db_models.LiteLLM_MCPUserCredentials | None": ...
|
||||
|
||||
async def upsert(
|
||||
self, *, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]]
|
||||
) -> "prisma_db_models.LiteLLM_MCPUserCredentials": ...
|
||||
|
||||
async def update(
|
||||
self, *, where: Mapping[str, object], data: Mapping[str, object]
|
||||
) -> "prisma_db_models.LiteLLM_MCPUserCredentials | None": ...
|
||||
|
||||
async def delete(self, *, where: Mapping[str, object]) -> "prisma_db_models.LiteLLM_MCPUserCredentials | None": ...
|
||||
|
||||
async def delete_many(self, *, where: Mapping[str, object]) -> int: ...
|
||||
|
||||
|
||||
class _MCPUserEnvVarsActions(Protocol):
|
||||
async def find_many(
|
||||
self, *, where: Mapping[str, object] | None = None
|
||||
) -> "Sequence[prisma_db_models.LiteLLM_MCPUserEnvVars]": ...
|
||||
|
||||
async def find_unique(self, *, where: Mapping[str, object]) -> "prisma_db_models.LiteLLM_MCPUserEnvVars | None": ...
|
||||
|
||||
async def upsert(
|
||||
self, *, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]]
|
||||
) -> "prisma_db_models.LiteLLM_MCPUserEnvVars": ...
|
||||
|
||||
async def update(
|
||||
self, *, where: Mapping[str, object], data: Mapping[str, object]
|
||||
) -> "prisma_db_models.LiteLLM_MCPUserEnvVars | None": ...
|
||||
|
||||
async def delete_many(self, *, where: Mapping[str, object]) -> int: ...
|
||||
|
||||
|
||||
class _MCPUserEnvVarsTransaction(Protocol):
|
||||
async def execute_raw(self, query: str, *args: object) -> int: ...
|
||||
|
||||
@property
|
||||
def litellm_mcpuserenvvars(self) -> _MCPUserEnvVarsActions: ...
|
||||
|
||||
|
||||
def _mcp_server_actions(prisma_client: PrismaClient) -> _MCPServerTableActions:
|
||||
table: Final[_MCPServerTableActions] = MCPServerRepository(prisma_client).table
|
||||
return table
|
||||
|
||||
|
||||
def _oauth_client_actions(prisma_client: PrismaClient) -> _MCPServerOAuthClientActions:
|
||||
table: Final[_MCPServerOAuthClientActions] = MCPServerOAuthClientRepository(prisma_client).table
|
||||
return table
|
||||
|
||||
|
||||
def _verification_token_actions(prisma_client: PrismaClient) -> _VerificationTokenActions:
|
||||
table: Final[_VerificationTokenActions] = VerificationTokenRepository(prisma_client).table
|
||||
return table
|
||||
|
||||
|
||||
def _team_actions(prisma_client: PrismaClient) -> _TeamTableActions:
|
||||
table: Final[_TeamTableActions] = TeamRepository(prisma_client).table
|
||||
return table
|
||||
|
||||
|
||||
def _user_env_var_tx(prisma_client: PrismaClient) -> AbstractAsyncContextManager[_MCPUserEnvVarsTransaction]:
|
||||
tx: Final[AbstractAsyncContextManager[_MCPUserEnvVarsTransaction]] = prisma_client.db.tx()
|
||||
return tx
|
||||
|
||||
|
||||
async def _db_find_mcp_server_rows(
|
||||
prisma_client: PrismaClient,
|
||||
where: "prisma_db_types.LiteLLM_MCPServerTableWhereInput | None" = None,
|
||||
) -> "list[prisma_db_models.LiteLLM_MCPServerTable]":
|
||||
rows: list[prisma_db_models.LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.find_many(
|
||||
where=where
|
||||
)
|
||||
return rows
|
||||
) -> "Sequence[prisma_db_models.LiteLLM_MCPServerTable]":
|
||||
return await _mcp_server_actions(prisma_client).find_many(where=where)
|
||||
|
||||
|
||||
async def _db_find_mcp_server_row(
|
||||
prisma_client: PrismaClient, server_id: str
|
||||
) -> "prisma_db_models.LiteLLM_MCPServerTable | None":
|
||||
row: prisma_db_models.LiteLLM_MCPServerTable | None = await MCPServerRepository(prisma_client).table.find_unique(
|
||||
where={"server_id": server_id}
|
||||
)
|
||||
return row
|
||||
return await _mcp_server_actions(prisma_client).find_unique(where={"server_id": server_id})
|
||||
|
||||
|
||||
async def _db_update_mcp_server_row(
|
||||
|
|
@ -458,28 +585,16 @@ async def _db_update_mcp_server_row(
|
|||
server_id: str,
|
||||
data: "prisma_db_types.LiteLLM_MCPServerTableUpdateInput",
|
||||
) -> "prisma_db_models.LiteLLM_MCPServerTable":
|
||||
row: Final[prisma_db_models.LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.update(
|
||||
where={"server_id": server_id},
|
||||
data=data,
|
||||
)
|
||||
return row
|
||||
return await _mcp_server_actions(prisma_client).update(where={"server_id": server_id}, data=data)
|
||||
|
||||
|
||||
def _user_credential_actions(
|
||||
prisma_client: PrismaClient,
|
||||
) -> "LiteLLM_MCPUserCredentialsActions[prisma_db_models.LiteLLM_MCPUserCredentials]":
|
||||
table: Final[LiteLLM_MCPUserCredentialsActions[prisma_db_models.LiteLLM_MCPUserCredentials]] = (
|
||||
MCPUserCredentialsRepository(prisma_client).table
|
||||
)
|
||||
def _user_credential_actions(prisma_client: PrismaClient) -> _MCPUserCredentialsActions:
|
||||
table: Final[_MCPUserCredentialsActions] = MCPUserCredentialsRepository(prisma_client).table
|
||||
return table
|
||||
|
||||
|
||||
def _user_env_var_actions(
|
||||
prisma_client: PrismaClient,
|
||||
) -> "LiteLLM_MCPUserEnvVarsActions[prisma_db_models.LiteLLM_MCPUserEnvVars]":
|
||||
table: Final[LiteLLM_MCPUserEnvVarsActions[prisma_db_models.LiteLLM_MCPUserEnvVars]] = (
|
||||
prisma_client.db.litellm_mcpuserenvvars
|
||||
)
|
||||
def _user_env_var_actions(prisma_client: PrismaClient) -> _MCPUserEnvVarsActions:
|
||||
table: Final[_MCPUserEnvVarsActions] = prisma_client.db.litellm_mcpuserenvvars
|
||||
return table
|
||||
|
||||
|
||||
|
|
@ -494,14 +609,14 @@ async def _db_find_user_credential_row(
|
|||
async def _db_find_user_credential_rows(
|
||||
prisma_client: PrismaClient,
|
||||
where: "prisma_db_types.LiteLLM_MCPUserCredentialsWhereInput | None" = None,
|
||||
) -> "list[prisma_db_models.LiteLLM_MCPUserCredentials]":
|
||||
) -> "Sequence[prisma_db_models.LiteLLM_MCPUserCredentials]":
|
||||
return await _user_credential_actions(prisma_client).find_many(where=where)
|
||||
|
||||
|
||||
async def _db_upsert_user_credential_row(
|
||||
prisma_client: PrismaClient, user_id: str, server_id: str, credential_b64: str
|
||||
) -> None:
|
||||
await MCPUserCredentialsRepository(prisma_client).table.upsert(
|
||||
await _user_credential_actions(prisma_client).upsert(
|
||||
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}},
|
||||
data={
|
||||
"create": {
|
||||
|
|
@ -517,7 +632,7 @@ async def _db_upsert_user_credential_row(
|
|||
async def _db_find_user_env_var_rows(
|
||||
prisma_client: PrismaClient,
|
||||
where: "prisma_db_types.LiteLLM_MCPUserEnvVarsWhereInput | None" = None,
|
||||
) -> "list[prisma_db_models.LiteLLM_MCPUserEnvVars]":
|
||||
) -> "Sequence[prisma_db_models.LiteLLM_MCPUserEnvVars]":
|
||||
return await _user_env_var_actions(prisma_client).find_many(where=where)
|
||||
|
||||
|
||||
|
|
@ -585,9 +700,7 @@ async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str]
|
|||
"""
|
||||
Returns the matching mcp servers from the db with the server_ids
|
||||
"""
|
||||
_mcp_servers: Final[list[prisma_db_models.LiteLLM_MCPServerTable]] = await MCPServerRepository(
|
||||
prisma_client
|
||||
).table.find_many(
|
||||
_mcp_servers: Final = await _mcp_server_actions(prisma_client).find_many(
|
||||
where={
|
||||
"server_id": {"in": server_ids},
|
||||
}
|
||||
|
|
@ -605,9 +718,7 @@ async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, toke
|
|||
"""
|
||||
Returns the mcp servers from the db for the verification token
|
||||
"""
|
||||
verification_token_record: prisma_db_models.LiteLLM_VerificationToken | None = await VerificationTokenRepository(
|
||||
prisma_client
|
||||
).table.find_unique(
|
||||
verification_token_record: Final = await _verification_token_actions(prisma_client).find_unique(
|
||||
where={
|
||||
"token": token,
|
||||
},
|
||||
|
|
@ -626,7 +737,7 @@ async def get_mcp_servers_by_team(prisma_client: PrismaClient, team_id: str) ->
|
|||
"""
|
||||
Returns the mcp servers from the db for the team id
|
||||
"""
|
||||
team_record: prisma_db_models.LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.find_unique(
|
||||
team_record: Final = await _team_actions(prisma_client).find_unique(
|
||||
where={
|
||||
"team_id": team_id,
|
||||
},
|
||||
|
|
@ -697,9 +808,9 @@ async def get_virtualkeys_for_mcp_server(
|
|||
"""
|
||||
Get all the virtual keys that have access to the mcp server
|
||||
"""
|
||||
virtual_keys: Final[list[prisma_db_models.LiteLLM_VerificationToken] | None] = await VerificationTokenRepository(
|
||||
virtual_keys: Final[list[prisma_db_models.LiteLLM_VerificationToken] | None] = await _verification_token_actions(
|
||||
prisma_client
|
||||
).table.find_many(
|
||||
).find_many(
|
||||
where={
|
||||
"mcp_servers": {"has": server_id},
|
||||
},
|
||||
|
|
@ -753,9 +864,9 @@ async def delete_mcp_server(
|
|||
if deleted_server is not None:
|
||||
credential_user_ids: list[str] = []
|
||||
try:
|
||||
credential_rows: Sequence[
|
||||
prisma_db_models.LiteLLM_MCPUserCredentials
|
||||
] = await prisma_client.db.litellm_mcpusercredentials.find_many(where={"server_id": server_id})
|
||||
credential_rows: Final[
|
||||
Sequence[prisma_db_models.LiteLLM_MCPUserCredentials]
|
||||
] = await _user_credential_actions(prisma_client).find_many(where={"server_id": server_id})
|
||||
credential_user_ids = [row.user_id for row in credential_rows]
|
||||
except Exception as e: # noqa: BLE001 - enumeration is best-effort; cached tokens expire by TTL
|
||||
verbose_proxy_logger.warning(
|
||||
|
|
@ -764,9 +875,9 @@ async def delete_mcp_server(
|
|||
e,
|
||||
)
|
||||
for model, label in (
|
||||
(prisma_client.db.litellm_mcpusercredentials, "credential"),
|
||||
(prisma_client.db.litellm_mcpuserenvvars, "env var"),
|
||||
(prisma_client.db.litellm_mcpserveroauthclient, "OAuth client"),
|
||||
(_user_credential_actions(prisma_client), "credential"),
|
||||
(_user_env_var_actions(prisma_client), "env var"),
|
||||
(_oauth_client_actions(prisma_client), "OAuth client"),
|
||||
):
|
||||
try:
|
||||
await model.delete_many(where={"server_id": server_id})
|
||||
|
|
@ -945,9 +1056,7 @@ async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, s
|
|||
LiteLLM_MCPServerTable row, so their dynamically registered client lives here keyed
|
||||
by server_id. The returned value is the raw credentials blob for
|
||||
``_get_persisted_dcr_credentials`` to parse."""
|
||||
row: Final[prisma_db_models.LiteLLM_MCPServerOAuthClient | None] = await MCPServerOAuthClientRepository(
|
||||
prisma_client
|
||||
).table.find_unique(where={"server_id": server_id})
|
||||
row: Final = await _oauth_client_actions(prisma_client).find_unique(where={"server_id": server_id})
|
||||
if row is None:
|
||||
return None
|
||||
return row.credentials
|
||||
|
|
@ -965,7 +1074,7 @@ async def upsert_mcp_server_oauth_client_credentials(
|
|||
|
||||
encrypted: Final = encrypt_credentials(credentials=MCPCredentials(**credentials), encryption_key=_get_salt_key())
|
||||
blob: Final = safe_dumps(encrypted)
|
||||
await MCPServerOAuthClientRepository(prisma_client).table.upsert(
|
||||
await _oauth_client_actions(prisma_client).upsert(
|
||||
where={"server_id": server_id},
|
||||
data={
|
||||
"create": {"server_id": server_id, "credentials": blob},
|
||||
|
|
@ -1012,21 +1121,19 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient,
|
|||
continue
|
||||
|
||||
update_data["updated_by"] = touched_by
|
||||
await MCPServerRepository(prisma_client).table.update(
|
||||
await _mcp_server_actions(prisma_client).update(
|
||||
where={"server_id": mcp_server.server_id},
|
||||
data=update_data,
|
||||
)
|
||||
updated += 1
|
||||
|
||||
oauth_clients: Final[list[prisma_db_models.LiteLLM_MCPServerOAuthClient]] = await MCPServerOAuthClientRepository(
|
||||
prisma_client
|
||||
).table.find_many()
|
||||
oauth_clients: Final = await _oauth_client_actions(prisma_client).find_many()
|
||||
oauth_updated = 0
|
||||
for oauth_client in oauth_clients:
|
||||
rotated_credentials = _reencrypt_mcp_credentials_blob(oauth_client.credentials, new_master_key)
|
||||
if rotated_credentials is None:
|
||||
continue
|
||||
await MCPServerOAuthClientRepository(prisma_client).table.update(
|
||||
await _oauth_client_actions(prisma_client).update(
|
||||
where={"server_id": oauth_client.server_id},
|
||||
data={"credentials": rotated_credentials},
|
||||
)
|
||||
|
|
@ -1347,7 +1454,7 @@ def mcp_oauth_token_identity(server: object) -> tuple[object, ...]:
|
|||
creds: Final = getattr(server, "credentials", None)
|
||||
if isinstance(creds, str):
|
||||
try:
|
||||
parsed: dict[str, object] | None = json.loads(creds)
|
||||
parsed: dict[str, object] | None = _json_loads_object(creds)
|
||||
except ValueError:
|
||||
parsed = None
|
||||
else:
|
||||
|
|
@ -1716,7 +1823,7 @@ async def get_mcp_submissions(
|
|||
along with a summary count breakdown by approval_status.
|
||||
Mirrors get_guardrail_submissions() from guardrail_endpoints.py.
|
||||
"""
|
||||
rows: list[prisma_db_models.LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.find_many(
|
||||
rows: Final = await _mcp_server_actions(prisma_client).find_many(
|
||||
where={"submitted_at": {"not": None}},
|
||||
order={"submitted_at": "desc"},
|
||||
take=500, # safety cap; paginate if needed in a future iteration
|
||||
|
|
@ -1757,12 +1864,11 @@ def _decode_user_env_vars(stored: str) -> dict[str, str]:
|
|||
"re-enter them rather than silently forwarding ciphertext"
|
||||
)
|
||||
return {}
|
||||
parsed: dict[str, object] | None
|
||||
try:
|
||||
parsed = json.loads(decrypted)
|
||||
parsed: Final = _json_loads_object(decrypted)
|
||||
except (ValueError, TypeError):
|
||||
return {}
|
||||
if not isinstance(parsed, dict):
|
||||
if parsed is None:
|
||||
return {}
|
||||
return {str(k): str(v) for k, v in parsed.items()}
|
||||
|
||||
|
|
@ -1818,9 +1924,9 @@ async def merge_user_env_vars(
|
|||
"big",
|
||||
signed=True,
|
||||
)
|
||||
async with prisma_client.db.tx() as tx:
|
||||
async with _user_env_var_tx(prisma_client) as tx:
|
||||
await tx.execute_raw("SELECT pg_advisory_xact_lock($1::bigint)", lock_key)
|
||||
row: Final[prisma_db_models.LiteLLM_MCPUserEnvVars | None] = await tx.litellm_mcpuserenvvars.find_unique(
|
||||
row: Final = await tx.litellm_mcpuserenvvars.find_unique(
|
||||
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}
|
||||
)
|
||||
existing: Final = _decode_user_env_vars(row.values_b64) if row is not None else {}
|
||||
|
|
|
|||
|
|
@ -11,10 +11,13 @@ MCP Spec Reference:
|
|||
"""
|
||||
|
||||
import typing
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Mapping, Sequence, Sized
|
||||
from typing import Any, Final, NamedTuple, Optional, Protocol, Union, runtime_checkable
|
||||
|
||||
if typing.TYPE_CHECKING:
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import TypeAlias
|
||||
|
||||
from fastapi import Request
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.shared.context import RequestContext
|
||||
|
|
@ -28,8 +31,12 @@ if typing.TYPE_CHECKING:
|
|||
ToolUseContent,
|
||||
)
|
||||
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
_AcompletionCallable: TypeAlias = Callable[..., Awaitable[ModelResponse | CustomStreamWrapper]]
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import TypeAdapter
|
||||
|
|
@ -92,7 +99,7 @@ def _resolve_model_from_preferences(
|
|||
if not available_model_names and litellm.model_list:
|
||||
for entry in litellm.model_list:
|
||||
if isinstance(entry, dict):
|
||||
name = entry.get("model_name")
|
||||
name: str | None = entry.get("model_name")
|
||||
if name:
|
||||
available_model_names.append(name)
|
||||
elif isinstance(entry, str):
|
||||
|
|
@ -439,9 +446,9 @@ def _convert_mcp_messages_to_openai(
|
|||
)
|
||||
|
||||
# Separate marker items from regular content parts
|
||||
tool_call_markers = []
|
||||
tool_result_markers = []
|
||||
regular_parts = []
|
||||
tool_call_markers: list[Mapping[str, object]] = []
|
||||
tool_result_markers: list[Mapping[str, object]] = []
|
||||
regular_parts: list[Mapping[str, object]] = []
|
||||
for part in converted_parts:
|
||||
marker = part.get("_marker_type") if isinstance(part, dict) else None
|
||||
if marker == "tool_use":
|
||||
|
|
@ -520,7 +527,7 @@ def _extract_text_parts(
|
|||
) -> str | None:
|
||||
"""Extract text parts from mixed content."""
|
||||
items: Final = content if isinstance(content, list) else [content]
|
||||
texts: Final = []
|
||||
texts: Final[list[str]] = []
|
||||
for item in items:
|
||||
if getattr(item, "type", None) == "text":
|
||||
texts.append(getattr(item, "text", ""))
|
||||
|
|
@ -1141,7 +1148,7 @@ async def _build_completion_kwargs(
|
|||
user_api_key_auth: "UserAPIKeyAuth",
|
||||
raw_headers: dict[str, str] | None,
|
||||
client_ip: str | None,
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
openai_messages: Final = _convert_mcp_messages_to_openai(
|
||||
messages=params.messages,
|
||||
system_prompt=params.systemPrompt,
|
||||
|
|
@ -1176,8 +1183,21 @@ async def _build_completion_kwargs(
|
|||
)
|
||||
|
||||
|
||||
async def _call_acompletion(
|
||||
acompletion: "_AcompletionCallable",
|
||||
completion_kwargs: dict[str, object],
|
||||
) -> "ModelResponse | CustomStreamWrapper":
|
||||
"""Invoke a completion entrypoint with the dynamically assembled proxy payload.
|
||||
|
||||
The payload's keys are only known at runtime (MCP params plus whatever
|
||||
``add_litellm_data_to_request`` and the pre-call hooks inject), so it is
|
||||
forwarded through a signature-erased callable.
|
||||
"""
|
||||
return await acompletion(**completion_kwargs)
|
||||
|
||||
|
||||
async def _run_guardrails_and_call_llm(
|
||||
completion_kwargs: dict[str, Any],
|
||||
completion_kwargs: dict[str, object],
|
||||
user_api_key_auth: "UserAPIKeyAuth",
|
||||
) -> Any:
|
||||
try:
|
||||
|
|
@ -1204,10 +1224,10 @@ async def _run_guardrails_and_call_llm(
|
|||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
if llm_router is not None:
|
||||
return await llm_router.acompletion(**completion_kwargs)
|
||||
return await litellm.acompletion(**completion_kwargs)
|
||||
return await _call_acompletion(llm_router.acompletion, completion_kwargs)
|
||||
return await _call_acompletion(litellm.acompletion, completion_kwargs)
|
||||
except ImportError:
|
||||
return await litellm.acompletion(**completion_kwargs)
|
||||
return await _call_acompletion(litellm.acompletion, completion_kwargs)
|
||||
|
||||
|
||||
async def handle_sampling_create_message(
|
||||
|
|
@ -1284,12 +1304,12 @@ async def handle_sampling_create_message(
|
|||
client_ip=client_ip,
|
||||
)
|
||||
|
||||
openai_messages: Final[Sequence[Mapping[str, object]]] = completion_kwargs["messages"]
|
||||
openai_messages: Final = completion_kwargs["messages"]
|
||||
openai_tools: Final = completion_kwargs.get("tools")
|
||||
verbose_logger.debug(
|
||||
"MCP sampling: calling litellm.acompletion with model=%s, num_messages=%d, has_tools=%s",
|
||||
model,
|
||||
len(openai_messages),
|
||||
len(openai_messages) if isinstance(openai_messages, Sized) else 0,
|
||||
bool(openai_tools),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -3,8 +3,9 @@ import os
|
|||
import sys
|
||||
import time
|
||||
import webbrowser
|
||||
from collections.abc import Mapping, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any, Final
|
||||
from typing import Any, Final, TypedDict
|
||||
from urllib.parse import urlencode
|
||||
|
||||
import click
|
||||
|
|
@ -18,6 +19,42 @@ from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh
|
|||
from .private_json import write_private_json
|
||||
|
||||
|
||||
class CliContext(TypedDict):
|
||||
base_url: str
|
||||
api_key: str | None
|
||||
base_url_explicit: bool
|
||||
|
||||
|
||||
class TeamRecord(TypedDict, total=False):
|
||||
team_id: str | None
|
||||
team_alias: str | None
|
||||
models: list[str]
|
||||
max_budget: float | None
|
||||
|
||||
|
||||
class CliSsoStartResponse(TypedDict):
|
||||
login_id: str
|
||||
poll_secret: str
|
||||
user_code: str
|
||||
|
||||
|
||||
class PollResponse(TypedDict, total=False):
|
||||
status: str
|
||||
key: str
|
||||
user_id: str
|
||||
team_id: str
|
||||
teams: list[str]
|
||||
team_details: list[dict[str, str]]
|
||||
requires_team_selection: bool
|
||||
|
||||
|
||||
class AuthResult(TypedDict):
|
||||
api_key: str
|
||||
user_id: str | None
|
||||
teams: list[str]
|
||||
team_id: str | None
|
||||
|
||||
|
||||
# Token storage utilities
|
||||
def get_token_file_path() -> str:
|
||||
"""Get the path to store the authentication token"""
|
||||
|
|
@ -27,7 +64,7 @@ def get_token_file_path() -> str:
|
|||
return str(config_dir / "token.json")
|
||||
|
||||
|
||||
def save_token(token_data: dict[str, Any]) -> None:
|
||||
def save_token(token_data: Mapping[str, object]) -> None:
|
||||
"""Save token data to file"""
|
||||
write_private_json(get_token_file_path(), token_data)
|
||||
|
||||
|
|
@ -65,7 +102,7 @@ def get_stored_api_key(expected_base_url: str | None = None) -> str | None:
|
|||
|
||||
|
||||
# Team selection utilities
|
||||
def display_teams_table(teams: list[dict[str, Any]]) -> None:
|
||||
def display_teams_table(teams: Sequence[TeamRecord]) -> None:
|
||||
"""Display teams in a formatted table"""
|
||||
console: Final = Console()
|
||||
|
||||
|
|
@ -249,10 +286,11 @@ def prompt_team_selection_fallback(
|
|||
|
||||
while True:
|
||||
try:
|
||||
choice = click.prompt(
|
||||
raw_choice: str = click.prompt(
|
||||
"\nSelect a team by entering the index number (or 'skip' to continue without a team)",
|
||||
type=str,
|
||||
).strip()
|
||||
)
|
||||
choice = raw_choice.strip()
|
||||
|
||||
if choice.lower() == "skip":
|
||||
return None
|
||||
|
|
@ -275,7 +313,7 @@ def prompt_team_selection_fallback(
|
|||
|
||||
def _response_error_detail(response: requests.Response) -> str | None:
|
||||
try:
|
||||
body: Final = response.json()
|
||||
body: Final[Mapping[str, object]] = response.json()
|
||||
except ValueError:
|
||||
return None
|
||||
detail: Final = body.get("detail") if isinstance(body, dict) else None
|
||||
|
|
@ -309,15 +347,16 @@ def _poll_for_ready_data(
|
|||
other_status_log_every: int = 10,
|
||||
http_error_log_every: int = 10,
|
||||
connection_error_log_every: int = 10,
|
||||
) -> dict[str, Any] | None:
|
||||
) -> PollResponse | None:
|
||||
for attempt in range(total_timeout // poll_interval):
|
||||
try:
|
||||
request_kwargs: dict[str, Any] = {"timeout": request_timeout}
|
||||
if headers is not None:
|
||||
request_kwargs["headers"] = headers
|
||||
response = requests.get(url, **request_kwargs)
|
||||
response = (
|
||||
requests.get(url, timeout=request_timeout)
|
||||
if headers is None
|
||||
else requests.get(url, timeout=request_timeout, headers=headers)
|
||||
)
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
data: PollResponse = response.json()
|
||||
status = data.get("status")
|
||||
if status == "ready":
|
||||
return data
|
||||
|
|
@ -341,7 +380,7 @@ def _poll_for_ready_data(
|
|||
return None
|
||||
|
||||
|
||||
def _normalize_teams(teams, team_details):
|
||||
def _normalize_teams(teams: Sequence[str], team_details: Sequence[Mapping[str, str]] | None) -> list[TeamRecord]:
|
||||
"""If team_details are a
|
||||
|
||||
Args:
|
||||
|
|
@ -365,7 +404,7 @@ def _normalize_teams(teams, team_details):
|
|||
return []
|
||||
|
||||
|
||||
def _start_cli_sso_flow(base_url: str) -> dict[str, Any]:
|
||||
def _start_cli_sso_flow(base_url: str) -> CliSsoStartResponse:
|
||||
start_url: Final = f"{base_url}/sso/cli/start"
|
||||
try:
|
||||
response: Final = requests.post(start_url, timeout=10)
|
||||
|
|
@ -389,7 +428,7 @@ def _start_cli_sso_flow(base_url: str) -> dict[str, Any]:
|
|||
)
|
||||
|
||||
try:
|
||||
data: Final = response.json()
|
||||
data: Final[CliSsoStartResponse] = response.json()
|
||||
except ValueError:
|
||||
content_type: Final = response.headers.get("content-type", "unknown")
|
||||
raise ValueError(
|
||||
|
|
@ -398,8 +437,10 @@ def _start_cli_sso_flow(base_url: str) -> dict[str, Any]:
|
|||
f"Response starts with: {response.text[:200]!r}"
|
||||
)
|
||||
|
||||
required_fields: Final = ("login_id", "poll_secret", "user_code")
|
||||
missing_fields: Final = tuple(field for field in required_fields if not isinstance(data.get(field), str))
|
||||
string_fields: Final = frozenset(name for name, value in data.items() if isinstance(value, str))
|
||||
missing_fields: Final = tuple(
|
||||
field for field in ("login_id", "poll_secret", "user_code") if field not in string_fields
|
||||
)
|
||||
if missing_fields:
|
||||
raise ValueError(
|
||||
f"The response from {start_url} is missing required field(s): {', '.join(missing_fields)}. "
|
||||
|
|
@ -412,7 +453,7 @@ def _get_cli_sso_poll_headers(poll_secret: str) -> dict[str, str]:
|
|||
return {"x-litellm-cli-poll-secret": poll_secret}
|
||||
|
||||
|
||||
def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> dict | None:
|
||||
def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> AuthResult | None:
|
||||
"""
|
||||
Poll the server for authentication completion and handle team selection.
|
||||
|
||||
|
|
@ -431,7 +472,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> di
|
|||
teams = data.get("teams", [])
|
||||
team_details: Final = data.get("team_details")
|
||||
user_id = data.get("user_id")
|
||||
normalized_teams: Final[list[dict[str, Any]]] = _normalize_teams(teams, team_details)
|
||||
normalized_teams: Final = _normalize_teams(teams, team_details)
|
||||
if not normalized_teams:
|
||||
click.echo("Warning: No teams available for selection.")
|
||||
return None
|
||||
|
|
@ -478,7 +519,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> di
|
|||
|
||||
|
||||
def _handle_team_selection_during_polling(
|
||||
base_url: str, key_id: str, poll_secret: str, teams: list[dict[str, Any]]
|
||||
base_url: str, key_id: str, poll_secret: str, teams: Sequence[TeamRecord]
|
||||
) -> str | None:
|
||||
"""
|
||||
Handle team selection and re-poll with selected team_id.
|
||||
|
|
@ -522,7 +563,7 @@ def _handle_team_selection_during_polling(
|
|||
return None
|
||||
|
||||
|
||||
def _render_and_prompt_for_team_selection(teams: list[dict[str, Any]]) -> str | None:
|
||||
def _render_and_prompt_for_team_selection(teams: Sequence[TeamRecord]) -> str | None:
|
||||
"""Render teams table and prompt user for a team selection.
|
||||
|
||||
Returns the selected team_id as a string, or None if selection was
|
||||
|
|
@ -546,10 +587,11 @@ def _render_and_prompt_for_team_selection(teams: list[dict[str, Any]]) -> str |
|
|||
# Simple selection
|
||||
while True:
|
||||
try:
|
||||
choice = click.prompt(
|
||||
raw_choice: str = click.prompt(
|
||||
"\nSelect a team by entering the index number (or 'skip' to use first team)",
|
||||
type=str,
|
||||
).strip()
|
||||
)
|
||||
choice = raw_choice.strip()
|
||||
|
||||
if choice.lower() == "skip":
|
||||
# Default to the first team's ID if the user skips an
|
||||
|
|
@ -582,7 +624,8 @@ def login(ctx: click.Context):
|
|||
from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER
|
||||
from litellm.proxy.client.cli.interface import show_commands
|
||||
|
||||
base_url: Final = ctx.obj["base_url"]
|
||||
cli_obj: Final[CliContext] = ctx.obj
|
||||
base_url: Final = cli_obj["base_url"]
|
||||
|
||||
try:
|
||||
cli_sso_flow: Final = _start_cli_sso_flow(base_url=base_url)
|
||||
|
|
@ -666,6 +709,7 @@ def print_token(ctx: click.Context):
|
|||
expires after `LITELLM_CLI_JWT_EXPIRATION_HOURS` (default 24h); once
|
||||
expired, run `lite login` again.
|
||||
"""
|
||||
cli_obj: Final[CliContext] = ctx.obj
|
||||
token_data: Final = load_token()
|
||||
if not token_data:
|
||||
click.echo("Not authenticated. Run 'lite login'.", err=True)
|
||||
|
|
@ -675,8 +719,8 @@ def print_token(ctx: click.Context):
|
|||
# explicitly pointed us at a server, trust whichever one `lite login`
|
||||
# actually issued this token for -- that's the whole point of not
|
||||
# needing a wrapper command.
|
||||
if ctx.obj.get("base_url_explicit"):
|
||||
base_url: Final = ctx.obj["base_url"]
|
||||
if cli_obj.get("base_url_explicit"):
|
||||
base_url: Final = cli_obj["base_url"]
|
||||
if token_data.get("base_url") != base_url.rstrip("/"):
|
||||
click.echo("Not authenticated for this server. Run 'lite login'.", err=True)
|
||||
sys.exit(1)
|
||||
|
|
|
|||
|
|
@ -10,9 +10,9 @@ POST /cache/settings - Save cache settings to database
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Coroutine, Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Final
|
||||
from typing import Any, Final, Protocol
|
||||
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException
|
||||
from pydantic import BaseModel, Field
|
||||
|
|
@ -63,6 +63,35 @@ _REDACTED_VALUE: Final = "***REDACTED***"
|
|||
_URL_OVERRIDDEN_CONNECTION_FIELDS: Final[frozenset] = frozenset({"host", "port", "db", "password", "username"})
|
||||
|
||||
|
||||
class _CacheConfigRow(Protocol):
|
||||
"""The single LiteLLM_CacheConfig row, narrowed to the column this module reads."""
|
||||
|
||||
cache_settings: object
|
||||
|
||||
|
||||
class _CacheConfigTable(Protocol):
|
||||
"""The prisma actions this module calls on the LiteLLM_CacheConfig table."""
|
||||
|
||||
def find_unique(self, *, where: Mapping[str, str]) -> Coroutine[None, None, _CacheConfigRow | None]: ...
|
||||
|
||||
def upsert(
|
||||
self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, str]]
|
||||
) -> Coroutine[None, None, _CacheConfigRow]: ...
|
||||
|
||||
|
||||
class _CacheConfigTableProvider(Protocol):
|
||||
@property
|
||||
def table(self) -> _CacheConfigTable: ...
|
||||
|
||||
|
||||
def _repository_table(repository: _CacheConfigTableProvider) -> _CacheConfigTable:
|
||||
return repository.table
|
||||
|
||||
|
||||
def _cache_config_table(prisma_client: object) -> _CacheConfigTable:
|
||||
return _repository_table(CacheConfigRepository(prisma_client))
|
||||
|
||||
|
||||
def _resolve_cache_url_precedence(settings: Mapping[str, object]) -> dict[str, Any]:
|
||||
"""Return cache settings with the url-vs-discrete-fields ambiguity resolved.
|
||||
|
||||
|
|
@ -197,7 +226,7 @@ def _saved_secret_is_reusable(incoming: Mapping[str, object], saved: Mapping[str
|
|||
return True
|
||||
|
||||
|
||||
def _merge_over_saved(incoming: Mapping[str, object], saved: Mapping[str, object]) -> dict[str, Any]:
|
||||
def _merge_over_saved(incoming: Mapping[str, object], saved: Mapping[str, object]) -> dict[str, object]:
|
||||
"""Keep the stored secret behind any credential the caller echoed back redacted or omitted.
|
||||
|
||||
GET returns credentials as the marker and the form never re-prefills a
|
||||
|
|
@ -339,7 +368,7 @@ class CacheSettingsManager:
|
|||
return normalized1 == normalized2
|
||||
|
||||
@staticmethod
|
||||
async def init_cache_settings_in_db(prisma_client, proxy_config):
|
||||
async def init_cache_settings_in_db(prisma_client: object, proxy_config):
|
||||
"""
|
||||
Initialize cache settings from database into the router on startup.
|
||||
Only reinitializes if cache params have changed.
|
||||
|
|
@ -347,18 +376,17 @@ class CacheSettingsManager:
|
|||
import json
|
||||
|
||||
try:
|
||||
cache_config: Final = await call_with_db_reconnect_retry(
|
||||
cache_config: Final[_CacheConfigRow | None] = await call_with_db_reconnect_retry(
|
||||
prisma_client,
|
||||
lambda: CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"}),
|
||||
lambda: _cache_config_table(prisma_client).find_unique(where={"id": "cache_config"}),
|
||||
reason="init_cache_settings_in_db_lookup_failure",
|
||||
)
|
||||
if cache_config is not None and cache_config.cache_settings:
|
||||
# Parse cache settings JSON
|
||||
cache_settings_json: Final = cache_config.cache_settings
|
||||
if isinstance(cache_settings_json, str):
|
||||
cache_settings_dict = json.loads(cache_settings_json)
|
||||
else:
|
||||
cache_settings_dict = cache_settings_json
|
||||
cache_settings_json: Final[object] = cache_config.cache_settings
|
||||
cache_settings_dict: Final[object] = (
|
||||
json.loads(cache_settings_json) if isinstance(cache_settings_json, str) else cache_settings_json
|
||||
)
|
||||
|
||||
# Decrypt cache settings
|
||||
decrypted_settings: Final = proxy_config._decrypt_db_variables(variables_dict=cache_settings_dict)
|
||||
|
|
@ -444,7 +472,9 @@ async def get_cache_settings(
|
|||
# Read the stored settings (decrypted); an env-only cache has none.
|
||||
stored: dict[str, object] = {}
|
||||
if prisma_client is not None:
|
||||
cache_config = await CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"})
|
||||
cache_config: Final[_CacheConfigRow | None] = await _cache_config_table(prisma_client).find_unique(
|
||||
where={"id": "cache_config"}
|
||||
)
|
||||
if cache_config is not None and cache_config.cache_settings:
|
||||
stored = proxy_config._decrypt_db_variables(
|
||||
variables_dict=_parse_stored_settings(cache_config.cache_settings)
|
||||
|
|
@ -511,7 +541,7 @@ async def test_cache_connection(
|
|||
saved_settings: dict[str, object] = {}
|
||||
if prisma_client is not None:
|
||||
try:
|
||||
existing_row: Final = await CacheConfigRepository(prisma_client).table.find_unique(
|
||||
existing_row: Final[_CacheConfigRow | None] = await _cache_config_table(prisma_client).find_unique(
|
||||
where={"id": "cache_config"}
|
||||
)
|
||||
if existing_row is not None and existing_row.cache_settings:
|
||||
|
|
@ -590,7 +620,9 @@ async def update_cache_settings(
|
|||
try:
|
||||
# Read the stored row first: its decrypted values back any credential the
|
||||
# caller echoed back redacted, and its key set drives the audit diff.
|
||||
existing_row: Final = await CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"})
|
||||
existing_row: Final[_CacheConfigRow | None] = await _cache_config_table(prisma_client).find_unique(
|
||||
where={"id": "cache_config"}
|
||||
)
|
||||
before_settings: dict[str, object] | None = None
|
||||
saved_settings: dict[str, object] = {}
|
||||
if existing_row is not None and existing_row.cache_settings:
|
||||
|
|
@ -606,7 +638,7 @@ async def update_cache_settings(
|
|||
encrypted_settings: Final = proxy_config._encrypt_env_variables(environment_variables=cache_settings)
|
||||
|
||||
# Save to database
|
||||
await CacheConfigRepository(prisma_client).table.upsert(
|
||||
await _cache_config_table(prisma_client).upsert(
|
||||
where={"id": "cache_config"},
|
||||
data={
|
||||
"create": {
|
||||
|
|
|
|||
|
|
@ -18,11 +18,14 @@ Scoping:
|
|||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Final
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.models.team import LiteLLM_TeamTable
|
||||
from litellm.proxy._types import (
|
||||
CommonProxyErrors,
|
||||
LitellmUserRoles,
|
||||
|
|
@ -40,10 +43,57 @@ from litellm.types.memory_management import (
|
|||
MemoryUpdateRequest,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
|
||||
def _serialize_metadata_for_prisma(metadata: Any) -> str:
|
||||
class _MemoryRecord(Protocol):
|
||||
memory_id: str
|
||||
key: str
|
||||
value: str
|
||||
user_id: str | None
|
||||
team_id: str | None
|
||||
created_at: datetime | None
|
||||
created_by: str | None
|
||||
updated_at: datetime | None
|
||||
updated_by: str | None
|
||||
|
||||
|
||||
class _MemoryTable(Protocol):
|
||||
async def create(self, data: Mapping[str, object]) -> _MemoryRecord: ...
|
||||
|
||||
async def count(self, where: Mapping[str, object]) -> int: ...
|
||||
|
||||
async def find_many(
|
||||
self,
|
||||
where: Mapping[str, object],
|
||||
order: Mapping[str, str],
|
||||
take: int,
|
||||
skip: int = 0,
|
||||
) -> list[_MemoryRecord]: ...
|
||||
|
||||
async def update(self, where: Mapping[str, str], data: Mapping[str, object]) -> _MemoryRecord: ...
|
||||
|
||||
async def delete(self, where: Mapping[str, str]) -> _MemoryRecord | None: ...
|
||||
|
||||
|
||||
class _TeamTable(Protocol):
|
||||
async def find_unique(self, where: Mapping[str, str]) -> LiteLLM_TeamTable | None: ...
|
||||
|
||||
|
||||
def _memory_table(prisma_client: "PrismaClient") -> _MemoryTable:
|
||||
table: Final[_MemoryTable] = MemoryRepository(prisma_client).table
|
||||
return table
|
||||
|
||||
|
||||
def _team_table(prisma_client: "PrismaClient") -> _TeamTable:
|
||||
table: Final[_TeamTable] = TeamRepository(prisma_client).table
|
||||
return table
|
||||
|
||||
|
||||
def _serialize_metadata_for_prisma(metadata: object) -> str:
|
||||
"""
|
||||
Encode a `metadata` payload for the `Json?` column.
|
||||
|
||||
|
|
@ -62,25 +112,25 @@ def _is_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:
|
|||
return user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
|
||||
|
||||
def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> dict | None:
|
||||
def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object] | None:
|
||||
"""
|
||||
Prisma `where` fragment restricting rows to those the caller can see.
|
||||
Returns None for admins (no restriction).
|
||||
"""
|
||||
if user_api_key_has_admin_view(user_api_key_dict):
|
||||
return None
|
||||
ors: Final[list[dict]] = []
|
||||
if user_api_key_dict.user_id:
|
||||
ors.append({"user_id": user_api_key_dict.user_id})
|
||||
if user_api_key_dict.team_id:
|
||||
ors.append({"team_id": user_api_key_dict.team_id})
|
||||
ors: Final[list[Mapping[str, str]]] = [
|
||||
{field: value}
|
||||
for field, value in (("user_id", user_api_key_dict.user_id), ("team_id", user_api_key_dict.team_id))
|
||||
if value
|
||||
]
|
||||
if not ors:
|
||||
# Caller has neither user_id nor team_id — match nothing.
|
||||
return {"memory_id": "__no_match__"}
|
||||
return {"OR": ors}
|
||||
|
||||
|
||||
def _row_to_model(row: Any) -> LiteLLM_MemoryRow:
|
||||
def _row_to_model(row: _MemoryRecord) -> LiteLLM_MemoryRow:
|
||||
return LiteLLM_MemoryRow(
|
||||
memory_id=row.memory_id,
|
||||
key=row.key,
|
||||
|
|
@ -95,7 +145,7 @@ def _row_to_model(row: Any) -> LiteLLM_MemoryRow:
|
|||
)
|
||||
|
||||
|
||||
def _require_prisma():
|
||||
def _require_prisma() -> "PrismaClient":
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
|
|
@ -113,7 +163,9 @@ def _internal_error(log_message: str, exc: Exception, default_detail: str) -> HT
|
|||
return HTTPException(status_code=500, detail=default_detail)
|
||||
|
||||
|
||||
async def _assert_write_access(prisma_client: Any, row: Any, user_api_key_dict: UserAPIKeyAuth) -> None:
|
||||
async def _assert_write_access(
|
||||
prisma_client: "PrismaClient", row: _MemoryRecord, user_api_key_dict: UserAPIKeyAuth
|
||||
) -> None:
|
||||
"""
|
||||
Enforce ownership for mutations (PUT/DELETE).
|
||||
|
||||
|
|
@ -135,8 +187,8 @@ async def _assert_write_access(prisma_client: Any, row: Any, user_api_key_dict:
|
|||
"""
|
||||
if _is_admin(user_api_key_dict):
|
||||
return
|
||||
row_user_id: Final = getattr(row, "user_id", None)
|
||||
row_team_id: Final = getattr(row, "team_id", None)
|
||||
row_user_id: Final[str | None] = getattr(row, "user_id", None)
|
||||
row_team_id: Final[str | None] = getattr(row, "team_id", None)
|
||||
|
||||
# Personal ownership.
|
||||
if row_user_id and row_user_id == user_api_key_dict.user_id:
|
||||
|
|
@ -153,7 +205,7 @@ async def _assert_write_access(prisma_client: Any, row: Any, user_api_key_dict:
|
|||
)
|
||||
|
||||
|
||||
async def _is_team_admin_for(prisma_client: Any, user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool:
|
||||
async def _is_team_admin_for(prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool:
|
||||
"""
|
||||
True if the caller is a team admin of `team_id`, or an org admin for the
|
||||
team's organization. Mirrors the auth pattern used by team-management
|
||||
|
|
@ -168,7 +220,7 @@ async def _is_team_admin_for(prisma_client: Any, user_api_key_dict: UserAPIKeyAu
|
|||
)
|
||||
|
||||
try:
|
||||
team_obj: Final = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id})
|
||||
team_obj: Final = await _team_table(prisma_client).find_unique(where={"team_id": team_id})
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error loading team for write-auth check (team_id=%s): %s", team_id, e)
|
||||
return False
|
||||
|
|
@ -261,7 +313,7 @@ def _resolve_scope(
|
|||
async def create_memory(
|
||||
body: MemoryCreateRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
) -> LiteLLM_MemoryRow:
|
||||
"""Create a new memory entry for the caller (or, for admins, any scope)."""
|
||||
prisma_client: Final = _require_prisma()
|
||||
user_id, team_id = _resolve_scope(user_api_key_dict, body.user_id, body.team_id)
|
||||
|
|
@ -269,19 +321,19 @@ async def create_memory(
|
|||
# `metadata` is a `Json?` column — prisma-client-python rejects raw
|
||||
# Python values, so JSON-encode any non-null payload and omit the field
|
||||
# entirely when None so the column defaults to SQL NULL.
|
||||
create_data: Final[dict] = {
|
||||
metadata: Final[object] = body.metadata
|
||||
create_data: Final[Mapping[str, object]] = {
|
||||
"key": body.key,
|
||||
"value": body.value,
|
||||
"user_id": user_id,
|
||||
"team_id": team_id,
|
||||
"created_by": user_api_key_dict.user_id,
|
||||
"updated_by": user_api_key_dict.user_id,
|
||||
**({"metadata": _serialize_metadata_for_prisma(metadata)} if metadata is not None else {}),
|
||||
}
|
||||
if body.metadata is not None:
|
||||
create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata)
|
||||
|
||||
try:
|
||||
row: Final = await MemoryRepository(prisma_client).table.create(data=create_data)
|
||||
row: Final = await _memory_table(prisma_client).create(data=create_data)
|
||||
except Exception as e:
|
||||
# Key is globally unique. Any duplicate → 409.
|
||||
if _is_unique_violation(e):
|
||||
|
|
@ -316,7 +368,7 @@ async def list_memory(
|
|||
page: int = Query(1, ge=1),
|
||||
page_size: int = Query(50, ge=1, le=500),
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
) -> MemoryListResponse:
|
||||
"""List memory entries visible to the caller."""
|
||||
prisma_client: Final = _require_prisma()
|
||||
|
||||
|
|
@ -325,24 +377,18 @@ async def list_memory(
|
|||
# top-level "AND" — safer than `dict.update` since future visibility
|
||||
# filters could grow an "OR" key that would clobber this one if merged
|
||||
# by key.
|
||||
key_filter: Final[dict] = {}
|
||||
if key_prefix is not None:
|
||||
key_filter["key"] = {"startsWith": key_prefix}
|
||||
elif key is not None:
|
||||
key_filter["key"] = key
|
||||
key_filter: Final[Mapping[str, object]] = (
|
||||
{"key": {"startsWith": key_prefix}} if key_prefix is not None else {"key": key} if key is not None else {}
|
||||
)
|
||||
|
||||
vis: Final = _visibility_filter(user_api_key_dict)
|
||||
where: dict
|
||||
if vis is None:
|
||||
where = key_filter
|
||||
elif not key_filter:
|
||||
where = vis
|
||||
else:
|
||||
where = {"AND": [key_filter, vis]}
|
||||
where: Final[Mapping[str, object]] = (
|
||||
key_filter if vis is None else vis if not key_filter else {"AND": [key_filter, vis]}
|
||||
)
|
||||
|
||||
try:
|
||||
total: Final = await MemoryRepository(prisma_client).table.count(where=where)
|
||||
rows: Final = await MemoryRepository(prisma_client).table.find_many(
|
||||
total: Final = await _memory_table(prisma_client).count(where=where)
|
||||
rows: Final = await _memory_table(prisma_client).find_many(
|
||||
where=where,
|
||||
order={"updated_at": "desc"},
|
||||
skip=(page - 1) * page_size,
|
||||
|
|
@ -354,12 +400,14 @@ async def list_memory(
|
|||
return MemoryListResponse(memories=[_row_to_model(r) for r in rows], total=total)
|
||||
|
||||
|
||||
async def _find_memory_for_caller(prisma_client: Any, key: str, user_api_key_dict: UserAPIKeyAuth) -> Any:
|
||||
async def _find_memory_for_caller(
|
||||
prisma_client: "PrismaClient", key: str, user_api_key_dict: UserAPIKeyAuth
|
||||
) -> _MemoryRecord:
|
||||
"""Look up a memory row by key, scoped to the caller's visibility."""
|
||||
key_filter: Final[dict] = {"key": key}
|
||||
key_filter: Final[Mapping[str, object]] = {"key": key}
|
||||
vis: Final = _visibility_filter(user_api_key_dict)
|
||||
where: Final[dict] = key_filter if vis is None else {"AND": [key_filter, vis]}
|
||||
rows = await MemoryRepository(prisma_client).table.find_many(where=where, take=1, order={"updated_at": "desc"})
|
||||
where: Final[Mapping[str, object]] = key_filter if vis is None else {"AND": [key_filter, vis]}
|
||||
rows = await _memory_table(prisma_client).find_many(where=where, take=1, order={"updated_at": "desc"})
|
||||
if not rows:
|
||||
raise HTTPException(status_code=404, detail=f"Memory with key '{key}' not found")
|
||||
return rows[0]
|
||||
|
|
@ -374,7 +422,7 @@ async def _find_memory_for_caller(prisma_client: Any, key: str, user_api_key_dic
|
|||
async def get_memory(
|
||||
key: str,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
) -> LiteLLM_MemoryRow:
|
||||
"""Get a single memory entry by key, scoped to the caller."""
|
||||
prisma_client: Final = _require_prisma()
|
||||
row: Final = await _find_memory_for_caller(prisma_client, key, user_api_key_dict)
|
||||
|
|
@ -391,7 +439,7 @@ async def upsert_memory(
|
|||
key: str,
|
||||
body: MemoryUpdateRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
) -> LiteLLM_MemoryRow:
|
||||
"""
|
||||
Upsert a memory entry by key within the caller's scope.
|
||||
|
||||
|
|
@ -414,20 +462,24 @@ async def upsert_memory(
|
|||
# `model_fields_set`), the column is preserved as-is.
|
||||
fields_sent: Final = body.model_fields_set
|
||||
metadata_in_payload: Final = "metadata" in fields_sent
|
||||
metadata: Final[object] = body.metadata
|
||||
|
||||
data: Final[dict] = {}
|
||||
if body.value is not None:
|
||||
data["value"] = body.value
|
||||
if metadata_in_payload:
|
||||
data["metadata"] = _serialize_metadata_for_prisma(body.metadata)
|
||||
if not data:
|
||||
value_field: Final[Mapping[str, object]] = {"value": body.value} if body.value is not None else {}
|
||||
metadata_field: Final[Mapping[str, object]] = (
|
||||
{"metadata": _serialize_metadata_for_prisma(metadata)} if metadata_in_payload else {}
|
||||
)
|
||||
if not value_field and not metadata_field:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Request body must include at least one of: value, metadata.",
|
||||
)
|
||||
data["updated_by"] = user_api_key_dict.user_id
|
||||
data: Final[Mapping[str, object]] = {
|
||||
**value_field,
|
||||
**metadata_field,
|
||||
"updated_by": user_api_key_dict.user_id,
|
||||
}
|
||||
|
||||
async def _find_existing() -> Any:
|
||||
async def _find_existing() -> _MemoryRecord | None:
|
||||
"""Return the caller-visible row for `key`, or None."""
|
||||
try:
|
||||
return await _find_memory_for_caller(prisma_client, key, user_api_key_dict)
|
||||
|
|
@ -444,7 +496,7 @@ async def upsert_memory(
|
|||
# their team) — otherwise a teammate could overwrite a personal
|
||||
# entry through the OR-based visibility filter.
|
||||
await _assert_write_access(prisma_client, existing, user_api_key_dict)
|
||||
row = await MemoryRepository(prisma_client).table.update(
|
||||
row = await _memory_table(prisma_client).update(
|
||||
where={"memory_id": existing.memory_id},
|
||||
data=data,
|
||||
)
|
||||
|
|
@ -459,18 +511,17 @@ async def upsert_memory(
|
|||
# Omit `metadata` when None so the column defaults to SQL NULL;
|
||||
# otherwise JSON-encode for Prisma — same pattern as
|
||||
# `create_memory` above.
|
||||
create_data: Final[dict] = {
|
||||
create_data: Final[Mapping[str, object]] = {
|
||||
"key": key,
|
||||
"value": body.value,
|
||||
"user_id": user_id,
|
||||
"team_id": team_id,
|
||||
"created_by": user_api_key_dict.user_id,
|
||||
"updated_by": user_api_key_dict.user_id,
|
||||
**({"metadata": _serialize_metadata_for_prisma(metadata)} if metadata is not None else {}),
|
||||
}
|
||||
if body.metadata is not None:
|
||||
create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata)
|
||||
try:
|
||||
row = await MemoryRepository(prisma_client).table.create(data=create_data)
|
||||
row = await _memory_table(prisma_client).create(data=create_data)
|
||||
except Exception as e:
|
||||
# Race: a concurrent PUT/POST created the row after our check.
|
||||
# Re-read and fall back to an update so the PUT stays idempotent
|
||||
|
|
@ -487,7 +538,7 @@ async def upsert_memory(
|
|||
)
|
||||
# Same write-authorization check as the non-race path.
|
||||
await _assert_write_access(prisma_client, existing_after_race, user_api_key_dict)
|
||||
row = await MemoryRepository(prisma_client).table.update(
|
||||
row = await _memory_table(prisma_client).update(
|
||||
where={"memory_id": existing_after_race.memory_id},
|
||||
data=data,
|
||||
)
|
||||
|
|
@ -508,14 +559,14 @@ async def upsert_memory(
|
|||
async def delete_memory(
|
||||
key: str,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
) -> MemoryDeleteResponse:
|
||||
"""Delete a memory entry by key, scoped to the caller."""
|
||||
prisma_client: Final = _require_prisma()
|
||||
row: Final = await _find_memory_for_caller(prisma_client, key, user_api_key_dict)
|
||||
# Visibility != write authority — see the upsert handler for the rationale.
|
||||
await _assert_write_access(prisma_client, row, user_api_key_dict)
|
||||
try:
|
||||
await MemoryRepository(prisma_client).table.delete(where={"memory_id": row.memory_id})
|
||||
await _memory_table(prisma_client).delete(where={"memory_id": row.memory_id})
|
||||
except Exception as e:
|
||||
raise _internal_error("Error deleting memory: %s", e, "Internal error deleting memory entry.")
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,8 @@ https://platform.openai.com/docs/api-reference/responses-streaming
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any, Final, cast
|
||||
from collections.abc import AsyncIterable, Callable, Mapping
|
||||
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, TypedDict
|
||||
|
||||
from fastapi import Request, Response
|
||||
|
||||
|
|
@ -18,28 +19,77 @@ from litellm._logging import verbose_proxy_logger
|
|||
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.router import Router
|
||||
from litellm.types.llms.openai import ResponsesAPIStatus
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
|
||||
_ContentEntry: TypeAlias = dict[str, object] | list[object] | str | int | float | bool | None
|
||||
|
||||
|
||||
class _StreamingBodyResponse(Protocol):
|
||||
body_iterator: AsyncIterable[str | bytes | memoryview]
|
||||
|
||||
|
||||
class _OutputItem(TypedDict, total=False):
|
||||
id: str
|
||||
content: list[_ContentEntry]
|
||||
|
||||
|
||||
class _TerminalResponse(TypedDict, total=False):
|
||||
status: ResponsesAPIStatus
|
||||
error: dict[str, object]
|
||||
usage: dict[str, object]
|
||||
reasoning: dict[str, object]
|
||||
tool_choice: str | dict[str, object]
|
||||
tools: list[dict[str, object]]
|
||||
model: str
|
||||
instructions: str
|
||||
temperature: float
|
||||
top_p: float
|
||||
max_output_tokens: int
|
||||
previous_response_id: str
|
||||
text: dict[str, object]
|
||||
truncation: str
|
||||
parallel_tool_calls: bool
|
||||
user: str
|
||||
store: bool
|
||||
incomplete_details: dict[str, object]
|
||||
output: list[_OutputItem]
|
||||
|
||||
|
||||
class _StreamEvent(TypedDict, total=False):
|
||||
type: str
|
||||
item: _OutputItem
|
||||
item_id: str
|
||||
part: dict[str, object]
|
||||
content_index: int
|
||||
delta: str
|
||||
response: _TerminalResponse
|
||||
|
||||
|
||||
async def background_streaming_task(
|
||||
polling_id: str,
|
||||
data: dict,
|
||||
data: dict[str, object],
|
||||
polling_handler: ResponsePollingHandler,
|
||||
request: Request,
|
||||
fastapi_response: Response,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
general_settings: dict,
|
||||
llm_router,
|
||||
proxy_config,
|
||||
proxy_logging_obj,
|
||||
select_data_generator,
|
||||
user_model,
|
||||
user_temperature,
|
||||
user_request_timeout,
|
||||
user_max_tokens,
|
||||
user_api_base,
|
||||
version,
|
||||
):
|
||||
general_settings: dict[str, object],
|
||||
llm_router: Router | None,
|
||||
proxy_config: "ProxyConfig",
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
select_data_generator: Callable[..., object] | None,
|
||||
user_model: str | None,
|
||||
user_temperature: float | None,
|
||||
user_request_timeout: float | None,
|
||||
user_max_tokens: int | None,
|
||||
user_api_base: str | None,
|
||||
version: str | None,
|
||||
) -> None:
|
||||
"""
|
||||
Background task to stream response and update cache
|
||||
|
||||
|
|
@ -69,7 +119,7 @@ async def background_streaming_task(
|
|||
# Make streaming request.
|
||||
# Pre-call checks (rate limits, guardrails, budget) were already run
|
||||
# before polling ID creation, so skip them here to avoid double-counting.
|
||||
response: Final = await processor.base_process_llm_request(
|
||||
response: Final[_StreamingBodyResponse] = await processor.base_process_llm_request(
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -91,8 +141,9 @@ async def background_streaming_task(
|
|||
|
||||
# Process streaming response following OpenAI events format
|
||||
# https://platform.openai.com/docs/api-reference/responses-streaming
|
||||
output_items: Final[dict[str, dict[str, Any]]] = {} # Track output items by ID
|
||||
accumulated_text: Final = {} # Track accumulated text deltas by (item_id, content_index)
|
||||
output_items: Final[dict[str, _OutputItem]] = {} # Track output items by ID
|
||||
# Track accumulated text deltas by (item_id, content_index)
|
||||
accumulated_text: Final[dict[tuple[str, int], str]] = {}
|
||||
|
||||
# ResponsesAPIResponse fields to extract from response.completed
|
||||
usage_data = None
|
||||
|
|
@ -121,7 +172,7 @@ async def background_streaming_task(
|
|||
None # Will be set by response.completed/failed/incomplete/cancelled
|
||||
)
|
||||
terminal_error = None
|
||||
_event_to_status: Final = {
|
||||
_event_to_status: Final[Mapping[str, ResponsesAPIStatus]] = {
|
||||
"response.completed": "completed",
|
||||
"response.failed": "failed",
|
||||
"response.incomplete": "incomplete",
|
||||
|
|
@ -162,7 +213,7 @@ async def background_streaming_task(
|
|||
break
|
||||
|
||||
try:
|
||||
event = json.loads(chunk_data)
|
||||
event: _StreamEvent = json.loads(chunk_data)
|
||||
event_type = event.get("type", "")
|
||||
|
||||
# Process different event types based on OpenAI streaming spec
|
||||
|
|
@ -181,9 +232,10 @@ async def background_streaming_task(
|
|||
|
||||
if item_id and item_id in output_items:
|
||||
# Update the output item with new content
|
||||
if "content" not in output_items[item_id]:
|
||||
output_items[item_id]["content"] = []
|
||||
output_items[item_id]["content"].append(content_part)
|
||||
added_to_item = output_items[item_id]
|
||||
if "content" not in added_to_item:
|
||||
added_to_item["content"] = []
|
||||
added_to_item["content"].append(content_part)
|
||||
state_dirty = True
|
||||
|
||||
elif event_type == "response.output_text.delta":
|
||||
|
|
@ -201,12 +253,14 @@ async def background_streaming_task(
|
|||
accumulated_text[key] += delta
|
||||
|
||||
# Update the content in output_items
|
||||
if "content" in output_items[item_id]:
|
||||
content_list = output_items[item_id]["content"]
|
||||
delta_item = output_items[item_id]
|
||||
if "content" in delta_item:
|
||||
content_list = delta_item["content"]
|
||||
if content_index < len(content_list):
|
||||
# Update existing content part with accumulated text
|
||||
if isinstance(content_list[content_index], dict):
|
||||
content_list[content_index]["text"] = accumulated_text[key]
|
||||
existing_part = content_list[content_index]
|
||||
if isinstance(existing_part, dict):
|
||||
existing_part["text"] = accumulated_text[key]
|
||||
state_dirty = True
|
||||
|
||||
elif event_type == "response.content_part.done":
|
||||
|
|
@ -217,8 +271,9 @@ async def background_streaming_task(
|
|||
|
||||
if item_id and item_id in output_items:
|
||||
# Update with final content from event
|
||||
if "content" in output_items[item_id]:
|
||||
content_list = output_items[item_id]["content"]
|
||||
done_item = output_items[item_id]
|
||||
if "content" in done_item:
|
||||
content_list = done_item["content"]
|
||||
if content_index < len(content_list):
|
||||
content_list[content_index] = content_part
|
||||
state_dirty = True
|
||||
|
|
@ -248,12 +303,9 @@ async def background_streaming_task(
|
|||
# Terminal event - extract all ResponsesAPIResponse fields
|
||||
# https://platform.openai.com/docs/api-reference/responses-streaming
|
||||
response_data = event.get("response", {})
|
||||
terminal_status = cast(
|
||||
ResponsesAPIStatus,
|
||||
response_data.get(
|
||||
"status",
|
||||
_event_to_status.get(event_type, "completed"),
|
||||
),
|
||||
terminal_status = response_data.get(
|
||||
"status",
|
||||
_event_to_status.get(event_type, "completed"),
|
||||
)
|
||||
|
||||
# Extract error for failed and incomplete responses
|
||||
|
|
|
|||
|
|
@ -17,7 +17,12 @@ from __future__ import annotations
|
|||
|
||||
import hashlib
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
from typing_extensions import NotRequired, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -39,6 +44,44 @@ if TYPE_CHECKING:
|
|||
from litellm.types.rag import RAGIngestOptions
|
||||
|
||||
|
||||
_NO_FILENAME_METADATA: Final[Mapping[str, str]] = MappingProxyType({})
|
||||
|
||||
|
||||
class S3VectorsVectorData(TypedDict):
|
||||
float32: Sequence[float]
|
||||
|
||||
|
||||
class S3VectorsRecord(TypedDict):
|
||||
key: str
|
||||
data: S3VectorsVectorData
|
||||
metadata: Mapping[str, str]
|
||||
|
||||
|
||||
class S3VectorsMetadataConfiguration(TypedDict):
|
||||
nonFilterableMetadataKeys: Sequence[str]
|
||||
|
||||
|
||||
class S3VectorsCreateIndexRequest(TypedDict):
|
||||
vectorBucketName: str
|
||||
indexName: str | None
|
||||
dataType: str
|
||||
dimension: int | None
|
||||
distanceMetric: str
|
||||
metadataConfiguration: NotRequired[S3VectorsMetadataConfiguration]
|
||||
|
||||
|
||||
class S3VectorsQueryResultMetadata(TypedDict, total=False):
|
||||
source_text: str
|
||||
|
||||
|
||||
class S3VectorsQueryResult(TypedDict, total=False):
|
||||
metadata: S3VectorsQueryResultMetadata
|
||||
|
||||
|
||||
class S3VectorsQueryResponse(TypedDict, total=False):
|
||||
vectors: Sequence[S3VectorsQueryResult]
|
||||
|
||||
|
||||
class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
|
||||
"""
|
||||
S3 Vectors RAG ingestion using httpx + AWS SigV4 signing.
|
||||
|
|
@ -57,6 +100,12 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
|
|||
- non_filterable_metadata_keys: List of metadata keys to exclude from filtering
|
||||
"""
|
||||
|
||||
vector_bucket_name: str
|
||||
index_name: str | None
|
||||
distance_metric: str
|
||||
non_filterable_metadata_keys: Sequence[str]
|
||||
dimension: int | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ingest_options: RAGIngestOptions,
|
||||
|
|
@ -78,7 +127,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
|
|||
self.dimension = self._get_dimension_from_config()
|
||||
|
||||
# Get AWS region using BaseAWSLLM method
|
||||
_aws_region: Final = self.vector_store_config.get("aws_region_name")
|
||||
_aws_region: Final[str | None] = self.vector_store_config.get("aws_region_name")
|
||||
self.aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(
|
||||
aws_region_name=str(_aws_region) if _aws_region else None
|
||||
)
|
||||
|
|
@ -166,7 +215,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
|
|||
url: str,
|
||||
data: str | None = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
) -> Any:
|
||||
) -> httpx.Response:
|
||||
"""
|
||||
Helper to sign and execute AWS API requests using httpx + SigV4.
|
||||
|
||||
|
|
@ -223,18 +272,24 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
|
|||
signed_headers: Final = dict(aws_request.headers.items())
|
||||
|
||||
# Make the request using specific method (pattern from s3_v2.py)
|
||||
method_upper: Final = method.upper()
|
||||
if method_upper == "PUT":
|
||||
response = await self.async_httpx_client.put(url, data=data, headers=signed_headers)
|
||||
elif method_upper == "POST":
|
||||
response = await self.async_httpx_client.post(url, data=data, headers=signed_headers)
|
||||
elif method_upper == "GET":
|
||||
response = await self.async_httpx_client.get(url, headers=signed_headers)
|
||||
else:
|
||||
raise ValueError(f"Unsupported HTTP method: {method}")
|
||||
response: Final = await self._execute_signed_request(method, url, data, signed_headers)
|
||||
if response is None:
|
||||
raise ValueError(f"No response from S3 Vectors for {method} {url}")
|
||||
|
||||
return response
|
||||
|
||||
async def _execute_signed_request(
|
||||
self, method: str, url: str, data: str | None, headers: dict[str, str]
|
||||
) -> httpx.Response | None:
|
||||
method_upper: Final = method.upper()
|
||||
if method_upper == "PUT":
|
||||
return await self.async_httpx_client.put(url, data=data, headers=headers)
|
||||
if method_upper == "POST":
|
||||
return await self.async_httpx_client.post(url, data=data, headers=headers)
|
||||
if method_upper == "GET":
|
||||
return await self.async_httpx_client.get(url, headers=headers)
|
||||
raise ValueError(f"Unsupported HTTP method: {method}")
|
||||
|
||||
async def _ensure_vector_bucket_exists(self):
|
||||
"""Create vector bucket if it doesn't exist using GetVectorBucket and CreateVectorBucket APIs."""
|
||||
verbose_logger.debug("Ensuring S3 vector bucket exists: %s", self.vector_bucket_name)
|
||||
|
|
@ -311,16 +366,24 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
|
|||
)
|
||||
|
||||
# Prepare index configuration per AWS API docs
|
||||
index_config: Final = {
|
||||
"vectorBucketName": self.vector_bucket_name,
|
||||
"indexName": self.index_name,
|
||||
"dataType": "float32",
|
||||
"dimension": self.dimension,
|
||||
"distanceMetric": self.distance_metric,
|
||||
}
|
||||
base_index_config: Final = S3VectorsCreateIndexRequest(
|
||||
vectorBucketName=self.vector_bucket_name,
|
||||
indexName=self.index_name,
|
||||
dataType="float32",
|
||||
dimension=self.dimension,
|
||||
distanceMetric=self.distance_metric,
|
||||
)
|
||||
|
||||
if self.non_filterable_metadata_keys:
|
||||
index_config["metadataConfiguration"] = {"nonFilterableMetadataKeys": self.non_filterable_metadata_keys}
|
||||
index_config: Final[S3VectorsCreateIndexRequest] = (
|
||||
{
|
||||
**base_index_config,
|
||||
"metadataConfiguration": S3VectorsMetadataConfiguration(
|
||||
nonFilterableMetadataKeys=self.non_filterable_metadata_keys
|
||||
),
|
||||
}
|
||||
if self.non_filterable_metadata_keys
|
||||
else base_index_config
|
||||
)
|
||||
|
||||
create_url: Final = f"https://s3vectors.{self.aws_region_name}.api.aws/CreateIndex"
|
||||
response = await self._sign_and_execute_request("POST", create_url, data=safe_dumps(index_config))
|
||||
|
|
@ -336,7 +399,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
|
|||
verbose_logger.exception("Error creating vector index: %s", e)
|
||||
raise
|
||||
|
||||
async def _put_vectors(self, vectors: list[dict[str, Any]]):
|
||||
async def _put_vectors(self, vectors: Sequence[S3VectorsRecord]):
|
||||
"""
|
||||
Call PutVectors API to store vectors in S3 Vectors.
|
||||
|
||||
|
|
@ -442,24 +505,19 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
|
|||
raise ValueError(error_msg)
|
||||
|
||||
# Prepare vectors for PutVectors API
|
||||
vectors: Final = []
|
||||
for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
|
||||
# Build metadata dict
|
||||
metadata: dict[str, str] = {
|
||||
"source_text": chunk, # Non-filterable (for reference)
|
||||
"chunk_index": str(i), # Filterable
|
||||
}
|
||||
|
||||
if filename:
|
||||
metadata["filename"] = filename # Filterable
|
||||
|
||||
vector_obj = {
|
||||
"key": f"{filename}_{i}" if filename else f"chunk_{i}",
|
||||
"data": {"float32": embedding},
|
||||
"metadata": metadata,
|
||||
}
|
||||
|
||||
vectors.append(vector_obj)
|
||||
filename_metadata: Final[Mapping[str, str]] = {"filename": filename} if filename else _NO_FILENAME_METADATA
|
||||
vectors: Final = tuple(
|
||||
S3VectorsRecord(
|
||||
key=f"{filename}_{i}" if filename else f"chunk_{i}",
|
||||
data=S3VectorsVectorData(float32=embedding),
|
||||
metadata={
|
||||
"source_text": chunk, # Non-filterable (for reference)
|
||||
"chunk_index": str(i), # Filterable
|
||||
**filename_metadata, # Filterable
|
||||
},
|
||||
)
|
||||
for i, (chunk, embedding) in enumerate(zip(chunks, embeddings))
|
||||
)
|
||||
|
||||
# Call PutVectors API
|
||||
await self._put_vectors(vectors)
|
||||
|
|
@ -468,7 +526,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
|
|||
vector_store_id: Final = f"{self.vector_bucket_name}:{self.index_name}"
|
||||
return vector_store_id, filename
|
||||
|
||||
async def query_vector_store(self, vector_store_id: str, query: str, top_k: int = 5) -> dict[str, Any] | None:
|
||||
async def query_vector_store(
|
||||
self, vector_store_id: str, query: str, top_k: int = 5
|
||||
) -> S3VectorsQueryResponse | None:
|
||||
"""
|
||||
Query S3 Vectors using QueryVectors API.
|
||||
|
||||
|
|
@ -507,12 +567,13 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
|
|||
response = await self._sign_and_execute_request("POST", url, data=safe_dumps(request_body))
|
||||
|
||||
if response.status_code == 200:
|
||||
results: Final = response.json()
|
||||
results: Final[S3VectorsQueryResponse] = response.json()
|
||||
verbose_logger.debug("Query returned %s results", len(results.get("vectors", [])))
|
||||
|
||||
# Check if query terms appear in results
|
||||
if results.get("vectors"):
|
||||
for result in results["vectors"]:
|
||||
matched_vectors: Final = results.get("vectors")
|
||||
if matched_vectors:
|
||||
for result in matched_vectors:
|
||||
metadata = result.get("metadata", {})
|
||||
source_text = metadata.get("source_text", "")
|
||||
if query.lower() in source_text.lower():
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
{
|
||||
"ANN001": {
|
||||
"limit": 3114
|
||||
"limit": 3101
|
||||
},
|
||||
"ANN002": {
|
||||
"limit": 71
|
||||
|
|
@ -9,10 +9,10 @@
|
|||
"limit": 834
|
||||
},
|
||||
"ANN201": {
|
||||
"limit": 2031
|
||||
"limit": 2025
|
||||
},
|
||||
"ANN202": {
|
||||
"limit": 865
|
||||
"limit": 863
|
||||
},
|
||||
"ANN204": {
|
||||
"limit": 713
|
||||
|
|
@ -24,7 +24,7 @@
|
|||
"limit": 133
|
||||
},
|
||||
"ANN401": {
|
||||
"limit": 1555
|
||||
"limit": 1543
|
||||
},
|
||||
"ASYNC230": {
|
||||
"limit": 11
|
||||
|
|
@ -201,7 +201,7 @@
|
|||
"limit": 58
|
||||
},
|
||||
"SIM102": {
|
||||
"limit": 322
|
||||
"limit": 321
|
||||
},
|
||||
"SIM103": {
|
||||
"limit": 119
|
||||
|
|
@ -234,7 +234,7 @@
|
|||
"limit": 5
|
||||
},
|
||||
"TID251": {
|
||||
"limit": 1238
|
||||
"limit": 1234
|
||||
},
|
||||
"TRY002": {
|
||||
"limit": 528
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
{
|
||||
"LIT001": {
|
||||
"limit": 23149
|
||||
"limit": 23144
|
||||
},
|
||||
"LIT002": {
|
||||
"limit": 27166
|
||||
|
|
@ -15,7 +15,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"LIT006": {
|
||||
"limit": 1086
|
||||
"limit": 1085
|
||||
},
|
||||
"LIT007": {
|
||||
"limit": 0
|
||||
|
|
@ -27,9 +27,9 @@
|
|||
"limit": 0
|
||||
},
|
||||
"LIT010": {
|
||||
"limit": 16758
|
||||
"limit": 16735
|
||||
},
|
||||
"LIT011": {
|
||||
"limit": 5598
|
||||
"limit": 5597
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue