329 lines
14 KiB
TypeScript
329 lines
14 KiB
TypeScript
import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
|
|
import { MindDB } from '../../src/mind/db.js';
|
|
import { FrameStore } from '../../src/mind/frames.js';
|
|
import { SessionStore } from '../../src/mind/sessions.js';
|
|
import { HybridSearch, rechunkAllFrames } from '../../src/mind/search.js';
|
|
import { MockEmbedder } from './helpers/mock-embedder.js';
|
|
|
|
/**
|
|
* D1 (oss-drift triage, 2026-06-11) — chunk-level retrieval lane, reverse-
|
|
* ported from OSS hive-mind "Phase 3b-3". FLAG-GATED: WAGGLE_CHUNK_RETRIEVAL=1
|
|
* is OPT-IN; default OFF must be byte-identical to pre-D1 behavior.
|
|
*/
|
|
|
|
const FLAG = 'WAGGLE_CHUNK_RETRIEVAL';
|
|
|
|
/** A single paragraph of `sentences` short sentences (~55 chars each). */
|
|
function para(topic: string, sentences: number): string {
|
|
return Array.from(
|
|
{ length: sentences },
|
|
(_, i) => `The ${topic} system processes record number ${i} every day.`
|
|
).join(' ');
|
|
}
|
|
|
|
/** Multi-paragraph content long enough to produce >= 2 chunks (default knobs). */
|
|
function longContent(topic: string): string {
|
|
return `${para(topic, 30)}\n\n${para(topic, 30)}`;
|
|
}
|
|
|
|
describe('HybridSearch — chunk-level retrieval lane (D1)', () => {
|
|
let db: MindDB;
|
|
let frames: FrameStore;
|
|
let sessions: SessionStore;
|
|
let search: HybridSearch;
|
|
let gopId: string;
|
|
let savedFlag: string | undefined;
|
|
|
|
beforeEach(() => {
|
|
savedFlag = process.env[FLAG];
|
|
// D1 flip (2026-06-12): the flag is now default-ON (unset = enabled), so
|
|
// "off" in tests must be the explicit kill switch '0', not deletion.
|
|
process.env[FLAG] = '0';
|
|
db = new MindDB(':memory:');
|
|
frames = new FrameStore(db);
|
|
sessions = new SessionStore(db);
|
|
search = new HybridSearch(db, new MockEmbedder());
|
|
gopId = sessions.create().gop_id;
|
|
});
|
|
|
|
afterEach(() => {
|
|
if (savedFlag === undefined) delete process.env[FLAG];
|
|
else process.env[FLAG] = savedFlag;
|
|
db.close();
|
|
});
|
|
|
|
function chunkRowCount(): number {
|
|
return (db.getDatabase().prepare('SELECT COUNT(*) AS n FROM memory_frame_chunks').get() as { n: number }).n;
|
|
}
|
|
|
|
function chunkVecCount(): number {
|
|
return (db.getDatabase().prepare('SELECT COUNT(*) AS n FROM memory_frame_chunks_vec').get() as { n: number }).n;
|
|
}
|
|
|
|
describe('flag OFF (default) — byte-identical to pre-D1', () => {
|
|
it('indexFrame writes ZERO chunk rows', async () => {
|
|
const f = frames.createIFrame(gopId, longContent('kubernetes'), 'normal', 'user_stated');
|
|
await search.indexFrame(f.id, f.content);
|
|
expect(chunkRowCount()).toBe(0);
|
|
expect(chunkVecCount()).toBe(0);
|
|
});
|
|
|
|
it('search() never touches the chunk lane and uses whole-frame vectors', async () => {
|
|
const f = frames.createIFrame(gopId, longContent('kubernetes'), 'normal', 'user_stated');
|
|
await search.indexFrame(f.id, f.content);
|
|
|
|
const chunkSpy = vi.spyOn(search, 'vectorSearchChunks');
|
|
const vecSpy = vi.spyOn(search, 'vectorSearch');
|
|
const results = await search.search('kubernetes system record', { limit: 5 });
|
|
|
|
expect(chunkSpy).not.toHaveBeenCalled();
|
|
expect(vecSpy).toHaveBeenCalledTimes(1);
|
|
expect(results.length).toBeGreaterThan(0);
|
|
expect(results[0].frame.id).toBe(f.id);
|
|
});
|
|
|
|
it('search() output is stable across calls (regression anchor)', async () => {
|
|
const a = frames.createIFrame(gopId, longContent('kubernetes'), 'normal', 'user_stated');
|
|
const b = frames.createIFrame(gopId, longContent('gardening'), 'normal', 'user_stated');
|
|
await search.indexFrame(a.id, a.content);
|
|
await search.indexFrame(b.id, b.content);
|
|
|
|
const r1 = await search.search('kubernetes system record', { limit: 5 });
|
|
const r2 = await search.search('kubernetes system record', { limit: 5 });
|
|
expect(r1.map(r => [r.frame.id, r.finalScore])).toEqual(r2.map(r => [r.frame.id, r.finalScore]));
|
|
});
|
|
});
|
|
|
|
describe('flag ON — chunk lane active', () => {
|
|
beforeEach(() => {
|
|
process.env[FLAG] = '1';
|
|
});
|
|
|
|
it('indexFrame also writes chunk rows + chunk vectors', async () => {
|
|
const f = frames.createIFrame(gopId, longContent('kubernetes'), 'normal', 'user_stated');
|
|
await search.indexFrame(f.id, f.content);
|
|
expect(chunkRowCount()).toBeGreaterThanOrEqual(2); // long content → multiple chunks
|
|
expect(chunkVecCount()).toBe(chunkRowCount());
|
|
});
|
|
|
|
it('indexFramesBatch also writes chunk rows for every frame', async () => {
|
|
const a = frames.createIFrame(gopId, longContent('kubernetes'), 'normal', 'user_stated');
|
|
const b = frames.createIFrame(gopId, 'short note about tomatoes', 'normal', 'user_stated');
|
|
await search.indexFramesBatch([
|
|
{ id: a.id, content: a.content },
|
|
{ id: b.id, content: b.content },
|
|
]);
|
|
const raw = db.getDatabase();
|
|
const perFrame = raw
|
|
.prepare('SELECT frame_id, COUNT(*) AS n FROM memory_frame_chunks GROUP BY frame_id')
|
|
.all() as Array<{ frame_id: number; n: number }>;
|
|
const byId = new Map(perFrame.map(r => [r.frame_id, r.n]));
|
|
expect(byId.get(a.id)).toBeGreaterThanOrEqual(2);
|
|
expect(byId.get(b.id)).toBe(1); // short content → single chunk
|
|
});
|
|
|
|
it('search() uses the chunk lane (whole-frame vectorSearch NOT called)', async () => {
|
|
const f = frames.createIFrame(gopId, longContent('kubernetes'), 'normal', 'user_stated');
|
|
await search.indexFrame(f.id, f.content);
|
|
|
|
const vecSpy = vi.spyOn(search, 'vectorSearch');
|
|
const results = await search.search('kubernetes system record', { limit: 5 });
|
|
|
|
expect(vecSpy).not.toHaveBeenCalled();
|
|
expect(results.length).toBeGreaterThan(0);
|
|
expect(results[0].frame.id).toBe(f.id);
|
|
});
|
|
|
|
it('vectorSearchChunks dedups to best-chunk-per-frame', async () => {
|
|
// Both of this frame's chunks match the query — the parent frame must
|
|
// appear exactly once in the returned ids.
|
|
const f = frames.createIFrame(gopId, longContent('kubernetes'), 'normal', 'user_stated');
|
|
const g = frames.createIFrame(gopId, longContent('gardening'), 'normal', 'user_stated');
|
|
await search.indexFrame(f.id, f.content);
|
|
await search.indexFrame(g.id, g.content);
|
|
expect(chunkRowCount()).toBeGreaterThanOrEqual(4);
|
|
|
|
const ids = await search.vectorSearchChunks('kubernetes system record', 10);
|
|
expect(ids).not.toBeNull();
|
|
const occurrences = (ids as number[]).filter(id => id === f.id).length;
|
|
expect(occurrences).toBe(1);
|
|
expect(new Set(ids as number[]).size).toBe((ids as number[]).length);
|
|
expect((ids as number[])[0]).toBe(f.id); // best-matching frame first
|
|
});
|
|
|
|
it('excludes deprecated chunk candidates before the KNN limit', async () => {
|
|
const live = frames.createIFrame(
|
|
gopId,
|
|
`${longContent('gardening')} One live kubernetes system record sentence.`,
|
|
'normal',
|
|
'user_stated',
|
|
);
|
|
const staleFrames = Array.from({ length: 13 }, (_, index) => frames.createIFrame(
|
|
gopId,
|
|
`${longContent('kubernetes')} Obsolete source ${index}.`,
|
|
'normal',
|
|
'user_stated',
|
|
));
|
|
await search.indexFramesBatch([
|
|
{ id: live.id, content: live.content },
|
|
...staleFrames.map((frame) => ({ id: frame.id, content: frame.content })),
|
|
]);
|
|
for (const stale of staleFrames) {
|
|
frames.update(stale.id, stale.content, 'deprecated');
|
|
}
|
|
const otherGop = sessions.create().gop_id;
|
|
const outOfScopeDecoy = frames.createIFrame(
|
|
otherGop,
|
|
longContent('kubernetes'),
|
|
'normal',
|
|
'user_stated',
|
|
);
|
|
await search.indexFrame(outOfScopeDecoy.id, outOfScopeDecoy.content);
|
|
const staleChunkCount = db.getDatabase().prepare(`
|
|
SELECT COUNT(*) AS n
|
|
FROM memory_frame_chunks c
|
|
JOIN memory_frames mf ON mf.id = c.frame_id
|
|
WHERE mf.importance = 'deprecated'
|
|
`).get() as { n: number };
|
|
expect(staleChunkCount.n).toBeGreaterThan(25);
|
|
|
|
const ids = await search.vectorSearchChunks(
|
|
'kubernetes system record',
|
|
1,
|
|
gopId,
|
|
true,
|
|
);
|
|
|
|
expect(ids).toEqual([live.id]);
|
|
await expect(search.vectorSearchChunks(
|
|
'kubernetes system record',
|
|
1,
|
|
undefined,
|
|
true,
|
|
)).resolves.toEqual([outOfScopeDecoy.id]);
|
|
});
|
|
|
|
it('falls back to whole-frame vectors when the chunk index is empty', async () => {
|
|
// Index with the flag OFF (explicit kill switch — default is ON) so no
|
|
// chunks are written…
|
|
process.env[FLAG] = '0';
|
|
const f = frames.createIFrame(gopId, longContent('kubernetes'), 'normal', 'user_stated');
|
|
await search.indexFrame(f.id, f.content);
|
|
expect(chunkRowCount()).toBe(0);
|
|
// …then search with the flag ON: chunk probe finds 0 rows → null → fallback.
|
|
process.env[FLAG] = '1';
|
|
|
|
const vecSpy = vi.spyOn(search, 'vectorSearch');
|
|
const results = await search.search('kubernetes system record', { limit: 5 });
|
|
|
|
expect(vecSpy).toHaveBeenCalledTimes(1);
|
|
expect(results.length).toBeGreaterThan(0);
|
|
expect(results[0].frame.id).toBe(f.id);
|
|
});
|
|
|
|
it('vectorSearchChunks honours gopId scoping', async () => {
|
|
const otherGop = sessions.create().gop_id;
|
|
const f = frames.createIFrame(gopId, longContent('kubernetes'), 'normal', 'user_stated');
|
|
// Distinct content: identical text would hit the D3 content-hash dedup
|
|
// and return the SAME frame (one frame can't be vector-indexed twice).
|
|
const g = frames.createIFrame(
|
|
otherGop,
|
|
`${longContent('kubernetes')} Extra kubernetes deployment sentence.`,
|
|
'normal',
|
|
'user_stated'
|
|
);
|
|
await search.indexFrame(f.id, f.content);
|
|
await search.indexFrame(g.id, g.content);
|
|
|
|
const ids = await search.vectorSearchChunks('kubernetes system record', 10, gopId);
|
|
expect(ids).not.toBeNull();
|
|
expect(ids).toContain(f.id);
|
|
expect(ids).not.toContain(g.id);
|
|
});
|
|
});
|
|
|
|
describe('indexChunksForFrame (flag-independent)', () => {
|
|
it('is callable with the flag OFF (backfill/eval path)', async () => {
|
|
const f = frames.createIFrame(gopId, longContent('kubernetes'), 'normal', 'user_stated');
|
|
const n = await search.indexChunksForFrame(f.id, f.content);
|
|
expect(n).toBeGreaterThanOrEqual(2);
|
|
expect(chunkRowCount()).toBe(n);
|
|
expect(chunkVecCount()).toBe(n);
|
|
});
|
|
|
|
it('replaces chunks on reindex (no orphaned vec rows)', async () => {
|
|
const f = frames.createIFrame(gopId, longContent('kubernetes'), 'normal', 'user_stated');
|
|
const first = await search.indexChunksForFrame(f.id, f.content);
|
|
expect(first).toBeGreaterThanOrEqual(2);
|
|
|
|
const second = await search.indexChunksForFrame(f.id, 'short replacement content');
|
|
expect(second).toBe(1);
|
|
expect(chunkRowCount()).toBe(1);
|
|
expect(chunkVecCount()).toBe(1);
|
|
const row = db.getDatabase()
|
|
.prepare('SELECT content FROM memory_frame_chunks WHERE frame_id = ?')
|
|
.get(f.id) as { content: string };
|
|
expect(row.content).toBe('short replacement content');
|
|
});
|
|
|
|
it('rejects invalid frame ids', async () => {
|
|
await expect(search.indexChunksForFrame(0, 'x')).rejects.toThrow('Invalid frame ID');
|
|
await expect(search.indexChunksForFrame(Number.NaN, 'x')).rejects.toThrow('Invalid frame ID');
|
|
});
|
|
});
|
|
|
|
describe('recreateVecTables (D1 extension)', () => {
|
|
it('drops + recreates the chunk vec table; chunk CONTENT rows survive', async () => {
|
|
const f = frames.createIFrame(gopId, longContent('kubernetes'), 'normal', 'user_stated');
|
|
await search.indexChunksForFrame(f.id, f.content);
|
|
const contentRows = chunkRowCount();
|
|
expect(contentRows).toBeGreaterThanOrEqual(2);
|
|
expect(chunkVecCount()).toBe(contentRows);
|
|
|
|
db.recreateVecTables(1024);
|
|
|
|
// Vectors discarded, content rows survive (they're re-derivable text,
|
|
// not vectors — rechunkAllFrames re-embeds them).
|
|
expect(chunkVecCount()).toBe(0);
|
|
expect(chunkRowCount()).toBe(contentRows);
|
|
// And the recreated table is writable again.
|
|
await search.indexChunksForFrame(f.id, f.content);
|
|
expect(chunkVecCount()).toBe(chunkRowCount());
|
|
});
|
|
});
|
|
|
|
describe('rechunkAllFrames (backfill helper)', () => {
|
|
it('populates chunks from existing frames, skipping deprecated', async () => {
|
|
const a = frames.createIFrame(gopId, longContent('kubernetes'), 'normal', 'user_stated');
|
|
const b = frames.createIFrame(gopId, 'short note about tomatoes', 'normal', 'user_stated');
|
|
const dep = frames.createIFrame(gopId, longContent('obsolete'), 'deprecated', 'user_stated');
|
|
expect(chunkRowCount()).toBe(0);
|
|
|
|
const result = await rechunkAllFrames(db, search);
|
|
|
|
expect(result.framesProcessed).toBe(2);
|
|
expect(result.framesFailed).toBe(0);
|
|
expect(result.chunksCreated).toBeGreaterThanOrEqual(3); // >=2 for a, 1 for b
|
|
expect(chunkRowCount()).toBe(result.chunksCreated);
|
|
expect(chunkVecCount()).toBe(result.chunksCreated);
|
|
|
|
const raw = db.getDatabase();
|
|
const frameIds = (raw.prepare('SELECT DISTINCT frame_id FROM memory_frame_chunks').all() as Array<{ frame_id: number }>)
|
|
.map(r => r.frame_id);
|
|
expect(frameIds).toContain(a.id);
|
|
expect(frameIds).toContain(b.id);
|
|
expect(frameIds).not.toContain(dep.id);
|
|
});
|
|
|
|
it('is idempotent — a second pass yields the same chunk counts', async () => {
|
|
const f = frames.createIFrame(gopId, longContent('kubernetes'), 'normal', 'user_stated');
|
|
void f;
|
|
const first = await rechunkAllFrames(db, search);
|
|
const second = await rechunkAllFrames(db, search);
|
|
expect(second.chunksCreated).toBe(first.chunksCreated);
|
|
expect(chunkRowCount()).toBe(first.chunksCreated);
|
|
expect(chunkVecCount()).toBe(first.chunksCreated);
|
|
});
|
|
});
|
|
});
|