Merge branch 'dev' of personal:logspace-ai/langflow into dev

langflow/__init__.py

@@ -0,0 +1 @@
+from langflow_backend.interface.loading import load_flow_from_json

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:

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)}"

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):