From d9868550e66fe8aaa7fff55a8e24b871ee51e3b1 Mon Sep 17 00:00:00 2001 From: blackhao <13851610112@163.com> Date: Fri, 5 Dec 2025 20:40:40 -0600 Subject: init: add project files and ignore secrets --- frontend/src/store/flowStore.ts | 405 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 405 insertions(+) create mode 100644 frontend/src/store/flowStore.ts (limited to 'frontend/src/store/flowStore.ts') diff --git a/frontend/src/store/flowStore.ts b/frontend/src/store/flowStore.ts new file mode 100644 index 0000000..6083a36 --- /dev/null +++ b/frontend/src/store/flowStore.ts @@ -0,0 +1,405 @@ +import { create } from 'zustand'; +import { + addEdge, + applyNodeChanges, + applyEdgeChanges, + type Connection, + type Edge, + type EdgeChange, + type Node, + type NodeChange, + type OnNodesChange, + type OnEdgesChange, + type OnConnect, + getIncomers, + getOutgoers +} from 'reactflow'; + +export type NodeStatus = 'idle' | 'loading' | 'success' | 'error'; + +export interface Message { + id?: string; + role: 'user' | 'assistant' | 'system'; + content: string; +} + +export interface Trace { + id: string; + sourceNodeId: string; + color: string; + messages: Message[]; +} + +export interface NodeData { + label: string; + model: string; + temperature: number; + apiKey?: string; + systemPrompt: string; + userPrompt: string; + mergeStrategy: 'raw' | 'smart'; + + // Traces logic + traces: Trace[]; // INCOMING Traces + outgoingTraces: Trace[]; // ALL Outgoing (inherited + self + forks) + forkedTraces: Trace[]; // Manually created forks from "New" handle + activeTraceIds: string[]; + + response: string; + status: NodeStatus; + inputs: number; + [key: string]: any; +} + +export type LLMNode = Node; + +interface FlowState { + nodes: LLMNode[]; + edges: Edge[]; + selectedNodeId: string | null; + + onNodesChange: OnNodesChange; + onEdgesChange: OnEdgesChange; + onConnect: OnConnect; + + addNode: (node: LLMNode) => void; + updateNodeData: (nodeId: string, data: Partial) => void; + setSelectedNode: (nodeId: string | null) => void; + + getActiveContext: (nodeId: string) => Message[]; + + propagateTraces: () => void; +} + +// Hash string to color +const getStableColor = (str: string) => { + let hash = 0; + for (let i = 0; i < str.length; i++) { + hash = str.charCodeAt(i) + ((hash << 5) - hash); + } + const hue = Math.abs(hash % 360); + return `hsl(${hue}, 70%, 60%)`; +}; + +const useFlowStore = create((set, get) => ({ + nodes: [], + edges: [], + selectedNodeId: null, + + onNodesChange: (changes: NodeChange[]) => { + set({ + nodes: applyNodeChanges(changes, get().nodes) as LLMNode[], + }); + }, + onEdgesChange: (changes: EdgeChange[]) => { + set({ + edges: applyEdgeChanges(changes, get().edges), + }); + get().propagateTraces(); + }, + onConnect: (connection: Connection) => { + const { nodes } = get(); + + // Check if connecting from "new-trace" handle + if (connection.sourceHandle === 'new-trace') { + // Logic: Create a new Forked Trace on the source node + const sourceNode = nodes.find(n => n.id === connection.source); + if (sourceNode) { + // Generate the content for this new trace (it's essentially the Self Trace of this node) + const myResponseMsg: Message[] = []; + if (sourceNode.data.userPrompt) myResponseMsg.push({ id: `${sourceNode.id}-u`, role: 'user', content: sourceNode.data.userPrompt }); + if (sourceNode.data.response) myResponseMsg.push({ id: `${sourceNode.id}-a`, role: 'assistant', content: sourceNode.data.response }); + + const newForkId = `trace-${sourceNode.id}-fork-${Date.now()}`; + const newForkTrace: Trace = { + id: newForkId, + sourceNodeId: sourceNode.id, + color: getStableColor(newForkId), // Unique color for this fork + messages: [...myResponseMsg] + }; + + // Update Source Node to include this fork + get().updateNodeData(sourceNode.id, { + forkedTraces: [...(sourceNode.data.forkedTraces || []), newForkTrace] + }); + + // Redirect connection to the new handle + // Note: We must wait for propagateTraces to render the new handle? + // ReactFlow might complain if handle doesn't exist yet. + // But since we updateNodeData synchronously (mostly), it might work. + // Let's use the new ID for the connection. + + set({ + edges: addEdge({ + ...connection, + sourceHandle: `trace-${newForkId}`, // Redirect! + style: { stroke: newForkTrace.color, strokeWidth: 2 } + }, get().edges), + }); + + // Trigger propagation to update downstream + setTimeout(() => get().propagateTraces(), 0); + return; + } + } + + // Normal connection + set({ + edges: addEdge({ + ...connection, + style: { stroke: '#888', strokeWidth: 2 } + }, get().edges), + }); + setTimeout(() => get().propagateTraces(), 0); + }, + + addNode: (node: LLMNode) => { + set((state) => ({ nodes: [...state.nodes, node] })); + setTimeout(() => get().propagateTraces(), 0); + }, + + updateNodeData: (nodeId: string, data: Partial) => { + set((state) => ({ + nodes: state.nodes.map((node) => { + if (node.id === nodeId) { + return { ...node, data: { ...node.data, ...data } }; + } + return node; + }), + })); + + if (data.response !== undefined || data.userPrompt !== undefined) { + get().propagateTraces(); + } + }, + + setSelectedNode: (nodeId: string | null) => { + set({ selectedNodeId: nodeId }); + }, + + getActiveContext: (nodeId: string) => { + const node = get().nodes.find(n => n.id === nodeId); + if (!node) return []; + + // The traces stored in node.data.traces are the INCOMING traces. + // If we select one, we want its history. + + const activeTraces = node.data.traces.filter(t => + node.data.activeTraceIds?.includes(t.id) + ); + + const contextMessages: Message[] = []; + activeTraces.forEach(t => { + contextMessages.push(...t.messages); + }); + + return contextMessages; + }, + + propagateTraces: () => { + const { nodes, edges } = get(); + + // We need to calculate traces for each node, AND update edge colors. + // Topological Sort + const inDegree = new Map(); + const graph = new Map(); + + nodes.forEach(node => { + inDegree.set(node.id, 0); + graph.set(node.id, []); + }); + + edges.forEach(edge => { + inDegree.set(edge.target, (inDegree.get(edge.target) || 0) + 1); + graph.get(edge.source)?.push(edge.target); + }); + + const topoQueue: string[] = []; + inDegree.forEach((count, id) => { + if (count === 0) topoQueue.push(id); + }); + + const sortedNodes: string[] = []; + while (topoQueue.length > 0) { + const u = topoQueue.shift()!; + sortedNodes.push(u); + + const children = graph.get(u) || []; + children.forEach(v => { + inDegree.set(v, (inDegree.get(v) || 0) - 1); + if (inDegree.get(v) === 0) { + topoQueue.push(v); + } + }); + } + + // Map: Traces LEAVING this node + const nodeOutgoingTraces = new Map(); + // Map: Traces ENTERING this node (to update NodeData) + const nodeIncomingTraces = new Map(); + + // Also track Edge updates (Color AND SourceHandle) + const updatedEdges = [...edges]; + let edgesChanged = false; + + // Iterate + sortedNodes.forEach(nodeId => { + const node = nodes.find(n => n.id === nodeId); + if (!node) return; + + // 1. Gather Incoming Traces + const incomingEdges = edges.filter(e => e.target === nodeId); + const myIncomingTraces: Trace[] = []; + + incomingEdges.forEach(edge => { + const parentOutgoing = nodeOutgoingTraces.get(edge.source) || []; + + // Find match based on Handle ID + // EXACT match first + // Since we removed 'new-trace' handle, we only look for exact trace matches. + let matchedTrace = parentOutgoing.find(t => edge.sourceHandle === `trace-${t.id}`); + + // If no exact match, try to find a "Semantic Match" (Auto-Reconnect) + // If edge.sourceHandle was 'trace-X', and now we have 'trace-X_Parent', that's a likely evolution. + if (!matchedTrace && edge.sourceHandle?.startsWith('trace-')) { + const oldId = edge.sourceHandle.replace('trace-', ''); + matchedTrace = parentOutgoing.find(t => t.id === `${oldId}_${edge.source}`); + } + + // Fallback: If still no match, and parent has traces, try to connect to the most logical one. + // If parent has only 1 trace, connect to it. + // This handles cases where edge.sourceHandle might be null or outdated. + if (!matchedTrace && parentOutgoing.length > 0) { + // If edge has no handle ID, default to the last generated trace (usually Self Trace) + if (!edge.sourceHandle) { + matchedTrace = parentOutgoing[parentOutgoing.length - 1]; + } + } + + if (matchedTrace) { + myIncomingTraces.push(matchedTrace); + + // Update Edge Visuals & Logical Connection + const edgeIndex = updatedEdges.findIndex(e => e.id === edge.id); + if (edgeIndex !== -1) { + const currentEdge = updatedEdges[edgeIndex]; + const newHandleId = `trace-${matchedTrace.id}`; + + // Check if we need to update + if (currentEdge.sourceHandle !== newHandleId || currentEdge.style?.stroke !== matchedTrace.color) { + updatedEdges[edgeIndex] = { + ...currentEdge, + sourceHandle: newHandleId, // Auto-update handle connection! + style: { ...currentEdge.style, stroke: matchedTrace.color, strokeWidth: 2 } + }; + edgesChanged = true; + } + } + } + }); + + // Deduplicate incoming traces by ID (in case multiple edges carry same trace) + const uniqueIncoming = Array.from(new Map(myIncomingTraces.map(t => [t.id, t])).values()); + nodeIncomingTraces.set(nodeId, uniqueIncoming); + + // 2. Generate Outgoing Traces + // Every incoming trace gets appended with this node's response. + // PLUS, we always generate a "Self Trace" (Start New) that starts here. + + const myResponseMsg: Message[] = []; + if (node.data.userPrompt) { + myResponseMsg.push({ + id: `${node.id}-user`, // Deterministic ID for stability + role: 'user', + content: node.data.userPrompt + }); + } + if (node.data.response) { + myResponseMsg.push({ + id: `${node.id}-assistant`, + role: 'assistant', + content: node.data.response + }); + } + + const myOutgoingTraces: Trace[] = []; + + // A. Pass-through traces (append history) + uniqueIncoming.forEach(t => { + // When a trace passes through a node and gets modified, it effectively becomes a NEW branch of that trace. + // We must append the current node ID to the trace ID to distinguish branches. + // e.g. Trace "root" -> passes Node A -> becomes "root_A" + // If it passes Node B -> becomes "root_B" + // Downstream Node D can then distinguish "root_A" from "root_B". + + // Match Logic: + // We need to find if this edge was PREVIOUSLY connected to a trace that has now evolved into 'newTrace'. + // The edge.sourceHandle might be the OLD ID. + // We need a heuristic: if edge.sourceHandle contains the ROOT ID of this trace, we assume it's a match. + // But this is risky if multiple branches exist. + + // Better heuristic: + // When we extend a trace t -> t_new (with id t.id + '_' + node.id), + // we record this evolution mapping. + + const newTraceId = `${t.id}_${node.id}`; + + myOutgoingTraces.push({ + ...t, + id: newTraceId, + messages: [...t.messages, ...myResponseMsg] + }); + }); + + // B. Self Trace (New Branch) -> This is the "Default" self trace (always there?) + // Actually, if we use Manual Forks, maybe we don't need an automatic self trace? + // Or maybe the "Default" self trace is just one of the outgoing ones. + // Let's keep it for compatibility if downstream picks it up automatically. + const selfTrace: Trace = { + id: `trace-${node.id}`, + sourceNodeId: node.id, + color: getStableColor(node.id), + messages: [...myResponseMsg] + }; + myOutgoingTraces.push(selfTrace); + + // C. Manual Forks + if (node.data.forkedTraces) { + // We need to keep them updated with the latest messages (if prompt changed) + // But keep their IDs and Colors stable. + const updatedForks = node.data.forkedTraces.map(fork => ({ + ...fork, + messages: [...myResponseMsg] // Re-sync messages + })); + myOutgoingTraces.push(...updatedForks); + } + + nodeOutgoingTraces.set(nodeId, myOutgoingTraces); + + // Update Node Data with INCOMING traces (for sidebar selection) + // We store uniqueIncoming in node.data.traces + // Note: We need to update the node in the `nodes` array, but we are inside the loop. + // We'll do a bulk set at the end. + }); + + // Bulk Update Store + set(state => ({ + edges: updatedEdges, + nodes: state.nodes.map(n => { + const traces = nodeIncomingTraces.get(n.id) || []; + const outTraces = nodeOutgoingTraces.get(n.id) || []; + return { + ...n, + data: { + ...n.data, + traces, + outgoingTraces: outTraces, + activeTraceIds: n.data.activeTraceIds + } + }; + }) + })); + } +})); + +export default useFlowStore; -- cgit v1.2.3