fix retriever signatures (#7097)

This commit is contained in:
Bagatur 2023-07-03 14:21:36 -06:00 committed by GitHub
parent 490fcf9d98
commit fd3f8efec7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 9 additions and 9 deletions

View File

@ -1,4 +1,4 @@
from typing import Any, Dict, List, Optional, cast from typing import Any, Dict, List, cast
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
@ -40,7 +40,7 @@ class LlamaIndexRetriever(BaseRetriever, BaseModel):
return docs return docs
async def _aget_relevant_documents( async def _aget_relevant_documents(
self, query: str, *, run_manager: Optional[AsyncCallbackManagerForRetrieverRun] self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
) -> List[Document]: ) -> List[Document]:
raise NotImplementedError("LlamaIndexRetriever does not support async") raise NotImplementedError("LlamaIndexRetriever does not support async")

View File

@ -117,7 +117,7 @@ class SelfQueryRetriever(BaseRetriever, BaseModel):
return docs return docs
async def _aget_relevant_documents( async def _aget_relevant_documents(
self, query: str, *, run_manager: Optional[AsyncCallbackManagerForRetrieverRun] self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
) -> List[Document]: ) -> List[Document]:
raise NotImplementedError raise NotImplementedError

View File

@ -407,7 +407,7 @@ class VectorStoreRetriever(BaseRetriever, BaseModel):
return values return values
def _get_relevant_documents( def _get_relevant_documents(
self, query: str, *, run_manager: Optional[CallbackManagerForRetrieverRun] self, query: str, *, run_manager: CallbackManagerForRetrieverRun
) -> List[Document]: ) -> List[Document]:
if self.search_type == "similarity": if self.search_type == "similarity":
docs = self.vectorstore.similarity_search(query, **self.search_kwargs) docs = self.vectorstore.similarity_search(query, **self.search_kwargs)
@ -427,7 +427,7 @@ class VectorStoreRetriever(BaseRetriever, BaseModel):
return docs return docs
async def _aget_relevant_documents( async def _aget_relevant_documents(
self, query: str, *, run_manager: Optional[AsyncCallbackManagerForRetrieverRun] self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
) -> List[Document]: ) -> List[Document]:
if self.search_type == "similarity": if self.search_type == "similarity":
docs = await self.vectorstore.asimilarity_search( docs = await self.vectorstore.asimilarity_search(

View File

@ -620,7 +620,7 @@ class RedisVectorStoreRetriever(VectorStoreRetriever, BaseModel):
return values return values
def _get_relevant_documents( def _get_relevant_documents(
self, query: str, *, run_manager: Optional[CallbackManagerForRetrieverRun] self, query: str, *, run_manager: CallbackManagerForRetrieverRun
) -> List[Document]: ) -> List[Document]:
if self.search_type == "similarity": if self.search_type == "similarity":
docs = self.vectorstore.similarity_search(query, k=self.k) docs = self.vectorstore.similarity_search(query, k=self.k)
@ -633,7 +633,7 @@ class RedisVectorStoreRetriever(VectorStoreRetriever, BaseModel):
return docs return docs
async def _aget_relevant_documents( async def _aget_relevant_documents(
self, query: str, *, run_manager: Optional[AsyncCallbackManagerForRetrieverRun] self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
) -> List[Document]: ) -> List[Document]:
raise NotImplementedError("RedisVectorStoreRetriever does not support async") raise NotImplementedError("RedisVectorStoreRetriever does not support async")

View File

@ -451,7 +451,7 @@ class SingleStoreDBRetriever(VectorStoreRetriever):
allowed_search_types: ClassVar[Collection[str]] = ("similarity",) allowed_search_types: ClassVar[Collection[str]] = ("similarity",)
def _get_relevant_documents( def _get_relevant_documents(
self, query: str, *, run_manager: Optional[CallbackManagerForRetrieverRun] self, query: str, *, run_manager: CallbackManagerForRetrieverRun
) -> List[Document]: ) -> List[Document]:
if self.search_type == "similarity": if self.search_type == "similarity":
docs = self.vectorstore.similarity_search(query, k=self.k) docs = self.vectorstore.similarity_search(query, k=self.k)
@ -460,7 +460,7 @@ class SingleStoreDBRetriever(VectorStoreRetriever):
return docs return docs
async def _aget_relevant_documents( async def _aget_relevant_documents(
self, query: str, *, run_manager: Optional[AsyncCallbackManagerForRetrieverRun] self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
) -> List[Document]: ) -> List[Document]:
raise NotImplementedError( raise NotImplementedError(
"SingleStoreDBVectorStoreRetriever does not support async" "SingleStoreDBVectorStoreRetriever does not support async"