mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
refactor(types): replace Any with real types across 20 backend files
Fifth round of basedpyright Any reduction. Every change is typing-only and leaves runtime behavior identical. Prisma table access now goes through the PrismaTableRepository and TableActions protocols the repo already has, instead of reading untyped attributes off prisma_client.db. Payload and parameter annotations move from dict[str, Any] to dict[str, object] or Mapping[str, object]. Guardrail constructors that took **kwargs: Any now take Unpack of a PEP 728 TypedDict, the same _CustomGuardrailOptions shape three other guardrails already use. Calls into the OpenAI and Azure assistants SDKs pass explicit keywords rather than splatting a dict, so the arguments are checked against the real SDK signatures.
This commit is contained in:
mateo-berri
parent
9dbfb060bd
commit
fa966ca2d0
20 changed files with 199 additions and 100 deletions
|
|
@ -1,5 +1,5 @@
|
|||
from collections.abc import Coroutine, Iterable
|
||||
from typing import Any, Final, Literal, TypedDict
|
||||
from typing import Final, Literal, TypedDict
|
||||
|
||||
import httpx
|
||||
from openai import AsyncAzureOpenAI, AzureOpenAI
|
||||
|
|
@ -715,7 +715,8 @@ class AzureAssistantsAPI(BaseAzureLLM):
|
|||
event_handler: AssistantEventHandler | None,
|
||||
litellm_params: dict | None = None,
|
||||
) -> AsyncAssistantStreamManager[AsyncAssistantEventHandler]:
|
||||
data: Final[dict[str, Any]] = {
|
||||
stream_fn: Final = client.beta.threads.runs.stream
|
||||
base_data: Final[_RunThreadStreamData] = {
|
||||
"thread_id": thread_id,
|
||||
"assistant_id": assistant_id,
|
||||
"additional_instructions": additional_instructions,
|
||||
|
|
@ -725,8 +726,8 @@ class AzureAssistantsAPI(BaseAzureLLM):
|
|||
"tools": tools,
|
||||
}
|
||||
if event_handler is not None:
|
||||
data["event_handler"] = event_handler
|
||||
return client.beta.threads.runs.stream(**data)
|
||||
return stream_fn(**base_data, event_handler=event_handler)
|
||||
return stream_fn(**base_data)
|
||||
|
||||
def run_thread_stream(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -124,7 +124,7 @@ class OllamaChatConfig(BaseConfig):
|
|||
setattr(self.__class__, key, value)
|
||||
|
||||
@classmethod
|
||||
def get_config(cls):
|
||||
def get_config(cls) -> dict[str, object]:
|
||||
return super().get_config()
|
||||
|
||||
def get_supported_openai_params(self, model: str):
|
||||
|
|
|
|||
|
|
@ -227,7 +227,7 @@ class OllamaConfig(BaseConfig):
|
|||
model: str,
|
||||
api_base: str | None = None,
|
||||
api_key: str | None = None,
|
||||
) -> Any:
|
||||
) -> dict[str, object] | None:
|
||||
"""
|
||||
curl http://localhost:11434/api/show -d '{
|
||||
"name": "mistral"
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import time
|
||||
import types
|
||||
from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional, cast
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -2754,7 +2754,12 @@ class OpenAIAssistantsAPI(BaseLLM):
|
|||
|
||||
message_thread: Final = await openai_client.beta.threads.create(**data)
|
||||
|
||||
return Thread(**message_thread.dict())
|
||||
return Thread(
|
||||
id=message_thread.id,
|
||||
created_at=message_thread.created_at,
|
||||
metadata=message_thread.metadata,
|
||||
object=message_thread.object,
|
||||
)
|
||||
|
||||
# fmt: off
|
||||
|
||||
|
|
@ -2840,7 +2845,12 @@ class OpenAIAssistantsAPI(BaseLLM):
|
|||
|
||||
message_thread: Final = openai_client.beta.threads.create(**data)
|
||||
|
||||
return Thread(**message_thread.dict())
|
||||
return Thread(
|
||||
id=message_thread.id,
|
||||
created_at=message_thread.created_at,
|
||||
metadata=message_thread.metadata,
|
||||
object=message_thread.object,
|
||||
)
|
||||
|
||||
async def async_get_thread(
|
||||
self,
|
||||
|
|
@ -2863,7 +2873,12 @@ class OpenAIAssistantsAPI(BaseLLM):
|
|||
|
||||
response: Final = await openai_client.beta.threads.retrieve(thread_id=thread_id)
|
||||
|
||||
return Thread(**response.dict())
|
||||
return Thread(
|
||||
id=response.id,
|
||||
created_at=response.created_at,
|
||||
metadata=response.metadata,
|
||||
object=response.object,
|
||||
)
|
||||
|
||||
# fmt: off
|
||||
|
||||
|
|
@ -2929,7 +2944,12 @@ class OpenAIAssistantsAPI(BaseLLM):
|
|||
|
||||
response: Final = openai_client.beta.threads.retrieve(thread_id=thread_id)
|
||||
|
||||
return Thread(**response.dict())
|
||||
return Thread(
|
||||
id=response.id,
|
||||
created_at=response.created_at,
|
||||
metadata=response.metadata,
|
||||
object=response.object,
|
||||
)
|
||||
|
||||
def delete_thread(self):
|
||||
pass
|
||||
|
|
@ -2986,18 +3006,27 @@ class OpenAIAssistantsAPI(BaseLLM):
|
|||
tools: Iterable[AssistantToolParam] | None,
|
||||
event_handler: AssistantEventHandler | None,
|
||||
) -> AsyncAssistantStreamManager[AsyncAssistantEventHandler]:
|
||||
data: Final[dict[str, Any]] = {
|
||||
"thread_id": thread_id,
|
||||
"assistant_id": assistant_id,
|
||||
"additional_instructions": additional_instructions,
|
||||
"instructions": instructions,
|
||||
"metadata": metadata,
|
||||
"model": model,
|
||||
"tools": tools,
|
||||
}
|
||||
runs_stream: Final = client.beta.threads.runs.stream
|
||||
if event_handler is not None:
|
||||
data["event_handler"] = event_handler
|
||||
return client.beta.threads.runs.stream(**data)
|
||||
return runs_stream(
|
||||
thread_id=thread_id,
|
||||
assistant_id=assistant_id,
|
||||
additional_instructions=additional_instructions,
|
||||
instructions=instructions,
|
||||
metadata=metadata,
|
||||
model=model,
|
||||
tools=tools,
|
||||
event_handler=event_handler,
|
||||
)
|
||||
return runs_stream(
|
||||
thread_id=thread_id,
|
||||
assistant_id=assistant_id,
|
||||
additional_instructions=additional_instructions,
|
||||
instructions=instructions,
|
||||
metadata=metadata,
|
||||
model=model,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
def run_thread_stream(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -280,7 +280,7 @@ class VertexFineTuningAPI(VertexLLM):
|
|||
vertex_location: str,
|
||||
vertex_credentials: str,
|
||||
request_route: str,
|
||||
):
|
||||
) -> object:
|
||||
_auth_header, vertex_project = await self._ensure_access_token_async(
|
||||
credentials=vertex_credentials,
|
||||
project_id=vertex_project,
|
||||
|
|
@ -341,5 +341,4 @@ class VertexFineTuningAPI(VertexLLM):
|
|||
f"Error creating fine tuning job. Status code: {response.status_code}. Response: {response.text}"
|
||||
)
|
||||
|
||||
response_json: Final = response.json()
|
||||
return response_json
|
||||
return response.json()
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import base64
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from io import BufferedRandom, BufferedReader, BytesIO
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
|
@ -47,11 +48,11 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
image_edit_optional_params: ImageEditOptionalRequestParams,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
supported_params: Final = self.get_supported_openai_params(model)
|
||||
filtered_params = {key: value for key, value in image_edit_optional_params.items() if key in supported_params}
|
||||
|
||||
mapped_params: Final[dict[str, Any]] = {}
|
||||
mapped_params: Final[dict[str, object]] = {}
|
||||
|
||||
# Map OpenAI parameters to Imagen format
|
||||
if "n" in filtered_params:
|
||||
|
|
@ -148,10 +149,10 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
model: str,
|
||||
prompt: str | None,
|
||||
image: FileTypes | None,
|
||||
image_edit_optional_request_params: dict[str, Any],
|
||||
image_edit_optional_request_params: Mapping[str, object],
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
headers: dict,
|
||||
) -> tuple[dict[str, Any], RequestFiles | None]:
|
||||
) -> tuple[dict[str, object], RequestFiles | None]:
|
||||
# Prepare reference images in the correct Imagen format
|
||||
if image is None:
|
||||
raise ValueError("Vertex AI Imagen image edit requires at least one reference image.")
|
||||
|
|
@ -182,14 +183,14 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
parameters["guidanceScale"] = 7.5 # Default guidance scale
|
||||
parameters["seed"] = None # Let Vertex AI choose random seed
|
||||
|
||||
request_body: Final[dict[str, Any]] = {
|
||||
request_body: Final[dict[str, object]] = {
|
||||
"instances": instances,
|
||||
"parameters": parameters,
|
||||
}
|
||||
|
||||
payload: Final[Any] = json.dumps(request_body)
|
||||
payload: Final = json.dumps(request_body)
|
||||
empty_files: Final = cast(RequestFiles, [])
|
||||
return cast(tuple[dict[str, Any], RequestFiles | None], (payload, empty_files))
|
||||
return cast(tuple[dict[str, object], RequestFiles | None], (payload, empty_files))
|
||||
|
||||
def transform_image_edit_response(
|
||||
self,
|
||||
|
|
@ -237,8 +238,8 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
def _prepare_reference_images(
|
||||
self,
|
||||
image: FileTypes | list[FileTypes],
|
||||
image_edit_optional_request_params: dict[str, Any],
|
||||
) -> list[dict[str, Any]]:
|
||||
image_edit_optional_request_params: Mapping[str, object],
|
||||
) -> list[dict[str, object]]:
|
||||
"""
|
||||
Prepare reference images in the correct Imagen API format
|
||||
"""
|
||||
|
|
@ -248,7 +249,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
else:
|
||||
images = [image]
|
||||
|
||||
reference_images: Final[list[dict[str, Any]]] = []
|
||||
reference_images: Final[list[dict[str, object]]] = []
|
||||
|
||||
for idx, img in enumerate(images):
|
||||
if img is None:
|
||||
|
|
@ -258,7 +259,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
base64_data = base64.b64encode(image_bytes).decode("utf-8")
|
||||
|
||||
# Create reference image structure
|
||||
reference_image = {
|
||||
reference_image: dict[str, object] = {
|
||||
"referenceType": "REFERENCE_TYPE_RAW",
|
||||
"referenceId": idx + 1,
|
||||
"referenceImage": {"bytesBase64Encoded": base64_data},
|
||||
|
|
@ -272,7 +273,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
mask_bytes: Final = self._read_all_bytes(mask_image)
|
||||
mask_base64: Final = base64.b64encode(mask_bytes).decode("utf-8")
|
||||
|
||||
mask_reference: Final = {
|
||||
mask_reference: Final[dict[str, object]] = {
|
||||
"referenceType": "REFERENCE_TYPE_MASK",
|
||||
"referenceId": len(reference_images) + 1,
|
||||
"referenceImage": {"bytesBase64Encoded": mask_base64},
|
||||
|
|
|
|||
|
|
@ -218,10 +218,10 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
|
|||
contents: Final = [{"role": "user", "parts": [{"text": prompt}]}]
|
||||
|
||||
# Prepare generation config
|
||||
generation_config: Final[dict[str, Any]] = {"responseModalities": ["IMAGE"]}
|
||||
generation_config: Final[dict[str, object]] = {"responseModalities": ["IMAGE"]}
|
||||
|
||||
# Seed from user-supplied imageConfig dict; flat params are overlaid for backward compat.
|
||||
image_config: Final[dict[str, Any]] = dict(optional_params.get("imageConfig") or {})
|
||||
image_config: Final[dict[str, object]] = dict(optional_params.get("imageConfig") or {})
|
||||
|
||||
if "aspectRatio" in optional_params:
|
||||
image_config["aspectRatio"] = optional_params["aspectRatio"]
|
||||
|
|
@ -242,7 +242,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
|
|||
elif "n" in optional_params:
|
||||
generation_config["candidateCount"] = optional_params["n"]
|
||||
|
||||
request_body: Final[dict[str, Any]] = {
|
||||
request_body: Final[dict[str, object]] = {
|
||||
"contents": contents,
|
||||
"generationConfig": generation_config,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ from litellm.repositories.table_repositories import (
|
|||
MCPServerOAuthClientRepository,
|
||||
MCPServerRepository,
|
||||
MCPUserCredentialsRepository,
|
||||
PrismaTableRepository,
|
||||
)
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.verification_token_repository import (
|
||||
|
|
@ -522,11 +523,14 @@ def _user_credential_actions(
|
|||
return table
|
||||
|
||||
|
||||
class _MCPUserEnvVarsRepository(PrismaTableRepository["prisma_db_models.LiteLLM_MCPUserEnvVars"]):
|
||||
table_name = "litellm_mcpuserenvvars"
|
||||
|
||||
|
||||
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
|
||||
return table
|
||||
return _MCPUserEnvVarsRepository(prisma_client).table
|
||||
|
||||
|
||||
async def _db_find_user_credential_row(
|
||||
|
|
|
|||
|
|
@ -48,7 +48,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 SpendLinkedTable
|
||||
from litellm.repositories.prisma_protocols import PrismaBatch, SpendLinkedTable
|
||||
from litellm.repositories.table_repositories import (
|
||||
EndUserRepository,
|
||||
ModelAccessGroupBudgetRepository,
|
||||
|
|
@ -435,6 +435,11 @@ class ResetBudgetJob:
|
|||
self.reset_settings: BudgetResetSettings = reset_settings or get_budget_reset_settings()
|
||||
self.pod_lock_manager: PodLockManager | None = pod_lock_manager
|
||||
|
||||
@property
|
||||
def _new_batch(self) -> Callable[[], PrismaBatch]:
|
||||
new_batch: Final[Callable[[], PrismaBatch]] = self.prisma_client.db.batch_
|
||||
return new_batch
|
||||
|
||||
async def _lease_is_held(self, lock_manager: PodLockManager) -> bool:
|
||||
"""True only when the lease is readable and someone holds it.
|
||||
|
||||
|
|
@ -721,7 +726,7 @@ class ResetBudgetJob:
|
|||
)
|
||||
|
||||
async def _commit_budget_cascade_once(self, cascade: _BudgetCascade) -> None:
|
||||
async with budget_cascade_unit_of_work(self.prisma_client.db.batch_) as uow:
|
||||
async with budget_cascade_unit_of_work(self._new_batch) as uow:
|
||||
_queue_budget_linked_resets(uow.team_memberships, cascade)
|
||||
_queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE)
|
||||
_queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE)
|
||||
|
|
@ -861,7 +866,7 @@ class ResetBudgetJob:
|
|||
)
|
||||
|
||||
async def _write_key_reset_updates_once(self, updated_keys: list[LiteLLM_VerificationToken]) -> None:
|
||||
async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow:
|
||||
async with spend_reset_unit_of_work(self._new_batch) as uow:
|
||||
for k in updated_keys:
|
||||
if k.token is None:
|
||||
continue
|
||||
|
|
@ -885,7 +890,7 @@ class ResetBudgetJob:
|
|||
)
|
||||
|
||||
async def _write_user_reset_updates_once(self, updated_users: list[LiteLLM_UserTable]) -> None:
|
||||
async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow:
|
||||
async with spend_reset_unit_of_work(self._new_batch) as uow:
|
||||
for u in updated_users:
|
||||
uow.users.queue_spend_reset(
|
||||
user_id=u.user_id,
|
||||
|
|
@ -907,7 +912,7 @@ class ResetBudgetJob:
|
|||
)
|
||||
|
||||
async def _write_team_reset_updates_once(self, updated_teams: list[LiteLLM_TeamTable]) -> None:
|
||||
async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow:
|
||||
async with spend_reset_unit_of_work(self._new_batch) as uow:
|
||||
for t in updated_teams:
|
||||
uow.teams.queue_spend_reset(
|
||||
team_id=t.team_id,
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal
|
|||
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
from typing_extensions import TypedDict, Unpack
|
||||
|
||||
from litellm import DualCache
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -111,6 +112,10 @@ class CiscoAIDefenseGuardrailAPIError(Exception):
|
|||
"""Raised when there is an error talking to the Cisco AI Defense API."""
|
||||
|
||||
|
||||
class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
|
||||
"""Base-class constructor options this guardrail forwards untouched to CustomGuardrail."""
|
||||
|
||||
|
||||
class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
|
||||
"""
|
||||
Cisco AI Defense guardrail integration.
|
||||
|
|
@ -144,7 +149,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
|
|||
on_flagged_action: str | None = None,
|
||||
fallback_on_error: str | None = None,
|
||||
timeout: float | None = None,
|
||||
**kwargs: Any,
|
||||
**kwargs: Unpack[_CustomGuardrailOptions],
|
||||
) -> None:
|
||||
resolved_api_key: Final = api_key or os.environ.get("CISCO_AI_DEFENSE_API_KEY")
|
||||
if not resolved_api_key:
|
||||
|
|
|
|||
|
|
@ -7,10 +7,10 @@
|
|||
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol
|
||||
|
||||
import httpx
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._version import version as litellm_version
|
||||
|
|
@ -56,7 +56,13 @@ class DeepKeepFirewallResponse(TypedDict):
|
|||
class _DeepKeepInitKwargsView(TypedDict):
|
||||
"""Typed read of the guardrail name carried in the untyped base-guardrail kwargs."""
|
||||
|
||||
guardrail_name: ReadOnly[str]
|
||||
guardrail_name: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
|
||||
"""Base-class constructor options this guardrail forwards untouched to CustomGuardrail."""
|
||||
|
||||
guardrail_name: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _DeepKeepMetadataSource(TypedDict, total=False):
|
||||
|
|
@ -110,7 +116,7 @@ class DeepKeepGuardrail(CustomGuardrail):
|
|||
firewall_id: str | None = None,
|
||||
unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed",
|
||||
extra_headers: Mapping[str, str] | list[str] | None = None,
|
||||
**kwargs: Any,
|
||||
**kwargs: Unpack[_CustomGuardrailOptions],
|
||||
):
|
||||
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
|
||||
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import time
|
|||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, TypeGuard
|
||||
from typing import TYPE_CHECKING, ClassVar, Final, Literal, TypeGuard
|
||||
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
|
|
@ -50,6 +50,7 @@ from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.anthropic_messages.transformation import BaseAnthropicMessagesConfig
|
||||
from litellm.types.guardrails import LitellmParams
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
|
||||
|
|
@ -878,9 +879,9 @@ class HeadroomGuardrail(CustomGuardrail):
|
|||
|
||||
async def async_pre_call_deployment_hook(
|
||||
self,
|
||||
kwargs: dict[str, Any],
|
||||
kwargs: dict[str, object],
|
||||
call_type: CallTypes | None,
|
||||
) -> dict[str, Any] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict
|
||||
) -> dict[str, object] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict
|
||||
base_result: Final = await super().async_pre_call_deployment_hook(kwargs, call_type)
|
||||
effective: Final = base_result if base_result is not None else kwargs
|
||||
if call_type not in _STREAM_CONVERTIBLE_CALL_TYPES:
|
||||
|
|
@ -897,7 +898,7 @@ class HeadroomGuardrail(CustomGuardrail):
|
|||
|
||||
async def async_should_run_agentic_loop(
|
||||
self,
|
||||
response: Any,
|
||||
response: object,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
tools: list[dict] | None,
|
||||
|
|
@ -919,8 +920,8 @@ class HeadroomGuardrail(CustomGuardrail):
|
|||
tools: dict,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
response: Any,
|
||||
anthropic_messages_provider_config: Any,
|
||||
response: object,
|
||||
anthropic_messages_provider_config: BaseAnthropicMessagesConfig | None,
|
||||
anthropic_messages_optional_request_params: dict,
|
||||
logging_obj: LiteLLMLoggingObj | None,
|
||||
stream: bool,
|
||||
|
|
|
|||
|
|
@ -360,7 +360,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
else:
|
||||
return {"modelResponseData": {"byteItem": {"byteDataType": file_type, "byteData": base64_data}}}
|
||||
|
||||
def _should_block_content(self, armor_response: Mapping[str, Any], allow_sanitization: bool = False) -> bool:
|
||||
def _should_block_content(self, armor_response: Mapping[str, object], allow_sanitization: bool = False) -> bool:
|
||||
"""Check if Model Armor response indicates content should be blocked, including both inspectResult and deidentifyResult."""
|
||||
for filt in self._filter_result_items(armor_response):
|
||||
# Check RAI, PI/Jailbreak, Malicious URI, CSAM, Virus scan as before
|
||||
|
|
@ -429,7 +429,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
return filter_results
|
||||
return []
|
||||
|
||||
def _has_deidentify_match(self, armor_response: Mapping[str, Any]) -> bool:
|
||||
def _has_deidentify_match(self, armor_response: Mapping[str, object]) -> bool:
|
||||
"""Whether an SDP de-identify filter matched, i.e. Model Armor owes this response a redaction."""
|
||||
for filter_entry in self._filter_result_items(armor_response):
|
||||
sdp = filter_entry.get("sdpFilterResult")
|
||||
|
|
@ -439,7 +439,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
|
||||
def _resolve_streaming_outcome(
|
||||
self,
|
||||
armor_response: Mapping[str, Any],
|
||||
armor_response: Mapping[str, object],
|
||||
assembled_response: object,
|
||||
content: str,
|
||||
) -> tuple[bool, str | None]:
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@ POST /auto_router/validate_complexity_router_config - Dry-run the complexity-rou
|
|||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from itertools import chain, groupby
|
||||
from operator import attrgetter
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Final, Protocol
|
||||
from uuid import uuid4
|
||||
|
|
@ -1094,6 +1093,10 @@ def _slices(rows: Sequence[_AttemptAggRow]) -> tuple[ShadowEvalSlice, ...]:
|
|||
)
|
||||
|
||||
|
||||
def _leg_group_id(leg: "_LegRow") -> str:
|
||||
return leg.group_id
|
||||
|
||||
|
||||
class _LegRow(BaseModel):
|
||||
"""One LiteLLM_ShadowEvalJob row, validated off the untyped prisma record. A row is
|
||||
one target's leg of a job; the legs of a job share group_id and identical config,
|
||||
|
|
@ -1598,10 +1601,7 @@ async def list_shadow_eval_jobs(
|
|||
or ()
|
||||
)
|
||||
by_group: Final[Mapping[str, tuple[_LegRow, ...]]] = MappingProxyType(
|
||||
{
|
||||
group_id: tuple(group)
|
||||
for group_id, group in groupby(sorted(legs, key=attrgetter("group_id")), key=attrgetter("group_id"))
|
||||
}
|
||||
{group_id: tuple(group) for group_id, group in groupby(sorted(legs, key=_leg_group_id), key=_leg_group_id)}
|
||||
)
|
||||
newest_first: Final = sorted(
|
||||
by_group, key=lambda group_id: max(leg.created_at for leg in by_group[group_id]), reverse=True
|
||||
|
|
|
|||
|
|
@ -31,11 +31,31 @@ from litellm.constants import (
|
|||
)
|
||||
from litellm.litellm_core_utils.ptu_pricing import ptu_terms
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
|
||||
from litellm.repositories.model_repository import ModelRepository
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import PrismaTableRepository
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
||||
class _DailyTeamSpendRepository(PrismaTableRepository["prisma_models.LiteLLM_DailyTeamSpend"]):
|
||||
table_name = "litellm_dailyteamspend"
|
||||
|
||||
|
||||
def _daily_team_spend_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_DailyTeamSpend]":
|
||||
"""The sentinel rows this rollup writes, reads back and prunes."""
|
||||
return _DailyTeamSpendRepository(prisma_client).table
|
||||
|
||||
|
||||
def _proxy_model_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_ProxyModelTable]":
|
||||
"""The stored deployments the rollup scans for PTU config."""
|
||||
return ModelRepository(prisma_client).table
|
||||
|
||||
|
||||
_HOURS_PER_DAY: Final = 24
|
||||
_PRUNE_ID_CHUNK_SIZE: Final = 5_000
|
||||
_UPSERT_ATTEMPTS: Final = 3
|
||||
|
|
@ -97,7 +117,7 @@ def _decode_model_info(raw: object) -> "Mapping[str, object] | None":
|
|||
"""
|
||||
if isinstance(raw, str):
|
||||
try:
|
||||
decoded: Final = json.loads(raw)
|
||||
decoded: Final[object] = json.loads(raw)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
return decoded if isinstance(decoded, dict) else None
|
||||
|
|
@ -240,7 +260,7 @@ async def _upsert_ptu_daily_row(
|
|||
}
|
||||
}
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
await prisma_client.db.litellm_dailyteamspend.upsert(
|
||||
await _daily_team_spend_table(prisma_client).upsert(
|
||||
where=where,
|
||||
data={ # mutable-ok: prisma upsert data payload
|
||||
"create": { # mutable-ok: prisma create payload
|
||||
|
|
@ -353,7 +373,7 @@ async def _load_ptu_models(prisma_client: "PrismaClient", *, router: object | No
|
|||
The router is handed in rather than read off the proxy module, so a run prices exactly
|
||||
the deployments its caller declares and nothing a co-resident process left behind.
|
||||
"""
|
||||
rows: Final = await prisma_client.db.litellm_proxymodeltable.find_many()
|
||||
rows: Final = await _proxy_model_table(prisma_client).find_many()
|
||||
db_ids: Final = frozenset(model_id for row in rows if (model_id := str(getattr(row, "model_id", "") or "")))
|
||||
config_records: Final = _config_deployments(router, owned_by_db=db_ids)
|
||||
models: Final = tuple(
|
||||
|
|
@ -503,7 +523,7 @@ async def _existing_sentinel_keys(
|
|||
survives a rename. Nothing here reads the display name.
|
||||
"""
|
||||
date_range: Final = {"gte": start.isoformat(), "lte": end.isoformat()} # mutable-ok: prisma range filter
|
||||
rows: Final = await prisma_client.db.litellm_dailyteamspend.find_many(
|
||||
rows: Final = await _daily_team_spend_table(prisma_client).find_many(
|
||||
where={"api_key": PTU_SENTINEL_API_KEY, "date": date_range} # mutable-ok: prisma find filter
|
||||
)
|
||||
return frozenset(
|
||||
|
|
@ -771,7 +791,7 @@ async def _prune_unrefreshed_sentinel_rows(
|
|||
)
|
||||
filters: Final = tuple(_prune_filter(date_str=date_str, cutoff=cutoff, chunk=chunk) for chunk in chunks)
|
||||
deletions: Final = tuple(
|
||||
[await prisma_client.db.litellm_dailyteamspend.delete_many(where=where) for where in filters]
|
||||
[await _daily_team_spend_table(prisma_client).delete_many(where=where) for where in filters]
|
||||
)
|
||||
deleted: Final = sum(deletions)
|
||||
if deleted:
|
||||
|
|
|
|||
|
|
@ -3401,7 +3401,7 @@ class _ConfigRow:
|
|||
|
||||
__slots__ = ("param_name", "param_value")
|
||||
|
||||
def __init__(self, param_name: str, param_value: Any) -> None:
|
||||
def __init__(self, param_name: str, param_value: object) -> None:
|
||||
self.param_name = param_name
|
||||
self.param_value = param_value
|
||||
|
||||
|
|
@ -3414,7 +3414,7 @@ def _pack_config_row(row: Any) -> dict[str, object]:
|
|||
return {"param_name": row.param_name, "param_value": row.param_value}
|
||||
|
||||
|
||||
def _unpack_config_row(cached: Any) -> _ConfigRow | None:
|
||||
def _unpack_config_row(cached: object) -> _ConfigRow | None:
|
||||
if cached is None or cached == _CONFIG_CACHE_MISS:
|
||||
return None
|
||||
if isinstance(cached, dict):
|
||||
|
|
@ -3557,6 +3557,7 @@ class PrismaClient:
|
|||
verbose_proxy_logger.debug("Creating Prisma Client..")
|
||||
try:
|
||||
from prisma import Prisma
|
||||
from prisma.types import DatasourceOverride
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("Failed to import Prisma client: %s", e)
|
||||
verbose_proxy_logger.error("This usually means 'prisma generate' hasn't been run yet.")
|
||||
|
|
@ -3607,11 +3608,11 @@ class PrismaClient:
|
|||
reader_token: Final = mint_database_token(token_auth, reader_iam_endpoint)
|
||||
read_replica_url = reader_iam_endpoint.build_url(reader_token)
|
||||
os.environ["DATABASE_URL_READ_REPLICA"] = read_replica_url
|
||||
reader_kwargs: Final[dict[str, Any]] = {"datasource": {"url": read_replica_url}}
|
||||
reader_datasource: Final = DatasourceOverride(url=read_replica_url)
|
||||
if http_client is not None:
|
||||
reader_prisma = Prisma(http=http_client, **reader_kwargs)
|
||||
reader_prisma = Prisma(http=http_client, datasource=reader_datasource)
|
||||
else:
|
||||
reader_prisma = Prisma(**reader_kwargs)
|
||||
reader_prisma = Prisma(datasource=reader_datasource)
|
||||
reader_wrapper: Final = PrismaWrapper(
|
||||
original_prisma=reader_prisma,
|
||||
token_auth=token_auth,
|
||||
|
|
|
|||
|
|
@ -53,7 +53,7 @@ bedrock_realtime: Final = BedrockRealtime()
|
|||
xai_realtime: Final = XAIRealtime()
|
||||
vertex_llm_base: Final = VertexBase()
|
||||
base_llm_http_handler = BaseLLMHTTPHandler()
|
||||
_EMPTY_MODEL_PARAMS: Final[Mapping[str, Any]] = MappingProxyType({})
|
||||
_EMPTY_MODEL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _with_resolved_session_model(session: dict[str, object], model_name: str) -> dict[str, object]:
|
||||
|
|
|
|||
|
|
@ -4,29 +4,23 @@ Model repository for database operations on LiteLLM_ProxyModelTable.
|
|||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from litellm.models.model import LiteLLM_ProxyModelTable
|
||||
from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
)
|
||||
from litellm.repositories.base_repository import BaseRepository
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import PrismaTableRepository
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
|
||||
class _PrismaModelDb(Protocol):
|
||||
@property
|
||||
def litellm_proxymodeltable(self) -> TableActions["prisma_models.LiteLLM_ProxyModelTable"]: ...
|
||||
|
||||
|
||||
class _PrismaClientView(Protocol):
|
||||
@property
|
||||
def db(self) -> _PrismaModelDb: ...
|
||||
class _ProxyModelTableRepository(PrismaTableRepository["prisma_models.LiteLLM_ProxyModelTable"]):
|
||||
table_name = "litellm_proxymodeltable"
|
||||
|
||||
|
||||
class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
|
||||
|
|
@ -38,11 +32,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
|
|||
|
||||
@property
|
||||
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",
|
||||
)
|
||||
return _ProxyModelTableRepository(self._prisma_client).table
|
||||
|
||||
@property
|
||||
def model_class(self) -> type[LiteLLM_ProxyModelTable]:
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ Team repository for database operations on LiteLLM_TeamTable.
|
|||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from types import TracebackType
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
|
@ -40,6 +41,36 @@ def _team_arrays(team: LiteLLM_TeamTable) -> _TeamArrays:
|
|||
return team
|
||||
|
||||
|
||||
class _TeamTables(Protocol):
|
||||
"""The two team tables this repository reads and writes."""
|
||||
|
||||
@property
|
||||
def litellm_teamtable(self) -> TableActions["prisma_models.LiteLLM_TeamTable"]: ...
|
||||
|
||||
@property
|
||||
def litellm_deletedteamtable(self) -> TableActions["prisma_models.LiteLLM_DeletedTeamTable"]: ...
|
||||
|
||||
|
||||
class _TeamTransactionManager(Protocol):
|
||||
async def __aenter__(self) -> _TeamTables: ...
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> bool | None: ...
|
||||
|
||||
|
||||
class _PrismaTeamDb(_TeamTables, Protocol):
|
||||
def tx(self) -> _TeamTransactionManager: ...
|
||||
|
||||
|
||||
class _PrismaClientView(Protocol):
|
||||
@property
|
||||
def db(self) -> _PrismaTeamDb: ...
|
||||
|
||||
|
||||
_MEMBERS_WITH_ROLES_ADAPTER: Final = TypeAdapter(list[Member])
|
||||
_JSON_ENCODED_TEAM_FIELDS: Final = (
|
||||
"metadata",
|
||||
|
|
@ -54,13 +85,18 @@ _JSON_ENCODED_TEAM_FIELDS: Final = (
|
|||
class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
|
||||
"""Repository for team database operations."""
|
||||
|
||||
@property
|
||||
def _db(self) -> _PrismaTeamDb:
|
||||
client: Final[_PrismaClientView] = self.prisma_client
|
||||
return client.db
|
||||
|
||||
@property
|
||||
def table(self) -> TableActions["prisma_models.LiteLLM_TeamTable"]:
|
||||
return self.prisma_client.db.litellm_teamtable
|
||||
return self._db.litellm_teamtable
|
||||
|
||||
@property
|
||||
def deleted_table(self) -> TableActions["prisma_models.LiteLLM_DeletedTeamTable"]:
|
||||
return self.prisma_client.db.litellm_deletedteamtable
|
||||
return self._db.litellm_deletedteamtable
|
||||
|
||||
@property
|
||||
def model_class(self) -> type[LiteLLM_TeamTable]:
|
||||
|
|
@ -256,7 +292,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
|
|||
archive_data["litellm_changed_by"] = litellm_changed_by
|
||||
archive_data["deleted_at"] = datetime.utcnow()
|
||||
|
||||
async with self.prisma_client.db.tx() as tx:
|
||||
async with self._db.tx() as tx:
|
||||
await tx.litellm_deletedteamtable.create(data=archive_data)
|
||||
await tx.litellm_teamtable.delete(where={"team_id": team_id})
|
||||
|
||||
|
|
|
|||
|
|
@ -266,7 +266,7 @@ def get_pre_routing_selection(kwargs: Mapping[str, object]) -> str | None:
|
|||
DISABLE_FALLBACKS_METADATA_KEY: Final = "_disable_fallbacks"
|
||||
|
||||
|
||||
def record_disable_fallbacks(request_kwargs: Mapping[str, Any] | None, disabled: bool) -> None:
|
||||
def record_disable_fallbacks(request_kwargs: Mapping[str, object] | None, disabled: bool) -> None:
|
||||
"""
|
||||
Write-or-clear the request's disable_fallbacks verdict into the router-internal metadata
|
||||
bucket. The wrapper pops the raw kwarg before any downstream frame runs, so the refusal
|
||||
|
|
@ -286,7 +286,7 @@ def record_disable_fallbacks(request_kwargs: Mapping[str, Any] | None, disabled:
|
|||
bucket.pop(DISABLE_FALLBACKS_METADATA_KEY, None)
|
||||
|
||||
|
||||
def fallbacks_disabled_for_request(kwargs: Mapping[str, Any]) -> bool:
|
||||
def fallbacks_disabled_for_request(kwargs: Mapping[str, object]) -> bool:
|
||||
"""True when this request opted out of fallbacks, read from the raw kwarg (pre-pop
|
||||
snapshots keep it) or the router-internal bucket the wrapper stamps after popping it."""
|
||||
if kwargs.get("disable_fallbacks") is True:
|
||||
|
|
@ -639,7 +639,7 @@ async def log_failure_fallback_event(original_model_group: str, kwargs: dict, or
|
|||
verbose_router_logger.error("Error in log_failure_fallback_event: %s", e)
|
||||
|
||||
|
||||
def _check_non_standard_fallback_format(fallbacks: list[Any] | None) -> bool:
|
||||
def _check_non_standard_fallback_format(fallbacks: Sequence[object] | None) -> bool:
|
||||
"""
|
||||
Checks if the fallbacks list is a list of strings or a list of dictionaries.
|
||||
|
||||
|
|
@ -653,8 +653,9 @@ def _check_non_standard_fallback_format(fallbacks: list[Any] | None) -> bool:
|
|||
return False
|
||||
if all(isinstance(item, str) for item in fallbacks):
|
||||
return True
|
||||
elif all(isinstance(item, dict) for item in fallbacks):
|
||||
for item in fallbacks:
|
||||
dict_entries: Final = tuple(item for item in fallbacks if isinstance(item, dict))
|
||||
if len(dict_entries) == len(fallbacks):
|
||||
for item in dict_entries:
|
||||
for key in LiteLLMParamsTypedDict.__annotations__:
|
||||
if key in item:
|
||||
# If the value is a list, it's likely a standard fallback model group mapping
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue