fix(decisions): preserve omitted probabilities (#48449) (#48524)

This commit is contained in:
ishaan-j-berri 2026-10-26 19:49:32 -07:00 • committed by GitHub
parent f2b59be6ea
commit ebe0871a28
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 104 additions and 7 deletions

View file

@ -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 {}),
)

View file

@ -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,
)

View file

@ -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
)

View file

@ -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)