From 12d9413860e96f1cf2b0d138466b34fcae53f925 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:37:58 -0700 Subject: [PATCH] refactor(passthrough): share the relay shape table and price FLUX 2 provider relays --- .../llms/azure/passthrough/transformation.py | 44 +++----------- .../azure_ai/passthrough/transformation.py | 36 ++++-------- .../base_llm/passthrough/transformation.py | 58 ++++++++++++++----- ...est_azure_ai_passthrough_transformation.py | 49 +++++++++++++--- 4 files changed, 106 insertions(+), 81 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index d7078d94a0b..bdb854bb87a 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,5 +1,4 @@ -from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional import httpx @@ -10,7 +9,8 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.base_llm.passthrough.transformation import ( BasePassthroughConfig, - relayed_json_object, + RelayShape, + logged_relay_shape, replace_path_segment, strip_leading_model_segment, ) @@ -22,6 +22,7 @@ from litellm.types.utils import CallTypes, EmbeddingResponse, ImageResponse if TYPE_CHECKING: from httpx import URL + from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse from litellm.types.utils import CostResponseTypes @@ -41,40 +42,13 @@ def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, obj return details.request_data.messages if details.request_data else None -@dataclass(frozen=True, slots=True) -class OpenAIRelayShape: - path_suffix: str - call_type: CallTypes - parse: Callable[[Mapping[str, object]], EmbeddingResponse | ImageResponse | ResponsesAPIResponse] - - OPENAI_RELAY_SHAPES: Final = ( - OpenAIRelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate), - OpenAIRelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate), - OpenAIRelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate), + RelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate), + RelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate), + RelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate), ) -def logged_openai_response( - httpx_response: Response, logging_obj: Logging, endpoint: str -) -> EmbeddingResponse | ImageResponse | ResponsesAPIResponse | None: - relayed_path: Final = f"/{endpoint.strip('/')}" - shape: Final = next( - (candidate for candidate in OPENAI_RELAY_SHAPES if relayed_path.endswith(candidate.path_suffix)), None - ) - body: Final = relayed_json_object(httpx_response) if shape else None - if shape is None or body is None: - return None - try: - parsed: Final = shape.parse(body) - except ValidationError: - return None - logging_obj.call_type = ( - shape.call_type.value - ) # rebind-ok: routes cost calculation to the relayed shape's pricing path - return parsed - - class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return bool(request_data.get("stream")) @@ -151,13 +125,13 @@ class AzurePassthroughConfig(BasePassthroughConfig): request_data: dict, logging_obj: Logging, endpoint: str, - ) -> Optional["CostResponseTypes | ResponsesAPIResponse"]: + ) -> Optional["LoggedRelayResponse"]: from litellm import encoding from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.utils import ModelResponse if "chat/completions" not in endpoint: - return logged_openai_response(httpx_response, logging_obj, endpoint) + return logged_relay_shape(OPENAI_RELAY_SHAPES, httpx_response, logging_obj, endpoint) openai_chat_config: Final = OpenAIGPTConfig() diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index c49e10566ed..e854874eaee 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -16,19 +16,20 @@ from litellm.llms.azure_ai.common_utils import ( from litellm.llms.azure_ai.ocr.common_utils import get_azure_ai_ocr_config from litellm.llms.base_llm.passthrough.transformation import ( BasePassthroughConfig, - relayed_json_object, + RelayShape, + logged_relay_shape, strip_leading_model_segment, ) from litellm.types.llms.openai import AllMessageValues from litellm.types.rerank import RerankResponse -from litellm.types.utils import CallTypes, StandardPassThroughResponseObject +from litellm.types.utils import CallTypes, ImageResponse, StandardPassThroughResponseObject if TYPE_CHECKING: from httpx import URL, Response from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse - from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse from litellm.types.utils import CostResponseTypes @@ -83,16 +84,10 @@ def relayed_body(httpx_response: Response) -> str | dict: return body if isinstance(body, dict) else httpx_response.text -def logged_rerank_response(httpx_response: Response, logging_obj: Logging, endpoint: str) -> RerankResponse | None: - body: Final = relayed_json_object(httpx_response) if f"/{endpoint.strip('/')}".endswith("/rerank") else None - if body is None: - return None - try: - rerank_response: Final = RerankResponse.model_validate(body) - except ValidationError: - return None - logging_obj.call_type = CallTypes.arerank.value # rebind-ok: routes cost calculation to the per-query rerank path - return rerank_response +FOUNDRY_RELAY_SHAPES: Final = ( + RelayShape("/rerank", CallTypes.arerank, RerankResponse.model_validate), + RelayShape("/providers/blackforestlabs/v1/flux-2-pro", CallTypes.aimage_generation, ImageResponse.model_validate), +) class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): @@ -148,14 +143,7 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): request_data: Mapping[str, object], logging_obj: Logging, endpoint: str, - ) -> ( - CostResponseTypes - | OCRResponse - | RerankResponse - | ResponsesAPIResponse - | StandardPassThroughResponseObject - | None - ): + ) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None: from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig chat_result: Final = AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict @@ -171,9 +159,9 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): ocr_result: Final = self.logged_ocr_response(model, httpx_response, logging_obj, endpoint) if ocr_result is not None: return ocr_result - rerank_result: Final = logged_rerank_response(httpx_response, logging_obj, endpoint) - if rerank_result is not None: - return rerank_result + foundry_result: Final = logged_relay_shape(FOUNDRY_RELAY_SHAPES, httpx_response, logging_obj, endpoint) + if foundry_result is not None: + return foundry_result return StandardPassThroughResponseObject(response=relayed_body(httpx_response)) def logged_ocr_response( diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index 975b74137db..93ec5a09e62 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -1,10 +1,15 @@ +from __future__ import annotations + import re from abc import abstractmethod -from collections.abc import Mapping -from typing import TYPE_CHECKING, Final, Optional, Union +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, TypeAlias from pydantic import TypeAdapter, ValidationError +from litellm.types.utils import CallTypes + from ..base_utils import BaseLLMModelInfo if TYPE_CHECKING: @@ -18,6 +23,8 @@ if TYPE_CHECKING: from ..chat.transformation import BaseLLMException from ..ocr.transformation import OCRResponse + LoggedRelayResponse: TypeAlias = CostResponseTypes | RerankResponse | ResponsesAPIResponse + RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) @@ -39,7 +46,7 @@ def replace_path_segment(endpoint: str, segment: str, replacement: str) -> str: return bounded_segment.sub(lambda _: replacement, endpoint) -def relayed_json_object(httpx_response: "Response") -> Mapping[str, object] | None: +def relayed_json_object(httpx_response: Response) -> Mapping[str, object] | None: if httpx_response.status_code != 200: return None try: @@ -48,6 +55,31 @@ def relayed_json_object(httpx_response: "Response") -> Mapping[str, object] | No return None +@dataclass(frozen=True, slots=True) +class RelayShape: + path_suffix: str + call_type: CallTypes + parse: Callable[[Mapping[str, object]], LoggedRelayResponse] + + +def logged_relay_shape( + shapes: Sequence[RelayShape], httpx_response: Response, logging_obj: LiteLLMLoggingObj, endpoint: str +) -> LoggedRelayResponse | None: + relayed_path: Final = f"/{endpoint.strip('/')}" + shape: Final = next((candidate for candidate in shapes if relayed_path.endswith(candidate.path_suffix)), None) + body: Final = relayed_json_object(httpx_response) if shape else None + if shape is None or body is None: + return None + try: + parsed: Final = shape.parse(body) + except ValidationError: + return None + logging_obj.call_type = ( + shape.call_type.value + ) # rebind-ok: routes cost calculation to the relayed shape's pricing path + return parsed + + class BasePassthroughConfig(BaseLLMModelInfo): @abstractmethod def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: @@ -60,7 +92,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): endpoint: str, base_target_url: str, request_query_params: Mapping[str, object] | None, - ) -> "URL": + ) -> URL: """ Helper function to add query params to the url Args: @@ -94,7 +126,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): endpoint: str, request_query_params: dict | None, litellm_params: dict, - ) -> tuple["URL", str]: + ) -> tuple[URL, str]: """ Get the complete url for the request Returns: @@ -124,9 +156,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): """ return headers, None - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, "Headers"] - ) -> "BaseLLMException": + def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException: from litellm.llms.base_llm.chat.transformation import BaseLLMException return BaseLLMException(status_code=status_code, message=error_message, headers=headers) @@ -135,23 +165,21 @@ class BasePassthroughConfig(BaseLLMModelInfo): self, model: str, custom_llm_provider: str, - httpx_response: "Response", + httpx_response: Response, request_data: dict, - logging_obj: "LiteLLMLoggingObj", + logging_obj: LiteLLMLoggingObj, endpoint: str, - ) -> Optional[ - "CostResponseTypes | OCRResponse | RerankResponse | ResponsesAPIResponse | StandardPassThroughResponseObject" - ]: + ) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None: pass def handle_logging_collected_chunks( self, all_chunks: list[str], - litellm_logging_obj: "LiteLLMLoggingObj", + litellm_logging_obj: LiteLLMLoggingObj, model: str, custom_llm_provider: str, endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> CostResponseTypes | None: return None def _convert_raw_bytes_to_str_lines(self, raw_bytes: list[bytes]) -> list[str]: diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 788aced50f8..8b6c76952d6 100644 --- a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -25,7 +25,9 @@ def clear_azure_ai_env(monkeypatch): def test_provider_config_manager_resolves_azure_ai_passthrough_config(): - config = ProviderConfigManager.get_provider_passthrough_config(model="Cohere-parse-v5", provider=LlmProviders.AZURE_AI) + config = ProviderConfigManager.get_provider_passthrough_config( + model="Cohere-parse-v5", provider=LlmProviders.AZURE_AI + ) assert isinstance(config, AzureAIPassthroughConfig) @@ -189,7 +191,10 @@ def test_no_credentials_at_all_raises(): [({"stream": True}, True), ({"stream": 1}, True), ({"stream": False}, False), ({}, False)], ) def test_is_streaming_request_reads_the_stream_flag(request_data, expected): - assert AzureAIPassthroughConfig().is_streaming_request(endpoint="models/chat/completions", request_data=request_data) is expected + assert ( + AzureAIPassthroughConfig().is_streaming_request(endpoint="models/chat/completions", request_data=request_data) + is expected + ) def _chat_completion_response() -> httpx.Response: @@ -266,7 +271,12 @@ def _relay_logging_obj(model: str, api_base: str) -> Logging: def _relay_logging_result( - config: AzureAIPassthroughConfig, model: str, native_path: str, body, api_base: str = FOUNDRY_BASE, status_code: int = 200 + config: AzureAIPassthroughConfig, + model: str, + native_path: str, + body, + api_base: str = FOUNDRY_BASE, + status_code: int = 200, ): relayed_url = f"{FOUNDRY_BASE}/{native_path}?api-version=2024-05-01-preview" logging_obj = _relay_logging_obj(model, api_base) @@ -355,7 +365,11 @@ def test_deployment_without_an_ocr_config_is_never_costed_as_ocr(): def test_accepted_ocr_job_without_a_result_body_is_not_costed(): result, logging_obj = _relay_logging_result( - AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", {"status": "running"}, status_code=202 + AzureAIPassthroughConfig(), + "mistral-document-ai-2512", + "providers/mistral/azure/ocr", + {"status": "running"}, + status_code=202, ) assert result == {"response": {"status": "running"}} @@ -364,7 +378,10 @@ def test_accepted_ocr_job_without_a_result_body_is_not_costed(): def test_unparseable_ocr_body_falls_back_to_the_passthrough_object(): result, logging_obj = _relay_logging_result( - AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", ["not", "an", "ocr", "body"] + AzureAIPassthroughConfig(), + "mistral-document-ai-2512", + "providers/mistral/azure/ocr", + ["not", "an", "ocr", "body"], ) assert result == {"response": '["not", "an", "ocr", "body"]'} @@ -423,6 +440,17 @@ def test_image_generation_relay_is_costed_per_image(): assert logging_obj._response_cost_calculator(result=result) == pytest.approx(per_image) +def test_flux_2_relay_through_the_provider_route_is_costed_per_image(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "FLUX.2-pro", "providers/blackforestlabs/v1/flux-2-pro", IMAGE_BODY + ) + per_image = litellm.get_model_info("azure_ai/FLUX.2-pro")["output_cost_per_image"] + + assert isinstance(result, ImageResponse) + assert logging_obj.call_type == "aimage_generation" + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(per_image) + + def test_rejected_rerank_relay_keeps_the_passthrough_object_and_call_type(): result, logging_obj = _relay_logging_result( AzureAIPassthroughConfig(), @@ -439,8 +467,15 @@ def test_rejected_rerank_relay_keeps_the_passthrough_object_and_call_type(): def test_streaming_chat_completion_chunks_are_costed_like_azure(): head = {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-5.4-mini"} chunks = [ - "data: " + json.dumps({**head, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}]}), - "data: " + json.dumps({**head, "choices": [], "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}}), + "data: " + + json.dumps( + { + **head, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + } + ), + "data: " + + json.dumps({**head, "choices": [], "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}}), "data: [DONE]", ]