jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Update RetrievalQAComponent parameters

Browse source
This commit is contained in:
anovazzi1 2024-01-10 17:35:35 -03:00
commit f9aba724b9
1 changed files with 8 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

18
src/backend/langflow/components/chains/RetrievalQA.py
View file

@ -1,13 +1,11 @@
from langflow import CustomComponent from langflow import CustomComponent
from langchain.chains import BaseRetrievalQA
from typing import Optional, Union, Callable from typing import Optional, Union, Callable
from langflow.field_typing import ( from langflow.field_typing import (
BaseCombineDocumentsChain,
BaseMemory, BaseMemory,
BaseRetriever, BaseRetriever)
) from langchain.chains.retrieval_qa.base import BaseRetrievalQA
from langchain.chains.combine_documents.base import BaseCombineDocumentsChain
class RetrievalQAComponent(CustomComponent): class RetrievalQAComponent(CustomComponent):
display_name = "RetrievalQA" display_name = "RetrievalQA"
description = "Chain for question-answering against an index." description = "Chain for question-answering against an index."
@ -17,8 +15,8 @@ class RetrievalQAComponent(CustomComponent):
"combine_documents_chain": {"display_name": "Combine Documents Chain"}, "combine_documents_chain": {"display_name": "Combine Documents Chain"},
"retriever": {"display_name": "Retriever"}, "retriever": {"display_name": "Retriever"},
"memory": {"display_name": "Memory", "required": False}, "memory": {"display_name": "Memory", "required": False},
"input_key": {"display_name": "Input Key"}, "input_key": {"display_name": "Input Key","advanced":True},
"output_key": {"display_name": "Output Key"}, "output_key": {"display_name": "Output Key","advanced":True},
"return_source_documents": {"display_name": "Return Source Documents"}, "return_source_documents": {"display_name": "Return Source Documents"},
} }
@ -27,9 +25,9 @@ class RetrievalQAComponent(CustomComponent):
combine_documents_chain: BaseCombineDocumentsChain, combine_documents_chain: BaseCombineDocumentsChain,
retriever: BaseRetriever, retriever: BaseRetriever,
memory: Optional[BaseMemory] = None, memory: Optional[BaseMemory] = None,
input_key: str = "query", input_key: Optional[str] = "query",
output_key: str = "result", output_key: Optional[str] = "result",
return_source_documents: bool = True, return_source_documents: Optional[bool] = True,
) -> Union[BaseRetrievalQA, Callable]: ) -> Union[BaseRetrievalQA, Callable]:
return BaseRetrievalQA( return BaseRetrievalQA(
combine_documents_chain=combine_documents_chain, combine_documents_chain=combine_documents_chain,