Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 8 commits into from Mar 20, 2024
4 changes: 2 additions & 2 deletions src/backend/langflow/api/v1/chat.py
Expand Up @@ -202,7 +202,7 @@ async def build_vertex_stream(
async def stream_vertex():
try:
if not session_id:
cache = chat_service.get_cache(flow_id)
cache = await chat_service.get_cache(flow_id)
if not cache:
# If there's no cache
raise ValueError(f"No cache found for {flow_id}.")
Expand Down Expand Up @@ -252,7 +252,7 @@ async def stream_vertex():
raise ValueError(f"No result found for vertex {vertex_id}")

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)}))
finally:
logger.debug("Closing stream")
Expand Down
8 changes: 7 additions & 1 deletion src/backend/langflow/graph/graph/base.py
Expand Up @@ -377,7 +377,10 @@ def update(self, other: "Graph") -> "Graph":

# Remove vertices that are not in the other graph
for vertex_id in removed_vertex_ids:
self.remove_vertex(vertex_id)
try:
self.remove_vertex(vertex_id)
except ValueError:
pass

# The order here matters because adding the vertex is required
# if any of them have edges that point to any of the new vertices
Expand Down Expand Up @@ -741,8 +744,11 @@ def _build_vertices(self) -> List[Vertex]:
vertex_data = vertex["data"]
vertex_type: str = vertex_data["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"])

vertex_instance = VertexClass(vertex, graph=self)
vertex_instance.set_top_level(self.top_level_vertices)
vertices.append(vertex_instance)
Expand Down
104 changes: 53 additions & 51 deletions src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
@@ -1,4 +1,4 @@
import _ from "lodash";
import _, { cloneDeep, set } from "lodash";
import { MouseEvent, useCallback, useEffect, useRef, useState } from "react";
import ReactFlow, {
Background,
Expand Down Expand Up @@ -32,6 +32,7 @@ import {
isValidConnection,
reconnectEdges,
scapeJSONParse,
updateIds,
validateSelection,
} from "../../../../utils/reactflowUtils";
import { getRandomName, isWrappedWithClass } from "../../../../utils/utils";
Expand Down Expand Up @@ -104,35 +105,40 @@ export default function Page({
) {
event.preventDefault();
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(
lastSelection!,
nodes,
edges,
clonedSelection!,
clonedNodes,
clonedEdges,
getRandomName()
);
const newGroupNode = generateNodeFromFlow(newFlow, getNodeId);
const newEdges = reconnectEdges(newGroupNode, removedEdges);
setNodes((oldNodes) => [
...oldNodes.filter(
(oldNodes) =>
!lastSelection?.nodes.some(
(selectionNode) => selectionNode.id === oldNodes.id
)
),
const newGroupNode = generateNodeFromFlow(
newFlow,
getNodeId
);
const newEdges = reconnectEdges(
newGroupNode,
]);
setEdges((oldEdges) => [
...oldEdges.filter(
(oldEdge) =>
!lastSelection!.nodes.some(
(selectionNode) =>
selectionNode.id === oldEdge.target ||
selectionNode.id === oldEdge.source
)
),
...newEdges,
]);
removedEdges
);
setNodes([...clonedNodes.filter(
(oldNodes) =>
!clonedSelection?.nodes.some(
(selectionNode) =>
selectionNode.id === oldNodes.id
)), newGroupNode])
setEdges([...clonedEdges.filter(
(oldEdge) =>
!clonedSelection!.nodes.some(
(selectionNode) =>
selectionNode.id === oldEdge.target ||
selectionNode.id === oldEdge.source
)), ...newEdges])
} else {
setErrorData({
title: INVALID_SELECTION_ERROR_ALERT,
Expand Down Expand Up @@ -431,7 +437,7 @@ export default function Page({
<div className="h-full w-full">
<div className="h-full w-full" ref={reactFlowWrapper}>
{Object.keys(templates).length > 0 &&
Object.keys(types).length > 0 ? (
Object.keys(types).length > 0 ? (
<div id="react-flow-id" className="h-full w-full">
<ReactFlow
nodes={nodes}
Expand Down Expand Up @@ -480,10 +486,14 @@ export default function Page({
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(
lastSelection!,
nodes,
edges,
clonedSelection!,
clonedNodes,
clonedEdges,
getRandomName()
);
const newGroupNode = generateNodeFromFlow(
Expand All @@ -494,27 +504,19 @@ export default function Page({
newGroupNode,
removedEdges
);
setNodes((oldNodes) => [
...oldNodes.filter(
(oldNodes) =>
!lastSelection?.nodes.some(
(selectionNode) =>
selectionNode.id === oldNodes.id
)
),
newGroupNode,
]);
setEdges((oldEdges) => [
...oldEdges.filter(
(oldEdge) =>
!lastSelection!.nodes.some(
(selectionNode) =>
selectionNode.id === oldEdge.target ||
selectionNode.id === oldEdge.source
)
),
...newEdges,
]);
setNodes([...clonedNodes.filter(
(oldNodes) =>
!clonedSelection?.nodes.some(
(selectionNode) =>
selectionNode.id === oldNodes.id
)), newGroupNode])
setEdges([...clonedEdges.filter(
(oldEdge) =>
!clonedSelection!.nodes.some(
(selectionNode) =>
selectionNode.id === oldEdge.target ||
selectionNode.id === oldEdge.source
)), ...newEdges])
} else {
setErrorData({
title: INVALID_SELECTION_ERROR_ALERT,
Expand Down
38 changes: 33 additions & 5 deletions src/frontend/src/stores/flowStore.ts
@@ -1,4 +1,4 @@
import { cloneDeep } from "lodash";
import { cloneDeep, zip } from "lodash";
import {
Edge,
EdgeChange,
Expand Down Expand Up @@ -26,6 +26,7 @@ import {
ChatOutputType,
FlowPoolObjectType,
FlowStoreType,
VertexLayerElementType,
chatInputType,
} from "../types/zustand/flow";
import { buildVertices } from "../utils/buildUtils";
Expand All @@ -36,6 +37,9 @@ import {
getNodeId,
scapeJSONParse,
scapedJSONStringfy,
updateEdgesIds,
updateIds,
updateProxyIdsOnTemplate,
validateNodes,
} from "../utils/reactflowUtils";
import { getInputsAndOutputs } from "../utils/storeUtils";
Expand Down Expand Up @@ -257,6 +261,14 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
let newId = getNodeId(node.data.type);
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
const newNode: NodeType = {
id: newId,
Expand Down Expand Up @@ -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
// 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)

// 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 = [
...get().verticesBuild!.verticesLayers,
vertexBuildData.next_vertices_ids,
nextVertices,
];
const newIds = [
...get().verticesBuild!.verticesIds,
Expand Down Expand Up @@ -508,12 +529,18 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
get().setIsBuilding(false);
},
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);
setErrorData({ list, title });
get().setIsBuilding(false);
},
onBuildStart: (idList) => {
onBuildStart: (elementList) => {
const idList = elementList
.map((element) => element.reference)
.filter(Boolean) as string[];
useFlowStore.getState().updateBuildStatus(idList, BuildStatus.BUILDING);
},
validateNodes: validateSubgraph,
Expand All @@ -531,7 +558,7 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
updateVerticesBuild: (
vertices: {
verticesIds: string[];
verticesLayers: string[][];
verticesLayers: VertexLayerElementType[][];
runId: string;
} | null
) => {
Expand Down Expand Up @@ -562,6 +589,7 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
},
updateBuildStatus: (nodeIdList: string[], status: BuildStatus) => {
const newFlowBuildStatus = { ...get().flowBuildStatus };
console.log("newFlowBuildStatus", newFlowBuildStatus);
nodeIdList.forEach((id) => {
newFlowBuildStatus[id] = {
status,
Expand Down
9 changes: 7 additions & 2 deletions src/frontend/src/types/zustand/flow/index.ts
Expand Up @@ -35,6 +35,11 @@ export type FlowPoolObjectType = {
buildId: string;
};

export type VertexLayerElementType = {
id: string;
reference?: string;
};

export type FlowPoolType = {
[key: string]: Array<FlowPoolObjectType>;
};
Expand Down Expand Up @@ -103,15 +108,15 @@ export type FlowStoreType = {
updateVerticesBuild: (
vertices: {
verticesIds: string[];
verticesLayers: string[][];
verticesLayers: VertexLayerElementType[][];
runId: string;
} | null
) => void;
addToVerticesBuild: (vertices: string[]) => void;
removeFromVerticesBuild: (vertices: string[]) => void;
verticesBuild: {
verticesIds: string[];
verticesLayers: string[][];
verticesLayers: VertexLayerElementType[][];
runId: string;
} | null;
updateBuildStatus: (nodeId: string[], status: BuildStatus) => void;
Expand Down