diff --git a/airbyte_cdk/sources/file_based/file_types/unstructured_parser.py b/airbyte_cdk/sources/file_based/file_types/unstructured_parser.py index f55675e0ad..0c7538ef86 100644 --- a/airbyte_cdk/sources/file_based/file_types/unstructured_parser.py +++ b/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", diff --git a/airbyte_cdk/sources/file_based/stream/default_file_based_stream.py b/airbyte_cdk/sources/file_based/stream/default_file_based_stream.py index 588c4a18e4..0ef94dfb10 100644 --- a/airbyte_cdk/sources/file_based/stream/default_file_based_stream.py +++ b/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): diff --git a/unit_tests/sources/file_based/file_types/test_unstructured_parser.py b/unit_tests/sources/file_based/file_types/test_unstructured_parser.py index 374a02eed4..fc1fbd3829 100644 --- a/unit_tests/sources/file_based/file_types/test_unstructured_parser.py +++ b/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", diff --git a/unit_tests/sources/file_based/scenarios/unstructured_scenarios.py b/unit_tests/sources/file_based/scenarios/unstructured_scenarios.py index c0db46e7a4..ae42b5f960 100644 --- a/unit_tests/sources/file_based/scenarios/unstructured_scenarios.py +++ b/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"], diff --git a/unit_tests/sources/file_based/stream/test_default_file_based_stream.py b/unit_tests/sources/file_based/stream/test_default_file_based_stream.py index 54394a36d8..aacf7d63cb 100644 --- a/unit_tests/sources/file_based/stream/test_default_file_based_stream.py +++ b/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}