From bd31b30c21ac21896fbe02cb18c2ada73483a1dd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:12:00 -0700 Subject: [PATCH] refactor(types): leave provider reply parsers for a follow-up with live provider calls Restores the Bedrock Stability, Bedrock KB, AgentCore, Replicate, Recraft, Aleph Alpha and Guardrails AI parsers to main, since no live provider call could prove the stricter schemas. Also types the cache delete route tests. --- .../providers/bedrock_agentcore/handler.py | 8 ++--- .../image_edit/stability_transformation.py | 18 ++-------- .../bedrock/vector_stores/transformation.py | 36 +++++-------------- .../llms/deprecated_providers/aleph_alpha.py | 19 +--------- .../llms/recraft/image_edit/transformation.py | 21 +---------- litellm/llms/replicate/chat/transformation.py | 24 ++----------- .../guardrails_ai/guardrails_ai.py | 25 ++++++------- .../test_stability_transformation.py | 33 ----------------- tests/unit/proxy/test_caching_routes.py | 17 +++++---- 9 files changed, 39 insertions(+), 162 deletions(-) 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"]