Pine
83bf19a3e5
后端： - tools.py：工具注册表（get_park_overview 园区实时数据 / query_companies 企业名录 / control_display 大屏控制 / get_time），POST /api/tools/exec 执行 - rag.py：园区知识库（park.md 分块 → dashscope text-embedding-v3 → numpy 余弦检索），/api/kb/brief 摘要注入 + /api/kb/retrieve 动态检索 - 知识索引 kb_index.json 入库，部署免首次构建 前端（VoiceAssistant）： - TOOL_DEFS 经 session.update 传给 LLM；toolcall 事件 → 后端执行 → sendToolOutput + requestResponse - instructions 拼接园区知识摘要；工具调用气泡提示 - 测试通过：4 个工具，检索「入驻政策」正确命中政策段落
103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
# -*- 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)
|