scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-27 06:13:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 97

feat(rerank): select decision methods automatically and bound admission

Browse source
This commit is contained in:
Ishaan Jaff 2026-10-20 22:18:46 -07:00
parent 931b03c0bc
commit 9b8e0715ea
15 changed files with 1087 additions and 194 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/caching/caching.py
View file

@ -42,7 +42,7 @@ from .redis_semantic_cache import RedisSemanticCache
from .s3_cache import S3Cache
RESPONSE_CACHE_TARGET: Final = "llm_response"
_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token", "rerank_mode", "decision_rerank"})
_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token"})
_DECODED_LITERAL: Final = TypeAdapter(object)

10
litellm/model_prices_and_context_window_backup.json
View file

@ -46044,6 +46044,7 @@
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"rerank_strategy": "multiple_choice",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/api-reference/decisions-post"
},
@ -46052,6 +46053,7 @@
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"rerank_strategy": "multiple_choice",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/getting-started/pricing"
},
@ -46060,6 +46062,7 @@
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"rerank_strategy": "multiple_choice",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
@ -46068,6 +46071,7 @@
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"rerank_strategy": "multiple_choice",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
@ -46076,6 +46080,7 @@
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"rerank_strategy": "multiple_choice",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
@ -46084,6 +46089,7 @@
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"rerank_strategy": "multiple_choice",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
@ -72235,6 +72241,7 @@
"litellm_provider": "typesafe",
"max_input_tokens": 65536,
"mode": "evaluation",
"rerank_strategy": "pointwise",
"output_cost_per_token": 0.0,
"source": "https://docs.typesafe.ai/models"
},
@ -72243,6 +72250,7 @@
"litellm_provider": "typesafe",
"max_input_tokens": 65536,
"mode": "evaluation",
"rerank_strategy": "pointwise",
"output_cost_per_token": 0.0,
"source": "https://docs.typesafe.ai/models",
"supports_box_questions": true
@ -72252,6 +72260,7 @@
"litellm_provider": "typesafe",
"max_input_tokens": 65536,
"mode": "evaluation",
"rerank_strategy": "pointwise",
"output_cost_per_token": 0.0,
"source": "https://docs.typesafe.ai/models"
},
@ -72260,6 +72269,7 @@
"litellm_provider": "typesafe",
"max_input_tokens": 65536,
"mode": "evaluation",
"rerank_strategy": "pointwise",
"output_cost_per_token": 0.0,
"source": "https://docs.typesafe.ai/models",
"supports_box_questions": true

266
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -7,6 +7,7 @@ import time
from collections.abc import AsyncIterator, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from itertools import chain
from types import MappingProxyType
from typing import Final, NoReturn, SupportsFloat, SupportsIndex, SupportsInt, cast
@ -42,6 +43,11 @@ from litellm.proxy.common_utils.user_api_key_cache import (
from litellm.proxy.spend_tracking.input_tokens import count_input_tokens, count_input_tokens_for_model
from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement, spend_counter_batch_scope
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.rerank_api.litellm_decisions_transformation.transformation import (
DecisionRerankRequest,
LiteLLMDecisionsRerankConfig,
)
from litellm.rerank_api.rerank_utils import get_decision_rerank_strategy
from litellm.router import Router
from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
from litellm.types.router import DeploymentTypedDict
@ -73,6 +79,7 @@ _COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = {
}
_CACHED_VALUE: Final = TypeAdapter(object)
_BUDGET_REQUEST: Final = TypeAdapter(Mapping[str, object])
_WINDOW_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
@ -273,12 +280,41 @@ async def reserve_budget_for_request(
raw_body=raw_body,
)
reservation_cost = estimate_request_max_cost(
request_body=request_body,
route=route,
llm_router=llm_router,
input_token_counts=input_token_counts,
decision_rerank: Final = "rerank" in route and any(
_decision_rerank_candidates(_BUDGET_REQUEST.validate_python(request_body), route, model, llm_router)
for model in _get_request_models(request_body, route, llm_router)
)
estimates: Final[tuple[float | None, float | None, Mapping[str, tuple[float, ...]]]] = (
await asyncio.to_thread(
_decision_rerank_reservation_estimates,
_BUDGET_REQUEST.validate_python(request_body),
route,
llm_router,
input_token_counts,
)
if decision_rerank
else (
estimate_request_max_cost(
request_body=request_body,
route=route,
llm_router=llm_router,
input_token_counts=input_token_counts,
),
estimate_request_input_cost(
request_body=request_body,
route=route,
llm_router=llm_router,
input_token_counts=input_token_counts,
),
{},
)
)
reservation_cost = estimates[0]
if decision_rerank and reservation_cost is None and fail_closed_budget_enforcement:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Budget enforcement unavailable: unable to estimate all decision rerank calls.",
)
# estimate_request_max_cost still returns None when the model is unknown
# to the cost map (no token-priced cost fields, e.g. image/audio routes).
# In that case we fall back to read-time enforcement only.
@ -309,19 +345,14 @@ async def reserve_budget_for_request(
if not applied_entries:
return None
input_cost: Final = estimate_request_input_cost(
request_body=request_body,
route=route,
llm_router=llm_router,
input_token_counts=input_token_counts,
)
budget_reservation: Final = {
"reserved_cost": reservation_cost,
"entries": applied_entries,
"finalized": False,
"callback_bound": False,
"input_cost": min(float(input_cost or 0.0), reservation_cost),
"input_cost": min(float(estimates[1] or 0.0), reservation_cost),
"input_tokens": max(input_token_counts.values(), default=None),
**({"decision_rerank_input_costs": estimates[2]} if decision_rerank else {}),
}
_start_reservation_lease_renewal(
budget_reservation=budget_reservation,
@ -1283,6 +1314,193 @@ def _coerce_datetime(value: object) -> datetime | None:
return None
@dataclass(frozen=True, slots=True)
class _DecisionRerankEstimate:
strategy: str
input_costs: tuple[float | None, ...]
max_cost: float | None
@property
def input_cost(self) -> float | None:
return (
None
if any(cost is None for cost in self.input_costs)
else sum(cost for cost in self.input_costs if cost is not None)
)
def _rerank_deployment_params(model: str, llm_router: Router | None) -> tuple[Mapping[str, object], ...]:
deployments: Final = (llm_router.get_model_list(model_name=model) or []) if llm_router is not None else []
return tuple(deployment["litellm_params"] for deployment in deployments) or ({},)
def _decision_rerank_candidates(
request_body: Mapping[str, object],
route: str,
model: str,
llm_router: Router | None,
) -> tuple[tuple[str, DecisionRerankRequest], ...]:
if "rerank" not in route:
return ()
return tuple(
(
str(params.get("model") or model),
LiteLLMDecisionsRerankConfig.validate_request(
{**params, **request_body},
model=model,
provider=str(params.get("custom_llm_provider") or ""),
strategy=strategy,
),
)
for params in _rerank_deployment_params(model, llm_router)
if (
strategy := get_decision_rerank_strategy(
str(params.get("model") or model),
custom_llm_provider=str(params.get("custom_llm_provider") or "") or None,
)
)
is not None
)
def _decision_rerank_inputs(request: DecisionRerankRequest, model: str) -> tuple[int | None, ...]:
count: Final = 1 if request.strategy == "multiple_choice" else len(request.documents)
return tuple(
count_input_tokens_for_model(
request_body={
"input": LiteLLMDecisionsRerankConfig.transform_rerank_request(request, index).model_dump_json()
},
model=model,
)
for index in range(count)
)
def _price_decision_rerank_inputs(
strategy: str,
model: str,
counts: Sequence[int | None],
model_info: Mapping[str, object],
) -> _DecisionRerankEstimate:
# Rerank does not forward caller output caps to the Decisions request.
maximums: Final = tuple(
_max_cost_for_cost_info(
request_body={},
route="/decisions",
model=model,
model_info=model_info,
input_tokens=tokens,
)
for tokens in counts
)
return _DecisionRerankEstimate(
strategy=strategy,
input_costs=tuple(
_input_cost_for_cost_info(
request_body={},
route="/decisions",
model=model,
model_info=model_info,
input_tokens=tokens,
)
for tokens in counts
),
max_cost=(
None if any(cost is None for cost in maximums) else sum(cost for cost in maximums if cost is not None)
),
)
def _decision_rerank_costs(
request_body: Mapping[str, object],
route: str,
model: str,
llm_router: Router | None,
) -> tuple[_DecisionRerankEstimate, ...]:
candidates: Final = _decision_rerank_candidates(request_body, route, model, llm_router)
if not candidates:
return ()
cost_infos: Final = _get_model_cost_infos(model=model, llm_router=llm_router)
inputs: Final = tuple(
(
backend_model,
request.strategy,
_decision_rerank_inputs(request, backend_model),
)
for backend_model, request in candidates
)
return tuple(
_price_decision_rerank_inputs(strategy=strategy, model=backend_model, counts=counts, model_info=model_info)
for backend_model, strategy, counts in inputs
for model_info in cost_infos # comprehension-ok: conservatively price every deployment/cost combination
) or (_DecisionRerankEstimate("pointwise", (None,), None),)
def _maximum_dispatch_costs(
candidates: Sequence[Sequence[float | None]],
) -> tuple[float, ...]:
return tuple(max(costs[index] or 0.0 for costs in candidates) for index in range(len(candidates[0])))
def _decision_rerank_reservation_estimates(
request_body: Mapping[str, object],
route: str,
llm_router: Router | None,
input_token_counts: Mapping[str, int],
) -> tuple[float | None, float | None, Mapping[str, tuple[float, ...]]]:
models: Final = _get_request_models(dict(request_body), route, llm_router)
decisions: Final = {
model: _decision_rerank_costs(request_body=request_body, route=route, model=model, llm_router=llm_router)
for model in models
}
all_estimates: Final = tuple(chain.from_iterable(decisions.values()))
if any(estimate.max_cost is None or estimate.input_cost is None for estimate in all_estimates):
return None, None, {}
maximums: Final = tuple(
_estimate_request_max_cost_for_model(
request_body=dict(request_body),
route=route,
model=model,
llm_router=llm_router,
input_tokens=input_token_counts.get(model),
decision_estimates=decisions[model],
)
for model in models
)
# A native sibling keeps its existing cancellation floor. The adapter replaces it before dispatch.
native_inputs: Final = tuple(
_estimate_request_input_cost_for_model(
request_body=dict(request_body),
route=route,
model=model,
llm_router=llm_router,
input_tokens=input_token_counts.get(model),
decision_estimates=(),
)
for model in models
if any(
get_decision_rerank_strategy(
str(params.get("model") or model),
custom_llm_provider=str(params.get("custom_llm_provider") or "") or None,
)
is None
for params in _rerank_deployment_params(model, llm_router)
)
)
strategies: Final = frozenset(estimate.strategy for estimate in all_estimates)
dispatch_costs: Final = {
strategy: _maximum_dispatch_costs(
tuple(estimate.input_costs for estimate in all_estimates if estimate.strategy == strategy)
)
for strategy in strategies
}
return (
max((cost for cost in maximums if cost is not None), default=None),
max((cost for cost in native_inputs if cost is not None), default=0.0),
dispatch_costs,
)
def estimate_request_max_cost(
request_body: dict,
route: str,
@ -1340,6 +1558,7 @@ def _estimate_request_input_cost_for_model(
model: str,
llm_router: Router | None,
input_tokens: int | None = None,
decision_estimates: Sequence[_DecisionRerankEstimate] | None = None,
) -> float | None:
estimates: Final = [
_input_cost_for_cost_info(
@ -1351,7 +1570,16 @@ def _estimate_request_input_cost_for_model(
)
for model_info in _get_model_cost_infos(model=model, llm_router=llm_router)
]
valid_estimates: Final = [estimate for estimate in estimates if estimate is not None]
transformed: Final = (
decision_estimates
if decision_estimates is not None
else _decision_rerank_costs(_BUDGET_REQUEST.validate_python(request_body), route, model, llm_router)
)
if any(estimate.input_cost is None for estimate in transformed):
return None
valid_estimates: Final = [
cost for cost in (*estimates, *(estimate.input_cost for estimate in transformed)) if cost is not None
]
return max(valid_estimates) if valid_estimates else None
@ -1388,6 +1616,7 @@ def _estimate_request_max_cost_for_model(
model: str,
llm_router: Router | None,
input_tokens: int | None = None,
decision_estimates: Sequence[_DecisionRerankEstimate] | None = None,
) -> float | None:
estimates: Final = [
_max_cost_for_cost_info(
@ -1399,7 +1628,16 @@ def _estimate_request_max_cost_for_model(
)
for model_info in _get_model_cost_infos(model=model, llm_router=llm_router)
]
valid_estimates: Final = [estimate for estimate in estimates if estimate is not None]
transformed: Final = (
decision_estimates
if decision_estimates is not None
else _decision_rerank_costs(_BUDGET_REQUEST.validate_python(request_body), route, model, llm_router)
)
if any(estimate.max_cost is None for estimate in transformed):
return None
valid_estimates: Final = [
cost for cost in (*estimates, *(estimate.max_cost for estimate in transformed)) if cost is not None
]
return max(valid_estimates) if valid_estimates else None

232
litellm/rerank_api/litellm_decisions_transformation/handler.py
View file

@ -2,14 +2,19 @@
import asyncio
import time
from collections.abc import Awaitable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor, wait
from dataclasses import dataclass, replace
from functools import partial
from typing import Final
from queue import Empty, SimpleQueue
from threading import Lock
from typing import Final, Literal
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm.decisions.call import DecisionsCall, OpenAIDecisionRequestFields, asend, prepare_call, send
from litellm.litellm_core_utils.core_helpers import budget_reservation_from_metadata
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.rerank_api.litellm_decisions_transformation.transformation import (
DecisionRerankRequest,
@ -20,7 +25,11 @@ from litellm.types.decisions import OpenAIDecisionResponse
from litellm.types.rerank import DecisionRerankResponse
from litellm.types.utils import Usage
_CHILD_OMIT: Final = frozenset({"litellm_logging_obj", "rerank_mode", "decision_rerank", "litellm_call_id"})
_METADATA: Final = TypeAdapter(Mapping[str, object])
_DISPATCH_COSTS: Final = TypeAdapter(Mapping[str, tuple[float, ...]])
_CHILD_OMIT: Final = frozenset({"litellm_logging_obj", "litellm_call_id"})
_MAX_CONCURRENCY: Final = 8
@dataclass(frozen=True, slots=True)
@ -30,6 +39,35 @@ class _DocumentCall:
call: DecisionsCall
def _dispatch_cost_recorder(
kwargs: Mapping[str, object], request: DecisionRerankRequest
) -> Callable[[int], None] | None:
try:
reservations: Final = (
budget_reservation_from_metadata(_METADATA.validate_python(kwargs.get(name) or {}))
for name in ("metadata", "litellm_metadata")
)
reservation: Final = next((item for item in reservations if item is not None), None)
if reservation is None:
return None
costs: Final = _DISPATCH_COSTS.validate_python(reservation.get("decision_rerank_input_costs") or {}).get(
request.strategy, ()
)
except ValidationError:
return None
if len(costs) != _call_count(request):
return None
reservation["input_cost"] = 0.0 # Only dispatched children belong in the cancellation floor.
lock: Final = Lock()
def record(index: int) -> None:
with lock:
current: Final = reservation.get("input_cost")
reservation["input_cost"] = (float(current) if isinstance(current, (float, int)) else 0.0) + costs[index]
return record
def _timeout(model: str, provider: str) -> litellm.Timeout:
return litellm.Timeout(
message="Decision reranking exceeded the request timeout",
@ -45,37 +83,63 @@ def _call_with_deadline(document: _DocumentCall, deadline: float) -> DecisionsCa
return replace(document.call, timeout=remaining)
def _send(document: _DocumentCall, deadline: float) -> DecisionRerankResult:
def _send(
document: _DocumentCall,
deadline: float,
record_dispatch: Callable[[int], None] | None,
) -> DecisionRerankResult:
call: Final = _call_with_deadline(document, deadline)
if isinstance(call, litellm.Timeout):
return DecisionRerankResult(document.index, document.text, error=call)
try:
if record_dispatch is not None:
record_dispatch(document.index)
response: Final = send(call)
if isinstance(response, OpenAIDecisionResponse):
return DecisionRerankResult(document.index, document.text, response=response)
return DecisionRerankResult(
document.index, document.text, error=_invalid_response(document, "expected a Decisions response")
document.index,
document.text,
error=_invalid_response(document, "expected a Decisions response"),
)
except Exception as error: # noqa: BLE001 # preserve mapped provider failures alongside completed billable calls
return DecisionRerankResult(document.index, document.text, error=error)
async def _asend(document: _DocumentCall, deadline: float, semaphore: asyncio.Semaphore) -> DecisionRerankResult:
async with semaphore:
call: Final = _call_with_deadline(document, deadline)
if isinstance(call, litellm.Timeout):
return DecisionRerankResult(document.index, document.text, error=call)
try:
response: Final = await asend(call)
if isinstance(response, OpenAIDecisionResponse):
return DecisionRerankResult(document.index, document.text, response=response)
return DecisionRerankResult(
document.index,
document.text,
error=_invalid_response(document, "expected a Decisions response"),
)
except Exception as error: # noqa: BLE001 # preserve mapped provider failures alongside completed billable calls
return DecisionRerankResult(document.index, document.text, error=error)
async def _asend(
document: _DocumentCall,
deadline: float,
cancellation_requested: asyncio.Event,
record_dispatch: Callable[[int], None] | None,
) -> DecisionRerankResult:
call: Final = _call_with_deadline(document, deadline)
if isinstance(call, litellm.Timeout):
return DecisionRerankResult(document.index, document.text, error=call)
try:
if record_dispatch is not None:
record_dispatch(document.index)
response: Final = await asend(call)
if isinstance(response, OpenAIDecisionResponse):
return DecisionRerankResult(document.index, document.text, response=response)
return DecisionRerankResult(
document.index,
document.text,
error=_invalid_response(document, "expected a Decisions response"),
)
except asyncio.CancelledError:
if cancellation_requested.is_set():
raise
return DecisionRerankResult(
document.index,
document.text,
error=litellm.APIConnectionError(
message="Decision rerank provider request was cancelled",
model=document.call.requested_model,
llm_provider=document.call.custom_llm_provider,
),
)
except Exception as error: # noqa: BLE001 # preserve mapped provider failures alongside completed billable calls
return DecisionRerankResult(document.index, document.text, error=error)
def _invalid_response(document: _DocumentCall, reason: str) -> litellm.InternalServerError:
@ -121,60 +185,89 @@ def _finish(
return response
def _call_count(request: DecisionRerankRequest) -> int:
return 1 if request.strategy == "multiple_choice" else len(request.documents)
def _rerank(
documents: tuple[_DocumentCall, ...],
prepare: Callable[[int], _DocumentCall],
request: DecisionRerankRequest,
logging_obj: Logging,
model: str,
provider: str,
deadline: float,
record_dispatch: Callable[[int], None] | None,
) -> DecisionRerankResponse:
deadline: Final = time.monotonic() + request.timeout
executor: Final = ThreadPoolExecutor(max_workers=request.decision_rerank.max_concurrency)
futures: Final = tuple(executor.submit(_send, document, deadline) for document in documents)
count: Final = _call_count(request)
workers: Final = min(count, _MAX_CONCURRENCY)
completed: Final[SimpleQueue[DecisionRerankResult]] = SimpleQueue()
def run(indices: range) -> None:
for index in indices:
if time.monotonic() >= deadline:
return
try:
completed.put(_send(prepare(index), deadline, record_dispatch))
except Exception as error: # noqa: BLE001 # preparation failures retain completed billable calls
completed.put(DecisionRerankResult(index, request.document_text(index), error=error))
executor: Final = ThreadPoolExecutor(max_workers=workers)
futures: Final = tuple(executor.submit(run, range(worker, count, workers)) for worker in range(workers))
try:
_, pending = wait(futures, timeout=max(0.0, deadline - time.monotonic()))
outcomes: Final = tuple(future.result() for future in futures if future.done() and not future.cancelled())
return _finish(outcomes, request, logging_obj, bool(pending), model, provider)
wait(futures, timeout=max(0.0, deadline - time.monotonic()))
outcomes: Final[list[DecisionRerankResult]] = [] # mutable-ok: retain completed results for partial billing
while True:
try:
outcomes.append(completed.get_nowait())
except Empty:
break
return _finish(outcomes, request, logging_obj, len(outcomes) < count, model, provider)
finally:
executor.shutdown(wait=False, cancel_futures=True)
async def _arerank(
documents: tuple[_DocumentCall, ...],
prepare: Callable[[int], _DocumentCall],
request: DecisionRerankRequest,
logging_obj: Logging,
model: str,
provider: str,
deadline: float,
record_dispatch: Callable[[int], None] | None,
) -> DecisionRerankResponse:
deadline: Final = time.monotonic() + request.timeout
semaphore: Final = asyncio.Semaphore(request.decision_rerank.max_concurrency)
tasks: Final = tuple(asyncio.create_task(_asend(document, deadline, semaphore)) for document in documents)
count: Final = _call_count(request)
workers: Final = min(count, _MAX_CONCURRENCY)
outcomes: Final[list[DecisionRerankResult]] = [] # mutable-ok: retain completed results for partial billing
cancellation_requested: Final = asyncio.Event()
async def run(indices: range) -> None:
for index in indices:
if time.monotonic() >= deadline:
return
try:
outcomes.append(
await _asend(
prepare(index),
deadline,
cancellation_requested,
record_dispatch,
)
)
except Exception as error: # noqa: BLE001 # preparation failures retain completed billable calls
outcomes.append(DecisionRerankResult(index, request.document_text(index), error=error))
tasks: Final = tuple(asyncio.create_task(run(range(worker, count, workers))) for worker in range(workers))
try:
_, pending = await asyncio.wait(tasks, timeout=max(0.0, deadline - time.monotonic()))
outcomes: Final = tuple(
DecisionRerankResult(
document.index,
document.text,
error=litellm.APIConnectionError(
message="Decision rerank provider request was cancelled",
model=model,
llm_provider=provider,
),
)
if task.cancelled()
else task.result()
for document, task in zip(documents, tasks, strict=True)
if task.done()
)
return _finish(outcomes, request, logging_obj, bool(pending), model, provider)
await asyncio.wait(tasks, timeout=max(0.0, deadline - time.monotonic()))
return _finish(outcomes, request, logging_obj, len(outcomes) < count, model, provider)
except asyncio.CancelledError:
outcomes_on_cancel: Final = tuple(task.result() for task in tasks if task.done() and not task.cancelled())
response_on_cancel, _ = LiteLLMDecisionsRerankConfig.transform_decisions_response(
outcomes_on_cancel, request, model, provider
outcomes, request, model, provider
)
_record_partial(response_on_cancel, logging_obj)
raise
finally:
cancellation_requested.set()
for task in tasks:
if not task.done():
task.cancel()
@ -217,10 +310,12 @@ def rerank_with_decisions(
api_base: str | None,
logging_obj: Logging,
is_async: bool,
strategy: Literal["pointwise", "multiple_choice"],
params: Mapping[str, object],
kwargs: Mapping[str, object],
) -> DecisionRerankResponse | Awaitable[DecisionRerankResponse]:
request: Final = LiteLLMDecisionsRerankConfig.validate_request(params, model, provider)
request: Final = LiteLLMDecisionsRerankConfig.validate_request(params, model, provider, strategy)
deadline: Final = time.monotonic() + request.timeout
logging_obj.update_from_kwargs( # pyright: ignore[reportUnknownMemberType] # Logging accepts legacy provider kwargs
kwargs=dict(kwargs),
model=model,
@ -239,25 +334,36 @@ def rerank_with_decisions(
api_base=api_base,
kwargs=kwargs,
)
documents: Final = (
(prepare(0),)
if request.decision_rerank.strategy == "multiple_choice"
else tuple(prepare(index) for index in range(len(request.documents)))
)
first: Final = prepare(0)
def prepare_next(index: int) -> _DocumentCall:
return first if index == 0 else prepare(index)
logging_obj.pre_call( # pyright: ignore[reportUnknownMemberType] # Logging accepts the logical rerank request
input=request.query,
api_key=documents[0].call.api_key,
api_key=first.call.api_key,
additional_args={
"complete_input_dict": {
"model": model,
**request.model_dump(
mode="json", exclude={"extra_headers", "safety_identifier", "timeout"}, exclude_none=True
mode="json",
exclude={"extra_headers", "safety_identifier", "timeout"},
exclude_none=True,
),
},
"api_base": documents[0].call.api_base,
"headers": dict(documents[0].call.headers),
"api_base": first.call.api_base,
"headers": dict(first.call.headers),
},
)
record_dispatch: Final = _dispatch_cost_recorder(kwargs, request)
if is_async:
return _arerank(documents, request, logging_obj, model, provider)
return _rerank(documents, request, logging_obj, model, provider)
return _arerank(
prepare_next,
request,
logging_obj,
model,
provider,
deadline,
record_dispatch,
)
return _rerank(prepare_next, request, logging_obj, model, provider, deadline, record_dispatch)

71
litellm/rerank_api/litellm_decisions_transformation/transformation.py
View file

@ -19,13 +19,16 @@ from litellm.types.decisions import (
OpenAIPredicateQuestion,
OpenAIRefusalAnswer,
)
from litellm.types.rerank import DecisionRerankConfig, DecisionRerankResponse, RerankResponseResult
from litellm.types.rerank import DecisionRerankResponse, RerankResponseResult
_RELEVANCE: Final = (
"Is the passage relevant to the query, meaning it contains information that answers, supports or contradicts it?"
)
_UNSUPPORTED: Final = frozenset(
{
"rerank_mode",
"decision_rerank",
"strategy",
"rank_fields",
"max_chunks_per_doc",
"max_tokens_per_doc",
@ -44,8 +47,7 @@ class _TextDocument(BaseModel):
class DecisionRerankRequest(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
rerank_mode: Literal["decisions"]
decision_rerank: DecisionRerankConfig = DecisionRerankConfig()
strategy: Literal["pointwise", "multiple_choice"] = Field(exclude=True)
query: str
documents: Annotated[tuple[str | _TextDocument, ...], Field(min_length=1)]
top_n: Annotated[StrictInt | None, Field(ge=1)] = None
@ -56,11 +58,8 @@ class DecisionRerankRequest(BaseModel):
@model_validator(mode="after")
def validate_choice_request(self) -> "DecisionRerankRequest":
if self.decision_rerank.strategy == "multiple_choice":
if not 2 <= len(self.documents) <= 255:
raise ValueError("multiple_choice requires between 2 and 255 documents")
if self.decision_rerank.refusal_policy != "error":
raise ValueError("multiple_choice requires refusal_policy='error'")
if self.strategy == "multiple_choice" and not 2 <= len(self.documents) <= 255:
raise ValueError("multiple_choice requires between 2 and 255 documents")
return self
def document_text(self, index: int) -> str:
@ -84,9 +83,7 @@ def _invalid_response(index: int, reason: str, model: str, provider: str) -> lit
)
def _score(
outcome: DecisionRerankResult, refusal_policy: Literal["error", "zero"], model: str, provider: str
) -> float | Exception:
def _score(outcome: DecisionRerankResult, model: str, provider: str) -> float | Exception:
if outcome.error is not None:
return outcome.error
response: Final = outcome.response
@ -94,8 +91,6 @@ def _score(
return _invalid_response(outcome.index, "expected one relevance answer", model, provider)
answer: Final = response.answers[0]
if isinstance(answer, OpenAIRefusalAnswer):
if refusal_policy == "zero":
return 0.0
return litellm.ContentPolicyViolationError(
message=f"Decision model refused to score document {outcome.index}",
model=model,
@ -104,7 +99,12 @@ def _score(
if not isinstance(answer, OpenAIPredicateAnswer) or answer.name != "relevance":
return _invalid_response(outcome.index, "expected the relevance predicate", model, provider)
if not math.isfinite(answer.probability) or not 0 <= answer.probability <= 1:
return _invalid_response(outcome.index, "probability must be finite and between zero and one", model, provider)
return _invalid_response(
outcome.index,
"probability must be finite and between zero and one",
model,
provider,
)
return answer.probability
@ -164,20 +164,30 @@ def _choice_scores(
def _scores(
outcomes: Sequence[DecisionRerankResult], request: DecisionRerankRequest, model: str, provider: str
outcomes: Sequence[DecisionRerankResult],
request: DecisionRerankRequest,
model: str,
provider: str,
) -> tuple[float | Exception, ...]:
if request.decision_rerank.strategy == "multiple_choice":
if request.strategy == "multiple_choice":
scores: Final = _choice_scores(outcomes[0], request, model, provider) if outcomes else ()
return (scores,) if isinstance(scores, Exception) else scores
return tuple(_score(outcome, request.decision_rerank.refusal_policy, model, provider) for outcome in outcomes)
return tuple(_score(outcome, model, provider) for outcome in outcomes)
def _ranked_results(
outcomes: Sequence[DecisionRerankResult], request: DecisionRerankRequest, scores: Sequence[float | Exception]
outcomes: Sequence[DecisionRerankResult],
request: DecisionRerankRequest,
scores: Sequence[float | Exception],
) -> tuple[RerankResponseResult, ...]:
if request.decision_rerank.strategy == "multiple_choice":
if request.strategy == "multiple_choice":
return tuple(
_result(index, request.document_text(index), score, bool(request.return_documents))
_result(
index,
request.document_text(index),
score,
bool(request.return_documents),
)
for index, score in enumerate(scores)
if isinstance(score, float)
)
@ -211,7 +221,12 @@ def _response(
class LiteLLMDecisionsRerankConfig:
@staticmethod
def validate_request(params: Mapping[str, object], model: str, provider: str) -> DecisionRerankRequest:
def validate_request(
params: Mapping[str, object],
model: str,
provider: str,
strategy: Literal["pointwise", "multiple_choice"],
) -> DecisionRerankRequest:
unsupported: Final = tuple(key for key in _UNSUPPORTED if params.get(key) is not None)
if unsupported:
raise litellm.UnsupportedParamsError(
@ -220,7 +235,7 @@ class LiteLLMDecisionsRerankConfig:
llm_provider=provider,
)
try:
return DecisionRerankRequest.model_validate(params)
return DecisionRerankRequest.model_validate({**params, "strategy": strategy})
except ValidationError as error:
raise litellm.BadRequestError(
message="Invalid decision rerank request: "
@ -231,7 +246,7 @@ class LiteLLMDecisionsRerankConfig:
@staticmethod
def transform_rerank_request(request: DecisionRerankRequest, index: int) -> OpenAIDecisionRequestBody:
if request.decision_rerank.strategy == "multiple_choice":
if request.strategy == "multiple_choice":
return OpenAIDecisionRequestBody(
input=json.dumps(
{
@ -252,13 +267,19 @@ class LiteLLMDecisionsRerankConfig:
),
)
return OpenAIDecisionRequestBody(
input=json.dumps({"query": request.query, "passage": request.document_text(index)}, ensure_ascii=False),
input=json.dumps(
{"query": request.query, "passage": request.document_text(index)},
ensure_ascii=False,
),
questions=(OpenAIPredicateQuestion(type="predicate", name="relevance", instructions=_RELEVANCE),),
)
@staticmethod
def transform_decisions_response(
outcomes: Sequence[DecisionRerankResult], request: DecisionRerankRequest, model: str, provider: str
outcomes: Sequence[DecisionRerankResult],
request: DecisionRerankRequest,
model: str,
provider: str,
) -> tuple[DecisionRerankResponse, Exception | None]:
"""Retain completed responses for billing even when another result fails validation."""
scores: Final = _scores(outcomes, request, model, provider)

21
litellm/rerank_api/main.py
View file

@ -15,7 +15,10 @@ from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.llms.together_ai.rerank.handler import TogetherAIRerank
from litellm.llms.watsonx.common_utils import IBMWatsonXMixin
from litellm.rerank_api.rerank_utils import get_optional_rerank_params
from litellm.rerank_api.rerank_utils import (
get_decision_rerank_strategy,
get_optional_rerank_params,
)
from litellm.secret_managers.main import get_secret, get_secret_str
from litellm.types.rerank import RerankResponse
from litellm.types.router import *
@ -158,9 +161,12 @@ def rerank(
api_key=optional_params.api_key,
)
decision_kwargs: Final = _DECISION_KWARGS.validate_python(kwargs)
if decision_kwargs.get("rerank_mode") is not None or decision_kwargs.get("decision_rerank") is not None:
from litellm.rerank_api.litellm_decisions_transformation.handler import rerank_with_decisions
decision_strategy: Final = get_decision_rerank_strategy(model, _custom_llm_provider)
if decision_strategy is not None:
decision_kwargs: Final = _DECISION_KWARGS.validate_python(kwargs)
from litellm.rerank_api.litellm_decisions_transformation.handler import (
rerank_with_decisions,
)
return rerank_with_decisions(
model=model,
@ -169,6 +175,7 @@ def rerank(
api_base=dynamic_api_base or optional_params.api_base,
logging_obj=litellm_logging_obj,
is_async=_is_async,
strategy=decision_strategy,
params=_DECISION_KWARGS.validate_python(
{
**decision_kwargs,
@ -179,9 +186,9 @@ def rerank(
"rank_fields": rank_fields,
"max_chunks_per_doc": max_chunks_per_doc,
"max_tokens_per_doc": max_tokens_per_doc,
"timeout": litellm.request_timeout
if optional_params.timeout is None
else optional_params.timeout,
"timeout": (
litellm.request_timeout if optional_params.timeout is None else optional_params.timeout
),
"extra_headers": decision_kwargs.get("extra_headers") or headers,
}
),

26
litellm/rerank_api/rerank_utils.py
View file

@ -1,8 +1,32 @@
from typing import Any, Final
from typing import Any, Final, Literal
import litellm
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
def get_decision_rerank_strategy(
model: str, custom_llm_provider: str | None = None
) -> Literal["pointwise", "multiple_choice"] | None:
"""Read the catalog method for the model, including a gateway's provider-prefixed model."""
try:
resolved_model, provider, _, _ = litellm.get_llm_provider(model=model, custom_llm_provider=custom_llm_provider)
except litellm.BadRequestError:
return None
candidates: Final = (f"{provider}/{resolved_model}", resolved_model, model)
catalog_model: Final = next(
(
name
for name in candidates
if litellm.model_cost.get(name, {}).get("rerank_strategy") is not None # pyright: ignore[reportUnknownMemberType] # check capability presence without invoking native provider metadata
),
None,
)
if catalog_model is None:
return None
strategy: Final = litellm.get_model_info(model=catalog_model).get("rerank_strategy")
return strategy if strategy in ("pointwise", "multiple_choice") else None
def get_optional_rerank_params(
rerank_provider_config: BaseRerankConfig,
model: str,

2
litellm/types/litellm_params.py
View file

@ -162,8 +162,6 @@ class DispatchOptions:
use_chat_completions_api: bool | None = None
use_in_pass_through: bool | None = None
allowed_openai_params: Sequence[str] | None = None
rerank_mode: Literal["decisions"] | None = None
decision_rerank: Mapping[str, object] | None = None
@dataclass(frozen=True, slots=True, kw_only=True)

12
litellm/types/rerank.py
View file

@ -5,23 +5,15 @@ https://docs.cohere.com/reference/rerank
"""
from collections.abc import Sequence
from typing import Annotated, Literal
from typing import Literal
from pydantic import ConfigDict, Field, PrivateAttr, StrictInt
from pydantic import ConfigDict, Field, PrivateAttr
from typing_extensions import ReadOnly, Required, TypedDict
from litellm.types.decisions import OpenAIDecisionResponse
from litellm.types.llms.base import LiteLLMBaseModel
class DecisionRerankConfig(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
strategy: Literal["pointwise", "multiple_choice"] = "pointwise"
max_concurrency: Annotated[StrictInt, Field(ge=1, le=32)] = 8
refusal_policy: Literal["error", "zero"] = "error"
class RerankRequest(LiteLLMBaseModel):
model: str
query: str

1
litellm/types/utils.py
View file

@ -203,6 +203,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
supports_audio_input: bool | None
supports_video_input: ReadOnly[bool | None]
supports_box_questions: ReadOnly[bool | None]
rerank_strategy: ReadOnly[Literal["pointwise", "multiple_choice"] | None]
supports_embedding_image_input: bool | None
supports_audio_output: bool | None
supports_pdf_input: bool | None

1
litellm/utils.py
View file

@ -6924,6 +6924,7 @@ def get_model_info_helper(
supports_audio_input=_model_info.get("supports_audio_input", None),
supports_video_input=_model_info.get("supports_video_input", None),
supports_box_questions=_model_info.get("supports_box_questions", None),
rerank_strategy=_model_info.get("rerank_strategy", None),
supports_audio_output=_model_info.get("supports_audio_output", None),
supports_pdf_input=_model_info.get("supports_pdf_input", None),
supports_embedding_image_input=_model_info.get("supports_embedding_image_input", None),

10
model_prices_and_context_window.json
View file

@ -46044,6 +46044,7 @@
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"rerank_strategy": "multiple_choice",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/api-reference/decisions-post"
},
@ -46052,6 +46053,7 @@
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"rerank_strategy": "multiple_choice",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/getting-started/pricing"
},
@ -46060,6 +46062,7 @@
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"rerank_strategy": "multiple_choice",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
@ -46068,6 +46071,7 @@
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"rerank_strategy": "multiple_choice",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
@ -46076,6 +46080,7 @@
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"rerank_strategy": "multiple_choice",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
@ -46084,6 +46089,7 @@
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"rerank_strategy": "multiple_choice",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
@ -72235,6 +72241,7 @@
"litellm_provider": "typesafe",
"max_input_tokens": 65536,
"mode": "evaluation",
"rerank_strategy": "pointwise",
"output_cost_per_token": 0.0,
"source": "https://docs.typesafe.ai/models"
},
@ -72243,6 +72250,7 @@
"litellm_provider": "typesafe",
"max_input_tokens": 65536,
"mode": "evaluation",
"rerank_strategy": "pointwise",
"output_cost_per_token": 0.0,
"source": "https://docs.typesafe.ai/models",
"supports_box_questions": true
@ -72252,6 +72260,7 @@
"litellm_provider": "typesafe",
"max_input_tokens": 65536,
"mode": "evaluation",
"rerank_strategy": "pointwise",
"output_cost_per_token": 0.0,
"source": "https://docs.typesafe.ai/models"
},
@ -72260,6 +72269,7 @@
"litellm_provider": "typesafe",
"max_input_tokens": 65536,
"mode": "evaluation",
"rerank_strategy": "pointwise",
"output_cost_per_token": 0.0,
"source": "https://docs.typesafe.ai/models",
"supports_box_questions": true

292
tests/unit/proxy/spend_tracking/test_budget_reservation.py
View file

@ -2,7 +2,7 @@ from __future__ import annotations
import math
from datetime import datetime, timedelta, timezone
from typing import Final
from typing import Final, Literal
import pytest
@ -426,3 +426,293 @@ async def test_reservation_holds_cost_against_cached_org_and_project_budgets(spe
)
def test_budget_window_start_reads_json_encoded_windows(window: str, expected: datetime | None) -> None:
assert get_budget_window_start(window) == expected
def _decision_rerank_model(strategy: str) -> str:
return "perplexity/pplx-decider-flash-v1.2" if strategy == "multiple_choice" else "typesafe/jev-1.14.0"
@pytest.fixture
def decision_rerank_catalog(monkeypatch):
for strategy in ("pointwise", "multiple_choice"):
model = _decision_rerank_model(strategy)
monkeypatch.setitem(
litellm.model_cost,
model,
{**litellm.model_cost[model], "rerank_strategy": strategy},
)
def _decision_rerank_router(strategy: str = "pointwise", *, model: str | None = None) -> Router:
return Router(
model_list=[
{
"model_name": "decision-reranker",
"litellm_params": {
"model": model or _decision_rerank_model(strategy),
"api_key": "test-key",
"input_cost_per_token": 0.001,
"output_cost_per_token": 0.01,
},
"model_info": {"max_output_tokens": 7},
}
]
)
def _decision_expected_costs(body: dict, strategy: Literal["pointwise", "multiple_choice"]) -> tuple[float, float]:
from litellm.rerank_api.litellm_decisions_transformation.transformation import (
LiteLLMDecisionsRerankConfig,
)
request: Final = LiteLLMDecisionsRerankConfig.validate_request(
body,
_decision_rerank_model(strategy),
"perplexity" if strategy == "multiple_choice" else "typesafe",
strategy=strategy,
)
count: Final = 1 if strategy == "multiple_choice" else len(body["documents"])
input_cost: Final = sum(
litellm.token_counter(
model=_decision_rerank_model(strategy),
text=LiteLLMDecisionsRerankConfig.transform_rerank_request(request, index).model_dump_json(),
)
* 0.001
for index in range(count)
)
return input_cost, input_cost + count * 7 * 0.01
@pytest.mark.parametrize("strategy", ["pointwise", "multiple_choice"])
def test_decision_rerank_alias_reserves_transformed_inputs_and_each_output_bound(strategy, decision_rerank_catalog):
from litellm.proxy.spend_tracking.budget_reservation import (
estimate_request_input_cost,
)
router: Final = _decision_rerank_router(strategy)
body: Final = {
"model": "decision-reranker",
"query": "repeat this query " * 30,
"documents": ["one", "two", "three"],
"max_tokens": 1,
}
input_cost, maximum = _decision_expected_costs(body, strategy)
# Cached raw-body token counts must not override the transformed call estimate.
cached: Final = {"decision-reranker": 1}
assert estimate_request_max_cost(body, "/rerank", router, cached) == pytest.approx(maximum)
assert estimate_request_input_cost(body, "/rerank", router, cached) == pytest.approx(input_cost)
def test_decision_rerank_request_cannot_override_catalog_strategy(
decision_rerank_catalog,
):
router: Final = _decision_rerank_router("multiple_choice")
body: Final = {
"model": "decision-reranker",
"query": "query",
"documents": ["one", "two", "three"],
"decision_rerank": {"strategy": "pointwise"},
}
with pytest.raises(litellm.UnsupportedParamsError):
estimate_request_max_cost(body, "/rerank", router)
def test_native_rerank_estimate_remains_input_only():
router: Final = _decision_rerank_router(model="cohere/rerank-v3.5")
body: Final = {
"model": "decision-reranker",
"query": "query",
"documents": ["one", "two"],
}
assert estimate_request_max_cost(body, "/rerank", router, {"decision-reranker": 10}) == pytest.approx(0.01)
@pytest.mark.asyncio
async def test_fail_closed_budget_rejects_aggregate_decision_fanout_before_admission(
spend_counter_cache, decision_rerank_catalog
):
router: Final = _decision_rerank_router()
body: Final = {
"model": "decision-reranker",
"query": "repeat this query " * 30,
"documents": ["one", "two", "three"],
}
_, required = _decision_expected_costs(body, "pointwise")
with pytest.raises(litellm.BudgetExceededError):
await reserve_budget_for_request(
request_body=body,
route="/rerank",
llm_router=router,
valid_token=UserAPIKeyAuth(token="decision-budget-test", max_budget=required / 2, spend=0.0),
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=UserApiKeyCache(),
proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()),
fail_closed_budget_enforcement=True,
)
@pytest.mark.asyncio
async def test_decision_rerank_cancellation_settles_only_dispatched_inputs(
spend_counter_cache, respx_mock, monkeypatch, decision_rerank_catalog
):
import asyncio
from litellm.proxy.spend_tracking.budget_reservation import (
release_budget_reservation_on_cancel,
)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
router: Final = _decision_rerank_router()
body: Final = {
"model": "decision-reranker",
"query": "query",
"documents": [f"document {index}" for index in range(12)],
}
reservation: Final = await reserve_budget_for_request(
request_body=body,
route="/rerank",
llm_router=router,
valid_token=UserAPIKeyAuth(token="cancel-decision-budget", max_budget=100, spend=0.0),
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=UserApiKeyCache(),
proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()),
fail_closed_budget_enforcement=True,
)
assert reservation is not None
assert reservation["input_cost"] == 0
started = []
eight_started = asyncio.Event()
async def score(request):
started.append(request)
if len(started) == 8:
eight_started.set()
await asyncio.Event().wait()
respx_mock.post("https://api.typesafe.ai/v1/systemone").mock(side_effect=score)
task = asyncio.create_task(router.arerank(**body, metadata={"user_api_key_budget_reservation": reservation}))
await asyncio.wait_for(eight_started.wait(), 5)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
expected: Final = sum(reservation["decision_rerank_input_costs"]["pointwise"][:8])
assert len(started) == 8
assert reservation["input_cost"] == pytest.approx(expected)
await release_budget_reservation_on_cancel(reservation)
assert spend_counter_cache.in_memory_cache.get_cache(key="spend:key:cancel-decision-budget") == pytest.approx(
expected
)
@pytest.mark.asyncio
async def test_unknown_decision_output_pricing_cannot_fall_back_to_raw_rerank_estimate(
spend_counter_cache, monkeypatch, decision_rerank_catalog
):
from fastapi import HTTPException
from litellm.proxy.spend_tracking import budget_reservation as budgets
monkeypatch.setattr(
budgets,
"_get_model_cost_infos",
lambda **kwargs: [{"input_cost_per_token": 0.001}],
)
body: Final = {
"model": "decision-reranker",
"query": "query",
"documents": ["one", "two"],
}
with pytest.raises(HTTPException) as error:
await reserve_budget_for_request(
request_body=body,
route="/rerank",
llm_router=_decision_rerank_router(),
valid_token=UserAPIKeyAuth(token="unknown-decision-budget", max_budget=100, spend=0.0),
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=UserApiKeyCache(),
proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()),
fail_closed_budget_enforcement=True,
)
assert error.value.status_code == 503
assert spend_counter_cache.in_memory_cache.get_cache(key="spend:key:unknown-decision-budget") is None
@pytest.mark.asyncio
async def test_cancellation_during_decision_estimation_does_not_reserve_counters(
spend_counter_cache, monkeypatch, decision_rerank_catalog
):
import asyncio
import threading
from litellm.proxy.spend_tracking import budget_reservation as budgets
started = threading.Event()
release = threading.Event()
original = budgets._decision_rerank_costs
def estimate(*args, **kwargs):
started.set()
release.wait(5)
return original(*args, **kwargs)
monkeypatch.setattr(budgets, "_decision_rerank_costs", estimate)
task = asyncio.create_task(
reserve_budget_for_request(
request_body={
"model": "decision-reranker",
"query": "query",
"documents": ["one", "two"],
},
route="/rerank",
llm_router=_decision_rerank_router(),
valid_token=UserAPIKeyAuth(token="estimate-cancel-budget", max_budget=100, spend=0.0),
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=UserApiKeyCache(),
proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()),
fail_closed_budget_enforcement=True,
)
)
try:
assert await asyncio.to_thread(started.wait, 5)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert spend_counter_cache.in_memory_cache.get_cache(key="spend:key:estimate-cancel-budget") is None
finally:
release.set()
def test_explicit_provider_deployment_reserves_each_decision_call(
decision_rerank_catalog,
):
router = Router(
model_list=[
{
"model_name": "decision-reranker",
"litellm_params": {
"model": "jev-1.14.0",
"custom_llm_provider": "typesafe",
"api_key": "test-key",
"input_cost_per_token": 0.001,
"output_cost_per_token": 0.01,
},
"model_info": {"max_output_tokens": 7},
}
]
)
body = {
"model": "decision-reranker",
"query": "query",
"documents": ["one", "two", "three"],
}
_, expected = _decision_expected_costs(body, "pointwise")
assert estimate_request_max_cost(body, "/rerank", router, {"decision-reranker": 1}) == pytest.approx(expected)

333
tests/unit/rerank_api/litellm_decisions_transformation/test_handler.py
View file

@ -1,6 +1,7 @@
import asyncio
import json
from typing import Final
from unittest.mock import patch
import httpx
import pytest
@ -14,6 +15,7 @@ from litellm.types.utils import CallTypes, CostBreakdown
from litellm.utils import return_raw_request
MODEL: Final = "typesafe/jevrank-unit-test"
CHOICE_MODEL: Final = "typesafe/choice-unit-test"
URL: Final = "https://api.typesafe.ai/v1/systemone"
RATE: Final = 0.0002
@ -29,6 +31,37 @@ def setup_model(monkeypatch: pytest.MonkeyPatch):
{
"litellm_provider": "typesafe",
"mode": "evaluation",
"rerank_strategy": "pointwise",
"input_cost_per_token": RATE,
"output_cost_per_token": 0,
},
)
for name in (
CHOICE_MODEL,
"perplexity/pplx-unit-test",
"perplexity/pplx-decider-v1-27b",
"openai/decision-unit-test",
"openai/test-decision",
):
monkeypatch.setitem(
litellm.model_cost,
name,
{
"litellm_provider": name.split("/")[0],
"mode": "evaluation",
"rerank_strategy": "multiple_choice",
"input_cost_per_token": RATE,
"output_cost_per_token": 0,
},
)
monkeypatch.setitem(
litellm.model_cost,
"openrouter/typesafe/jev-1.14",
{
"litellm_provider": "openrouter",
"mode": "evaluation",
"rerank_strategy": "pointwise",
"input_cost_per_token": RATE,
"output_cost_per_token": 0,
},
@ -46,15 +79,18 @@ def response(probability=0.8, *, tokens=10):
)
async def rerank(sync, **kwargs):
async def rerank(sync, *, strategy="pointwise", concurrency=8, **kwargs):
params = {
"model": MODEL,
"model": CHOICE_MODEL if strategy == "multiple_choice" else MODEL,
"api_key": "test-key",
"rerank_mode": "decisions",
"query": "warranty",
**kwargs,
}
return litellm.rerank(**params) if sync else await litellm.arerank(**params)
with patch(
"litellm.rerank_api.litellm_decisions_transformation.handler._MAX_CONCURRENCY",
concurrency,
):
return litellm.rerank(**params) if sync else await litellm.arerank(**params)
@pytest.mark.parametrize("sync", [True, False])
@ -98,10 +134,6 @@ async def test_single_text_document_and_top_n_above_length(sync, respx_mock):
{"documents": []},
{"documents": [{"text": "one", "other": "field"}]},
{"documents": ["one"], "top_n": 0},
{"documents": ["one"], "decision_rerank": {"max_concurrency": 0}},
{"documents": ["one"], "decision_rerank": {"max_concurrency": True}},
{"documents": ["one"], "decision_rerank": {"strategy": "multiple_choice"}},
{"documents": ["one"], "decision_rerank": {"refusal_policy": "ignore"}},
{"documents": ["one"], "timeout": 0},
],
)
@ -129,8 +161,7 @@ async def test_invalid_upstream_probability_is_an_error(sync, probability, respx
await rerank(sync, documents=["one"])
@pytest.mark.parametrize("policy", ["error", "zero"])
async def test_explicit_refusal_policy_preserves_other_documents(policy, respx_mock):
async def test_refusal_fails_the_whole_request(respx_mock):
def score(request):
state = json.loads(json.loads(request.content)["state"])
if state["passage"] == "refused":
@ -145,25 +176,8 @@ async def test_explicit_refusal_policy_preserves_other_documents(policy, respx_m
return response()
respx_mock.post(URL).mock(side_effect=score)
if policy == "error":
with pytest.raises(litellm.ContentPolicyViolationError):
await rerank(
False,
documents=["refused", "allowed"],
decision_rerank={"refusal_policy": policy},
)
else:
result = await rerank(
False,
documents=["refused", "allowed"],
decision_rerank={"refusal_policy": policy},
)
assert result.results == [
{"index": 1, "relevance_score": 0.8},
{"index": 0, "relevance_score": 0.0},
]
assert result.meta["refused_document_indices"] == [0]
assert litellm.completion_cost(completion_response=result) == pytest.approx(20 * RATE)
with pytest.raises(litellm.ContentPolicyViolationError):
await rerank(False, documents=["refused", "allowed"])
@pytest.mark.parametrize(
@ -177,7 +191,7 @@ async def test_explicit_refusal_policy_preserves_other_documents(policy, respx_m
async def test_outages_are_never_zero_scores(status, exception, respx_mock):
respx_mock.post(URL).mock(return_value=httpx.Response(status, json={"error": "provider failure"}))
with pytest.raises(exception):
await rerank(False, documents=["one"], decision_rerank={"refusal_policy": "zero"})
await rerank(False, documents=["one"])
async def test_async_concurrency_and_out_of_order_completion(respx_mock):
@ -203,7 +217,7 @@ async def test_async_concurrency_and_out_of_order_completion(respx_mock):
rerank(
False,
documents=["first", "second", "third"],
decision_rerank={"max_concurrency": 2},
concurrency=2,
)
)
await asyncio.wait_for(first_started.wait(), 10)
@ -231,7 +245,7 @@ async def test_cancellation_cancels_inflight_requests_and_does_not_start_queued_
cancelled.set()
route = respx_mock.post(URL).mock(side_effect=score)
task = asyncio.create_task(rerank(False, documents=["one", "two"], decision_rerank={"max_concurrency": 1}))
task = asyncio.create_task(rerank(False, documents=["one", "two"], concurrency=1))
await asyncio.wait_for(started.wait(), 10)
task.cancel()
with pytest.raises(asyncio.CancelledError):
@ -272,7 +286,7 @@ async def test_success_is_billed_once_and_failed_request_retains_completed_usage
assert len(logger.events) == 1
assert logger.events[0]["response_cost"] == pytest.approx(20 * RATE)
with pytest.raises(litellm.ServiceUnavailableError):
await rerank(False, documents=["one", "failed"], decision_rerank={"max_concurrency": 1})
await rerank(False, documents=["one", "failed"], concurrency=1)
await flush_logging()
assert len(logger.events) == 2
assert logger.events[1]["response_cost"] == pytest.approx(10 * RATE)
@ -287,7 +301,6 @@ async def test_custom_token_prices_apply_to_decision_usage(respx_mock):
"litellm_params": {
"model": MODEL,
"api_key": "test-key",
"rerank_mode": "decisions",
},
"model_info": {
"mode": "rerank",
@ -301,16 +314,7 @@ async def test_custom_token_prices_apply_to_decision_usage(respx_mock):
assert result.hidden_params["response_cost"] == pytest.approx(20 * 0.004)
def test_cache_separates_native_and_decision_rerank_and_refusal_policy():
cache = Cache(type="local")
common = {"model": MODEL, "query": "query", "documents": ["one"]}
native = cache.get_cache_key(**common)
strict = cache.get_cache_key(**common, rerank_mode="decisions", decision_rerank={"refusal_policy": "error"})
zero = cache.get_cache_key(**common, rerank_mode="decisions", decision_rerank={"refusal_policy": "zero"})
assert len({native, strict, zero}) == 3
async def test_router_model_alias_keeps_rerank_configuration(respx_mock):
async def test_router_model_alias_uses_catalog_without_mode_override(respx_mock):
route = respx_mock.post(URL).mock(return_value=response())
router = litellm.Router(
model_list=[
@ -319,10 +323,7 @@ async def test_router_model_alias_keeps_rerank_configuration(respx_mock):
"litellm_params": {
"model": MODEL,
"api_key": "test-key",
"rerank_mode": "decisions",
"decision_rerank": {"strategy": "pointwise", "max_concurrency": 1},
},
"model_info": {"mode": "rerank"},
}
]
)
@ -364,7 +365,7 @@ async def test_multiple_choice_ranks_one_whole_set_and_bills_once(sync, respx_mo
result = await rerank(
sync,
documents=["first", "second", "third"],
decision_rerank={"strategy": "multiple_choice"},
strategy="multiple_choice",
top_n=2,
return_documents=True,
)
@ -392,14 +393,14 @@ async def test_multiple_choice_rejects_malformed_distributions(probabilities, re
await rerank(
False,
documents=["first", "second"],
decision_rerank={"strategy": "multiple_choice"},
strategy="multiple_choice",
)
@pytest.mark.parametrize(
"model,url",
[
(MODEL, URL),
(CHOICE_MODEL, URL),
("perplexity/pplx-unit-test", "https://api.perplexity.ai/v1/decisions"),
("openai/decision-unit-test", "https://api.openai.com/v1/decisions"),
],
@ -445,7 +446,7 @@ async def test_multiple_choice_uses_existing_decisions_normalization(model, url,
False,
model=model,
documents=["one", "two"],
decision_rerank={"strategy": "multiple_choice"},
strategy="multiple_choice",
)
native_scores = tuple(item.probability for item in native.answers[0].probabilities)
assert native_scores == (1, 0)
@ -462,17 +463,17 @@ async def test_multiple_choice_rejects_excess_mass_at_maximum_candidate_count(
await rerank(
False,
documents=[str(i) for i in range(255)],
decision_rerank={"strategy": "multiple_choice"},
strategy="multiple_choice",
)
@pytest.mark.parametrize("count,policy", [(1, "error"), (256, "error"), (2, "zero")])
async def test_multiple_choice_rejects_unsupported_size_or_refusal_policy(count, policy, respx_mock):
@pytest.mark.parametrize("count", [1, 256])
async def test_multiple_choice_rejects_unsupported_size(count, respx_mock):
with pytest.raises(litellm.BadRequestError):
await rerank(
False,
documents=[str(i) for i in range(count)],
decision_rerank={"strategy": "multiple_choice", "refusal_policy": policy},
strategy="multiple_choice",
)
assert not respx_mock.calls
@ -483,7 +484,7 @@ async def test_multiple_choice_refusal_is_an_error(respx_mock):
await rerank(
False,
documents=["one", "two"],
decision_rerank={"strategy": "multiple_choice"},
strategy="multiple_choice",
)
@ -498,10 +499,7 @@ async def test_perplexity_multiple_choice_router_preset(respx_mock):
"litellm_params": {
"model": "perplexity/pplx-decider-v1-27b",
"api_key": "test-key",
"rerank_mode": "decisions",
"decision_rerank": {"strategy": "multiple_choice"},
},
"model_info": {"mode": "rerank"},
}
]
)
@ -552,7 +550,7 @@ async def test_custom_transport_configuration_and_untrusted_text(strategy, respx
await rerank(
False,
documents=documents,
decision_rerank={"strategy": strategy},
strategy=strategy,
api_base="https://custom.example",
api_key="explicit-key",
extra_headers={"x-custom": "custom-value"},
@ -585,7 +583,7 @@ async def test_deadline_bounds_fanout_and_never_starts_queued_requests(sync, res
sync,
documents=["one", "two", "three"],
timeout=0.05,
decision_rerank={"max_concurrency": 1},
concurrency=1,
)
assert time.monotonic() - start < 0.75
assert len(started) == 1
@ -639,7 +637,7 @@ async def test_openai_choice_invalid_normalized_scores_rejected(probabilities, r
False,
model="openai/test-decision",
documents=["one", "two"],
decision_rerank={"strategy": "multiple_choice"},
strategy="multiple_choice",
)
@ -667,7 +665,7 @@ async def test_openai_choice_duplicate_labels_use_native_last_value(respx_mock):
False,
model="openai/test-decision",
documents=["one", "two"],
decision_rerank={"strategy": "multiple_choice"},
strategy="multiple_choice",
)
assert result.results == [
{"index": 0, "relevance_score": 0.7},
@ -711,7 +709,7 @@ async def test_invalid_options_never_echo_credentials(respx_mock):
False,
api_key="sensitive-test-key",
documents=["one", "two"],
decision_rerank={"strategy": "multiple_choice", "refusal_policy": "zero"},
top_n=0,
)
assert "sensitive-test-key" not in str(caught.value)
assert not respx_mock.calls
@ -743,9 +741,7 @@ def test_raw_request_inspection_never_sends_decision_requests(strategy, respx_mo
raw = return_raw_request(
CallTypes.rerank,
{
"model": MODEL,
"rerank_mode": "decisions",
"decision_rerank": {"strategy": strategy},
"model": CHOICE_MODEL if strategy == "multiple_choice" else MODEL,
"query": "warranty",
"documents": ["one", "two"],
},
@ -768,7 +764,7 @@ async def test_fixed_margin_uses_the_pricing_provider_for_aggregate_and_breakdow
model_list=[
{
"model_name": "jevrank",
"litellm_params": {"model": MODEL, "api_key": "test-key", "rerank_mode": "decisions"},
"litellm_params": {"model": MODEL, "api_key": "test-key"},
"model_info": {
"mode": "rerank",
"pricing_provider": "openai",
@ -847,10 +843,209 @@ async def test_cancelled_child_fails_whole_request_and_bills_completed_usage(res
await rerank(
False,
documents=["completed", "cancelled"],
decision_rerank={"max_concurrency": 1},
concurrency=1,
)
await flush_logging()
assert len(logger.events) == 1
assert logger.events[0]["status"] == "failure"
assert logger.events[0]["prompt_tokens"] == 10
assert logger.events[0]["response_cost"] == pytest.approx(10 * RATE)
@pytest.mark.parametrize("sync", [True, False])
async def test_request_preparation_is_bounded_and_stops_at_deadline(sync, respx_mock, monkeypatch):
import threading
from litellm.rerank_api.litellm_decisions_transformation import handler
prepared = []
started = []
release = threading.Event()
original = handler._prepare_document
def prepare(index, *args, **kwargs):
prepared.append(index)
return original(index, *args, **kwargs)
def sync_score(request):
started.append(request)
release.wait(5)
return response()
async def async_score(request):
started.append(request)
await asyncio.Event().wait()
monkeypatch.setattr(handler, "_prepare_document", prepare)
respx_mock.post(URL).mock(side_effect=sync_score if sync else async_score)
try:
with pytest.raises(litellm.Timeout):
await rerank(
sync,
query="long query " * 1000,
documents=[""] * 100,
timeout=0.05,
concurrency=2,
)
assert len(prepared) == 2
assert len(started) == 2
finally:
release.set()
await asyncio.sleep(0.02)
assert len(prepared) == 2
@pytest.mark.parametrize("sync", [True, False])
async def test_preparation_time_counts_toward_request_deadline(sync, respx_mock, monkeypatch):
import time
from litellm.rerank_api.litellm_decisions_transformation import handler
original = handler._prepare_document
prepared = []
def prepare(index, *args, **kwargs):
prepared.append(index)
time.sleep(0.03)
return original(index, *args, **kwargs)
monkeypatch.setattr(handler, "_prepare_document", prepare)
route = respx_mock.post(URL).mock(return_value=response())
with pytest.raises(litellm.Timeout):
await rerank(sync, documents=["one", "two"], timeout=0.01)
assert prepared == [0]
assert route.call_count == 0
@pytest.mark.parametrize("sync", [True, False])
async def test_late_preparation_failure_preserves_completed_usage(sync, respx_mock, monkeypatch):
import threading
from litellm.rerank_api.litellm_decisions_transformation import handler
logger = RecordingLogger()
logged = threading.Event()
def record(kwargs, response_obj, start_time, end_time):
logger.events.append(kwargs["standard_logging_object"])
logged.set()
if sync:
monkeypatch.setattr(logger, "log_failure_event", record)
monkeypatch.setattr(litellm, "callbacks", [logger])
original = handler._prepare_document
def prepare(index, *args, **kwargs):
if index == 1:
raise litellm.BadRequestError("preparation failed", model=MODEL, llm_provider="typesafe")
return original(index, *args, **kwargs)
monkeypatch.setattr(handler, "_prepare_document", prepare)
route = respx_mock.post(URL).mock(return_value=response())
with pytest.raises(litellm.BadRequestError, match="preparation failed"):
await rerank(sync, documents=["completed", "invalid"], concurrency=1)
if sync:
assert await asyncio.to_thread(logged.wait, 5)
else:
await flush_logging()
assert route.call_count == 1
assert len(logger.events) == 1
assert logger.events[0]["status"] == "failure"
assert logger.events[0]["prompt_tokens"] == 10
assert logger.events[0]["response_cost"] == pytest.approx(10 * RATE)
@pytest.mark.parametrize("mode", ["sync", "async", "router"])
async def test_gateway_prefix_uses_underlying_catalog_model_without_configuration(mode, respx_mock, monkeypatch):
route = respx_mock.post("https://gateway.example/v1/decisions").respond(
json={
"model": "jevrank-unit-test",
"answers": [{"type": "predicate", "name": "relevance", "probability": 0.8}],
"usage": {"input_tokens": 10, "output_tokens": 0, "total_tokens": 10},
}
)
params = {
"model": f"openai/{MODEL}",
"api_base": "https://gateway.example/v1",
"api_key": "test-key",
"query": "warranty",
"documents": ["one", "two"],
"return_documents": False,
}
# Gateway-specific pricing must not hide the underlying catalog's capability.
monkeypatch.setitem(
litellm.model_cost,
params["model"],
{
"litellm_provider": "openai",
"mode": "evaluation",
"input_cost_per_token": RATE,
"output_cost_per_token": 0,
},
)
if mode == "router":
router = litellm.Router(
model_list=[
{
"model_name": "gateway-reranker",
"litellm_params": {name: params[name] for name in ("model", "api_base", "api_key")},
}
]
)
result = await router.arerank(
model="gateway-reranker",
query="warranty",
documents=["one", "two"],
return_documents=False,
)
else:
result = litellm.rerank(**params) if mode == "sync" else await litellm.arerank(**params)
assert result.results == [
{"index": 0, "relevance_score": 0.8},
{"index": 1, "relevance_score": 0.8},
]
assert route.call_count == 2
@pytest.mark.parametrize(
"parameter,value",
[
("rerank_mode", "decisions"),
("decision_rerank", {"strategy": "multiple_choice"}),
("strategy", "multiple_choice"),
],
)
async def test_removed_rerank_configuration_is_rejected(parameter, value, respx_mock):
with pytest.raises(litellm.UnsupportedParamsError):
await litellm.arerank(
model=MODEL,
api_key="test-key",
query="warranty",
documents=["one", "two"],
**{parameter: value},
)
assert not respx_mock.calls
def test_catalog_selects_methods_and_leaves_native_rerank_unchanged(monkeypatch):
from pathlib import Path
from litellm.rerank_api.rerank_utils import get_decision_rerank_strategy
catalog = json.loads((Path(__file__).parents[4] / "model_prices_and_context_window.json").read_text())
selected = {name: info for name, info in catalog.items() if info.get("rerank_strategy") is not None}
assert selected
for name, info in selected.items():
monkeypatch.setitem(litellm.model_cost, name, info)
expected = "pointwise" if info["litellm_provider"] == "typesafe" else "multiple_choice"
assert get_decision_rerank_strategy(name) == expected
assert get_decision_rerank_strategy(f"openai/{name}") == expected
assert get_decision_rerank_strategy("cohere/rerank-v3.5") is None
assert get_decision_rerank_strategy("not-a-catalog-model") is None
def test_native_rerank_does_not_invoke_model_metadata_lookup(monkeypatch):
from litellm.rerank_api.rerank_utils import get_decision_rerank_strategy
with patch.object(litellm, "get_model_info", side_effect=AssertionError("native metadata lookup")):
assert get_decision_rerank_strategy("cohere/rerank-v3.5") is None

2
tests/unit/rerank_api/litellm_decisions_transformation/test_transformation.py
View file

@ -10,7 +10,7 @@ from litellm.types.decisions import OpenAIDecisionResponse, OpenAIDecisionUsage,
def test_invalid_score_keeps_completed_responses_and_usage_for_failure_billing():
request: Final = DecisionRerankRequest(rerank_mode="decisions", query="warranty", documents=("relevant", "invalid"))
request: Final = DecisionRerankRequest(strategy="pointwise", query="warranty", documents=("relevant", "invalid"))
responses: Final = tuple(
OpenAIDecisionResponse(
model="decision-unit-test",