feat: Adding TruncatedEmbeddingStep (2) and document metadata to Qdrant (#183)

Author: malteos

tests/data/embedded.csv

@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: 2025 Deutsche Telekom AG (opensource@telekom.de)
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from wurzel.steps.data import EmbeddingResult
+
+
+def test_load_from_csv_converts_dict_columns(tmp_path):
+    csv = tmp_path / "embedded.csv"
+    csv.write_text(
+        'text,vector,url,keywords,embedding_input_text,metadata\n"foo","[0.1, 0.2]","","https://example.com","kw","{\'foo\': \'bar\'}"\n',
+        encoding="utf-8",
+    )
+
+    df = EmbeddingResult.load_from_path(csv)
+
+    metadata = df["metadata"].iloc[0]
+    assert isinstance(metadata, dict)
+    assert metadata["foo"] == "bar"

tests/datacontract/dataframe_model_test.py

@@ -100,8 +100,17 @@ def test_split_markdown_document(Splitter):
     )
 
     result = Splitter.split_markdown_document(contract)
-    assert len(result) > 1
-    assert "TV HD Recorder Fehlerbehebun" in result[-1].md
+    assert len(result) == 5, "Splitter produce invalid number of chunks"
+    assert "TV HD Recorder Fehlerbehebun" in result[-1].md, "Invalid chunk content"
+
+    # Check metadata
+    expected_hash = "1b5098dbc4584f019bb00cbbb42a36ef27e908b216f40e09ae77f30ca1cddc2f"  # pragma: allowlist secret
+    assert result[0].metadata["source_sha256_hash"] == expected_hash, "Invalid source hash"
+
+    assert result[0].metadata["source_sha256_hash"] == result[-1].metadata["source_sha256_hash"], "Source hashes are not the same"
+    assert result[0].metadata["chunks_count"] == 5, "Chunk metadata is invalid"
+    assert result[0].metadata["chunk_index"] == 0, "Chunk metadata is invalid"
+    assert result[-1].metadata["chunk_index"] == 4, "Chunk metadata is invalid"
 
 
 def test_sentence_splitter(Splitter: SemanticSplitter):

tests/splitter_test.py

@@ -25,8 +25,8 @@
         pytest.param(
             DataFrame[EmbeddingResult](
                 [
-                    {"text": "a", "url": "url", "vector": [0.1], "keywords": "kw"},
-                    {"text": "b", "url": "url", "vector": [0.1], "keywords": "kw"},
+                    {"text": "a", "url": "url", "vector": [0.1], "keywords": "kw", "embedding_input_text": "a", "metadata": {"foo": "bar"}},
+                    {"text": "b", "url": "url", "vector": [0.1], "keywords": "kw", "embedding_input_text": "b", "metadata": {"foo": "bar"}},
                 ]
             ),
             id="DataFrame",
@@ -63,14 +63,14 @@ def test_unsorted(run_num, expected):
 def test_unsorted_df():
     unsorted = DataFrame[EmbeddingResult](
         [
-            {"text": "b", "url": "url", "vector": [0.1], "keywords": "kw"},
-            {"text": "a", "url": "url", "vector": [0.1], "keywords": "kw"},
+            {"text": "b", "url": "url", "vector": [0.1], "keywords": "kw", "embedding_input_text": "b", "metadata": {"foo": "bar"}},
+            {"text": "a", "url": "url", "vector": [0.1], "keywords": "kw", "embedding_input_text": "a", "metadata": {"foo": "bar"}},
         ]
     )
     sort = DataFrame[EmbeddingResult](
         [
-            {"text": "a", "url": "url", "vector": [0.1], "keywords": "kw"},
-            {"text": "b", "url": "url", "vector": [0.1], "keywords": "kw"},
+            {"text": "a", "url": "url", "vector": [0.1], "keywords": "kw", "embedding_input_text": "a", "metadata": {"foo": "bar"}},
+            {"text": "b", "url": "url", "vector": [0.1], "keywords": "kw", "embedding_input_text": "b", "metadata": {"foo": "bar"}},
         ]
     )
     assert not sort.equals(unsorted), "sanity check"

tests/step_executor/sort_test.py

@@ -2,6 +2,10 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import shutil
+from pathlib import Path
+
+import numpy as np
 import pytest
 import requests_mock
 
@@ -32,3 +36,61 @@ def embedding_service_mock():
         m.post("/embed", text=POST_RESULT_EMBEDDING_STR)
         m.get("/info", text=GET_RESULT_INFO)
         yield
+
+
+@pytest.fixture(scope="module")
+def mock_embedding():
+    """A pytest fixture that provides a mock embedding class for testing.
+
+    Overrides the `_select_embedding` method of the `EmbeddingStep` class
+    to return an instance of the mock embedding class, which generates
+    a fixed-size random vector upon calling `embed_query`.
+
+    Returns:
+    -------
+    MockEmbedding
+        An instance of the mock embedding class.
+
+    """
+
+    class MockEmbedding:
+        def embed_query(self, _: str) -> list[float]:
+            """Simulates embedding of a query string into a fixed-size random vector.
+
+            Parameters
+            ----------
+            _ : str
+                The input query string (ignored in this mock implementation).
+
+            Returns:
+            -------
+            np.ndarray
+                A random vector of size 768.
+
+            """
+            return list(np.random.random(768))
+
+    def mock_func(*args, **kwargs):
+        return MockEmbedding()
+
+    return mock_func
+
+
+@pytest.fixture
+def default_embedding_data(tmp_path):
+    mock_file = Path("tests/data/markdown.json")
+    input_folder = tmp_path / "input"
+    input_folder.mkdir()
+    shutil.copy(mock_file, input_folder)
+    output_folder = tmp_path / "out"
+    return (input_folder, output_folder)
+
+
+@pytest.fixture
+def splitter_tokenizer_model():
+    return "gpt-3.5-turbo"
+
+
+@pytest.fixture
+def sentence_splitter_model():
+    return "de_core_news_sm"

tests/steps/embedding/conftest.py

@@ -7,7 +7,6 @@
 import shutil
 from pathlib import Path
 
-import numpy as np
 import pytest
 
 from wurzel.utils import HAS_LANGCHAIN_CORE, HAS_REQUESTS, HAS_SPACY, HAS_TIKTOKEN
@@ -21,59 +20,8 @@
 from wurzel.steps.embedding.huggingface import HuggingFaceInferenceAPIEmbeddings
 from wurzel.steps.embedding.step_multivector import EmbeddingMultiVectorStep
 
-SPLITTER_TOKENIZER_MODEL = "gpt-3.5-turbo"
-SENTENCE_SPLITTER_MODEL = "de_core_news_sm"
 
-
-@pytest.fixture(scope="module")
-def mock_embedding():
-    """A pytest fixture that provides a mock embedding class for testing.
-
-    Overrides the `_select_embedding` method of the `EmbeddingStep` class
-    to return an instance of the mock embedding class, which generates
-    a fixed-size random vector upon calling `embed_query`.
-
-    Returns:
-    -------
-    MockEmbedding
-        An instance of the mock embedding class.
-
-    """
-
-    class MockEmbedding:
-        def embed_query(self, _: str) -> list[float]:
-            """Simulates embedding of a query string into a fixed-size random vector.
-
-            Parameters
-            ----------
-            _ : str
-                The input query string (ignored in this mock implementation).
-
-            Returns:
-            -------
-            np.ndarray
-                A random vector of size 768.
-
-            """
-            return list(np.random.random(768))
-
-    def mock_func(*args, **kwargs):
-        return MockEmbedding()
-
-    return mock_func
-
-
-@pytest.fixture
-def default_embedding_data(tmp_path):
-    mock_file = Path("tests/data/markdown.json")
-    input_folder = tmp_path / "input"
-    input_folder.mkdir()
-    shutil.copy(mock_file, input_folder)
-    output_folder = tmp_path / "out"
-    return (input_folder, output_folder)
-
-
-def test_embedding_step(mock_embedding, default_embedding_data, env):
+def test_embedding_step(mock_embedding, default_embedding_data, env, splitter_tokenizer_model, sentence_splitter_model):
     """Tests the execution of the `EmbeddingStep` with a mock input file.
 
     Parameters
@@ -92,8 +40,8 @@ def test_embedding_step(mock_embedding, default_embedding_data, env):
     env.set("EMBEDDINGSTEP__TOKEN_COUNT_MIN", "64")
     env.set("EMBEDDINGSTEP__TOKEN_COUNT_MAX", "256")
     env.set("EMBEDDINGSTEP__TOKEN_COUNT_BUFFER", "32")
-    env.set("EMBEDDINGSTEP__TOKENIZER_MODEL", SPLITTER_TOKENIZER_MODEL)
-    env.set("EMBEDDINGSTEP__SENTENCE_SPLITTER_MODEL", SENTENCE_SPLITTER_MODEL)
+    env.set("EMBEDDINGSTEP__TOKENIZER_MODEL", splitter_tokenizer_model)
+    env.set("EMBEDDINGSTEP__SENTENCE_SPLITTER_MODEL", sentence_splitter_model)
 
     EmbeddingStep._select_embedding = mock_embedding
     input_folder, output_folder = default_embedding_data
@@ -152,15 +100,17 @@ def _select_embedding(*args, **kwargs) -> HuggingFaceInferenceAPIEmbeddings:
     assert sf.value.message.endswith(EXPECTED_EXCEPTION)
 
 
-def test_embedding_step_log_statistics(mock_embedding, default_embedding_data, env, caplog):
+def test_embedding_step_log_statistics(
+    mock_embedding, default_embedding_data, env, caplog, splitter_tokenizer_model, sentence_splitter_model
+):
     """Tests the logging of descriptive statistics in the `EmbeddingStep` with a mock input file."""
     env.set("EMBEDDINGSTEP__API", "https://example-embedding.com/embed")
     env.set("EMBEDDINGSTEP__NUM_THREADS", "1")  # Ensure deterministic behavior with single thread
     env.set("EMBEDDINGSTEP__TOKEN_COUNT_MIN", "64")
     env.set("EMBEDDINGSTEP__TOKEN_COUNT_MAX", "256")
     env.set("EMBEDDINGSTEP__TOKEN_COUNT_BUFFER", "32")
-    env.set("EMBEDDINGSTEP__TOKENIZER_MODEL", SPLITTER_TOKENIZER_MODEL)
-    env.set("EMBEDDINGSTEP__SENTENCE_SPLITTER_MODEL", SENTENCE_SPLITTER_MODEL)
+    env.set("EMBEDDINGSTEP__TOKENIZER_MODEL", splitter_tokenizer_model)
+    env.set("EMBEDDINGSTEP__SENTENCE_SPLITTER_MODEL", sentence_splitter_model)
 
     EmbeddingStep._select_embedding = mock_embedding
     input_folder, output_folder = default_embedding_data

tests/steps/embedding/e2e_test.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: 2025 Deutsche Telekom AG (opensource@telekom.de)
+#
+# SPDX-License-Identifier: Apache-2.0
+import pytest
+
+from wurzel.utils import HAS_LANGCHAIN_CORE, HAS_REQUESTS, HAS_SPACY, HAS_TIKTOKEN
+
+if not HAS_LANGCHAIN_CORE or not HAS_REQUESTS or not HAS_SPACY or not HAS_TIKTOKEN:
+    pytest.skip("Embedding dependencies (langchain-core, requests, spacy, tiktoken) are not available", allow_module_level=True)
+
+from wurzel.step_executor import BaseStepExecutor
+from wurzel.steps.embedding.step import TruncatedEmbeddingStep
+
+
+@pytest.mark.parametrize(
+    "token_count_max,mean_text_length",
+    [
+        (99999, 959.4),
+        (9999, 959.4),
+        (256, 491.1),
+        (128, 309.8),
+        (32, 103.4),
+    ],
+)
+def test_truncated_embedding_step(
+    token_count_max, mean_text_length, mock_embedding, default_embedding_data, env, splitter_tokenizer_model, sentence_splitter_model
+):
+    """Tests the execution of the `TruncatedEmbeddingStep` with a mock input file and check total output count and mean length of texts.
+
+    Parameters
+    ----------
+    mock_embedding : MockEmbedding
+        The mock embedding fixture.
+    tmp_path : pathlib.Path
+        A pytest fixture that provides a temporary directory unique to the test invocation.
+
+    Asserts
+    -------
+    Asserts that the `embedding.csv` file is created in the output folder.
+
+    """
+    env.set("TRUNCATEDEMBEDDINGSTEP__API", "https://example-embedding.com/embed")
+    env.set("TRUNCATEDEMBEDDINGSTEP__TOKEN_COUNT_MIN", "64")
+    env.set("TRUNCATEDEMBEDDINGSTEP__TOKEN_COUNT_MAX", str(token_count_max))
+    env.set("TRUNCATEDEMBEDDINGSTEP__TOKEN_COUNT_BUFFER", "32")
+    env.set("TRUNCATEDEMBEDDINGSTEP__TOKENIZER_MODEL", splitter_tokenizer_model)
+    env.set("TRUNCATEDEMBEDDINGSTEP__SENTENCE_SPLITTER_MODEL", sentence_splitter_model)
+
+    TruncatedEmbeddingStep._select_embedding = mock_embedding
+    input_folder, output_folder = default_embedding_data
+    step_res = BaseStepExecutor(dont_encapsulate=False).execute_step(TruncatedEmbeddingStep, [input_folder], output_folder)
+    assert output_folder.is_dir()
+    assert len(list(output_folder.glob("*"))) > 0
+
+    step_output, step_report = step_res[0]
+
+    assert len(step_output) == 7, "Step outputs have wrong count."
+    assert step_report.results == 7, "Step report has wrong count of outputs."
+
+    assert step_output.embedding_input_text.str.len().mean() == pytest.approx(mean_text_length, abs=0.1), (
+        "Invalid mean length of embedding_input_text"
+    )

tests/steps/embedding/truncated_embedding_step_test.py

@@ -31,7 +31,17 @@ def test_qdrant_connector_first(input_output_folder: tuple[Path, Path], dummy_co
     input_file = input_path / "qdrant_at.csv"
     output_file = output_path / "QdrantConnectorStep"
     shutil.copy("./tests/data/embedded.csv", input_file)
-    BaseStepExecutor().execute_step(QdrantConnectorStep, {input_path}, output_file)
+
+    with BaseStepExecutor() as ex:
+        step_res = ex(QdrantConnectorStep, {input_path}, output_file)
+
+        step_output, step_report = step_res[0]
+
+        # Validate step output
+        assert step_report.results == 2, "Invalid step results"
+
+        assert step_output["collection"][1] == "dummy_v1", "Invalid step output in collection name"
+        assert step_output["metadata"][0]["foo"] == "bar", "Invalid step output in metadata"
 
 
 def test_qdrant_connector_has_previous(input_output_folder: tuple[Path, Path], dummy_collection):

tests/steps/qdrant/e2e_test.py

@@ -52,10 +52,28 @@ def load_from_path(cls, path: Path, *args) -> Self:
         """Switch case to find the matching file ending."""
         import pandas as pd  # pylint: disable=import-outside-toplevel
 
+        # Load CSV from path
         read_data = pd.read_csv(path.open(encoding="utf-8"))
+
+        def _literal_eval_or_passthrough(value):
+            """Convert stringified literals to Python objects because pandas keeps CSV cells as strings."""
+            if not isinstance(value, str):
+                return value
+            stripped = value.strip()
+            if stripped == "":
+                return None
+            try:
+                return literal_eval(stripped)
+            except (ValueError, SyntaxError):
+                return value
+
+        # Iterate over coluns and load data
         for key, atr in cls.to_schema().columns.items():
-            if atr.dtype.type is list:
-                read_data[key] = read_data[key].apply(literal_eval)
+            if key not in read_data.columns:
+                continue
+            if atr.dtype.type in {list, dict}:
+                read_data[key] = read_data[key].apply(_literal_eval_or_passthrough)
+
         return patyp.DataFrame[cls](read_data)
 
 

wurzel/datacontract/datacontract.py

@@ -12,7 +12,10 @@
 if HAS_LANGCHAIN_CORE and HAS_REQUESTS:
     try:
         from .embedding import *  # noqa: F403 Allow importing Step classes
-        from .embedding import EmbeddingStep  # noqa: F401
+        from .embedding import (
+            EmbeddingStep,  # noqa: F401
+            TruncatedEmbeddingStep,  # noqa: F401
+        )
     except ImportError:
         pass