114 lines
3.6 KiB
Python
114 lines
3.6 KiB
Python
import json
|
|
from langchain import LLMChain, OpenAI
|
|
import pytest
|
|
from pathlib import Path
|
|
from langflow import load_flow_from_json
|
|
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)
|
|
assert loaded is not None
|
|
|
|
|
|
def test_extract_json():
|
|
with open(EXAMPLE_JSON_PATH, "r") as f:
|
|
flow_graph = json.load(f)
|
|
data_graph = flow_graph["data"]
|
|
extracted = extract_json(data_graph)
|
|
assert extracted is not None
|
|
assert isinstance(extracted, dict)
|
|
|
|
|
|
def test_get_root_node():
|
|
with open(EXAMPLE_JSON_PATH, "r") as f:
|
|
flow_graph = json.load(f)
|
|
data_graph = flow_graph["data"]
|
|
nodes = data_graph["nodes"]
|
|
edges = data_graph["edges"]
|
|
root = get_root_node(nodes, edges)
|
|
assert root is not None
|
|
assert "id" in root
|
|
assert "data" in root
|
|
|
|
|
|
def test_build_json():
|
|
with open(EXAMPLE_JSON_PATH, "r") as f:
|
|
flow_graph = json.load(f)
|
|
data_graph = flow_graph["data"]
|
|
nodes = data_graph["nodes"]
|
|
edges = data_graph["edges"]
|
|
root = get_root_node(nodes, edges)
|
|
built_json = build_json(root, nodes, edges)
|
|
assert built_json is not None
|
|
assert isinstance(built_json, dict)
|
|
|
|
|
|
def test_build_json_missing_child():
|
|
with open(EXAMPLE_JSON_PATH, "r") as f:
|
|
flow_graph = json.load(f)
|
|
data_graph = flow_graph["data"]
|
|
nodes = data_graph["nodes"]
|
|
edges = data_graph["edges"]
|
|
|
|
# Modify nodes to create a missing required child scenario
|
|
for node in nodes:
|
|
if "data" in node and "node" in node["data"]:
|
|
for key, value in node["data"]["node"]["template"].items():
|
|
if isinstance(value, dict) and "required" in value:
|
|
value["required"] = True
|
|
|
|
root = get_root_node(nodes, edges)
|
|
with pytest.raises(ValueError):
|
|
build_json(root, nodes, edges)
|
|
|
|
|
|
def test_build_json_no_nodes():
|
|
with pytest.raises(TypeError):
|
|
build_json(None, [], [])
|
|
|
|
|
|
def test_build_json_invalid_edge():
|
|
with open(EXAMPLE_JSON_PATH, "r") as f:
|
|
flow_graph = json.load(f)
|
|
data_graph = flow_graph["data"]
|
|
nodes = data_graph["nodes"]
|
|
edges = data_graph["edges"]
|
|
# Modify edges to create an invalid edge scenario
|
|
for edge in edges:
|
|
edge["source"] = "invalid_id"
|
|
|
|
root = get_root_node(nodes, edges)
|
|
with pytest.raises(ValueError):
|
|
build_json(root, nodes, edges)
|
|
|
|
|
|
def test_load_langchain_type_from_config():
|
|
with open(EXAMPLE_JSON_PATH, "r") as f:
|
|
flow_graph = json.load(f)
|
|
data_graph = flow_graph["data"]
|
|
extracted = extract_json(data_graph)
|
|
|
|
agent_config = extracted.copy()
|
|
agent_type = "AgentExecutor" # Replace with the actual agent type in the JSON
|
|
|
|
invalid_config = extracted.copy()
|
|
invalid_config["_type"] = "invalid_type"
|
|
|
|
agent_loaded = load_langchain_type_from_config(agent_config)
|
|
assert agent_loaded is not None
|
|
assert agent_loaded.__class__.__name__ == agent_type
|
|
assert hasattr(agent_loaded.agent, "llm_chain")
|
|
assert isinstance(
|
|
agent_loaded.agent.llm_chain, LLMChain
|
|
) # Replace Chain with the appropriate class
|
|
assert hasattr(agent_loaded.agent.llm_chain, "llm")
|
|
assert isinstance(agent_loaded.agent.llm_chain.llm, OpenAI)
|
|
|
|
with pytest.raises(ValueError):
|
|
load_langchain_type_from_config(invalid_config)
|