From 39e7a92f4564140ad582e4096ca84954c30a16da Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sat, 30 Mar 2024 17:32:44 -0300 Subject: [PATCH 1/2] Refactor RetrievalQAComponent and RetrievalQAWithSourcesChainComponent --- .../langflow/components/chains/RetrievalQA.py | 20 ++++++++++++------- .../chains/RetrievalQAWithSourcesChain.py | 3 ++- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/src/backend/base/langflow/components/chains/RetrievalQA.py b/src/backend/base/langflow/components/chains/RetrievalQA.py index 4ac514ec5..7b6ba49e4 100644 --- a/src/backend/base/langflow/components/chains/RetrievalQA.py +++ b/src/backend/base/langflow/components/chains/RetrievalQA.py @@ -1,11 +1,11 @@ from typing import Optional -from langchain.chains.combine_documents.base import BaseCombineDocumentsChain from langchain.chains.retrieval_qa.base import RetrievalQA from langchain_core.documents import Document -from langflow.field_typing import BaseMemory, BaseRetriever, Text +from langflow.field_typing import BaseLanguageModel, BaseMemory, BaseRetriever, Text from langflow.interface.custom.custom_component import CustomComponent +from langflow.schema.schema import Record class RetrievalQAComponent(CustomComponent): @@ -14,7 +14,8 @@ class RetrievalQAComponent(CustomComponent): def build_config(self): return { - "combine_documents_chain": {"display_name": "Combine Documents Chain"}, + "llm": {"display_name": "LLM"}, + "chain_type": {"display_name": "Chain Type", "options": ["Stuff", "Map Reduce", "Refine", "Map Rerank"]}, "retriever": {"display_name": "Retriever"}, "memory": {"display_name": "Memory", "required": False}, "input_key": {"display_name": "Input Key", "advanced": True}, @@ -22,13 +23,14 @@ class RetrievalQAComponent(CustomComponent): "return_source_documents": {"display_name": "Return Source Documents"}, "input_value": { "display_name": "Input", - "input_types": ["Text", "Document"], + "input_types": ["Record", "Document"], }, } def build( self, - combine_documents_chain: BaseCombineDocumentsChain, + llm: BaseLanguageModel, + chain_type: str, retriever: BaseRetriever, input_value: str = "", memory: Optional[BaseMemory] = None, @@ -36,8 +38,10 @@ class RetrievalQAComponent(CustomComponent): output_key: str = "result", return_source_documents: bool = True, ) -> Text: - runnable = RetrievalQA( - combine_documents_chain=combine_documents_chain, + chain_type = chain_type.lower().replace(" ", "_") + runnable = RetrievalQA.from_chain_type( + llm=llm, + chain_type=chain_type, retriever=retriever, memory=memory, input_key=input_key, @@ -46,6 +50,8 @@ class RetrievalQAComponent(CustomComponent): ) if isinstance(input_value, Document): input_value = input_value.page_content + if isinstance(input_value, Record): + input_value = input_value.get_text() self.status = runnable result = runnable.invoke({input_key: input_value}) result = result.content if hasattr(result, "content") else result diff --git a/src/backend/base/langflow/components/chains/RetrievalQAWithSourcesChain.py b/src/backend/base/langflow/components/chains/RetrievalQAWithSourcesChain.py index 0c98a16ec..75a9131f5 100644 --- a/src/backend/base/langflow/components/chains/RetrievalQAWithSourcesChain.py +++ b/src/backend/base/langflow/components/chains/RetrievalQAWithSourcesChain.py @@ -16,7 +16,7 @@ class RetrievalQAWithSourcesChainComponent(CustomComponent): "llm": {"display_name": "LLM"}, "chain_type": { "display_name": "Chain Type", - "options": ["stuff", "map_reduce", "map_rerank", "refine"], + "options": ["Stuff", "Map Reduce", "Refine", "Map Rerank"], "info": "The type of chain to use to combined Documents.", }, "memory": {"display_name": "Memory"}, @@ -37,6 +37,7 @@ class RetrievalQAWithSourcesChainComponent(CustomComponent): memory: Optional[BaseMemory] = None, return_source_documents: Optional[bool] = True, ) -> Text: + chain_type = chain_type.lower().replace(" ", "_") runnable = RetrievalQAWithSourcesChain.from_chain_type( llm=llm, chain_type=chain_type, From 4428c36bfa58926bebd47fad889cb2b1c34bf043 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sat, 30 Mar 2024 17:33:00 -0300 Subject: [PATCH 2/2] Refactor SubFlowComponent display name --- src/backend/base/langflow/components/experimental/SubFlow.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/backend/base/langflow/components/experimental/SubFlow.py b/src/backend/base/langflow/components/experimental/SubFlow.py index af309fa88..72e07c8e8 100644 --- a/src/backend/base/langflow/components/experimental/SubFlow.py +++ b/src/backend/base/langflow/components/experimental/SubFlow.py @@ -1,19 +1,19 @@ from typing import Any, List, Optional -from langflow.helpers.flow import get_flow_inputs from loguru import logger from langflow.custom import CustomComponent from langflow.graph.graph.base import Graph from langflow.graph.schema import ResultData, RunOutputs from langflow.graph.vertex.base import Vertex +from langflow.helpers.flow import get_flow_inputs from langflow.schema import Record from langflow.schema.dotdict import dotdict from langflow.template.field.base import TemplateField class SubFlowComponent(CustomComponent): - display_name = "SubFlow" + display_name = "Sub Flow" description = "Dynamically Generates a Component from a Flow. The output is a list of records with keys 'result' and 'message'." beta: bool = True field_order = ["flow_name"]