From 2625b60f30c341e845a987cc413faf3883b9d0b6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 20 Oct 2026 21:34:10 -0700 Subject: [PATCH] refactor(rerank): reuse the unchanged Decisions transport --- litellm/cost_calculator.py | 54 +++++-- litellm/decisions/call.py | 22 +-- litellm/rerank_api/decision_rerank.py | 25 ++-- litellm/responses/mcp/round_billing.py | 9 +- tests/unit/rerank_api/test_decision_rerank.py | 141 +++++++++++++++++- 5 files changed, 195 insertions(+), 56 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 1de8f094877..37804888c58 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -5,7 +5,7 @@ import math import time from collections.abc import Callable, Mapping, Sequence from datetime import datetime -from functools import lru_cache, partial, reduce +from functools import lru_cache, partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast @@ -1501,6 +1501,43 @@ def _response_time_ms_for_cost( return 0.0 +def _decision_rerank_breakdown( + breakdowns: Sequence[CostBreakdown | None], repeated_fixed: float +) -> CostBreakdown | None: + complete: Final = tuple(item for item in breakdowns if item is not None) + if not complete or len(complete) != len(breakdowns): + return None + monetary_fields: Final = ( + "input_cost", + "cache_read_cost", + "cache_creation_cost", + "output_cost", + "reasoning_cost", + "total_cost", + "tool_usage_cost", + "guardrail_cost", + "original_cost", + "discount_amount", + "margin_total_amount", + ) + totals: Final = { + key: sum(value for item in complete if isinstance(value := item.get(key), (int, float))) + for key in monetary_fields + if any(key in item for item in complete) + } + additional_keys: Final = frozenset[str]().union(*(item.get("additional_costs", {}) for item in complete)) + return TypeAdapter(CostBreakdown).validate_python( + { + **complete[-1], + **totals, + "additional_costs": { + key: sum(item.get("additional_costs", {}).get(key, 0.0) for item in complete) for key in additional_keys + }, + **{key: totals[key] - repeated_fixed for key in ("total_cost", "margin_total_amount") if key in totals}, + } + ) + + def _decision_rerank_cost( responses: Sequence[OpenAIDecisionResponse], price: Callable[[OpenAIDecisionResponse], float], @@ -1508,8 +1545,6 @@ def _decision_rerank_cost( provider: str | None, custom_pricing: bool, ) -> float: - from litellm.responses.mcp.round_billing import summed_cost_breakdown - def price_child( response: OpenAIDecisionResponse, ) -> tuple[float, CostBreakdown | None, float]: @@ -1531,18 +1566,9 @@ def _decision_rerank_cost( prices: Final = tuple(price_child(response) for response in responses) repeated_fixed: Final = sum(item[2] for item in prices) - max((item[2] for item in prices), default=0.0) - def merge_breakdown( - accumulated: CostBreakdown | None, - item: tuple[OpenAIDecisionResponse, tuple[float, CostBreakdown | None, float]], - ) -> CostBreakdown | None: - return summed_cost_breakdown(responses[0], accumulated, item[0], item[1][1]) - if logging_obj is not None and responses: - logging_obj.cost_breakdown = reduce( # rebind-ok: the outer request logs the aggregate rather than the last child - merge_breakdown, - zip(responses[1:], prices[1:], strict=True), - prices[0][1], - ) + combined: Final = _decision_rerank_breakdown(tuple(item[1] for item in prices), repeated_fixed) + logging_obj.cost_breakdown = combined # rebind-ok: log the aggregate rather than the last child return sum(item[0] for item in prices) - repeated_fixed diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index 2538c8ee390..abbeccc10b3 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -210,12 +210,6 @@ class DecisionsCall: caller_questions: tuple[DecisionsIRQuestion, ...] -@dataclass(frozen=True, slots=True) -class DecisionsResult: - ir_response: DecisionsIRResponse - response: DecisionsResponse | OpenAIDecisionResponse - - def supported_decisions_providers() -> tuple[str, ...]: return tuple( provider.value @@ -756,7 +750,7 @@ def _map_upstream_exception(error: Exception, call: DecisionsCall) -> Exception: ) -async def asend_result(call: DecisionsCall) -> DecisionsResult: +async def asend(call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse: supported: Final = _checked_image_count(_checked_image_types(call)) inlined: Final = _checked_image_sizes(_checked_image_types(await _async_with_remote_images_inlined(supported))) body: Final = _request_body(inlined) @@ -777,10 +771,10 @@ async def asend_result(call: DecisionsCall) -> DecisionsResult: ) except Exception as error: raise _map_upstream_exception(error, inlined) from error - return DecisionsResult(ir_response=response, response=_format_response(response, inlined)) + return _format_response(response, inlined) -def send_result(call: DecisionsCall) -> DecisionsResult: +def send(call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse: supported: Final = _checked_image_count(_checked_image_types(call)) inlined: Final = _checked_image_sizes(_checked_image_types(_with_remote_images_inlined(supported))) body: Final = _request_body(inlined) @@ -801,15 +795,7 @@ def send_result(call: DecisionsCall) -> DecisionsResult: ) except Exception as error: raise _map_upstream_exception(error, inlined) from error - return DecisionsResult(ir_response=response, response=_format_response(response, inlined)) - - -async def asend(call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse: - return (await asend_result(call)).response - - -def send(call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse: - return send_result(call).response + return _format_response(response, inlined) def reject_other_format(*, model: str, kwargs: Mapping[str, object], field_name: str, other_function: str) -> None: diff --git a/litellm/rerank_api/decision_rerank.py b/litellm/rerank_api/decision_rerank.py index 40d82ddad1c..f20ea13b884 100644 --- a/litellm/rerank_api/decision_rerank.py +++ b/litellm/rerank_api/decision_rerank.py @@ -21,15 +21,13 @@ import litellm from litellm.decisions.call import ( DecisionsCall, OpenAIDecisionRequestFields, - asend_result, + asend, prepare_call, - send_result, + send, ) from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.decisions import ( - DecisionsIRChoiceAnswer, - DecisionsIRRefusal, - DecisionsIRResponse, + OpenAIChoiceAnswer, OpenAIDecisionResponse, OpenAIPredicateAnswer, OpenAIRefusalAnswer, @@ -96,7 +94,6 @@ class _DocumentCall: class _Outcome: document: _DocumentCall response: OpenAIDecisionResponse | None = None - ir_response: DecisionsIRResponse | None = None error: Exception | None = None @@ -120,10 +117,9 @@ def _send(document: _DocumentCall, deadline: float) -> _Outcome: if isinstance(call, litellm.Timeout): return _Outcome(document, error=call) try: - result: Final = send_result(call) - response: Final = result.response + response: Final = send(call) if isinstance(response, OpenAIDecisionResponse): - return _Outcome(document, response=response, ir_response=result.ir_response) + return _Outcome(document, response=response) return _Outcome(document, error=_invalid_response(document, "expected a Decisions response")) except Exception as error: # noqa: BLE001 # preserve mapped provider failures alongside completed billable calls return _Outcome(document, error=error) @@ -135,10 +131,9 @@ async def _asend(document: _DocumentCall, deadline: float, semaphore: asyncio.Se if isinstance(call, litellm.Timeout): return _Outcome(document, error=call) try: - result: Final = await asend_result(call) - response: Final = result.response + response: Final = await asend(call) if isinstance(response, OpenAIDecisionResponse): - return _Outcome(document, response=response, ir_response=result.ir_response) + return _Outcome(document, response=response) return _Outcome( document, error=_invalid_response(document, "expected a Decisions response"), @@ -194,17 +189,17 @@ def _is_refusal(response: OpenAIDecisionResponse) -> bool: def _choice_scores(outcome: _Outcome, request: _Request) -> tuple[float, ...] | Exception: if outcome.error is not None: return outcome.error - response: Final = outcome.ir_response + response: Final = outcome.response if response is None or len(response.answers) != 1: return _invalid_response(outcome.document, "expected one choice answer") answer: Final = response.answers[0] - if isinstance(answer, DecisionsIRRefusal): + if isinstance(answer, OpenAIRefusalAnswer): return litellm.ContentPolicyViolationError( message="Decision model refused to rank the candidate set", model=outcome.document.call.requested_model, llm_provider=outcome.document.call.custom_llm_provider, ) - if not isinstance(answer, DecisionsIRChoiceAnswer): + if not isinstance(answer, OpenAIChoiceAnswer) or answer.name != "relevance": return _invalid_response(outcome.document, "expected the relevance choice") labels: Final = tuple(f"p{index}" for index in range(len(request.documents))) probabilities: Final = {item.value: item.probability for item in answer.probabilities} diff --git a/litellm/responses/mcp/round_billing.py b/litellm/responses/mcp/round_billing.py index bc343872e24..93769fbe82a 100644 --- a/litellm/responses/mcp/round_billing.py +++ b/litellm/responses/mcp/round_billing.py @@ -10,7 +10,6 @@ from pydantic_core import to_jsonable_python from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.responses.utils import ResponseAPILoggingUtils -from litellm.types.decisions import OpenAIDecisionResponse from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.utils import CostBreakdown, ModelResponse, Usage @@ -90,9 +89,9 @@ def _usage_of_both_rounds( def _repeated_fixed_margin( - first: ResponsesAPIResponse | ModelResponse | OpenAIDecisionResponse, + first: ResponsesAPIResponse | ModelResponse, first_breakdown: CostBreakdown | None, - final: ResponsesAPIResponse | ModelResponse | OpenAIDecisionResponse, + final: ResponsesAPIResponse | ModelResponse, final_breakdown: CostBreakdown | None, ) -> float: if first_breakdown is None or final_breakdown is None or first_breakdown is final_breakdown: @@ -118,9 +117,9 @@ def breakdown_less_repeated_fixed_margin( def summed_cost_breakdown( - first: ResponsesAPIResponse | ModelResponse | OpenAIDecisionResponse, + first: ResponsesAPIResponse | ModelResponse, first_breakdown: CostBreakdown | None, - final: ResponsesAPIResponse | ModelResponse | OpenAIDecisionResponse, + final: ResponsesAPIResponse | ModelResponse, final_breakdown: CostBreakdown | None, ) -> CostBreakdown | None: """The breakdown of the charged rounds, or None when a charged round was priced without its own breakdown, since a diff --git a/tests/unit/rerank_api/test_decision_rerank.py b/tests/unit/rerank_api/test_decision_rerank.py index 271ab327727..54558afe483 100644 --- a/tests/unit/rerank_api/test_decision_rerank.py +++ b/tests/unit/rerank_api/test_decision_rerank.py @@ -7,9 +7,10 @@ import pytest import litellm from litellm.caching.caching import Cache +from litellm.cost_calculator import _decision_rerank_breakdown from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.types.utils import CallTypes +from litellm.types.utils import CallTypes, CostBreakdown from litellm.utils import return_raw_request MODEL: Final = "typesafe/jevrank-unit-test" @@ -379,9 +380,7 @@ async def test_multiple_choice_ranks_one_whole_set_and_bills_once(sync, respx_mo @pytest.mark.parametrize( "probabilities", [ - {"p0": 1}, {"p0": 0.5, "wrong": 0.5}, - {"p0": 0.5, "p1": 0.5, "extra": 0}, {"p0": 0.8, "p1": 0.8}, {"p0": -0.1, "p1": 1.1}, {"p0": 0, "p1": 0}, @@ -397,6 +396,64 @@ async def test_multiple_choice_rejects_malformed_distributions(probabilities, re ) +@pytest.mark.parametrize( + "model,url", + [ + (MODEL, URL), + ("perplexity/pplx-unit-test", "https://api.perplexity.ai/v1/decisions"), + ("openai/decision-unit-test", "https://api.openai.com/v1/decisions"), + ], +) +@pytest.mark.parametrize("probabilities", [{"p0": 1}, {"p0": 1, "p1": 0, "extra": 0}]) +async def test_multiple_choice_uses_existing_decisions_normalization(model, url, probabilities, respx_mock): + if model.startswith("openai/"): + upstream = httpx.Response( + 200, + json={ + "model": "decision-unit-test", + "answers": [ + { + "type": "choice", + "name": "relevance", + "choice": "p0", + "confidence": 1, + "probabilities": [ + {"value": label, "probability": score} for label, score in probabilities.items() + ], + } + ], + "usage": {"input_tokens": 20, "output_tokens": 0, "total_tokens": 20}, + }, + ) + else: + upstream = choice_response(probabilities) + route = respx_mock.post(url).mock(return_value=upstream) + native = await litellm.adecisions( + model=model, + api_key="test-key", + input="warranty", + questions=[ + { + "type": "choice", + "name": "relevance", + "instructions": "Which passage is most relevant?", + "choices": [{"value": "p0"}, {"value": "p1"}], + } + ], + ) + result = await rerank( + False, + model=model, + documents=["one", "two"], + decision_rerank={"strategy": "multiple_choice"}, + ) + native_scores = tuple(item.probability for item in native.answers[0].probabilities) + assert native_scores == (1, 0) + assert tuple(item["relevance_score"] for item in result.results) == native_scores + assert tuple(item["index"] for item in result.results) == (0, 1) + assert route.call_count == 2 + + async def test_multiple_choice_rejects_excess_mass_at_maximum_candidate_count( respx_mock, ): @@ -560,7 +617,7 @@ async def test_direct_token_price_override_is_respected(respx_mock): ], ], ) -async def test_openai_choice_duplicate_labels_and_nonfinite_scores_rejected(probabilities, respx_mock): +async def test_openai_choice_invalid_normalized_scores_rejected(probabilities, respx_mock): payload = { "model": "test-decision", "answers": [ @@ -586,6 +643,38 @@ async def test_openai_choice_duplicate_labels_and_nonfinite_scores_rejected(prob ) +async def test_openai_choice_duplicate_labels_use_native_last_value(respx_mock): + respx_mock.post("https://api.openai.com/v1/decisions").respond( + json={ + "model": "test-decision", + "answers": [ + { + "type": "choice", + "name": "relevance", + "choice": "p0", + "confidence": 0.7, + "probabilities": [ + {"value": "p0", "probability": 0.2}, + {"value": "p0", "probability": 0.7}, + {"value": "p1", "probability": 0.3}, + ], + } + ], + "usage": {"input_tokens": 20, "output_tokens": 0, "total_tokens": 20}, + } + ) + result = await rerank( + False, + model="openai/test-decision", + documents=["one", "two"], + decision_rerank={"strategy": "multiple_choice"}, + ) + assert result.results == [ + {"index": 0, "relevance_score": 0.7}, + {"index": 1, "relevance_score": 0.3}, + ] + + async def test_success_cache_hit_reuses_aggregate_response_without_new_provider_calls(respx_mock, monkeypatch): monkeypatch.setattr(litellm, "cache", Cache(type="local")) route = respx_mock.post(URL).mock(return_value=response()) @@ -699,6 +788,50 @@ async def test_fixed_margin_uses_the_pricing_provider_for_aggregate_and_breakdow assert logger.events[0]["cost_breakdown"]["margin_fixed_amount"] == pytest.approx(0.02) +def test_aggregate_breakdown_adds_charges_and_preserves_pricing_metadata(): + first: Final[CostBreakdown] = { + "input_cost": 0.02, + "total_cost": 0.034, + "additional_costs": {"request": 0.004}, + "discount_percent": 0.1, + "margin_percent": 0.2, + "margin_fixed_amount": 0.01, + "margin_total_amount": 0.014, + "service_tier": "default", + "data_residency": "us", + "vertex_location": None, + } + second: Final[CostBreakdown] = { + **first, + "input_cost": 0.04, + "total_cost": 0.058, + "additional_costs": {"request": 0.004, "other": 0.002}, + "margin_total_amount": 0.018, + } + actual = _decision_rerank_breakdown((first, second), repeated_fixed=0.01) + assert actual is not None + assert actual["input_cost"] == pytest.approx(0.06) + assert actual["total_cost"] == pytest.approx(0.082) + assert actual["margin_total_amount"] == pytest.approx(0.022) + assert actual["additional_costs"] == {"request": pytest.approx(0.008), "other": pytest.approx(0.002)} + for key in ( + "discount_percent", + "margin_percent", + "margin_fixed_amount", + "service_tier", + "data_residency", + "vertex_location", + ): + assert actual[key] == second[key] + + +@pytest.mark.parametrize("missing_first", [True, False]) +def test_aggregate_breakdown_does_not_publish_partial_totals(missing_first): + priced: Final[CostBreakdown] = {"input_cost": 0.02, "total_cost": 0.02} + breakdowns = (None, priced) if missing_first else (priced, None) + assert _decision_rerank_breakdown(breakdowns, repeated_fixed=0) is None + + async def test_cancelled_child_fails_whole_request_and_bills_completed_usage(respx_mock, monkeypatch): logger = RecordingLogger() monkeypatch.setattr(litellm, "callbacks", [logger])