287 lines
11 KiB
Python
287 lines
11 KiB
Python
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)
|