From acb542ab59692468f4819c5cf42a9ebe2270a90d Mon Sep 17 00:00:00 2001 From: Gabriel Almeida Date: Thu, 6 Apr 2023 10:17:12 -0300 Subject: [PATCH 1/4] refact: more maintable build_nodes --- src/backend/langflow/graph/graph.py | 57 ++++++++++++----------------- 1 file changed, 24 insertions(+), 33 deletions(-) diff --git a/src/backend/langflow/graph/graph.py b/src/backend/langflow/graph/graph.py index 76cab071d..1e4555fc8 100644 --- a/src/backend/langflow/graph/graph.py +++ b/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 ( @@ -108,6 +108,26 @@ class Graph: edges.append(Edge(source, target)) return edges + def _get_node_class(self, node_type: str, node_lc_type: str) -> Type[Node]: + node_type_map = { + **{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 node_type_map: + return node_type_map[node_type] + if node_lc_type in node_type_map: + return node_type_map[node_lc_type] + if node_type in FILE_TOOLS: + return FileToolNode + return Node + def _build_nodes(self) -> List[Node]: nodes: List[Node] = [] for node in self._nodes: @@ -115,38 +135,9 @@ class Graph: 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]: From d921983a00670e5b0c5c2a3bee084a7296b988ea Mon Sep 17 00:00:00 2001 From: Gabriel Almeida Date: Thu, 6 Apr 2023 15:40:16 -0300 Subject: [PATCH 2/4] fix: check for specific nodes first --- src/backend/langflow/graph/graph.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/backend/langflow/graph/graph.py b/src/backend/langflow/graph/graph.py index 1e4555fc8..81c488881 100644 --- a/src/backend/langflow/graph/graph.py +++ b/src/backend/langflow/graph/graph.py @@ -120,12 +120,12 @@ class Graph: **{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] - if node_type in FILE_TOOLS: - return FileToolNode return Node def _build_nodes(self) -> List[Node]: From 16276f18dd126a908bc57617f80b50c6b4ab3d2e Mon Sep 17 00:00:00 2001 From: Gabriel Almeida Date: Thu, 6 Apr 2023 15:46:20 -0300 Subject: [PATCH 3/4] fix linting --- src/backend/langflow/api/validate.py | 1 - src/backend/langflow/graph/graph.py | 3 +-- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/src/backend/langflow/api/validate.py b/src/backend/langflow/api/validate.py index 15600f9f7..a60bcc506 100644 --- a/src/backend/langflow/api/validate.py +++ b/src/backend/langflow/api/validate.py @@ -7,7 +7,6 @@ from langflow.api.base import ( 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 diff --git a/src/backend/langflow/graph/graph.py b/src/backend/langflow/graph/graph.py index 81c488881..57b9082bb 100644 --- a/src/backend/langflow/graph/graph.py +++ b/src/backend/langflow/graph/graph.py @@ -19,7 +19,6 @@ 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 @@ -109,7 +108,7 @@ class Graph: return edges def _get_node_class(self, node_type: str, node_lc_type: str) -> Type[Node]: - node_type_map = { + 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()}, From 3a4bfc2e366d0b30bf4f666225b7b8a6aa5f119f Mon Sep 17 00:00:00 2001 From: Gabriel Almeida Date: Thu, 6 Apr 2023 16:18:13 -0300 Subject: [PATCH 4/4] formatting --- src/backend/langflow/api/base.py | 3 ++- src/backend/langflow/graph/graph.py | 2 +- tests/test_prompts_template.py | 4 ++-- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/backend/langflow/api/base.py b/src/backend/langflow/api/base.py index 4ae313fd3..084e04d65 100644 --- a/src/backend/langflow/api/base.py +++ b/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 diff --git a/src/backend/langflow/graph/graph.py b/src/backend/langflow/graph/graph.py index 57b9082bb..84b94a6a5 100644 --- a/src/backend/langflow/graph/graph.py +++ b/src/backend/langflow/graph/graph.py @@ -15,12 +15,12 @@ from langflow.graph.nodes import ( 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.wrappers.base import wrapper_creator -from langflow.interface.memories.base import memory_creator from langflow.utils import payload diff --git a/tests/test_prompts_template.py b/tests/test_prompts_template.py index 24729cb63..caa30821c 100644 --- a/tests/test_prompts_template.py +++ b/tests/test_prompts_template.py @@ -167,7 +167,7 @@ def test_zero_shot_prompt(client: TestClient): "placeholder": "", "show": True, "multiline": True, - "value": "Answer the following questions as best you can. You have access to the following tools:", # noqa: E501 + "value": "Answer the following questions as best you can. You have access to the following tools:", # noqa: E501 "password": False, "name": "prefix", "type": "str", @@ -189,7 +189,7 @@ def test_zero_shot_prompt(client: TestClient): "placeholder": "", "show": False, "multiline": False, - "value": "Use the following format:\n\nQuestion: the input question you must answer\nThought: you should always think about what to do\nAction: the action to take, should be one of [{tool_names}]\nAction Input: the input to the action\nObservation: the result of the action\n... (this Thought/Action/Action Input/Observation can repeat N times)\nThought: I now know the final answer\nFinal Answer: the final answer to the original input question", # noqa: E501 + "value": "Use the following format:\n\nQuestion: the input question you must answer\nThought: you should always think about what to do\nAction: the action to take, should be one of [{tool_names}]\nAction Input: the input to the action\nObservation: the result of the action\n... (this Thought/Action/Action Input/Observation can repeat N times)\nThought: I now know the final answer\nFinal Answer: the final answer to the original input question", # noqa: E501 "password": False, "name": "format_instructions", "type": "str",