120 lines
4.5 KiB
Python
120 lines
4.5 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
RAG 知识库：文档分块 + jieba 分词 + BM25 检索（本地零依赖）
|
||
升级路径：接入商用 embedding + pgvector/Qdrant 做向量检索（见 README）
|
||
"""
|
||
import math
|
||
import re
|
||
import jieba
|
||
import db
|
||
|
||
CHUNK_SIZE = 400 # 每块目标字数
|
||
CHUNK_OVERLAP = 80 # 块间重叠
|
||
|
||
|
||
def segment(text):
|
||
"""jieba 分词，去停用词/单字/空白"""
|
||
text = re.sub(r'[\s，。！？、；：""''（）【】《》·…—0-9a-zA-Z-]+', ' ', text)
|
||
words = [w for w in jieba.cut(text) if len(w.strip()) > 1 and not w.isspace()]
|
||
return words
|
||
|
||
def split_chunks(content, size=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
|
||
"""按段落聚合切块，避免从句子中间切断"""
|
||
paras = [p.strip() for p in re.split(r'\n+', content) if p.strip()]
|
||
chunks, buf, buf_len = [], '', 0
|
||
for p in paras:
|
||
if buf_len + len(p) > size and buf:
|
||
chunks.append(buf)
|
||
tail = buf[-overlap:] if overlap else ''
|
||
buf, buf_len = tail + p, len(tail) + len(p)
|
||
else:
|
||
buf += ('\n' if buf else '') + p
|
||
buf_len += len(p)
|
||
if buf:
|
||
chunks.append(buf)
|
||
return chunks or [content[:size]]
|
||
|
||
|
||
def rebuild_document(doc_id):
|
||
"""重新分块索引文档"""
|
||
doc = db.q('SELECT * FROM documents WHERE id=?', (doc_id,), one=True)
|
||
if not doc:
|
||
return 0
|
||
db.w('DELETE FROM doc_chunks WHERE document_id=?', (doc_id,))
|
||
chunks = split_chunks(doc['content'])
|
||
for i, c in enumerate(chunks):
|
||
db.w('INSERT INTO doc_chunks (document_id, idx, content, tokens) VALUES (?,?,?,?)',
|
||
(doc_id, i, c, len(segment(c))))
|
||
db.w('UPDATE documents SET chunk_size=?, updated_at=? WHERE id=?', (len(chunks), db.now(), doc_id))
|
||
return len(chunks)
|
||
|
||
|
||
class BM25Index:
|
||
"""轻量 BM25：按查询词 IDF 加权打分"""
|
||
|
||
def __init__(self, chunks):
|
||
# chunks: [{id, document_id, content, tokens}]
|
||
self.chunks = chunks
|
||
self.doc_len = [max(c['tokens'], 1) for c in chunks]
|
||
self.avg_len = sum(self.doc_len) / max(len(self.doc_len), 1)
|
||
self.N = len(chunks)
|
||
self.k1, self.b = 1.5, 0.75
|
||
# 倒排：term -> set of chunk idx
|
||
self.inv = {}
|
||
self.df = {}
|
||
for i, c in enumerate(chunks):
|
||
seen = set()
|
||
for w in segment(c['content']):
|
||
if w in seen:
|
||
continue
|
||
seen.add(w)
|
||
self.inv.setdefault(w, []).append(i)
|
||
for w, lst in self.inv.items():
|
||
self.df[w] = len(lst)
|
||
|
||
def search(self, query, top_k=5):
|
||
q_terms = [w for w in segment(query)]
|
||
if not q_terms or not self.N:
|
||
return []
|
||
scores = {}
|
||
for w in q_terms:
|
||
postings = self.inv.get(w, [])
|
||
if not postings:
|
||
continue
|
||
idf = math.log(1 + (self.N - self.df[w] + 0.5) / (self.df[w] + 0.5))
|
||
for idx in postings:
|
||
tf = sum(1 for x in segment(self.chunks[idx]['content']) if x == w)
|
||
denom = tf + self.k1 * (1 - self.b + self.b * self.doc_len[idx] / self.avg_len)
|
||
scores[idx] = scores.get(idx, 0) + idf * (tf * (self.k1 + 1)) / denom
|
||
ranked = sorted(scores.items(), key=lambda x: -x[1])[:top_k]
|
||
return [{'chunk_id': self.chunks[i]['id'], 'document_id': self.chunks[i]['document_id'],
|
||
'content': self.chunks[i]['content'], 'score': round(s, 4)}
|
||
for i, s in ranked]
|
||
|
||
|
||
def search_project(project_id, query, top_k=5):
|
||
"""在项目知识库中检索，返回 (片段列表, 是否命中)"""
|
||
chunks = db.q(
|
||
'SELECT c.id, c.document_id, c.content, c.tokens FROM doc_chunks c '
|
||
'JOIN documents d ON d.id=c.document_id WHERE d.project_id=?',
|
||
(project_id,))
|
||
if not chunks:
|
||
return [], False
|
||
idx = BM25Index(chunks)
|
||
return idx.search(query, top_k), True
|
||
|
||
|
||
def build_context(project_id, query, top_k=4):
|
||
"""生成注入提示词的检索上下文（含来源标注）"""
|
||
hits, hit = search_project(project_id, query, top_k)
|
||
if not hit or not hits:
|
||
return '', []
|
||
parts, refs = [], []
|
||
for h in hits:
|
||
doc = db.q('SELECT name FROM documents WHERE id=?', (h['document_id'],), one=True)
|
||
name = doc['name'] if doc else f'文档#{h["document_id"]}'
|
||
parts.append(f'【来源：{name}】\n{h["content"]}')
|
||
refs.append(f'{name}#块{h["chunk_id"]}')
|
||
ctx = '以下是项目知识库中的相关资料，回答时请优先参考：\n\n' + '\n\n---\n\n'.join(parts)
|
||
return ctx, refs
|