feat(reedy): Phase 1A — MVP retrieval primitives (#4293)
* feat(db): add reedy schema migration with Tantivy FTS + lazy embeddings Registers a new `reedy` migration set bound to reedy.db. Creates reedy_book_meta + reedy_book_chunks with a Tantivy FTS index on chunks.text (ngram tokenizer) and a per-book position index used by BookRetriever. The vector embeddings table is intentionally NOT created here — the indexer creates it lazily on first index so the vector32(<dim>) column matches the active embedding model. Tests cover the migration applies cleanly, is idempotent, and that the FTS index is queryable. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * feat(reedy): retrieval primitives — DB, chunker, indexer, retriever, lookupPassage tool Wires the MVP retrieval pipeline behind reedy.db, all under src/services/reedy/ and with no integration into the existing AI module yet (Phase 1B will do that). - ReedyDb wrapper over DatabaseService: book-meta CRUD, lazy embeddings table at the active model's dim, bulk chunk + embedding writes via batch(), hybridSearch (brute-force cosine + Tantivy FTS + reciprocal-rank fusion with 3× per-path over-fetch), per-book and global wipe. Internal write queue serializes batch() calls so Turso's single-writer transaction guard doesn't trip when BookIndexer runs across books in parallel. - CfiChunker: TreeWalker over the section's DOM, ~maxChunkSize windows with paragraph > sentence > word break-points, full epubcfi(/6/N!/…) anchors, round-trip verified via CFI.toRange before each chunk lands. - BookIndexer: per-book mutex, lazy embeddings-table creation, model.batchSize embedding batches with dim assertion, terminal status transitions (indexed | empty_index | failed). Re-indexing clears prior chunks via a new ReedyDb.clearBookChunks helper. - BookRetriever: status-typed results (ok | not_indexed | empty_index | stale_index | degraded). Embedding has a 5s wall-clock budget; on timeout it falls through to FTS-only with status=degraded. - lookupPassage Vercel ai-SDK tool: Zod-validated query/topK, per-turn composite-key dedupe, parallel-call serialization, 10s per-turn budget, 6000-char result clamp, status-with-hint passthrough, and a separate serializeForModel that wraps each passage in <retrieved trust="untrusted"> with XML-escaped content so book text cannot escape the envelope. 59 unit tests; pnpm test + pnpm lint clean. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Huang Xin
@@ -0,0 +1,104 @@
|
||||
import { describe, it, expect, beforeEach, afterEach } from 'vitest';
|
||||
import { NodeDatabaseService } from '@/services/database/nodeDatabaseService';
|
||||
import { DatabaseService } from '@/types/database';
|
||||
import { migrate } from '@/services/database/migrate';
|
||||
import { getMigrations } from '@/services/database/migrations';
|
||||
|
||||
/**
|
||||
* Verify the Reedy schema migration applies cleanly against a real Turso
|
||||
* (in-memory) SQLite database with the same `experimental: ['index_method']`
|
||||
* opt that production opens reedy.db with.
|
||||
*
|
||||
* Per plan §M1.1 the embeddings table is created lazily by BookIndexer on
|
||||
* first index (so its dim can match the active embedding model). The
|
||||
* migration must NOT create it; this test guards that contract.
|
||||
*/
|
||||
describe('Reedy migration', () => {
|
||||
let db: DatabaseService;
|
||||
|
||||
beforeEach(async () => {
|
||||
db = await NodeDatabaseService.open(':memory:', { experimental: ['index_method'] });
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
await db.close();
|
||||
});
|
||||
|
||||
async function listObjects(type: 'table' | 'index'): Promise<string[]> {
|
||||
const rows = await db.select<{ name: string }>(
|
||||
`SELECT name FROM sqlite_master WHERE type = '${type}' AND name NOT LIKE 'sqlite_%'`,
|
||||
);
|
||||
return rows.map((r) => r.name);
|
||||
}
|
||||
|
||||
it('registers a non-empty migration set under the "reedy" schema', () => {
|
||||
const reedyMigrations = getMigrations('reedy');
|
||||
expect(reedyMigrations.length).toBeGreaterThan(0);
|
||||
for (const m of reedyMigrations) {
|
||||
expect(m.name).toMatch(/^\d{10}_/);
|
||||
expect(typeof m.sql).toBe('string');
|
||||
expect(m.sql.length).toBeGreaterThan(0);
|
||||
}
|
||||
});
|
||||
|
||||
it('creates reedy_book_meta and reedy_book_chunks tables', async () => {
|
||||
await migrate(db, getMigrations('reedy'));
|
||||
const tables = await listObjects('table');
|
||||
expect(tables).toContain('reedy_book_meta');
|
||||
expect(tables).toContain('reedy_book_chunks');
|
||||
});
|
||||
|
||||
it('does NOT create the embeddings table at migration time (lazy)', async () => {
|
||||
await migrate(db, getMigrations('reedy'));
|
||||
const tables = await listObjects('table');
|
||||
expect(tables).not.toContain('reedy_book_chunk_embeddings');
|
||||
});
|
||||
|
||||
it('creates idx_chunks_book_position index on (book_hash, position_index)', async () => {
|
||||
await migrate(db, getMigrations('reedy'));
|
||||
const indexes = await listObjects('index');
|
||||
expect(indexes).toContain('idx_chunks_book_position');
|
||||
});
|
||||
|
||||
it('creates an FTS index over reedy_book_chunks.text that is queryable', async () => {
|
||||
await migrate(db, getMigrations('reedy'));
|
||||
|
||||
await db.execute(
|
||||
`INSERT INTO reedy_book_chunks
|
||||
(id, book_hash, section_index, chapter_title, start_cfi, end_cfi, position_index, text, token_count)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
['c1', 'bk1', 0, 'Ch1', '/6/4!/4/2,/1:0,/1:10', '/6/4!/4/2,/1:10,/1:20', 0, 'alpha bravo', 2],
|
||||
);
|
||||
await db.execute(
|
||||
`INSERT INTO reedy_book_chunks
|
||||
(id, book_hash, section_index, chapter_title, start_cfi, end_cfi, position_index, text, token_count)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
[
|
||||
'c2',
|
||||
'bk1',
|
||||
0,
|
||||
'Ch1',
|
||||
'/6/4!/4/4,/1:0,/1:10',
|
||||
'/6/4!/4/4,/1:10,/1:20',
|
||||
1,
|
||||
'charlie delta',
|
||||
2,
|
||||
],
|
||||
);
|
||||
|
||||
const matches = await db.select<{ id: string; text: string }>(
|
||||
"SELECT id, text FROM reedy_book_chunks WHERE fts_match(text, 'alpha')",
|
||||
);
|
||||
expect(matches).toHaveLength(1);
|
||||
expect(matches[0]!.id).toBe('c1');
|
||||
});
|
||||
|
||||
it('is idempotent — running twice does not error and PRAGMA user_version stays at target', async () => {
|
||||
const reedyMigrations = getMigrations('reedy');
|
||||
await migrate(db, reedyMigrations);
|
||||
await migrate(db, reedyMigrations);
|
||||
|
||||
const version = await db.select<{ user_version: number }>('PRAGMA user_version');
|
||||
expect(version[0]!.user_version).toBe(reedyMigrations.length);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,227 @@
|
||||
/**
|
||||
* @vitest-environment jsdom
|
||||
*/
|
||||
import { describe, it, expect, beforeEach, afterEach } from 'vitest';
|
||||
import { NodeDatabaseService } from '@/services/database/nodeDatabaseService';
|
||||
import { DatabaseService } from '@/types/database';
|
||||
import { migrate } from '@/services/database/migrate';
|
||||
import { getMigrations } from '@/services/database/migrations';
|
||||
import { ReedyDb } from '@/services/reedy/db/ReedyDb';
|
||||
import { BookIndexer } from '@/services/reedy/retrieval/BookIndexer';
|
||||
import type { EmbeddingModel } from '@/services/reedy/models/EmbeddingModel';
|
||||
import type { BookDoc, SectionItem } from '@/libs/document';
|
||||
|
||||
const DIM = 4;
|
||||
|
||||
function fakeModel(overrides: Partial<EmbeddingModel> = {}): EmbeddingModel {
|
||||
return {
|
||||
id: 'fake-model',
|
||||
dim: DIM,
|
||||
batchSize: 2,
|
||||
async embed(texts) {
|
||||
// Deterministic embedding: char-code sums over four buckets, normalized.
|
||||
return texts.map((t) => {
|
||||
const v = [0, 0, 0, 0];
|
||||
for (let i = 0; i < t.length; i++) {
|
||||
v[i % 4]! += t.charCodeAt(i);
|
||||
}
|
||||
const norm = Math.sqrt(v.reduce((s, x) => s + x * x, 0)) || 1;
|
||||
return v.map((x) => x / norm);
|
||||
});
|
||||
},
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
function section(id: string, html: string): SectionItem {
|
||||
return {
|
||||
id,
|
||||
cfi: '',
|
||||
size: html.length,
|
||||
linear: 'yes',
|
||||
async createDocument() {
|
||||
return new DOMParser().parseFromString(
|
||||
`<!DOCTYPE html><html><body>${html}</body></html>`,
|
||||
'text/html',
|
||||
);
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function fakeBook(sections: SectionItem[]): BookDoc {
|
||||
return {
|
||||
metadata: { title: 'T', author: 'A', language: 'en' },
|
||||
rendition: {},
|
||||
dir: 'ltr',
|
||||
sections,
|
||||
splitTOCHref: () => [],
|
||||
async getCover() {
|
||||
return null;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
describe('BookIndexer', () => {
|
||||
let svc: DatabaseService;
|
||||
let reedy: ReedyDb;
|
||||
let indexer: BookIndexer;
|
||||
|
||||
beforeEach(async () => {
|
||||
svc = await NodeDatabaseService.open(':memory:', { experimental: ['index_method'] });
|
||||
await migrate(svc, getMigrations('reedy'));
|
||||
reedy = new ReedyDb(svc);
|
||||
indexer = new BookIndexer(reedy);
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
await svc.close();
|
||||
});
|
||||
|
||||
it('happy path: chunks all sections, writes embeddings, lands status=indexed', async () => {
|
||||
const book = fakeBook([
|
||||
section('s0', '<p>Alpha bravo charlie delta echo.</p>'),
|
||||
section('s1', '<p>Foxtrot golf hotel india juliet.</p>'),
|
||||
]);
|
||||
const model = fakeModel();
|
||||
|
||||
await indexer.indexBook(book, 'bk-happy', model);
|
||||
|
||||
const meta = await reedy.getBookMeta('bk-happy');
|
||||
expect(meta?.indexingStatus).toBe('indexed');
|
||||
expect(meta?.chunkCount).toBeGreaterThan(0);
|
||||
expect(meta?.embeddingModel).toBe('fake-model');
|
||||
expect(meta?.embeddingDim).toBe(DIM);
|
||||
expect(meta?.indexedAt).toBeGreaterThan(0);
|
||||
|
||||
const chunkRows = await svc.select<{ c: number }>(
|
||||
"SELECT COUNT(*) as c FROM reedy_book_chunks WHERE book_hash = 'bk-happy'",
|
||||
);
|
||||
expect(chunkRows[0]!.c).toBeGreaterThan(0);
|
||||
|
||||
const embRows = await svc.select<{ c: number }>(
|
||||
"SELECT COUNT(*) as c FROM reedy_book_chunk_embeddings WHERE book_hash = 'bk-happy'",
|
||||
);
|
||||
expect(embRows[0]!.c).toBe(chunkRows[0]!.c);
|
||||
});
|
||||
|
||||
it('image-only book lands status=empty_index with chunk_count=0', async () => {
|
||||
const book = fakeBook([section('s0', '<img src="cover.png" alt=""/>')]);
|
||||
await indexer.indexBook(book, 'bk-empty', fakeModel());
|
||||
|
||||
const meta = await reedy.getBookMeta('bk-empty');
|
||||
expect(meta?.indexingStatus).toBe('empty_index');
|
||||
expect(meta?.chunkCount).toBe(0);
|
||||
});
|
||||
|
||||
it('respects model.batchSize when calling embed (multiple smaller batches)', async () => {
|
||||
const calls: number[] = [];
|
||||
const model: EmbeddingModel = {
|
||||
id: 'batch-model',
|
||||
dim: DIM,
|
||||
batchSize: 2,
|
||||
async embed(texts) {
|
||||
calls.push(texts.length);
|
||||
return texts.map(() => [1, 0, 0, 0]);
|
||||
},
|
||||
};
|
||||
// 5 small paragraphs → CfiChunker emits 5 chunks (each well below maxChunkSize)
|
||||
const html = Array.from({ length: 5 }, (_, i) => `<p>Para ${i} text here.</p>`).join('');
|
||||
const book = fakeBook([section('s0', html)]);
|
||||
|
||||
await indexer.indexBook(book, 'bk-batch', model, {
|
||||
chunkOptions: { maxChunkSize: 30, minChunkSize: 5, overlapSize: 0, breakSearchRange: 5 },
|
||||
});
|
||||
|
||||
expect(calls.length).toBeGreaterThan(1);
|
||||
for (const n of calls) expect(n).toBeLessThanOrEqual(2);
|
||||
});
|
||||
|
||||
it('on embed failure lands status=failed with the error message and does not throw past indexBook', async () => {
|
||||
const model: EmbeddingModel = {
|
||||
id: 'broken-model',
|
||||
dim: DIM,
|
||||
async embed() {
|
||||
throw new Error('embedding gateway down');
|
||||
},
|
||||
};
|
||||
const book = fakeBook([section('s0', '<p>Some real text content.</p>')]);
|
||||
|
||||
await expect(indexer.indexBook(book, 'bk-fail', model)).rejects.toThrow(/gateway down/);
|
||||
|
||||
const meta = await reedy.getBookMeta('bk-fail');
|
||||
expect(meta?.indexingStatus).toBe('failed');
|
||||
expect(meta?.error).toContain('gateway down');
|
||||
});
|
||||
|
||||
it('rejects an embedding model whose returned vector length differs from model.dim', async () => {
|
||||
const wrongDimModel: EmbeddingModel = {
|
||||
id: 'wrong-dim',
|
||||
dim: DIM,
|
||||
async embed(texts) {
|
||||
// Claims dim=DIM but actually returns dim=DIM+1
|
||||
return texts.map(() => Array.from({ length: DIM + 1 }, () => 0));
|
||||
},
|
||||
};
|
||||
const book = fakeBook([section('s0', '<p>Real chunkable content text body.</p>')]);
|
||||
|
||||
await expect(indexer.indexBook(book, 'bk-bad-dim', wrongDimModel)).rejects.toThrow(/dim/);
|
||||
const meta = await reedy.getBookMeta('bk-bad-dim');
|
||||
expect(meta?.indexingStatus).toBe('failed');
|
||||
});
|
||||
|
||||
it('serializes concurrent indexBook calls for the same book (mutex)', async () => {
|
||||
let active = 0;
|
||||
let maxActive = 0;
|
||||
const model: EmbeddingModel = {
|
||||
id: 'mtx',
|
||||
dim: DIM,
|
||||
async embed(texts) {
|
||||
active++;
|
||||
maxActive = Math.max(maxActive, active);
|
||||
await new Promise((r) => setTimeout(r, 5));
|
||||
active--;
|
||||
return texts.map(() => [1, 0, 0, 0]);
|
||||
},
|
||||
};
|
||||
const make = () =>
|
||||
fakeBook([section('s0', '<p>Content text for serialization test goes here.</p>')]);
|
||||
|
||||
await Promise.all([
|
||||
indexer.indexBook(make(), 'bk-mtx', model),
|
||||
indexer.indexBook(make(), 'bk-mtx', model),
|
||||
indexer.indexBook(make(), 'bk-mtx', model),
|
||||
]);
|
||||
|
||||
// With the mutex, only one indexBook is ever embedding for bk-mtx at a time.
|
||||
expect(maxActive).toBe(1);
|
||||
|
||||
// After all three resolve the row is in a terminal state (indexed/empty).
|
||||
const meta = await reedy.getBookMeta('bk-mtx');
|
||||
expect(meta?.indexingStatus).toBe('indexed');
|
||||
});
|
||||
|
||||
it('does NOT serialize indexBook calls across different books', async () => {
|
||||
let active = 0;
|
||||
let maxActive = 0;
|
||||
const model: EmbeddingModel = {
|
||||
id: 'cross',
|
||||
dim: DIM,
|
||||
async embed(texts) {
|
||||
active++;
|
||||
maxActive = Math.max(maxActive, active);
|
||||
await new Promise((r) => setTimeout(r, 10));
|
||||
active--;
|
||||
return texts.map(() => [1, 0, 0, 0]);
|
||||
},
|
||||
};
|
||||
const make = () => fakeBook([section('s0', '<p>Cross-book parallelism check.</p>')]);
|
||||
|
||||
await Promise.all([
|
||||
indexer.indexBook(make(), 'bookA', model),
|
||||
indexer.indexBook(make(), 'bookB', model),
|
||||
indexer.indexBook(make(), 'bookC', model),
|
||||
]);
|
||||
|
||||
expect(maxActive).toBeGreaterThan(1);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,302 @@
|
||||
import { describe, it, expect, beforeEach, afterEach } from 'vitest';
|
||||
import { NodeDatabaseService } from '@/services/database/nodeDatabaseService';
|
||||
import { DatabaseService } from '@/types/database';
|
||||
import { migrate } from '@/services/database/migrate';
|
||||
import { getMigrations } from '@/services/database/migrations';
|
||||
import { ReedyDb } from '@/services/reedy/db/ReedyDb';
|
||||
import { BookRetriever } from '@/services/reedy/retrieval/BookRetriever';
|
||||
import type { EmbeddingModel } from '@/services/reedy/models/EmbeddingModel';
|
||||
import type { ChunkRow, EmbeddingRow } from '@/services/reedy/db/types';
|
||||
|
||||
const DIM = 4;
|
||||
|
||||
function unitVec(values: number[]): number[] {
|
||||
const norm = Math.sqrt(values.reduce((s, v) => s + v * v, 0)) || 1;
|
||||
return values.map((v) => v / norm);
|
||||
}
|
||||
|
||||
function fakeModel(
|
||||
opts: { id?: string; embedFn?: (texts: string[]) => Promise<number[][]> } = {},
|
||||
): EmbeddingModel {
|
||||
return {
|
||||
id: opts.id ?? 'fake-model',
|
||||
dim: DIM,
|
||||
embed: opts.embedFn ?? (async (texts) => texts.map(() => unitVec([1, 0, 0, 0]))),
|
||||
};
|
||||
}
|
||||
|
||||
function chunk(id: string, bookHash: string, pos: number, text: string): ChunkRow {
|
||||
return {
|
||||
id,
|
||||
bookHash,
|
||||
sectionIndex: 0,
|
||||
chapterTitle: 'Ch1',
|
||||
startCfi: `epubcfi(/6/2!/4/${pos * 2 + 2},/1:0,/1:10)`,
|
||||
endCfi: `epubcfi(/6/2!/4/${pos * 2 + 2},/1:10,/1:20)`,
|
||||
positionIndex: pos,
|
||||
text,
|
||||
tokenCount: text.split(/\s+/).length,
|
||||
};
|
||||
}
|
||||
|
||||
describe('BookRetriever', () => {
|
||||
let svc: DatabaseService;
|
||||
let reedy: ReedyDb;
|
||||
let retriever: BookRetriever;
|
||||
|
||||
beforeEach(async () => {
|
||||
svc = await NodeDatabaseService.open(':memory:', { experimental: ['index_method'] });
|
||||
await migrate(svc, getMigrations('reedy'));
|
||||
reedy = new ReedyDb(svc);
|
||||
retriever = new BookRetriever(reedy);
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
await svc.close();
|
||||
});
|
||||
|
||||
// -------------------------------------------------------------------------
|
||||
// status: not_indexed
|
||||
// -------------------------------------------------------------------------
|
||||
|
||||
it('returns not_indexed when the book has no meta row', async () => {
|
||||
const res = await retriever.search({
|
||||
bookHash: 'unknown',
|
||||
query: 'whatever',
|
||||
k: 5,
|
||||
activeEmbeddingModel: fakeModel(),
|
||||
});
|
||||
expect(res.status).toBe('not_indexed');
|
||||
expect(res.passages).toEqual([]);
|
||||
});
|
||||
|
||||
it('returns not_indexed while indexing is still in progress', async () => {
|
||||
await reedy.upsertBookMeta({
|
||||
bookHash: 'bk1',
|
||||
indexingStatus: 'indexing',
|
||||
chunkCount: 0,
|
||||
embeddingModel: 'fake-model',
|
||||
embeddingDim: DIM,
|
||||
indexedAt: null,
|
||||
error: null,
|
||||
});
|
||||
const res = await retriever.search({
|
||||
bookHash: 'bk1',
|
||||
query: 'q',
|
||||
k: 5,
|
||||
activeEmbeddingModel: fakeModel(),
|
||||
});
|
||||
expect(res.status).toBe('not_indexed');
|
||||
});
|
||||
|
||||
it('returns not_indexed when the prior index failed', async () => {
|
||||
await reedy.upsertBookMeta({
|
||||
bookHash: 'bk1',
|
||||
indexingStatus: 'failed',
|
||||
chunkCount: 0,
|
||||
embeddingModel: 'fake-model',
|
||||
embeddingDim: DIM,
|
||||
indexedAt: null,
|
||||
error: 'gateway down',
|
||||
});
|
||||
const res = await retriever.search({
|
||||
bookHash: 'bk1',
|
||||
query: 'q',
|
||||
k: 5,
|
||||
activeEmbeddingModel: fakeModel(),
|
||||
});
|
||||
expect(res.status).toBe('not_indexed');
|
||||
});
|
||||
|
||||
// -------------------------------------------------------------------------
|
||||
// status: empty_index
|
||||
// -------------------------------------------------------------------------
|
||||
|
||||
it('returns empty_index for an image-only book that was indexed with zero chunks', async () => {
|
||||
await reedy.upsertBookMeta({
|
||||
bookHash: 'bk1',
|
||||
indexingStatus: 'empty_index',
|
||||
chunkCount: 0,
|
||||
embeddingModel: 'fake-model',
|
||||
embeddingDim: DIM,
|
||||
indexedAt: Date.now(),
|
||||
error: null,
|
||||
});
|
||||
const res = await retriever.search({
|
||||
bookHash: 'bk1',
|
||||
query: 'q',
|
||||
k: 5,
|
||||
activeEmbeddingModel: fakeModel(),
|
||||
});
|
||||
expect(res.status).toBe('empty_index');
|
||||
expect(res.passages).toEqual([]);
|
||||
});
|
||||
|
||||
// -------------------------------------------------------------------------
|
||||
// status: stale_index
|
||||
// -------------------------------------------------------------------------
|
||||
|
||||
it('returns stale_index when the active model differs from the one used to index', async () => {
|
||||
await reedy.upsertBookMeta({
|
||||
bookHash: 'bk1',
|
||||
indexingStatus: 'indexed',
|
||||
chunkCount: 5,
|
||||
embeddingModel: 'nomic-embed-text',
|
||||
embeddingDim: DIM,
|
||||
indexedAt: Date.now(),
|
||||
error: null,
|
||||
});
|
||||
const res = await retriever.search({
|
||||
bookHash: 'bk1',
|
||||
query: 'q',
|
||||
k: 5,
|
||||
activeEmbeddingModel: fakeModel({ id: 'text-embedding-3-small' }),
|
||||
});
|
||||
expect(res.status).toBe('stale_index');
|
||||
expect(res.reason).toMatch(/text-embedding-3-small/);
|
||||
expect(res.reason).toMatch(/nomic-embed-text/);
|
||||
});
|
||||
|
||||
// -------------------------------------------------------------------------
|
||||
// happy path + per-book isolation
|
||||
// -------------------------------------------------------------------------
|
||||
|
||||
describe('with indexed data', () => {
|
||||
beforeEach(async () => {
|
||||
await reedy.upsertBookMeta({
|
||||
bookHash: 'bookA',
|
||||
indexingStatus: 'indexed',
|
||||
chunkCount: 4,
|
||||
embeddingModel: 'fake-model',
|
||||
embeddingDim: DIM,
|
||||
indexedAt: Date.now(),
|
||||
error: null,
|
||||
});
|
||||
await reedy.upsertBookMeta({
|
||||
bookHash: 'bookB',
|
||||
indexingStatus: 'indexed',
|
||||
chunkCount: 1,
|
||||
embeddingModel: 'fake-model',
|
||||
embeddingDim: DIM,
|
||||
indexedAt: Date.now(),
|
||||
error: null,
|
||||
});
|
||||
await reedy.ensureEmbeddingsTable(DIM);
|
||||
await reedy.insertChunks([
|
||||
chunk('a0', 'bookA', 0, 'alpha bravo charlie introduction'),
|
||||
chunk('a1', 'bookA', 1, 'delta echo foxtrot middle chapter'),
|
||||
chunk('a2', 'bookA', 2, 'golf hotel india later passages'),
|
||||
chunk('a3', 'bookA', 3, 'juliet kilo lima final wrap-up'),
|
||||
chunk('b1', 'bookB', 0, 'alpha bravo charlie introduction'),
|
||||
]);
|
||||
const embs: EmbeddingRow[] = [
|
||||
{ chunkId: 'a0', bookHash: 'bookA', embedding: unitVec([1, 0, 0, 0]) },
|
||||
{ chunkId: 'a1', bookHash: 'bookA', embedding: unitVec([0, 1, 0, 0]) },
|
||||
{ chunkId: 'a2', bookHash: 'bookA', embedding: unitVec([0, 0, 1, 0]) },
|
||||
{ chunkId: 'a3', bookHash: 'bookA', embedding: unitVec([0, 0, 0, 1]) },
|
||||
{ chunkId: 'b1', bookHash: 'bookB', embedding: unitVec([1, 0, 0, 0]) },
|
||||
];
|
||||
await reedy.insertEmbeddings(embs);
|
||||
});
|
||||
|
||||
it('returns ok with passages strictly from the requested book (T3 isolation)', async () => {
|
||||
const res = await retriever.search({
|
||||
bookHash: 'bookA',
|
||||
query: 'introduction',
|
||||
k: 5,
|
||||
activeEmbeddingModel: fakeModel({
|
||||
embedFn: async (texts) => texts.map(() => unitVec([1, 0, 0, 0])),
|
||||
}),
|
||||
});
|
||||
expect(res.status).toBe('ok');
|
||||
expect(res.passages.length).toBeGreaterThan(0);
|
||||
for (const p of res.passages) expect(p.bookHash).toBe('bookA');
|
||||
});
|
||||
|
||||
it('exposes start_cfi and end_cfi on each passage so the UI can navigate', async () => {
|
||||
const res = await retriever.search({
|
||||
bookHash: 'bookA',
|
||||
query: 'introduction',
|
||||
k: 2,
|
||||
activeEmbeddingModel: fakeModel({
|
||||
embedFn: async (texts) => texts.map(() => unitVec([1, 0, 0, 0])),
|
||||
}),
|
||||
});
|
||||
expect(res.passages.length).toBeGreaterThan(0);
|
||||
for (const p of res.passages) {
|
||||
expect(p.cfi).toMatch(/^epubcfi\(/);
|
||||
expect(p.endCfi).toMatch(/^epubcfi\(/);
|
||||
expect(p.text).toBeTruthy();
|
||||
}
|
||||
});
|
||||
|
||||
it('exact-quote query ranks the FTS-aligned chunk first (T2 FTS dominance)', async () => {
|
||||
const res = await retriever.search({
|
||||
bookHash: 'bookA',
|
||||
query: 'wrap-up',
|
||||
k: 3,
|
||||
activeEmbeddingModel: fakeModel({
|
||||
// Vector points elsewhere so FTS has to win
|
||||
embedFn: async (texts) => texts.map(() => unitVec([1, 0, 0, 0])),
|
||||
}),
|
||||
});
|
||||
expect(res.status).toBe('ok');
|
||||
expect(res.passages[0]!.id).toBe('a3');
|
||||
});
|
||||
|
||||
it('paraphrase query without a lexical match ranks by vector similarity (T2 vector dominance)', async () => {
|
||||
const res = await retriever.search({
|
||||
bookHash: 'bookA',
|
||||
query: 'something completely orthogonal to chunk text',
|
||||
k: 3,
|
||||
activeEmbeddingModel: fakeModel({
|
||||
embedFn: async (texts) => texts.map(() => unitVec([0, 0, 1, 0])),
|
||||
}),
|
||||
});
|
||||
expect(res.status).toBe('ok');
|
||||
expect(res.passages[0]!.id).toBe('a2');
|
||||
});
|
||||
|
||||
it('drops passages above spoilerBoundPosition', async () => {
|
||||
const res = await retriever.search({
|
||||
bookHash: 'bookA',
|
||||
query: 'final',
|
||||
k: 5,
|
||||
spoilerBoundPosition: 1,
|
||||
activeEmbeddingModel: fakeModel({
|
||||
embedFn: async (texts) => texts.map(() => unitVec([0, 0, 0, 1])),
|
||||
}),
|
||||
});
|
||||
for (const p of res.passages) expect(p.positionIndex).toBeLessThanOrEqual(1);
|
||||
});
|
||||
|
||||
it('falls back to FTS-only with status=degraded when the embedding call times out', async () => {
|
||||
const slowModel: EmbeddingModel = {
|
||||
id: 'fake-model',
|
||||
dim: DIM,
|
||||
embed: async (_texts, opts) => {
|
||||
await new Promise((resolve, reject) => {
|
||||
const t = setTimeout(resolve, 50);
|
||||
opts?.signal?.addEventListener('abort', () => {
|
||||
clearTimeout(t);
|
||||
reject(new DOMException('aborted', 'AbortError'));
|
||||
});
|
||||
});
|
||||
return [unitVec([1, 0, 0, 0])];
|
||||
},
|
||||
};
|
||||
const res = await retriever.search({
|
||||
bookHash: 'bookA',
|
||||
query: 'introduction',
|
||||
k: 3,
|
||||
activeEmbeddingModel: slowModel,
|
||||
embeddingTimeoutMs: 10,
|
||||
});
|
||||
expect(res.status).toBe('degraded');
|
||||
expect(res.reason).toMatch(/timeout|abort/i);
|
||||
// FTS still finds the lexical match for "introduction"
|
||||
expect(res.passages.length).toBeGreaterThan(0);
|
||||
expect(res.passages[0]!.id).toBe('a0');
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,148 @@
|
||||
/**
|
||||
* @vitest-environment jsdom
|
||||
*/
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import * as CFI from 'foliate-js/epubcfi.js';
|
||||
import { chunkSection } from '@/services/reedy/retrieval/CfiChunker';
|
||||
|
||||
function makeDoc(bodyHtml: string): Document {
|
||||
return new DOMParser().parseFromString(
|
||||
`<!DOCTYPE html><html><body>${bodyHtml}</body></html>`,
|
||||
'text/html',
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Strip the wrapping `epubcfi(...)` and the leading `/6/N!` spine step that
|
||||
* CfiChunker prepends, so we can hand the inner path to CFI.toRange against
|
||||
* the section document.
|
||||
*/
|
||||
function innerCfi(stored: string): string {
|
||||
const m = stored.match(/^epubcfi\((.+)\)$/);
|
||||
if (!m) throw new Error(`malformed CFI: ${stored}`);
|
||||
const inner = m[1]!;
|
||||
const spineSplit = inner.indexOf('!');
|
||||
return spineSplit >= 0 ? `epubcfi(${inner.slice(spineSplit + 1)})` : `epubcfi(${inner})`;
|
||||
}
|
||||
|
||||
function normalizeWs(s: string): string {
|
||||
return s.replace(/\s+/g, ' ').trim();
|
||||
}
|
||||
|
||||
describe('CfiChunker', () => {
|
||||
it('returns no chunks for an image-only document with no extractable text', () => {
|
||||
const doc = makeDoc('<img src="cover.png" alt="" />');
|
||||
const chunks = chunkSection(doc, 0, 'Front', 'bk1');
|
||||
expect(chunks).toEqual([]);
|
||||
});
|
||||
|
||||
it('returns no chunks when the body is empty', () => {
|
||||
const doc = makeDoc('');
|
||||
const chunks = chunkSection(doc, 0, 'Empty', 'bk1');
|
||||
expect(chunks).toEqual([]);
|
||||
});
|
||||
|
||||
it('produces a single chunk for a short section that fits under the size limit', () => {
|
||||
const doc = makeDoc('<p>Hello world.</p>');
|
||||
const chunks = chunkSection(doc, 0, 'Greeting', 'bk1');
|
||||
expect(chunks).toHaveLength(1);
|
||||
expect(chunks[0]!.text).toContain('Hello world');
|
||||
expect(chunks[0]!.bookHash).toBe('bk1');
|
||||
expect(chunks[0]!.sectionIndex).toBe(0);
|
||||
expect(chunks[0]!.chapterTitle).toBe('Greeting');
|
||||
expect(chunks[0]!.positionIndex).toBe(0);
|
||||
});
|
||||
|
||||
it('prepends the EPUB spine prefix /6/{(idx+1)*2}! to stored CFIs', () => {
|
||||
const doc = makeDoc('<p>Hello world.</p>');
|
||||
const chunks = chunkSection(doc, 3, 'Ch4', 'bk1');
|
||||
expect(chunks[0]!.startCfi).toMatch(/^epubcfi\(\/6\/8!/);
|
||||
expect(chunks[0]!.endCfi).toMatch(/^epubcfi\(\/6\/8!/);
|
||||
});
|
||||
|
||||
it('generates CFIs that round-trip via CFI.toRange and resolve to the chunk text', () => {
|
||||
const doc = makeDoc(
|
||||
'<p id="p1">First paragraph here.</p><p id="p2">Second has <em>emphasis</em> and more text.</p>',
|
||||
);
|
||||
const chunks = chunkSection(doc, 0, 'Roundtrip', 'bk1');
|
||||
expect(chunks.length).toBeGreaterThan(0);
|
||||
for (const c of chunks) {
|
||||
const parts = CFI.parse(innerCfi(c.startCfi));
|
||||
const range = CFI.toRange(doc, parts);
|
||||
expect(range, `toRange returned null for ${c.startCfi}`).not.toBeNull();
|
||||
// The resolved range's start position should fall inside the chunk text.
|
||||
const resolvedText = range!.toString();
|
||||
// We don't require an exact equality here because toRange returns a
|
||||
// collapsed start range; instead verify the first 8 chars of the chunk
|
||||
// align with the text starting at the resolved start position.
|
||||
const chunkHead = normalizeWs(c.text).slice(0, 8);
|
||||
if (chunkHead.length > 0) {
|
||||
// Build a fresh range from the start CFI through end CFI and compare.
|
||||
const endParts = CFI.parse(innerCfi(c.endCfi));
|
||||
const endRange = CFI.toRange(doc, endParts);
|
||||
expect(endRange).not.toBeNull();
|
||||
const fullRange = doc.createRange();
|
||||
fullRange.setStart(range!.startContainer, range!.startOffset);
|
||||
fullRange.setEnd(endRange!.startContainer, endRange!.startOffset);
|
||||
expect(normalizeWs(fullRange.toString())).toBe(normalizeWs(c.text));
|
||||
} else {
|
||||
expect(resolvedText).toBeDefined();
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
it('skips text inside <script>, <style>, and <noscript>', () => {
|
||||
const doc = makeDoc(
|
||||
'<p>Visible text here.</p>' +
|
||||
'<script>alert("hidden");</script>' +
|
||||
'<style>.x{color:red}</style>' +
|
||||
'<noscript>no js</noscript>' +
|
||||
'<p>Another visible paragraph.</p>',
|
||||
);
|
||||
const chunks = chunkSection(doc, 0, 'Filtering', 'bk1');
|
||||
const allText = chunks.map((c) => c.text).join(' ');
|
||||
expect(allText).not.toContain('alert(');
|
||||
expect(allText).not.toContain('color:red');
|
||||
expect(allText).not.toContain('no js');
|
||||
expect(allText).toContain('Visible text here');
|
||||
expect(allText).toContain('Another visible paragraph');
|
||||
});
|
||||
|
||||
it('splits long content into multiple chunks with monotonically increasing position_index', () => {
|
||||
const paragraph = 'Lorem ipsum dolor sit amet, consectetur adipiscing elit. '.repeat(40);
|
||||
const doc = makeDoc(`<p>${paragraph}</p><p>${paragraph}</p>`);
|
||||
const chunks = chunkSection(doc, 0, 'Long', 'bk1', {
|
||||
maxChunkSize: 300,
|
||||
minChunkSize: 50,
|
||||
overlapSize: 30,
|
||||
});
|
||||
expect(chunks.length).toBeGreaterThan(1);
|
||||
for (let i = 1; i < chunks.length; i++) {
|
||||
expect(chunks[i]!.positionIndex).toBe(chunks[i - 1]!.positionIndex + 1);
|
||||
}
|
||||
for (const c of chunks) {
|
||||
expect(c.text.length).toBeLessThanOrEqual(400); // maxChunkSize + breakSlack
|
||||
}
|
||||
});
|
||||
|
||||
it('produces token_count approximating whitespace-separated word count', () => {
|
||||
const doc = makeDoc('<p>one two three four five</p>');
|
||||
const chunks = chunkSection(doc, 0, 'Tokens', 'bk1');
|
||||
expect(chunks).toHaveLength(1);
|
||||
expect(chunks[0]!.tokenCount).toBe(5);
|
||||
});
|
||||
|
||||
it('assigns deterministic ids that include book hash, section, and position', () => {
|
||||
const doc = makeDoc('<p>A</p><p>B</p>');
|
||||
const chunks = chunkSection(doc, 2, 'Det', 'hashXYZ', {
|
||||
maxChunkSize: 1,
|
||||
minChunkSize: 1,
|
||||
overlapSize: 0,
|
||||
});
|
||||
for (let i = 0; i < chunks.length; i++) {
|
||||
expect(chunks[i]!.id).toContain('hashXYZ');
|
||||
expect(chunks[i]!.id).toContain('2');
|
||||
expect(chunks[i]!.id).toContain(`${i}`);
|
||||
}
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,299 @@
|
||||
import { describe, it, expect, beforeEach, afterEach } from 'vitest';
|
||||
import { NodeDatabaseService } from '@/services/database/nodeDatabaseService';
|
||||
import { DatabaseService } from '@/types/database';
|
||||
import { migrate } from '@/services/database/migrate';
|
||||
import { getMigrations } from '@/services/database/migrations';
|
||||
import { ReedyDb } from '@/services/reedy/db/ReedyDb';
|
||||
import type { ChunkRow, EmbeddingRow } from '@/services/reedy/db/types';
|
||||
|
||||
const DIM = 4;
|
||||
|
||||
function chunk(id: string, bookHash: string, pos: number, text: string): ChunkRow {
|
||||
return {
|
||||
id,
|
||||
bookHash,
|
||||
sectionIndex: 0,
|
||||
chapterTitle: 'Ch1',
|
||||
startCfi: `/6/4!/4/${pos * 2 + 2},/1:0,/1:10`,
|
||||
endCfi: `/6/4!/4/${pos * 2 + 2},/1:10,/1:20`,
|
||||
positionIndex: pos,
|
||||
text,
|
||||
tokenCount: text.split(/\s+/).length,
|
||||
};
|
||||
}
|
||||
|
||||
function unitVec(values: number[]): number[] {
|
||||
const norm = Math.sqrt(values.reduce((s, v) => s + v * v, 0));
|
||||
return values.map((v) => v / norm);
|
||||
}
|
||||
|
||||
describe('ReedyDb', () => {
|
||||
let svc: DatabaseService;
|
||||
let reedy: ReedyDb;
|
||||
|
||||
beforeEach(async () => {
|
||||
svc = await NodeDatabaseService.open(':memory:', { experimental: ['index_method'] });
|
||||
await migrate(svc, getMigrations('reedy'));
|
||||
reedy = new ReedyDb(svc);
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
await svc.close();
|
||||
});
|
||||
|
||||
describe('book meta', () => {
|
||||
it('getBookMeta returns null when no row exists', async () => {
|
||||
const meta = await reedy.getBookMeta('missing');
|
||||
expect(meta).toBeNull();
|
||||
});
|
||||
|
||||
it('upsertBookMeta then getBookMeta round-trips all fields', async () => {
|
||||
await reedy.upsertBookMeta({
|
||||
bookHash: 'bk1',
|
||||
indexingStatus: 'indexed',
|
||||
chunkCount: 42,
|
||||
embeddingModel: 'nomic-embed-text',
|
||||
embeddingDim: DIM,
|
||||
indexedAt: 1700000000,
|
||||
error: null,
|
||||
});
|
||||
|
||||
const meta = await reedy.getBookMeta('bk1');
|
||||
expect(meta).toEqual({
|
||||
bookHash: 'bk1',
|
||||
indexingStatus: 'indexed',
|
||||
chunkCount: 42,
|
||||
embeddingModel: 'nomic-embed-text',
|
||||
embeddingDim: DIM,
|
||||
indexedAt: 1700000000,
|
||||
error: null,
|
||||
});
|
||||
});
|
||||
|
||||
it('setIndexingStatus preserves untouched fields (partial update)', async () => {
|
||||
await reedy.upsertBookMeta({
|
||||
bookHash: 'bk1',
|
||||
indexingStatus: 'indexing',
|
||||
chunkCount: 0,
|
||||
embeddingModel: 'nomic-embed-text',
|
||||
embeddingDim: DIM,
|
||||
indexedAt: null,
|
||||
error: null,
|
||||
});
|
||||
|
||||
await reedy.setIndexingStatus('bk1', 'failed', { error: 'embed timeout' });
|
||||
|
||||
const meta = await reedy.getBookMeta('bk1');
|
||||
expect(meta?.indexingStatus).toBe('failed');
|
||||
expect(meta?.error).toBe('embed timeout');
|
||||
expect(meta?.embeddingModel).toBe('nomic-embed-text');
|
||||
expect(meta?.embeddingDim).toBe(DIM);
|
||||
});
|
||||
});
|
||||
|
||||
describe('ensureEmbeddingsTable', () => {
|
||||
it('creates the embeddings table on first call', async () => {
|
||||
await reedy.ensureEmbeddingsTable(DIM);
|
||||
const tables = await svc.select<{ name: string }>(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name='reedy_book_chunk_embeddings'",
|
||||
);
|
||||
expect(tables).toHaveLength(1);
|
||||
});
|
||||
|
||||
it('is idempotent across repeated calls with the same dim', async () => {
|
||||
await reedy.ensureEmbeddingsTable(DIM);
|
||||
await reedy.ensureEmbeddingsTable(DIM);
|
||||
await reedy.ensureEmbeddingsTable(DIM);
|
||||
const tables = await svc.select<{ name: string }>(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name='reedy_book_chunk_embeddings'",
|
||||
);
|
||||
expect(tables).toHaveLength(1);
|
||||
});
|
||||
|
||||
it('throws when called with a different dim than the existing table', async () => {
|
||||
await reedy.ensureEmbeddingsTable(DIM);
|
||||
await expect(reedy.ensureEmbeddingsTable(DIM + 1)).rejects.toThrow(/dim/);
|
||||
});
|
||||
});
|
||||
|
||||
describe('chunk + embedding writes', () => {
|
||||
beforeEach(async () => {
|
||||
await reedy.ensureEmbeddingsTable(DIM);
|
||||
});
|
||||
|
||||
it('insertChunks writes multiple rows in one batch', async () => {
|
||||
const chunks = [
|
||||
chunk('c1', 'bk1', 0, 'alpha bravo'),
|
||||
chunk('c2', 'bk1', 1, 'charlie delta'),
|
||||
chunk('c3', 'bk1', 2, "let's go — apostrophe & ampersand"),
|
||||
];
|
||||
await reedy.insertChunks(chunks);
|
||||
|
||||
const rows = await svc.select<{ id: string; text: string }>(
|
||||
'SELECT id, text FROM reedy_book_chunks ORDER BY position_index',
|
||||
);
|
||||
expect(rows).toHaveLength(3);
|
||||
expect(rows.map((r) => r.text)).toEqual(chunks.map((c) => c.text));
|
||||
});
|
||||
|
||||
it('insertEmbeddings writes vectors that round-trip via vector_extract', async () => {
|
||||
await reedy.insertChunks([chunk('c1', 'bk1', 0, 'alpha')]);
|
||||
const emb: EmbeddingRow = {
|
||||
chunkId: 'c1',
|
||||
bookHash: 'bk1',
|
||||
embedding: unitVec([1, 2, 3, 4]),
|
||||
};
|
||||
await reedy.insertEmbeddings([emb]);
|
||||
|
||||
const rows = await svc.select<{ extracted: string }>(
|
||||
"SELECT vector_extract(embedding) AS extracted FROM reedy_book_chunk_embeddings WHERE chunk_id = 'c1'",
|
||||
);
|
||||
expect(rows).toHaveLength(1);
|
||||
const parsed = JSON.parse(rows[0]!.extracted) as number[];
|
||||
expect(parsed).toHaveLength(DIM);
|
||||
for (let i = 0; i < DIM; i++) {
|
||||
expect(parsed[i]).toBeCloseTo(emb.embedding[i]!, 4);
|
||||
}
|
||||
});
|
||||
|
||||
it('insertEmbeddings throws when embedding length does not match dim', async () => {
|
||||
await reedy.insertChunks([chunk('c1', 'bk1', 0, 'alpha')]);
|
||||
await expect(
|
||||
reedy.insertEmbeddings([{ chunkId: 'c1', bookHash: 'bk1', embedding: [1, 2, 3] }]),
|
||||
).rejects.toThrow(/dim/);
|
||||
});
|
||||
});
|
||||
|
||||
describe('dropBookData / wipeAllData', () => {
|
||||
beforeEach(async () => {
|
||||
await reedy.ensureEmbeddingsTable(DIM);
|
||||
await reedy.insertChunks([
|
||||
chunk('a1', 'bookA', 0, 'A one'),
|
||||
chunk('a2', 'bookA', 1, 'A two'),
|
||||
chunk('b1', 'bookB', 0, 'B one'),
|
||||
]);
|
||||
await reedy.insertEmbeddings([
|
||||
{ chunkId: 'a1', bookHash: 'bookA', embedding: unitVec([1, 0, 0, 0]) },
|
||||
{ chunkId: 'a2', bookHash: 'bookA', embedding: unitVec([0, 1, 0, 0]) },
|
||||
{ chunkId: 'b1', bookHash: 'bookB', embedding: unitVec([0, 0, 1, 0]) },
|
||||
]);
|
||||
});
|
||||
|
||||
it('dropBookData removes only the targeted book’s chunks and embeddings', async () => {
|
||||
await reedy.dropBookData('bookA');
|
||||
|
||||
const chunks = await svc.select<{ id: string }>(
|
||||
'SELECT id FROM reedy_book_chunks ORDER BY id',
|
||||
);
|
||||
expect(chunks.map((c) => c.id)).toEqual(['b1']);
|
||||
|
||||
const embs = await svc.select<{ chunk_id: string }>(
|
||||
'SELECT chunk_id FROM reedy_book_chunk_embeddings ORDER BY chunk_id',
|
||||
);
|
||||
expect(embs.map((e) => e.chunk_id)).toEqual(['b1']);
|
||||
});
|
||||
|
||||
it('wipeAllData clears chunks, embeddings, and meta across every book', async () => {
|
||||
await reedy.upsertBookMeta({
|
||||
bookHash: 'bookA',
|
||||
indexingStatus: 'indexed',
|
||||
chunkCount: 2,
|
||||
embeddingModel: 'nomic-embed-text',
|
||||
embeddingDim: DIM,
|
||||
indexedAt: 1700000000,
|
||||
error: null,
|
||||
});
|
||||
|
||||
await reedy.wipeAllData();
|
||||
|
||||
const chunks = await svc.select('SELECT id FROM reedy_book_chunks');
|
||||
const meta = await svc.select('SELECT book_hash FROM reedy_book_meta');
|
||||
// wipeAllData DROPS the embeddings table so ensureEmbeddingsTable(newDim)
|
||||
// can recreate it with a different vector32 width. Check it's gone.
|
||||
const embTable = await svc.select<{ name: string }>(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name='reedy_book_chunk_embeddings'",
|
||||
);
|
||||
expect(chunks).toHaveLength(0);
|
||||
expect(meta).toHaveLength(0);
|
||||
expect(embTable).toHaveLength(0);
|
||||
});
|
||||
|
||||
it('wipeAllData lets ensureEmbeddingsTable recreate the table at a new dim', async () => {
|
||||
await reedy.wipeAllData();
|
||||
await reedy.ensureEmbeddingsTable(DIM + 4);
|
||||
|
||||
const rows = await svc.select<{ sql: string }>(
|
||||
"SELECT sql FROM sqlite_master WHERE type='table' AND name='reedy_book_chunk_embeddings'",
|
||||
);
|
||||
expect(rows[0]!.sql).toMatch(/vector32\s*\(\s*8\s*\)/i);
|
||||
});
|
||||
});
|
||||
|
||||
describe('hybridSearch', () => {
|
||||
beforeEach(async () => {
|
||||
await reedy.ensureEmbeddingsTable(DIM);
|
||||
// bookA: 4 chunks at increasing positions, embeddings on different axes.
|
||||
// bookB: 1 chunk with an embedding very close to bookA's first to test
|
||||
// that book filtering works.
|
||||
const chunks = [
|
||||
chunk('a1', 'bookA', 0, 'apple banana cherry'),
|
||||
chunk('a2', 'bookA', 1, 'date elderberry fig'),
|
||||
chunk('a3', 'bookA', 2, 'grape honeydew imbe'),
|
||||
chunk('a4', 'bookA', 3, 'jackfruit kiwi lemon'),
|
||||
chunk('b1', 'bookB', 0, 'apple banana cherry'),
|
||||
];
|
||||
const embs: EmbeddingRow[] = [
|
||||
{ chunkId: 'a1', bookHash: 'bookA', embedding: unitVec([1, 0, 0, 0]) },
|
||||
{ chunkId: 'a2', bookHash: 'bookA', embedding: unitVec([0, 1, 0, 0]) },
|
||||
{ chunkId: 'a3', bookHash: 'bookA', embedding: unitVec([0, 0, 1, 0]) },
|
||||
{ chunkId: 'a4', bookHash: 'bookA', embedding: unitVec([0, 0, 0, 1]) },
|
||||
{ chunkId: 'b1', bookHash: 'bookB', embedding: unitVec([1, 0, 0, 0]) },
|
||||
];
|
||||
await reedy.insertChunks(chunks);
|
||||
await reedy.insertEmbeddings(embs);
|
||||
});
|
||||
|
||||
it('filters strictly by book_hash (no cross-book bleed)', async () => {
|
||||
const res = await reedy.hybridSearch({
|
||||
bookHash: 'bookA',
|
||||
queryText: 'apple',
|
||||
queryEmbedding: unitVec([1, 0, 0, 0]),
|
||||
k: 5,
|
||||
});
|
||||
for (const r of res) expect(r.bookHash).toBe('bookA');
|
||||
});
|
||||
|
||||
it('ranks the vector-aligned chunk highly when query embedding points at it', async () => {
|
||||
const res = await reedy.hybridSearch({
|
||||
bookHash: 'bookA',
|
||||
queryText: 'something unrelated',
|
||||
queryEmbedding: unitVec([0, 0, 1, 0]),
|
||||
k: 5,
|
||||
});
|
||||
expect(res.length).toBeGreaterThan(0);
|
||||
expect(res[0]!.id).toBe('a3');
|
||||
});
|
||||
|
||||
it('ranks the FTS-matched chunk highly when query text matches that chunk', async () => {
|
||||
const res = await reedy.hybridSearch({
|
||||
bookHash: 'bookA',
|
||||
queryText: 'jackfruit',
|
||||
queryEmbedding: unitVec([1, 0, 0, 0]), // misaligned vector
|
||||
k: 5,
|
||||
});
|
||||
expect(res.length).toBeGreaterThan(0);
|
||||
expect(res[0]!.id).toBe('a4');
|
||||
});
|
||||
|
||||
it('drops chunks with position_index > spoilerBoundPosition', async () => {
|
||||
const res = await reedy.hybridSearch({
|
||||
bookHash: 'bookA',
|
||||
queryText: 'apple banana',
|
||||
queryEmbedding: unitVec([0, 0, 0, 1]), // would otherwise surface a4
|
||||
k: 5,
|
||||
spoilerBoundPosition: 1,
|
||||
});
|
||||
for (const r of res) expect(r.positionIndex).toBeLessThanOrEqual(1);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,269 @@
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import {
|
||||
buildLookupTool,
|
||||
createTurnState,
|
||||
lookupInputSchema,
|
||||
serializeForModel,
|
||||
type LookupToolResult,
|
||||
} from '@/services/reedy/tools/lookupPassage';
|
||||
import type { BookRetriever, RetrieverResult } from '@/services/reedy/retrieval/BookRetriever';
|
||||
import type { EmbeddingModel } from '@/services/reedy/models/EmbeddingModel';
|
||||
|
||||
const model: EmbeddingModel = {
|
||||
id: 'fake-model',
|
||||
dim: 4,
|
||||
async embed(texts) {
|
||||
return texts.map(() => [1, 0, 0, 0]);
|
||||
},
|
||||
};
|
||||
|
||||
function passage(id: string, text: string, position = 0): RetrieverResult['passages'][number] {
|
||||
return {
|
||||
id,
|
||||
bookHash: 'bk1',
|
||||
cfi: `epubcfi(/6/2!/4/${position * 2 + 2})`,
|
||||
endCfi: `epubcfi(/6/2!/4/${position * 2 + 4})`,
|
||||
chapterTitle: 'Ch1',
|
||||
text,
|
||||
positionIndex: position,
|
||||
score: 0.5,
|
||||
};
|
||||
}
|
||||
|
||||
function fakeRetriever(impl: (args: { query: string }) => Promise<RetrieverResult>): BookRetriever {
|
||||
return { search: vi.fn((args) => impl({ query: args.query })) } as unknown as BookRetriever;
|
||||
}
|
||||
|
||||
async function runExecute(
|
||||
tool: ReturnType<typeof buildLookupTool>,
|
||||
input: unknown,
|
||||
): Promise<LookupToolResult> {
|
||||
if (!tool.execute) throw new Error('tool execute missing');
|
||||
return tool.execute(input as { query: string; topK: number }, {
|
||||
toolCallId: 'tc1',
|
||||
messages: [],
|
||||
}) as Promise<LookupToolResult>;
|
||||
}
|
||||
|
||||
describe('buildLookupTool', () => {
|
||||
let turnState: ReturnType<typeof createTurnState>;
|
||||
|
||||
beforeEach(() => {
|
||||
turnState = createTurnState();
|
||||
});
|
||||
|
||||
describe('schema validation', () => {
|
||||
it('rejects empty query', () => {
|
||||
expect(lookupInputSchema.safeParse({ query: '', topK: 3 }).success).toBe(false);
|
||||
});
|
||||
|
||||
it('rejects oversized query (>500 chars)', () => {
|
||||
expect(lookupInputSchema.safeParse({ query: 'x'.repeat(501), topK: 3 }).success).toBe(false);
|
||||
});
|
||||
|
||||
it('rejects topK > 5', () => {
|
||||
expect(lookupInputSchema.safeParse({ query: 'q', topK: 6 }).success).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('dedupe', () => {
|
||||
it('returns cached: true on a second identical call within the same turn', async () => {
|
||||
let calls = 0;
|
||||
const retriever = fakeRetriever(async () => {
|
||||
calls++;
|
||||
return { passages: [passage('p1', 'hello')], status: 'ok' };
|
||||
});
|
||||
const tool = buildLookupTool({
|
||||
bookHash: 'bk1',
|
||||
retriever,
|
||||
activeEmbeddingModel: model,
|
||||
turnState,
|
||||
});
|
||||
const a = await runExecute(tool, { query: 'foo', topK: 3 });
|
||||
const b = await runExecute(tool, { query: ' Foo ', topK: 3 }); // case + whitespace normalized
|
||||
expect(calls).toBe(1);
|
||||
expect(a.cached).toBeFalsy();
|
||||
expect(b.cached).toBe(true);
|
||||
});
|
||||
|
||||
it('does NOT dedupe when topK differs', async () => {
|
||||
let calls = 0;
|
||||
const retriever = fakeRetriever(async () => {
|
||||
calls++;
|
||||
return { passages: [passage('p1', 'hello')], status: 'ok' };
|
||||
});
|
||||
const tool = buildLookupTool({
|
||||
bookHash: 'bk1',
|
||||
retriever,
|
||||
activeEmbeddingModel: model,
|
||||
turnState,
|
||||
});
|
||||
await runExecute(tool, { query: 'foo', topK: 3 });
|
||||
await runExecute(tool, { query: 'foo', topK: 5 });
|
||||
expect(calls).toBe(2);
|
||||
});
|
||||
|
||||
it('does NOT dedupe when spoilerBoundPosition differs', async () => {
|
||||
let calls = 0;
|
||||
const retriever = fakeRetriever(async () => {
|
||||
calls++;
|
||||
return { passages: [passage('p1', 'hello')], status: 'ok' };
|
||||
});
|
||||
const toolA = buildLookupTool({
|
||||
bookHash: 'bk1',
|
||||
retriever,
|
||||
activeEmbeddingModel: model,
|
||||
turnState,
|
||||
spoilerBoundPosition: 5,
|
||||
});
|
||||
const toolB = buildLookupTool({
|
||||
bookHash: 'bk1',
|
||||
retriever,
|
||||
activeEmbeddingModel: model,
|
||||
turnState,
|
||||
spoilerBoundPosition: 10,
|
||||
});
|
||||
await runExecute(toolA, { query: 'foo', topK: 3 });
|
||||
await runExecute(toolB, { query: 'foo', topK: 3 });
|
||||
expect(calls).toBe(2);
|
||||
});
|
||||
});
|
||||
|
||||
describe('parallel-call serialization', () => {
|
||||
it('serializes concurrent executes through turnState.pendingChain', async () => {
|
||||
let active = 0;
|
||||
let maxActive = 0;
|
||||
const retriever = fakeRetriever(async () => {
|
||||
active++;
|
||||
maxActive = Math.max(maxActive, active);
|
||||
await new Promise((r) => setTimeout(r, 10));
|
||||
active--;
|
||||
return { passages: [passage('p1', 'x')], status: 'ok' };
|
||||
});
|
||||
const tool = buildLookupTool({
|
||||
bookHash: 'bk1',
|
||||
retriever,
|
||||
activeEmbeddingModel: model,
|
||||
turnState,
|
||||
});
|
||||
await Promise.all([
|
||||
runExecute(tool, { query: 'a', topK: 3 }),
|
||||
runExecute(tool, { query: 'b', topK: 3 }),
|
||||
runExecute(tool, { query: 'c', topK: 3 }),
|
||||
]);
|
||||
expect(maxActive).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe('wall-clock budget', () => {
|
||||
it('returns status=budget_exceeded once turnState.totalToolMs exceeds 10000', async () => {
|
||||
const retriever = fakeRetriever(async () => ({
|
||||
passages: [passage('p1', 'x')],
|
||||
status: 'ok',
|
||||
}));
|
||||
const tool = buildLookupTool({
|
||||
bookHash: 'bk1',
|
||||
retriever,
|
||||
activeEmbeddingModel: model,
|
||||
turnState,
|
||||
});
|
||||
// Pre-load the budget so the next call is over the limit.
|
||||
turnState.totalToolMs = 10001;
|
||||
const res = await runExecute(tool, { query: 'q', topK: 3 });
|
||||
expect(res.status).toBe('budget_exceeded');
|
||||
expect(res.passages).toEqual([]);
|
||||
expect(res.hint).toBeTruthy();
|
||||
});
|
||||
});
|
||||
|
||||
describe('result-size clamp', () => {
|
||||
it('drops the lowest-ranked passages until total chars ≤ 6000', async () => {
|
||||
const huge = 'x'.repeat(2000);
|
||||
const retriever = fakeRetriever(async () => ({
|
||||
passages: [
|
||||
passage('p0', huge, 0),
|
||||
passage('p1', huge, 1),
|
||||
passage('p2', huge, 2),
|
||||
passage('p3', huge, 3),
|
||||
passage('p4', huge, 4),
|
||||
],
|
||||
status: 'ok',
|
||||
}));
|
||||
const tool = buildLookupTool({
|
||||
bookHash: 'bk1',
|
||||
retriever,
|
||||
activeEmbeddingModel: model,
|
||||
turnState,
|
||||
});
|
||||
const res = await runExecute(tool, { query: 'q', topK: 5 });
|
||||
const total = res.passages.reduce((s, p) => s + p.text.length, 0);
|
||||
expect(total).toBeLessThanOrEqual(6000);
|
||||
expect(res.truncated).toBe(true);
|
||||
// Lowest-ranked passages drop first — p0 (highest rank) should survive.
|
||||
expect(res.passages[0]!.cfi).toContain('epubcfi(');
|
||||
});
|
||||
});
|
||||
|
||||
describe('status passthrough', () => {
|
||||
it.each([
|
||||
'not_indexed',
|
||||
'empty_index',
|
||||
'stale_index',
|
||||
'degraded',
|
||||
] as const)('forwards status=%s with a human-readable hint and empty passages', async (status) => {
|
||||
const retriever = fakeRetriever(async () => ({
|
||||
passages: [],
|
||||
status,
|
||||
reason: status === 'stale_index' ? 'model changed' : undefined,
|
||||
}));
|
||||
const tool = buildLookupTool({
|
||||
bookHash: 'bk1',
|
||||
retriever,
|
||||
activeEmbeddingModel: model,
|
||||
turnState,
|
||||
});
|
||||
const res = await runExecute(tool, { query: 'q', topK: 3 });
|
||||
expect(res.status).toBe(status);
|
||||
expect(res.passages).toEqual([]);
|
||||
expect(res.hint).toBeTruthy();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('serializeForModel (XML envelope)', () => {
|
||||
it('wraps passage text in a <retrieved> envelope with the CFI attribute', () => {
|
||||
const out = serializeForModel({
|
||||
cfi: 'epubcfi(/6/2!/4/2)',
|
||||
chapter: 'Ch1',
|
||||
text: 'plain content',
|
||||
});
|
||||
expect(out).toMatch(/^<retrieved /);
|
||||
expect(out).toContain('cfi="epubcfi(/6/2!/4/2)"');
|
||||
expect(out).toContain('trust="untrusted"');
|
||||
expect(out).toContain('plain content');
|
||||
expect(out).toMatch(/<\/retrieved>$/);
|
||||
});
|
||||
|
||||
it('XML-escapes literal </retrieved>, &, <, > in book text', () => {
|
||||
const out = serializeForModel({
|
||||
cfi: 'epubcfi(/6/2!/4/2)',
|
||||
text: 'evil </retrieved> & <script>alert()</script>',
|
||||
});
|
||||
// Escaped form of the user-supplied close tag must be present...
|
||||
expect(out).toContain('</retrieved>');
|
||||
expect(out).toContain('&');
|
||||
expect(out).toContain('<script>');
|
||||
expect(out).toContain('</script>');
|
||||
// ...and the only literal </retrieved> is the envelope's own closing tag.
|
||||
const closes = out.match(/<\/retrieved>/g) ?? [];
|
||||
expect(closes).toHaveLength(1);
|
||||
});
|
||||
|
||||
it('XML-escapes quotes in the cfi attribute value to keep the envelope well-formed', () => {
|
||||
const out = serializeForModel({
|
||||
cfi: 'epubcfi(/6/2!/4/2[id"with"quote])',
|
||||
text: 'ok',
|
||||
});
|
||||
expect(out).toContain('"');
|
||||
});
|
||||
});
|
||||
@@ -31,6 +31,47 @@ const migrations: Record<SchemaType, MigrationEntry[]> = {
|
||||
`,
|
||||
},
|
||||
],
|
||||
// The embeddings table is created lazily by BookIndexer because its
|
||||
// vector32(<dim>) column needs the active embedding model's dim, which
|
||||
// isn't known at migration time. Tantivy FTS lives on the chunks.text
|
||||
// column directly (no virtual table). Writers MUST DELETE+INSERT chunk
|
||||
// rows rather than UPDATE — Tantivy 0.25→0.26 has a known WASM-only
|
||||
// UPDATE regression (see fts-tests.ts:306 FIXME). MVP indexing is
|
||||
// write-once per book so this is naturally satisfied.
|
||||
reedy: [
|
||||
{
|
||||
name: '2026052601_reedy_init',
|
||||
sql: `
|
||||
CREATE TABLE IF NOT EXISTS reedy_book_meta (
|
||||
book_hash TEXT PRIMARY KEY,
|
||||
indexing_status TEXT NOT NULL,
|
||||
chunk_count INTEGER NOT NULL DEFAULT 0,
|
||||
embedding_model TEXT NOT NULL,
|
||||
embedding_dim INTEGER NOT NULL,
|
||||
indexed_at INTEGER,
|
||||
error TEXT
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS reedy_book_chunks (
|
||||
id TEXT PRIMARY KEY,
|
||||
book_hash TEXT NOT NULL,
|
||||
section_index INTEGER NOT NULL,
|
||||
chapter_title TEXT,
|
||||
start_cfi TEXT NOT NULL,
|
||||
end_cfi TEXT NOT NULL,
|
||||
position_index INTEGER NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
token_count INTEGER NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_chunks_book_position
|
||||
ON reedy_book_chunks (book_hash, position_index);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_chunks_fts
|
||||
ON reedy_book_chunks USING fts (text) WITH (tokenizer = 'ngram');
|
||||
`,
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
export function getMigrations(schema: SchemaType): MigrationEntry[] {
|
||||
|
||||
@@ -0,0 +1,405 @@
|
||||
import type { DatabaseService } from '@/types/database';
|
||||
import type { BookMeta, ChunkRow, EmbeddingRow, IndexingStatus, ScoredChunk } from './types';
|
||||
|
||||
/**
|
||||
* Typed wrapper around a Turso DatabaseService opened against reedy.db.
|
||||
*
|
||||
* MVP scope: single embedding model locked per database lifetime, single
|
||||
* global `reedy_book_chunk_embeddings` table created lazily by the first
|
||||
* `ensureEmbeddingsTable(dim)` call. Multi-model routing lives only in
|
||||
* Appendix A of the plan.
|
||||
*
|
||||
* Multi-row writes go through DatabaseService.batch() per plan §M1.2. Because
|
||||
* batch() takes raw SQL strings (no parameter binding), text values are
|
||||
* inline-quoted via {@link sqlQuote}; this is safe for SQLite which only
|
||||
* honours `''` as an escape sequence inside single-quoted strings.
|
||||
*/
|
||||
export class ReedyDb {
|
||||
/**
|
||||
* Serializes every DB-mutating call that goes through `batch()` because the
|
||||
* underlying Turso connection only allows one BEGIN/COMMIT at a time. This
|
||||
* lets BookIndexer run embedding requests in parallel across books while
|
||||
* the writes themselves still go through one at a time.
|
||||
*/
|
||||
private writeQueue: Promise<unknown> = Promise.resolve();
|
||||
|
||||
constructor(private readonly db: DatabaseService) {}
|
||||
|
||||
private enqueue<T>(fn: () => Promise<T>): Promise<T> {
|
||||
const next = this.writeQueue.then(fn, fn);
|
||||
// Swallow the value on the queue so a failure in one write doesn't
|
||||
// poison every subsequent write, while still letting the caller see it.
|
||||
this.writeQueue = next.catch(() => undefined);
|
||||
return next;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// book meta
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async upsertBookMeta(meta: BookMeta): Promise<void> {
|
||||
await this.enqueue(() =>
|
||||
this.db.execute(
|
||||
`INSERT INTO reedy_book_meta
|
||||
(book_hash, indexing_status, chunk_count, embedding_model, embedding_dim, indexed_at, error)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(book_hash) DO UPDATE SET
|
||||
indexing_status = excluded.indexing_status,
|
||||
chunk_count = excluded.chunk_count,
|
||||
embedding_model = excluded.embedding_model,
|
||||
embedding_dim = excluded.embedding_dim,
|
||||
indexed_at = excluded.indexed_at,
|
||||
error = excluded.error`,
|
||||
[
|
||||
meta.bookHash,
|
||||
meta.indexingStatus,
|
||||
meta.chunkCount,
|
||||
meta.embeddingModel,
|
||||
meta.embeddingDim,
|
||||
meta.indexedAt,
|
||||
meta.error,
|
||||
],
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
async getBookMeta(bookHash: string): Promise<BookMeta | null> {
|
||||
const rows = await this.db.select<{
|
||||
book_hash: string;
|
||||
indexing_status: string;
|
||||
chunk_count: number;
|
||||
embedding_model: string;
|
||||
embedding_dim: number;
|
||||
indexed_at: number | null;
|
||||
error: string | null;
|
||||
}>('SELECT * FROM reedy_book_meta WHERE book_hash = ?', [bookHash]);
|
||||
const row = rows[0];
|
||||
if (!row) return null;
|
||||
return {
|
||||
bookHash: row.book_hash,
|
||||
indexingStatus: row.indexing_status as IndexingStatus,
|
||||
chunkCount: row.chunk_count,
|
||||
embeddingModel: row.embedding_model,
|
||||
embeddingDim: row.embedding_dim,
|
||||
indexedAt: row.indexed_at,
|
||||
error: row.error,
|
||||
};
|
||||
}
|
||||
|
||||
async setIndexingStatus(
|
||||
bookHash: string,
|
||||
status: IndexingStatus,
|
||||
partial?: Partial<Pick<BookMeta, 'chunkCount' | 'indexedAt' | 'error'>>,
|
||||
): Promise<void> {
|
||||
const sets: string[] = ['indexing_status = ?'];
|
||||
const params: unknown[] = [status];
|
||||
if (partial?.chunkCount !== undefined) {
|
||||
sets.push('chunk_count = ?');
|
||||
params.push(partial.chunkCount);
|
||||
}
|
||||
if (partial?.indexedAt !== undefined) {
|
||||
sets.push('indexed_at = ?');
|
||||
params.push(partial.indexedAt);
|
||||
}
|
||||
if ('error' in (partial ?? {})) {
|
||||
sets.push('error = ?');
|
||||
params.push(partial!.error ?? null);
|
||||
}
|
||||
params.push(bookHash);
|
||||
await this.enqueue(() =>
|
||||
this.db.execute(`UPDATE reedy_book_meta SET ${sets.join(', ')} WHERE book_hash = ?`, params),
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// embeddings table lifecycle
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/**
|
||||
* Idempotently create the embeddings table with `vector32(<dim>)`. If the
|
||||
* table exists with a different dim, throws — the MVP locks one embedding
|
||||
* model per database, so a dim mismatch signals a misuse (e.g. the active
|
||||
* model changed but `wipeAllData` wasn't called).
|
||||
*/
|
||||
async ensureEmbeddingsTable(dim: number): Promise<void> {
|
||||
if (!Number.isInteger(dim) || dim <= 0) {
|
||||
throw new Error(`ensureEmbeddingsTable: dim must be a positive integer, got ${dim}`);
|
||||
}
|
||||
// Enqueue the whole check-then-create so concurrent indexers don't race
|
||||
// between the SELECT and the CREATE.
|
||||
await this.enqueue(async () => {
|
||||
const existing = await this.db.select<{ sql: string | null }>(
|
||||
"SELECT sql FROM sqlite_master WHERE type='table' AND name='reedy_book_chunk_embeddings'",
|
||||
);
|
||||
if (existing.length === 0) {
|
||||
// Inline `dim` is a validated positive integer above — safe to interpolate.
|
||||
await this.db.batch([
|
||||
`CREATE TABLE reedy_book_chunk_embeddings (
|
||||
chunk_id TEXT PRIMARY KEY REFERENCES reedy_book_chunks(id) ON DELETE CASCADE,
|
||||
book_hash TEXT NOT NULL,
|
||||
embedding vector32(${dim})
|
||||
)`,
|
||||
'CREATE INDEX idx_embeddings_book ON reedy_book_chunk_embeddings(book_hash)',
|
||||
]);
|
||||
return;
|
||||
}
|
||||
const m = existing[0]!.sql?.match(/vector32\s*\(\s*(\d+)\s*\)/i);
|
||||
const existingDim = m ? parseInt(m[1]!, 10) : NaN;
|
||||
if (existingDim !== dim) {
|
||||
throw new Error(
|
||||
`ensureEmbeddingsTable: dim mismatch — existing table is vector32(${existingDim}), requested vector32(${dim}). ` +
|
||||
`Switching embedding models requires wipeAllData() first.`,
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// bulk writes
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async insertChunks(chunks: ChunkRow[]): Promise<void> {
|
||||
if (chunks.length === 0) return;
|
||||
const stmts = chunks.map(
|
||||
(c) =>
|
||||
`INSERT INTO reedy_book_chunks
|
||||
(id, book_hash, section_index, chapter_title, start_cfi, end_cfi, position_index, text, token_count)
|
||||
VALUES (${sqlQuote(c.id)}, ${sqlQuote(c.bookHash)}, ${c.sectionIndex}, ${sqlQuoteNullable(c.chapterTitle)}, ${sqlQuote(c.startCfi)}, ${sqlQuote(c.endCfi)}, ${c.positionIndex}, ${sqlQuote(c.text)}, ${c.tokenCount})`,
|
||||
);
|
||||
await this.enqueue(() => this.db.batch(stmts));
|
||||
}
|
||||
|
||||
/**
|
||||
* Insert embedding rows. Asserts every row's vector matches the existing
|
||||
* table's dim (queried once from sqlite_master) before issuing SQL.
|
||||
*/
|
||||
async insertEmbeddings(rows: EmbeddingRow[]): Promise<void> {
|
||||
if (rows.length === 0) return;
|
||||
const dimRows = await this.db.select<{ sql: string | null }>(
|
||||
"SELECT sql FROM sqlite_master WHERE type='table' AND name='reedy_book_chunk_embeddings'",
|
||||
);
|
||||
const sql = dimRows[0]?.sql;
|
||||
const m = sql?.match(/vector32\s*\(\s*(\d+)\s*\)/i);
|
||||
if (!m) {
|
||||
throw new Error(
|
||||
'insertEmbeddings: reedy_book_chunk_embeddings does not exist — call ensureEmbeddingsTable(dim) first.',
|
||||
);
|
||||
}
|
||||
const dim = parseInt(m[1]!, 10);
|
||||
for (const r of rows) {
|
||||
if (r.embedding.length !== dim) {
|
||||
throw new Error(
|
||||
`insertEmbeddings: embedding for chunk ${r.chunkId} has length ${r.embedding.length}, expected dim ${dim}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
const stmts = rows.map(
|
||||
(r) =>
|
||||
`INSERT INTO reedy_book_chunk_embeddings (chunk_id, book_hash, embedding)
|
||||
VALUES (${sqlQuote(r.chunkId)}, ${sqlQuote(r.bookHash)}, vector32(${sqlQuote(serializeVector(r.embedding))}))`,
|
||||
);
|
||||
await this.enqueue(() => this.db.batch(stmts));
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// data lifecycle
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/**
|
||||
* Clear a single book's chunks and embeddings but leave the meta row alone.
|
||||
* Used by BookIndexer when re-indexing — the meta row has just been set to
|
||||
* 'indexing' and must be preserved.
|
||||
*/
|
||||
async clearBookChunks(bookHash: string): Promise<void> {
|
||||
await this.enqueue(async () => {
|
||||
// Embeddings reference chunks via ON DELETE CASCADE, but the embeddings
|
||||
// table may not exist yet on first index — guard with sqlite_master.
|
||||
const has = await this.db.select<{ name: string }>(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name='reedy_book_chunk_embeddings'",
|
||||
);
|
||||
if (has.length > 0) {
|
||||
await this.db.execute('DELETE FROM reedy_book_chunk_embeddings WHERE book_hash = ?', [
|
||||
bookHash,
|
||||
]);
|
||||
}
|
||||
await this.db.execute('DELETE FROM reedy_book_chunks WHERE book_hash = ?', [bookHash]);
|
||||
});
|
||||
}
|
||||
|
||||
async dropBookData(bookHash: string): Promise<void> {
|
||||
await this.enqueue(async () => {
|
||||
await this.db.execute('DELETE FROM reedy_book_chunk_embeddings WHERE book_hash = ?', [
|
||||
bookHash,
|
||||
]);
|
||||
await this.db.execute('DELETE FROM reedy_book_chunks WHERE book_hash = ?', [bookHash]);
|
||||
await this.db.execute('DELETE FROM reedy_book_meta WHERE book_hash = ?', [bookHash]);
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Wipe every Reedy-managed row across the database. Used when the user
|
||||
* switches embedding models — the lazy embeddings table keeps its
|
||||
* existing vector32(<old-dim>) shape until something writes to it again,
|
||||
* but the next ensureEmbeddingsTable(<new-dim>) call will succeed
|
||||
* because we DROP the table here as well.
|
||||
*/
|
||||
async wipeAllData(): Promise<void> {
|
||||
await this.enqueue(async () => {
|
||||
await this.db.execute('DELETE FROM reedy_book_meta');
|
||||
// Drop embeddings table so a future ensureEmbeddingsTable(newDim) is free
|
||||
// to recreate it with a different vector32 width.
|
||||
await this.db.execute('DROP TABLE IF EXISTS reedy_book_chunk_embeddings');
|
||||
await this.db.execute('DELETE FROM reedy_book_chunks');
|
||||
});
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// hybrid search (brute-force cosine + Tantivy FTS + RRF)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async hybridSearch(args: {
|
||||
bookHash: string;
|
||||
queryText: string;
|
||||
queryEmbedding: number[];
|
||||
k: number;
|
||||
spoilerBoundPosition?: number;
|
||||
}): Promise<ScoredChunk[]> {
|
||||
const { bookHash, queryText, queryEmbedding, k, spoilerBoundPosition } = args;
|
||||
if (k <= 0) return [];
|
||||
|
||||
const spoilerClause = spoilerBoundPosition !== undefined ? ' AND c.position_index <= ?' : '';
|
||||
const spoilerParam: unknown[] =
|
||||
spoilerBoundPosition !== undefined ? [spoilerBoundPosition] : [];
|
||||
|
||||
// Over-fetch from each path so a chunk surfaced by one path still has a
|
||||
// chance to gain RRF mass from the other. Standard hybrid-search pattern.
|
||||
const fetchK = Math.max(k, k * RRF_FETCH_MULTIPLIER);
|
||||
|
||||
// Vector path — brute-force cosine. Turso has no native vector index
|
||||
// module so this is O(n) per book; sub-ms at MVP corpus sizes (see
|
||||
// bench/vector-retrieval.bench.ts).
|
||||
const vectorRows = await this.db.select<ScoredChunkRowSql>(
|
||||
`SELECT c.id, c.book_hash, c.section_index, c.chapter_title,
|
||||
c.start_cfi, c.end_cfi, c.position_index, c.text, c.token_count,
|
||||
vector_distance_cos(e.embedding, vector32(?)) AS metric
|
||||
FROM reedy_book_chunk_embeddings e
|
||||
JOIN reedy_book_chunks c ON c.id = e.chunk_id
|
||||
WHERE e.book_hash = ?${spoilerClause}
|
||||
ORDER BY metric ASC
|
||||
LIMIT ?`,
|
||||
[serializeVector(queryEmbedding), bookHash, ...spoilerParam, fetchK],
|
||||
);
|
||||
|
||||
// FTS path — Tantivy BM25 over the chunks.text column.
|
||||
let ftsRows: ScoredChunkRowSql[] = [];
|
||||
if (queryText.trim().length > 0) {
|
||||
try {
|
||||
ftsRows = await this.db.select<ScoredChunkRowSql>(
|
||||
`SELECT c.id, c.book_hash, c.section_index, c.chapter_title,
|
||||
c.start_cfi, c.end_cfi, c.position_index, c.text, c.token_count,
|
||||
fts_score(c.text, ?) AS metric
|
||||
FROM reedy_book_chunks c
|
||||
WHERE fts_match(c.text, ?) AND c.book_hash = ?${spoilerClause}
|
||||
ORDER BY metric DESC
|
||||
LIMIT ?`,
|
||||
[queryText, queryText, bookHash, ...spoilerParam, fetchK],
|
||||
);
|
||||
} catch {
|
||||
// FTS index may legitimately be empty (no chunks yet) or the query
|
||||
// may be malformed for Tantivy. The vector path is the primary
|
||||
// signal; FTS is purely a lexical booster.
|
||||
ftsRows = [];
|
||||
}
|
||||
}
|
||||
|
||||
return reciprocalRankFusion(vectorRows, ftsRows, k);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
interface ScoredChunkRowSql {
|
||||
id: string;
|
||||
book_hash: string;
|
||||
section_index: number;
|
||||
chapter_title: string | null;
|
||||
start_cfi: string;
|
||||
end_cfi: string;
|
||||
position_index: number;
|
||||
text: string;
|
||||
token_count: number;
|
||||
metric: number;
|
||||
[key: string]: unknown;
|
||||
}
|
||||
|
||||
function rowToChunk(row: ScoredChunkRowSql): Omit<ScoredChunk, 'score' | 'vectorRank' | 'ftsRank'> {
|
||||
return {
|
||||
id: row.id,
|
||||
bookHash: row.book_hash,
|
||||
sectionIndex: row.section_index,
|
||||
chapterTitle: row.chapter_title,
|
||||
startCfi: row.start_cfi,
|
||||
endCfi: row.end_cfi,
|
||||
positionIndex: row.position_index,
|
||||
text: row.text,
|
||||
tokenCount: row.token_count,
|
||||
};
|
||||
}
|
||||
|
||||
// RRF constant per the Cormack/Clarke/Buettcher paper; dampens single-path
|
||||
// dominance so a chunk surfaced by both lists outranks a chunk near the top
|
||||
// of only one list.
|
||||
const RRF_K = 60;
|
||||
// Per-path over-fetch multiplier so a chunk that won on FTS but tied on
|
||||
// vector (or vice-versa) still gets a vector rank contributing to its
|
||||
// RRF score. Standard 3-5× hybrid-search rule of thumb.
|
||||
const RRF_FETCH_MULTIPLIER = 3;
|
||||
|
||||
function reciprocalRankFusion(
|
||||
vectorRows: ScoredChunkRowSql[],
|
||||
ftsRows: ScoredChunkRowSql[],
|
||||
topK: number,
|
||||
): ScoredChunk[] {
|
||||
const merged = new Map<string, ScoredChunk>();
|
||||
for (let i = 0; i < vectorRows.length; i++) {
|
||||
const row = vectorRows[i]!;
|
||||
const rank = i + 1;
|
||||
merged.set(row.id, {
|
||||
...rowToChunk(row),
|
||||
score: 1 / (RRF_K + rank),
|
||||
vectorRank: rank,
|
||||
ftsRank: null,
|
||||
});
|
||||
}
|
||||
for (let i = 0; i < ftsRows.length; i++) {
|
||||
const row = ftsRows[i]!;
|
||||
const rank = i + 1;
|
||||
const existing = merged.get(row.id);
|
||||
if (existing) {
|
||||
existing.score += 1 / (RRF_K + rank);
|
||||
existing.ftsRank = rank;
|
||||
} else {
|
||||
merged.set(row.id, {
|
||||
...rowToChunk(row),
|
||||
score: 1 / (RRF_K + rank),
|
||||
vectorRank: null,
|
||||
ftsRank: rank,
|
||||
});
|
||||
}
|
||||
}
|
||||
return [...merged.values()].sort((a, b) => b.score - a.score).slice(0, topK);
|
||||
}
|
||||
|
||||
function serializeVector(v: number[]): string {
|
||||
return JSON.stringify(v);
|
||||
}
|
||||
|
||||
function sqlQuote(s: string): string {
|
||||
return `'${s.replace(/'/g, "''")}'`;
|
||||
}
|
||||
|
||||
function sqlQuoteNullable(s: string | null): string {
|
||||
return s === null ? 'NULL' : sqlQuote(s);
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
/**
|
||||
* Public types for the Reedy retrieval layer. Kept narrow on purpose — the
|
||||
* MVP locks one embedding model per database lifetime, so we don't need
|
||||
* per-model routing types here. See plan §M1.2 and Appendix A for the
|
||||
* deferred multi-model story.
|
||||
*/
|
||||
|
||||
export type IndexingStatus = 'not_indexed' | 'indexing' | 'indexed' | 'failed' | 'empty_index';
|
||||
|
||||
export interface BookMeta {
|
||||
bookHash: string;
|
||||
indexingStatus: IndexingStatus;
|
||||
chunkCount: number;
|
||||
embeddingModel: string;
|
||||
embeddingDim: number;
|
||||
indexedAt: number | null;
|
||||
error: string | null;
|
||||
}
|
||||
|
||||
export interface ChunkRow {
|
||||
id: string;
|
||||
bookHash: string;
|
||||
sectionIndex: number;
|
||||
chapterTitle: string | null;
|
||||
startCfi: string;
|
||||
endCfi: string;
|
||||
positionIndex: number;
|
||||
text: string;
|
||||
tokenCount: number;
|
||||
}
|
||||
|
||||
export interface EmbeddingRow {
|
||||
chunkId: string;
|
||||
bookHash: string;
|
||||
embedding: number[];
|
||||
}
|
||||
|
||||
/**
|
||||
* A chunk returned by hybridSearch, annotated with the RRF-fused score and
|
||||
* which retrieval paths surfaced it. Per-path ranks are 1-indexed; `null`
|
||||
* means the path didn't surface this chunk in its top-K.
|
||||
*/
|
||||
export interface ScoredChunk extends ChunkRow {
|
||||
score: number;
|
||||
vectorRank: number | null;
|
||||
ftsRank: number | null;
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
/**
|
||||
* Minimal embedding-model interface the Reedy retrieval layer talks to.
|
||||
*
|
||||
* MVP scope: just enough surface for BookIndexer and BookRetriever to drive
|
||||
* indexing + query embedding. The actual provider plumbing (Ollama,
|
||||
* AIGateway, OpenRouter, ...) lives in `src/services/ai/` and is adapted to
|
||||
* this interface by the M1.7 ReedyBackend so we don't have to re-implement
|
||||
* provider transports here.
|
||||
*/
|
||||
export interface EmbeddingModel {
|
||||
/** Stable identifier — matches the `embedding_model` column in reedy_book_meta. */
|
||||
readonly id: string;
|
||||
/** Vector width. Must match the `vector32(<dim>)` column once the lazy embeddings table exists. */
|
||||
readonly dim: number;
|
||||
/**
|
||||
* Batch size hint for indexing. Ollama and local engines typically prefer
|
||||
* small batches (4); hosted providers (AIGateway, OpenAI) accept larger
|
||||
* batches (16+). BookIndexer respects this; embedding-time backpressure is
|
||||
* the model's responsibility.
|
||||
*/
|
||||
readonly batchSize?: number;
|
||||
embed(texts: string[], opts?: { signal?: AbortSignal }): Promise<number[][]>;
|
||||
}
|
||||
@@ -0,0 +1,179 @@
|
||||
import type { BookDoc } from '@/libs/document';
|
||||
import { ReedyDb } from '../db/ReedyDb';
|
||||
import type { ChunkRow, EmbeddingRow } from '../db/types';
|
||||
import type { EmbeddingModel } from '../models/EmbeddingModel';
|
||||
import { chunkSection, type ChunkOptions } from './CfiChunker';
|
||||
|
||||
const DEFAULT_BATCH_SIZE = 16;
|
||||
|
||||
export interface IndexBookOptions {
|
||||
/** Override CfiChunker tuning. Falls back to chunker defaults. */
|
||||
chunkOptions?: Partial<ChunkOptions>;
|
||||
/** Optional callback for progress reporting. Phases: 'chunking' | 'embedding'. */
|
||||
onProgress?: (event: { phase: 'chunking' | 'embedding'; current: number; total: number }) => void;
|
||||
/** Optional chapter-title resolver; defaults to `Section ${i + 1}`. */
|
||||
getChapterTitle?: (sectionIndex: number) => string | null;
|
||||
/** AbortSignal honoured by the embedding model. */
|
||||
signal?: AbortSignal;
|
||||
}
|
||||
|
||||
/**
|
||||
* Orchestrates one book's indexing pipeline:
|
||||
* 1. mutex per book so concurrent calls serialize
|
||||
* 2. chunk every section via CfiChunker
|
||||
* 3. lazy-create the embeddings table at the active model's dim
|
||||
* 4. embed in model-sized batches and insert
|
||||
* 5. land a terminal status (indexed | empty_index | failed) on reedy_book_meta
|
||||
*
|
||||
* The caller decides when to call this (settings panel "Index this book"
|
||||
* button, library-import hook, etc). Failures throw — the caller surfaces
|
||||
* the error to the user — but the meta row is updated to 'failed' first so
|
||||
* subsequent BookRetriever calls return a useful status.
|
||||
*/
|
||||
export class BookIndexer {
|
||||
private readonly inflight = new Map<string, Promise<void>>();
|
||||
|
||||
constructor(private readonly reedy: ReedyDb) {}
|
||||
|
||||
async indexBook(
|
||||
bookDoc: BookDoc,
|
||||
bookHash: string,
|
||||
model: EmbeddingModel,
|
||||
options: IndexBookOptions = {},
|
||||
): Promise<void> {
|
||||
// Chain on whatever's already in flight for this book so concurrent
|
||||
// callers serialize. We register the chained promise synchronously
|
||||
// (before any await) so a second concurrent caller sees the chain even
|
||||
// if no prior run had been registered when we entered.
|
||||
const prior = this.inflight.get(bookHash);
|
||||
const promise = (prior ? prior.catch(() => undefined) : Promise.resolve()).then(() =>
|
||||
this.runIndex(bookDoc, bookHash, model, options),
|
||||
);
|
||||
this.inflight.set(bookHash, promise);
|
||||
try {
|
||||
await promise;
|
||||
} finally {
|
||||
// Only clear if this is still the tail of the chain — a subsequent
|
||||
// caller may have appended after us and we shouldn't drop that.
|
||||
if (this.inflight.get(bookHash) === promise) {
|
||||
this.inflight.delete(bookHash);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private async runIndex(
|
||||
bookDoc: BookDoc,
|
||||
bookHash: string,
|
||||
model: EmbeddingModel,
|
||||
options: IndexBookOptions,
|
||||
): Promise<void> {
|
||||
await this.reedy.upsertBookMeta({
|
||||
bookHash,
|
||||
indexingStatus: 'indexing',
|
||||
chunkCount: 0,
|
||||
embeddingModel: model.id,
|
||||
embeddingDim: model.dim,
|
||||
indexedAt: null,
|
||||
error: null,
|
||||
});
|
||||
// Re-indexing must replace, not duplicate — drop any prior chunks +
|
||||
// embeddings for this book before writing the new ones. The meta row
|
||||
// upserted above is preserved (clearBookChunks only touches the chunk
|
||||
// and embedding tables).
|
||||
await this.reedy.clearBookChunks(bookHash);
|
||||
|
||||
try {
|
||||
const chunks = await this.collectChunks(bookDoc, bookHash, options);
|
||||
|
||||
if (chunks.length === 0) {
|
||||
await this.reedy.setIndexingStatus(bookHash, 'empty_index', {
|
||||
chunkCount: 0,
|
||||
indexedAt: Date.now(),
|
||||
error: null,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
await this.reedy.ensureEmbeddingsTable(model.dim);
|
||||
await this.reedy.insertChunks(chunks);
|
||||
|
||||
await this.embedAndStore(chunks, model, options);
|
||||
|
||||
await this.reedy.setIndexingStatus(bookHash, 'indexed', {
|
||||
chunkCount: chunks.length,
|
||||
indexedAt: Date.now(),
|
||||
error: null,
|
||||
});
|
||||
} catch (err) {
|
||||
const message = err instanceof Error ? err.message : String(err);
|
||||
await this.reedy.setIndexingStatus(bookHash, 'failed', { error: message });
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
private async collectChunks(
|
||||
bookDoc: BookDoc,
|
||||
bookHash: string,
|
||||
options: IndexBookOptions,
|
||||
): Promise<ChunkRow[]> {
|
||||
const all: ChunkRow[] = [];
|
||||
const sections = bookDoc.sections;
|
||||
for (let i = 0; i < sections.length; i++) {
|
||||
options.onProgress?.({ phase: 'chunking', current: i, total: sections.length });
|
||||
const section = sections[i]!;
|
||||
let doc: Document;
|
||||
try {
|
||||
doc = await section.createDocument();
|
||||
} catch (err) {
|
||||
console.warn('[Reedy] section createDocument failed', { sectionIndex: i, err });
|
||||
continue;
|
||||
}
|
||||
const title = options.getChapterTitle?.(i) ?? `Section ${i + 1}`;
|
||||
const sectionChunks = chunkSection(doc, i, title, bookHash, options.chunkOptions);
|
||||
// Rewrite the position index to be monotonic across the whole book —
|
||||
// CfiChunker numbers within a section, the indexer needs a global order.
|
||||
for (const c of sectionChunks) {
|
||||
all.push({ ...c, positionIndex: all.length, id: `${bookHash}-${all.length}` });
|
||||
}
|
||||
}
|
||||
options.onProgress?.({ phase: 'chunking', current: sections.length, total: sections.length });
|
||||
return all;
|
||||
}
|
||||
|
||||
private async embedAndStore(
|
||||
chunks: ChunkRow[],
|
||||
model: EmbeddingModel,
|
||||
options: IndexBookOptions,
|
||||
): Promise<void> {
|
||||
const batchSize = Math.max(1, model.batchSize ?? DEFAULT_BATCH_SIZE);
|
||||
const total = chunks.length;
|
||||
let done = 0;
|
||||
for (let i = 0; i < total; i += batchSize) {
|
||||
if (options.signal?.aborted) {
|
||||
throw new Error('indexing aborted');
|
||||
}
|
||||
const batch = chunks.slice(i, i + batchSize);
|
||||
const vectors = await model.embed(
|
||||
batch.map((c) => c.text),
|
||||
{ signal: options.signal },
|
||||
);
|
||||
if (vectors.length !== batch.length) {
|
||||
throw new Error(
|
||||
`embedding model returned ${vectors.length} vectors for ${batch.length} inputs`,
|
||||
);
|
||||
}
|
||||
const rows: EmbeddingRow[] = batch.map((c, j) => {
|
||||
const v = vectors[j]!;
|
||||
if (v.length !== model.dim) {
|
||||
throw new Error(
|
||||
`embedding for chunk ${c.id} has length ${v.length}, expected dim ${model.dim}`,
|
||||
);
|
||||
}
|
||||
return { chunkId: c.id, bookHash: c.bookHash, embedding: v };
|
||||
});
|
||||
await this.reedy.insertEmbeddings(rows);
|
||||
done += batch.length;
|
||||
options.onProgress?.({ phase: 'embedding', current: done, total });
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,153 @@
|
||||
import type { ReedyDb } from '../db/ReedyDb';
|
||||
import type { EmbeddingModel } from '../models/EmbeddingModel';
|
||||
|
||||
/**
|
||||
* Status the retriever reports back so the lookupPassage tool can phrase the
|
||||
* model's response, and so the UI can offer an appropriate next action
|
||||
* (e.g. "Index this book", "Re-index with new model"). `budget_exceeded` is
|
||||
* intentionally NOT part of this union — that status is surfaced by the
|
||||
* lookupPassage tool layer when it refuses to call the retriever again.
|
||||
*/
|
||||
export type RetrieverStatus = 'ok' | 'not_indexed' | 'empty_index' | 'stale_index' | 'degraded';
|
||||
|
||||
export interface RetrievedChunk {
|
||||
id: string;
|
||||
bookHash: string;
|
||||
/** The chunk's start CFI — the navigable anchor handed to the UI. */
|
||||
cfi: string;
|
||||
/** End CFI, useful for highlighting or future tool operations. */
|
||||
endCfi: string;
|
||||
chapterTitle: string | null;
|
||||
text: string;
|
||||
positionIndex: number;
|
||||
/** Fused RRF score; informational only. */
|
||||
score: number;
|
||||
}
|
||||
|
||||
export interface RetrieverResult {
|
||||
passages: RetrievedChunk[];
|
||||
status: RetrieverStatus;
|
||||
/** Human-readable reason for non-`ok` statuses; surfaced to the user via the model. */
|
||||
reason?: string;
|
||||
}
|
||||
|
||||
export interface RetrieveArgs {
|
||||
bookHash: string;
|
||||
query: string;
|
||||
k: number;
|
||||
spoilerBoundPosition?: number;
|
||||
activeEmbeddingModel: EmbeddingModel;
|
||||
/** Query embedding wall-clock budget. @default 5000 */
|
||||
embeddingTimeoutMs?: number;
|
||||
}
|
||||
|
||||
const DEFAULT_EMBEDDING_TIMEOUT_MS = 5000;
|
||||
|
||||
/**
|
||||
* Per plan §M1.5 — wraps ReedyDb.hybridSearch with status detection and
|
||||
* graceful degradation. The retriever:
|
||||
*
|
||||
* 1. checks reedy_book_meta → returns `not_indexed` / `empty_index` /
|
||||
* `stale_index` without touching the chunks/embeddings tables;
|
||||
* 2. embeds the user's query with a wall-clock budget; on timeout it
|
||||
* reports `degraded` and falls through to FTS-only fusion;
|
||||
* 3. calls hybridSearch (vector cosine + Tantivy FTS + RRF) filtered by
|
||||
* bookHash and spoilerBoundPosition;
|
||||
* 4. shapes ScoredChunk → RetrievedChunk for the tool layer.
|
||||
*/
|
||||
export class BookRetriever {
|
||||
constructor(private readonly reedy: ReedyDb) {}
|
||||
|
||||
async search(args: RetrieveArgs): Promise<RetrieverResult> {
|
||||
const meta = await this.reedy.getBookMeta(args.bookHash);
|
||||
if (!meta) {
|
||||
return { passages: [], status: 'not_indexed' };
|
||||
}
|
||||
// Check indexing_status FIRST so 'indexing' / 'failed' rows (which start
|
||||
// life with chunk_count=0) report not_indexed instead of empty_index.
|
||||
if (meta.indexingStatus === 'empty_index') {
|
||||
return { passages: [], status: 'empty_index' };
|
||||
}
|
||||
if (meta.indexingStatus !== 'indexed') {
|
||||
// 'indexing' / 'failed' fall here — no usable corpus yet.
|
||||
return { passages: [], status: 'not_indexed' };
|
||||
}
|
||||
if (meta.chunkCount === 0) {
|
||||
// Indexed but zero chunks — same shape as empty_index.
|
||||
return { passages: [], status: 'empty_index' };
|
||||
}
|
||||
if (meta.embeddingModel !== args.activeEmbeddingModel.id) {
|
||||
return {
|
||||
passages: [],
|
||||
status: 'stale_index',
|
||||
reason: `${args.activeEmbeddingModel.id} is selected but this book was indexed with ${meta.embeddingModel}; re-index required`,
|
||||
};
|
||||
}
|
||||
|
||||
const timeoutMs = args.embeddingTimeoutMs ?? DEFAULT_EMBEDDING_TIMEOUT_MS;
|
||||
const { embedding, degraded, reason } = await embedQueryWithTimeout(
|
||||
args.activeEmbeddingModel,
|
||||
args.query,
|
||||
timeoutMs,
|
||||
);
|
||||
|
||||
const scored = await this.reedy.hybridSearch({
|
||||
bookHash: args.bookHash,
|
||||
queryText: args.query,
|
||||
// When the embedding times out we fall back to FTS-only by passing a
|
||||
// zero vector — vector_distance_cos will produce uniform distances and
|
||||
// contribute nothing useful to the RRF; FTS still ranks meaningfully.
|
||||
queryEmbedding: embedding ?? new Array(args.activeEmbeddingModel.dim).fill(0),
|
||||
k: args.k,
|
||||
spoilerBoundPosition: args.spoilerBoundPosition,
|
||||
});
|
||||
|
||||
const passages: RetrievedChunk[] = scored.map((s) => ({
|
||||
id: s.id,
|
||||
bookHash: s.bookHash,
|
||||
cfi: s.startCfi,
|
||||
endCfi: s.endCfi,
|
||||
chapterTitle: s.chapterTitle,
|
||||
text: s.text,
|
||||
positionIndex: s.positionIndex,
|
||||
score: s.score,
|
||||
}));
|
||||
|
||||
if (degraded) {
|
||||
return { passages, status: 'degraded', reason };
|
||||
}
|
||||
return { passages, status: 'ok' };
|
||||
}
|
||||
}
|
||||
|
||||
async function embedQueryWithTimeout(
|
||||
model: EmbeddingModel,
|
||||
query: string,
|
||||
timeoutMs: number,
|
||||
): Promise<{ embedding: number[] | null; degraded: boolean; reason?: string }> {
|
||||
const controller = new AbortController();
|
||||
const timer = setTimeout(() => controller.abort(), timeoutMs);
|
||||
try {
|
||||
const result = await model.embed([query], { signal: controller.signal });
|
||||
const v = result[0];
|
||||
if (!v || v.length !== model.dim) {
|
||||
return {
|
||||
embedding: null,
|
||||
degraded: true,
|
||||
reason: `embedding_dim_mismatch: model returned ${v?.length ?? 'no'} values, expected ${model.dim}`,
|
||||
};
|
||||
}
|
||||
return { embedding: v, degraded: false };
|
||||
} catch (err) {
|
||||
const message = err instanceof Error ? err.message : String(err);
|
||||
return {
|
||||
embedding: null,
|
||||
degraded: true,
|
||||
reason: controller.signal.aborted
|
||||
? `embedding_timeout after ${timeoutMs}ms`
|
||||
: `embedding_failed: ${message}`,
|
||||
};
|
||||
} finally {
|
||||
clearTimeout(timer);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,276 @@
|
||||
import * as CFI from 'foliate-js/epubcfi.js';
|
||||
import type { ChunkRow } from '../db/types';
|
||||
|
||||
/**
|
||||
* CFI-aware chunker. Walks an EPUB section's DOM via TreeWalker, accumulates
|
||||
* text from <body>, and slices it into ~maxChunkSize windows with paragraph >
|
||||
* sentence > word break-points. Each chunk carries the full epubcfi(/6/N!/…)
|
||||
* range for its first and last character positions so the retriever can hand
|
||||
* back navigable anchors.
|
||||
*
|
||||
* MVP scope (per plan §M1.3):
|
||||
* - Plain text only; image-only sections produce zero chunks (callers handle
|
||||
* that via the BookIndexer's `empty_index` status).
|
||||
* - Skips `<script>`, `<style>`, `<noscript>` and any node marked with the
|
||||
* `cfi-inert` class.
|
||||
* - Verifies every generated CFI round-trips via CFI.toRange against the
|
||||
* section document. Mismatches are dropped with a console warning rather
|
||||
* than silently writing bad data.
|
||||
*/
|
||||
|
||||
export interface ChunkOptions {
|
||||
/** Target chunk size in characters. */
|
||||
maxChunkSize: number;
|
||||
/** Minimum acceptable chunk size; smaller tails are merged into the prior chunk. */
|
||||
minChunkSize: number;
|
||||
/** Characters of overlap between adjacent chunks (re-emitted from the end of the prior chunk). */
|
||||
overlapSize: number;
|
||||
/** Maximum chars to search left/right of the target boundary for a break-point. */
|
||||
breakSearchRange: number;
|
||||
}
|
||||
|
||||
const DEFAULT_OPTIONS: ChunkOptions = {
|
||||
maxChunkSize: 500,
|
||||
minChunkSize: 100,
|
||||
overlapSize: 50,
|
||||
breakSearchRange: 50,
|
||||
};
|
||||
|
||||
const SKIP_TAGS = new Set(['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEMPLATE']);
|
||||
|
||||
interface TextSlice {
|
||||
node: Text;
|
||||
/** Cumulative character offset of this text node's first char within the section's flat string. */
|
||||
cumStart: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Walk the document body collecting text nodes alongside their position in
|
||||
* the flat concatenated string we use for break-point detection.
|
||||
*/
|
||||
function collectTextNodes(doc: Document): { slices: TextSlice[]; flatText: string } {
|
||||
const body = doc.body ?? doc.documentElement;
|
||||
if (!body) return { slices: [], flatText: '' };
|
||||
|
||||
const walker = doc.createTreeWalker(body, NodeFilter.SHOW_TEXT, {
|
||||
acceptNode(node) {
|
||||
let p: Node | null = node.parentNode;
|
||||
while (p && p.nodeType === 1) {
|
||||
const el = p as Element;
|
||||
if (SKIP_TAGS.has(el.tagName)) return NodeFilter.FILTER_REJECT;
|
||||
if (el.classList?.contains('cfi-inert')) return NodeFilter.FILTER_REJECT;
|
||||
p = p.parentNode;
|
||||
}
|
||||
return (node.nodeValue ?? '').length > 0 ? NodeFilter.FILTER_ACCEPT : NodeFilter.FILTER_SKIP;
|
||||
},
|
||||
});
|
||||
|
||||
const slices: TextSlice[] = [];
|
||||
const parts: string[] = [];
|
||||
let cum = 0;
|
||||
let n: Node | null = walker.nextNode();
|
||||
while (n) {
|
||||
const text = (n as Text).nodeValue ?? '';
|
||||
slices.push({ node: n as Text, cumStart: cum });
|
||||
parts.push(text);
|
||||
cum += text.length;
|
||||
n = walker.nextNode();
|
||||
}
|
||||
return { slices, flatText: parts.join('') };
|
||||
}
|
||||
|
||||
/**
|
||||
* Map a cumulative character offset to a (text node, offset-within-node) pair.
|
||||
* Caller guarantees `0 <= offset <= flatText.length`.
|
||||
*/
|
||||
function offsetToNode(slices: TextSlice[], offset: number): { node: Text; offset: number } | null {
|
||||
if (slices.length === 0) return null;
|
||||
// Binary search for the slice whose cumStart <= offset < next cumStart.
|
||||
let lo = 0;
|
||||
let hi = slices.length - 1;
|
||||
while (lo < hi) {
|
||||
const mid = (lo + hi + 1) >>> 1;
|
||||
if (slices[mid]!.cumStart <= offset) lo = mid;
|
||||
else hi = mid - 1;
|
||||
}
|
||||
const slice = slices[lo]!;
|
||||
const within = offset - slice.cumStart;
|
||||
const nodeLen = (slice.node.nodeValue ?? '').length;
|
||||
// Clamp to node length so the very last position resolves to end of last node.
|
||||
return { node: slice.node, offset: Math.min(within, nodeLen) };
|
||||
}
|
||||
|
||||
function findBreakPoint(text: string, targetPos: number, searchRange: number): number {
|
||||
const start = Math.max(0, targetPos - searchRange);
|
||||
const end = Math.min(text.length, targetPos + searchRange);
|
||||
const window = text.slice(start, end);
|
||||
|
||||
// Prefer paragraph break, then sentence terminator + space, then word break.
|
||||
const paragraphBreak = window.lastIndexOf('\n\n');
|
||||
if (paragraphBreak !== -1 && paragraphBreak > searchRange / 2) {
|
||||
return start + paragraphBreak + 2;
|
||||
}
|
||||
const sentenceBreak = window.lastIndexOf('. ');
|
||||
if (sentenceBreak !== -1 && sentenceBreak > searchRange / 2) {
|
||||
return start + sentenceBreak + 2;
|
||||
}
|
||||
const wordBreak = window.lastIndexOf(' ');
|
||||
if (wordBreak !== -1) {
|
||||
return start + wordBreak + 1;
|
||||
}
|
||||
return targetPos;
|
||||
}
|
||||
|
||||
function composeSectionCfi(innerCfiWrapped: string, sectionIndex: number): string {
|
||||
// fromRange returns "epubcfi(/4/2[p1],/1:0,/1:5)" — unwrap, then prepend
|
||||
// the spine itemref step "/6/{(sectionIndex+1)*2}!" matching the pattern
|
||||
// foliate-js uses for full document CFIs (see src/utils/xcfi.ts).
|
||||
const m = innerCfiWrapped.match(/^epubcfi\((.+)\)$/);
|
||||
if (!m) return innerCfiWrapped;
|
||||
const spineStep = (sectionIndex + 1) * 2;
|
||||
return `epubcfi(/6/${spineStep}!${m[1]!})`;
|
||||
}
|
||||
|
||||
function tokenCount(text: string): number {
|
||||
const trimmed = text.trim();
|
||||
if (trimmed.length === 0) return 0;
|
||||
return trimmed.split(/\s+/).length;
|
||||
}
|
||||
|
||||
export function chunkSection(
|
||||
doc: Document,
|
||||
sectionIndex: number,
|
||||
chapterTitle: string,
|
||||
bookHash: string,
|
||||
options?: Partial<ChunkOptions>,
|
||||
): ChunkRow[] {
|
||||
const opts: ChunkOptions = { ...DEFAULT_OPTIONS, ...options };
|
||||
const { slices, flatText } = collectTextNodes(doc);
|
||||
if (flatText.trim().length === 0 || slices.length === 0) return [];
|
||||
|
||||
const totalLen = flatText.length;
|
||||
// Below the minimum chunk size, emit the whole section as one chunk so very
|
||||
// short sections (a single paragraph, a back-cover blurb) still get indexed.
|
||||
if (totalLen < opts.minChunkSize) {
|
||||
return buildChunks(
|
||||
[{ start: 0, end: totalLen }],
|
||||
flatText,
|
||||
slices,
|
||||
doc,
|
||||
sectionIndex,
|
||||
chapterTitle,
|
||||
bookHash,
|
||||
);
|
||||
}
|
||||
|
||||
const windows: Array<{ start: number; end: number }> = [];
|
||||
let cursor = 0;
|
||||
while (cursor < totalLen) {
|
||||
const targetEnd = cursor + opts.maxChunkSize;
|
||||
if (targetEnd >= totalLen) {
|
||||
windows.push({ start: cursor, end: totalLen });
|
||||
break;
|
||||
}
|
||||
const snappedEnd = findBreakPoint(flatText, targetEnd, opts.breakSearchRange);
|
||||
// Guarantee forward progress even if the breakpoint search returns <= cursor.
|
||||
const end = snappedEnd > cursor ? snappedEnd : Math.min(totalLen, cursor + opts.maxChunkSize);
|
||||
windows.push({ start: cursor, end });
|
||||
cursor = end > opts.overlapSize ? end - opts.overlapSize : end;
|
||||
if (cursor >= totalLen) break;
|
||||
}
|
||||
|
||||
return buildChunks(windows, flatText, slices, doc, sectionIndex, chapterTitle, bookHash);
|
||||
}
|
||||
|
||||
function buildChunks(
|
||||
windows: Array<{ start: number; end: number }>,
|
||||
flatText: string,
|
||||
slices: TextSlice[],
|
||||
doc: Document,
|
||||
sectionIndex: number,
|
||||
chapterTitle: string,
|
||||
bookHash: string,
|
||||
): ChunkRow[] {
|
||||
const out: ChunkRow[] = [];
|
||||
let position = 0;
|
||||
for (const w of windows) {
|
||||
const sliceText = flatText.slice(w.start, w.end).trim();
|
||||
if (sliceText.length === 0) continue;
|
||||
const startPair = offsetToNode(slices, w.start);
|
||||
// For the end position we want the END of the chunk character, not the
|
||||
// start, so step one past the last char (clamped to total length).
|
||||
const endPair = offsetToNode(slices, Math.min(flatText.length, w.end));
|
||||
if (!startPair || !endPair) continue;
|
||||
|
||||
let range: Range;
|
||||
try {
|
||||
range = doc.createRange();
|
||||
range.setStart(startPair.node, startPair.offset);
|
||||
range.setEnd(endPair.node, endPair.offset);
|
||||
} catch (err) {
|
||||
console.warn('[Reedy] chunk_cfi_mismatch: failed to build range', err);
|
||||
continue;
|
||||
}
|
||||
|
||||
let startInner: string;
|
||||
let endInner: string;
|
||||
try {
|
||||
const startCollapsed = doc.createRange();
|
||||
startCollapsed.setStart(startPair.node, startPair.offset);
|
||||
startCollapsed.collapse(true);
|
||||
const endCollapsed = doc.createRange();
|
||||
endCollapsed.setStart(endPair.node, endPair.offset);
|
||||
endCollapsed.collapse(true);
|
||||
startInner = CFI.fromRange(startCollapsed);
|
||||
endInner = CFI.fromRange(endCollapsed);
|
||||
} catch (err) {
|
||||
console.warn('[Reedy] chunk_cfi_mismatch: fromRange threw', err);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Round-trip verification: parsing the generated CFI must resolve to a
|
||||
// range whose start position equals the original position. We don't
|
||||
// require the resolved text to match exactly because toRange of a
|
||||
// collapsed CFI returns a zero-length range — we just need a valid node
|
||||
// reference.
|
||||
if (!verifyRoundTrip(doc, startInner, startPair) || !verifyRoundTrip(doc, endInner, endPair)) {
|
||||
console.warn('[Reedy] chunk_cfi_mismatch: CFI failed round-trip verification', {
|
||||
sectionIndex,
|
||||
position,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
out.push({
|
||||
id: `${bookHash}-${sectionIndex}-${position}`,
|
||||
bookHash,
|
||||
sectionIndex,
|
||||
chapterTitle,
|
||||
startCfi: composeSectionCfi(startInner, sectionIndex),
|
||||
endCfi: composeSectionCfi(endInner, sectionIndex),
|
||||
positionIndex: position,
|
||||
text: sliceText,
|
||||
tokenCount: tokenCount(sliceText),
|
||||
});
|
||||
position++;
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
function verifyRoundTrip(
|
||||
doc: Document,
|
||||
innerCfiWrapped: string,
|
||||
expected: { node: Text; offset: number },
|
||||
): boolean {
|
||||
try {
|
||||
const parts = CFI.parse(innerCfiWrapped);
|
||||
const resolved = CFI.toRange(doc, parts);
|
||||
if (!resolved) return false;
|
||||
// Loose match: same text node, offset within ±1 (CFI normalization can
|
||||
// collapse a zero-length character difference at node boundaries).
|
||||
if (resolved.startContainer !== expected.node) return false;
|
||||
return Math.abs(resolved.startOffset - expected.offset) <= 1;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,258 @@
|
||||
import { tool } from 'ai';
|
||||
import { z } from 'zod';
|
||||
import type { BookRetriever, RetrieverStatus } from '../retrieval/BookRetriever';
|
||||
import type { EmbeddingModel } from '../models/EmbeddingModel';
|
||||
|
||||
/**
|
||||
* Statuses the tool may return to the model. Mirrors RetrieverStatus from
|
||||
* BookRetriever plus the tool-only `budget_exceeded` flag that fires when
|
||||
* the assistant turn has already spent its per-turn retrieval wall-clock.
|
||||
*/
|
||||
export type LookupToolStatus = RetrieverStatus | 'budget_exceeded';
|
||||
|
||||
export interface LookupPassage {
|
||||
cfi: string;
|
||||
endCfi: string;
|
||||
chapter?: string;
|
||||
text: string;
|
||||
}
|
||||
|
||||
export interface LookupToolResult {
|
||||
passages: LookupPassage[];
|
||||
status: LookupToolStatus;
|
||||
/** True when the result came from this turn's dedupe cache. */
|
||||
cached?: boolean;
|
||||
/** True when result-size clamping dropped passages to stay under 6000 chars. */
|
||||
truncated?: boolean;
|
||||
/** Human-readable next-step hint for non-`ok` statuses. */
|
||||
hint?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Per-turn state shared across every lookupPassage invocation in one assistant
|
||||
* turn. Holds (a) the dedupe cache keyed on the composite request shape, and
|
||||
* (b) the parallel-call serialization chain so concurrent tool dispatches
|
||||
* mutate `totalToolMs` and `cache` in a consistent order.
|
||||
*/
|
||||
export interface LookupTurnState {
|
||||
totalToolMs: number;
|
||||
cache: Map<string, LookupToolResult>;
|
||||
pendingChain: Promise<void>;
|
||||
}
|
||||
|
||||
export function createTurnState(): LookupTurnState {
|
||||
return { totalToolMs: 0, cache: new Map(), pendingChain: Promise.resolve() };
|
||||
}
|
||||
|
||||
const MAX_QUERY_CHARS = 500;
|
||||
const MAX_TOP_K = 5;
|
||||
const PER_TURN_BUDGET_MS = 10_000;
|
||||
const RESULT_SIZE_CAP_CHARS = 6_000;
|
||||
|
||||
/**
|
||||
* Exported so tests and prospective callers can pre-validate input without
|
||||
* going through the Tool wrapper (the Tool's `inputSchema` becomes a
|
||||
* provider-utils FlexibleSchema with no `.safeParse`).
|
||||
*/
|
||||
export const lookupInputSchema = z.object({
|
||||
query: z.string().min(1).max(MAX_QUERY_CHARS),
|
||||
topK: z.number().int().min(1).max(MAX_TOP_K).default(MAX_TOP_K),
|
||||
});
|
||||
|
||||
export interface BuildLookupToolArgs {
|
||||
bookHash: string;
|
||||
retriever: BookRetriever;
|
||||
activeEmbeddingModel: EmbeddingModel;
|
||||
turnState: LookupTurnState;
|
||||
/** Optional position cap for spoiler-free retrieval. */
|
||||
spoilerBoundPosition?: number;
|
||||
/**
|
||||
* Optional sink for telemetry — wired by M1.9 to record `tool_called`,
|
||||
* `tool_returned_empty`, etc. Kept as a callback so the tool factory has
|
||||
* no direct dependency on the metrics module.
|
||||
*/
|
||||
onEvent?: (event: { type: string; payload?: Record<string, unknown> }) => void;
|
||||
}
|
||||
|
||||
/**
|
||||
* Construct the Vercel `ai`-SDK Tool factory used by ReedyBackend (M1.7).
|
||||
*
|
||||
* Behaviour mandated by plan §M1.6:
|
||||
* - Zod-validated input (`{ query, topK }`).
|
||||
* - Per-turn dedupe via a composite key over query + topK + spoiler +
|
||||
* active model id.
|
||||
* - Parallel-call serialization so concurrent tool dispatches mutate
|
||||
* shared state in order.
|
||||
* - 10s per-turn wall-clock budget; over-budget calls short-circuit with
|
||||
* `status: 'budget_exceeded'` so the model finalizes its answer.
|
||||
* - Result-size clamp at 6000 chars; lowest-ranked passages drop first.
|
||||
* - Status passthrough (`not_indexed`, `empty_index`, `stale_index`,
|
||||
* `degraded`) with human-readable hints the model can repeat verbatim.
|
||||
*
|
||||
* Trust markers (XML envelope + escape) are produced by `serializeForModel`,
|
||||
* not the tool itself — the tool returns the structured result; the M1.7
|
||||
* prompt builder wraps each passage at the system-message boundary.
|
||||
*/
|
||||
export function buildLookupTool(args: BuildLookupToolArgs) {
|
||||
const { bookHash, retriever, activeEmbeddingModel, turnState, spoilerBoundPosition, onEvent } =
|
||||
args;
|
||||
|
||||
return tool({
|
||||
description:
|
||||
"Look up passages from the user's currently open book by semantic + lexical search. " +
|
||||
'Returns up to topK passages with CFI anchors the UI uses to navigate. ' +
|
||||
'Call this whenever the user asks about the book content. ' +
|
||||
"If status != 'ok', use the hint to phrase the user-visible reply.",
|
||||
inputSchema: lookupInputSchema,
|
||||
async execute({ query, topK }) {
|
||||
// Chain on the prior call so concurrent dispatches serialize.
|
||||
const chained = turnState.pendingChain.then(
|
||||
() => doExecute({ query, topK }),
|
||||
() => doExecute({ query, topK }),
|
||||
);
|
||||
turnState.pendingChain = chained.then(
|
||||
() => undefined,
|
||||
() => undefined,
|
||||
);
|
||||
return chained;
|
||||
},
|
||||
});
|
||||
|
||||
async function doExecute({
|
||||
query,
|
||||
topK,
|
||||
}: {
|
||||
query: string;
|
||||
topK: number;
|
||||
}): Promise<LookupToolResult> {
|
||||
const cacheKey = JSON.stringify({
|
||||
q: query.trim().toLowerCase(),
|
||||
k: topK,
|
||||
sb: spoilerBoundPosition ?? null,
|
||||
m: activeEmbeddingModel.id,
|
||||
b: bookHash,
|
||||
});
|
||||
const cached = turnState.cache.get(cacheKey);
|
||||
if (cached) {
|
||||
onEvent?.({ type: 'tool_call_cached', payload: { query_length: query.length } });
|
||||
return { ...cached, cached: true };
|
||||
}
|
||||
|
||||
if (turnState.totalToolMs > PER_TURN_BUDGET_MS) {
|
||||
const result: LookupToolResult = {
|
||||
passages: [],
|
||||
status: 'budget_exceeded',
|
||||
hint: 'Per-turn retrieval budget exhausted; do not call lookupPassage again this turn — finalize the answer with what you already have.',
|
||||
};
|
||||
onEvent?.({ type: 'budget_exceeded' });
|
||||
// Don't cache budget_exceeded — caller might want to retry next turn.
|
||||
return result;
|
||||
}
|
||||
|
||||
onEvent?.({
|
||||
type: 'tool_called',
|
||||
payload: { tool: 'lookupPassage', query_length: query.length },
|
||||
});
|
||||
const t0 = Date.now();
|
||||
const retrieved = await retriever.search({
|
||||
bookHash,
|
||||
query,
|
||||
k: topK,
|
||||
spoilerBoundPosition,
|
||||
activeEmbeddingModel,
|
||||
});
|
||||
turnState.totalToolMs += Date.now() - t0;
|
||||
|
||||
const passages: LookupPassage[] = retrieved.passages.map((p) => ({
|
||||
cfi: p.cfi,
|
||||
endCfi: p.endCfi,
|
||||
chapter: p.chapterTitle ?? undefined,
|
||||
text: p.text,
|
||||
}));
|
||||
|
||||
const { clamped, truncated } = clampToCharCap(passages, RESULT_SIZE_CAP_CHARS);
|
||||
|
||||
if (clamped.length === 0 && retrieved.status === 'ok') {
|
||||
onEvent?.({ type: 'tool_returned_empty' });
|
||||
}
|
||||
if (retrieved.status === 'stale_index') {
|
||||
onEvent?.({ type: 'tool_returned_stale' });
|
||||
}
|
||||
|
||||
const result: LookupToolResult = {
|
||||
passages: clamped,
|
||||
status: retrieved.status,
|
||||
truncated: truncated || undefined,
|
||||
hint: retrieved.status === 'ok' ? undefined : hintFor(retrieved.status, retrieved.reason),
|
||||
};
|
||||
turnState.cache.set(cacheKey, result);
|
||||
return result;
|
||||
}
|
||||
}
|
||||
|
||||
function clampToCharCap(
|
||||
passages: LookupPassage[],
|
||||
cap: number,
|
||||
): { clamped: LookupPassage[]; truncated: boolean } {
|
||||
let total = 0;
|
||||
for (const p of passages) total += p.text.length;
|
||||
if (total <= cap) return { clamped: passages, truncated: false };
|
||||
// Drop from the end (lowest RRF rank) until under cap.
|
||||
const clamped = [...passages];
|
||||
while (clamped.length > 0 && total > cap) {
|
||||
const dropped = clamped.pop()!;
|
||||
total -= dropped.text.length;
|
||||
}
|
||||
return { clamped, truncated: true };
|
||||
}
|
||||
|
||||
function hintFor(status: LookupToolStatus, reason?: string): string {
|
||||
switch (status) {
|
||||
case 'not_indexed':
|
||||
return "This book hasn't been indexed yet. Tell the user to open AI settings and click 'Index this book'.";
|
||||
case 'empty_index':
|
||||
return 'This book contains no extractable text (image-only PDF or scanned book). Tell the user Reedy cannot answer questions about its content.';
|
||||
case 'stale_index':
|
||||
return reason
|
||||
? `${reason}. Tell the user to re-index the book from settings.`
|
||||
: 'The active embedding model differs from the one this book was indexed with. Tell the user to re-index.';
|
||||
case 'degraded':
|
||||
return reason
|
||||
? `Vector search unavailable (${reason}). Answer with what you got and mention that results are text-match only.`
|
||||
: 'Vector search was temporarily unavailable; results are from text matching only.';
|
||||
case 'budget_exceeded':
|
||||
return 'Per-turn retrieval budget exhausted; finalize the answer with what you already have.';
|
||||
case 'ok':
|
||||
return '';
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Wrap a passage for inclusion in the assistant's system prompt. Used by the
|
||||
* M1.7 prompt builder, not by the tool layer — the tool returns the
|
||||
* structured result and the adapter decides where (if anywhere) to inline
|
||||
* the envelope text.
|
||||
*
|
||||
* Per plan §M1.6 / Codex F7: book text containing literal `</retrieved>`,
|
||||
* `&`, `<`, `>` is XML-escaped so the model cannot mistake it for a closing
|
||||
* tag. The opener uses `trust="untrusted"` to remind the model these are
|
||||
* data, not instructions.
|
||||
*/
|
||||
export function serializeForModel(passage: {
|
||||
cfi: string;
|
||||
chapter?: string;
|
||||
text: string;
|
||||
}): string {
|
||||
const escapedText = xmlEscape(passage.text);
|
||||
const escapedCfi = xmlAttrEscape(passage.cfi);
|
||||
const chapterAttr = passage.chapter ? ` chapter="${xmlAttrEscape(passage.chapter)}"` : '';
|
||||
return `<retrieved trust="untrusted" cfi="${escapedCfi}"${chapterAttr}>${escapedText}</retrieved>`;
|
||||
}
|
||||
|
||||
function xmlEscape(s: string): string {
|
||||
return s.replace(/&/g, '&').replace(/</g, '<').replace(/>/g, '>');
|
||||
}
|
||||
|
||||
function xmlAttrEscape(s: string): string {
|
||||
return xmlEscape(s).replace(/"/g, '"');
|
||||
}
|
||||
Reference in New Issue
Block a user