diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index 1e6da5f2f43..e36d455a20f 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -9,8 +9,6 @@ import json from collections.abc import AsyncIterator, Mapping from typing import Final -from pydantic import ConfigDict, TypeAdapter - from litellm._logging import verbose_logger from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import ( BedrockAgentCoreA2ATransformation, @@ -19,8 +17,6 @@ from litellm.llms.bedrock.base_aws_llm import run_aws_signing from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider -_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) - class BedrockAgentCoreA2AHandler: """ @@ -70,7 +66,7 @@ class BedrockAgentCoreA2AHandler: data=body, ) response.raise_for_status() - response_data: Final = _JSON_OBJECT.validate_python(response.json(), strict=True) + response_data: Final = response.json() if "error" in response_data: verbose_logger.warning("BedrockAgentCore A2A: Agent returned error: %s", response_data["error"]) @@ -129,7 +125,7 @@ class BedrockAgentCoreA2AHandler: "BedrockAgentCore A2A streaming: received JSON instead of SSE, yielding as single event" ) response_body: Final = await response.aread() - response_data: Final = _JSON_OBJECT.validate_python(json.loads(response_body), strict=True) + response_data: Final = json.loads(response_body) yield response_data else: # SSE stream — parse data: lines diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py index fd2fff8fc0b..d0ce5895970 100644 --- a/litellm/llms/bedrock/image_edit/stability_transformation.py +++ b/litellm/llms/bedrock/image_edit/stability_transformation.py @@ -22,13 +22,10 @@ API Reference: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parame """ import base64 -from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles -from pydantic import ConfigDict, TypeAdapter, with_config -from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.bedrock.common_utils import BedrockError @@ -48,15 +45,6 @@ else: LiteLLMLoggingObj = Any -@with_config(ConfigDict(extra="allow", strict=True)) -class _StabilityResponse(TypedDict): - finish_reasons: ReadOnly[NotRequired[Sequence[object] | None]] - images: ReadOnly[NotRequired[Sequence[str | None] | None]] - - -_STABILITY_RESPONSE: Final = TypeAdapter(_StabilityResponse) - - class BedrockStabilityImageEditConfig(BaseImageEditConfig): """ Configuration for Bedrock Stability AI image edit. @@ -327,10 +315,8 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): headers=raw_response.headers, ) - payload: Final = _STABILITY_RESPONSE.validate_python(response_data) - # Check finish_reasons - finish_reasons: Final = payload.get("finish_reasons", []) + finish_reasons: Final = response_data.get("finish_reasons", []) if finish_reasons and finish_reasons[0]: raise self.get_error_class( error_message=f"Bedrock Stability error: {finish_reasons[0]}", @@ -343,7 +329,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): model_response.data = [] # Extract images from response - images: Final = payload.get("images", []) + images: Final = response_data.get("images", []) if images: for image_b64 in images: if image_b64: diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 2d558e3b8d3..96d1da19bcc 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -1,11 +1,9 @@ -from collections.abc import Mapping, Sequence +from collections.abc import Mapping from copy import deepcopy from typing import TYPE_CHECKING, Any, Final, cast from urllib.parse import urlparse import httpx -from pydantic import ConfigDict, TypeAdapter, with_config -from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base @@ -13,6 +11,8 @@ from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreCon from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.integrations.rag.bedrock_knowledgebase import ( + BedrockKBContent, + BedrockKBResponse, BedrockKBRetrievalConfiguration, BedrockKBRetrievalQuery, BedrockKBUserContext, @@ -34,26 +34,6 @@ else: LiteLLMLoggingObj = Any -@with_config(ConfigDict(extra="allow", strict=True)) -class _KBContent(TypedDict): - text: ReadOnly[NotRequired[str | None]] - - -@with_config(ConfigDict(extra="allow", strict=True)) -class _KBRetrievalResult(TypedDict): - content: ReadOnly[NotRequired[_KBContent | None]] - score: ReadOnly[NotRequired[int | float | None]] - metadata: ReadOnly[NotRequired[Mapping[str, object] | None]] - - -@with_config(ConfigDict(extra="allow", strict=True)) -class _KBResponse(TypedDict): - retrievalResults: ReadOnly[NotRequired[Sequence[_KBRetrievalResult] | None]] - - -_KB_RESPONSE: Final = TypeAdapter(_KBResponse) - - class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): """Vector store configuration for AWS Bedrock Knowledge Bases.""" @@ -304,7 +284,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): api_key=api_key, ) - def _get_file_id_from_metadata(self, metadata: Mapping[str, Any]) -> str: + def _get_file_id_from_metadata(self, metadata: dict[str, Any]) -> str: """ Extract file_id from Bedrock KB metadata. Uses source URI if available, otherwise generates a fallback ID. @@ -316,7 +296,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): chunk_id: Final = metadata.get("x-amz-bedrock-kb-chunk-id", "unknown") if metadata else "unknown" return f"bedrock-kb-{chunk_id}" - def _get_filename_from_metadata(self, metadata: Mapping[str, Any]) -> str: + def _get_filename_from_metadata(self, metadata: dict[str, Any]) -> str: """ Extract filename from Bedrock KB metadata. Tries to extract filename from source URI, falls back to domain name or data source ID. @@ -338,7 +318,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): data_source_id: Final = metadata.get("x-amz-bedrock-kb-data-source-id", "unknown") if metadata else "unknown" return f"bedrock-kb-document-{data_source_id}" - def _get_attributes_from_metadata(self, metadata: Mapping[str, object]) -> dict[str, object]: + def _get_attributes_from_metadata(self, metadata: dict[str, object]) -> dict[str, object]: """ Extract all attributes from Bedrock KB metadata. Returns a copy of the metadata dictionary. @@ -351,10 +331,10 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj ) -> VectorStoreSearchResponse: try: - response_data: Final = _KB_RESPONSE.validate_python(response.json()) + response_data: Final = BedrockKBResponse(**response.json()) results: Final[list[VectorStoreSearchResult]] = [] for item in response_data.get("retrievalResults", []) or []: - content: _KBContent | None = item.get("content") + content: BedrockKBContent | None = item.get("content") text = content.get("text") if content else None if text is None: continue diff --git a/litellm/llms/deprecated_providers/aleph_alpha.py b/litellm/llms/deprecated_providers/aleph_alpha.py index 3f65f2997fd..cb56717993e 100644 --- a/litellm/llms/deprecated_providers/aleph_alpha.py +++ b/litellm/llms/deprecated_providers/aleph_alpha.py @@ -5,27 +5,11 @@ from collections.abc import Callable, Mapping from typing import Final import httpx -from pydantic import ConfigDict, TypeAdapter, with_config -from typing_extensions import ReadOnly, TypedDict import litellm from litellm.utils import Choices, Message, ModelResponse, Usage -@with_config(ConfigDict(extra="allow", strict=True)) -class _AlephAlphaCompletion(TypedDict): - completion: ReadOnly[str] - finish_reason: ReadOnly[str | None] - - -@with_config(ConfigDict(extra="allow", strict=True)) -class _AlephAlphaResponse(TypedDict): - completions: ReadOnly[list[_AlephAlphaCompletion]] - - -_ALEPH_ALPHA_RESPONSE: Final = TypeAdapter(_AlephAlphaResponse) - - class AlephAlphaError(Exception): def __init__(self, status_code, message): self.status_code = status_code @@ -273,9 +257,8 @@ def completion( ) else: try: - validated_response: Final = _ALEPH_ALPHA_RESPONSE.validate_python(completion_response) choices_list: Final = [] - for idx, item in enumerate(validated_response["completions"]): + for idx, item in enumerate(completion_response["completions"]): if len(item["completion"]) > 0: message_obj = Message(content=item["completion"]) else: diff --git a/litellm/llms/recraft/image_edit/transformation.py b/litellm/llms/recraft/image_edit/transformation.py index 3230f2e6d07..74bfa52c5e0 100644 --- a/litellm/llms/recraft/image_edit/transformation.py +++ b/litellm/llms/recraft/image_edit/transformation.py @@ -1,11 +1,8 @@ -from collections.abc import Sequence from io import BufferedReader from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles -from pydantic import ConfigDict, TypeAdapter, with_config -from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.images.utils import ImageEditRequestUtils from litellm.litellm_core_utils.url_utils import join_api_base @@ -25,21 +22,6 @@ else: LiteLLMLoggingObj = Any -@with_config(ConfigDict(extra="allow", strict=True)) -class _RecraftImageData(TypedDict): - url: ReadOnly[NotRequired[str | None]] - b64_json: ReadOnly[NotRequired[str | None]] - - -@with_config(ConfigDict(extra="allow", strict=True)) -class _RecraftImageEditResponse(TypedDict): - data: ReadOnly[Sequence[_RecraftImageData]] - - -_JSON_BODY: Final = TypeAdapter(object, config=ConfigDict(hide_input_in_errors=True)) -_RECRAFT_IMAGE_EDIT_RESPONSE: Final = TypeAdapter(_RecraftImageEditResponse) - - class RecraftImageEditConfig(BaseImageEditConfig): DEFAULT_BASE_URL: str = "https://external.api.recraft.ai" IMAGE_EDIT_ENDPOINT: str = "v1/images/imageToImage" @@ -170,14 +152,13 @@ class RecraftImageEditConfig(BaseImageEditConfig): ) -> ImageResponse: model_response: Final = ImageResponse() try: - raw_response_data: Final = _JSON_BODY.validate_python(raw_response.json()) + response_data: Final = raw_response.json() except Exception as e: raise self.get_error_class( error_message=f"Error transforming image edit response: {e}", status_code=raw_response.status_code, headers=raw_response.headers, ) - response_data: Final = _RECRAFT_IMAGE_EDIT_RESPONSE.validate_python(raw_response_data) if not model_response.data: model_response.data = [] diff --git a/litellm/llms/replicate/chat/transformation.py b/litellm/llms/replicate/chat/transformation.py index 72fd6ce37b6..62dfac43bd7 100644 --- a/litellm/llms/replicate/chat/transformation.py +++ b/litellm/llms/replicate/chat/transformation.py @@ -3,8 +3,7 @@ from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final import httpx -from pydantic import ConfigDict, TypeAdapter, with_config -from typing_extensions import NotRequired, ReadOnly, TypedDict +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.constants import REPLICATE_MODEL_NAME_WITH_ID_LENGTH @@ -34,20 +33,6 @@ _TEXTS: Final = TypeAdapter(Iterable[str], config=ConfigDict(strict=True, hide_i _JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) -@with_config(ConfigDict(extra="allow", strict=True)) -class _ReplicatePredictionUrls(TypedDict): - get: ReadOnly[NotRequired[str | None]] - - -@with_config(ConfigDict(extra="allow", strict=True)) -class _ReplicatePrediction(TypedDict): - urls: ReadOnly[NotRequired[_ReplicatePredictionUrls]] - - -_JSON_BODY: Final = TypeAdapter(object, config=ConfigDict(hide_input_in_errors=True)) -_REPLICATE_PREDICTION: Final = TypeAdapter(_ReplicatePrediction) - - class ReplicateConfig(BaseConfig): """ Reference: https://replicate.com/meta/llama-2-70b-chat/api @@ -307,11 +292,8 @@ class ReplicateConfig(BaseConfig): "urls":{"cancel":"https://api.replicate.com/v1/predictions/gqsmqmp1pdrj00cknr08dgmvb4/cancel","get":"https://api.replicate.com/v1/predictions/gqsmqmp1pdrj00cknr08dgmvb4","stream":"https://stream-b.svc.rno2.c.replicate.net/v1/streams/eot4gbydowuin4snhncydwxt57dfwgsc3w3snycx5nid7oef7jga"} } """ - response_json: Final = _JSON_BODY.validate_python(response.json()) - prediction_urls: Final = _REPLICATE_PREDICTION.validate_python(response_json).get("urls") - prediction_url: Final = ( - prediction_urls["get"] if prediction_urls is not None and "get" in prediction_urls else None - ) + response_json: Final = response.json() + prediction_url: Final = response_json.get("urls", {}).get("get") if prediction_url is None: raise ReplicateError( status_code=400, diff --git a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py index 4877c166026..cf6592a3e58 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py @@ -7,12 +7,9 @@ import json import os -from collections.abc import Sequence -from typing import TYPE_CHECKING, Final, Literal +from typing import TYPE_CHECKING, Final, Literal, TypedDict from fastapi import HTTPException -from pydantic import ConfigDict, TypeAdapter, with_config -from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -37,17 +34,17 @@ class GuardrailsAIResponse(TypedDict): validationPassed: bool -@with_config(ConfigDict(extra="allow", strict=True)) -class _PreCallOutput(TypedDict): - data: ReadOnly[NotRequired[Sequence[str]]] +class InferenceData(TypedDict): + name: str + shape: list[int] + data: list + datatype: str -@with_config(ConfigDict(extra="allow", strict=True)) -class _PreCallResponse(TypedDict): - outputs: ReadOnly[NotRequired[Sequence[_PreCallOutput]]] - - -_PRE_CALL_RESPONSE: Final = TypeAdapter(_PreCallResponse) +class GuardrailsAIResponsePreCall(TypedDict): + modelname: str + modelversion: str + outputs: list[InferenceData] class GuardrailsAI(CustomGuardrail): @@ -133,7 +130,7 @@ class GuardrailsAI(CustomGuardrail): }, ) - _json_response: Final = _PRE_CALL_RESPONSE.validate_python(response.json()) + _json_response: Final = GuardrailsAIResponsePreCall(**response.json()) response = _json_response.get("outputs", [])[0].get("data", [])[0] return response diff --git a/tests/unit/llms/bedrock/image_edit/test_stability_transformation.py b/tests/unit/llms/bedrock/image_edit/test_stability_transformation.py index d761deabbe4..6ca70324162 100644 --- a/tests/unit/llms/bedrock/image_edit/test_stability_transformation.py +++ b/tests/unit/llms/bedrock/image_edit/test_stability_transformation.py @@ -1,8 +1,3 @@ -from unittest.mock import MagicMock - -import httpx -import pytest - from litellm.llms.bedrock.image_edit.stability_transformation import BedrockStabilityImageEditConfig @@ -17,31 +12,3 @@ def test_transform_image_edit_request_returns_the_json_body_and_no_files(): ) assert request == ({"output_format": "png", "prompt": "add a red hat", "image": "iVBORy1ieXRlcw=="}, {}) - - -@pytest.mark.parametrize( - ("body", "expected_images"), - [ - ({"images": ["aW1hZ2U="], "finish_reasons": [None], "seeds": [42]}, ["aW1hZ2U="]), - ({"images": ["aW1hZ2U="], "finish_reasons": None}, ["aW1hZ2U="]), - ({"images": None, "finish_reasons": None}, []), - ({}, []), - ], -) -def test_transform_image_edit_response_returns_each_image_and_tolerates_null_lists(body, expected_images): - response = BedrockStabilityImageEditConfig().transform_image_edit_response( - model="stability.stable-image-inpaint-v1:0", - raw_response=httpx.Response(200, json=body), - logging_obj=MagicMock(), - ) - - assert [image.b64_json for image in response.data or []] == expected_images - - -def test_transform_image_edit_response_raises_the_first_finish_reason(): - with pytest.raises(Exception, match="Bedrock Stability error: Filter reason: prompt"): - BedrockStabilityImageEditConfig().transform_image_edit_response( - model="stability.stable-image-inpaint-v1:0", - raw_response=httpx.Response(200, json={"images": [None], "finish_reasons": ["Filter reason: prompt"]}), - logging_obj=MagicMock(), - ) diff --git a/tests/unit/proxy/test_caching_routes.py b/tests/unit/proxy/test_caching_routes.py index 0190b9390e8..5acc2513fd1 100644 --- a/tests/unit/proxy/test_caching_routes.py +++ b/tests/unit/proxy/test_caching_routes.py @@ -1,8 +1,11 @@ import json +from typing import Final +from unittest.mock import MagicMock import pytest from fastapi.testclient import TestClient from pydantic import ValidationError +from pytest_mock import MockerFixture @@ -53,8 +56,8 @@ def mock_redis_failure(mocker): @pytest.fixture -def mock_redis_delete(mocker): - mock_cache = mocker.MagicMock() +def mock_redis_delete(mocker: MockerFixture) -> MagicMock: + mock_cache: Final = mocker.MagicMock() mock_cache.type = "redis" mock_cache.delete_cache_keys = mocker.AsyncMock() mocker.patch.object(litellm, "cache", mock_cache) @@ -62,16 +65,18 @@ def mock_redis_delete(mocker): @pytest.mark.parametrize("keys", [["a-key", "b-key"], [12345, 1.5]]) -def test_cache_delete_passes_string_and_numeric_keys_to_the_cache(mock_redis_delete, keys): - response = client.post("/cache/delete", headers={"Authorization": "Bearer sk-9876"}, json={"keys": keys}) +def test_cache_delete_passes_string_and_numeric_keys_to_the_cache( + mock_redis_delete: MagicMock, keys: list[str | int | float] +) -> None: + response: Final = client.post("/cache/delete", headers={"Authorization": "Bearer sk-9876"}, json={"keys": keys}) assert response.status_code == 200 assert response.json() == {"status": "success"} mock_redis_delete.delete_cache_keys.assert_awaited_once_with(keys=keys) -def test_cache_delete_rejects_keys_that_are_not_a_list_with_a_400_that_hides_them(mock_redis_delete): - response = client.post("/cache/delete", headers={"Authorization": "Bearer sk-9876"}, json={"keys": "sk-secret-key"}) +def test_cache_delete_rejects_keys_that_are_not_a_list_with_a_400_that_hides_them(mock_redis_delete: MagicMock) -> None: + response: Final = client.post("/cache/delete", headers={"Authorization": "Bearer sk-9876"}, json={"keys": "sk-secret-key"}) assert response.status_code == 400 assert response.json()["detail"][0]["loc"] == ["keys"]