refactor(rerank): validate raw decision results without changing OpenAI formatting

litellm/decisions/call.py

@@ -208,7 +208,12 @@ class DecisionsCall:
     timeout: float | httpx.Timeout | None
     litellm_params: Mapping[str, object] = field(repr=False)
     caller_questions: tuple[DecisionsIRQuestion, ...]
-    preserve_choice_probabilities: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class DecisionsResult:
+    ir_response: DecisionsIRResponse
+    response: DecisionsResponse | OpenAIDecisionResponse
 
 
 def supported_decisions_providers() -> tuple[str, ...]:
@@ -728,12 +733,7 @@ def _formatted_response(
         case DecisionsRequestBody():
             return ir_to_systemone_response(caller_response, caller_request)
         case OpenAIDecisionRequestBody():
-            return ir_to_openai_response(
-                caller_response,
-                caller_request,
-                call.requested_model,
-                preserve_choice_probabilities=call.preserve_choice_probabilities,
-            )
+            return ir_to_openai_response(caller_response, caller_request, call.requested_model)
         case _:
             assert_never(call.request)
 
@@ -756,7 +756,7 @@ def _map_upstream_exception(error: Exception, call: DecisionsCall) -> Exception:
             )
 
 
-async def asend(call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse:
+async def asend_result(call: DecisionsCall) -> DecisionsResult:
     supported: Final = _checked_image_count(_checked_image_types(call))
     inlined: Final = _checked_image_sizes(_checked_image_types(await _async_with_remote_images_inlined(supported)))
     body: Final = _request_body(inlined)
@@ -777,10 +777,10 @@ async def asend(call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionRespon
         )
     except Exception as error:
         raise _map_upstream_exception(error, inlined) from error
-    return _format_response(response, inlined)
+    return DecisionsResult(ir_response=response, response=_format_response(response, inlined))
 
 
-def send(call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse:
+def send_result(call: DecisionsCall) -> DecisionsResult:
     supported: Final = _checked_image_count(_checked_image_types(call))
     inlined: Final = _checked_image_sizes(_checked_image_types(_with_remote_images_inlined(supported)))
     body: Final = _request_body(inlined)
@@ -801,7 +801,15 @@ def send(call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse:
         )
     except Exception as error:
         raise _map_upstream_exception(error, inlined) from error
-    return _format_response(response, inlined)
+    return DecisionsResult(ir_response=response, response=_format_response(response, inlined))
+
+
+async def asend(call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse:
+    return (await asend_result(call)).response
+
+
+def send(call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse:
+    return send_result(call).response
 
 
 def reject_other_format(*, model: str, kwargs: Mapping[str, object], field_name: str, other_function: str) -> None:

litellm/llms/openai/decisions/transformation.py

@@ -264,18 +264,12 @@ def openai_response_to_ir(response: OpenAIDecisionResponse, request: DecisionsIR
     )
 
 
-def _openai_choice_answer(
-    question: DecisionsIRChoiceQuestion, answer: DecisionsIRChoiceAnswer, preserve_choice_probabilities: bool = False
-) -> OpenAIChoiceAnswer:
+def _openai_choice_answer(question: DecisionsIRChoiceQuestion, answer: DecisionsIRChoiceAnswer) -> OpenAIChoiceAnswer:
     probabilities: Final = {systemone_choice_key(item.value): item.probability for item in answer.probabilities}
     return OpenAIChoiceAnswer(
         name=question.name,
         choice=answer.choice,
         probabilities=tuple(
-            OpenAIChoiceProbability(value=item.value, probability=item.probability) for item in answer.probabilities
-        )
-        if preserve_choice_probabilities
-        else tuple(
             OpenAIChoiceProbability(
                 value=option.value, probability=probabilities.get(systemone_choice_key(option.value), 0.0)
             )
@@ -300,14 +294,12 @@ def _openai_score_answer(question: DecisionsIRScoreQuestion, answer: DecisionsIR
     )
 
 
-def _openai_answer(
-    question: DecisionsIRQuestion, answer: DecisionsIRAnswer, preserve_choice_probabilities: bool = False
-) -> OpenAIDecisionAnswer:
+def _openai_answer(question: DecisionsIRQuestion, answer: DecisionsIRAnswer) -> OpenAIDecisionAnswer:
     match question, answer:
         case DecisionsIRPredicateQuestion(), DecisionsIRPredicateAnswer():
             return OpenAIPredicateAnswer(name=question.name, probability=answer.probability)
         case DecisionsIRChoiceQuestion(), DecisionsIRChoiceAnswer():
-            return _openai_choice_answer(question, answer, preserve_choice_probabilities)
+            return _openai_choice_answer(question, answer)
         case DecisionsIRScoreQuestion(), DecisionsIRScoreAnswer():
             return _openai_score_answer(question, answer)
         case DecisionsIRBoxQuestion(), DecisionsIRBoxAnswer():
@@ -317,16 +309,12 @@ def _openai_answer(
 
 
 def ir_to_openai_response(
-    response: DecisionsIRResponse,
-    request: DecisionsIRRequest,
-    requested_model: str,
-    *,
-    preserve_choice_probabilities: bool = False,
+    response: DecisionsIRResponse, request: DecisionsIRRequest, requested_model: str
 ) -> OpenAIDecisionResponse:
     return OpenAIDecisionResponse(
         model=response.model or requested_model,
         answers=tuple(
-            _openai_answer(question, answer, preserve_choice_probabilities)
+            _openai_answer(question, answer)
             for question, answer in zip(request.questions, response.answers, strict=True)
         ),
         usage=OpenAIDecisionUsage(

litellm/rerank_api/decision_rerank.py

@@ -21,13 +21,15 @@ import litellm
 from litellm.decisions.call import (
     DecisionsCall,
     OpenAIDecisionRequestFields,
-    asend,
+    asend_result,
     prepare_call,
-    send,
+    send_result,
 )
 from litellm.litellm_core_utils.litellm_logging import Logging
 from litellm.types.decisions import (
-    OpenAIChoiceAnswer,
+    DecisionsIRChoiceAnswer,
+    DecisionsIRRefusal,
+    DecisionsIRResponse,
     OpenAIDecisionResponse,
     OpenAIPredicateAnswer,
     OpenAIRefusalAnswer,
@@ -94,6 +96,7 @@ class _DocumentCall:
 class _Outcome:
     document: _DocumentCall
     response: OpenAIDecisionResponse | None = None
+    ir_response: DecisionsIRResponse | None = None
     error: Exception | None = None
 
 
@@ -117,9 +120,10 @@ def _send(document: _DocumentCall, deadline: float) -> _Outcome:
     if isinstance(call, litellm.Timeout):
         return _Outcome(document, error=call)
     try:
-        response: Final = send(call)
+        result: Final = send_result(call)
+        response: Final = result.response
         if isinstance(response, OpenAIDecisionResponse):
-            return _Outcome(document, response=response)
+            return _Outcome(document, response=response, ir_response=result.ir_response)
         return _Outcome(document, error=_invalid_response(document, "expected a Decisions response"))
     except Exception as error:  # noqa: BLE001  # preserve mapped provider failures alongside completed billable calls
         return _Outcome(document, error=error)
@@ -131,9 +135,10 @@ async def _asend(document: _DocumentCall, deadline: float, semaphore: asyncio.Se
         if isinstance(call, litellm.Timeout):
             return _Outcome(document, error=call)
         try:
-            response: Final = await asend(call)
+            result: Final = await asend_result(call)
+            response: Final = result.response
             if isinstance(response, OpenAIDecisionResponse):
-                return _Outcome(document, response=response)
+                return _Outcome(document, response=response, ir_response=result.ir_response)
             return _Outcome(
                 document,
                 error=_invalid_response(document, "expected a Decisions response"),
@@ -189,17 +194,17 @@ def _is_refusal(response: OpenAIDecisionResponse) -> bool:
 def _choice_scores(outcome: _Outcome, request: _Request) -> tuple[float, ...] | Exception:
     if outcome.error is not None:
         return outcome.error
-    response: Final = outcome.response
+    response: Final = outcome.ir_response
     if response is None or len(response.answers) != 1:
         return _invalid_response(outcome.document, "expected one choice answer")
     answer: Final = response.answers[0]
-    if isinstance(answer, OpenAIRefusalAnswer):
+    if isinstance(answer, DecisionsIRRefusal):
         return litellm.ContentPolicyViolationError(
             message="Decision model refused to rank the candidate set",
             model=outcome.document.call.requested_model,
             llm_provider=outcome.document.call.custom_llm_provider,
         )
-    if not isinstance(answer, OpenAIChoiceAnswer) or answer.name != "relevance":
+    if not isinstance(answer, DecisionsIRChoiceAnswer):
         return _invalid_response(outcome.document, "expected the relevance choice")
     labels: Final = tuple(f"p{index}" for index in range(len(request.documents)))
     probabilities: Final = {item.value: item.probability for item in answer.probabilities}
@@ -418,19 +423,16 @@ def _prepare_choice(document: _DocumentCall, request: _Request) -> _DocumentCall
     call: Final = document.call
     return replace(
         document,
-        call=replace(
-            prepare_call(
-                model=call.model,
-                fields=fields,
-                safety_identifier=request.safety_identifier,
-                api_key=call.api_key,
-                api_base=call.api_base,
-                timeout=request.timeout,
-                custom_llm_provider=call.custom_llm_provider,
-                extra_headers=call.headers,
-                kwargs=call.litellm_params,
-            ),
-            preserve_choice_probabilities=True,
+        call=prepare_call(
+            model=call.model,
+            fields=fields,
+            safety_identifier=request.safety_identifier,
+            api_key=call.api_key,
+            api_base=call.api_base,
+            timeout=request.timeout,
+            custom_llm_provider=call.custom_llm_provider,
+            extra_headers=call.headers,
+            kwargs=call.litellm_params,
         ),
     )