scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

refactor(rerank): reuse the unchanged Decisions transport

Browse source
This commit is contained in:
Ishaan Jaff 2026-10-20 21:34:10 -07:00
parent bc1c45c7ea
commit 2625b60f30
No known key found for this signature in database
5 changed files with 195 additions and 56 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

54
litellm/cost_calculator.py
View file

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

22
litellm/decisions/call.py
View file

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

25
litellm/rerank_api/decision_rerank.py
View file

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

9
litellm/responses/mcp/round_billing.py
View file

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

141
tests/unit/rerank_api/test_decision_rerank.py
View file

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