mirror of
https://github.com/hellodigua/ChatLab.git
synced 2026-10-26 02:04:37 +08:00
fix(ai): semantic index core reliability and evidence text
Prevent stale semantic index writes by checking write guard before persisting chunks; make rebuild detection explicit on app restart via reactivate() path; preserve full evidence snippet text without truncation; add regression test for same-second chunk tie-breaking in mapMessageToChunk.
This commit is contained in:
1 parent
88dea48ec6
commit
abd05981c6
11 files changed
+323
-53
No files matched your search
@@ -33,12 +33,15 @@ export const CHUNK_VECTOR_INDEX_TABLE = `
|
||||
`
|
||||
|
||||
/**
|
||||
* 范围映射索引:FTS message_id -> chunk 的 O(log n) 查询依赖前导列顺序,
|
||||
* 配合 `start_message_id <= ? ORDER BY start_message_id DESC LIMIT 1`。
|
||||
* 范围映射索引:
|
||||
* - idx_chunk_range 保留 message_id 领先的旧查询路径和既有数据库兼容性。
|
||||
* - idx_chunk_ts_range 支撑当前 `start_ts, start_message_id` 组合查找,避免按会话扫描排序。
|
||||
*/
|
||||
export const CHUNK_VECTOR_INDEX_INDEXES = `
|
||||
CREATE INDEX IF NOT EXISTS idx_chunk_range
|
||||
ON chunk_vector_index(db_path_hash, model_id, strategy_id, start_message_id, end_message_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_chunk_ts_range
|
||||
ON chunk_vector_index(db_path_hash, model_id, strategy_id, start_ts, start_message_id);
|
||||
`
|
||||
|
||||
export const EMBEDDING_INDEX_SCHEMA = CHUNK_VECTOR_INDEX_TABLE + CHUNK_VECTOR_INDEX_INDEXES
|
||||
|
||||
@@ -84,10 +84,15 @@ function setup(opts?: { embedDelayMs?: number }) {
|
||||
return { service, chatDbPath, dir, db, authProfiles, getEmbedCount: () => embedTextCount }
|
||||
}
|
||||
|
||||
test('enable builds index to completion and reports status', async () => {
|
||||
const { service, db } = setup()
|
||||
async function enableAndBuild(service: SemanticIndexService): Promise<void> {
|
||||
service.enable(SESSION_ID)
|
||||
service.build(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
}
|
||||
|
||||
test('explicit build indexes an enabled session to completion and reports status', async () => {
|
||||
const { service, db } = setup()
|
||||
await enableAndBuild(service)
|
||||
|
||||
const status = service.status(SESSION_ID)!
|
||||
assert.equal(status.enabled, true)
|
||||
@@ -100,10 +105,25 @@ test('enable builds index to completion and reports status', async () => {
|
||||
db.close()
|
||||
})
|
||||
|
||||
test('enable only marks a session enabled without starting a build', async () => {
|
||||
const { service, db } = setup({ embedDelayMs: 25 })
|
||||
|
||||
service.enable(SESSION_ID)
|
||||
|
||||
const status = service.status(SESSION_ID)!
|
||||
assert.equal(status.enabled, true)
|
||||
assert.equal(status.indexStatus, 'idle')
|
||||
assert.equal(status.queued, false)
|
||||
assert.equal(status.running, false)
|
||||
assert.equal(status.chunkCount, 0)
|
||||
|
||||
service.close()
|
||||
db.close()
|
||||
})
|
||||
|
||||
test('search returns evidence blocks for an enabled completed index', async () => {
|
||||
const { service, db } = setup()
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
|
||||
const result = await service.search(SESSION_ID, '项目排期')
|
||||
assert.equal(result.available, true)
|
||||
@@ -115,8 +135,7 @@ test('search returns evidence blocks for an enabled completed index', async () =
|
||||
|
||||
test('changing model identity marks needsRebuild and blocks search until rebuilt', async () => {
|
||||
const { service, db } = setup()
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
|
||||
// 切换到不同模型身份(改用 API)-> 身份变化 -> 需重建
|
||||
service.setConfig(
|
||||
@@ -144,8 +163,7 @@ test('changing model identity marks needsRebuild and blocks search until rebuilt
|
||||
|
||||
test('changing chunker identity marks needsRebuild and blocks search until rebuilt', async () => {
|
||||
const { service, dir, db } = setup()
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
assert.equal(service.status(SESSION_ID)!.needsRebuild, false)
|
||||
assert.equal(service.canSearch(SESSION_ID), true)
|
||||
|
||||
@@ -174,8 +192,7 @@ test('changing chunker identity marks needsRebuild and blocks search until rebui
|
||||
|
||||
test('buildAllPending rebuilds a stale index back to searchable', async () => {
|
||||
const { service, dir, db, getEmbedCount } = setup()
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
const afterFirstBuild = getEmbedCount()
|
||||
assert.ok(afterFirstBuild > 0)
|
||||
|
||||
@@ -200,10 +217,9 @@ test('buildAllPending rebuilds a stale index back to searchable', async () => {
|
||||
db.close()
|
||||
})
|
||||
|
||||
test('re-enabling a stale index rebuilds instead of resuming stale chunks', async () => {
|
||||
test('re-enabling a stale index keeps it pending until explicit rebuild', async () => {
|
||||
const { service, dir, db, getEmbedCount } = setup()
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
const afterFirstBuild = getEmbedCount()
|
||||
|
||||
const external = openBetterSqliteDatabase(path.join(dir, 'embedding_index.db'))
|
||||
@@ -212,11 +228,17 @@ test('re-enabling a stale index rebuilds instead of resuming stale chunks', asyn
|
||||
.run('stale-cfg', computeDbPathHash(SESSION_ID))
|
||||
external.close()
|
||||
|
||||
// 重新启用旧 stale 索引:必须走 rebuild(重新嵌入),而不是按旧游标续建跳过
|
||||
// 重新启用旧 stale 索引:只恢复 enabled,不自动消耗算力/API 额度。
|
||||
service.enable(SESSION_ID)
|
||||
|
||||
assert.equal(getEmbedCount(), afterFirstBuild)
|
||||
assert.equal(service.status(SESSION_ID)!.needsRebuild, true)
|
||||
assert.equal(service.canSearch(SESSION_ID), false)
|
||||
|
||||
service.buildAllPending()
|
||||
await service.whenIdle()
|
||||
|
||||
assert.ok(getEmbedCount() > afterFirstBuild, 're-enabling a stale index should rebuild, not resume-skip')
|
||||
assert.ok(getEmbedCount() > afterFirstBuild, 'explicit pending rebuild should re-embed chunks')
|
||||
assert.equal(service.status(SESSION_ID)!.needsRebuild, false)
|
||||
assert.equal(service.canSearch(SESSION_ID), true)
|
||||
service.close()
|
||||
@@ -225,8 +247,7 @@ test('re-enabling a stale index rebuilds instead of resuming stale chunks', asyn
|
||||
|
||||
test('disable then cleanup removes the index and blocks search', async () => {
|
||||
const { service, db } = setup()
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
|
||||
service.disable(SESSION_ID)
|
||||
const disabledStatus = service.status(SESSION_ID)!
|
||||
@@ -243,15 +264,13 @@ test('disable then cleanup removes the index and blocks search', async () => {
|
||||
|
||||
test('re-enabling after cleanup rebuilds from scratch', async () => {
|
||||
const { service, db, getEmbedCount } = setup()
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
const afterFirstBuild = getEmbedCount()
|
||||
|
||||
service.disable(SESSION_ID)
|
||||
service.cleanupUnused()
|
||||
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
|
||||
assert.ok(getEmbedCount() > afterFirstBuild, 're-enabling a cleaned index should embed chunks again')
|
||||
assert.equal(service.status(SESSION_ID)!.indexStatus, 'completed')
|
||||
@@ -264,8 +283,7 @@ test('re-enabling after cleanup rebuilds from scratch', async () => {
|
||||
|
||||
test('rebuild wipes and rebuilds the index', async () => {
|
||||
const { service, db } = setup()
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
const before = service.status(SESSION_ID)!.chunkCount
|
||||
|
||||
service.rebuild(SESSION_ID)
|
||||
@@ -297,6 +315,23 @@ test('setConfig stores API key by reference only and hasApiKey reflects it', asy
|
||||
db.close()
|
||||
})
|
||||
|
||||
test('keyless local Ollama API config can run without an auth profile', async () => {
|
||||
const { service, db } = setup()
|
||||
|
||||
service.setConfig({
|
||||
version: 1,
|
||||
mode: 'api',
|
||||
local: { modelId: QWEN3_PROFILE.modelId },
|
||||
api: { baseUrl: 'http://localhost:11434/v1', model: 'nomic-embed-text' },
|
||||
})
|
||||
|
||||
assert.equal(service.hasApiKey(), false)
|
||||
assert.equal(service.isConfigured(), true)
|
||||
assert.equal(service.canRun(), true)
|
||||
service.close()
|
||||
db.close()
|
||||
})
|
||||
|
||||
test('canSearch reflects availability: disabled/needs-rebuild/empty are false', async () => {
|
||||
const { service, db } = setup()
|
||||
|
||||
@@ -304,7 +339,10 @@ test('canSearch reflects availability: disabled/needs-rebuild/empty are false',
|
||||
assert.equal(service.canSearch(SESSION_ID), false)
|
||||
|
||||
service.enable(SESSION_ID)
|
||||
// 未完成建立前没有 chunk
|
||||
// 未显式建立前没有 chunk
|
||||
assert.equal(service.canSearch(SESSION_ID), false)
|
||||
|
||||
service.build(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
assert.equal(service.canSearch(SESSION_ID), true)
|
||||
|
||||
@@ -333,8 +371,7 @@ test('canSearch reflects availability: disabled/needs-rebuild/empty are false',
|
||||
|
||||
test('searchForTool desensitizes + anonymizes via pipeline, never leaks raw', async () => {
|
||||
const { service, db } = setup()
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
|
||||
const result = await service.searchForTool(SESSION_ID, '项目排期', {
|
||||
preprocessConfig: {
|
||||
@@ -372,6 +409,9 @@ test('searchForTool desensitizes + anonymizes via pipeline, never leaks raw', as
|
||||
assert.equal(s.snippet.includes('项目'), false)
|
||||
assert.equal(s.snippet.includes('张三'), false)
|
||||
assert.equal(s.snippet.includes('[Name Map]'), false)
|
||||
assert.equal(s.text?.includes('项目'), false)
|
||||
assert.equal(s.text?.includes('张三'), false)
|
||||
assert.equal(s.text?.includes('[Name Map]'), false)
|
||||
assert.ok(typeof s.startMessageId === 'number')
|
||||
assert.ok(Array.isArray(s.chunkIds))
|
||||
}
|
||||
@@ -390,8 +430,7 @@ test('searchForTool desensitizes + anonymizes via pipeline, never leaks raw', as
|
||||
|
||||
test('disabling the global switch blocks search and re-enabling keeps data usable', async () => {
|
||||
const { service, db } = setup()
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
assert.equal(service.canSearch(SESSION_ID), true)
|
||||
|
||||
// 关闭全局开关:不暴露工具 + 检索返回 disabled(已建索引数据保留)
|
||||
@@ -416,8 +455,7 @@ test('searchForTool is unavailable when disabled or empty query', async () => {
|
||||
assert.equal(disabled.available, false)
|
||||
assert.equal(disabled.reason, 'disabled')
|
||||
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
const empty = await service.searchForTool(SESSION_ID, ' ')
|
||||
assert.equal(empty.available, false)
|
||||
assert.equal(empty.reason, 'empty-query')
|
||||
@@ -427,8 +465,7 @@ test('searchForTool is unavailable when disabled or empty query', async () => {
|
||||
|
||||
test('searchForTool uses configured default max results when caller omits it', async () => {
|
||||
const { service, db } = setup()
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
service.setConfig({
|
||||
version: 1,
|
||||
mode: 'local',
|
||||
@@ -446,8 +483,7 @@ test('searchForTool uses configured default max results when caller omits it', a
|
||||
|
||||
test('searchForTool honors maxResultTokens as evidence budget ceiling', async () => {
|
||||
const { service, db } = setup()
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
|
||||
const tiny = await service.searchForTool(SESSION_ID, '项目排期', { maxResults: 10, maxResultTokens: 10 })
|
||||
const big = await service.searchForTool(SESSION_ID, '项目排期', { maxResults: 10, maxResultTokens: 100000 })
|
||||
@@ -463,8 +499,7 @@ test('searchForTool honors maxResultTokens as evidence budget ceiling', async ()
|
||||
|
||||
test('searchForTool applies timeFilter to exclude out-of-range chunks', async () => {
|
||||
const { service, db } = setup()
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
|
||||
// 不带时间过滤:应有命中
|
||||
const all = await service.searchForTool(SESSION_ID, '项目排期', { maxResults: 10 })
|
||||
@@ -485,8 +520,7 @@ test('searchForTool applies timeFilter to exclude out-of-range chunks', async ()
|
||||
|
||||
test('recover marks stale running as paused without auto-resuming', async () => {
|
||||
const { service, dir, db } = setup()
|
||||
service.enable(SESSION_ID)
|
||||
await service.whenIdle()
|
||||
await enableAndBuild(service)
|
||||
|
||||
// 模拟崩溃:另一连接把状态置为 running
|
||||
const external = new SemanticIndexStateStore(path.join(dir, 'embedding_index.db'))
|
||||
|
||||
@@ -27,7 +27,12 @@ import {
|
||||
STRATEGY_ID,
|
||||
} from './chunker-config'
|
||||
import type { ChunkSource } from './chunker'
|
||||
import { SemanticIndexConfigStore, type SemanticIndexConfig, type SemanticIndexConfigInput } from './config'
|
||||
import {
|
||||
isKeylessSemanticIndexApiBaseUrl,
|
||||
SemanticIndexConfigStore,
|
||||
type SemanticIndexConfig,
|
||||
type SemanticIndexConfigInput,
|
||||
} from './config'
|
||||
import { createEmbedder, type EmbedderFactoryDeps } from './embedder-factory'
|
||||
import type { EmbeddingProvider } from './embedding/types'
|
||||
import { EmbeddingIndexStore, type LoadSqliteVec } from './store'
|
||||
@@ -94,6 +99,7 @@ export interface SemanticSearchToolSource {
|
||||
score: number
|
||||
chunkIds: string[]
|
||||
snippet: string
|
||||
text?: string
|
||||
startTime?: string
|
||||
endTime?: string
|
||||
}
|
||||
@@ -229,7 +235,7 @@ export class SemanticIndexService {
|
||||
canRun(): boolean {
|
||||
if (!this.configStore.canRun()) return false
|
||||
const cfg = this.configStore.get()
|
||||
return cfg.mode !== 'api' || this.hasApiKey()
|
||||
return cfg.mode !== 'api' || isKeylessSemanticIndexApiBaseUrl(cfg.api?.baseUrl) || this.hasApiKey()
|
||||
}
|
||||
|
||||
/** 当前 API 模式是否已配置可用的 API Key(不返回明文,仅布尔) */
|
||||
@@ -323,7 +329,6 @@ export class SemanticIndexService {
|
||||
} else {
|
||||
this.stateStore.enable(this.enableParams(dbPath))
|
||||
}
|
||||
this.queue.enqueue({ type: stale ? 'rebuild' : 'build', dbPathHash: hash })
|
||||
}
|
||||
|
||||
disable(sessionId: string): void {
|
||||
@@ -609,6 +614,7 @@ export class SemanticIndexService {
|
||||
score: block.score,
|
||||
chunkIds: block.chunkIds,
|
||||
snippet,
|
||||
text: safeText,
|
||||
startTime,
|
||||
endTime,
|
||||
})
|
||||
@@ -656,11 +662,13 @@ export class SemanticIndexService {
|
||||
const activeSessionIds = new Set(this.sessionAdapter.listSessionIds())
|
||||
for (const state of this.stateStore.listAll()) {
|
||||
if (state.cleanupStatus !== 'pending' && !activeSessionIds.has(sessionIdFromDbPath(state.dbPath))) {
|
||||
this.queue.cancel(state.dbPathHash)
|
||||
this.stateStore.disable(state.dbPathHash)
|
||||
}
|
||||
}
|
||||
|
||||
for (const state of this.stateStore.listPendingCleanup()) {
|
||||
this.queue.cancel(state.dbPathHash)
|
||||
this.store.deleteByDbPathHash(state.dbPathHash)
|
||||
this.stateStore.resetProgress(state.dbPathHash)
|
||||
this.stateStore.setCleanupStatus(state.dbPathHash, 'done')
|
||||
|
||||
@@ -58,6 +58,27 @@ test('deleteByDbPathHash removes only the target conversation across models and
|
||||
store.close()
|
||||
})
|
||||
|
||||
test('deleteByModelFromPosition removes only chunks at and after a chat-order boundary', () => {
|
||||
const store = makeStore()
|
||||
store.insertChunk(record('a1', 'A', 'm1', 4, 1), [1, 0, 0, 0])
|
||||
store.insertChunk(record('a2', 'A', 'm1', 4, 3), [0, 1, 0, 0])
|
||||
store.insertChunk(record('a3', 'A', 'm1', 4, 5), [0, 0, 1, 0])
|
||||
store.insertChunk(record('other-model', 'A', 'm2', 4, 3), [0, 0, 0, 1])
|
||||
store.insertChunk(record('other-db', 'B', 'm1', 4, 3), [1, 1, 0, 0])
|
||||
|
||||
const removed = store.deleteByModelFromPosition({ dbPathHash: 'A', modelId: 'm1', startTs: 3000, startMessageId: 3 })
|
||||
assert.equal(removed, 2)
|
||||
assert.equal(store.countChunks('A', 'm1'), 1)
|
||||
assert.equal(store.countChunks('A', 'm2'), 1)
|
||||
assert.equal(store.countChunks('B', 'm1'), 1)
|
||||
assert.equal(store.getChunkById('a1')?.chunkId, 'a1')
|
||||
assert.equal(store.getChunkById('a2'), null)
|
||||
assert.equal(store.getChunkById('a3'), null)
|
||||
|
||||
assert.equal(store.queryDense({ dbPathHash: 'A', modelId: 'm1', dim: 4, embedding: [0, 1, 0, 0], k: 10 }).length, 1)
|
||||
store.close()
|
||||
})
|
||||
|
||||
test('listDbPathHashes returns distinct conversation hashes', () => {
|
||||
const store = makeStore()
|
||||
store.insertChunk(record('a1', 'A', 'm', 4, 1), [1, 0, 0, 0])
|
||||
|
||||
@@ -3,6 +3,7 @@ import fs from 'node:fs'
|
||||
import os from 'node:os'
|
||||
import path from 'node:path'
|
||||
import test from 'node:test'
|
||||
import Database from 'better-sqlite3'
|
||||
import { EmbeddingIndexStore } from './store'
|
||||
import type { ChunkRecord } from './types'
|
||||
|
||||
@@ -170,6 +171,21 @@ test('mapMessageToChunk breaks ties when multiple chunks share the same start_ts
|
||||
store.close()
|
||||
})
|
||||
|
||||
test('schema includes timestamp-leading index for message-to-chunk lookup', () => {
|
||||
const dbPath = makeTempDbPath()
|
||||
const store = new EmbeddingIndexStore(dbPath)
|
||||
store.close()
|
||||
|
||||
const db = new Database(dbPath, { readonly: true })
|
||||
const indexes = db.pragma('index_list(chunk_vector_index)') as { name: string }[]
|
||||
const timestampIndex = indexes.find((index) => index.name === 'idx_chunk_ts_range')
|
||||
assert.ok(timestampIndex, 'expected idx_chunk_ts_range to exist')
|
||||
|
||||
const columns = (db.pragma(`index_info(${timestampIndex.name})`) as { name: string }[]).map((column) => column.name)
|
||||
assert.deepEqual(columns, ['db_path_hash', 'model_id', 'strategy_id', 'start_ts', 'start_message_id'])
|
||||
db.close()
|
||||
})
|
||||
|
||||
test('data persists across store reopen', () => {
|
||||
const dbPath = makeTempDbPath()
|
||||
const store = new EmbeddingIndexStore(dbPath)
|
||||
|
||||
@@ -245,6 +245,50 @@ export class EmbeddingIndexStore {
|
||||
})()
|
||||
}
|
||||
|
||||
/** 删除某聊天库 + 模型中从指定聊天顺序位置开始的 chunks,返回删除的元数据行数 */
|
||||
deleteByModelFromPosition(params: {
|
||||
dbPathHash: string
|
||||
modelId: string
|
||||
startTs: number
|
||||
startMessageId: number
|
||||
}): number {
|
||||
return this.db.transaction(() => {
|
||||
const rows = this.db
|
||||
.prepare(
|
||||
`SELECT rowid AS rowid, dim FROM chunk_vector_index
|
||||
WHERE db_path_hash = ? AND model_id = ?
|
||||
AND (start_ts > CAST(? AS INTEGER)
|
||||
OR (start_ts = CAST(? AS INTEGER) AND start_message_id >= CAST(? AS INTEGER)))`
|
||||
)
|
||||
.all(params.dbPathHash, params.modelId, params.startTs, params.startTs, params.startMessageId) as Array<{
|
||||
rowid: number
|
||||
dim: number
|
||||
}>
|
||||
|
||||
const delStmts = new Map<number, Database.Statement>()
|
||||
for (const { rowid, dim } of rows) {
|
||||
const table = vecTableName(dim)
|
||||
if (!this.tableExists(table)) continue
|
||||
let stmt = delStmts.get(dim)
|
||||
if (!stmt) {
|
||||
stmt = this.db.prepare(`DELETE FROM ${table} WHERE vector_id = CAST(? AS INTEGER)`)
|
||||
delStmts.set(dim, stmt)
|
||||
}
|
||||
stmt.run(rowid)
|
||||
}
|
||||
|
||||
const res = this.db
|
||||
.prepare(
|
||||
`DELETE FROM chunk_vector_index
|
||||
WHERE db_path_hash = ? AND model_id = ?
|
||||
AND (start_ts > CAST(? AS INTEGER)
|
||||
OR (start_ts = CAST(? AS INTEGER) AND start_message_id >= CAST(? AS INTEGER)))`
|
||||
)
|
||||
.run(params.dbPathHash, params.modelId, params.startTs, params.startTs, params.startMessageId)
|
||||
return res.changes
|
||||
})()
|
||||
}
|
||||
|
||||
/** 列出存储中存在 chunk 的所有 db_path_hash */
|
||||
listDbPathHashes(): string[] {
|
||||
const rows = this.db.prepare('SELECT DISTINCT db_path_hash FROM chunk_vector_index').all() as Array<{
|
||||
|
||||
@@ -53,11 +53,13 @@ class FakeEmbedder implements EmbeddingProvider {
|
||||
readonly maxTokens = 1000
|
||||
calls = 0
|
||||
failAtCall?: number
|
||||
afterEmbed?: () => void
|
||||
|
||||
async embedDocuments(texts: string[]): Promise<Float32Array[]> {
|
||||
return texts.map(() => {
|
||||
this.calls++
|
||||
if (this.failAtCall && this.calls >= this.failAtCall) throw new Error('boom')
|
||||
this.afterEmbed?.()
|
||||
return new Float32Array([1, 0, 0, 0])
|
||||
})
|
||||
}
|
||||
@@ -86,6 +88,12 @@ function countStored(store: EmbeddingIndexStore): number {
|
||||
return store.queryDense({ dbPathHash: DB_HASH, modelId: MODEL, dim: 4, embedding: [1, 0, 0, 0], k: 1000 }).length
|
||||
}
|
||||
|
||||
function storedParentIds(store: EmbeddingIndexStore): string[] {
|
||||
return store
|
||||
.queryDense({ dbPathHash: DB_HASH, modelId: MODEL, dim: 4, embedding: [1, 0, 0, 0], k: 1000 })
|
||||
.map((result) => result.record.parentId)
|
||||
}
|
||||
|
||||
test('full warmup writes all chunks and marks completed', async () => {
|
||||
const { store, stateStore, source } = setup(makeMessages(8))
|
||||
const embedder = new FakeEmbedder()
|
||||
@@ -110,7 +118,7 @@ test('pause then resume completes without duplicate inserts', async () => {
|
||||
const embedder = new FakeEmbedder()
|
||||
|
||||
let calls = 0
|
||||
const pauseOnThird: StopSignal = () => (++calls >= 3 ? 'paused' : null)
|
||||
const pauseAfterTwoWrittenChunks: StopSignal = () => (++calls >= 5 ? 'paused' : null)
|
||||
const paused = await runWarmup({
|
||||
dbPathHash: DB_HASH,
|
||||
modelId: MODEL,
|
||||
@@ -119,7 +127,7 @@ test('pause then resume completes without duplicate inserts', async () => {
|
||||
stateStore,
|
||||
source,
|
||||
config,
|
||||
checkStop: pauseOnThird,
|
||||
checkStop: pauseAfterTwoWrittenChunks,
|
||||
})
|
||||
|
||||
assert.equal(paused.status, 'paused')
|
||||
@@ -159,6 +167,34 @@ test('cancellation stops immediately and records cancelled state', async () => {
|
||||
stateStore.close()
|
||||
})
|
||||
|
||||
test('cancellation after embedding does not write the just-finished chunk', async () => {
|
||||
const { store, stateStore, source } = setup(makeMessages(8))
|
||||
const embedder = new FakeEmbedder()
|
||||
|
||||
let cancelled = false
|
||||
embedder.afterEmbed = () => {
|
||||
cancelled = true
|
||||
}
|
||||
|
||||
const result = await runWarmup({
|
||||
dbPathHash: DB_HASH,
|
||||
modelId: MODEL,
|
||||
embedder,
|
||||
store,
|
||||
stateStore,
|
||||
source,
|
||||
config,
|
||||
checkStop: () => (cancelled ? 'cancelled' : null),
|
||||
})
|
||||
|
||||
assert.equal(result.status, 'cancelled')
|
||||
assert.equal(result.chunksWritten, 0)
|
||||
assert.equal(countStored(store), 0)
|
||||
assert.equal(stateStore.getState(DB_HASH)!.indexStatus, 'cancelled')
|
||||
store.close()
|
||||
stateStore.close()
|
||||
})
|
||||
|
||||
test('embedding failure marks failed and preserves partial progress', async () => {
|
||||
const { store, stateStore, source } = setup(makeMessages(8))
|
||||
const embedder = new FakeEmbedder()
|
||||
@@ -234,3 +270,43 @@ test('backfilled older messages trigger full re-index', async () => {
|
||||
store.close()
|
||||
stateStore.close()
|
||||
})
|
||||
|
||||
test('append-only warmup replaces stale chunks from the extended tail parent', async () => {
|
||||
const original = makeMessages(4)
|
||||
const { store, stateStore } = setup(original)
|
||||
const embedder = new FakeEmbedder()
|
||||
|
||||
await runWarmup({
|
||||
dbPathHash: DB_HASH,
|
||||
modelId: MODEL,
|
||||
embedder,
|
||||
store,
|
||||
stateStore,
|
||||
source: makeSource(original),
|
||||
config,
|
||||
})
|
||||
assert.equal(countStored(store), 2)
|
||||
assert.ok(storedParentIds(store).every((id) => id.includes('parent:1:4:')))
|
||||
|
||||
const appended = [
|
||||
...original,
|
||||
...makeMessages(2).map((message) => ({ ...message, id: message.id + 4, ts: message.ts + 4 * 0.1 * MINUTE })),
|
||||
]
|
||||
const result = await runWarmup({
|
||||
dbPathHash: DB_HASH,
|
||||
modelId: MODEL,
|
||||
embedder: new FakeEmbedder(),
|
||||
store,
|
||||
stateStore,
|
||||
source: makeSource(appended),
|
||||
config,
|
||||
})
|
||||
|
||||
assert.equal(result.status, 'completed')
|
||||
assert.equal(result.chunksWritten, 3)
|
||||
assert.equal(countStored(store), 3)
|
||||
assert.ok(storedParentIds(store).every((id) => id.includes('parent:1:6:')))
|
||||
assert.equal(stateStore.getState(DB_HASH)!.chunkCount, 3)
|
||||
store.close()
|
||||
stateStore.close()
|
||||
})
|
||||
@@ -9,7 +9,8 @@
|
||||
* - chunk 是纯函数且廉价,每次运行对全量消息重新 chunk(保证 parent 边界稳定)。
|
||||
* - 通过业务状态的 lastIndexedMessageId 游标跳过已写入 chunk,实现断点续跑。
|
||||
* - 每写入一个 chunk 即更新游标,保证崩溃后续跑不重复写入(chunk_id UNIQUE)。
|
||||
* - embedding 是瓶颈;每 chunk 前检查停止信号,暂停/取消可快速响应。
|
||||
* - embedding 是瓶颈;每 chunk 前检查暂停/取消,embedding 返回后再次检查取消,
|
||||
* 避免在清理已取消索引后写入刚完成的向量。
|
||||
*
|
||||
* 依赖全部注入,便于单测(fake source/embedder + 真实内存级 SQLite store)。
|
||||
*/
|
||||
@@ -58,13 +59,13 @@ export async function runWarmup(options: WarmupRunnerOptions): Promise<WarmupRes
|
||||
try {
|
||||
const messages = source.readAllMessages()
|
||||
const total = source.countMessages()
|
||||
const savedState = stateStore.getState(dbPathHash)
|
||||
stateStore.updateProgress(dbPathHash, { indexStatus: 'running', totalMessages: total, error: null })
|
||||
|
||||
// 消息 id -> 流位置,用于进度计数(消息按 ts 排序,id 未必单调)
|
||||
const streamIndexById = new Map<number, number>()
|
||||
messages.forEach((m, i) => streamIndexById.set(m.id, i))
|
||||
|
||||
const savedState = stateStore.getState(dbPathHash)
|
||||
const resumeMessageId = savedState?.lastIndexedMessageId ?? null
|
||||
const rawResumeIndex = resumeMessageId !== null ? (streamIndexById.get(resumeMessageId) ?? -1) : -1
|
||||
|
||||
@@ -75,16 +76,42 @@ export async function runWarmup(options: WarmupRunnerOptions): Promise<WarmupRes
|
||||
if (isNonAppendOnly) {
|
||||
store.deleteByDbPathHash(dbPathHash)
|
||||
}
|
||||
const resumeIndex = isNonAppendOnly ? -1 : rawResumeIndex
|
||||
|
||||
const { chunks, chunkerConfigHash } = chunkMessages({ messages, source: source.getSource(), config })
|
||||
|
||||
const chunkRanges = chunks.map((chunk) => ({
|
||||
chunk,
|
||||
startIndex: streamIndexById.get(chunk.startMessageId) ?? -1,
|
||||
endIndex: streamIndexById.get(chunk.endMessageId) ?? -1,
|
||||
}))
|
||||
|
||||
let resumeIndex = isNonAppendOnly ? -1 : rawResumeIndex
|
||||
const hasAppendedMessages = !isNonAppendOnly && rawResumeIndex >= 0 && total > (savedState?.totalMessages ?? 0)
|
||||
if (hasAppendedMessages) {
|
||||
const cursorRange =
|
||||
chunkRanges.find((range) => range.startIndex <= rawResumeIndex && rawResumeIndex <= range.endIndex) ??
|
||||
chunkRanges.filter((range) => range.endIndex <= rawResumeIndex).sort((a, b) => b.endIndex - a.endIndex)[0]
|
||||
|
||||
if (cursorRange) {
|
||||
const parentRanges = chunkRanges.filter((range) => range.chunk.parentId === cursorRange.chunk.parentId)
|
||||
const rewindIndex = Math.min(...parentRanges.map((range) => range.startIndex).filter((index) => index >= 0))
|
||||
const rewindMessage = messages[rewindIndex]
|
||||
if (rewindMessage) {
|
||||
store.deleteByModelFromPosition({
|
||||
dbPathHash,
|
||||
modelId,
|
||||
startTs: rewindMessage.ts,
|
||||
startMessageId: rewindMessage.id,
|
||||
})
|
||||
resumeIndex = rewindIndex - 1
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 续跑时统计已写入 chunk 数,保证 chunkCount 连续
|
||||
let storedChunkCount = stateStore.getState(dbPathHash)?.chunkCount ?? 0
|
||||
let storedChunkCount = store.countChunks(dbPathHash, modelId)
|
||||
if (resumeIndex < 0) storedChunkCount = 0
|
||||
|
||||
for (const chunk of chunks) {
|
||||
const chunkEndIndex = streamIndexById.get(chunk.endMessageId) ?? -1
|
||||
for (const { chunk, endIndex: chunkEndIndex } of chunkRanges) {
|
||||
if (chunkEndIndex <= resumeIndex) continue
|
||||
|
||||
const stop = checkStop?.()
|
||||
@@ -94,6 +121,12 @@ export async function runWarmup(options: WarmupRunnerOptions): Promise<WarmupRes
|
||||
}
|
||||
|
||||
const [vector] = await embedder.embedDocuments([chunk.embeddingInput])
|
||||
const stopAfterEmbedding = checkStop?.()
|
||||
if (stopAfterEmbedding === 'cancelled') {
|
||||
stateStore.setIndexStatus(dbPathHash, 'cancelled')
|
||||
return { status: 'cancelled', chunksWritten }
|
||||
}
|
||||
|
||||
const record: ChunkRecord = {
|
||||
chunkId: composeChunkId(dbPathHash, modelId, chunk.localChunkId),
|
||||
dbPathHash,
|
||||
|
||||
Reference in new issue
Block a user