Hook up get_documents_from_query_results

googleapis · Jan 30, 2025 · 039b1bd · 039b1bd
1 parent 5f01339
commit 039b1bd
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/samples/search_ann.py b/samples/search_ann.py
@@ -56,7 +56,7 @@ def use_case():
                 nullable_column=True,
                 num_branches=1000,
                 tree_depth=3,
-                index_type=distance_strategy,
+                distance_type=distance_strategy,
                 num_leaves=100000,
             ),
         ],
@@ -140,6 +140,10 @@ def clear_and_insert_data(tx):
         id_column="categoryId",
         embedding_service=embeddings,
         embedding_column="productDescriptionEmbedding",
+        query_parameters=QueryParameters(
+            algorithm=QueryParameters.NearestNeighborsAlgorithm.APPROXIMATE_NEAREST_NEIGHBOR,
+            distance_strategy=distance_strategy,
+        ),
         skip_not_nullable_columns=True,
     )
     results = vec_store.search_by_ANN(

diff --git a/src/langchain_google_spanner/vector_store.py b/src/langchain_google_spanner/vector_store.py
@@ -1053,7 +1053,7 @@ def search_by_ANN(
         embedding_column_is_nullable: bool = False,
         ascending: bool = True,
         return_columns: List[str] = None,
-    ) -> List[Any]:
+    ) -> List[Document]:
         sql = SpannerVectorStore._query_ANN(
             self._table_name,
             index_name,
@@ -1074,7 +1074,12 @@ def search_by_ANN(
         ) as snapshot:
             print("search by ANN sql", sql)
             results = snapshot.execute_sql(sql=sql)
-            return list(results)
+            column_order_map = {
+                value: index for index, value in enumerate(self._columns_to_insert)
+            }
+            return self._get_documents_from_query_results(
+                list(results), column_order_map
+            )
 
     @staticmethod
     def _query_ANN(