From 89286c64563b757a3974f4875ecba2c50caa05d1 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 26 Jun 2023 10:52:53 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A8=20refactor(utils.py):=20extract=20?= =?UTF-8?q?build=5Finput=5Fkeys=5Fresponse=20function=20to=20improve=20cod?= =?UTF-8?q?e=20reusability=20=F0=9F=9A=80=20feat(chat.py):=20use=20build?= =?UTF-8?q?=5Finput=5Fkeys=5Fresponse=20function=20to=20build=20input=20ke?= =?UTF-8?q?ys=20response=20The=20build=5Finput=5Fkeys=5Fresponse=20functio?= =?UTF-8?q?n=20was=20extracted=20from=20the=20remove=5Fapi=5Fkeys=20functi?= =?UTF-8?q?on=20to=20improve=20code=20reusability.=20It=20builds=20the=20i?= =?UTF-8?q?nput=20keys=20response=20object,=20which=20contains=20the=20inp?= =?UTF-8?q?ut=20keys=20and=20memory=20keys=20of=20a=20langchain=20object.?= =?UTF-8?q?=20The=20chat.py=20file=20now=20uses=20this=20function=20to=20b?= =?UTF-8?q?uild=20the=20input=20keys=20response=20object=20instead=20of=20?= =?UTF-8?q?building=20it=20manually.=20This=20improves=20code=20readabilit?= =?UTF-8?q?y=20and=20reduces=20code=20duplication.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/langflow/api/utils.py | 23 +++++++++++++++++++++++ src/backend/langflow/api/v1/chat.py | 5 ++--- 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/src/backend/langflow/api/utils.py b/src/backend/langflow/api/utils.py index bfd9e3da5..cd7daa5d2 100644 --- a/src/backend/langflow/api/utils.py +++ b/src/backend/langflow/api/utils.py @@ -22,3 +22,26 @@ def remove_api_keys(flow: dict): value["value"] = None return flow + + +def build_input_keys_response(langchain_object): + """Build the input keys response.""" + input_keys_response = { + "input_keys": langchain_object.input_keys, + "memory_keys": [], + } + # If the object has memory, that memory will have a memory_variables attribute + # memory variables should be removed from the input keys + if hasattr(langchain_object, "memory") and hasattr( + langchain_object.memory, "memory_variables" + ): + # Remove memory variables from input keys + input_keys_response["input_keys"] = [ + key + for key in input_keys_response["input_keys"] + if key not in langchain_object.memory.memory_variables + ] + # Add memory variables to memory_keys + input_keys_response["memory_keys"] = langchain_object.memory.memory_variables + + return input_keys_response diff --git a/src/backend/langflow/api/v1/chat.py b/src/backend/langflow/api/v1/chat.py index 666c5d3cd..5054bf107 100644 --- a/src/backend/langflow/api/v1/chat.py +++ b/src/backend/langflow/api/v1/chat.py @@ -6,6 +6,7 @@ from fastapi import ( status, ) from fastapi.responses import StreamingResponse +from langflow.api.utils import build_input_keys_response from langflow.api.v1.schemas import BuiltResponse, InitResponse, StreamData from langflow.chat.manager import ChatManager @@ -123,9 +124,7 @@ async def stream_build(flow_id: str): langchain_object = graph.build() # Now we need to check the input_keys to send them to the client if hasattr(langchain_object, "input_keys"): - input_keys_response = { - "input_keys": langchain_object.input_keys, - } + input_keys_response = build_input_keys_response(langchain_object) yield str(StreamData(event="input_keys", data=input_keys_response)) chat_manager.set_cache(flow_id, langchain_object)