Merge remote-tracking branch 'origin/litellm_internal_staging' into shadow-eval-pre-adoption

.github/workflows/check-schema-sync.yml

@@ -10,6 +10,10 @@ on:
 permissions:
   contents: read
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   check-sync:
     name: Verify schema.prisma copies match root

.github/workflows/conventional-commits.yml

@@ -14,6 +14,10 @@ on:
 permissions:
   pull-requests: read
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   lint-pr-title:
     name: Validate PR title

.github/workflows/guard-fork-dependencies.yml

@@ -15,6 +15,10 @@ on:
 
 permissions: {}
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   guard:
     name: Block fork dependency changes

.github/workflows/helm_unit_test.yml

@@ -9,6 +9,10 @@ on:
 permissions:
   contents: read
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   unit-test:
     runs-on: ubuntu-latest

.github/workflows/test-linting.yml

@@ -11,6 +11,10 @@ on:
 permissions:
   contents: read
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   lint:
     runs-on: ubuntu-latest

.github/workflows/test-litellm-ui-build.yml

@@ -10,6 +10,10 @@ on:
       - litellm_oss_staging
       - "litellm_**"
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   build-ui:
     runs-on: ubuntu-latest

.github/workflows/test-litellm-ui-lint.yml

@@ -10,6 +10,10 @@ on:
       - litellm_oss_staging
       - "litellm_**"
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   frontend-lint:
     runs-on: ubuntu-latest

.github/workflows/test-mcp.yml

@@ -11,6 +11,10 @@ on:
 permissions:
   contents: read
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   test:
     runs-on: ubuntu-latest

.github/workflows/test-model-map.yaml

@@ -11,6 +11,10 @@ on:
 permissions:
   contents: read
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   validate-model-prices-json:
     runs-on: ubuntu-latest

.github/workflows/test-unit-proxy-endpoints.yml

@@ -38,6 +38,8 @@ jobs:
         tests/test_litellm/proxy/google_endpoints
         tests/test_litellm/proxy/openai_files_endpoint
         tests/test_litellm/proxy/batches_endpoints
+        tests/test_litellm/proxy/fine_tuning_endpoints
+        tests/test_litellm/proxy/vector_store_files_endpoints
         tests/test_litellm/proxy/video_endpoints
         tests/test_litellm/proxy/response_api_endpoints
         tests/test_litellm/proxy/image_endpoints

basedpyright-code-budget.json

@@ -1,9 +1,9 @@
 {
   "reportAny": {
-    "limit": 28842
+    "limit": 27731
   },
   "reportArgumentType": {
-    "limit": 2634
+    "limit": 2626
   },
   "reportAssignmentType": {
     "limit": 329
@@ -12,7 +12,7 @@
     "limit": 514
   },
   "reportCallIssue": {
-    "limit": 117
+    "limit": 116
   },
   "reportConstantRedefinition": {
     "limit": 40
@@ -24,7 +24,7 @@
     "limit": 19
   },
   "reportExplicitAny": {
-    "limit": 9103
+    "limit": 8807
   },
   "reportFunctionMemberAccess": {
     "limit": 7
@@ -54,10 +54,10 @@
     "limit": 0
   },
   "reportMissingParameterType": {
-    "limit": 5843
+    "limit": 5835
   },
   "reportMissingTypeArgument": {
-    "limit": 15816
+    "limit": 15790
   },
   "reportMissingTypeStubs": {
     "limit": 40
@@ -72,7 +72,7 @@
     "limit": 0
   },
   "reportOptionalMemberAccess": {
-    "limit": 1078
+    "limit": 1077
   },
   "reportOptionalOperand": {
     "limit": 0
@@ -90,28 +90,28 @@
     "limit": 8
   },
   "reportReturnType": {
-    "limit": 218
+    "limit": 217
   },
   "reportTypedDictNotRequiredAccess": {
-    "limit": 27
+    "limit": 26
   },
   "reportUndefinedVariable": {
     "limit": 0
   },
   "reportUnknownArgumentType": {
-    "limit": 45098
+    "limit": 45063
   },
   "reportUnknownLambdaType": {
     "limit": 113
   },
   "reportUnknownMemberType": {
-    "limit": 39826
+    "limit": 39773
   },
   "reportUnknownParameterType": {
-    "limit": 20237
+    "limit": 20207
   },
   "reportUnknownVariableType": {
-    "limit": 31371
+    "limit": 31281
   },
   "reportUnnecessaryCast": {
     "limit": 122
@@ -123,7 +123,7 @@
     "limit": 5
   },
   "reportUnnecessaryIsInstance": {
-    "limit": 864
+    "limit": 862
   },
   "reportUntypedBaseClass": {
     "limit": 0

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -3,9 +3,21 @@
 
 import base64
 import json
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Final,
+    List,
+    Literal,
+    Optional,
+    Protocol,
+    TypedDict,
+    Union,
+    cast,
+)
 from uuid import NAMESPACE_URL, uuid5
 
 from fastapi import HTTPException
@@ -98,33 +110,76 @@ def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> Optional[LiteLLMB
     try:
         batch_obj: Final = LiteLLMBatch.model_validate(_decode_json_blob(row.file_object))
     except Exception as e:
-        verbose_logger.warning(
-            f"Failed to parse batch object {row.unified_object_id}: {_sanitized_parse_error(e)}"
-        )
+        verbose_logger.warning(f"Failed to parse batch object {row.unified_object_id}: {_sanitized_parse_error(e)}")
         return None
     batch_obj.id = row.unified_object_id
     return batch_obj
 
 
-def _parse_managed_file_object(
-    raw_file_object: object, unified_file_id: str
-) -> Optional[OpenAIFileObject]:
+def _parse_managed_file_object(raw_file_object: object, unified_file_id: str) -> Optional[OpenAIFileObject]:
     if raw_file_object is None:
         return None
     try:
         return OpenAIFileObject.model_validate(raw_file_object)
     except Exception as e:
-        verbose_logger.warning(
-            f"Failed to parse managed file object {unified_file_id}: {_sanitized_parse_error(e)}"
-        )
+        verbose_logger.warning(f"Failed to parse managed file object {unified_file_id}: {_sanitized_parse_error(e)}")
         return None
 
 
+class _ManagedFileRow(Protocol):
+    unified_file_id: str
+    file_object: OpenAIFileObject
+    storage_backend: Optional[str]
+    storage_url: Optional[str]
+    created_by: Optional[str]
+    team_id: Optional[str]
+
+    def model_dump(self) -> Mapping[str, object]: ...
+
+
+class _ManagedFileTableActions(Protocol):
+    async def find_first(self, where: Mapping[str, object]) -> Optional[_ManagedFileRow]: ...
+
+    async def find_many(self, where: Mapping[str, object]) -> Sequence[_ManagedFileRow]: ...
+
+    async def upsert(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]]) -> _ManagedFileRow: ...
+
+    async def delete(self, where: Mapping[str, str]) -> Optional[_ManagedFileRow]: ...
+
+
+class _ManagedObjectTableActions(Protocol):
+    async def find_first(self, where: Mapping[str, object]) -> "Optional[PrismaManagedObjectRow]": ...
+
+    async def find_many(
+        self,
+        where: Mapping[str, object],
+        take: int,
+        order: Union[Mapping[str, str], Sequence[Mapping[str, str]]],
+        cursor: Mapping[str, str] = ...,
+        skip: int = ...,
+    ) -> "Sequence[PrismaManagedObjectRow]": ...
+
+    async def upsert(
+        self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]]
+    ) -> "PrismaManagedObjectRow": ...
+
+
+class _CursorPageArgs(TypedDict, total=False):
+    cursor: Mapping[str, str]
+    skip: int
+
+
+def _managed_file_table(prisma_client: PrismaClient) -> _ManagedFileTableActions:
+    return prisma_client.db.litellm_managedfiletable
+
+
+def _managed_object_table(prisma_client: PrismaClient) -> _ManagedObjectTableActions:
+    return prisma_client.db.litellm_managedobjecttable
+
+
 class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
     # Class variables or attributes
-    def __init__(
-        self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient
-    ):
+    def __init__(self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient):
         self.internal_usage_cache = internal_usage_cache
         self.prisma_client = prisma_client
 
@@ -143,9 +198,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         model_mappings: Dict[str, str],
         user_api_key_dict: UserAPIKeyAuth,
     ) -> None:
-        verbose_logger.info(
-            f"Storing LiteLLM Managed File object with id={file_id} in cache"
-        )
+        verbose_logger.info(f"Storing LiteLLM Managed File object with id={file_id} in cache")
         if file_object is not None:
             litellm_managed_file_object = LiteLLM_ManagedFileTable(
                 unified_file_id=file_id,
@@ -196,13 +249,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 f"storage_url={db_data.get('storage_url')}"
             )
 
-        result = await self.prisma_client.db.litellm_managedfiletable.upsert(
+        result = await _managed_file_table(self.prisma_client).upsert(
             where={"unified_file_id": file_id},
             data={"create": db_data, "update": update_data},
         )
-        verbose_logger.debug(
-            f"LiteLLM Managed File object with id={file_id} stored in db: {result}"
-        )
+        verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}")
 
     async def store_unified_object_id(
         self,
@@ -213,9 +264,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         file_purpose: Literal["batch", "fine-tune", "response"],
         user_api_key_dict: UserAPIKeyAuth,
     ) -> None:
-        verbose_logger.info(
-            f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache"
-        )
+        verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache")
         litellm_managed_object = LiteLLM_ManagedObjectTable(
             unified_object_id=unified_object_id,
             model_object_id=model_object_id,
@@ -228,7 +277,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             litellm_parent_otel_span=litellm_parent_otel_span,
         )
 
-        await self.prisma_client.db.litellm_managedobjecttable.upsert(
+        await _managed_object_table(self.prisma_client).upsert(
             where={"unified_object_id": unified_object_id},
             data={
                 "create": {
@@ -265,9 +314,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             return LiteLLM_ManagedFileTable.model_validate(result)
 
         ## CHECK DB
-        db_object = await self.prisma_client.db.litellm_managedfiletable.find_first(
-            where={"unified_file_id": file_id}
-        )
+        db_object = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id})
 
         if db_object:
             return LiteLLM_ManagedFileTable.model_validate(db_object.model_dump())
@@ -277,9 +324,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         self, file_id: str, litellm_parent_otel_span: Optional[Span] = None
     ) -> OpenAIFileObject:
         ## get old value
-        initial_value = await self.prisma_client.db.litellm_managedfiletable.find_first(
-            where={"unified_file_id": file_id}
-        )
+        initial_value = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id})
         if initial_value is None:
             raise Exception(f"LiteLLM Managed File object with id={file_id} not found")
         ## delete old value
@@ -288,15 +333,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             value=None,
             litellm_parent_otel_span=litellm_parent_otel_span,
         )
-        await self.prisma_client.db.litellm_managedfiletable.delete(
-            where={"unified_file_id": file_id}
-        )
+        await _managed_file_table(self.prisma_client).delete(where={"unified_file_id": file_id})
         return initial_value.file_object
 
-    async def can_user_call_unified_file_id(
-        self, unified_file_id: str, user_api_key_dict: UserAPIKeyAuth
-    ) -> bool:
-        managed_file = await self.prisma_client.db.litellm_managedfiletable.find_first(
+    async def can_user_call_unified_file_id(self, unified_file_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool:
+        managed_file = await _managed_file_table(self.prisma_client).find_first(
             where={"unified_file_id": unified_file_id}
         )
 
@@ -311,13 +352,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             detail=f"File not found: {unified_file_id}",
         )
 
-    async def can_user_call_unified_object_id(
-        self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth
-    ) -> bool:
-        managed_object = (
-            await self.prisma_client.db.litellm_managedobjecttable.find_first(
-                where={"unified_object_id": unified_object_id}
-            )
+    async def can_user_call_unified_object_id(self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool:
+        managed_object = await _managed_object_table(self.prisma_client).find_first(
+            where={"unified_object_id": unified_object_id}
         )
 
         if managed_object:
@@ -339,34 +376,28 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         provider: Optional[str] = None,
         target_model_names: Optional[str] = None,
         llm_router: Optional[Router] = None,
-    ) -> Dict[str, Any]:
+    ) -> Dict[str, object]:
         # Provider filtering is not supported for managed batches
         # This is because the encoded object ids stored in the managed objects table do not contain the provider information
         # To support provider filtering, we would need to store the provider information in the encoded object ids
         if provider:
-            raise Exception(
-                "Filtering by 'provider' is not supported when using managed batches."
-            )
+            raise Exception("Filtering by 'provider' is not supported when using managed batches.")
 
         # Model name filtering is not supported for managed batches
         # This is because the encoded object ids stored in the managed objects table do not contain the model name
         # A hash of the model name + litellm_params for the model name is encoded as the model id. This is not sufficient to reliably map the target model names to the model ids.
         if target_model_names:
-            raise Exception(
-                "Filtering by 'target_model_names' is not supported when using managed batches."
-            )
+            raise Exception("Filtering by 'target_model_names' is not supported when using managed batches.")
 
         owner_filter = build_owner_filter(user_api_key_dict)
         if owner_filter is None:
             return build_list_page([])
 
-        where_clause: Dict[str, Any] = {"file_purpose": "batch", **owner_filter}
+        where_clause: Dict[str, object] = {"file_purpose": "batch", **owner_filter}
 
         if after:
-            cursor_row = (
-                await self.prisma_client.db.litellm_managedobjecttable.find_first(
-                    where={**where_clause, "unified_object_id": after}
-                )
+            cursor_row = await _managed_object_table(self.prisma_client).find_first(
+                where={**where_clause, "unified_object_id": after}
             )
             if cursor_row is None:
                 raise HTTPException(
@@ -375,11 +406,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 )
 
         page_size: Final = min(limit or 20, 100)
-        cursor_args: Dict[str, Any] = (
-            {"cursor": {"unified_object_id": after}, "skip": 1} if after else {}
-        )
+        cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": after}, "skip": 1} if after else {}
 
-        batches = await self.prisma_client.db.litellm_managedobjecttable.find_many(
+        batches = await _managed_object_table(self.prisma_client).find_many(
             where=where_clause,
             take=page_size + 1,
             order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
@@ -389,9 +418,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         has_more = len(batches) > page_size
 
         parsed_rows: Final = tuple(
-            (row, batch_obj)
-            for row in batches[:page_size]
-            if (batch_obj := _parse_managed_batch_row(row)) is not None
+            (row, batch_obj) for row in batches[:page_size] if (batch_obj := _parse_managed_batch_row(row)) is not None
         )
         unified_id_by_raw_id: Final = await map_raw_file_ids_to_unified(
             raw_file_ids=frozenset(
@@ -432,14 +459,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 verbose_proxy_logger=verbose_logger,
                 user_api_key_dict=user_api_key_dict,
                 db_batch_object=row,
-                unified_batch_id=_is_base64_encoded_unified_file_id(
-                    row.unified_object_id
-                ),
+                unified_batch_id=_is_base64_encoded_unified_file_id(row.unified_object_id),
             )
         except Exception as e:
-            verbose_logger.warning(
-                f"Failed to resolve managed file ids for batch {row.unified_object_id}: {e}"
-            )
+            verbose_logger.warning(f"Failed to resolve managed file ids for batch {row.unified_object_id}: {e}")
             return None
         return batch_obj
 
@@ -458,7 +481,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         if owner_filter is None:
             return []
 
-        file_ids = await self.prisma_client.db.litellm_managedfiletable.find_many(
+        file_ids = await _managed_file_table(self.prisma_client).find_many(
             where={
                 **owner_filter,
                 "flat_model_file_ids": {"hasSome": model_object_ids},
@@ -467,27 +490,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         return [
             parsed_file_object.model_copy(update={"id": row.unified_file_id})
             for row in file_ids
-            if (
-                parsed_file_object := _parse_managed_file_object(
-                    row.file_object, row.unified_file_id
-                )
-            )
-            is not None
+            if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None
         ]
 
-    async def check_managed_file_id_access(
-        self, data: Dict, user_api_key_dict: UserAPIKeyAuth
-    ) -> bool:
+    async def check_managed_file_id_access(self, data: Dict, user_api_key_dict: UserAPIKeyAuth) -> bool:
         retrieve_file_id = cast(Optional[str], data.get("file_id"))
-        potential_file_id = (
-            _is_base64_encoded_unified_file_id(retrieve_file_id)
-            if retrieve_file_id
-            else False
-        )
+        potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False
         if potential_file_id and retrieve_file_id:
-            if await self.can_user_call_unified_file_id(
-                retrieve_file_id, user_api_key_dict
-            ):
+            if await self.can_user_call_unified_file_id(retrieve_file_id, user_api_key_dict):
                 return True
             else:
                 raise HTTPException(
@@ -496,9 +506,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 )
         return False
 
-    async def check_file_ids_access(
-        self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth
-    ) -> None:
+    async def check_file_ids_access(self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth) -> None:
         """
         Check if the user has access to a list of file IDs.
         Only checks managed (unified) file IDs.
@@ -513,9 +521,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         for file_id in file_ids:
             is_unified_file_id = _is_base64_encoded_unified_file_id(file_id)
             if is_unified_file_id:
-                if not await self.can_user_call_unified_file_id(
-                    file_id, user_api_key_dict
-                ):
+                if not await self.can_user_call_unified_file_id(file_id, user_api_key_dict):
                     raise HTTPException(
                         status_code=403,
                         detail=f"User {user_api_key_dict.user_id} does not have access to the file {file_id}",
@@ -543,10 +549,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
 
         ### HANDLE TRANSFORMATIONS ###
         # Check both completion and acompletion call types
-        is_completion_call = (
-            call_type == CallTypes.completion.value
-            or call_type == CallTypes.acompletion.value
-        )
+        is_completion_call = call_type == CallTypes.completion.value or call_type == CallTypes.acompletion.value
 
         if is_completion_call:
             messages = data.get("messages")
@@ -559,9 +562,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
 
                     # Check if any files are stored in storage backends and need base64 conversion
                     # This is needed for Vertex AI/Gemini which requires base64 content
-                    is_vertex_ai = model and (
-                        "vertex_ai" in model or "gemini" in model.lower()
-                    )
+                    is_vertex_ai = model and ("vertex_ai" in model or "gemini" in model.lower())
                     if is_vertex_ai:
                         await self._convert_storage_files_to_base64(
                             messages=messages,
@@ -573,10 +574,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                         file_ids, user_api_key_dict.parent_otel_span
                     )
                     data["model_file_id_mapping"] = model_file_id_mapping
-        elif (
-            call_type == CallTypes.aresponses.value
-            or call_type == CallTypes.responses.value
-        ):
+        elif call_type == CallTypes.aresponses.value or call_type == CallTypes.responses.value:
             # Handle managed files in responses API input and tools
             file_ids = []
 
@@ -603,23 +601,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             if tools:
                 unified_vs_ids = self.get_vector_store_ids_from_file_search_tools(tools)
                 if unified_vs_ids:
-                    await self.check_vector_store_ids_access(
-                        unified_vs_ids, user_api_key_dict
-                    )
+                    await self.check_vector_store_ids_access(unified_vs_ids, user_api_key_dict)
         elif call_type == CallTypes.afile_content.value:
             retrieve_file_id = cast(Optional[str], data.get("file_id"))
-            potential_file_id = (
-                _is_base64_encoded_unified_file_id(retrieve_file_id)
-                if retrieve_file_id
-                else False
-            )
+            potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False
             if potential_file_id and "llm_output_file_id," in potential_file_id:
                 model_id = self.get_model_id_from_unified_file_id(potential_file_id)
                 if model_id:
                     data["model"] = model_id
-                    data["file_id"] = self.get_output_file_id_from_unified_file_id(
-                        potential_file_id
-                    )
+                    data["file_id"] = self.get_output_file_id_from_unified_file_id(potential_file_id)
         elif call_type == CallTypes.acreate_batch.value:
             input_file_id = cast(Optional[str], data.get("input_file_id"))
             if input_file_id:
@@ -636,10 +626,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         ):
             accessor_key: Optional[str] = None
             retrieve_object_id: Optional[str] = None
-            if (
-                call_type == CallTypes.aretrieve_batch.value
-                or call_type == CallTypes.acancel_batch.value
-            ):
+            if call_type == CallTypes.aretrieve_batch.value or call_type == CallTypes.acancel_batch.value:
                 accessor_key = "batch_id"
             elif (
                 call_type == CallTypes.acancel_fine_tuning_job.value
@@ -651,32 +638,24 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 retrieve_object_id = cast(Optional[str], data.get(accessor_key))
 
             potential_llm_object_id = (
-                _is_base64_encoded_unified_file_id(retrieve_object_id)
-                if retrieve_object_id
-                else False
+                _is_base64_encoded_unified_file_id(retrieve_object_id) if retrieve_object_id else False
             )
             if potential_llm_object_id and retrieve_object_id:
                 ## VALIDATE USER HAS ACCESS TO THE OBJECT ##
-                if not await self.can_user_call_unified_object_id(
-                    retrieve_object_id, user_api_key_dict
-                ):
+                if not await self.can_user_call_unified_object_id(retrieve_object_id, user_api_key_dict):
                     raise HTTPException(
                         status_code=403,
                         detail=f"User {user_api_key_dict.user_id} does not have access to the object {retrieve_object_id}",
                     )
 
                 ## for managed batch id - get the model id
-                potential_model_id = get_model_id_from_unified_batch_id(
-                    potential_llm_object_id
-                )
+                potential_model_id = get_model_id_from_unified_batch_id(potential_llm_object_id)
                 if potential_model_id is None:
                     raise Exception(
                         f"LiteLLM Managed {accessor_key} with id={retrieve_object_id} is invalid - does not contain encoded model_id."
                     )
                 data["model"] = potential_model_id
-                data[accessor_key] = get_batch_id_from_unified_batch_id(
-                    potential_llm_object_id
-                )
+                data[accessor_key] = get_batch_id_from_unified_batch_id(potential_llm_object_id)
         elif call_type == CallTypes.acreate_fine_tuning_job.value:
             input_file_id = cast(Optional[str], data.get("training_file"))
             if input_file_id:
@@ -732,24 +711,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
 
         if accessor_key:
             input_file_id = cast(Optional[str], kwargs.get(accessor_key))
-            model_file_id_mapping = cast(
-                Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping")
-            )
+            model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping"))
             # model_info may be at top-level or nested under litellm_metadata
             # (batch/file operations use litellm_metadata)
             model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None))
             if model_id is None:
                 model_id = cast(
                     Optional[str],
-                    kwargs.get("litellm_metadata", {})
-                    .get("model_info", {})
-                    .get("id", None),
+                    kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None),
                 )
             mapped_file_id: Optional[str] = None
             if input_file_id and model_file_id_mapping and model_id:
-                mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get(
-                    model_id, None
-                )
+                mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get(model_id, None)
             if mapped_file_id:
                 kwargs[accessor_key] = mapped_file_id
 
@@ -775,9 +748,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                                 file_ids.append(file_id)
         return file_ids
 
-    def get_file_ids_from_responses_input(
-        self, input: Union[str, List[Dict[str, Any]]]
-    ) -> List[str]:
+    def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, Any]]]) -> List[str]:
         """
         Gets file ids from responses API input.
 
@@ -809,19 +780,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             content = item.get("content")
             if isinstance(content, list):
                 for content_item in content:
-                    if (
-                        isinstance(content_item, dict)
-                        and content_item.get("type") == "input_file"
-                    ):
+                    if isinstance(content_item, dict) and content_item.get("type") == "input_file":
                         file_id = content_item.get("file_id")
                         if file_id:
                             file_ids.append(file_id)
 
         return file_ids
 
-    def get_file_ids_from_responses_tools(
-        self, tools: List[Dict[str, Any]]
-    ) -> List[str]:
+    def get_file_ids_from_responses_tools(self, tools: List[Dict[str, object]]) -> List[str]:
         """
         Gets file ids from responses API tools parameter.
 
@@ -854,9 +820,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
 
         return file_ids
 
-    def get_vector_store_ids_from_file_search_tools(
-        self, tools: List[Dict[str, Any]]
-    ) -> List[str]:
+    def get_vector_store_ids_from_file_search_tools(self, tools: List[Dict[str, object]]) -> List[str]:
         """
         Extract unified vector_store_ids from file_search tools.
 
@@ -949,9 +913,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                     ),
                 )
 
-    async def get_model_file_id_mapping(
-        self, file_ids: List[str], litellm_parent_otel_span: Span
-    ) -> dict:
+    async def get_model_file_id_mapping(self, file_ids: List[str], litellm_parent_otel_span: Span) -> dict:
         """
         Get model-specific file IDs for a list of proxy file IDs.
         Returns a dictionary mapping litellm_proxy/ file_id -> model_id -> model_file_id
@@ -981,9 +943,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             # Get all cache keys matching the pattern file_id:*
             for file_id in litellm_managed_file_ids:
                 # Search for any cache key starting with this file_id
-                unified_file_object = await self.get_unified_file_id(
-                    file_id, litellm_parent_otel_span
-                )
+                unified_file_object = await self.get_unified_file_id(file_id, litellm_parent_otel_span)
 
                 if unified_file_object:
                     file_id_mapping[file_id] = unified_file_object.model_mappings
@@ -1001,9 +961,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             raise Exception("LLM Router not initialized. Ensure models added to proxy.")
         responses = []
         for model in target_model_names_list:
-            individual_response = await llm_router.acreate_file(
-                model=model, **_create_file_request
-            )
+            individual_response = await llm_router.acreate_file(model=model, **_create_file_request)
             responses.append(individual_response)
 
         return responses
@@ -1034,9 +992,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         model_mappings: Dict[str, str] = {}
 
         for file_object in responses:
-            model_file_id_mapping = file_object._hidden_params.get(
-                "model_file_id_mapping"
-            )
+            model_file_id_mapping = file_object._hidden_params.get("model_file_id_mapping")
             if model_file_id_mapping and isinstance(model_file_id_mapping, dict):
                 model_mappings.update(model_file_id_mapping)
 
@@ -1051,17 +1007,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         # Emit Prometheus metrics for managed file creation
         prom_logger = self._get_prometheus_logger()
         if prom_logger:
-            first_model = (
-                target_model_names_list[0] if target_model_names_list else None
-            )
+            first_model = target_model_names_list[0] if target_model_names_list else None
             first_provider = ""
             if responses:
-                first_provider = (
-                    getattr(responses[0], "_hidden_params", {}).get(
-                        "custom_llm_provider"
-                    )
-                    or ""
-                )
+                first_provider = getattr(responses[0], "_hidden_params", {}).get("custom_llm_provider") or ""
             prom_logger.record_managed_file_created(
                 model=first_model or "",
                 api_provider=first_provider,
@@ -1104,9 +1053,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         )
 
         # Convert to URL-safe base64 and strip padding
-        base64_unified_file_id = (
-            base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=")
-        )
+        base64_unified_file_id = base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=")
 
         ## CREATE RESPONSE OBJECT
 
@@ -1123,46 +1070,26 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
 
         return response
 
-    def get_unified_generic_response_id(
-        self, model_id: str, generic_response_id: str
-    ) -> str:
-        unified_generic_response_id = (
-            SpecialEnums.LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR.value.format(
-                model_id, generic_response_id
-            )
-        )
-        return (
-            base64.urlsafe_b64encode(unified_generic_response_id.encode())
-            .decode()
-            .rstrip("=")
+    def get_unified_generic_response_id(self, model_id: str, generic_response_id: str) -> str:
+        unified_generic_response_id = SpecialEnums.LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR.value.format(
+            model_id, generic_response_id
         )
+        return base64.urlsafe_b64encode(unified_generic_response_id.encode()).decode().rstrip("=")
 
     def get_unified_batch_id(self, batch_id: str, model_id: str) -> str:
-        unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(
-            model_id, batch_id
-        )
+        unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id)
         return base64.urlsafe_b64encode(unified_batch_id.encode()).decode().rstrip("=")
 
-    def get_unified_output_file_id(
-        self, output_file_id: str, model_id: str, model_name: Optional[str]
-    ) -> str:
-        deterministic_uuid: Final = uuid5(
-            uuid5(NAMESPACE_URL, model_id), output_file_id
-        )
-        unified_output_file_id = (
-            SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
-                "application/json",
-                str(deterministic_uuid),
-                model_name or "",
-                output_file_id,
-                model_id,
-            )
-        )
-        return (
-            base64.urlsafe_b64encode(unified_output_file_id.encode())
-            .decode()
-            .rstrip("=")
+    def get_unified_output_file_id(self, output_file_id: str, model_id: str, model_name: Optional[str]) -> str:
+        deterministic_uuid: Final = uuid5(uuid5(NAMESPACE_URL, model_id), output_file_id)
+        unified_output_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
+            "application/json",
+            str(deterministic_uuid),
+            model_name or "",
+            output_file_id,
+            model_id,
         )
+        return base64.urlsafe_b64encode(unified_output_file_id.encode()).decode().rstrip("=")
 
     def get_model_id_from_unified_file_id(self, file_id: str) -> str:
         return file_id.split("llm_output_file_model_id,")[1].split(";")[0]
@@ -1170,59 +1097,39 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
     def get_output_file_id_from_unified_file_id(self, file_id: str) -> str:
         marker = "llm_output_file_id,"
         if marker not in file_id:
-            raise ValueError(
-                f"Unified id does not contain {marker!r}: {file_id[:80]!r}"
-            )
+            raise ValueError(f"Unified id does not contain {marker!r}: {file_id[:80]!r}")
         return file_id.split(marker, 1)[1].split(";")[0]
 
     async def async_post_call_success_hook(
         self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
-    ) -> Any:
+    ) -> LLMResponseTypes:
         if isinstance(response, LiteLLMBatch):
             ## Check if unified_file_id is in the response
-            unified_file_id = response._hidden_params.get(
-                "unified_file_id"
-            )  # managed file id
-            unified_batch_id = response._hidden_params.get(
-                "unified_batch_id"
-            )  # managed batch id
+            unified_file_id = response._hidden_params.get("unified_file_id")  # managed file id
+            unified_batch_id = response._hidden_params.get("unified_batch_id")  # managed batch id
             model_id = cast(Optional[str], response._hidden_params.get("model_id"))
             model_name = cast(Optional[str], response._hidden_params.get("model_name"))
 
             resolved_model_name = resolve_managed_output_file_model_name(
-                unified_input_file_id=unified_file_id
-                if isinstance(unified_file_id, str)
-                else response.input_file_id,
+                unified_input_file_id=unified_file_id if isinstance(unified_file_id, str) else response.input_file_id,
                 fallback_model_name=model_name,
             )
             original_response_id = response.id
 
             if (unified_batch_id or unified_file_id) and model_id:
-                response.id = self.get_unified_batch_id(
-                    batch_id=response.id, model_id=model_id
-                )
+                response.id = self.get_unified_batch_id(batch_id=response.id, model_id=model_id)
 
                 # Handle both output_file_id and error_file_id
                 for file_attr in ["output_file_id", "error_file_id"]:
                     file_id_value = getattr(response, file_attr, None)
                     if file_id_value and model_id:
-                        decoded_output_file_id = _is_base64_encoded_unified_file_id(
-                            file_id_value
-                        )
-                        if (
-                            decoded_output_file_id
-                            and "llm_output_file_id," in decoded_output_file_id
-                        ):
-                            provider_file_id = (
-                                self.get_output_file_id_from_unified_file_id(
-                                    decoded_output_file_id
-                                )
-                            )
+                        decoded_output_file_id = _is_base64_encoded_unified_file_id(file_id_value)
+                        if decoded_output_file_id and "llm_output_file_id," in decoded_output_file_id:
+                            provider_file_id = self.get_output_file_id_from_unified_file_id(decoded_output_file_id)
                             unified_file_id = file_id_value
                         elif decoded_output_file_id:
                             verbose_logger.warning(
-                                f"Skipping {file_attr}={file_id_value!r}: "
-                                "unified id is not a managed file output id"
+                                f"Skipping {file_attr}={file_id_value!r}: unified id is not a managed file output id"
                             )
                             continue
                         else:
@@ -1241,23 +1148,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                             # Import module and use getattr for better testability with mocks
                             import litellm.proxy.proxy_server as proxy_server_module
 
-                            _llm_router = getattr(
-                                proxy_server_module, "llm_router", None
-                            )
+                            _llm_router = getattr(proxy_server_module, "llm_router", None)
                             if _llm_router is not None and model_id:
-                                _creds = (
-                                    _llm_router.get_deployment_credentials_with_provider(
-                                        model_id
-                                    )
-                                    or {}
-                                )
+                                _creds = _llm_router.get_deployment_credentials_with_provider(model_id) or {}
                                 file_object = await litellm.afile_retrieve(
                                     file_id=provider_file_id,
                                     **_creds,
                                 )
                             else:
                                 file_object = await litellm.afile_retrieve(
-                                    custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai",  # type: ignore[arg-type]
+                                    custom_llm_provider=model_name.split("/")[0]
+                                    if model_name and "/" in model_name
+                                    else "openai",  # type: ignore[arg-type]
                                     file_id=provider_file_id,
                                 )
                             verbose_logger.debug(
@@ -1311,9 +1213,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
 
         elif isinstance(response, LiteLLMFineTuningJob):
             ## Check if unified_file_id is in the response
-            unified_file_id = response._hidden_params.get(
-                "unified_file_id"
-            )  # managed file id
+            unified_file_id = response._hidden_params.get("unified_file_id")  # managed file id
             unified_finetuning_job_id = response._hidden_params.get(
                 "unified_finetuning_job_id"
             )  # managed finetuning job id
@@ -1321,9 +1221,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             model_name = cast(Optional[str], response._hidden_params.get("model_name"))
             original_response_id = response.id
             if (unified_file_id or unified_finetuning_job_id) and model_id:
-                response.id = self.get_unified_generic_response_id(
-                    model_id=model_id, generic_response_id=response.id
-                )
+                response.id = self.get_unified_generic_response_id(model_id=model_id, generic_response_id=response.id)
             await self.store_unified_object_id(
                 unified_object_id=response.id,
                 file_object=response,
@@ -1338,9 +1236,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             """
             ## check if file object
             if hasattr(response, "data") and isinstance(response.data, list):
-                if all(
-                    isinstance(file_object, FileObject) for file_object in response.data
-                ):
+                if all(isinstance(file_object, FileObject) for file_object in response.data):
                     ## Get all file id's
                     ## Check which file id's were created by the user
                     ## Filter the response to only include the files created by the user
@@ -1349,9 +1245,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                         file_object.id
                         for file_object in cast(List[FileObject], response.data)  # type: ignore
                     ]
-                    user_created_file_ids = await self.get_user_created_file_ids(
-                        user_api_key_dict, file_ids
-                    )
+                    user_created_file_ids = await self.get_user_created_file_ids(user_api_key_dict, file_ids)
                     ## Filter the response to only include the files created by the user
                     response.data = user_created_file_ids  # type: ignore
                     return response
@@ -1359,11 +1253,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         return response
 
     async def afile_retrieve(
-        self, file_id: str, litellm_parent_otel_span: Optional[Span], llm_router=None
+        self, file_id: str, litellm_parent_otel_span: Optional[Span], llm_router: Optional[Router] = None
     ) -> OpenAIFileObject:
-        stored_file_object = await self.get_unified_file_id(
-            file_id, litellm_parent_otel_span
-        )
+        stored_file_object = await self.get_unified_file_id(file_id, litellm_parent_otel_span)
 
         # Case 1 : This is not a managed file
         if not stored_file_object:
@@ -1386,21 +1278,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             )
 
         try:
-            model_id, model_file_id = next(
-                iter(stored_file_object.model_mappings.items())
-            )
-            credentials = (
-                llm_router.get_deployment_credentials_with_provider(model_id) or {}
-            )
-            response = await litellm.afile_retrieve(
-                file_id=model_file_id, **credentials
-            )
+            model_id, model_file_id = next(iter(stored_file_object.model_mappings.items()))
+            credentials = llm_router.get_deployment_credentials_with_provider(model_id) or {}
+            response = await litellm.afile_retrieve(file_id=model_file_id, **credentials)
             response.id = file_id  # Replace with unified ID
             return response
         except Exception as e:
-            raise Exception(
-                f"Failed to retrieve file {file_id} from provider: {str(e)}"
-            ) from e
+            raise Exception(f"Failed to retrieve file {file_id} from provider: {str(e)}") from e
 
     async def afile_list(
         self,
@@ -1437,12 +1321,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
 
             return False
         except Exception as e:
-            verbose_logger.warning(
-                f"Error checking batch polling configuration: {e}. Assuming disabled."
-            )
+            verbose_logger.warning(f"Error checking batch polling configuration: {e}. Assuming disabled.")
             return False
 
-    async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, Any]]:
+    async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, object]]:
         """
         Find batches that reference this file and still need cost tracking.
         Find batches that are in non-terminal state and have not yet been processed by CheckBatchCost.
@@ -1458,9 +1340,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
 
         # Get model-specific file IDs for this unified file ID if it's a managed file
         try:
-            model_file_id_mapping = await self.get_model_file_id_mapping(
-                [file_id], litellm_parent_otel_span=None
-            )
+            model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span=None)
 
             if model_file_id_mapping and file_id in model_file_id_mapping:
                 # Add all provider file IDs for this unified file
@@ -1468,8 +1348,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 file_ids_to_check.extend(provider_file_ids)
         except Exception as e:
             verbose_logger.debug(
-                f"Could not get model file ID mapping for {file_id}: {e}. "
-                f"Will only check unified file ID."
+                f"Could not get model file ID mapping for {file_id}: {e}. Will only check unified file ID."
             )
         MAX_MATCHES_TO_RETURN = 10
 
@@ -1487,11 +1366,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         for batch in batches:
             try:
                 # Parse the batch file_object to check for file references
-                batch_data = (
-                    json.loads(batch.file_object)
-                    if isinstance(batch.file_object, str)
-                    else batch.file_object
-                )
+                batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object
 
                 # Extract file IDs from batch
                 # Batches typically reference the unified file ID in input_file_id
@@ -1500,9 +1375,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 output_file_id = batch_data.get("output_file_id")
                 error_file_id = batch_data.get("error_file_id")
 
-                referenced_file_ids = [
-                    fid for fid in [input_file_id, output_file_id, error_file_id] if fid
-                ]
+                referenced_file_ids = [fid for fid in [input_file_id, output_file_id, error_file_id] if fid]
 
                 # Check if any referenced file ID matches the file we're trying to delete
                 if any(ref_id in file_ids_to_check for ref_id in referenced_file_ids):
@@ -1514,9 +1387,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                         }
                     )
             except Exception as e:
-                verbose_logger.warning(
-                    f"Error parsing batch object {batch.unified_object_id}: {e}"
-                )
+                verbose_logger.warning(f"Error parsing batch object {batch.unified_object_id}: {e}")
                 continue
 
         return referencing_batches
@@ -1545,21 +1416,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
 
         if referencing_batches:
             # File is referenced by non-terminal batches and polling is enabled
-            MAX_BATCHES_IN_ERROR = (
-                5  # Limit batches shown in error message for readability
-            )
+            MAX_BATCHES_IN_ERROR = 5  # Limit batches shown in error message for readability
 
             # Show up to MAX_BATCHES_IN_ERROR in the error message
             batches_to_show = referencing_batches[:MAX_BATCHES_IN_ERROR]
-            batch_statuses = [
-                f"{b['batch_id']}: {b['status']}" for b in batches_to_show
-            ]
+            batch_statuses = [f"{b['batch_id']}: {b['status']}" for b in batches_to_show]
 
             # Determine the count message
             count_message = f"{len(referencing_batches)}"
-            if (
-                len(referencing_batches) >= 10
-            ):  # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file
+            if len(referencing_batches) >= 10:  # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file
                 count_message = "10+"
 
             error_message = (
@@ -1600,23 +1465,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         await self._check_file_deletion_allowed(file_id)
 
         # file_id = convert_b64_uid_to_unified_uid(file_id)
-        model_file_id_mapping = await self.get_model_file_id_mapping(
-            [file_id], litellm_parent_otel_span
-        )
+        model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span)
 
         delete_response = None
         specific_model_file_id_mapping = model_file_id_mapping.get(file_id)
         if specific_model_file_id_mapping:
             # Remove conflicting keys from data to avoid duplicate keyword arguments
-            filtered_data = {
-                k: v for k, v in data.items() if k not in ("model", "file_id")
-            }
+            filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")}
             for model_id, model_file_id in specific_model_file_id_mapping.items():
                 delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data)  # type: ignore
 
-        stored_file_object = await self.delete_unified_file_id(
-            file_id, litellm_parent_otel_span
-        )
+        stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span)
 
         # Record successful deletion metric only on actual success
         if stored_file_object or delete_response:
@@ -1643,9 +1502,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         Get the content of a file from first model that has it
         """
         model_file_id_mapping = data.pop("model_file_id_mapping", None)
-        model_file_id_mapping = (
-            model_file_id_mapping
-            or await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span)
+        model_file_id_mapping = model_file_id_mapping or await self.get_model_file_id_mapping(
+            [file_id], litellm_parent_otel_span
         )
 
         specific_model_file_id_mapping = model_file_id_mapping.get(file_id)
@@ -1658,13 +1516,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                     # against the deployment's configured bucket, which they only
                     # trust from this immutable server-side snapshot, never from
                     # request params.
-                    credentials = llm_router.get_deployment_credentials_with_provider(
-                        model_id=model_id
-                    )
+                    credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
                     if credentials is not None:
-                        data["_litellm_internal_model_credentials"] = cast(
-                            Dict, MappingProxyType(dict(credentials))
-                        )
+                        data["_litellm_internal_model_credentials"] = cast(Dict, MappingProxyType(dict(credentials)))
                     else:
                         data.pop("_litellm_internal_model_credentials", None)
                     return await llm_router.afile_content(model=model_id, file_id=provider_file_id, **data)  # type: ignore
@@ -1699,9 +1553,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             # Check database for storage backend info
             # IMPORTANT: The database stores the base64 encoded unified_file_id (not the decoded version)
             # So we query with the original file_id (which is base64 encoded)
-            db_file = await self.prisma_client.db.litellm_managedfiletable.find_first(
-                where={"unified_file_id": file_id}
-            )
+            db_file = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id})
 
             if not db_file or not db_file.storage_backend or not db_file.storage_url:
                 continue
@@ -1727,22 +1579,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 file_content = await storage_backend.download_file(storage_url)
 
                 # Determine content type from file object
-                content_type = self._get_content_type_from_file_object(
-                    db_file.file_object
-                )
+                content_type = self._get_content_type_from_file_object(db_file.file_object)
 
                 # Convert to base64
                 base64_data = base64.b64encode(file_content).decode("utf-8")
                 base64_data_uri = f"data:{content_type};base64,{base64_data}"
 
                 # Update messages to use base64 instead of file_id
-                self._update_messages_with_base64_data(
-                    messages, file_id, base64_data_uri, content_type
-                )
+                self._update_messages_with_base64_data(messages, file_id, base64_data_uri, content_type)
             except Exception as e:
-                verbose_logger.exception(
-                    f"Error converting file {file_id} from storage backend to base64: {str(e)}"
-                )
+                verbose_logger.exception(f"Error converting file {file_id} from storage backend to base64: {str(e)}")
                 # Continue with other files even if one fails
                 continue
 

litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py

@@ -6,16 +6,33 @@ This module provides fake streaming by converting non-streaming responses into s
 """
 
 import asyncio
-from collections.abc import AsyncIterator
-from typing import Any, Final, cast
+from collections.abc import AsyncIterator, Mapping, Sequence
+from typing import Any, Final, Protocol, cast, runtime_checkable
 from uuid import uuid4
 
+from pydantic import TypeAdapter
+
 from litellm._logging import verbose_logger
 from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
     get_async_httpx_client,
 )
 
+_ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object])
+_STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
+_LIST_ADAPTER: Final = TypeAdapter(list[object])
+_TEXT_ADAPTER: Final = TypeAdapter(str)
+
+
+@runtime_checkable
+class _SupportsModelDump(Protocol):
+    def model_dump(self, *, mode: str, exclude_none: bool) -> Mapping[str, object]: ...
+
+
+@runtime_checkable
+class _SupportsPydanticDict(Protocol):
+    def dict(self, *, exclude_none: bool) -> Mapping[str, object]: ...
+
 
 class PydanticAITransformation:
     """
@@ -28,7 +45,7 @@ class PydanticAITransformation:
     """
 
     @staticmethod
-    def _remove_none_values(obj: Any) -> Any:
+    def _remove_none_values(obj: object) -> object:
         """
         Recursively remove None values from a dict/list structure.
 
@@ -42,14 +59,18 @@ class PydanticAITransformation:
             Cleaned object with None values removed
         """
         if isinstance(obj, dict):
-            return {k: PydanticAITransformation._remove_none_values(v) for k, v in obj.items() if v is not None}
+            typed_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python(obj)
+            return {k: PydanticAITransformation._remove_none_values(v) for k, v in typed_dict.items() if v is not None}
         elif isinstance(obj, list):
-            return [PydanticAITransformation._remove_none_values(item) for item in obj if item is not None]
+            typed_list: Final = _LIST_ADAPTER.validate_python(obj)
+            return [PydanticAITransformation._remove_none_values(item) for item in typed_list if item is not None]
         else:
             return obj
 
     @staticmethod
-    def _params_to_dict(params: Any) -> dict[str, Any]:
+    def _params_to_dict(
+        params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
+    ) -> Mapping[str, object]:
         """
         Convert params to a dict, handling Pydantic models.
 
@@ -59,10 +80,10 @@ class PydanticAITransformation:
         Returns:
             Dict representation of params
         """
-        if hasattr(params, "model_dump"):
+        if isinstance(params, _SupportsModelDump):
             # Pydantic v2 model
             return params.model_dump(mode="python", exclude_none=True)
-        elif hasattr(params, "dict"):
+        elif isinstance(params, _SupportsPydanticDict):
             # Pydantic v1 model
             return params.dict(exclude_none=True)
         elif isinstance(params, dict):
@@ -75,12 +96,12 @@ class PydanticAITransformation:
     async def _poll_for_completion(
         client: AsyncHTTPHandler,
         endpoint: str,
-        task_id: str,
+        task_id: object,
         request_id: str,
         max_attempts: int = 30,
         poll_interval: float = 0.5,
         agent_extra_headers: dict[str, str] | None = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """
         Poll for task completion using tasks/get method.
 
@@ -112,10 +133,10 @@ class PydanticAITransformation:
                 },
             )
             response.raise_for_status()
-            poll_data = response.json()
+            poll_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json())
 
-            result = poll_data.get("result", {})
-            status = result.get("status", {})
+            result = _STR_KEY_DICT_ADAPTER.validate_python(poll_data.get("result", {}))
+            status = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {}))
             state = status.get("state", "")
 
             verbose_logger.debug("Pydantic AI: Poll attempt %s/%s, state=%s", attempt + 1, max_attempts, state)
@@ -133,10 +154,10 @@ class PydanticAITransformation:
     async def _send_and_poll_raw(
         api_base: str,
         request_id: str,
-        params: Any,
+        params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
         timeout: float = 60.0,
         agent_extra_headers: dict[str, str] | None = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """
         Send a request to Pydantic AI agent and return the raw task response.
 
@@ -153,14 +174,16 @@ class PydanticAITransformation:
             Raw Pydantic AI task response (with history/artifacts)
         """
         # Convert params to dict if it's a Pydantic model
-        params_dict = PydanticAITransformation._params_to_dict(params)
-
         # Remove None values - FastA2A doesn't accept null for optional fields
-        params_dict = PydanticAITransformation._remove_none_values(params_dict)
+        params_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python(
+            PydanticAITransformation._remove_none_values(PydanticAITransformation._params_to_dict(params))
+        )
 
         # Ensure the message has 'kind': 'message' as required by FastA2A/Pydantic AI
         if "message" in params_dict:
-            params_dict["message"]["kind"] = "message"
+            message_value: Final = _ANY_KEY_DICT_ADAPTER.validate_python(params_dict["message"])
+            message_value["kind"] = "message"
+            params_dict["message"] = message_value
 
         # Build A2A JSON-RPC request using message/send method for FastA2A compatibility
         a2a_request: Final = {
@@ -189,11 +212,11 @@ class PydanticAITransformation:
             },
         )
         response.raise_for_status()
-        response_data = response.json()
+        response_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json())
 
         # Check if task is already completed
-        result: Final = response_data.get("result", {})
-        status: Final = result.get("status", {})
+        result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {}))
+        status: Final = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {}))
         state: Final = status.get("state", "")
 
         if state != "completed":
@@ -217,10 +240,10 @@ class PydanticAITransformation:
     async def send_non_streaming_request(
         api_base: str,
         request_id: str,
-        params: Any,
+        params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
         timeout: float = 60.0,
         agent_extra_headers: dict[str, str] | None = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """
         Send a non-streaming A2A request to Pydantic AI agent and wait for completion.
 
@@ -253,10 +276,10 @@ class PydanticAITransformation:
     async def send_and_get_raw_response(
         api_base: str,
         request_id: str,
-        params: Any,
+        params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
         timeout: float = 60.0,
         agent_extra_headers: dict[str, str] | None = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """
         Send a request to Pydantic AI agent and return the raw task response.
 
@@ -282,9 +305,9 @@ class PydanticAITransformation:
 
     @staticmethod
     def _transform_to_a2a_response(
-        response_data: dict[str, Any],
+        response_data: Mapping[str, object],
         request_id: str,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """
         Transform Pydantic AI task response to standard A2A non-streaming format.
 
@@ -328,7 +351,7 @@ class PydanticAITransformation:
         }
 
     @staticmethod
-    def _extract_response_text(response_data: dict[str, Any]) -> tuple[str, str, list]:
+    def _extract_response_text(response_data: Mapping[str, object]) -> tuple[object, object, Sequence[object]]:
         """
         Extract response text from completed task response.
 
@@ -342,52 +365,53 @@ class PydanticAITransformation:
         Returns:
             Tuple of (full_text, message_id, parts)
         """
-        result: Final = response_data.get("result", {})
+        result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {}))
 
         # Try to extract from artifacts first (preferred for results)
         artifacts: Final = result.get("artifacts", [])
         if artifacts:
-            for artifact in artifacts:
-                parts = artifact.get("parts", [])
+            for artifact in _LIST_ADAPTER.validate_python(artifacts):
+                parts = _LIST_ADAPTER.validate_python(_STR_KEY_DICT_ADAPTER.validate_python(artifact).get("parts", []))
                 for part in parts:
-                    if part.get("kind") == "text":
-                        text = part.get("text", "")
+                    if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text":
+                        text = part_dict.get("text", "")
                         if text:
                             return text, str(uuid4()), parts
 
         # Fall back to history - get the last agent message
-        history: Final = result.get("history", [])
+        history: Final = _LIST_ADAPTER.validate_python(result.get("history", []))
         for msg in reversed(history):
-            if msg.get("role") == "agent":
-                parts = msg.get("parts", [])
-                message_id = msg.get("messageId", str(uuid4()))
+            if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "agent":
+                parts = _LIST_ADAPTER.validate_python(msg_dict.get("parts", []))
+                message_id = msg_dict.get("messageId", str(uuid4()))
                 full_text = ""
                 for part in parts:
-                    if part.get("kind") == "text":
-                        full_text += part.get("text", "")
+                    if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text":
+                        full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", ""))
                 if full_text:
                     return full_text, message_id, parts
 
         # Fall back to message field (original format)
         message: Final = result.get("message", {})
         if message:
-            parts = message.get("parts", [])
-            message_id = message.get("messageId", str(uuid4()))
+            message_dict: Final = _STR_KEY_DICT_ADAPTER.validate_python(message)
+            parts = _LIST_ADAPTER.validate_python(message_dict.get("parts", []))
+            message_id = message_dict.get("messageId", str(uuid4()))
             full_text = ""
             for part in parts:
-                if part.get("kind") == "text":
-                    full_text += part.get("text", "")
+                if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text":
+                    full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", ""))
             return full_text, message_id, parts
 
         return "", str(uuid4()), []
 
     @staticmethod
     async def fake_streaming_from_response(
-        response_data: dict[str, Any],
+        response_data: Mapping[str, object],
         request_id: str,
         chunk_size: int = 50,
         delay_ms: int = 10,
-    ) -> AsyncIterator[dict[str, Any]]:
+    ) -> AsyncIterator[dict[str, object]]:
         """
         Convert a non-streaming A2A response into fake streaming chunks.
 
@@ -410,12 +434,12 @@ class PydanticAITransformation:
         full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data)
 
         # Extract input message from raw response for history
-        result: Final = response_data.get("result", {})
-        history: Final = result.get("history", [])
-        input_message = {}
+        result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {}))
+        history: Final = _LIST_ADAPTER.validate_python(result.get("history", []))
+        input_message = _STR_KEY_DICT_ADAPTER.validate_python({})
         for msg in history:
-            if msg.get("role") == "user":
-                input_message = msg
+            if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "user":
+                input_message = msg_dict
                 break
 
         # Generate IDs for streaming events
@@ -426,45 +450,49 @@ class PydanticAITransformation:
 
         # 1. Emit initial task event (kind: "task", status: "submitted")
         # Format matches A2ACompletionBridgeTransformation.create_task_event
-        task_event: Final = {
-            "jsonrpc": "2.0",
-            "id": request_id,
-            "result": {
-                "contextId": context_id,
-                "history": [
-                    {
-                        "contextId": context_id,
-                        "kind": "message",
-                        "messageId": input_message_id,
-                        "parts": input_message.get("parts", [{"kind": "text", "text": ""}]),
-                        "role": "user",
-                        "taskId": task_id,
-                    }
-                ],
-                "id": task_id,
-                "kind": "task",
-                "status": {
-                    "state": "submitted",
+        task_event: Final = _STR_KEY_DICT_ADAPTER.validate_python(
+            {
+                "jsonrpc": "2.0",
+                "id": request_id,
+                "result": {
+                    "contextId": context_id,
+                    "history": [
+                        {
+                            "contextId": context_id,
+                            "kind": "message",
+                            "messageId": input_message_id,
+                            "parts": input_message.get("parts", [{"kind": "text", "text": ""}]),
+                            "role": "user",
+                            "taskId": task_id,
+                        }
+                    ],
+                    "id": task_id,
+                    "kind": "task",
+                    "status": {
+                        "state": "submitted",
+                    },
                 },
-            },
-        }
+            }
+        )
         yield task_event
 
         # 2. Emit status update (kind: "status-update", status: "working")
         # Format matches A2ACompletionBridgeTransformation.create_status_update_event
-        working_event: Final = {
-            "jsonrpc": "2.0",
-            "id": request_id,
-            "result": {
-                "contextId": context_id,
-                "final": False,
-                "kind": "status-update",
-                "status": {
-                    "state": "working",
+        working_event: Final = _STR_KEY_DICT_ADAPTER.validate_python(
+            {
+                "jsonrpc": "2.0",
+                "id": request_id,
+                "result": {
+                    "contextId": context_id,
+                    "final": False,
+                    "kind": "status-update",
+                    "status": {
+                        "state": "working",
+                    },
+                    "taskId": task_id,
                 },
-                "taskId": task_id,
-            },
-        }
+            }
+        )
         yield working_event
 
         # Small delay to simulate processing
@@ -473,29 +501,32 @@ class PydanticAITransformation:
         # 3. Emit artifact update chunks (kind: "artifact-update")
         # Format matches A2ACompletionBridgeTransformation.create_artifact_update_event
         if full_text:
+            full_text_str: Final = _TEXT_ADAPTER.validate_python(full_text)
             # Split text into chunks
-            for i in range(0, len(full_text), chunk_size):
-                chunk_text = full_text[i : i + chunk_size]
-                is_last_chunk = (i + chunk_size) >= len(full_text)
+            for i in range(0, len(full_text_str), chunk_size):
+                chunk_text = full_text_str[i : i + chunk_size]
+                is_last_chunk = (i + chunk_size) >= len(full_text_str)
 
-                artifact_event = {
-                    "jsonrpc": "2.0",
-                    "id": request_id,
-                    "result": {
-                        "contextId": context_id,
-                        "kind": "artifact-update",
-                        "taskId": task_id,
-                        "artifact": {
-                            "artifactId": artifact_id,
-                            "parts": [
-                                {
-                                    "kind": "text",
-                                    "text": chunk_text,
-                                }
-                            ],
+                artifact_event = _STR_KEY_DICT_ADAPTER.validate_python(
+                    {
+                        "jsonrpc": "2.0",
+                        "id": request_id,
+                        "result": {
+                            "contextId": context_id,
+                            "kind": "artifact-update",
+                            "taskId": task_id,
+                            "artifact": {
+                                "artifactId": artifact_id,
+                                "parts": [
+                                    {
+                                        "kind": "text",
+                                        "text": chunk_text,
+                                    }
+                                ],
+                            },
                         },
-                    },
-                }
+                    }
+                )
                 yield artifact_event
 
                 # Add delay between chunks (except for last chunk)
@@ -503,19 +534,21 @@ class PydanticAITransformation:
                     await asyncio.sleep(delay_ms / 1000.0)
 
         # 4. Emit final status update (kind: "status-update", status: "completed", final: true)
-        completed_event: Final = {
-            "jsonrpc": "2.0",
-            "id": request_id,
-            "result": {
-                "contextId": context_id,
-                "final": True,
-                "kind": "status-update",
-                "status": {
-                    "state": "completed",
+        completed_event: Final = _STR_KEY_DICT_ADAPTER.validate_python(
+            {
+                "jsonrpc": "2.0",
+                "id": request_id,
+                "result": {
+                    "contextId": context_id,
+                    "final": True,
+                    "kind": "status-update",
+                    "status": {
+                        "state": "completed",
+                    },
+                    "taskId": task_id,
                 },
-                "taskId": task_id,
-            },
-        }
+            }
+        )
         yield completed_event
 
         verbose_logger.info("Pydantic AI: Fake streaming completed for request_id=%s", request_id)

litellm/integrations/websearch_interception/handler.py

@@ -42,7 +42,7 @@ from litellm.types.integrations.websearch_interception import (
     WebSearchInterceptionConfig,
 )
 from litellm.types.llms.openai import AllMessageValues
-from litellm.types.utils import CallTypes, LlmProviders
+from litellm.types.utils import AgenticLoopParams, CallTypes, LlmProviders
 from litellm.utils import ProviderConfigManager
 
 if TYPE_CHECKING:
@@ -265,7 +265,7 @@ class WebSearchInterceptionLogger(CustomLogger):
             return None
 
         # Check if request has tools with native web_search
-        tools: Final = kwargs.get("tools")
+        tools: Final[Sequence[dict[str, object]] | None] = kwargs.get("tools")
         if not tools:
             return None
 
@@ -314,7 +314,9 @@ class WebSearchInterceptionLogger(CustomLogger):
 
         return kwargs
 
-    def _convert_responses_tools(self, kwargs: Mapping[str, object], tools: list[dict[str, object]]) -> dict | None:
+    def _convert_responses_tools(
+        self, kwargs: Mapping[str, object], tools: Sequence[dict[str, object]]
+    ) -> dict[str, object] | None:
         """Convert Responses API web search tools to the LiteLLM standard function tool."""
         if not any(is_web_search_tool_responses(tool) for tool in tools):
             return None
@@ -379,7 +381,7 @@ class WebSearchInterceptionLogger(CustomLogger):
         )
 
     @staticmethod
-    def _tool_name(tool: dict[str, Any]) -> str | None:
+    def _tool_name(tool: Mapping[str, object]) -> object:
         """Effective tool name, handling OpenAI ``function`` wrapper shape."""
         fn: Final = tool.get("function")
         if tool.get("type") == "function" and isinstance(fn, dict):
@@ -1271,7 +1273,7 @@ class WebSearchInterceptionLogger(CustomLogger):
         kwargs_for_followup: Final = self._prepare_followup_kwargs(kwargs)
 
         if logging_obj is not None:
-            agentic_params: Final = logging_obj.model_call_details.get("agentic_loop_params", {})
+            agentic_params: Final[AgenticLoopParams] = logging_obj.model_call_details.get("agentic_loop_params", {})
             full_model_name = agentic_params.get("model", model)
         verbose_logger.debug(
             "WebSearchInterception: Built anthropic request patch [call_id=%s model=%s messages=%d searches=%d]",

litellm/litellm_core_utils/realtime_streaming.py

@@ -1,5 +1,6 @@
 import asyncio
 import json
+from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final, Protocol, cast
 
 import litellm
@@ -20,11 +21,24 @@ from .litellm_logging import Logging as LiteLLMLogging
 if TYPE_CHECKING:
     from websockets.asyncio.client import ClientConnection
 
+    from litellm.types.guardrails import GuardrailEventHooks
+
     CLIENT_CONNECTION_CLASS = ClientConnection
 else:
     CLIENT_CONNECTION_CLASS = Any
 
 
+class _ClientWebSocketExceptions(Protocol):
+    ConnectionClosed: type[Exception]
+
+
+class _ClientWebSocket(Protocol):
+    exceptions: _ClientWebSocketExceptions
+
+    async def send_text(self, data: str) -> None: ...
+    async def receive_text(self) -> str: ...
+
+
 class RealtimeEventNormalizer(Protocol):
     def should_drop(self, event: object) -> bool: ...
     def normalize(self, event: dict) -> dict: ...
@@ -48,13 +62,13 @@ class RealTimeStreaming:
         logging_obj: LiteLLMLogging,
         provider_config: BaseRealtimeConfig | None = None,
         model: str = "",
-        user_api_key_dict: Any | None = None,
+        user_api_key_dict: object | None = None,
         request_data: dict | None = None,
         backend_uses_beta_protocol: bool | None = None,
         force_transcription_model: str | None = None,
         event_normalizer: RealtimeEventNormalizer | None = None,
     ):
-        self.websocket = websocket
+        self.websocket: _ClientWebSocket = websocket
         self.backend_ws = backend_ws
         self.logging_obj = logging_obj
         self.messages: list[OpenAIRealtimeEvents] = []
@@ -127,7 +141,7 @@ class RealTimeStreaming:
         ]
     )
     _CLIENT_AUDIO_BUFFER_COMMIT_TYPES = frozenset(["input_audio_buffer.commit", "input_audio_buffer.end"])
-    _AUDIO_FORMAT_MAP: dict[str, dict[str, Any]] = {
+    _AUDIO_FORMAT_MAP: dict[str, dict[str, str | int]] = {
         "pcm16": {"type": "audio/pcm", "rate": 24000},
         "g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000},
         "g711_alaw": {"type": "audio/G711-alaw", "rate": 8000},
@@ -281,6 +295,7 @@ class RealTimeStreaming:
             if event_obj.get("type") != "response.done":
                 return
             response: Final = cast(dict[str, Any], event_obj.get("response", {}))
+            item: Mapping[str, object]
             for item in response.get("output", []):
                 if item.get("type") == "function_call":
                     self.tool_calls.append(
@@ -384,7 +399,7 @@ class RealTimeStreaming:
             return message
 
         try:
-            message_obj: Final = json.loads(message)
+            message_obj: Final[Mapping[str, object]] = json.loads(message)
         except (json.JSONDecodeError, TypeError):
             return message
 
@@ -487,7 +502,7 @@ class RealTimeStreaming:
         if self._backend_setup_complete and not self._flushing_pending_messages_until_setup:
             return False
         try:
-            msg_obj: Final = json.loads(message)
+            msg_obj: Final[Mapping[str, object]] = json.loads(message)
         except (json.JSONDecodeError, TypeError):
             return False
         return msg_obj.get("type") in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES
@@ -555,7 +570,7 @@ class RealTimeStreaming:
     def _event_to_client_json(self, event: dict) -> str:
         return json.dumps(self._normalize_event_for_ga_client(event))
 
-    async def _send_event_to_client(self, event: Any, event_str: str) -> bool:
+    async def _send_event_to_client(self, event: object, event_str: str) -> bool:
         if self._should_drop_event_from_client(event):
             return False
         if isinstance(event, dict):
@@ -595,12 +610,12 @@ class RealTimeStreaming:
 
     def _make_disable_auto_response_message(self) -> str:
         """Return a session.update that disables VAD auto-response."""
-        turn_detection: Final[dict[str, Any]] = {
+        turn_detection: Final[dict[str, str | bool]] = {
             "type": "server_vad",
             "create_response": False,
         }
         if self._backend_uses_beta_protocol:
-            session: dict[str, Any] = {"turn_detection": turn_detection}
+            session: dict[str, object] = {"turn_detection": turn_detection}
         else:
             session = {
                 "type": "realtime",
@@ -654,7 +669,7 @@ class RealTimeStreaming:
 
     def _has_realtime_guardrails_for_event_hooks(
         self,
-        event_hooks: list[Any],
+        event_hooks: Sequence["GuardrailEventHooks"],
     ) -> bool:
         """Return True if any callback would run for one of ``event_hooks``."""
         from litellm.integrations.custom_guardrail import CustomGuardrail
@@ -699,7 +714,7 @@ class RealTimeStreaming:
         transcript: str,
         item_id: str | None = None,
         pre_block_backend_message: str | None = None,
-        event_hooks: list[Any] | None = None,
+        event_hooks: Sequence["GuardrailEventHooks"] | None = None,
     ) -> bool:
         """
         Run registered guardrails on realtime text (transcript, user message, tool output).
@@ -753,7 +768,7 @@ class RealTimeStreaming:
                     raise
                 # Extract the human-readable error from the detail dict (HTTPException)
                 # or fall back to str(e) for plain ValueError.
-                detail = getattr(e, "detail", None)
+                detail: object | None = getattr(e, "detail", None)
                 if isinstance(detail, dict):
                     safe_msg = detail.get("error") or str(e)
                 elif detail is not None:
@@ -826,7 +841,7 @@ class RealTimeStreaming:
                 return True
         return False
 
-    async def _handle_provider_config_message(self, raw_response) -> None:
+    async def _handle_provider_config_message(self, raw_response: str) -> None:
         """Process a backend message when a provider_config is set (transformed path)."""
         returned_object: Final = self.provider_config.transform_realtime_response(
             raw_response,
@@ -910,7 +925,7 @@ class RealTimeStreaming:
             await self._send_event_to_client(event, event_str)
 
     @staticmethod
-    def _parse_backend_event(raw_response: str) -> dict | None:
+    def _parse_backend_event(raw_response: str) -> dict[str, object] | None:
         """Parse a backend frame once. Returns None for non-JSON or non-object frames."""
         try:
             event: Final = json.loads(raw_response)
@@ -1020,7 +1035,7 @@ class RealTimeStreaming:
         objects and any test doubles that expose a .scope dict.
         """
         try:
-            headers: Final = websocket.scope.get("headers", [])
+            headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = websocket.scope.get("headers", [])
             for name, value in headers:
                 if isinstance(name, bytes):
                     name = name.decode("latin-1")
@@ -1071,9 +1086,9 @@ class RealTimeStreaming:
                     session["output_modalities"] = ["text"]
 
         # 3-7. Lift flat audio fields into the nested audio object
-        audio: Final[dict[str, Any]] = {}
-        inp: Final[dict[str, Any]] = {}
-        out: Final[dict[str, Any]] = {}
+        audio: Final[dict[str, object]] = {}
+        inp: Final[dict[str, object]] = {}
+        out: Final[dict[str, object]] = {}
 
         # voice → audio.output.voice
         if "voice" in session:
@@ -1190,7 +1205,7 @@ class RealTimeStreaming:
                             # model; check them with the same guardrail used for
                             # user text so an attacker cannot smuggle blocked
                             # content into a function_call_output.
-                            output = item.get("output", "")
+                            output: object = item.get("output", "")
                             output_text = output if isinstance(output, str) else json.dumps(output)
                             if output_text:
                                 # Build the sanitized function_call_output up
@@ -1241,7 +1256,7 @@ class RealTimeStreaming:
                                     # interaction turn.
                                     continue
                         elif item.get("role") == "user":
-                            content_list = item.get("content", [])
+                            content_list: Sequence[object] = item.get("content", [])
                             texts = [
                                 c.get("text", "")
                                 for c in content_list
@@ -1280,7 +1295,7 @@ class RealTimeStreaming:
                         and not self._guardrail_turn_detection_update_sent
                         and self._has_audio_transcription_guardrails()
                     ):
-                        session = msg_obj.setdefault("session", {})
+                        session: object = msg_obj.setdefault("session", {})
                         if isinstance(session, dict):
                             existing_td = session.get("turn_detection")
                             if not isinstance(existing_td, dict):

litellm/litellm_core_utils/streaming_chunk_builder_utils.py

@@ -3,7 +3,7 @@ import time
 from collections.abc import Iterator, Mapping, Sequence
 from itertools import groupby
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, Union, cast
+from typing import TYPE_CHECKING, Any, Final, TypedDict, Union, cast
 
 from litellm._logging import verbose_logger
 from litellm.types.llms.openai import (
@@ -30,6 +30,7 @@ from litellm.types.utils import (
 from litellm.utils import print_verbose, token_counter
 
 if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging
     from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import (
         UsagePerChunk,
     )
@@ -39,6 +40,60 @@ if TYPE_CHECKING:
     )
 
 
+class _ThinkingBlockFragment(TypedDict, total=False):
+    type: str | None
+    data: str | None
+    thinking: str | None
+    signature: str | None
+
+
+class _ThinkingDelta(TypedDict, total=False):
+    thinking_blocks: Sequence[_ThinkingBlockFragment]
+
+
+class _ThinkingChoice(TypedDict, total=False):
+    delta: _ThinkingDelta
+
+
+class _ThinkingChunk(TypedDict):
+    choices: Sequence[_ThinkingChoice]
+
+
+class _ContentChoice(TypedDict, total=False):
+    delta: Mapping[str, str | None]
+
+
+class _ContentChunk(TypedDict):
+    choices: Sequence[_ContentChoice]
+
+
+class _AudioDelta(TypedDict, total=False):
+    audio: ChatCompletionAudioDelta | None
+
+
+class _AudioChoice(TypedDict, total=False):
+    delta: _AudioDelta
+
+
+class _AudioChunk(TypedDict):
+    choices: Sequence[_AudioChoice]
+
+
+class _UsageBearingChunk(TypedDict, total=False):
+    usage: Usage | None
+    _hidden_params: Mapping[str, str]
+
+
+class _UsageSummary(TypedDict):
+    prompt_tokens: int | None
+    completion_tokens: int | None
+    cache_creation_input_tokens: int | None
+    cache_read_input_tokens: int | None
+    completion_tokens_details: CompletionTokensDetails | None
+    prompt_tokens_details: PromptTokensDetailsWrapper | None
+    cost: float | None
+
+
 def capture_cache_creation_token_details(
     prompt_tokens_details: PromptTokensDetailsWrapper | None,
     current: CacheCreationTokenDetails | None,
@@ -78,7 +133,7 @@ class ChunkProcessor:
             return []
 
         first_chunk: Final = chunks[0]
-        first_hidden_params: dict[str, Any] = {}
+        first_hidden_params: dict[str, object] = {}
         if isinstance(first_chunk, dict):
             candidate = first_chunk.get("_hidden_params", {})
             if isinstance(candidate, dict):
@@ -115,8 +170,8 @@ class ChunkProcessor:
     @staticmethod
     def apply_provider_assembled_streaming_metadata(
         response: ModelResponse,
-        chunks: list[Any],
-        logging_obj: Any | None = None,
+        chunks: list[object],
+        logging_obj: "Logging | None" = None,
     ) -> None:
         if not chunks:
             return
@@ -456,7 +511,7 @@ class ChunkProcessor:
         )
 
     def get_combined_content(
-        self, chunks: list[dict[str, Any]], delta_key: str = "content"
+        self, chunks: Sequence["_ContentChunk"], delta_key: str = "content"
     ) -> ChatCompletionAssistantContentValue:
         content_list: Final[list[str]] = []
         for chunk in chunks:
@@ -475,7 +530,7 @@ class ChunkProcessor:
         return combined_content
 
     def get_combined_thinking_content(
-        self, chunks: list[dict[str, Any]]
+        self, chunks: Sequence["_ThinkingChunk"]
     ) -> list[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]] | None:
         from litellm.types.llms.openai import (
             ChatCompletionRedactedThinkingBlock,
@@ -532,10 +587,10 @@ class ChunkProcessor:
             return thinking_blocks
         return None
 
-    def get_combined_reasoning_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAssistantContentValue:
+    def get_combined_reasoning_content(self, chunks: Sequence["_ContentChunk"]) -> ChatCompletionAssistantContentValue:
         return self.get_combined_content(chunks, delta_key="reasoning_content")
 
-    def get_combined_audio_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAudioResponse:
+    def get_combined_audio_content(self, chunks: Sequence["_AudioChunk"]) -> ChatCompletionAudioResponse:
         base64_data_list: Final[list[str]] = []
         transcript_list: Final[list[str]] = []
         expires_at: int | None = None
@@ -544,7 +599,7 @@ class ChunkProcessor:
         for chunk in chunks:
             choices = chunk["choices"]
             for choice in choices:
-                delta = choice.get("delta") or {}
+                delta: _AudioDelta = choice.get("delta") or {}
                 audio: ChatCompletionAudioDelta | None = delta.get("audio")
                 if audio is not None:
                     for k, v in audio.items():
@@ -565,7 +620,7 @@ class ChunkProcessor:
             id=id,
         )
 
-    def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> dict:
+    def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> "_UsageSummary":
         prompt_tokens = 0
         completion_tokens = 0
         ## anthropic prompt caching information ##
@@ -623,8 +678,8 @@ class ChunkProcessor:
         return reasoning_tokens
 
     @staticmethod
-    def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse | ModelResponseStream) -> Usage | None:
-        usage_chunk: Usage | dict[str, Any] | None = None
+    def _extract_usage_chunk(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Usage | None:
+        usage_chunk: Usage | None = None
         if hasattr(chunk, "usage") and chunk.usage is not None:
             usage_chunk = chunk.usage
         elif "usage" in chunk:
@@ -640,7 +695,7 @@ class ChunkProcessor:
 
     def _calculate_usage_per_chunk(
         self,
-        chunks: list[dict[str, Any] | ModelResponse],
+        chunks: Sequence["_UsageBearingChunk | ModelResponse"],
     ) -> "UsagePerChunk":
         from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import (
             UsagePerChunk,
@@ -721,13 +776,7 @@ class ChunkProcessor:
                         "web_search_requests",
                     )
 
-                prompt_tokens_details = (
-                    cast(
-                        PromptTokensDetailsWrapper | None,
-                        usage_chunk_dict["prompt_tokens_details"],
-                    )
-                    or prompt_tokens_details
-                )
+                prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"] or prompt_tokens_details
 
                 cache_creation_token_details = capture_cache_creation_token_details(
                     prompt_tokens_details, cache_creation_token_details
@@ -758,7 +807,7 @@ class ChunkProcessor:
 
     @staticmethod
     def _reset_anthropic_cursor_completion_tokens(
-        chunks: list[dict[str, Any] | ModelResponse],
+        chunks: Sequence["_UsageBearingChunk | ModelResponse"],
         completion_tokens: int,
         completion_usage_updates: int,
     ) -> int:
@@ -797,7 +846,7 @@ class ChunkProcessor:
 
     def calculate_usage(
         self,
-        chunks: list[dict[str, Any] | ModelResponse],
+        chunks: Sequence["_UsageBearingChunk | ModelResponse"],
         model: str,
         completion_output: str,
         messages: list | None = None,
@@ -851,8 +900,8 @@ class ChunkProcessor:
             setattr(returned_usage, "cache_read_input_tokens", cache_read_input_tokens)  # for anthropic
         if completion_tokens_details is not None:
             if isinstance(completion_tokens_details, CompletionTokensDetails):
-                returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper(
-                    **completion_tokens_details.model_dump()
+                returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper.model_validate(
+                    completion_tokens_details.model_dump()
                 )
             else:
                 returned_usage.completion_tokens_details = completion_tokens_details

litellm/llms/anthropic/experimental_pass_through/adapters/handler.py

@@ -1,8 +1,9 @@
-from collections.abc import AsyncIterator, Coroutine, Iterator
+from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping
 from typing import (
     TYPE_CHECKING,
     Any,
     Final,
+    TypeAlias,
     cast,
 )
 
@@ -33,8 +34,12 @@ if TYPE_CHECKING:
 # Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge.
 ANTHROPIC_ONLY_REQUEST_KEYS: Final[frozenset[str]] = frozenset({"output_config"})
 
+_AnthropicMessages: TypeAlias = "list[dict[str, object]]"
+_AnthropicSystem: TypeAlias = "str | list[dict[str, object]] | None"
+_ContextManagementSpec: TypeAlias = "dict[str, object] | list[dict[str, object]] | None"
 
-def _messages_have_compaction_block(messages: list[dict]) -> bool:
+
+def _messages_have_compaction_block(messages: _AnthropicMessages) -> bool:
     """Return True when any message carries a ``compaction`` content block."""
     for msg in messages:
         content = msg.get("content")
@@ -54,8 +59,10 @@ def _proxy_router_fallback() -> "Router | None":
     return _proxy_router
 
 
-def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None:
-    """Return ``kwargs["litellm_metadata"]`` when it's a dict; ``None`` otherwise.
+def _extract_proxy_litellm_metadata(
+    kwargs: Mapping[str, object],
+) -> "tuple[dict[str, object], UserAPIKeyAuth | None] | tuple[None, None]":
+    """Return ``(kwargs["litellm_metadata"], its user_api_key_auth)`` when it's a dict; ``(None, None)`` otherwise.
 
     The proxy attaches its auth/spend-attribution fields (``user_api_key``,
     ``user_api_key_team_id``, ``litellm_call_id``, the full ``UserAPIKeyAuth``
@@ -68,18 +75,19 @@ def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] |
     """
     litellm_metadata: Final = kwargs.get("litellm_metadata")
     if not isinstance(litellm_metadata, dict):
-        return None
-    return litellm_metadata
+        return None, None
+    user_api_key_auth: Final[UserAPIKeyAuth | None] = litellm_metadata.get("user_api_key_auth")
+    return litellm_metadata, user_api_key_auth
 
 
 async def _prepare_context_managed_request(
     *,
     model: str,
-    messages: list[dict],
-    tools: list[dict] | None,
-    system: Any | None,
-    context_management_spec: Any,
-    litellm_metadata: dict | None,
+    messages: _AnthropicMessages,
+    tools: list[dict[str, object]] | None,
+    system: _AnthropicSystem,
+    context_management_spec: _ContextManagementSpec,
+    litellm_metadata: dict[str, object] | None,
     additional_drop_params: list[str] | None,
     llm_router: "Router | None",
     user_api_key_auth: "UserAPIKeyAuth | None" = None,
@@ -102,11 +110,11 @@ async def _prepare_context_managed_request(
 
     if polyfill_will_run:
         history_result: PolyfillResult | None = None
-        working_messages: list[dict] = messages
-        working_system: Any | None = system
+        working_messages: _AnthropicMessages = messages
+        working_system: _AnthropicSystem = system
     else:
         history_result = apply_client_compaction_block_history(
-            messages=cast(list[dict[str, Any]], messages),
+            messages=messages,
             system=system,
         )
         working_messages = history_result.messages if history_result is not None else messages
@@ -136,7 +144,7 @@ async def _prepare_context_managed_request(
     # to non-Anthropic backends that would reject them.
     if polyfill_will_run and history_result is None:
         history_result = apply_client_compaction_block_history(
-            messages=cast(list[dict[str, Any]], messages),
+            messages=messages,
             system=system,
         )
     return history_result
@@ -144,7 +152,7 @@ async def _prepare_context_managed_request(
 
 def _polyfill_will_run(
     *,
-    context_management_spec: Any,
+    context_management_spec: _ContextManagementSpec,
     additional_drop_params: list[str] | None,
 ) -> bool:
     """Return True when ``compact_20260112`` will run via the polyfill dispatcher.
@@ -171,7 +179,7 @@ def _polyfill_will_run(
 
 def _spec_has_non_compact_edits(
     *,
-    context_management_spec: Any,
+    context_management_spec: _ContextManagementSpec,
     additional_drop_params: list[str] | None,
 ) -> bool:
     """Return True when the spec includes edits other than ``compact_20260112``.
@@ -209,9 +217,9 @@ def _context_management_explicitly_dropped(additional_drop_params: list[str] | N
 
 def _normalize_spec_edits(
     *,
-    context_management_spec: Any,
+    context_management_spec: _ContextManagementSpec,
     additional_drop_params: list[str] | None,
-) -> list[dict[str, Any]] | None:
+) -> list[dict[str, object]] | None:
     """Return the normalized ``edits`` list, or ``None`` if the polyfill won't run.
 
     Delegates spec-shape normalization to the dispatcher's ``_normalize_spec``
@@ -236,11 +244,11 @@ def _normalize_spec_edits(
 async def _run_polyfill_if_enabled(
     *,
     model: str,
-    messages: list[dict],
-    tools: list[dict] | None,
-    system: Any | None,
-    context_management_spec: Any,
-    litellm_metadata: dict | None,
+    messages: _AnthropicMessages,
+    tools: list[dict[str, object]] | None,
+    system: _AnthropicSystem,
+    context_management_spec: _ContextManagementSpec,
+    litellm_metadata: dict[str, object] | None,
     additional_drop_params: list[str] | None,
     llm_router: "Router | None",
     user_api_key_auth: "UserAPIKeyAuth | None" = None,
@@ -306,7 +314,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
     def _route_openai_thinking_to_responses_api_if_needed(
         completion_kwargs: dict[str, Any],
         *,
-        thinking: dict[str, Any] | None,
+        thinking: Mapping[str, object] | None,
     ) -> None:
         """
         When users call `litellm.anthropic.messages.*` with a non-Anthropic model and
@@ -407,12 +415,12 @@ class LiteLLMMessagesToCompletionTransformationHandler:
     def _prepare_completion_kwargs(
         *,
         max_tokens: int,
-        messages: list[dict],
+        messages: _AnthropicMessages,
         model: str,
         metadata: dict | None = None,
         stop_sequences: list[str] | None = None,
         stream: bool | None = False,
-        system: str | list[dict[str, Any]] | None = None,
+        system: _AnthropicSystem = None,
         temperature: float | None = None,
         thinking: dict | None = None,
         tool_choice: dict | None = None,
@@ -420,7 +428,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
         top_k: int | None = None,
         top_p: float | None = None,
         output_format: dict | None = None,
-        extra_kwargs: dict[str, Any] | None = None,
+        extra_kwargs: Mapping[str, object] | None = None,
     ) -> tuple[dict[str, Any], dict[str, str]]:
         """Prepare kwargs for litellm.completion/acompletion.
 
@@ -433,7 +441,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
             Logging as LiteLLMLoggingObject,
         )
 
-        request_data: Final = {
+        request_data: Final[dict[str, object]] = {
             "model": model,
             "messages": messages,
             "max_tokens": max_tokens,
@@ -528,7 +536,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
     @staticmethod
     async def async_anthropic_messages_handler(
         max_tokens: int,
-        messages: list[dict],
+        messages: _AnthropicMessages,
         model: str,
         metadata: dict | None = None,
         stop_sequences: list[str] | None = None,
@@ -537,7 +545,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
         temperature: float | None = None,
         thinking: dict | None = None,
         tool_choice: dict | None = None,
-        tools: list[dict] | None = None,
+        tools: list[dict[str, object]] | None = None,
         top_k: int | None = None,
         top_p: float | None = None,
         output_format: dict | None = None,
@@ -551,10 +559,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
             requested_router if requested_router is not None else _proxy_router_fallback()
         )
 
-        proxy_litellm_metadata: Final = _extract_proxy_litellm_metadata(kwargs)
-        user_api_key_auth: Final[UserAPIKeyAuth | None] = (
-            proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None
-        )
+        proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs)
 
         polyfill_result: Final = await _prepare_context_managed_request(
             model=model,
@@ -618,7 +623,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
     @staticmethod
     def anthropic_messages_handler(
         max_tokens: int,
-        messages: list[dict],
+        messages: _AnthropicMessages,
         model: str,
         metadata: dict | None = None,
         stop_sequences: list[str] | None = None,
@@ -627,7 +632,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
         temperature: float | None = None,
         thinking: dict | None = None,
         tool_choice: dict | None = None,
-        tools: list[dict] | None = None,
+        tools: list[dict[str, object]] | None = None,
         top_k: int | None = None,
         top_p: float | None = None,
         output_format: dict | None = None,
@@ -688,10 +693,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
         if context_management is None and not _messages_have_compaction_block(messages):
             polyfill_result: PolyfillResult | None = None
         else:
-            proxy_litellm_metadata: Final = _extract_proxy_litellm_metadata(kwargs)
-            user_api_key_auth: Final[UserAPIKeyAuth | None] = (
-                proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None
-            )
+            proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs)
             polyfill_result = run_async_function(
                 _prepare_context_managed_request,
                 model=model,

litellm/llms/azure/assistants.py

@@ -1,8 +1,9 @@
 from collections.abc import Coroutine, Iterable
-from typing import Any, Final, Literal
+from typing import Any, Final, Literal, TypedDict
 
 import httpx
 from openai import AsyncAzureOpenAI, AzureOpenAI
+from openai.types.shared_params.metadata import Metadata
 from typing_extensions import overload
 
 from ...types.llms.openai import (
@@ -22,6 +23,16 @@ from ...types.llms.openai import (
 from .common_utils import BaseAzureLLM
 
 
+class _RunThreadStreamData(TypedDict):
+    thread_id: str
+    assistant_id: str
+    additional_instructions: str | None
+    instructions: str | None
+    metadata: Metadata | None
+    model: str | None
+    tools: Iterable[AssistantToolParam] | None
+
+
 class AzureAssistantsAPI(BaseAzureLLM):
     def __init__(self) -> None:
         super().__init__()
@@ -212,9 +223,9 @@ class AzureAssistantsAPI(BaseAzureLLM):
         response_obj: OpenAIMessage | None = None
         if getattr(thread_message, "status", None) is None:
             thread_message.status = "completed"
-            response_obj = OpenAIMessage(**thread_message.dict())
+            response_obj = OpenAIMessage.model_validate(thread_message.dict())
         else:
-            response_obj = OpenAIMessage(**thread_message.dict())
+            response_obj = OpenAIMessage.model_validate(thread_message.dict())
         return response_obj
 
     # fmt: off
@@ -301,9 +312,9 @@ class AzureAssistantsAPI(BaseAzureLLM):
         response_obj: OpenAIMessage | None = None
         if getattr(thread_message, "status", None) is None:
             thread_message.status = "completed"
-            response_obj = OpenAIMessage(**thread_message.dict())
+            response_obj = OpenAIMessage.model_validate(thread_message.dict())
         else:
-            response_obj = OpenAIMessage(**thread_message.dict())
+            response_obj = OpenAIMessage.model_validate(thread_message.dict())
         return response_obj
 
     async def async_get_messages(
@@ -443,7 +454,7 @@ class AzureAssistantsAPI(BaseAzureLLM):
 
         message_thread: Final = await openai_client.beta.threads.create(**data)
 
-        return Thread(**message_thread.dict())
+        return Thread.model_validate(message_thread.dict())
 
     # fmt: off
 
@@ -539,7 +550,7 @@ class AzureAssistantsAPI(BaseAzureLLM):
 
         message_thread: Final = azure_openai_client.beta.threads.create(**data)
 
-        return Thread(**message_thread.dict())
+        return Thread.model_validate(message_thread.dict())
 
     async def async_get_thread(
         self,
@@ -566,7 +577,7 @@ class AzureAssistantsAPI(BaseAzureLLM):
 
         response: Final = await openai_client.beta.threads.retrieve(thread_id=thread_id)
 
-        return Thread(**response.dict())
+        return Thread.model_validate(response.dict())
 
     # fmt: off
 
@@ -642,7 +653,7 @@ class AzureAssistantsAPI(BaseAzureLLM):
 
         response: Final = openai_client.beta.threads.retrieve(thread_id=thread_id)
 
-        return Thread(**response.dict())
+        return Thread.model_validate(response.dict())
 
     # def delete_thread(self):
     #     pass
@@ -730,7 +741,8 @@ class AzureAssistantsAPI(BaseAzureLLM):
         event_handler: AssistantEventHandler | None,
         litellm_params: dict | None = None,
     ) -> AssistantStreamManager[AssistantEventHandler]:
-        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,
@@ -740,8 +752,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)
 
     # fmt: off
 

litellm/llms/vertex_ai/batches/handler.py

@@ -14,7 +14,7 @@ from litellm.llms.custom_httpx.http_handler import (
     _get_httpx_client,
     get_async_httpx_client,
 )
-from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
+from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url
 from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
 from litellm.types.llms.openai import CreateBatchRequest
 from litellm.types.llms.vertex_ai import (
@@ -98,9 +98,6 @@ class VertexAIBatchPrediction(VertexLLM):
             data=json.dumps(vertex_batch_request),
         )
 
-        if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
-
         _json_response: Final = response.json()
         vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(
             response=_json_response
@@ -130,8 +127,6 @@ class VertexAIBatchPrediction(VertexLLM):
                 error_body[:1000],
             )
             raise
-        if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
 
         _json_response: Final = response.json()
         vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(
@@ -243,7 +238,9 @@ class VertexAIBatchPrediction(VertexLLM):
         )
 
         if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
+            raise VertexAIError(
+                status_code=response.status_code, message=f"Error: {response.status_code} {response.text}"
+            )
 
         _json_response: Final = response.json()
         vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(
@@ -293,7 +290,9 @@ class VertexAIBatchPrediction(VertexLLM):
             headers=headers,
         )
         if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
+            raise VertexAIError(
+                status_code=response.status_code, message=f"Error: {response.status_code} {response.text}"
+            )
 
         _json_response: Final = response.json()
         vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(
@@ -366,7 +365,9 @@ class VertexAIBatchPrediction(VertexLLM):
         )
 
         if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
+            raise VertexAIError(
+                status_code=response.status_code, message=f"Error: {response.status_code} {response.text}"
+            )
 
         _json_response: Final = response.json()
         vertex_batch_response: Final = (
@@ -391,7 +392,9 @@ class VertexAIBatchPrediction(VertexLLM):
             params=params,
         )
         if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
+            raise VertexAIError(
+                status_code=response.status_code, message=f"Error: {response.status_code} {response.text}"
+            )
 
         _json_response: Final = response.json()
         vertex_batch_response: Final = (
@@ -461,7 +464,7 @@ class VertexAIBatchPrediction(VertexLLM):
 
         sync_handler: Final = _get_httpx_client()
         try:
-            response: Final = sync_handler.post(
+            sync_handler.post(
                 url=api_base,
                 headers=headers,
                 data=json.dumps({}),
@@ -475,9 +478,6 @@ class VertexAIBatchPrediction(VertexLLM):
             )
             raise
 
-        if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
-
         # HTTPHandler.get() does not accept a timeout parameter
         retrieve_response: Final = sync_handler.get(
             url=retrieve_api_base,
@@ -489,7 +489,10 @@ class VertexAIBatchPrediction(VertexLLM):
                 retrieve_response.status_code,
                 retrieve_response.text[:1000],
             )
-            raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}")
+            raise VertexAIError(
+                status_code=retrieve_response.status_code,
+                message=f"Error: {retrieve_response.status_code} {retrieve_response.text}",
+            )
 
         _json_response: Final = retrieve_response.json()
         vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(
@@ -508,7 +511,7 @@ class VertexAIBatchPrediction(VertexLLM):
             llm_provider=litellm.LlmProviders.VERTEX_AI,
         )
         try:
-            response: Final = await client.post(
+            await client.post(
                 url=api_base,
                 headers=headers,
                 data=json.dumps({}),
@@ -521,8 +524,6 @@ class VertexAIBatchPrediction(VertexLLM):
                 e.response.text[:1000],
             )
             raise
-        if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
 
         # AsyncHTTPHandler.get() does not accept a timeout parameter
         retrieve_response: Final = await client.get(
@@ -535,7 +536,10 @@ class VertexAIBatchPrediction(VertexLLM):
                 retrieve_response.status_code,
                 retrieve_response.text[:1000],
             )
-            raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}")
+            raise VertexAIError(
+                status_code=retrieve_response.status_code,
+                message=f"Error: {retrieve_response.status_code} {retrieve_response.text}",
+            )
 
         _json_response: Final = retrieve_response.json()
         vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(

litellm/llms/vertex_ai/batches/transformation.py

@@ -1,7 +1,9 @@
 from typing import Any, Final
+from urllib.parse import unquote
 
 from litellm._uuid import uuid
 from litellm.llms.vertex_ai.common_utils import (
+    VertexAIError,
     _convert_vertex_datetime_to_openai_datetime,
 )
 from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest
@@ -199,16 +201,40 @@ class VertexAIBatchTransformation:
 
         gcs_file_uri format: gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8
         returns: "publishers/google/models/gemini-1.5-flash-001"
+
+        Raises a 400 `VertexAIError` when the uri carries no parseable model path.
         """
-        from urllib.parse import unquote
-
-        decoded_uri: Final = unquote(gcs_file_uri)
-
-        model_path: Final = decoded_uri.split("publishers/")[1]
-        parts: Final = model_path.split("/")
-        model: Final = f"publishers/{'/'.join(parts[:3])}"
+        model: Final = cls._parse_model_from_gcs_file(gcs_file_uri)
+        if model is None:
+            raise VertexAIError(
+                status_code=400,
+                message=(
+                    "Vertex AI batch creation requires the model to be part of `input_file_id`, but "
+                    f"'{gcs_file_uri}' contains no 'publishers/<publisher>/models/<model>' path segment. "
+                    "Either upload the input file through LiteLLM (POST /v1/files with "
+                    "custom_llm_provider=vertex_ai), which encodes the model into the returned file id, or "
+                    "pass a uri of the form "
+                    "gs://<bucket>/<prefix>/publishers/<publisher>/models/<model>/<file>"
+                ),
+            )
         return model
 
+    @classmethod
+    def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None:
+        """
+        Returns the `publishers/<publisher>/models/<model>` path from a gcs uri, or None if the uri
+        does not contain one.
+        """
+        _, separator, model_path = unquote(gcs_file_uri).partition("publishers/")
+        if not separator:
+            return None
+
+        parts: Final = model_path.split("/")
+        if len(parts) < 3 or parts[1] != "models" or not parts[2]:
+            return None
+
+        return f"publishers/{'/'.join(parts[:3])}"
+
     @classmethod
     def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: str | None) -> bool:
         """
@@ -216,7 +242,11 @@ class VertexAIBatchTransformation:
         LiteLLM-managed unified file id) with a `publishers/` model path that
         `_get_model_from_gcs_file` can parse.
         """
-        return input_file_id is not None and input_file_id.startswith("gs://") and "publishers/" in input_file_id
+        return (
+            input_file_id is not None
+            and input_file_id.startswith("gs://")
+            and cls._parse_model_from_gcs_file(input_file_id) is not None
+        )
 
     @classmethod
     def get_bare_model_name_from_gcs_file(cls, gcs_file_uri: str) -> str:

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -109,7 +109,7 @@ if MCP_AVAILABLE:
     ############ MCP Server REST API Routes #################
     async def _safe_fire_mcp_tool_call_logging(
         logging_obj: Any | None,
-        result: Any,
+        result: "CallToolResult",
         start_time: datetime,
         end_time: datetime,
         user_api_key_auth: UserAPIKeyAuth | None = None,

litellm/proxy/_experimental/mcp_server/server.py

@@ -13,9 +13,9 @@ import time
 import traceback
 import types
 import uuid
-from collections.abc import AsyncIterator, Callable, Mapping
+from collections.abc import AsyncIterator, Callable, Mapping, Sequence
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Final, cast
+from typing import TYPE_CHECKING, Any, Final, Protocol
 
 import httpx
 from fastapi import FastAPI, HTTPException
@@ -145,7 +145,7 @@ try:
     )
 
     # Robust auth lookup keyed by session_object.
-    _session_obj_auth_storage: "weakref.WeakKeyDictionary[Any, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary()
+    _session_obj_auth_storage: "weakref.WeakKeyDictionary[object, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary()
 except ImportError as e:
     verbose_logger.debug("MCP module not found: %s", e)
     MCP_AVAILABLE = False
@@ -493,14 +493,14 @@ if MCP_AVAILABLE:
     def _gateway_create_initialization_options(
         self,
         notification_options: NotificationOptions | None = None,
-        experimental_capabilities: dict[str, dict[str, Any]] | None = None,
+        experimental_capabilities: dict[str, dict[str, object]] | None = None,
     ) -> InitializationOptions:
         opts: Final = Server.create_initialization_options(
             self,
             notification_options=notification_options,
             experimental_capabilities=experimental_capabilities or {},
         )
-        updates: Final[dict[str, Any]] = {}
+        updates: Final[dict[str, str]] = {}
         merged: Final = _mcp_gateway_initialize_instructions.get()
         if merged is not None:
             updates["instructions"] = merged
@@ -549,6 +549,17 @@ if MCP_AVAILABLE:
     _stateful_session_locks: Final[dict[str, asyncio.Lock]] = {}
     _stateful_session_active_request_counts: Final[dict[str, int]] = {}
 
+    class _TerminableTransport(Protocol):
+        async def terminate(self) -> None: ...
+
+    class _TransportRegistry(Protocol):
+        def __contains__(self, session_id: object, /) -> bool: ...
+
+        def pop(self, session_id: str, default: None, /) -> "_TerminableTransport | None": ...
+
+    def _stateful_server_instances() -> _TransportRegistry:
+        return getattr(session_manager_stateful, "_server_instances", {})
+
     def _remove_stateful_session_tracking(session_id: str) -> None:
         _stateful_session_auth_contexts.pop(session_id, None)
         _stateful_session_auth_context_last_seen.pop(session_id, None)
@@ -578,8 +589,8 @@ if MCP_AVAILABLE:
     ) -> None:
         """Terminate expired stateful sessions and drop their auth contexts."""
         now = time.monotonic() if now is None else now
-        server_instances: Final = getattr(session_manager_stateful, "_server_instances", {})
-        expired_session_ids: Final = []
+        server_instances: Final = _stateful_server_instances()
+        expired_session_ids: Final[list[str]] = []
         for session_id, last_seen in _stateful_session_auth_context_last_seen.items():
             if _stateful_session_active_request_counts.get(session_id, 0) > 0:
                 continue
@@ -619,7 +630,7 @@ if MCP_AVAILABLE:
         session may proceed, or ``False`` when the caller is already at the cap
         with every session in flight (the new ``initialize`` should be rejected).
         """
-        server_instances: Final = getattr(session_manager_stateful, "_server_instances", {})
+        server_instances: Final = _stateful_server_instances()
 
         def _owned_live_session_ids() -> list[str]:
             return [
@@ -778,7 +789,7 @@ if MCP_AVAILABLE:
                     get_virtual_tool_definitions,
                 )
 
-                return [Tool(**d) for d in get_virtual_tool_definitions()]
+                return [Tool.model_validate(d) for d in get_virtual_tool_definitions()]
 
             # Get mcp_servers from context variable
             verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools")
@@ -847,7 +858,7 @@ if MCP_AVAILABLE:
 
     async def _build_virtual_call_logging_obj(
         name: str,
-        arguments: dict[str, Any],
+        arguments: dict[str, object],
         user_api_key_auth: UserAPIKeyAuth,
     ) -> LiteLLMLoggingObj | None:
         """Run the pre-call pipeline (guardrails + logging setup) for a virtual
@@ -885,7 +896,7 @@ if MCP_AVAILABLE:
 
     async def _dispatch_virtual_mcp_tool(
         name: str,
-        arguments: dict[str, Any] | None,
+        arguments: dict[str, object] | None,
         user_api_key_auth: UserAPIKeyAuth | None,
         client_ip: str | None,
         mcp_servers: list[str] | None = None,
@@ -957,7 +968,7 @@ if MCP_AVAILABLE:
         )
 
     @server.call_tool()
-    async def mcp_server_tool_call(name: str, arguments: dict[str, Any] | None) -> CallToolResult:
+    async def mcp_server_tool_call(name: str, arguments: dict[str, object] | None) -> CallToolResult:
         """
         Call a specific tool with the provided arguments
         Args:
@@ -1621,7 +1632,7 @@ if MCP_AVAILABLE:
     async def _get_user_oauth_extra_headers_from_db(
         server: MCPServer,
         user_api_key_auth: UserAPIKeyAuth | None,
-        prefetched_creds: dict[str, dict[str, Any]] | None = None,
+        prefetched_creds: 'Mapping[str, "OAuthCredentialPayload"] | None' = None,
     ) -> dict[str, str] | None:
         """Stored OAuth2 token for (user, server) as an ``Authorization: Bearer`` header, or None.
 
@@ -1646,7 +1657,7 @@ if MCP_AVAILABLE:
 
         Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops.
         """
-        user_id: Final = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None
+        user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None
         if not user_id:
             return {}
         try:
@@ -1871,7 +1882,7 @@ if MCP_AVAILABLE:
 
         list_tools_start_time: Final = datetime.now()
         litellm_logging_obj: LiteLLMLoggingObj | None = None
-        list_tools_request_data: dict[str, Any] = {}
+        list_tools_request_data: dict[str, object] = {}
 
         if log_list_tools_to_spendlogs:
             # This is intentionally minimal: only async_success_handler / post_call_failure_hook
@@ -1879,7 +1890,7 @@ if MCP_AVAILABLE:
             list_tools_call_id: Final = str(uuid.uuid4())
             # Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool)
             effective_litellm_trace_id: Final = litellm_trace_id or get_chain_id_from_headers(raw_headers)
-            spend_logs_metadata: Final[dict[str, Any]] = {
+            spend_logs_metadata: Final[dict[str, object]] = {
                 "mcp_operation": "list_tools",
             }
             if isinstance(list_tools_log_source, str):
@@ -2615,7 +2626,7 @@ if MCP_AVAILABLE:
 
     async def execute_mcp_tool(
         name: str,
-        arguments: dict[str, Any],
+        arguments: dict[str, object],
         allowed_mcp_servers: list[MCPServer],
         start_time: datetime,
         user_api_key_auth: UserAPIKeyAuth | None = None,
@@ -2882,7 +2893,7 @@ if MCP_AVAILABLE:
                 _request_auth_header.reset(_auth_token)
                 _request_extra_headers.reset(_extra_token)
                 _request_resolved_auth_headers.reset(_resolved_token)
-            response = CallToolResult(content=cast(Any, local_content), isError=False)
+            response = CallToolResult(content=local_content, isError=False)
 
         # Try managed MCP server tool (the name is bare; the prefix boundary was
         # already resolved above against this server's registered prefixes)
@@ -2956,7 +2967,7 @@ if MCP_AVAILABLE:
                     arguments = hook_result["arguments"]  # pyright: ignore[reportAny]  # hook returns untyped args
 
             local_content = await _handle_local_mcp_tool(original_tool_name, arguments)
-            response = CallToolResult(content=cast(Any, local_content), isError=False)
+            response = CallToolResult(content=local_content, isError=False)
 
         return await _run_post_mcp_call_guardrails(
             result=response,
@@ -3003,7 +3014,7 @@ if MCP_AVAILABLE:
 
     async def _fire_mcp_tool_call_logging(
         logging_obj: LiteLLMLoggingObj,
-        result: Any,
+        result: CallToolResult,
         start_time: datetime,
         end_time: datetime,
         user_api_key_auth: UserAPIKeyAuth | None = None,
@@ -3070,7 +3081,7 @@ if MCP_AVAILABLE:
     @client
     async def call_mcp_tool(
         name: str,
-        arguments: dict[str, Any] | None = None,
+        arguments: dict[str, object] | None = None,
         user_api_key_auth: UserAPIKeyAuth | None = None,
         mcp_auth_header: str | None = None,
         mcp_servers: list[str] | None = None,
@@ -3161,7 +3172,7 @@ if MCP_AVAILABLE:
 
     async def mcp_get_prompt(
         name: str,
-        arguments: dict[str, Any] | None = None,
+        arguments: dict[str, object] | None = None,
         user_api_key_auth: UserAPIKeyAuth | None = None,
         mcp_auth_header: str | None = None,
         mcp_servers: list[str] | None = None,
@@ -3262,7 +3273,7 @@ if MCP_AVAILABLE:
 
     def _get_standard_logging_mcp_tool_call(
         name: str,
-        arguments: dict[str, Any],
+        arguments: dict[str, object],
         server_name: str | None,
         session_id: str | None = None,
     ) -> StandardLoggingMCPToolCall:
@@ -3291,13 +3302,13 @@ if MCP_AVAILABLE:
     async def _handle_managed_mcp_tool(
         server_name: str,
         name: str,
-        arguments: dict[str, Any],
+        arguments: dict[str, object],
         user_api_key_auth: UserAPIKeyAuth | None = None,
         mcp_auth_header: str | None = None,
         mcp_server_auth_headers: dict[str, dict[str, str]] | None = None,
         oauth2_headers: dict[str, str] | None = None,
         raw_headers: dict[str, str] | None = None,
-        litellm_logging_obj: Any | None = None,
+        litellm_logging_obj: LiteLLMLoggingObj | None = None,
         host_progress_callback: Callable | None = None,
     ) -> CallToolResult:
         """Handle tool execution for managed server tools"""
@@ -3320,7 +3331,7 @@ if MCP_AVAILABLE:
         return call_tool_result
 
     async def _handle_local_mcp_tool(
-        name: str, arguments: dict[str, Any]
+        name: str, arguments: dict[str, object]
     ) -> list[TextContent | ImageContent | EmbeddedResource]:
         """
         Handle tool execution for local registry tools
@@ -3426,7 +3437,8 @@ if MCP_AVAILABLE:
         Extract mcp-session-id from ASGI scope headers.
         Returns None if not present.
         """
-        for header_name, header_value in scope.get("headers", []):
+        scope_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", [])
+        for header_name, header_value in scope_headers:
             name = header_name if isinstance(header_name, bytes) else header_name.encode()
             if name.lower() == b"mcp-session-id":
                 return header_value.decode() if isinstance(header_value, bytes) else str(header_value)
@@ -3528,7 +3540,7 @@ if MCP_AVAILABLE:
             if message.get("type") != "http.request":
                 break
 
-            body = message.get("body", b"") or b""
+            body: bytes = message.get("body", b"") or b""
             if body:
                 # Only retain up to the remaining peek budget for sniffing.
                 # The full ``message`` is already in memory (delivered by
@@ -3571,9 +3583,9 @@ if MCP_AVAILABLE:
         Fixes https://github.com/BerriAI/litellm/issues/20992
         """
         _mcp_session_header: Final = b"mcp-session-id"
-        _headers: Final = scope.get("headers", [])
+        _headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", [])
 
-        def _normalize_header_name(header_name: Any) -> bytes | None:
+        def _normalize_header_name(header_name: object) -> bytes | None:
             if isinstance(header_name, bytes):
                 return header_name.lower()
             if isinstance(header_name, str):
@@ -3902,7 +3914,8 @@ if MCP_AVAILABLE:
 
     def _get_authorization_header_from_scope(scope: Scope) -> str | None:
         """First ``Authorization`` header value in the ASGI scope, or None."""
-        for key, value in scope.get("headers", []):
+        scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", [])
+        for key, value in scope_headers:
             if key.lower() == b"authorization":
                 return value.decode("latin-1")
         return None
@@ -3921,7 +3934,8 @@ if MCP_AVAILABLE:
         ``MCPRequestHandler.process_mcp_request``), and forwarding it upstream
         would leak the proxy key to a third-party MCP server.
         """
-        has_litellm_key_header: Final = any(key.lower() == b"x-litellm-api-key" for key, _ in scope.get("headers", []))
+        scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", [])
+        has_litellm_key_header: Final = any(key.lower() == b"x-litellm-api-key" for key, _ in scope_headers)
         if not has_litellm_key_header:
             return None
         return _get_authorization_header_from_scope(scope)
@@ -4115,7 +4129,7 @@ if MCP_AVAILABLE:
     async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None:
         """Handle MCP requests through StreamableHTTP."""
         try:
-            path: Final = scope.get("path", "")
+            path: Final[str] = scope.get("path", "")
             (
                 user_api_key_auth,
                 mcp_auth_header,
@@ -4135,7 +4149,8 @@ if MCP_AVAILABLE:
             )
 
             # Strip any client-supplied x-mcp-toolset-id to prevent forgery.
-            scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"]
+            scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", [])
+            scope["headers"] = [(k, v) for k, v in scope_headers if k.lower() != b"x-mcp-toolset-id"]
 
             # Apply toolset scope if set server-side via ContextVar (set by
             # /toolset/{name}/mcp and /{name}/mcp route handlers in proxy_server.py).
@@ -4436,7 +4451,7 @@ if MCP_AVAILABLE:
     async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None:
         """Handle MCP requests through SSE."""
         try:
-            path: Final = scope.get("path", "")
+            path: Final[str] = scope.get("path", "")
             (
                 user_api_key_auth,
                 mcp_auth_header,
@@ -4456,7 +4471,8 @@ if MCP_AVAILABLE:
             )
 
             # Strip any client-supplied x-mcp-toolset-id to prevent forgery.
-            scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"]
+            scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", [])
+            scope["headers"] = [(k, v) for k, v in scope_headers if k.lower() != b"x-mcp-toolset-id"]
 
             # Apply toolset scope if set server-side via ContextVar so the
             # downstream probe list matches the fully-authorized server set
@@ -4680,7 +4696,8 @@ if MCP_AVAILABLE:
     ) -> Send:
         async def wrapped_send(message: Message) -> None:
             if message.get("type") == "http.response.start":
-                for key, value in message.get("headers", []):
+                response_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = message.get("headers", [])
+                for key, value in response_headers:
                     header_name = key if isinstance(key, bytes) else str(key).encode()
                     if header_name.lower() == b"mcp-session-id":
                         session_id = value.decode() if isinstance(value, bytes) else str(value)

litellm/proxy/batches_endpoints/endpoints.py

@@ -36,6 +36,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
     get_original_file_id,
     prepare_data_with_credentials,
     update_batch_in_database,
+    validate_managed_id_requirement,
 )
 from litellm.proxy.utils import handle_exception_on_proxy, is_known_model
 from litellm.repositories.table_repositories import ManagedFileRepository
@@ -176,6 +177,12 @@ async def create_batch(
             }
 
         input_file_id: Final = _create_batch_data.get("input_file_id", None)
+        await validate_managed_id_requirement(
+            resource_id=input_file_id,
+            resource_kind="file",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
         unified_file_id: str | Literal[False] = False
 
         model_from_file_id = None
@@ -392,6 +399,12 @@ async def retrieve_batch(
 
     data: dict = {}
     try:
+        await validate_managed_id_requirement(
+            resource_id=batch_id,
+            resource_kind="batch",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
         model_from_id: Final = decode_model_from_file_id(batch_id)
         _retrieve_batch_request: Final = RetrieveBatchRequest(
             batch_id=batch_id,
@@ -840,6 +853,13 @@ async def cancel_batch(
 
     data: dict = {}
     try:
+        await validate_managed_id_requirement(
+            resource_id=batch_id,
+            resource_kind="batch",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
+
         # Check for encoded batch ID with model info
         model_from_id: Final = decode_model_from_file_id(batch_id)
 

litellm/proxy/fine_tuning_endpoints/endpoints.py

@@ -17,6 +17,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
 from litellm.proxy.openai_files_endpoints.common_utils import (
     _is_base64_encoded_unified_file_id,
+    validate_managed_id_requirement,
 )
 from litellm.proxy.utils import handle_exception_on_proxy
 from litellm.types.utils import LiteLLMFineTuningJob
@@ -134,6 +135,18 @@ async def create_fine_tuning_job(
         ## CHECK IF MANAGED FILE ID
         unified_file_id: str | Literal[False] = False
         training_file: Final = fine_tuning_request.training_file
+        await validate_managed_id_requirement(
+            resource_id=training_file,
+            resource_kind="file",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
+        await validate_managed_id_requirement(
+            resource_id=fine_tuning_request.validation_file,
+            resource_kind="file",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
         response: LiteLLMFineTuningJob | None = None
         if training_file:
             unified_file_id = _is_base64_encoded_unified_file_id(training_file)
@@ -246,6 +259,12 @@ async def retrieve_fine_tuning_job(
     try:
         if premium_user is not True:
             raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}")
+        await validate_managed_id_requirement(
+            resource_id=fine_tuning_job_id,
+            resource_kind="fine-tuning job",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
         # Include original request and headers in the data
         base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data)
         (
@@ -513,6 +532,12 @@ async def cancel_fine_tuning_job(
     try:
         if premium_user is not True:
             raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}")
+        await validate_managed_id_requirement(
+            resource_id=fine_tuning_job_id,
+            resource_kind="fine-tuning job",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
         # Include original request and headers in the data
         base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data)
         (

litellm/proxy/guardrails/guardrail_endpoints.py

@@ -6,10 +6,10 @@ import concurrent.futures
 import inspect
 import json
 import os
-from collections.abc import Mapping, Sequence
+from collections.abc import Awaitable, Callable, Mapping, Sequence
 from datetime import datetime, timezone
 from types import UnionType
-from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, Union, cast, get_args, get_origin
+from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union, cast, get_args, get_origin
 from urllib.parse import urlparse
 
 from fastapi import APIRouter, Depends, HTTPException, Request
@@ -54,8 +54,8 @@ from litellm.types.guardrails import (
 if TYPE_CHECKING:
     from types import CodeType
 
-    from prisma.actions import LiteLLM_GuardrailsTableActions
     from prisma.models import LiteLLM_GuardrailsTable
+    from pydantic.fields import FieldInfo
 
     from litellm.proxy.utils import PrismaClient
 
@@ -65,24 +65,44 @@ router: Final = APIRouter()
 GUARDRAIL_REGISTRY: Final = GuardrailRegistry()
 
 
-def _guardrails_table(prisma_client: "PrismaClient") -> "LiteLLM_GuardrailsTableActions[LiteLLM_GuardrailsTable]":
-    table: Final[LiteLLM_GuardrailsTableActions[LiteLLM_GuardrailsTable]] = GuardrailsRepository(prisma_client).table
+class _GuardrailsTableActions(Protocol):
+    async def create(self, data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable": ...
+
+    async def delete(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ...
+
+    async def find_unique(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ...
+
+    async def find_many(
+        self, where: Mapping[str, object], order: Mapping[str, str]
+    ) -> "Sequence[LiteLLM_GuardrailsTable]": ...
+
+    async def update(
+        self, where: Mapping[str, object], data: Mapping[str, object]
+    ) -> "LiteLLM_GuardrailsTable | None": ...
+
+
+def _as_str_object_mapping(mapping: Mapping[str, object]) -> Mapping[str, object]:
+    return mapping
+
+
+def _guardrails_table(prisma_client: "PrismaClient") -> _GuardrailsTableActions:
+    table: Final[_GuardrailsTableActions] = GuardrailsRepository(prisma_client).table
     return table
 
 
 async def _create_guardrail_row(prisma_client: "PrismaClient", data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable":
-    row: Final[LiteLLM_GuardrailsTable] = await GuardrailsRepository(prisma_client).table.create(data=data)
+    row: Final = await _guardrails_table(prisma_client).create(data=data)
     return row
 
 
 async def _delete_guardrail_row(prisma_client: "PrismaClient", where: Mapping[str, object]) -> None:
-    await GuardrailsRepository(prisma_client).table.delete(where=where)
+    await _guardrails_table(prisma_client).delete(where=where)
 
 
 async def _find_team_guardrail_rows(
     prisma_client: "PrismaClient", where: Mapping[str, object]
 ) -> "Sequence[LiteLLM_GuardrailsTable]":
-    rows: Final[Sequence[LiteLLM_GuardrailsTable]] = await GuardrailsRepository(prisma_client).table.find_many(
+    rows: Final = await _guardrails_table(prisma_client).find_many(
         where=where,
         order={"created_at": "desc"},
     )
@@ -499,10 +519,12 @@ async def update_guardrail(
         if existing_guardrail is None:
             raise HTTPException(status_code=404, detail=f"Guardrail with ID {guardrail_id} not found")
 
-        result: Final = await GUARDRAIL_REGISTRY.update_guardrail_in_db(
-            guardrail_id=guardrail_id,
-            guardrail=request.guardrail,
-            prisma_client=prisma_client,
+        result: Final = _as_str_object_mapping(
+            await GUARDRAIL_REGISTRY.update_guardrail_in_db(
+                guardrail_id=guardrail_id,
+                guardrail=request.guardrail,
+                prisma_client=prisma_client,
+            )
         )
 
         guardrail_name: Final = result.get("guardrail_name", "Unknown")
@@ -613,7 +635,7 @@ class RegisterGuardrailRequest(BaseModel):
     """Request body for POST /guardrails/register. Follows Generic Guardrail API config."""
 
     guardrail_name: str
-    litellm_params: dict[str, Any]  # guardrail, mode, api_base required; api_key, headers, etc. optional
+    litellm_params: dict[str, object]  # guardrail, mode, api_base required; api_key, headers, etc. optional
     guardrail_info: dict[str, object] | None = None
     team_id: str | None = None
 
@@ -1172,12 +1194,14 @@ async def patch_guardrail(
         )
 
         # Update litellm_params if default_on is provided or pii_entities_config is provided
-        litellm_params = LitellmParams(**dict(existing_guardrail.get("litellm_params", {})))
+        existing_litellm_params: Final = _as_str_object_mapping(dict(existing_guardrail.get("litellm_params", {})))
+        litellm_params = LitellmParams(**existing_litellm_params)
         if request.litellm_params is not None:
             requested_litellm_params: Final = request.litellm_params.model_dump(exclude_unset=True)
             litellm_params_dict: Final = litellm_params.model_dump(exclude_unset=True)
             litellm_params_dict.update(requested_litellm_params)
-            litellm_params = LitellmParams(**litellm_params_dict)
+            merged_litellm_params: Final = _as_str_object_mapping(litellm_params_dict)
+            litellm_params = LitellmParams(**merged_litellm_params)
 
         # Update guardrail_info if provided
         guardrail_info: Final = (
@@ -1193,10 +1217,12 @@ async def patch_guardrail(
             litellm_params=litellm_params,
             guardrail_info=guardrail_info,
         )
-        result: Final = await GUARDRAIL_REGISTRY.update_guardrail_in_db(
-            guardrail_id=guardrail_id,
-            guardrail=guardrail,
-            prisma_client=prisma_client,
+        result: Final = _as_str_object_mapping(
+            await GUARDRAIL_REGISTRY.update_guardrail_in_db(
+                guardrail_id=guardrail_id,
+                guardrail=guardrail,
+                prisma_client=prisma_client,
+            )
         )
 
         guardrail_name = result.get("guardrail_name", "Unknown")
@@ -1552,31 +1578,46 @@ async def validate_blocked_words_file(request: dict[str, str]):
         return {"valid": False, "error": f"Validation error: {e}"}
 
 
-def _get_field_type_from_annotation(field_annotation: Any) -> str:
+def _dunder_origin(annotation: object) -> object:
+    origin: Final[object] = getattr(annotation, "__origin__", None)
+    return origin
+
+
+def _dunder_name(annotation: object) -> object:
+    name: Final[object] = getattr(annotation, "__name__", None)
+    return name
+
+
+def _dunder_args(annotation: object) -> tuple[object, ...]:
+    args: Final[tuple[object, ...]] = getattr(annotation, "__args__", ())
+    return args
+
+
+def _get_field_type_from_annotation(field_annotation: object) -> str:
     """
     Convert a Python type annotation to a UI-friendly type string
     """
     # Handle Union types (like Optional[T])
     if get_origin(field_annotation) is Union or get_origin(field_annotation) is UnionType:
         # For Optional[T], get the non-None type
-        args: Final = get_args(field_annotation)
+        args: Final[tuple[object, ...]] = get_args(field_annotation)
         non_none_args: Final = [arg for arg in args if arg is not type(None)]
         if non_none_args:
             field_annotation = non_none_args[0]
 
     # Handle List types
-    if hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is list:
+    if hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is list:
         return "array"
 
     # Handle Dict types
-    if hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is dict:
+    if hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is dict:
         return "dict"
 
     # Handle Literal types
     if hasattr(field_annotation, "__origin__") and hasattr(field_annotation, "__args__"):
         # Check for Literal types (Python 3.8+)
-        origin: Final = field_annotation.__origin__
-        if hasattr(origin, "__name__") and origin.__name__ == "Literal":
+        origin: Final = _dunder_origin(field_annotation)
+        if hasattr(origin, "__name__") and _dunder_name(origin) == "Literal":
             return "select"  # For dropdown/select inputs
 
     # Handle basic types
@@ -1595,66 +1636,66 @@ def _get_field_type_from_annotation(field_annotation: Any) -> str:
     return "string"
 
 
-def _extract_literal_values(annotation: Any) -> list[str]:
+def _extract_literal_values(annotation: object) -> Sequence[object]:
     """
     Extract literal values from a Literal type annotation
     """
     if hasattr(annotation, "__origin__") and hasattr(annotation, "__args__"):
-        origin: Final = annotation.__origin__
-        if hasattr(origin, "__name__") and origin.__name__ == "Literal":
-            return list(annotation.__args__)
+        origin: Final = _dunder_origin(annotation)
+        if hasattr(origin, "__name__") and _dunder_name(origin) == "Literal":
+            return list(_dunder_args(annotation))
     return []
 
 
-def _get_dict_key_options(field_annotation: Any) -> list[str] | None:
+def _get_dict_key_options(field_annotation: object) -> Sequence[object] | None:
     """
     Extract key options from Dict[Literal[...], T] types
     """
     if (
         hasattr(field_annotation, "__origin__")
-        and field_annotation.__origin__ is dict
+        and _dunder_origin(field_annotation) is dict
         and hasattr(field_annotation, "__args__")
     ):
-        args: Final = field_annotation.__args__
+        args: Final = _dunder_args(field_annotation)
         if len(args) >= 2:
             key_type: Final = args[0]
             return _extract_literal_values(key_type)
     return None
 
 
-def _get_dict_value_type(field_annotation: Any) -> str:
+def _get_dict_value_type(field_annotation: object) -> str:
     """
     Get the value type from Dict[K, V] types
     """
     if (
         hasattr(field_annotation, "__origin__")
-        and field_annotation.__origin__ is dict
+        and _dunder_origin(field_annotation) is dict
         and hasattr(field_annotation, "__args__")
     ):
-        args: Final = field_annotation.__args__
+        args: Final = _dunder_args(field_annotation)
         if len(args) >= 2:
             value_type: Final = args[1]
             return _get_field_type_from_annotation(value_type)
     return "string"
 
 
-def _get_list_element_options(field_annotation: Any) -> list[str] | None:
+def _get_list_element_options(field_annotation: object) -> Sequence[object] | None:
     """
     Extract element options from List[Literal[...]] types
     """
     if (
         hasattr(field_annotation, "__origin__")
-        and field_annotation.__origin__ is list
+        and _dunder_origin(field_annotation) is list
         and hasattr(field_annotation, "__args__")
     ):
-        args: Final = field_annotation.__args__
+        args: Final = _dunder_args(field_annotation)
         if len(args) >= 1:
             element_type: Final = args[0]
             return _extract_literal_values(element_type)
     return None
 
 
-def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool:
+def _should_skip_optional_params(field_name: str, field_annotation: object) -> bool:
     """Check if optional_params field should be skipped (not meaningfully overridden)."""
     if field_name != "optional_params":
         return False
@@ -1664,12 +1705,12 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool
 
     # Check if the annotation is still a generic TypeVar (not specialized)
     if isinstance(field_annotation, TypeVar) or (
-        hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is TypeVar
+        hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is TypeVar
     ):
         return True
 
     # Also skip if it's a generic type that wasn't specialized
-    if hasattr(field_annotation, "__name__") and field_annotation.__name__ in (
+    if hasattr(field_annotation, "__name__") and _dunder_name(field_annotation) in (
         "T",
         "TypeVar",
     ):
@@ -1677,18 +1718,18 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool
 
     # Handle Optional[T] where T is still a TypeVar
     if hasattr(field_annotation, "__args__"):
-        non_none_args: Final = [arg for arg in field_annotation.__args__ if arg is not type(None)]
+        non_none_args: Final = [arg for arg in _dunder_args(field_annotation) if arg is not type(None)]
         if non_none_args and isinstance(non_none_args[0], TypeVar):
             return True
 
     return False
 
 
-def _unwrap_optional_type(field_annotation: Any) -> Any:
+def _unwrap_optional_type(field_annotation: object) -> object:
     """Unwrap Optional types to get the actual type."""
     if get_origin(field_annotation) is Union or get_origin(field_annotation) is UnionType:
         # For Optional[BaseModel], get the non-None type
-        args: Final = get_args(field_annotation)
+        args: Final[tuple[object, ...]] = get_args(field_annotation)
         non_none_args: Final = [arg for arg in args if arg is not type(None)]
         if non_none_args:
             return non_none_args[0]
@@ -1696,20 +1737,20 @@ def _unwrap_optional_type(field_annotation: Any) -> Any:
 
 
 def _build_field_dict(
-    field: Any,
-    field_annotation: Any,
+    field: "FieldInfo",
+    field_annotation: object,
     description: str,
     required: bool,
-) -> dict[str, Any]:
+) -> dict[str, object]:
     """Build field dictionary for non-nested fields."""
     # Determine the field type from annotation
     field_type = _get_field_type_from_annotation(field_annotation)
 
     # Check for custom UI type override
-    field_json_schema_extra: Final = getattr(field, "json_schema_extra", {})
+    field_json_schema_extra: Final[Mapping[str, object]] = getattr(field, "json_schema_extra", {})
     if field_json_schema_extra and "ui_type" in field_json_schema_extra:
         ui_type: Final = field_json_schema_extra["ui_type"]
-        field_type = ui_type.value if hasattr(ui_type, "value") else ui_type
+        field_type = getattr(ui_type, "value", ui_type)
     elif field_json_schema_extra and "type" in field_json_schema_extra:
         field_type = field_json_schema_extra["type"]
 
@@ -1748,8 +1789,9 @@ def _build_field_dict(
             field_dict["options"] = literal_options
 
     # Add default value if it exists
-    if field.default is not None and field.default is not ...:
-        field_dict["default_value"] = field.default
+    field_default: Final[object] = getattr(field, "default", None)
+    if field_default is not None and field_default is not ...:
+        field_dict["default_value"] = field_default
 
     # Copy min, max, step from json_schema_extra for number/percentage inputs
     if field_json_schema_extra:
@@ -1763,7 +1805,7 @@ def _build_field_dict(
 def _extract_fields_recursive(
     model: type[BaseModel],
     depth: int = 0,
-) -> dict[str, Any]:
+) -> dict[str, object]:
     # Check if we've exceeded the maximum recursion depth
     if depth > DEFAULT_MAX_RECURSE_DEPTH:
         raise HTTPException(
@@ -1817,7 +1859,7 @@ def _extract_fields_recursive(
     return fields
 
 
-def _get_fields_from_model(model_class: type[BaseModel]) -> dict[str, Any]:
+def _get_fields_from_model(model_class: type[BaseModel]) -> dict[str, object]:
     """
     Get the fields from a Pydantic model as a nested dictionary structure
     """
@@ -2141,7 +2183,26 @@ def _resolve_guardrail_input_type(active_guardrail: CustomGuardrail, input_type:
     return "response" if input_type == "response" else "request"
 
 
-def _patch_logging_obj_for_guardrail(litellm_logging_obj: Any, request: ApplyGuardrailRequest) -> None:
+class _GuardrailLoggingObj(Protocol):
+    call_type: str
+    model_call_details: dict[str, object]
+
+    @property
+    def update_messages(self) -> "Callable[..., object]": ...
+
+    @property
+    def async_success_handler(self) -> "Callable[..., Awaitable[object]]": ...
+
+    @property
+    def success_handler(self) -> "Callable[..., object]": ...
+
+
+class _GuardrailProxyLogging(Protocol):
+    @property
+    def post_call_success_hook(self) -> "Callable[..., Awaitable[object]]": ...
+
+
+def _patch_logging_obj_for_guardrail(litellm_logging_obj: _GuardrailLoggingObj, request: ApplyGuardrailRequest) -> None:
     """Configure the logging object so Langfuse/OTEL extract input and output correctly."""
     litellm_logging_obj.call_type = "pass_through_endpoint"
     litellm_logging_obj.model_call_details["call_type"] = "pass_through_endpoint"
@@ -2151,8 +2212,8 @@ def _patch_logging_obj_for_guardrail(litellm_logging_obj: Any, request: ApplyGua
 
 
 async def _emit_guardrail_success_logs(
-    proxy_logging_obj: Any,
-    litellm_logging_obj: Any,
+    proxy_logging_obj: _GuardrailProxyLogging,
+    litellm_logging_obj: _GuardrailLoggingObj | None,
     data: dict,
     user_api_key_dict: UserAPIKeyAuth,
     response: ApplyGuardrailResponse,

litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py

@@ -19,7 +19,7 @@ request is sent with the ``X-Cisco-AI-Defense-API-Key`` header.
 
 import json
 import os
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Mapping, Sequence
 from dataclasses import dataclass, replace
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, Literal
@@ -94,13 +94,13 @@ class _CiscoVerdict:
     is_safe: bool | None
     classifications: list[str]
     severity: str | None
-    rules: list[dict[str, Any]]
+    rules: list[dict[str, object]]
     explanation: str | None
     event_id: str | None
     action: str | None = None
     sanitized_text: str | None = None
-    sanitized_messages: list[dict[str, Any]] | None = None
-    sanitized_mcp_arguments: dict[str, Any] | None = None
+    sanitized_messages: list[dict[str, object]] | None = None
+    sanitized_mcp_arguments: dict[str, object] | None = None
 
 
 class CiscoAIDefenseGuardrailMissingSecrets(Exception):
@@ -136,7 +136,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         api_base: str | None = None,
         inspection_type: str | None = None,
         inspect_path: str | None = None,
-        enabled_rules: list[dict[str, Any]] | None = None,
+        enabled_rules: Sequence[object] | None = None,
         integration_profile_id: str | None = None,
         integration_profile_version: str | None = None,
         integration_tenant_id: str | None = None,
@@ -415,7 +415,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
     async def async_post_call_streaming_iterator_hook(
         self,
         user_api_key_dict: UserAPIKeyAuth,
-        response: AsyncIterator[Any],
+        response: AsyncIterator[object],
         request_data: dict,
     ):
         """Buffer and inspect streaming chat output before delivery."""
@@ -437,7 +437,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
             self.guardrail_name,
         )
 
-        all_chunks: Final[list[Any]] = []
+        all_chunks: Final[list[object]] = []
         try:
             async for chunk in response:
                 all_chunks.append(chunk)
@@ -497,7 +497,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
                 response_obj=assembled,
             )
         except HTTPException as exc:
-            error_obj: dict[str, Any] = self._http_exception_to_error_obj(exc)
+            error_obj: dict[str, object] = self._http_exception_to_error_obj(exc)
             verbose_proxy_logger.warning(
                 "Cisco AI Defense guardrail (%s): streaming response "
                 "blocked — emitting SSE error event instead of "
@@ -531,7 +531,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
             for chunk in all_chunks:
                 yield chunk
 
-    def _build_block_payload(self, context: _ScanContext, verdict: _CiscoVerdict) -> dict[str, Any]:
+    def _build_block_payload(self, context: _ScanContext, verdict: _CiscoVerdict) -> dict[str, object]:
         """Canonical block payload used across all four block paths.
 
         Same dict is the ``HTTPException.detail`` for chat / MCP request
@@ -555,34 +555,34 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
             "event_id": verdict.event_id,
         }
 
-    def _http_exception_to_error_obj(self, exc: HTTPException) -> dict[str, Any]:
+    def _http_exception_to_error_obj(self, exc: HTTPException) -> dict[str, object]:
         """Wrap an ``HTTPException`` detail into the SSE ``error`` payload.
 
         For Cisco's own blocks the detail is already the canonical block
         payload, so this is a near-passthrough that just adds ``code``
         / ``guardrail`` defaults for non-Cisco / unstructured details.
         """
-        error_obj: dict[str, Any] = dict(exc.detail) if isinstance(exc.detail, dict) else {"message": str(exc.detail)}
+        error_obj: dict[str, object] = {**exc.detail} if isinstance(exc.detail, dict) else {"message": str(exc.detail)}
         error_obj.setdefault("message", error_obj.get("error", "Guardrail block"))
         error_obj.setdefault("code", exc.status_code)
         error_obj.setdefault("guardrail", self.guardrail_name)
         return error_obj
 
     @classmethod
-    def _streaming_content_was_modified(cls, original_chunks: list[Any], assembled: ModelResponse) -> bool:
+    def _streaming_content_was_modified(cls, original_chunks: Sequence[object], assembled: ModelResponse) -> bool:
         """Decide whether redact changed content or tool/function arguments."""
         original_text: Final = cls._extract_streaming_chunk_scan_text(original_chunks)
         assembled_text: Final = " ".join(m.get("content", "") for m in cls._extract_response_messages(assembled))
         return original_text != assembled_text
 
     @classmethod
-    def _extract_streaming_chunk_scan_text(cls, chunks: list[Any]) -> str:
+    def _extract_streaming_chunk_scan_text(cls, chunks: Sequence[object]) -> str:
         original_text = ""
         argument_text = ""
         for chunk in chunks:
             choices = getattr(chunk, "choices", None) or []
             for c in choices:
-                delta = getattr(c, "delta", None)
+                delta: object | None = getattr(c, "delta", None)
                 if delta is None:
                     continue
                 text = getattr(delta, "content", None)
@@ -595,7 +595,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
                     args = cls._extract_tool_call_arguments(tc)
                     if args:
                         argument_text += args
-                fc = getattr(delta, "function_call", None)
+                fc: object | None = getattr(delta, "function_call", None)
                 if fc is not None:
                     args = cls._extract_function_call_arguments(fc)
                     if args:
@@ -673,7 +673,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         allow, WARNING for intervened/redacted, ERROR is left for
         upstream API failures.
         """
-        fields: Final[dict[str, Any]] = {
+        fields: Final[dict[str, object]] = {
             "guardrail": self.guardrail_name,
             "surface": context.surface,
             "direction": context.direction,
@@ -752,7 +752,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         user_api_key_dict: UserAPIKeyAuth,
         direction: str = "input",
         response_obj: object = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         url: Final = f"{self.api_base}{self.inspect_path}"
         payload: Final = self._build_chat_payload(messages, request_data, user_api_key_dict)
         start_time: Final = datetime.now()
@@ -784,7 +784,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         messages: list[dict[str, str]],
         request_data: dict,
         user_api_key_dict: UserAPIKeyAuth,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         return {
             "messages": messages,
             "metadata": self._build_metadata(request_data, user_api_key_dict),
@@ -798,9 +798,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
     async def _post_inspection(
         self,
         url: str,
-        payload: dict[str, Any],
+        payload: dict[str, object],
         surface: str,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         headers: Final = self._build_headers()
         verbose_proxy_logger.debug(
             "Cisco AI Defense guardrail: posting %s inspection to %s",
@@ -856,8 +856,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         self,
         request_data: dict,
         user_api_key_dict: UserAPIKeyAuth,
-    ) -> dict[str, Any]:
-        metadata: Final[dict[str, Any]] = {}
+    ) -> dict[str, object]:
+        metadata: Final[dict[str, object]] = {}
 
         user: Final = request_data.get("user") or getattr(user_api_key_dict, "user_id", None)
         if user:
@@ -884,8 +884,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
 
         return metadata
 
-    def _build_config(self) -> dict[str, Any]:
-        config: Final[dict[str, Any]] = {}
+    def _build_config(self) -> dict[str, object]:
+        config: Final[dict[str, object]] = {}
         if self.enabled_rules:
             config["enabled_rules"] = self.enabled_rules
         if self.integration_profile_id:
@@ -899,7 +899,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         return config
 
     @staticmethod
-    def _normalize_rule(rule: object) -> dict[str, Any]:
+    def _normalize_rule(rule: object) -> dict[str, object]:
         """Coerce a user-supplied rule into the wire-shape dict Cisco expects.
 
         Accepts ``str``, ``dict``, and Pydantic model inputs.
@@ -922,7 +922,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
                     rule = dumped
 
         if isinstance(rule, dict):
-            normalized: Final[dict[str, Any]] = {}
+            normalized: Final[dict[str, object]] = {}
             rule_name: Final = rule.get("rule_name")
             if rule_name:
                 normalized["rule_name"] = rule_name
@@ -950,7 +950,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         context: _ScanContext,
         start_time: datetime,
         response_obj: object = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """Parse, log, and (optionally) raise/redact on the Cisco verdict.
 
         ``context.direction`` is ``"input"`` for request scans and ``"output"``
@@ -1119,10 +1119,10 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
     @classmethod
     def _sanitize_response_for_logging(
         cls,
-        inspect_response: dict[str, Any],
+        inspect_response: Mapping[str, object],
         surface: str,
         action: str | None = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """Drop bulky / privacy-sensitive fields, recursing into nested dicts.
 
         MCP verdicts are commonly nested under ``result``, so a
@@ -1138,9 +1138,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         return sanitized
 
     @classmethod
-    def _strip_sensitive_keys(cls, d: dict[str, Any]) -> dict[str, Any]:
+    def _strip_sensitive_keys(cls, d: Mapping[str, object]) -> dict[str, object]:
         """Recursively strip privacy-sensitive keys from a verdict dict."""
-        out: Final[dict[str, Any]] = {}
+        out: Final[dict[str, object]] = {}
         for key, value in d.items():
             if key.startswith("_") or key in cls._REDACTED_LOG_KEYS:
                 continue
@@ -1222,8 +1222,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
 
     @staticmethod
     def _extract_jsonrpc_error(
-        inspect_response: dict[str, Any],
-    ) -> dict[str, Any] | None:
+        inspect_response: Mapping[str, object],
+    ) -> dict[str, object] | None:
         """Detect a JSON-RPC error envelope inside an HTTP 200 response.
 
         The Cisco Inspect API can return ``{"error": {...}}`` (or nest one
@@ -1270,7 +1270,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
 
     @staticmethod
     def _extract_sanitized_text(
-        inspect_response: dict[str, Any],
+        inspect_response: Mapping[str, object],
     ) -> str | None:
         """Pull ``sanitized_text`` (or camelCase variant) off the verdict."""
         for key in ("sanitized_text", "sanitizedText"):
@@ -1287,8 +1287,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
 
     @staticmethod
     def _extract_sanitized_messages(
-        inspect_response: dict[str, Any],
-    ) -> list[dict[str, Any]] | None:
+        inspect_response: Mapping[str, object],
+    ) -> list[dict[str, object]] | None:
         """Pull a sanitized OpenAI-format messages array off the verdict.
 
         Cisco can return the rewrite under several keys; we accept any of
@@ -1354,7 +1354,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
     def _redact_mcp_input(
         request_data: dict,
         sanitized_text: str | None,
-        sanitized_mcp_arguments: dict[str, Any] | None,
+        sanitized_mcp_arguments: dict[str, object] | None,
     ) -> bool:
         """Rewrite MCP request arguments in all locations the proxy reads."""
         if sanitized_mcp_arguments is not None:
@@ -1388,7 +1388,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         self,
         request_data: dict,
         sanitized_text: str | None,
-        sanitized_messages: list[dict[str, Any]] | None,
+        sanitized_messages: list[dict[str, object]] | None,
     ) -> bool:
         """Rewrite chat request input (``messages`` or ``input``)."""
         if sanitized_messages and self._extract_tool_definition_text(request_data):
@@ -1444,7 +1444,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         cls,
         request_data: dict,
         sanitized_text: str | None,
-        sanitized_messages: list[dict[str, Any]] | None,
+        sanitized_messages: list[dict[str, object]] | None,
     ) -> bool:
         if sanitized_messages:
             instruction_text: Final = cls._instruction_text_from_messages(sanitized_messages)
@@ -1457,7 +1457,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         return False
 
     @classmethod
-    def _instruction_text_from_messages(cls, messages: list[dict[str, Any]]) -> str | None:
+    def _instruction_text_from_messages(cls, messages: list[dict[str, object]]) -> str | None:
         for message in messages:
             if not isinstance(message, dict):
                 continue
@@ -1468,7 +1468,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         return None
 
     @classmethod
-    def _non_instruction_messages(cls, messages: list[dict[str, Any]] | None) -> list[dict[str, Any]] | None:
+    def _non_instruction_messages(cls, messages: list[dict[str, object]] | None) -> list[dict[str, object]] | None:
         if messages is None:
             return None
         return [
@@ -1499,7 +1499,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         self,
         response_obj: object,
         sanitized_text: str | None,
-        sanitized_messages: list[dict[str, Any]] | None,
+        sanitized_messages: list[dict[str, object]] | None,
     ) -> bool:
         """Rewrite chat response (``ModelResponse`` or ``ResponsesAPIResponse``)."""
         if response_obj is None:
@@ -1526,7 +1526,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
     def _redact_model_response_choices(
         choices: list,
         sanitized_text: str | None,
-        sanitized_messages: list[dict[str, Any]] | None,
+        sanitized_messages: list[dict[str, object]] | None,
     ) -> bool:
         """Redact every returned choice, including tool-call/reasoning fields."""
         if sanitized_messages:
@@ -1570,7 +1570,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
     def _redact_text_completion_choices(
         choices: list,
         sanitized_text: str | None,
-        sanitized_messages: list[dict[str, Any]] | None,
+        sanitized_messages: list[dict[str, object]] | None,
     ) -> bool:
         """Rewrite ``/v1/completions`` text choices after Cisco redaction."""
         replacement = sanitized_text
@@ -1638,7 +1638,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         self,
         output_items: list,
         sanitized_text: str | None,
-        sanitized_messages: list[dict[str, Any]] | None,
+        sanitized_messages: list[dict[str, object]] | None,
     ) -> bool:
         replacement_text: str | None = sanitized_text
         if not replacement_text and sanitized_messages:
@@ -1672,14 +1672,14 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
 
     @staticmethod
     def _sanitized_messages_to_responses_input(
-        sanitized_messages: list[dict[str, Any]],
-    ) -> list[dict[str, Any]] | None:
+        sanitized_messages: list[dict[str, object]],
+    ) -> list[dict[str, object]] | None:
         """Convert chat-shape sanitized_messages to Responses API ``input``.
 
         Returns ``None`` if nothing usable could be converted, so the
         caller falls back to ``on_flagged_action``.
         """
-        out: Final[list[dict[str, Any]]] = []
+        out: Final[list[dict[str, object]]] = []
         for m in sanitized_messages:
             if not isinstance(m, dict):
                 continue
@@ -1764,7 +1764,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         start_time: datetime | None = None,
         surface: str = "chat",
         direction: str = "input",
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         verbose_proxy_logger.error(
             "Cisco AI Defense guardrail (%s): API communication failed: %s",
             surface,
@@ -2060,7 +2060,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         return getattr(obj, key, None)
 
     @classmethod
-    def _field_list(cls, obj: object, key: str) -> list[Any]:
+    def _field_list(cls, obj: object, key: str) -> list[object]:
         value: Final = cls._field(obj, key)
         return value if isinstance(value, list) else []
 

litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py

@@ -8,8 +8,8 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint
 
 import copy
 import json
-from collections.abc import AsyncGenerator
-from typing import TYPE_CHECKING, Any, Final
+from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Final, Protocol
 
 from fastapi import HTTPException
 
@@ -34,6 +34,9 @@ if TYPE_CHECKING:
     # Imported lazily at runtime (inside the streaming hook) to avoid a
     # module-level cyclic import with litellm.integrations.custom_guardrail.
     from litellm.integrations.custom_guardrail import ModifyResponseException
+    from litellm.llms.base_llm.guardrail_translation.base_translation import (
+        BaseTranslation,
+    )
 
 # Call types that use NDJSON streaming (A2A); guardrail HTTPException is emitted as in-stream error
 A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message)
@@ -41,12 +44,35 @@ A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message)
 GUARDRAIL_NAME: Final = "unified_llm_guardrails"
 
 
+class _EndpointTranslation(Protocol):
+    @property
+    def process_input_messages(self) -> "Callable[..., Awaitable[dict[str, object]]]": ...
+
+    @property
+    def process_output_response(self) -> "Callable[..., Awaitable[object]]": ...
+
+    @property
+    def process_output_streaming_response(self) -> "Callable[..., Awaitable[object]]": ...
+
+    @property
+    def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ...
+
+
+def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTranslation:
+    return translation
+
+
+def _chunk_choices(item: object) -> Sequence[object]:
+    choices: Final[Sequence[object]] = getattr(item, "choices", None) or []
+    return choices
+
+
 class _StreamTerminated(Exception):
     """Internal signal that the incremental transform stream has already emitted
     its terminal chunks (block message or in-stream error) and must stop."""
 
 
-def _get_a2a_request_id(responses_so_far: list[Any], request_data: dict) -> str | None:
+def _get_a2a_request_id(responses_so_far: Sequence[object], request_data: dict) -> str | None:
     """Get JSON-RPC request id from first A2A chunk or request body for in-stream error reporting."""
     for item in responses_so_far:
         if isinstance(item, dict) and "id" in item:
@@ -138,7 +164,9 @@ class UnifiedLLMGuardrails(CustomLogger):
         except ValueError:
             return data  # handle unmapped call types
 
-        endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
+        endpoint_translation: Final = _as_endpoint_translation(
+            endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
+        )
 
         _ensure_litellm_metadata(data, user_api_key_dict)
 
@@ -156,7 +184,7 @@ class UnifiedLLMGuardrails(CustomLogger):
 
     async def async_moderation_hook(
         self, data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral
-    ) -> Any:
+    ) -> object:
         """
         Runs in parallel to LLM API call
         Runs on only Input
@@ -187,7 +215,9 @@ class UnifiedLLMGuardrails(CustomLogger):
         if call_type is not None and CallTypes(call_type) not in endpoint_guardrail_translation_mappings:
             return data
 
-        endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
+        endpoint_translation: Final = _as_endpoint_translation(
+            endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
+        )
 
         _ensure_litellm_metadata(data, user_api_key_dict)
 
@@ -202,7 +232,7 @@ class UnifiedLLMGuardrails(CustomLogger):
         data: dict,
         user_api_key_dict: UserAPIKeyAuth,
         response,
-    ) -> Any:
+    ) -> object:
         """
         Runs on response from LLM API call
 
@@ -271,7 +301,9 @@ class UnifiedLLMGuardrails(CustomLogger):
             )
             return response
 
-        endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
+        endpoint_translation: Final = _as_endpoint_translation(
+            endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
+        )
 
         try:
             response = await endpoint_translation.process_output_response(
@@ -299,10 +331,10 @@ class UnifiedLLMGuardrails(CustomLogger):
     async def _handle_streaming_block(
         self,
         exc: "ModifyResponseException",
-        endpoint_translation: Any,
+        endpoint_translation: _EndpointTranslation,
         stream_started: bool,
-        responses_so_far: list[Any],
-    ) -> AsyncGenerator[Any, None]:
+        responses_so_far: Sequence[object],
+    ) -> AsyncGenerator[object, None]:
         """
         Terminate a streamed response cleanly when a guardrail blocks it.
 
@@ -323,7 +355,7 @@ class UnifiedLLMGuardrails(CustomLogger):
     @staticmethod
     def _resolve_transform_call_type(
         user_api_key_dict: UserAPIKeyAuth,
-        mappings: dict,
+        mappings: Mapping[CallTypes, type["BaseTranslation"]],
     ) -> str | None:
         """Resolve the call type for the incremental_diff path, or None if the
         route is unresolvable / unsupported.
@@ -356,9 +388,9 @@ class UnifiedLLMGuardrails(CustomLogger):
         self,
         exc: HTTPException,
         call_type: str | None,
-        responses_so_far: list[Any],
+        responses_so_far: Sequence[object],
         request_data: dict,
-    ) -> AsyncGenerator[Any, None]:
+    ) -> AsyncGenerator[object, None]:
         """Surface a mid-stream HTTPException. For A2A (NDJSON) call types the
         response has already started, so emit an in-stream JSON-RPC error chunk;
         otherwise re-raise so the proxy can report it.
@@ -387,7 +419,7 @@ class UnifiedLLMGuardrails(CustomLogger):
     def _build_transform_chunk(
         self,
         *,
-        reference_chunk: Any,
+        reference_chunk: object,
         mutated_text_per_choice: dict[int, str],
         emitted_text_per_choice: dict[int, str],
         holdback_per_choice: dict[int, int],
@@ -500,18 +532,18 @@ class UnifiedLLMGuardrails(CustomLogger):
     async def _emit_transform_round(
         self,
         *,
-        endpoint_translation: Any,
+        endpoint_translation: _EndpointTranslation,
         guardrail_to_apply: CustomGuardrail,
         request_data: dict,
         user_api_key_dict: UserAPIKeyAuth,
         call_type: str,
-        reference_chunk: Any,
-        responses_so_far: list[Any],
-        responses_yielded: list[Any],
+        reference_chunk: object,
+        responses_so_far: Sequence[object],
+        responses_yielded: list[object],
         emitted_text_per_choice: dict[int, str],
         finish_reason_per_choice: dict[int, str | None],
         is_final: bool,
-    ) -> AsyncGenerator[Any, None]:
+    ) -> AsyncGenerator[object, None]:
         """Run one guardrail processing round and emit the resulting diff chunk.
 
         Raises ``_StreamTerminated`` (after emitting the terminal block message or
@@ -564,14 +596,14 @@ class UnifiedLLMGuardrails(CustomLogger):
         self,
         *,
         guardrail_to_apply: CustomGuardrail,
-        response: Any,
+        response: AsyncIterable[object],
         request_data: dict,
         user_api_key_dict: UserAPIKeyAuth,
         call_type: str,
         sampling_rate: int,
         end_of_stream_only: bool,
-        mappings: dict,
-    ) -> AsyncGenerator[Any, None]:
+        mappings: Mapping[CallTypes, type["BaseTranslation"]],
+    ) -> AsyncGenerator[object, None]:
         """Emit guardrail text transformations as new deltas on the stream.
 
         Raw chunks are withheld and accumulated; on each sampled processing round
@@ -580,15 +612,15 @@ class UnifiedLLMGuardrails(CustomLogger):
         synthetic chunk. A BLOCK terminates the stream via the shared block
         handler; an underflow surfaces as an HTTPException.
         """
-        endpoint_translation: Final = mappings[CallTypes(call_type)]()
-        responses_so_far: Final[list[Any]] = []
-        responses_yielded: Final[list[Any]] = []
+        endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]())
+        responses_so_far: Final[list[object]] = []
+        responses_yielded: Final[list[object]] = []
         emitted_text_per_choice: Final[dict[int, str]] = {}
         finish_reason_per_choice: Final[dict[int, str | None]] = {}
         chunk_counter = 0
-        last_chunk: Any | None = None
+        last_chunk: object | None = None
 
-        def _round(reference_chunk: Any, is_final: bool) -> AsyncGenerator[Any, None]:
+        def _round(reference_chunk: object, is_final: bool) -> AsyncGenerator[object, None]:
             return self._emit_transform_round(
                 endpoint_translation=endpoint_translation,
                 guardrail_to_apply=guardrail_to_apply,
@@ -694,13 +726,13 @@ class UnifiedLLMGuardrails(CustomLogger):
     async def _inspect_full_response_for_block(
         self,
         *,
-        endpoint_translation: Any,
+        endpoint_translation: _EndpointTranslation,
         guardrail_to_apply: CustomGuardrail,
         request_data: dict,
         user_api_key_dict: UserAPIKeyAuth,
-        responses_so_far: list[Any],
-        responses_yielded: list[Any],
-    ) -> AsyncGenerator[Any, None]:
+        responses_so_far: Sequence[object],
+        responses_yielded: Sequence[object],
+    ) -> AsyncGenerator[object, None]:
         """Run the block-only guardrail inspection over the full assembled
         response (text + tool calls) so nothing bypasses the block decision.
 
@@ -734,17 +766,17 @@ class UnifiedLLMGuardrails(CustomLogger):
             raise _StreamTerminated()
 
     @staticmethod
-    def _chunk_has_tool_calls(item: Any) -> bool:
-        for choice in getattr(item, "choices", None) or []:
+    def _chunk_has_tool_calls(item: object) -> bool:
+        for choice in _chunk_choices(item):
             delta = getattr(choice, "delta", None)
             if getattr(delta, "tool_calls", None):
                 return True
         return False
 
     @staticmethod
-    def _chunk_carries_text(item: Any) -> bool:
+    def _chunk_carries_text(item: object) -> bool:
         """True if any choice in this chunk has non-empty string ``delta.content``."""
-        for choice in getattr(item, "choices", None) or []:
+        for choice in _chunk_choices(item):
             delta = getattr(choice, "delta", None)
             content = getattr(delta, "content", None)
             if isinstance(content, str) and content != "":
@@ -753,7 +785,7 @@ class UnifiedLLMGuardrails(CustomLogger):
 
     @staticmethod
     def _tool_call_passthrough_chunk(
-        item: Any,
+        item: object,
         finish_reason_per_choice: "dict[int, str | None] | None" = None,
     ) -> ModelResponseStream:
         """Copy of a chunk carrying tool calls with all text content stripped.
@@ -772,7 +804,7 @@ class UnifiedLLMGuardrails(CustomLogger):
         redaction purpose.
         """
         synthetic_choices: Final[list[StreamingChoices]] = []
-        for choice in getattr(item, "choices", None) or []:
+        for choice in _chunk_choices(item):
             delta = getattr(choice, "delta", None)
             idx = getattr(choice, "index", 0) or 0
             original_finish = getattr(choice, "finish_reason", None)
@@ -801,15 +833,15 @@ class UnifiedLLMGuardrails(CustomLogger):
         )
 
     @staticmethod
-    def _record_finish_reasons(item: Any, finish_reason_per_choice: dict[int, str | None]) -> None:
-        for choice in getattr(item, "choices", None) or []:
+    def _record_finish_reasons(item: object, finish_reason_per_choice: dict[int, str | None]) -> None:
+        for choice in _chunk_choices(item):
             finish_reason = getattr(choice, "finish_reason", None)
             if finish_reason is not None:
                 finish_reason_per_choice[getattr(choice, "index", 0) or 0] = finish_reason
 
     @staticmethod
-    def _chunk_has_finish_reason(item: Any) -> bool:
-        choices: Final = getattr(item, "choices", None) or []
+    def _chunk_has_finish_reason(item: object) -> bool:
+        choices: Final = _chunk_choices(item)
         return any(getattr(choice, "finish_reason", None) is not None for choice in choices)
 
     async def async_post_call_streaming_iterator_hook(
@@ -845,22 +877,22 @@ class UnifiedLLMGuardrails(CustomLogger):
         # Get streaming configuration. Resolution order (later wins): default
         # < guardrail attribute < guardrail_config dict < this callback's
         # optional_params.
-        def _streaming_flag(name: str, default: Any) -> Any:
+        def _streaming_flag(name: str, default: object) -> Any:
             value = default
             if guardrail_to_apply is not None:
                 value = getattr(guardrail_to_apply, name, value)
-                config: Final = getattr(guardrail_to_apply, "guardrail_config", {})
+                config: Final[Mapping[str, object]] = getattr(guardrail_to_apply, "guardrail_config", {})
                 if isinstance(config, dict):
                     value = config.get(name, value)
             return self.optional_params.get(name, value)
 
-        sampling_rate: Final = _streaming_flag("streaming_sampling_rate", 5)
+        sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5)
         # Only apply the guardrail at end of stream (not per chunk).
-        end_of_stream_only = _streaming_flag("streaming_end_of_stream_only", False)
+        end_of_stream_only: bool = _streaming_flag("streaming_end_of_stream_only", False)
         # "block_only" (default) drops guardrail text rewrites on the streaming
         # path; "incremental_diff" emits them as synthetic deltas (see
         # _run_incremental_transform_stream).
-        streaming_transform_mode: Final = _streaming_flag("streaming_transform_mode", "block_only")
+        streaming_transform_mode: Final[str] = _streaming_flag("streaming_transform_mode", "block_only")
         # Withhold every chunk until end-of-stream moderation passes, then
         # release the original chunks (clean) or only the block message
         # (blocked) -- moderating the whole response *before* any content
@@ -868,7 +900,9 @@ class UnifiedLLMGuardrails(CustomLogger):
         # release the original chunks are replayed as-is, so a
         # content-rewriting guardrail (e.g. PII masking) would leak
         # unredacted content. Guarded below via mask_response_content.
-        buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", buffer_until_moderated_default)
+        buffer_until_moderated: bool = _streaming_flag(
+            "streaming_buffer_until_moderated", buffer_until_moderated_default
+        )
 
         if (
             buffer_until_moderated
@@ -939,9 +973,9 @@ class UnifiedLLMGuardrails(CustomLogger):
         # Infer call type from first chunk
         call_type = None
         chunk_counter = 0
-        responses_so_far: Final[list[Any]] = []
-        responses_yielded: Final[list[Any]] = []
-        pending_end_of_stream_items: Final[list[Any]] = []
+        responses_so_far: Final[list[object]] = []
+        responses_yielded: Final[list[object]] = []
+        pending_end_of_stream_items: Final[list[object]] = []
         # Whether any real response chunk has been forwarded to the client.
         # Drives how a block terminates the stream: continue the in-progress
         # message (True) vs emit a standalone block message (False, buffered).

litellm/proxy/hooks/litellm_skills/main.py

@@ -26,7 +26,8 @@ Usage:
 
 import base64
 import json
-from typing import Any, Final
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Final
 
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.caching import DualCache
@@ -36,7 +37,10 @@ from litellm.llms.litellm_proxy.skills.prompt_injection import (
     SkillPromptInjectionHandler,
 )
 from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth
-from litellm.types.utils import CallTypes, CallTypesLiteral
+from litellm.types.utils import CallTypes, CallTypesLiteral, LLMResponseTypes
+
+if TYPE_CHECKING:
+    from litellm.llms.litellm_proxy.skills.sandbox_executor import SkillsSandboxExecutor
 
 
 class SkillsInjectionHook(CustomLogger):
@@ -99,7 +103,7 @@ class SkillsInjectionHook(CustomLogger):
         verbose_proxy_logger.debug("SkillsInjectionHook: Processing %s skills", len(skills))
 
         litellm_skills: Final[list[LiteLLM_SkillsTable]] = []
-        anthropic_skills: Final[list[dict[str, Any]]] = []
+        anthropic_skills: Final[list[dict[str, object]]] = []
 
         # Separate skills by prefix
         for skill in skills:
@@ -324,9 +328,9 @@ class SkillsInjectionHook(CustomLogger):
     async def async_post_call_success_deployment_hook(
         self,
         request_data: dict,
-        response: Any,
+        response: LLMResponseTypes,
         call_type: CallTypes | None,
-    ) -> Any | None:
+    ) -> LLMResponseTypes | None:
         """
         Post-call hook to handle automatic code execution.
 
@@ -372,7 +376,7 @@ class SkillsInjectionHook(CustomLogger):
         # Check if any tool call needs execution (litellm_code_execution or skill tool)
         has_executable_tool = False
         for tc in tool_calls:
-            tool_name = tc.get("name", "")
+            tool_name: str = tc.get("name", "")
             # Execute if it's litellm_code_execution OR a skill tool (litellm_skill_xxx)
             if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value or tool_name.startswith(LITELLM_SKILL_ID_PREFIX):
                 has_executable_tool = True
@@ -441,7 +445,7 @@ class SkillsInjectionHook(CustomLogger):
         data: dict,
         response: Any,
         skill_files: dict[str, bytes],
-    ) -> Any:
+    ) -> LLMResponseTypes | None:
         """
         Execute the code execution loop for messages API (Anthropic format).
 
@@ -466,7 +470,7 @@ class SkillsInjectionHook(CustomLogger):
         max_tokens: Final = data.get("max_tokens", 4096)
 
         executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout)
-        generated_files: Final[list[dict[str, Any]]] = []
+        generated_files: Final[list[dict[str, object]]] = []
         current_response = response
 
         for iteration in range(self.max_iterations):
@@ -511,9 +515,9 @@ class SkillsInjectionHook(CustomLogger):
             # Process tool calls
             tool_results = []
             for tc in tool_calls:
-                tool_name = tc.get("name", "")
+                tool_name: str = tc.get("name", "")
                 tool_id = tc.get("id", "")
-                tool_input = tc.get("input", {})
+                tool_input: Mapping[str, str] = tc.get("input", {})
 
                 # Execute if it's litellm_code_execution OR a skill tool
                 if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value:
@@ -561,8 +565,8 @@ class SkillsInjectionHook(CustomLogger):
         self,
         code: str,
         skill_files: dict[str, bytes],
-        executor: Any,
-        generated_files: list[dict[str, Any]],
+        executor: "SkillsSandboxExecutor",
+        generated_files: list[dict[str, object]],
     ) -> str:
         """Execute code in sandbox and return result string."""
         try:
@@ -574,7 +578,8 @@ class SkillsInjectionHook(CustomLogger):
 
             # Collect generated files
             if exec_result.get("files"):
-                for f in exec_result["files"]:
+                files: Final[Sequence[Mapping[str, str]]] = exec_result["files"]
+                for f in files:
                     generated_files.append(
                         {
                             "name": f["name"],
@@ -595,10 +600,10 @@ class SkillsInjectionHook(CustomLogger):
     async def _execute_skill_tool(
         self,
         tool_name: str,
-        tool_input: dict[str, Any],
+        tool_input: Mapping[str, str],
         skill_files: dict[str, bytes],
-        executor: Any,
-        generated_files: list[dict[str, Any]],
+        executor: "SkillsSandboxExecutor",
+        generated_files: list[dict[str, object]],
     ) -> str:
         """Execute a skill tool by generating and running code based on skill content."""
         # Generate code based on available skill modules
@@ -670,7 +675,7 @@ print('No executable skill module found')
         data: dict,
         response: Any,
         skill_files: dict[str, bytes],
-    ) -> Any:
+    ) -> LLMResponseTypes:
         """
         Execute the code execution loop until model gives final response.
 
@@ -704,7 +709,7 @@ print('No executable skill module found')
         kwargs: Final = {k: v for k, v in data.items() if k not in _EXCLUDED_ACOMPLETION_KEYS}
 
         executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout)
-        generated_files: Final[list[dict[str, Any]]] = []
+        generated_files: Final[list[dict[str, object]]] = []
         current_response: Any = response
 
         for iteration in range(self.max_iterations):
@@ -713,7 +718,7 @@ print('No executable skill module found')
             stop_reason = current_response.choices[0].finish_reason
 
             # Build assistant message for conversation history
-            assistant_msg_dict: dict[str, Any] = {
+            assistant_msg_dict: dict[str, object] = {
                 "role": "assistant",
                 "content": assistant_message.content,
             }
@@ -781,13 +786,13 @@ print('No executable skill module found')
         self,
         tool_call: Any,
         skill_files: dict[str, bytes],
-        executor: Any,
-        generated_files: list[dict[str, Any]],
+        executor: "SkillsSandboxExecutor",
+        generated_files: list[dict[str, object]],
     ) -> str:
         """Execute a litellm_code_execution tool call and return result string."""
         try:
             args: Final = json.loads(tool_call.function.arguments)
-            code: Final = args.get("code", "")
+            code: Final[str] = args.get("code", "")
 
             verbose_proxy_logger.debug("SkillsInjectionHook: Executing code (%s chars)", len(code))
 
@@ -802,7 +807,8 @@ print('No executable skill module found')
             # Collect generated files
             if exec_result.get("files"):
                 tool_result += "\n\nGenerated files:"
-                for f in exec_result["files"]:
+                files: Final[Sequence[Mapping[str, str]]] = exec_result["files"]
+                for f in files:
                     file_content = base64.b64decode(f["content_base64"])
                     generated_files.append(
                         {
@@ -830,8 +836,8 @@ print('No executable skill module found')
     def _attach_files_to_response(
         self,
         response: Any,
-        generated_files: list[dict[str, Any]],
-    ) -> Any:
+        generated_files: list[dict[str, object]],
+    ) -> LLMResponseTypes:
         """
         Attach generated files to the response object.
 
@@ -841,11 +847,13 @@ print('No executable skill module found')
         if not generated_files:
             return response
 
+        raw_response: Final = response
+
         # Handle dict response (Anthropic/messages API format)
         if isinstance(response, dict):
             response["_litellm_generated_files"] = generated_files
             verbose_proxy_logger.debug("SkillsInjectionHook: Attached %s files to dict response", len(generated_files))
-            return response
+            return raw_response
 
         # Handle object response (OpenAI format)
         try:

litellm/proxy/management_endpoints/internal_user_endpoints.py

@@ -17,7 +17,7 @@ import json
 import traceback
 from collections.abc import Mapping, Sequence
 from datetime import datetime, timezone
-from typing import Any, Final, cast
+from typing import Any, Final, Literal, cast
 
 import fastapi
 from fastapi import APIRouter, Depends, Header, HTTPException, Request, status

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -18,7 +18,7 @@ import os
 import re
 import secrets
 import traceback
-from collections.abc import Callable, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Mapping, Sequence
 from datetime import datetime, timedelta, timezone
 from typing import Any, Final, Literal, Optional, Protocol, TypeVar, cast
 
@@ -171,8 +171,12 @@ class _PrismaTableActions(Protocol[_PrismaRowT]):
 
     async def count(self, *, where: Mapping[str, object] | None = None) -> int: ...
 
+    async def create(self, *, data: Mapping[str, object]) -> _PrismaRowT: ...
+
     async def create_many(self, *, data: Sequence[Mapping[str, object]]) -> int: ...
 
+    async def delete_many(self, *, where: Mapping[str, object] | None = None) -> int: ...
+
     async def update(
         self,
         *,
@@ -181,6 +185,10 @@ class _PrismaTableActions(Protocol[_PrismaRowT]):
     ) -> _PrismaRowT | None: ...
 
 
+class _TxTables(Protocol):
+    litellm_proxymodeltable: _PrismaTableActions[object]
+
+
 def _prisma_table(
     repository: BaseRepository[_RepositoryModelT],
 ) -> _PrismaTableActions[_RepositoryModelT]:
@@ -1650,9 +1658,12 @@ async def generate_key_fn(
                 detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"},
             )
 
-        if user_custom_key_generate is not None:
-            if inspect.iscoroutinefunction(user_custom_key_generate):
-                result: Final = await user_custom_key_generate(data)
+        custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = (
+            user_custom_key_generate
+        )
+        if custom_key_generate_hook is not None:
+            if inspect.iscoroutinefunction(custom_key_generate_hook):
+                result: Final = await custom_key_generate_hook(data)
             else:
                 raise ValueError("user_custom_key_generate must be a coroutine")
             decision: Final = result.get("decision", True)
@@ -1847,9 +1858,10 @@ async def generate_service_account_key_fn(
 
     verbose_proxy_logger.debug("entered /key/generate")
 
-    if user_custom_key_generate is not None:
-        if inspect.iscoroutinefunction(user_custom_key_generate):
-            result: Final = await user_custom_key_generate(data)
+    custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_generate
+    if custom_key_generate_hook is not None:
+        if inspect.iscoroutinefunction(custom_key_generate_hook):
+            result: Final = await custom_key_generate_hook(data)
         else:
             raise ValueError("user_custom_key_generate must be a coroutine")
         decision: Final = result.get("decision", True)
@@ -1918,7 +1930,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_
             )
         casted_metadata[reserved_field] = existing_value
 
-    data_json: Final = data.model_dump(exclude_unset=True, exclude_none=True)
+    data_json: Final[Mapping[str, object]] = data.model_dump(exclude_unset=True, exclude_none=True)
 
     try:
         for k, v in data_json.items():
@@ -2179,7 +2191,7 @@ async def _process_single_key_update(
     llm_router: Router | None,
     user_custom_key_update: Callable | None = None,
     existing_key_row: LiteLLM_VerificationToken | None = None,
-) -> dict[str, Any]:
+) -> dict[str, object]:
     """
     Process a single key update with all validations and checks.
 
@@ -2722,9 +2734,10 @@ async def update_key_fn(
         )
 
         # Custom key update hook
-        if user_custom_key_update is not None:
-            if inspect.iscoroutinefunction(user_custom_key_update):
-                result: Final = await user_custom_key_update(data)
+        custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_update
+        if custom_key_update_hook is not None:
+            if inspect.iscoroutinefunction(custom_key_update_hook):
+                result: Final = await custom_key_update_hook(data)
             else:
                 raise ValueError("user_custom_key_update must be a coroutine")
             decision: Final = result.get("decision", True)
@@ -4089,10 +4102,11 @@ async def delete_verification_tokens(
     failed_tokens: list = []
     try:
         if prisma_client:
-            tokens = [_hash_token_if_needed(token=key) for key in tokens]
-            _keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = await VerificationTokenRepository(
-                prisma_client
-            ).table.find_many(where={"token": {"in": tokens}})
+            hashed_tokens: Final[list[str]] = [_hash_token_if_needed(token=key) for key in tokens]
+            tokens = hashed_tokens
+            _keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = await _prisma_table(
+                VerificationTokenRepository(prisma_client)
+            ).find_many(where={"token": {"in": hashed_tokens}})
 
             if len(_keys_being_deleted) == 0:
                 raise HTTPException(
@@ -4291,7 +4305,7 @@ async def _rotate_master_key(
     if models:
         decrypted_models: Final = proxy_config.decrypt_model_list_from_db(new_models=models)
         verbose_proxy_logger.debug("ABLE TO DECRYPT MODELS - len(decrypted_models): %s", len(decrypted_models))
-        new_models: Final = []
+        new_models: Final[list[dict[str, object]]] = []
         for model in decrypted_models:
             new_model = await _add_model_to_db(
                 model_params=Deployment(**model),
@@ -4306,7 +4320,8 @@ async def _rotate_master_key(
                 _dumped["model_info"] = prisma.Json(_dumped["model_info"])
                 new_models.append(_dumped)
         verbose_proxy_logger.debug("Resetting proxy model table")
-        async with prisma_client.db.tx() as tx:
+        async with prisma_client.db.tx() as tx_ctx:
+            tx: Final[_TxTables] = tx_ctx
             await tx.litellm_proxymodeltable.delete_many()
             verbose_proxy_logger.debug("Creating %s models", len(new_models))
             await tx.litellm_proxymodeltable.create_many(
@@ -4630,7 +4645,7 @@ async def _execute_virtual_key_regeneration(
             _validate_key_alias_format(key_alias=new_key_alias)
         verbose_proxy_logger.debug("non_default_values: %s", non_default_values)
     update_data.update(non_default_values)
-    update_data = prisma_client.jsonify_object(data=update_data)
+    jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data)
 
     # If grace period set, insert deprecated key so old key remains valid
     await _insert_deprecated_key(
@@ -4642,9 +4657,9 @@ async def _execute_virtual_key_regeneration(
 
     updated_token: Final = await VerificationTokenRepository(prisma_client).table.update(
         where={"token": hashed_api_key},
-        data=update_data,
+        data=jsonified_update_data,
     )
-    updated_token_dict: Final = dict(updated_token) if updated_token is not None else {}
+    updated_token_dict: Final[dict[str, object]] = dict(updated_token) if updated_token is not None else {}
     updated_token_dict["key"] = new_token
     updated_token_dict["token_id"] = updated_token_dict.pop("token")
 
@@ -5589,7 +5604,7 @@ async def key_aliases(
         where_sql: Final = " AND ".join(where_parts)
 
         count_sql: Final = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}'
-        count_rows: Final = await prisma_client.db.query_raw(count_sql, *query_params)
+        count_rows: Final[Sequence[Mapping[str, int]]] = await prisma_client.db.query_raw(count_sql, *query_params)
         total_count: Final = int(count_rows[0]["count"]) if count_rows else 0
 
         aliases_params: Final = query_params + [size, (page - 1) * size]
@@ -5602,7 +5617,7 @@ async def key_aliases(
             f" ORDER BY key_alias ASC"
             f" LIMIT ${limit_idx} OFFSET ${offset_idx}"
         )
-        alias_rows: Final = await prisma_client.db.query_raw(aliases_sql, *aliases_params)
+        alias_rows: Final[Sequence[Mapping[str, str]]] = await prisma_client.db.query_raw(aliases_sql, *aliases_params)
         aliases: Final[list[str]] = [row["key_alias"] for row in alias_rows if row.get("key_alias")]
 
         total_pages: Final = -(-total_count // size) if total_count > 0 else 0
@@ -5695,7 +5710,7 @@ def _build_key_filter_conditions(
     agent_id: str | None = None,
     use_substring_matching: bool = False,
     expires_filter: str | None = None,
-) -> dict[str, str | dict[str, Any] | list[dict[str, Any]]]:
+) -> Mapping[str, object]:
     """Build filter conditions for key listing.
 
     Visibility rules:
@@ -5707,14 +5722,14 @@ def _build_key_filter_conditions(
       so former members cannot see service accounts they created after leaving.
     """
     # Prepare filter conditions
-    where: dict[str, str | dict[str, Any] | list[dict[str, Any]]] = {}
+    where: dict[str, object] = {}
     where.update(_get_condition_to_filter_out_ui_session_tokens())
 
     # Build the OR conditions for user's keys and admin team keys
-    or_conditions: Final[list[dict[str, Any]]] = []
+    or_conditions: Final[list[dict[str, object]]] = []
 
     # Base conditions for user's own keys
-    user_condition: Final[dict[str, Any]] = {}
+    user_condition: Final[dict[str, object]] = {}
     if user_id and isinstance(user_id, str):
         if use_substring_matching:
             user_condition["user_id"] = {
@@ -5784,7 +5799,7 @@ def _build_key_filter_conditions(
 
     # Apply team_id, project_id and access_group_id as global AND filters so they
     # narrow results across all visibility conditions (own keys, team keys, etc.)
-    global_filters: tuple[dict[str, Any], ...] = (
+    global_filters: Final[tuple[dict[str, object], ...]] = (
         *(
             (
                 {"key_alias": {"contains": key_alias, "mode": "insensitive"}}
@@ -5805,7 +5820,7 @@ def _build_key_filter_conditions(
             else ()
         ),
     )
-    combined_where = {"AND": [where, *global_filters]} if global_filters else where
+    combined_where: Final[Mapping[str, object]] = {"AND": [where, *global_filters]} if global_filters else where
     verbose_proxy_logger.debug("Filter conditions: %s", combined_where)
     return combined_where
 
@@ -5986,7 +6001,7 @@ async def _list_key_helper(
     )
 
 
-def _get_condition_to_filter_out_ui_session_tokens() -> dict[str, Any]:
+def _get_condition_to_filter_out_ui_session_tokens() -> Mapping[str, object]:
     """
     Condition to filter out UI session tokens
     """
@@ -6395,7 +6410,7 @@ async def _can_user_query_key_info(
 async def test_key_logging(
     user_api_key_dict: UserAPIKeyAuth,
     request: Request,
-    key_logging: list[dict[str, Any]],
+    key_logging: Sequence[Mapping[str, str]],
 ) -> LoggingCallbackStatus:
     """
     Test the key-based logging

litellm/proxy/management_endpoints/model_management_endpoints.py

@@ -13,9 +13,9 @@ model/{model_id}/update - PATCH endpoint for model update.
 import asyncio
 import datetime
 import json
-from collections.abc import Mapping, Sequence
+from collections.abc import Awaitable, Mapping, Sequence
 from json import JSONDecodeError
-from typing import Any, Final, Literal, cast
+from typing import Final, Literal, Protocol, cast
 
 from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
 from pydantic import BaseModel, ConfigDict, Field, ValidationError
@@ -78,9 +78,7 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import
 from litellm.types.router import (
     SPECIAL_MODEL_INFO_PARAMS,
     Deployment,
-    DeploymentTypedDict,
     GenericLiteLLMParams,
-    LiteLLMParamsTypedDict,
     updateDeployment,
 )
 from litellm.utils import get_utc_datetime
@@ -104,10 +102,80 @@ class UpdatePublicModelGroupsRequest(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
 
+class _ProxyModelRow(Protocol):
+    model_id: str
+    model_name: str
+    model_info: Mapping[str, object] | None
+
+    def model_dump_json(self, *, exclude_none: bool = False) -> str: ...
+
+
+class _ProxyModelTable(Protocol):
+    def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ...
+
+    def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ...
+
+    def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[_ProxyModelRow]: ...
+
+    def delete(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ...
+
+    def delete_many(self, *, where: Mapping[str, object]) -> Awaitable[int]: ...
+
+
+class _TxModelTables(Protocol):
+    litellm_proxymodeltable: _ProxyModelTable
+
+
+class _TeamRow(Protocol):
+    models: Sequence[str]
+
+    def model_dump(self) -> Mapping[str, object]: ...
+
+
+class _TeamTable(Protocol):
+    def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_TeamRow | None]: ...
+
+    def update(
+        self, *, where: Mapping[str, object], data: Mapping[str, object], include: Mapping[str, bool]
+    ) -> Awaitable[LiteLLM_TeamTable]: ...
+
+
+class _TeamIdRef(Protocol):
+    team_id: str
+
+
+class _ModelAliasRow(Protocol):
+    id: int
+    model_aliases: dict[str, str]
+    team: _TeamIdRef | None
+
+
+class _ModelAliasTable(Protocol):
+    def find_many(self, *, include: Mapping[str, bool]) -> Awaitable[Sequence[_ModelAliasRow]]: ...
+
+    def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[object]: ...
+
+
+def _proxy_model_table(prisma_client: PrismaClient) -> _ProxyModelTable:
+    return ModelRepository(prisma_client).table
+
+
+def _repo_team_table(prisma_client: PrismaClient) -> _TeamTable:
+    return TeamRepository(prisma_client).table
+
+
+def _db_team_table(prisma_client: PrismaClient) -> _TeamTable:
+    return prisma_client.db.litellm_teamtable
+
+
+def _model_alias_table(prisma_client: PrismaClient) -> _ModelAliasTable:
+    return ModelTableRepository(prisma_client).table
+
+
 async def get_db_model(model_id: str, prisma_client: PrismaClient) -> Deployment | None:
     db_model: Final = cast(
         BaseModel | None,
-        await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id}),
+        await _proxy_model_table(prisma_client).find_unique(where={"model_id": model_id}),
     )
 
     if not db_model:
@@ -166,14 +234,9 @@ def _raise_on_strategy_router_write_violation(
 
 
 def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel:
-    merged_deployment_dict: Final = DeploymentTypedDict(
-        model_name=db_model.model_name,
-        litellm_params=LiteLLMParamsTypedDict(**db_model.litellm_params.model_dump(exclude_none=True)),
-        model_info=db_model.model_info.model_dump(exclude_none=True),
-    )
-    # update model name
-    if updated_patch.model_name:
-        merged_deployment_dict["model_name"] = updated_patch.model_name
+    merged_model_name: Final = updated_patch.model_name or db_model.model_name
+    merged_litellm_params: Final = db_model.litellm_params.model_dump(exclude_none=True)
+    merged_model_info: Final = db_model.model_info.model_dump(exclude_none=True)
 
     # update litellm params
     if updated_patch.litellm_params:
@@ -182,13 +245,11 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
             k: encrypt_value_helper(v) for k, v in updated_patch.litellm_params.model_dump(exclude_none=True).items()
         }
 
-        merged_deployment_dict["litellm_params"].update(encrypted_params)
+        merged_litellm_params.update(encrypted_params)
 
     # update model info
     if updated_patch.model_info:
-        if "model_info" not in merged_deployment_dict:
-            merged_deployment_dict["model_info"] = {}
-        merged_deployment_dict["model_info"].update(updated_patch.model_info.model_dump(exclude_none=True))
+        merged_model_info.update(updated_patch.model_info.model_dump(exclude_none=True))
 
     # Honor explicit-null clears LAST, after both merges, so a model_info blob the UI
     # passes through (which today re-sends the OLD pricing on every save) cannot
@@ -202,29 +263,25 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
     if updated_patch.litellm_params:
         for field in updated_patch.litellm_params.model_fields_set:
             if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.litellm_params, field) is None:
-                merged_deployment_dict["litellm_params"].pop(field, None)
-                merged_deployment_dict.get("model_info", {}).pop(field, None)
+                merged_litellm_params.pop(field, None)
+                merged_model_info.pop(field, None)
     if updated_patch.model_info:
         for field in updated_patch.model_info.model_fields_set:
             if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None:
-                merged_deployment_dict["model_info"].pop(field, None)
-                merged_deployment_dict.get("litellm_params", {}).pop(field, None)
+                merged_model_info.pop(field, None)
+                merged_litellm_params.pop(field, None)
 
     # convert to prisma compatible format
 
-    prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel()
-    if "model_name" in merged_deployment_dict:
-        prisma_compatible_model_dict["model_name"] = merged_deployment_dict["model_name"]
+    for key, value in merged_model_info.items():
+        if isinstance(value, datetime.datetime):
+            merged_model_info[key] = value.isoformat()
 
-    if "litellm_params" in merged_deployment_dict:
-        prisma_compatible_model_dict["litellm_params"] = json.dumps(merged_deployment_dict["litellm_params"])
-
-    if "model_info" in merged_deployment_dict:
-        model_info: Final = merged_deployment_dict["model_info"]
-        for key, value in model_info.items():
-            if isinstance(value, datetime.datetime):
-                model_info[key] = value.isoformat()
-        prisma_compatible_model_dict["model_info"] = json.dumps(model_info)
+    prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel(
+        model_name=merged_model_name,
+        litellm_params=json.dumps(merged_litellm_params),
+        model_info=json.dumps(merged_model_info),
+    )
 
     if updated_patch.blocked is not None:
         prisma_compatible_model_dict["blocked"] = updated_patch.blocked
@@ -338,7 +395,7 @@ async def patch_model(
         update_data["updated_at"] = cast(str, get_utc_datetime())
 
         # Perform partial update
-        updated_model: Final = await ModelRepository(prisma_client).table.update(
+        updated_model: Final = await _proxy_model_table(prisma_client).update(
             where={"model_id": model_id},
             data=update_data,
         )
@@ -769,8 +826,8 @@ async def _setup_new_team_model_assignment(
 
 
 async def _get_team_deployments(
-    team_id: str, prisma_client: PrismaClient, table: Any | None = None
-) -> list[LiteLLM_ProxyModelTable]:
+    team_id: str, prisma_client: PrismaClient, table: _ProxyModelTable | None = None
+) -> Sequence[_ProxyModelRow]:
     """
     Fetch all deployments for a given team_id from the database.
 
@@ -785,7 +842,7 @@ async def _get_team_deployments(
     existing transaction.
     """
     prefix: Final = f"model_name_{team_id}_"
-    table = table or ModelRepository(prisma_client).table
+    table = table or _proxy_model_table(prisma_client)
     response: Final = await table.find_many(
         where={
             "model_name": {"startswith": prefix},
@@ -806,7 +863,7 @@ async def _get_team_deployments(
 async def delete_team_models(
     team_ids: list[str],
     prisma_client: PrismaClient,
-    llm_router: Any | None,
+    llm_router: Router | None,
 ) -> list[str]:
     """
     Delete every BYOK model owned by the given teams, from the DB and the router.
@@ -820,7 +877,8 @@ async def delete_team_models(
     Returns the model_ids that were deleted.
     """
     deleted_model_ids: Final[list[str]] = []
-    async with prisma_client.db.tx() as tx:
+    async with prisma_client.db.tx() as tx_ctx:
+        tx: Final[_TxModelTables] = tx_ctx
         for team_id in team_ids:
             rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable)
             model_ids = [row.model_id for row in rows]
@@ -920,11 +978,11 @@ async def _remove_unbacked_team_models(
     if not names_to_remove:
         return
 
-    existing_team_row: Final = await prisma_client.db.litellm_teamtable.find_unique(where={"team_id": team_id})
+    existing_team_row: Final = await _db_team_table(prisma_client).find_unique(where={"team_id": team_id})
     if existing_team_row is None:
         return
 
-    updated_team_row: Final[LiteLLM_TeamTable] = await prisma_client.db.litellm_teamtable.update(
+    updated_team_row: Final[LiteLLM_TeamTable] = await _db_team_table(prisma_client).update(
         where={"team_id": team_id},
         data={"models": [model for model in existing_team_row.models if model not in names_to_remove]},
         include={"object_permission": True},
@@ -953,7 +1011,7 @@ async def _update_existing_team_model_assignment(
     """
 
     def _get_team_public_model_name(
-        model_info: dict | str | None,
+        model_info: object,
     ) -> str | None:
         parsed: Final = model_info_as_mapping(model_info)
         if parsed is None:
@@ -1062,7 +1120,7 @@ class ModelManagementAuthChecks:
                 detail={"error": CommonProxyErrors.not_premium_user.value},
             )
 
-        _existing_team_row: Final = await TeamRepository(prisma_client).table.find_unique(
+        _existing_team_row: Final = await _repo_team_table(prisma_client).find_unique(
             where={"team_id": model_params.model_info.team_id}
         )
 
@@ -1091,7 +1149,7 @@ class ModelManagementAuthChecks:
     ) -> Literal[True]:
         ## Check team model auth
         if model_params.model_info is not None and model_params.model_info.team_id is not None:
-            team_obj_row: Final = await TeamRepository(prisma_client).table.find_unique(
+            team_obj_row: Final = await _repo_team_table(prisma_client).find_unique(
                 where={"team_id": model_params.model_info.team_id}
             )
             if team_obj_row is None:
@@ -1192,7 +1250,7 @@ async def delete_model(
             - store keys separately
             """
             # encrypt litellm params #
-            result: Final = await ModelRepository(prisma_client).table.delete(where={"model_id": model_info.id})
+            result: Final = await _proxy_model_table(prisma_client).delete(where={"model_id": model_info.id})
 
             if result is None:
                 raise HTTPException(
@@ -1265,9 +1323,9 @@ async def delete_team_model_alias(
     Returns:
     - List of team id + model alias pairs that were removed
     """
-    team_model_aliases: Final = await ModelTableRepository(prisma_client).table.find_many(include={"team": True})
+    team_model_aliases: Final = await _model_alias_table(prisma_client).find_many(include={"team": True})
     tasks: Final = []
-    removed_model_aliases: Final = []
+    removed_model_aliases: Final[list[tuple[str, str]]] = []
     for team_model_alias in team_model_aliases:
         model_aliases = team_model_alias.model_aliases  # {"alias": "public model name"}
         id = team_model_alias.id
@@ -1278,7 +1336,7 @@ async def delete_team_model_alias(
                 removed_model_aliases.append((team_model_alias.team.team_id, key))
             del model_aliases[key]
             tasks.append(
-                ModelTableRepository(prisma_client).table.update(
+                _model_alias_table(prisma_client).update(
                     where={"id": id},
                     data={"model_aliases": json.dumps(model_aliases)},
                 )
@@ -1492,7 +1550,7 @@ async def update_model(
                 },
             )
 
-        _model_id = None
+        _model_id: str | None = None
         _model_info: Final = getattr(model_params, "model_info", None)
         if _model_info is None:
             raise Exception("model_info not provided")
@@ -1551,11 +1609,11 @@ async def update_model(
                 else:
                     pass
 
-            _data: Final[dict] = {
+            _data: Final[dict[str, str]] = {
                 "litellm_params": json.dumps(merged_dictionary),
                 "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
             }
-            model_response: Final = await ModelRepository(prisma_client).table.update(
+            model_response: Final = await _proxy_model_table(prisma_client).update(
                 where={"model_id": _model_id},
                 data=_data,
             )

litellm/proxy/management_endpoints/team_endpoints.py

@@ -15,11 +15,11 @@ import math
 import traceback
 from collections.abc import Mapping, Sequence
 from datetime import datetime, timezone
-from typing import Annotated, Final, Protocol, TypeVar, cast
+from typing import Annotated, Final, Protocol, TypedDict, TypeVar, cast
 
 import fastapi
 from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
-from pydantic import BaseModel
+from pydantic import BaseModel, JsonValue
 
 import litellm
 from litellm._logging import verbose_proxy_logger
@@ -29,6 +29,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.proxy._types import (
     UI_TEAM_ID,
     BlockTeamRequest,
+    BudgetNewRequest,
     CommonProxyErrors,
     DeleteTeamRequest,
     LiteLLM_AccessGroupTable,
@@ -156,6 +157,15 @@ router: Final = APIRouter()
 _DbRecordT = TypeVar("_DbRecordT")
 
 
+class _TeamIdKeyCount(TypedDict):
+    team_id: int
+
+
+class _TeamIdGroupRow(TypedDict):
+    team_id: str
+    _count: _TeamIdKeyCount
+
+
 class _PrismaTableActions(Protocol[_DbRecordT]):
     async def find_unique(
         self,
@@ -220,59 +230,127 @@ class _PrismaTableActions(Protocol[_DbRecordT]):
         where: Mapping[str, object] | None = None,
     ) -> int: ...
 
+    async def group_by(
+        self,
+        by: Sequence[str],
+        where: Mapping[str, object] | None = None,
+        count: Mapping[str, bool] | None = None,
+    ) -> Sequence[_TeamIdGroupRow]: ...
+
+
+class _HasTableActions(Protocol[_DbRecordT]):
+    @property
+    def table(self) -> "_PrismaTableActions[_DbRecordT]": ...
+
+
+def _typed_table(
+    repo: "_HasTableActions[_DbRecordT]", record_type: type[_DbRecordT]
+) -> "_PrismaTableActions[_DbRecordT]":
+    return repo.table
+
+
+def _as_object(value: object) -> object:
+    return value
+
+
+def _nullable(value: _DbRecordT | None) -> _DbRecordT | None:
+    return value
+
+
+class _UserIdRow(Protocol):
+    @property
+    def user_id(self) -> str | None: ...
+
+
+class _HasUserIdTable(Protocol):
+    @property
+    def table(self) -> "_PrismaTableActions[_UserIdRow]": ...
+
+
+def _user_id_rows_db(repo: "_HasUserIdTable") -> "_PrismaTableActions[_UserIdRow]":
+    return repo.table
+
+
+class _RawTeamRow(Protocol):
+    @property
+    def members_with_roles(self) -> Sequence[Mapping[str, object]] | None: ...
+
+
+class _HasRawTeamTable(Protocol):
+    @property
+    def table(self) -> "_PrismaTableActions[_RawTeamRow]": ...
+
+
+def _raw_team_db(repo: "_HasRawTeamTable") -> "_PrismaTableActions[_RawTeamRow]":
+    return repo.table
+
+
+class _BudgetWriteCall(Protocol):
+    async def __call__(
+        self, budget_obj: BudgetNewRequest, user_api_key_dict: UserAPIKeyAuth
+    ) -> LiteLLM_BudgetTableFull: ...
+
+
+def _as_budget_write(fn: "_BudgetWriteCall") -> "_BudgetWriteCall":
+    return fn
+
+
+class _TeamFindManyArgs(TypedDict, total=False):
+    take: int
+    skip: int
+    order: Mapping[str, str]
+    cursor: Mapping[str, object]
+
+
+class _TeamUiViewFilters(TypedDict, total=False):
+    team_id: Mapping[str, str]
+    team_alias: Mapping[str, str]
+
+
+class _TeamIdInFilter(TypedDict, total=False):
+    team_id: Mapping[str, Sequence[str]]
+
 
 def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]":
-    team_table: Final[_PrismaTableActions[LiteLLM_TeamTable]] = TeamRepository(prisma_client).table
-    return team_table
+    return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable)
 
 
 def _team_membership_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamMembership]":
-    membership_table: Final[_PrismaTableActions[LiteLLM_TeamMembership]] = TeamMembershipRepository(prisma_client).table
-    return membership_table
+    return _typed_table(TeamMembershipRepository(prisma_client), LiteLLM_TeamMembership)
 
 
 def _user_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_UserTable]":
-    user_table: Final[_PrismaTableActions[LiteLLM_UserTable]] = UserRepository(prisma_client).table
-    return user_table
+    return _typed_table(UserRepository(prisma_client), LiteLLM_UserTable)
 
 
 def _model_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_ModelTable]":
-    model_table: Final[_PrismaTableActions[LiteLLM_ModelTable]] = ModelTableRepository(prisma_client).table
-    return model_table
+    return _typed_table(ModelTableRepository(prisma_client), LiteLLM_ModelTable)
 
 
 def _org_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_OrganizationTable]":
-    org_table: Final[_PrismaTableActions[LiteLLM_OrganizationTable]] = OrganizationRepository(prisma_client).table
-    return org_table
+    return _typed_table(OrganizationRepository(prisma_client), LiteLLM_OrganizationTable)
 
 
 def _org_membership_db(
     prisma_client: PrismaClient | None,
 ) -> "_PrismaTableActions[LiteLLM_OrganizationMembershipTable]":
-    org_membership_table: _PrismaTableActions[LiteLLM_OrganizationMembershipTable] = OrganizationMembershipRepository(
-        prisma_client
-    ).table
-    return org_membership_table
+    return _typed_table(OrganizationMembershipRepository(prisma_client), LiteLLM_OrganizationMembershipTable)
 
 
 def _budget_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_BudgetTableFull]":
-    budget_table: Final[_PrismaTableActions[LiteLLM_BudgetTableFull]] = BudgetRepository(prisma_client).table
-    return budget_table
+    return _typed_table(BudgetRepository(prisma_client), LiteLLM_BudgetTableFull)
 
 
 def _deleted_team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_DeletedTeamTable]":
-    deleted_team_table: _PrismaTableActions[LiteLLM_DeletedTeamTable] = DeletedTeamRepository(prisma_client).table
-    return deleted_team_table
+    return _typed_table(DeletedTeamRepository(prisma_client), LiteLLM_DeletedTeamTable)
 
 
 def _access_group_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_AccessGroupTable]":
-    access_group_table: _PrismaTableActions[LiteLLM_AccessGroupTable] = AccessGroupRepository(prisma_client).table
-    return access_group_table
+    return _typed_table(AccessGroupRepository(prisma_client), LiteLLM_AccessGroupTable)
 
 
 def _tokens_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_VerificationToken]":
-    tokens_table: _PrismaTableActions[LiteLLM_VerificationToken] = VerificationTokenRepository(prisma_client).table
-    return tokens_table
+    return _typed_table(VerificationTokenRepository(prisma_client), LiteLLM_VerificationToken)
 
 
 def _sanitize_for_log(value: object) -> str:
@@ -408,7 +486,7 @@ class TeamMemberBudgetHandler:
         if team_member_budget_duration is not None:
             budget_request.budget_duration = team_member_budget_duration
 
-        team_member_budget_table: Final = await new_budget(
+        team_member_budget_table: Final = await _as_budget_write(new_budget)(
             budget_obj=budget_request,
             user_api_key_dict=user_api_key_dict,
         )
@@ -456,7 +534,7 @@ class TeamMemberBudgetHandler:
             if team_member_budget_duration is not None:
                 budget_request.budget_duration = team_member_budget_duration
 
-            budget_row: Final = await update_budget(
+            budget_row: Final = await _as_budget_write(update_budget)(
                 budget_obj=budget_request,
                 user_api_key_dict=user_api_key_dict,
             )
@@ -571,7 +649,7 @@ class TeamMemberBudgetHandler:
                 )
 
         if missing:
-            await TeamMembershipRepository(prisma_client).table.create_many(
+            await _team_membership_db(prisma_client).create_many(
                 data=missing,
                 skip_duplicates=True,  # safety net against concurrent races
             )
@@ -1407,9 +1485,10 @@ async def new_team(
             complete_team_data_dict["metadata"] = encrypt_callback_vars(complete_team_data_dict["metadata"])
 
         complete_team_data_dict = prisma_client.jsonify_team_object(db_data=complete_team_data_dict)
+        team_creation_data: Final[Mapping[str, object]] = complete_team_data_dict
 
-        team_row: Final[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.create(
-            data=complete_team_data_dict,
+        team_row: Final[LiteLLM_TeamTable] = await _team_db(prisma_client).create(
+            data=team_creation_data,
             include={"litellm_model_table": True},
         )
 
@@ -1856,7 +1935,7 @@ async def update_team(
                 detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"},
             )
 
-        existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id})
+        existing_team_row = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id})
 
         if existing_team_row is None:
             raise HTTPException(
@@ -1884,7 +1963,7 @@ async def update_team(
                     )
 
         if data.max_budget is not None:
-            existing_soft_budget: Final = getattr(existing_team_row, "soft_budget", None)
+            existing_soft_budget: Final[object] = _as_object(getattr(existing_team_row, "soft_budget", None))
             soft_budget_to_check: Final = data.soft_budget if data.soft_budget is not None else existing_soft_budget
             if soft_budget_to_check is not None and isinstance(soft_budget_to_check, (int, float)):
                 if data.max_budget <= soft_budget_to_check:
@@ -1943,7 +2022,7 @@ async def update_team(
             data.organization_id = None
 
         # check org team limits - if updating team that belongs to an org
-        org_id_to_check: Final = (
+        org_id_to_check: Final[object] = _as_object(
             data.organization_id if data.organization_id is not None else existing_team_row.organization_id
         )
         if org_id_to_check is not None and isinstance(org_id_to_check, str) and prisma_client is not None:
@@ -1976,7 +2055,7 @@ async def update_team(
             TeamMemberBudgetHandler.strip_system_managed_metadata_keys(updated_kv["metadata"])
 
         if "metadata" in updated_kv:
-            stored_metadata: Final = (
+            stored_metadata: Final[Mapping[str, JsonValue] | None] = (
                 {  # mutable-ok: the validator payload's isinstance guard requires a plain dict
                     key: value
                     for key, value in existing_team_row.metadata.items()
@@ -2079,16 +2158,19 @@ async def update_team(
             updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"])
 
         updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv)
-        team_row: Final[LiteLLM_TeamTable | None] = await TeamRepository(prisma_client).table.update(
-            where={"team_id": data.team_id},
-            data=updated_kv,
-            # `object_permission` is included so `_refresh_cached_team`
-            # doesn't write a cached team with the relation nulled out —
-            # see team_model_add for the full rationale.
-            include={
-                "litellm_model_table": True,
-                "object_permission": True,
-            },
+        team_update_data: Final[Mapping[str, object]] = updated_kv
+        team_row: Final[LiteLLM_TeamTable | None] = _nullable(
+            await _team_db(prisma_client).update(
+                where={"team_id": data.team_id},
+                data=team_update_data,
+                # `object_permission` is included so `_refresh_cached_team`
+                # doesn't write a cached team with the relation nulled out —
+                # see team_model_add for the full rationale.
+                include={
+                    "litellm_model_table": True,
+                    "object_permission": True,
+                },
+            )
         )
 
         if team_row is None or team_row.team_id is None:
@@ -2603,7 +2685,7 @@ async def _resolve_existing_member_user_ids(
     if not requested_user_ids:
         return frozenset()
 
-    found: Final = await UserRepository(prisma_client).table.find_many(
+    found: Final = await _user_id_rows_db(UserRepository(prisma_client)).find_many(
         where={  # mutable-ok: Prisma query filters are dict-shaped
             "user_id": {  # mutable-ok: Prisma query filters are dict-shaped
                 "in": sorted(requested_user_ids)
@@ -3098,7 +3180,9 @@ async def team_member_delete(
         key_val["user_id"] = data.user_id
     elif data.user_email is not None:
         key_val["user_email"] = data.user_email
-    existing_user_rows: Final = await UserRepository(prisma_client).table.find_many(where=key_val)
+    existing_user_rows: Final[Sequence[LiteLLM_UserTable] | None] = await UserRepository(prisma_client).table.find_many(
+        where=key_val
+    )
 
     if existing_user_rows is not None and (isinstance(existing_user_rows, list) and len(existing_user_rows) > 0):
         for existing_user in existing_user_rows:
@@ -3106,7 +3190,7 @@ async def team_member_delete(
             if data.team_id in existing_user.teams:
                 team_list = existing_user.teams
                 team_list.remove(data.team_id)
-                await UserRepository(prisma_client).table.update(
+                await _user_db(prisma_client).update(
                     where={
                         "user_id": existing_user.user_id,
                     },
@@ -3114,7 +3198,7 @@ async def team_member_delete(
                 )
 
     # Also clean up any existing team membership rows for this user and team
-    user_ids_to_delete: Final = set()
+    user_ids_to_delete: Final = set[str]()
     if data.user_id is not None:
         user_ids_to_delete.add(data.user_id)
     if existing_user_rows is not None and isinstance(existing_user_rows, list):
@@ -3123,9 +3207,7 @@ async def team_member_delete(
                 user_ids_to_delete.add(existing_user.user_id)
 
     for _uid in user_ids_to_delete:
-        await TeamMembershipRepository(prisma_client).table.delete_many(
-            where={"team_id": data.team_id, "user_id": _uid}
-        )
+        await _team_membership_db(prisma_client).delete_many(where={"team_id": data.team_id, "user_id": _uid})
 
     ## DELETE KEYS CREATED BY USER FOR THIS TEAM
     if user_ids_to_delete:
@@ -3134,9 +3216,7 @@ async def team_member_delete(
         )
 
         # Fetch keys before deletion to persist them
-        keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await VerificationTokenRepository(
-            prisma_client
-        ).table.find_many(
+        keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many(
             where={
                 "user_id": {"in": list(user_ids_to_delete)},
                 "team_id": data.team_id,
@@ -3151,7 +3231,7 @@ async def team_member_delete(
                 litellm_changed_by=None,
             )
 
-        await VerificationTokenRepository(prisma_client).table.delete_many(
+        await _tokens_db(prisma_client).delete_many(
             where={
                 "user_id": {"in": list(user_ids_to_delete)},
                 "team_id": data.team_id,
@@ -3311,7 +3391,7 @@ async def team_member_update(
 
     ### upsert new budget
     budget_patch: Final = _build_member_budget_patch(data)
-    async with prisma_client.db.tx() as tx:
+    async with prisma_client.tx() as tx:
         await _upsert_budget_and_membership(
             tx=tx,
             team_id=data.team_id,
@@ -3654,7 +3734,7 @@ async def delete_team(
         _persist_deleted_verification_tokens,
     )
 
-    keys_to_delete: list[LiteLLM_VerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many(
+    keys_to_delete: list[LiteLLM_VerificationToken] = await _tokens_db(prisma_client).find_many(
         where={"team_id": {"in": data.team_ids}}
     )
 
@@ -4469,7 +4549,7 @@ async def _get_keys_count_by_team(
     if not page_team_ids:
         return {}
 
-    grouped: Final = await VerificationTokenRepository(prisma_client).table.group_by(
+    grouped: Final = await _tokens_db(prisma_client).group_by(
         by=["team_id"],
         where={"team_id": {"in": page_team_ids}},
         count={"team_id": True},
@@ -4786,7 +4866,7 @@ async def _authorize_and_filter_teams(
 
     if allowed_org_ids is not None:
         # Org admin: query DB for teams in their orgs
-        org_teams: Final = await TeamRepository(prisma_client).table.find_many(
+        org_teams: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many(
             where={"organization_id": {"in": allowed_org_ids}},
             include={"litellm_model_table": True},
         )
@@ -4800,7 +4880,9 @@ async def _authorize_and_filter_teams(
         ]
     elif user_id:
         # Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays)
-        response: Final = await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True})
+        response: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many(
+            include={"litellm_model_table": True}
+        )
         return [
             team
             for team in response
@@ -4808,7 +4890,7 @@ async def _authorize_and_filter_teams(
         ]
     else:
         # Proxy admin: all teams
-        return list(await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True}))
+        return list(await _raw_team_db(TeamRepository(prisma_client)).find_many(include={"litellm_model_table": True}))
 
 
 @router.get("/team/list", tags=["team management"], dependencies=[Depends(user_api_key_auth)])
@@ -4860,7 +4942,7 @@ async def list_team(
                 _team_memberships.append(tm)
 
         # add all keys that belong to the team
-        keys = await VerificationTokenRepository(prisma_client).table.find_many(where={"team_id": team.team_id})
+        keys = await _tokens_db(prisma_client).find_many(where={"team_id": team.team_id})
 
         try:
             returned_responses.append(
@@ -4911,7 +4993,7 @@ async def get_paginated_teams(
         total_count: Final = await _team_db(prisma_client).count()
 
         # Get paginated teams
-        teams: Final = await TeamRepository(prisma_client).table.find_many(
+        teams: Final = await _team_db(prisma_client).find_many(
             skip=skip,
             take=page_size,
             order={"team_alias": "asc"},  # Sort by team_alias
@@ -4961,7 +5043,7 @@ async def ui_view_teams(
         skip: Final = (page - 1) * page_size
 
         # Build where conditions based on provided parameters
-        where_conditions: Final = {}
+        where_conditions: Final[_TeamUiViewFilters] = {}
 
         if team_id:
             where_conditions["team_id"] = {
@@ -4976,7 +5058,7 @@ async def ui_view_teams(
             }
 
         # Query users with pagination and filters
-        teams: Final = await TeamRepository(prisma_client).table.find_many(
+        teams: Final = await _team_db(prisma_client).find_many(
             where=where_conditions,
             skip=skip,
             take=page_size,
@@ -5166,13 +5248,13 @@ async def team_model_delete(
         )
 
     # Get current models list
-    current_models: Final = team_obj.models or []
+    current_models: Final[Sequence[str]] = team_obj.models or []
 
     # Remove specified models
     updated_models: Final = [m for m in current_models if m not in data.models]
 
     # Update team. See team_model_add for the rationale on `include`.
-    updated_team: Final = await TeamRepository(prisma_client).table.update(
+    updated_team: Final = await _team_db(prisma_client).update(
         where={"team_id": data.team_id},
         data={"models": updated_models},
         include={"object_permission": True},
@@ -5425,7 +5507,7 @@ async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissi
     BATCH_SIZE: Final = 500
 
     while True:
-        find_args: dict = {
+        find_args: _TeamFindManyArgs = {
             "take": BATCH_SIZE,
             "order": {"team_id": "asc"},
         }
@@ -5433,7 +5515,7 @@ async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissi
             find_args["cursor"] = {"team_id": cursor}
             find_args["skip"] = 1
 
-        teams = await TeamRepository(prisma_client).table.find_many(**find_args)
+        teams = await _team_db(prisma_client).find_many(**find_args)
 
         if not teams:
             break
@@ -5528,11 +5610,11 @@ async def get_team_daily_activity(
                     )
 
     ## Fetch team aliases and check team admin status
-    where_condition: Final = {}
+    where_condition: Final[_TeamIdInFilter] = {}
     if team_ids_list:
         where_condition["team_id"] = {"in": list(team_ids_list)}
-    team_aliases: Final = await TeamRepository(prisma_client).table.find_many(where=where_condition)
-    team_alias_metadata: Final = {t.team_id: {"team_alias": t.team_alias} for t in team_aliases}
+    team_aliases: Final = await _team_db(prisma_client).find_many(where=where_condition)
+    team_alias_metadata: Final = {t.team_id: {"team_alias": _as_object(t.team_alias)} for t in team_aliases}
 
     # Check if user is team admin or has /team/daily/activity permission
     # If not, filter by user's API keys.

litellm/proxy/management_endpoints/ui_sso.py

@@ -16,9 +16,22 @@ import json
 import os
 import re
 import secrets
+from collections.abc import Mapping, Sequence
 from copy import deepcopy
 from html import escape
-from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Final,
+    Literal,
+    NoReturn,
+    Optional,
+    Protocol,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
 from urllib.parse import parse_qs, urlencode, urlparse
 
 if TYPE_CHECKING:
@@ -155,6 +168,102 @@ _CLI_SSO_SECRET_KEY_FRAGMENTS: Final = frozenset(
     }
 )
 
+_DbRecordT: Final = TypeVar("_DbRecordT", covariant=True)
+
+
+class _PrismaTableActions(Protocol[_DbRecordT]):
+    async def find_unique(
+        self,
+        where: Mapping[str, object],
+    ) -> _DbRecordT | None: ...
+
+    async def find_first(
+        self,
+        where: Mapping[str, object] | None = None,
+    ) -> _DbRecordT | None: ...
+
+    async def find_many(
+        self,
+        where: Mapping[str, object] | None = None,
+    ) -> Sequence[_DbRecordT]: ...
+
+    async def update(
+        self,
+        where: Mapping[str, object],
+        data: Mapping[str, object],
+    ) -> _DbRecordT: ...
+
+    async def update_many(
+        self,
+        where: Mapping[str, object],
+        data: Mapping[str, object],
+    ) -> int: ...
+
+
+class _UserMetadataRow(Protocol):
+    @property
+    def metadata(self) -> Mapping[str, object] | None: ...
+
+
+class _HasUserMetadataTable(Protocol):
+    @property
+    def table(self) -> "_PrismaTableActions[_UserMetadataRow]": ...
+
+
+def _user_meta_db(repo: "_HasUserMetadataTable") -> "_PrismaTableActions[_UserMetadataRow]":
+    return repo.table
+
+
+class _SsoConfigRow(Protocol):
+    @property
+    def sso_settings(self) -> Mapping[str, object] | None: ...
+
+
+class _HasSsoConfigTable(Protocol):
+    @property
+    def table(self) -> "_PrismaTableActions[_SsoConfigRow]": ...
+
+
+def _sso_config_db(repo: "_HasSsoConfigTable") -> "_PrismaTableActions[_SsoConfigRow]":
+    return repo.table
+
+
+class _TeamDetailRow(Protocol):
+    def model_dump(self) -> Mapping[str, object]: ...
+
+
+class _HasTeamDetailTable(Protocol):
+    @property
+    def table(self) -> "_PrismaTableActions[_TeamDetailRow]": ...
+
+
+def _team_detail_db(repo: "_HasTeamDetailTable") -> "_PrismaTableActions[_TeamDetailRow]":
+    return repo.table
+
+
+class _CustomSsoCall(Protocol):
+    async def __call__(self, sso_response: object) -> SSOUserDefinedValues | None: ...
+
+
+class _ServicePrincipalAssignment(Protocol):
+    def get(self, key: str) -> str: ...
+
+
+class _ServicePrincipalPage(Protocol):
+    @overload
+    def get(
+        self,
+        key: Literal["value"],
+        default: Sequence["_ServicePrincipalAssignment"],
+    ) -> Sequence["_ServicePrincipalAssignment"]: ...
+
+    @overload
+    def get(self, key: Literal["@odata.nextLink"]) -> str | None: ...
+
+
+def _as_object(value: object) -> object:
+    return value
+
 
 def _hash_cli_sso_secret(secret: str) -> str:
     return hashlib.sha256(secret.encode("utf-8")).hexdigest()
@@ -256,7 +365,7 @@ def _get_cli_sso_flow_or_raise(login_id: str | None, cache: DualCache) -> dict:
         flow = cache.get_cache(key=cache_key)
     if isinstance(flow, str):
         try:
-            flow = json.loads(flow)
+            flow = _as_object(json.loads(flow))
         except ValueError:
             flow = None
     if not isinstance(flow, dict) or "poll_secret_hash" not in flow:
@@ -421,7 +530,7 @@ def _flatten_cli_sso_metadata_for_poll(
 
 def build_cli_sso_attribution_metadata(
     result: CustomOpenID | OpenID | dict,
-) -> dict[str, Any]:
+) -> dict[str, object]:
     """
     Build allowlisted, non-secret scalar attribution metadata from an SSO result.
 
@@ -432,7 +541,7 @@ def build_cli_sso_attribution_metadata(
     if not claim_map:
         return {}
 
-    metadata: Final[dict[str, Any]] = {}
+    metadata: Final[dict[str, object]] = {}
     for source_claim, dest_key in claim_map:
         if not _is_safe_cli_sso_metadata_dest_key(dest_key):
             verbose_proxy_logger.debug("Skipping unsafe CLI SSO metadata destination key: %s", dest_key)
@@ -474,14 +583,14 @@ def _merge_cli_sso_attribution_metadata(
 async def _persist_cli_sso_user_metadata(
     prisma_client: PrismaClient,
     user_id: str,
-    attribution_metadata: dict[str, Any],
+    attribution_metadata: dict[str, object],
 ) -> None:
     if not attribution_metadata:
         return
 
     try:
-        user_row: Final = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id})
-        existing_metadata: dict[str, Any] = {}
+        user_row: Final = await _user_meta_db(UserRepository(prisma_client)).find_unique(where={"user_id": user_id})
+        existing_metadata: dict[str, object] = {}
         if user_row is not None:
             row_metadata: Final = user_row.metadata
             if isinstance(row_metadata, dict):
@@ -491,7 +600,7 @@ async def _persist_cli_sso_user_metadata(
             existing_metadata=existing_metadata,
             attribution_metadata=attribution_metadata,
         )
-        await UserRepository(prisma_client).table.update_many(
+        await _user_meta_db(UserRepository(prisma_client)).update_many(
             where={"user_id": user_id},
             data={"metadata": merged_metadata},
         )
@@ -1104,7 +1213,7 @@ def generic_response_convertor(
                 )
 
     # Build extra_fields dict from GENERIC_USER_EXTRA_ATTRIBUTES if specified
-    extra_fields: dict[str, Any] | None = None
+    extra_fields: dict[str, object] | None = None
     if generic_user_extra_attributes:
         extra_fields = {}
         for attr_name in generic_user_extra_attributes.split(","):
@@ -1193,7 +1302,9 @@ async def _setup_team_mappings() -> Optional["TeamMappings"]:
 
         prisma_client: Final = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy")
 
-        sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"})
+        sso_db_record: Final = await _sso_config_db(SSOConfigRepository(prisma_client)).find_unique(
+            where={"id": "sso_config"}
+        )
 
         if sso_db_record and sso_db_record.sso_settings:
             sso_settings_dict: Final = dict(sso_db_record.sso_settings)
@@ -1225,7 +1336,9 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]:
 
         prisma_client: Final = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy")
 
-        sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"})
+        sso_db_record: Final = await _sso_config_db(SSOConfigRepository(prisma_client)).find_unique(
+            where={"id": "sso_config"}
+        )
 
         if sso_db_record and sso_db_record.sso_settings:
             sso_settings_dict: Final = dict(sso_db_record.sso_settings)
@@ -1273,7 +1386,7 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]:
     return role_mappings
 
 
-def _parse_generic_sso_headers() -> dict:
+def _parse_generic_sso_headers() -> dict[str, str]:
     """Parse comma-separated GENERIC_SSO_HEADERS env var into a dict."""
     raw: Final = os.getenv("GENERIC_SSO_HEADERS", None)
     if raw is None:
@@ -1677,7 +1790,7 @@ def _build_sso_user_update_data(
     result: Union["CustomOpenID", OpenID, dict] | None,
     user_email: str | None,
     user_id: str | None,
-) -> dict:
+) -> dict[str, object]:
     """
     Build the update data dictionary for SSO user upsert.
 
@@ -1689,7 +1802,7 @@ def _build_sso_user_update_data(
     Returns:
         dict: Update data containing user_email and optionally user_role if valid
     """
-    update_data: Final[dict] = {"user_email": normalize_email(user_email)}
+    update_data: Final[dict[str, object]] = {"user_email": normalize_email(user_email)}
 
     # Get SSO role from result and include if valid
     sso_role: Final = getattr(result, "user_role", None)
@@ -1740,7 +1853,7 @@ async def _sync_user_role_from_jwt_role_map(
 
     # Update existing DB record if role differs
     if user_info is not None and user_info.user_role != mapped_role.value:
-        await UserRepository(prisma_client).table.update(
+        await _user_meta_db(UserRepository(prisma_client)).update(
             where={"user_id": user_info.user_id},
             data={"user_role": mapped_role.value},
         )
@@ -1796,7 +1909,7 @@ async def check_and_update_if_proxy_admin_id(user_role: str, user_id: str, prism
             return user_role
 
         if prisma_client:
-            await UserRepository(prisma_client).table.update(
+            await _user_meta_db(UserRepository(prisma_client)).update(
                 where={"user_id": user_id},
                 data={"user_role": LitellmUserRoles.PROXY_ADMIN.value},
             )
@@ -2016,10 +2129,11 @@ async def _build_cli_sso_user_defined_values(
 ) -> SSOUserDefinedValues | None:
     from litellm.proxy.proxy_server import user_custom_sso
 
+    custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso
     user_id: Final = parsed_openid_result.get("user_id")
-    if user_custom_sso is not None:
-        if inspect.iscoroutinefunction(user_custom_sso):
-            return await user_custom_sso(result)
+    if custom_sso_handler is not None:
+        if inspect.iscoroutinefunction(custom_sso_handler):
+            return await custom_sso_handler(result)
         raise ValueError("user_custom_sso must be a coroutine function")
     if user_id is None:
         return None
@@ -2035,12 +2149,14 @@ async def _build_cli_sso_user_defined_values(
 
 async def _fetch_cli_sso_team_details(
     prisma_client: PrismaClient,
-    teams: list[str],
-) -> list[dict[str, Any]]:
-    team_details: Final[list[dict[str, Any]]] = []
+    teams: Sequence[str],
+) -> list[dict[str, object]]:
+    team_details: Final[list[dict[str, object]]] = []
     try:
         if teams:
-            prisma_teams: Final = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": teams}})
+            prisma_teams: Final = await _team_detail_db(TeamRepository(prisma_client)).find_many(
+                where={"team_id": {"in": teams}}
+            )
             for team_row in prisma_teams:
                 team_dict = team_row.model_dump()
                 team_details.append(
@@ -2257,12 +2373,12 @@ async def cli_poll_key(
                 verbose_proxy_logger.info("Returning teams list for user %s to select from: %s", user_id, user_teams)
                 # Best-effort construction of team_details if it wasn't
                 # already cached for some reason.
-                team_details_response: list[dict[str, Any]] | None = None
+                team_details_response: list[dict[str, object]] | None = None
                 if isinstance(user_team_details, list) and user_team_details:
                     team_details_response = user_team_details
                 elif user_teams:
                     team_details_response = [{"team_id": t, "team_alias": None} for t in user_teams]
-                poll_response: dict[str, Any] = {
+                poll_response: dict[str, object] = {
                     "status": "ready",
                     "user_id": user_id,
                     "teams": user_teams,
@@ -2997,7 +3113,9 @@ class SSOAuthenticationHandler:
                     user_id=user_id,
                 )
 
-                await UserRepository(prisma_client).table.update_many(where={"user_id": user_id}, data=update_data)
+                await _user_meta_db(UserRepository(prisma_client)).update_many(
+                    where={"user_id": user_id}, data=update_data
+                )
             else:
                 verbose_proxy_logger.info("user not in DB, inserting user into LiteLLM DB")
                 # user not in DB, insert User into LiteLLM DB
@@ -3089,7 +3207,9 @@ class SSOAuthenticationHandler:
                 code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             )
         try:
-            team_obj: Final = await TeamRepository(prisma_client).table.find_first(where={"team_id": litellm_team_id})
+            team_obj: Final = await _team_detail_db(TeamRepository(prisma_client)).find_first(
+                where={"team_id": litellm_team_id}
+            )
             verbose_proxy_logger.debug("Team object: %s", team_obj)
 
             # only create a new team if it doesn't exist
@@ -3278,9 +3398,10 @@ class SSOAuthenticationHandler:
         # But if it is, we want their models preferences
         user_defined_values: SSOUserDefinedValues | None = None
 
-        if user_custom_sso is not None:
-            if inspect.iscoroutinefunction(user_custom_sso):
-                user_defined_values = await user_custom_sso(result)
+        custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso
+        if custom_sso_handler is not None:
+            if inspect.iscoroutinefunction(custom_sso_handler):
+                user_defined_values = await custom_sso_handler(result)
             else:
                 raise ValueError("user_custom_sso must be a coroutine function")
         elif user_id is not None:
@@ -3448,7 +3569,7 @@ class SSOAuthenticationHandler:
             dict: Token exchange parameters
         """
         # Prepare token exchange parameters (may add code_verifier: str later)
-        token_params: Final[dict[str, Any]] = {"include_client_id": generic_include_client_id}
+        token_params: Final[dict[str, object]] = {"include_client_id": generic_include_client_id}
 
         # Retrieve PKCE code_verifier if PKCE was used in authorization.
         # Gate on GENERIC_CLIENT_USE_PKCE to avoid an unnecessary Redis round-trip
@@ -3663,7 +3784,7 @@ class SSOAuthenticationHandler:
         access_token string. Raises ProxyException on any validation failure.
         """
         try:
-            token_response_raw: Final = response.json()
+            token_response_raw: Final[object] = _as_object(response.json())
         except Exception as json_err:
             verbose_proxy_logger.error(
                 "Failed to parse token response as JSON: %s. Body: %s",
@@ -4253,7 +4374,7 @@ class MicrosoftSSOHandler:
 
         while next_link is not None and page_count < MicrosoftSSOHandler.MAX_GRAPH_API_PAGES:
             response = await async_client.get(next_link, headers=headers)
-            response_json = response.json()
+            response_json: _ServicePrincipalPage = response.json()
             verbose_proxy_logger.debug("Response from service principal app role assigned to: %s", response_json)
 
             for _object in response_json.get("value", []):

litellm/proxy/openai_files_endpoints/common_utils.py

@@ -4,7 +4,7 @@ import re
 from collections.abc import Mapping
 from dataclasses import dataclass, field
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Final, Literal, Optional
+from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, runtime_checkable
 
 from litellm.repositories.table_repositories import (
     ManagedFileRepository,
@@ -22,6 +22,21 @@ if TYPE_CHECKING:
     from litellm.types.utils import LiteLLMBatch
 
 
+@runtime_checkable
+class ManagedResourceAccessChecker(Protocol):
+    async def can_user_call_unified_file_id(
+        self,
+        unified_file_id: str,
+        user_api_key_dict: "UserAPIKeyAuth",
+    ) -> bool: ...
+
+    async def can_user_call_unified_object_id(
+        self,
+        unified_object_id: str,
+        user_api_key_dict: "UserAPIKeyAuth",
+    ) -> bool: ...
+
+
 def _is_base64_encoded_unified_file_id(b64_uid: str) -> str | Literal[False]:
     # Ensure b64_uid is a string and not a mock object
     if not isinstance(b64_uid, str):
@@ -881,6 +896,65 @@ def validate_managed_files_requirement(
         )
 
 
+async def validate_managed_id_requirement(
+    resource_id: str | None,
+    resource_kind: Literal["file", "batch", "fine-tuning job"],
+    user_api_key_dict: "UserAPIKeyAuth",
+    managed_files_obj: object | None,
+) -> None:
+    """
+    Enforce proxy-level managed resources on every route that accepts a provider-issued id
+    when ``litellm.require_managed_files`` is enabled, and authenticate managed ids against
+    the caller's stored ownership record.
+
+    Ownership is only recorded for LiteLLM managed ids, so a raw provider id is forwarded to the
+    provider under shared credentials without any tenant check; knowing another tenant's provider
+    id would be enough to read, reuse, or destroy the object behind it.
+
+    Raises:
+        HTTPException: 400 for a raw id, 403 for an inaccessible managed id, or 500 when
+            ownership validation is unavailable.
+    """
+    from fastapi import HTTPException
+
+    import litellm
+
+    if litellm.require_managed_files is not True:
+        return
+
+    if not resource_id:
+        return
+
+    if not _is_base64_encoded_unified_file_id(resource_id):
+        raise HTTPException(
+            status_code=400,
+            detail=(
+                f"Raw provider {resource_kind} ids cannot be used when require_managed_files is enabled in "
+                f"litellm_settings. Use the LiteLLM managed {resource_kind} id returned when the "
+                f"{resource_kind} was created."
+            ),
+        )
+
+    if not isinstance(managed_files_obj, ManagedResourceAccessChecker):
+        raise HTTPException(
+            status_code=500,
+            detail="Managed resource ownership validation is unavailable.",
+        )
+
+    can_access: Final = (
+        await managed_files_obj.can_user_call_unified_file_id(resource_id, user_api_key_dict)
+        if resource_kind == "file"
+        else await managed_files_obj.can_user_call_unified_object_id(resource_id, user_api_key_dict)
+    )
+    if can_access:
+        return
+
+    raise HTTPException(
+        status_code=403,
+        detail=f"The caller does not have access to this managed {resource_kind} id.",
+    )
+
+
 def _extract_model_param(request: "Request", request_body: dict) -> str | None:
     """
     Extract model parameter from request.

litellm/proxy/openai_files_endpoints/files_endpoints.py

@@ -50,6 +50,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
     handle_model_based_routing,
     prepare_data_with_credentials,
     validate_managed_files_requirement,
+    validate_managed_id_requirement,
 )
 from litellm.proxy.utils import ProxyLogging, is_known_model
 from litellm.repositories.table_repositories import ManagedFileRepository
@@ -612,6 +613,13 @@ async def get_file_content(
 
     data: dict = {"file_id": file_id}
     try:
+        await validate_managed_id_requirement(
+            resource_id=file_id,
+            resource_kind="file",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
+
         # Include original request and headers in the data
         base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data)
         (
@@ -908,6 +916,13 @@ async def get_file(
 
     data: dict = {"file_id": file_id}
     try:
+        await validate_managed_id_requirement(
+            resource_id=file_id,
+            resource_kind="file",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
+
         custom_llm_provider: Final = (
             provider
             or get_custom_llm_provider_from_request_headers(request=request)
@@ -1098,6 +1113,13 @@ async def delete_file(
 
     data: dict = {"file_id": file_id}
     try:
+        await validate_managed_id_requirement(
+            resource_id=file_id,
+            resource_kind="file",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
+
         custom_llm_provider: Final = (
             provider
             or get_custom_llm_provider_from_request_headers(request=request)

litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py

@@ -4,11 +4,11 @@ import json
 import os
 from collections import Counter
 from collections.abc import Mapping
-from typing import Any, Final
+from typing import Any, Final, Protocol, TypeVar
 from urllib.parse import urlparse
 
 from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile
-from pydantic import ConfigDict, ValidationError, create_model
+from pydantic import ConfigDict, JsonValue, ValidationError, create_model
 from pydantic.fields import FieldInfo
 
 import litellm
@@ -36,6 +36,73 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
 
 router: Final = APIRouter()
 
+_DbRecordT: Final = TypeVar("_DbRecordT", covariant=True)
+
+
+class _PrismaTableActions(Protocol[_DbRecordT]):
+    async def find_unique(self, where: Mapping[str, object]) -> _DbRecordT | None: ...
+
+    async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _DbRecordT: ...
+
+    async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> _DbRecordT: ...
+
+
+class _SsoSettingsMappingRow(Protocol):
+    @property
+    def sso_settings(self) -> Mapping[str, object] | None: ...
+
+
+class _HasSsoSettingsMappingTable(Protocol):
+    @property
+    def table(self) -> _PrismaTableActions[_SsoSettingsMappingRow]: ...
+
+
+def _sso_settings_mapping_db(repo: _HasSsoSettingsMappingTable) -> _PrismaTableActions[_SsoSettingsMappingRow]:
+    return repo.table
+
+
+class _StoredSsoSettingsRow(Protocol):
+    @property
+    def sso_settings(self) -> object: ...
+
+
+class _HasStoredSsoSettingsTable(Protocol):
+    @property
+    def table(self) -> _PrismaTableActions[_StoredSsoSettingsRow]: ...
+
+
+def _stored_sso_settings_db(repo: _HasStoredSsoSettingsTable) -> _PrismaTableActions[_StoredSsoSettingsRow]:
+    return repo.table
+
+
+class _UiSettingsRow(Protocol):
+    @property
+    def ui_settings(self) -> str | Mapping[str, JsonValue] | None: ...
+
+
+class _HasUiSettingsTable(Protocol):
+    @property
+    def table(self) -> _PrismaTableActions[_UiSettingsRow]: ...
+
+
+def _ui_settings_db(repo: _HasUiSettingsTable) -> _PrismaTableActions[_UiSettingsRow]:
+    return repo.table
+
+
+class _ConfigParamRow(Protocol):
+    @property
+    def param_value(self) -> str | Mapping[str, object] | None: ...
+
+
+class _HasConfigParamTable(Protocol):
+    @property
+    def table(self) -> _PrismaTableActions[_ConfigParamRow]: ...
+
+
+def _config_param_db(repo: _HasConfigParamTable) -> _PrismaTableActions[_ConfigParamRow]:
+    return repo.table
+
+
 # Maps each UIThemeConfig field to the env var the UI branding path reads it
 # from. /update/ui_theme_settings writes both the stored ui_theme_config and
 # these env vars, so /get/ui_theme_settings resolves the same env vars to
@@ -54,7 +121,7 @@ def _is_public_http_url(value: str | None) -> bool:
     return parsed.scheme in ("http", "https") and bool(parsed.netloc)
 
 
-def _resolve_ui_theme_field(stored_values: Mapping[str, Any], field_name: str) -> str | None:
+def _resolve_ui_theme_field(stored_values: Mapping[str, object], field_name: str) -> str | None:
     """Resolve one UI theme field to the value the branding path actually uses.
 
     The stored ui_theme_config wins; a field absent or blank there falls back to
@@ -263,7 +330,7 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS: Final = [
 # include generics like ``Optional[int]`` / ``List[str]`` that are not
 # instances of ``type`` — so tightening this to ``type`` would reject
 # valid inputs.
-_EXTRA_UI_SETTINGS_FIELDS: Final[dict[str, tuple[Any, FieldInfo]]] = {}
+_EXTRA_UI_SETTINGS_FIELDS: Final[dict[str, tuple[object, FieldInfo]]] = {}
 
 # Settings OSS knows about as enterprise-gated. If a caller sends one of
 # these keys and no extension package has registered it, the PATCH
@@ -275,7 +342,7 @@ _ENTERPRISE_ONLY_UI_SETTINGS: Final[set[str]] = {"enable_projects_ui"}
 _EFFECTIVE_UI_SETTINGS_CLASS: type[UISettings] | None = None
 
 
-def register_extra_ui_setting(name: str, annotation: Any, field: FieldInfo) -> None:
+def register_extra_ui_setting(name: str, annotation: object, field: FieldInfo) -> None:
     """Register an additional UI settings field contributed by an extension package.
 
     ``field`` must be a ``FieldInfo`` instance — construct it directly
@@ -470,7 +537,7 @@ async def delete_allowed_ip(
 
 async def _get_settings_with_schema(
     settings_key: str,
-    settings_class: Any,
+    settings_class: type[BaseModel],
     config: dict,
 ) -> dict:
     """
@@ -842,7 +909,9 @@ async def get_sso_settings():
     # Resolve the effective SSO config: the stored row wins, else the process
     # environment, else each field's default. Unlike the legacy read path this
     # does not write os.environ; a GET has no business mutating the environment.
-    sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"})
+    sso_db_record: Final = await _sso_settings_mapping_db(SSOConfigRepository(prisma_client)).find_unique(
+        where={"id": "sso_config"}
+    )
     sso_db_settings: Final = dict(sso_db_record.sso_settings) if sso_db_record and sso_db_record.sso_settings else None
     resolved: Final = resolve_sso_config(sso_db_settings, os.environ)
 
@@ -914,8 +983,10 @@ async def update_sso_settings(
     # before-snapshot has the same shape as after_value, and rely on
     # create_config_audit_log's secret-name redaction to mask the
     # *_client_secret fields before the audit row is written.
-    existing_sso_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"})
-    before_sso_data: dict[str, Any] | None = None
+    existing_sso_record: Final = await _stored_sso_settings_db(SSOConfigRepository(prisma_client)).find_unique(
+        where={"id": "sso_config"}
+    )
+    before_sso_data: dict[str, JsonValue] | None = None
     if existing_sso_record and existing_sso_record.sso_settings:
         stored = existing_sso_record.sso_settings
         if isinstance(stored, str):
@@ -948,7 +1019,7 @@ async def update_sso_settings(
     encrypted_sso_data: Final = proxy_config._encrypt_env_variables(environment_variables=sso_data)
 
     # Save to dedicated SSO table
-    await SSOConfigRepository(prisma_client).table.upsert(
+    await _stored_sso_settings_db(SSOConfigRepository(prisma_client)).upsert(
         where={"id": "sso_config"},
         data={
             "create": {
@@ -974,7 +1045,7 @@ async def update_sso_settings(
 
     # Remove SSO-related env vars from config.environment_variables
     try:
-        env_var_entry: Final = await ConfigRepository(prisma_client).table.find_unique(
+        env_var_entry: Final = await _config_param_db(ConfigRepository(prisma_client)).find_unique(
             where={"param_name": "environment_variables"}
         )
 
@@ -982,7 +1053,7 @@ async def update_sso_settings(
         if env_var_entry is not None:
             if env_var_entry.param_value is not None:
                 if isinstance(env_var_entry.param_value, str):
-                    environment_variables = json.loads(env_var_entry.param_value)
+                    environment_variables: Mapping[str, object] = json.loads(env_var_entry.param_value)
                 else:
                     environment_variables = dict(env_var_entry.param_value)
             else:
@@ -993,7 +1064,7 @@ async def update_sso_settings(
                 key: value for key, value in environment_variables.items() if key not in env_vars_to_remove
             }
 
-            await ConfigRepository(prisma_client).table.update(
+            await _config_param_db(ConfigRepository(prisma_client)).update(
                 where={"param_name": "environment_variables"},
                 data={
                     "param_value": json.dumps(filtered_env_vars, default=str),
@@ -1239,8 +1310,10 @@ async def get_ui_settings_cached() -> dict[str, Any]:
     if prisma_client is None:
         return {}
 
-    db_record: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"})
-    ui_settings: dict[str, Any] = {}
+    db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique(
+        where={"id": "ui_settings"}
+    )
+    ui_settings: dict[str, JsonValue] = {}
     if db_record and db_record.ui_settings:
         raw: Final = db_record.ui_settings
         ui_settings = json.loads(raw) if isinstance(raw, str) else dict(raw)
@@ -1272,9 +1345,11 @@ async def get_ui_settings():
             detail={"error": "Database not connected. Please connect a database."},
         )
 
-    ui_settings: dict[str, Any] = {}
+    ui_settings: Mapping[str, JsonValue] = {}
 
-    db_record: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"})
+    db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique(
+        where={"id": "ui_settings"}
+    )
 
     if db_record and db_record.ui_settings:
         ui_settings_json: Final = db_record.ui_settings
@@ -1300,7 +1375,7 @@ async def get_ui_settings():
     await user_api_key_cache.async_set_cache(key=UI_SETTINGS_CACHE_KEY, value=ui_settings, ttl=UI_SETTINGS_CACHE_TTL)
 
     # Build config-like object for schema helper
-    config: Final[dict[str, Any]] = {"litellm_settings": {"ui_settings": ui_settings}}
+    config: Final[dict[str, object]] = {"litellm_settings": {"ui_settings": ui_settings}}
 
     return await _get_settings_with_schema(
         settings_key="ui_settings",
@@ -1315,7 +1390,7 @@ async def get_ui_settings():
     dependencies=[Depends(user_api_key_auth)],
 )
 async def update_ui_settings(
-    settings_body: dict[str, Any] = Body(...),
+    settings_body: dict[str, object] = Body(...),
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
 ):
     """
@@ -1352,7 +1427,7 @@ async def update_ui_settings(
         raise HTTPException(status_code=422, detail=e.errors())
 
     # Only include fields the caller actually sent (not Pydantic defaults).
-    settings_dict: Final = settings.model_dump(exclude_unset=True)
+    settings_dict: Final[Mapping[str, JsonValue]] = settings.model_dump(exclude_unset=True)
 
     # Reject enterprise-only settings up front so the caller gets a clear
     # signal instead of a silent drop.
@@ -1373,15 +1448,17 @@ async def update_ui_settings(
 
     # Merge with existing persisted settings so a partial PATCH doesn't
     # overwrite fields the caller didn't send.
-    existing: dict = {}
-    db_existing: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"})
+    existing: dict[str, JsonValue] = {}
+    db_existing: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique(
+        where={"id": "ui_settings"}
+    )
     if db_existing and db_existing.ui_settings:
         raw: Final = db_existing.ui_settings
         existing = json.loads(raw) if isinstance(raw, str) else dict(raw)
 
     ui_settings: Final = {**existing, **incoming}
 
-    await UISettingsRepository(prisma_client).table.upsert(
+    await _ui_settings_db(UISettingsRepository(prisma_client)).upsert(
         where={"id": "ui_settings"},
         data={
             "create": {

litellm/proxy/vector_store_endpoints/management_endpoints.py

@@ -10,10 +10,17 @@ All /vector_store management endpoints
 
 import copy
 import json
-from typing import Any, Final
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, Final, Protocol
 
 from fastapi import APIRouter, Depends, HTTPException
 
+if TYPE_CHECKING:
+    from prisma.models import LiteLLM_ManagedVectorStoresTable as _VectorStoreRow
+
+    from litellm.proxy.utils import PrismaClient
+    from litellm.router import Router
+
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.in_memory_cache import InMemoryCache
@@ -43,6 +50,25 @@ from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
 
 router: Final = APIRouter()
 
+
+class _VectorStoreTableActions(Protocol):
+    async def find_unique(self, where: Mapping[str, str]) -> "_VectorStoreRow | None": ...
+
+    async def create(self, data: Mapping[str, object]) -> "_VectorStoreRow": ...
+
+    async def update(self, where: Mapping[str, str], data: Mapping[str, object]) -> "_VectorStoreRow": ...
+
+    async def delete(self, where: Mapping[str, str]) -> "_VectorStoreRow | None": ...
+
+
+def _vector_store_table(prisma_client: "PrismaClient") -> _VectorStoreTableActions:
+    return ManagedVectorStoresRepository(prisma_client).table
+
+
+def _row_to_vector_store(row: "_VectorStoreRow") -> LiteLLM_ManagedVectorStore:
+    return LiteLLM_ManagedVectorStore(**row.model_dump())
+
+
 _LITELLM_PARAMS_MASKER: Final = SensitiveDataMasker()
 
 
@@ -117,22 +143,20 @@ def _redact_sensitive_litellm_params(litellm_params: Any, _depth: int = 0) -> An
 async def _fetch_and_authorize_vector_store(
     vector_store_id: str,
     user_api_key_dict: UserAPIKeyAuth,
-    prisma_client: Any,
+    prisma_client: "PrismaClient",
 ) -> "LiteLLM_ManagedVectorStore":
     """
     Look up a vector store by id and confirm the caller can access it.
     Raises HTTPException(404) on miss and HTTPException(403) on access
     denial.
     """
-    row: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique(
-        where={"vector_store_id": vector_store_id}
-    )
+    row: Final = await _vector_store_table(prisma_client).find_unique(where={"vector_store_id": vector_store_id})
     if row is None:
         raise HTTPException(
             status_code=404,
             detail=f"Vector store with ID {vector_store_id} not found",
         )
-    typed: Final = LiteLLM_ManagedVectorStore(**row.model_dump())
+    typed: Final = _row_to_vector_store(row)
     if not await _check_vector_store_access(typed, user_api_key_dict):
         raise HTTPException(
             status_code=403,
@@ -141,7 +165,7 @@ async def _fetch_and_authorize_vector_store(
     return typed
 
 
-def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> dict[str, Any] | None:
+def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> dict[str, object] | None:
     """
     Resolve embedding config from router's config-defined models.
 
@@ -177,7 +201,7 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d
                 litellm_params = deployment.litellm_params
 
                 # Build embedding config from model params
-                embedding_config: dict[str, Any] = {}
+                embedding_config: dict[str, object] = {}
 
                 # Extract api_key
                 api_key = getattr(litellm_params, "api_key", None)
@@ -217,7 +241,9 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d
     return None
 
 
-async def _resolve_embedding_config_from_db(embedding_model: str, prisma_client) -> dict[str, Any] | None:
+async def _resolve_embedding_config_from_db(
+    embedding_model: str, prisma_client: "PrismaClient"
+) -> dict[str, object] | None:
     """
     Resolve embedding config from database model configuration.
 
@@ -307,7 +333,9 @@ async def _resolve_embedding_config_from_db(embedding_model: str, prisma_client)
     return None
 
 
-async def _resolve_embedding_config(embedding_model: str, prisma_client, llm_router=None) -> dict[str, Any] | None:
+async def _resolve_embedding_config(
+    embedding_model: str, prisma_client: "PrismaClient | None", llm_router: "Router | None" = None
+) -> dict[str, object] | None:
     """
     Resolve embedding config from either router (config-defined) or database models.
 
@@ -388,7 +416,7 @@ async def _check_vector_store_access(
 async def create_vector_store_in_db(
     vector_store_id: str,
     custom_llm_provider: str,
-    prisma_client,
+    prisma_client: "PrismaClient | None",
     vector_store_name: str | None = None,
     vector_store_description: str | None = None,
     vector_store_metadata: dict | None = None,
@@ -417,7 +445,7 @@ async def create_vector_store_in_db(
         raise HTTPException(status_code=500, detail="Database not connected")
 
     # Check if vector store already exists
-    existing_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique(
+    existing_vector_store: Final = await _vector_store_table(prisma_client).find_unique(
         where={"vector_store_id": vector_store_id}
     )
     if existing_vector_store is not None:
@@ -427,7 +455,7 @@ async def create_vector_store_in_db(
         )
 
     # Prepare data for database
-    data_to_create: Final[dict[str, Any]] = {
+    data_to_create: Final[dict[str, object]] = {
         "vector_store_id": vector_store_id,
         "custom_llm_provider": custom_llm_provider,
     }
@@ -463,9 +491,9 @@ async def create_vector_store_in_db(
         data_to_create["litellm_params"] = safe_dumps({})
 
     # Create in database
-    _new_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.create(data=data_to_create)
+    _new_vector_store: Final = await _vector_store_table(prisma_client).create(data=data_to_create)
 
-    new_vector_store: Final[LiteLLM_ManagedVectorStore] = LiteLLM_ManagedVectorStore(**_new_vector_store.model_dump())
+    new_vector_store: Final[LiteLLM_ManagedVectorStore] = _row_to_vector_store(_new_vector_store)
 
     # Add vector store to registry
     if litellm.vector_store_registry is not None:
@@ -682,12 +710,12 @@ async def delete_vector_store(
         memory_vector_store_exists = False
         vector_store_to_check = None
 
-        existing_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique(
+        existing_vector_store: Final = await _vector_store_table(prisma_client).find_unique(
             where={"vector_store_id": data.vector_store_id}
         )
         if existing_vector_store is not None:
             db_vector_store_exists = True
-            vector_store_to_check = LiteLLM_ManagedVectorStore(**existing_vector_store.model_dump())
+            vector_store_to_check = _row_to_vector_store(existing_vector_store)
 
         # Check in-memory registry
         if litellm.vector_store_registry is not None:
@@ -715,9 +743,7 @@ async def delete_vector_store(
 
         # Delete from database if exists
         if db_vector_store_exists:
-            await ManagedVectorStoresRepository(prisma_client).table.delete(
-                where={"vector_store_id": data.vector_store_id}
-            )
+            await _vector_store_table(prisma_client).delete(where={"vector_store_id": data.vector_store_id})
 
         # Delete from in-memory registry if exists
         if memory_vector_store_exists and litellm.vector_store_registry is not None:
@@ -829,7 +855,7 @@ async def update_vector_store(
 
     try:
         update_data: Final = data.model_dump(exclude_unset=True)
-        vector_store_id: Final = update_data.pop("vector_store_id")
+        vector_store_id: Final[str] = update_data.pop("vector_store_id")
 
         # Per-store access control: anyone authenticated who passes the
         # premium-feature gate could otherwise update *any* vector store —
@@ -857,12 +883,12 @@ async def update_vector_store(
             update_data["litellm_params"] = safe_dumps(litellm_params_dict)
 
         # Update in database
-        updated: Final = await ManagedVectorStoresRepository(prisma_client).table.update(
+        updated: Final = await _vector_store_table(prisma_client).update(
             where={"vector_store_id": vector_store_id},
             data=update_data,
         )
 
-        updated_vs: Final = LiteLLM_ManagedVectorStore(**updated.model_dump())
+        updated_vs: Final = _row_to_vector_store(updated)
 
         # Immediately update in-memory registry to keep it in sync
         if litellm.vector_store_registry is not None:

litellm/proxy/vector_store_files_endpoints/endpoints.py

@@ -30,10 +30,12 @@ if TYPE_CHECKING:
 router: Final = APIRouter()
 
 
-def _update_request_data_with_managed_file_id(
+async def _update_request_data_with_managed_file_id(
     data: dict,
     file_id: str,
     request: Request,
+    user_api_key_dict: UserAPIKeyAuth,
+    managed_files_obj: object | None,
     llm_router: Optional["Router"] = None,
 ) -> tuple[dict, str | None]:
     """
@@ -65,6 +67,16 @@ def _update_request_data_with_managed_file_id(
         is_base64_encoded_unified_id,
         parse_unified_id,
     )
+    from litellm.proxy.openai_files_endpoints.common_utils import (
+        validate_managed_id_requirement,
+    )
+
+    await validate_managed_id_requirement(
+        resource_id=file_id,
+        resource_kind="file",
+        user_api_key_dict=user_api_key_dict,
+        managed_files_obj=managed_files_obj,
+    )
 
     # First, check if this is a unified managed file ID (base64 encoded)
     decoded_id: Final = is_base64_encoded_unified_id(file_id)
@@ -509,8 +521,13 @@ async def vector_store_file_create(
     # Handle managed file IDs if present in request body
     original_managed_file_id = None
     if "file_id" in data:
-        data, original_managed_file_id = _update_request_data_with_managed_file_id(
-            data=data, file_id=data["file_id"], request=request, llm_router=llm_router
+        data, original_managed_file_id = await _update_request_data_with_managed_file_id(
+            data=data,
+            file_id=data["file_id"],
+            request=request,
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+            llm_router=llm_router,
         )
 
     # Then handle managed vector store IDs
@@ -707,8 +724,13 @@ async def vector_store_file_retrieve(
     )
 
     # Handle managed file IDs first
-    data, original_managed_file_id = _update_request_data_with_managed_file_id(
-        data=data, file_id=file_id, request=request, llm_router=llm_router
+    data, original_managed_file_id = await _update_request_data_with_managed_file_id(
+        data=data,
+        file_id=file_id,
+        request=request,
+        user_api_key_dict=user_api_key_dict,
+        managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        llm_router=llm_router,
     )
 
     # Then handle managed vector store IDs
@@ -809,8 +831,13 @@ async def vector_store_file_content(
     )
 
     # Handle managed file IDs first
-    data, original_managed_file_id = _update_request_data_with_managed_file_id(
-        data=data, file_id=file_id, request=request, llm_router=llm_router
+    data, original_managed_file_id = await _update_request_data_with_managed_file_id(
+        data=data,
+        file_id=file_id,
+        request=request,
+        user_api_key_dict=user_api_key_dict,
+        managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        llm_router=llm_router,
     )
 
     # Then handle managed vector store IDs
@@ -911,8 +938,13 @@ async def vector_store_file_update(
     )
 
     # Handle managed file IDs first
-    data, original_managed_file_id = _update_request_data_with_managed_file_id(
-        data=data, file_id=file_id, request=request, llm_router=llm_router
+    data, original_managed_file_id = await _update_request_data_with_managed_file_id(
+        data=data,
+        file_id=file_id,
+        request=request,
+        user_api_key_dict=user_api_key_dict,
+        managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        llm_router=llm_router,
     )
 
     # Then handle managed vector store IDs
@@ -1013,8 +1045,13 @@ async def vector_store_file_delete(
     )
 
     # Handle managed file IDs first
-    data, original_managed_file_id = _update_request_data_with_managed_file_id(
-        data=data, file_id=file_id, request=request, llm_router=llm_router
+    data, original_managed_file_id = await _update_request_data_with_managed_file_id(
+        data=data,
+        file_id=file_id,
+        request=request,
+        user_api_key_dict=user_api_key_dict,
+        managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        llm_router=llm_router,
     )
 
     # Then handle managed vector store IDs

litellm/responses/litellm_completion_transformation/transformation.py

@@ -4,8 +4,8 @@ Handles transforming from Responses API -> LiteLLM completion  (Chat Completion
 
 import json
 import re
-from collections.abc import Sequence
-from typing import Any, Final, Literal, cast
+from collections.abc import Iterator, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable
 
 from openai.types.chat.chat_completion_named_tool_choice_param import (
     ChatCompletionNamedToolChoiceParam,
@@ -16,6 +16,7 @@ from openai.types.chat.chat_completion_named_tool_choice_param import (
 from openai.types.responses import ResponseFunctionToolCall
 from openai.types.responses.response_create_params import ResponseInputParam
 from openai.types.responses.tool_param import FunctionToolParam
+from pydantic import TypeAdapter
 from typing_extensions import TypedDict
 
 from litellm._logging import verbose_logger
@@ -78,9 +79,35 @@ from .custom_tools import (
     unwrap_custom_tool_arguments,
 )
 
+if TYPE_CHECKING:
+    from openai.types.responses.response_apply_patch_tool_call import (
+        ResponseApplyPatchToolCall,
+    )
+
 ########### Initialize Classes used for Responses API  ###########
 TOOL_CALLS_CACHE: Final = InMemoryCache()
 
+_ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object])
+_STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
+_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
+_DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]])
+_TEXT_ADAPTER: Final = TypeAdapter(str)
+
+
+@runtime_checkable
+class _SupportsIter(Protocol):
+    def __iter__(self) -> Iterator[object]: ...
+
+
+@runtime_checkable
+class _HasToolCalls(Protocol):
+    tool_calls: object
+
+
+@runtime_checkable
+class _HasId(Protocol):
+    id: object
+
 
 class ChatCompletionSession(TypedDict, total=False):
     messages: list[
@@ -205,7 +232,7 @@ class LiteLLMCompletionResponsesConfig:
         responses_api_request: ResponsesAPIOptionalRequestParams,
         custom_llm_provider: str | None = None,
         stream: bool | None = None,
-        extra_headers: dict[str, Any] | None = None,
+        extra_headers: Mapping[str, object] | None = None,
         **kwargs,
     ) -> dict:
         """
@@ -462,7 +489,9 @@ class LiteLLMCompletionResponsesConfig:
                     if not chat_completion_messages:
                         continue
 
-                    deduped_in_place: list[Any] = []
+                    deduped_in_place: list[
+                        AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage
+                    ] = []
                     for m in chat_completion_messages:
                         role = ""
                         if isinstance(m, dict):
@@ -472,7 +501,7 @@ class LiteLLMCompletionResponsesConfig:
 
                         # Drop assistant tool_calls wrappers if we already have this call_id
                         if role == "assistant":
-                            tool_calls: Any = (
+                            tool_calls: object = (
                                 m.get("tool_calls") if isinstance(m, dict) else getattr(m, "tool_calls", None)
                             )
                             call_id = ""
@@ -534,7 +563,7 @@ class LiteLLMCompletionResponsesConfig:
             call_id = ""
 
             if role == "assistant":
-                tool_calls: Any = None
+                tool_calls: object = None
                 if isinstance(tool_call_message, dict):
                     tool_calls = tool_call_message.get("tool_calls")
                 else:
@@ -578,7 +607,16 @@ class LiteLLMCompletionResponsesConfig:
         return False
 
     @staticmethod
-    def _find_previous_assistant_idx(messages: list[Any], current_idx: int) -> int | None:
+    def _find_previous_assistant_idx(
+        messages: Sequence[
+            AllMessageValues
+            | GenericChatCompletionMessage
+            | ChatCompletionResponseMessage
+            | ChatCompletionMessageToolCall
+            | Message
+        ],
+        current_idx: int,
+    ) -> int | None:
         """Find the index of the previous assistant message."""
         for j in range(current_idx - 1, -1, -1):
             if messages[j].get("role") == "assistant":
@@ -586,7 +624,18 @@ class LiteLLMCompletionResponsesConfig:
         return None
 
     @staticmethod
-    def _recover_tool_call_id_from_assistant(assistant_message: Any, message: Any) -> str:
+    def _recover_tool_call_id_from_assistant(
+        assistant_message: AllMessageValues
+        | GenericChatCompletionMessage
+        | ChatCompletionResponseMessage
+        | ChatCompletionMessageToolCall
+        | Message,
+        message: AllMessageValues
+        | GenericChatCompletionMessage
+        | ChatCompletionResponseMessage
+        | ChatCompletionMessageToolCall
+        | Message,
+    ) -> str:
         """Try to recover empty tool_call_id from assistant message's tool_calls."""
         tool_calls_raw: Final = (
             assistant_message.get("tool_calls")
@@ -594,17 +643,23 @@ class LiteLLMCompletionResponsesConfig:
             else getattr(assistant_message, "tool_calls", None)
         )
         if tool_calls_raw and isinstance(tool_calls_raw, list) and len(tool_calls_raw) > 0:
-            first_tool_call: Final = tool_calls_raw[0]
+            first_tool_call: Final = _OBJECT_LIST_ADAPTER.validate_python(tool_calls_raw)[0]
             if isinstance(first_tool_call, dict):
-                tool_call_id_raw = first_tool_call.get("id", "")
+                tool_call_id_raw = _ANY_KEY_DICT_ADAPTER.validate_python(first_tool_call).get("id", "")
                 return str(tool_call_id_raw) if tool_call_id_raw is not None else ""
-            elif hasattr(first_tool_call, "id"):
-                tool_call_id_raw = getattr(first_tool_call, "id", None)
+            elif isinstance(first_tool_call, _HasId):
+                tool_call_id_raw = first_tool_call.id
                 return str(tool_call_id_raw) if tool_call_id_raw is not None else ""
         return ""
 
     @staticmethod
-    def _get_tool_calls_list(assistant_message: Any) -> list[Any]:
+    def _get_tool_calls_list(
+        assistant_message: AllMessageValues
+        | GenericChatCompletionMessage
+        | ChatCompletionResponseMessage
+        | ChatCompletionMessageToolCall
+        | Message,
+    ) -> Sequence[object]:
         """Extract tool_calls as a list from assistant message."""
         tool_calls_raw: Final = (
             assistant_message.get("tool_calls")
@@ -614,18 +669,18 @@ class LiteLLMCompletionResponsesConfig:
         if tool_calls_raw is None:
             return []
         if isinstance(tool_calls_raw, list):
-            return tool_calls_raw
-        if hasattr(tool_calls_raw, "__iter__") and not isinstance(tool_calls_raw, (str, bytes)):
+            return _OBJECT_LIST_ADAPTER.validate_python(tool_calls_raw)
+        if isinstance(tool_calls_raw, _SupportsIter) and not isinstance(tool_calls_raw, (str, bytes)):
             return list(tool_calls_raw)
         return []
 
     @staticmethod
-    def _check_tool_call_exists(tool_calls: list[Any], tool_call_id: str) -> bool:
+    def _check_tool_call_exists(tool_calls: Sequence[object], tool_call_id: str) -> bool:
         """Check if a tool_call with the given ID exists in the list."""
         for tool_call in tool_calls:
-            tool_call_id_to_check: str | None = None
+            tool_call_id_to_check: object = None
             if isinstance(tool_call, dict):
-                tool_call_id_to_check = tool_call.get("id")
+                tool_call_id_to_check = _ANY_KEY_DICT_ADAPTER.validate_python(tool_call).get("id")
             elif hasattr(tool_call, "id"):
                 tool_call_id_to_check = getattr(tool_call, "id", None)
             if tool_call_id_to_check == tool_call_id:
@@ -633,12 +688,13 @@ class LiteLLMCompletionResponsesConfig:
         return False
 
     @staticmethod
-    def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: list[Any]) -> dict[str, Any] | None:
+    def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: Sequence[object]) -> dict[str, object] | None:
         """Reconstruct a minimal tool_call definition from tools list."""
         for tool in tools:
             if isinstance(tool, dict):
-                tool_function = tool.get("function") or {}
-                tool_name = tool_function.get("name") or tool.get("name") or ""
+                tool_map = _ANY_KEY_DICT_ADAPTER.validate_python(tool)
+                tool_function = _ANY_KEY_DICT_ADAPTER.validate_python(tool_map.get("function") or {})
+                tool_name = tool_function.get("name") or tool_map.get("name") or ""
                 if tool_name:
                     return {
                         "id": tool_call_id,
@@ -651,7 +707,7 @@ class LiteLLMCompletionResponsesConfig:
         return None
 
     @staticmethod
-    def _get_mapping_or_attr_value(obj: Any, key: str, default: Any = None) -> Any:
+    def _get_mapping_or_attr_value(obj: object, key: str, default: object = None) -> object:
         """
         Safely read a field from dict-like or attribute-based objects.
         """
@@ -659,7 +715,7 @@ class LiteLLMCompletionResponsesConfig:
             return default
 
         if isinstance(obj, dict):
-            return obj.get(key, default)
+            return _ANY_KEY_DICT_ADAPTER.validate_python(obj).get(key, default)
 
         getter: Final = getattr(obj, "get", None)
         if callable(getter):
@@ -672,13 +728,13 @@ class LiteLLMCompletionResponsesConfig:
 
     @staticmethod
     def _create_tool_call_chunk(
-        tool_use_definition: dict[str, Any], tool_call_id: str, index: int
+        tool_use_definition: Mapping[object, object], tool_call_id: str, index: int
     ) -> ChatCompletionToolCallChunk:
         """Create a ChatCompletionToolCallChunk from tool_use_definition."""
         function_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "function")
         function_name_raw: Final = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "name")
         function_arguments_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "arguments")
-        function: Final[dict[str, Any]] = {
+        function: Final[dict[str, object]] = {
             "name": function_name_raw or "",
             "arguments": function_arguments_raw or "{}",
         }
@@ -697,7 +753,7 @@ class LiteLLMCompletionResponsesConfig:
         )
 
     @staticmethod
-    def _normalize_tool_use_definition(tool_use_definition: Any, tool_call_id: str) -> dict[str, Any] | None:
+    def _normalize_tool_use_definition(tool_use_definition: object, tool_call_id: str) -> dict[object, object] | None:
         """
         Normalize cached tool_call definitions to a dict-like shape consumed by _create_tool_call_chunk.
         """
@@ -705,7 +761,7 @@ class LiteLLMCompletionResponsesConfig:
             return None
 
         if isinstance(tool_use_definition, dict):
-            normalized_definition: dict[str, Any] = dict(tool_use_definition)
+            normalized_definition: dict[object, object] = _ANY_KEY_DICT_ADAPTER.validate_python(tool_use_definition)
         else:
             tool_use_id_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "id")
             tool_use_type_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "type")
@@ -738,7 +794,7 @@ class LiteLLMCompletionResponsesConfig:
         return normalized_definition
 
     @staticmethod
-    def _add_tool_call_to_assistant(assistant_message: Any, tool_call_chunk: ChatCompletionToolCallChunk) -> None:
+    def _add_tool_call_to_assistant(assistant_message: object, tool_call_chunk: ChatCompletionToolCallChunk) -> None:
         """Add a tool_call to an assistant message."""
         if isinstance(assistant_message, dict):
             prev_assistant_dict: Final = cast(dict[str, Any], assistant_message)
@@ -747,7 +803,7 @@ class LiteLLMCompletionResponsesConfig:
             tool_calls_list: Final = prev_assistant_dict["tool_calls"]
             if isinstance(tool_calls_list, list):
                 tool_calls_list.append(tool_call_chunk)
-        elif hasattr(assistant_message, "tool_calls"):
+        elif isinstance(assistant_message, _HasToolCalls):
             if assistant_message.tool_calls is None:
                 assistant_message.tool_calls = []
             if isinstance(assistant_message.tool_calls, list):
@@ -762,7 +818,7 @@ class LiteLLMCompletionResponsesConfig:
             | ChatCompletionMessageToolCall
             | Message
         ],
-        tools: list[Any] | None = None,
+        tools: Sequence[object] | None = None,
     ) -> list[
         AllMessageValues
         | GenericChatCompletionMessage
@@ -851,7 +907,7 @@ class LiteLLMCompletionResponsesConfig:
                 tool_calls = LiteLLMCompletionResponsesConfig._get_tool_calls_list(prev_assistant)
 
                 if not LiteLLMCompletionResponsesConfig._check_tool_call_exists(tool_calls, tool_call_id):
-                    _tool_use_definition = TOOL_CALLS_CACHE.get_cache(key=tool_call_id)
+                    _tool_use_definition: object = TOOL_CALLS_CACHE.get_cache(key=tool_call_id)
 
                     if not _tool_use_definition and tools:
                         _tool_use_definition = LiteLLMCompletionResponsesConfig._reconstruct_tool_call_from_tools(
@@ -908,7 +964,7 @@ class LiteLLMCompletionResponsesConfig:
                 function_call=input_item
             )
         else:
-            content: Final = input_item.get("content")
+            content: Final[object] = input_item.get("content")
             # Handle None content: Responses API allows None content, but GenericChatCompletionMessage requires content
             # Since guardrails skip None content anyway, we return empty list to exclude it from structured messages
             if content is None:
@@ -923,7 +979,7 @@ class LiteLLMCompletionResponsesConfig:
             ]
 
     @staticmethod
-    def _is_input_item_tool_call_output(input_item: Any) -> bool:
+    def _is_input_item_tool_call_output(input_item: Mapping[str, object]) -> bool:
         """
         Check if the input item is a tool call output
         """
@@ -936,7 +992,7 @@ class LiteLLMCompletionResponsesConfig:
         ]
 
     @staticmethod
-    def _is_input_item_function_call(input_item: Any) -> bool:
+    def _is_input_item_function_call(input_item: Mapping[str, object]) -> bool:
         """
         Check if the input item is a function call or custom tool call.
         Both need to be reconstructed as assistant tool_calls for Chat
@@ -946,7 +1002,7 @@ class LiteLLMCompletionResponsesConfig:
 
     @staticmethod
     def _transform_responses_api_tool_call_output_to_chat_completion_message(
-        tool_call_output: dict[str, Any],
+        tool_call_output: Mapping[str, object],
     ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]:
         """
         ChatCompletionToolMessage is used to indicate the output from a tool call
@@ -958,7 +1014,7 @@ class LiteLLMCompletionResponsesConfig:
             return []
 
         def _normalize_function_call_output_to_tool_content(
-            output: Any,
+            output: object,
         ) -> Any:
             """
             Normalize Responses API function_call_output.output into a shape that downstream
@@ -981,7 +1037,7 @@ class LiteLLMCompletionResponsesConfig:
 
             # Some adapters represent tool output as a list of "input_*" parts
             if isinstance(output, list):
-                normalized_blocks: Final[list[dict[str, Any]]] = []
+                normalized_blocks: Final[list[dict[str, object]]] = []
                 text_acc: Final[list[str]] = []
                 for part in output:
                     if not isinstance(part, dict):
@@ -1082,7 +1138,7 @@ class LiteLLMCompletionResponsesConfig:
 
     @staticmethod
     def _transform_responses_api_function_call_to_chat_completion_message(
-        function_call: dict[str, Any],
+        function_call: Mapping[str, str],
     ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]:
         """
         Transform a Responses API function_call into a Chat Completion message with tool calls
@@ -1127,7 +1183,7 @@ class LiteLLMCompletionResponsesConfig:
         return [chat_completion_response_message]
 
     @staticmethod
-    def _resolve_file_id(item: dict[str, Any]) -> str | None:
+    def _resolve_file_id(item: Mapping[str, object]) -> object:
         """
         Return the effective file_id for a Responses API input_file item.
         Explicit file_id takes precedence; file_url is used as fallback so
@@ -1136,7 +1192,7 @@ class LiteLLMCompletionResponsesConfig:
         return item.get("file_id") or item.get("file_url") or None
 
     @staticmethod
-    def _transform_input_file_item_to_file_item(item: dict[str, Any]) -> dict[str, Any]:
+    def _transform_input_file_item_to_file_item(item: Mapping[str, object]) -> dict[str, object]:
         """
         Transform a Responses API input_file item to a Chat Completion file item
 
@@ -1146,21 +1202,21 @@ class LiteLLMCompletionResponsesConfig:
         Returns:
             Dictionary with transformed file structure for Chat Completion
         """
-        file_dict: Final[dict[str, Any]] = {}
+        file_dict: Final[dict[str, object]] = {}
         file_id: Final = LiteLLMCompletionResponsesConfig._resolve_file_id(item)
         if file_id:
             file_dict["file_id"] = file_id
         if item.get("file_data"):
             file_dict["file_data"] = item["file_data"]
 
-        new_item: Final[dict[str, Any]] = {"type": "file", "file": file_dict}
+        new_item: Final[dict[str, object]] = {"type": "file", "file": file_dict}
         if "cache_control" in item:
             new_item["cache_control"] = item["cache_control"]
         return new_item
 
     @staticmethod
     def _transform_input_image_item_to_image_item(
-        item: dict[str, Any],
+        item: Mapping[str, str],
     ) -> ChatCompletionImageObject:
         """
         Transform a Responses API input_image item to a Chat Completion image item
@@ -1173,8 +1229,8 @@ class LiteLLMCompletionResponsesConfig:
 
     @staticmethod
     def _transform_responses_api_content_to_chat_completion_content(
-        content: Any,
-    ) -> str | list[str | dict[str, Any]]:
+        content: object,
+    ) -> str | list[str | dict[str, object]]:
         """
         Transform a Responses API content into a Chat Completion content
 
@@ -1188,7 +1244,7 @@ class LiteLLMCompletionResponsesConfig:
         elif isinstance(content, str):
             return content
         elif isinstance(content, list):
-            content_list: Final[list[str | dict[str, Any]]] = []
+            content_list: Final[list[str | dict[str, object]]] = []
             for item in content:
                 if isinstance(item, str):
                     content_list.append(item)
@@ -1198,8 +1254,8 @@ class LiteLLMCompletionResponsesConfig:
                             LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item)
                         )
                     elif item.get("type") == "input_image":
-                        image_block = dict(
-                            LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item)
+                        image_block = _STR_KEY_DICT_ADAPTER.validate_python(
+                            dict(LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item))
                         )
                         if "cache_control" in item:
                             image_block["cache_control"] = item["cache_control"]
@@ -1209,7 +1265,7 @@ class LiteLLMCompletionResponsesConfig:
                         text_value = item.get("text")
                         if text_value is None:
                             continue
-                        content_block: dict[str, Any] = {
+                        content_block: dict[str, object] = {
                             "type": LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type(
                                 item.get("type") or "text"
                             ),
@@ -1299,7 +1355,7 @@ class LiteLLMCompletionResponsesConfig:
                 parameters = dict(typed_tool.get("parameters", {}) or {})
                 if not parameters or "type" not in parameters:
                     parameters["type"] = "object"
-                chat_completion_tool: dict[str, Any] = {
+                chat_completion_tool: dict[str, object] = {
                     "type": "function",
                     "function": {
                         "name": typed_tool.get("name") or "",
@@ -1340,7 +1396,7 @@ class LiteLLMCompletionResponsesConfig:
     @staticmethod
     def transform_chat_completion_tool_params_to_responses_api_tools(
         chat_completion_tools: list[ChatCompletionToolParam | OpenAIMcpServerTool] | None,
-    ) -> list[dict[str, Any]]:
+    ) -> list[dict[str, object]]:
         """
         Transform Chat Completion tool params (e.g. from guardrail output) back to
         Responses API request tool format. Inverse of
@@ -1348,7 +1404,7 @@ class LiteLLMCompletionResponsesConfig:
         """
         if chat_completion_tools is None or not chat_completion_tools:
             return []
-        result: Final[list[dict[str, Any]]] = []
+        result: Final[list[dict[str, object]]] = []
         for tool in chat_completion_tools:
             if not isinstance(tool, dict):
                 result.append(tool)
@@ -1358,7 +1414,7 @@ class LiteLLMCompletionResponsesConfig:
                 parameters = dict(fn.get("parameters", {}) or {})
                 if not parameters or "type" not in parameters:
                     parameters["type"] = "object"
-                responses_tool: dict[str, Any] = {
+                responses_tool: dict[str, object] = {
                     "type": "function",
                     "name": fn.get("name") or "",
                     "description": fn.get("description") or "",
@@ -1510,7 +1566,7 @@ class LiteLLMCompletionResponsesConfig:
     def convert_response_function_tool_call_to_chat_completion_tool_call(
         tool_call_item: Any,
         index: int = 0,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """
         Convert ResponseFunctionToolCall to ChatCompletionToolCallChunk format.
 
@@ -1536,7 +1592,7 @@ class LiteLLMCompletionResponsesConfig:
                     else (dict(provider_fields) if hasattr(provider_fields, "__dict__") else {})
                 )
 
-        function_dict: Final[dict[str, Any]] = {
+        function_dict: Final[dict[str, object]] = {
             "name": tool_call_item.name,
             "arguments": tool_call_item.arguments,
         }
@@ -1544,7 +1600,7 @@ class LiteLLMCompletionResponsesConfig:
         if provider_specific_fields:
             function_dict["provider_specific_fields"] = provider_specific_fields
 
-        tool_call_dict: Final[dict[str, Any]] = {
+        tool_call_dict: Final[dict[str, object]] = {
             "id": LiteLLMCompletionResponsesConfig._tool_call_id_from_responses_item(
                 getattr(tool_call_item, "id", None),
                 getattr(tool_call_item, "call_id", None),
@@ -1561,9 +1617,9 @@ class LiteLLMCompletionResponsesConfig:
 
     @staticmethod
     def convert_apply_patch_tool_call_to_chat_completion_tool_call(
-        tool_call_item: Any,
+        tool_call_item: "ResponseApplyPatchToolCall",
         index: int = 0,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """
         Convert ResponseApplyPatchToolCall to ChatCompletionToolCallChunk format.
 
@@ -1581,7 +1637,7 @@ class LiteLLMCompletionResponsesConfig:
         import json
 
         operation_dict: Final = tool_call_item.operation.model_dump()
-        tool_call_dict: Final[dict[str, Any]] = {
+        tool_call_dict: Final[dict[str, object]] = {
             "id": tool_call_item.call_id,
             "function": {
                 "name": "apply_patch",
@@ -1795,9 +1851,11 @@ class LiteLLMCompletionResponsesConfig:
         if not images:
             return image_generation_items
 
-        for idx, image_item in enumerate(images):
+        for idx, image_item in enumerate(_DICT_ITEMS_LIST_ADAPTER.validate_python(images)):
             # Extract base64 from data URL
-            image_url = image_item.get("image_url", {}).get("url", "")
+            image_url = _TEXT_ADAPTER.validate_python(
+                _ANY_KEY_DICT_ADAPTER.validate_python(image_item.get("image_url", {})).get("url", "")
+            )
             base64_data = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(image_url)
 
             if base64_data:
@@ -2048,8 +2106,8 @@ class LiteLLMCompletionResponsesConfig:
 
     @staticmethod
     def _transform_text_format_to_response_format(
-        text_param: dict[str, Any] | Any,
-    ) -> dict[str, Any] | None:
+        text_param: object,
+    ) -> dict[str, object] | None:
         """
         Transform Responses API text.format parameter to Chat Completion response_format parameter.
 

litellm/responses/streaming_iterator.py

@@ -5,7 +5,7 @@ import json
 import time
 import traceback
 import uuid
-from collections.abc import Mapping
+from collections.abc import Awaitable, Callable, Mapping
 from datetime import datetime
 from functools import lru_cache
 from types import MappingProxyType
@@ -33,6 +33,7 @@ from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfi
 from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils
 from litellm.types.llms.openai import (
     PART_UNION_TYPES,
+    ResponseAPIUsage,
     ResponsesAPIResponse,
     ResponsesAPIStreamEvents,
     ResponsesAPIStreamingResponse,
@@ -112,7 +113,7 @@ _ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType(
 
 
 def _error_event_fields(error_obj: object) -> tuple[str, str | None, str | None]:
-    if isinstance(error_obj, dict):
+    if _is_json_object(error_obj):
         raw_message = error_obj.get("message")
         raw_type = error_obj.get("type")
         raw_code = error_obj.get("code")
@@ -243,7 +244,9 @@ class BaseResponsesAPIStreamingIterator:
                 # Using getattr(..., "response") alone is unsafe with Mocks: they synthesize a
                 # truthy child Mock for any attribute, which breaks tests and is wrong on stream.
                 if "response" in parsed_chunk:
-                    response_object: Final = getattr(openai_responses_api_chunk, "response", None)
+                    response_object: Final[ResponsesAPIResponse | None] = getattr(
+                        openai_responses_api_chunk, "response", None
+                    )
                     if response_object is not None:
                         response: Final = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id(
                             responses_api_response=response_object,
@@ -279,7 +282,9 @@ class BaseResponsesAPIStreamingIterator:
                             model_id=_stream_model_id,
                         )
                 elif _event_type == ResponsesAPIStreamEvents.CONTENT_PART_DONE:
-                    _part: Final = getattr(openai_responses_api_chunk, "part", None)
+                    _part: Final[PART_UNION_TYPES | Mapping[str, object] | None] = getattr(
+                        openai_responses_api_chunk, "part", None
+                    )
                     if _part is not None:
                         if isinstance(_part, dict):
                             ResponsesAPIRequestUtils._encode_container_ids_in_annotations(
@@ -302,7 +307,7 @@ class BaseResponsesAPIStreamingIterator:
                         openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
                         openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
                     ):
-                        item: Final = getattr(openai_responses_api_chunk, "item", None)
+                        item: Final[object | None] = getattr(openai_responses_api_chunk, "item", None)
                         if item:
                             encrypted_content: Final = getattr(item, "encrypted_content", None)
                             if encrypted_content and isinstance(encrypted_content, str):
@@ -324,9 +329,11 @@ class BaseResponsesAPIStreamingIterator:
                     self.completed_response = openai_responses_api_chunk
                     # Add cost to usage object if include_cost_in_streaming_usage is True
                     if litellm.include_cost_in_streaming_usage and self.logging_obj is not None:
-                        response_obj: Final[Any | None] = getattr(openai_responses_api_chunk, "response", None)
+                        response_obj: Final[ResponsesAPIResponse | None] = getattr(
+                            openai_responses_api_chunk, "response", None
+                        )
                         if response_obj:
-                            usage_obj: Final[Any | None] = getattr(response_obj, "usage", None)
+                            usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None)
                             if usage_obj is not None:
                                 try:
                                     cost: float | None = self.logging_obj._response_cost_calculator(result=response_obj)
@@ -414,7 +421,9 @@ class BaseResponsesAPIStreamingIterator:
         async_failure_handler / failure_handler so logging integrations correctly
         record the call as failed.
         """
-        response_obj: Final = getattr(self.completed_response, "response", None) if self.completed_response else None
+        response_obj: Final[ResponsesAPIResponse | None] = (
+            getattr(self.completed_response, "response", None) if self.completed_response else None
+        )
         error_info: Final = getattr(response_obj, "error", None) if response_obj else None
         error_message, error_type, error_code = _error_event_fields(error_info)
         self._record_failed_response_usage(response_obj)
@@ -429,7 +438,7 @@ class BaseResponsesAPIStreamingIterator:
     def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None:
         if response_obj is None or self.logging_obj is None:
             return
-        usage_obj: Final = getattr(response_obj, "usage", None)
+        usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None)
         if usage_obj is None:
             return
         try:
@@ -506,7 +515,7 @@ class BaseResponsesAPIStreamingIterator:
             return
 
         request_kwargs = getattr(caching_handler, "request_kwargs", None)
-        if not isinstance(request_kwargs, dict) or request_kwargs.get("stream") is not True:
+        if not _is_json_object(request_kwargs) or request_kwargs.get("stream") is not True:
             return
         request_kwargs = request_kwargs.copy()
         preset_cache_key = getattr(caching_handler, "preset_cache_key", None)
@@ -606,7 +615,7 @@ class BaseResponsesAPIStreamingIterator:
         if self.completed_response is None:
             return
 
-        request_payload: Final[dict[str, Any]] = {}
+        request_payload: Final[dict[str, object]] = {}
         if isinstance(self.request_data, dict):
             request_payload.update(self.request_data)
         try:
@@ -695,11 +704,15 @@ class BaseResponsesAPIStreamingIterator:
             pass
 
 
-async def call_post_streaming_hooks_for_testing(iterator, chunk):
+async def call_post_streaming_hooks_for_testing(
+    iterator: object, chunk: ResponsesAPIStreamingResponse
+) -> ResponsesAPIStreamingResponse:
     """
     Module-level helper for tests to ensure hooks can be invoked even if the iterator is wrapped.
     """
-    hook_fn: Final = getattr(iterator, "_call_post_streaming_deployment_hook", None)
+    hook_fn: Final[Callable[[ResponsesAPIStreamingResponse], Awaitable[ResponsesAPIStreamingResponse]] | None] = (
+        getattr(iterator, "_call_post_streaming_deployment_hook", None)
+    )
     if hook_fn is None:
         return chunk
     return await hook_fn(chunk)
@@ -1019,7 +1032,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
 def _dump_response_object(obj: Any) -> dict[str, Any]:
     if hasattr(obj, "model_dump"):
         return obj.model_dump()
-    if isinstance(obj, dict):
+    if _is_json_object(obj):
         return obj
     return {}
 
@@ -1684,7 +1697,7 @@ class ResponsesWebSocketStreaming:
             return response_str
 
         try:
-            evt_obj: Final = json.loads(response_str)
+            evt_obj: Final[Mapping[str, object]] = json.loads(response_str)
         except (json.JSONDecodeError, TypeError):
             return response_str
 
@@ -1925,7 +1938,7 @@ class ManagedResponsesWebSocketHandler:
 
     @staticmethod
     def _extract_output_messages(
-        completed_event: dict[str, Any],
+        completed_event: dict[str, object],
     ) -> list[dict[str, object]]:
         """
         Convert the output items in a ``response.completed`` event into
@@ -2065,7 +2078,7 @@ class ManagedResponsesWebSocketHandler:
           Flat:   {"type": "response.create", "input": [...], "model": "...", ...}
         """
         nested: Final = msg_obj.get("response")
-        response_params: Final[dict[str, Any]] = (
+        response_params: Final[dict[str, object]] = (
             nested if _is_json_object(nested) and nested else {k: v for k, v in msg_obj.items() if k != "type"}
         )
         return {
@@ -2076,7 +2089,7 @@ class ManagedResponsesWebSocketHandler:
 
     def _apply_history(
         self,
-        call_kwargs: dict[str, Any],
+        call_kwargs: dict[str, object],
         previous_response_id: str | None,
         current_messages: list[dict[str, object]],
         prior_history: list[dict[str, object]],
@@ -2129,7 +2142,7 @@ class ManagedResponsesWebSocketHandler:
             return False
         return event_provider == self._connection_provider
 
-    def _inject_credentials(self, call_kwargs: dict[str, Any], model: str | None = None) -> None:
+    def _inject_credentials(self, call_kwargs: dict[str, object], model: str | None = None) -> None:
         """Inject connection-level credentials and metadata into call_kwargs."""
         if self.api_key is not None:
             call_kwargs["api_key"] = self.api_key

litellm/utils.py

@@ -5769,7 +5769,7 @@ def json_schema_type(python_type_name: str):
     return python_to_json_schema_types.get(python_type_name, "string")
 
 
-def function_to_dict(input_function) -> dict:  # noqa: C901
+def function_to_dict(input_function) -> dict:
     """Using type hints and numpy-styled docstring,
     produce a dictionary usable for OpenAI function calling
 

ruff-strict-budget.json

@@ -1,6 +1,6 @@
 {
   "ANN001": {
-    "limit": 3121
+    "limit": 3114
   },
   "ANN002": {
     "limit": 71
@@ -9,7 +9,7 @@
     "limit": 834
   },
   "ANN201": {
-    "limit": 2032
+    "limit": 2031
   },
   "ANN202": {
     "limit": 865
@@ -24,7 +24,7 @@
     "limit": 133
   },
   "ANN401": {
-    "limit": 1630
+    "limit": 1555
   },
   "ASYNC230": {
     "limit": 11
@@ -56,9 +56,6 @@
   "B026": {
     "limit": 3
   },
-  "B033": {
-    "limit": 0
-  },
   "BLE001": {
     "limit": 2924
   },
@@ -81,7 +78,7 @@
     "limit": 1
   },
   "C901": {
-    "limit": 315
+    "limit": 314
   },
   "D419": {
     "limit": 6
@@ -113,18 +110,6 @@
   "F401": {
     "limit": 17
   },
-  "FURB136": {
-    "limit": 0
-  },
-  "FURB168": {
-    "limit": 0
-  },
-  "FURB188": {
-    "limit": 0
-  },
-  "I001": {
-    "limit": 0
-  },
   "LOG015": {
     "limit": 5
   },
@@ -137,18 +122,9 @@
   "PERF401": {
     "limit": 12
   },
-  "PERF402": {
-    "limit": 0
-  },
   "PERF403": {
     "limit": 34
   },
-  "PIE790": {
-    "limit": 0
-  },
-  "PIE800": {
-    "limit": 0
-  },
   "PIE804": {
     "limit": 18
   },
@@ -158,9 +134,6 @@
   "PLC0206": {
     "limit": 26
   },
-  "PLC0208": {
-    "limit": 0
-  },
   "PLC0414": {
     "limit": 46
   },
@@ -170,24 +143,12 @@
   "PLR0206": {
     "limit": 1
   },
-  "PLR0402": {
-    "limit": 0
-  },
   "PLR1704": {
     "limit": 3
   },
-  "PLR1711": {
-    "limit": 0
-  },
   "PLR1714": {
     "limit": 257
   },
-  "PLR1730": {
-    "limit": 0
-  },
-  "PLR2044": {
-    "limit": 0
-  },
   "PLW0127": {
     "limit": 57
   },
@@ -206,27 +167,12 @@
   "PLW1510": {
     "limit": 2
   },
-  "PYI030": {
-    "limit": 0
-  },
   "PYI036": {
     "limit": 3
   },
-  "PYI041": {
-    "limit": 0
-  },
-  "PYI064": {
-    "limit": 0
-  },
-  "RET501": {
-    "limit": 0
-  },
   "RET504": {
     "limit": 177
   },
-  "RUF010": {
-    "limit": 0
-  },
   "RUF012": {
     "limit": 241
   },
@@ -236,23 +182,14 @@
   "RUF019": {
     "limit": 38
   },
-  "RUF022": {
-    "limit": 0
-  },
-  "RUF023": {
-    "limit": 0
-  },
   "RUF046": {
     "limit": 4
   },
-  "RUF051": {
-    "limit": 0
-  },
   "RUF059": {
     "limit": 67
   },
   "RUF100": {
-    "limit": 100
+    "limit": 0
   },
   "S110": {
     "limit": 218
@@ -272,18 +209,12 @@
   "SIM113": {
     "limit": 3
   },
-  "SIM114": {
-    "limit": 0
-  },
   "SIM115": {
     "limit": 2
   },
   "SIM117": {
     "limit": 7
   },
-  "SIM118": {
-    "limit": 0
-  },
   "SIM201": {
     "limit": 1
   },
@@ -302,11 +233,8 @@
   "TC004": {
     "limit": 5
   },
-  "TC005": {
-    "limit": 0
-  },
   "TID251": {
-    "limit": 1240
+    "limit": 1238
   },
   "TRY002": {
     "limit": 528
@@ -323,46 +251,13 @@
   "TRY300": {
     "limit": 860
   },
-  "UP006": {
-    "limit": 0
-  },
-  "UP007": {
-    "limit": 0
-  },
-  "UP008": {
-    "limit": 0
-  },
-  "UP012": {
-    "limit": 0
-  },
-  "UP018": {
-    "limit": 0
-  },
-  "UP024": {
-    "limit": 0
-  },
   "UP028": {
     "limit": 2
   },
   "UP031": {
     "limit": 2
   },
-  "UP032": {
-    "limit": 0
-  },
-  "UP034": {
-    "limit": 0
-  },
-  "UP035": {
-    "limit": 0
-  },
   "UP036": {
     "limit": 1
-  },
-  "UP037": {
-    "limit": 0
-  },
-  "UP045": {
-    "limit": 0
   }
 }

ruff-strict.toml

@@ -4,6 +4,17 @@ extend = "ruff.toml"
 preview = true
 select = ["ANN", "ASYNC230", "B004", "B006", "B008", "B009", "B010", "B018", "B019", "B021", "B026", "B033", "BLE", "C401", "C404", "C405", "C408", "C414", "C419", "C901", "D419", "DTZ001", "DTZ003", "DTZ005", "DTZ006", "DTZ007", "DTZ011", "EXE001", "EXE002", "F401", "FURB136", "FURB168", "FURB188", "I001", "LOG015", "N999", "PERF102", "PERF401", "PERF402", "PERF403", "PIE790", "PIE800", "PIE804", "PIE810", "PLC0206", "PLC0208", "PLC0414", "PLR0124", "PLR0206", "PLR0402", "PLR1704", "PLR1711", "PLR1714", "PLR1730", "PLR2044", "PLW0127", "PLW0133", "PLW0602", "PLW0603", "PLW1508", "PLW1510", "PYI030", "PYI036", "PYI041", "PYI064", "RET501", "RET504", "RUF010", "RUF012", "RUF015", "RUF019", "RUF022", "RUF023", "RUF046", "RUF051", "RUF059", "RUF100", "S110", "S112", "SIM101", "SIM102", "SIM103", "SIM113", "SIM114", "SIM115", "SIM117", "SIM118", "SIM201", "SIM210", "SIM211", "SIM222", "SIM401", "TC004", "TC005", "TID251", "TRY002", "TRY004", "TRY201", "TRY203", "TRY300", "UP006", "UP007", "UP008", "UP012", "UP018", "UP024", "UP028", "UP031", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045"]
 extend-select = []
+# Overrides the inherited list: rules this gate enforces itself must NOT be external here,
+# so this config's RUF100 flags their stale `# noqa` directives. What remains external is
+# only what other tooling enforces: every base ruff.toml rule this select list doesn't
+# re-enable (all of the default E/F families plus T20/PGH004/RUF008/RUF009, minus the
+# strict-selected F401 and RUF100; F4 is split out so stale F401 noqas stay detectable),
+# plus upstream litellm's ruff config.
+external = [
+    "T20", "PGH004", "RUF008", "RUF009", "E4", "E7", "E9",
+    "F402", "F404", "F406", "F407", "F5", "F6", "F7", "F8", "F9",
+    "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405",
+]
 
 [lint.mccabe]
 max-complexity = 15

ruff.toml

@@ -1,12 +1,26 @@
 lint.ignore = ["F405", "E402", "F403"]
-lint.extend-select = ["T20", "PGH004", "RUF008", "RUF009", "RUF100"]
+# The second group is the strict gate's graduates: rules the codebase already has zero
+# violations of, so they hard-fail here instead of being ratcheted in ruff-strict-budget.json.
+# That gives editors and `ruff check --fix` the diagnostic, which the gate script cannot.
+lint.extend-select = [
+    "T20", "PGH004", "RUF008", "RUF009", "RUF100",
+    "B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790", "PIE800", "PLC0208",
+    "PLR0402", "PLR1711", "PLR1730", "PLR2044", "PYI030", "PYI041", "PYI064", "RET501", "RUF010",
+    "RUF022", "RUF023", "RUF051", "SIM114", "SIM118", "TC005", "UP006", "UP007", "UP008", "UP012",
+    "UP018", "UP024", "UP032", "UP034", "UP035", "UP037", "UP045",
+]
 # RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip
 # `# noqa` directives that protect rules enforced elsewhere. List those codes as external
 # so RUF100 leaves their directives alone: the strict gate (ruff-strict.toml) and upstream
 # litellm's own ruff config both rely on suppressions this config can't see.
 lint.external = [
-    # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml)
-    "C901", "TID251",
+    # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml).
+    # Family entries whose every strict rule graduated into extend-select above (FURB), and
+    # standalone graduated codes (I001, RUF010, RUF022, RUF023, RUF051), are dropped so this
+    # config's RUF100 polices their directives itself.
+    "ANN", "ASYNC230", "B", "C", "D419", "DTZ", "EXE", "LOG015", "N999", "PERF",
+    "PIE", "PL", "PYI", "RET", "RUF012", "RUF015", "RUF019",
+    "RUF046", "RUF059", "S110", "S112", "SIM", "TC", "TID251", "TRY", "UP",
     # Enforced by upstream litellm's ruff config, but not run in this repo's CI
     "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405",
 ]

scripts/budget_ratchet_check.py

@@ -10,7 +10,10 @@ content at the merge-base with the target branch and fails (exits 1, red) if:
   * a rule was dropped from a budget (its ceiling effectively became infinite), or
   * an entire budget file was deleted.
 
-New rules and lowered/equal limits are fine.
+New rules and lowered/equal limits are fine. So is a rule that graduated: once a
+paired config (ruff.toml for the ruff-strict budget) selects the rule outright it
+hard-fails at the first violation, which is stricter than any ceiling the budget
+could hold, so dropping its entry tightens the guard rather than removing it.
 
 This is deliberately NOT a gating check. It should turn the run red so that a
 loosening is impossible to miss in review, but it must stay OUT of the
@@ -30,7 +33,9 @@ import argparse
 import json
 import subprocess
 import sys
+import tomllib
 from pathlib import Path
+from types import MappingProxyType
 from typing import NamedTuple
 
 REPO_ROOT = Path(__file__).resolve().parent.parent
@@ -40,6 +45,7 @@ DEFAULT_BUDGETS: tuple[str, ...] = (
     "type-discipline-budget.json",
     "basedpyright-code-budget.json",
 )
+GRADUATION_CONFIGS = MappingProxyType({"ruff-strict-budget.json": "ruff.toml"})
 
 
 class Regression(NamedTuple):
@@ -106,24 +112,57 @@ def _limits(budget: dict) -> dict[str, int]:
     }
 
 
+def selectors_hard_failed_by(lint: dict) -> tuple[str, ...]:
+    """A ruff `[lint]` table's selected codes, minus anything `ignore` turns back off.
+
+    `lint.ignore` wins over `lint.extend-select` in ruff, so an ignored code is not
+    actually enforced and must not count as a graduation.
+    """
+    ignored = tuple(lint.get("ignore", ()))
+    return tuple(
+        selector
+        for selector in lint.get("extend-select", ())
+        if not (ignored and selector.startswith(ignored))
+    )
+
+
+def graduated_selectors(rel: str) -> tuple[str, ...]:
+    """Selectors the budget's paired ruff config hard-fails, so its ceiling is moot."""
+    config = GRADUATION_CONFIGS.get(rel)
+    if config is None or not (REPO_ROOT / config).exists():
+        return ()
+    return selectors_hard_failed_by(
+        tomllib.loads((REPO_ROOT / config).read_text()).get("lint", {})
+    )
+
+
 def _regression_detail(
     rule: str,
     base_limits: dict[str, int],
     head_limits: dict[str, int],
+    graduated: tuple[str, ...],
 ) -> str | None:
-    """Why `rule` regressed vs base, or None when it held flat or fell.
+    """Why `rule` regressed vs base, or None when it held flat, fell, or graduated.
 
-    A dropped rule is terminal; otherwise the only loosening left is a raised limit.
+    A dropped rule is terminal unless it graduated; otherwise the only loosening
+    left is a raised limit.
     """
     base_limit = base_limits[rule]
     if rule not in head_limits:
+        if graduated and rule.startswith(graduated):
+            return None
         return f"rule dropped (limit {base_limit} -> removed)"
     if head_limits[rule] > base_limit:
         return f"limit raised {base_limit} -> {head_limits[rule]}"
     return None
 
 
-def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regression]:
+def regressions_for(
+    rel: str,
+    base: dict | None,
+    head: dict | None,
+    graduated: tuple[str, ...] = (),
+) -> list[Regression]:
     if base is None:
         return []  # new budget file: nothing to ratchet against yet
     if head is None:
@@ -133,7 +172,7 @@ def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regr
     return [
         Regression(rel, rule, detail)
         for rule in sorted(base_limits)
-        if (detail := _regression_detail(rule, base_limits, head_limits)) is not None
+        if (detail := _regression_detail(rule, base_limits, head_limits, graduated)) is not None
     ]
 
 
@@ -164,7 +203,7 @@ def main() -> int:
             print(f"skip {rel}: new file (no base at {args.base} to ratchet against)")
             continue
         checked.append(rel)
-        regressions.extend(regressions_for(rel, base, head))
+        regressions.extend(regressions_for(rel, base, head, graduated_selectors(rel)))
 
     if regressions:
         print(

tests/e2e/proxy_client.py

@@ -76,15 +76,12 @@ from transport import HttpTransport, SplitTransport, Transport
 
 RowsPredicate = Callable[[list[SpendLogRow]], bool]
 
-# After /model/new, the control-plane writer reloads itself immediately, but every
-# other gateway worker (and peer pod) only picks the model up on its add_deployment
-# job. That job runs every proxy_config_reload_interval_seconds (product default 30).
-# A single /v1/models hit can land on a hot worker while the next /chat hits a cold
-# one ("Invalid model name"). Wait for first listing within MODEL_SERVABLE_TIMEOUT,
-# then require continuous listing for MODEL_SERVABLE_DB_SYNC_SECONDS (the default
-# reload interval) so every worker has had a chance to sync from the DB.
+# After /model/new, poll data-plane /v1/models until the model is listed (or fail).
+# Bound by MODEL_SERVABLE_TIMEOUT so a stuck reload does not burn the spend
+# poll_timeout (120s). Return on first listing; settle_propagation owns the separate
+# wait that lets every worker and replica reload before the caller uses the model.
 MODEL_SERVABLE_TIMEOUT = 40.0
-MODEL_SERVABLE_DB_SYNC_SECONDS = 30.0
+MODEL_SERVABLE_DB_SYNC_SECONDS = 0.0
 MODEL_SERVABLE_INTERVAL = 2.0
 # Cap each /v1/models poll so one slow request cannot outlast the remaining budget.
 MODEL_SERVABLE_REQUEST_TIMEOUT = 5.0
@@ -310,10 +307,7 @@ class ProxyClient:
         return model_id
 
     def _await_model_servable(self, model_name: str) -> None:
-        """Block until the data plane lists `model_name` long enough for DB sync.
-
-        Fails if first listing misses model_servable_timeout, or if continuous listing
-        for model_servable_db_sync_seconds never holds (multi-worker / peer reload)."""
+        """Block until the data plane lists `model_name`, or fail at model_servable_timeout."""
         outcome = await_servable(
             lambda poll_timeout: self.transport.get(
                 "/v1/models",

tests/test_litellm/conftest.py

@@ -19,6 +19,7 @@ sys.path.insert(
 import asyncio
 
 import litellm
+from litellm import router as litellm_router_module
 from litellm import utils as litellm_utils_module
 from litellm._logging import ALL_LOGGERS
 from litellm.litellm_core_utils.prompt_templates import (
@@ -244,6 +245,8 @@ def isolate_litellm_state():
         for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items()
     }
 
+    original_live_routers = set(litellm_router_module._live_routers)
+
     # Store LiteLLM logger state. Some tests reconfigure handlers/propagation for
     # JSON logging and do not restore them, which breaks later caplog-based tests.
     logger_state = {}
@@ -313,6 +316,11 @@ def isolate_litellm_state():
     litellm_utils_module._runtime_registered_model_cost.clear()
     litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost)
 
+    for _router in tuple(litellm_router_module._live_routers):
+        litellm_router_module._live_routers.discard(_router)
+    for _router in original_live_routers:
+        litellm_router_module._live_routers.add(_router)
+
     # Restore logger configuration mutated by logging-focused tests.
     for logger in ALL_LOGGERS:
         original_logger_state = logger_state.get(logger.name)

tests/test_litellm/llms/vertex_ai/batches/test_handler.py

@@ -5,8 +5,10 @@ The handler is HTTP/auth glue around the (separately-tested) pure
 ``VertexAIBatchTransformation``. Each public method (create / retrieve / list /
 cancel) resolves a Vertex access token + URL, branches on ``_is_async``
 (returning the coroutine in the async case, doing the sync HTTP call otherwise),
-checks the HTTP status, and parses the JSON into ``LiteLLMBatch`` (or the OpenAI
-list shape).
+and parses the JSON into ``LiteLLMBatch`` (or the OpenAI list shape). POST-backed
+calls rely on the client's ``raise_for_status`` (non-2xx surfaces as
+``httpx.HTTPStatusError``); GET-backed calls return without raising, so the
+handler checks their status codes itself.
 
 We mock only true I/O / auth seams:
   * ``_ensure_access_token`` - the Vertex credential seam. Returns a fixed
@@ -20,7 +22,7 @@ We mock only true I/O / auth seams:
     what URL/headers/body, and that the response is parsed into the litellm
     type. Sibling seams are asserted NOT called where relevant.
 
-The ``_is_async`` branch, status-code error paths, and the cancel
+The ``_is_async`` branch, the error paths, and the cancel
 retrieve-after-cancel sequencing run for real.
 """
 
@@ -40,6 +42,7 @@ sys.path.insert(0, os.path.abspath("../../../../.."))
 from litellm.llms.vertex_ai.batches.handler import (  # noqa: E402
     VertexAIBatchPrediction,
 )
+from litellm.llms.vertex_ai.common_utils import VertexAIError  # noqa: E402
 from litellm.types.utils import LiteLLMBatch  # noqa: E402
 
 HMOD = "litellm.llms.vertex_ai.batches.handler"
@@ -178,13 +181,19 @@ def test_create_batch_async_returns_coroutine_and_uses_async_client():
     sync_client.post.assert_not_called()
 
 
-def test_create_batch_sync_non_200_raises():
+def test_create_batch_sync_httpstatuserror_propagates():
+    """``HTTPHandler.post`` raises for non-2xx via ``raise_for_status``; the
+    sync create path must surface that error, not swallow it."""
     h = _make_handler()
     client = MagicMock()
-    client.post.return_value = _http_response(status_code=500)
+    request = httpx.Request("POST", "https://x/batchPredictionJobs")
+    err_response = httpx.Response(status_code=500, request=request, text="boom")
+    client.post.side_effect = httpx.HTTPStatusError(
+        "boom", request=request, response=err_response
+    )
 
     with patch(f"{HMOD}._get_httpx_client", return_value=client):
-        with pytest.raises(Exception, match="Error: 500"):
+        with pytest.raises(httpx.HTTPStatusError):
             h.create_batch(
                 _is_async=False,
                 create_batch_data=CREATE_DATA,
@@ -197,27 +206,27 @@ def test_create_batch_sync_non_200_raises():
             )
 
 
-def test_create_batch_async_non_200_raises():
+def test_create_batch_input_file_id_without_model_raises_400_before_post():
+    """A gs:// uri with no publishers/<publisher>/models/<model> path is a 400, not a bare 500."""
     h = _make_handler()
-    async_client = MagicMock()
-    async_client.post = AsyncMock(return_value=_http_response(status_code=403))
+    client = MagicMock()
 
-    with (
-        patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()),
-        patch(f"{HMOD}.get_async_httpx_client", return_value=async_client),
-    ):
-        coro = h.create_batch(
-            _is_async=True,
-            create_batch_data=CREATE_DATA,
-            api_base=None,
-            vertex_credentials=None,
-            vertex_project=PROJECT,
-            vertex_location=LOCATION,
-            timeout=600.0,
-            max_retries=None,
-        )
-        with pytest.raises(Exception, match="Error: 403"):
-            _run(coro)
+    with patch(f"{HMOD}._get_httpx_client", return_value=client):
+        with pytest.raises(VertexAIError) as exc_info:
+            h.create_batch(
+                _is_async=False,
+                create_batch_data={"input_file_id": "gs://bucket/batch-input.jsonl"},
+                api_base=None,
+                vertex_credentials=None,
+                vertex_project=PROJECT,
+                vertex_location=LOCATION,
+                timeout=600.0,
+                max_retries=None,
+            )
+
+    assert exc_info.value.status_code == 400
+    assert "gs://bucket/batch-input.jsonl" in str(exc_info.value)
+    client.post.assert_not_called()
 
 
 # =========================================================================== #
@@ -292,7 +301,7 @@ def test_retrieve_batch_sync_non_200_raises():
         patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()),
         patch(f"{HMOD}.safe_get", return_value=_http_response(status_code=404)),
     ):
-        with pytest.raises(Exception, match="Error: 404"):
+        with pytest.raises(VertexAIError, match="Error: 404"):
             h.retrieve_batch(
                 _is_async=False,
                 batch_id=BATCH_ID,
@@ -438,7 +447,7 @@ def test_list_batches_sync_non_200_raises():
     client.get.return_value = _http_response(status_code=500)
 
     with patch(f"{HMOD}._get_httpx_client", return_value=client):
-        with pytest.raises(Exception, match="Error: 500"):
+        with pytest.raises(VertexAIError, match="Error: 500"):
             h.list_batches(
                 _is_async=False,
                 after=None,
@@ -524,27 +533,6 @@ def test_cancel_batch_async_returns_coroutine_posts_then_retrieves():
     assert post_kwargs["url"].endswith(":cancel")
 
 
-def test_cancel_batch_sync_cancel_post_non_200_raises():
-    h = _make_handler()
-    client = MagicMock()
-    client.post.return_value = _http_response(status_code=500)
-
-    with patch(f"{HMOD}._get_httpx_client", return_value=client):
-        with pytest.raises(Exception, match="Error: 500"):
-            h.cancel_batch(
-                _is_async=False,
-                batch_id=BATCH_ID,
-                api_base=None,
-                vertex_credentials=None,
-                vertex_project=PROJECT,
-                vertex_location=LOCATION,
-                timeout=600.0,
-                max_retries=None,
-            )
-    # cancel POST failed -> retrieve GET must never fire
-    client.get.assert_not_called()
-
-
 def test_cancel_batch_sync_retrieve_non_200_raises():
     h = _make_handler()
     client = MagicMock()
@@ -552,7 +540,7 @@ def test_cancel_batch_sync_retrieve_non_200_raises():
     client.get.return_value = _http_response(status_code=404)
 
     with patch(f"{HMOD}._get_httpx_client", return_value=client):
-        with pytest.raises(Exception, match="Error: 404"):
+        with pytest.raises(VertexAIError, match="Error: 404"):
             h.cancel_batch(
                 _is_async=False,
                 batch_id=BATCH_ID,
@@ -672,7 +660,7 @@ def test_async_retrieve_batch_non_200_raises():
             timeout=600.0,
             max_retries=None,
         )
-        with pytest.raises(Exception, match="Error: 500"):
+        with pytest.raises(VertexAIError, match="Error: 500"):
             _run(coro)
 
 
@@ -726,7 +714,7 @@ def test_async_list_batches_non_200_raises():
             timeout=600.0,
             max_retries=None,
         )
-        with pytest.raises(Exception, match="Error: 500"):
+        with pytest.raises(VertexAIError, match="Error: 500"):
             _run(coro)
 
 
@@ -761,28 +749,6 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200():
             _run(coro)
     async_client.get.assert_not_awaited()
 
-    # (a2) cancel POST returns a plain non-200 (no exception) -> raises
-    async_client_post500 = MagicMock()
-    async_client_post500.post = AsyncMock(return_value=_http_response(status_code=500))
-    async_client_post500.get = AsyncMock()
-    with (
-        patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()),
-        patch(f"{HMOD}.get_async_httpx_client", return_value=async_client_post500),
-    ):
-        coro = h.cancel_batch(
-            _is_async=True,
-            batch_id=BATCH_ID,
-            api_base=None,
-            vertex_credentials=None,
-            vertex_project=PROJECT,
-            vertex_location=LOCATION,
-            timeout=600.0,
-            max_retries=None,
-        )
-        with pytest.raises(Exception, match="Error: 500"):
-            _run(coro)
-    async_client_post500.get.assert_not_awaited()
-
     # (b) retrieve-after-cancel returns non-200
     async_client2 = MagicMock()
     async_client2.post = AsyncMock(return_value=_http_response(json_body={}))
@@ -801,5 +767,5 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200():
             timeout=600.0,
             max_retries=None,
         )
-        with pytest.raises(Exception, match="Error: 404"):
+        with pytest.raises(VertexAIError, match="Error: 404"):
             _run(coro)

tests/test_litellm/llms/vertex_ai/batches/test_transformation.py

@@ -25,6 +25,7 @@ from litellm.llms.vertex_ai.batches.transformation import (  # noqa: E402
     VertexAIBatchTransformation,
 )
 from litellm.llms.vertex_ai.common_utils import (  # noqa: E402
+    VertexAIError,
     _convert_vertex_datetime_to_openai_datetime,
 )
 from litellm.types.utils import LiteLLMBatch  # noqa: E402
@@ -69,6 +70,24 @@ def test_transform_openai_request_missing_input_file_id_raises():
         T.transform_openai_batch_request_to_vertex_ai_batch_request({})
 
 
+@pytest.mark.parametrize(
+    "input_file_id",
+    [
+        "gs://bucket/no-model-here.jsonl",
+        "gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid",
+        "gs://bucket/publishers/google/models",
+        "gs://bucket/publishers/google/models//file-uuid",
+    ],
+)
+def test_transform_openai_request_unparseable_model_raises_400(input_file_id: str):
+    """An input_file_id with no parseable model path is a client error, not an IndexError -> 500."""
+    with pytest.raises(VertexAIError) as exc_info:
+        T.transform_openai_batch_request_to_vertex_ai_batch_request({"input_file_id": input_file_id})
+
+    assert exc_info.value.status_code == 400
+    assert input_file_id in str(exc_info.value)
+
+
 # =========================================================================== #
 # transform_vertex_ai_batch_response_to_openai_batch_response
 # =========================================================================== #
@@ -299,9 +318,29 @@ def test_get_model_from_gcs_file_url_encoded():
     assert T._get_model_from_gcs_file(encoded) == "publishers/google/models/gemini-1.5-flash-001"
 
 
-def test_get_model_from_gcs_file_no_publishers_raises():
-    with pytest.raises(IndexError):
+def test_get_model_from_gcs_file_no_publishers_raises_400():
+    with pytest.raises(VertexAIError) as exc_info:
         T._get_model_from_gcs_file("gs://bucket/no-model-here.jsonl")
+    assert exc_info.value.status_code == 400
+
+
+# =========================================================================== #
+# is_unmanaged_gcs_batch_input_file_id
+# =========================================================================== #
+
+
+@pytest.mark.parametrize(
+    "input_file_id, expected",
+    [
+        (INPUT_FILE, True),
+        (None, False),
+        ("file-abc123", False),
+        ("gs://bucket/no-model-here.jsonl", False),
+        ("gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", False),
+    ],
+)
+def test_is_unmanaged_gcs_batch_input_file_id(input_file_id, expected):
+    assert T.is_unmanaged_gcs_batch_input_file_id(input_file_id) is expected
 
 
 # =========================================================================== #

tests/test_litellm/proxy/batches_endpoints/test_endpoints.py

@@ -2253,3 +2253,176 @@ async def test_cancel__provider_only_resolves_named_vertex_credentials(cancel_ha
         "vertex_location": "us-central1",
         "vertex_credentials": "/creds/customer-sa.json",
     }
+
+
+# =========================================================================== #
+# require_managed_files - raw provider ids must not reach the provider.        #
+#                                                                              #
+# Ownership rows only exist for LiteLLM managed ids. A raw provider id sent to #
+# these routes is forwarded under the shared provider credentials with no      #
+# tenant check, so any caller who learns another tenant's id can read its      #
+# batch, reuse its file as batch input, or cancel its job. These lock the      #
+# guard on every batches route that accepts a caller-supplied id.              #
+# =========================================================================== #
+
+
+def _unified_batch_id(model_id: str = "azure/gpt-4o", batch_id: str = "batch-provider-id") -> str:
+    import base64
+
+    from litellm.types.utils import SpecialEnums
+
+    unified = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id)
+    return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=")
+
+
+def _unified_file_id() -> str:
+    import base64
+
+    from litellm.types.utils import SpecialEnums
+
+    unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
+        "application/json", "managed-id", "gpt-4o-mini", "file-provider-id", "gpt-4o-mini-id"
+    )
+    return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=")
+
+
+@dataclass(frozen=True)
+class ManagedResourceAccessCheckerStub:
+    file_access: bool = True
+    object_access: bool = True
+
+    async def can_user_call_unified_file_id(
+        self,
+        unified_file_id: str,
+        user_api_key_dict: UserAPIKeyAuth,
+    ) -> bool:
+        return self.file_access
+
+    async def can_user_call_unified_object_id(
+        self,
+        unified_object_id: str,
+        user_api_key_dict: UserAPIKeyAuth,
+    ) -> bool:
+        return self.object_access
+
+
+@pytest.mark.asyncio
+async def test_create__raw_input_file_id_rejected_when_managed_files_required(harness):
+    set_body(
+        harness,
+        {
+            "input_file_id": "file-victim-abc123",
+            "endpoint": "/v1/chat/completions",
+            "completion_window": "24h",
+        },
+    )
+
+    with patch.object(litellm, "require_managed_files", True):
+        with pytest.raises(ProxyException) as exc:
+            await call_create(harness)
+
+    assert exc.value.code == "400"
+    harness.litellm_acreate.assert_not_called()
+    harness.router_acreate.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_create__model_encoded_input_file_id_rejected_when_managed_files_required(harness):
+    """A model-encoded id is client-forgeable and has no ownership row, so it is
+    not a managed file id and must be rejected like any other raw id."""
+    set_body(
+        harness,
+        {
+            "input_file_id": AZURE_FILE_ID,
+            "endpoint": "/v1/chat/completions",
+            "completion_window": "24h",
+        },
+    )
+
+    with patch.object(litellm, "require_managed_files", True):
+        with pytest.raises(ProxyException) as exc:
+            await call_create(harness)
+
+    assert exc.value.code == "400"
+    harness.litellm_acreate.assert_not_called()
+    harness.router_acreate.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_create__raw_input_file_id_allowed_when_managed_files_not_required(harness):
+    set_body(
+        harness,
+        {
+            "input_file_id": "file-victim-abc123",
+            "endpoint": "/v1/chat/completions",
+            "completion_window": "24h",
+        },
+    )
+
+    with patch.object(litellm, "require_managed_files", False):
+        await call_create(harness)
+
+    assert harness.acreate_kwargs()["input_file_id"] == "file-victim-abc123"
+
+
+@pytest.mark.asyncio
+async def test_create__other_teams_unified_input_file_id_rejected(harness):
+    set_body(
+        harness,
+        {
+            "input_file_id": _unified_file_id(),
+            "endpoint": "/v1/chat/completions",
+            "completion_window": "24h",
+        },
+    )
+    harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub(file_access=False)
+
+    with patch.object(litellm, "require_managed_files", True):
+        with pytest.raises(ProxyException) as exc:
+            await call_create(harness)
+
+    assert exc.value.code == "403"
+    harness.litellm_acreate.assert_not_called()
+    harness.router_acreate.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_retrieve__raw_batch_id_rejected_when_managed_files_required(retrieve_harness):
+    with patch.object(litellm, "require_managed_files", True):
+        with pytest.raises(ProxyException) as exc:
+            await call_retrieve(retrieve_harness, "batch-victim-abc123")
+
+    assert exc.value.code == "400"
+    retrieve_harness.litellm_aretrieve.assert_not_called()
+    retrieve_harness.router_aretrieve.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_retrieve__unified_batch_id_allowed_when_managed_files_required(retrieve_harness):
+    retrieve_harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub()
+
+    with patch.object(litellm, "require_managed_files", True):
+        await call_retrieve(retrieve_harness, _unified_batch_id())
+
+    assert retrieve_harness.router_aretrieve.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_cancel__raw_batch_id_rejected_when_managed_files_required(cancel_harness):
+    with patch.object(litellm, "require_managed_files", True):
+        with pytest.raises(ProxyException) as exc:
+            await call_cancel(cancel_harness, "batch-victim-abc123")
+
+    assert exc.value.code == "400"
+    cancel_harness.litellm_acancel.assert_not_called()
+    cancel_harness.router_acancel.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_cancel__unified_batch_id_allowed_when_managed_files_required(cancel_harness):
+    cancel_harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub()
+
+    with patch.object(litellm, "require_managed_files", True):
+        await call_cancel(cancel_harness, _unified_batch_id())
+
+    assert cancel_harness.router_acancel.call_count == 1

tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py

@@ -0,0 +1,278 @@
+"""
+require_managed_files enforcement for litellm/proxy/fine_tuning_endpoints/endpoints.py
+
+Ownership rows only exist for LiteLLM managed ids. A raw provider id sent to these
+routes is forwarded to the provider under the shared proxy credentials with no tenant
+check, so any caller who learns another tenant's file id can train on it, and any
+caller who learns another tenant's job id can read or cancel it.
+
+Each test asserts BOTH that the request is rejected AND that every downstream provider
+seam stayed untouched, so a guard that raises after the provider call would still fail.
+"""
+
+import base64
+import os
+import sys
+from contextlib import ExitStack
+from dataclasses import dataclass
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../../../.."))
+
+from fastapi import Response
+
+import litellm
+import litellm.proxy.fine_tuning_endpoints.endpoints as endpoints
+import litellm.proxy.proxy_server as proxy_server
+from litellm.proxy._types import ProxyException, UserAPIKeyAuth
+from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+from litellm.proxy.utils import ProxyLogging
+from litellm.router import Router
+from litellm.types.llms.openai import LiteLLMFineTuningJobCreate
+from litellm.types.utils import LiteLLMFineTuningJob, SpecialEnums
+
+RAW_FILE_ID = "file-victim-abc123"
+RAW_JOB_ID = "ftjob-victim-abc123"
+
+
+def _unified_file_id() -> str:
+    unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
+        "application/json", "victim-unified-id", "gpt-4o-mini", RAW_FILE_ID, "gpt-4o-mini-id"
+    )
+    return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=")
+
+
+def _unified_job_id() -> str:
+    unified = SpecialEnums.LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR.value.format("gpt-4o-mini-id", RAW_JOB_ID)
+    return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=")
+
+
+def _job() -> LiteLLMFineTuningJob:
+    job = LiteLLMFineTuningJob(
+        id=RAW_JOB_ID,
+        created_at=1234567890,
+        fine_tuned_model=None,
+        finished_at=None,
+        hyperparameters={"n_epochs": 1},
+        model="gpt-4o-mini",
+        object="fine_tuning.job",
+        organization_id="org-test",
+        result_files=[],
+        seed=0,
+        status="running",
+        trained_tokens=None,
+        training_file=RAW_FILE_ID,
+        validation_file=None,
+    )
+    job._hidden_params = {}
+    return job
+
+
+class FakeRequest:
+    def __init__(self):
+        self.headers = {}
+        self.query_params = {}
+
+    async def json(self):
+        return {}
+
+
+@dataclass(frozen=True)
+class ManagedResourceAccessCheckerStub:
+    file_access: bool = True
+    object_access: bool = True
+
+    async def can_user_call_unified_file_id(
+        self,
+        unified_file_id: str,
+        user_api_key_dict: UserAPIKeyAuth,
+    ) -> bool:
+        return self.file_access
+
+    async def can_user_call_unified_object_id(
+        self,
+        unified_object_id: str,
+        user_api_key_dict: UserAPIKeyAuth,
+    ) -> bool:
+        return self.object_access
+
+
+class Seams:
+    def __init__(self, router: MagicMock, litellm_calls: dict[str, AsyncMock], logging: MagicMock):
+        self.router = router
+        self.litellm_calls = litellm_calls
+        self.logging = logging
+
+    def assert_no_provider_call(self) -> None:
+        for name, mock in self.litellm_calls.items():
+            assert mock.call_count == 0, f"litellm.{name} was called"
+        for name in ("acreate_fine_tuning_job", "aretrieve_fine_tuning_job", "acancel_fine_tuning_job"):
+            assert getattr(self.router, name).call_count == 0, f"router.{name} was called"
+
+
+@pytest.fixture
+def seams():
+    logging = MagicMock(spec=ProxyLogging)
+    logging.post_call_success_hook = AsyncMock(side_effect=lambda **kw: kw["response"])
+    logging.post_call_failure_hook = AsyncMock()
+    logging.update_request_status = AsyncMock()
+    logging.get_proxy_hook = MagicMock(return_value=None)
+
+    router = MagicMock(spec=Router)
+    router.acreate_fine_tuning_job = AsyncMock(return_value=_job())
+    router.aretrieve_fine_tuning_job = AsyncMock(return_value=_job())
+    router.acancel_fine_tuning_job = AsyncMock(return_value=_job())
+
+    litellm_calls = {
+        name: AsyncMock(return_value=_job())
+        for name in ("acreate_fine_tuning_job", "aretrieve_fine_tuning_job", "acancel_fine_tuning_job")
+    }
+
+    with ExitStack() as stack:
+        stack.enter_context(
+            patch.object(
+                ProxyBaseLLMRequestProcessing,
+                "common_processing_pre_call_logic",
+                AsyncMock(side_effect=lambda self=None, **kw: (self.data if self else {}, MagicMock())),
+            )
+        )
+        stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", MagicMock(return_value={})))
+        for name, mock in litellm_calls.items():
+            stack.enter_context(patch.object(litellm, name, mock))
+        stack.enter_context(patch.object(proxy_server, "llm_router", router))
+        stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging))
+        stack.enter_context(patch.object(proxy_server, "premium_user", True))
+        stack.enter_context(patch.object(proxy_server, "general_settings", {}))
+        stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock()))
+        stack.enter_context(patch.object(proxy_server, "version", "test-version"))
+        stack.enter_context(patch.object(endpoints, "fine_tuning_config", [{"custom_llm_provider": "openai"}]))
+        yield Seams(router=router, litellm_calls=litellm_calls, logging=logging)
+
+
+async def _create(training_file: str, validation_file: str | None = None):
+    return await endpoints.create_fine_tuning_job(
+        request=FakeRequest(),
+        fastapi_response=Response(),
+        fine_tuning_request=LiteLLMFineTuningJobCreate(
+            model="gpt-4o-mini",
+            training_file=training_file,
+            validation_file=validation_file,
+            custom_llm_provider="openai",
+        ),
+        user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
+    )
+
+
+async def _retrieve(job_id: str):
+    return await endpoints.retrieve_fine_tuning_job(
+        request=FakeRequest(),
+        fastapi_response=Response(),
+        fine_tuning_job_id=job_id,
+        custom_llm_provider="openai",
+        user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
+    )
+
+
+async def _cancel(job_id: str):
+    return await endpoints.cancel_fine_tuning_job(
+        request=FakeRequest(),
+        fastapi_response=Response(),
+        fine_tuning_job_id=job_id,
+        user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
+    )
+
+
+@pytest.mark.asyncio
+async def test_create__raw_training_file_rejected_when_managed_files_required(seams):
+    with patch.object(litellm, "require_managed_files", True):
+        with pytest.raises(ProxyException) as exc:
+            await _create(RAW_FILE_ID)
+
+    assert exc.value.code == "400"
+    seams.assert_no_provider_call()
+
+
+@pytest.mark.asyncio
+async def test_create__raw_validation_file_rejected_when_managed_files_required(seams):
+    """The validation file is uploaded and readable exactly like the training file,
+    so a managed training_file must not smuggle a raw validation_file past the guard."""
+    seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub()
+
+    with patch.object(litellm, "require_managed_files", True):
+        with pytest.raises(ProxyException) as exc:
+            await _create(_unified_file_id(), validation_file=RAW_FILE_ID)
+
+    assert exc.value.code == "400"
+    seams.assert_no_provider_call()
+
+
+@pytest.mark.asyncio
+async def test_create__unified_training_file_allowed_when_managed_files_required(seams):
+    seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub()
+
+    with patch.object(litellm, "require_managed_files", True):
+        await _create(_unified_file_id())
+
+    assert seams.router.acreate_fine_tuning_job.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_create__other_teams_unified_training_file_rejected(seams):
+    seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub(file_access=False)
+
+    with patch.object(litellm, "require_managed_files", True):
+        with pytest.raises(ProxyException) as exc:
+            await _create(_unified_file_id())
+
+    assert exc.value.code == "403"
+    seams.assert_no_provider_call()
+
+
+@pytest.mark.asyncio
+async def test_create__raw_training_file_allowed_when_managed_files_not_required(seams):
+    with patch.object(litellm, "require_managed_files", False):
+        await _create(RAW_FILE_ID)
+
+    assert seams.litellm_calls["acreate_fine_tuning_job"].call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_retrieve__raw_job_id_rejected_when_managed_files_required(seams):
+    with patch.object(litellm, "require_managed_files", True):
+        with pytest.raises(ProxyException) as exc:
+            await _retrieve(RAW_JOB_ID)
+
+    assert exc.value.code == "400"
+    seams.assert_no_provider_call()
+
+
+@pytest.mark.asyncio
+async def test_retrieve__unified_job_id_allowed_when_managed_files_required(seams):
+    seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub()
+
+    with patch.object(litellm, "require_managed_files", True):
+        await _retrieve(_unified_job_id())
+
+    assert seams.router.aretrieve_fine_tuning_job.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_cancel__raw_job_id_rejected_when_managed_files_required(seams):
+    with patch.object(litellm, "require_managed_files", True):
+        with pytest.raises(ProxyException) as exc:
+            await _cancel(RAW_JOB_ID)
+
+    assert exc.value.code == "400"
+    seams.assert_no_provider_call()
+
+
+@pytest.mark.asyncio
+async def test_cancel__unified_job_id_allowed_when_managed_files_required(seams):
+    seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub()
+
+    with patch.object(litellm, "require_managed_files", True):
+        await _cancel(_unified_job_id())
+
+    assert seams.router.acancel_fine_tuning_job.call_count == 1

tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py

@@ -3051,3 +3051,147 @@ def test_list_files_key_allowed_openai_model_still_resolves_team_credentials(
         mocker, monkeypatch, _team_openai_plus_global_anthropic_router(), ["team-gpt"]
     )
     assert captured_kwargs.get("api_key") == "team-openai-key"
+
+
+@pytest.mark.parametrize(
+    "http_method, url, patched_litellm_call",
+    [
+        ("get", "/v1/files/file-victim-abc123", "litellm.afile_retrieve"),
+        ("get", "/v1/files/file-victim-abc123/content", "litellm.afile_content"),
+        ("delete", "/v1/files/file-victim-abc123", "litellm.afile_delete"),
+    ],
+)
+def test_require_managed_files_rejects_raw_provider_file_id(
+    mocker: MockerFixture,
+    monkeypatch,
+    llm_router: Router,
+    http_method: str,
+    url: str,
+    patched_litellm_call: str,
+):
+    import litellm.proxy.proxy_server as ps
+    from litellm.proxy._types import LitellmUserRoles
+
+    monkeypatch.setattr("litellm.require_managed_files", True)
+    monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router)
+    setup_proxy_logging_object(monkeypatch, llm_router)
+
+    mock_call = mocker.patch(patched_litellm_call, new=mocker.AsyncMock())
+
+    app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
+        user_role=LitellmUserRoles.INTERNAL_USER, user_id="attacker-user"
+    )
+
+    try:
+        response = getattr(client, http_method)(
+            url, headers={"Authorization": "Bearer test-key"}
+        )
+    finally:
+        app.dependency_overrides.pop(ps.user_api_key_auth, None)
+        monkeypatch.setattr("litellm.require_managed_files", False)
+
+    assert response.status_code == 400, response.text
+    mock_call.assert_not_called()
+
+
+def _unified_managed_file_id() -> str:
+    import base64
+
+    from litellm.types.utils import SpecialEnums
+
+    unified_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
+        "application/json", "victim-unified-id", "gpt-3.5-turbo", "file-victim-abc123", "gpt-3.5-turbo-id"
+    )
+    return base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=")
+
+
+class _ManagedResourceAccessCheckerStub:
+    async def can_user_call_unified_file_id(
+        self,
+        unified_file_id: str,
+        user_api_key_dict: UserAPIKeyAuth,
+    ) -> bool:
+        return True
+
+    async def can_user_call_unified_object_id(
+        self,
+        unified_object_id: str,
+        user_api_key_dict: UserAPIKeyAuth,
+    ) -> bool:
+        return True
+
+
+@pytest.mark.asyncio
+async def test_require_managed_files_allows_owned_unified_managed_file_id(monkeypatch):
+    from litellm.proxy.openai_files_endpoints.common_utils import (
+        validate_managed_id_requirement,
+    )
+
+    monkeypatch.setattr("litellm.require_managed_files", True)
+
+    await validate_managed_id_requirement(
+        resource_id=_unified_managed_file_id(),
+        resource_kind="file",
+        user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="owner-user"),
+        managed_files_obj=_ManagedResourceAccessCheckerStub(),
+    )
+
+
+@pytest.mark.asyncio
+async def test_managed_file_id_requirement_is_opt_in(monkeypatch):
+    from litellm.proxy.openai_files_endpoints.common_utils import (
+        validate_managed_id_requirement,
+    )
+
+    monkeypatch.setattr("litellm.require_managed_files", False)
+
+    await validate_managed_id_requirement(
+        resource_id="file-victim-abc123",
+        resource_kind="file",
+        user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
+        managed_files_obj=None,
+    )
+
+
+def test_raw_provider_file_id_retrieve_allowed_when_managed_files_not_required(
+    mocker: MockerFixture, monkeypatch, llm_router: Router
+):
+    import litellm.proxy.proxy_server as ps
+    from litellm.proxy._types import LitellmUserRoles
+
+    monkeypatch.setattr("litellm.require_managed_files", False)
+    monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router)
+    setup_proxy_logging_object(monkeypatch, llm_router)
+
+    mock_retrieve = mocker.patch(
+        "litellm.afile_retrieve",
+        new=mocker.AsyncMock(
+            return_value=OpenAIFileObject(
+                id="file-victim-abc123",
+                object="file",
+                bytes=3,
+                created_at=1234567890,
+                filename="test.txt",
+                purpose="user_data",
+                status="uploaded",
+            )
+        ),
+    )
+
+    app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
+        user_role=LitellmUserRoles.INTERNAL_USER, user_id="some-user"
+    )
+
+    try:
+        response = client.get(
+            "/v1/files/file-victim-abc123", headers={"Authorization": "Bearer test-key"}
+        )
+    finally:
+        app.dependency_overrides.pop(ps.user_api_key_auth, None)
+
+    assert response.status_code == 200, response.text
+    mock_retrieve.assert_called_once()

tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py

@@ -0,0 +1,146 @@
+"""
+require_managed_files enforcement for litellm/proxy/vector_store_files_endpoints/endpoints.py
+
+Every vector-store file route (create, retrieve, content, update, delete) resolves its
+caller-supplied file id through _update_request_data_with_managed_file_id before the
+provider call, so the guard lives there once and covers all five.
+
+A raw or forged managed-looking file id has no ownership row, so without the guard it
+is attached to a vector store or read back under shared provider credentials.
+"""
+
+import base64
+import os
+import sys
+from dataclasses import dataclass
+from typing import Literal
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../../../.."))
+
+from fastapi import HTTPException
+
+import litellm
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.vector_store_files_endpoints.endpoints import (
+    _update_request_data_with_managed_file_id,
+)
+from litellm.types.utils import SpecialEnums
+
+RAW_FILE_ID = "file-victim-abc123"
+CALLER = UserAPIKeyAuth(api_key="sk-test", user_id="attacker-user", team_id="team-b")
+
+
+@dataclass(frozen=True)
+class ManagedResourceAccessCheckerStub:
+    file_access: Literal["allow", "deny", "missing"]
+
+    async def can_user_call_unified_file_id(
+        self,
+        unified_file_id: str,
+        user_api_key_dict: UserAPIKeyAuth,
+    ) -> bool:
+        if self.file_access == "missing":
+            raise HTTPException(status_code=404, detail=f"File not found: {unified_file_id}")
+        return self.file_access == "allow"
+
+    async def can_user_call_unified_object_id(
+        self,
+        unified_object_id: str,
+        user_api_key_dict: UserAPIKeyAuth,
+    ) -> bool:
+        return False
+
+
+def _unified_file_id() -> str:
+    unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
+        "application/json", "victim-unified-id", "gpt-4o-mini", RAW_FILE_ID, "gpt-4o-mini-id"
+    )
+    return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=")
+
+
+async def _resolve(
+    file_id: str,
+    file_access: Literal["allow", "deny", "missing"] = "allow",
+):
+    return await _update_request_data_with_managed_file_id(
+        data={"vector_store_id": "vs-test", "file_id": file_id},
+        file_id=file_id,
+        request=MagicMock(headers={}, query_params={}),
+        user_api_key_dict=CALLER,
+        managed_files_obj=ManagedResourceAccessCheckerStub(file_access=file_access),
+        llm_router=None,
+    )
+
+
+@pytest.mark.asyncio
+async def test_raw_file_id_rejected_when_managed_files_required():
+    with patch.object(litellm, "require_managed_files", True):
+        with pytest.raises(HTTPException) as exc:
+            await _resolve(RAW_FILE_ID)
+
+    assert exc.value.status_code == 400
+
+
+@pytest.mark.asyncio
+async def test_model_encoded_file_id_rejected_when_managed_files_required():
+    """encode_file_id_with_model output is client-forgeable and carries no ownership
+    row, so it is not a managed file id."""
+    from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model
+
+    encoded = encode_file_id_with_model(RAW_FILE_ID, "gpt-4o-mini", id_type="file")
+
+    with patch.object(litellm, "require_managed_files", True):
+        with pytest.raises(HTTPException) as exc:
+            await _resolve(encoded)
+
+    assert exc.value.status_code == 400
+
+
+@pytest.mark.asyncio
+async def test_forged_unified_file_id_rejected_without_ownership_record():
+    forged_id = _unified_file_id()
+    data = {"vector_store_id": "vs-test", "file_id": forged_id}
+
+    with patch.object(litellm, "require_managed_files", True):
+        with pytest.raises(HTTPException) as exc:
+            await _update_request_data_with_managed_file_id(
+                data=data,
+                file_id=forged_id,
+                request=MagicMock(headers={}, query_params={}),
+                user_api_key_dict=CALLER,
+                managed_files_obj=ManagedResourceAccessCheckerStub(file_access="missing"),
+                llm_router=None,
+            )
+
+    assert exc.value.status_code == 404
+    assert data["file_id"] == forged_id
+
+
+@pytest.mark.asyncio
+async def test_other_teams_unified_file_id_rejected():
+    with patch.object(litellm, "require_managed_files", True):
+        with pytest.raises(HTTPException) as exc:
+            await _resolve(_unified_file_id(), file_access="deny")
+
+    assert exc.value.status_code == 403
+
+
+@pytest.mark.asyncio
+async def test_owned_unified_file_id_allowed_when_managed_files_required():
+    with patch.object(litellm, "require_managed_files", True):
+        data, original = await _resolve(_unified_file_id())
+
+    assert original == _unified_file_id()
+    assert data["file_id"] == RAW_FILE_ID
+
+
+@pytest.mark.asyncio
+async def test_raw_file_id_allowed_when_managed_files_not_required():
+    with patch.object(litellm, "require_managed_files", False):
+        data, original = await _resolve(RAW_FILE_ID)
+
+    assert original is None
+    assert data["file_id"] == RAW_FILE_ID

tests/test_litellm/test_budget_ratchet_check.py

@@ -68,6 +68,50 @@ def test_new_rule_in_head_is_clean():
     assert ratchet.regressions_for("b.json", {}, {"new-rule": _spec_of(5)}) == []
 
 
+def test_dropped_rule_that_graduated_to_a_hard_failing_config_is_clean():
+    base = {"UP006": _spec_of(0)}
+    assert ratchet.regressions_for("b.json", base, {}, graduated=("UP006",)) == []
+
+
+def test_graduation_matches_by_prefix_like_ruff_selectors_do():
+    base = {"ANN202": _spec_of(865)}
+    assert ratchet.regressions_for("b.json", base, {}, graduated=("ANN",)) == []
+
+
+def test_an_unrelated_graduation_does_not_excuse_a_dropped_rule():
+    base = {"C901": _spec_of(3)}
+    regs = ratchet.regressions_for("b.json", base, {}, graduated=("UP006", "SIM118"))
+    assert [r.rule for r in regs] == ["C901"]
+    assert "dropped" in regs[0].detail
+
+
+def test_graduation_never_excuses_a_raised_limit():
+    base = {"UP006": _spec_of(0)}
+    regs = ratchet.regressions_for("b.json", base, {"UP006": _spec_of(7)}, graduated=("UP006",))
+    assert [r.rule for r in regs] == ["UP006"]
+    assert "0 -> 7" in regs[0].detail
+
+
+def test_graduated_selectors_come_from_the_paired_ruff_config():
+    selectors = ratchet.graduated_selectors("ruff-strict-budget.json")
+    assert "UP006" in selectors
+    assert "ANN" not in selectors
+
+
+def test_budgets_without_a_paired_config_can_never_graduate():
+    assert ratchet.graduated_selectors("type-discipline-budget.json") == ()
+    assert ratchet.graduated_selectors("basedpyright-code-budget.json") == ()
+
+
+def test_a_selector_the_config_also_ignores_does_not_count_as_graduated():
+    lint = {"ignore": ["UP006"], "extend-select": ["UP006", "SIM118"]}
+    assert ratchet.selectors_hard_failed_by(lint) == ("SIM118",)
+
+
+def test_selectors_hard_failed_by_reads_a_config_with_no_ignore_list():
+    assert ratchet.selectors_hard_failed_by({"extend-select": ["UP006"]}) == ("UP006",)
+
+
 def test_deleted_budget_file_is_a_regression():
     regs = ratchet.regressions_for("b.json", {"LIT006": _spec_of(1)}, None)
     assert [r.rule for r in regs] == ["*"]

tests/test_litellm/test_conftest_isolation.py

@@ -1,9 +1,15 @@
 import litellm
+from litellm import Router
+from litellm import router as litellm_router_module
 from litellm import utils as litellm_utils_module
 
 CANARY_MODEL = "conftest-isolation-canary-model"
 
 
+class _CanaryRouterHolder:
+    router: Router | None = None
+
+
 def test_register_model_ledger_entry_is_scoped_to_this_test():
     litellm.register_model({CANARY_MODEL: {"litellm_provider": "openai", "input_cost_per_token": 0.001}})
     assert CANARY_MODEL in litellm_utils_module._runtime_registered_model_cost
@@ -11,3 +17,20 @@ def test_register_model_ledger_entry_is_scoped_to_this_test():
 
 def test_register_model_ledger_entry_was_rolled_back():
     assert CANARY_MODEL not in litellm_utils_module._runtime_registered_model_cost
+
+
+def test_live_router_membership_is_scoped_to_this_test():
+    _CanaryRouterHolder.router = Router(
+        model_list=[
+            {
+                "model_name": "conftest-isolation-canary-router",
+                "litellm_params": {"model": "openai/conftest-isolation-canary-backend", "api_key": "sk-canary"},
+            }
+        ]
+    )
+    assert _CanaryRouterHolder.router in litellm_router_module._live_routers
+
+
+def test_live_router_membership_was_rolled_back():
+    assert _CanaryRouterHolder.router is not None
+    assert _CanaryRouterHolder.router not in litellm_router_module._live_routers

tests/test_litellm/test_ruff_strict_gate.py

@@ -1,16 +1,24 @@
 import importlib.util
+import json
+import re
+import shutil
 import subprocess
+import sys
+import tomllib
 from pathlib import Path
 
 import pytest
 
-_MODULE_PATH = Path(__file__).resolve().parents[2] / "scripts" / "ruff_strict_gate.py"
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+_MODULE_PATH = _REPO_ROOT / "scripts" / "ruff_strict_gate.py"
 _spec = importlib.util.spec_from_file_location("ruff_strict_gate", _MODULE_PATH)
 gate = importlib.util.module_from_spec(_spec)
 _spec.loader.exec_module(gate)
 
 Violation = gate.Violation
 
+_ENABLED_BY_RUFF_DEFAULTS = frozenset({"F401"})
+
 
 def rule(name, limit):
     return {name: {"limit": limit}}
@@ -151,3 +159,213 @@ def test_base_point_mid_merge_advances_to_the_merged_in_base_tip(tmp_path):
     repo, _, base_tip = _branched_repo(tmp_path)
     _git(repo, "merge", "--no-commit", "--no-ff", "main")
     assert gate.resolve_base_point("main", cwd=repo) == base_tip
+
+
+def _lint_section(config_name: str) -> dict:
+    return tomllib.loads((_REPO_ROOT / config_name).read_text())["lint"]
+
+
+def _base_external() -> tuple[str, ...]:
+    return tuple(_lint_section("ruff.toml")["external"])
+
+
+def _strict_external() -> tuple[str, ...]:
+    return tuple(_lint_section("ruff-strict.toml")["external"])
+
+
+def _strict_selected() -> frozenset:
+    return frozenset(_lint_section("ruff-strict.toml")["select"])
+
+
+def _prefix_covered(code: str, prefixes: tuple[str, ...]) -> bool:
+    return any(code.startswith(prefix) for prefix in prefixes)
+
+
+def _selected_by_the_normal_config() -> frozenset:
+    return frozenset(_lint_section("ruff.toml")["extend-select"]) | _ENABLED_BY_RUFF_DEFAULTS
+
+
+def _budgeted_rules() -> frozenset:
+    return frozenset(json.loads((_REPO_ROOT / "ruff-strict-budget.json").read_text()))
+
+
+def _ruff_binary() -> str | None:
+    beside_interpreter = Path(sys.executable).with_name("ruff")
+    return str(beside_interpreter) if beside_interpreter.exists() else shutil.which("ruff")
+
+
+_RUFF = _ruff_binary()
+_needs_ruff = pytest.mark.skipif(_RUFF is None, reason="ruff is not installed in this environment")
+
+
+def _ruff_output_for_noqa(code: str, *extra_args: str) -> str:
+    proc = subprocess.run(
+        [
+            _RUFF,
+            "check",
+            "--no-cache",
+            "--stdin-filename",
+            "litellm/types/_external_probe.py",
+            *extra_args,
+            "-",
+        ],
+        cwd=_REPO_ROOT,
+        input=f"def _probe(x: int):  # noqa: {code}\n    return x\n",
+        capture_output=True,
+        text=True,
+    )
+    return proc.stdout
+
+
+def test_every_strict_gate_rule_is_protected_from_base_ruf100():
+    unprotected = frozenset(
+        selector
+        for selector in _strict_selected()
+        if not _prefix_covered(selector, _base_external())
+        and selector not in _selected_by_the_normal_config()
+    )
+    assert unprotected == frozenset(), (
+        f"`ruff check` deletes any `# noqa` naming {sorted(unprotected)} as unused, so suppressing "
+        "one of those strict-gate rules breaks lint. Cover them in ruff.toml's lint.external or "
+        "enable them in its lint.extend-select."
+    )
+
+
+def test_every_selected_rule_keeps_stale_noqa_detection_somewhere():
+    policed_by_strict = frozenset(
+        selector
+        for selector in _strict_selected()
+        if not _prefix_covered(selector, _strict_external())
+    )
+    policed_by_base = frozenset(
+        selector
+        for selector in _selected_by_the_normal_config()
+        if not _prefix_covered(selector, _base_external())
+    )
+    shadowed = (
+        _strict_selected() | _selected_by_the_normal_config()
+    ) - policed_by_strict - policed_by_base
+    assert shadowed == frozenset(), (
+        f"no config's RUF100 can ever report a stale `# noqa` for {sorted(shadowed)}: every config "
+        "that selects each of them also shadows it with an external entry. Narrow the external "
+        "entry in ruff.toml or ruff-strict.toml."
+    )
+
+
+_BASE_OWNED_FAMILY = re.compile(r"E[479]\d+|F\d+|T20\d+")
+_BASE_OWNED_SINGLES = frozenset({"PGH004", "RUF008", "RUF009", "RUF100"})
+
+
+@pytest.fixture(scope="module")
+def all_ruff_rule_codes() -> frozenset:
+    listing = subprocess.run(
+        [_RUFF, "rule", "--all", "--output-format", "json"],
+        capture_output=True,
+        text=True,
+    )
+    assert listing.returncode == 0, listing.stderr
+    return frozenset(
+        entry["code"] for entry in json.loads(listing.stdout) if "Removed" not in entry["status"]
+    )
+
+
+@_needs_ruff
+def test_every_base_owned_rule_is_external_or_selected_in_the_strict_config(all_ruff_rule_codes):
+    base_owned = frozenset(
+        code
+        for code in all_ruff_rule_codes
+        if _BASE_OWNED_FAMILY.fullmatch(code) or code in _BASE_OWNED_SINGLES
+    )
+    stranded = frozenset(
+        code
+        for code in base_owned
+        if code not in _strict_selected() and not _prefix_covered(code, _strict_external())
+    )
+    assert stranded == frozenset(), (
+        f"the strict gate's RUF100 reads a valid `# noqa` for {sorted(stranded)} as unused, the "
+        "spurious-breach trap ruff-strict.toml's external override exists to prevent. Cover them "
+        "there."
+    )
+    double_booked = frozenset(
+        code
+        for code in base_owned
+        if code in _strict_selected() and _prefix_covered(code, _strict_external())
+    )
+    assert double_booked == frozenset(), (
+        f"{sorted(double_booked)} are selected by the strict config yet shadowed by its external "
+        "list, so their stale suppressions can never be reported. Narrow the external entry in "
+        "ruff-strict.toml."
+    )
+
+
+def test_every_budgeted_rule_is_one_the_gate_actually_measures():
+    selectors = tuple(_lint_section("ruff-strict.toml")["select"])
+    unmeasured = frozenset(code for code in _budgeted_rules() if not code.startswith(selectors))
+    assert unmeasured == frozenset(), (
+        f"the gate never counts {sorted(unmeasured)}, so their ceilings are dead config that reads "
+        "as coverage. Either select them in ruff-strict.toml or drop them from the budget."
+    )
+
+
+@_needs_ruff
+def test_every_strict_selected_rule_is_budgeted_or_hard_failed_by_the_base_config(all_ruff_rule_codes):
+    strict_enabled = frozenset(
+        code
+        for code in all_ruff_rule_codes
+        if code.startswith(tuple(_lint_section("ruff-strict.toml")["select"]))
+    )
+    base_hard_failed = tuple(_lint_section("ruff.toml")["extend-select"])
+    unpoliced = frozenset(
+        code
+        for code in strict_enabled
+        if code not in _budgeted_rules()
+        and not code.startswith(base_hard_failed)
+        and code not in _ENABLED_BY_RUFF_DEFAULTS
+    )
+    assert unpoliced == frozenset(), (
+        f"nothing enforces {sorted(unpoliced)}: the gate skips rules missing from the budget, and "
+        "the base config does not hard-fail them. Re-add a budget ceiling or graduate them into "
+        "ruff.toml's lint.extend-select."
+    )
+
+
+@_needs_ruff
+def test_a_noqa_for_a_strict_gate_rule_survives_the_normal_ruff_run():
+    assert "RUF100" not in _ruff_output_for_noqa("ANN202")
+
+
+@_needs_ruff
+def test_the_external_list_is_what_saves_that_noqa():
+    assert "RUF100" in _ruff_output_for_noqa("ANN202", "--config", "lint.external=[]")
+
+
+@_needs_ruff
+def test_a_stale_noqa_for_a_locally_enabled_rule_is_still_reported():
+    assert "RUF100" in _ruff_output_for_noqa("F401")
+
+
+def _ruff_output_for_source(source: str) -> str:
+    proc = subprocess.run(
+        [_RUFF, "check", "--no-cache", "--stdin-filename", "litellm/types/_graduate_probe.py", "-"],
+        cwd=_REPO_ROOT,
+        input=source,
+        capture_output=True,
+        text=True,
+    )
+    return proc.stdout
+
+
+_DEPRECATED_TYPING_ALIAS = "from typing import List  # noqa: UP035\n\n\ndef _probe(x: List[int]) -> None: ...\n"
+
+
+@_needs_ruff
+def test_a_graduated_rule_now_fails_the_normal_ruff_run_instead_of_waiting_for_the_gate():
+    assert "UP006" in _ruff_output_for_source(_DEPRECATED_TYPING_ALIAS)
+
+
+@_needs_ruff
+def test_a_graduated_rule_can_still_be_suppressed_without_tripping_unused_noqa():
+    suppressed = _DEPRECATED_TYPING_ALIAS.replace("...\n", "...  # noqa: UP006\n")
+    output = _ruff_output_for_source(suppressed)
+    assert "UP006" not in output
+    assert "RUF100" not in output

type-discipline-budget.json

@@ -1,9 +1,9 @@
 {
   "LIT001": {
-    "limit": 23235
+    "limit": 23149
   },
   "LIT002": {
-    "limit": 27176
+    "limit": 27166
   },
   "LIT003": {
     "limit": 269
@@ -15,7 +15,7 @@
     "limit": 0
   },
   "LIT006": {
-    "limit": 1091
+    "limit": 1086
   },
   "LIT007": {
     "limit": 0
@@ -27,7 +27,7 @@
     "limit": 0
   },
   "LIT010": {
-    "limit": 16769
+    "limit": 16760
   },
   "LIT011": {
     "limit": 5598