fix: union types validation in element enricher (#499)

Author: micpst

packages/ragbits-document-search/CHANGELOG.md

@@ -2,12 +2,13 @@
 
 ## Unreleased
 
+- fix union types validation in element enricher (#499)
+
 ## 0.13.0 (2025-04-02)
 
 ### Changed
 
 - ragbits-core updated to version v0.13.0
-
 - DocumentSearch.ingest now raises IngestExecutionError when any errors are encountered during ingestion.
 
 ## 0.12.0 (2025-03-25)

packages/ragbits-document-search/src/ragbits/document_search/ingestion/enrichers/base.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
-from types import ModuleType
-from typing import ClassVar, Generic, TypeVar
+from types import ModuleType, UnionType
+from typing import ClassVar, Generic, TypeVar, get_args, get_origin
 
 from ragbits.core.utils.config_handling import WithConstructionConfig
 from ragbits.document_search.documents.element import Element
@@ -47,5 +47,18 @@ def validate_element_type(cls, element_type: type[Element]) -> None:
         Raises:
             EnricherElementNotSupportedError: If the element type is not supported.
         """
-        if element_type != cls.__orig_bases__[0].__args__[0]:  # type: ignore
-            raise EnricherElementNotSupportedError(enricher_name=cls.__name__, element_type=element_type)
+        expected_element_type = cls.__orig_bases__[0].__args__[0]  # type: ignore
+
+        # Check if expected_element_type is a Union and if element_type is in that Union
+        if (
+            (origin := get_origin(expected_element_type))
+            and origin == UnionType
+            and element_type in get_args(expected_element_type)
+        ):
+            return
+
+        # Check if element_type matches expected_element_type exactly
+        if element_type == expected_element_type:
+            return
+
+        raise EnricherElementNotSupportedError(enricher_name=cls.__name__, element_type=element_type)

packages/ragbits-document-search/tests/unit/test_element_enrichers.py

@@ -11,16 +11,27 @@
 from ragbits.document_search.ingestion.enrichers.image import ImageDescriberPrompt, ImageElementEnricher
 
 
-def test_enricher_validates_supported_element_types_passes() -> None:
+def test_enricher_validates_supported_element_types() -> None:
     ImageElementEnricher.validate_element_type(ImageElement)
 
+    with pytest.raises(EnricherElementNotSupportedError):
+        ImageElementEnricher.validate_element_type(TextElement)
+
 
-def test_enricher_validates_supported_document_types_fails() -> None:
+def test_enricher_validates_supported_document_union_types() -> None:
     class CustomElement(Element):
+        @property
+        def text_representation(self) -> str:
+            return ""
+
+    class CustomElementEnricher(ElementEnricher[CustomElement | TextElement]):
         pass
 
+    CustomElementEnricher.validate_element_type(CustomElement)
+    CustomElementEnricher.validate_element_type(TextElement)
+
     with pytest.raises(EnricherElementNotSupportedError):
-        ImageElementEnricher.validate_element_type(CustomElement)  # type: ignore
+        CustomElementEnricher.validate_element_type(ImageElement)
 
 
 @pytest.mark.parametrize(