Initial open-source release

packages/context/src/wiki/sqlite-knowledge-index.ts

@@ -0,0 +1,276 @@
+import { mkdirSync } from 'node:fs';
+import { dirname } from 'node:path';
+import Database from 'better-sqlite3';
+import { buildKnowledgeSearchText } from './knowledge-search-text.js';
+import type { LocalKnowledgeScope } from './local-knowledge.js';
+
+export interface SqliteKnowledgeIndexOptions {
+  dbPath: string;
+}
+
+export interface SqliteKnowledgeIndexPage {
+  path: string;
+  key: string;
+  scope: LocalKnowledgeScope;
+  summary: string;
+  content: string;
+  tags: string[];
+  embedding?: number[] | null;
+}
+
+export interface SqliteKnowledgeIndexSearchResult {
+  path: string;
+  score: number;
+}
+
+export interface WikiSqliteLaneCandidate {
+  id: string;
+  path: string;
+  rank: number;
+  rawScore: number;
+}
+
+export interface ExistingKnowledgeIndexPage {
+  searchText: string;
+  embedding: number[] | null;
+}
+
+interface SearchRow {
+  path: string;
+  rank: number;
+}
+
+type IndexedPageRow = {
+  path: string;
+  embedding_json: string | null;
+};
+
+function cosineSimilarity(left: number[], right: number[]): number {
+  if (left.length === 0 || left.length !== right.length) {
+    return 0;
+  }
+  let dot = 0;
+  let leftNorm = 0;
+  let rightNorm = 0;
+  for (let i = 0; i < left.length; i++) {
+    const l = left[i] ?? 0;
+    const r = right[i] ?? 0;
+    dot += l * r;
+    leftNorm += l * l;
+    rightNorm += r * r;
+  }
+  if (leftNorm === 0 || rightNorm === 0) {
+    return 0;
+  }
+  return dot / (Math.sqrt(leftNorm) * Math.sqrt(rightNorm));
+}
+
+function scoreFromRank(rank: number): number {
+  return Number((1 / (1 + Math.abs(rank))).toFixed(6));
+}
+
+function parseEmbedding(raw: string | null): number[] | null {
+  if (!raw) {
+    return null;
+  }
+  try {
+    const embedding = JSON.parse(raw) as unknown;
+    return Array.isArray(embedding) && embedding.every((value) => typeof value === 'number') ? embedding : null;
+  } catch {
+    return null;
+  }
+}
+
+function normalizeFtsQuery(query: string): string {
+  const terms = query
+    .toLowerCase()
+    .split(/[^a-z0-9_]+/u)
+    .map((term) => term.trim())
+    .filter(Boolean);
+
+  return [...new Set(terms)].map((term) => `"${term.replaceAll('"', '""')}"`).join(' OR ');
+}
+
+export class SqliteKnowledgeIndex {
+  private readonly db: Database.Database;
+
+  constructor(options: SqliteKnowledgeIndexOptions) {
+    mkdirSync(dirname(options.dbPath), { recursive: true });
+    this.db = new Database(options.dbPath);
+    this.db.pragma('journal_mode = WAL');
+    this.db.pragma('foreign_keys = ON');
+    this.db.exec(`
+      CREATE TABLE IF NOT EXISTS knowledge_pages (
+        path TEXT PRIMARY KEY,
+        key TEXT NOT NULL,
+        scope TEXT NOT NULL,
+        summary TEXT NOT NULL,
+        content TEXT NOT NULL,
+        tags TEXT NOT NULL,
+        search_text TEXT NOT NULL,
+        embedding_json TEXT
+      );
+
+      CREATE VIRTUAL TABLE IF NOT EXISTS knowledge_pages_fts USING fts5(
+        path UNINDEXED,
+        key,
+        summary,
+        content,
+        tags
+      );
+    `);
+    const columns = this.db.prepare('PRAGMA table_info(knowledge_pages)').all() as Array<{ name: string }>;
+    const columnNames = new Set(columns.map((column) => column.name));
+    if (!columnNames.has('search_text')) {
+      this.db.exec("ALTER TABLE knowledge_pages ADD COLUMN search_text TEXT NOT NULL DEFAULT ''");
+    }
+    if (!columnNames.has('embedding_json')) {
+      this.db.exec('ALTER TABLE knowledge_pages ADD COLUMN embedding_json TEXT');
+    }
+  }
+
+  sync(pages: SqliteKnowledgeIndexPage[]): void {
+    const keepPaths = pages.map((page) => page.path);
+    const clearPages =
+      keepPaths.length === 0
+        ? this.db.prepare('DELETE FROM knowledge_pages')
+        : this.db.prepare(`DELETE FROM knowledge_pages WHERE path NOT IN (${keepPaths.map(() => '?').join(', ')})`);
+    const clearFts =
+      keepPaths.length === 0
+        ? this.db.prepare('DELETE FROM knowledge_pages_fts')
+        : this.db.prepare(`DELETE FROM knowledge_pages_fts WHERE path NOT IN (${keepPaths.map(() => '?').join(', ')})`);
+    const upsertPage = this.db.prepare(`
+      INSERT INTO knowledge_pages (path, key, scope, summary, content, tags, search_text, embedding_json)
+      VALUES (@path, @key, @scope, @summary, @content, @tags, @searchText, @embeddingJson)
+      ON CONFLICT(path) DO UPDATE SET
+        key = excluded.key,
+        scope = excluded.scope,
+        summary = excluded.summary,
+        content = excluded.content,
+        tags = excluded.tags,
+        search_text = excluded.search_text,
+        embedding_json = excluded.embedding_json
+    `);
+    const deleteFts = this.db.prepare('DELETE FROM knowledge_pages_fts WHERE path = @path');
+    const insertFts = this.db.prepare(`
+      INSERT INTO knowledge_pages_fts (path, key, summary, content, tags)
+      VALUES (@path, @key, @summary, @content, @tags)
+    `);
+
+    const transaction = this.db.transaction((items: SqliteKnowledgeIndexPage[]) => {
+      clearPages.run(...keepPaths);
+      clearFts.run(...keepPaths);
+      for (const page of items) {
+        const searchText = buildKnowledgeSearchText(page.key, page.summary, page.content, page.tags);
+        const row = {
+          path: page.path,
+          key: page.key,
+          scope: page.scope,
+          summary: page.summary,
+          content: searchText,
+          tags: page.tags.join(' '),
+          searchText,
+          embeddingJson: page.embedding ? JSON.stringify(page.embedding) : null,
+        };
+        upsertPage.run(row);
+        deleteFts.run(row);
+        insertFts.run(row);
+      }
+    });
+
+    transaction(pages);
+  }
+
+  rebuild(pages: SqliteKnowledgeIndexPage[]): void {
+    this.sync(pages);
+  }
+
+  getExistingPages(): Map<string, ExistingKnowledgeIndexPage> {
+    const rows = this.db
+      .prepare(
+        `
+        SELECT path, search_text, embedding_json
+        FROM knowledge_pages
+        ORDER BY path ASC
+      `,
+      )
+      .all() as Array<{ path: string; search_text: string; embedding_json: string | null }>;
+
+    return new Map(
+      rows.map((row) => [
+        row.path,
+        {
+          searchText: row.search_text,
+          embedding: parseEmbedding(row.embedding_json),
+        },
+      ]),
+    );
+  }
+
+  searchLexicalCandidates(input: { queryText: string; limit: number }): WikiSqliteLaneCandidate[] {
+    const ftsQuery = normalizeFtsQuery(input.queryText);
+    if (!ftsQuery) {
+      return [];
+    }
+
+    const rows = this.db
+      .prepare(
+        `
+        SELECT path, bm25(knowledge_pages_fts) AS rank
+        FROM knowledge_pages_fts
+        WHERE knowledge_pages_fts MATCH ?
+        ORDER BY rank ASC, path ASC
+        LIMIT ?
+      `,
+      )
+      .all(ftsQuery, Math.max(1, input.limit)) as SearchRow[];
+
+    return rows.map((row, index) => ({
+      id: row.path,
+      path: row.path,
+      rank: index + 1,
+      rawScore: Number(row.rank),
+    }));
+  }
+
+  searchSemanticCandidates(input: { queryEmbedding: number[]; limit: number }): WikiSqliteLaneCandidate[] {
+    const rows = this.db
+      .prepare(
+        `
+        SELECT path, embedding_json
+        FROM knowledge_pages
+        ORDER BY path ASC
+      `,
+      )
+      .all() as IndexedPageRow[];
+
+    return rows
+      .flatMap((row) => {
+        if (!row.embedding_json) {
+          return [];
+        }
+        const embedding = parseEmbedding(row.embedding_json);
+        if (!embedding) {
+          return [];
+        }
+        return [
+          {
+            id: row.path,
+            path: row.path,
+            rank: 0,
+            rawScore: cosineSimilarity(input.queryEmbedding, embedding),
+          },
+        ];
+      })
+      .sort((left, right) => right.rawScore - left.rawScore || left.path.localeCompare(right.path))
+      .slice(0, Math.max(1, input.limit))
+      .map((candidate, index) => ({ ...candidate, rank: index + 1 }));
+  }
+
+  search(query: string, limit: number): SqliteKnowledgeIndexSearchResult[] {
+    return this.searchLexicalCandidates({ queryText: query, limit }).map((row) => ({
+      path: row.path,
+      score: scoreFromRank(row.rawScore),
+    }));
+  }
+}