Harrison/typing of return (#4685)

Co-authored-by: OlajideOgun <37077640+OlajideOgun@users.noreply.github.com>
This commit is contained in:
Harrison Chase 2023-05-14 18:25:50 -07:00 committed by GitHub
parent c48f1301ee
commit ed8207b2fb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 4 additions and 5 deletions

View File

@ -329,7 +329,7 @@ class VectorStore(ABC):
"""Return VectorStore initialized from texts and embeddings."""
raise NotImplementedError
def as_retriever(self, **kwargs: Any) -> BaseRetriever:
def as_retriever(self, **kwargs: Any) -> VectorStoreRetriever:
return VectorStoreRetriever(vectorstore=self, **kwargs)

View File

@ -23,9 +23,8 @@ from pydantic import BaseModel, root_validator
from langchain.docstore.document import Document
from langchain.embeddings.base import Embeddings
from langchain.schema import BaseRetriever
from langchain.utils import get_from_dict_or_env
from langchain.vectorstores.base import VectorStore
from langchain.vectorstores.base import VectorStore, VectorStoreRetriever
logger = logging.getLogger(__name__)
@ -544,11 +543,11 @@ class Redis(VectorStore):
**kwargs,
)
def as_retriever(self, **kwargs: Any) -> BaseRetriever:
def as_retriever(self, **kwargs: Any) -> RedisVectorStoreRetriever:
return RedisVectorStoreRetriever(vectorstore=self, **kwargs)
class RedisVectorStoreRetriever(BaseRetriever, BaseModel):
class RedisVectorStoreRetriever(VectorStoreRetriever, BaseModel):
vectorstore: Redis
search_type: str = "similarity"
k: int = 4