jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

refact: more maintable build_nodes

Browse source
This commit is contained in:
Gabriel Almeida 2023-04-06 10:17:12 -03:00
commit acb542ab59
1 changed files with 24 additions and 33 deletions
Download patch file Download diff file Expand all files Collapse all files

57
src/backend/langflow/graph/graph.py
View file

@ -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.base import Edge, Node
from langflow.graph.nodes import ( from langflow.graph.nodes import (
@ -108,6 +108,26 @@ class Graph:
edges.append(Edge(source, target)) edges.append(Edge(source, target))
return edges 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]: def _build_nodes(self) -> List[Node]:
nodes: List[Node] = [] nodes: List[Node] = []
for node in self._nodes: for node in self._nodes:
@ -115,38 +135,9 @@ class Graph:
node_type: str = node_data["type"] # type: ignore node_type: str = node_data["type"] # type: ignore
node_lc_type: str = node_data["node"]["template"]["_type"] # type: ignore node_lc_type: str = node_data["node"]["template"]["_type"] # type: ignore
if node_type in prompt_creator.to_list(): NodeClass = self._get_node_class(node_type, node_lc_type)
nodes.append(PromptNode(node)) nodes.append(NodeClass(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))
return nodes return nodes
def get_children_by_node_type(self, node: Node, node_type: str) -> List[Node]: def get_children_by_node_type(self, node: Node, node_type: str) -> List[Node]: