From 13db79da46f41b9775e781ac0df9c1aa14c29f9a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 22:38:41 -0700 Subject: [PATCH] fix(azure_ai): price streamed Responses relays from their terminal event A streamed Responses API relay handed the success handler a bare ResponsesAPIResponse, which the streaming assembly step drops, so the relay never reached the spend callbacks. Hand it the terminal response.completed event instead, which the assembly step already converts, and cover the whole flush path with a regression test that fails on the previous tip. --- .../llms/azure/passthrough/transformation.py | 11 +-- .../base_llm/passthrough/transformation.py | 4 +- .../llms/openai/responses/transformation.py | 9 +- litellm/types/llms/openai.py | 3 + .../test_azure_passthrough_transformation.py | 35 ++++++-- ...est_azure_ai_passthrough_transformation.py | 90 +++++++++++++------ 6 files changed, 107 insertions(+), 45 deletions(-) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index ab7fe412628..85647ae02c9 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -15,7 +15,7 @@ from litellm.llms.base_llm.passthrough.transformation import ( strip_leading_model_segment, ) from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse +from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse, ResponsesTerminalEvent from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CallTypes, EmbeddingResponse, ImageResponse @@ -50,18 +50,19 @@ OPENAI_RELAY_SHAPES: Final = ( ) -def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> ResponsesAPIResponse | None: +def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> ResponsesTerminalEvent | None: + """A streaming logging object assembles the logged response from the terminal event, not from its body.""" from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig - terminal_response: Final = OpenAIResponsesAPIConfig.parse_terminal_response_from_stream_chunks( + terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks( all_chunks=list(all_chunks) ) - if terminal_response is None: + if terminal_event is None: return None logging_obj.call_type = ( RESPONSES_RELAY_SHAPE.call_type.value ) # rebind-ok: routes cost calculation to the relayed shape's pricing path - return terminal_response + return terminal_event class AzurePassthroughConfig(BasePassthroughConfig): diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index dc857dfc808..20180c5cfa2 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -16,14 +16,14 @@ if TYPE_CHECKING: from httpx import URL, Headers, Response from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesTerminalEvent from litellm.types.rerank import RerankResponse from litellm.types.utils import CostResponseTypes, StandardPassThroughResponseObject from ..chat.transformation import BaseLLMException from ..ocr.transformation import OCRResponse - LoggedRelayResponse: TypeAlias = CostResponseTypes | RerankResponse | ResponsesAPIResponse + LoggedRelayResponse: TypeAlias = CostResponseTypes | RerankResponse | ResponsesAPIResponse | ResponsesTerminalEvent RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 926de3e8854..cde399065fd 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -620,15 +620,20 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return event_pydantic_model.model_construct(**parsed_chunk) @staticmethod - def parse_terminal_response_from_stream_chunks(all_chunks: list[str]) -> ResponsesAPIResponse | None: + def parse_terminal_event_from_stream_chunks(all_chunks: list[str]) -> ResponsesTerminalEvent | None: for chunk_str in reversed(all_chunks): for event_model in (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent): try: - return event_model.model_validate_json(chunk_str.removeprefix("data: ")).response + return event_model.model_validate_json(chunk_str.removeprefix("data: ")) except ValueError: continue return None + @staticmethod + def parse_terminal_response_from_stream_chunks(all_chunks: list[str]) -> ResponsesAPIResponse | None: + terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks) + return None if terminal_event is None else terminal_event.response + @staticmethod def get_event_model_class(event_type: str) -> type[BaseLiteLLMOpenAIResponseObject]: """ diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index b6da9490e01..b7c4371f32f 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1564,6 +1564,9 @@ class ResponseIncompleteEvent(BaseLiteLLMOpenAIResponseObject): response: ResponsesAPIResponse +ResponsesTerminalEvent: TypeAlias = ResponseCompletedEvent | ResponseIncompleteEvent | ResponseFailedEvent + + class ResponsePartAddedEvent(BaseLiteLLMOpenAIResponseObject): type: Literal[ResponsesAPIStreamEvents.RESPONSE_PART_ADDED] item_id: str diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 42dd65ea661..6e85b0cdca6 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -9,7 +9,7 @@ import litellm from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse from litellm.types.utils import EmbeddingResponse, ModelResponse @@ -103,7 +103,9 @@ def _relay_logging_result(model: str, endpoint: str, body, status_code: int = 20 status_code=status_code, headers={"content-type": "application/json"}, content=json.dumps(body).encode("utf-8"), - request=httpx.Request("POST", f"https://my-resource.openai.azure.com/{endpoint}?api-version=2025-04-01-preview"), + request=httpx.Request( + "POST", f"https://my-resource.openai.azure.com/{endpoint}?api-version=2025-04-01-preview" + ), ) result = AzurePassthroughConfig().logging_non_streaming_response( model=model, @@ -291,7 +293,11 @@ def _azure_responses_stream_chunks(terminal_event: str | None = "response.comple "response.output_text.delta", {"type": "response.output_text.delta", "sequence_number": 1, "item_id": "msg_1", "delta": "hi"}, ), - ] + ([(terminal_event, {"type": terminal_event, "sequence_number": 2, "response": RESPONSES_BODY})] if terminal_event else []) + ] + ( + [(terminal_event, {"type": terminal_event, "sequence_number": 2, "response": RESPONSES_BODY})] + if terminal_event + else [] + ) return [line for name, payload in events for line in (f"event: {name}", _sse_line(payload))] @@ -307,10 +313,10 @@ def test_azure_passthrough_streaming_responses_chunks_are_costed_per_token(): ) info = litellm.get_model_info("azure/gpt-4.1-mini") - assert isinstance(response, ResponsesAPIResponse) - assert response.usage.input_tokens == 1000 + assert isinstance(response, ResponseCompletedEvent) + assert response.response.usage.input_tokens == 1000 assert logging_obj.call_type == "aresponses" - assert logging_obj._response_cost_calculator(result=response) == pytest.approx( + assert logging_obj._response_cost_calculator(result=response.response) == pytest.approx( 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] ) @@ -373,7 +379,10 @@ def test_azure_passthrough_url_strips_the_leading_router_model_segment(): litellm_params={}, ) - assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + assert ( + str(url) + == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + ) def test_azure_passthrough_url_rewrites_the_model_group_only_as_a_whole_segment(): @@ -386,7 +395,10 @@ def test_azure_passthrough_url_rewrites_the_model_group_only_as_a_whole_segment( litellm_params={"litellm_metadata": {"model_group": "gpt"}}, ) - assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + assert ( + str(url) + == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + ) @pytest.mark.parametrize( @@ -394,4 +406,9 @@ def test_azure_passthrough_url_rewrites_the_model_group_only_as_a_whole_segment( [({"stream": True}, True), ({"stream": 1}, True), ({"stream": False}, False), ({}, False)], ) def test_azure_passthrough_is_streaming_request_reads_the_stream_flag(request_data, expected): - assert AzurePassthroughConfig().is_streaming_request(endpoint="openai/deployments/x/chat/completions", request_data=request_data) is expected + assert ( + AzurePassthroughConfig().is_streaming_request( + endpoint="openai/deployments/x/chat/completions", request_data=request_data + ) + is expected + ) 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 0b23da984e8..064be954518 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 @@ -6,6 +6,7 @@ import httpx import pytest import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure_ai.passthrough.transformation import AzureAIPassthroughConfig from litellm.llms.base_llm.ocr.transformation import OCRResponse @@ -14,6 +15,36 @@ from litellm.types.utils import EmbeddingResponse, ImageResponse, LlmProviders, from litellm.utils import ProviderConfigManager FOUNDRY_BASE = "https://my-resource.services.ai.azure.com" +RESPONSES_COMPLETED_EVENT = { + "type": "response.completed", + "sequence_number": 2, + "response": { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.4-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, + }, +} + + +class _SpendProbe(CustomLogger): + logged_call_type: str | None = None + logged_cost: float | None = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged_call_type = kwargs["call_type"] + self.logged_cost = kwargs["response_cost"] @pytest.fixture(autouse=True) @@ -87,7 +118,9 @@ def test_api_base_that_already_ends_in_models_is_cut_back_to_the_foundry_root(): def test_full_url_api_base_that_already_ends_with_the_native_path_is_not_doubled(): - model_router_url = "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router/chat/completions" + model_router_url = ( + "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router/chat/completions" + ) url, base = AzureAIPassthroughConfig().get_complete_url( api_base=f"{model_router_url}?api-version=2025-01-01-preview", @@ -267,21 +300,29 @@ def test_non_chat_relay_with_a_non_json_body_logs_the_raw_text(): assert _non_chat_logging_result(b"page one", "text/plain") == {"response": "page one"} -def _relay_logging_obj(model: str, api_base: str) -> Logging: +def _relay_logging_obj( + model: str, + api_base: str, + stream: bool = False, + callbacks: list[CustomLogger] | None = None, + endpoint: str = "", +) -> Logging: logging_obj = Logging( model=model, messages=[], - stream=False, + stream=stream, call_type="allm_passthrough_route", start_time=datetime.now(), litellm_call_id="call-1", function_id="fn-1", + dynamic_async_success_callbacks=callbacks, ) logging_obj.update_environment_variables( model=model, litellm_params={"api_base": api_base, "custom_llm_provider": "azure_ai"}, optional_params={}, custom_llm_provider="azure_ai", + endpoint=endpoint, ) return logging_obj @@ -509,31 +550,10 @@ def test_streaming_chat_completion_chunks_are_costed_like_azure(): def test_streaming_responses_chunks_through_a_router_relay_are_costed_like_azure(): - completed = { - "type": "response.completed", - "sequence_number": 2, - "response": { - "id": "resp_1", - "object": "response", - "created_at": 1, - "status": "completed", - "model": "gpt-5.4-mini", - "output": [ - { - "type": "message", - "id": "msg_1", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "hi", "annotations": []}], - } - ], - "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, - }, - } logging_obj = _relay_logging_obj("gpt-5.4-mini", FOUNDRY_BASE) response = AzureAIPassthroughConfig().handle_logging_collected_chunks( - all_chunks=["event: response.completed", "data: " + json.dumps(completed)], + all_chunks=["event: response.completed", "data: " + json.dumps(RESPONSES_COMPLETED_EVENT)], litellm_logging_obj=logging_obj, model="gpt-5.4-mini", custom_llm_provider="azure_ai", @@ -542,8 +562,24 @@ def test_streaming_responses_chunks_through_a_router_relay_are_costed_like_azure info = litellm.get_model_info("azure_ai/gpt-5.4-mini") assert response is not None - assert response.usage.output_tokens == 100 + assert response.response.usage.output_tokens == 100 assert logging_obj.call_type == "aresponses" - assert logging_obj._response_cost_calculator(result=response) == pytest.approx( + assert logging_obj._response_cost_calculator(result=response.response) == pytest.approx( 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] ) + + +async def test_streaming_responses_relay_flush_reaches_the_success_callbacks_with_a_price(): + probe = _SpendProbe() + logging_obj = _relay_logging_obj( + "gpt-5.4-mini", FOUNDRY_BASE, stream=True, callbacks=[probe], endpoint="gpt/openai/responses" + ) + stream = "event: response.completed\ndata: " + json.dumps(RESPONSES_COMPLETED_EVENT) + "\n\n" + + await logging_obj.async_flush_passthrough_collected_chunks( + raw_bytes=[stream.encode()], provider_config=AzureAIPassthroughConfig() + ) + info = litellm.get_model_info("azure_ai/gpt-5.4-mini") + + assert probe.logged_call_type == "allm_passthrough_route" + assert probe.logged_cost == pytest.approx(1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"])