515 lines
18 KiB
Python
515 lines
18 KiB
Python
import copy
|
|
import json
|
|
import os
|
|
import pickle
|
|
from pathlib import Path
|
|
from typing import Type, Union
|
|
|
|
import pytest
|
|
from langchain.agents import AgentExecutor
|
|
from langchain.chains.base import Chain
|
|
from langchain.llms.fake import FakeListLLM
|
|
from langflow.graph import Graph
|
|
from langflow.graph.edge.base import Edge
|
|
from langflow.graph.graph.utils import (
|
|
find_last_node,
|
|
process_flow,
|
|
set_new_target_handle,
|
|
ungroup_node,
|
|
update_source_handle,
|
|
update_target_handle,
|
|
update_template,
|
|
)
|
|
from langflow.graph.vertex.base import Vertex
|
|
from langflow.graph.vertex.types import FileToolVertex, LLMVertex, ToolkitVertex
|
|
from langflow.processing.process import get_result_and_thought
|
|
from langflow.utils.payload import get_root_vertex
|
|
|
|
# Test cases for the graph module
|
|
|
|
# now we have three types of graph:
|
|
# BASIC_EXAMPLE_PATH, COMPLEX_EXAMPLE_PATH, OPENAPI_EXAMPLE_PATH
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_template():
|
|
return {
|
|
"field1": {"proxy": {"field": "some_field", "id": "node1"}},
|
|
"field2": {"proxy": {"field": "other_field", "id": "node2"}},
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_nodes():
|
|
return [
|
|
{
|
|
"id": "node1",
|
|
"data": {"node": {"template": {"some_field": {"show": True, "advanced": False, "name": "Name1"}}}},
|
|
},
|
|
{
|
|
"id": "node2",
|
|
"data": {
|
|
"node": {
|
|
"template": {
|
|
"other_field": {
|
|
"show": False,
|
|
"advanced": True,
|
|
"display_name": "DisplayName2",
|
|
}
|
|
}
|
|
}
|
|
},
|
|
},
|
|
{
|
|
"id": "node3",
|
|
"data": {"node": {"template": {"unrelated_field": {"show": True, "advanced": True}}}},
|
|
},
|
|
]
|
|
|
|
|
|
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.vertices if isinstance(node, node_type)), None)
|
|
|
|
|
|
def test_graph_structure(basic_graph):
|
|
assert isinstance(basic_graph, Graph)
|
|
assert len(basic_graph.vertices) > 0
|
|
assert len(basic_graph.edges) > 0
|
|
for node in basic_graph.vertices:
|
|
assert isinstance(node, Vertex)
|
|
for edge in basic_graph.edges:
|
|
assert isinstance(edge, Edge)
|
|
source_vertex = basic_graph.get_vertex(edge.source_id)
|
|
target_vertex = basic_graph.get_vertex(edge.target_id)
|
|
assert source_vertex in basic_graph.vertices
|
|
assert target_vertex in basic_graph.vertices
|
|
|
|
|
|
def test_circular_dependencies(basic_graph):
|
|
assert isinstance(basic_graph, Graph)
|
|
|
|
def check_circular(node, visited):
|
|
visited.add(node)
|
|
neighbors = basic_graph.get_vertices_with_target(node)
|
|
for neighbor in neighbors:
|
|
if neighbor in visited:
|
|
return True
|
|
if check_circular(neighbor, visited.copy()):
|
|
return True
|
|
return False
|
|
|
|
for node in basic_graph.vertices:
|
|
assert not check_circular(node, set())
|
|
|
|
|
|
def test_invalid_node_types():
|
|
graph_data = {
|
|
"nodes": [
|
|
{
|
|
"id": "1",
|
|
"data": {
|
|
"node": {
|
|
"base_classes": ["BaseClass"],
|
|
"template": {
|
|
"_type": "InvalidNodeType",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
],
|
|
"edges": [],
|
|
}
|
|
with pytest.raises(Exception):
|
|
Graph(graph_data["nodes"], graph_data["edges"])
|
|
|
|
|
|
def test_get_vertices_with_target(basic_graph):
|
|
"""Test getting connected nodes"""
|
|
assert isinstance(basic_graph, Graph)
|
|
# Get root node
|
|
root = get_root_vertex(basic_graph)
|
|
assert root is not None
|
|
connected_nodes = basic_graph.get_vertices_with_target(root.id)
|
|
assert connected_nodes is not None
|
|
|
|
|
|
def test_get_node_neighbors_basic(basic_graph):
|
|
"""Test getting node neighbors"""
|
|
|
|
assert isinstance(basic_graph, Graph)
|
|
# Get root node
|
|
root = get_root_vertex(basic_graph)
|
|
assert root is not None
|
|
neighbors = basic_graph.get_vertex_neighbors(root)
|
|
assert neighbors is not None
|
|
assert isinstance(neighbors, dict)
|
|
# Root Node is an Agent, it requires an LLMChain and tools
|
|
# We need to check if there is a Chain in the one of the neighbors'
|
|
# data attribute in the type key
|
|
assert any("ConversationBufferMemory" in neighbor.data["type"] for neighbor, val in neighbors.items() if val)
|
|
|
|
assert any("OpenAI" in neighbor.data["type"] for neighbor, val in neighbors.items() if val)
|
|
|
|
|
|
def test_get_node(basic_graph):
|
|
"""Test getting a single node"""
|
|
node_id = basic_graph.vertices[0].id
|
|
node = basic_graph.get_vertex(node_id)
|
|
assert isinstance(node, Vertex)
|
|
assert node.id == node_id
|
|
|
|
|
|
def test_build_nodes(basic_graph):
|
|
"""Test building nodes"""
|
|
|
|
assert len(basic_graph.vertices) == len(basic_graph._vertices)
|
|
for node in basic_graph.vertices:
|
|
assert isinstance(node, Vertex)
|
|
|
|
|
|
def test_build_edges(basic_graph):
|
|
"""Test building edges"""
|
|
assert len(basic_graph.edges) == len(basic_graph._edges)
|
|
for edge in basic_graph.edges:
|
|
assert isinstance(edge, Edge)
|
|
assert isinstance(edge.source_id, str)
|
|
assert isinstance(edge.target_id, str)
|
|
|
|
|
|
def test_get_root_vertex(client, basic_graph, complex_graph):
|
|
"""Test getting root node"""
|
|
assert isinstance(basic_graph, Graph)
|
|
root = get_root_vertex(basic_graph)
|
|
assert root is not None
|
|
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_vertex(complex_graph)
|
|
assert root is not None
|
|
assert isinstance(root, Vertex)
|
|
assert root.data["type"] == "ZeroShotAgent"
|
|
|
|
|
|
def test_validate_edges(basic_graph):
|
|
"""Test validating edges"""
|
|
|
|
assert isinstance(basic_graph, Graph)
|
|
# all edges should be valid
|
|
assert all(edge.valid for edge in basic_graph.edges)
|
|
|
|
|
|
def test_matched_type(basic_graph):
|
|
"""Test matched type attribute in Edge"""
|
|
assert isinstance(basic_graph, Graph)
|
|
# all edges should be valid
|
|
assert all(edge.valid for edge in basic_graph.edges)
|
|
# all edges should have a matched_type attribute
|
|
assert all(hasattr(edge, "matched_type") for edge in basic_graph.edges)
|
|
# The matched_type attribute should be in the source_types attr
|
|
assert all(edge.matched_type in edge.source_types for edge in basic_graph.edges)
|
|
|
|
|
|
def test_build_params(basic_graph):
|
|
"""Test building params"""
|
|
|
|
assert isinstance(basic_graph, Graph)
|
|
# all edges should be valid
|
|
assert all(edge.valid for edge in basic_graph.edges)
|
|
# all edges should have a matched_type attribute
|
|
assert all(hasattr(edge, "matched_type") for edge in basic_graph.edges)
|
|
# The matched_type attribute should be in the source_types attr
|
|
assert all(edge.matched_type in edge.source_types for edge in basic_graph.edges)
|
|
# Get the root node
|
|
root = get_root_vertex(basic_graph)
|
|
# Root node is a TimeTravelGuideChain
|
|
# which requires an llm and memory
|
|
assert root is not None
|
|
assert isinstance(root.params, dict)
|
|
assert "llm" in root.params
|
|
assert "memory" in root.params
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_build(basic_graph):
|
|
"""Test Node's build method"""
|
|
await assert_agent_was_built(basic_graph)
|
|
|
|
|
|
async def assert_agent_was_built(graph):
|
|
"""Assert that the agent was built"""
|
|
assert isinstance(graph, Graph)
|
|
# Now we test the build method
|
|
# Build the Agent
|
|
result = await graph.build()
|
|
# The agent should be a AgentExecutor
|
|
assert isinstance(result, Chain)
|
|
|
|
|
|
def test_llm_node_build(basic_graph):
|
|
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
|
|
|
|
|
|
def test_toolkit_node_build(client, openapi_graph):
|
|
# Write a file to the disk
|
|
file_path = "api-with-examples.yaml"
|
|
with open(file_path, "w") as f:
|
|
f.write("openapi: 3.0.0")
|
|
|
|
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
|
|
# Remove the file
|
|
os.remove(file_path)
|
|
assert not Path(file_path).exists()
|
|
|
|
|
|
def test_file_tool_node_build(client, openapi_graph):
|
|
file_path = "api-with-examples.yaml"
|
|
with open(file_path, "w") as f:
|
|
f.write("openapi: 3.0.0")
|
|
|
|
assert Path(file_path).exists()
|
|
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
|
|
# Remove the file
|
|
os.remove(file_path)
|
|
assert not Path(file_path).exists()
|
|
|
|
|
|
# def test_wrapper_node_build(openapi_graph):
|
|
# 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
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_result_and_thought(basic_graph):
|
|
"""Test the get_result_and_thought method"""
|
|
responses = [
|
|
"Final Answer: I am a response",
|
|
]
|
|
message = {"input": "Hello"}
|
|
# Find the node that is an LLMNode and change the
|
|
# _built_object to a FakeListLLM
|
|
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
|
|
langchain_object = await basic_graph.build()
|
|
# assert all nodes are built
|
|
assert all(node._built for node in basic_graph.vertices)
|
|
# now build again and check if FakeListLLM was used
|
|
|
|
# Get the result and thought
|
|
result = get_result_and_thought(langchain_object, message)
|
|
assert isinstance(result, dict)
|
|
|
|
|
|
def test_find_last_node(grouped_chat_json_flow):
|
|
grouped_chat_data = json.loads(grouped_chat_json_flow).get("data")
|
|
nodes, edges = grouped_chat_data["nodes"], grouped_chat_data["edges"]
|
|
last_node = find_last_node(nodes, edges)
|
|
assert last_node is not None # Replace with the actual expected value
|
|
assert last_node["id"] == "LLMChain-pimAb" # Replace with the actual expected value
|
|
|
|
|
|
def test_ungroup_node(grouped_chat_json_flow):
|
|
grouped_chat_data = json.loads(grouped_chat_json_flow).get("data")
|
|
group_node = grouped_chat_data["nodes"][2] # Assuming the first node is a group node
|
|
base_flow = copy.deepcopy(grouped_chat_data)
|
|
ungroup_node(group_node["data"], base_flow)
|
|
# after ungroup_node is called, the base_flow and grouped_chat_data should be different
|
|
assert base_flow != grouped_chat_data
|
|
# assert node 2 is not a group node anymore
|
|
assert base_flow["nodes"][2]["data"]["node"].get("flow") is None
|
|
# assert the edges are updated
|
|
assert len(base_flow["edges"]) > len(grouped_chat_data["edges"])
|
|
assert base_flow["edges"][0]["source"] == "ConversationBufferMemory-kUMif"
|
|
assert base_flow["edges"][0]["target"] == "LLMChain-2P369"
|
|
assert base_flow["edges"][1]["source"] == "PromptTemplate-Wjk4g"
|
|
assert base_flow["edges"][1]["target"] == "LLMChain-2P369"
|
|
assert base_flow["edges"][2]["source"] == "ChatOpenAI-rUJ1b"
|
|
assert base_flow["edges"][2]["target"] == "LLMChain-2P369"
|
|
|
|
|
|
def test_process_flow(grouped_chat_json_flow):
|
|
grouped_chat_data = json.loads(grouped_chat_json_flow).get("data")
|
|
|
|
processed_flow = process_flow(grouped_chat_data)
|
|
assert processed_flow is not None
|
|
assert isinstance(processed_flow, dict)
|
|
assert "nodes" in processed_flow
|
|
assert "edges" in processed_flow
|
|
|
|
|
|
def test_process_flow_one_group(one_grouped_chat_json_flow):
|
|
grouped_chat_data = json.loads(one_grouped_chat_json_flow).get("data")
|
|
# There should be only one node
|
|
assert len(grouped_chat_data["nodes"]) == 1
|
|
# Get the node, it should be a group node
|
|
group_node = grouped_chat_data["nodes"][0]
|
|
node_data = group_node["data"]["node"]
|
|
assert node_data.get("flow") is not None
|
|
template_data = node_data["template"]
|
|
assert any("openai_api_key" in key for key in template_data.keys())
|
|
# Get the openai_api_key dict
|
|
openai_api_key = next(
|
|
(template_data[key] for key in template_data.keys() if "openai_api_key" in key),
|
|
None,
|
|
)
|
|
assert openai_api_key is not None
|
|
assert openai_api_key["value"] == "test"
|
|
|
|
processed_flow = process_flow(grouped_chat_data)
|
|
assert processed_flow is not None
|
|
assert isinstance(processed_flow, dict)
|
|
assert "nodes" in processed_flow
|
|
assert "edges" in processed_flow
|
|
|
|
# Now get the node that has ChatOpenAI in its id
|
|
chat_openai_node = next((node for node in processed_flow["nodes"] if "ChatOpenAI" in node["id"]), None)
|
|
assert chat_openai_node is not None
|
|
assert chat_openai_node["data"]["node"]["template"]["openai_api_key"]["value"] == "test"
|
|
|
|
|
|
def test_process_flow_vector_store_grouped(vector_store_grouped_json_flow):
|
|
grouped_chat_data = json.loads(vector_store_grouped_json_flow).get("data")
|
|
nodes = grouped_chat_data["nodes"]
|
|
assert len(nodes) == 4
|
|
# There are two group nodes in this flow
|
|
# One of them is inside the other totalling 7 nodes
|
|
# 4 nodes grouped, one of these turns into 1 normal node and 1 group node
|
|
# This group node has 2 nodes inside it
|
|
|
|
processed_flow = process_flow(grouped_chat_data)
|
|
assert processed_flow is not None
|
|
processed_nodes = processed_flow["nodes"]
|
|
assert len(processed_nodes) == 7
|
|
assert isinstance(processed_flow, dict)
|
|
assert "nodes" in processed_flow
|
|
assert "edges" in processed_flow
|
|
edges = processed_flow["edges"]
|
|
# Expected keywords in source and target fields
|
|
expected_keywords = [
|
|
{"source": "VectorStoreInfo", "target": "VectorStoreAgent"},
|
|
{"source": "ChatOpenAI", "target": "VectorStoreAgent"},
|
|
{"source": "OpenAIEmbeddings", "target": "Chroma"},
|
|
{"source": "Chroma", "target": "VectorStoreInfo"},
|
|
{"source": "WebBaseLoader", "target": "RecursiveCharacterTextSplitter"},
|
|
{"source": "RecursiveCharacterTextSplitter", "target": "Chroma"},
|
|
]
|
|
|
|
for idx, expected_keyword in enumerate(expected_keywords):
|
|
for key, value in expected_keyword.items():
|
|
assert (
|
|
value in edges[idx][key].split("-")[0]
|
|
), f"Edge {idx}, key {key} expected to contain {value} but got {edges[idx][key]}"
|
|
|
|
|
|
def test_update_template(sample_template, sample_nodes):
|
|
# Making a deep copy to keep original sample_nodes unchanged
|
|
nodes_copy = copy.deepcopy(sample_nodes)
|
|
update_template(sample_template, nodes_copy)
|
|
|
|
# Now, validate the updates.
|
|
node1_updated = next((n for n in nodes_copy if n["id"] == "node1"), None)
|
|
node2_updated = next((n for n in nodes_copy if n["id"] == "node2"), None)
|
|
node3_updated = next((n for n in nodes_copy if n["id"] == "node3"), None)
|
|
|
|
assert node1_updated["data"]["node"]["template"]["some_field"]["show"] is True
|
|
assert node1_updated["data"]["node"]["template"]["some_field"]["advanced"] is False
|
|
assert node1_updated["data"]["node"]["template"]["some_field"]["display_name"] == "Name1"
|
|
|
|
assert node2_updated["data"]["node"]["template"]["other_field"]["show"] is False
|
|
assert node2_updated["data"]["node"]["template"]["other_field"]["advanced"] is True
|
|
assert node2_updated["data"]["node"]["template"]["other_field"]["display_name"] == "DisplayName2"
|
|
|
|
# Ensure node3 remains unchanged
|
|
assert node3_updated == sample_nodes[2]
|
|
|
|
|
|
# Test `update_target_handle`
|
|
def test_update_target_handle_proxy():
|
|
new_edge = {
|
|
"data": {
|
|
"targetHandle": {
|
|
"type": "some_type",
|
|
"proxy": {"id": "some_id", "field": ""},
|
|
}
|
|
}
|
|
}
|
|
g_nodes = [{"id": "some_id", "data": {"node": {"flow": None}}}]
|
|
group_node_id = "group_id"
|
|
updated_edge = update_target_handle(new_edge, g_nodes, group_node_id)
|
|
assert updated_edge["data"]["targetHandle"] == new_edge["data"]["targetHandle"]
|
|
|
|
|
|
# Test `set_new_target_handle`
|
|
def test_set_new_target_handle():
|
|
proxy_id = "proxy_id"
|
|
new_edge = {"target": None, "data": {"targetHandle": {}}}
|
|
target_handle = {"type": "type_1", "proxy": {"field": "field_1"}}
|
|
node = {
|
|
"data": {
|
|
"node": {
|
|
"flow": True,
|
|
"template": {"field_1": {"proxy": {"field": "new_field", "id": "new_id"}}},
|
|
}
|
|
}
|
|
}
|
|
set_new_target_handle(proxy_id, new_edge, target_handle, node)
|
|
assert new_edge["target"] == "proxy_id"
|
|
assert new_edge["data"]["targetHandle"]["fieldName"] == "field_1"
|
|
assert new_edge["data"]["targetHandle"]["proxy"] == {
|
|
"field": "new_field",
|
|
"id": "new_id",
|
|
}
|
|
|
|
|
|
# Test `update_source_handle`
|
|
def test_update_source_handle():
|
|
new_edge = {"source": None, "data": {"sourceHandle": {"id": None}}}
|
|
flow_data = {
|
|
"nodes": [{"id": "some_node"}, {"id": "last_node"}],
|
|
"edges": [{"source": "some_node"}],
|
|
}
|
|
updated_edge = update_source_handle(new_edge, flow_data["nodes"], flow_data["edges"])
|
|
assert updated_edge["source"] == "last_node"
|
|
assert updated_edge["data"]["sourceHandle"]["id"] == "last_node"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pickle_graph(json_vector_store):
|
|
loaded_json = json.loads(json_vector_store)
|
|
graph = Graph.from_payload(loaded_json)
|
|
assert isinstance(graph, Graph)
|
|
first_result = await graph.build()
|
|
assert isinstance(first_result, AgentExecutor)
|
|
pickled = pickle.dumps(graph)
|
|
assert pickled is not None
|
|
unpickled = pickle.loads(pickled)
|
|
assert unpickled is not None
|
|
result = await unpickled.build()
|
|
assert isinstance(result, AgentExecutor)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pickle_each_vertex(json_vector_store):
|
|
loaded_json = json.loads(json_vector_store)
|
|
graph = Graph.from_payload(loaded_json)
|
|
assert isinstance(graph, Graph)
|
|
for vertex in graph.vertices:
|
|
await vertex.build()
|
|
pickled = pickle.dumps(vertex)
|
|
assert pickled is not None
|
|
unpickled = pickle.loads(pickled)
|
|
assert unpickled is not None
|