Fix bugs, add tests, and refactor code (#138)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langflow"
-version = "0.0.54"
+version = "0.0.55"
 description = "A Python package with a built-in web application"
 authors = ["Logspace <[email protected]>"]
 maintainers = [

src/backend/langflow/api/base.py

@@ -1,6 +1,7 @@
-from langflow.graph.utils import extract_input_variables_from_prompt
 from pydantic import BaseModel, validator
 
+from langflow.graph.utils import extract_input_variables_from_prompt
+
 
 class Code(BaseModel):
     code: str

src/backend/langflow/api/validate.py

@@ -7,7 +7,6 @@
     PromptValidationResponse,
     validate_prompt,
 )
-from langflow.graph.utils import extract_input_variables_from_prompt
 from langflow.utils.logger import logger
 from langflow.utils.validate import validate_code
 

src/backend/langflow/config.yaml

@@ -24,15 +24,15 @@ prompts:
 
 llms:
   - OpenAI
-  - AzureOpenAI
+  # - AzureOpenAI
   - ChatOpenAI
   - HuggingFaceHub
 
 tools:
-  # - Search
+  - Search
   - PAL-MATH
   - Calculator
-  # - Serper Search
+  - Serper Search
   - Tool
   - PythonFunction
   - JsonSpec

src/backend/langflow/graph/base.py

@@ -4,6 +4,7 @@
 #   - Build each inner agent first, then build the outer agent
 
 import types
+import warnings
 from copy import deepcopy
 from typing import Any, Dict, List, Optional
 
@@ -119,7 +120,13 @@ def _build_params(self):
                     params[key] = edges[0].source
 
             elif value["required"] or value.get("value"):
-                params[key] = value["value"]
+                # If value does not have value this still passes
+                # but then gives a keyError
+                # so we need to check if value has value
+                new_value = value.get("value")
+                if new_value is None:
+                    warnings.warn(f"Value for {key} in {self.node_type} is None. ")
+                params[key] = new_value
 
         # Add _type to params
         self.params = params

src/backend/langflow/graph/constants.py

@@ -1 +1 @@
-DIRECT_TYPES = ["str", "bool", "code", "int", "float", "Any"]
+DIRECT_TYPES = ["str", "bool", "code", "int", "float", "Any", "prompt"]

src/backend/langflow/graph/graph.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Union
+from typing import Dict, List, Type, Union
 
 from langflow.graph.base import Edge, Node
 from langflow.graph.nodes import (
@@ -15,13 +15,12 @@
 from langflow.interface.agents.base import agent_creator
 from langflow.interface.chains.base import chain_creator
 from langflow.interface.llms.base import llm_creator
+from langflow.interface.memories.base import memory_creator
 from langflow.interface.prompts.base import prompt_creator
 from langflow.interface.toolkits.base import toolkits_creator
 from langflow.interface.tools.base import tool_creator
 from langflow.interface.tools.constants import FILE_TOOLS
-from langflow.interface.tools.util import get_tools_dict
 from langflow.interface.wrappers.base import wrapper_creator
-from langflow.interface.memories.base import memory_creator
 from langflow.utils import payload
 
 
@@ -108,45 +107,36 @@ def _build_edges(self) -> List[Edge]:
             edges.append(Edge(source, target))
         return edges
 
+    def _get_node_class(self, node_type: str, node_lc_type: str) -> Type[Node]:
+        node_type_map: Dict[str, Type[Node]] = {
+            **{t: PromptNode for t in prompt_creator.to_list()},
+            **{t: AgentNode for t in agent_creator.to_list()},
+            **{t: ChainNode for t in chain_creator.to_list()},
+            **{t: ToolNode for t in tool_creator.to_list()},
+            **{t: ToolkitNode for t in toolkits_creator.to_list()},
+            **{t: WrapperNode for t in wrapper_creator.to_list()},
+            **{t: LLMNode for t in llm_creator.to_list()},
+            **{t: MemoryNode for t in memory_creator.to_list()},
+        }
+
+        if node_type in FILE_TOOLS:
+            return FileToolNode
+        if node_type in node_type_map:
+            return node_type_map[node_type]
+        if node_lc_type in node_type_map:
+            return node_type_map[node_lc_type]
+        return Node
+
     def _build_nodes(self) -> List[Node]:
         nodes: List[Node] = []
         for node in self._nodes:
             node_data = node["data"]
             node_type: str = node_data["type"]  # type: ignore
             node_lc_type: str = node_data["node"]["template"]["_type"]  # type: ignore
 
-            if node_type in prompt_creator.to_list():
-                nodes.append(PromptNode(node))
-            elif (
-                node_type in agent_creator.to_list()
-                or node_lc_type in agent_creator.to_list()
-            ):
-                nodes.append(AgentNode(node))
-            elif node_type in chain_creator.to_list():
-                nodes.append(ChainNode(node))
-            elif (
-                node_type in tool_creator.to_list()
-                or node_lc_type in get_tools_dict().keys()
-            ):
-                if node_type in FILE_TOOLS:
-                    nodes.append(FileToolNode(node))
-                nodes.append(ToolNode(node))
-            elif node_type in toolkits_creator.to_list():
-                nodes.append(ToolkitNode(node))
-            elif node_type in wrapper_creator.to_list():
-                nodes.append(WrapperNode(node))
-            elif (
-                node_type in llm_creator.to_list()
-                or node_lc_type in llm_creator.to_list()
-            ):
-                nodes.append(LLMNode(node))
-            elif (
-                node_type in memory_creator.to_list()
-                or node_lc_type in memory_creator.to_list()
-            ):
-                nodes.append(MemoryNode(node))
-            else:
-                nodes.append(Node(node))
+            NodeClass = self._get_node_class(node_type, node_lc_type)
+            nodes.append(NodeClass(node))
+
         return nodes
 
     def get_children_by_node_type(self, node: Node, node_type: str) -> List[Node]:

src/backend/langflow/graph/nodes.py

@@ -55,7 +55,10 @@ def build(
         tools: Optional[Union[List[Node], List[ToolNode]]] = None,
     ) -> Any:
         if not self._built or force:
-            if "input_variables" not in self.params:
+            if (
+                "input_variables" not in self.params
+                or self.params["input_variables"] is None
+            ):
                 self.params["input_variables"] = []
             # Check if it is a ZeroShotPrompt and needs a tool
             if "ShotPrompt" in self.node_type:
@@ -75,7 +78,6 @@ def build(
             for param in prompt_params:
                 prompt_text = self.params[param]
                 variables = extract_input_variables_from_prompt(prompt_text)
-
                 self.params["input_variables"].extend(variables)
             self.params["input_variables"] = list(set(self.params["input_variables"]))
 

src/backend/langflow/interface/chains/base.py

@@ -1,9 +1,10 @@
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Type
 
 from langflow.custom.customs import get_custom_nodes
 from langflow.interface.base import LangChainTypeCreator
 from langflow.interface.custom_lists import chain_type_to_cls_dict
 from langflow.settings import settings
+from langflow.template.nodes import ChainFrontendNode
 from langflow.utils.util import build_template_from_class
 
 # Assuming necessary imports for Field, Template, and FrontendNode classes
@@ -12,6 +13,10 @@
 class ChainCreator(LangChainTypeCreator):
     type_name: str = "chains"
 
+    @property
+    def frontend_node_class(self) -> Type[ChainFrontendNode]:
+        return ChainFrontendNode
+
     @property
     def type_to_loader_dict(self) -> Dict:
         if self.type_dict is None:

src/backend/langflow/interface/chains/custom.py

@@ -19,21 +19,21 @@ class BaseCustomChain(ConversationChain):
 
     template: Optional[str]
 
-    ai_prefix_key: Optional[str]
+    ai_prefix_value: Optional[str]
     """Field to use as the ai_prefix. It needs to be set and has to be in the template"""
 
     @root_validator(pre=False)
     def build_template(cls, values):
         format_dict = {}
         input_variables = extract_input_variables_from_prompt(values["template"])
 
-        if values.get("ai_prefix_key", None) is None:
-            values["ai_prefix_key"] = values["memory"].ai_prefix
+        if values.get("ai_prefix_value", None) is None:
+            values["ai_prefix_value"] = values["memory"].ai_prefix
 
         for key in input_variables:
             new_value = values.get(key, f"{{{key}}}")
             format_dict[key] = new_value
-            if key == values.get("ai_prefix_key", None):
+            if key == values.get("ai_prefix_value", None):
                 values["memory"].ai_prefix = new_value
 
         values["template"] = values["template"].format(**format_dict)
@@ -62,7 +62,7 @@ class SeriesCharacterChain(BaseCustomChain):
 Human: {input}
 {character}:"""
     memory: BaseMemory = Field(default_factory=ConversationBufferMemory)
-    ai_prefix_key: Optional[str] = "character"
+    ai_prefix_value: Optional[str] = "character"
     """Default memory store."""
 
 

src/backend/langflow/interface/tools/base.py

@@ -73,25 +73,36 @@ def get_signature(self, name: str) -> Optional[Dict]:
         base_classes = ["Tool"]
         all_tools = {}
         for tool in self.type_to_loader_dict.keys():
-            if tool_params := get_tool_params(get_tool_by_name(tool)):
+            tool_fcn = get_tool_by_name(tool)
+            if tool_params := get_tool_params(tool_fcn):
                 tool_name = tool_params.get("name") or str(tool)
-                all_tools[tool_name] = {"type": tool, "params": tool_params}
+                all_tools[tool_name] = {
+                    "type": tool,
+                    "params": tool_params,
+                    "fcn": tool_fcn,
+                }
 
         # Raise error if name is not in tools
         if name not in all_tools.keys():
             raise ValueError("Tool not found")
 
         tool_type: str = all_tools[name]["type"]  # type: ignore
 
-        if tool_type in _BASE_TOOLS:
+        if all_tools[tool_type]["fcn"] in _BASE_TOOLS.values():
             params = []
-        elif tool_type in _LLM_TOOLS:
+        elif all_tools[tool_type]["fcn"] in _LLM_TOOLS.values():
             params = ["llm"]
-        elif tool_type in _EXTRA_LLM_TOOLS:
-            _, extra_keys = _EXTRA_LLM_TOOLS[tool_type]
+        elif all_tools[tool_type]["fcn"] in [
+            val[0] for val in _EXTRA_LLM_TOOLS.values()
+        ]:
+            n_dict = {val[0]: val[1] for val in _EXTRA_LLM_TOOLS.values()}
+            extra_keys = n_dict[all_tools[tool_type]["fcn"]]
             params = ["llm"] + extra_keys
-        elif tool_type in _EXTRA_OPTIONAL_TOOLS:
-            _, extra_keys = _EXTRA_OPTIONAL_TOOLS[tool_type]
+        elif all_tools[tool_type]["fcn"] in [
+            val[0] for val in _EXTRA_OPTIONAL_TOOLS.values()
+        ]:
+            n_dict = {val[0]: val[1] for val in _EXTRA_OPTIONAL_TOOLS.values()}  # type: ignore
+            extra_keys = n_dict[all_tools[tool_type]["fcn"]]
             params = extra_keys
         elif tool_type == "Tool":
             params = ["name", "description", "func"]
@@ -104,7 +115,6 @@ def get_signature(self, name: str) -> Optional[Dict]:
         elif tool_type in FILE_TOOLS:
             params = all_tools[name]["params"]  # type: ignore
             base_classes += [name]
-
         else:
             params = []
 

src/backend/langflow/template/base.py

@@ -179,12 +179,13 @@ def format_field(field: TemplateField, name: Optional[str] = None) -> None:
             (field.required and key not in ["input_variables"])
             or key in FORCE_SHOW_FIELDS
             or "api" in key
-            or "key" in key
+            or ("key" in key and "input" not in key and "output" not in key)
         )
 
         # Add password field
-        field.password = any(
-            text in key.lower() for text in {"password", "token", "api", "key"}
+        field.password = (
+            any(text in key.lower() for text in {"password", "token", "api", "key"})
+            and field.show
         )
 
         # Add multline

src/backend/langflow/template/nodes.py

@@ -289,3 +289,13 @@ def format_field(field: TemplateField, name: Optional[str] = None) -> None:
             field.field_type = "int"
             field.value = 10
             field.display_name = "Memory Size"
+
+
+class ChainFrontendNode(FrontendNode):
+    @staticmethod
+    def format_field(field: TemplateField, name: Optional[str] = None) -> None:
+        FrontendNode.format_field(field, name)
+
+        if "key" in field.name:
+            field.password = False
+            field.show = False