refactor: change Node to Vertex

This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-06-01 15:37:11 -03:00
commit 0cef4bd818
10 changed files with 164 additions and 158 deletions

View file

@ -9,7 +9,7 @@ from langflow.api.base import (
PromptValidationResponse,
validate_prompt,
)
from langflow.graph.node.types import VectorStoreNode
from langflow.graph.vertex.types import VectorStoreVertex
from langflow.interface.run import build_graph
from langflow.utils.logger import logger
from langflow.utils.validate import validate_code
@ -49,7 +49,7 @@ def post_validate_node(node_id: str, data: dict):
node = graph.get_node(node_id)
if node is None:
raise ValueError(f"Node {node_id} not found")
if not isinstance(node, VectorStoreNode):
if not isinstance(node, VectorStoreVertex):
node.build()
return json.dumps({"valid": True, "params": str(node._built_object_repr())})
except Exception as e:

View file

@ -1,35 +1,35 @@
from langflow.graph.edge.base import Edge
from langflow.graph.graph.base import Graph
from langflow.graph.node.base import Node
from langflow.graph.node.types import (
AgentNode,
ChainNode,
DocumentLoaderNode,
EmbeddingNode,
LLMNode,
MemoryNode,
PromptNode,
TextSplitterNode,
ToolNode,
ToolkitNode,
VectorStoreNode,
WrapperNode,
from langflow.graph.vertex.base import Vertex
from langflow.graph.vertex.types import (
AgentVertex,
ChainVertex,
DocumentLoaderVertex,
EmbeddingVertex,
LLMVertex,
MemoryVertex,
PromptVertex,
TextSplitterVertex,
ToolVertex,
ToolkitVertex,
VectorStoreVertex,
WrapperVertex,
)
__all__ = [
"Graph",
"Node",
"Vertex",
"Edge",
"AgentNode",
"ChainNode",
"DocumentLoaderNode",
"EmbeddingNode",
"LLMNode",
"MemoryNode",
"PromptNode",
"TextSplitterNode",
"ToolNode",
"ToolkitNode",
"VectorStoreNode",
"WrapperNode",
"AgentVertex",
"ChainVertex",
"DocumentLoaderVertex",
"EmbeddingVertex",
"LLMVertex",
"MemoryVertex",
"PromptVertex",
"TextSplitterVertex",
"ToolVertex",
"ToolkitVertex",
"VectorStoreVertex",
"WrapperVertex",
]

View file

@ -2,13 +2,13 @@ from langflow.utils.logger import logger
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from langflow.graph.node.base import Node
from langflow.graph.vertex.base import Vertex
class Edge:
def __init__(self, source: "Node", target: "Node"):
self.source: "Node" = source
self.target: "Node" = target
def __init__(self, source: "Vertex", target: "Vertex"):
self.source: "Vertex" = source
self.target: "Vertex" = target
self.validate_edge()
def validate_edge(self) -> None:
@ -41,7 +41,7 @@ class Edge:
logger.debug(self.target_reqs)
if no_matched_type:
raise ValueError(
f"Edge between {self.source.node_type} and {self.target.node_type} "
f"Edge between {self.source.vertex_type} and {self.target.vertex_type} "
f"has no matched type"
)

View file

@ -1,12 +1,12 @@
from typing import Dict, List, Type, Union
from langflow.graph.edge.base import Edge
from langflow.graph.graph.constants import NODE_TYPE_MAP
from langflow.graph.node.base import Node
from langflow.graph.node.types import (
FileToolNode,
LLMNode,
ToolkitNode,
from langflow.graph.graph.constants import VERTEX_TYPE_MAP
from langflow.graph.vertex.base import Vertex
from langflow.graph.vertex.types import (
FileToolVertex,
LLMVertex,
ToolkitVertex,
)
from langflow.interface.tools.constants import FILE_TOOLS
from langflow.utils import payload
@ -26,7 +26,7 @@ class Graph:
def _build_graph(self) -> None:
"""Builds the graph from the nodes and edges."""
self.nodes = self._build_nodes()
self.nodes = self._build_vertices()
self.edges = self._build_edges()
for edge in self.edges:
edge.source.add_edge(edge)
@ -43,12 +43,12 @@ class Graph:
llm_node = None
for node in self.nodes:
node._build_params()
if isinstance(node, LLMNode):
if isinstance(node, LLMVertex):
llm_node = node
if llm_node:
for node in self.nodes:
if isinstance(node, ToolkitNode):
if isinstance(node, ToolkitVertex):
node.params["llm"] = llm_node
def _remove_invalid_nodes(self) -> None:
@ -60,23 +60,23 @@ class Graph:
or (len(self.nodes) == 1 and len(self.edges) == 0)
]
def _validate_node(self, node: Node) -> bool:
def _validate_node(self, node: Vertex) -> bool:
"""Validates a node."""
# All nodes that do not have edges are invalid
return len(node.edges) > 0
def get_node(self, node_id: str) -> Union[None, Node]:
def get_node(self, node_id: str) -> Union[None, Vertex]:
"""Returns a node by id."""
return next((node for node in self.nodes if node.id == node_id), None)
def get_nodes_with_target(self, node: Node) -> List[Node]:
def get_nodes_with_target(self, node: Vertex) -> List[Vertex]:
"""Returns the nodes connected to a node."""
connected_nodes: List[Node] = [
connected_nodes: List[Vertex] = [
edge.source for edge in self.edges if edge.target == node
]
return connected_nodes
def build(self) -> List[Node]:
def build(self) -> List[Vertex]:
"""Builds the graph."""
# Get root node
root_node = payload.get_root_node(self)
@ -84,9 +84,9 @@ class Graph:
raise ValueError("No root node found")
return root_node.build()
def get_node_neighbors(self, node: Node) -> Dict[Node, int]:
def get_node_neighbors(self, node: Vertex) -> Dict[Vertex, int]:
"""Returns the neighbors of a node."""
neighbors: Dict[Node, int] = {}
neighbors: Dict[Vertex, int] = {}
for edge in self.edges:
if edge.source == node:
neighbor = edge.target
@ -117,28 +117,30 @@ class Graph:
edges.append(Edge(source, target))
return edges
def _get_node_class(self, node_type: str, node_lc_type: str) -> Type[Node]:
def _get_vertex_class(self, node_type: str, node_lc_type: str) -> Type[Vertex]:
"""Returns the node class based on the node type."""
if node_type in FILE_TOOLS:
return FileToolNode
if node_type in NODE_TYPE_MAP:
return NODE_TYPE_MAP[node_type]
return NODE_TYPE_MAP[node_lc_type] if node_lc_type in NODE_TYPE_MAP else Node
return FileToolVertex
if node_type in VERTEX_TYPE_MAP:
return VERTEX_TYPE_MAP[node_type]
return (
VERTEX_TYPE_MAP[node_lc_type] if node_lc_type in VERTEX_TYPE_MAP else Vertex
)
def _build_nodes(self) -> List[Node]:
"""Builds the nodes of the graph."""
nodes: List[Node] = []
def _build_vertices(self) -> List[Vertex]:
"""Builds the vertices of the graph."""
nodes: List[Vertex] = []
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
NodeClass = self._get_node_class(node_type, node_lc_type)
nodes.append(NodeClass(node))
VertexClass = self._get_vertex_class(node_type, node_lc_type)
nodes.append(VertexClass(node))
return nodes
def get_children_by_node_type(self, node: Node, node_type: str) -> List[Node]:
def get_children_by_node_type(self, node: Vertex, node_type: str) -> List[Vertex]:
"""Returns the children of a node based on the node type."""
children = []
node_types = [node.data["type"]]

View file

@ -1,17 +1,17 @@
from langflow.graph.node.base import Node
from langflow.graph.node.types import (
AgentNode,
ChainNode,
DocumentLoaderNode,
EmbeddingNode,
LLMNode,
MemoryNode,
PromptNode,
TextSplitterNode,
ToolNode,
ToolkitNode,
VectorStoreNode,
WrapperNode,
from langflow.graph.vertex.base import Vertex
from langflow.graph.vertex.types import (
AgentVertex,
ChainVertex,
DocumentLoaderVertex,
EmbeddingVertex,
LLMVertex,
MemoryVertex,
PromptVertex,
TextSplitterVertex,
ToolVertex,
ToolkitVertex,
VectorStoreVertex,
WrapperVertex,
)
from langflow.interface.agents.base import agent_creator
from langflow.interface.chains.base import chain_creator
@ -33,17 +33,17 @@ from typing import Dict, Type
DIRECT_TYPES = ["str", "bool", "code", "int", "float", "Any", "prompt"]
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()},
**{t: EmbeddingNode for t in embedding_creator.to_list()},
**{t: VectorStoreNode for t in vectorstore_creator.to_list()},
**{t: DocumentLoaderNode for t in documentloader_creator.to_list()},
**{t: TextSplitterNode for t in textsplitter_creator.to_list()},
VERTEX_TYPE_MAP: Dict[str, Type[Vertex]] = {
**{t: PromptVertex for t in prompt_creator.to_list()},
**{t: AgentVertex for t in agent_creator.to_list()},
**{t: ChainVertex for t in chain_creator.to_list()},
**{t: ToolVertex for t in tool_creator.to_list()},
**{t: ToolkitVertex for t in toolkits_creator.to_list()},
**{t: WrapperVertex for t in wrapper_creator.to_list()},
**{t: LLMVertex for t in llm_creator.to_list()},
**{t: MemoryVertex for t in memory_creator.to_list()},
**{t: EmbeddingVertex for t in embedding_creator.to_list()},
**{t: VectorStoreVertex for t in vectorstore_creator.to_list()},
**{t: DocumentLoaderVertex for t in documentloader_creator.to_list()},
**{t: TextSplitterVertex for t in textsplitter_creator.to_list()},
}

View file

@ -1,5 +1,5 @@
from langflow.cache import base as cache_utils
from langflow.graph.node.constants import DIRECT_TYPES
from langflow.graph.vertex.constants import DIRECT_TYPES
from langflow.interface import loading
from langflow.interface.listing import ALL_TYPES_DICT
from langflow.utils.logger import logger
@ -17,7 +17,7 @@ if TYPE_CHECKING:
from langflow.graph.edge.base import Edge
class Node:
class Vertex:
def __init__(self, data: Dict, base_type: Optional[str] = None) -> None:
self.id: str = data["id"]
self._data = data
@ -48,12 +48,12 @@ class Node:
]
template_dict = self.data["node"]["template"]
self.node_type = (
self.vertex_type = (
self.data["type"] if "Tool" not in self.output else template_dict["_type"]
)
if self.base_type is None:
for base_type, value in ALL_TYPES_DICT.items():
if self.node_type in value:
if self.vertex_type in value:
self.base_type = base_type
break
@ -113,7 +113,7 @@ class Node:
if value["required"] and not edges:
# If a required parameter is not found, raise an error
raise ValueError(
f"Required input {key} for module {self.node_type} not found"
f"Required input {key} for module {self.vertex_type} not found"
)
elif value["list"]:
# If this is a list parameter, append all sources to a list
@ -128,7 +128,7 @@ class Node:
# 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. ")
warnings.warn(f"Value for {key} in {self.vertex_type} is None. ")
if value.get("type") == "int":
with contextlib.suppress(TypeError, ValueError):
new_value = int(new_value) # type: ignore
@ -148,12 +148,12 @@ class Node:
# and continue
# Another aspect is that the node_type is the class that we need to import
# and instantiate with these built params
logger.debug(f"Building {self.node_type}")
logger.debug(f"Building {self.vertex_type}")
# Build each node in the params dict
for key, value in self.params.copy().items():
# Check if Node or list of Nodes and not self
# to avoid recursion
if isinstance(value, Node):
if isinstance(value, Vertex):
if value == self:
del self.params[key]
continue
@ -177,7 +177,7 @@ class Node:
self.params[key] = result
elif isinstance(value, list) and all(
isinstance(node, Node) for node in value
isinstance(node, Vertex) for node in value
):
self.params[key] = []
for node in value:
@ -193,17 +193,17 @@ class Node:
try:
self._built_object = loading.instantiate_class(
node_type=self.node_type,
node_type=self.vertex_type,
base_type=self.base_type,
params=self.params,
)
except Exception as exc:
raise ValueError(
f"Error building node {self.node_type}: {str(exc)}"
f"Error building node {self.vertex_type}: {str(exc)}"
) from exc
if self._built_object is None:
raise ValueError(f"Node type {self.node_type} not found")
raise ValueError(f"Node type {self.vertex_type} not found")
self._built = True
@ -220,7 +220,7 @@ class Node:
return f"Node(id={self.id}, data={self.data})"
def __eq__(self, __o: object) -> bool:
return self.id == __o.id if isinstance(__o, Node) else False
return self.id == __o.id if isinstance(__o, Vertex) else False
def __hash__(self) -> int:
return id(self)

View file

@ -1,22 +1,22 @@
from typing import Any, Dict, List, Optional, Union
from langflow.graph.node.base import Node
from langflow.graph.vertex.base import Vertex
from langflow.graph.utils import extract_input_variables_from_prompt
class AgentNode(Node):
class AgentVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="agents")
self.tools: List[ToolNode] = []
self.chains: List[ChainNode] = []
self.tools: List[ToolVertex] = []
self.chains: List[ChainVertex] = []
def _set_tools_and_chains(self) -> None:
for edge in self.edges:
source_node = edge.source
if isinstance(source_node, ToolNode):
if isinstance(source_node, ToolVertex):
self.tools.append(source_node)
elif isinstance(source_node, ChainNode):
elif isinstance(source_node, ChainVertex):
self.chains.append(source_node)
def build(self, force: bool = False) -> Any:
@ -33,24 +33,28 @@ class AgentNode(Node):
self._build()
#! Cannot deepcopy VectorStore, VectorStoreRouter, or SQL agents
if self.node_type in ["VectorStoreAgent", "VectorStoreRouterAgent", "SQLAgent"]:
if self.vertex_type in [
"VectorStoreAgent",
"VectorStoreRouterAgent",
"SQLAgent",
]:
return self._built_object
return self._built_object
class ToolNode(Node):
class ToolVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="tools")
class PromptNode(Node):
class PromptVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="prompts")
def build(
self,
force: bool = False,
tools: Optional[Union[List[Node], List[ToolNode]]] = None,
tools: Optional[Union[List[Vertex], List[ToolVertex]]] = None,
) -> Any:
if not self._built or force:
if (
@ -59,7 +63,7 @@ class PromptNode(Node):
):
self.params["input_variables"] = []
# Check if it is a ZeroShotPrompt and needs a tool
if "ShotPrompt" in self.node_type:
if "ShotPrompt" in self.vertex_type:
tools = (
[tool_node.build() for tool_node in tools]
if tools is not None
@ -83,31 +87,31 @@ class PromptNode(Node):
return self._built_object
class ChainNode(Node):
class ChainVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="chains")
def build(
self,
force: bool = False,
tools: Optional[Union[List[Node], List[ToolNode]]] = None,
tools: Optional[Union[List[Vertex], List[ToolVertex]]] = None,
) -> Any:
if not self._built or force:
# Check if the chain requires a PromptNode
for key, value in self.params.items():
if isinstance(value, PromptNode):
if isinstance(value, PromptVertex):
# Build the PromptNode, passing the tools if available
self.params[key] = value.build(tools=tools, force=force)
self._build()
#! Cannot deepcopy SQLDatabaseChain
if self.node_type in ["SQLDatabaseChain"]:
if self.vertex_type in ["SQLDatabaseChain"]:
return self._built_object
return self._built_object
class LLMNode(Node):
class LLMVertex(Vertex):
built_node_type = None
class_built_object = None
@ -117,28 +121,28 @@ class LLMNode(Node):
def build(self, force: bool = False) -> Any:
# LLM is different because some models might take up too much memory
# or time to load. So we only load them when we need them.ß
if self.node_type == self.built_node_type:
if self.vertex_type == self.built_node_type:
return self.class_built_object
if not self._built or force:
self._build()
self.built_node_type = self.node_type
self.built_node_type = self.vertex_type
self.class_built_object = self._built_object
# Avoid deepcopying the LLM
# that are loaded from a file
return self._built_object
class ToolkitNode(Node):
class ToolkitVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="toolkits")
class FileToolNode(ToolNode):
class FileToolVertex(ToolVertex):
def __init__(self, data: Dict):
super().__init__(data)
class WrapperNode(Node):
class WrapperVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="wrappers")
@ -150,7 +154,7 @@ class WrapperNode(Node):
return self._built_object
class DocumentLoaderNode(Node):
class DocumentLoaderVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="documentloaders")
@ -158,17 +162,17 @@ class DocumentLoaderNode(Node):
# This built_object is a list of documents. Maybe we should
# show how many documents are in the list?
if self._built_object:
return f"""{self.node_type}({len(self._built_object)} documents)
return f"""{self.vertex_type}({len(self._built_object)} documents)
Documents: {self._built_object[:3]}..."""
return f"{self.node_type}()"
return f"{self.vertex_type}()"
class EmbeddingNode(Node):
class EmbeddingVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="embeddings")
class VectorStoreNode(Node):
class VectorStoreVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="vectorstores")
@ -176,12 +180,12 @@ class VectorStoreNode(Node):
return "Vector stores can take time to build. It will build on the first query."
class MemoryNode(Node):
class MemoryVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="memory")
class TextSplitterNode(Node):
class TextSplitterVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="textsplitters")
@ -189,5 +193,5 @@ class TextSplitterNode(Node):
# This built_object is a list of documents. Maybe we should
# show how many documents are in the list?
if self._built_object:
return f"""{self.node_type}({len(self._built_object)} documents)\nDocuments: {self._built_object[:3]}..."""
return f"{self.node_type}()"
return f"""{self.vertex_type}({len(self._built_object)} documents)\nDocuments: {self._built_object[:3]}..."""
return f"{self.vertex_type}()"