obsidian_ollama/src/structured-memory.ts

// src/structured-memory.ts

import {
  StructuredMemoryData,
  StructuredMemoryConfig,
  ConversationSummary,
  UserPreference,
  LearnedFact,
  OllamaMessage,
} from './types';

export function createDefaultStructuredMemoryData(): StructuredMemoryData {
  return {
    conversationSummaries: [],
    userPreferences: [],
    learnedFacts: [],
  };
}

/**
 * Manages the agent's structured memory: conversation summaries,
 * user preferences, and learned facts. Persists in plugin data JSON.
 */
export class StructuredMemoryManager {
  private data: StructuredMemoryData;
  private config: StructuredMemoryConfig;

  constructor(config: StructuredMemoryConfig, initialData?: StructuredMemoryData) {
    this.config = config;
    this.data = initialData ?? createDefaultStructuredMemoryData();
  }

  /**
   * Replace the in-memory data (e.g., after loading from disk).
   */
  loadData(data: StructuredMemoryData): void {
    this.data = {
      conversationSummaries: data.conversationSummaries ?? [],
      userPreferences: data.userPreferences ?? [],
      learnedFacts: data.learnedFacts ?? [],
    };
  }

  /**
   * Get a serializable copy of the current memory data.
   */
  getData(): StructuredMemoryData {
    return {
      conversationSummaries: [...this.data.conversationSummaries],
      userPreferences: [...this.data.userPreferences],
      learnedFacts: [...this.data.learnedFacts],
    };
  }

  /**
   * Update the config (e.g., when settings change).
   */
  updateConfig(config: StructuredMemoryConfig): void {
    this.config = config;
    this.enforceLimits();
  }

  /**
   * Add a conversation summary, keeping the newest within maxSummaries.
   */
  addConversationSummary(summary: ConversationSummary): void {
    if (!this.config.enabled) return;
    this.data.conversationSummaries.push(summary);
    this.enforceLimits();
  }

  getConversationSummaries(): ConversationSummary[] {
    return [...this.data.conversationSummaries];
  }

  clearConversationSummaries(): void {
    this.data.conversationSummaries = [];
  }

  /**
   * Add or update a user preference. If the key already exists, update it.
   */
  addUserPreference(preference: UserPreference): void {
    if (!this.config.enabled) return;
    const existingIndex = this.data.userPreferences.findIndex((p) => p.key === preference.key);
    if (existingIndex >= 0) {
      this.data.userPreferences[existingIndex] = preference;
    } else {
      this.data.userPreferences.push(preference);
    }
    this.enforceLimits();
  }

  getUserPreference(key: string): UserPreference | undefined {
    return this.data.userPreferences.find((p) => p.key === key);
  }

  getUserPreferences(): UserPreference[] {
    return [...this.data.userPreferences];
  }

  removeUserPreference(key: string): void {
    this.data.userPreferences = this.data.userPreferences.filter((p) => p.key !== key);
  }

  clearUserPreferences(): void {
    this.data.userPreferences = [];
  }

  /**
   * Add a learned fact, deduplicating by content (case-insensitive).
   */
  addLearnedFact(fact: LearnedFact): void {
    if (!this.config.enabled) return;
    const normalizedContent = fact.content.trim().toLowerCase();
    const existingIndex = this.data.learnedFacts.findIndex(
      (f) => f.content.trim().toLowerCase() === normalizedContent
    );
    if (existingIndex >= 0) {
      // Update confidence and timestamp if duplicate
      this.data.learnedFacts[existingIndex] = {
        ...fact,
        timestamp: Date.now(),
        confidence: Math.max(fact.confidence, this.data.learnedFacts[existingIndex].confidence),
      };
    } else {
      this.data.learnedFacts.push(fact);
    }
    this.enforceLimits();
  }

  getLearnedFacts(): LearnedFact[] {
    return [...this.data.learnedFacts];
  }

  getLearnedFactsByCategory(category: LearnedFact['category']): LearnedFact[] {
    return this.data.learnedFacts.filter((f) => f.category === category);
  }

  clearLearnedFacts(): void {
    this.data.learnedFacts = [];
  }

  clearAll(): void {
    this.data = createDefaultStructuredMemoryData();
  }

  /**
   * Build a context string from stored memory for injection into the system prompt.
   * Returns an empty string if memory is disabled or empty.
   */
  buildMemoryContext(): string {
    if (!this.config.enabled) return '';

    const parts: string[] = [];

    const summaries = this.data.conversationSummaries;
    if (summaries.length > 0) {
      parts.push('## Past Conversations');
      for (const s of summaries.slice(-3)) {
        parts.push(`- ${s.topic}: ${s.summary}`);
      }
    }

    const preferences = this.data.userPreferences;
    if (preferences.length > 0) {
      parts.push('## User Preferences');
      for (const p of preferences) {
        parts.push(`- ${p.key}: ${p.value}`);
      }
    }

    const facts = this.data.learnedFacts;
    if (facts.length > 0) {
      parts.push('## Learned Facts');
      for (const f of facts.filter((fact) => fact.confidence >= 0.5).slice(-10)) {
        parts.push(`- ${f.content}`);
      }
    }

    if (parts.length === 0) return '';
    return 'The following is remembered context from past sessions:\n' + parts.join('\n');
  }

  /**
   * Extract likely user preferences from a message using lightweight regex heuristics.
   */
  extractPreferencesFromMessage(message: string): UserPreference[] {
    if (!this.config.enabled) return [];

    const preferences: UserPreference[] = [];
    const patterns = [
      { regex: /i(?:'d| would)?\s+prefer\s+(?:that\s+)?(.+?)(?:\.|$)/i, keyPrefix: 'preference' },
      { regex: /i\s+(?:like|love|enjoy)\s+(.+?)(?:\.|$)/i, keyPrefix: 'preference' },
      { regex: /i\s+(?:dislike|hate|avoid)\s+(.+?)(?:\.|$)/i, keyPrefix: 'preference' },
      { regex: /(?:always|never)\s+(.+?)(?:\.|$)/i, keyPrefix: 'preference' },
      {
        regex: /my\s+(?:favorite|preferred)\s+(\w+)\s+(?:is|are)\s+(.+?)(?:\.|$)/i,
        keyPrefix: 'favorite',
      },
    ];

    for (const { regex, keyPrefix } of patterns) {
      const match = regex.exec(message);
      if (match) {
        const value = match[match.length - 1].trim();
        const key =
          value.length > 30
            ? `${keyPrefix}-${Date.now()}`
            : `${keyPrefix}-${value.toLowerCase().replace(/\s+/g, '-')}`;
        preferences.push({
          key,
          value,
          timestamp: Date.now(),
          source: 'inferred',
        });
      }
    }

    return preferences;
  }

  /**
   * Extract likely facts from a message using lightweight regex heuristics.
   */
  extractFactsFromMessage(message: string): LearnedFact[] {
    if (!this.config.enabled) return [];

    const facts: LearnedFact[] = [];

    // Vault structure patterns
    const folderPattern = /(\/[^\s]+\/(?:[^\s/]+\/)*)/g;
    const folderMatches = message.matchAll(folderPattern);
    for (const match of folderMatches) {
      facts.push({
        id: crypto.randomUUID?.() ?? `fact-${Date.now()}-${Math.random()}`,
        timestamp: Date.now(),
        content: `The vault contains a folder at ${match[1]}.`,
        category: 'vault_structure',
        confidence: 0.6,
      });
    }

    // Topic patterns ("X is a Y")
    const topicPattern = /(\w+(?:\s+\w+){0,5})\s+is\s+(?:a|an|the)\s+(.+?)(?:\.|$)/gi;
    const topicMatches = message.matchAll(topicPattern);
    for (const match of topicMatches) {
      const subject = match[1].trim();
      const predicate = match[2].trim();
      if (subject.length > 2 && predicate.length > 2) {
        facts.push({
          id: crypto.randomUUID?.() ?? `fact-${Date.now()}-${Math.random()}`,
          timestamp: Date.now(),
          content: `${subject} is ${predicate}.`,
          category: 'topic',
          confidence: 0.5,
        });
      }
    }

    return facts;
  }

  /**
   * Generate a simple topic string from a conversation by looking at the first user message.
   */
  summarizeConversation(messages: OllamaMessage[]): { topic: string; keyPoints: string[] } {
    const firstUser = messages.find((m) => m.role === 'user');
    const topic = firstUser
      ? firstUser.content.slice(0, 60).replace(/\n/g, ' ')
      : 'Untitled conversation';

    const keyPoints: string[] = [];
    for (const msg of messages) {
      if (msg.role === 'assistant' && msg.content) {
        const sentences = msg.content
          .split(/[.!?]+/)
          .map((s) => s.trim())
          .filter((s) => s.length > 10 && s.length < 120);
        keyPoints.push(...sentences.slice(0, 2));
      }
      if (keyPoints.length >= 3) break;
    }

    return { topic, keyPoints };
  }

  private enforceLimits(): void {
    if (this.data.conversationSummaries.length > this.config.maxSummaries) {
      this.data.conversationSummaries = this.data.conversationSummaries.slice(
        -this.config.maxSummaries
      );
    }
    if (this.data.userPreferences.length > this.config.maxPreferences) {
      // Keep most recently updated preferences
      const sorted = [...this.data.userPreferences].sort((a, b) => b.timestamp - a.timestamp);
      this.data.userPreferences = sorted.slice(0, this.config.maxPreferences);
    }
    if (this.data.learnedFacts.length > this.config.maxFacts) {
      // Keep highest-confidence facts
      const sorted = [...this.data.learnedFacts].sort((a, b) => b.confidence - a.confidence);
      this.data.learnedFacts = sorted.slice(0, this.config.maxFacts);
    }
  }
}