Pine
064b06ecf0
迁入 app/park：llm(对话)、rag(双路向量知识库)、tools/ai_tools(智能体工具)、 asr(语音识别)、s2s_bridge(实时语音桥)、vision_yolo/vision_llm(人脸/多模态)、 knowledge/*.md、vendor/s2s-cloud(s2s 云化栈)；routers 补 /api/ai|kb|asr|vision| s2s|tools 端点。智能体提示词/企业名录改读 park_config 主数据源。pyproject 加重依赖 (dashscope/numpy/openai/torch/transformers/ultralytics/websockets/soundfile/scipy/ nltk/jinja2)。TestClient 冒烟：ai/chat(无 Key 走本地规则)、kb、s2s、tools、display 均 200。
461 lines
18 KiB
Python
461 lines
18 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""园区知识库 —— 轻量本地 RAG（embedding + numpy 余弦检索，双路索引）
|
||
|
||
升级为「细粒度分块 + LLM 生成问答对 + 双路向量索引（原句 + 问题）+ 余弦检索返回原文」：
|
||
- 知识源：backend/knowledge/park.md（按 H2 节 → 原子事实两级分块）
|
||
- 向量化：dashscope text-embedding-v4（云端 API，索引构建一次后缓存，构建时 L2 归一化）
|
||
- 问答：复用 llm._chat_once 对每块生成 3-5 个口语化问句，缓存到 kb_qa.json（无 key 自动单路降级）
|
||
- 检索：query 同时比对原句向量与问题向量，按块去重，返回 top_k 个原文块
|
||
- 提供：
|
||
- retrieve(query, top_k) → 检索命中的原文块（对话动态注入）
|
||
- full_kb() → 整库原文（小库全量注入，供语音 instructions）
|
||
- brief() → 前 N 字摘要（兼容 /api/kb/brief 与检索无命中回退）
|
||
"""
|
||
import hashlib
|
||
import json
|
||
import logging
|
||
import os
|
||
import re
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
|
||
log = logging.getLogger("dpm.rag")
|
||
|
||
|
||
def _trunc(s, n=300):
|
||
"""日志用截断：超长内容只留前 n 字并标注总长。"""
|
||
s = str(s)
|
||
return s if len(s) <= n else s[:n] + f"...（共{len(s)}字）"
|
||
|
||
|
||
KB_DIR = Path(__file__).resolve().parent / "knowledge"
|
||
KB_GLOB = "*.md" # 知识库 = knowledge/ 下全部 markdown（park.md + opc.md 等）
|
||
META_FILE = KB_DIR / "kb_meta.json" # 元数据（JSON，可读）：model/hash/chunks/questions/embed_texts/embed_owner
|
||
VECTOR_FILE = KB_DIR / "kb_vectors.npy" # 向量矩阵（二进制 numpy float32，快载小体积）
|
||
LEGACY_INDEX_FILE = KB_DIR / "kb_index.json" # 旧版单文件索引（迁移后删除）
|
||
QA_FILE = KB_DIR / "kb_qa.json"
|
||
EMBED_MODEL = "text-embedding-v4"
|
||
EMBED_BATCH = 10 # DashScope embedding 单批上限
|
||
QA_BATCH = 5 # 每批几块生成问答
|
||
CHUNK_MAX = 240 # 单块超过此长度按句拆
|
||
|
||
_index = None # 缓存 {chunks, questions, embed_texts, embed_owner, embeddings(归一化 np.ndarray)}
|
||
|
||
|
||
def _kb_text() -> str:
|
||
"""拼接 knowledge/ 下全部 markdown 内容（按文件名排序，稳定顺序）。"""
|
||
parts = []
|
||
for f in sorted(KB_DIR.glob(KB_GLOB)):
|
||
try:
|
||
parts.append(f.read_text("utf-8"))
|
||
except OSError as e: # noqa: BLE001
|
||
log.warning("rag: 跳过知识文件 %s: %s", f.name, e)
|
||
return "\n\n".join(parts)
|
||
|
||
|
||
def _data_hash() -> str:
|
||
"""全部知识文件内容签名，用于索引/问答缓存失效判断。"""
|
||
return hashlib.sha1(_kb_text().encode("utf-8")).hexdigest()
|
||
|
||
|
||
def _split_long(text: str, limit: int = CHUNK_MAX) -> list[str]:
|
||
"""超长块按句号/分号/顿号等边界拆成多块，保留语义完整。"""
|
||
if len(text) <= limit:
|
||
return [text]
|
||
parts: list[str] = []
|
||
cur = ""
|
||
for seg in re.split(r"(?<=[。！？；、])", text):
|
||
if len(cur) + len(seg) > limit and cur:
|
||
parts.append(cur.strip())
|
||
cur = ""
|
||
cur += seg
|
||
if cur.strip():
|
||
parts.append(cur.strip())
|
||
return [p for p in parts if len(p) > 8]
|
||
|
||
|
||
def _chunk_md(text: str) -> list[str]:
|
||
"""两级切分：H2 节标题作上下文前缀 + 节内原子事实。
|
||
|
||
企业条目（加粗公司名行 `**...｜公司名**` + 负责人/简介列表）合并为单个块，
|
||
保证「企业名称 + 企业介绍」作为一个整体可被按名召回。
|
||
"""
|
||
chunks: list[str] = []
|
||
section = ""
|
||
para: list[str] = []
|
||
company: list[str] | None = None # 正在累积的单个企业块
|
||
|
||
def flush_para():
|
||
nonlocal para
|
||
if para:
|
||
t = "\n".join(para).strip()
|
||
if len(t) > 8:
|
||
chunks.extend(_split_long(t))
|
||
para = []
|
||
|
||
def flush_company():
|
||
nonlocal company
|
||
if company:
|
||
t = "\n".join(company).strip()
|
||
if len(t) > 8:
|
||
chunks.append(t) # 企业块不拆，保证名称+介绍同块
|
||
company = None
|
||
|
||
def add_fact(content: str):
|
||
full = f"{section}\n{content}" if section else content
|
||
if len(full) > 8:
|
||
chunks.extend(_split_long(full))
|
||
|
||
for line in text.splitlines():
|
||
s = line.rstrip()
|
||
if s.startswith("# "): # H1：不产块
|
||
flush_para(); flush_company()
|
||
continue
|
||
if s.startswith("## "): # H2 小节：切换上下文前缀
|
||
flush_para(); flush_company()
|
||
section = s[3:].strip()
|
||
continue
|
||
if s.startswith("###"): # H3 子节：上下文（如分区），flush 企业
|
||
flush_para(); flush_company()
|
||
para.append(s.strip())
|
||
continue
|
||
stripped = s.strip()
|
||
if not stripped:
|
||
continue
|
||
# 企业条目：加粗公司名行作为单个企业块的起点
|
||
if stripped.startswith("**") and "｜" in stripped:
|
||
flush_para(); flush_company()
|
||
company = [stripped]
|
||
continue
|
||
if company is not None:
|
||
# 企业块的跟随行（负责人/简介等列表项）
|
||
if stripped.startswith("- ") or stripped.startswith("问：") or re.match(r"^\d+[\.、]\s", stripped):
|
||
company.append(stripped)
|
||
else:
|
||
flush_company()
|
||
para.append(stripped)
|
||
continue
|
||
if stripped.startswith("- "): # 列表项 = 原子事实
|
||
flush_para(); add_fact(stripped)
|
||
elif re.match(r"^\d+[\.、]\s", stripped): # 编号步
|
||
flush_para(); add_fact(stripped)
|
||
elif stripped.startswith("问："): # 单条问答
|
||
flush_para(); add_fact(stripped)
|
||
else: # 普通段落：累积后整段成块（超长再拆）
|
||
para.append(stripped)
|
||
flush_para()
|
||
flush_company()
|
||
return chunks
|
||
|
||
|
||
def _company_name(chunk: str) -> str:
|
||
"""从企业块提取企业名称：`**加速1｜云南派音人工智能科技有限公司**` → `云南派音人工智能科技有限公司`。"""
|
||
m = re.match(r"^\*\*[^｜\n]*｜([^*]+)\*\*", chunk) or re.match(r"^\*\*([^*]+)\*\*", chunk)
|
||
return m.group(1).strip() if m else ""
|
||
|
||
|
||
def _embed(texts: list[str]) -> list[list[float]]:
|
||
"""dashscope text-embedding-v4 批量向量化（每批 ≤10 条，按 text_index 排序拼接）。"""
|
||
import dashscope
|
||
|
||
api_key = os.environ.get("DASHSCOPE_API_KEY") or ""
|
||
out: list[list[float]] = []
|
||
for i in range(0, len(texts), EMBED_BATCH):
|
||
batch = texts[i:i + EMBED_BATCH]
|
||
resp = dashscope.TextEmbedding.call(
|
||
model=EMBED_MODEL,
|
||
input=batch,
|
||
api_key=api_key,
|
||
)
|
||
if resp.status_code != 200:
|
||
raise RuntimeError(
|
||
f"embedding API {resp.status_code}: {resp.code} {resp.message}"
|
||
)
|
||
embs = sorted(
|
||
resp.output["embeddings"],
|
||
key=lambda e: e.get("text_index", e.get("index", 0)),
|
||
)
|
||
out.extend(e["embedding"] for e in embs)
|
||
log.info("rag: embedding 完成 %d 条（%s，每批 %d）", len(texts), EMBED_MODEL, EMBED_BATCH)
|
||
return out
|
||
|
||
|
||
def _normalize(vectors: list[list[float]]) -> np.ndarray:
|
||
"""转为 float32 np.ndarray 并做 L2 行归一化（点积即余弦相似度）。"""
|
||
arr = np.asarray(vectors, dtype=np.float32)
|
||
norms = np.linalg.norm(arr, axis=1, keepdims=True)
|
||
norms[norms == 0] = 1.0
|
||
return arr / norms
|
||
|
||
|
||
# ── LLM 问答生成 ────────────────────────────────────────────────
|
||
|
||
def _parse_json_array(text: str) -> list:
|
||
"""解析模型返回的 JSON 数组：剥 markdown 围栏 → json.loads → 截取首尾方括号兜底。"""
|
||
text = (text or "").strip()
|
||
if text.startswith("```"):
|
||
lines = [l for l in text.splitlines() if not l.strip().startswith("```")]
|
||
text = "\n".join(lines).strip()
|
||
for candidate in (text,):
|
||
try:
|
||
return json.loads(candidate)
|
||
except Exception:
|
||
pass
|
||
i = text.find("[")
|
||
j = text.rfind("]")
|
||
if i != -1 and j > i:
|
||
try:
|
||
return json.loads(text[i:j + 1])
|
||
except Exception:
|
||
pass
|
||
return []
|
||
|
||
|
||
def _qa_for_batch(chunks: list[str]) -> list[list[str]]:
|
||
"""对一批分块，用 LLM 生成每块的 3-5 个口语化问句（输出 JSON 数组，与块一一对应）。"""
|
||
from .llm import _chat_once
|
||
|
||
numbered = "\n\n".join(f"[{i}]\n{c}" for i, c in enumerate(chunks))
|
||
|
||
def build_prompt():
|
||
return (
|
||
"下面是园区知识库的若干分块，每块标有 [编号]。"
|
||
"请为每一块生成 3-5 个用户可能会问的自然口语化中文问句，"
|
||
"覆盖该块全部信息点与常见同义问法。"
|
||
f"严格只输出一个 JSON 数组，数组长度必须为 {len(chunks)}，"
|
||
"第 i 个元素是第 i 块的问题字符串数组。不要输出任何其他文字或解释。\n\n"
|
||
+ numbered
|
||
)
|
||
|
||
for attempt in range(2):
|
||
try:
|
||
data = _chat_once([{"role": "user", "content": build_prompt()}], with_tools=False)
|
||
content = data["choices"][0]["message"].get("content") or ""
|
||
arr = _parse_json_array(content)
|
||
if isinstance(arr, list) and arr:
|
||
out: list[list[str]] = []
|
||
for i in range(len(chunks)):
|
||
qs = arr[i] if i < len(arr) and isinstance(arr[i], list) else []
|
||
out.append([str(q).strip() for q in qs if str(q).strip()][:5])
|
||
return out
|
||
except Exception as e: # noqa: BLE001
|
||
log.warning("rag: QA 批生成失败（尝试 %d/%d）: %s", attempt + 1, 2, e)
|
||
return [[] for _ in chunks]
|
||
|
||
|
||
def _first_line(chunk: str) -> str:
|
||
"""取块首行（去标题），用作无问句时的兜底检索问题。"""
|
||
for ln in chunk.splitlines():
|
||
ln = ln.strip()
|
||
if ln and not ln.startswith("#"):
|
||
return ln
|
||
return chunk[:60]
|
||
|
||
|
||
def _ensure_questions(chunks: list[str], questions) -> list[list[str]]:
|
||
"""确保每块都有至少一个检索问题：未生成问句的块，用块首句兜底（保证全量覆盖）。"""
|
||
qs: list[list[str]] = [list(q or []) for q in (questions or [])]
|
||
if len(qs) < len(chunks):
|
||
qs.extend([] for _ in range(len(chunks) - len(qs)))
|
||
for i, c in enumerate(chunks):
|
||
if not qs[i]:
|
||
qs[i] = [_first_line(c)]
|
||
return qs
|
||
|
||
|
||
def _load_questions(chunks: list[str]) -> list[list[str]]:
|
||
"""加载/生成问答缓存；返回与 chunks 对齐的问题数组（保证每块 ≥1 个问题）。无 key 时单路回退。"""
|
||
from .config import settings
|
||
|
||
if not settings.DASHSCOPE_API_KEY:
|
||
log.warning("rag: 无 DASHSCOPE_API_KEY，跳过问答生成（单路检索）")
|
||
return _ensure_questions(chunks, [[] for _ in chunks])
|
||
h = _data_hash()
|
||
if QA_FILE.exists():
|
||
try:
|
||
data = json.loads(QA_FILE.read_text("utf-8"))
|
||
if data.get("data_hash") == h and data.get("chunks") == chunks:
|
||
log.info("rag: 问答缓存已加载（%d 块）", len(chunks))
|
||
return _ensure_questions(chunks, data["questions"])
|
||
except Exception: # noqa: BLE001
|
||
log.warning("rag: 问答缓存损坏，重建")
|
||
questions: list[list[str]] = []
|
||
total = len(chunks)
|
||
for i in range(0, total, QA_BATCH):
|
||
questions.extend(_qa_for_batch(chunks[i:i + QA_BATCH]))
|
||
log.info("rag: 问答生成进度 %d/%d", min(i + QA_BATCH, total), total)
|
||
questions = _ensure_questions(chunks, questions)
|
||
QA_FILE.write_text(
|
||
json.dumps({"data_hash": h, "chunks": chunks, "questions": questions}, ensure_ascii=False),
|
||
"utf-8",
|
||
)
|
||
log.info("rag: 问答缓存已生成（%d 块）-> %s", len(chunks), QA_FILE)
|
||
return questions
|
||
|
||
|
||
# ── 双路向量索引 ────────────────────────────────────────────────
|
||
|
||
def _migrate_legacy():
|
||
"""把旧版单文件 kb_index.json（向量内嵌在 JSON）就地转为 kb_vectors.npy + kb_meta.json，再删除旧文件。"""
|
||
if not LEGACY_INDEX_FILE.exists() or META_FILE.exists():
|
||
return
|
||
try:
|
||
d = json.loads(LEGACY_INDEX_FILE.read_text("utf-8"))
|
||
np.save(VECTOR_FILE, np.asarray(d["embeddings"], dtype=np.float32))
|
||
META_FILE.write_text(
|
||
json.dumps({
|
||
"model": d.get("model"),
|
||
"data_hash": d.get("data_hash"),
|
||
"chunks": d["chunks"],
|
||
"questions": d.get("questions", []),
|
||
"embed_texts": d["embed_texts"],
|
||
"embed_owner": d["embed_owner"],
|
||
}, ensure_ascii=False),
|
||
"utf-8",
|
||
)
|
||
LEGACY_INDEX_FILE.unlink()
|
||
log.info("rag: 已迁移旧 kb_index.json -> kb_vectors.npy + kb_meta.json（无重建，未调 API）")
|
||
except Exception as e: # noqa: BLE001
|
||
log.warning("rag: 旧索引迁移失败（将重建）: %s", e)
|
||
|
||
|
||
def _load_index():
|
||
"""构建/加载双路向量索引（原句 + 生成问题）。向量存二进制 .npy，元数据存 JSON。"""
|
||
global _index
|
||
if _index is not None:
|
||
return _index
|
||
_migrate_legacy()
|
||
h = _data_hash()
|
||
if META_FILE.exists() and VECTOR_FILE.exists():
|
||
try:
|
||
data = json.loads(META_FILE.read_text("utf-8"))
|
||
valid = (
|
||
data.get("model") == EMBED_MODEL
|
||
and data.get("data_hash") == h
|
||
and len(data.get("embed_texts", [])) == len(data.get("embed_owner", []))
|
||
)
|
||
if valid:
|
||
vectors = np.load(VECTOR_FILE) # 已归一化 float32
|
||
if len(vectors) == len(data["embed_texts"]):
|
||
_index = {
|
||
"chunks": data["chunks"],
|
||
"questions": data.get("questions", []),
|
||
"embed_texts": data["embed_texts"],
|
||
"embed_owner": data["embed_owner"],
|
||
"embeddings": vectors,
|
||
}
|
||
log.info("rag: 双路知识索引已加载（%d 块 → %d 向量, %s）",
|
||
len(_index["chunks"]), len(_index["embed_texts"]), EMBED_MODEL)
|
||
return _index
|
||
log.warning("rag: 索引失效（model/hash/长度不符），重建")
|
||
except Exception: # noqa: BLE001
|
||
log.warning("rag: 索引损坏，重建")
|
||
# 构建
|
||
chunks = _chunk_md(_kb_text())
|
||
questions = _load_questions(chunks)
|
||
embed_texts: list[str] = []
|
||
embed_owner: list[int] = []
|
||
for i, c in enumerate(chunks):
|
||
embed_texts.append(c)
|
||
embed_owner.append(i)
|
||
# 企业块：把「企业名称」作为确定性问题加入索引，保证按名称可召回（名称为问、介绍为答）
|
||
cname = _company_name(c)
|
||
if cname:
|
||
embed_texts.append(cname)
|
||
embed_owner.append(i)
|
||
for q in questions[i]:
|
||
embed_texts.append(q)
|
||
embed_owner.append(i)
|
||
log.info("rag: 构建双路索引（%d 块 → %d 向量，%s embedding 中…）",
|
||
len(chunks), len(embed_texts), EMBED_MODEL)
|
||
arr = _normalize(_embed(embed_texts))
|
||
_index = {
|
||
"chunks": chunks,
|
||
"questions": questions,
|
||
"embed_texts": embed_texts,
|
||
"embed_owner": embed_owner,
|
||
"embeddings": arr,
|
||
}
|
||
np.save(VECTOR_FILE, arr)
|
||
META_FILE.write_text(
|
||
json.dumps({
|
||
"model": EMBED_MODEL,
|
||
"data_hash": h,
|
||
"chunks": chunks,
|
||
"questions": questions,
|
||
"embed_texts": embed_texts,
|
||
"embed_owner": embed_owner,
|
||
}, ensure_ascii=False),
|
||
"utf-8",
|
||
)
|
||
log.info("rag: 双路索引构建完成（%d 块, dim=%d）-> %s + %s",
|
||
len(chunks), arr.shape[1], VECTOR_FILE, META_FILE)
|
||
return _index
|
||
|
||
|
||
# ── 对外接口 ────────────────────────────────────────────────────
|
||
|
||
def retrieve(query: str, top_k: int = 3, threshold: float = 0.3) -> list[str]:
|
||
"""双路检索：query 同时比对原句与问题向量，按块去重，返回 top_k 个原文块。"""
|
||
if not query or not query.strip():
|
||
return []
|
||
idx = _load_index()
|
||
q = _normalize([_embed([query])[0]])[0]
|
||
sims = np.dot(idx["embeddings"], q)
|
||
order = np.argsort(-sims)
|
||
best: dict[int, float] = {} # owner -> 该块最高分
|
||
for e in order:
|
||
owner = idx["embed_owner"][int(e)]
|
||
sc = float(sims[int(e)])
|
||
if sc < threshold:
|
||
break
|
||
if owner not in best:
|
||
best[owner] = sc
|
||
if len(best) >= top_k:
|
||
break
|
||
ordered = sorted(best.items(), key=lambda kv: kv[1], reverse=True)
|
||
hits = [idx["chunks"][o] for o, _ in ordered]
|
||
log.info("rag: retrieve() query=%s top_k=%d threshold=%.2f -> 命中 %d 段",
|
||
_trunc(query, 120), top_k, threshold, len(hits))
|
||
for j, (o, sc) in enumerate(ordered):
|
||
log.info("rag: [%d] %.4f %s", j, sc, _trunc(idx["chunks"][o], 300))
|
||
return hits
|
||
|
||
|
||
def brief(max_chars: int = 1200) -> str:
|
||
"""固定知识摘要（前 N 字；兼容 /api/kb/brief 与检索无命中回退）。"""
|
||
idx = _load_index()
|
||
out = []
|
||
total = 0
|
||
for c in idx["chunks"]:
|
||
if total + len(c) > max_chars:
|
||
break
|
||
out.append(c)
|
||
total += len(c)
|
||
text = "\n\n".join(out)
|
||
log.info("rag: brief() 知识库摘要 %d 字 -> %s", total, _trunc(text, 300))
|
||
return text
|
||
|
||
|
||
def full_kb() -> str:
|
||
"""整库原文（小库全量注入，用于语音 instructions，修复尾部知识丢失）。"""
|
||
idx = _load_index()
|
||
text = "\n\n".join(idx["chunks"])
|
||
log.info("rag: full_kb() 全库 %d 字 / %d 块", len(text), len(idx["chunks"]))
|
||
return text
|
||
|
||
|
||
def build_instructions() -> str:
|
||
"""语音助手的完整基础提示词 = 配置提示词 + 整库知识（服务端统一组装）。"""
|
||
from .config import settings
|
||
base = settings.S2S_INSTRUCTIONS.strip()
|
||
try:
|
||
kb = full_kb()
|
||
if kb:
|
||
base += f"\n\n【园区知识库】\n{kb}"
|
||
except Exception as e: # noqa: BLE001
|
||
log.warning("rag: 知识库拼接失败（仅用基础提示词）: %s", e)
|
||
log.info("rag: build_instructions() 语音系统提示词=%s", _trunc(base, 600))
|
||
return base
|