feat: add dump and dumps methods to Graph (#3202)
* feat(utils.py): add escape_json_dump function to escape JSON strings for Edge dictionaries * refactor(Output): streamline add_types method to prevent duplicate entries in types list for improved type management * feat(data.py): add classmethod decorator to validate_data for enhanced validation logic when checking data types * feat(setup.py): implement retry logic for loading starter projects to enhance robustness against JSON decode errors * fix(input_mixin.py): improve model_config formatting and update field_type alias for clarity and consistency in field definitions * feat(types.py): refactor vertex constructors to use NodeData and add input/output methods for better component interaction * feat(schema.py): add NodeData and Position TypedDicts for improved type safety and structure in vertex data handling * feat(base.py): update Vertex to use NodeData type and add to_data method for better data management and access * refactor(schema.py): update TargetHandle and SourceHandle models to include model_config attribute * Add TypedDict classes for graph schema serialization in `schema.py` * Refactor `Edge` class to improve handle validation and data handling - Consolidated imports and removed redundant `BaseModel` definitions for `SourceHandle` and `TargetHandle`. - Added `valid_handles`, `target_param`, and `_target_handle` attributes to `Edge` class. - Enhanced handle validation logic to distinguish between dictionary and string types. - Introduced `to_data` method to return edge data. - Updated attribute names to follow consistent naming conventions (`base_classes`, `input_types`, `field_name`). * Refactor `Edge` class to improve handle validation and data handling * Refactor: Standardize attribute naming and add `to_data` method in Edge class - Renamed attributes to use snake_case consistently (`baseClasses` to `base_classes`, `inputTypes` to `input_types`, `fieldName` to `field_name`). - Added `to_data` method to return `_data` attribute. - Updated validation methods to use new attribute names. * Refactor: Update Edge class to consistently use snake_case for attributes and improve validation logic for handles * Refactor: Change node argument type in add_node and _create_vertex methods to NodeData for better type safety and clarity * Refactor: Implement JSON serialization for graph data with `dumps` and `dump` methods, enhancing data export capabilities * Refactor: Add pytest fixtures for ingestion and RAG graphs, enhance test structure for better clarity and organization * Refactor: Add pytest fixtures for memory_chatbot_graph tests and improve test structure * Refactor: Remove unused methods in ComponentVertex class to streamline code and improve readability * Refactor: Remove unnecessary line in ComponentVertex class to enhance code clarity and maintainability * Refactor: Update import path for DefaultPromptField to improve code organization and maintainability in api_utils.py * Refactor: Update import path for DefaultPromptField to enhance code organization and maintainability in prompt.py * fix: Remove fixture in test_memory_chatbot.py that blocked db setup * Refactor: Add durations path for unit tests to improve test reporting * Refactor: Add splitting algorithm option for unit tests * Add async option to Makefile for unit tests and update GitHub Actions workflow - Introduced `async` variable in Makefile to conditionally run unit tests with or without parallel execution. - Updated `unit_tests` target in Makefile to handle `async` flag. - Modified GitHub Actions workflow to set `async=false` for unit tests.
This commit is contained in:
Gabriel Luiz Freitas Almeida
GitHub
parent
8ece1cadf6
commit
bb1bc5c2df
19 changed files with 797 additions and 107 deletions
|
|
@ -1,5 +1,7 @@
|
|||
from collections import deque
|
||||
|
||||
import pytest
|
||||
|
||||
from langflow.components.helpers.Memory import MemoryComponent
|
||||
from langflow.components.inputs.ChatInput import ChatInput
|
||||
from langflow.components.models.OpenAIModel import OpenAIModelComponent
|
||||
|
|
@ -7,9 +9,11 @@ from langflow.components.outputs.ChatOutput import ChatOutput
|
|||
from langflow.components.prompts.Prompt import PromptComponent
|
||||
from langflow.graph import Graph
|
||||
from langflow.graph.graph.constants import Finish
|
||||
from langflow.graph.graph.schema import GraphDump
|
||||
|
||||
|
||||
def test_memory_chatbot():
|
||||
@pytest.fixture
|
||||
def memory_chatbot_graph():
|
||||
session_id = "test_session_id"
|
||||
template = """{context}
|
||||
|
||||
|
|
@ -32,10 +36,87 @@ AI: """
|
|||
chat_output.set(input_value=openai_component.text_response)
|
||||
|
||||
graph = Graph(chat_input, chat_output)
|
||||
return graph
|
||||
|
||||
|
||||
def test_memory_chatbot(memory_chatbot_graph):
|
||||
# Now we run step by step
|
||||
expected_order = deque(["chat_input", "chat_memory", "prompt", "openai", "chat_output"])
|
||||
for step in expected_order:
|
||||
result = graph.step()
|
||||
result = memory_chatbot_graph.step()
|
||||
if isinstance(result, Finish):
|
||||
break
|
||||
assert step == result.vertex.id
|
||||
|
||||
|
||||
def test_memory_chatbot_dump_structure(memory_chatbot_graph: Graph):
|
||||
# Now we run step by step
|
||||
graph_dict = memory_chatbot_graph.dump(
|
||||
name="Memory Chatbot", description="A memory chatbot", endpoint_name="membot"
|
||||
)
|
||||
assert isinstance(graph_dict, dict)
|
||||
# Test structure
|
||||
assert "data" in graph_dict
|
||||
assert "is_component" in graph_dict
|
||||
|
||||
data_dict = graph_dict["data"]
|
||||
assert "nodes" in data_dict
|
||||
assert "edges" in data_dict
|
||||
assert "description" in graph_dict
|
||||
assert "endpoint_name" in graph_dict
|
||||
|
||||
# Test data
|
||||
nodes = data_dict["nodes"]
|
||||
edges = data_dict["edges"]
|
||||
description = graph_dict["description"]
|
||||
endpoint_name = graph_dict["endpoint_name"]
|
||||
|
||||
assert len(nodes) == 5
|
||||
assert len(edges) == 4
|
||||
assert description is not None
|
||||
assert endpoint_name is not None
|
||||
|
||||
|
||||
def test_memory_chatbot_dump_components_and_edges(memory_chatbot_graph: Graph):
|
||||
# Check all components and edges were dumped correctly
|
||||
graph_dict: GraphDump = memory_chatbot_graph.dump(
|
||||
name="Memory Chatbot", description="A memory chatbot", endpoint_name="membot"
|
||||
)
|
||||
|
||||
data_dict = graph_dict["data"]
|
||||
nodes = data_dict["nodes"]
|
||||
edges = data_dict["edges"]
|
||||
|
||||
# sort the nodes by id
|
||||
nodes = sorted(nodes, key=lambda x: x["id"])
|
||||
|
||||
# Check each node
|
||||
assert nodes[0]["data"]["type"] == "ChatInput"
|
||||
assert nodes[0]["id"] == "chat_input"
|
||||
|
||||
assert nodes[1]["data"]["type"] == "MemoryComponent"
|
||||
assert nodes[1]["id"] == "chat_memory"
|
||||
|
||||
assert nodes[2]["data"]["type"] == "ChatOutput"
|
||||
assert nodes[2]["id"] == "chat_output"
|
||||
|
||||
assert nodes[3]["data"]["type"] == "OpenAIModelComponent"
|
||||
assert nodes[3]["id"] == "openai"
|
||||
|
||||
assert nodes[4]["data"]["type"] == "PromptComponent"
|
||||
assert nodes[4]["id"] == "prompt"
|
||||
|
||||
# Check edges
|
||||
expected_edges = [
|
||||
("chat_input", "prompt"),
|
||||
("chat_memory", "prompt"),
|
||||
("prompt", "openai"),
|
||||
("openai", "chat_output"),
|
||||
]
|
||||
|
||||
assert len(edges) == len(expected_edges)
|
||||
|
||||
for edge in edges:
|
||||
source = edge["source"]
|
||||
target = edge["target"]
|
||||
assert (source, target) in expected_edges, edge
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
from textwrap import dedent
|
||||
|
||||
import pytest
|
||||
|
||||
from langflow.components.data.File import FileComponent
|
||||
from langflow.components.embeddings.OpenAIEmbeddings import OpenAIEmbeddingsComponent
|
||||
from langflow.components.helpers.ParseData import ParseDataComponent
|
||||
|
|
@ -14,10 +16,17 @@ from langflow.graph.graph.constants import Finish
|
|||
from langflow.schema.data import Data
|
||||
|
||||
|
||||
def test_vector_store_rag():
|
||||
@pytest.fixture
|
||||
def client():
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ingestion_graph():
|
||||
# Ingestion Graph
|
||||
file_component = FileComponent(_id="file-123")
|
||||
file_component.set(path="test.txt")
|
||||
file_component.set_output_value("data", Data(text="This is a test file."))
|
||||
text_splitter = SplitTextComponent(_id="text-splitter-123")
|
||||
text_splitter.set(data_inputs=file_component.load_file)
|
||||
openai_embeddings = OpenAIEmbeddingsComponent(_id="openai-embeddings-123")
|
||||
|
|
@ -31,8 +40,18 @@ def test_vector_store_rag():
|
|||
api_endpoint="https://astra.example.com",
|
||||
token="token",
|
||||
)
|
||||
vector_store.set_output_value("vector_store", "mock_vector_store")
|
||||
vector_store.set_output_value("base_retriever", "mock_retriever")
|
||||
vector_store.set_output_value("search_results", [Data(text="This is a test file.")])
|
||||
|
||||
ingestion_graph = Graph(file_component, vector_store)
|
||||
return ingestion_graph
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rag_graph():
|
||||
# RAG Graph
|
||||
openai_embeddings = OpenAIEmbeddingsComponent(_id="openai-embeddings-124")
|
||||
chat_input = ChatInput(_id="chatinput-123")
|
||||
chat_input.get_output("message").value = "What is the meaning of life?"
|
||||
rag_vector_store = AstraVectorStoreComponent(_id="rag-vector-store-123")
|
||||
|
|
@ -69,21 +88,160 @@ def test_vector_store_rag():
|
|||
chat_output.set(input_value=openai_component.text_response)
|
||||
|
||||
graph = Graph(start=chat_input, end=chat_output)
|
||||
assert graph is not None
|
||||
ids = [
|
||||
return graph
|
||||
|
||||
|
||||
def test_vector_store_rag(ingestion_graph, rag_graph):
|
||||
assert ingestion_graph is not None
|
||||
ingestion_ids = [
|
||||
"file-123",
|
||||
"text-splitter-123",
|
||||
"openai-embeddings-123",
|
||||
"vector-store-123",
|
||||
]
|
||||
assert rag_graph is not None
|
||||
rag_ids = [
|
||||
"chatinput-123",
|
||||
"chatoutput-123",
|
||||
"openai-123",
|
||||
"parse-data-123",
|
||||
"prompt-123",
|
||||
"rag-vector-store-123",
|
||||
"openai-embeddings-123",
|
||||
"openai-embeddings-124",
|
||||
]
|
||||
results = []
|
||||
for result in graph.start():
|
||||
results.append(result)
|
||||
for ids, graph, len_results in zip([ingestion_ids, rag_ids], [ingestion_graph, rag_graph], [5, 8]):
|
||||
results = []
|
||||
for result in graph.start():
|
||||
results.append(result)
|
||||
|
||||
assert len(results) == 8
|
||||
vids = [result.vertex.id for result in results if hasattr(result, "vertex")]
|
||||
assert all(vid in ids for vid in vids), f"Diff: {set(vids) - set(ids)}"
|
||||
assert results[-1] == Finish()
|
||||
assert len(results) == len_results
|
||||
vids = [result.vertex.id for result in results if hasattr(result, "vertex")]
|
||||
assert all(vid in ids for vid in vids), f"Diff: {set(vids) - set(ids)}"
|
||||
assert results[-1] == Finish()
|
||||
|
||||
|
||||
def test_vector_store_rag_dump_components_and_edges(ingestion_graph, rag_graph):
|
||||
# Test ingestion graph components and edges
|
||||
ingestion_graph_dump = ingestion_graph.dump(
|
||||
name="Ingestion Graph", description="Graph for data ingestion", endpoint_name="ingestion"
|
||||
)
|
||||
|
||||
ingestion_data = ingestion_graph_dump["data"]
|
||||
ingestion_nodes = ingestion_data["nodes"]
|
||||
ingestion_edges = ingestion_data["edges"]
|
||||
|
||||
# Sort nodes by id to check components
|
||||
ingestion_nodes = sorted(ingestion_nodes, key=lambda x: x["id"])
|
||||
|
||||
# Check components in the ingestion graph
|
||||
assert ingestion_nodes[0]["data"]["type"] == "FileComponent"
|
||||
assert ingestion_nodes[0]["id"] == "file-123"
|
||||
|
||||
assert ingestion_nodes[1]["data"]["type"] == "OpenAIEmbeddingsComponent"
|
||||
assert ingestion_nodes[1]["id"] == "openai-embeddings-123"
|
||||
|
||||
assert ingestion_nodes[2]["data"]["type"] == "SplitTextComponent"
|
||||
assert ingestion_nodes[2]["id"] == "text-splitter-123"
|
||||
|
||||
assert ingestion_nodes[3]["data"]["type"] == "AstraVectorStoreComponent"
|
||||
assert ingestion_nodes[3]["id"] == "vector-store-123"
|
||||
|
||||
# Check edges in the ingestion graph
|
||||
expected_ingestion_edges = [
|
||||
("file-123", "text-splitter-123"),
|
||||
("text-splitter-123", "vector-store-123"),
|
||||
("openai-embeddings-123", "vector-store-123"),
|
||||
]
|
||||
assert len(ingestion_edges) == len(expected_ingestion_edges)
|
||||
|
||||
for edge in ingestion_edges:
|
||||
source = edge["source"]
|
||||
target = edge["target"]
|
||||
assert (source, target) in expected_ingestion_edges, edge
|
||||
|
||||
# Test RAG graph components and edges
|
||||
rag_graph_dump = rag_graph.dump(
|
||||
name="RAG Graph", description="Graph for Retrieval-Augmented Generation", endpoint_name="rag"
|
||||
)
|
||||
|
||||
rag_data = rag_graph_dump["data"]
|
||||
rag_nodes = rag_data["nodes"]
|
||||
rag_edges = rag_data["edges"]
|
||||
|
||||
# Sort nodes by id to check components
|
||||
rag_nodes = sorted(rag_nodes, key=lambda x: x["id"])
|
||||
|
||||
# Check components in the RAG graph
|
||||
assert rag_nodes[0]["data"]["type"] == "ChatInput"
|
||||
assert rag_nodes[0]["id"] == "chatinput-123"
|
||||
|
||||
assert rag_nodes[1]["data"]["type"] == "ChatOutput"
|
||||
assert rag_nodes[1]["id"] == "chatoutput-123"
|
||||
|
||||
assert rag_nodes[2]["data"]["type"] == "OpenAIModelComponent"
|
||||
assert rag_nodes[2]["id"] == "openai-123"
|
||||
|
||||
assert rag_nodes[3]["data"]["type"] == "OpenAIEmbeddingsComponent"
|
||||
assert rag_nodes[3]["id"] == "openai-embeddings-124"
|
||||
|
||||
assert rag_nodes[4]["data"]["type"] == "ParseDataComponent"
|
||||
assert rag_nodes[4]["id"] == "parse-data-123"
|
||||
|
||||
assert rag_nodes[5]["data"]["type"] == "PromptComponent"
|
||||
assert rag_nodes[5]["id"] == "prompt-123"
|
||||
|
||||
assert rag_nodes[6]["data"]["type"] == "AstraVectorStoreComponent"
|
||||
assert rag_nodes[6]["id"] == "rag-vector-store-123"
|
||||
|
||||
# Check edges in the RAG graph
|
||||
expected_rag_edges = [
|
||||
("chatinput-123", "rag-vector-store-123"),
|
||||
("openai-embeddings-124", "rag-vector-store-123"),
|
||||
("chatinput-123", "prompt-123"),
|
||||
("rag-vector-store-123", "parse-data-123"),
|
||||
("parse-data-123", "prompt-123"),
|
||||
("prompt-123", "openai-123"),
|
||||
("openai-123", "chatoutput-123"),
|
||||
]
|
||||
assert len(rag_edges) == len(expected_rag_edges), rag_edges
|
||||
|
||||
for edge in rag_edges:
|
||||
source = edge["source"]
|
||||
target = edge["target"]
|
||||
assert (source, target) in expected_rag_edges, f"Edge {source} -> {target} not found"
|
||||
|
||||
|
||||
def test_vector_store_rag_dump(ingestion_graph, rag_graph):
|
||||
# Test ingestion graph dump
|
||||
ingestion_graph_dump = ingestion_graph.dump(
|
||||
name="Ingestion Graph", description="Graph for data ingestion", endpoint_name="ingestion"
|
||||
)
|
||||
assert isinstance(ingestion_graph_dump, dict)
|
||||
|
||||
ingestion_data = ingestion_graph_dump["data"]
|
||||
assert "nodes" in ingestion_data
|
||||
assert "edges" in ingestion_data
|
||||
assert "description" in ingestion_graph_dump
|
||||
assert "endpoint_name" in ingestion_graph_dump
|
||||
|
||||
ingestion_nodes = ingestion_data["nodes"]
|
||||
ingestion_edges = ingestion_data["edges"]
|
||||
assert len(ingestion_nodes) == 4 # There are 4 components in the ingestion graph
|
||||
assert len(ingestion_edges) == 3 # There are 3 connections between components
|
||||
|
||||
# Test RAG graph dump
|
||||
rag_graph_dump = rag_graph.dump(
|
||||
name="RAG Graph", description="Graph for Retrieval-Augmented Generation", endpoint_name="rag"
|
||||
)
|
||||
assert isinstance(rag_graph_dump, dict)
|
||||
|
||||
rag_data = rag_graph_dump["data"]
|
||||
assert "nodes" in rag_data
|
||||
assert "edges" in rag_data
|
||||
assert "description" in rag_graph_dump
|
||||
assert "endpoint_name" in rag_graph_dump
|
||||
|
||||
rag_nodes = rag_data["nodes"]
|
||||
rag_edges = rag_data["edges"]
|
||||
assert len(rag_nodes) == 7 # There are 7 components in the RAG graph
|
||||
assert len(rag_edges) == 7 # There are 7 connections between components
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue