129 lines
3.7 KiB
TypeScript
129 lines
3.7 KiB
TypeScript
import fs from 'node:fs';
|
||
import path from 'node:path';
|
||
|
||
/** 知识库中的一个文档块 */
|
||
interface Chunk {
|
||
/** 来源文件名 */
|
||
source: string;
|
||
/** 块内文本 */
|
||
content: string;
|
||
}
|
||
|
||
/**
|
||
* 轻量知识库：从 knowledge/ 目录加载 .md 文件，按 ## 标题分块，
|
||
* 提供基于关键词匹配的检索能力。
|
||
*
|
||
* 用法：
|
||
* const kb = new KnowledgeBase('./knowledge');
|
||
* await kb.load();
|
||
* const results = kb.search('Agent 工具调用');
|
||
*/
|
||
export class KnowledgeBase {
|
||
private chunks: Chunk[] = [];
|
||
private knowledgeDir: string;
|
||
|
||
constructor(knowledgeDir: string) {
|
||
this.knowledgeDir = knowledgeDir;
|
||
}
|
||
|
||
/** 加载 knowledgeDir 下所有 .md 文件并分块 */
|
||
async load(): Promise<void> {
|
||
this.chunks = [];
|
||
|
||
if (!fs.existsSync(this.knowledgeDir)) {
|
||
console.warn(`知识库目录不存在: ${this.knowledgeDir}`);
|
||
return;
|
||
}
|
||
|
||
const files = fs
|
||
.readdirSync(this.knowledgeDir)
|
||
.filter((f) => f.endsWith('.md'));
|
||
|
||
for (const file of files) {
|
||
const filePath = path.join(this.knowledgeDir, file);
|
||
const raw = fs.readFileSync(filePath, 'utf-8');
|
||
const fileChunks = this.splitChunks(raw, file);
|
||
this.chunks.push(...fileChunks);
|
||
}
|
||
|
||
console.log(`知识库已加载: ${this.chunks.length} 个块，来自 ${files.length} 个文件`);
|
||
}
|
||
|
||
/** 按 ## 标题将文档拆分为块 */
|
||
private splitChunks(raw: string, source: string): Chunk[] {
|
||
const blocks = raw.split(/(?=^## )/m);
|
||
return blocks
|
||
.map((b) => b.trim())
|
||
.filter(Boolean)
|
||
.map((content) => ({ source, content }));
|
||
}
|
||
|
||
/**
|
||
* 基于关键词匹配搜索，返回相关块（按相关性降序）
|
||
*
|
||
* 算法：将查询分词，统计每个块命中关键词的次数，
|
||
* 同时给标题匹配额外加权。
|
||
*/
|
||
search(query: string, topK: number = 3): Chunk[] {
|
||
const keywords = this.tokenize(query);
|
||
if (keywords.length === 0) return [];
|
||
|
||
const scored = this.chunks.map((chunk) => {
|
||
const lower = chunk.content.toLowerCase();
|
||
let score = 0;
|
||
for (const kw of keywords) {
|
||
// 标题行命中加权 ×3
|
||
const headlineRegex = /^## .+$/gm;
|
||
let match: RegExpExecArray | null;
|
||
while ((match = headlineRegex.exec(chunk.content)) !== null) {
|
||
if (match[0].toLowerCase().includes(kw)) {
|
||
score += 3;
|
||
}
|
||
}
|
||
// 正文命中
|
||
const count = (lower.match(new RegExp(this.escapeRegex(kw), 'gi')) || []).length;
|
||
score += count;
|
||
}
|
||
// 标题匹配额外加分
|
||
return { chunk, score };
|
||
});
|
||
|
||
return scored
|
||
.filter((s) => s.score > 0)
|
||
.sort((a, b) => b.score - a.score)
|
||
.slice(0, topK)
|
||
.map((s) => s.chunk);
|
||
}
|
||
|
||
/** 中文 + 英文简单分词 */
|
||
private tokenize(text: string): string[] {
|
||
// 按空白/标点拆分，过滤长度 ≤1 的词
|
||
return text
|
||
.split(/[\s,，。.！!？?：:；;、]+/)
|
||
.map((t) => t.toLowerCase().trim())
|
||
.filter((t) => t.length > 1);
|
||
}
|
||
|
||
private escapeRegex(s: string): string {
|
||
return s.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
|
||
}
|
||
|
||
/** 获取知识库摘要（供 Agent 概览） */
|
||
summary(): string {
|
||
if (this.chunks.length === 0) return '知识库为空';
|
||
const sources = [...new Set(this.chunks.map((c) => c.source))];
|
||
const titles = this.chunks
|
||
.map((c) => {
|
||
const m = c.content.match(/^## (.+)$/m);
|
||
return m ? ` - ${m[1]} (${c.source})` : null;
|
||
})
|
||
.filter(Boolean);
|
||
return `知识库包含 ${sources.length} 个文档:\n${titles.join('\n')}`;
|
||
}
|
||
}
|
||
|
||
/** 全局单例 */
|
||
export const knowledgeBase = new KnowledgeBase(
|
||
path.join(process.cwd(), 'knowledge'),
|
||
);
|