jyapayne/langflow
Watch
1
0
Fork
You've already forked langflow
 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

refactor: change Node to Vertex

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-06-01 15:37:11 -03:00
commit 0cef4bd818
10 changed files with 165 additions and 159 deletions
Download patch file Download diff file Expand all files Collapse all files

4
src/backend/langflow/api/validate.py
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:

54
src/backend/langflow/graph/__init__.py
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",
]

10
src/backend/langflow/graph/edge/base.py
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"
)

56
src/backend/langflow/graph/graph/base.py
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"]]

54
src/backend/langflow/graph/graph/constants.py
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()},
}

0
src/backend/langflow/graph/vertex/__init__.py Normal file
View file

229
src/backend/langflow/graph/vertex/base.py Normal file
View file

@ -0,0 +1,229 @@
from langflow.cache import base as cache_utils
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
from langflow.utils.util import sync_to_async
import contextlib
import inspect
import types
import warnings
from typing import Any, Dict, List, Optional
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from langflow.graph.edge.base import Edge
class Vertex:
def __init__(self, data: Dict, base_type: Optional[str] = None) -> None:
self.id: str = data["id"]
self._data = data
self.edges: List["Edge"] = []
self.base_type: Optional[str] = base_type
self._parse_data()
self._built_object = None
self._built = False
def _parse_data(self) -> None:
self.data = self._data["data"]
self.output = self.data["node"]["base_classes"]
template_dicts = {
key: value
for key, value in self.data["node"]["template"].items()
if isinstance(value, dict)
}
self.required_inputs = [
template_dicts[key]["type"]
for key, value in template_dicts.items()
if value["required"]
]
self.optional_inputs = [
template_dicts[key]["type"]
for key, value in template_dicts.items()
if not value["required"]
]
template_dict = self.data["node"]["template"]
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.vertex_type in value:
self.base_type = base_type
break
def _build_params(self):
# Some params are required, some are optional
# but most importantly, some params are python base classes
# like str and others are LangChain objects like LLMChain, BasePromptTemplate
# so we need to be able to distinguish between the two
# The dicts with "type" == "str" are the ones that are python base classes
# and most likely have a "value" key
# So for each key besides "_type" in the template dict, we have a dict
# with a "type" key. If the type is not "str", then we need to get the
# edge that connects to that node and get the Node with the required data
# and use that as the value for the param
# If the type is "str", then we need to get the value of the "value" key
# and use that as the value for the param
template_dict = {
key: value
for key, value in self.data["node"]["template"].items()
if isinstance(value, dict)
}
params = {}
for key, value in template_dict.items():
if key == "_type":
continue
# If the type is not transformable to a python base class
# then we need to get the edge that connects to this node
if value.get("type") == "file":
# Load the type in value.get('suffixes') using
# what is inside value.get('content')
# value.get('value') is the file name
file_name = value.get("value")
content = value.get("content")
type_to_load = value.get("suffixes")
file_path = cache_utils.save_binary_file(
content=content, file_name=file_name, accepted_types=type_to_load
)
params[key] = file_path
elif value.get("type") not in DIRECT_TYPES:
# Get the edge that connects to this node
edges = [
edge
for edge in self.edges
if edge.target == self and edge.matched_type in value["type"]
]
# Get the output of the node that the edge connects to
# if the value['list'] is True, then there will be more
# than one time setting to params[key]
# so we need to append to a list if it exists
# or create a new list if it doesn't
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.vertex_type} not found"
)
elif value["list"]:
# If this is a list parameter, append all sources to a list
params[key] = [edge.source for edge in edges]
elif edges:
# If a single parameter is found, use its source
params[key] = edges[0].source
elif value["required"] or value.get("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.vertex_type} is None. ")
if value.get("type") == "int":
with contextlib.suppress(TypeError, ValueError):
new_value = int(new_value) # type: ignore
params[key] = new_value
# Add _type to params
self.params = params
def _build(self):
# The params dict is used to build the module
# it contains values and keys that point to nodes which
# have their own params dict
# When build is called, we iterate through the params dict
# and if the value is a node, we call build on that node
# and use the output of that build as the value for the param
# if the value is not a node, then we use the value as the param
# 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.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, Vertex):
if value == self:
del self.params[key]
continue
result = value.build()
# If the key is "func", then we need to use the run method
if key == "func":
if not isinstance(result, types.FunctionType):
# func can be
# PythonFunction(code='\ndef upper_case(text: str) -> str:\n return text.upper()\n')
# so we need to check if there is an attribute called run
if hasattr(result, "run"):
result = result.run # type: ignore
elif hasattr(result, "get_function"):
result = result.get_function() # type: ignore
elif inspect.iscoroutinefunction(result):
self.params["coroutine"] = result
else:
# turn result which is a function into a coroutine
# so that it can be awaited
self.params["coroutine"] = sync_to_async(result)
self.params[key] = result
elif isinstance(value, list) and all(
isinstance(node, Vertex) for node in value
):
self.params[key] = []
for node in value:
built = node.build()
if isinstance(built, list):
self.params[key].extend(built)
else:
self.params[key].append(built)
# Get the class from LANGCHAIN_TYPES_DICT
# and instantiate it with the params
# and return the instance
try:
self._built_object = loading.instantiate_class(
node_type=self.vertex_type,
base_type=self.base_type,
params=self.params,
)
except Exception as exc:
raise ValueError(
f"Error building node {self.vertex_type}: {str(exc)}"
) from exc
if self._built_object is None:
raise ValueError(f"Node type {self.vertex_type} not found")
self._built = True
def build(self, force: bool = False) -> Any:
if not self._built or force:
self._build()
return self._built_object
def add_edge(self, edge: "Edge") -> None:
self.edges.append(edge)
def __repr__(self) -> str:
return f"Node(id={self.id}, data={self.data})"
def __eq__(self, __o: object) -> bool:
return self.id == __o.id if isinstance(__o, Vertex) else False
def __hash__(self) -> int:
return id(self)
def _built_object_repr(self):
return repr(self._built_object)

1
src/backend/langflow/graph/vertex/constants.py Normal file
View file

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

197
src/backend/langflow/graph/vertex/types.py Normal file
View file

@ -0,0 +1,197 @@
from typing import Any, Dict, List, Optional, Union
from langflow.graph.vertex.base import Vertex
from langflow.graph.utils import extract_input_variables_from_prompt
class AgentVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="agents")
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, ToolVertex):
self.tools.append(source_node)
elif isinstance(source_node, ChainVertex):
self.chains.append(source_node)
def build(self, force: bool = False) -> Any:
if not self._built or force:
self._set_tools_and_chains()
# First, build the tools
for tool_node in self.tools:
tool_node.build()
# Next, build the chains and the rest
for chain_node in self.chains:
chain_node.build(tools=self.tools)
self._build()
#! Cannot deepcopy VectorStore, VectorStoreRouter, or SQL agents
if self.vertex_type in [
"VectorStoreAgent",
"VectorStoreRouterAgent",
"SQLAgent",
]:
return self._built_object
return self._built_object
class ToolVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="tools")
class PromptVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="prompts")
def build(
self,
force: bool = False,
tools: Optional[Union[List[Vertex], List[ToolVertex]]] = None,
) -> Any:
if not self._built or force:
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.vertex_type:
tools = (
[tool_node.build() for tool_node in tools]
if tools is not None
else []
)
self.params["tools"] = tools
prompt_params = [
key
for key, value in self.params.items()
if isinstance(value, str) and key != "format_instructions"
]
else:
prompt_params = ["template"]
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"]))
self._build()
return self._built_object
class ChainVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="chains")
def build(
self,
force: bool = False,
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, 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.vertex_type in ["SQLDatabaseChain"]:
return self._built_object
return self._built_object
class LLMVertex(Vertex):
built_node_type = None
class_built_object = None
def __init__(self, data: Dict):
super().__init__(data, base_type="llms")
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.vertex_type == self.built_node_type:
return self.class_built_object
if not self._built or force:
self._build()
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 ToolkitVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="toolkits")
class FileToolVertex(ToolVertex):
def __init__(self, data: Dict):
super().__init__(data)
class WrapperVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="wrappers")
def build(self, force: bool = False) -> Any:
if not self._built or force:
if "headers" in self.params:
self.params["headers"] = eval(self.params["headers"])
self._build()
return self._built_object
class DocumentLoaderVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="documentloaders")
def _built_object_repr(self):
# 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.vertex_type}({len(self._built_object)} documents)
Documents: {self._built_object[:3]}..."""
return f"{self.vertex_type}()"
class EmbeddingVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="embeddings")
class VectorStoreVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="vectorstores")
def _built_object_repr(self):
return "Vector stores can take time to build. It will build on the first query."
class MemoryVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="memory")
class TextSplitterVertex(Vertex):
def __init__(self, data: Dict):
super().__init__(data, base_type="textsplitters")
def _built_object_repr(self):
# 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.vertex_type}({len(self._built_object)} documents)\nDocuments: {self._built_object[:3]}..."""
return f"{self.vertex_type}()"

54
tests/test_graph.py
View file

@ -1,20 +1,20 @@
from typing import Type, Union
from langflow.graph.edge.base import Edge
from langflow.graph.node.base import Node
from langflow.graph.vertex.base import Vertex
import pytest
from langchain.chains.base import Chain
from langchain.llms.fake import FakeListLLM
from langflow.graph import Graph
from langflow.graph.node.types import (
AgentNode,
ChainNode,
FileToolNode,
LLMNode,
PromptNode,
ToolkitNode,
ToolNode,
WrapperNode,
from langflow.graph.vertex.types import (
AgentVertex,
ChainVertex,
FileToolVertex,
LLMVertex,
PromptVertex,
ToolkitVertex,
ToolVertex,
WrapperVertex,
)
from langflow.interface.run import get_result_and_thought
from langflow.utils.payload import get_root_node
@ -25,7 +25,7 @@ from langflow.utils.payload import get_root_node
# BASIC_EXAMPLE_PATH, COMPLEX_EXAMPLE_PATH, OPENAPI_EXAMPLE_PATH
def get_node_by_type(graph, node_type: Type[Node]) -> Union[Node, None]:
def get_node_by_type(graph, node_type: Type[Vertex]) -> Union[Vertex, None]:
"""Get a node by type"""
return next((node for node in graph.nodes if isinstance(node, node_type)), None)
@ -35,7 +35,7 @@ def test_graph_structure(basic_graph):
assert len(basic_graph.nodes) > 0
assert len(basic_graph.edges) > 0
for node in basic_graph.nodes:
assert isinstance(node, Node)
assert isinstance(node, Vertex)
for edge in basic_graph.edges:
assert isinstance(edge, Edge)
assert edge.source in basic_graph.nodes
@ -165,7 +165,7 @@ def test_get_node(basic_graph):
"""Test getting a single node"""
node_id = basic_graph.nodes[0].id
node = basic_graph.get_node(node_id)
assert isinstance(node, Node)
assert isinstance(node, Vertex)
assert node.id == node_id
@ -174,7 +174,7 @@ def test_build_nodes(basic_graph):
assert len(basic_graph.nodes) == len(basic_graph._nodes)
for node in basic_graph.nodes:
assert isinstance(node, Node)
assert isinstance(node, Vertex)
def test_build_edges(basic_graph):
@ -182,8 +182,8 @@ def test_build_edges(basic_graph):
assert len(basic_graph.edges) == len(basic_graph._edges)
for edge in basic_graph.edges:
assert isinstance(edge, Edge)
assert isinstance(edge.source, Node)
assert isinstance(edge.target, Node)
assert isinstance(edge.source, Vertex)
assert isinstance(edge.target, Vertex)
def test_get_root_node(basic_graph, complex_graph):
@ -191,13 +191,13 @@ def test_get_root_node(basic_graph, complex_graph):
assert isinstance(basic_graph, Graph)
root = get_root_node(basic_graph)
assert root is not None
assert isinstance(root, Node)
assert isinstance(root, Vertex)
assert root.data["type"] == "TimeTravelGuideChain"
# For complex example, the root node is a ZeroShotAgent too
assert isinstance(complex_graph, Graph)
root = get_root_node(complex_graph)
assert root is not None
assert isinstance(root, Node)
assert isinstance(root, Vertex)
assert root.data["type"] == "ZeroShotAgent"
@ -257,14 +257,14 @@ def assert_agent_was_built(graph):
def test_agent_node_build(complex_graph):
agent_node = get_node_by_type(complex_graph, AgentNode)
agent_node = get_node_by_type(complex_graph, AgentVertex)
assert agent_node is not None
built_object = agent_node.build()
assert built_object is not None
def test_tool_node_build(complex_graph):
tool_node = get_node_by_type(complex_graph, ToolNode)
tool_node = get_node_by_type(complex_graph, ToolVertex)
assert tool_node is not None
built_object = tool_node.build()
assert built_object is not None
@ -272,7 +272,7 @@ def test_tool_node_build(complex_graph):
def test_chain_node_build(complex_graph):
chain_node = get_node_by_type(complex_graph, ChainNode)
chain_node = get_node_by_type(complex_graph, ChainVertex)
assert chain_node is not None
built_object = chain_node.build()
assert built_object is not None
@ -280,7 +280,7 @@ def test_chain_node_build(complex_graph):
def test_prompt_node_build(complex_graph):
prompt_node = get_node_by_type(complex_graph, PromptNode)
prompt_node = get_node_by_type(complex_graph, PromptVertex)
assert prompt_node is not None
built_object = prompt_node.build()
assert built_object is not None
@ -288,7 +288,7 @@ def test_prompt_node_build(complex_graph):
def test_llm_node_build(basic_graph):
llm_node = get_node_by_type(basic_graph, LLMNode)
llm_node = get_node_by_type(basic_graph, LLMVertex)
assert llm_node is not None
built_object = llm_node.build()
assert built_object is not None
@ -296,7 +296,7 @@ def test_llm_node_build(basic_graph):
def test_toolkit_node_build(openapi_graph):
toolkit_node = get_node_by_type(openapi_graph, ToolkitNode)
toolkit_node = get_node_by_type(openapi_graph, ToolkitVertex)
assert toolkit_node is not None
built_object = toolkit_node.build()
assert built_object is not None
@ -304,7 +304,7 @@ def test_toolkit_node_build(openapi_graph):
def test_file_tool_node_build(openapi_graph):
file_tool_node = get_node_by_type(openapi_graph, FileToolNode)
file_tool_node = get_node_by_type(openapi_graph, FileToolVertex)
assert file_tool_node is not None
built_object = file_tool_node.build()
assert built_object is not None
@ -312,7 +312,7 @@ def test_file_tool_node_build(openapi_graph):
def test_wrapper_node_build(openapi_graph):
wrapper_node = get_node_by_type(openapi_graph, WrapperNode)
wrapper_node = get_node_by_type(openapi_graph, WrapperVertex)
assert wrapper_node is not None
built_object = wrapper_node.build()
assert built_object is not None
@ -327,7 +327,7 @@ def test_get_result_and_thought(basic_graph):
message = "Hello"
# Find the node that is an LLMNode and change the
# _built_object to a FakeListLLM
llm_node = get_node_by_type(basic_graph, LLMNode)
llm_node = get_node_by_type(basic_graph, LLMVertex)
assert llm_node is not None
llm_node._built_object = FakeListLLM(responses=responses)
llm_node._built = True