added filter kwarg to VectorStoreIndexWrapper query and query_with_so… (#8844)

- Description: added filter to query methods in VectorStoreIndexWrapper
for filtering by metadata (i.e. search_kwargs)
- Tag maintainer: @rlancemartin, @eyurtsev

Updated the doc snippet on this topic as well. It took me a long while
to figure out how to filter the vectorstore by filename, so this might
help someone else out.

---------

Co-authored-by: Bagatur <baskaryan@gmail.com>
This commit is contained in:
Chris Pappalardo
2023-08-08 10:10:45 -07:00
committed by GitHub
parent 4a63533216
commit beab637f04
2 changed files with 31 additions and 5 deletions

View File

@@ -1,4 +1,4 @@
from typing import Any, List, Optional, Type
from typing import Any, Dict, List, Optional, Type
from pydantic import BaseModel, Extra, Field
@@ -31,22 +31,32 @@ class VectorStoreIndexWrapper(BaseModel):
arbitrary_types_allowed = True
def query(
self, question: str, llm: Optional[BaseLanguageModel] = None, **kwargs: Any
self,
question: str,
llm: Optional[BaseLanguageModel] = None,
retriever_kwargs: Optional[Dict[str, Any]] = None,
**kwargs: Any
) -> str:
"""Query the vectorstore."""
llm = llm or OpenAI(temperature=0)
retriever_kwargs = retriever_kwargs or {}
chain = RetrievalQA.from_chain_type(
llm, retriever=self.vectorstore.as_retriever(), **kwargs
llm, retriever=self.vectorstore.as_retriever(**retriever_kwargs), **kwargs
)
return chain.run(question)
def query_with_sources(
self, question: str, llm: Optional[BaseLanguageModel] = None, **kwargs: Any
self,
question: str,
llm: Optional[BaseLanguageModel] = None,
retriever_kwargs: Optional[Dict[str, Any]] = None,
**kwargs: Any
) -> dict:
"""Query the vectorstore and get back sources."""
llm = llm or OpenAI(temperature=0)
retriever_kwargs = retriever_kwargs or {}
chain = RetrievalQAWithSourcesChain.from_chain_type(
llm, retriever=self.vectorstore.as_retriever(), **kwargs
llm, retriever=self.vectorstore.as_retriever(**retriever_kwargs), **kwargs
)
return chain({chain.question_key: question})