diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index dbcb73be6e8..2538c8ee390 100644 --- a/litellm/decisions/call.py +++ b/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: diff --git a/litellm/llms/openai/decisions/transformation.py b/litellm/llms/openai/decisions/transformation.py index 2e01b800e1d..56f348881de 100644 --- a/litellm/llms/openai/decisions/transformation.py +++ b/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( diff --git a/litellm/rerank_api/decision_rerank.py b/litellm/rerank_api/decision_rerank.py index b23475cfb31..be0c4fd62c9 100644 --- a/litellm/rerank_api/decision_rerank.py +++ b/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, ), )