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

Fix import statements and formatting issues

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-01-25 13:03:42 -03:00
commit a6a69a0f3a
24 changed files with 924 additions and 527 deletions
Download patch file Download diff file Expand all files Collapse all files

226
tests/test_graph.py
View file

@ -1,31 +1,32 @@
import copy
import json
import os
import pickle
from pathlib import Path
import pickle
from typing import Type, Union
import pytest
from langflow.graph.edge.base import Edge
from langflow.graph.vertex.base import Vertex
from langchain.agents import AgentExecutor
import pytest
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.vertex.types import (
FileToolVertex,
LLMVertex,
ToolkitVertex,
)
from langflow.processing.process import get_result_and_thought
from langflow.utils.payload import get_root_node
from langflow.graph.graph.utils import (
find_last_node,
process_flow,
set_new_target_handle,
ungroup_node,
process_flow,
update_source_handle,
update_target_handle,
update_template,
)
from langflow.graph.utils import UnbuiltObject
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
@ -46,7 +47,13 @@ def sample_nodes():
return [
{
"id": "node1",
"data": {"node": {"template": {"some_field": {"show": True, "advanced": False, "name": "Name1"}}}},
"data": {
"node": {
"template": {
"some_field": {"show": True, "advanced": False, "name": "Name1"}
}
}
},
},
{
"id": "node2",
@ -64,7 +71,11 @@ def sample_nodes():
},
{
"id": "node3",
"data": {"node": {"template": {"unrelated_field": {"show": True, "advanced": True}}}},
"data": {
"node": {
"template": {"unrelated_field": {"show": True, "advanced": True}}
}
},
},
]
@ -82,8 +93,8 @@ def test_graph_structure(basic_graph):
assert isinstance(node, Vertex)
for edge in basic_graph.edges:
assert isinstance(edge, Edge)
assert edge.source_id in basic_graph.vertex_map.keys()
assert edge.target_id in basic_graph.vertex_map.keys()
assert edge.source in basic_graph.vertices
assert edge.target in basic_graph.vertices
def test_circular_dependencies(basic_graph):
@ -91,7 +102,7 @@ def test_circular_dependencies(basic_graph):
def check_circular(node, visited):
visited.add(node)
neighbors = basic_graph.get_vertices_with_target(node)
neighbors = basic_graph.get_nodes_with_target(node)
for neighbor in neighbors:
if neighbor in visited:
return True
@ -124,13 +135,13 @@ def test_invalid_node_types():
Graph(graph_data["nodes"], graph_data["edges"])
def test_get_vertices_with_target(basic_graph):
def test_get_nodes_with_target(basic_graph):
"""Test getting connected nodes"""
assert isinstance(basic_graph, Graph)
# Get root node
root = get_root_vertex(basic_graph)
root = get_root_node(basic_graph)
assert root is not None
connected_nodes = basic_graph.get_vertices_with_target(root.id)
connected_nodes = basic_graph.get_nodes_with_target(root)
assert connected_nodes is not None
@ -139,66 +150,23 @@ def test_get_node_neighbors_basic(basic_graph):
assert isinstance(basic_graph, Graph)
# Get root node
root = get_root_vertex(basic_graph)
root = get_root_node(basic_graph)
assert root is not None
neighbors = basic_graph.get_vertex_neighbors(root)
neighbors = basic_graph.get_node_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(
"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_neighbors_complex(complex_graph):
# """Test getting node neighbors"""
# assert isinstance(complex_graph, Graph)
# # Get root node
# root = get_root_node(complex_graph)
# assert root is not None
# neighbors = complex_graph.get_nodes_with_target(root)
# assert neighbors is not None
# # Neighbors should be a list of nodes
# assert isinstance(neighbors, list)
# # 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'
# assert any("Chain" in neighbor.data["type"] for neighbor in neighbors)
# # assert Tool is in the neighbors
# assert any("Tool" in neighbor.data["type"] for neighbor in neighbors)
# # Now on to the Chain's neighbors
# chain = next(neighbor for neighbor in neighbors if "Chain" in neighbor.data["type"])
# chain_neighbors = complex_graph.get_nodes_with_target(chain)
# assert chain_neighbors is not None
# # Check if there is a LLM in the chain's neighbors
# assert any("OpenAI" in neighbor.data["type"] for neighbor in chain_neighbors)
# # Chain should have a Prompt as a neighbor
# assert any("Prompt" in neighbor.data["type"] for neighbor in chain_neighbors)
# # Now on to the Tool's neighbors
# tool = next(neighbor for neighbor in neighbors if "Tool" in neighbor.data["type"])
# tool_neighbors = complex_graph.get_nodes_with_target(tool)
# assert tool_neighbors is not None
# # Check if there is an Agent in the tool's neighbors
# assert any("Agent" in neighbor.data["type"] for neighbor in tool_neighbors)
# # This Agent has a Tool that has a PythonFunction as func
# agent = next(
# neighbor for neighbor in tool_neighbors if "Agent" in neighbor.data["type"]
# )
# agent_neighbors = complex_graph.get_nodes_with_target(agent)
# assert agent_neighbors is not None
# # Check if there is a Tool in the agent's neighbors
# assert any("Tool" in neighbor.data["type"] for neighbor in agent_neighbors)
# # This Tool has a PythonFunction as func
# tool = next(
# neighbor for neighbor in agent_neighbors if "Tool" in neighbor.data["type"]
# )
# tool_neighbors = complex_graph.get_nodes_with_target(tool)
# assert tool_neighbors is not None
# # Check if there is a PythonFunction in the tool's neighbors
# assert any(
# "PythonFunctionTool" in neighbor.data["type"] for neighbor in tool_neighbors
# )
assert any(
"OpenAI" in neighbor.data["type"] for neighbor, val in neighbors.items() if val
)
def test_get_node(basic_graph):
@ -212,7 +180,7 @@ def test_get_node(basic_graph):
def test_build_nodes(basic_graph):
"""Test building nodes"""
assert len(basic_graph.vertices) == len(basic_graph._vertices)
assert len(basic_graph.vertices) == len(basic_graph._nodes)
for node in basic_graph.vertices:
assert isinstance(node, Vertex)
@ -222,21 +190,20 @@ 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_id, str)
assert isinstance(edge.target_id, str)
assert isinstance(edge.source, Vertex)
assert isinstance(edge.target, Vertex)
def test_get_root_vertex(client, basic_graph, complex_graph):
def test_get_root_node(client, basic_graph, complex_graph):
"""Test getting root node"""
assert isinstance(basic_graph, Graph)
root = get_root_vertex(basic_graph)
root = get_root_node(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)
root = get_root_node(complex_graph)
assert root is not None
assert isinstance(root, Vertex)
assert root.data["type"] == "ZeroShotAgent"
@ -272,7 +239,7 @@ def test_build_params(basic_graph):
# 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 = get_root_node(basic_graph)
# Root node is a TimeTravelGuideChain
# which requires an llm and memory
assert root is not None
@ -281,32 +248,29 @@ def test_build_params(basic_graph):
assert "memory" in root.params
@pytest.mark.asyncio
async def test_build(basic_graph):
def test_build(basic_graph):
"""Test Node's build method"""
await assert_agent_was_built(basic_graph)
assert_agent_was_built(basic_graph)
async def assert_agent_was_built(graph):
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()
result = graph.build()
# The agent should be a AgentExecutor
assert isinstance(result, Chain)
@pytest.mark.asyncio
async def test_llm_node_build(basic_graph):
def test_llm_node_build(basic_graph):
llm_node = get_node_by_type(basic_graph, LLMVertex)
assert llm_node is not None
built_object = await llm_node.build()
assert built_object is not UnbuiltObject()
built_object = llm_node.build()
assert built_object is not None
@pytest.mark.asyncio
async def test_toolkit_node_build(client, openapi_graph):
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:
@ -314,31 +278,36 @@ async def test_toolkit_node_build(client, openapi_graph):
toolkit_node = get_node_by_type(openapi_graph, ToolkitVertex)
assert toolkit_node is not None
built_object = await toolkit_node.build()
assert built_object is not UnbuiltObject
built_object = toolkit_node.build()
assert built_object is not None
# Remove the file
os.remove(file_path)
assert not Path(file_path).exists()
@pytest.mark.asyncio
async def test_file_tool_node_build(client, openapi_graph):
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 UnbuiltObject and file_tool_node is not None
built_object = await file_tool_node.build()
assert built_object is not UnbuiltObject
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()
@pytest.mark.asyncio
async def test_get_result_and_thought(basic_graph):
# 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
def test_get_result_and_thought(basic_graph):
"""Test the get_result_and_thought method"""
responses = [
"Final Answer: I am a response",
@ -350,7 +319,7 @@ async def test_get_result_and_thought(basic_graph):
assert llm_node is not None
llm_node._built_object = FakeListLLM(responses=responses)
llm_node._built = True
langchain_object = await basic_graph.build()
langchain_object = 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
@ -370,7 +339,9 @@ def test_find_last_node(grouped_chat_json_flow):
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
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
@ -422,9 +393,14 @@ def test_process_flow_one_group(one_grouped_chat_json_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)
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"
assert (
chat_openai_node["data"]["node"]["template"]["openai_api_key"]["value"]
== "test"
)
def test_process_flow_vector_store_grouped(vector_store_grouped_json_flow):
@ -471,15 +447,19 @@ def test_update_template(sample_template, sample_nodes):
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 is not 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 (
node1_updated["data"]["node"]["template"]["some_field"]["display_name"]
== "Name1"
)
assert node2_updated is not None
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"
assert (
node2_updated["data"]["node"]["template"]["other_field"]["display_name"]
== "DisplayName2"
)
# Ensure node3 remains unchanged
assert node3_updated == sample_nodes[2]
@ -510,7 +490,9 @@ def test_set_new_target_handle():
"data": {
"node": {
"flow": True,
"template": {"field_1": {"proxy": {"field": "new_field", "id": "new_id"}}},
"template": {
"field_1": {"proxy": {"field": "new_field", "id": "new_id"}}
},
}
}
}
@ -530,34 +512,34 @@ def test_update_source_handle():
"nodes": [{"id": "some_node"}, {"id": "last_node"}],
"edges": [{"source": "some_node"}],
}
updated_edge = update_source_handle(new_edge, flow_data["nodes"], flow_data["edges"])
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):
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()
first_result = graph.build()
assert isinstance(first_result, AgentExecutor)
pickled = pickle.dumps(graph)
assert pickled is not UnbuiltObject
assert pickled is not None
unpickled = pickle.loads(pickled)
assert unpickled is not UnbuiltObject
result = await unpickled.build()
assert unpickled is not None
result = unpickled.build()
assert isinstance(result, AgentExecutor)
@pytest.mark.asyncio
async def test_pickle_each_vertex(json_vector_store):
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()
for vertex in graph.nodes:
vertex.build()
pickled = pickle.dumps(vertex)
assert pickled is not UnbuiltObject
assert pickled is not None
unpickled = pickle.loads(pickled)
assert unpickled is not UnbuiltObject
assert unpickled is not None