ai-worker-platform/kb.py

# -*- coding: utf-8 -*-
"""
V3.5.2 全局知识库
=================
- kb_documents / kb_chunks：文档 + 分块（jieba 分词，BM25 风格检索）
- 基本功能：增删改查、上传（txt/md/pdf）、全文检索、上下文注入（对话可选）
"""
import json
import re
import os

import db
import config

KB_UPLOAD_DIR = os.path.join(config.DATA_DIR, 'kb_uploads')


def _tok(text):
    """jieba 分词（去停用字、只留长度>=2 的 token）"""
    try:
        import jieba
        toks = []
        for t in jieba.cut_for_search((text or '').lower()):
            t = t.strip()
            if len(t) >= 2 and not t.isdigit():
                toks.append(t)
        return toks
    except Exception:
        return [w for w in re.findall(r'[\u4e00-\u9fff]{2,}|[a-zA-Z0-9_]{2,}', (text or '').lower())]


def _chunks(content, size=400, overlap=60):
    """把文档切成小块（按段落聚合 + 超长硬切 + 前后重叠）"""
    content = content or ''
    paras = [p for p in re.split(r'\n+', content) if p.strip()]
    blocks, buf = [], ''
    for p in paras:
        if buf and len(buf) + len(p) > size:
            blocks.append(buf)
            buf = ''
        buf = (buf + '\n' + p) if buf else p
    if buf:
        blocks.append(buf)
    out = []
    for b in blocks:
        while len(b) > size:
            out.append(b[:size])
            b = b[size - overlap:]
        if b:
            out.append(b)
    return out or ['']


def rebuild_chunks(doc_id):
    doc = db.q('SELECT * FROM kb_documents WHERE id=?', (doc_id,), one=True)
    if not doc:
        return 0
    db.w('DELETE FROM kb_chunks WHERE doc_id=?', (doc_id,))
    n = 0
    for i, c in enumerate(_chunks(doc.get('content') or '')):
        db.w('INSERT INTO kb_chunks (doc_id, idx, content, tokens) VALUES (?,?,?,?)',
             (doc_id, i, c, json.dumps(_tok(c))))
        n += 1
    return n


def search(q, top_k=6):
    """BM25 风格检索：返回 [{doc_id,title,content,score}]（按命中 token 数 + IDF 加权）"""
    q_tokens = _tok(q)
    if not q_tokens:
        return []
    rows = db.q('SELECT * FROM kb_chunks ORDER BY doc_id, idx')
    if not rows:
        return []
    docs = {d['id']: d for d in db.q('SELECT id,title FROM kb_documents')}
    df = {}
    for c in rows:
        for t in set(json.loads(c['tokens'] or '[]')):
            df[t] = df.get(t, 0) + 1
    n_docs = max(1, len(set(r['doc_id'] for r in rows)))
    scored = []
    for c in rows:
        toks = json.loads(c['tokens'] or '[]')
        tf = {}
        for t in toks:
            tf[t] = tf.get(t, 0) + 1
        score = 0.0
        for t in q_tokens:
            if t in tf:
                score += (1 + tf[t]) * max(0.1, (n_docs - df.get(t, 0) + 0.5) / (df.get(t, 0) + 0.5))
        if score > 0:
            scored.append({'doc_id': c['doc_id'], 'idx': c['idx'],
                           'content': c['content'], 'score': round(score, 3),
                           'title': docs.get(c['doc_id'], {}).get('title', '')})
    scored.sort(key=lambda x: -x['score'])
    return scored[:top_k]


def build_context(q, top_k=4):
    """把检索结果拼成可注入的上下文（用于对话/任务），返回 (ctx, hits)"""
    hits = search(q, top_k)
    if not hits:
        return '', []
    parts = []
    for i, h in enumerate(hits):
        parts.append(f"[{i + 1}]《{h['title']}》\n{h['content'][:900]}")
    ctx = ('以下是与你问题相关的【知识库参考】资料（可据此回答）：\n' + '\n\n'.join(parts) + '\n\n----\n')
    return ctx, hits


def extract_text(filename, raw):
    """按扩展名抽取文本：txt/md/html/csv/json；pdf 用 pypdf（有则装）。返回 (text, ok)"""
    ext = os.path.splitext(filename)[1].lower()
    name = filename or 'doc'
    if ext in ('.txt', '.md', '.markdown', '.html', '.htm', '.csv', '.json', '.log', '.py', '.js', '.css'):
        for enc in ('utf-8', 'gbk', 'utf-8-sig'):
            try:
                return raw.decode(enc), True
            except Exception:
                continue
        return raw.decode('utf-8', errors='ignore'), True
    if ext == '.pdf':
        try:
            from pypdf import PdfReader
            import io
            reader = PdfReader(io.BytesIO(raw))
            text = '\n'.join((pg.extract_text() or '') for pg in reader.pages)
            return text, bool(text.strip())
        except Exception:
            return '', False
    if ext in ('.docx',):
        try:
            import io
            from docx import Document
            doc = Document(io.BytesIO(raw))
            text = '\n'.join(p.text for p in doc.paragraphs)
            return text, bool(text.strip())
        except Exception:
            return '', False
    return '', False