scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

refactor(passthrough): share the relay shape table and price FLUX 2 provider relays

Browse source
This commit is contained in:
mateo-berri 2026-09-07 18:37:58 -07:00
parent 17b7003592
commit 12d9413860
4 changed files with 106 additions and 81 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
litellm/llms/azure/passthrough/transformation.py
View file

@ -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()

36
litellm/llms/azure_ai/passthrough/transformation.py
View file

@ -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(

58
litellm/llms/base_llm/passthrough/transformation.py
View file

@ -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]:

49
tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py
View file

@ -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]",
]