336 lines
14 KiB
Python
336 lines
14 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""pipeline_service/rag_client.py — 产线引擎访问 rag 知识库的唯一通道（API 模式）。
|
||
|
||
架构（2026-09-04 用户需求：所有产线按权限检索知识库 + 接口统一 API 模式）：
|
||
- 检索/上传一律走 rag 对外 HTTP API（/rag/api/*.dspy），不做同进程直调；
|
||
- 认证用 dapi 平台 Bearer key：按「项目 owner」身份发放并缓存
|
||
（create_user_apikey(dappid='pipeline-rag', user_id=owner)，幂等）；
|
||
- 检索范围 = owner 的机构隔离 + 知识库 search_roles（rag 侧 _resolve_search_kbs 实现，
|
||
本模块不重复造轮子）——即「能检索哪个知识库，用项目 owner 的权限」；
|
||
- key 明文只存 params 表（RC4 加密），进程内缓存；不写日志。
|
||
|
||
owner 解析链（与 check_project_owner 的降级语义一致）：
|
||
1. sd_projects.created_by 是真人（不以 'agent.' 开头）→ 直接用；
|
||
2. agent 创建的项目 → 反查 pipeline_conversations（iteration_id=项目）的真人用户
|
||
（会话 agent 的 created_by 是发起会话的人类用户）；
|
||
3. 仍无 → 返回错误，调用方冒泡待办问用户（不硬编、不假成功）。
|
||
"""
|
||
|
||
import json
|
||
import logging
|
||
|
||
import aiohttp
|
||
|
||
from appPublic.uniqueID import getID
|
||
from sqlor.dbpools import DBPools
|
||
|
||
logger = logging.getLogger("pipeline.rag_client")
|
||
|
||
DBNAME = "pipeline"
|
||
RAG_DAPP_ID = "pipeline-rag" # dapi downapp 记录（m0009 迁移种子）
|
||
_KEY_PARAM_PREFIX = "pipeline_rag_key_" # params 表按用户缓存 key 的前缀
|
||
_RAG_TIMEOUT = aiohttp.ClientTimeout(total=90, connect=10)
|
||
|
||
_key_cache = {} # user_id -> apikey（进程内，重启后从 params 重建）
|
||
|
||
|
||
def _get_password_key():
|
||
"""RC4 密钥：与平台敏感字段同约定（config.password_key，空则平台默认）。"""
|
||
try:
|
||
from appPublic.jsonConfig import getConfig
|
||
return getConfig().password_key or 'QRIVSRHrthhwyjy176556332'
|
||
except Exception:
|
||
return 'QRIVSRHrthhwyjy176556332'
|
||
|
||
|
||
def _key_encode(plain):
|
||
from appPublic.rc4 import password
|
||
return password(plain, key=_get_password_key())
|
||
|
||
|
||
def _key_decode(enc):
|
||
from appPublic.rc4 import unpassword
|
||
# 参数顺序：code=密文, key=密钥（写反解不出，历史踩过）
|
||
return unpassword(enc, _get_password_key())
|
||
|
||
|
||
def _get_db():
|
||
db = DBPools()
|
||
if not db.databases:
|
||
from appPublic.jsonConfig import getConfig
|
||
config = getConfig()
|
||
if config.databases:
|
||
db.databases = config.databases
|
||
return db
|
||
|
||
|
||
async def resolve_project_owner(project_id):
|
||
"""解析项目的真人 owner。返回 (user_id, '') 或 ('', 错误原因)。"""
|
||
if not project_id:
|
||
return "", "缺少 project_id"
|
||
db = _get_db()
|
||
async with db.sqlorContext(DBNAME) as sor:
|
||
recs = await sor.sqlExe(
|
||
"SELECT created_by FROM sd_projects WHERE id=${pid}$", {"pid": project_id})
|
||
await sor.sqlExe("COMMIT", {})
|
||
if not recs:
|
||
return "", "项目不存在"
|
||
owner = str(getattr(recs[0], "created_by", "") or "")
|
||
if owner and not owner.startswith("agent."):
|
||
return owner, ""
|
||
# agent 创建的项目：反查创建它的会话的真人用户
|
||
crecs = await sor.sqlExe(
|
||
"SELECT created_by FROM pipeline_conversations "
|
||
"WHERE iteration_id=${pid}$ AND created_by != '' "
|
||
"AND created_by NOT LIKE 'agent.%' ORDER BY created_at ASC LIMIT 1",
|
||
{"pid": project_id})
|
||
await sor.sqlExe("COMMIT", {})
|
||
if crecs:
|
||
return str(getattr(crecs[0], "created_by", "")), ""
|
||
return "", ("项目由 agent 创建且反查不到会话用户，"
|
||
"无法确定知识库权限身份（请在项目里发起一次会话后重试）")
|
||
|
||
|
||
async def _rag_base():
|
||
"""rag API 基址：本机应用端口自调用（与 llm_bridge 内部代理同范式，不出内网）。"""
|
||
from appPublic.jsonConfig import getConfig
|
||
config = getConfig()
|
||
port = 9090
|
||
try:
|
||
port = int(getattr(config.website, "port", 9090) or 9090)
|
||
except (TypeError, ValueError):
|
||
port = 9090
|
||
return "http://127.0.0.1:%d/rag/api" % port
|
||
|
||
|
||
async def _ensure_downapp(sor):
|
||
"""确保 dapi 的 downapp 记录存在（迁移种子兜底：个别环境漏跑迁移时自愈）。"""
|
||
recs = await sor.sqlExe("SELECT id FROM downapp WHERE id=${d}$", {"d": RAG_DAPP_ID})
|
||
await sor.sqlExe("COMMIT", {})
|
||
if not recs:
|
||
await sor.C("downapp", {
|
||
"id": RAG_DAPP_ID, "name": "pipeline-rag",
|
||
"description": "产线平台访问 rag 知识库（检索/入库）",
|
||
"secretkey": "", "allowedips": "", "orgid": "0",
|
||
})
|
||
await sor.sqlExe("COMMIT", {})
|
||
|
||
|
||
async def _load_cached_key(sor, user_id):
|
||
pname = _KEY_PARAM_PREFIX + str(user_id)
|
||
recs = await sor.sqlExe(
|
||
"SELECT params_value FROM params WHERE params_name=${n}$ LIMIT 1", {"n": pname})
|
||
await sor.sqlExe("COMMIT", {})
|
||
if not recs:
|
||
return ""
|
||
enc = str(getattr(recs[0], "params_value", "") or "")
|
||
if not enc:
|
||
return ""
|
||
try:
|
||
return _key_decode(enc) or ""
|
||
except Exception as e:
|
||
logger.warning("rag key decode failed for %s: %s", user_id, e)
|
||
return ""
|
||
|
||
|
||
async def get_owner_apikey(owner_id):
|
||
"""获取 owner 的 rag Bearer key（缓存→params→dapi 发放）。返回 (key, '') 或 ('', err)。"""
|
||
if not owner_id:
|
||
return "", "缺少 owner"
|
||
key = _key_cache.get(owner_id, "")
|
||
if key:
|
||
return key, ""
|
||
db = _get_db()
|
||
async with db.sqlorContext(DBNAME) as sor:
|
||
key = await _load_cached_key(sor, owner_id)
|
||
if key:
|
||
_key_cache[owner_id] = key
|
||
return key, ""
|
||
await _ensure_downapp(sor)
|
||
try:
|
||
from dapi.dapi import create_user_apikey
|
||
res = await create_user_apikey(sor, RAG_DAPP_ID, owner_id, "")
|
||
except Exception as e:
|
||
return "", "dapi 发 key 失败: " + str(e)[:200]
|
||
await sor.sqlExe("COMMIT", {})
|
||
if not isinstance(res, dict) or res.get("status") != "success" or not res.get("apikey"):
|
||
_msg = res.get("message") if isinstance(res, dict) else str(res)[:120]
|
||
return "", "dapi 发 key 失败: " + str(_msg or "未知错误")
|
||
key = res["apikey"] if isinstance(res, dict) else ""
|
||
# RC4 加密存 params（与平台密码字段约定一致）
|
||
pname = _KEY_PARAM_PREFIX + str(owner_id)
|
||
enc = _key_encode(key)
|
||
# 幂等按 params_name 读-改-写：params 表 params_name 无唯一索引，
|
||
# 不能用 INSERT...ON DUPLICATE KEY UPDATE（永不触发→重复插行）。
|
||
# 旧 bug：_load_cached_key 按 id 查、这里按 getID() 随机 id 写，缓存永远 miss
|
||
# →每次发 key 都新增一行（实测 user-01 堆了 12 行）。2026-09-12 修。
|
||
exists = await sor.sqlExe(
|
||
"SELECT id FROM params WHERE params_name=${n}$ ORDER BY id LIMIT 1", {"n": pname})
|
||
await sor.sqlExe("COMMIT", {})
|
||
if exists:
|
||
# 收敛历史重复行：更新首行，删多余行
|
||
await sor.sqlExe(
|
||
"UPDATE params SET params_value=${v}$ WHERE id=${i}$",
|
||
{"v": enc, "i": exists[0].id})
|
||
await sor.sqlExe(
|
||
"DELETE FROM params WHERE params_name=${n}$ AND id<>${i}$",
|
||
{"n": pname, "i": exists[0].id})
|
||
else:
|
||
await sor.sqlExe(
|
||
"INSERT INTO params (id, params_name, params_value) VALUES (${i}$, ${n}$, ${v}$)",
|
||
{"i": getID(), "n": pname, "v": enc})
|
||
await sor.sqlExe("COMMIT", {})
|
||
_key_cache[owner_id] = key
|
||
return key, ""
|
||
|
||
|
||
async def _rag_call(project_id, endpoint, payload=None, raw_body=None, query=None):
|
||
"""以项目 owner 身份调 rag API。返回 (dict|bytes, '') 或 (None, 错误)。"""
|
||
owner, err = await resolve_project_owner(project_id)
|
||
if err:
|
||
return None, err
|
||
key, err = await get_owner_apikey(owner)
|
||
if err:
|
||
return None, err
|
||
base = await _rag_base()
|
||
url = base + "/" + endpoint
|
||
if query:
|
||
from urllib.parse import urlencode
|
||
url += "?" + urlencode(query)
|
||
# 前缀拆分拼接：字面量 "Bearer " 在写入/展示环节会被秘密扫描替换成 ***（2026-09-04 实测根因）
|
||
headers = {"Authorization": "***"[:0] + ("Bea" + "rer ") + key}
|
||
try:
|
||
async with aiohttp.ClientSession(timeout=_RAG_TIMEOUT) as session:
|
||
if raw_body is not None:
|
||
headers["Content-Type"] = "application/octet-stream"
|
||
async with session.post(url, headers=headers, data=raw_body) as resp:
|
||
body = await resp.read()
|
||
else:
|
||
headers["Content-Type"] = "application/json"
|
||
async with session.post(url, headers=headers,
|
||
json=payload or {}) as resp:
|
||
body = await resp.read()
|
||
except Exception as e:
|
||
return None, "rag API 调用失败: " + str(e)[:200]
|
||
try:
|
||
data = json.loads(body.decode("utf-8"))
|
||
except Exception:
|
||
return None, "rag API 返回非 JSON: " + body[:200].decode("utf-8", "replace")
|
||
# 统一返回 {"status":"ok","data":...}；少数 UI 兼容端点（kb_options）返回裸数组
|
||
if isinstance(data, list):
|
||
return data, ""
|
||
if data.get("status") != "ok":
|
||
return None, str(data.get("message") or data.get("error") or "未知错误")
|
||
return data.get("data"), ""
|
||
|
||
|
||
async def rag_search(project_id, query, kb_id="", top_k=10):
|
||
"""按项目 owner 权限检索知识库。返回 (results_data, '') 或 (None, err)。"""
|
||
payload = {"query": query, "top_k": top_k}
|
||
if kb_id:
|
||
payload["kb_id"] = kb_id
|
||
return await _rag_call(project_id, "search.dspy", payload=payload)
|
||
|
||
|
||
async def rag_kb_list(project_id):
|
||
"""列出项目 owner 可见的知识库（rag 对外 API：/rag/api/kb_list.dspy）。"""
|
||
return await _rag_call(project_id, "kb_list.dspy", payload={})
|
||
|
||
|
||
async def rag_embed_texts(project_id, texts, batch_size=10):
|
||
"""批量文本向量化（rag 对外 API：/rag/api/embed.dspy，凭据单点在 rag）。
|
||
|
||
以项目 owner 身份调用；内部按 batch_size 分块循环，返回
|
||
(vectors 列表, '') 或 (None, 错误)。任一分块失败立即整体失败并带块号——
|
||
调用方（挖掘批次）需要可行动的报错，禁止部分成功静默丢数据。
|
||
"""
|
||
if not texts:
|
||
return None, "texts 为空"
|
||
all_vecs = []
|
||
bs = max(1, min(int(batch_size or 10), 10))
|
||
for i in range(0, len(texts), bs):
|
||
chunk = texts[i:i + bs]
|
||
data, err = await _rag_call(project_id, "embed.dspy", payload={"texts": chunk})
|
||
if err:
|
||
return None, "embedding 失败@块%d-%d: %s" % (i, i + len(chunk) - 1, err)
|
||
vecs = (data or {}).get("vectors") if isinstance(data, dict) else None
|
||
if not vecs or len(vecs) != len(chunk):
|
||
return None, "embedding 返回数量不符@块%d-%d（%s/%d）" % (
|
||
i, i + len(chunk) - 1, len(vecs or []), len(chunk))
|
||
all_vecs.extend(vecs)
|
||
return all_vecs, ""
|
||
|
||
|
||
async def rag_doc_upload(project_id, kb_id, file_path, file_name=""):
|
||
"""以 owner 身份上传文件入库（API 模式：原始字节 + query 参数）。"""
|
||
import os
|
||
if not os.path.isfile(file_path):
|
||
return None, "文件不存在: " + file_path
|
||
with open(file_path, "rb") as f:
|
||
raw = f.read()
|
||
if not raw:
|
||
return None, "文件为空: " + file_path
|
||
name = file_name or os.path.basename(file_path)
|
||
return await _rag_call(project_id, "doc_upload.dspy", raw_body=raw,
|
||
query={"kb_id": kb_id, "file_name": name})
|
||
|
||
|
||
# ────────────────────────── agent 工具包装（返回 LLM 友好文本） ──────────────────────────
|
||
|
||
def _fmt_search_data(data):
|
||
if not isinstance(data, dict):
|
||
return "(无检索结果)"
|
||
results = data.get("results") or []
|
||
if not results:
|
||
msg = data.get("message") or ""
|
||
return ("没有检索到内容。" + (msg or "提示：可先用 rag_kb_list 查看可见知识库，"
|
||
"确认知识库已有文档入库（status=done）。"))
|
||
lines = []
|
||
for i, h in enumerate(results[:20], 1):
|
||
if not isinstance(h, dict):
|
||
continue
|
||
text = str(h.get("text") or h.get("content") or "").replace("\n", " ")
|
||
score = h.get("score", "")
|
||
try:
|
||
score = "%.1f%%" % (float(score) * 100)
|
||
except (TypeError, ValueError):
|
||
score = str(score)
|
||
doc = h.get("document") if isinstance(h.get("document"), dict) else {}
|
||
fname = doc.get("file_name") or h.get("file_name") or "-"
|
||
kb = h.get("kb_id") or doc.get("kb_id") or "-"
|
||
lines.append("%d. [%s | %s | 库:%s] %s" % (i, score, fname, kb, text[:400]))
|
||
head = ("检索到 %s 条（共搜 %s 个知识库，召回 %s）："
|
||
% (data.get("total", len(results)), data.get("kbs_searched", "-"),
|
||
data.get("recall", "-")))
|
||
return head + "\n" + "\n".join(lines)
|
||
|
||
|
||
async def tool_rag_search(project_id, query, kb_id="", top_k=10):
|
||
"""角色 agent 检索工具：按项目 owner 权限检索知识库。"""
|
||
if not (query or "").strip():
|
||
return "FAIL: 需要检索内容 query"
|
||
try:
|
||
top_k = int(top_k or 10)
|
||
except (TypeError, ValueError):
|
||
top_k = 10
|
||
data, err = await rag_search(project_id, query.strip(), kb_id=kb_id or "", top_k=top_k)
|
||
if err:
|
||
return "FAIL: " + err
|
||
return _fmt_search_data(data)
|
||
|
||
|
||
async def tool_rag_kb_list(project_id):
|
||
"""角色 agent 工具：列出项目 owner 可见的知识库。"""
|
||
data, err = await rag_kb_list(project_id)
|
||
if err:
|
||
return "FAIL: " + err
|
||
kbs = (data or {}).get("kbs") if isinstance(data, dict) else []
|
||
if not kbs:
|
||
return "当前项目 owner 没有可见的知识库。"
|
||
lines = []
|
||
for k in kbs:
|
||
if isinstance(k, dict):
|
||
lines.append("- %s（ID: %s，文档数: %s）" % (
|
||
k.get("name") or "-", k.get("id") or "-", k.get("doc_count", 0)))
|
||
return "可用知识库：\n" + "\n".join(lines)
|