fix: expand chunk IDs in remove() when chunking is enabled

src/db/sqlite-fts.ts

@@ -207,6 +207,11 @@ export async function sqliteFts(config: SqliteFtsConfig = {}): Promise<SearchPro
       }
     },
 
+    async listIds() {
+      const rows = db.prepare('SELECT id FROM documents_meta').all() as { id: string }[]
+      return rows.map(r => r.id)
+    },
+
     async clear() {
       db.exec('DELETE FROM documents_fts')
       db.exec('DELETE FROM documents_meta')

src/retriv.ts

@@ -220,9 +220,25 @@ export async function createRetriv(options: RetrivOptions): Promise<SearchProvid
     },
 
     async remove(ids: string[]) {
+      let removeIds = ids
+      if (chunker) {
+        const lister = drivers.find(d => d.listIds)
+        if (lister) {
+          const allIds = await lister.listIds!()
+          const idSet = new Set(ids)
+          removeIds = allIds.filter((id) => {
+            if (idSet.has(id))
+              return true
+            const sep = id.indexOf('#chunk-')
+            return sep >= 0 && idSet.has(id.substring(0, sep))
+          })
+        }
+      }
       const results = await Promise.all(
-        drivers.filter(d => d.remove).map(d => d.remove!(ids)),
+        drivers.filter(d => d.remove).map(d => d.remove!(removeIds)),
       )
+      for (const id of ids)
+        parentDocs.delete(id)
       return { count: results[0]?.count ?? 0 }
     },
 

test/retriv.test.ts

@@ -130,6 +130,26 @@ describe('createRetriv', () => {
     await retriv.close?.()
   })
 
+  it('removes chunked documents by parent id', async () => {
+    const retriv = await createRetriv({
+      driver: sqliteFts({ path: ':memory:' }),
+      chunking: markdownChunker({ chunkSize: 20, chunkOverlap: 0 }),
+    })
+
+    await retriv.index([
+      { id: 'doc1', content: 'First part.\n\nSecond part.\n\nThird part.' },
+      { id: 'doc2', content: 'Keep this.\n\nStill here.\n\nNot removed.' },
+    ])
+
+    await retriv.remove?.(['doc1'])
+
+    const doc1Results = await retriv.search('part', { limit: 10 })
+    const doc2Results = await retriv.search('keep', { limit: 10 })
+
+    expect(doc1Results.every(r => !r.id.startsWith('doc1'))).toBe(true)
+    expect(doc2Results.length).toBeGreaterThanOrEqual(1)
+  })
+
   it('extracts snippets with highlights', async () => {
     const retriv = await createRetriv({
       driver: sqliteFts({ path: ':memory:' }),