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

Add "retriever" input to the build method

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-26 16:17:05 -03:00
commit d7f1a5cc8c
1 changed files with 3 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

4
src/backend/langflow/components/chains/RetrievalQAWithSourcesChain.py
View file

@ -21,10 +21,12 @@ class RetrievalQAWithSourcesChainComponent(CustomComponent):
}, },
"memory": {"display_name": "Memory"}, "memory": {"display_name": "Memory"},
"return_source_documents": {"display_name": "Return Source Documents"}, "return_source_documents": {"display_name": "Return Source Documents"},
"retriever": {"display_name": "Retriever"},
} }
def build( def build(
self, self,
inputs: str,
retriever: BaseRetriever, retriever: BaseRetriever,
llm: BaseLanguageModel, llm: BaseLanguageModel,
chain_type: str, chain_type: str,
@ -50,7 +52,7 @@ class RetrievalQAWithSourcesChainComponent(CustomComponent):
references_str = "" references_str = ""
if return_source_documents: if return_source_documents:
references_str = self.create_references_from_records(records) references_str = self.create_references_from_records(records)
result_str = result.get("result") result_str = result.get("answer")
final_result = "\n".join([result_str, references_str]) final_result = "\n".join([result_str, references_str])
self.status = final_result self.status = final_result
return final_result return final_result