From f9aba724b941ca04689657cafed961d80ca099d4 Mon Sep 17 00:00:00 2001 From: anovazzi1 Date: Wed, 10 Jan 2024 17:35:35 -0300 Subject: [PATCH] Update RetrievalQAComponent parameters --- .../langflow/components/chains/RetrievalQA.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/src/backend/langflow/components/chains/RetrievalQA.py b/src/backend/langflow/components/chains/RetrievalQA.py index 9ed188bf5..4f0abb264 100644 --- a/src/backend/langflow/components/chains/RetrievalQA.py +++ b/src/backend/langflow/components/chains/RetrievalQA.py @@ -1,13 +1,11 @@ from langflow import CustomComponent -from langchain.chains import BaseRetrievalQA from typing import Optional, Union, Callable from langflow.field_typing import ( - BaseCombineDocumentsChain, BaseMemory, - BaseRetriever, -) - + BaseRetriever) +from langchain.chains.retrieval_qa.base import BaseRetrievalQA +from langchain.chains.combine_documents.base import BaseCombineDocumentsChain class RetrievalQAComponent(CustomComponent): display_name = "RetrievalQA" description = "Chain for question-answering against an index." @@ -17,8 +15,8 @@ class RetrievalQAComponent(CustomComponent): "combine_documents_chain": {"display_name": "Combine Documents Chain"}, "retriever": {"display_name": "Retriever"}, "memory": {"display_name": "Memory", "required": False}, - "input_key": {"display_name": "Input Key"}, - "output_key": {"display_name": "Output Key"}, + "input_key": {"display_name": "Input Key","advanced":True}, + "output_key": {"display_name": "Output Key","advanced":True}, "return_source_documents": {"display_name": "Return Source Documents"}, } @@ -27,9 +25,9 @@ class RetrievalQAComponent(CustomComponent): combine_documents_chain: BaseCombineDocumentsChain, retriever: BaseRetriever, memory: Optional[BaseMemory] = None, - input_key: str = "query", - output_key: str = "result", - return_source_documents: bool = True, + input_key: Optional[str] = "query", + output_key: Optional[str] = "result", + return_source_documents: Optional[bool] = True, ) -> Union[BaseRetrievalQA, Callable]: return BaseRetrievalQA( combine_documents_chain=combine_documents_chain,