diff --git a/src/backend/base/langflow/graph/edge/base.py b/src/backend/base/langflow/graph/edge/base.py index f3520e72cf1..80aa5519bb6 100644 --- a/src/backend/base/langflow/graph/edge/base.py +++ b/src/backend/base/langflow/graph/edge/base.py @@ -87,6 +87,13 @@ def __setstate__(self, state): self.target_param = state["target_param"] self.source_handle = state.get("source_handle") self.target_handle = state.get("target_handle") + self._source_handle = state.get("_source_handle") + self._target_handle = state.get("_target_handle") + self._data = state.get("_data") + self.valid_handles = state.get("valid_handles") + self.source_types = state.get("source_types") + self.target_reqs = state.get("target_reqs") + self.matched_type = state.get("matched_type") def validate_edge(self, source, target) -> None: # If the self.source_handle has base_classes, then we are using the legacy diff --git a/src/backend/base/langflow/graph/graph/base.py b/src/backend/base/langflow/graph/graph/base.py index 901f4288810..f9772fe325c 100644 --- a/src/backend/base/langflow/graph/graph/base.py +++ b/src/backend/base/langflow/graph/graph/base.py @@ -1,4 +1,5 @@ import asyncio +import copy import json import uuid from collections import defaultdict, deque @@ -100,6 +101,29 @@ def __init__( if (start is not None and end is None) or (start is None and end is not None): raise ValueError("You must provide both input and output components") + def __add__(self, other): + if not isinstance(other, Graph): + raise TypeError("Can only add Graph objects") + # Add the vertices and edges from the other graph to this graph + new_instance = copy.deepcopy(self) + for vertex in other.vertices: + # This updates the edges as well + new_instance.add_vertex(vertex) + new_instance.build_graph_maps(new_instance.edges) + new_instance.define_vertices_lists() + return new_instance + + def __iadd__(self, other): + if not isinstance(other, Graph): + raise TypeError("Can only add Graph objects") + # Add the vertices and edges from the other graph to this graph + for vertex in other.vertices: + # This updates the edges as well + self.add_vertex(vertex) + self.build_graph_maps(self.edges) + self.define_vertices_lists() + return self + def dumps( self, name: Optional[str] = None, @@ -779,6 +803,14 @@ def __getstate__(self): "vertices_to_run": self.vertices_to_run, "stop_vertex": self.stop_vertex, "vertex_map": self.vertex_map, + "_run_queue": self._run_queue, + "_first_layer": self._first_layer, + "_vertices": self._vertices, + "_edges": self._edges, + "_is_input_vertices": self._is_input_vertices, + "_is_output_vertices": self._is_output_vertices, + "_has_session_id_vertices": self._has_session_id_vertices, + "_sorted_vertices_layers": self._sorted_vertices_layers, } def __setstate__(self, state): @@ -1450,6 +1482,7 @@ def _create_vertex(self, frontend_data: NodeData): return vertex_instance def prepare(self, stop_component_id: Optional[str] = None, start_component_id: Optional[str] = None): + self.initialize() if stop_component_id and start_component_id: raise ValueError("You can only provide one of stop_component_id or start_component_id") self.validate_stream() diff --git a/src/backend/tests/unit/initial_setup/starter_projects/test_vector_store_rag.py b/src/backend/tests/unit/initial_setup/starter_projects/test_vector_store_rag.py index a005b97deeb..587c65779a2 100644 --- a/src/backend/tests/unit/initial_setup/starter_projects/test_vector_store_rag.py +++ b/src/backend/tests/unit/initial_setup/starter_projects/test_vector_store_rag.py @@ -1,3 +1,4 @@ +import copy from textwrap import dedent import pytest @@ -211,6 +212,73 @@ def test_vector_store_rag_dump_components_and_edges(ingestion_graph, rag_graph): assert (source, target) in expected_rag_edges, f"Edge {source} -> {target} not found" +def test_vector_store_rag_add(ingestion_graph, rag_graph): + ingestion_graph_copy = copy.deepcopy(ingestion_graph) + rag_graph_copy = copy.deepcopy(rag_graph) + ingestion_graph_copy += rag_graph_copy + + assert ( + len(ingestion_graph_copy.vertices) == len(ingestion_graph.vertices) + len(rag_graph.vertices) + ), f"Vertices mismatch: {len(ingestion_graph_copy.vertices)} != {len(ingestion_graph.vertices)} + {len(rag_graph.vertices)}" + assert len(ingestion_graph_copy.edges) == len(ingestion_graph.edges) + len( + rag_graph.edges + ), f"Edges mismatch: {len(ingestion_graph_copy.edges)} != {len(ingestion_graph.edges)} + {len(rag_graph.edges)}" + + combined_graph_dump = ingestion_graph_copy.dump( + name="Combined Graph", description="Graph for data ingestion and RAG", endpoint_name="combined" + ) + + combined_data = combined_graph_dump["data"] + combined_nodes = combined_data["nodes"] + combined_edges = combined_data["edges"] + + # Sort nodes by id to check components + combined_nodes = sorted(combined_nodes, key=lambda x: x["id"]) + + # Expected components in the combined graph (both ingestion and RAG nodes) + expected_nodes = sorted( + [ + {"id": "file-123", "type": "File"}, + {"id": "openai-embeddings-123", "type": "OpenAIEmbeddings"}, + {"id": "text-splitter-123", "type": "SplitText"}, + {"id": "vector-store-123", "type": "AstraDB"}, + {"id": "chatinput-123", "type": "ChatInput"}, + {"id": "chatoutput-123", "type": "ChatOutput"}, + {"id": "openai-123", "type": "OpenAIModel"}, + {"id": "openai-embeddings-124", "type": "OpenAIEmbeddings"}, + {"id": "parse-data-123", "type": "ParseData"}, + {"id": "prompt-123", "type": "Prompt"}, + {"id": "rag-vector-store-123", "type": "AstraDB"}, + ], + key=lambda x: x["id"], + ) + + for expected_node, combined_node in zip(expected_nodes, combined_nodes): + assert combined_node["data"]["type"] == expected_node["type"] + assert combined_node["id"] == expected_node["id"] + + # Expected edges in the combined graph (both ingestion and RAG edges) + expected_combined_edges = [ + ("file-123", "text-splitter-123"), + ("text-splitter-123", "vector-store-123"), + ("openai-embeddings-123", "vector-store-123"), + ("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(combined_edges) == len(expected_combined_edges), combined_edges + + for edge in combined_edges: + source = edge["source"] + target = edge["target"] + assert (source, target) in expected_combined_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( diff --git a/src/frontend/package-lock.json b/src/frontend/package-lock.json index db0e15228cd..f150e6b429d 100644 --- a/src/frontend/package-lock.json +++ b/src/frontend/package-lock.json @@ -1079,6 +1079,7 @@ }, "node_modules/@clack/prompts/node_modules/is-unicode-supported": { "version": "1.3.0", + "extraneous": true, "inBundle": true, "license": "MIT", "engines": {