obsidian_ollama/src/vault-indexer.ts

// src/vault-indexer.ts

import { VaultIndexEntry } from './types';
import { Logger } from './utils';

interface Cache {
  get(key: string): Promise<string | null>;
  put(key: string, value: string): Promise<void>;
  clear(): Promise<void>;
}

class InMemoryCache implements Cache {
  private store: Map<string, string>;

  constructor() {
    this.store = new Map();
  }

  get(key: string): Promise<string | null> {
    return Promise.resolve(this.store.get(key) || null);
  }

  put(key: string, value: string): Promise<void> {
    this.store.set(key, value);
    return Promise.resolve();
  }

  clear(): Promise<void> {
    this.store.clear();
    return Promise.resolve();
  }
}

interface Frontmatter {
  title?: string;
  tags?: string;
}

interface VaultFile {
  basename: string;
  path: string;
}

interface VaultLike {
  getMarkdownFiles(): VaultFile[];
  read(file: VaultFile): Promise<string>;
}

interface TokenizedContent {
  tokens: string[];
  headings: string[];
  frontmatter: Frontmatter;
  firstParagraph?: string;
}

interface ScoreResult {
  score: number;
  matchedFields: string[];
}

class VaultIndexer {
  private vault: VaultLike | null = null;
  private cache?: Cache;

  // Define weights for scoring
  private readonly SCORING_WEIGHTS = {
    HEADING: 5,
    FRONTMATTER_TITLE: 3,
    FRONTMATTER_TAGS: 2.5,
    FIRST_PARAGRAPH: 1.5,
    TOKEN: 1,
  };

  constructor(vault: VaultLike, cache?: Cache) {
    this.vault = vault;
    this.cache = cache;
  }

  async searchVault(query: string, limit: number = 5): Promise<VaultIndexEntry[]> {
    if (!query || !query.trim()) {
      return [];
    }

    if (!this.vault) {
      throw new Error('Vault-like object not provided to VaultIndexer');
    }

    const cacheKey = `query:${query.trim()}:limit:${limit}`;
    if (this.cache) {
      let cachedResults;
      try {
        cachedResults = await this.cache.get(cacheKey);
      } catch {
        // Ignore cache retrieval errors and continue with normal processing
        cachedResults = null;
      }
      if (cachedResults) {
        try {
          const parsedResults = JSON.parse(cachedResults) as VaultIndexEntry[];
          return parsedResults.slice(0, limit);
        } catch {
          // Ignore cache parse errors and continue with normal processing
        }
      }
    }

    const queryTokens = this.tokenize(query.trim());
    const vault = this.vault;
    const allFiles = vault.getMarkdownFiles();
    const results = await this.processFilesInBatches(vault, allFiles, queryTokens);

    const filteredResults = results.sort((a, b) => b.score - a.score).slice(0, limit);

    if (this.cache) {
      try {
        await this.cache.put(cacheKey, JSON.stringify(filteredResults));
      } catch (error) {
        Logger.warn(
          `Failed to cache results for query "${query}": ${error instanceof Error ? error.message : String(error)}`,
          'vault-indexer'
        );
      }
    }

    return filteredResults;
  }

  private async processFilesInBatches(
    vault: VaultLike,
    files: VaultFile[],
    queryTokens: string[]
  ): Promise<VaultIndexEntry[]> {
    const batchSize = 10;
    const results: VaultIndexEntry[] = [];
    const seenPaths = new Set<string>();

    for (let i = 0; i < files.length; i += batchSize) {
      const batch = files.slice(i, i + batchSize);
      const batchResults = await Promise.all(
        batch.map(async (file) => {
          try {
            const content = await vault.read(file);
            const tokenized = this.tokenizeContent(content);
            const scoreResult = this.calculateWeightedScore(tokenized, queryTokens, file);
            if (scoreResult.score > 0) {
              const entry: VaultIndexEntry = {
                path: file.path,
                title: file.basename.replace(/\.md$/, ''),
                content: content.substring(0, 500),
                score: scoreResult.score,
              };
              if (!seenPaths.has(entry.path)) {
                seenPaths.add(entry.path);
                return entry;
              }
              return null;
            }
            return null;
          } catch (error) {
            Logger.warn(
              `Failed to process ${file.path}: ${error instanceof Error ? error.message : String(error)}`,
              'vault-indexer'
            );
            return null;
          }
        })
      );

      const validResults = batchResults.filter(
        (result): result is NonNullable<typeof result> => result !== null
      );
      results.push(...validResults);
    }

    return results;
  }

  private tokenize(text: string): string[] {
    const stopWords = new Set([
      'the',
      'a',
      'an',
      'and',
      'or',
      'but',
      'is',
      'are',
      'was',
      'were',
      'in',
      'on',
      'at',
      'to',
      'of',
      'for',
      'with',
      'as',
      'by',
      'it',
      'its',
      'that',
      'this',
      'these',
      'those',
    ]);
    return text
      .toLowerCase()
      .split(/\W+/)
      .filter((token) => token.length > 1 && !stopWords.has(token));
  }

  private tokenizeContent(content: string): TokenizedContent {
    const tokens: string[] = [];
    const headings: string[] = [];
    const frontmatter: Frontmatter = {};
    let firstParagraph: string | undefined;

    const frontmatterMatch = content.match(/^---(.*?)---/s);
    if (frontmatterMatch) {
      try {
        const frontmatterContent = frontmatterMatch[1];
        const lines = frontmatterContent.trim().split('\n');
        for (const line of lines) {
          const [key, ...valueParts] = line.split(':');
          if (!key) continue;
          const value = valueParts.join(':').trim();
          if (key.trim() === 'title') {
            if (value) {
              frontmatter.title = value;
            }
          } else if (key.trim() === 'tags') {
            if (value) {
              frontmatter.tags = value;
            }
          }
        }
      } catch {
        Logger.warn('Failed to parse frontmatter', 'vault-indexer');
      }
    }

    const headingMatches = content.match(/^#{1,6} (.*?)$/gm);
    if (headingMatches) {
      headings.push(...headingMatches.map((h: string) => h.replace(/^#{1,6} /, '')));
    }

    const paragraphMatch = content.match(/^([^#]*?)(?=#|\n\n|$)/s);
    if (paragraphMatch) {
      firstParagraph = paragraphMatch[1].trim();
    }

    const allText = content
      .replace(/^---.*?---/s, '')
      .replace(/^#.*?$/gm, '')
      .replace(/```.*?```/gs, '')
      .replace(/`.*?`/g, '')
      .replace(/\[.*?\]\(.*?\)/g, '');
    tokens.push(...this.tokenize(allText));

    return { tokens, headings, frontmatter, firstParagraph };
  }

  private calculateWeightedScore(
    tokenized: TokenizedContent,
    queryTokens: string[],
    file?: VaultFile
  ): ScoreResult {
    let totalScore = 0;
    const matchedTokens: Set<string> = new Set<string>();

    for (const queryToken of queryTokens) {
      let tokenScore = 0;
      const stemmed = this.stemToken(queryToken);
      let matched = false;

      if (
        tokenized.frontmatter?.title &&
        this.exactMatch(tokenized.frontmatter.title, queryToken)
      ) {
        tokenScore += this.SCORING_WEIGHTS.FRONTMATTER_TITLE;
        matched = true;
      } else if (
        file &&
        file.basename &&
        this.exactMatch(file.basename.replace(/\.md$/, ''), queryToken)
      ) {
        tokenScore += this.SCORING_WEIGHTS.FRONTMATTER_TITLE;
        matched = true;
      }

      if (tokenized.frontmatter?.tags && this.exactMatch(tokenized.frontmatter.tags, queryToken)) {
        tokenScore += this.SCORING_WEIGHTS.FRONTMATTER_TAGS;
        matched = true;
      }

      if (tokenized.headings.some((heading: string) => heading.toLowerCase().includes(stemmed))) {
        tokenScore += this.SCORING_WEIGHTS.HEADING;
        matched = true;
      }

      if (tokenized.firstParagraph && tokenized.firstParagraph.toLowerCase().includes(stemmed)) {
        tokenScore += this.SCORING_WEIGHTS.FIRST_PARAGRAPH;
        matched = true;
      }

      if (tokenized.tokens.includes(stemmed)) {
        tokenScore += this.SCORING_WEIGHTS.TOKEN;
        matched = true;
      }

      if (matched) {
        totalScore += tokenScore;
        matchedTokens.add(queryToken);
      }
    }

    return {
      score: totalScore,
      matchedFields: Array.from(matchedTokens),
    };
  }

  private stemToken(token: string): string {
    // Improved stemmer that handles edge cases
    //
    // Limitations:
    // - Simple suffix removal (e.g., 's', 'ed', 'ing') may lead to over-stemming or incorrect roots.
    // - For example, stemming "mice" results in "mic", which is incorrect.
    // - Consider using a more robust NLP library if the plugin environment permits.
    //
    if (token.length <= 3) return token; // Don't stem very short tokens
    if (token.endsWith('s')) return token.slice(0, -1);
    if (token.endsWith('ed') && token.length > 4) return token.slice(0, -2); // Don't stem 3-letter words ending in ed
    if (token.endsWith('ing') && token.length > 5) return token.slice(0, -3); // Don't stem 4-letter words ending in ing
    return token;
  }

  private exactMatch(content: string, token: string): boolean {
    const stemmedToken = this.stemToken(token);
    return content.toLowerCase().includes(stemmedToken);
  }
}

export { VaultIndexer, Cache, InMemoryCache };

// Convenience method to create a VaultIndexer with an in-memory cache
export function createVaultIndexerWithCache(vault: VaultLike): VaultIndexer {
  return new VaultIndexer(vault, new InMemoryCache());
}