diff --git a/domain/data_transformation/guarded_column_classifier.py b/domain/data_transformation/guarded_column_classifier.py new file mode 100644 index 000000000..20f4f3ac4 --- /dev/null +++ b/domain/data_transformation/guarded_column_classifier.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from enum import Enum +from typing import Callable, Optional, TypeVar + +from domain.data_transformation.column_classifier import ColumnClassifier + +E = TypeVar("E", bound=Enum) + + +class GuardedColumnClassifier(ColumnClassifier[E]): + """A ``ColumnClassifier`` that resolves the descriptions a deterministic guard + is certain about, and delegates the rest to a fallback classifier. + + The ``guard`` maps a raw description to a category member when it recognises it + deterministically (e.g. a party-ceiling roof marker — #1376), else ``None``. + Guard hits never reach the fallback, so an unreliable classifier (the LLM) + cannot override a description the guard is sure of — and the LLM is not billed + for it. Every description still appears in the result (guarded or fallen-back). + """ + + def __init__( + self, + guard: Callable[[str], Optional[E]], + fallback: ColumnClassifier[E], + ) -> None: + self._guard = guard + self._fallback = fallback + + def classify(self, descriptions: set[str]) -> dict[str, E]: + raise NotImplementedError diff --git a/tests/domain/data_transformation/__init__.py b/tests/domain/data_transformation/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/domain/data_transformation/test_guarded_column_classifier.py b/tests/domain/data_transformation/test_guarded_column_classifier.py new file mode 100644 index 000000000..a6ec38b3a --- /dev/null +++ b/tests/domain/data_transformation/test_guarded_column_classifier.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from enum import Enum +from typing import Optional + +from domain.data_transformation.column_classifier import ColumnClassifier +from domain.data_transformation.guarded_column_classifier import ( + GuardedColumnClassifier, +) + + +class _Category(Enum): + GUARDED = "guarded" + FALLBACK = "fallback" + UNKNOWN = "unknown" + + +class _RecordingFallback(ColumnClassifier[_Category]): + """Stands in for the LLM: records what it was asked and maps everything it + sees to FALLBACK, so the test can see which descriptions reached it.""" + + def __init__(self) -> None: + self.asked: set[str] = set() + + def classify(self, descriptions: set[str]) -> dict[str, _Category]: + self.asked = set(descriptions) + return {d: _Category.FALLBACK for d in descriptions} + + +def _guard(description: str) -> Optional[_Category]: + return _Category.GUARDED if description == "marker" else None + + +def test_guard_hits_win_and_misses_fall_through_to_the_fallback() -> None: + # Arrange + fallback = _RecordingFallback() + classifier = GuardedColumnClassifier(guard=_guard, fallback=fallback) + + # Act + result = classifier.classify({"marker", "other"}) + + # Assert — the guarded description takes the guard's member and never reaches + # the fallback; the unrecognised one is resolved by the fallback. + assert result == {"marker": _Category.GUARDED, "other": _Category.FALLBACK} + assert fallback.asked == {"other"}