291 lines
9.7 KiB
JavaScript
291 lines
9.7 KiB
JavaScript
const fs = require('fs/promises');
|
|
const os = require('os');
|
|
const path = require('path');
|
|
const { mockClient } = require('aws-sdk-client-mock');
|
|
const { BedrockRuntimeClient, InvokeModelCommand } = require('@aws-sdk/client-bedrock-runtime');
|
|
const { chunkBook } = require('../src/chunk');
|
|
const { applyCorrections } = require('../src/corrections');
|
|
const { embedMany } = require('../src/embed');
|
|
const { LocalFsStorage } = require('../src/storage');
|
|
const { buildIndexBody } = require('../src/indexMapping');
|
|
const { bulkWrite, deleteStale, verifyDocs, indexBook } = require('../src/indexer');
|
|
const { migrateRun, swapAlias, cutover, rollback } = require('../src/migrate');
|
|
|
|
const bedrock = mockClient(BedrockRuntimeClient);
|
|
|
|
function paragraph(words) {
|
|
return Array.from({ length: words }, (_, index) => `w${index}`).join(' ');
|
|
}
|
|
|
|
describe('chunking', () => {
|
|
const chapters = [
|
|
{ chapterNo: 1, title: 'One', startPage: 1, endPage: 1 },
|
|
{ chapterNo: 2, title: 'Two', startPage: 2, endPage: 2 },
|
|
];
|
|
const pages = [
|
|
{ page: 1, text: `${paragraph(80)}\n\n${paragraph(80)}` },
|
|
{ page: 2, text: `${paragraph(40)}\n\n1.1 Harbor notes\n\n${paragraph(40)}` },
|
|
];
|
|
|
|
it('stays inside a chapter, covers pages, and links neighbors', () => {
|
|
const chunks = chunkBook({ bookId: 5, pages, chapters, targetTokens: 30, overlapTokens: 8 });
|
|
expect(chunks.length).toBeGreaterThan(2);
|
|
for (const chunk of chunks) {
|
|
expect(chunk.pageStart).toBe(chunk.chapterNo);
|
|
expect(chunk.pageEnd).toBe(chunk.chapterNo);
|
|
expect(chunk.text).not.toMatch(/w0[\s\S]*Chapter/);
|
|
}
|
|
const chapter2 = chunks.filter((chunk) => chunk.chapterNo === 2);
|
|
expect(chapter2.some((chunk) => chunk.sectionTitle.includes('Harbor'))).toBe(true);
|
|
expect(chunks[0].chunkId).toBe('5:chunk:00001');
|
|
expect(chunks[0].prevChunkId).toBeNull();
|
|
expect(chunks[0].nextChunkId).toBe(chunks[1].chunkId);
|
|
expect(chunks[chunks.length - 1].nextChunkId).toBeNull();
|
|
const again = chunkBook({ bookId: 5, pages, chapters, targetTokens: 30, overlapTokens: 8 });
|
|
expect(again.map((chunk) => chunk.chunkId)).toEqual(chunks.map((chunk) => chunk.chunkId));
|
|
});
|
|
|
|
it('applies corrections before the chunk text is built', () => {
|
|
const corrected = applyCorrections({
|
|
pages,
|
|
chapters,
|
|
corrections: {
|
|
textReplacements: [{ page: 1, find: 'w0', replace: 'FIXED' }],
|
|
dropPages: [2],
|
|
chapterOverrides: [{ title: 'Only', startPage: 1 }],
|
|
metadata: { title: 'Given Title', author: 'Given Author' },
|
|
},
|
|
pageCount: 2,
|
|
});
|
|
expect(corrected.pages).toHaveLength(1);
|
|
expect(corrected.pages[0].text).toContain('FIXED');
|
|
expect(corrected.chapters[0].title).toBe('Only');
|
|
expect(corrected.metadata.author).toBe('Given Author');
|
|
const chunks = chunkBook({
|
|
bookId: 5,
|
|
pages: corrected.pages,
|
|
chapters: corrected.chapters,
|
|
targetTokens: 30,
|
|
overlapTokens: 0,
|
|
});
|
|
expect(chunks.every((chunk) => chunk.chapterNo !== 2)).toBe(true);
|
|
expect(chunks.some((chunk) => chunk.text.includes('FIXED'))).toBe(true);
|
|
});
|
|
});
|
|
|
|
describe('embedding cache', () => {
|
|
it('does not call the model again for the same text', async () => {
|
|
bedrock.reset();
|
|
const vector = Array.from({ length: 4 }, () => 0.1);
|
|
bedrock.on(InvokeModelCommand).resolves({
|
|
body: new TextEncoder().encode(JSON.stringify({ embedding: vector })),
|
|
});
|
|
const dir = await fs.mkdtemp(path.join(os.tmpdir(), 'embed-'));
|
|
const storage = new LocalFsStorage(dir);
|
|
const client = new BedrockRuntimeClient({ region: 'ap-south-1' });
|
|
const first = await embedMany({
|
|
client,
|
|
modelId: 'amazon.titan-embed-text-v2:0',
|
|
dimensions: 4,
|
|
texts: ['alpha', 'beta'],
|
|
storage,
|
|
bookId: 1,
|
|
concurrency: 2,
|
|
});
|
|
const second = await embedMany({
|
|
client,
|
|
modelId: 'amazon.titan-embed-text-v2:0',
|
|
dimensions: 4,
|
|
texts: ['alpha', 'beta'],
|
|
storage,
|
|
bookId: 1,
|
|
concurrency: 2,
|
|
});
|
|
expect(first.calls).toBe(2);
|
|
expect(second.calls).toBe(0);
|
|
expect(second.vectors).toEqual(first.vectors);
|
|
expect(bedrock.commandCalls(InvokeModelCommand)).toHaveLength(2);
|
|
await fs.rm(dir, { recursive: true, force: true });
|
|
});
|
|
});
|
|
|
|
describe('index mapping', () => {
|
|
it('matches the OpenSearch 2.7 lucene cosine mapping', () => {
|
|
const body = buildIndexBody(1024);
|
|
expect(body.settings).toEqual({ index: { knn: true } });
|
|
expect(body.mappings.properties.embedding).toEqual({
|
|
type: 'knn_vector',
|
|
dimension: 1024,
|
|
method: { name: 'hnsw', space_type: 'cosinesimil', engine: 'lucene' },
|
|
});
|
|
expect(Object.keys(body.mappings.properties).sort()).toEqual(
|
|
[
|
|
'book_id',
|
|
'chapter_no',
|
|
'chapter_title',
|
|
'chunk_id',
|
|
'contextual_text',
|
|
'created_at',
|
|
'embedding',
|
|
'ingest_version',
|
|
'language',
|
|
'level',
|
|
'next_chunk_id',
|
|
'ordinal',
|
|
'page_end',
|
|
'page_start',
|
|
'prev_chunk_id',
|
|
'section_title',
|
|
'text',
|
|
'token_count',
|
|
].sort(),
|
|
);
|
|
expect(body.mappings.properties.text.fields.english.analyzer).toBe('english');
|
|
});
|
|
});
|
|
|
|
function doc(id, source) {
|
|
return { _id: id, level: source.level, _source: source };
|
|
}
|
|
|
|
describe('indexer', () => {
|
|
it('retries failed bulk items and deletes only stale versions of the same book', async () => {
|
|
const calls = [];
|
|
const client = {
|
|
async bulk(params) {
|
|
calls.push(params);
|
|
if (calls.length === 1) {
|
|
return {
|
|
errors: true,
|
|
items: [{ index: { _id: 'a', status: 201 } }, { index: { _id: 'b', status: 429, error: { type: 'rejected' } } }],
|
|
};
|
|
}
|
|
return { errors: false, items: [{ index: { _id: 'b', status: 201 } }] };
|
|
},
|
|
async deleteByQuery(params) {
|
|
calls.push(params);
|
|
},
|
|
};
|
|
const docs = [doc('a', { level: 'chunk' }), doc('b', { level: 'chunk' })];
|
|
await bulkWrite(client, 'books_current_v2', docs);
|
|
expect(calls[1].body.map((row) => row.index && row.index._id).filter(Boolean)).toEqual(['b']);
|
|
await deleteStale(client, 'books_current_v2', 9, 'v9');
|
|
expect(calls[2].body.query).toEqual({
|
|
bool: {
|
|
must: [{ term: { book_id: '9' } }],
|
|
must_not: [{ term: { ingest_version: 'v9' } }],
|
|
},
|
|
});
|
|
});
|
|
|
|
it('fails verification when a page is missing or a vector is the wrong length', () => {
|
|
const vector = [0.1, 0.2];
|
|
const good = {
|
|
level: 'chunk',
|
|
chunk_id: '1:chunk:00001',
|
|
chapter_no: 1,
|
|
page_start: 1,
|
|
page_end: 1,
|
|
embedding: vector,
|
|
};
|
|
expect(() => verifyDocs([good, { level: 'book', chapter_no: 0, embedding: vector }], { dimensions: 2, pages: [1, 2] })).toThrow(/page 2/);
|
|
expect(() =>
|
|
verifyDocs(
|
|
[{ ...good, embedding: [0.1] }, { level: 'book', chapter_no: 0, embedding: vector }],
|
|
{ dimensions: 2, pages: [1] },
|
|
),
|
|
).toThrow(/bad embedding/);
|
|
});
|
|
|
|
it('does not call OpenSearch in local mode', async () => {
|
|
const vector = [0.2, 0.3];
|
|
const docs = [
|
|
doc('1:chunk:00001', {
|
|
level: 'chunk',
|
|
chunk_id: '1:chunk:00001',
|
|
chapter_no: 1,
|
|
page_start: 1,
|
|
page_end: 1,
|
|
embedding: vector,
|
|
}),
|
|
doc('1:book', { level: 'book', chapter_no: 0, embedding: vector }),
|
|
];
|
|
const result = await indexBook({
|
|
bookId: 1,
|
|
docs,
|
|
client: null,
|
|
config: { embedDimensions: 2, ingestVersion: 'v1', indexAlias: 'books_current', indexVersion: 'v2' },
|
|
pages: [1],
|
|
logger: { info() {} },
|
|
});
|
|
expect(result.chunkCount).toBe(1);
|
|
expect(result.indexName).toBeNull();
|
|
});
|
|
});
|
|
|
|
describe('migration', () => {
|
|
it('dry-run makes no rebuild calls', async () => {
|
|
let calls = 0;
|
|
const results = await migrateRun({
|
|
bookIds: ['3', '4'],
|
|
dryRun: true,
|
|
rebuild: async () => {
|
|
calls += 1;
|
|
},
|
|
});
|
|
expect(calls).toBe(0);
|
|
expect(results.map((row) => row.status)).toEqual(['skipped', 'skipped']);
|
|
});
|
|
|
|
it('continues after one book fails', async () => {
|
|
const results = await migrateRun({
|
|
bookIds: ['3', '4'],
|
|
dryRun: false,
|
|
rebuild: async (bookId) => {
|
|
if (bookId === '3') throw new Error('missing pdf');
|
|
},
|
|
});
|
|
expect(results[0]).toMatchObject({ status: 'failed', error: 'missing pdf' });
|
|
expect(results[1].status).toBe('success');
|
|
});
|
|
|
|
it('swaps the alias in one request and rollback restores the previous target', async () => {
|
|
const calls = [];
|
|
const client = {
|
|
async indices() {},
|
|
};
|
|
client.indices = {
|
|
async updateAliases(params) {
|
|
calls.push(params);
|
|
},
|
|
async getAlias() {
|
|
return { books_current_v1: { aliases: { books_current: {} } } };
|
|
},
|
|
};
|
|
const actions = await swapAlias(client, { alias: 'books_current', addIndex: 'books_current_v2', removeIndex: 'books_current_v1' });
|
|
expect(calls).toHaveLength(1);
|
|
expect(actions).toEqual([
|
|
{ remove: { index: 'books_current_v1', alias: 'books_current' } },
|
|
{ add: { index: 'books_current_v2', alias: 'books_current' } },
|
|
]);
|
|
|
|
const state = {
|
|
value: null,
|
|
async get() {
|
|
return this.value;
|
|
},
|
|
async put(value) {
|
|
this.value = value;
|
|
},
|
|
};
|
|
await cutover({ client, alias: 'books_current', toIndex: 'books_current_v2', state });
|
|
expect(calls).toHaveLength(2);
|
|
await rollback({ client, state });
|
|
expect(calls).toHaveLength(3);
|
|
expect(calls[2].body.actions).toEqual([
|
|
{ remove: { index: 'books_current_v2', alias: 'books_current' } },
|
|
{ add: { index: 'books_current_v1', alias: 'books_current' } },
|
|
]);
|
|
});
|
|
});
|