@disco_trooper/apple-notes-mcp 1.2.0 → 1.3.0

package/src/db/lancedb.test.ts

@@ -1,6 +1,6 @@
 import { describe, it, expect, beforeEach, afterEach } from "vitest";
-import { LanceDBStore } from "./lancedb.js";
-import type { NoteRecord } from "./lancedb.js";
+import { LanceDBStore, ChunkStore } from "./lancedb.js";
+import type { NoteRecord, ChunkRecord } from "./lancedb.js";
 import * as fs from "node:fs";
 import * as path from "node:path";
 
@@ -28,6 +28,9 @@ describe("LanceDBStore", () => {
     created: new Date().toISOString(),
     indexed_at: new Date().toISOString(),
     vector: Array(384).fill(0.1),
+    // LanceDB requires at least one element to infer the list type
+    tags: ["test-tag"],
+    outlinks: ["test-link"],
   });
 
   describe("index and getByTitle", () => {
@@ -184,3 +187,207 @@ describe("LanceDBStore", () => {
     });
   });
 });
+
+describe("ChunkStore", () => {
+  let chunkStore: ChunkStore;
+  let testDbPath: string;
+
+  beforeEach(() => {
+    testDbPath = path.join("/tmp", `lancedb-chunk-test-${Date.now()}`);
+    chunkStore = new ChunkStore(testDbPath);
+  });
+
+  afterEach(() => {
+    if (fs.existsSync(testDbPath)) {
+      fs.rmSync(testDbPath, { recursive: true, force: true });
+    }
+  });
+
+  const createTestChunk = (
+    noteId: string,
+    chunkIndex: number,
+    totalChunks: number,
+    content?: string
+  ): ChunkRecord => ({
+    chunk_id: `${noteId}_chunk_${chunkIndex}`,
+    note_id: noteId,
+    note_title: `Note ${noteId}`,
+    folder: "Test",
+    chunk_index: chunkIndex,
+    total_chunks: totalChunks,
+    content: content ?? `Chunk ${chunkIndex} content for note ${noteId}`,
+    vector: Array(384).fill(0.1),
+    created: new Date().toISOString(),
+    modified: new Date().toISOString(),
+    indexed_at: new Date().toISOString(),
+    tags: ["test-tag"],
+    outlinks: ["test-link"],
+  });
+
+  describe("indexChunks", () => {
+    it("indexes chunks and allows retrieval", async () => {
+      const chunks = [
+        createTestChunk("note-1", 0, 2),
+        createTestChunk("note-1", 1, 2),
+        createTestChunk("note-2", 0, 1),
+      ];
+
+      await chunkStore.indexChunks(chunks);
+      const count = await chunkStore.count();
+
+      expect(count).toBe(3);
+    });
+
+    it("handles empty chunks array", async () => {
+      await chunkStore.indexChunks([]);
+      const count = await chunkStore.count();
+      expect(count).toBe(0);
+    });
+
+    it("handles chunks with empty tags and outlinks", async () => {
+      const chunks = [
+        { ...createTestChunk("note-1", 0, 1), tags: [], outlinks: [] },
+      ];
+
+      await chunkStore.indexChunks(chunks);
+      const count = await chunkStore.count();
+      expect(count).toBe(1);
+    });
+  });
+
+  describe("searchChunks", () => {
+    it("returns results based on vector similarity", async () => {
+      const chunks = [
+        createTestChunk("note-1", 0, 2),
+        createTestChunk("note-1", 1, 2),
+      ];
+      await chunkStore.indexChunks(chunks);
+
+      const queryVector = Array(384).fill(0.1);
+      const results = await chunkStore.searchChunks(queryVector, 2);
+
+      expect(results).toHaveLength(2);
+      expect(results[0]).toHaveProperty("chunk_id");
+      expect(results[0]).toHaveProperty("note_id");
+      expect(results[0]).toHaveProperty("score");
+    });
+  });
+
+  describe("searchChunksFTS", () => {
+    it("returns results matching query text", async () => {
+      const chunks = [
+        createTestChunk("note-1", 0, 1, "Meeting notes about project planning"),
+        createTestChunk("note-2", 0, 1, "Shopping list for groceries"),
+      ];
+      await chunkStore.indexChunks(chunks);
+      await chunkStore.rebuildFtsIndex();
+
+      const results = await chunkStore.searchChunksFTS("Meeting", 10);
+      expect(results.length).toBeGreaterThanOrEqual(1);
+      expect(results[0].content).toContain("Meeting");
+    });
+
+    it("returns empty array for no matches", async () => {
+      const chunks = [createTestChunk("note-1", 0, 1)];
+      await chunkStore.indexChunks(chunks);
+      await chunkStore.rebuildFtsIndex();
+
+      const results = await chunkStore.searchChunksFTS("nonexistentquery12345", 10);
+      expect(results).toHaveLength(0);
+    });
+  });
+
+  describe("getChunksByNoteId", () => {
+    it("returns all chunks for a note sorted by chunk_index", async () => {
+      const chunks = [
+        createTestChunk("note-1", 2, 3),
+        createTestChunk("note-1", 0, 3),
+        createTestChunk("note-1", 1, 3),
+        createTestChunk("note-2", 0, 1),
+      ];
+      await chunkStore.indexChunks(chunks);
+
+      const noteChunks = await chunkStore.getChunksByNoteId("note-1");
+
+      expect(noteChunks).toHaveLength(3);
+      expect(noteChunks[0].chunk_index).toBe(0);
+      expect(noteChunks[1].chunk_index).toBe(1);
+      expect(noteChunks[2].chunk_index).toBe(2);
+    });
+
+    it("returns empty array for non-existent note", async () => {
+      await chunkStore.indexChunks([createTestChunk("note-1", 0, 1)]);
+
+      const chunks = await chunkStore.getChunksByNoteId("non-existent");
+      expect(chunks).toHaveLength(0);
+    });
+  });
+
+  describe("deleteNoteChunks", () => {
+    it("deletes all chunks for a note", async () => {
+      const chunks = [
+        createTestChunk("note-1", 0, 2),
+        createTestChunk("note-1", 1, 2),
+        createTestChunk("note-2", 0, 1),
+      ];
+      await chunkStore.indexChunks(chunks);
+
+      await chunkStore.deleteNoteChunks("note-1");
+
+      const remaining = await chunkStore.count();
+      expect(remaining).toBe(1);
+
+      const note1Chunks = await chunkStore.getChunksByNoteId("note-1");
+      expect(note1Chunks).toHaveLength(0);
+
+      const note2Chunks = await chunkStore.getChunksByNoteId("note-2");
+      expect(note2Chunks).toHaveLength(1);
+    });
+
+    it("does not throw when deleting non-existent note chunks", async () => {
+      await chunkStore.indexChunks([createTestChunk("note-1", 0, 1)]);
+      await expect(chunkStore.deleteNoteChunks("non-existent")).resolves.not.toThrow();
+    });
+  });
+
+  describe("count", () => {
+    it("returns correct count", async () => {
+      const chunks = [
+        createTestChunk("note-1", 0, 2),
+        createTestChunk("note-1", 1, 2),
+        createTestChunk("note-2", 0, 1),
+      ];
+      await chunkStore.indexChunks(chunks);
+
+      expect(await chunkStore.count()).toBe(3);
+    });
+
+    it("returns 0 for empty store", async () => {
+      expect(await chunkStore.count()).toBe(0);
+    });
+  });
+
+  describe("clear", () => {
+    it("removes all chunks", async () => {
+      await chunkStore.indexChunks([
+        createTestChunk("note-1", 0, 1),
+        createTestChunk("note-2", 0, 1),
+      ]);
+      expect(await chunkStore.count()).toBe(2);
+
+      await chunkStore.clear();
+      expect(await chunkStore.count()).toBe(0);
+    });
+  });
+
+  describe("rebuildFtsIndex", () => {
+    it("rebuilds FTS index without error", async () => {
+      await chunkStore.indexChunks([
+        createTestChunk("note-1", 0, 1),
+        createTestChunk("note-2", 0, 1),
+      ]);
+
+      await expect(chunkStore.rebuildFtsIndex()).resolves.not.toThrow();
+    });
+  });
+});

package/src/db/lancedb.ts

@@ -1,9 +1,8 @@
 import * as lancedb from "@lancedb/lancedb";
-import path from "node:path";
-import os from "node:os";
 import { validateTitle, escapeForFilter } from "./validation.js";
 import type { DBSearchResult as SearchResult } from "../types/index.js";
 import { createDebugLogger } from "../utils/debug.js";
+import { getDataDir } from "../config/paths.js";
 
 // Schema for stored notes
 export interface NoteRecord {
@@ -15,9 +14,30 @@ export interface NoteRecord {
   created: string;      // ISO date
   modified: string;     // ISO date
   indexed_at: string;   // ISO date - when embedding was generated
+  // Knowledge Graph fields
+  tags: string[];       // Extracted #hashtags (without #)
+  outlinks: string[];   // Extracted [[wiki-links]] titles
   [key: string]: unknown; // Index signature for LanceDB compatibility
 }
 
+// Schema for chunked notes (Parent Document Retriever pattern)
+export interface ChunkRecord {
+  chunk_id: string;      // `${note_id}_chunk_${index}`
+  note_id: string;       // Parent note Apple ID
+  note_title: string;    // For display and deduplication
+  folder: string;
+  chunk_index: number;   // 0, 1, 2...
+  total_chunks: number;  // Total chunks in this note
+  content: string;       // Chunk content
+  vector: number[];
+  created: string;       // ISO date (from parent)
+  modified: string;      // ISO date (from parent)
+  indexed_at: string;    // ISO date
+  tags: string[];        // From parent note
+  outlinks: string[];    // From parent note
+  [key: string]: unknown; // Index signature for LanceDB
+}
+
 // SearchResult is imported from ../types/index.js as DBSearchResult
 export type { SearchResult };
 
@@ -61,7 +81,7 @@ export class LanceDBStore implements VectorStore {
   private readonly tableName = "notes";
 
   constructor(dataDir?: string) {
-    this.dbPath = dataDir || path.join(os.homedir(), ".apple-notes-mcp", "data");
+    this.dbPath = dataDir || getDataDir();
   }
 
   private async ensureConnection(): Promise<lancedb.Connection> {
@@ -103,9 +123,71 @@ export class LanceDBStore implements VectorStore {
       debug("Table drop skipped (table may not exist):", error);
     }
 
+    // Arrow type inference requires the FIRST record to have non-empty arrays.
+    // Strategy: Reorder records so the first has non-empty tags/outlinks,
+    // or add placeholders that stay in the data (filtered on read).
+    const processedRecords = records.map((r) => ({
+      ...r,
+      tags: r.tags ?? [],
+      outlinks: r.outlinks ?? [],
+    }));
+
+    // Track if we added placeholders
+    let addedTagPlaceholder = false;
+    let addedOutlinkPlaceholder = false;
+
+    if (processedRecords.length > 0) {
+      // Ensure FIRST record has non-empty tags for type inference
+      if (processedRecords[0].tags.length === 0) {
+        // Try to find a record with tags and swap
+        const tagIdx = processedRecords.findIndex(r => r.tags.length > 0);
+        if (tagIdx > 0) {
+          // Swap first record with the one that has tags
+          [processedRecords[0], processedRecords[tagIdx]] =
+            [processedRecords[tagIdx], processedRecords[0]];
+        } else {
+          // No record has tags - add placeholder
+          processedRecords[0].tags = ["__type_placeholder__"];
+          addedTagPlaceholder = true;
+        }
+      }
+
+      // Ensure FIRST record has non-empty outlinks for type inference
+      if (processedRecords[0].outlinks.length === 0) {
+        // Try to find a record with outlinks and copy its structure
+        const outlinkIdx = processedRecords.findIndex(r => r.outlinks.length > 0);
+        if (outlinkIdx === -1) {
+          // No record has outlinks - add placeholder
+          processedRecords[0].outlinks = ["__type_placeholder__"];
+          addedOutlinkPlaceholder = true;
+        } else {
+          // Copy first outlink to first record temporarily, then remove
+          processedRecords[0].outlinks = ["__type_placeholder__"];
+          addedOutlinkPlaceholder = true;
+        }
+      }
+    }
+
+    debug(`Creating table with ${processedRecords.length} records (tag placeholder: ${addedTagPlaceholder}, outlink placeholder: ${addedOutlinkPlaceholder})`);
+
     // Create new table with records
-    debug(`Creating table with ${records.length} records`);
-    this.table = await db.createTable(this.tableName, records);
+    this.table = await db.createTable(this.tableName, processedRecords);
+
+    // Remove placeholders by deleting and re-inserting the first record
+    if (addedTagPlaceholder || addedOutlinkPlaceholder) {
+      const firstRecord = processedRecords[0];
+      const cleanRecord = {
+        ...firstRecord,
+        tags: addedTagPlaceholder ? [] : firstRecord.tags,
+        outlinks: addedOutlinkPlaceholder ? [] : firstRecord.outlinks,
+      };
+
+      // Delete the record with placeholders
+      await this.table.delete(`id = '${escapeForFilter(firstRecord.id)}'`);
+      // Re-insert without placeholders
+      await this.table.add([cleanRecord]);
+      debug("Removed type inference placeholders via delete+insert");
+    }
 
     // Create FTS index for hybrid search
     debug("Creating FTS index on content");
@@ -217,15 +299,18 @@ export class LanceDBStore implements VectorStore {
 
     const results = await table.query().toArray();
 
-    return results.map((row) => ({
+    return results.map((row): NoteRecord => ({
       id: (row.id as string) ?? "",
       title: row.title as string,
       content: row.content as string,
-      vector: row.vector as number[],
+      vector: Array.isArray(row.vector) ? row.vector : Array.from(row.vector as Iterable<number>),
       folder: row.folder as string,
       created: row.created as string,
       modified: row.modified as string,
       indexed_at: row.indexed_at as string,
+      // Arrow Vectors need explicit conversion to JS arrays
+      tags: Array.isArray(row.tags) ? row.tags : Array.from(row.tags as Iterable<string>),
+      outlinks: Array.isArray(row.outlinks) ? row.outlinks : Array.from(row.outlinks as Iterable<string>),
     }));
   }
 
@@ -270,3 +355,256 @@ export function getVectorStore(): VectorStore {
   }
   return storeInstance;
 }
+
+// Search result type for chunks
+export interface ChunkSearchResult {
+  chunk_id: string;
+  note_id: string;
+  note_title: string;
+  folder: string;
+  chunk_index: number;
+  total_chunks: number;
+  content: string;
+  modified: string;
+  score: number;
+}
+
+/**
+ * Convert a chunk database row to a ChunkSearchResult with rank-based score.
+ */
+function rowToChunkSearchResult(row: Record<string, unknown>, index: number): ChunkSearchResult {
+  return {
+    chunk_id: row.chunk_id as string,
+    note_id: row.note_id as string,
+    note_title: row.note_title as string,
+    folder: row.folder as string,
+    chunk_index: row.chunk_index as number,
+    total_chunks: row.total_chunks as number,
+    content: row.content as string,
+    modified: row.modified as string,
+    score: 1 / (1 + index),
+  };
+}
+
+// ChunkStore for Parent Document Retriever pattern
+export class ChunkStore {
+  private db: lancedb.Connection | null = null;
+  private table: lancedb.Table | null = null;
+  private readonly dbPath: string;
+  private readonly tableName = "chunks";
+
+  constructor(dataDir?: string) {
+    this.dbPath = dataDir || getDataDir();
+  }
+
+  private async ensureConnection(): Promise<lancedb.Connection> {
+    if (!this.db) {
+      debug(`ChunkStore: Connecting to LanceDB at ${this.dbPath}`);
+      this.db = await lancedb.connect(this.dbPath);
+    }
+    return this.db;
+  }
+
+  private async ensureTable(): Promise<lancedb.Table> {
+    if (!this.table) {
+      const db = await this.ensureConnection();
+      try {
+        this.table = await db.openTable(this.tableName);
+        debug(`ChunkStore: Opened existing table: ${this.tableName}`);
+      } catch (error) {
+        debug(`ChunkStore: Table ${this.tableName} not found. Error:`, error);
+        throw new Error("Chunk index not found. Run index-notes first.");
+      }
+    }
+    return this.table;
+  }
+
+  async indexChunks(chunks: ChunkRecord[]): Promise<void> {
+    if (chunks.length === 0) {
+      debug("ChunkStore: No chunks to index");
+      return;
+    }
+
+    const db = await this.ensureConnection();
+
+    // Drop existing table if exists
+    try {
+      await db.dropTable(this.tableName);
+      debug(`ChunkStore: Dropped existing table: ${this.tableName}`);
+    } catch (error) {
+      debug("ChunkStore: Table drop skipped (table may not exist):", error);
+    }
+
+    // Arrow type inference requires the FIRST record to have non-empty arrays.
+    // Same strategy as LanceDBStore
+    const processedChunks = chunks.map((c) => ({
+      ...c,
+      tags: c.tags ?? [],
+      outlinks: c.outlinks ?? [],
+    }));
+
+    let addedTagPlaceholder = false;
+    let addedOutlinkPlaceholder = false;
+
+    if (processedChunks.length > 0) {
+      // Ensure FIRST chunk has non-empty tags for type inference
+      if (processedChunks[0].tags.length === 0) {
+        const tagIdx = processedChunks.findIndex(c => c.tags.length > 0);
+        if (tagIdx > 0) {
+          [processedChunks[0], processedChunks[tagIdx]] =
+            [processedChunks[tagIdx], processedChunks[0]];
+        } else {
+          processedChunks[0].tags = ["__type_placeholder__"];
+          addedTagPlaceholder = true;
+        }
+      }
+
+      // Ensure FIRST chunk has non-empty outlinks for type inference
+      if (processedChunks[0].outlinks.length === 0) {
+        const outlinkIdx = processedChunks.findIndex(c => c.outlinks.length > 0);
+        if (outlinkIdx === -1) {
+          processedChunks[0].outlinks = ["__type_placeholder__"];
+          addedOutlinkPlaceholder = true;
+        } else {
+          processedChunks[0].outlinks = ["__type_placeholder__"];
+          addedOutlinkPlaceholder = true;
+        }
+      }
+    }
+
+    debug(`ChunkStore: Creating table with ${processedChunks.length} chunks (tag placeholder: ${addedTagPlaceholder}, outlink placeholder: ${addedOutlinkPlaceholder})`);
+
+    // Create new table with chunks
+    this.table = await db.createTable(this.tableName, processedChunks);
+
+    // Remove placeholders by deleting and re-inserting the first chunk
+    if (addedTagPlaceholder || addedOutlinkPlaceholder) {
+      const firstChunk = processedChunks[0];
+      const cleanChunk = {
+        ...firstChunk,
+        tags: addedTagPlaceholder ? [] : firstChunk.tags,
+        outlinks: addedOutlinkPlaceholder ? [] : firstChunk.outlinks,
+      };
+
+      // Delete the chunk with placeholders
+      await this.table.delete(`chunk_id = '${escapeForFilter(firstChunk.chunk_id)}'`);
+      // Re-insert without placeholders
+      await this.table.add([cleanChunk]);
+      debug("ChunkStore: Removed type inference placeholders via delete+insert");
+    }
+
+    // Create FTS index for hybrid search
+    debug("ChunkStore: Creating FTS index on content");
+    await this.table.createIndex("content", {
+      config: lancedb.Index.fts(),
+      replace: true,
+    });
+
+    debug(`ChunkStore: Indexed ${chunks.length} chunks`);
+  }
+
+  async searchChunks(queryVector: number[], limit: number): Promise<ChunkSearchResult[]> {
+    const table = await this.ensureTable();
+
+    const results = await table
+      .search(queryVector)
+      .limit(limit)
+      .toArray();
+
+    return results.map(rowToChunkSearchResult);
+  }
+
+  async searchChunksFTS(query: string, limit: number): Promise<ChunkSearchResult[]> {
+    const table = await this.ensureTable();
+
+    try {
+      const results = await table
+        .query()
+        .fullTextSearch(query)
+        .limit(limit)
+        .toArray();
+
+      return results.map(rowToChunkSearchResult);
+    } catch (error) {
+      debug("ChunkStore: FTS search failed, returning empty results. Error:", error);
+      return [];
+    }
+  }
+
+  async getChunksByNoteId(noteId: string): Promise<ChunkRecord[]> {
+    const table = await this.ensureTable();
+    const escapedNoteId = escapeForFilter(noteId);
+
+    const results = await table
+      .query()
+      .where(`note_id = '${escapedNoteId}'`)
+      .toArray();
+
+    // Convert and sort by chunk_index
+    const chunks = results.map((row): ChunkRecord => ({
+      chunk_id: row.chunk_id as string,
+      note_id: row.note_id as string,
+      note_title: row.note_title as string,
+      folder: row.folder as string,
+      chunk_index: row.chunk_index as number,
+      total_chunks: row.total_chunks as number,
+      content: row.content as string,
+      vector: Array.isArray(row.vector) ? row.vector : Array.from(row.vector as Iterable<number>),
+      created: row.created as string,
+      modified: row.modified as string,
+      indexed_at: row.indexed_at as string,
+      tags: Array.isArray(row.tags) ? row.tags : Array.from(row.tags as Iterable<string>),
+      outlinks: Array.isArray(row.outlinks) ? row.outlinks : Array.from(row.outlinks as Iterable<string>),
+    }));
+
+    return chunks.sort((a, b) => a.chunk_index - b.chunk_index);
+  }
+
+  async deleteNoteChunks(noteId: string): Promise<void> {
+    const table = await this.ensureTable();
+    const escapedNoteId = escapeForFilter(noteId);
+    await table.delete(`note_id = '${escapedNoteId}'`);
+    debug(`ChunkStore: Deleted chunks for note: ${noteId}`);
+  }
+
+  async count(): Promise<number> {
+    try {
+      const table = await this.ensureTable();
+      return await table.countRows();
+    } catch (error) {
+      debug("ChunkStore: Count failed (table may not exist):", error);
+      return 0;
+    }
+  }
+
+  async clear(): Promise<void> {
+    const db = await this.ensureConnection();
+    try {
+      await db.dropTable(this.tableName);
+      this.table = null;
+      debug("ChunkStore: Cleared table");
+    } catch (error) {
+      debug("ChunkStore: Clear skipped (table may not exist):", error);
+    }
+  }
+
+  async rebuildFtsIndex(): Promise<void> {
+    const table = await this.ensureTable();
+    debug("ChunkStore: Rebuilding FTS index on content");
+    await table.createIndex("content", {
+      config: lancedb.Index.fts(),
+      replace: true,
+    });
+    debug("ChunkStore: FTS index rebuilt");
+  }
+}
+
+// Singleton instance for ChunkStore
+let chunkStoreInstance: ChunkStore | null = null;
+
+export function getChunkStore(): ChunkStore {
+  if (!chunkStoreInstance) {
+    chunkStoreInstance = new ChunkStore();
+  }
+  return chunkStoreInstance;
+}