feat: added tests for graph

This commit is contained in:
Gabriel Almeida 2023-03-24 11:21:29 -03:00
commit 269c06e9c1
6 changed files with 859 additions and 27 deletions

View file

@ -2,7 +2,7 @@ from typing import Dict, List, Union
class Node:
def __init__(self, data: Dict[str, Union[str, Dict[str, Union[str, List[str]]]]]):
def __init__(self, data: Dict):
self.id: str = data["id"]
self._data = data
self.edges: List[Edge] = []
@ -62,26 +62,36 @@ class Graph:
connected_nodes.append(edge.source)
return connected_nodes
def get_node_neighbors(self, node_id: str) -> Dict[str, int]:
neighbors: Dict[str, int] = {}
def get_node_neighbors(self, node: Node) -> Dict[str, int]:
neighbors: Dict[Node, int] = {}
for edge in self.edges:
if edge.source.id == node_id:
neighbor_id = edge.target.id
if neighbor_id not in neighbors:
neighbors[neighbor_id] = 0
neighbors[neighbor_id] += 1
elif edge.target.id == node_id:
neighbor_id = edge.source.id
if neighbor_id not in neighbors:
neighbors[neighbor_id] = 0
neighbors[neighbor_id] += 1
if edge.source == node:
neighbor = edge.target
if neighbor not in neighbors:
neighbors[neighbor] = 0
neighbors[neighbor] += 1
elif edge.target == node:
neighbor = edge.source
if neighbor not in neighbors:
neighbors[neighbor] = 0
neighbors[neighbor] += 1
return neighbors
def _build_edges(self) -> List[Edge]:
return [
Edge(self.get_node(edge["source"]), self.get_node(edge["target"]))
for edge in self._edges
]
# Edge takes two nodes as arguments, so we need to build the nodes first
# and then build the edges
# if we can't find a node, we raise an error
edges: List[Edge] = []
for edge in self._edges:
source = self.get_node(edge["source"])
target = self.get_node(edge["target"])
if source is None:
raise ValueError(f"Source node {edge['source']} not found")
if target is None:
raise ValueError(f"Target node {edge['target']} not found")
edges.append(Edge(source, target))
return edges
def _build_nodes(self) -> List[Node]:
return [Node(node) for node in self._nodes]

11
tests/conftest.py Normal file
View file

@ -0,0 +1,11 @@
from pathlib import Path
import pytest
def pytest_configure():
pytest.BASIC_EXAMPLE_PATH = (
Path(__file__).parent.absolute() / "data" / "basic_example.json"
)
pytest.COMPLEX_EXAMPLE_PATH = (
Path(__file__).parent.absolute() / "data" / "complex_example.json"
)

View file

@ -0,0 +1,741 @@
{
"name": "New Flow ",
"id": "0",
"data": {
"nodes": [
{
"width": 384,
"height": 477,
"id": "dndnode_7",
"type": "genericNode",
"position": {
"x": -211.61829328351757,
"y": 132.6841414309356
},
"data": {
"type": "OpenAI",
"node": {
"template": {
"_type": "openai",
"cache": {
"type": "bool",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": null
},
"verbose": {
"type": "bool",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": null
},
"client": {
"type": "Any",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": null
},
"model_name": {
"type": "str",
"required": false,
"placeholder": "",
"list": false,
"show": true,
"password": false,
"multiline": false,
"value": "text-davinci-003",
"options": [
"text-davinci-003",
"text-davinci-002"
]
},
"temperature": {
"type": "float",
"required": false,
"placeholder": "",
"list": false,
"show": true,
"password": false,
"multiline": false,
"value": 0.7
},
"max_tokens": {
"type": "int",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": true,
"multiline": false,
"value": 256
},
"top_p": {
"type": "float",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": 1
},
"frequency_penalty": {
"type": "float",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": 0
},
"presence_penalty": {
"type": "float",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": 0
},
"n": {
"type": "int",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": 1
},
"best_of": {
"type": "int",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": 1
},
"model_kwargs": {
"type": "dict[str, Any]",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": null
},
"openai_api_key": {
"type": "str",
"required": false,
"placeholder": "",
"list": false,
"show": true,
"password": true,
"multiline": false,
"value": null
},
"batch_size": {
"type": "int",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": 20
},
"request_timeout": {
"type": "Union[float, Tuple[float, float], NoneType]",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": null
},
"logit_bias": {
"type": "dict[str, float]",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": null
},
"max_retries": {
"type": "int",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": 6
},
"streaming": {
"type": "bool",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": false
}
},
"description": "Generic OpenAI class that uses model name.",
"base_classes": [
"BaseOpenAI",
"BaseLLM",
"BaseLanguageModel"
]
},
"id": "dndnode_7",
"value": null
},
"selected": false,
"positionAbsolute": {
"x": -211.61829328351757,
"y": 132.6841414309356
},
"dragging": false
},
{
"width": 384,
"height": 351,
"id": "dndnode_8",
"type": "genericNode",
"position": {
"x": 429.5817067164825,
"y": 552.6841414309356
},
"data": {
"type": "PAL-MATH",
"node": {
"template": {
"llm": {
"type": "BaseLLM",
"required": true,
"list": false,
"show": true
},
"_type": "pal-math"
},
"name": "PAL-MATH",
"description": "A language model that is really good at solving complex word math problems. Input should be a fully worded hard word math problem.",
"base_classes": [
"Tool"
]
},
"id": "dndnode_8",
"value": null
},
"selected": false,
"positionAbsolute": {
"x": 429.5817067164825,
"y": 552.6841414309356
},
"dragging": false
},
{
"width": 384,
"height": 529,
"id": "dndnode_9",
"type": "genericNode",
"position": {
"x": -238.41829328351753,
"y": 734.6841414309356
},
"data": {
"type": "ZeroShotPrompt",
"node": {
"template": {
"_type": "zero_shot",
"prefix": {
"type": "str",
"required": false,
"placeholder": "",
"list": false,
"show": true,
"multiline": true,
"value": "Answer the following questions as best you can. You have access to the following tools:"
},
"suffix": {
"type": "str",
"required": true,
"placeholder": "",
"list": false,
"show": true,
"multiline": true,
"value": "Begin!\n\nQuestion: {input}\nThought:{agent_scratchpad}"
},
"format_instructions": {
"type": "str",
"required": false,
"placeholder": "",
"list": false,
"show": true,
"multiline": true,
"value": "Use the following format:\n\nQuestion: the input question you must answer\nThought: you should always think about what to do\nAction: the action to take, should be one of [{tool_names}]\nAction Input: the input to the action\nObservation: the result of the action\n... (this Thought/Action/Action Input/Observation can repeat N times)\nThought: I now know the final answer\nFinal Answer: the final answer to the original input question"
}
},
"description": "Prompt template for Zero Shot Agent.",
"base_classes": [
"BasePromptTemplate"
]
},
"id": "dndnode_9",
"value": null
},
"selected": false,
"positionAbsolute": {
"x": -238.41829328351753,
"y": 734.6841414309356
},
"dragging": false
},
{
"width": 384,
"height": 351,
"id": "dndnode_34",
"type": "genericNode",
"position": {
"x": 931.431035174925,
"y": 285.11234969236074
},
"data": {
"type": "ZeroShotAgent",
"node": {
"template": {
"_type": "zero-shot-react-description",
"llm_chain": {
"type": "LLMChain",
"required": true,
"placeholder": "",
"list": false,
"show": true,
"password": false,
"multiline": false
},
"allowed_tools": {
"type": "Tool",
"required": false,
"placeholder": "",
"list": true,
"show": true,
"password": false,
"multiline": false,
"value": null
},
"return_values": {
"type": "str",
"required": false,
"placeholder": "",
"list": true,
"show": false,
"password": false,
"multiline": false,
"value": [
"output"
]
}
},
"description": "Agent for the MRKL chain.",
"base_classes": [
"Agent",
"function"
]
},
"id": "dndnode_34",
"value": null
},
"selected": false,
"positionAbsolute": {
"x": 931.431035174925,
"y": 285.11234969236074
},
"dragging": false
},
{
"width": 384,
"height": 523,
"id": "dndnode_41",
"type": "genericNode",
"position": {
"x": 1456.9923517899285,
"y": 180.9267225160853
},
"data": {
"type": "BaseTool",
"node": {
"template": {
"name": {
"type": "str",
"required": true,
"list": false,
"show": true,
"placeholder": "",
"value": ""
},
"description": {
"type": "str",
"required": true,
"list": false,
"show": true,
"placeholder": "",
"value": ""
},
"func": {
"type": "function",
"required": true,
"list": false,
"show": true,
"value": "",
"multiline": true
},
"_type": "BaseTool"
},
"name": "PAL-MATH",
"description": "A language model that is really good at solving complex word math problems. Input should be a fully worded hard word math problem.",
"base_classes": [
"Tool"
]
},
"id": "dndnode_41",
"value": null
},
"selected": false,
"positionAbsolute": {
"x": 1456.9923517899285,
"y": 180.9267225160853
},
"dragging": false
},
{
"width": 384,
"height": 351,
"id": "dndnode_42",
"type": "genericNode",
"position": {
"x": 1994.12084226133,
"y": -48.749412190849014
},
"data": {
"type": "ZeroShotAgent",
"node": {
"template": {
"_type": "zero-shot-react-description",
"llm_chain": {
"type": "LLMChain",
"required": true,
"placeholder": "",
"list": false,
"show": true,
"password": false,
"multiline": false
},
"allowed_tools": {
"type": "Tool",
"required": false,
"placeholder": "",
"list": true,
"show": true,
"password": false,
"multiline": false,
"value": null
},
"return_values": {
"type": "str",
"required": false,
"placeholder": "",
"list": true,
"show": false,
"password": false,
"multiline": false,
"value": [
"output"
]
}
},
"description": "Agent for the MRKL chain.",
"base_classes": [
"Agent",
"function"
]
},
"id": "dndnode_42",
"value": null
},
"selected": false,
"positionAbsolute": {
"x": 1994.12084226133,
"y": -48.749412190849014
},
"dragging": false
},
{
"width": 384,
"height": 391,
"id": "dndnode_43",
"type": "genericNode",
"position": {
"x": 394.51027997668166,
"y": 45.75513998834094
},
"data": {
"type": "LLMChain",
"node": {
"template": {
"_type": "llm_chain",
"memory": {
"type": "BaseMemory",
"required": false,
"placeholder": "",
"list": false,
"show": true,
"password": false,
"multiline": false,
"value": null
},
"verbose": {
"type": "bool",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": false
},
"prompt": {
"type": "BasePromptTemplate",
"required": true,
"placeholder": "",
"list": false,
"show": true,
"password": false,
"multiline": false
},
"llm": {
"type": "BaseLanguageModel",
"required": true,
"placeholder": "",
"list": false,
"show": true,
"password": false,
"multiline": false
},
"output_key": {
"type": "str",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": true,
"multiline": false,
"value": "text"
}
},
"description": "Chain to run queries against LLMs.",
"base_classes": [
"Chain"
]
},
"id": "dndnode_43",
"value": null
},
"selected": false,
"positionAbsolute": {
"x": 394.51027997668166,
"y": 45.75513998834094
},
"dragging": false
},
{
"width": 384,
"height": 391,
"id": "dndnode_44",
"type": "genericNode",
"position": {
"x": 1404.5102799766814,
"y": -328.24486001165906
},
"data": {
"type": "LLMChain",
"node": {
"template": {
"_type": "llm_chain",
"memory": {
"type": "BaseMemory",
"required": false,
"placeholder": "",
"list": false,
"show": true,
"password": false,
"multiline": false,
"value": null
},
"verbose": {
"type": "bool",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": false,
"multiline": false,
"value": false
},
"prompt": {
"type": "BasePromptTemplate",
"required": true,
"placeholder": "",
"list": false,
"show": true,
"password": false,
"multiline": false
},
"llm": {
"type": "BaseLanguageModel",
"required": true,
"placeholder": "",
"list": false,
"show": true,
"password": false,
"multiline": false
},
"output_key": {
"type": "str",
"required": false,
"placeholder": "",
"list": false,
"show": false,
"password": true,
"multiline": false,
"value": "text"
}
},
"description": "Chain to run queries against LLMs.",
"base_classes": [
"Chain"
]
},
"id": "dndnode_44",
"value": null
},
"selected": false,
"positionAbsolute": {
"x": 1404.5102799766814,
"y": -328.24486001165906
},
"dragging": false
}
],
"edges": [
{
"source": "dndnode_7",
"sourceHandle": "OpenAI|dndnode_7|BaseOpenAI|BaseLLM|BaseLanguageModel",
"target": "dndnode_8",
"targetHandle": "BaseLLM|llm|dndnode_8",
"className": "animate-pulse",
"id": "reactflow__edge-dndnode_7OpenAI|dndnode_7|BaseOpenAI|BaseLLM|BaseLanguageModel-dndnode_8BaseLLM|llm|dndnode_8"
},
{
"source": "dndnode_8",
"sourceHandle": "PAL-MATH|dndnode_8|Tool",
"target": "dndnode_34",
"targetHandle": "Tool|allowed_tools|dndnode_34",
"className": "animate-pulse",
"id": "reactflow__edge-dndnode_8PAL-MATH|dndnode_8|Tool-dndnode_34Tool|allowed_tools|dndnode_34"
},
{
"source": "dndnode_34",
"sourceHandle": "ZeroShotAgent|dndnode_34|Agent|function",
"target": "dndnode_41",
"targetHandle": "function|func|dndnode_41",
"className": "animate-pulse",
"id": "reactflow__edge-dndnode_34ZeroShotAgent|dndnode_34|Agent|function-dndnode_41function|func|dndnode_41"
},
{
"source": "dndnode_41",
"sourceHandle": "BaseTool|dndnode_41|Tool",
"target": "dndnode_42",
"targetHandle": "Tool|allowed_tools|dndnode_42",
"className": "animate-pulse",
"id": "reactflow__edge-dndnode_41BaseTool|dndnode_41|Tool-dndnode_42Tool|allowed_tools|dndnode_42"
},
{
"source": "dndnode_9",
"sourceHandle": "ZeroShotPrompt|dndnode_9|BasePromptTemplate",
"target": "dndnode_43",
"targetHandle": "BasePromptTemplate|prompt|dndnode_43",
"className": "animate-pulse",
"id": "reactflow__edge-dndnode_9ZeroShotPrompt|dndnode_9|BasePromptTemplate-dndnode_43BasePromptTemplate|prompt|dndnode_43"
},
{
"source": "dndnode_7",
"sourceHandle": "OpenAI|dndnode_7|BaseOpenAI|BaseLLM|BaseLanguageModel",
"target": "dndnode_43",
"targetHandle": "BaseLanguageModel|llm|dndnode_43",
"className": "animate-pulse",
"id": "reactflow__edge-dndnode_7OpenAI|dndnode_7|BaseOpenAI|BaseLLM|BaseLanguageModel-dndnode_43BaseLanguageModel|llm|dndnode_43"
},
{
"source": "dndnode_43",
"sourceHandle": "LLMChain|dndnode_43|Chain",
"target": "dndnode_34",
"targetHandle": "LLMChain|llm_chain|dndnode_34",
"className": "animate-pulse",
"id": "reactflow__edge-dndnode_43LLMChain|dndnode_43|Chain-dndnode_34LLMChain|llm_chain|dndnode_34"
},
{
"source": "dndnode_44",
"sourceHandle": "LLMChain|dndnode_44|Chain",
"target": "dndnode_42",
"targetHandle": "LLMChain|llm_chain|dndnode_42",
"className": "animate-pulse",
"id": "reactflow__edge-dndnode_44LLMChain|dndnode_44|Chain-dndnode_42LLMChain|llm_chain|dndnode_42"
},
{
"source": "dndnode_9",
"sourceHandle": "ZeroShotPrompt|dndnode_9|BasePromptTemplate",
"target": "dndnode_44",
"targetHandle": "BasePromptTemplate|prompt|dndnode_44",
"className": "animate-pulse",
"id": "reactflow__edge-dndnode_9ZeroShotPrompt|dndnode_9|BasePromptTemplate-dndnode_44BasePromptTemplate|prompt|dndnode_44"
},
{
"source": "dndnode_7",
"sourceHandle": "OpenAI|dndnode_7|BaseOpenAI|BaseLLM|BaseLanguageModel",
"target": "dndnode_44",
"targetHandle": "BaseLanguageModel|llm|dndnode_44",
"className": "animate-pulse",
"id": "reactflow__edge-dndnode_7OpenAI|dndnode_7|BaseOpenAI|BaseLLM|BaseLanguageModel-dndnode_44BaseLanguageModel|llm|dndnode_44"
}
],
"viewport": {
"x": -12.25513998834083,
"y": 135.5224300058294,
"zoom": 0.5
}
},
"chat": []
}

72
tests/test_graph.py Normal file
View file

@ -0,0 +1,72 @@
import json
from langflow.utils.graph import Graph
import pytest
from langflow.utils.payload import get_root_node
# Test cases for the graph module
def get_graph(basic=True):
"""Get a graph from a json file"""
path = pytest.BASIC_EXAMPLE_PATH if basic else pytest.COMPLEX_EXAMPLE_PATH
with open(path, "r") as f:
flow_graph = json.load(f)
data_graph = flow_graph["data"]
nodes = data_graph["nodes"]
edges = data_graph["edges"]
return Graph(nodes, edges)
def test_get_connected_nodes():
"""Test getting connected nodes"""
graph = get_graph()
assert isinstance(graph, Graph)
# Get root node
root = get_root_node(graph)
assert root is not None
connected_nodes = graph.get_connected_nodes(root)
assert connected_nodes is not None
def test_get_node_neighbors():
"""Test getting node neighbors"""
graph = get_graph(basic=True)
assert isinstance(graph, Graph)
# Get root node
root = get_root_node(graph)
assert root is not None
neighbors = 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(
"Chain" in neighbor.data["type"] for neighbor, val in neighbors.items() if val
)
# assert Serper Search is in the neighbors
assert any(
"Serper" in neighbor.data["type"] for neighbor, val in neighbors.items() if val
)
# Now on to the Chain's neighbors
chain = next(
neighbor
for neighbor, val in neighbors.items()
if "Chain" in neighbor.data["type"] and val
)
chain_neighbors = graph.get_node_neighbors(chain)
assert chain_neighbors is not None
assert isinstance(chain_neighbors, dict)
# Check if there is a LLM in the chain's neighbors
assert any(
"OpenAI" in neighbor.data["type"]
for neighbor, val in chain_neighbors.items()
if val
)
# Chain should have a Prompt as a neighbor
assert any(
"Prompt" in neighbor.data["type"]
for neighbor, val in chain_neighbors.items()
if val
)

View file

@ -8,17 +8,15 @@ from langflow.interface.loading import extract_json
from langflow.utils.payload import get_root_node, build_json
from langflow.interface.loading import load_langchain_type_from_config
EXAMPLE_JSON_PATH = Path(__file__).parent.absolute() / "data" / "example_flow.json"
def test_load_flow_from_json():
"""Test loading a flow from a json file"""
loaded = load_flow_from_json(EXAMPLE_JSON_PATH)
loaded = load_flow_from_json(pytest.EXAMPLE_JSON_PATH)
assert loaded is not None
def test_extract_json():
with open(EXAMPLE_JSON_PATH, "r") as f:
with open(pytest.EXAMPLE_JSON_PATH, "r") as f:
flow_graph = json.load(f)
data_graph = flow_graph["data"]
extracted = extract_json(data_graph)
@ -27,7 +25,7 @@ def test_extract_json():
def test_get_root_node():
with open(EXAMPLE_JSON_PATH, "r") as f:
with open(pytest.EXAMPLE_JSON_PATH, "r") as f:
flow_graph = json.load(f)
data_graph = flow_graph["data"]
nodes = data_graph["nodes"]
@ -40,7 +38,7 @@ def test_get_root_node():
def test_build_json():
with open(EXAMPLE_JSON_PATH, "r") as f:
with open(pytest.EXAMPLE_JSON_PATH, "r") as f:
flow_graph = json.load(f)
data_graph = flow_graph["data"]
nodes = data_graph["nodes"]
@ -53,7 +51,7 @@ def test_build_json():
def test_build_json_missing_child():
with open(EXAMPLE_JSON_PATH, "r") as f:
with open(pytest.EXAMPLE_JSON_PATH, "r") as f:
flow_graph = json.load(f)
data_graph = flow_graph["data"]
nodes = data_graph["nodes"]
@ -78,7 +76,7 @@ def test_build_json_no_nodes():
def test_build_json_invalid_edge():
with open(EXAMPLE_JSON_PATH, "r") as f:
with open(pytest.EXAMPLE_JSON_PATH, "r") as f:
flow_graph = json.load(f)
data_graph = flow_graph["data"]
nodes = data_graph["nodes"]
@ -87,14 +85,14 @@ def test_build_json_invalid_edge():
for edge in edges:
edge["source"] = "invalid_id"
with pytest.raises(AttributeError):
with pytest.raises(ValueError):
graph = Graph(nodes, edges)
root = get_root_node(graph)
build_json(root, nodes, edges)
def test_load_langchain_type_from_config():
with open(EXAMPLE_JSON_PATH, "r") as f:
with open(pytest.EXAMPLE_JSON_PATH, "r") as f:
flow_graph = json.load(f)
data_graph = flow_graph["data"]
extracted = extract_json(data_graph)