add async vector operations in VectorStore base class (#2535)

not currently implemented by any subclasses
This commit is contained in:
Ankush Gola 2023-04-07 16:22:14 +02:00 committed by GitHub
parent 481de8df7f
commit 6dbd29e440
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -32,6 +32,15 @@ class VectorStore(ABC):
List of ids from adding the texts into the vectorstore. List of ids from adding the texts into the vectorstore.
""" """
async def aadd_texts(
self,
texts: Iterable[str],
metadatas: Optional[List[dict]] = None,
**kwargs: Any,
) -> List[str]:
"""Run more texts through the embeddings and add to the vectorstore."""
raise NotImplementedError
def add_documents(self, documents: List[Document], **kwargs: Any) -> List[str]: def add_documents(self, documents: List[Document], **kwargs: Any) -> List[str]:
"""Run more documents through the embeddings and add to the vectorstore. """Run more documents through the embeddings and add to the vectorstore.
@ -47,12 +56,33 @@ class VectorStore(ABC):
metadatas = [doc.metadata for doc in documents] metadatas = [doc.metadata for doc in documents]
return self.add_texts(texts, metadatas, **kwargs) return self.add_texts(texts, metadatas, **kwargs)
async def aadd_documents(
self, documents: List[Document], **kwargs: Any
) -> List[str]:
"""Run more documents through the embeddings and add to the vectorstore.
Args:
documents (List[Document]: Documents to add to the vectorstore.
Returns:
List[str]: List of IDs of the added texts.
"""
texts = [doc.page_content for doc in documents]
metadatas = [doc.metadata for doc in documents]
return await self.aadd_texts(texts, metadatas, **kwargs)
@abstractmethod @abstractmethod
def similarity_search( def similarity_search(
self, query: str, k: int = 4, **kwargs: Any self, query: str, k: int = 4, **kwargs: Any
) -> List[Document]: ) -> List[Document]:
"""Return docs most similar to query.""" """Return docs most similar to query."""
async def asimilarity_search(
self, query: str, k: int = 4, **kwargs: Any
) -> List[Document]:
"""Return docs most similar to query."""
raise NotImplementedError
def similarity_search_by_vector( def similarity_search_by_vector(
self, embedding: List[float], k: int = 4, **kwargs: Any self, embedding: List[float], k: int = 4, **kwargs: Any
) -> List[Document]: ) -> List[Document]:
@ -67,6 +97,12 @@ class VectorStore(ABC):
""" """
raise NotImplementedError raise NotImplementedError
async def asimilarity_search_by_vector(
self, embedding: List[float], k: int = 4, **kwargs: Any
) -> List[Document]:
"""Return docs most similar to embedding vector."""
raise NotImplementedError
def max_marginal_relevance_search( def max_marginal_relevance_search(
self, query: str, k: int = 4, fetch_k: int = 20 self, query: str, k: int = 4, fetch_k: int = 20
) -> List[Document]: ) -> List[Document]:
@ -85,6 +121,12 @@ class VectorStore(ABC):
""" """
raise NotImplementedError raise NotImplementedError
async def amax_marginal_relevance_search(
self, query: str, k: int = 4, fetch_k: int = 20
) -> List[Document]:
"""Return docs selected using the maximal marginal relevance."""
raise NotImplementedError
def max_marginal_relevance_search_by_vector( def max_marginal_relevance_search_by_vector(
self, embedding: List[float], k: int = 4, fetch_k: int = 20 self, embedding: List[float], k: int = 4, fetch_k: int = 20
) -> List[Document]: ) -> List[Document]:
@ -103,6 +145,12 @@ class VectorStore(ABC):
""" """
raise NotImplementedError raise NotImplementedError
async def amax_marginal_relevance_search_by_vector(
self, embedding: List[float], k: int = 4, fetch_k: int = 20
) -> List[Document]:
"""Return docs selected using the maximal marginal relevance."""
raise NotImplementedError
@classmethod @classmethod
def from_documents( def from_documents(
cls, cls,
@ -115,6 +163,18 @@ class VectorStore(ABC):
metadatas = [d.metadata for d in documents] metadatas = [d.metadata for d in documents]
return cls.from_texts(texts, embedding, metadatas=metadatas, **kwargs) return cls.from_texts(texts, embedding, metadatas=metadatas, **kwargs)
@classmethod
async def afrom_documents(
cls,
documents: List[Document],
embedding: Embeddings,
**kwargs: Any,
) -> VectorStore:
"""Return VectorStore initialized from documents and embeddings."""
texts = [d.page_content for d in documents]
metadatas = [d.metadata for d in documents]
return await cls.afrom_texts(texts, embedding, metadatas=metadatas, **kwargs)
@classmethod @classmethod
@abstractmethod @abstractmethod
def from_texts( def from_texts(
@ -126,6 +186,17 @@ class VectorStore(ABC):
) -> VectorStore: ) -> VectorStore:
"""Return VectorStore initialized from texts and embeddings.""" """Return VectorStore initialized from texts and embeddings."""
@classmethod
async def afrom_texts(
cls,
texts: List[str],
embedding: Embeddings,
metadatas: Optional[List[dict]] = None,
**kwargs: Any,
) -> VectorStore:
"""Return VectorStore initialized from texts and embeddings."""
raise NotImplementedError
def as_retriever(self, **kwargs: Any) -> BaseRetriever: def as_retriever(self, **kwargs: Any) -> BaseRetriever:
return VectorStoreRetriever(vectorstore=self, **kwargs) return VectorStoreRetriever(vectorstore=self, **kwargs)
@ -161,4 +232,14 @@ class VectorStoreRetriever(BaseRetriever, BaseModel):
return docs return docs
async def aget_relevant_documents(self, query: str) -> List[Document]: async def aget_relevant_documents(self, query: str) -> List[Document]:
raise NotImplementedError("VectorStoreRetriever does not support async") if self.search_type == "similarity":
docs = await self.vectorstore.asimilarity_search(
query, **self.search_kwargs
)
elif self.search_type == "mmr":
docs = await self.vectorstore.amax_marginal_relevance_search(
query, **self.search_kwargs
)
else:
raise ValueError(f"search_type of {self.search_type} not allowed.")
return docs