diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 32b8eb3d4d0..0385f7a96e7 100644 --- a/basedpyright-code-budget.json +++ b/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 diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index d8318962633..f0914240f79 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/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 @@ async def store_unified_file_id( 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 @@ async def store_unified_file_id( 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 @@ async def store_unified_object_id( 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 @@ async def store_unified_object_id( 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 @@ async def get_unified_file_id( 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 @@ async def delete_unified_file_id( 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 @@ async def delete_unified_file_id( 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 @@ async def can_user_call_unified_file_id( 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 @@ async def list_user_batches( 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 @@ async def list_user_batches( ) 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 @@ async def list_user_batches( 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 @@ async def _resolve_listed_batch( 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 @@ async def get_user_created_file_ids( 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 @@ async def get_user_created_file_ids( 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 @@ async def check_managed_file_id_access( ) 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 @@ async def check_file_ids_access( 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 @@ async def async_pre_call_hook( ### 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 @@ async def async_pre_call_hook( # 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 @@ async def async_pre_call_hook( 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 @@ async def async_pre_call_hook( 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 @@ async def async_pre_call_hook( ): 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 @@ async def async_pre_call_hook( 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 @@ async def async_pre_call_deployment_hook( 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 @@ def get_file_ids_from_messages(self, messages: List[AllMessageValues]) -> List[s 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 @@ def get_file_ids_from_responses_input( 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 @@ def get_file_ids_from_responses_tools( 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 @@ async def check_vector_store_ids_access( ), ) - 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 @@ async def get_model_file_id_mapping( # 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 @@ async def create_file_for_each_model( 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 @@ async def acreate_file( 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 @@ async def acreate_file( # 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 @@ async def return_unified_file_id( ) # 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 @@ async def return_unified_file_id( 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 @@ def get_model_id_from_unified_file_id(self, file_id: str) -> str: 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 @@ async def async_post_call_success_hook( # 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 @@ async def async_post_call_success_hook( 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 @@ async def async_post_call_success_hook( 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 @@ async def async_post_call_success_hook( """ ## 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 @@ async def async_post_call_success_hook( 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 @@ async def async_post_call_success_hook( 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 @@ async def afile_retrieve( ) 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 @@ def _is_batch_polling_enabled(self) -> bool: 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 @@ async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, An # 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 @@ async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, An 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 @@ async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, An 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 @@ async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, An 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 @@ async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, An } ) 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 @@ async def _check_file_deletion_allowed(self, file_id: str) -> None: 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 @@ async def afile_delete( 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 @@ async def afile_content( 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 @@ async def afile_content( # 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 @@ async def _convert_storage_files_to_base64( # 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 @@ async def _convert_storage_files_to_base64( 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 diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 339da998d56..024e8c179c2 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -6,16 +6,33 @@ """ 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 @@ def _remove_none_values(obj: Any) -> Any: 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 @@ def _params_to_dict(params: Any) -> dict[str, Any]: 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 @@ def _params_to_dict(params: Any) -> dict[str, Any]: 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 @@ async def _poll_for_completion( }, ) 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 @@ async def _poll_for_completion( 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 @@ async def _send_and_poll_raw( 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 @@ async def _send_and_poll_raw( }, ) 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 @@ async def _send_and_poll_raw( 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 @@ async def send_non_streaming_request( 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 @@ async def send_and_get_raw_response( @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 @@ def _transform_to_a2a_response( } @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 @@ def _extract_response_text(response_data: dict[str, Any]) -> tuple[str, str, lis 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 @@ async def fake_streaming_from_response( 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 @@ async def fake_streaming_from_response( # 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 @@ async def fake_streaming_from_response( # 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) - - 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, - } - ], + 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 = _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 @@ async def fake_streaming_from_response( 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) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 9748db2dcd2..7abbf0c96e5 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -42,7 +42,7 @@ 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 @@ async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type 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 @@ async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type 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 @@ def from_config_yaml(cls, config: WebSearchInterceptionConfig) -> "WebSearchInte ) @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 @@ async def _build_anthropic_request_patch( 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]", diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 858d10df53b..d68bdc4a250 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/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 @@ 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 @@ def __init__( 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 @@ def __init__( ] ) _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 @@ def _collect_tool_calls_from_response_done(self, event_obj: dict | OpenAIRealtim 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 @@ def _enforce_transcription_session_model(self, message: str) -> str: 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 @@ def _should_buffer_client_message_until_setup(self, message: str) -> bool: 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 @@ def _normalize_event_for_ga_client(self, event: dict) -> dict: 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 @@ def _cache_session_configuration_request(self, transformed_message: str) -> None 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 @@ def _maybe_inject_guardrail_auto_response_disable(self, setup_message: str) -> s 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 @@ async def run_realtime_guardrails( 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 @@ async def run_realtime_guardrails( 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 @@ async def run_realtime_guardrails( 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 @@ async def _handle_provider_config_message(self, raw_response) -> None: 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 @@ def _detect_beta_header(websocket: Any) -> bool: 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 @@ def _remap_beta_session_to_ga(session: dict) -> dict: 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 @@ async def client_ack_messages(self): # 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 @@ async def client_ack_messages(self): # 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 @@ async def client_ack_messages(self): 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): diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 3f967e29002..886ba6a3a18 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -3,7 +3,7 @@ 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.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 @@ ) +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 @@ def _sort_chunks(self, chunks: list) -> list: 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 @@ def update_model_response_with_hidden_params( @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 @@ def get_combined_function_call_content(self, function_call_chunks: list[dict[str ) 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 @@ def get_combined_content( 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 @@ def _flush_thinking_block() -> None: 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 @@ def get_combined_audio_content(self, chunks: list[dict[str, Any]]) -> ChatComple 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 @@ def get_combined_audio_content(self, chunks: list[dict[str, Any]]) -> ChatComple 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 @@ def count_reasoning_tokens(self, response: ModelResponse) -> int | None: 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 @@ def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse | ModelResponseSt 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 @@ def _calculate_usage_per_chunk( "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 @@ def _calculate_usage_per_chunk( @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 @@ def _reset_anthropic_cursor_completion_tokens( 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 @@ def calculate_usage( 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 diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index a9751489473..36f3e875a7e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/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 @@ # 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 @@ def _normalize_reasoning_effort( 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 @@ def _prepare_completion_kwargs( 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 @@ def _prepare_completion_kwargs( Logging as LiteLLMLoggingObject, ) - request_data: Final = { + request_data: Final[dict[str, object]] = { "model": model, "messages": messages, "max_tokens": max_tokens, @@ -528,7 +536,7 @@ def _prepare_completion_kwargs( @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 @@ async def async_anthropic_messages_handler( 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 @@ async def async_anthropic_messages_handler( 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 @@ async def async_anthropic_messages_handler( @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 @@ def anthropic_messages_handler( 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 @@ def anthropic_messages_handler( 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, diff --git a/litellm/llms/azure/assistants.py b/litellm/llms/azure/assistants.py index 671e4633af4..f7b419405ac 100644 --- a/litellm/llms/azure/assistants.py +++ b/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 .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 @@ async def a_add_message( 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 @@ def add_message( 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 @@ async def async_create_thread( 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 @@ def create_thread( 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 @@ async def async_get_thread( 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 @@ def get_thread( 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 @@ def run_thread_stream( 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 @@ def run_thread_stream( "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 diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 76618e0f742..e285feb77ee 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -109,7 +109,7 @@ def _connection_error_message(exc: BaseException) -> str: ############ 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, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 1c6ad84ddb4..49a1f1314f0 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -13,9 +13,9 @@ 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 @@ def _write_byok_cred_cache(user_id: str, server_id: str, credential: str | None) ) # 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 @@ def _normalize_resource_contents(contents: list) -> list[ReadResourceContents]: 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 @@ def _gateway_create_initialization_options( _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 @@ async def _purge_expired_stateful_session_auth_contexts( ) -> 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 @@ async def _enforce_stateful_session_cap_for_owner(owner: str) -> bool: 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 @@ async def handle_list_tools() -> "ListToolsResult | list[Tool]": 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 @@ async def forward_progress(progress: float, total: float | None): 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 @@ async def _build_virtual_call_logging_obj( 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 @@ async def _dispatch_virtual_mcp_tool( ) @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 @@ def _client_has_passthrough_authorization( 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 @@ async def _prefetch_oauth_creds_for_user( 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 @@ async def _get_tools_from_mcp_servers( 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 @@ async def _get_tools_from_mcp_servers( 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 @@ async def _check_byok_credential( 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 @@ async def execute_mcp_tool( _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 @@ async def execute_mcp_tool( 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 @@ async def _run_post_mcp_call_guardrails( 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 @@ async def _fire_mcp_tool_call_logging( @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 @@ async def call_mcp_tool( 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 @@ async def mcp_read_resource( 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 @@ def _get_standard_logging_mcp_tool_call( 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 @@ async def _handle_managed_mcp_tool( 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 @@ def _get_session_id_from_scope(scope: Scope) -> str | None: 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 @@ async def _read_request_body_for_routing( 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 @@ async def _handle_stale_mcp_session( 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 @@ async def _raise_preemptive_401_for_unauthenticated_servers( 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 @@ def _get_forwarded_auth_from_scope(scope: Scope) -> str | None: ``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 @@ async def _check_passthrough_upstream_auth( 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 @@ async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) ) # 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 @@ async def _dispatch() -> None: 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 @@ async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None: ) # 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 @@ def _wrap_send_with_stateful_session_auth_context( ) -> 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) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 761d8aabc8a..b68d4a68b79 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -6,10 +6,10 @@ 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 @@ 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 @@ 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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py index 068a3ecf31b..facb822d00d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py @@ -19,7 +19,7 @@ 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 @@ def __init__( 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 @@ async def async_post_call_success_hook( 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 @@ async def async_post_call_streaming_iterator_hook( 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 @@ async def async_post_call_streaming_iterator_hook( 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 @@ async def async_post_call_streaming_iterator_hook( 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 @@ def _build_block_payload(self, context: _ScanContext, verdict: _CiscoVerdict) -> "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 @@ def _extract_streaming_chunk_scan_text(cls, chunks: list[Any]) -> str: 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 @@ def _log_decision( 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 @@ async def _inspect_chat( 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 @@ def _build_chat_payload( 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 @@ def _build_chat_payload( 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 @@ def _build_metadata( 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 @@ def _build_metadata( 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 @@ def _build_config(self) -> dict[str, Any]: 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 @@ def _normalize_rule(rule: object) -> dict[str, Any]: 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 @@ def _finalize_inspection( 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 @@ def _stash_verdict_on_request(request_data: dict, context: _ScanContext, verdict @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 @@ def _sanitize_response_for_logging( 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 @@ def _unwrap_verdict_envelope(cls, inspect_response: dict[str, Any]) -> dict[str, @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 @@ def _normalize_action(raw_action: str) -> str: @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 @@ def _extract_sanitized_text( @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 @@ def _apply_redaction( 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 @@ def _redact_chat_input( 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 @@ def _redact_responses_instructions( 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 @@ def _redact_responses_instructions( 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 @@ def _instruction_text_from_messages(cls, messages: list[dict[str, Any]]) -> str 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 @@ def _redact_chat_output( 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 @@ def _redact_chat_output( 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 @@ def _redact_model_response_choices( 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 @@ def _redact_responses_api_output( 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 @@ def _redact_responses_api_output( @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 @@ def _handle_api_error( 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 @@ def _field(obj: object, key: str) -> object: 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 [] diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index c29da89b15f..5bbb01c6c8e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -8,8 +8,8 @@ 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 @@ # 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 @@ 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 @@ async def async_pre_call_hook( 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 @@ async def async_pre_call_hook( 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 @@ async def async_moderation_hook( 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 @@ async def async_post_call_success_hook( data: dict, user_api_key_dict: UserAPIKeyAuth, response, - ) -> Any: + ) -> object: """ Runs on response from LLM API call @@ -271,7 +301,9 @@ async def async_post_call_success_hook( ) 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 @@ async def async_post_call_success_hook( 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 @@ async def _handle_streaming_block( @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 @@ async def _emit_streaming_http_error( 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 @@ async def _emit_streaming_http_error( 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 @@ def _build_transform_chunk( 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 @@ async def _run_incremental_transform_stream( 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 @@ async def _run_incremental_transform_stream( 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 @@ def _round(reference_chunk: Any, is_final: bool) -> AsyncGenerator[Any, None]: 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 @@ async def _inspect_full_response_for_block( 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 @@ def _chunk_carries_text(item: Any) -> bool: @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 @@ def _tool_call_passthrough_chunk( 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 @@ def _tool_call_passthrough_chunk( ) @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 @@ async def async_post_call_streaming_iterator_hook( # 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 @@ def _streaming_flag(name: str, default: Any) -> Any: # 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 @@ def _streaming_flag(name: str, default: Any) -> Any: # 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). diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index dd61cad15a1..a64ed764a67 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -26,7 +26,8 @@ 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 @@ 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 @@ async def async_pre_call_hook( 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 @@ def _is_anthropic_model(self, model: str) -> bool: 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 @@ async def async_post_call_success_deployment_hook( # 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 @@ async def _execute_code_loop_messages_api( 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 @@ async def _execute_code_loop_messages_api( 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 @@ async def _execute_code_loop_messages_api( # 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 @@ async def _execute_code( 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 @@ async def _execute_code( # 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 @@ async def _execute_code( 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 @@ async def _execute_code_loop( 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 @@ async def _execute_code_loop( 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 @@ async def _execute_code_loop( 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 @@ async def _execute_code_tool( 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 @@ async def _execute_code_tool( # 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 @@ async def _execute_code_tool( 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 @@ def _attach_files_to_response( 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: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f2cb1124fa0..2de1d177b33 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -18,7 +18,7 @@ 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 @@ async def find_many( 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 @@ async def update( ) -> _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 diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 71407c89813..c2087005863 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -13,9 +13,9 @@ 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.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"] - - if "litellm_params" in merged_deployment_dict: - prisma_compatible_model_dict["litellm_params"] = json.dumps(merged_deployment_dict["litellm_params"]) + for key, value in merged_model_info.items(): + if isinstance(value, datetime.datetime): + merged_model_info[key] = value.isoformat() - 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 @@ async def allow_team_model_action( 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 @@ async def can_user_make_model_call( ) -> 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, ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 0b99879f9fe..97f494c51de 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -15,11 +15,11 @@ 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.proxy._types import ( UI_TEAM_ID, BlockTeamRequest, + BudgetNewRequest, CommonProxyErrors, DeleteTeamRequest, LiteLLM_AccessGroupTable, @@ -156,6 +157,15 @@ _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 @@ async def count( 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 @@ async def create_team_member_budget_table( 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 @@ async def upsert_team_member_budget_table( 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 @@ async def backfill_team_member_budget_entries( ) 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. diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 44abc56713f..a2c50590dd5 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -16,9 +16,22 @@ 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 @@ } ) +_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 @@ async def upsert_sso_user( 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 @@ async def create_litellm_team_from_sso_group( 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 @@ async def get_redirect_response_from_openid( # 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 @@ async def prepare_token_exchange_parameters( 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 @@ def _validate_token_response(response: "httpx.Response") -> dict: 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 @@ async def get_group_ids_from_service_principal( 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", []): diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 382df608a0c..08bb8698cac 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -4,11 +4,11 @@ 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 @@ 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 @@ class UISettingsResponse(SettingsResponse): # 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 @@ class UISettingsResponse(SettingsResponse): _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": { diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 6b40281198c..2b037bef795 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -10,10 +10,17 @@ 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 @@ 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: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index f54023836e5..b2d065ea23b 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -4,8 +4,8 @@ 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.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 @@ 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 @@ def transform_responses_api_request_to_chat_completion_request( 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 @@ def _transform_response_input_param_to_chat_completion_message( 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 @@ def _transform_response_input_param_to_chat_completion_message( # 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 @@ def _deduplicate_tool_call_output_messages( 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 @@ def _ensure_tool_call_output_has_corresponding_tool_call( 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 @@ def _find_previous_assistant_idx(messages: list[Any], current_idx: int) -> int | 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 @@ def _recover_tool_call_id_from_assistant(assistant_message: Any, message: Any) - 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 @@ def _get_tool_calls_list(assistant_message: Any) -> list[Any]: 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 @@ def _check_tool_call_exists(tool_calls: list[Any], tool_call_id: str) -> bool: 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 @@ def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: list[Any]) -> di 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 @@ def _get_mapping_or_attr_value(obj: Any, key: str, default: Any = None) -> Any: 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 @@ def _get_mapping_or_attr_value(obj: Any, key: str, default: Any = None) -> Any: @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 @@ def _create_tool_call_chunk( ) @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 @@ def _normalize_tool_use_definition(tool_use_definition: Any, tool_call_id: str) 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 @@ def _normalize_tool_use_definition(tool_use_definition: Any, tool_call_id: str) 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 @@ def _add_tool_call_to_assistant(assistant_message: Any, tool_call_chunk: ChatCom 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 @@ def _ensure_tool_results_have_corresponding_tool_calls( | ChatCompletionMessageToolCall | Message ], - tools: list[Any] | None = None, + tools: Sequence[object] | None = None, ) -> list[ AllMessageValues | GenericChatCompletionMessage @@ -851,7 +907,7 @@ def _ensure_tool_results_have_corresponding_tool_calls( 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 @@ def _transform_responses_api_input_item_to_chat_completion_message( 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 @@ def _transform_responses_api_input_item_to_chat_completion_message( ] @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 @@ def _is_input_item_tool_call_output(input_item: Any) -> bool: ] @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 @@ def _is_input_item_function_call(input_item: Any) -> bool: @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 @@ def _transform_responses_api_tool_call_output_to_chat_completion_message( 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 @@ def _normalize_function_call_output_to_tool_content( # 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 @@ def _normalize_function_call_output_to_tool_content( @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 @@ def _transform_responses_api_function_call_to_chat_completion_message( 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 @@ def _resolve_file_id(item: dict[str, Any]) -> str | None: 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 @@ def _transform_input_file_item_to_file_item(item: dict[str, Any]) -> dict[str, A 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 @@ def _transform_input_image_item_to_image_item( @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 @@ def _transform_responses_api_content_to_chat_completion_content( 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 @@ def _transform_responses_api_content_to_chat_completion_content( 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 @@ def _transform_responses_api_content_to_chat_completion_content( 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 @@ def transform_responses_api_tools_to_chat_completion_tools( 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 @@ def transform_responses_api_tools_to_chat_completion_tools( @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 @@ def transform_chat_completion_tool_params_to_responses_api_tools( """ 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 @@ def transform_chat_completion_tool_params_to_responses_api_tools( 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 @@ def _tool_call_id_from_responses_item(item_id: str | None, call_id: str | None) 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 @@ def convert_response_function_tool_call_to_chat_completion_tool_call( 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 @@ def convert_response_function_tool_call_to_chat_completion_tool_call( 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 @@ def convert_response_function_tool_call_to_chat_completion_tool_call( @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 @@ def convert_apply_patch_tool_call_to_chat_completion_tool_call( 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 @@ def _extract_image_generation_output_items( 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 @@ def _transform_chat_completion_usage_to_responses_usage( @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. diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 820839fc6bf..2e1e1a44594 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -5,7 +5,7 @@ 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.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PART_UNION_TYPES, + ResponseAPIUsage, ResponsesAPIResponse, ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, @@ -112,7 +113,7 @@ def _log_background_task_failure(task: asyncio.Task[object], *, task_name: str) 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 @@ def _process_chunk(self, chunk: str) -> ResponsesAPIStreamingResponse | None: # 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 @@ def _process_chunk(self, chunk: str) -> ResponsesAPIStreamingResponse | None: 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 @@ def _process_chunk(self, chunk: str) -> ResponsesAPIStreamingResponse | None: 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 @@ def _process_chunk(self, chunk: str) -> ResponsesAPIStreamingResponse | None: 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 @@ def _handle_logging_failed_response(self): 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 @@ def _handle_logging_failed_response(self): 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 @@ def _persist_completed_response_to_cache(self, *, is_async: bool) -> None: 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 @@ def _run_post_success_hooks(self, end_time: datetime): 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 @@ def _handle_failure(self, exception: Exception): 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 @@ def __next__(self) -> ResponsesAPIStreamingResponse: 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 @@ async def _mask_response_completed(self, response_str: str) -> str: 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 @@ def _extract_response_id(completed_event: dict[str, object]) -> str | None: @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 @@ def _build_base_call_kwargs(msg_obj: dict[str, object]) -> dict[str, Any]: 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 @@ def _build_base_call_kwargs(msg_obj: dict[str, object]) -> dict[str, Any]: 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 @@ def _same_provider(self, model: str | None) -> bool: 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 diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index a8af4eabb3f..68a6451e273 100644 --- a/ruff-strict-budget.json +++ b/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 @@ -81,7 +81,7 @@ "limit": 1 }, "C901": { - "limit": 315 + "limit": 314 }, "D419": { "limit": 6 @@ -306,7 +306,7 @@ "limit": 0 }, "TID251": { - "limit": 1240 + "limit": 1238 }, "TRY002": { "limit": 528 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 0a0cfe9a617..3a670bc7345 100644 --- a/type-discipline-budget.json +++ b/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