From ebe0871a28bbd585b73988c78edebf675985b9de Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Mon, 26 Oct 2026 19:49:32 -0700 Subject: [PATCH] fix(decisions): preserve omitted probabilities (#48449) (#48524) --- .../llms/base_llm/decisions/transformation.py | 3 +- .../llms/openai/decisions/transformation.py | 10 ++-- .../test_base_decisions_transformation.py | 46 ++++++++++++++++ .../test_openai_decisions_transformation.py | 52 +++++++++++++++++++ 4 files changed, 104 insertions(+), 7 deletions(-) diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index 14c6a2d2a66..b3df4654a6b 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -351,9 +351,10 @@ def _ir_score_answer(question: DecisionsIRScoreQuestion, answer: ScoreAnswer) -> DecisionsIRScoreProbability( value=index, label=answer.legend.get(str(index), level.label), - probability=answer.probabilities.get(str(index), 0.0), + probability=answer.probabilities[str(index)], ) for index, level in enumerate(question.levels) + if str(index) in answer.probabilities ), extra=MappingProxyType(answer.model_extra or {}), ) diff --git a/litellm/llms/openai/decisions/transformation.py b/litellm/llms/openai/decisions/transformation.py index 25d403a0dff..4518fe6a77e 100644 --- a/litellm/llms/openai/decisions/transformation.py +++ b/litellm/llms/openai/decisions/transformation.py @@ -270,10 +270,9 @@ def _openai_choice_answer(question: DecisionsIRChoiceQuestion, answer: Decisions name=question.name, choice=answer.choice, probabilities=tuple( - OpenAIChoiceProbability( - value=option.value, probability=probabilities.get(systemone_choice_key(option.value), 0.0) - ) + OpenAIChoiceProbability(value=option.value, probability=probabilities[systemone_choice_key(option.value)]) for option in question.choices + if systemone_choice_key(option.value) in probabilities ), confidence=answer.confidence, ) @@ -285,10 +284,9 @@ def _openai_score_answer(question: DecisionsIRScoreQuestion, answer: DecisionsIR name=question.name, score=answer.score, probabilities=tuple( - OpenAIScoreProbability( - value=index, label=decisions_text(level.label), probability=probabilities.get(index, 0.0) - ) + OpenAIScoreProbability(value=index, label=decisions_text(level.label), probability=probabilities[index]) for index, level in enumerate(question.levels) + if index in probabilities ), confidence=answer.confidence, ) 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 3aebf56ee73..798cbe5fe00 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 @@ -355,3 +355,49 @@ def test_invalid_provider_response_is_an_api_error(payload: object, provider: st with pytest.raises(BaseLLMException, match="unexpected response") as error: config.transform_decisions_response("drex-v2.0", provider, httpx.Response(200, json=payload), request) assert error.value.status_code == 500 + + +@pytest.mark.parametrize( + ("probabilities", "expected"), + ( + ({}, ()), + ({"1": 0.8}, ((1, "medium", 0.8),)), + ({"0": 0.0, "1": 0.8}, ((0, "low", 0.0), (1, "medium", 0.8))), + ({"2": 0.0}, ((2, "high", 0.0),)), + ), + ids=("empty", "missing", "explicit_zero", "only_zero"), +) +def test_score_probabilities_preserve_missing_entries_and_explicit_zeros( + probabilities: Mapping[str, float], expected: tuple[tuple[int, str, float], ...] +) -> None: + ir: Final = systemone_request_to_ir( + _SYSTEMONE_BODY.validate_python( + { + "state": "Classify the fixture.", + "questions": {"score": {"type": "score", "criteria": ["low", "medium", "high"]}}, + } + ) + ) + payload: Final = { + "model": "fixture", + "answers": { + "score": { + "type": "score", + "score": 1, + "confidence": 0.8, + "legend": {"0": "low", "1": "medium", "2": "high"}, + "probabilities": probabilities, + } + }, + "usage": {"input_tokens": 0, "output_tokens": 0}, + } + + response: Final = parse_systemone_response(payload, ir) + answer: Final = response.answers[0] + + assert isinstance(answer, DecisionsIRScoreAnswer) + assert tuple((item.value, item.label, item.probability) for item in answer.probabilities) == expected + assert ( + ir_to_systemone_response(response, ir).model_dump(mode="json")["answers"]["score"]["probabilities"] + == probabilities + ) diff --git a/tests/unit/llms/openai/decisions/test_openai_decisions_transformation.py b/tests/unit/llms/openai/decisions/test_openai_decisions_transformation.py index 2a07462fe81..eddfe578256 100644 --- a/tests/unit/llms/openai/decisions/test_openai_decisions_transformation.py +++ b/tests/unit/llms/openai/decisions/test_openai_decisions_transformation.py @@ -299,3 +299,55 @@ def test_an_openai_response_becomes_systemone_answers_with_the_callers_score_lab } assert response.usage is not None assert (response.usage.cached_tokens, response.usage.cache_write_tokens) == (256, 64) + + +@pytest.mark.parametrize( + ("choice_probabilities", "score_probabilities"), + ( + ((), ()), + ( + ({"value": True, "probability": 0.8},), + ({"value": 1, "label": "major", "probability": 0.8},), + ), + ( + ( + {"value": True, "probability": 0.0}, + {"value": "escalate", "probability": 0.8}, + ), + ( + {"value": 0, "label": "minor", "probability": 0.0}, + {"value": 1, "label": "major", "probability": 0.8}, + ), + ), + ( + ({"value": True, "probability": 0.0},), + ({"value": 1, "label": "major", "probability": 0.0},), + ), + ), + ids=("empty", "missing", "explicit_zero", "only_zero"), +) +def test_openai_probabilities_preserve_missing_entries_and_explicit_zeros( + choice_probabilities: tuple[Mapping[str, object], ...], + score_probabilities: tuple[Mapping[str, object], ...], +) -> None: + ir: Final = _openai_ir(_OPENAI_REQUEST) + payload: Final = { + **_OPENAI_RESPONSE, + "answers": [ + _PREDICATE_ANSWER, + {**_CHOICE_ANSWER, "probabilities": choice_probabilities}, + { + "type": "score", + "name": "severity", + "score": 1, + "confidence": 0.8, + "probabilities": score_probabilities, + }, + _REFUSAL_ANSWER, + ], + } + + response: Final = ir_to_openai_response(OpenAIDecisionsConfig().parse_response(payload, ir), ir, "requested") + + assert response.answers[1].model_dump(mode="json")["probabilities"] == list(choice_probabilities) + assert response.answers[2].model_dump(mode="json")["probabilities"] == list(score_probabilities)