refactor: add graph utility tests and refactor sorting methods (#5538)
* refactor: turn sorting methods into functions in a separate module - Added `layered_topological_sort` function to perform layered topological sorting of graph vertices, accommodating cycles and input vertex checks. - Introduced `refine_layers` function to ensure proper dependency ordering among vertices. - Implemented helper functions for sorting layers by dependency and filtering vertices based on predecessors. - Enhanced utility functions to support better graph traversal and layer management. This update improves the graph processing capabilities, allowing for more efficient handling of complex graph structures. * feat(tests): enhance graph utility tests with cycle detection and sorting functionality - Added a new fixture `graph_with_loop` to simulate a graph containing cycles for testing purposes. - Improved the `test_large_graph_efficiency` to validate cycle detection in large graphs. - Introduced multiple tests for sorting vertices in graphs with cycles, ensuring correct order and handling of input vertices. - Enhanced assertions to provide clearer error messages for failed tests, improving debugging experience. These changes strengthen the testing framework for graph utilities, ensuring robust handling of complex graph structures. * refactor(graph): remove unused parent_node_map from Graph class initialization - Eliminated the `parent_node_map` parameter from the Graph class constructor, streamlining the graph initialization process. - This change enhances code clarity and reduces unnecessary complexity in graph management. This update contributes to cleaner and more maintainable graph-related code. * refactor(graph): optimize dependency sorting and vertex filtering - Improved the `_max_dependency_index` function by utilizing `index_map.get()` for cleaner code and better handling of missing successors. - Enhanced the `_sort_single_layer_by_dependency` function with a caching mechanism to avoid redundant calculations, improving performance during vertex sorting. - Updated `filter_vertices_up_to_vertex` to use a set for `vertices_ids`, optimizing membership checks and enhancing efficiency in vertex filtering. These changes contribute to more efficient graph processing and improved code readability. * chore: remove unused 'parent_node_map' parameter * [autofix.ci] apply automated fixes * fix: replace old method call with a new func * test: enhance assertions for file existence in webhook tests * refactor(graph): enhance component ID retrieval and chat input sorting - Updated `find_start_component_id` to accept an optional `is_webhook` parameter, allowing for dynamic priority input selection based on the flow type. - Improved `sort_chat_inputs_first` to handle chat input positioning more efficiently, ensuring only one chat input exists and adjusting its position within the layers as needed. - These changes enhance the flexibility and efficiency of graph processing, particularly for webhook flows. * test(graph): update assertions in sort_chat_inputs_first test for accuracy - Modified assertions in the `test_chat_inputs_at_start` function to reflect the correct expected output of the `sort_chat_inputs_first` utility. - Adjusted the expected length and order of the result to ensure accurate validation of chat input sorting functionality. These changes enhance the reliability of the test suite for graph utilities, ensuring that the sorting logic is correctly validated. * test(chat): update assertion in consume_and_assert_stream for accurate ID validation - Modified the assertion in the `consume_and_assert_stream` function to include an additional expected ID, ensuring the test accurately reflects the current output of the chat endpoint. - This change enhances the reliability of the test suite by validating the correct behavior of the chat input sorting functionality. * test(endpoints): update assertion in test_get_vertices for accurate ID validation - Modified the assertion in the `test_get_vertices` function to include an additional expected ID, "Webhook", alongside "ChatInput". - This change ensures the test accurately reflects the current output of the endpoint, enhancing the reliability of the test suite for endpoint functionality. --------- Co-authored-by: italojohnny <italojohnnydosanjos@gmail.com> Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Gabriel Luiz Freitas Almeida
GitHub
parent
319380462e
commit
9c23759c7d
8 changed files with 1940 additions and 946 deletions
|
|
@ -36,6 +36,24 @@ def graph():
|
|||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def graph_with_loop():
|
||||
return {
|
||||
"Playlist Extractor": {"successors": ["Loop"], "predecessors": []},
|
||||
"Loop": {
|
||||
"successors": ["Parse Data 1", "Parse Data 2"],
|
||||
"predecessors": ["Playlist Extractor", "YouTube Transcripts"],
|
||||
},
|
||||
"Parse Data 1": {"successors": ["YouTube Transcripts"], "predecessors": ["Loop"]},
|
||||
"Parse Data 2": {"successors": ["Message to Data"], "predecessors": ["Loop"]},
|
||||
"YouTube Transcripts": {"successors": ["Loop"], "predecessors": ["Parse Data 1"]},
|
||||
"Message to Data": {"successors": ["Split Text"], "predecessors": ["Parse Data 2"]},
|
||||
"Split Text": {"successors": ["Chroma DB"], "predecessors": ["Message to Data"]},
|
||||
"OpenAI Embeddings": {"successors": ["Chroma DB"], "predecessors": []},
|
||||
"Chroma DB": {"successors": [], "predecessors": ["Split Text", "OpenAI Embeddings"]},
|
||||
}
|
||||
|
||||
|
||||
def test_get_successors_a(graph):
|
||||
vertex_id = "A"
|
||||
|
||||
|
|
@ -202,10 +220,22 @@ class TestFindCycleEdge:
|
|||
|
||||
# Handles large graphs efficiently
|
||||
def test_large_graph_efficiency(self):
|
||||
entry_point = "0"
|
||||
edges = [(str(i), str(i + 1)) for i in range(1000)] + [("999", "0")]
|
||||
entry_point = "A"
|
||||
# Create a graph with 50 nodes that definitely contains cycles
|
||||
base_edges = [(chr(65 + i), chr(65 + (i + 1) % 26)) for i in range(25)]
|
||||
cycle_edges = [(chr(65 + i), chr(65 + (i - 2) % 26)) for i in range(2, 25, 3)]
|
||||
edges = base_edges + cycle_edges
|
||||
|
||||
result = utils.find_cycle_edge(entry_point, edges)
|
||||
assert result == ("999", "0")
|
||||
|
||||
assert result is not None, (
|
||||
"No cycle was found, but the graph should contain cycles.\n"
|
||||
f"Entry point: {entry_point}\n"
|
||||
f"Number of edges: {len(edges)}"
|
||||
)
|
||||
assert isinstance(result, tuple), f"Expected result to be a tuple, but got {type(result)}"
|
||||
assert len(result) == 2, f"Expected tuple of length 2, but got length {len(result)}"
|
||||
assert all(isinstance(x, str) for x in result), "Expected both elements to be strings"
|
||||
|
||||
# Manages graphs with duplicate edges
|
||||
def test_duplicate_edges(self):
|
||||
|
|
@ -444,3 +474,379 @@ class TestFindCycleVertices:
|
|||
expected_output = ["router", "chat_input", "concatenate"]
|
||||
result = utils.find_cycle_vertices(edges)
|
||||
assert sorted(result) == sorted(expected_output)
|
||||
|
||||
|
||||
def test_chat_inputs_at_start():
|
||||
vertices_layers = [["ChatInput1", "B"], ["C"], ["D"]]
|
||||
|
||||
def get_vertex_predecessors(vertex_id: str) -> list[str]: # noqa: ARG001
|
||||
return []
|
||||
|
||||
result = utils.sort_chat_inputs_first(vertices_layers, get_vertex_predecessors)
|
||||
assert len(result) == 3 # [chat_input] + original 3 layers
|
||||
assert result[0] == ["ChatInput1", "B"]
|
||||
assert result[1] == ["C"] # Original second layer
|
||||
assert result[2] == ["D"] # Original third layer
|
||||
|
||||
# Test that multiple chat inputs raise an error
|
||||
vertices_layers_multiple = [["ChatInput1", "B"], ["ChatInput2", "C"], ["D"]]
|
||||
with pytest.raises(ValueError, match="Only one chat input is allowed in the graph"):
|
||||
utils.sort_chat_inputs_first(vertices_layers_multiple, get_vertex_predecessors)
|
||||
|
||||
|
||||
def test_get_sorted_vertices_simple():
|
||||
# Simple graph with chat input
|
||||
vertices_ids = ["ChatInput1", "B", "C", "D"]
|
||||
cycle_vertices = set()
|
||||
graph_dict = {
|
||||
"ChatInput1": {"successors": ["B"], "predecessors": []},
|
||||
"B": {"successors": ["C"], "predecessors": ["ChatInput1"]},
|
||||
"C": {"successors": ["D"], "predecessors": ["B"]},
|
||||
"D": {"successors": [], "predecessors": ["C"]},
|
||||
}
|
||||
in_degree_map = {"ChatInput1": 0, "B": 1, "C": 1, "D": 1}
|
||||
successor_map = {"ChatInput1": ["B"], "B": ["C"], "C": ["D"], "D": []}
|
||||
predecessor_map = {"ChatInput1": [], "B": ["ChatInput1"], "C": ["B"], "D": ["C"]}
|
||||
|
||||
def is_input_vertex(vertex_id: str) -> bool:
|
||||
return vertex_id == "ChatInput1"
|
||||
|
||||
def get_vertex_predecessors(vertex_id: str) -> list[str]:
|
||||
return predecessor_map[vertex_id]
|
||||
|
||||
def get_vertex_successors(vertex_id: str) -> list[str]:
|
||||
return successor_map[vertex_id]
|
||||
|
||||
first_layer, remaining_layers = utils.get_sorted_vertices(
|
||||
vertices_ids=vertices_ids,
|
||||
cycle_vertices=cycle_vertices,
|
||||
stop_component_id=None,
|
||||
start_component_id=None,
|
||||
graph_dict=graph_dict,
|
||||
in_degree_map=in_degree_map,
|
||||
successor_map=successor_map,
|
||||
predecessor_map=predecessor_map,
|
||||
is_input_vertex=is_input_vertex,
|
||||
get_vertex_predecessors=get_vertex_predecessors,
|
||||
get_vertex_successors=get_vertex_successors,
|
||||
is_cyclic=False,
|
||||
)
|
||||
|
||||
assert first_layer == ["ChatInput1"]
|
||||
assert len(remaining_layers) == 3
|
||||
assert remaining_layers[0] == ["B"]
|
||||
assert remaining_layers[1] == ["C"]
|
||||
assert remaining_layers[2] == ["D"]
|
||||
|
||||
|
||||
def test_get_sorted_vertices_with_cycle():
|
||||
# Graph with a cycle
|
||||
vertices_ids = ["A", "B", "C"]
|
||||
cycle_vertices = {"A", "B", "C"}
|
||||
graph_dict = {
|
||||
"A": {"successors": ["B"], "predecessors": ["C"]},
|
||||
"B": {"successors": ["C"], "predecessors": ["A"]},
|
||||
"C": {"successors": ["A"], "predecessors": ["B"]},
|
||||
}
|
||||
in_degree_map = {"A": 1, "B": 1, "C": 1}
|
||||
successor_map = {"A": ["B"], "B": ["C"], "C": ["A"]}
|
||||
predecessor_map = {"A": ["C"], "B": ["A"], "C": ["B"]}
|
||||
|
||||
def is_input_vertex(vertex_id: str) -> bool: # noqa: ARG001
|
||||
return False
|
||||
|
||||
def get_vertex_predecessors(vertex_id: str) -> list[str]:
|
||||
return predecessor_map[vertex_id]
|
||||
|
||||
def get_vertex_successors(vertex_id: str) -> list[str]:
|
||||
return successor_map[vertex_id]
|
||||
|
||||
# Test with stop_component_id in cycle
|
||||
first_layer, remaining_layers = utils.get_sorted_vertices(
|
||||
vertices_ids=vertices_ids,
|
||||
cycle_vertices=cycle_vertices,
|
||||
stop_component_id="B",
|
||||
start_component_id=None,
|
||||
graph_dict=graph_dict,
|
||||
in_degree_map=in_degree_map,
|
||||
successor_map=successor_map,
|
||||
predecessor_map=predecessor_map,
|
||||
is_input_vertex=is_input_vertex,
|
||||
get_vertex_predecessors=get_vertex_predecessors,
|
||||
get_vertex_successors=get_vertex_successors,
|
||||
is_cyclic=True,
|
||||
)
|
||||
|
||||
# When there's a cycle and stop_component_id is in the cycle,
|
||||
# stop_component_id becomes start_component_id
|
||||
assert first_layer == ["B"]
|
||||
assert len(remaining_layers) == 2
|
||||
assert remaining_layers[0] == ["C"]
|
||||
assert remaining_layers[1] == ["A"]
|
||||
|
||||
|
||||
def test_get_sorted_vertices_with_stop():
|
||||
# Graph with a stop component
|
||||
vertices_ids = ["A", "B", "C", "D", "E"]
|
||||
cycle_vertices = set()
|
||||
graph_dict = {
|
||||
"A": {"successors": ["B"], "predecessors": []},
|
||||
"B": {"successors": ["C"], "predecessors": ["A"]},
|
||||
"C": {"successors": ["D"], "predecessors": ["B"]},
|
||||
"D": {"successors": ["E"], "predecessors": ["C"]},
|
||||
"E": {"successors": [], "predecessors": ["D"]},
|
||||
}
|
||||
in_degree_map = {"A": 0, "B": 1, "C": 1, "D": 1, "E": 1}
|
||||
successor_map = {"A": ["B"], "B": ["C"], "C": ["D"], "D": ["E"], "E": []}
|
||||
predecessor_map = {"A": [], "B": ["A"], "C": ["B"], "D": ["C"], "E": ["D"]}
|
||||
|
||||
def is_input_vertex(vertex_id: str) -> bool:
|
||||
return vertex_id == "A"
|
||||
|
||||
def get_vertex_predecessors(vertex_id: str) -> list[str]:
|
||||
return predecessor_map[vertex_id]
|
||||
|
||||
def get_vertex_successors(vertex_id: str) -> list[str]:
|
||||
return successor_map[vertex_id]
|
||||
|
||||
first_layer, remaining_layers = utils.get_sorted_vertices(
|
||||
vertices_ids=vertices_ids,
|
||||
cycle_vertices=cycle_vertices,
|
||||
stop_component_id="C",
|
||||
start_component_id=None,
|
||||
graph_dict=graph_dict,
|
||||
in_degree_map=in_degree_map,
|
||||
successor_map=successor_map,
|
||||
predecessor_map=predecessor_map,
|
||||
is_input_vertex=is_input_vertex,
|
||||
get_vertex_predecessors=get_vertex_predecessors,
|
||||
get_vertex_successors=get_vertex_successors,
|
||||
is_cyclic=False,
|
||||
)
|
||||
|
||||
assert first_layer == ["A"]
|
||||
assert len(remaining_layers) == 2
|
||||
assert remaining_layers[0] == ["B"]
|
||||
assert remaining_layers[1] == ["C"]
|
||||
|
||||
|
||||
def test_get_sorted_vertices_with_complex_cycle(graph_with_loop):
|
||||
# Convert the graph structure to the format needed by get_sorted_vertices
|
||||
vertices_ids = list(graph_with_loop.keys())
|
||||
cycle_vertices = {"Loop", "Parse Data 1", "YouTube Transcripts"} # Known cycle in the graph
|
||||
graph_dict = graph_with_loop
|
||||
|
||||
# Build in_degree_map from predecessors
|
||||
in_degree_map = {vertex: len(data["predecessors"]) for vertex, data in graph_with_loop.items()}
|
||||
|
||||
# Build successor and predecessor maps
|
||||
successor_map = {vertex: data["successors"] for vertex, data in graph_with_loop.items()}
|
||||
predecessor_map = {vertex: data["predecessors"] for vertex, data in graph_with_loop.items()}
|
||||
|
||||
def is_input_vertex(vertex_id: str) -> bool:
|
||||
# Only Playlist Extractor is an input vertex
|
||||
return vertex_id == "Playlist Extractor"
|
||||
|
||||
def get_vertex_predecessors(vertex_id: str) -> list[str]:
|
||||
return predecessor_map[vertex_id]
|
||||
|
||||
def get_vertex_successors(vertex_id: str) -> list[str]:
|
||||
return successor_map[vertex_id]
|
||||
|
||||
# Test with the cycle
|
||||
first_layer, remaining_layers = utils.get_sorted_vertices(
|
||||
vertices_ids=vertices_ids,
|
||||
cycle_vertices=cycle_vertices,
|
||||
stop_component_id=None,
|
||||
start_component_id=None,
|
||||
graph_dict=graph_dict,
|
||||
in_degree_map=in_degree_map,
|
||||
successor_map=successor_map,
|
||||
predecessor_map=predecessor_map,
|
||||
is_input_vertex=is_input_vertex,
|
||||
get_vertex_predecessors=get_vertex_predecessors,
|
||||
get_vertex_successors=get_vertex_successors,
|
||||
is_cyclic=True,
|
||||
)
|
||||
|
||||
# When is_cyclic is True and start_vertex_id is provided:
|
||||
# 1. The first layer will contain vertices with no predecessors and vertices that are part of the cycle
|
||||
# 2. This is because the cycle vertices are treated as having no dependencies in the initial sort
|
||||
assert (
|
||||
"OpenAI Embeddings" in first_layer
|
||||
), "Vertex with no predecessors 'OpenAI Embeddings' should be in first layer"
|
||||
assert "Playlist Extractor" in first_layer, "Input vertex 'Playlist Extractor' should be in first layer"
|
||||
assert (
|
||||
len(first_layer) == 2
|
||||
), f"First layer should contain exactly 4 vertices, got {len(first_layer)}: {first_layer}"
|
||||
|
||||
# Verify that the remaining layers contain the rest of the vertices in the correct order
|
||||
# The graph structure shows:
|
||||
# Loop -> Parse Data 2 -> Message to Data -> Split Text -> Chroma DB
|
||||
# OpenAI Embeddings -> Chroma DB
|
||||
vertex_to_layer = {}
|
||||
for i, layer in enumerate(remaining_layers):
|
||||
for vertex in layer:
|
||||
vertex_to_layer[vertex] = i
|
||||
|
||||
# Verify that vertices appear in the correct order
|
||||
assert "Loop" in vertex_to_layer, "Vertex 'Loop' should be present in remaining layers"
|
||||
assert "Parse Data 2" in vertex_to_layer, "Vertex 'Parse Data 2' should be present in remaining layers"
|
||||
assert "Message to Data" in vertex_to_layer, "Vertex 'Message to Data' should be present in remaining layers"
|
||||
assert "Chroma DB" in vertex_to_layer, "Vertex 'Chroma DB' should be present in remaining layers"
|
||||
|
||||
# Verify the dependencies are respected
|
||||
# Note: Due to the cycle and the way layered_topological_sort works,
|
||||
# some vertices might appear in earlier layers than expected
|
||||
# What's important is that the dependencies are respected within the non-cycle components
|
||||
assert vertex_to_layer["Parse Data 2"] <= vertex_to_layer["Message to Data"], (
|
||||
f"'Parse Data 2' (layer {vertex_to_layer['Parse Data 2']}) should appear in same or earlier layer than "
|
||||
f"'Message to Data' (layer {vertex_to_layer['Message to Data']})"
|
||||
)
|
||||
|
||||
|
||||
def test_get_sorted_vertices_with_stop_at_chroma(graph_with_loop):
|
||||
# Convert the graph structure to the format needed by get_sorted_vertices
|
||||
vertices_ids = list(graph_with_loop.keys())
|
||||
cycle_vertices = {"Loop", "Parse Data 1", "YouTube Transcripts"} # Known cycle in the graph
|
||||
graph_dict = graph_with_loop
|
||||
|
||||
# Build in_degree_map from predecessors
|
||||
in_degree_map = {vertex: len(data["predecessors"]) for vertex, data in graph_with_loop.items()}
|
||||
|
||||
# Build successor and predecessor maps
|
||||
successor_map = {vertex: data["successors"] for vertex, data in graph_with_loop.items()}
|
||||
predecessor_map = {vertex: data["predecessors"] for vertex, data in graph_with_loop.items()}
|
||||
|
||||
def is_input_vertex(vertex_id: str) -> bool:
|
||||
# Only Playlist Extractor is an input vertex
|
||||
return vertex_id == "Playlist Extractor"
|
||||
|
||||
def get_vertex_predecessors(vertex_id: str) -> list[str]:
|
||||
return predecessor_map[vertex_id]
|
||||
|
||||
def get_vertex_successors(vertex_id: str) -> list[str]:
|
||||
return successor_map[vertex_id]
|
||||
|
||||
# Test with ChromaDB as stop component
|
||||
first_layer, remaining_layers = utils.get_sorted_vertices(
|
||||
vertices_ids=vertices_ids,
|
||||
cycle_vertices=cycle_vertices,
|
||||
stop_component_id="Chroma DB", # Stop at ChromaDB
|
||||
start_component_id=None,
|
||||
graph_dict=graph_dict,
|
||||
in_degree_map=in_degree_map,
|
||||
successor_map=successor_map,
|
||||
predecessor_map=predecessor_map,
|
||||
is_input_vertex=is_input_vertex,
|
||||
get_vertex_predecessors=get_vertex_predecessors,
|
||||
get_vertex_successors=get_vertex_successors,
|
||||
is_cyclic=True,
|
||||
)
|
||||
|
||||
# When is_cyclic is True and we have a stop component:
|
||||
# 1. The first layer will contain vertices with no predecessors and vertices that are part of the cycle
|
||||
# 2. This is because the cycle vertices are treated as having no dependencies in the initial sort
|
||||
assert (
|
||||
"OpenAI Embeddings" in first_layer
|
||||
), "Vertex with no predecessors 'OpenAI Embeddings' should be in first layer"
|
||||
assert "Playlist Extractor" in first_layer, "Input vertex 'Playlist Extractor' should be in first layer"
|
||||
|
||||
assert (
|
||||
len(first_layer) == 2
|
||||
), f"First layer should contain exactly 4 vertices, got {len(first_layer)}: {first_layer}"
|
||||
|
||||
# Verify that the remaining layers contain the rest of the vertices in the correct order
|
||||
# The graph structure shows:
|
||||
# Loop -> Parse Data 2 -> Message to Data -> Split Text -> Chroma DB
|
||||
# OpenAI Embeddings -> Chroma DB
|
||||
vertex_to_layer = {}
|
||||
for i, layer in enumerate(remaining_layers):
|
||||
for vertex in layer:
|
||||
vertex_to_layer[vertex] = i
|
||||
|
||||
# Verify that vertices appear in the correct order
|
||||
assert "Loop" in vertex_to_layer, "Vertex 'Loop' should be present in remaining layers"
|
||||
assert "Parse Data 2" in vertex_to_layer, "Vertex 'Parse Data 2' should be present in remaining layers"
|
||||
assert "Message to Data" in vertex_to_layer, "Vertex 'Message to Data' should be present in remaining layers"
|
||||
assert "Chroma DB" in vertex_to_layer, "Vertex 'Chroma DB' should be present in remaining layers"
|
||||
|
||||
# Verify that dependencies are respected
|
||||
assert vertex_to_layer["Parse Data 2"] <= vertex_to_layer["Message to Data"], (
|
||||
f"'Parse Data 2' (layer {vertex_to_layer['Parse Data 2']}) should appear in same or earlier layer than "
|
||||
f"'Message to Data' (layer {vertex_to_layer['Message to Data']})"
|
||||
)
|
||||
|
||||
# When a vertex is marked as a stop component, it will appear in layer 0
|
||||
# of the remaining layers. This is because the algorithm stops at this vertex.
|
||||
assert vertex_to_layer["Chroma DB"] == 5, (
|
||||
f"Stop component 'Chroma DB' should be in layer 5, "
|
||||
f"but was found in layer {vertex_to_layer['Chroma DB']}. "
|
||||
f"Remaining layers: {remaining_layers}"
|
||||
)
|
||||
|
||||
|
||||
def test_get_sorted_vertices_exact_sequence(graph_with_loop):
|
||||
# Convert the graph structure to the format needed by get_sorted_vertices
|
||||
vertices_ids = list(graph_with_loop.keys())
|
||||
cycle_vertices = {"Loop", "Parse Data 1", "YouTube Transcripts"} # Known cycle in the graph
|
||||
graph_dict = graph_with_loop
|
||||
|
||||
# Build in_degree_map from predecessors
|
||||
in_degree_map = {vertex: len(data["predecessors"]) for vertex, data in graph_with_loop.items()}
|
||||
|
||||
# Build successor and predecessor maps
|
||||
successor_map = {vertex: data["successors"] for vertex, data in graph_with_loop.items()}
|
||||
predecessor_map = {vertex: data["predecessors"] for vertex, data in graph_with_loop.items()}
|
||||
|
||||
def is_input_vertex(vertex_id: str) -> bool:
|
||||
# Only Playlist Extractor is an input vertex
|
||||
return vertex_id == "Playlist Extractor"
|
||||
|
||||
def get_vertex_predecessors(vertex_id: str) -> list[str]:
|
||||
return predecessor_map[vertex_id]
|
||||
|
||||
def get_vertex_successors(vertex_id: str) -> list[str]:
|
||||
return successor_map[vertex_id]
|
||||
|
||||
# Test with the cycle
|
||||
first_layer, remaining_layers = utils.get_sorted_vertices(
|
||||
vertices_ids=vertices_ids,
|
||||
cycle_vertices=cycle_vertices,
|
||||
stop_component_id=None,
|
||||
start_component_id=None,
|
||||
graph_dict=graph_dict,
|
||||
in_degree_map=in_degree_map,
|
||||
successor_map=successor_map,
|
||||
predecessor_map=predecessor_map,
|
||||
is_input_vertex=is_input_vertex,
|
||||
get_vertex_predecessors=get_vertex_predecessors,
|
||||
get_vertex_successors=get_vertex_successors,
|
||||
is_cyclic=True,
|
||||
)
|
||||
|
||||
# Convert layers to a flat sequence
|
||||
sequence = []
|
||||
sequence.extend(sorted(first_layer))
|
||||
for layer in remaining_layers:
|
||||
sequence.extend(sorted(layer))
|
||||
|
||||
# Expected sequence
|
||||
expected_sequence = [
|
||||
"OpenAI Embeddings",
|
||||
"Playlist Extractor",
|
||||
"YouTube Transcripts",
|
||||
"Loop",
|
||||
"Parse Data 1",
|
||||
"Parse Data 2",
|
||||
"Message to Data",
|
||||
"Split Text",
|
||||
"Chroma DB",
|
||||
]
|
||||
|
||||
# Check each vertex appears in the correct order
|
||||
assert sequence == expected_sequence, f"Sequence: {sequence}"
|
||||
# Verify the exact sequence
|
||||
assert len(sequence) == len(expected_sequence), (
|
||||
f"Expected sequence length {len(expected_sequence)}, " f"but got {len(sequence)}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -73,7 +73,7 @@ async def consume_and_assert_stream(r):
|
|||
assert parsed["event"] == "vertices_sorted"
|
||||
ids = parsed["data"]["ids"]
|
||||
ids.sort()
|
||||
assert ids == ["ChatInput-CIGht"]
|
||||
assert ids == ["ChatInput-CIGht", "Memory-amN4Z"]
|
||||
|
||||
to_run = parsed["data"]["to_run"]
|
||||
to_run.sort()
|
||||
|
|
|
|||
|
|
@ -257,7 +257,7 @@ async def test_get_vertices(client, added_flow_webhook_test, logged_in_headers):
|
|||
# The important part is before the - (ConversationBufferMemory, PromptTemplate, ChatOpenAI, LLMChain)
|
||||
ids = [_id.split("-")[0] for _id in response.json()["ids"]]
|
||||
|
||||
assert set(ids) == {"ChatInput"}
|
||||
assert set(ids) == {"ChatInput", "Webhook"}
|
||||
|
||||
|
||||
async def test_build_vertex_invalid_flow_id(client, logged_in_headers):
|
||||
|
|
|
|||
|
|
@ -23,15 +23,16 @@ async def test_webhook_endpoint(client, added_webhook_test):
|
|||
|
||||
response = await client.post(endpoint, json=payload)
|
||||
assert response.status_code == 202
|
||||
assert await file_path.exists()
|
||||
|
||||
assert not await file_path.exists()
|
||||
# Wait a few seconds for the file to be created
|
||||
assert await file_path.exists(), f"File {file_path} does not exist"
|
||||
file_does_not_exist = not await file_path.exists()
|
||||
assert file_does_not_exist, f"File {file_path} still exists"
|
||||
|
||||
# Send an invalid payload
|
||||
payload = {"invalid_key": "invalid_value"}
|
||||
response = await client.post(endpoint, json=payload)
|
||||
assert response.status_code == 202
|
||||
assert not await file_path.exists()
|
||||
assert not await file_path.exists(), f"File {file_path} should not exist"
|
||||
|
||||
|
||||
async def test_webhook_flow_on_run_endpoint(client, added_webhook_test, created_api_key):
|
||||
|
|
@ -50,7 +51,6 @@ async def test_webhook_with_random_payload(client, added_webhook_test):
|
|||
endpoint_name = added_webhook_test["endpoint_name"]
|
||||
endpoint = f"api/v1/webhook/{endpoint_name}"
|
||||
# Just test that "Random Payload" returns 202
|
||||
# returns 202
|
||||
response = await client.post(
|
||||
endpoint,
|
||||
json="Random Payload",
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue