for ChatVectorDBChain, add top_k_docs_for_context to allow control how many chunks of context will be retrieved (#1155)

given that we allow user define chunk size, think it would be useful for
user to define how many chunks of context will be retrieved.
searx-api
Ji 1 year ago committed by GitHub
parent 955c89fccb
commit ed37fbaeff
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -32,6 +32,7 @@ class ChatVectorDBChain(Chain, BaseModel):
question_generator: LLMChain
output_key: str = "answer"
return_source_documents: bool = False
top_k_docs_for_context: int = 4
"""Return the source documents."""
@property
@ -88,7 +89,9 @@ class ChatVectorDBChain(Chain, BaseModel):
)
else:
new_question = question
docs = self.vectorstore.similarity_search(new_question, k=4, **vectordbkwargs)
docs = self.vectorstore.similarity_search(
new_question, k=self.top_k_docs_for_context, **vectordbkwargs
)
new_inputs = inputs.copy()
new_inputs["question"] = new_question
new_inputs["chat_history"] = chat_history_str
@ -109,7 +112,9 @@ class ChatVectorDBChain(Chain, BaseModel):
else:
new_question = question
# TODO: This blocks the event loop, but it's not clear how to avoid it.
docs = self.vectorstore.similarity_search(new_question, k=4, **vectordbkwargs)
docs = self.vectorstore.similarity_search(
new_question, k=self.top_k_docs_for_context, **vectordbkwargs
)
new_inputs = inputs.copy()
new_inputs["question"] = new_question
new_inputs["chat_history"] = chat_history_str

Loading…
Cancel
Save