DPM/backend/app/rag.py

# -*- coding: utf-8 -*-
"""园区知识库 —— 轻量本地 RAG（embedding + numpy 余弦检索）

- 知识源：backend/knowledge/park.md（按标题/段落分块）
- 向量化：dashscope text-embedding-v3（云端 API，索引构建一次后缓存）
- 检索：numpy 余弦 top-k（资料量小，暴力检索足够）
- 提供：
  - retrieve(query, top_k) → 相关段落（供对话动态注入）
  - brief() → 固定知识摘要（供前端拼进 instructions，静态知识一次注入）
"""
import json
import logging
import os
from pathlib import Path

import numpy as np

log = logging.getLogger("dpm.rag")

KB_FILE = Path(__file__).resolve().parent.parent / "knowledge" / "park.md"
INDEX_FILE = Path(__file__).resolve().parent.parent / "knowledge" / "kb_index.json"
EMBED_MODEL = "text-embedding-v3"

_index = None  # 缓存 {chunks, embeddings}


def _chunk_md(text: str) -> list[str]:
    """按标题/段落分块"""
    chunks: list[str] = []
    cur: list[str] = []
    for line in text.splitlines():
        if line.startswith("#"):
            if cur:
                chunks.append("\n".join(cur).strip())
            cur = [line]
        else:
            cur.append(line)
    if cur:
        chunks.append("\n".join(cur).strip())
    return [c for c in chunks if len(c) > 10]


def _embed(texts: list[str]) -> list[list[float]]:
    """dashscope text-embedding-v3 批量向量化"""
    import dashscope

    api_key = os.environ.get("DASHSCOPE_API_KEY") or ""
    resp = dashscope.TextEmbedding.call(
        model=EMBED_MODEL,
        input=texts,
        api_key=api_key,
    )
    if resp.status_code != 200:
        raise RuntimeError(f"embedding API {resp.status_code}: {resp.code} {resp.message}")
    out = [e["embedding"] for e in resp.output["embeddings"]]
    return out


def _load_index():
    """懒加载索引；缺失时构建并缓存"""
    global _index
    if _index is not None:
        return _index
    if INDEX_FILE.exists():
        try:
            _index = json.loads(INDEX_FILE.read_text("utf-8"))
            log.info("rag: 知识索引已加载（%d 块）", len(_index["chunks"]))
            return _index
        except Exception:  # noqa: BLE001
            log.warning("rag: 索引损坏，重建")
    # 构建
    chunks = _chunk_md(KB_FILE.read_text("utf-8"))
    log.info("rag: 构建知识索引（%d 块，embedding 中…）", len(chunks))
    embs = _embed(chunks)
    _index = {"chunks": chunks, "embeddings": embs}
    INDEX_FILE.write_text(json.dumps(_index, ensure_ascii=False), "utf-8")
    log.info("rag: 知识索引构建完成 -> %s", INDEX_FILE)
    return _index


def retrieve(query: str, top_k: int = 3) -> list[str]:
    """检索与 query 最相关的知识段落"""
    if not query or not query.strip():
        return []
    idx = _load_index()
    q = _embed([query])[0]
    sims = np.dot(np.array(idx["embeddings"]), np.array(q))
    top = np.argsort(-sims)[:top_k]
    return [idx["chunks"][int(i)] for i in top if float(sims[int(i)]) > 0.3]


def brief(max_chars: int = 1200) -> str:
    """固定知识摘要（前端拼进 instructions，静态知识一次注入）"""
    idx = _load_index()
    out = []
    total = 0
    for c in idx["chunks"]:
        if total + len(c) > max_chars:
            break
        out.append(c)
        total += len(c)
    return "\n\n".join(out)


def build_instructions() -> str:
    """语音助手的完整基础提示词 = 配置提示词 + 园区知识库摘要（服务端统一组装）"""
    from .config import settings
    base = settings.S2S_INSTRUCTIONS.strip()
    try:
        kb = brief()
        if kb:
            base += f"\n\n【园区知识库】\n{kb}"
    except Exception as e:  # noqa: BLE001
        log.warning("rag: 知识库拼接失败（仅用基础提示词）: %s", e)
    return base