181 lines
6.9 KiB
JavaScript
181 lines
6.9 KiB
JavaScript
'use strict';
|
||
|
||
const fs = require('fs');
|
||
const path = require('path');
|
||
const crypto = require('crypto');
|
||
const { parseFrontmatter, splitChunks, createParentChildChunks } = require('./kb-chunker');
|
||
|
||
// 新 bridge 独立 KB 索引器：读本地 kb/*.md，生成 parent/child + embedding，
|
||
// staging 表校验后原子切换 kb_docs，旧表保留为 kb_docs_v<timestamp>。
|
||
function createKbIndexer({ pool, embeddingClient, kbDir, dim = 1024, logger = console }) {
|
||
const DIM = Number(dim) || 1024;
|
||
const KB_DIR = kbDir || path.resolve(__dirname, '..', 'kb');
|
||
|
||
async function ensureBase() {
|
||
await pool.query('CREATE EXTENSION IF NOT EXISTS vector');
|
||
await pool.query(`
|
||
CREATE TABLE IF NOT EXISTS kb_sync_runs (
|
||
id SERIAL PRIMARY KEY,
|
||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||
source_hash TEXT,
|
||
parent_count INT,
|
||
child_count INT,
|
||
backup_table TEXT,
|
||
note TEXT
|
||
)
|
||
`);
|
||
}
|
||
|
||
function listFiles() {
|
||
return fs.readdirSync(KB_DIR)
|
||
.filter((f) => f.endsWith('.md') && /^\d+/.test(f))
|
||
.sort();
|
||
}
|
||
|
||
function sourceHash(files) {
|
||
const h = crypto.createHash('sha256');
|
||
for (const f of files) {
|
||
h.update(f);
|
||
h.update(fs.readFileSync(path.join(KB_DIR, f)));
|
||
}
|
||
return h.digest('hex');
|
||
}
|
||
|
||
async function createStaging() {
|
||
await pool.query('DROP TABLE IF EXISTS kb_docs_staging CASCADE');
|
||
await pool.query(`
|
||
CREATE TABLE kb_docs_staging (
|
||
id SERIAL PRIMARY KEY,
|
||
source_file TEXT,
|
||
title TEXT,
|
||
category TEXT,
|
||
section TEXT,
|
||
question TEXT,
|
||
answer TEXT,
|
||
keywords TEXT[],
|
||
source_urls TEXT[],
|
||
text_vec vector(${DIM}),
|
||
parent_id INTEGER,
|
||
chunk_type TEXT DEFAULT 'parent',
|
||
keywords_text TEXT,
|
||
text_tsv tsvector
|
||
)
|
||
`);
|
||
}
|
||
|
||
async function insertDoc(table, doc, parentId) {
|
||
const text = [doc.title, doc.category, doc.section, doc.question, doc.answer, (doc.keywords || []).join(' ')].filter(Boolean).join(' ');
|
||
const r = await pool.query(
|
||
`INSERT INTO ${table}
|
||
(source_file,title,category,section,question,answer,keywords,source_urls,text_vec,parent_id,chunk_type,keywords_text,text_tsv)
|
||
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9::vector,$10,$11,$12,to_tsvector('english',$13))
|
||
RETURNING id`,
|
||
[doc.source_file, doc.title, doc.category, doc.section, doc.question, doc.answer,
|
||
doc.keywords || [], doc.source_urls || [], embeddingClient.vecStr(doc.vec), parentId || null,
|
||
doc.chunk_type, (doc.keywords || []).join(' '), text]
|
||
);
|
||
return r.rows[0].id;
|
||
}
|
||
|
||
async function indexFile(table, file) {
|
||
const raw = fs.readFileSync(path.join(KB_DIR, file), 'utf8');
|
||
const { meta, body } = parseFrontmatter(raw);
|
||
const chunks = splitChunks(body);
|
||
let parents = 0;
|
||
let children = 0;
|
||
for (const chunk of chunks) {
|
||
const { parent, children: childChunks } = createParentChildChunks(chunk, meta, file);
|
||
parent.vec = await embeddingClient.embedQuery(parent.embText);
|
||
const parentId = await insertDoc(table, parent, null);
|
||
parents += 1;
|
||
for (const child of childChunks) {
|
||
child.vec = await embeddingClient.embedQuery(child.embText);
|
||
await insertDoc(table, child, parentId);
|
||
children += 1;
|
||
}
|
||
}
|
||
return { parents, children };
|
||
}
|
||
|
||
async function createStagingIndexes() {
|
||
await pool.query('CREATE INDEX kb_docs_staging_vec_idx ON kb_docs_staging USING hnsw (text_vec vector_cosine_ops)');
|
||
await pool.query('CREATE INDEX kb_docs_staging_tsv_idx ON kb_docs_staging USING GIN (text_tsv)');
|
||
await pool.query('CREATE INDEX kb_docs_staging_parent_idx ON kb_docs_staging (parent_id) WHERE parent_id IS NOT NULL');
|
||
}
|
||
|
||
async function atomicSwitch() {
|
||
const client = await pool.connect();
|
||
try {
|
||
await client.query('BEGIN');
|
||
const has = await client.query("SELECT to_regclass('public.kb_docs') IS NOT NULL AS exists");
|
||
let backup = null;
|
||
if (has.rows[0].exists) {
|
||
const stamp = new Date().toISOString().replace(/[-:TZ.]/g, '').slice(0, 14);
|
||
backup = `kb_docs_v${stamp}`;
|
||
let suffix = 0;
|
||
while ((await client.query('SELECT to_regclass($1) IS NOT NULL AS exists', [`public.${backup}`])).rows[0].exists) {
|
||
suffix += 1;
|
||
backup = `kb_docs_v${stamp}_${suffix}`;
|
||
}
|
||
await client.query(`ALTER TABLE kb_docs RENAME TO ${backup}`);
|
||
await client.query(`ALTER INDEX IF EXISTS kb_docs_vec_idx RENAME TO ${backup}_vec_idx`);
|
||
await client.query(`ALTER INDEX IF EXISTS kb_docs_tsv_idx RENAME TO ${backup}_tsv_idx`);
|
||
await client.query(`ALTER INDEX IF EXISTS kb_docs_parent_idx RENAME TO ${backup}_parent_idx`);
|
||
}
|
||
await client.query('ALTER TABLE kb_docs_staging RENAME TO kb_docs');
|
||
await client.query('ALTER INDEX IF EXISTS kb_docs_staging_vec_idx RENAME TO kb_docs_vec_idx');
|
||
await client.query('ALTER INDEX IF EXISTS kb_docs_staging_tsv_idx RENAME TO kb_docs_tsv_idx');
|
||
await client.query('ALTER INDEX IF EXISTS kb_docs_staging_parent_idx RENAME TO kb_docs_parent_idx');
|
||
await client.query('COMMIT');
|
||
return backup;
|
||
} catch (err) {
|
||
await client.query('ROLLBACK');
|
||
throw err;
|
||
} finally {
|
||
client.release();
|
||
}
|
||
}
|
||
|
||
async function reindex({ note = 'kb sync' } = {}) {
|
||
const lockClient = await pool.connect();
|
||
const lock = await lockClient.query("SELECT pg_try_advisory_lock(hashtext('bridge_kb_sync')) AS locked");
|
||
if (!lock.rows[0].locked) {
|
||
lockClient.release();
|
||
throw new Error('another KB sync is running');
|
||
}
|
||
try {
|
||
await ensureBase();
|
||
const files = listFiles();
|
||
if (!files.length) throw new Error(`no KB markdown files in ${KB_DIR}`);
|
||
const hash = sourceHash(files);
|
||
await createStaging();
|
||
let parents = 0;
|
||
let children = 0;
|
||
for (const file of files) {
|
||
const res = await indexFile('kb_docs_staging', file);
|
||
parents += res.parents;
|
||
children += res.children;
|
||
}
|
||
if (parents === 0) throw new Error('indexed 0 parent chunks');
|
||
await createStagingIndexes();
|
||
const backup = await atomicSwitch();
|
||
await pool.query(
|
||
'INSERT INTO kb_sync_runs (source_hash,parent_count,child_count,backup_table,note) VALUES ($1,$2,$3,$4,$5)',
|
||
[hash, parents, children, backup, note]
|
||
);
|
||
logger.log(JSON.stringify({ ts: new Date().toISOString(), kb_sync: 'done', parents, children, backup, files: files.length }));
|
||
return { parents, children, backup, sourceHash: hash, files: files.length };
|
||
} catch (err) {
|
||
await pool.query('DROP TABLE IF EXISTS kb_docs_staging CASCADE').catch(() => {});
|
||
throw err;
|
||
} finally {
|
||
await lockClient.query("SELECT pg_advisory_unlock(hashtext('bridge_kb_sync'))").catch(() => {});
|
||
lockClient.release();
|
||
}
|
||
}
|
||
|
||
return { reindex, listFiles, sourceHash, KB_DIR };
|
||
}
|
||
|
||
module.exports = { createKbIndexer };
|