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

Refactor imports and update function names

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-01-25 16:17:40 -03:00
commit abe4925cc3
8 changed files with 156 additions and 314 deletions
Download patch file Download diff file Expand all files Collapse all files

136
tests/test_graph.py
View file

@ -1,32 +1,29 @@
import copy
import json
import os
from pathlib import Path
import pickle
from pathlib import Path
from typing import Type, Union
from langflow.graph.edge.base import Edge
from langflow.graph.vertex.base import Vertex
from langchain.agents import AgentExecutor
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.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.edge.base import Edge
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.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
@ -47,13 +44,7 @@ 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",
@ -71,11 +62,7 @@ def sample_nodes():
},
{
"id": "node3",
"data": {
"node": {
"template": {"unrelated_field": {"show": True, "advanced": True}}
}
},
"data": {"node": {"template": {"unrelated_field": {"show": True, "advanced": True}}}},
},
]
@ -93,8 +80,10 @@ def test_graph_structure(basic_graph):
assert isinstance(node, Vertex)
for edge in basic_graph.edges:
assert isinstance(edge, Edge)
assert edge.source in basic_graph.vertices
assert edge.target in basic_graph.vertices
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):
@ -102,7 +91,7 @@ def test_circular_dependencies(basic_graph):
def check_circular(node, visited):
visited.add(node)
neighbors = basic_graph.get_nodes_with_target(node)
neighbors = basic_graph.get_vertices_with_target(node)
for neighbor in neighbors:
if neighbor in visited:
return True
@ -135,13 +124,13 @@ def test_invalid_node_types():
Graph(graph_data["nodes"], graph_data["edges"])
def test_get_nodes_with_target(basic_graph):
def test_get_vertices_with_target(basic_graph):
"""Test getting connected nodes"""
assert isinstance(basic_graph, Graph)
# Get root node
root = get_root_node(basic_graph)
root = get_root_vertex(basic_graph)
assert root is not None
connected_nodes = basic_graph.get_nodes_with_target(root)
connected_nodes = basic_graph.get_vertices_with_target(root.id)
assert connected_nodes is not None
@ -150,23 +139,17 @@ def test_get_node_neighbors_basic(basic_graph):
assert isinstance(basic_graph, Graph)
# Get root node
root = get_root_node(basic_graph)
root = get_root_vertex(basic_graph)
assert root is not None
neighbors = basic_graph.get_node_neighbors(root)
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("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
)
assert any("OpenAI" in neighbor.data["type"] for neighbor, val in neighbors.items() if val)
def test_get_node(basic_graph):
@ -180,7 +163,7 @@ def test_get_node(basic_graph):
def test_build_nodes(basic_graph):
"""Test building nodes"""
assert len(basic_graph.vertices) == len(basic_graph._nodes)
assert len(basic_graph.vertices) == len(basic_graph._vertices)
for node in basic_graph.vertices:
assert isinstance(node, Vertex)
@ -190,20 +173,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, Vertex)
assert isinstance(edge.target, Vertex)
assert isinstance(edge.source_id, str)
assert isinstance(edge.target_id, str)
def test_get_root_node(client, basic_graph, complex_graph):
def test_get_root_vertex(client, basic_graph, complex_graph):
"""Test getting root node"""
assert isinstance(basic_graph, Graph)
root = get_root_node(basic_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_node(complex_graph)
root = get_root_vertex(complex_graph)
assert root is not None
assert isinstance(root, Vertex)
assert root.data["type"] == "ZeroShotAgent"
@ -239,7 +222,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_node(basic_graph)
root = get_root_vertex(basic_graph)
# Root node is a TimeTravelGuideChain
# which requires an llm and memory
assert root is not None
@ -248,17 +231,18 @@ def test_build_params(basic_graph):
assert "memory" in root.params
def test_build(basic_graph):
@pytest.mark.asyncio
async def test_build(basic_graph):
"""Test Node's build method"""
assert_agent_was_built(basic_graph)
await assert_agent_was_built(basic_graph)
def assert_agent_was_built(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 = graph.build()
result = await graph.build()
# The agent should be a AgentExecutor
assert isinstance(result, Chain)
@ -307,7 +291,8 @@ def test_file_tool_node_build(client, openapi_graph):
# assert built_object is not None
def test_get_result_and_thought(basic_graph):
@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",
@ -319,7 +304,7 @@ 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 = basic_graph.build()
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
@ -339,9 +324,7 @@ 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
@ -393,14 +376,9 @@ 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):
@ -449,17 +427,11 @@ def test_update_template(sample_template, sample_nodes):
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["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]
@ -490,9 +462,7 @@ 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"}}},
}
}
}
@ -512,33 +482,33 @@ 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"
def test_pickle_graph(json_vector_store):
@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 = graph.build()
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 = unpickled.build()
result = await unpickled.build()
assert isinstance(result, AgentExecutor)
def test_pickle_each_vertex(json_vector_store):
@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.nodes:
vertex.build()
for vertex in graph.vertices:
await vertex.build()
pickled = pickle.dumps(vertex)
assert pickled is not None
unpickled = pickle.loads(pickled)