make metadata column in vecstore inputs optional (#7913)

GitOrigin-RevId: 33c6857542c21351ce9f26967a5984c72e9ddbb9

Author: berkecanrizai

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

@@ -406,3 +406,32 @@ def fake_embeddings_model(x: str) -> list[float]:
     (query_result,) = val.as_list()  # extract the single match
     assert isinstance(query_result, dict)
     assert query_result["text"]  # just check if some text was returned
+
+
+def test_docstore_on_table_without_metadata():
+    @pw.udf
+    def fake_embeddings_model(x: str) -> list[float]:
+        return [1.0, 1.0, 0.0]
+
+    docs = pw.debug.table_from_rows(
+        schema=pw.schema_from_types(data=bytes),
+        rows=[("test".encode("utf-8"),)],
+    )
+
+    vector_server = VectorStoreServer(
+        docs,
+        embedder=fake_embeddings_model,
+    )
+
+    retrieve_queries = pw.debug.table_from_rows(
+        schema=vector_server.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"] == "test"  # just check if some text was returned

python/pathway/xpacks/llm/vector_store.py

@@ -12,8 +12,9 @@
 import json
 import logging
 import threading
+import warnings
 from collections.abc import Callable, Coroutine
-from typing import TYPE_CHECKING, TypeAlias, cast
+from typing import TYPE_CHECKING, Iterable, TypeAlias, cast
 
 import jmespath
 import requests
@@ -206,18 +207,39 @@ def generic_transformer(x: str) -> list[tuple[str, dict]]:
             **kwargs,
         )
 
+    def _clean_tables(
+        self, docs: pw.Table | Iterable[pw.Table]
+    ) -> tuple[pw.Table, ...]:
+        if isinstance(docs, pw.Table):
+            docs = [docs]
+
+        def _clean_table(doc: pw.Table) -> pw.Table:
+            if "_metadata" not in doc.column_names():
+                warnings.warn(
+                    f"`_metadata` column is not present in Table {doc}. Filtering will not work for this Table"
+                )
+                doc = doc.with_columns(_metadata=dict())
+
+            return doc.select(pw.this.data, pw.this._metadata)
+
+        return tuple([_clean_table(doc) for doc in docs])
+
     def _build_graph(self) -> dict:
         """
         Builds the pathway computation graph for indexing documents and serving queries.
         """
         docs_s = self.docs
+
         if not docs_s:
             raise ValueError(
                 """Please provide at least one data source, e.g. read files from disk:
 
 pw.io.fs.read('./sample_docs', format='binary', mode='static', with_metadata=True)
 """
             )
+
+        docs_s = self._clean_tables(docs_s)
+
         if len(docs_s) == 1:
             (docs,) = docs_s
         else: