feat(file-based): declare Unstructured Document Parser content as semantic-search indexable

airbyte_cdk/sources/file_based/file_types/unstructured_parser.py

@@ -43,6 +43,28 @@
 unstructured_partition_docx = None
 unstructured_partition_pptx = None
 
+SEMANTIC_SEARCH_ANNOTATION: Dict[str, Any] = {
+    "content_type": "markdown",
+    "samples": [
+        {
+            "name": "chunk",
+            "windowed": True,
+            "sampling": {"sample_type": "whole", "unit_label": "chunk"},
+        }
+    ],
+    "windowing": {"context_max_chars": 2048, "context_boundary": "whole_unit"},
+    "embedding": {"model": "text-embedding-3-small"},
+    "metadata": [
+        {"name": "document_key", "path": "/document_key", "type": "string"},
+        {"name": "_ab_source_file_url", "path": "/_ab_source_file_url", "type": "string"},
+        {
+            "name": "_ab_source_file_last_modified",
+            "path": "/_ab_source_file_last_modified",
+            "type": "string",
+        },
+    ],
+}
+
 AIRBYTE_NLTK_DATA_DIR = "/airbyte/nltk_data"
 TMP_NLTK_DATA_DIR = "/tmp/nltk_data"
 
@@ -151,6 +173,7 @@ async def infer_schema(
                 "content": {
                     "type": "string",
                     "description": "Content of the file as markdown. Might be null if the file could not be parsed",
+                    "x-airbyte-semantic-search": SEMANTIC_SEARCH_ANNOTATION,  # type: ignore[dict-item]
                 },
                 "document_key": {
                     "type": "string",

airbyte_cdk/sources/file_based/stream/default_file_based_stream.py

@@ -341,13 +341,19 @@ def infer_schema(self, files: List[RemoteFile]) -> Mapping[str, Any]:
         loop = asyncio.get_event_loop()
         schema = loop.run_until_complete(self._infer_schema(files))
         # as infer schema returns a Mapping that is assumed to be immutable, we need to create a deepcopy to avoid modifying the reference
-        return self._fill_nulls(deepcopy(schema))
+        return self._fill_nulls(deepcopy(schema), is_properties_map=True)
 
     @staticmethod
-    def _fill_nulls(schema: Mapping[str, Any]) -> Mapping[str, Any]:
+    def _fill_nulls(
+        schema: Mapping[str, Any], is_properties_map: bool = False
+    ) -> Mapping[str, Any]:
         if isinstance(schema, dict):
             for k, v in schema.items():
-                if k == "type":
+                if is_properties_map:
+                    DefaultFileBasedStream._fill_nulls(v)
+                elif isinstance(k, str) and k.startswith("x-"):
+                    continue
+                elif k == "type":
                     if isinstance(v, list):
                         if "null" not in v:
                             schema[k] = ["null"] + v
@@ -356,6 +362,8 @@ def _fill_nulls(schema: Mapping[str, Any]) -> Mapping[str, Any]:
                             schema[k] = ["null", v]
                         else:
                             DefaultFileBasedStream._fill_nulls(v)
+                elif k == "properties":
+                    DefaultFileBasedStream._fill_nulls(v, is_properties_map=True)
                 else:
                     DefaultFileBasedStream._fill_nulls(v)
         elif isinstance(schema, list):

unit_tests/sources/file_based/file_types/test_unstructured_parser.py

@@ -21,6 +21,7 @@
 )
 from airbyte_cdk.sources.file_based.exceptions import RecordParseError
 from airbyte_cdk.sources.file_based.file_types import UnstructuredParser
+from airbyte_cdk.sources.file_based.file_types.unstructured_parser import SEMANTIC_SEARCH_ANNOTATION
 from airbyte_cdk.sources.file_based.remote_file import RemoteFile
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
@@ -96,6 +97,7 @@ def test_infer_schema(mock_detect_filetype, filetype, format_config, raises):
             "content": {
                 "type": "string",
                 "description": "Content of the file as markdown. Might be null if the file could not be parsed",
+                "x-airbyte-semantic-search": SEMANTIC_SEARCH_ANNOTATION,
             },
             "document_key": {
                 "type": "string",

unit_tests/sources/file_based/scenarios/unstructured_scenarios.py

@@ -5,6 +5,7 @@
 
 import nltk
 
+from airbyte_cdk.sources.file_based.file_types.unstructured_parser import SEMANTIC_SEARCH_ANNOTATION
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 from unit_tests.sources.file_based.scenarios.file_based_source_builder import FileBasedSourceBuilder
 from unit_tests.sources.file_based.scenarios.scenario_builder import TestScenarioBuilder
@@ -19,6 +20,7 @@
         "content": {
             "type": ["null", "string"],
             "description": "Content of the file as markdown. Might be null if the file could not be parsed",
+            "x-airbyte-semantic-search": SEMANTIC_SEARCH_ANNOTATION,
         },
         "document_key": {
             "type": ["null", "string"],

unit_tests/sources/file_based/stream/test_default_file_based_stream.py

@@ -65,8 +65,8 @@ class MockFormat:
             id="simple-schema-already-has-null",
         ),
         pytest.param(
-            {"properties": {"type": "string"}},
-            {"properties": {"type": ["null", "string"]}},
+            {"properties": {"prop": {"type": "string"}}},
+            {"properties": {"prop": {"type": ["null", "string"]}}},
             id="nested-schema",
         ),
         pytest.param(
@@ -88,6 +88,29 @@ def test_fill_nulls(input_schema: Mapping[str, Any], expected_output: Mapping[st
     assert DefaultFileBasedStream._fill_nulls(input_schema) == expected_output
 
 
+def test_fill_nulls_properties_map_only_exempts_x_keys_on_schema_objects() -> None:
+    properties = {
+        "x-status": {"type": "string"},
+        "content": {
+            "type": "string",
+            "x-airbyte-semantic-search": {"metadata": [{"type": "string"}]},
+        },
+        "nested": {"type": "object", "properties": {"x-inner": {"type": "string"}}},
+    }
+
+    assert DefaultFileBasedStream._fill_nulls(properties, is_properties_map=True) == {
+        "x-status": {"type": ["null", "string"]},
+        "content": {
+            "type": ["null", "string"],
+            "x-airbyte-semantic-search": {"metadata": [{"type": "string"}]},
+        },
+        "nested": {
+            "type": ["null", "object"],
+            "properties": {"x-inner": {"type": ["null", "string"]}},
+        },
+    }
+
+
 class DefaultFileBasedStreamTest(unittest.TestCase):
     _NOW = datetime(2022, 10, 22, tzinfo=timezone.utc)
     _A_RECORD = {"a_record": 1}