From 29765950630f8680a6d7be16009d8542b38c8cc8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 7 Oct 2026 19:04:39 -0700 Subject: [PATCH] fix(decisions): keep provider fields inside System One answers Backport of #45214 to stable/1.104.x. Cherry-picked from 1348de864f (litellm_openai_decisions_provider). (cherry picked from commit 1348de864fed181306fc1d67a5af3d3aa37aa8f0) --- .../llms/base_llm/decisions/transformation.py | 36 ++++++++++++------- litellm/types/decisions.py | 3 ++ .../test_base_decisions_transformation.py | 4 ++- 3 files changed, 29 insertions(+), 14 deletions(-) diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index 23a00d60ba0..ab827ae5dcf 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -183,6 +183,7 @@ def _ir_choice_answer(question: DecisionsIRChoiceQuestion, answer: ChoiceAnswer) DecisionsIRChoiceProbability(value=typed_values.get(key, key), probability=probability) for key, probability in answer.probabilities.items() ), + extra=MappingProxyType(answer.model_extra or {}), ) @@ -198,13 +199,14 @@ def _ir_score_answer(question: DecisionsIRScoreQuestion, answer: ScoreAnswer) -> ) for index, level in enumerate(question.levels) ), + extra=MappingProxyType(answer.model_extra or {}), ) def _ir_answer(question: DecisionsIRQuestion, answer: DecisionAnswer | None) -> DecisionsIRAnswer: match question, answer: case DecisionsIRPredicateQuestion(), NoulAnswer(): - return DecisionsIRPredicateAnswer(probability=answer.noul) + return DecisionsIRPredicateAnswer(probability=answer.noul, extra=MappingProxyType(answer.model_extra or {})) case DecisionsIRChoiceQuestion(), ChoiceAnswer(): return _ir_choice_answer(question, answer) case DecisionsIRScoreQuestion(), ScoreAnswer(): @@ -237,21 +239,29 @@ def parse_systemone_response(payload: object, request: DecisionsIRRequest) -> De def _systemone_answer(answer: DecisionsIRAnswer) -> DecisionAnswer | None: match answer: case DecisionsIRPredicateAnswer(): - return NoulAnswer(type="noul", noul=answer.probability) + return NoulAnswer.model_validate({**answer.extra, "type": "noul", "noul": answer.probability}) case DecisionsIRChoiceAnswer(): - return ChoiceAnswer( - type="choice", - choice=systemone_choice_key(answer.choice), - confidence=answer.confidence, - probabilities={systemone_choice_key(item.value): item.probability for item in answer.probabilities}, + return ChoiceAnswer.model_validate( + { + **answer.extra, + "type": "choice", + "choice": systemone_choice_key(answer.choice), + "confidence": answer.confidence, + "probabilities": { + systemone_choice_key(item.value): item.probability for item in answer.probabilities + }, + } ) case DecisionsIRScoreAnswer(): - return ScoreAnswer( - type="score", - score=answer.score, - confidence=answer.confidence, - legend={str(item.value): item.label for item in answer.probabilities}, - probabilities={str(item.value): item.probability for item in answer.probabilities}, + return ScoreAnswer.model_validate( + { + **answer.extra, + "type": "score", + "score": answer.score, + "confidence": answer.confidence, + "legend": {str(item.value): item.label for item in answer.probabilities}, + "probabilities": {str(item.value): item.probability for item in answer.probabilities}, + } ) case DecisionsIRRefusal(): return None diff --git a/litellm/types/decisions.py b/litellm/types/decisions.py index ba3c93a7d41..5fbbd905123 100644 --- a/litellm/types/decisions.py +++ b/litellm/types/decisions.py @@ -378,6 +378,7 @@ class DecisionsIRRequest: @dataclass(frozen=True, slots=True) class DecisionsIRPredicateAnswer: probability: float + extra: Mapping[str, object] = field(default=_NO_EXTRA) @dataclass(frozen=True, slots=True) @@ -391,6 +392,7 @@ class DecisionsIRChoiceAnswer: choice: str | bool confidence: float probabilities: tuple[DecisionsIRChoiceProbability, ...] + extra: Mapping[str, object] = field(default=_NO_EXTRA) @dataclass(frozen=True, slots=True) @@ -405,6 +407,7 @@ class DecisionsIRScoreAnswer: score: float confidence: float probabilities: tuple[DecisionsIRScoreProbability, ...] + extra: Mapping[str, object] = field(default=_NO_EXTRA) @dataclass(frozen=True, slots=True) diff --git a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py index 9c4d8a27260..a452a621d2e 100644 --- a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py +++ b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py @@ -136,12 +136,13 @@ def test_a_systemone_response_reaches_the_caller_unchanged_with_provider_extras( payload: Final = { "model": "jev-latest", "answers": { - "damaged": {"type": "noul", "noul": 0.95}, + "damaged": {"type": "noul", "noul": 0.95, "rationale": "crack visible"}, "action": { "type": "choice", "choice": "refund", "confidence": 0.8, "probabilities": {"refund": 0.9, "escalate": 0.1}, + "calibrated": True, }, "severity": { "type": "score", @@ -149,6 +150,7 @@ def test_a_systemone_response_reaches_the_caller_unchanged_with_provider_extras( "confidence": 0.6, "legend": {"0": "minor", "1": {"label": "major"}}, "probabilities": {"0": 0.6, "1": 0.4}, + "raw_logits": [0.1, 0.2], }, }, "usage": {"input_tokens": 383, "output_tokens": 2, "cost": 0.25},