fix LanceDB retrieval (#195)

Quansight · Nov 10, 2023 · 8d14f87 · 8d14f87
1 parent dc497ef
commit 8d14f87
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/ragna/source_storages/_lancedb.py b/ragna/source_storages/_lancedb.py
@@ -99,7 +99,10 @@ def retrieve(
         # retrieving to few sources and needed to query again.
         limit = int(num_tokens * 2 / chunk_size)
         results = (
-            table.search(vector_column_name=self._VECTOR_COLUMN_NAME)
+            table.search(
+                self._embedding_function([prompt])[0],
+                vector_column_name=self._VECTOR_COLUMN_NAME,
+            )
             .limit(limit)
             .to_arrow()
         )