114 lines
4.0 KiB
JavaScript
114 lines
4.0 KiB
JavaScript
import { strict as assert } from 'node:assert'
|
|
import { rm } from 'node:fs/promises'
|
|
import { join } from 'node:path'
|
|
import { test } from 'node:test'
|
|
|
|
import { RecallIndex } from '../lib/store.js'
|
|
import { makeStateDir } from './helpers.js'
|
|
|
|
const IDENTITY = { model: 'test-model', dimensions: 4, apiBase: 'http://stub' }
|
|
|
|
/** A chunk with a hand-written vector, so assertions are exact. */
|
|
function chunk(startLine, vector, symbol) {
|
|
return {
|
|
startLine,
|
|
endLine: startLine + 2,
|
|
...(symbol === undefined ? {} : { symbol }),
|
|
text: `line ${startLine}`,
|
|
vec: Float32Array.from(vector),
|
|
}
|
|
}
|
|
|
|
async function withIndex(t, work) {
|
|
const dir = await makeStateDir()
|
|
const index = new RecallIndex(join(dir, 'index.sqlite'))
|
|
t.after(async () => {
|
|
index.close()
|
|
await rm(dir, { recursive: true, force: true })
|
|
})
|
|
index.open(IDENTITY)
|
|
return work(index)
|
|
}
|
|
|
|
test('replaceFile round-trips chunks and their vectors', async (t) => {
|
|
await withIndex(t, (index) => {
|
|
index.replaceFile('a.ts', 'hash-a', 100, [chunk(1, [1, 0, 0, 0], 'alpha'), chunk(10, [0, 1, 0, 0], 'beta')])
|
|
const stats = index.stats()
|
|
assert.equal(stats.files, 1)
|
|
assert.equal(stats.chunks, 2)
|
|
|
|
const matrix = index.matrix()
|
|
assert.equal(matrix.dimensions, 4)
|
|
assert.equal(matrix.rows.length, 2)
|
|
assert.deepEqual(Array.from(matrix.vectors.slice(0, 4)), [1, 0, 0, 0])
|
|
assert.equal(matrix.rows[0].symbol, 'alpha')
|
|
assert.equal(matrix.rows[0].path, 'a.ts')
|
|
})
|
|
})
|
|
|
|
test('replacing a file drops its previous chunks rather than appending', async (t) => {
|
|
await withIndex(t, (index) => {
|
|
index.replaceFile('a.ts', 'hash-1', 10, [chunk(1, [1, 0, 0, 0])])
|
|
index.replaceFile('a.ts', 'hash-2', 20, [chunk(1, [0, 1, 0, 0]), chunk(5, [0, 0, 1, 0])])
|
|
assert.equal(index.stats().chunks, 2)
|
|
assert.equal(index.fileHashes().get('a.ts'), 'hash-2')
|
|
})
|
|
})
|
|
|
|
test('prune forgets files that disappeared from disk', async (t) => {
|
|
await withIndex(t, (index) => {
|
|
index.replaceFile('a.ts', 'ha', 1, [chunk(1, [1, 0, 0, 0])])
|
|
index.replaceFile('b.ts', 'hb', 1, [chunk(1, [0, 1, 0, 0])])
|
|
assert.equal(index.prune(['a.ts']), 1)
|
|
assert.equal(index.stats().files, 1)
|
|
assert.equal(index.matrix().rows[0].path, 'b.ts')
|
|
assert.equal(index.prune([]), 0)
|
|
})
|
|
})
|
|
|
|
test('changing the embedding model discards the index instead of mixing spaces', async (t) => {
|
|
await withIndex(t, (index) => {
|
|
index.replaceFile('a.ts', 'ha', 1, [chunk(1, [1, 0, 0, 0])])
|
|
assert.equal(index.stats().chunks, 1)
|
|
|
|
const outcome = index.recheck({ model: 'other-model', dimensions: 4, apiBase: 'http://stub' })
|
|
assert.match(outcome.discarded, /embedding model changed/)
|
|
assert.equal(index.stats().chunks, 0)
|
|
assert.equal(index.stats().files, 0)
|
|
})
|
|
})
|
|
|
|
test('changing the embedding width discards the index', async (t) => {
|
|
await withIndex(t, (index) => {
|
|
index.replaceFile('a.ts', 'ha', 1, [chunk(1, [1, 0, 0, 0])])
|
|
const outcome = index.recheck({ model: 'test-model', dimensions: 8, apiBase: 'http://stub' })
|
|
assert.match(outcome.discarded, /embedding width changed/)
|
|
assert.equal(index.stats().chunks, 0)
|
|
})
|
|
})
|
|
|
|
test('recheck is a no-op when nothing about the identity moved', async (t) => {
|
|
await withIndex(t, (index) => {
|
|
index.replaceFile('a.ts', 'ha', 1, [chunk(1, [1, 0, 0, 0])])
|
|
assert.deepEqual(index.recheck(IDENTITY), {})
|
|
assert.equal(index.stats().chunks, 1)
|
|
})
|
|
})
|
|
|
|
test('a write invalidates the cached matrix', async (t) => {
|
|
await withIndex(t, (index) => {
|
|
index.replaceFile('a.ts', 'ha', 1, [chunk(1, [1, 0, 0, 0])])
|
|
assert.equal(index.matrix().rows.length, 1)
|
|
index.replaceFile('b.ts', 'hb', 1, [chunk(1, [0, 1, 0, 0])])
|
|
assert.equal(index.matrix().rows.length, 2, 'the matrix must not serve a stale cache')
|
|
})
|
|
})
|
|
|
|
test('an index with no chunks reports a zero-width matrix', async (t) => {
|
|
await withIndex(t, (index) => {
|
|
const matrix = index.matrix()
|
|
assert.equal(matrix.rows.length, 0)
|
|
assert.equal(matrix.dimensions, 0)
|
|
})
|
|
})
|