obsidian_ollama/src/conversation-state.ts

// src/conversation-state.ts

import type { OllamaMessage } from './types';

export interface ConversationState {
  shortTermContext: OllamaMessage[];
  mediumTermContext: OllamaMessage[];
  longTermContext: OllamaMessage[];
}

const DEFAULT_SYSTEM_PROMPT = `You are an assistant that can help answer questions using the contents of a vault.
When a user asks for information about their vault, you MUST call the search_vault_files or read_vault_file tool to find the answer.
Do not say you will search or read files — immediately emit the tool_call.`;

export class ConversationStateManager {
  private shortTermContext: OllamaMessage[] = [];
  private mediumTermContext: OllamaMessage[] = [];
  private longTermContext: OllamaMessage[] = [];
  private maxShortTermTurns: number = 10;
  private maxMediumTermMessages: number = 20;

  constructor(initialSystemPrompt?: string) {
    // Initialize with default system context
    this.longTermContext = [
      {
        role: 'system',
        content: initialSystemPrompt ?? DEFAULT_SYSTEM_PROMPT,
      },
    ];
  }

  /**
   * Updates the short-term context with a new message
   * @param message The message to add to short-term context
   */
  updateShortTermContext(message: OllamaMessage): void {
    // Add new message
    this.shortTermContext.push(message);

    // Limit to max turns
    if (this.shortTermContext.length > this.maxShortTermTurns) {
      this.shortTermContext = this.shortTermContext.slice(-this.maxShortTermTurns);
    }
  }

  /**
   * Updates the medium-term context with a new message
   * @param message The message to add to medium-term context
   */
  updateMediumTermContext(message: OllamaMessage): void {
    // Add new message
    this.mediumTermContext.push(message);

    // Limit to max messages
    if (this.mediumTermContext.length > this.maxMediumTermMessages) {
      this.mediumTermContext = this.mediumTermContext.slice(-this.maxMediumTermMessages);
    }
  }

  /**
   * Sets the user's persona or core knowledge as long-term context
   * @param personaContent The persona or core knowledge content
   */
  setSystemPrompt(systemPrompt: string): void {
    // Replace all existing system messages with the new system prompt
    this.longTermContext = this.longTermContext.filter((msg) => msg.role !== 'system');
    this.longTermContext.unshift({
      role: 'system',
      content: systemPrompt,
    });
  }

  setPersona(personaContent: string): void {
    // Replace all existing system messages with the new persona
    this.longTermContext = this.longTermContext.filter((msg) => msg.role !== 'system');
    this.longTermContext.push({
      role: 'system',
      content: personaContent,
    });
  }

  /**
   * Gets the combined conversation context for the current turn
   * @param userMessage The user's current message
   * @returns Complete conversation context with all three layers
   */
  getConversationContext(_userMessage: string): ConversationState {
    return {
      shortTermContext: this.shortTermContext,
      mediumTermContext: this.mediumTermContext,
      longTermContext: this.longTermContext,
    };
  }

  /**
   * Gets the complete messages array for sending to the LLM
   * @param userMessage The user's current message
   * @returns Complete message array for the LLM
   */
  getCompleteMessages(userMessage: string): OllamaMessage[] {
    const userMessageWithContext: OllamaMessage = {
      role: 'user',
      content: userMessage,
    };

    // Build messages in the proper order:
    // 1. Long-term context (user persona, system instructions)
    // 2. Medium-term context (session knowledge base query results)
    // 3. Short-term context (last N turns)
    // 4. Current user message
    return [
      ...this.longTermContext,
      ...this.mediumTermContext,
      ...this.shortTermContext,
      userMessageWithContext,
    ];
  }

  /**
   * Clears all conversation context
   */
  clear(systemPrompt?: string): void {
    this.shortTermContext = [];
    this.mediumTermContext = [];
    this.longTermContext = [
      {
        role: 'system',
        content: systemPrompt ?? DEFAULT_SYSTEM_PROMPT,
      },
    ];
  }

  /**
   * Sets the medium-term context from a knowledge base query result
   * @param queryResult The result from a knowledge base query
   */
  setMediumTermContextFromQuery(queryResult: string): void {
    // Clear previous medium-term context
    this.mediumTermContext = [];

    // Add the query result as context
    if (queryResult.trim()) {
      this.mediumTermContext.push({
        role: 'system',
        content: `Knowledge base results for current query:\n${queryResult}`,
      });
    }
  }

  /**
   * Gets the current short-term context
   */
  getShortTermContext(): OllamaMessage[] {
    return [...this.shortTermContext];
  }

  /**
   * Gets the current medium-term context
   */
  getMediumTermContext(): OllamaMessage[] {
    return [...this.mediumTermContext];
  }

  /**
   * Gets the current long-term context
   */
  getLongTermContext(): OllamaMessage[] {
    return [...this.longTermContext];
  }
}