add similarity search with distance score

Author: juanfe88

src/langchain_google_firestore/vectorstores.py

@@ -240,6 +240,7 @@ def _similarity_search(
         query: List[float],
         k: int = DEFAULT_TOP_K,
         filters: Optional[BaseFilter] = None,
+        with_scores: Optional[bool] = False,
     ) -> List[DocumentSnapshot]:
         _filters = filters or self.filters
 
@@ -253,6 +254,7 @@ def _similarity_search(
             query_vector=Vector(query),
             distance_measure=self.distance_strategy,
             limit=k,
+            distance_result_field= 'distance' if with_scores else None
         )
 
         return results.get()
@@ -413,6 +415,14 @@ def max_marginal_relevance_search_by_vector(
         )
         return [convert_firestore_document(doc_results[i]) for i in mmr_doc_indexes]
 
+    def similarity_search_with_score(self, query, k = 4,filters: Optional[BaseFilter] = None, **kwargs):
+        docs = self._similarity_search(
+            self.embedding_service.embed_query(query), k, filters=filters,with_scores=True
+        )
+        return [
+            (convert_firestore_document(doc, page_content_fields=[self.content_field]),doc.to_dict()["distance"])
+            for doc in docs
+        ]
     @classmethod
     def from_texts(
         cls: Type[FirestoreVectorStore],

tests/test_vectorstores.py

@@ -495,6 +495,76 @@ def test_firestore_max_marginal_relevance_by_vector(
     test_case.assertEqual(len(results), k)
 
 
+def test_firestore_similarity_search_with_score(
+    test_case: TestCase,
+    test_collection: str,
+    client,
+    embeddings: FakeEmbeddings,
+):
+    """
+    An end-to-end test for similarity search with score in FirestoreVectorStore.
+    """
+
+    # Create FirestoreVectorStore instance
+    firestore_store = FirestoreVectorStore(test_collection, embeddings, client=client)
+
+    texts = ["test1", "test2"]
+    k = 2
+
+    # Add vectors to Firestore
+    firestore_store.add_texts(texts, ids=["1", "2"])
+
+    # Perform similarity search with score
+    results = firestore_store.similarity_search_with_score("test1", k)
+
+    # Verify that the search results are as expected
+    test_case.assertEqual(len(results), k)
+    
+    # Check that each result is a tuple with a Document and a score
+    for result in results:
+        test_case.assertTrue(isinstance(result, tuple))
+        test_case.assertEqual(len(result), 2)
+        test_case.assertTrue(isinstance(result[0], Document))
+        test_case.assertTrue(isinstance(result[1], float))
+
+
+def test_firestore_similarity_search_with_score_with_filters(
+    test_case: TestCase,
+    test_collection: str,
+    client: firestore.Client,
+    embeddings: FakeEmbeddings,
+):
+    """
+    An end-to-end test for similarity search with score in FirestoreVectorStore with filters.
+    Requires an index on the filter field in Firestore.
+    """
+
+    # Create FirestoreVectorStore instance
+    firestore_store = FirestoreVectorStore(test_collection, embeddings, client=client)
+
+    # Add vectors to Firestore
+    firestore_store.add_texts(
+        ["test1", "test2"],
+        ids=["1", "2"],
+        metadatas=[{"foo": "bar"}, {"foo": "baz"}],
+    )
+
+    # Perform similarity search with score and filter
+    results = firestore_store.similarity_search_with_score(
+        "test1", k=2, filters=FieldFilter("metadata.foo", "==", "bar")
+    )
+
+    # Verify that the search results are as expected with the filter applied
+    test_case.assertEqual(len(results), 1)
+    
+    # Check that the result is a tuple with a Document and a score
+    doc, score = results[0]
+    test_case.assertTrue(isinstance(doc, Document))
+    test_case.assertTrue(isinstance(score, float))
+    test_case.assertEqual(doc.page_content, "test1")
+    test_case.assertEqual(doc.metadata["metadata"]["foo"], "bar")
+
+
 def test_firestore_from_texts(
     test_case: TestCase,
     test_collection: str,