scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 100

fix(azure_ai): price streamed Responses relays from their terminal event

Browse source 
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.
This commit is contained in:
mateo-berri 2026-09-07 22:38:41 -07:00
parent 9666a21cf0
commit 13db79da46
6 changed files with 107 additions and 45 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

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

9
litellm/llms/openai/responses/transformation.py
View file

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

3
litellm/types/llms/openai.py
View file

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

35
tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py
View file

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

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

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