From 3eb8d52c8c202e99e3fcaee02f921f76f8c61a4a Mon Sep 17 00:00:00 2001 From: Enias Cailliau Date: Mon, 17 Apr 2023 19:13:53 +0100 Subject: [PATCH] Add similarity search with score (#37) * Hotfix * Lint * Add similarity_search_with_score * Clean up codebase --- .../vectorstores/steamship_vector_store.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/steamship_langchain/vectorstores/steamship_vector_store.py b/src/steamship_langchain/vectorstores/steamship_vector_store.py index 3477d73..85fe68a 100644 --- a/src/steamship_langchain/vectorstores/steamship_vector_store.py +++ b/src/steamship_langchain/vectorstores/steamship_vector_store.py @@ -1,6 +1,6 @@ import uuid from itertools import zip_longest -from typing import Any, Iterable, List, Optional +from typing import Any, Iterable, List, Optional, Tuple from langchain.docstore.document import Document from langchain.text_splitter import TextSplitter @@ -161,6 +161,15 @@ def similarity_search(self, query: str, k: int = 4, **kwargs: Any) -> List[Docum for item in search_results.output.items ] + def similarity_search_with_score(self, query: str, k: int = 4) -> List[Tuple[Document, float]]: + search_results = self.index.search(query, k=k) + search_results.wait() + docs = [] + for item in search_results.output.items: + doc = Document(page_content=item.tag.text, metadata=item.tag.value) + docs.append((doc, item.score)) + return docs + def similarity_search_by_vector( self, embedding: List[float], k: int = 4, **kwargs: Any ) -> List[Document]: