github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-06-22 06:39:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

enable mmr search (#807)

Browse Source
This commit is contained in:
Harrison Chase 2023-01-30 14:48:24 -08:00 committed by GitHub
parent f46f1d28af
commit b3916f74a7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 21 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
langchain/chains/vector_db_qa/base.py
View File

@ -41,6 +41,8 @@ class VectorDBQA(Chain, BaseModel):
"""Return the source documents."""
search_kwargs: Dict[str, Any] = Field(default_factory=dict)
"""Extra search args."""
search_type: str = "similarity"
"""Search type to use over vectorstore. `similarity` or `mmr`."""
class Config:
"""Configuration for this pydantic object."""
@ -90,6 +92,15 @@ class VectorDBQA(Chain, BaseModel):
values["combine_documents_chain"] = combine_documents_chain
return values
@root_validator()
def validate_search_type(cls, values: Dict) -> Dict:
"""Validate search type."""
if "search_type" in values:
search_type = values["search_type"]
if search_type not in ("similarity", "mmr"):
raise ValueError(f"search_type of {search_type} not allowed.")
return values
@classmethod
def from_llm(
cls, llm: BaseLLM, prompt: PromptTemplate = PROMPT, **kwargs: Any
@ -129,9 +140,16 @@ class VectorDBQA(Chain, BaseModel):
"""
question = inputs[self.input_key]
docs = self.vectorstore.similarity_search(
question, k=self.k, **self.search_kwargs
)
if self.search_type == "similarity":
docs = self.vectorstore.similarity_search(
question, k=self.k, **self.search_kwargs
)
elif self.search_type == "mmr":
docs = self.vectorstore.max_marginal_relevance_search(
question, k=self.k, **self.search_kwargs
)
else:
raise ValueError(f"search_type of {self.search_type} not allowed.")
answer, _ = self.combine_documents_chain.combine_docs(docs, question=question)
if self.return_source_documents: