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

Refactor RetrievalQAComponent and RetrievalQAWithSourcesChainComponent

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-30 17:32:44 -03:00
commit 39e7a92f45
2 changed files with 15 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

20
src/backend/base/langflow/components/chains/RetrievalQA.py
View file

@ -1,11 +1,11 @@
from typing import Optional from typing import Optional
from langchain.chains.combine_documents.base import BaseCombineDocumentsChain
from langchain.chains.retrieval_qa.base import RetrievalQA from langchain.chains.retrieval_qa.base import RetrievalQA
from langchain_core.documents import Document 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.interface.custom.custom_component import CustomComponent
from langflow.schema.schema import Record
class RetrievalQAComponent(CustomComponent): class RetrievalQAComponent(CustomComponent):
@ -14,7 +14,8 @@ class RetrievalQAComponent(CustomComponent):
def build_config(self): def build_config(self):
return { 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"}, "retriever": {"display_name": "Retriever"},
"memory": {"display_name": "Memory", "required": False}, "memory": {"display_name": "Memory", "required": False},
"input_key": {"display_name": "Input Key", "advanced": True}, "input_key": {"display_name": "Input Key", "advanced": True},
@ -22,13 +23,14 @@ class RetrievalQAComponent(CustomComponent):
"return_source_documents": {"display_name": "Return Source Documents"}, "return_source_documents": {"display_name": "Return Source Documents"},
"input_value": { "input_value": {
"display_name": "Input", "display_name": "Input",
"input_types": ["Text", "Document"], "input_types": ["Record", "Document"],
}, },
} }
def build( def build(
self, self,
combine_documents_chain: BaseCombineDocumentsChain, llm: BaseLanguageModel,
chain_type: str,
retriever: BaseRetriever, retriever: BaseRetriever,
input_value: str = "", input_value: str = "",
memory: Optional[BaseMemory] = None, memory: Optional[BaseMemory] = None,
@ -36,8 +38,10 @@ class RetrievalQAComponent(CustomComponent):
output_key: str = "result", output_key: str = "result",
return_source_documents: bool = True, return_source_documents: bool = True,
) -> Text: ) -> Text:
runnable = RetrievalQA( chain_type = chain_type.lower().replace(" ", "_")
combine_documents_chain=combine_documents_chain, runnable = RetrievalQA.from_chain_type(
llm=llm,
chain_type=chain_type,
retriever=retriever, retriever=retriever,
memory=memory, memory=memory,
input_key=input_key, input_key=input_key,
@ -46,6 +50,8 @@ class RetrievalQAComponent(CustomComponent):
) )
if isinstance(input_value, Document): if isinstance(input_value, Document):
input_value = input_value.page_content input_value = input_value.page_content
if isinstance(input_value, Record):
input_value = input_value.get_text()
self.status = runnable self.status = runnable
result = runnable.invoke({input_key: input_value}) result = runnable.invoke({input_key: input_value})
result = result.content if hasattr(result, "content") else result result = result.content if hasattr(result, "content") else result

3
src/backend/base/langflow/components/chains/RetrievalQAWithSourcesChain.py
View file

@ -16,7 +16,7 @@ class RetrievalQAWithSourcesChainComponent(CustomComponent):
"llm": {"display_name": "LLM"}, "llm": {"display_name": "LLM"},
"chain_type": { "chain_type": {
"display_name": "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.", "info": "The type of chain to use to combined Documents.",
}, },
"memory": {"display_name": "Memory"}, "memory": {"display_name": "Memory"},
@ -37,6 +37,7 @@ class RetrievalQAWithSourcesChainComponent(CustomComponent):
memory: Optional[BaseMemory] = None, memory: Optional[BaseMemory] = None,
return_source_documents: Optional[bool] = True, return_source_documents: Optional[bool] = True,
) -> Text: ) -> Text:
chain_type = chain_type.lower().replace(" ", "_")
runnable = RetrievalQAWithSourcesChain.from_chain_type( runnable = RetrievalQAWithSourcesChain.from_chain_type(
llm=llm, llm=llm,
chain_type=chain_type, chain_type=chain_type,