Fix doc_post_processors in document_store (#8276)

Co-authored-by: berkecanrizai <63911408+berkecanrizai@users.noreply.github.com>
GitOrigin-RevId: 7099adbc821b5f7f6bac29a81d983fcef30b895f

Author: 2 people

CHANGELOG.md

@@ -13,10 +13,11 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ### Changed
 - **BREAKING**: Changed the argument in `DoclingParser` from `parse_images` (bool) into `image_parsing_strategy` (Literal["llm"] | None)
+- **BREAKING**: `doc_post_processors` argument in the `pw.xpacks.llm.document_store.DocumentStore` now longer accepts `pw.UDF`.
 - Better error messages when using `pathway spawn` with multiple workers. Now error messages are printed only from the worker experiencing the error directly.
 
-### Removed
-
+### Fixed
+- `doc_post_processors` argument in the `pw.xpacks.llm.document_store.DocumentStore` had no effect. This is now fixed.
 
 ## [0.19.0] - 2025-02-20
 

python/pathway/xpacks/llm/__init__.py

@@ -2,6 +2,7 @@
 from ._typing import Doc, DocTransformer, DocTransformerCallable  # isort: skip
 
 from . import (
+    document_store,
     embedders,
     llms,
     parsers,
@@ -13,6 +14,7 @@
 )
 
 __all__ = [
+    "document_store",
     "embedders",
     "llms",
     "parsers",

python/pathway/xpacks/llm/_utils.py

@@ -98,3 +98,12 @@ def _to_dict(element: dict | pw.Json):
         return element.as_dict()
     else:
         return element
+
+
+def _wrap_doc_post_processor(fun: Callable[[str, dict], tuple[str, dict]]) -> pw.UDF:
+    @pw.udf
+    def wrapper(text: str, metadata: pw.Json) -> tuple[str, dict]:
+        metadata_dict = metadata.as_dict()
+        return fun(text, metadata_dict)
+
+    return wrapper

python/pathway/xpacks/llm/document_store.py

@@ -19,6 +19,7 @@
 from pathway.stdlib.indexing.data_index import _SCORE, DataIndex
 from pathway.stdlib.indexing.retrievers import AbstractRetrieverFactory
 from pathway.stdlib.ml.classifiers import _knn_lsh
+from pathway.xpacks.llm._utils import _wrap_doc_post_processor
 from pathway.xpacks.llm.utils import combine_metadata
 
 from ._utils import _wrap_udf
@@ -46,7 +47,8 @@ class DocumentStore:
         parser: callable that parses file contents into a list of documents.
         splitter: callable that splits long documents.
         doc_post_processors: optional list of callables that modify parsed files and metadata.
-            any callable takes two arguments (text: str, metadata: dict) and returns them as a tuple.
+            Each doc_post_processor is a Callable that takes two arguments
+            (text: str, metadata: dict) and returns them as a tuple.
     """
 
     def __init__(
@@ -56,7 +58,7 @@ def __init__(
         parser: Callable[[bytes], list[tuple[str, dict]]] | pw.UDF | None = None,
         splitter: Callable[[str], list[tuple[str, dict]]] | pw.UDF | None = None,
         doc_post_processors: (
-            list[Callable[[str, dict], tuple[str, dict]] | pw.UDF] | None
+            list[Callable[[str, dict], tuple[str, dict]]] | None
         ) = None,
     ):
         self.docs = docs
@@ -67,7 +69,9 @@ def __init__(
             else pathway.xpacks.llm.parsers.Utf8Parser()
         )
         self.doc_post_processors: list[pw.UDF] = (
-            [_wrap_udf(p) for p in doc_post_processors] if doc_post_processors else []
+            [_wrap_doc_post_processor(p) for p in doc_post_processors]
+            if doc_post_processors
+            else []
         )
         self.splitter: pw.UDF = (
             _wrap_udf(splitter)
@@ -247,6 +251,21 @@ def apply_processor(
         # `metadata` column: old_meta_dict -> old_meta_dict | new_meta_dict
         return combine_metadata(processed_docs)
 
+    @pw.table_transformer
+    def apply_doc_post_processor(
+        self, table: pw.Table, processor: pw.UDF
+    ) -> pw.Table[_DocumentSchema]:
+
+        processed_docs: pw.Table[DocumentStore._DocumentSchema] = table.select(
+            text=processor(pw.this.text, pw.this.metadata)
+            # some processors might split document into multiple parts so we flatten the results
+            # metadata will be propagated to all new rows
+        ).select(text=pw.this.text[0], metadata=pw.this.text[1])
+        # combine_metadata will transform our columns as follows:
+        # `text` column: tuple[str, new_meta_dict] -> str
+        # `metadata` column: new_meta_dict
+        return processed_docs
+
     def build_pipeline(self):
 
         cleaned_tables = self._clean_tables(self.docs)
@@ -272,13 +291,13 @@ def build_pipeline(self):
         # POST PROCESSING
         self.post_processed_docs = self.parsed_docs
         for post_processor in self.doc_post_processors:
-            self.post_processed_docs = self.apply_processor(
+            self.post_processed_docs = self.apply_doc_post_processor(
                 self.post_processed_docs, post_processor
             )
 
         # CHUNKING
         self.chunked_docs: pw.Table[DocumentStore._DocumentSchema] = (
-            self.apply_processor(self.parsed_docs, self.splitter)
+            self.apply_processor(self.post_processed_docs, self.splitter)
         )
 
         # INDEXING

python/pathway/xpacks/llm/tests/test_document_store.py

@@ -19,6 +19,7 @@
 from pathway.tests.utils import assert_table_equality
 from pathway.xpacks.llm.document_store import DocumentStore
 from pathway.xpacks.llm.servers import DocumentStoreServer
+from pathway.xpacks.llm.tests import mocks
 
 
 class DebugStatsInputSchema(DocumentStore.StatisticsQuerySchema):
@@ -663,3 +664,84 @@ def fake_embeddings_model(x: str) -> list[float]:
     (val,) = rows["result"].values()
     assert isinstance(val, pw.Json)
     assert len(val.as_list()) == 2
+
+
+def test_docstore_post_processor():
+
+    def add_baz(text: str, metadata: dict) -> tuple:
+        return (text + "baz", metadata)
+
+    docs = pw.debug.table_from_rows(
+        schema=pw.schema_from_types(data=bytes, _metadata=dict),
+        rows=[
+            (
+                "test".encode("utf-8"),
+                {"foo": "bar"},
+            )
+        ],
+    )
+
+    index_factory = BruteForceKnnFactory(
+        dimensions=3,
+        reserved_space=10,
+        embedder=mocks.fake_embeddings_model,
+        metric=BruteForceKnnMetricKind.COS,
+    )
+
+    vector_server = DocumentStore(
+        docs, retriever_factory=index_factory, doc_post_processors=[add_baz]
+    )
+
+    retrieve_queries = pw.debug.table_from_rows(
+        schema=DocumentStore.RetrieveQuerySchema,
+        rows=[("Foo", 1, None, None)],
+    )
+
+    retrieve_outputs = vector_server.retrieve_query(retrieve_queries)
+    _, rows = pw.debug.table_to_dicts(retrieve_outputs)
+    (val,) = rows["result"].values()
+    assert isinstance(val, pw.Json)
+    (query_result,) = val.as_list()  # extract the single match
+    assert isinstance(query_result, dict)
+    assert query_result["text"] == "testbaz"
+
+
+def test_docstore_metadata_post_processor():
+
+    def add_id(text: str, metadata: dict) -> tuple:
+        metadata["id"] = 1
+        return (text, metadata)
+
+    docs = pw.debug.table_from_rows(
+        schema=pw.schema_from_types(data=bytes, _metadata=dict),
+        rows=[
+            (
+                "test".encode("utf-8"),
+                {"foo": "bar"},
+            )
+        ],
+    )
+
+    index_factory = BruteForceKnnFactory(
+        dimensions=3,
+        reserved_space=10,
+        embedder=mocks.fake_embeddings_model,
+        metric=BruteForceKnnMetricKind.COS,
+    )
+
+    vector_server = DocumentStore(
+        docs, retriever_factory=index_factory, doc_post_processors=[add_id]
+    )
+
+    retrieve_queries = pw.debug.table_from_rows(
+        schema=DocumentStore.RetrieveQuerySchema,
+        rows=[("Foo", 1, None, None)],
+    )
+
+    retrieve_outputs = vector_server.retrieve_query(retrieve_queries)
+    _, rows = pw.debug.table_to_dicts(retrieve_outputs)
+    (val,) = rows["result"].values()
+    assert isinstance(val, pw.Json)
+    (query_result,) = val.as_list()  # extract the single match
+    assert isinstance(query_result, dict)
+    assert query_result["metadata"]["id"] == 1