@boundless-oss/atlas/src/modules/memory-management/document-memory.ts

Atlas - MCP Server for comprehensive startup project management

import { ConfigManager } from '../../config/config-manager.js';
import { DocumentContext } from './types.js';
import { promises as fs } from 'fs';
import path from 'path';

export class DocumentMemory {
  private configManager: ConfigManager;
  private documentsPath: string;
  private documents: Map<string, DocumentContext>;

  constructor(configManager: ConfigManager) {
    this.configManager = configManager;
    this.documentsPath = '';
    this.documents = new Map();
  }

  async initialize(): Promise<void> {
    const storageManager = this.configManager.getStorageManager();
    const location = await storageManager.getStorageLocation();
    
    this.documentsPath = path.join(location.data, 'memory', 'documents.json');
    
    // Ensure directory exists
    await fs.mkdir(path.dirname(this.documentsPath), { recursive: true });
    
    // Load existing document knowledge
    await this.loadDocuments();
  }

  async updateKnowledge(context: DocumentContext): Promise<void> {
    this.documents.set(context.documentPath, context);
    await this.saveDocuments();
  }

  async getDocumentContext(documentPath: string): Promise<DocumentContext | undefined> {
    return this.documents.get(documentPath);
  }

  async searchDocuments(query: string): Promise<DocumentContext[]> {
    const results: DocumentContext[] = [];
    const queryLower = query.toLowerCase();

    for (const [, doc] of this.documents) {
      let relevance = 0;

      // Check summary
      if (doc.summary.toLowerCase().includes(queryLower)) {
        relevance += 0.4;
      }

      // Check key points
      const keyPointMatch = doc.keyPoints.some(point => 
        point.toLowerCase().includes(queryLower)
      );
      if (keyPointMatch) {
        relevance += 0.3;
      }

      // Check topics
      const topicMatch = doc.topics.some(topic => 
        topic.toLowerCase().includes(queryLower)
      );
      if (topicMatch) {
        relevance += 0.2;
      }

      // Check document path
      if (doc.documentPath.toLowerCase().includes(queryLower)) {
        relevance += 0.1;
      }

      if (relevance > 0.2) {
        results.push(doc);
      }
    }

    // Sort by relevance (importance and recency)
    results.sort((a, b) => {
      const aScore = this.getImportanceScore(a.importance) + this.getRecencyScore(a.lastUpdated);
      const bScore = this.getImportanceScore(b.importance) + this.getRecencyScore(b.lastUpdated);
      return bScore - aScore;
    });

    return results;
  }

  async getRelatedDocuments(documentPath: string): Promise<DocumentContext[]> {
    const targetDoc = this.documents.get(documentPath);
    if (!targetDoc) return [];

    const related: { doc: DocumentContext; score: number }[] = [];

    for (const [path, doc] of this.documents) {
      if (path === documentPath) continue;

      let score = 0;

      // Check if documents reference each other
      if (targetDoc.relatedFiles.includes(path) || doc.relatedFiles.includes(documentPath)) {
        score += 0.5;
      }

      // Check for shared topics
      const sharedTopics = targetDoc.topics.filter(topic => doc.topics.includes(topic));
      score += sharedTopics.length * 0.2;

      // Check for similar key points
      const similarPoints = targetDoc.keyPoints.filter(point => 
        doc.keyPoints.some(otherPoint => 
          this.calculateSimilarity(point, otherPoint) > 0.3
        )
      );
      score += similarPoints.length * 0.1;

      if (score > 0.2) {
        related.push({ doc, score });
      }
    }

    return related
      .sort((a, b) => b.score - a.score)
      .slice(0, 5)
      .map(item => item.doc);
  }

  async generateTopics(content: string): Promise<string[]> {
    // Simple topic extraction based on keywords
    const topics: string[] = [];
    const contentLower = content.toLowerCase();

    const topicKeywords = {
      'api': ['api', 'endpoint', 'rest', 'graphql', 'swagger'],
      'authentication': ['auth', 'login', 'token', 'jwt', 'oauth'],
      'database': ['database', 'sql', 'mongodb', 'postgres', 'schema'],
      'testing': ['test', 'testing', 'jest', 'mocha', 'cypress'],
      'deployment': ['deploy', 'deployment', 'docker', 'kubernetes', 'ci/cd'],
      'frontend': ['react', 'vue', 'angular', 'frontend', 'ui', 'component'],
      'backend': ['server', 'backend', 'express', 'node', 'api'],
      'security': ['security', 'encryption', 'vulnerability', 'secure'],
      'performance': ['performance', 'optimization', 'speed', 'cache'],
      'configuration': ['config', 'configuration', 'settings', 'environment'],
    };

    for (const [topic, keywords] of Object.entries(topicKeywords)) {
      if (keywords.some(keyword => contentLower.includes(keyword))) {
        topics.push(topic);
      }
    }

    return topics;
  }

  async getDocumentStats(): Promise<{
    totalDocuments: number;
    byImportance: Record<string, number>;
    topTopics: Array<{ name: string; count: number }>;
    recentlyUpdated: DocumentContext[];
  }> {
    const documents = Array.from(this.documents.values());
    
    const byImportance: Record<string, number> = {};
    const topicCounts: Record<string, number> = {};

    documents.forEach(doc => {
      byImportance[doc.importance] = (byImportance[doc.importance] || 0) + 1;
      
      doc.topics.forEach(topic => {
        topicCounts[topic] = (topicCounts[topic] || 0) + 1;
      });
    });

    const topTopics = Object.entries(topicCounts)
      .sort(([, a], [, b]) => b - a)
      .slice(0, 10)
      .map(([name, count]) => ({ name, count }));

    const recentlyUpdated = documents
      .sort((a, b) => new Date(b.lastUpdated).getTime() - new Date(a.lastUpdated).getTime())
      .slice(0, 5);

    return {
      totalDocuments: documents.length,
      byImportance,
      topTopics,
      recentlyUpdated,
    };
  }

  private async loadDocuments(): Promise<void> {
    try {
      const data = await fs.readFile(this.documentsPath, 'utf-8');
      const documentsArray: DocumentContext[] = JSON.parse(data);
      
      this.documents.clear();
      for (const doc of documentsArray) {
        this.documents.set(doc.documentPath, doc);
      }
    } catch (error) {
      // File doesn't exist or is invalid, start with empty documents
      this.documents.clear();
    }
  }

  private async saveDocuments(): Promise<void> {
    const documentsArray = Array.from(this.documents.values());
    await fs.writeFile(this.documentsPath, JSON.stringify(documentsArray, null, 2));
  }

  private getImportanceScore(importance: string): number {
    const scores = { low: 1, medium: 2, high: 3, critical: 4 };
    return scores[importance as keyof typeof scores] || 2;
  }

  private getRecencyScore(timestamp: string): number {
    const ageInDays = (Date.now() - new Date(timestamp).getTime()) / (1000 * 60 * 60 * 24);
    return Math.max(0, 1 - (ageInDays * 0.01)); // Decay over time
  }

  private calculateSimilarity(text1: string, text2: string): number {
    const words1 = text1.toLowerCase().split(/\s+/);
    const words2 = text2.toLowerCase().split(/\s+/);
    
    const commonWords = words1.filter(word => words2.includes(word));
    const totalWords = new Set([...words1, ...words2]).size;
    
    return commonWords.length / totalWords;
  }
}