diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 80f1f8f5eca..8b66c27b930 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3481,7 +3481,7 @@ class Logging(LiteLLMLoggingBaseClass): start_time=start_time, end_time=end_time, print_verbose=print_verbose, - level=LogfireLevel.INFO.value, + level=LogfireLevel.INFO, ) if callback == "lunary" and lunaryLogger is not None: @@ -4442,7 +4442,7 @@ class Logging(LiteLLMLoggingBaseClass): response_obj=result, start_time=start_time, end_time=end_time, - level=LogfireLevel.ERROR.value, + level=LogfireLevel.ERROR, print_verbose=print_verbose, ) diff --git a/litellm/litellm_core_utils/safe_json_dumps.py b/litellm/litellm_core_utils/safe_json_dumps.py index 63242a580e7..24c401281ef 100644 --- a/litellm/litellm_core_utils/safe_json_dumps.py +++ b/litellm/litellm_core_utils/safe_json_dumps.py @@ -34,7 +34,7 @@ def safe_json_structure( def _transform(key: str | None, value: str) -> str: return value if value_transform is None else value_transform(key, value) - def _serialize(obj: object, seen: set[int], depth: int, key: str | None = None) -> Any: + def _serialize(obj: object, seen: set[int], depth: int, key: str | None = None) -> object: # Check for maximum depth. if depth > max_depth: return "MaxDepthExceeded" @@ -48,32 +48,32 @@ def safe_json_structure( if id(obj) in seen: return "CircularReference Detected" seen.add(id(obj)) - result: dict | list | tuple | set | str if isinstance(obj, dict): - result = {} - for k, v in obj.items(): - if isinstance(k, (str)): - clean_k = k.replace("\x00", "") if "\x00" in k else k - result[clean_k] = _serialize(v, seen, depth + 1, clean_k) + mapping_result: Final = { + (clean_k := k.replace("\x00", "") if "\x00" in k else k): _serialize(v, seen, depth + 1, clean_k) + for k, v in obj.items() + if isinstance(k, str) + } seen.remove(id(obj)) - return result + return mapping_result elif isinstance(obj, list): - result = [_serialize(item, seen, depth + 1, key) for item in obj] + list_result: Final = [_serialize(item, seen, depth + 1, key) for item in obj] seen.remove(id(obj)) - return result + return list_result elif isinstance(obj, tuple): - result = tuple(_serialize(item, seen, depth + 1, key) for item in obj) + tuple_result: Final = tuple(_serialize(item, seen, depth + 1, key) for item in obj) seen.remove(id(obj)) - return result + return tuple_result elif isinstance(obj, set): - result = sorted([_serialize(item, seen, depth + 1, key) for item in obj]) + sortable_values: Final[list[Any]] = [_serialize(item, seen, depth + 1, key) for item in obj] + set_result: Final = sorted(sortable_values) seen.remove(id(obj)) - return result + return set_result elif isinstance(obj, BaseModel): dumped: Final = obj.model_dump() - result = _serialize(dumped, seen, depth + 1, key) + model_result: Final = _serialize(dumped, seen, depth + 1, key) seen.remove(id(obj)) - return result + return model_result else: # Fall back to string conversion for non-serializable objects. try: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py index f7566ac4923..dc14ef005b5 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py @@ -10,6 +10,8 @@ import time from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -34,6 +36,15 @@ else: LiteLLMLoggingObj = Any +@with_config(ConfigDict(extra="allow", strict=True)) +class _PegasusResponse(TypedDict): + message: ReadOnly[NotRequired[str | None]] + finishReason: ReadOnly[NotRequired[str | None]] + + +_PEGASUS_RESPONSE: Final = TypeAdapter(_PegasusResponse) + + class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): """ Handles transforming OpenAI-style requests into Bedrock InvokeModel requests for @@ -207,7 +218,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): ModelResponse with choices[0].message.content and finish_reason """ try: - completion_response: Final = raw_response.json() + completion_response: Final = _PEGASUS_RESPONSE.validate_python(raw_response.json()) except Exception as e: raise BedrockError( message=f"Error parsing response: {raw_response.text}, error: {e}", @@ -224,7 +235,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): message_content: Final = completion_response.get("message", "") # Extract finish reason and map to LiteLLM format - finish_reason_raw: Final = completion_response.get("finishReason", "stop") + finish_reason_raw: Final = completion_response.get("finishReason") or "stop" finish_reason: Final = map_finish_reason(finish_reason_raw) # Set the response content diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 51b5a99272e..852179e8caa 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -1,11 +1,12 @@ import copy import json import time -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast, get_args import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -48,6 +49,27 @@ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM _GUARDRAIL_CONFIG_VALIDATOR: Final["TypeAdapter[GuardrailConfigBlock]"] = TypeAdapter(GuardrailConfigBlock) + +@with_config(ConfigDict(extra="allow", strict=True)) +class _CohereGeneration(TypedDict): + text: ReadOnly[str | None] + finish_reason: ReadOnly[str | None] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _CohereInvokeResponse(TypedDict): + text: NotRequired[ReadOnly[str | None]] + generations: NotRequired[ReadOnly[Sequence[_CohereGeneration]]] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _LlamaInvokeResponse(TypedDict): + generation: ReadOnly[str | None] + + +_COHERE_INVOKE_RESPONSE: Final = TypeAdapter(_CohereInvokeResponse) +_LLAMA_INVOKE_RESPONSE: Final = TypeAdapter(_LlamaInvokeResponse) + _GUARDRAIL_CONFIG_EXPECTED_FORMAT: Final = ( "{'guardrailIdentifier': str, 'guardrailVersion': str, 'trace': 'enabled'|'disabled'|'enabled_full'}" ) @@ -310,12 +332,13 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): outputText: str | None = None try: if provider == "cohere": - if "text" in completion_response: - outputText = completion_response["text"] - elif "generations" in completion_response: - outputText = completion_response["generations"][0]["text"] + cohere_response: Final = _COHERE_INVOKE_RESPONSE.validate_python(completion_response) + if "text" in cohere_response: + outputText = cohere_response["text"] + elif "generations" in cohere_response: + outputText = cohere_response["generations"][0]["text"] model_response.choices[0].finish_reason = map_finish_reason( - completion_response["generations"][0]["finish_reason"] + cohere_response["generations"][0]["finish_reason"] or "stop" ) elif provider == "anthropic": return litellm.AmazonAnthropicClaudeConfig().transform_response( @@ -361,7 +384,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): elif provider == "ai21": outputText = completion_response.get("completions")[0].get("data").get("text") elif provider == "meta" or provider == "llama" or provider == "deepseek_r1": - outputText = completion_response["generation"] + outputText = _LLAMA_INVOKE_RESPONSE.validate_python(completion_response)["generation"] elif provider == "mistral": outputText = litellm.AmazonMistralConfig.get_outputText(completion_response, model_response) else: # amazon titan diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 9cb38822881..952d8858f78 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -53,6 +53,7 @@ from litellm.types.realtime import ( RealtimeResponseTypedDict, ) +_JSON_OBJECT_MESSAGE: Final = TypeAdapter(dict[str, JsonValue], config=ConfigDict(strict=True)) _JSON_VALUE: Final = TypeAdapter(object) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) @@ -856,7 +857,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): List of Bedrock format messages (JSON strings) """ try: - json_message: Final[dict[str, JsonValue]] = json.loads(message) + json_message: Final = _JSON_OBJECT_MESSAGE.validate_python(json.loads(message)) except json.JSONDecodeError: verbose_logger.warning("Invalid JSON message: %s", message[:200]) return [] diff --git a/litellm/llms/dashscope/embed/transformation.py b/litellm/llms/dashscope/embed/transformation.py index d0c2b4dff0c..35132346261 100644 --- a/litellm/llms/dashscope/embed/transformation.py +++ b/litellm/llms/dashscope/embed/transformation.py @@ -14,6 +14,8 @@ Docs - https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import ensure_api_route @@ -28,6 +30,23 @@ from ..common_utils import DashScopeError DEFAULT_API_BASE: Final = "https://dashscope.aliyuncs.com/compatible-mode/v1" +@with_config(ConfigDict(extra="allow", strict=True)) +class _DashScopeUsage(TypedDict, total=False): + prompt_tokens: NotRequired[ReadOnly[int]] + total_tokens: NotRequired[ReadOnly[int]] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _DashScopeEmbeddingResponse(TypedDict, total=False): + data: NotRequired[ReadOnly[list[object]]] + model: NotRequired[ReadOnly[str | None]] + usage: NotRequired[ReadOnly[_DashScopeUsage | None]] + id: NotRequired[ReadOnly[object]] + + +_DASHSCOPE_EMBEDDING_RESPONSE: Final = TypeAdapter(_DashScopeEmbeddingResponse) + + class DashScopeEmbeddingConfig(BaseEmbeddingConfig): """ Reference: https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api @@ -152,11 +171,12 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): message=message, ) + parsed_response: Final = _DASHSCOPE_EMBEDDING_RESPONSE.validate_python(response_json) model_response.object = "list" - model_response.data = response_json.get("data", []) - model_response.model = response_json.get("model", model) + model_response.data = parsed_response.get("data", []) + model_response.model = parsed_response.get("model", model) - usage: Final = response_json.get("usage") or {} + usage: Final = parsed_response.get("usage") or {} prompt_tokens: Final = usage.get("prompt_tokens", 0) total_tokens: Final = usage.get("total_tokens", prompt_tokens) setattr( @@ -169,8 +189,8 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): ), ) - if "id" in response_json: - setattr(model_response, "id", response_json["id"]) + if "id" in parsed_response: + setattr(model_response, "id", parsed_response["id"]) return model_response diff --git a/litellm/llms/databricks/streaming_utils.py b/litellm/llms/databricks/streaming_utils.py index eca87f9dd5f..cc85012b803 100644 --- a/litellm/llms/databricks/streaming_utils.py +++ b/litellm/llms/databricks/streaming_utils.py @@ -1,6 +1,8 @@ import json from typing import Final +from pydantic import ConfigDict, TypeAdapter + import litellm from litellm import verbose_logger from litellm.types.llms.openai import ( @@ -23,6 +25,9 @@ def _usage_block(usage: Usage) -> ChatCompletionUsageBlock: ) +_JSON_CHUNK: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + + class ModelResponseIterator: def __init__(self, streaming_response, sync_stream: bool): self.streaming_response = streaming_response @@ -99,7 +104,7 @@ class ModelResponseIterator: self.response_iterator = self.streaming_response chunk: Final = self._next_sse_data() try: - json_chunk: Final = json.loads(chunk) + json_chunk: Final = _JSON_CHUNK.validate_python(json.loads(chunk)) return self.chunk_parser(chunk=json_chunk) except ValueError as e: verbose_logger.debug( @@ -126,7 +131,7 @@ class ModelResponseIterator: if chunk == "[DONE]": raise StopAsyncIteration try: - json_chunk: Final = json.loads(chunk) + json_chunk: Final = _JSON_CHUNK.validate_python(json.loads(chunk)) return self.chunk_parser(chunk=json_chunk) except ValueError as e: verbose_logger.debug( diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index d43e203d7a1..f37c6e88896 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -12,6 +12,8 @@ from types import MappingProxyType from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.caching.caching import InMemoryCache @@ -37,6 +39,17 @@ _token_cache: Final = InMemoryCache() _NO_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) +@with_config(ConfigDict(extra="allow", strict=True)) +class _GigaChatTokenResponse(TypedDict): + tok: NotRequired[ReadOnly[str | None]] + access_token: NotRequired[ReadOnly[str | None]] + exp: NotRequired[ReadOnly[int | float | str | None]] + expires_at: NotRequired[ReadOnly[int | float | str | None]] + + +_TOKEN_RESPONSE: Final = TypeAdapter(_GigaChatTokenResponse) + + class GigaChatAuthError(BaseLLMException): """GigaChat authentication error.""" @@ -230,7 +243,7 @@ async def _request_token_async( def _parse_token_response(response: httpx.Response) -> tuple[str, int]: """Parse OAuth token response.""" - data: Final = response.json() + data: Final = _TOKEN_RESPONSE.validate_python(response.json()) # GigaChat returns either 'tok'/'exp' or 'access_token'/'expires_at' access_token: Final = data.get("tok") or data.get("access_token") diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index 080fc3d24c9..b7e4cc6868f 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -5,7 +5,7 @@ This requires websockets, and is currently only supported on LiteLLM Proxy. """ import ssl -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Final, cast from litellm._logging import redact_string, verbose_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES @@ -23,6 +23,9 @@ from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..common_utils import is_openai_backed_api_base from ..openai import OpenAIChatCompletion +if TYPE_CHECKING: + from fastapi import WebSocket + class OpenAIRealtime(OpenAIChatCompletion): """ @@ -131,7 +134,7 @@ class OpenAIRealtime(OpenAIChatCompletion): async def async_realtime( self, model: str, - websocket: Any, + websocket: "WebSocket", logging_obj: LiteLLMLogging, api_base: str | None = None, api_key: str | None = None, diff --git a/litellm/llms/perplexity/embedding/transformation.py b/litellm/llms/perplexity/embedding/transformation.py index 7b4da933976..147d344795d 100644 --- a/litellm/llms/perplexity/embedding/transformation.py +++ b/litellm/llms/perplexity/embedding/transformation.py @@ -12,12 +12,14 @@ This module decodes them into float arrays for OpenAI-compatible responses. """ import base64 +import builtins import struct -from collections.abc import Mapping -from typing import Final +from collections.abc import Mapping, Sequence +from typing import Final, Literal import httpx -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import ReadOnly, TypedDict from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import join_api_base @@ -31,6 +33,17 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i _INT: Final = TypeAdapter(int) +@with_config(ConfigDict(extra="allow", strict=True)) +class _PerplexityEmbeddingResponse(TypedDict, total=False): + model: ReadOnly[str | None] + object: ReadOnly[Literal["list"]] + data: ReadOnly[Sequence[builtins.object] | Mapping[str, builtins.object] | Literal[""]] + usage: ReadOnly[builtins.object] + + +_EMBEDDING_RESPONSE_ADAPTER: Final = TypeAdapter(_PerplexityEmbeddingResponse) + + class PerplexityEmbeddingError(BaseLLMException): def __init__( self, @@ -153,14 +166,18 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict = {}, ) -> EmbeddingResponse: try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final[object] = raw_response.json() except Exception: raise PerplexityEmbeddingError(message=raw_response.text, status_code=raw_response.status_code) - model_response.model = raw_response_json.get("model", model) - model_response.object = raw_response_json.get("object", "list") + response_data: Final[_PerplexityEmbeddingResponse] = _EMBEDDING_RESPONSE_ADAPTER.validate_python( + raw_response_json + ) - raw_data: Final = raw_response_json.get("data", []) + model_response.model = response_data.get("model", model) + model_response.object = response_data.get("object", "list") + + raw_data: Final = response_data.get("data", []) decoded_data: Final[list[dict[str, object]]] = [] for item in raw_data: decoded_item = dict(_JSON_OBJECT.validate_python(item)) @@ -168,7 +185,7 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig): decoded_data.append(decoded_item) model_response.data = decoded_data - usage_data: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {})) + usage_data: Final = _JSON_OBJECT.validate_python(response_data.get("usage", {})) usage: Final = Usage( prompt_tokens=_INT.validate_python( usage_data.get("prompt_tokens", 0) or usage_data.get("total_tokens", 0) or 0 diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py index 9fa86ec30f5..ec2fd16f93e 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py @@ -9,10 +9,11 @@ The actual message transformation reuses OpenAIGPTConfig since Gemma uses OpenAI """ from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Final, cast import httpx -from pydantic import ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError +from typing_extensions import NotRequired, ReadOnly, TypeAliasType, TypedDict from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import ( @@ -32,6 +33,17 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +class _PredictionsEnvelope(TypedDict): + predictions: NotRequired[ReadOnly[dict[str, object]]] + + +_PredictionsEnvelopeType = TypeAliasType("_PredictionsEnvelopeType", _PredictionsEnvelope) +_PREDICTIONS_ENVELOPE: Final[TypeAdapter[_PredictionsEnvelope]] = TypeAdapter( + _PredictionsEnvelopeType, + config=ConfigDict(strict=True, extra="allow"), +) + + def parse_vertex_gemma_container_error(predictions: object) -> VertexGemmaContainerError | None: try: return VertexGemmaContainerError.model_validate(predictions) @@ -142,7 +154,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): def _unwrap_predictions_response( self, - response_json: dict[str, Any], + response_json: _PredictionsEnvelope, ) -> dict[str, object]: """ Unwrap the Vertex Gemma predictions format to OpenAI format. @@ -354,7 +366,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): message=f"Request failed: {response.text}", ) - response_json: Final = response.json() + response_json: Final = _PREDICTIONS_ENVELOPE.validate_python(response.json()) # Unwrap predictions to get OpenAI-compatible response openai_response: Final = self._unwrap_predictions_response(response_json) @@ -449,7 +461,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): message=f"Request failed: {response.text}", ) - response_json: Final = response.json() + response_json: Final = _PREDICTIONS_ENVELOPE.validate_python(response.json()) # Unwrap predictions to get OpenAI-compatible response openai_response: Final = self._unwrap_predictions_response(response_json) diff --git a/litellm/llms/volcengine/responses/transformation.py b/litellm/llms/volcengine/responses/transformation.py index 747ea21e4fb..4076ed06fc2 100644 --- a/litellm/llms/volcengine/responses/transformation.py +++ b/litellm/llms/volcengine/responses/transformation.py @@ -3,6 +3,7 @@ from types import UnionType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, Union, get_args, get_origin import httpx +from pydantic import ConfigDict, TypeAdapter from pydantic import fields as pyd_fields import litellm @@ -37,6 +38,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_RESPONSE_BODY: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + class _EventModelClass(Protocol): @property @@ -420,7 +423,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): @staticmethod def _parsed_response_body(raw_response: httpx.Response) -> dict[str, object]: - return raw_response.json() + return _RESPONSE_BODY.validate_python(raw_response.json()) @staticmethod def _annotation_origin(annotation: object) -> object: diff --git a/litellm/proxy/client/models.py b/litellm/proxy/client/models.py index 1361a3e6364..859322f061d 100644 --- a/litellm/proxy/client/models.py +++ b/litellm/proxy/client/models.py @@ -3,9 +3,12 @@ from collections.abc import Mapping from typing import Any, Final import requests +from pydantic import ConfigDict, TypeAdapter from .exceptions import NotFoundError, UnauthorizedError +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + class ModelsManagementClient: def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): @@ -109,7 +112,7 @@ class ModelsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -145,7 +148,7 @@ class ModelsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -287,7 +290,7 @@ class ModelsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index 98be35079c3..16b3c45f6f8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -169,7 +169,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): from litellm.types.llms.openai import OpenAIModerationResponse - return OpenAIModerationResponse(**response.json()) + return OpenAIModerationResponse.model_validate(response.json()) def _check_moderation_result(self, moderation_response: "OpenAIModerationResponse") -> None: """ diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 601d741c189..d4528bc85e8 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -691,7 +691,7 @@ def _metadata_string_tags(metadata_objects: Iterable[Mapping[str, object]]) -> t def _strip_untrusted_request_header_controls( - headers: Any, + headers: object, *, allow_client_message_redaction_opt_out: bool = False, ) -> None: @@ -1695,7 +1695,7 @@ class LiteLLMProxyRequestSetup: @staticmethod def add_internal_user_from_user_mapping( - general_settings: dict | None, + general_settings: Mapping[str, object] | None, user_api_key_dict: UserAPIKeyAuth, headers: dict, ) -> UserAPIKeyAuth: @@ -1714,7 +1714,7 @@ class LiteLLMProxyRequestSetup: return user_api_key_dict @staticmethod - def get_user_from_headers(headers: dict, general_settings: dict | None = None) -> str | None: + def get_user_from_headers(headers: dict, general_settings: Mapping[str, object] | None = None) -> str | None: """ Get the user from the specified header if `general_settings.user_header_name` is set. """ @@ -1735,7 +1735,9 @@ class LiteLLMProxyRequestSetup: return user @staticmethod - def get_openai_org_id_from_headers(headers: dict, general_settings: dict | None = None) -> str | None: + def get_openai_org_id_from_headers( + headers: dict, general_settings: Mapping[str, object] | None = None + ) -> str | None: """ Get the OpenAI Org ID from the headers. """ @@ -1826,7 +1828,7 @@ class LiteLLMProxyRequestSetup: headers: dict, request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, - general_settings: dict[str, object] | None = None, + general_settings: Mapping[str, object] | None = None, ) -> LitellmDataForBackendLLMCall: """ - Adds user from headers diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 7a843aecc0c..4f48f3c57b9 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -231,9 +231,12 @@ def _shadow_eval_attempts(prisma_client: "PrismaClient") -> _ShadowEvalAttemptTa return _ShadowEvalAttemptRepository(prisma_client).table +_RAW_ROWS: Final = TypeAdapter(Sequence[dict[str, object]], config=ConfigDict(strict=True)) + + async def _query_raw(prisma_client: "PrismaClient", query: str, *args: object) -> Sequence[Mapping[str, object]]: async with db_span("auto_router_report_query", sql_relation(query)): - return await prisma_client.db.query_raw(query, *args) + return _RAW_ROWS.validate_python(await prisma_client.db.query_raw(query, *args)) async def _execute_raw(prisma_client: "PrismaClient", query: str, *args: object) -> int: diff --git a/litellm/proxy/management_endpoints/router_settings_endpoints.py b/litellm/proxy/management_endpoints/router_settings_endpoints.py index 1f4ca442d20..41c4b5cfaf7 100644 --- a/litellm/proxy/management_endpoints/router_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/router_settings_endpoints.py @@ -10,7 +10,7 @@ GET /router/fields - Get router settings field definitions without values (for U import inspect from collections.abc import Mapping from types import MappingProxyType -from typing import Any, Final, get_args +from typing import Final, get_args from fastapi import APIRouter, Depends from pydantic import Field @@ -32,7 +32,7 @@ router: Final = APIRouter() class RouterSettingsResponse(LiteLLMBaseModel): fields: list[RouterSettingsField] = Field(description="List of all configurable router settings with metadata") - current_values: dict[str, Any] = Field(description="Current values of router settings") + current_values: dict[str, object] = Field(description="Current values of router settings") routing_strategy_descriptions: dict[str, str] = Field(description="Descriptions for each routing strategy option") routing_group_strategies: tuple[str, ...] = Field( description="Strategies supported when constructing a routing group" diff --git a/litellm/proxy/management_helpers/team_roster_sync.py b/litellm/proxy/management_helpers/team_roster_sync.py index dc0ba2e47d4..08be455a55d 100644 --- a/litellm/proxy/management_helpers/team_roster_sync.py +++ b/litellm/proxy/management_helpers/team_roster_sync.py @@ -40,6 +40,7 @@ from litellm.proxy.management_helpers.access_group_team_sync import ( ) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.chunked_in import delete_many_in, find_many_in, update_many_in +from litellm.repositories.team_repository import TeamRepository if TYPE_CHECKING: from prisma import Prisma @@ -144,7 +145,7 @@ def _changes_nothing(plan: RosterPlan) -> bool: async def _unchanged_roster(prisma_client: PrismaClient, team_id: str) -> RosterSync | TeamGone: - team_row: Final = await prisma_client.writer_db.litellm_teamtable.find_unique(where={"team_id": team_id}) + team_row: Final = await TeamRepository(prisma_client, use_writer=True).table.find_unique(where={"team_id": team_id}) if team_row is None: return TeamGone(team_id=team_id) team: Final = LiteLLM_TeamTable.model_validate(team_row.model_dump()) diff --git a/litellm/proxy/mcp_server/discoverable_endpoints.py b/litellm/proxy/mcp_server/discoverable_endpoints.py index eb12b15cdae..07a267d6caa 100644 --- a/litellm/proxy/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/mcp_server/discoverable_endpoints.py @@ -300,12 +300,13 @@ def decode_state_hash(encrypted_state: str) -> dict: if decrypted_json is None: raise ValueError("Failed to decrypt state parameter") - state_data: Final = json.loads(decrypted_json) + state_data: Final = _DECODED_STATE.validate_python(json.loads(decrypted_json)) return state_data _MANAGED_AUTH_CODE_PREFIX: Final = "llm_scope_" _DECODED_JSON: Final = TypeAdapter(object) +_DECODED_STATE: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) class _ManagedAuthorizationCode(LiteLLMBaseModel): diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 650dcae3239..e017eee3d27 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -30,6 +30,7 @@ from pydantic import TypeAdapter from starlette.datastructures import UploadFile as StarletteUploadFile from starlette.routing import BaseRoute, Route from starlette.websockets import WebSocketState +from typing_extensions import ReadOnly from websockets.asyncio.client import connect from websockets.exceptions import ( ConnectionClosedError, @@ -173,8 +174,17 @@ def _is_mergeable_client_metadata(key: str, value: object) -> bool: return not key.startswith("user_api_key_") and (key != "model_info" or isinstance(value, Mapping)) +class _RegisteredPassThroughRoute(TypedDict): + endpoint_id: ReadOnly[str] + path: ReadOnly[str | None] + type: ReadOnly[str] + methods: ReadOnly[str | bool | list[str] | tuple[str, ...] | Mapping[str, object]] + auth: ReadOnly[bool] + passthrough_params: ReadOnly[Mapping[str, object]] + + # Global registry to track registered pass-through routes and prevent memory leaks -_registered_pass_through_routes: Final[dict[str, dict[str, str | bool | list[str] | Mapping[str, object]]]] = {} +_registered_pass_through_routes: Final[dict[str, _RegisteredPassThroughRoute]] = {} def get_response_body(response: httpx.Response) -> dict | None: @@ -2239,7 +2249,7 @@ def create_pass_through_route( status_code=status.HTTP_405_METHOD_NOT_ALLOWED, detail=f"Method {request.method} is not allowed for pass-through endpoint {path}.", ) - target_params: Final = { + target_params: Final[dict[str, object]] = { "target": target, "custom_headers": custom_headers, "forward_headers": _forward_headers, @@ -2271,7 +2281,9 @@ def create_pass_through_route( # Ensure custom_headers is a dict. Botocore returns a HeadersDict # for SigV4-prepared requests, which is a Mapping but not a dict. - headers_dict: Final = dict(param_custom_headers) if isinstance(param_custom_headers, Mapping) else {} + headers_dict: Final[dict[str, object]] = ( + {**param_custom_headers} if isinstance(param_custom_headers, Mapping) else {} + ) # Ensure query_params and custom_body are dicts or None final_query_params: Final = query_params_data if isinstance(query_params_data, dict) else {} @@ -3104,9 +3116,7 @@ class SafeRouteAdder: return f"{path}/{{subpath:path}}" return None - def registered_route_methods( - route_info: Mapping[str, str | bool | Sequence[str] | Mapping[str, object]], - ) -> tuple[tuple[str, str], ...]: + def registered_route_methods(route_info: _RegisteredPassThroughRoute) -> tuple[tuple[str, str], ...]: starlette_path: Final = registered_route_path(route_info) methods: Final = route_info.get("methods") if starlette_path is None or not isinstance(methods, (list, tuple)): @@ -3462,7 +3472,7 @@ class InitPassThroughEndpointHelpers: return (comparison_route, *forwarded) @staticmethod - def get_registered_pass_through_route(route: str, method: str | None = None) -> dict[str, Any] | None: + def get_registered_pass_through_route(route: str, method: str | None = None) -> _RegisteredPassThroughRoute | None: """Get passthrough params for a given route and optionally filter by HTTP method""" comparison_route: Final = InitPassThroughEndpointHelpers._route_for_registry_lookup(route) for key in _registered_pass_through_routes: diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 5442c0f3f81..8fd4fbfcb8b 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -6,7 +6,7 @@ from datetime import datetime, timedelta from types import MappingProxyType from typing import Final, Literal, TypeVar -from pydantic import TypeAdapter +from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -155,6 +155,7 @@ class _DailyUserSpendOwnerRow(LiteLLMBaseModel): last_owner: str | None = None +_RAW_ROWS: Final = TypeAdapter(Sequence[dict[str, object]], config=ConfigDict(strict=True)) _TOKEN_DIGEST_ROWS: Final = TypeAdapter(tuple[_TokenDigestRow, ...]) _SPEND_LOG_DIGEST_ROWS: Final = TypeAdapter(tuple[_SpendLogDigestRow, ...]) _DAILY_USER_SPEND_OWNER_ROWS: Final = TypeAdapter(tuple[_DailyUserSpendOwnerRow, ...]) @@ -194,7 +195,7 @@ async def _rows_within_the_statement_timeout( await transaction.execute_raw(_SPEND_LOG_STATEMENT_TIMEOUT_SQL) for setting in planner_settings: await transaction.execute_raw(setting) - return await transaction.query_raw(sql, *params) + return _RAW_ROWS.validate_python(await transaction.query_raw(sql, *params)) async def _reverse_hash_key_metadata( diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index 98d12e497da..de8566e6fdc 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -21,6 +21,8 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -41,6 +43,15 @@ if TYPE_CHECKING: from botocore.awsrequest import HTTPHeaders +@with_config(ConfigDict(extra="allow", strict=True)) +class _SecretStringResponse(TypedDict): + SecretString: ReadOnly[str] + + +_SECRET_STRING_RESPONSE: Final = TypeAdapter(_SecretStringResponse) +_AWS_SECRET_MANAGER_RESPONSE: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + + class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): def __init__( self, @@ -154,7 +165,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() - return response.json()["SecretString"] + return _SECRET_STRING_RESPONSE.validate_python(response.json())["SecretString"] except httpx.TimeoutException: raise ValueError("Timeout error occurred") except Exception as e: @@ -204,7 +215,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = sync_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) - return response.json()["SecretString"] + return _SECRET_STRING_RESPONSE.validate_python(response.json())["SecretString"] except httpx.TimeoutException: raise ValueError("Timeout error occurred") except httpx.HTTPStatusError as e: @@ -424,7 +435,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() - return response.json() + return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json()) except httpx.HTTPStatusError as err: raise ValueError(f"HTTP error occurred: {err.response.text}") except httpx.TimeoutException: @@ -481,7 +492,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() - return response.json() + return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json()) except httpx.HTTPStatusError as err: raise ValueError(f"HTTP error occurred: {err.response.text}") except httpx.TimeoutException: @@ -533,7 +544,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() - return response.json() + return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json()) except httpx.HTTPStatusError as err: raise ValueError(f"HTTP error occurred: {err.response.text}") except httpx.TimeoutException: @@ -613,7 +624,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() - return response.json() + return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json()) except httpx.HTTPStatusError as err: raise ValueError(f"HTTP error occurred: {err.response.text}") except httpx.TimeoutException: diff --git a/litellm/types/integrations/rag/bedrock_knowledgebase.py b/litellm/types/integrations/rag/bedrock_knowledgebase.py index 7156d8101e1..7288eca03b4 100644 --- a/litellm/types/integrations/rag/bedrock_knowledgebase.py +++ b/litellm/types/integrations/rag/bedrock_knowledgebase.py @@ -1,4 +1,4 @@ -from typing import Any, Literal +from typing import Literal from typing_extensions import ReadOnly, TypedDict @@ -40,7 +40,7 @@ class BedrockKBRetrievalResult(TypedDict, total=False): content: BedrockKBContent | None location: BedrockKBLocation | None score: float | None - metadata: dict[str, Any] | None + metadata: dict[str, object] | None class BedrockKBResponse(TypedDict, total=False): @@ -87,21 +87,21 @@ class BedrockKBMetadataConfiguration(TypedDict, total=False): class BedrockKBModelConfiguration(TypedDict, total=False): """Model configuration for reranking.""" - additionalModelRequestFields: dict[str, Any] | None + additionalModelRequestFields: dict[str, object] | None modelArn: str | None class BedrockKBRerankingConfiguration(TypedDict, total=False): """Configuration for reranking in vector search.""" - bedrockRerankingConfiguration: dict[str, Any] | None # This could be further typed if needed + bedrockRerankingConfiguration: dict[str, object] | None # This could be further typed if needed type: str | None class BedrockKBVectorSearchConfiguration(TypedDict, total=False): """Configuration for vector search.""" - filter: dict[str, Any] | None + filter: dict[str, object] | None implicitFilterConfiguration: BedrockKBImplicitFilterConfiguration | None numberOfResults: int | None overrideSearchType: str | None diff --git a/litellm/types/llms/bedrock_invoke_agents.py b/litellm/types/llms/bedrock_invoke_agents.py index 62f3c18ca7a..b334c986d6b 100644 --- a/litellm/types/llms/bedrock_invoke_agents.py +++ b/litellm/types/llms/bedrock_invoke_agents.py @@ -4,7 +4,7 @@ Type definitions for AWS Bedrock Invoke Agent API responses. https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_InvokeAgent.html """ -from typing import Any, Final +from typing import Final from typing_extensions import TypedDict @@ -39,7 +39,7 @@ class InvokeAgentModelInvocationInput(TypedDict, total=False): """Model invocation input details.""" foundationModel: str | None - inferenceConfiguration: dict[str, Any] | None + inferenceConfiguration: dict[str, object] | None text: str | None traceId: str | None type: str | None @@ -49,9 +49,9 @@ class InvokeAgentModelInvocationOutput(TypedDict, total=False): """Model invocation output details.""" metadata: InvokeAgentMetadata | None - parsedResponse: dict[str, Any] | None - rawResponse: dict[str, Any] | None - reasoningContent: dict[str, Any] | None + parsedResponse: dict[str, object] | None + rawResponse: dict[str, object] | None + reasoningContent: dict[str, object] | None traceId: str | None diff --git a/litellm/types/llms/gemini.py b/litellm/types/llms/gemini.py index c300c26bcd3..2ad81e64e05 100644 --- a/litellm/types/llms/gemini.py +++ b/litellm/types/llms/gemini.py @@ -1,5 +1,6 @@ +from collections.abc import Mapping from enum import Enum -from typing import Any, Literal +from typing import Literal from typing_extensions import Required, TypedDict @@ -235,7 +236,7 @@ class GeminiVideoGenerationInstance(TypedDict, total=False): """Instance data for Gemini video generation request""" prompt: Required[str] - image: dict[str, Any] + image: Mapping[str, object] class GeminiVideoGenerationParameters(LiteLLMBaseModel): @@ -263,7 +264,7 @@ class GeminiVideoGenerationParameters(LiteLLMBaseModel): negativePrompt: str | None = None """Text describing what not to include in the video.""" - lastFrame: Any | None = None + lastFrame: object | None = None """ The final image for interpolation video to transition. Should be used with the 'image' parameter. @@ -275,7 +276,7 @@ class GeminiVideoGenerationParameters(LiteLLMBaseModel): Only supported in Veo 3.1 (list of VideoGenerationReferenceImage objects). """ - video: Any | None = None + video: object | None = None """ Video to be used for video extension (Video object). Only supported in Veo 3.1 & Veo 3 Fast. @@ -354,5 +355,5 @@ class GeminiLongRunningOperationResponse(LiteLLMBaseModel): response: GeminiOperationResponse | None = None """Response object when operation is complete""" - error: dict[str, Any] | None = None + error: dict[str, object] | None = None """Error details if operation failed""" diff --git a/litellm/types/llms/langgraph.py b/litellm/types/llms/langgraph.py index 7e34329992d..83f7bb090d1 100644 --- a/litellm/types/llms/langgraph.py +++ b/litellm/types/llms/langgraph.py @@ -4,7 +4,8 @@ Type definitions for LangGraph API. LangGraph provides a streaming and non-streaming API for running agents. """ -from typing import Any, Literal +from collections.abc import Mapping +from typing import Literal from typing_extensions import TypedDict @@ -29,8 +30,8 @@ class LangGraphRequest(TypedDict, total=False): assistant_id: str input: LangGraphInput stream_mode: str | None - config: dict[str, Any] | None - metadata: dict[str, Any] | None + config: Mapping[str, object] | None + metadata: Mapping[str, object] | None # Response Types - Streaming @@ -38,7 +39,7 @@ class LangGraphStreamEvent(TypedDict, total=False): """Single event in a LangGraph stream response.""" event: str - data: Any + data: object # Response Types - Non-streaming @@ -55,7 +56,7 @@ class LangGraphResponse(TypedDict, total=False): """Non-streaming response structure from LangGraph.""" messages: list[LangGraphResponseMessage] - values: dict[str, Any] + values: dict[str, object] # Parsed response for internal use diff --git a/litellm/types/mcp_server/tool_registry.py b/litellm/types/mcp_server/tool_registry.py index 738d37c9503..080e27b6ed3 100644 --- a/litellm/types/mcp_server/tool_registry.py +++ b/litellm/types/mcp_server/tool_registry.py @@ -1,5 +1,5 @@ from collections.abc import Callable -from typing import Any, ClassVar +from typing import ClassVar from pydantic import ConfigDict, Field @@ -10,7 +10,7 @@ class MCPTool(LiteLLMBaseModel): model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True) name: str description: str - input_schema: dict[str, Any] + input_schema: dict[str, object] handler: Callable server_id: str | None = Field(default=None, frozen=True) @@ -18,18 +18,18 @@ class MCPTool(LiteLLMBaseModel): class ToolSchema(LiteLLMBaseModel): name: str description: str - inputSchema: dict[str, Any] + inputSchema: dict[str, object] class ListToolsResponse(LiteLLMBaseModel): tools: list[ToolSchema] nextCursor: str | None = None - _meta: dict[str, Any] | None = None + _meta: dict[str, object] | None = None class CallToolRequest(LiteLLMBaseModel): method: str = "tools/call" - params: dict[str, Any] + params: dict[str, object] class ContentItem(LiteLLMBaseModel): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py b/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py index d1eba1dd4a9..f936c9d85cf 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py @@ -1,5 +1,3 @@ -from typing import Any - from pydantic import Field from typing_extensions import TypedDict @@ -20,7 +18,7 @@ class EnkryptAIPolicyViolationDetail(TypedDict, total=False): class EnkryptAIPIIDetail(TypedDict, total=False): """Details for PII detection.""" - pii: dict[str, Any] + pii: dict[str, object] class EnkryptAIToxicityDetail(TypedDict, total=False): @@ -85,8 +83,8 @@ class EnkryptAIResponseDetails(TypedDict, total=False): toxicity: EnkryptAIToxicityDetail | None keyword_detected: EnkryptAIKeywordDetail | None bias: EnkryptAIBiasDetail | None - prompt_injection: dict[str, Any] | None - jailbreak: dict[str, Any] | None + prompt_injection: dict[str, object] | None + jailbreak: dict[str, object] | None class EnkryptAIResponse(TypedDict, total=False): @@ -107,7 +105,7 @@ class EnkryptAIProcessedResult(TypedDict): | EnkryptAIToxicityDetail | EnkryptAIKeywordDetail | EnkryptAIBiasDetail - | dict[str, Any], + | dict[str, object], ] diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py index 1890c68b30d..c36644ba536 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -242,6 +242,33 @@ def test_transform_response_hands_json_mode_to_nova(): assert json.loads(result.choices[0].message.content) == {"city": "Paris", "temperature": 21} +@pytest.mark.parametrize(("finish_reason", "expected"), [("MAX_TOKENS", "length"), (None, "stop")]) +def test_transform_response_accepts_null_cohere_finish_reason(finish_reason: str | None, expected: str) -> None: + from litellm.types.utils import ModelResponse + + raw_response: Final = httpx.Response( + 200, + json={"generations": [{"text": "Hello there", "finish_reason": finish_reason}]}, + request=httpx.Request("POST", "https://bedrock"), + ) + + result: Final = AmazonInvokeConfig().transform_response( + model="invoke/cohere.command-text-v14", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key=None, + ) + + assert result.choices[0].message.content == "Hello there" + assert result.choices[0].finish_reason == expected + + def _stream_invoke_completion_with_spied_client(**kwargs: object) -> tuple[MagicMock, MagicMock]: mock_response = MagicMock() mock_response.status_code = 200 diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py index d90e9933b88..5b3cb2e0c92 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py @@ -1,6 +1,13 @@ +from typing import Final +from unittest.mock import Mock + +import httpx +import pytest + from litellm.llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation import ( AmazonTwelveLabsPegasusConfig, ) +from litellm.types.utils import ModelResponse def _make_messages() -> list[dict]: @@ -84,3 +91,24 @@ def test_transform_request_includes_s3_media(): s3_location = request["mediaSource"]["s3Location"] assert s3_location["uri"] == "s3://test-bucket/video.mp4" assert s3_location["bucketOwner"] == "123456789012" + + +@pytest.mark.parametrize( + ("finish_reason", "expected"), + [("length", "length"), (None, "stop")], +) +def test_transform_response_accepts_null_finish_reason(finish_reason: str | None, expected: str) -> None: + result: Final = AmazonTwelveLabsPegasusConfig().transform_response( + model="twelvelabs.pegasus-1-2-v1:0", + raw_response=httpx.Response(200, json={"message": "A cat jumps.", "finishReason": finish_reason}), + model_response=ModelResponse(), + logging_obj=Mock(), + request_data={}, + messages=_make_messages(), + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result.choices[0].message.content == "A cat jumps." + assert result.choices[0].finish_reason == expected