V1: DAG多任务编排 + AI规划WBS + RAG知识库 + 预算告警中心 + 开放API + 飞书/企微/邮件通知

2026-08-12 00:48:52 +08:00
parent e046770bad
commit 29874e6ca9
13 changed files with 1402 additions and 111 deletions
@@ -0,0 +1,119 @@
+# -*- 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