Update node and edge IDs, refactor node ID generation, and fix async bug in build_vertex_stream function (#1546)

* Update node and edge IDs in PageComponent and reactflowUtils

* Merge remote-tracking branch 'origin/zustand/io/migration' into fixGroup

* Refactor node ID generation and update node IDs in selection

* Update flowStore.ts and reactflowUtils.ts

* Fix async bug in build_vertex_stream function

* Add check for missing id in vertex data

* Fix exception message for missing vertex id

* Update code: Added VertexLayerElementType type and modified updateIds function
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-20 20:06:25 -03:00 • committed by GitHub
commit a6625bbad5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 152 additions and 84 deletions

View file

@ -202,7 +202,7 @@ async def build_vertex_stream(
async def stream_vertex(): async def stream_vertex():
try: try:
if not session_id: if not session_id:
cache = chat_service.get_cache(flow_id) cache = await chat_service.get_cache(flow_id)
if not cache: if not cache:
# If there's no cache # If there's no cache
raise ValueError(f"No cache found for {flow_id}.") raise ValueError(f"No cache found for {flow_id}.")
@ -252,7 +252,7 @@ async def build_vertex_stream(
raise ValueError(f"No result found for vertex {vertex_id}") raise ValueError(f"No result found for vertex {vertex_id}")
except Exception as exc: except Exception as exc:
logger.error(f"Error building vertex: {exc}") logger.exception(f"Error building vertex: {exc}")
yield str(StreamData(event="error", data={"error": str(exc)})) yield str(StreamData(event="error", data={"error": str(exc)}))
finally: finally:
logger.debug("Closing stream") logger.debug("Closing stream")

View file

@ -377,7 +377,10 @@ class Graph:
# Remove vertices that are not in the other graph # Remove vertices that are not in the other graph
for vertex_id in removed_vertex_ids: for vertex_id in removed_vertex_ids:
try:
self.remove_vertex(vertex_id) self.remove_vertex(vertex_id)
except ValueError:
pass
# The order here matters because adding the vertex is required # The order here matters because adding the vertex is required
# if any of them have edges that point to any of the new vertices # if any of them have edges that point to any of the new vertices
@ -741,8 +744,11 @@ class Graph:
vertex_data = vertex["data"] vertex_data = vertex["data"]
vertex_type: str = vertex_data["type"] # type: ignore vertex_type: str = vertex_data["type"] # type: ignore
vertex_base_type: str = vertex_data["node"]["template"]["_type"] # type: ignore vertex_base_type: str = vertex_data["node"]["template"]["_type"] # type: ignore
if "id" not in vertex_data:
raise ValueError(f"Vertex data for {vertex_data['display_name']} does not contain an id")
VertexClass = self._get_vertex_class(vertex_type, vertex_base_type, vertex_data["id"]) VertexClass = self._get_vertex_class(vertex_type, vertex_base_type, vertex_data["id"])
vertex_instance = VertexClass(vertex, graph=self) vertex_instance = VertexClass(vertex, graph=self)
vertex_instance.set_top_level(self.top_level_vertices) vertex_instance.set_top_level(self.top_level_vertices)
vertices.append(vertex_instance) vertices.append(vertex_instance)

View file

@ -1,4 +1,4 @@
import _ from "lodash"; import _, { cloneDeep, set } from "lodash";
import { MouseEvent, useCallback, useEffect, useRef, useState } from "react"; import { MouseEvent, useCallback, useEffect, useRef, useState } from "react";
import ReactFlow, { import ReactFlow, {
Background, Background,
@ -32,6 +32,7 @@ import {
isValidConnection, isValidConnection,
reconnectEdges, reconnectEdges,
scapeJSONParse, scapeJSONParse,
updateIds,
validateSelection, validateSelection,
} from "../../../../utils/reactflowUtils"; } from "../../../../utils/reactflowUtils";
import { getRandomName, isWrappedWithClass } from "../../../../utils/utils"; import { getRandomName, isWrappedWithClass } from "../../../../utils/utils";
@ -104,35 +105,40 @@ export default function Page({
) { ) {
event.preventDefault(); event.preventDefault();
takeSnapshot(); takeSnapshot();
if (validateSelection(lastSelection!, edges).length === 0) { if (
validateSelection(lastSelection!, edges).length === 0
) {
const clonedNodes = cloneDeep(nodes)
const clonedEdges = cloneDeep(edges)
const clonedSelection = cloneDeep(lastSelection)
updateIds({ nodes: clonedNodes, edges: clonedEdges }, clonedSelection!)
const { newFlow, removedEdges } = generateFlow( const { newFlow, removedEdges } = generateFlow(
lastSelection!, clonedSelection!,
nodes, clonedNodes,
edges, clonedEdges,
getRandomName() getRandomName()
); );
const newGroupNode = generateNodeFromFlow(newFlow, getNodeId); const newGroupNode = generateNodeFromFlow(
const newEdges = reconnectEdges(newGroupNode, removedEdges); newFlow,
setNodes((oldNodes) => [ getNodeId
...oldNodes.filter( );
(oldNodes) => const newEdges = reconnectEdges(
!lastSelection?.nodes.some(
(selectionNode) => selectionNode.id === oldNodes.id
)
),
newGroupNode, newGroupNode,
]); removedEdges
setEdges((oldEdges) => [ );
...oldEdges.filter( setNodes([...clonedNodes.filter(
(oldNodes) =>
!clonedSelection?.nodes.some(
(selectionNode) =>
selectionNode.id === oldNodes.id
)), newGroupNode])
setEdges([...clonedEdges.filter(
(oldEdge) => (oldEdge) =>
!lastSelection!.nodes.some( !clonedSelection!.nodes.some(
(selectionNode) => (selectionNode) =>
selectionNode.id === oldEdge.target || selectionNode.id === oldEdge.target ||
selectionNode.id === oldEdge.source selectionNode.id === oldEdge.source
) )), ...newEdges])
),
...newEdges,
]);
} else { } else {
setErrorData({ setErrorData({
title: INVALID_SELECTION_ERROR_ALERT, title: INVALID_SELECTION_ERROR_ALERT,
@ -480,10 +486,14 @@ export default function Page({
if ( if (
validateSelection(lastSelection!, edges).length === 0 validateSelection(lastSelection!, edges).length === 0
) { ) {
const clonedNodes = cloneDeep(nodes)
const clonedEdges = cloneDeep(edges)
const clonedSelection = cloneDeep(lastSelection)
updateIds({ nodes: clonedNodes, edges: clonedEdges }, clonedSelection!)
const { newFlow, removedEdges } = generateFlow( const { newFlow, removedEdges } = generateFlow(
lastSelection!, clonedSelection!,
nodes, clonedNodes,
edges, clonedEdges,
getRandomName() getRandomName()
); );
const newGroupNode = generateNodeFromFlow( const newGroupNode = generateNodeFromFlow(
@ -494,27 +504,19 @@ export default function Page({
newGroupNode, newGroupNode,
removedEdges removedEdges
); );
setNodes((oldNodes) => [ setNodes([...clonedNodes.filter(
...oldNodes.filter(
(oldNodes) => (oldNodes) =>
!lastSelection?.nodes.some( !clonedSelection?.nodes.some(
(selectionNode) => (selectionNode) =>
selectionNode.id === oldNodes.id selectionNode.id === oldNodes.id
) )), newGroupNode])
), setEdges([...clonedEdges.filter(
newGroupNode,
]);
setEdges((oldEdges) => [
...oldEdges.filter(
(oldEdge) => (oldEdge) =>
!lastSelection!.nodes.some( !clonedSelection!.nodes.some(
(selectionNode) => (selectionNode) =>
selectionNode.id === oldEdge.target || selectionNode.id === oldEdge.target ||
selectionNode.id === oldEdge.source selectionNode.id === oldEdge.source
) )), ...newEdges])
),
...newEdges,
]);
} else { } else {
setErrorData({ setErrorData({
title: INVALID_SELECTION_ERROR_ALERT, title: INVALID_SELECTION_ERROR_ALERT,

View file

@ -1,4 +1,4 @@
import { cloneDeep } from "lodash"; import { cloneDeep, zip } from "lodash";
import { import {
Edge, Edge,
EdgeChange, EdgeChange,
@ -26,6 +26,7 @@ import {
ChatOutputType, ChatOutputType,
FlowPoolObjectType, FlowPoolObjectType,
FlowStoreType, FlowStoreType,
VertexLayerElementType,
chatInputType, chatInputType,
} from "../types/zustand/flow"; } from "../types/zustand/flow";
import { buildVertices } from "../utils/buildUtils"; import { buildVertices } from "../utils/buildUtils";
@ -36,6 +37,9 @@ import {
getNodeId, getNodeId,
scapeJSONParse, scapeJSONParse,
scapedJSONStringfy, scapedJSONStringfy,
updateEdgesIds,
updateIds,
updateProxyIdsOnTemplate,
validateNodes, validateNodes,
} from "../utils/reactflowUtils"; } from "../utils/reactflowUtils";
import { getInputsAndOutputs } from "../utils/storeUtils"; import { getInputsAndOutputs } from "../utils/storeUtils";
@ -257,6 +261,14 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
let newId = getNodeId(node.data.type); let newId = getNodeId(node.data.type);
idsMap[node.id] = newId; idsMap[node.id] = newId;
if (node.data.node!.flow) {
let newFlow = node.data.node!.flow;
const idsMap = updateIds(newFlow.data!);
updateProxyIdsOnTemplate(node.data.node!.template, idsMap);
let flowEdges = selection.edges;
updateEdgesIds(flowEdges, idsMap);
}
// Create a new node object // Create a new node object
const newNode: NodeType = { const newNode: NodeType = {
id: newId, id: newId,
@ -459,9 +471,18 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
// verticesLayers is a list of list of vertices ids, where each list is a layer of vertices // verticesLayers is a list of list of vertices ids, where each list is a layer of vertices
// we want to add a new layer (next_vertices_ids) to the list of layers (verticesLayers) // we want to add a new layer (next_vertices_ids) to the list of layers (verticesLayers)
// and the values of next_vertices_ids to the list of vertices ids (verticesIds) // and the values of next_vertices_ids to the list of vertices ids (verticesIds)
// const nextVertices will be the zip of vertexBuildData.next_vertices_ids and
// vertexBuildData.top_level_vertices
// the VertexLayerElementType as {id: next_vertices_id, layer: top_level_vertex}
const nextVertices: VertexLayerElementType[] = zip(
vertexBuildData.next_vertices_ids,
vertexBuildData.top_level_vertices
).map(([id, reference]) => ({ id: id!, reference }));
const newLayers = [ const newLayers = [
...get().verticesBuild!.verticesLayers, ...get().verticesBuild!.verticesLayers,
vertexBuildData.next_vertices_ids, nextVertices,
]; ];
const newIds = [ const newIds = [
...get().verticesBuild!.verticesIds, ...get().verticesBuild!.verticesIds,
@ -508,12 +529,18 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
get().setIsBuilding(false); get().setIsBuilding(false);
}, },
onBuildUpdate: handleBuildUpdate, onBuildUpdate: handleBuildUpdate,
onBuildError: (title, list, idList) => { onBuildError: (title: string, list: string[], elementList) => {
const idList = elementList
.map((element) => element.id)
.filter(Boolean) as string[];
useFlowStore.getState().updateBuildStatus(idList, BuildStatus.BUILT); useFlowStore.getState().updateBuildStatus(idList, BuildStatus.BUILT);
setErrorData({ list, title }); setErrorData({ list, title });
get().setIsBuilding(false); get().setIsBuilding(false);
}, },
onBuildStart: (idList) => { onBuildStart: (elementList) => {
const idList = elementList
.map((element) => element.reference)
.filter(Boolean) as string[];
useFlowStore.getState().updateBuildStatus(idList, BuildStatus.BUILDING); useFlowStore.getState().updateBuildStatus(idList, BuildStatus.BUILDING);
}, },
validateNodes: validateSubgraph, validateNodes: validateSubgraph,
@ -531,7 +558,7 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
updateVerticesBuild: ( updateVerticesBuild: (
vertices: { vertices: {
verticesIds: string[]; verticesIds: string[];
verticesLayers: string[][]; verticesLayers: VertexLayerElementType[][];
runId: string; runId: string;
} | null } | null
) => { ) => {
@ -562,6 +589,7 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
}, },
updateBuildStatus: (nodeIdList: string[], status: BuildStatus) => { updateBuildStatus: (nodeIdList: string[], status: BuildStatus) => {
const newFlowBuildStatus = { ...get().flowBuildStatus }; const newFlowBuildStatus = { ...get().flowBuildStatus };
console.log("newFlowBuildStatus", newFlowBuildStatus);
nodeIdList.forEach((id) => { nodeIdList.forEach((id) => {
newFlowBuildStatus[id] = { newFlowBuildStatus[id] = {
status, status,

View file

@ -35,6 +35,11 @@ export type FlowPoolObjectType = {
buildId: string; buildId: string;
}; };
export type VertexLayerElementType = {
id: string;
reference?: string;
};
export type FlowPoolType = { export type FlowPoolType = {
[key: string]: Array<FlowPoolObjectType>; [key: string]: Array<FlowPoolObjectType>;
}; };
@ -103,7 +108,7 @@ export type FlowStoreType = {
updateVerticesBuild: ( updateVerticesBuild: (
vertices: { vertices: {
verticesIds: string[]; verticesIds: string[];
verticesLayers: string[][]; verticesLayers: VertexLayerElementType[][];
runId: string; runId: string;
} | null } | null
) => void; ) => void;
@ -111,7 +116,7 @@ export type FlowStoreType = {
removeFromVerticesBuild: (vertices: string[]) => void; removeFromVerticesBuild: (vertices: string[]) => void;
verticesBuild: { verticesBuild: {
verticesIds: string[]; verticesIds: string[];
verticesLayers: string[][]; verticesLayers: VertexLayerElementType[][];
runId: string; runId: string;
} | null; } | null;
updateBuildStatus: (nodeId: string[], status: BuildStatus) => void; updateBuildStatus: (nodeId: string[], status: BuildStatus) => void;

View file

@ -4,6 +4,7 @@ import { getVerticesOrder, postBuildVertex } from "../controllers/API";
import useAlertStore from "../stores/alertStore"; import useAlertStore from "../stores/alertStore";
import useFlowStore from "../stores/flowStore"; import useFlowStore from "../stores/flowStore";
import { VertexBuildTypeAPI } from "../types/api"; import { VertexBuildTypeAPI } from "../types/api";
import { VertexLayerElementType } from "../types/zustand/flow";
type BuildVerticesParams = { type BuildVerticesParams = {
flowId: string; // Assuming FlowType is the type for your flow flowId: string; // Assuming FlowType is the type for your flow
@ -17,8 +18,8 @@ type BuildVerticesParams = {
buildId: string buildId: string
) => void; // Replace any with the actual type if it's not any ) => void; // Replace any with the actual type if it's not any
onBuildComplete?: (allNodesValid: boolean) => void; onBuildComplete?: (allNodesValid: boolean) => void;
onBuildError?: (title, list, idList: string[]) => void; onBuildError?: (title, list, idList: VertexLayerElementType[]) => void;
onBuildStart?: (idList: string[]) => void; onBuildStart?: (idList: VertexLayerElementType[]) => void;
validateNodes?: (nodes: string[]) => void; validateNodes?: (nodes: string[]) => void;
}; };
@ -49,7 +50,7 @@ export async function updateVerticesOrder(
startNodeId?: string | null, startNodeId?: string | null,
stopNodeId?: string | null stopNodeId?: string | null
): Promise<{ ): Promise<{
verticesLayers: string[][]; verticesLayers: VertexLayerElementType[][];
verticesIds: string[]; verticesIds: string[];
runId: string; runId: string;
}> { }> {
@ -67,7 +68,14 @@ export async function updateVerticesOrder(
useFlowStore.getState().setIsBuilding(false); useFlowStore.getState().setIsBuilding(false);
throw new Error("Invalid nodes"); throw new Error("Invalid nodes");
} }
let verticesLayers: Array<Array<string>> = [orderResponse.data.ids]; // orderResponse.data.ids,
// for each id we need to build the VertexLayerElementType object as
// {id: id, reference: id}
let verticesLayers: Array<Array<VertexLayerElementType>> =
orderResponse.data.ids.map((id: string) => {
return [{ id: id, reference: id }];
});
const runId = orderResponse.data.run_id; const runId = orderResponse.data.run_id;
// if (nodeId) { // if (nodeId) {
// for (let i = 0; i < verticesOrder.length; i += 1) { // for (let i = 0; i < verticesOrder.length; i += 1) {
@ -161,17 +169,17 @@ export async function buildVertices({
if (onBuildStart) onBuildStart(currentLayer); if (onBuildStart) onBuildStart(currentLayer);
// Build each vertex in the current layer // Build each vertex in the current layer
await Promise.all( await Promise.all(
currentLayer.map(async (vertexId) => { currentLayer.map(async (element) => {
// Check if id is in the list of inactive nodes // Check if id is in the list of inactive nodes
if ( if (
!useFlowStore !useFlowStore
.getState() .getState()
.verticesBuild?.verticesIds.includes(vertexId) && .verticesBuild?.verticesIds.includes(element.id) &&
onBuildUpdate onBuildUpdate
) { ) {
// If it is, skip building and set the state to inactive // If it is, skip building and set the state to inactive
onBuildUpdate( onBuildUpdate(
getInactiveVertexData(vertexId), getInactiveVertexData(element.id),
BuildStatus.INACTIVE, BuildStatus.INACTIVE,
runId runId
); );
@ -182,7 +190,7 @@ export async function buildVertices({
// Build the vertex // Build the vertex
await buildVertex({ await buildVertex({
flowId, flowId,
id: vertexId, id: element.id,
input_value, input_value,
onBuildUpdate: (data: VertexBuildTypeAPI, status: BuildStatus) => { onBuildUpdate: (data: VertexBuildTypeAPI, status: BuildStatus) => {
if (onBuildUpdate) onBuildUpdate(data, status, runId); if (onBuildUpdate) onBuildUpdate(data, status, runId);
@ -227,7 +235,7 @@ async function buildVertex({
id: string; id: string;
input_value: string; input_value: string;
onBuildUpdate?: (data: any, status: BuildStatus) => void; onBuildUpdate?: (data: any, status: BuildStatus) => void;
onBuildError?: (title, list, idList: string[]) => void; onBuildError?: (title, list, idList: VertexLayerElementType[]) => void;
verticesIds: string[]; verticesIds: string[];
buildResults: boolean[]; buildResults: boolean[];
stopBuild: () => void; stopBuild: () => void;
@ -241,7 +249,7 @@ async function buildVertex({
onBuildError!( onBuildError!(
"Error Building Component", "Error Building Component",
[buildData.params], [buildData.params],
verticesIds verticesIds.map((id) => ({ id }))
); );
stopBuild(); stopBuild();
} }
@ -252,7 +260,7 @@ async function buildVertex({
onBuildError!( onBuildError!(
"Error Building Component", "Error Building Component",
[(error as AxiosError<any>).response?.data?.detail ?? "Unknown Error"], [(error as AxiosError<any>).response?.data?.detail ?? "Unknown Error"],
verticesIds verticesIds.map((id) => ({ id }))
); );
stopBuild(); stopBuild();
} }

View file

@ -204,23 +204,33 @@ export const processDataFromFlow = (flow: FlowType, refreshIds = true) => {
return data; return data;
}; };
export function updateIds(newFlow: ReactFlowJsonObject) { export function updateIds(
{ edges, nodes }: { edges: Edge[]; nodes: Node[] },
selection?: { edges: Edge[]; nodes: Node[] }
) {
let idsMap = {}; let idsMap = {};
const selectionIds = selection?.nodes.map((n) => n.id);
if (newFlow.nodes) if (nodes) {
newFlow.nodes.forEach((node: NodeType) => { nodes.forEach((node: NodeType) => {
// Generate a unique node ID // Generate a unique node ID
let newId = getNodeId( let newId = getNodeId(node.data.type);
node.data.node?.flow ? "GroupNode" : node.data.type if (selection && !selectionIds?.includes(node.id)) {
); newId = node.id;
}
idsMap[node.id] = newId; idsMap[node.id] = newId;
node.id = newId; node.id = newId;
node.data.id = newId; node.data.id = newId;
// Add the new node to the list of nodes in state // Add the new node to the list of nodes in state
}); });
selection?.nodes.forEach((sNode: NodeType) => {
if (newFlow.edges) let newId = idsMap[sNode.id];
newFlow.edges.forEach((edge: Edge) => { sNode.id = newId;
sNode.data.id = newId;
});
}
const concatedEdges = [...edges, ...(selection?.edges ?? [])];
if (concatedEdges)
concatedEdges.forEach((edge: Edge) => {
edge.source = idsMap[edge.source]; edge.source = idsMap[edge.source];
edge.target = idsMap[edge.target]; edge.target = idsMap[edge.target];
const sourceHandleObject: sourceHandleType = scapeJSONParse( const sourceHandleObject: sourceHandleType = scapeJSONParse(
@ -955,7 +965,7 @@ export function connectedInputNodesOnHandle(
return connectedNodes; return connectedNodes;
} }
function updateProxyIdsOnTemplate( export function updateProxyIdsOnTemplate(
template: APITemplateType, template: APITemplateType,
idsMap: { [key: string]: string } idsMap: { [key: string]: string }
) { ) {
@ -966,12 +976,16 @@ function updateProxyIdsOnTemplate(
}); });
} }
function updateEdgesIds(edges: Edge[], idsMap: { [key: string]: string }) { export function updateEdgesIds(
edges: Edge[],
idsMap: { [key: string]: string }
) {
edges.forEach((edge) => { edges.forEach((edge) => {
let targetHandle: targetHandleType = edge.data.targetHandle; let targetHandle: targetHandleType = edge.data.targetHandle;
if (targetHandle.proxy && idsMap[targetHandle.proxy!.id]) { if (targetHandle.proxy && idsMap[targetHandle.proxy!.id]) {
targetHandle.proxy!.id = idsMap[targetHandle.proxy!.id]; targetHandle.proxy!.id = idsMap[targetHandle.proxy!.id];
} }
console.log("edge", edge);
edge.data.targetHandle = targetHandle; edge.data.targetHandle = targetHandle;
edge.targetHandle = scapedJSONStringfy(targetHandle); edge.targetHandle = scapedJSONStringfy(targetHandle);
}); });

View file

@ -705,3 +705,8 @@ export function sortFields(a, b, fieldOrder) {
// You might want to sort them alphabetically or in another specific manner // You might want to sort them alphabetically or in another specific manner
return a.localeCompare(b); return a.localeCompare(b);
} }
export function freezeObject(obj: any) {
if(!obj) return obj;
return(JSON.parse(JSON.stringify(obj)));
}