From 9b8e0715ea6ea9f34a0141e452ccdfc7645c02cb Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 20 Oct 2026 22:18:46 -0700 Subject: [PATCH] feat(rerank): select decision methods automatically and bound admission --- litellm/caching/caching.py | 2 +- ...odel_prices_and_context_window_backup.json | 10 + .../spend_tracking/budget_reservation.py | 266 +++++++++++++- .../handler.py | 232 ++++++++---- .../transformation.py | 71 ++-- litellm/rerank_api/main.py | 21 +- litellm/rerank_api/rerank_utils.py | 26 +- litellm/types/litellm_params.py | 2 - litellm/types/rerank.py | 12 +- litellm/types/utils.py | 1 + litellm/utils.py | 1 + model_prices_and_context_window.json | 10 + .../spend_tracking/test_budget_reservation.py | 292 ++++++++++++++- .../test_handler.py | 333 ++++++++++++++---- .../test_transformation.py | 2 +- 15 files changed, 1087 insertions(+), 194 deletions(-) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 90536056067..5f5b8210096 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -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) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3d89af9d691..f8849494a8d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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 diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 803df9f83ca..acf87e678d7 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -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 diff --git a/litellm/rerank_api/litellm_decisions_transformation/handler.py b/litellm/rerank_api/litellm_decisions_transformation/handler.py index 38c3f1cf45c..e8191dda4da 100644 --- a/litellm/rerank_api/litellm_decisions_transformation/handler.py +++ b/litellm/rerank_api/litellm_decisions_transformation/handler.py @@ -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) diff --git a/litellm/rerank_api/litellm_decisions_transformation/transformation.py b/litellm/rerank_api/litellm_decisions_transformation/transformation.py index c98fb8a5a79..b4df011e8e7 100644 --- a/litellm/rerank_api/litellm_decisions_transformation/transformation.py +++ b/litellm/rerank_api/litellm_decisions_transformation/transformation.py @@ -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) diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 4a010d5ab54..96bb2fe104a 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -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, } ), diff --git a/litellm/rerank_api/rerank_utils.py b/litellm/rerank_api/rerank_utils.py index f365ae749d1..b700fe7bf3f 100644 --- a/litellm/rerank_api/rerank_utils.py +++ b/litellm/rerank_api/rerank_utils.py @@ -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, diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index d0235e4b01a..0ed156524be 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -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) diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py index 7ae2993b16b..71691b8db82 100644 --- a/litellm/types/rerank.py +++ b/litellm/types/rerank.py @@ -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 diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5939a350970..b633bcab087 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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 diff --git a/litellm/utils.py b/litellm/utils.py index 1a114ffe9f0..3172cc1ed0c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3d89af9d691..f8849494a8d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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 diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 024fe5229be..0816466fc60 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -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) diff --git a/tests/unit/rerank_api/litellm_decisions_transformation/test_handler.py b/tests/unit/rerank_api/litellm_decisions_transformation/test_handler.py index 54558afe483..2773167eb2e 100644 --- a/tests/unit/rerank_api/litellm_decisions_transformation/test_handler.py +++ b/tests/unit/rerank_api/litellm_decisions_transformation/test_handler.py @@ -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 diff --git a/tests/unit/rerank_api/litellm_decisions_transformation/test_transformation.py b/tests/unit/rerank_api/litellm_decisions_transformation/test_transformation.py index 6f309b52cf3..3d7fba68615 100644 --- a/tests/unit/rerank_api/litellm_decisions_transformation/test_transformation.py +++ b/tests/unit/rerank_api/litellm_decisions_transformation/test_transformation.py @@ -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",