feat: add BookMate AI reading workspace and full EPUB reader

2026-10-24 03:12:20 +08:00
parent eaeb45fc85
commit c7b0cfb92b
60 changed files with 8996 additions and 191 deletions
@@ -0,0 +1,286 @@
+import re
+
+import jieba
+from sqlalchemy import select, or_
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.core.database import async_session
+from app.models.book import Book, Chapter, Segment
+from app.rag.vector import (
+    embed_texts, rerank_documents, load_vectors, vector_search,
+    rrf_fuse, ensure_embeddings, get_status, RERANK_ENABLED,
+)
+
+CITATION_RE = re.compile(r"\[\[ref:(\d+)\]\]")
+
+# 保守停用词：只过滤纯功能词，避免误删有检索价值的实词
+STOPWORDS = set(
+    "的 了 呢 吗 吧 啊 呀 嘛 哦 啦 哈 么 这 那 是 在 和 与 或 之 对 从 有 就 都 也 很 还 又 而 却 但 但是 然而 因此 所以 因为 如果 我们 你们 他们 它们 她们 什么 怎么 为什么 如何 这个 那个 一些 一个 本书 这本书 内容 主题 大概 比如 例如 应该 可能 可以 时候 里面 当中 其中 一些 部分 相关 关于 请问 告诉 解释 介绍 总结 分析".split()
+)
+
+
+def _tokenize(query: str) -> list[str]:
+    """分词：jieba 词 + 整句（精确子串）+ 双字 bigram，过滤停用词与纯符号。"""
+    tokens: set[str] = set()
+    for w in jieba.cut(query):
+        w = w.strip().lower()
+        if not w or w in STOPWORDS or re.fullmatch(r"[\s\W_]+", w):
+            continue
+        if len(w) >= 2:
+            tokens.add(w)
+    q = query.strip().lower()
+    if len(q) >= 4 and q not in STOPWORDS:
+        tokens.add(q)
+    # 中文连续串生成 bigram，增强长句召回（SQL LIKE 无法分词匹配）
+    for part in re.findall(r"[一-鿿]+", q):
+        if len(part) >= 4:
+            tokens.update(part[i:i + 2] for i in range(len(part) - 1))
+    # 过滤 bigram 中出现的停用字组合（如「的」开头的）
+    tokens = {t for t in tokens if not any(c in STOPWORDS for c in t) or len(t) >= 4}
+    return list(tokens)
+
+
+def match_chapters(chapters: list[Chapter], query: str, tokens: list[str] | None = None) -> list[Chapter]:
+    """章节标题匹配：查询词命中标题，或标题词命中查询。按目录顺序返回。"""
+    if not chapters:
+        return []
+    if tokens is None:
+        tokens = _tokenize(query)
+    q = query.lower()
+    hits = []
+    for ch in chapters:
+        title = (ch.title or "").strip().lower()
+        if not title or len(title) < 2:
+            continue
+        if any(w and w in title for w in tokens):
+            hits.append(ch)
+            continue
+        title_tokens = [
+            w.strip() for w in jieba.cut(title)
+            if len(w.strip()) >= 2 and w.strip() not in STOPWORDS
+        ]
+        if any(tw in q for tw in title_tokens):
+            hits.append(ch)
+    hits.sort(key=lambda c: c.order_idx)
+    return hits
+
+
+async def _keyword_candidates(db: AsyncSession, book_id: str, tokens: list[str], limit: int) -> list[Segment]:
+    """关键词候选：SQL LIKE 召回（含整句子串），按命中词数 + 长词加权打分。"""
+    if not tokens:
+        return []
+    ordered = sorted(tokens, key=len, reverse=True)[:8]
+    conds = [Segment.text.contains(w) for w in ordered]
+    q = (
+        select(Segment)
+        .where(Segment.book_id == book_id, or_(*conds))
+        .order_by(Segment.order_idx)
+        .limit(limit * 4)
+    )
+    result = await db.execute(q)
+    segs = list(result.scalars().all())
+    scored = []
+    for s in segs:
+        text = s.text.lower()
+        score = sum(3 if len(w) >= 4 else 1 for w in ordered if w in text)
+        scored.append((score, s))
+    scored.sort(key=lambda x: -x[0])
+    return [s for _, s in scored[:limit]]
+
+
+async def _empty_fallback(db: AsyncSession, book_id: str, chapters: list[Chapter], top_k: int) -> list[Segment]:
+    """空检索兜底：返回目录最靠前的章节的开头段（全书开头），让 AI 至少有所依据。"""
+    result = await db.execute(
+        select(Segment)
+        .where(Segment.book_id == book_id)
+        .order_by(Segment.order_idx)
+        .limit(top_k * 3)
+    )
+    segs = [s for s in result.scalars().all() if (s.text or "").strip()]
+    return segs[:top_k]
+
+
+async def search_segments(
+    db: AsyncSession,
+    book_id: str,
+    query: str,
+    top_k: int = 15,
+    chapters: list[Chapter] | None = None,
+) -> list[Segment]:
+    """混合检索：jieba 关键词 + API 向量余弦 → RRF 融合 → rerank 精排 → top_k。
+
+    未配置 / 未生成向量 / API 失败时自动降级为纯关键词检索，绝不空手返回。
+    """
+    tokens = _tokenize(query)
+
+    # 懒补向量（后台任务没跑完 / 重启后缺失时）
+    if get_status(book_id) != "done":
+        # 用独立会话执行：ensure_embeddings 失败时的 rollback 会过期会话内已加载的实例，
+        # 不能让调用方持有的 book/chapters 一起被过期（异步惰性访问会抛 MissingGreenlet）
+        async with async_session() as embed_session:
+            status = await ensure_embeddings(embed_session, book_id)
+        if status == "running":
+            # 后台正在生成，本次先用关键词，避免阻塞
+            status = "failed"
+
+    # 1. 关键词候选
+    kw_candidates = await _keyword_candidates(db, book_id, tokens, top_k * 3)
+
+    # 2. 向量候选
+    vec_candidates: list[Segment] = []
+    if tokens:
+        qvec = await embed_texts([query])
+        if qvec and qvec[0]:
+            mat, vec_segments = await load_vectors(db, book_id)
+            if mat is not None and len(mat) > 0:
+                idxs = vector_search(qvec[0], mat, top_k * 3)
+                vec_candidates = [vec_segments[i] for i in idxs]
+
+    # 3. RRF 融合
+    if kw_candidates and vec_candidates:
+        by_id = {s.id: s for s in kw_candidates + vec_candidates}
+        fused = rrf_fuse(
+            [s.id for s in kw_candidates],
+            [s.id for s in vec_candidates],
+        )
+        merged = [by_id[sid] for sid in fused if sid in by_id]
+    elif vec_candidates:
+        merged = vec_candidates
+    else:
+        merged = kw_candidates
+
+    # 4. rerank 精排（对融合后的 top50 段打分）
+    if merged and RERANK_ENABLED:
+        docs = [s.text[:800] for s in merged[:50]]
+        idxs = await rerank_documents(query, docs, top_n=top_k)
+        if idxs is not None:
+            merged = [merged[i] for i in idxs if i < len(merged)]
+
+    # 5. 兜底：完全没召回时给全书开头段，而不是空着
+    if not merged:
+        merged = await _empty_fallback(db, book_id, chapters or [], top_k)
+
+    return merged[:top_k]
+
+
+def _chapter_full_text(chapter_segments: list[Segment], budget_tokens: int = 6000) -> str:
+    """整章注入：按 token 预算截取命中章节的原文。"""
+    parts = []
+    used = 0
+    for seg in chapter_segments:
+        t = seg.text
+        cost = estimate_tokens(t)
+        if parts and used + cost > budget_tokens:
+            break
+        parts.append(t)
+        used += cost
+    return "\n\n".join(parts)
+
+
+def build_chat_prompt(
+    book: Book,
+    chapters: list[Chapter],
+    segments: list[Segment],
+    history: list[dict],
+    query: str,
+    chapter_segments: list[Segment] | None = None,
+) -> list[dict]:
+    system = f"""# 你是 BookMate，一个帮助用户阅读书籍的 AI 助手。
+
+## 当前书籍
+- 书名：《{book.title or '未知'}》
+- 作者：{book.author or '未知'}
+- 总字数：约 {book.char_count:,} 字
+
+## 目录结构
+"""
+    for ch in chapters[:40]:
+        indent = "  " * (ch.level - 1)
+        system += f"{indent}- {ch.title or '(无标题)'}\n"
+    if len(chapters) > 40:
+        system += f"  ...（共 {len(chapters)} 章）\n"
+
+    # 命中章节的整章原文：全局性问题（如「第3章讲了什么」）直接给全文
+    chapter_text = _chapter_full_text(chapter_segments or [])
+    if chapter_text:
+        system += "\n## 与问题直接相关的整章原文\n"
+        system += chapter_text + "\n"
+
+    if segments:
+        system += "\n## 检索到的相关原文段落\n"
+        for i, seg in enumerate(segments, 1):
+            text_preview = seg.text[:800].replace("\n", " ")
+            system += f"\n[[ref:{i}]] {text_preview}\n"
+    else:
+        system += (
+            "\n## 检索说明\n"
+            "本次没有检索到与问题直接匹配的原文段落。请基于上方目录结构、"
+            "书名与你的常识，指出书中哪些章节最可能涉及该主题，"
+            "并明确说明「这是基于目录的推断」，不要编造具体原文内容。\n"
+        )
+
+    system += """
+## 回答要求
+1. 优先基于「检索到的相关原文段落」和「整章原文」回答，不要臆造书里没有的内容
+2. 引用原文时必须带 [[ref:N]] 标记（N 对应上方段落编号）；同一段落只需引用一次，不要反复出现相同的 [[ref:N]]；整个回答的引用控制在 5 条左右，只引用你真正依据的段落
+3. 如果检索内容不足以回答，坦诚说明，并给出基于目录的建议
+4. 用中文回答
+5. 当用户要求「思维导图 / 脑图 / mind map」时，输出 mermaid mindmap 语法，放在 ```mermaid 代码块里。格式：第一行写 `mindmap`，根节点写 `root((主题))`，子节点通过缩进表示层级，例如：
+```mermaid
+mindmap
+  root((主题))
+    一级分支
+      二级分支
+```
+不要把其他解释文字塞进代码块，代码块外可以有简短说明。"""
+
+    messages = [{"role": "system", "content": system}]
+    if history:
+        messages.extend(history[-10:])
+    messages.append({"role": "user", "content": query})
+    return messages
+
+
+def _trim_messages(messages: list[dict], budget_tokens: int) -> list[dict]:
+    """内容感知裁剪：保留 system（含检索原文）与最新提问，从最旧历史开始丢弃。
+
+    替代原先粗糙的 `messages[:1] + messages[-(len//2):]`，避免裁掉检索上下文。
+    """
+    if estimate_tokens("".join(m["content"] for m in messages)) <= budget_tokens:
+        return messages
+
+    system = messages[0]
+    tail = messages[-1]
+    used = estimate_tokens(system["content"] + tail["content"])
+    kept = []
+    for m in reversed(messages[1:-1]):
+        cost = estimate_tokens(m["content"])
+        if used + cost <= budget_tokens:
+            kept.append(m)
+            used += cost
+        else:
+            break
+    kept.reverse()
+    return [system] + kept + [tail]
+
+
+def parse_citations(content: str, segments: list[Segment]) -> dict:
+    citations = []
+    for m in CITATION_RE.finditer(content):
+        idx = int(m.group(1)) - 1
+        if 0 <= idx < len(segments):
+            seg = segments[idx]
+            citations.append({
+                "segment_id": seg.id,
+                "cfi": seg.cfi,
+                "excerpt": seg.text[:200],
+            })
+    clean_text = CITATION_RE.sub("", content)
+    return {"text": clean_text, "citations": citations}
+
+
+def estimate_tokens(text: str) -> int:
+    """粗略估算 token 数：中文每字约 1.5 token，英文每词约 1.3 token"""
+    chinese_chars = len(re.findall(r"[一-鿿]", text))
+    english_words = len(re.findall(r"[a-zA-Z]+", text))
+    return int(chinese_chars * 1.5 + english_words * 1.3 + (len(text) - chinese_chars - len(re.findall(r"[a-zA-Z]", text))) * 0.5)