bgeEmbedder.ts 的源代码：
/**
* BgeEmbeddingFunction —— BGE 中文 embedding 适配器
* ============================================================
*
* 继承官方 @chroma-core/default-embed 的 DefaultEmbeddingFunction，
* 补一个 generateForQueries 方法处理 BGE 查询前缀。
*
* 【BGE 的前缀规则】（非对称检索模型特性）
* - 文档端（入库）：不加前缀，原文向量化
* - 查询端（检索）：加官方前缀 "为这个句子生成表示以用于检索相关文章："
* 不加前缀，模型不知道当前是"查询"，检索质量打折。
*
* 【chromadb 的自动分流机制】（源码：chromadb.mjs:3976-3979 的 embed() 方法）
* - collection.add() 调 generate()（文档端，不加前缀）
* - collection.query() 优先调 generateForQueries()（查询端，加前缀）
* 所以只要实现 generateForQueries，chromadb 内部自动分流，调用方无感。
*
* 【一致性约束】
* 写入端和查询端必须用同一个模型名 + dtype，否则向量空间不匹配。
* 换模型 = 必须重建整个库。
*/
import { DefaultEmbeddingFunction } from '@chroma-core/default-embed'
/** BGE 官方查询前缀（中文版，必须一字不差） */
export const BGE_QUERY_PREFIX = '为这个句子生成表示以用于检索相关文章：'
export class BgeEmbeddingFunction extends DefaultEmbeddingFunction {
/**
* 查询路径：chromadb query 时自动调用，加 BGE 前缀。
* 给每条 query 拼上前缀，再交给父类的 generate（它已封装好向量化全流程）。
*/
async generateForQueries(texts: string[]): Promise<number[][]> {
return super.generate(texts.map(t => BGE_QUERY_PREFIX + t))
}
// generate()（文档端）直接继承父类，不加前缀
}
loadNovelsFromDir.ts 的源代码：
import { readdirSync, readFileSync } from 'node:fs'
import { join, basename, extname } from 'node:path'
export interface NovelDocument {
id: string
source: string
text: string
}
/**
* 读取目录下所有 .txt 文件，组装成 NovelDocument 数组
* - id：文件名去掉 .txt 扩展名（同时作为 chunk id 前缀和 metadata.book）
* - source：原始文件名（含 .txt）
* - text：文件全文，utf-8 解码
*
* 仅处理 .txt，其他文件忽略；目录为空时返回空数组。
*/
export function loadNovelsFromDir(dir: string): NovelDocument[] {
const entries = readdirSync(dir, { withFileTypes: true })
const docs: NovelDocument[] = []
for (const entry of entries) {
// 只处理文件，跳过子目录
if (!entry.isFile()) continue
if (extname(entry.name).toLowerCase() !== '.txt') continue
const fullPath = join(dir, entry.name)
const text = readFileSync(fullPath, 'utf-8')
const id = basename(entry.name, extname(entry.name))
docs.push({ id, source: entry.name, text })
}
return docs
}
simpleChroma.ts 的源代码：
/**
* SimpleChroma —— 自封装的 Chroma LangChain 向量库集成（生产级）
* ============================================================
*
* 【对齐的官方 VectorStore 契约】（不继承抽象类，但接口对齐——鸭子类型）
* - addDocuments：写入（支持增量去重）
* - similaritySearch / similaritySearchWithScore：检索
* - delete：按 ids 或 filter 删除
* - asRetriever：转 Runnable，能接 LCEL 链
* 额外补了 count（查总数，增量逻辑要用）
*
* 【核心设计：embedding 模型一致性】
* 写入和查询必须用同一个 embedding 模型 + 同一个 dtype。
* 本类接受任意 EmbeddingFunction（chromadb 的接口），只要写入端和查询端传同一个实例。
* 配套的 BgeEmbeddingFunction 见本文件底部（或 8.3/8.4 各自有一份）。
*
* 【增量去重】
* addDocuments 支持传 ids 参数。如果传了，会先查集合里已存在的 id，跳过已存在的，
* 只写入新的。这样重复运行不会产生重复 chunk（解决"每次跑都新增"的问题）。
*/
import { ChromaClient, type Collection, type Where } from 'chromadb'
import { Document } from '@langchain/core/documents'
import { Runnable, RunnableLambda } from '@langchain/core/runnables'
// ============================================================================
// 配置类型
// ============================================================================
/** SimpleChroma 构造配置 */
export interface SimpleChromaConfig {
/** Chroma 服务 host */
host: string
/** Chroma 服务 port */
port: number
/** Basic Auth 凭据（明文，内部会转 base64） */
username: string
password: string
/** 集合名 */
collectionName: string
/** embedding 函数（写入和查询必须用同一个实例） */
embedder: {
generate(texts: string[]): Promise<number[][]>
generateForQueries?(texts: string[]): Promise<number[][]>
}
/** 是否用 ssl（默认 false） */
ssl?: boolean
}
/** addDocuments 的选项 */
export interface AddDocumentsOptions {
/**
* 每个 doc 的 id。如果传了，会做增量去重（已存在的 id 跳过）。
* 如果不传，自动用 `chunk_${timestamp}_${i}` 生成（不去重，每次都新增）。
*/
ids?: string[]
}
/** delete 的参数：ids 和 filter 至少传一个 */
export interface DeleteParams {
/** 按 id 删除（精确删除指定 chunk） */
ids?: string[]
/** 按 metadata 条件删除（Chroma Where 语法，如 { source: 'xxx.txt' } 或 { $or: [...] }） */
filter?: Where
}
// ============================================================================
// SimpleChroma 主类
// ============================================================================
/**
* 自封装的 Chroma LangChain 向量库集成。
*
* 用法：
* ```ts
* const store = new SimpleChroma({
* host: '192.168.31.249', port: 7777,
* username: 'admin', password: '123456',
* collectionName: 'my_kb',
* embedder: new BgeEmbeddingFunction({ modelName: 'Xenova/bge-small-zh-v1.5', dtype: 'fp16' }),
* })
*
* // 建库（增量去重）
* await store.addDocuments(docs, { ids: docs.map((_, i) => `chunk_${i}`) })
*
* // 检索
* const hits = await store.similaritySearchWithScore('问题', 3)
*
* // 删除
* await store.delete({ filter: { source: 'xxx.txt' } })
*
* // 接 LCEL 链
* const retriever = store.asRetriever(3)
* const chain = retriever.pipe(...)
* ```
*/
export class SimpleChroma {
private client: ChromaClient
private collectionName: string
private embedder: SimpleChromaConfig['embedder']
private collectionCache?: Collection
constructor(config: SimpleChromaConfig) {
this.collectionName = config.collectionName
this.embedder = config.embedder
const auth = Buffer.from(`${config.username}:${config.password}`).toString('base64')
this.client = new ChromaClient({
ssl: config.ssl ?? false,
host: config.host,
port: config.port,
headers: { Authorization: `Basic ${auth}` },
})
}
/** 拿到（或创建）集合，结果缓存避免重复请求 */
private async getCollection(): Promise<Collection> {
if (this.collectionCache) return this.collectionCache
try {
this.collectionCache = await this.client.getCollection({ name: this.collectionName })
} catch {
this.collectionCache = await this.client.getOrCreateCollection({ name: this.collectionName })
}
return this.collectionCache
}
/** 重置集合缓存（delete 等操作后集合引用可能失效时调用） */
private resetCollectionCache() {
this.collectionCache = undefined
}
/**
* 查询集合中已存在的 id 集合（用于增量去重）。
* 注意：这里拉所有 id 到内存。如果集合非常大（百万级），需改用更分批的策略。
*/
private async getExistingIds(): Promise<Set<string>> {
const collection = await this.getCollection()
const result = await collection.get()
return new Set(result.ids ?? [])
}
/**
* ✅ 官方契约：写入文档（建库用）。
*
* @param docs LangChain Document 数组
* @param options.ids 可选，传了会增量去重（已存在的 id 跳过）
* @returns 实际写入的 id 数组（已跳过已存在的）
*/
async addDocuments(docs: Document[], options?: AddDocumentsOptions): Promise<string[]> {
const collection = await this.getCollection()
// 确定 id：传了用传的，没传自动生成（不去重）
const useDedup = !!options?.ids
const ids = options?.ids ?? docs.map((_, i) => `chunk_${Date.now()}_${i}`)
// 增量去重：过滤掉已存在的 id
let pendingIndices: number[]
let pendingIds: string[]
if (useDedup) {
const existingIds = await this.getExistingIds()
pendingIndices = []
pendingIds = []
for (let i = 0; i < ids.length; i++) {
if (!existingIds.has(ids[i])) {
pendingIndices.push(i)
pendingIds.push(ids[i])
}
}
if (pendingIndices.length === 0) {
return [] // 全部已存在，无需写入
}
} else {
pendingIndices = docs.map((_, i) => i)
pendingIds = ids
}
const pendingDocs = pendingIndices.map(i => docs[i])
const texts = pendingDocs.map(d => d.pageContent)
const vectors = await this.embedder.generate(texts)
// 过滤 metadata：Chroma 只支持 string/number/boolean 基本类型
// RecursiveCharacterTextSplitter 会加 loc（对象，记录行号），要剔除
const cleanMetadatas = pendingDocs.map(d => filterBasicMetadata(d.metadata))
await collection.add({
ids: pendingIds,
documents: texts,
embeddings: vectors,
metadatas: cleanMetadatas,
})
return pendingIds
}
/**
* ✅ 官方契约：检索（带距离）。距离越小越相关。
* 查询端如果 embedder 实现了 generateForQueries，会调用它（BGE 用来加查询前缀）。
*/
async similaritySearchWithScore(query: string, k: number): Promise<[Document, number][]> {
const collection = await this.getCollection()
// 查询端：优先用 generateForQueries（BGE 加前缀），没有就 fallback 到 generate
const queryVectors = this.embedder.generateForQueries
? await this.embedder.generateForQueries([query])
: await this.embedder.generate([query])
const queryVector = queryVectors[0]
const result = await collection.query({
queryEmbeddings: [queryVector],
nResults: k,
})
// chromadb v3 返回 QueryResult 类：documents/metadatas/distances 都是二维数组（取 [0]）
const docs = result.documents[0] ?? []
const metas = result.metadatas[0] ?? []
const dists = result.distances[0] ?? []
const hits: [Document, number][] = []
for (let i = 0; i < docs.length; i++) {
const text = docs[i]
if (text == null) continue
hits.push([
new Document({
pageContent: text,
metadata: (metas[i] as Record<string, unknown>) ?? {},
}),
dists[i] ?? Infinity,
])
}
return hits
}
/** ✅ 官方契约：检索（不带距离），返回纯 Document 数组。 */
async similaritySearch(query: string, k: number): Promise<Document[]> {
const hits = await this.similaritySearchWithScore(query, k)
return hits.map(([doc]) => doc)
}
/**
* ✅ 官方契约：删除文档。支持按 ids 或按 filter 删除，至少传一个。
*/
async delete(params: DeleteParams): Promise<void> {
const collection = await this.getCollection()
if (params.ids && params.ids.length > 0) {
await collection.delete({ ids: params.ids })
} else if (params.filter) {
await collection.delete({ where: params.filter })
} else {
throw new Error('delete 至少要传 ids 或 filter 之一')
}
}
/**
* ✅ 官方契约：转 Runnable（对齐 asRetriever）。
* 用 RunnableLambda 包一层 similaritySearch，返回的 Runnable 能 pipe 进 LCEL 链。
* 输入：查询字符串；输出：Document[]。
*/
asRetriever(k = 3): Runnable<string, Document[]> {
const self = this
return RunnableLambda.from(async (query: string) => {
return self.similaritySearch(query, k)
})
}
/**
* 查询集合中的文档总数。
* 增量逻辑判断、监控统计常用。
*/
async count(): Promise<number> {
const collection = await this.getCollection()
return await collection.count()
}
/**
* 删除整个集合并重置缓存（慎用，不可恢复）。
* 用于"重建知识库"场景：先 dropCollection，再重新 addDocuments。
*/
async dropCollection(): Promise<void> {
try {
await this.client.deleteCollection({ name: this.collectionName })
} catch {
// 集合本来就不存在，忽略
}
this.resetCollectionCache()
}
}
// ============================================================================
// 辅助函数
// ============================================================================
/** 过滤 metadata：只保留 string/number/boolean 基本类型（Chroma 的要求） */
function filterBasicMetadata(meta?: Record<string, unknown>): Record<string, string | number | boolean> {
const clean: Record<string, string | number | boolean> = {}
if (!meta) return clean
for (const [k, v] of Object.entries(meta)) {
if (typeof v === 'string' || typeof v === 'number' || typeof v === 'boolean') {
clean[k] = v
}
}
return clean
}
// ============================================================================
// 配套：BgeEmbeddingFunction（BGE 中文 embedding 适配器）
// ============================================================================
/**
* BGE embedding 适配器：继承官方 DefaultEmbeddingFunction，补一个 generateForQueries。
* - generate() 继承父类：文档端用，不加前缀
* - generateForQueries()：查询端用，加 BGE 官方前缀
*
* 必须和写入端用同一个模型名 + dtype，否则向量空间不匹配。
* 建库和查询共用同一个实例。
*/
export { BgeEmbeddingFunction } from './bgeEmbedder'
splitText.ts 的源代码：
interface SplitTextOptions {
/** 每个块的最大字符数（非严格，边界模式可能略少） */
chunkSize?: number
/** 相邻两块之间的重叠字符数，用于保证跨块上下文不丢失 */
overlap?: number
/** 是否在切割时尽量回退到句子/词边界，保证语义单元的完整性 */
respectBoundary?: boolean
/** 边界模式下，允许 end 最多往回退的比例（相对于 chunkSize），防止过度回退产生极小碎块，默认 0.5 */
maxBacktrackRatio?: number
}
/**
* 将长文本切分成多个带有重叠的文本块，专为 RAG（检索增强生成）和向量化场景设计
*
* 核心机制：
* 1. 滑动窗口 + 重叠（overlap），保证上下文连续性
* 2. 支持语义边界回退（句号、换行），避免切断句子
* 3. 严格限制回退范围（maxBacktrackRatio），防止碎块过多导致效率低下
* 4. 完整的边界安全（Unicode 代理对保护、除零/死循环防御）
*/
export function splitText(text: string, options: SplitTextOptions = {}): string[] {
const {
chunkSize = 500,
overlap = 50,
respectBoundary = false,
maxBacktrackRatio = 0.5,
} = options
// --- 1. 防御性校验：彻底杜绝死循环与非法入参 ---
if (!Number.isFinite(chunkSize) || chunkSize <= 0) {
throw new Error(`chunkSize 必须为正数，当前为 ${chunkSize}`)
}
if (!Number.isFinite(overlap) || overlap < 0 || overlap >= chunkSize) {
throw new Error(`overlap 必须满足 0 <= overlap < chunkSize，当前 overlap=${overlap}`)
}
if (text.length === 0) return []
const chunks: string[] = []
// 最大允许回退的字符数（用于边界模式下限制搜索范围）
const maxBacktrack = Math.floor(chunkSize * maxBacktrackRatio)
// 非边界模式下的固定步长
const fixedStep = chunkSize - overlap
let start = 0
while (start < text.length) {
// --- 2. 确定初始结束位置（硬切） ---
let end = Math.min(start + chunkSize, text.length)
// --- 3. Unicode 安全切割：避免在 UTF-16 代理对（如 Emoji）中间截断 ---
// 如果 end 指向低位代理（Low Surrogate），则回退一格，保证不拆散字符
if (end < text.length && isLowSurrogate(text.charCodeAt(end))) {
end -= 1
}
// --- 4. 语义边界回退（可选） ---
if (respectBoundary && end < text.length) {
// 设定回退的底线：至少要比 start 大（避免原地踏步），且不能退过 maxBacktrack 范围
const floor = Math.max(start + 1, end - maxBacktrack)
// 在 (start, end] 范围内寻找最优的边界（句号、换行、空格等）
const boundary = findBoundary(text, start, end)
// 如果找到的边界在底线之上，则采用，否则宁可硬切（保底策略）
if (boundary >= floor) {
end = boundary
}
// 如果 boundary < floor，说明最优边界离当前 end 太远了，强行回退会导致块太小，
// 此时保留硬切的 end，避免产生大量无意义的碎块。
}
// --- 5. 截取并存储当前块 ---
chunks.push(text.slice(start, end))
// --- 6. 判断是否已经到达末尾，跳出循环 ---
if (end >= text.length) break
// --- 7. 计算下一次循环的起始位置（核心策略） ---
if (respectBoundary) {
// 【关键修复】当开启边界模式时，end 可能向左回退。
// 为了确保 overlap 精确，下一块的起点必须基于实际的 end 反推。
// 这样即使 end 回退了，下一块的开头也会相应提前，保留足够的重叠文本。
const nextStart = end - overlap
// 兜底保护：万一 overlap 为 0 或极端情况导致 nextStart 不大于 start，
// 强制向前推进至少 1 个字符，避免死循环。
start = nextStart > start ? nextStart : start + 1
} else {
// 非边界模式下，按固定步长滑动（普通硬切模式）
start += fixedStep
}
}
return chunks
}
/**
* 判断某个 Unicode 码点是否为 UTF-16 低位代理（Low Surrogate）
* 用于检测是否切中了 Emoji 或特殊字符的内部，防止乱码
*/
function isLowSurrogate(code: number): boolean {
return code >= 0xdc00 && code <= 0xdfff
}
/**
* 在 (start, end] 区间内，从 end 开始向前寻找最近的语义边界
* 优先级：换行/句子终结符 > 空格/逗号等分隔符
*
* 注意：搜索范围包含 end 位置本身的字符，修正了 V2 版本遗漏 end 处标点的 bug
*/
function findBoundary(text: string, start: number, end: number): number {
// 限定搜索上界，防止 end 等于 text.length 时越界
const upperBound = Math.min(end, text.length - 1)
// 第一优先级：段落/句子终结符（中英文标点、换行）
for (let i = upperBound; i > start; i--) {
if (/[\n。！？!?；;]/.test(text[i])) {
// 返回边界的下一个索引（将标点符号保留在当前块中，或者归入下一块开头？此处返回 i+1，
// 意味着标点符号被归入当前块末尾，符合大多数场景的阅读习惯）
return i + 1
}
}
// 第二优先级：词/词组分隔符（空格、逗号、顿号）
for (let i = upperBound; i > start; i--) {
if (/[\s,，、]/.test(text[i])) {
return i + 1
}
}
// 找不到合适边界：返回原 end，保持硬切
return end
}
💬 评论 0
还没有评论,快来抢沙发吧~