zclaw_openfang/desktop/src/store/memoryGraphStore.ts

/**
 * MemoryGraphStore - 记忆图谱状态管理
 *
 * 管理记忆图谱可视化的状态，包括节点、边、布局和交互。
 */

import { create } from 'zustand';
import { persist } from 'zustand/middleware';
import {
  intelligenceClient,
  type MemoryEntry,
  type MemoryType,
} from '../lib/intelligence-client';

export type { MemoryType };

// === Types ===

export interface GraphNode {
  id: string;
  type: MemoryType;
  label: string;
  x: number;
  y: number;
  vx: number;
  vy: number;
  importance: number;
  accessCount: number;
  createdAt: string;
  isHighlighted: boolean;
  isSelected: boolean;
}

export interface GraphEdge {
  id: string;
  source: string;
  target: string;
  type: 'reference' | 'related' | 'derived';
  strength: number;
}

export interface GraphFilter {
  types: MemoryType[];
  minImportance: number;
  dateRange: {
    start?: string;
    end?: string;
  };
  searchQuery: string;
}

export interface GraphLayout {
  width: number;
  height: number;
  zoom: number;
  offsetX: number;
  offsetY: number;
}

interface MemoryGraphState {
  nodes: GraphNode[];
  edges: GraphEdge[];
  isLoading: boolean;
  error: string | null;
  filter: GraphFilter;
  layout: GraphLayout;
  selectedNodeId: string | null;
  hoveredNodeId: string | null;
  showLabels: boolean;
  simulationRunning: boolean;
}

interface MemoryGraphActions {
  loadGraph: (agentId: string) => Promise<void>;
  setFilter: (filter: Partial<GraphFilter>) => void;
  resetFilter: () => void;
  setLayout: (layout: Partial<GraphLayout>) => void;
  selectNode: (nodeId: string | null) => void;
  hoverNode: (nodeId: string | null) => void;
  toggleLabels: () => void;
  startSimulation: () => void;
  stopSimulation: () => void;
  updateNodePositions: (updates: Array<{ id: string; x: number; y: number }>) => void;
  highlightSearch: (query: string) => void;
  clearHighlight: () => void;
  exportAsImage: () => Promise<Blob | null>;
  getFilteredNodes: () => GraphNode[];
  getFilteredEdges: () => GraphEdge[];
}

const DEFAULT_FILTER: GraphFilter = {
  types: ['fact', 'preference', 'lesson', 'context', 'task'],
  minImportance: 0,
  dateRange: {},
  searchQuery: '',
};

const DEFAULT_LAYOUT: GraphLayout = {
  width: 800,
  height: 600,
  zoom: 1,
  offsetX: 0,
  offsetY: 0,
};

export type MemoryGraphStore = MemoryGraphState & MemoryGraphActions;

// === Helper Functions ===

function memoryToNode(memory: MemoryEntry, index: number, total: number): GraphNode {
  // 使用圆形布局初始位置
  const angle = (index / total) * 2 * Math.PI;
  const radius = 200;

  return {
    id: memory.id,
    type: memory.type,
    label: memory.content.slice(0, 50) + (memory.content.length > 50 ? '...' : ''),
    x: 400 + radius * Math.cos(angle),
    y: 300 + radius * Math.sin(angle),
    vx: 0,
    vy: 0,
    importance: memory.importance,
    accessCount: memory.accessCount,
    createdAt: memory.createdAt,
    isHighlighted: false,
    isSelected: false,
  };
}

function findRelatedMemories(memories: MemoryEntry[]): GraphEdge[] {
  const edges: GraphEdge[] = [];

  // 简单的关联算法：基于共同标签和关键词
  for (let i = 0; i < memories.length; i++) {
    for (let j = i + 1; j < memories.length; j++) {
      const m1 = memories[i];
      const m2 = memories[j];

      // 检查共同标签
      const commonTags = m1.tags.filter(t => m2.tags.includes(t));
      if (commonTags.length > 0) {
        edges.push({
          id: `edge-${m1.id}-${m2.id}`,
          source: m1.id,
          target: m2.id,
          type: 'related',
          strength: commonTags.length * 0.3,
        });
      }

      // 同类型记忆关联
      if (m1.type === m2.type) {
        const existingEdge = edges.find(
          e => e.source === m1.id && e.target === m2.id
        );
        if (!existingEdge) {
          edges.push({
            id: `edge-${m1.id}-${m2.id}-type`,
            source: m1.id,
            target: m2.id,
            type: 'derived',
            strength: 0.1,
          });
        }
      }
    }
  }

  return edges;
}

export const useMemoryGraphStore = create<MemoryGraphStore>()(
  persist(
    (set, get) => ({
      nodes: [],
      edges: [],
      isLoading: false,
      error: null,
      filter: DEFAULT_FILTER,
      layout: DEFAULT_LAYOUT,
      selectedNodeId: null,
      hoveredNodeId: null,
      showLabels: true,
      simulationRunning: false,

      loadGraph: async (agentId: string) => {
        set({ isLoading: true, error: null });

        try {
          const memories = await intelligenceClient.memory.search({
            agentId,
            limit: 200,
          });

          const nodes = memories.map((m, i) => memoryToNode(m, i, memories.length));
          const edges = findRelatedMemories(memories);

          set({
            nodes,
            edges,
            isLoading: false,
          });
        } catch (err) {
          set({
            isLoading: false,
            error: err instanceof Error ? err.message : '加载图谱失败',
          });
        }
      },

      setFilter: (filter) => {
        set(state => ({
          filter: { ...state.filter, ...filter },
        }));
      },

      resetFilter: () => {
        set({ filter: DEFAULT_FILTER });
      },

      setLayout: (layout) => {
        set(state => ({
          layout: { ...state.layout, ...layout },
        }));
      },

      selectNode: (nodeId) => {
        set(state => ({
          selectedNodeId: nodeId,
          nodes: state.nodes.map(n => ({
            ...n,
            isSelected: n.id === nodeId,
          })),
        }));
      },

      hoverNode: (nodeId) => {
        set(state => ({
          hoveredNodeId: nodeId,
          nodes: state.nodes.map(n => ({
            ...n,
            isHighlighted: nodeId ? n.id === nodeId : n.isHighlighted,
          })),
        }));
      },

      toggleLabels: () => {
        set(state => ({ showLabels: !state.showLabels }));
      },

      startSimulation: () => {
        set({ simulationRunning: true });
      },

      stopSimulation: () => {
        set({ simulationRunning: false });
      },

      updateNodePositions: (updates) => {
        set(state => ({
          nodes: state.nodes.map(node => {
            const update = updates.find(u => u.id === node.id);
            return update ? { ...node, x: update.x, y: update.y } : node;
          }),
        }));
      },

      highlightSearch: (query) => {
        const lowerQuery = query.toLowerCase();
        set(state => ({
          filter: { ...state.filter, searchQuery: query },
          nodes: state.nodes.map(n => ({
            ...n,
            isHighlighted: query ? n.label.toLowerCase().includes(lowerQuery) : false,
          })),
        }));
      },

      clearHighlight: () => {
        set(state => ({
          nodes: state.nodes.map(n => ({ ...n, isHighlighted: false })),
        }));
      },

      exportAsImage: async () => {
        // SVG 导出逻辑在组件中实现
        return null;
      },

      getFilteredNodes: () => {
        const { nodes, filter } = get();
        return nodes.filter(n => {
          if (!filter.types.includes(n.type)) return false;
          if (n.importance < filter.minImportance) return false;
          if (filter.dateRange.start && n.createdAt < filter.dateRange.start) return false;
          if (filter.dateRange.end && n.createdAt > filter.dateRange.end) return false;
          if (filter.searchQuery) {
            return n.label.toLowerCase().includes(filter.searchQuery.toLowerCase());
          }
          return true;
        });
      },

      getFilteredEdges: () => {
        const { edges } = get();
        const filteredNodes = get().getFilteredNodes();
        const nodeIds = new Set(filteredNodes.map(n => n.id));

        return edges.filter(e => nodeIds.has(e.source) && nodeIds.has(e.target));
      },
    }),
    {
      name: 'zclaw-memory-graph',
      partialize: (state) => ({
        filter: state.filter,
        layout: state.layout,
        showLabels: state.showLabels,
      }),
    }
  )
);