From eaa2623fc6d9e7dbd84d7aa0612abb3b426f47a3 Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Thu, 9 Mar 2023 15:00:36 -0300 Subject: [PATCH] fix: loading from json works now --- langflow/__init__.py | 1 + .../langflow_backend/interface/loading.py | 21 +++++++++---- .../backend/langflow_backend/interface/run.py | 31 ++++++++----------- .../backend/langflow_backend/utils/payload.py | 14 ++++----- 4 files changed, 35 insertions(+), 32 deletions(-) diff --git a/langflow/__init__.py b/langflow/__init__.py index e69de29bb..1ece2c9f9 100644 --- a/langflow/__init__.py +++ b/langflow/__init__.py @@ -0,0 +1 @@ +from langflow_backend.interface.loading import load_flow_from_json diff --git a/langflow/backend/langflow_backend/interface/loading.py b/langflow/backend/langflow_backend/interface/loading.py index f23aa3599..4b4a779ad 100644 --- a/langflow/backend/langflow_backend/interface/loading.py +++ b/langflow/backend/langflow_backend/interface/loading.py @@ -12,13 +12,22 @@ def load_flow_from_json(path: str): with open(path, "r") as f: flow_graph = json.load(f) data_graph = flow_graph["data"] - edges = data_graph["edges"] - nodes = replace_zero_shot_prompt_with_prompt_template(data_graph["nodes"]) - root = payload.get_root_node(data_graph) - extracted_json = payload.build_json(root, nodes, edges) + extracted_json = extract_json(data_graph) return load_langchain_type_from_config(config=extracted_json) +def extract_json(data_graph): + nodes = data_graph["nodes"] + # Substitute ZeroShotPrompt with PromptTemplate + nodes = replace_zero_shot_prompt_with_prompt_template(nodes) + # Add input variables + nodes = payload.extract_input_variables(nodes) + # Nodes, edges and root node + edges = data_graph["edges"] + root = payload.get_root_node(nodes, edges) + return payload.build_json(root, nodes, edges) + + def replace_zero_shot_prompt_with_prompt_template(nodes): """Replace ZeroShotPrompt with PromptTemplate""" for node in nodes: @@ -40,9 +49,9 @@ def load_langchain_type_from_config(config: Dict[str, Any]): # Get type list type_list = get_type_list() if config["_type"] in type_list["agents"]: - return load_agent_executor_from_config(config).run + return load_agent_executor_from_config(config) elif config["_type"] in type_list["chains"]: - return load_chain_from_config(config).run + return load_chain_from_config(config) elif config["_type"] in type_list["llms"]: return load_llm_from_config(config) else: diff --git a/langflow/backend/langflow_backend/interface/run.py b/langflow/backend/langflow_backend/interface/run.py index 541663236..667679966 100644 --- a/langflow/backend/langflow_backend/interface/run.py +++ b/langflow/backend/langflow_backend/interface/run.py @@ -2,11 +2,7 @@ import contextlib import io import re from typing import Any, Dict -from langflow_backend.interface.loading import ( - load_langchain_type_from_config, - replace_zero_shot_prompt_with_prompt_template, -) -from langflow_backend.utils import payload +from langflow_backend.interface import loading def process_data_graph(data_graph: Dict[str, Any]): @@ -14,17 +10,10 @@ def process_data_graph(data_graph: Dict[str, Any]): Process data graph by extracting input variables and replacing ZeroShotPrompt with PromptTemplate,then run the graph and return the result and thought. """ - nodes = data_graph["nodes"] - # Substitute ZeroShotPrompt with PromptTemplate - nodes = replace_zero_shot_prompt_with_prompt_template(nodes) - # Add input variables - data_graph = payload.extract_input_variables(data_graph) - # Nodes, edges and root node - message = data_graph["message"] - edges = data_graph["edges"] - root = payload.get_root_node(data_graph) - extracted_json = payload.build_json(root, nodes, edges) + extracted_json = loading.extract_json(data_graph) + + message = data_graph["message"] # Process json result, thought = get_result_and_thought(extracted_json, message) @@ -42,11 +31,17 @@ def process_data_graph(data_graph: Dict[str, Any]): def get_result_and_thought(extracted_json: Dict[str, Any], message: str): """Get result and thought from extracted json""" - # Get type list try: - loaded = load_langchain_type_from_config(config=extracted_json) + loaded_langchain = loading.load_langchain_type_from_config( + config=extracted_json + ) with io.StringIO() as output_buffer, contextlib.redirect_stdout(output_buffer): - result = loaded(message) + result = loaded_langchain(message) + result = ( + result.get(loaded_langchain.output_keys[0]) + if isinstance(result, dict) + else result + ) thought = output_buffer.getvalue() except Exception as e: result = f"Error: {str(e)}" diff --git a/langflow/backend/langflow_backend/utils/payload.py b/langflow/backend/langflow_backend/utils/payload.py index 4d183ba82..e7d38139e 100644 --- a/langflow/backend/langflow_backend/utils/payload.py +++ b/langflow/backend/langflow_backend/utils/payload.py @@ -2,12 +2,12 @@ import contextlib import re -def extract_input_variables(data): +def extract_input_variables(nodes): """ Extracts input variables from the template and adds them to the input_variables field. """ - for node in data["nodes"]: + for node in nodes: with contextlib.suppress(Exception): if "input_variables" in node["data"]["node"]["template"]: if node["data"]["node"]["template"]["_type"] == "prompt": @@ -24,17 +24,15 @@ def extract_input_variables(data): else: variables = [] node["data"]["node"]["template"]["input_variables"]["value"] = variables - return data + return nodes -def get_root_node(data): +def get_root_node(nodes, edges): """ Returns the root node of the template. """ - incoming_edges = {edge["source"] for edge in data["edges"]} - return next( - (node for node in data["nodes"] if node["id"] not in incoming_edges), None - ) + incoming_edges = {edge["source"] for edge in edges} + return next((node for node in nodes if node["id"] not in incoming_edges), None) def build_json(root, nodes, edges):