@boundless-oss/atlas/src/modules/local-ai/graph-memory.ts

Atlas - MCP Server for comprehensive startup project management

/**
 * Graph-based memory system for code intelligence
 * Explores different database approaches for agent memory
 */

export interface GraphNode {
  id: string;
  type: 'file' | 'function' | 'class' | 'variable' | 'import' | 'test' | 'commit';
  name: string;
  properties: Record<string, any>;
  timestamp: Date;
}

export interface GraphEdge {
  from: string;
  to: string;
  type: 'calls' | 'imports' | 'extends' | 'implements' | 'tests' | 'modifies' | 'depends_on';
  properties?: Record<string, any>;
  timestamp: Date;
}

export interface MemoryQuery {
  type: 'graph' | 'temporal' | 'semantic';
  query: string;
  timeRange?: { start: Date; end: Date };
  depth?: number;
}

/**
 * In-memory graph database for code relationships
 * Lightweight alternative to Neo4j for local use
 */
export class CodeGraphMemory {
  private nodes: Map<string, GraphNode> = new Map();
  private edges: Map<string, GraphEdge[]> = new Map();
  private temporalIndex: Map<number, string[]> = new Map(); // timestamp -> node IDs
  
  /**
   * Add a node to the graph
   */
  addNode(node: GraphNode): void {
    this.nodes.set(node.id, node);
    
    // Add to temporal index
    const timestamp = node.timestamp.getTime();
    if (!this.temporalIndex.has(timestamp)) {
      this.temporalIndex.set(timestamp, []);
    }
    this.temporalIndex.get(timestamp)!.push(node.id);
  }
  
  /**
   * Add an edge between nodes
   */
  addEdge(edge: GraphEdge): void {
    if (!this.edges.has(edge.from)) {
      this.edges.set(edge.from, []);
    }
    this.edges.get(edge.from)!.push(edge);
  }
  
  /**
   * Find all nodes connected to a given node
   */
  getConnectedNodes(nodeId: string, edgeType?: string, depth: number = 1): GraphNode[] {
    const visited = new Set<string>();
    const result: GraphNode[] = [];
    
    const traverse = (id: string, currentDepth: number) => {
      if (visited.has(id) || currentDepth > depth) return;
      visited.add(id);
      
      const node = this.nodes.get(id);
      if (node && currentDepth > 0) {
        result.push(node);
      }
      
      const edges = this.edges.get(id) || [];
      for (const edge of edges) {
        if (!edgeType || edge.type === edgeType) {
          traverse(edge.to, currentDepth + 1);
        }
      }
    };
    
    traverse(nodeId, 0);
    return result;
  }
  
  /**
   * Find nodes modified within a time range
   */
  getNodesInTimeRange(start: Date, end: Date): GraphNode[] {
    const startTime = start.getTime();
    const endTime = end.getTime();
    const result: GraphNode[] = [];
    
    for (const [timestamp, nodeIds] of this.temporalIndex) {
      if (timestamp >= startTime && timestamp <= endTime) {
        for (const nodeId of nodeIds) {
          const node = this.nodes.get(nodeId);
          if (node) result.push(node);
        }
      }
    }
    
    return result;
  }
  
  /**
   * Find dependency chains
   */
  findDependencyChain(fromId: string, toId: string): GraphNode[][] {
    const paths: GraphNode[][] = [];
    const visited = new Set<string>();
    
    const dfs = (current: string, path: GraphNode[]) => {
      if (current === toId) {
        paths.push([...path]);
        return;
      }
      
      if (visited.has(current)) return;
      visited.add(current);
      
      const edges = this.edges.get(current) || [];
      for (const edge of edges) {
        if (edge.type === 'depends_on' || edge.type === 'imports') {
          const nextNode = this.nodes.get(edge.to);
          if (nextNode) {
            dfs(edge.to, [...path, nextNode]);
          }
        }
      }
      
      visited.delete(current);
    };
    
    const startNode = this.nodes.get(fromId);
    if (startNode) {
      dfs(fromId, [startNode]);
    }
    
    return paths;
  }
  
  /**
   * Detect circular dependencies
   */
  detectCycles(): string[][] {
    const cycles: string[][] = [];
    const visited = new Set<string>();
    const recursionStack = new Set<string>();
    
    const dfs = (nodeId: string, path: string[]) => {
      visited.add(nodeId);
      recursionStack.add(nodeId);
      
      const edges = this.edges.get(nodeId) || [];
      for (const edge of edges) {
        if (!visited.has(edge.to)) {
          dfs(edge.to, [...path, edge.to]);
        } else if (recursionStack.has(edge.to)) {
          // Found a cycle
          const cycleStart = path.indexOf(edge.to);
          if (cycleStart !== -1) {
            cycles.push(path.slice(cycleStart));
          }
        }
      }
      
      recursionStack.delete(nodeId);
    };
    
    for (const nodeId of this.nodes.keys()) {
      if (!visited.has(nodeId)) {
        dfs(nodeId, [nodeId]);
      }
    }
    
    return cycles;
  }
}

/**
 * Document-based memory using embeddings and metadata
 */
export class DocumentMemory {
  private documents: Map<string, {
    content: string;
    metadata: Record<string, any>;
    embedding?: number[];
    timestamp: Date;
  }> = new Map();
  
  addDocument(id: string, content: string, metadata: Record<string, any>): void {
    this.documents.set(id, {
      content,
      metadata,
      timestamp: new Date(),
    });
  }
  
  searchByMetadata(query: Record<string, any>): Array<[string, any]> {
    const results: Array<[string, any]> = [];
    
    for (const [id, doc] of this.documents) {
      let matches = true;
      for (const [key, value] of Object.entries(query)) {
        if (doc.metadata[key] !== value) {
          matches = false;
          break;
        }
      }
      if (matches) {
        results.push([id, doc]);
      }
    }
    
    return results;
  }
}

/**
 * Time-series memory for tracking code evolution
 */
export class TimeSeriesMemory {
  private events: Array<{
    timestamp: Date;
    type: string;
    data: any;
    tags: string[];
  }> = [];
  
  addEvent(type: string, data: any, tags: string[] = []): void {
    this.events.push({
      timestamp: new Date(),
      type,
      data,
      tags,
    });
  }
  
  queryByTimeRange(start: Date, end: Date, type?: string): any[] {
    return this.events.filter(event => {
      const inRange = event.timestamp >= start && event.timestamp <= end;
      const matchesType = !type || event.type === type;
      return inRange && matchesType;
    });
  }
  
  getEventPatterns(windowSize: number = 3600000): Record<string, number> {
    const patterns: Record<string, number> = {};
    const now = Date.now();
    
    for (const event of this.events) {
      const age = now - event.timestamp.getTime();
      if (age <= windowSize) {
        patterns[event.type] = (patterns[event.type] || 0) + 1;
      }
    }
    
    return patterns;
  }
}

/**
 * Unified memory interface combining all approaches
 */
export class HybridCodeMemory {
  private graph: CodeGraphMemory;
  private documents: DocumentMemory;
  private timeSeries: TimeSeriesMemory;
  
  constructor() {
    this.graph = new CodeGraphMemory();
    this.documents = new DocumentMemory();
    this.timeSeries = new TimeSeriesMemory();
  }
  
  /**
   * Record a code change with full context
   */
  recordCodeChange(change: {
    file: string;
    type: 'create' | 'modify' | 'delete';
    functions?: string[];
    imports?: string[];
    tests?: string[];
  }): void {
    // Add to graph
    const nodeId = `file:${change.file}:${Date.now()}`;
    this.graph.addNode({
      id: nodeId,
      type: 'file',
      name: change.file,
      properties: { changeType: change.type },
      timestamp: new Date(),
    });
    
    // Create relationships
    if (change.imports) {
      for (const imp of change.imports) {
        this.graph.addEdge({
          from: nodeId,
          to: `import:${imp}`,
          type: 'imports',
          timestamp: new Date(),
        });
      }
    }
    
    // Add to time series
    this.timeSeries.addEvent('code_change', change, [change.type, change.file]);
    
    // Add to documents
    this.documents.addDocument(nodeId, JSON.stringify(change), {
      file: change.file,
      type: change.type,
      timestamp: new Date(),
    });
  }
  
  /**
   * Query memory with context awareness
   */
  query(query: MemoryQuery): any {
    switch (query.type) {
      case 'graph':
        // Use graph for relationship queries
        return this.graph.getConnectedNodes(query.query, undefined, query.depth);
        
      case 'temporal':
        // Use time series for time-based queries
        if (query.timeRange) {
          return this.timeSeries.queryByTimeRange(
            query.timeRange.start,
            query.timeRange.end
          );
        }
        break;
        
      case 'semantic':
        // Use document store for content queries
        return this.documents.searchByMetadata({ type: query.query });
    }
  }
  
  /**
   * Get coding patterns and insights
   */
  getInsights(): {
    recentPatterns: Record<string, number>;
    dependencies: string[][];
    hotspots: string[];
  } {
    return {
      recentPatterns: this.timeSeries.getEventPatterns(),
      dependencies: this.graph.detectCycles(),
      hotspots: this.findHotspots(),
    };
  }
  
  private findHotspots(): string[] {
    const changeFrequency: Record<string, number> = {};
    const recentEvents = this.timeSeries.queryByTimeRange(
      new Date(Date.now() - 7 * 24 * 60 * 60 * 1000), // Last week
      new Date()
    );
    
    for (const event of recentEvents) {
      if (event.data.file) {
        changeFrequency[event.data.file] = (changeFrequency[event.data.file] || 0) + 1;
      }
    }
    
    return Object.entries(changeFrequency)
      .sort((a, b) => b[1] - a[1])
      .slice(0, 10)
      .map(([file]) => file);
  }
}