fix(decisions): keep provider fields inside System One answers

This commit is contained in:
mateo-berri 2026-10-07 19:03:18 -07:00
parent 0dd207fde1
commit 1348de864f
3 changed files with 29 additions and 14 deletions

View file

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

View file

@ -382,6 +382,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)
@ -395,6 +396,7 @@ class DecisionsIRChoiceAnswer:
choice: str | bool
confidence: float
probabilities: tuple[DecisionsIRChoiceProbability, ...]
extra: Mapping[str, object] = field(default=_NO_EXTRA)
@dataclass(frozen=True, slots=True)
@ -409,6 +411,7 @@ class DecisionsIRScoreAnswer:
score: float
confidence: float
probabilities: tuple[DecisionsIRScoreProbability, ...]
extra: Mapping[str, object] = field(default=_NO_EXTRA)
@dataclass(frozen=True, slots=True)

View file

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