Update build_and_cache_graph function to handle graph data from database and from input data (#1750)

* Update build_and_cache_graph function to handle graph data from database and from input data

* Update get_vertices endpoint to handle graph data from input

* Update build_vertex function to handle graph data from database

* Update build_vertex function to handle graph data from input

* Update FlowDataRequest schema to include optional viewport field

src/backend/base/langflow/api/utils.py

@@ -201,7 +201,7 @@ def format_elapsed_time(elapsed_time: float) -> str:
         return f"{minutes} {minutes_unit}, {seconds} {seconds_unit}"
 
 
-async def build_and_cache_graph(
+async def build_and_cache_graph_from_db(
     flow_id: str,
     session: Session,
     chat_service: "ChatService",
@@ -220,6 +220,17 @@ async def build_and_cache_graph(
     return graph
 
 
+async def build_and_cache_graph_from_data(
+    flow_id: str,
+    chat_service: "ChatService",
+    graph_data: dict,
+):  # -> Graph | Any:
+    """Build and cache the graph."""
+    graph = Graph.from_payload(graph_data, flow_id)
+    await chat_service.set_cache(flow_id, graph)
+    return graph
+
+
 def format_syntax_error_message(exc: SyntaxError) -> str:
     """Format a SyntaxError message for returning to the frontend."""
     if exc.text is None:

src/backend/base/langflow/api/v1/chat.py

@@ -8,13 +8,15 @@ from fastapi.responses import StreamingResponse
 from loguru import logger
 
 from langflow.api.utils import (
-    build_and_cache_graph,
+    build_and_cache_graph_from_data,
+    build_and_cache_graph_from_db,
     format_elapsed_time,
     format_exception_message,
     get_top_level_vertices,
     parse_exception,
 )
 from langflow.api.v1.schemas import (
+    FlowDataRequest,
     InputValueRequest,
     ResultDataResponse,
     StreamData,
@@ -52,6 +54,7 @@ async def try_running_celery_task(vertex, user_id):
 @router.get("/build/{flow_id}/vertices", response_model=VerticesOrderResponse)
 async def get_vertices(
     flow_id: str,
+    data: Optional[FlowDataRequest] = None,
     stop_component_id: Optional[str] = None,
     start_component_id: Optional[str] = None,
     chat_service: "ChatService" = Depends(get_chat_service),
@@ -76,9 +79,14 @@ async def get_vertices(
     try:
         # First, we need to check if the flow_id is in the cache
         graph = None
-        if cache := await chat_service.get_cache(flow_id):
-            graph = cache.get("result")
-        graph = await build_and_cache_graph(flow_id, session, chat_service, graph)
+        if not data:
+            if cache := await chat_service.get_cache(flow_id):
+                graph = cache.get("result")
+            graph = await build_and_cache_graph_from_db(
+                flow_id=flow_id, session=session, chat_service=chat_service, graph=graph
+            )
+        else:
+            graph = await build_and_cache_graph_from_data(flow_id=flow_id, data=data, chat_service=chat_service)
         if stop_component_id or start_component_id:
             try:
                 first_layer = graph.sort_vertices(stop_component_id, start_component_id)
@@ -144,7 +152,9 @@ async def build_vertex(
         if not cache:
             # If there's no cache
             logger.warning(f"No cache found for {flow_id}. Building graph starting at {vertex_id}")
-            graph = await build_and_cache_graph(flow_id=flow_id, session=next(get_session()), chat_service=chat_service)
+            graph = await build_and_cache_graph_from_db(
+                flow_id=flow_id, session=next(get_session()), chat_service=chat_service
+            )
         else:
             graph = cache.get("result")
         result_data_response = ResultDataResponse(results={})

src/backend/base/langflow/api/v1/schemas.py

@@ -294,3 +294,15 @@ class SimplifiedAPIRequest(BaseModel):
     )
     tweaks: Optional[Tweaks] = Field(default=None, description="The tweaks")
     session_id: Optional[str] = Field(default=None, description="The session id")
+
+
+# (alias) type ReactFlowJsonObject<NodeData = any, EdgeData = any> = {
+#     nodes: Node<NodeData>[];
+#     edges: Edge<EdgeData>[];
+#     viewport: Viewport;
+# }
+# import ReactFlowJsonObject
+class FlowDataRequest(BaseModel):
+    nodes: List[dict]
+    edges: List[dict]
+    viewport: Optional[dict] = None