184 lines
7.9 KiB
JavaScript
184 lines
7.9 KiB
JavaScript
import { strict as assert } from 'node:assert'
|
|
import { test } from 'node:test'
|
|
|
|
import { excerptOf, queryTerms, rankChunks, splitIdentifier } from '../lib/rank.js'
|
|
import { hashEmbed } from './helpers.js'
|
|
|
|
const DIMENSIONS = 128
|
|
|
|
/** Build a scoring matrix from a list of chunks. */
|
|
function makeMatrix(chunks) {
|
|
const vectors = new Float32Array(chunks.length * DIMENSIONS)
|
|
chunks.forEach((chunk, at) => {
|
|
vectors.set(hashEmbed(`${chunk.path} ${chunk.symbol ?? ''} ${chunk.text}`, DIMENSIONS), at * DIMENSIONS)
|
|
})
|
|
return {
|
|
vectors,
|
|
dimensions: DIMENSIONS,
|
|
rows: chunks.map((chunk, at) => ({
|
|
id: at + 1,
|
|
path: chunk.path,
|
|
startLine: chunk.startLine,
|
|
endLine: chunk.endLine,
|
|
...(chunk.symbol === undefined ? {} : { symbol: chunk.symbol }),
|
|
text: chunk.text,
|
|
})),
|
|
}
|
|
}
|
|
|
|
test('splitIdentifier breaks camelCase and snake_case alike', () => {
|
|
assert.deepEqual(splitIdentifier('getUserById'), ['get', 'User', 'By', 'Id'])
|
|
assert.deepEqual(splitIdentifier('get_user_by_id'), ['get', 'user', 'by', 'id'])
|
|
assert.deepEqual(splitIdentifier('HTTPResponseParser'), ['HTTP', 'Response', 'Parser'])
|
|
})
|
|
|
|
test('queryTerms keeps identifier sub-words and drops stopwords', () => {
|
|
const terms = queryTerms('where is the getUserById function')
|
|
assert.ok(terms.includes('get'))
|
|
assert.ok(terms.includes('user'))
|
|
assert.ok(!terms.includes('the'), 'stopwords must not rank anything')
|
|
})
|
|
|
|
test('a named identifier wins even against a topically similar decoy', () => {
|
|
const matrix = makeMatrix([
|
|
{
|
|
path: 'src/services/auth/session.ts',
|
|
symbol: 'issueSession',
|
|
startLine: 1,
|
|
endLine: 10,
|
|
text: 'export function issueSession(userId) { return { userId, expiresAt: Date.now() } }',
|
|
},
|
|
{
|
|
path: 'src/services/auth/other.ts',
|
|
symbol: 'revokeSession',
|
|
startLine: 1,
|
|
endLine: 10,
|
|
text: 'export function revokeSession(record) { record.revoked = true; return record }',
|
|
},
|
|
{ path: 'src/weather/forecast.ts', symbol: 'average', startLine: 1, endLine: 5, text: 'export function average(days) { return days.length }' },
|
|
])
|
|
const query = 'issueSession'
|
|
const ranked = rankChunks({ matrix, queryVector: hashEmbed(query, DIMENSIONS), query, limit: 3 })
|
|
assert.equal(ranked.results[0].path, 'src/services/auth/session.ts')
|
|
assert.equal(ranked.results[0].symbol, 'issueSession')
|
|
})
|
|
|
|
test('a behavioural query still finds the right file without naming it', () => {
|
|
const matrix = makeMatrix([
|
|
{
|
|
path: 'src/services/auth/session.ts',
|
|
symbol: 'issueSession',
|
|
startLine: 1,
|
|
endLine: 6,
|
|
text: 'export function issueSession(userId) { const expiresAt = Date.now() + ttl; return { userId, expiresAt } }',
|
|
},
|
|
{ path: 'src/weather/forecast.ts', symbol: 'averageTemperature', startLine: 1, endLine: 6, text: 'export function averageTemperature(days) { return days.reduce(sum, 0) }' },
|
|
])
|
|
const query = 'session expiry userId expiresAt'
|
|
const ranked = rankChunks({ matrix, queryVector: hashEmbed(query, DIMENSIONS), query, limit: 2 })
|
|
assert.equal(ranked.results[0].path, 'src/services/auth/session.ts')
|
|
})
|
|
|
|
test('at most perFile chunks come from one file', () => {
|
|
const chunks = Array.from({ length: 6 }, (_, at) => ({
|
|
path: 'src/big.ts',
|
|
symbol: `fn_${at}`,
|
|
startLine: at * 10 + 1,
|
|
endLine: at * 10 + 9,
|
|
text: `export function fn_${at}(handler) { return handler(${at}) }`,
|
|
}))
|
|
chunks.push({ path: 'src/other.ts', symbol: 'fn_x', startLine: 1, endLine: 5, text: 'export function fn_x(handler) { return handler }' })
|
|
const matrix = makeMatrix(chunks)
|
|
const query = 'fn handler'
|
|
const ranked = rankChunks({ matrix, queryVector: hashEmbed(query, DIMENSIONS), query, limit: 6, perFile: 2 })
|
|
const fromBig = ranked.results.filter((result) => result.path === 'src/big.ts')
|
|
assert.ok(fromBig.length <= 2, `expected at most 2 from one file, got ${fromBig.length}`)
|
|
})
|
|
|
|
test('excerptOf keeps the head and reports truncation', () => {
|
|
const text = Array.from({ length: 100 }, (_, at) => `line ${at}`).join('\n')
|
|
const short = excerptOf(text, 200)
|
|
assert.equal(short.truncated, false)
|
|
assert.equal(short.excerpt, text)
|
|
const long = excerptOf(text, 10)
|
|
assert.equal(long.truncated, true)
|
|
assert.equal(long.excerpt.split('\n').length, 10)
|
|
assert.equal(long.excerpt.split('\n')[0], 'line 0')
|
|
})
|
|
|
|
test('rankChunks handles an empty matrix without inventing results', () => {
|
|
const ranked = rankChunks({
|
|
matrix: { vectors: new Float32Array(0), rows: [], dimensions: 0 },
|
|
queryVector: new Float32Array(0),
|
|
query: 'anything',
|
|
})
|
|
assert.deepEqual(ranked.results, [])
|
|
assert.equal(ranked.vectorBest, 0)
|
|
})
|
|
|
|
test('ranking is deterministic for equal scores', () => {
|
|
const chunks = Array.from({ length: 4 }, (_, at) => ({
|
|
path: `src/f${at}.ts`,
|
|
startLine: 1,
|
|
endLine: 2,
|
|
text: 'export const same = 1',
|
|
}))
|
|
const matrix = makeMatrix(chunks)
|
|
const query = 'same'
|
|
const first = rankChunks({ matrix, queryVector: hashEmbed(query, DIMENSIONS), query, limit: 4 })
|
|
const second = rankChunks({ matrix, queryVector: hashEmbed(query, DIMENSIONS), query, limit: 4 })
|
|
assert.deepEqual(
|
|
first.results.map((result) => result.path),
|
|
second.results.map((result) => result.path),
|
|
)
|
|
})
|
|
test('a chunk that fits is returned whole rather than cut in half', () => {
|
|
// The default cap equals the chunk cap, so a chunk — already a semantic unit —
|
|
// is not truncated. Cutting below it is what forces a follow-up read.
|
|
const body = Array.from({ length: 60 }, (_, at) => ` const line_${at} = compute(${at})`).join('\n')
|
|
const text = `export function handler() {\n${body}\n}`
|
|
const matrix = makeMatrix([{ path: 'src/a.ts', symbol: 'handler', startLine: 1, endLine: 62, text }])
|
|
const ranked = rankChunks({ matrix, queryVector: hashEmbed('handler', DIMENSIONS), query: 'handler', limit: 1 })
|
|
assert.equal(ranked.results[0].truncated, false, 'a 62-line chunk should arrive whole')
|
|
assert.equal(ranked.results[0].excerpt, text)
|
|
})
|
|
|
|
test('a per-excerpt cap still truncates, and says so', () => {
|
|
const text = Array.from({ length: 60 }, (_, at) => `line ${at}`).join('\n')
|
|
const matrix = makeMatrix([{ path: 'src/a.ts', startLine: 1, endLine: 60, text }])
|
|
const ranked = rankChunks({ matrix, queryVector: hashEmbed('line', DIMENSIONS), query: 'line', limit: 1, excerptLines: 10 })
|
|
assert.equal(ranked.results[0].truncated, true)
|
|
assert.equal(ranked.results[0].excerpt.split('\n').length, 10)
|
|
})
|
|
|
|
test('the total line budget bounds the whole set, allocating by rank', () => {
|
|
const chunks = Array.from({ length: 6 }, (_, at) => ({
|
|
path: `src/f${at}.ts`,
|
|
symbol: `fn_${at}`,
|
|
startLine: 1,
|
|
endLine: 40,
|
|
text: Array.from({ length: 40 }, (_, line) => `fn_${at} line ${line} handler`).join('\n'),
|
|
}))
|
|
const matrix = makeMatrix(chunks)
|
|
const query = 'fn handler'
|
|
const ranked = rankChunks({ matrix, queryVector: hashEmbed(query, DIMENSIONS), query, limit: 6, excerptLines: 40, totalLines: 90 })
|
|
|
|
const total = ranked.results.reduce((sum, result) => sum + result.excerpt.split('\n').length, 0)
|
|
assert.ok(total <= 90 + 40, `budget should bind, total was ${total}`)
|
|
assert.ok(ranked.results.length < 6, 'the budget should have stopped the set early')
|
|
assert.ok(ranked.results.length >= 1, 'at least one result must survive any budget')
|
|
})
|
|
|
|
test('at least one result is returned even with no budget at all', () => {
|
|
const chunks = Array.from({ length: 3 }, (_, at) => ({
|
|
path: `src/f${at}.ts`,
|
|
startLine: 1,
|
|
endLine: 30,
|
|
text: Array.from({ length: 30 }, (_, line) => `handler line ${line}`).join('\n'),
|
|
}))
|
|
const matrix = makeMatrix(chunks)
|
|
const ranked = rankChunks({ matrix, queryVector: hashEmbed('handler', DIMENSIONS), query: 'handler', limit: 3, totalLines: 0 })
|
|
assert.equal(ranked.results.length, 1, 'an empty budget must still answer with something')
|
|
assert.equal(ranked.results[0].excerpt.split('\n').length, 1, 'squeezed to the minimum, not dropped')
|
|
})
|