diff --git a/infrastructure/chatgpt/chatgpt_property_type_classifier.py b/infrastructure/chatgpt/chatgpt_property_type_classifier.py index d4f0c060..75ec1556 100644 --- a/infrastructure/chatgpt/chatgpt_property_type_classifier.py +++ b/infrastructure/chatgpt/chatgpt_property_type_classifier.py @@ -33,6 +33,14 @@ class ChatGptPropertyTypeClassifier(PropertyTypeClassifier): ) raw: dict[str, Any] = json.loads(reply) return { - description: PropertyType(raw[description]) + description: self._to_property_type(raw[description]) for description in descriptions } + + @staticmethod + def _to_property_type(value: Any) -> PropertyType: + """Map a reply value to a PropertyType, defaulting to UNKNOWN.""" + try: + return PropertyType(value) + except ValueError: + return PropertyType.UNKNOWN diff --git a/tests/infrastructure/chatgpt/test_chatgpt_property_type_classifier.py b/tests/infrastructure/chatgpt/test_chatgpt_property_type_classifier.py index 8c697eb2..d4801154 100644 --- a/tests/infrastructure/chatgpt/test_chatgpt_property_type_classifier.py +++ b/tests/infrastructure/chatgpt/test_chatgpt_property_type_classifier.py @@ -31,3 +31,15 @@ def test_classifies_description_into_property_type() -> None: # Assert assert result == {"semi-detached": PropertyType.HOUSE} + + +def test_unrecognised_category_maps_to_unknown() -> None: + # Arrange + chat_gpt = _FakeChatGPT(reply='{"garden shed": "Shed"}') + classifier = ChatGptPropertyTypeClassifier(chat_gpt) + + # Act + result = classifier.classify({"garden shed"}) + + # Assert + assert result == {"garden shed": PropertyType.UNKNOWN}