diff --git a/libs/schemas/reasoning.py b/libs/schemas/reasoning.py index 5008cf3..a658af7 100644 --- a/libs/schemas/reasoning.py +++ b/libs/schemas/reasoning.py @@ -1,7 +1,6 @@ from __future__ import annotations from typing import Literal, Optional -from pydantic import BaseModel, Field -import json +from pydantic import BaseModel, Field, computed_field class ReasoningResult(BaseModel): @@ -16,24 +15,22 @@ class ReasoningResult(BaseModel): severity_score: float = Field(0.0, ge=0.0, le=1.0) alert_id: Optional[str] = None + @computed_field @property def confidence_tier(self) -> Literal["high", "medium", "low"]: - if self.confidence >= 0.75: return "high" - if self.confidence >= 0.50: return "medium" + if self.confidence >= 0.75: + return "high" + if self.confidence >= 0.50: + return "medium" return "low" + @computed_field @property def is_actionable(self) -> bool: return self.label == "Suspicious" and self.confidence >= 0.65 - def model_dump_json(self, **kw) -> str: - d = self.model_dump(**kw) - d["confidence_tier"] = self.confidence_tier - d["is_actionable"] = self.is_actionable - return json.dumps(d, default=str) - class GroundingResult(BaseModel): grounded: bool invented_label: Optional[str] = None - checked_caption: str = "" + checked_caption: str = "" \ No newline at end of file diff --git a/tests/test_reasoning_schema.py b/tests/test_reasoning_schema.py new file mode 100644 index 0000000..8b18763 --- /dev/null +++ b/tests/test_reasoning_schema.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +import pytest +from libs.schemas.reasoning import ReasoningResult + + +@pytest.mark.parametrize( + "confidence, expected_tier", + [ + (0.49, "low"), + (0.50, "medium"), + (0.74, "medium"), + (0.75, "high"), + (1.0, "high"), + ], +) +def test_confidence_tier_boundaries(confidence: float, expected_tier: str) -> None: + result = ReasoningResult(track_id=1, label="Suspicious", confidence=confidence, reason="test") + assert result.confidence_tier == expected_tier + + +@pytest.mark.parametrize( + "label, confidence, expected", + [ + ("Suspicious", 0.65, True), + ("Suspicious", 0.64, False), + ("Normal", 0.99, False), + ("Normal", 0.65, False), + ], +) +def test_is_actionable(label: str, confidence: float, expected: bool) -> None: + result = ReasoningResult(track_id=1, label=label, confidence=confidence, reason="test") + assert result.is_actionable is expected + + +def test_model_dump_json_mode_includes_computed_fields() -> None: + result = ReasoningResult(track_id=1, label="Suspicious", confidence=0.88, reason="test") + payload = result.model_dump(mode="json") + + assert payload["confidence_tier"] == "high" + assert payload["is_actionable"] is True \ No newline at end of file