From 0387de2dd22b53e2047b79cf7ec0e88ffcd41514 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 26 Feb 2024 09:51:41 -0300 Subject: [PATCH] Refactor RetrievalQAComponent to include source documents in the returned result --- src/backend/langflow/components/chains/RetrievalQA.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/backend/langflow/components/chains/RetrievalQA.py b/src/backend/langflow/components/chains/RetrievalQA.py index bfa3b2e11..2fe31353e 100644 --- a/src/backend/langflow/components/chains/RetrievalQA.py +++ b/src/backend/langflow/components/chains/RetrievalQA.py @@ -3,6 +3,7 @@ from typing import Callable, Optional, Union from langchain.chains.combine_documents.base import BaseCombineDocumentsChain from langchain.chains.retrieval_qa.base import BaseRetrievalQA, RetrievalQA from langchain_core.documents import Document + from langflow import CustomComponent from langflow.field_typing import BaseMemory, BaseRetriever, Text @@ -47,4 +48,11 @@ class RetrievalQAComponent(CustomComponent): result = result.content if hasattr(result, "content") else result # Result is a dict with keys "query", "result" and "source_documents" # for now we just return the result - return result.get("result") + records = self.to_records(result.get("source_documents")) + references_str = "" + if return_source_documents: + references_str = self.create_references_from_records(records) + result_str = result.get("result") + final_result = "\n".join([result_str, references_str]) + self.status = final_result + return final_result