554 lines
20 KiB
Python
554 lines
20 KiB
Python
from rag.uapi_service import APIService
|
||
from sqlor.dbpools import DBPools
|
||
from appPublic.log import debug, error, info
|
||
import time
|
||
import traceback
|
||
import json
|
||
import math
|
||
import uuid
|
||
import os
|
||
from rag.service_opts import get_service_params, sor_get_service_params, sor_get_embedding_mode, get_embedding_mode
|
||
from rag.rag_operations import RagOperations
|
||
from langchain_core.documents import Document
|
||
|
||
REAL_PHOTO_ROOT = "/home/wangmeihua/kyrag/files"
|
||
|
||
helptext = """kyrag API:
|
||
|
||
1. 得到kdb表:
|
||
path: /v1/get_kdbs
|
||
headers: {
|
||
"Content-Type": "application/json"
|
||
}
|
||
response:
|
||
[{"id": "1", "name": "textdb", "description": "文本数据库"}, {"id": "testkdb", "name": "testkdb", "description": ""}, {"id": "Vdtbt3qBfocteit1HIxVH", "name": "trm", "description": ""}]
|
||
|
||
2. 向量检索文本块:
|
||
path: /v1/fusedsearch
|
||
headers: {
|
||
"Content-Type": "application/json"
|
||
}
|
||
data: {
|
||
"query": "什么是知识抽取。",
|
||
"fiids":["1"],
|
||
"limit":5
|
||
}
|
||
|
||
3、docs文档
|
||
path: /v1/docs
|
||
|
||
4. 纯文本插入接口:
|
||
path: /v1/textinsert
|
||
headers: {
|
||
"Content-Type": "application/json"
|
||
}
|
||
data: {
|
||
"text": "要插入的纯文本内容",
|
||
"fiid": "知识库ID",
|
||
"db_type": "数据库类型（如 textdb）"
|
||
}
|
||
response: {
|
||
"status": "success",
|
||
"userid": "用户组织ID",
|
||
"collection_name": "ragdb_{dbtype}",
|
||
"message": "文本成功嵌入并处理三元组",
|
||
"status_code": 200,
|
||
"unique_triples": ["提取的三元组列表"],
|
||
"timings": {"total": 0.123}
|
||
}
|
||
error response: {
|
||
"status": "error",
|
||
"userid": "用户组织ID",
|
||
"collection_name": "ragdb_{dbtype}",
|
||
"message": "错误信息",
|
||
"status_code": 400,
|
||
"timings": {"total": 0.123}
|
||
}
|
||
|
||
4. 添加用户消息到记忆:
|
||
path: /v1/add_user_messages
|
||
headers: {
|
||
"Content-Type": "application/json"
|
||
}
|
||
data: {
|
||
"orgid": "用户组织ID",
|
||
"messages": [{"role": "user", "content": "消息内容"}, ...]
|
||
}
|
||
|
||
5. 获取用户所有记忆:
|
||
path: /v1/get_user_memories
|
||
headers: {
|
||
"Content-Type": "application/json"
|
||
}
|
||
data: {
|
||
"orgid": "用户组织ID",
|
||
"limit": 10 # 可选，默认为 10
|
||
}
|
||
"""
|
||
|
||
async def docs(request, params_kw, *params, **kw):
|
||
return helptext
|
||
|
||
async def get_kdbs(request, params_kw, *params, **kw):
|
||
"""返回 kdb 表的全部内容，返回 JSON"""
|
||
f = kw.get('get_userorgid')
|
||
orgid = await f()
|
||
debug(f"orgid: {orgid}，{f=}")
|
||
debug(f"params_kw: {params_kw}")
|
||
db = DBPools()
|
||
dbname = kw.get('get_module_dbname')('rag')
|
||
sql_opts = """
|
||
SELECT id, name, description
|
||
FROM kdb
|
||
WHERE orgid = ${orgid}$
|
||
"""
|
||
try:
|
||
async with db.sqlorContext(dbname) as sor:
|
||
opts_result = await sor.sqlExe(sql_opts, {"orgid": orgid})
|
||
if not opts_result:
|
||
error("未找到 kdb 表记录")
|
||
return json.dumps({"status": "error", "message": "未找到记录"})
|
||
return json.dumps(opts_result, ensure_ascii=False)
|
||
except Exception as e:
|
||
error(f"查询 kdb 表失败: {str(e)}, 堆栈: {traceback.format_exc()}")
|
||
return json.dumps({"status": "error", "message": str(e)})
|
||
|
||
except Exception as e:
|
||
error(f"列出用户文件失败: {str(e)}, 堆栈: {traceback.format_exc()}")
|
||
result.update({
|
||
"status": "error",
|
||
"files_by_knowledge_base": {},
|
||
"message": f"列出用户文件失败: {str(e)}",
|
||
"status_code": 400
|
||
})
|
||
return result
|
||
|
||
async def fusedsearch(request, params_kw, *params):
|
||
"""
|
||
融合搜索，调用服务化端点
|
||
|
||
"""
|
||
kw = request._run_ns
|
||
f = kw.get('get_userorgid')
|
||
orgid = await f()
|
||
debug(f"orgid: {orgid}，{f=}")
|
||
f = kw.get('get_user')
|
||
userid = await f()
|
||
debug(f"params_kw: {params_kw}")
|
||
# orgid = "04J6VbxLqB_9RPMcgOv_8"
|
||
# userid = "04J6VbxLqB_9RPMcgOv_8"
|
||
query = params_kw.get('query', '').strip()
|
||
img_path = params_kw.get('image')
|
||
if isinstance(img_path, str):
|
||
img_path = img_path.strip()
|
||
relative_part = img_path.lstrip("/")
|
||
real_img_path = os.path.join(REAL_PHOTO_ROOT, relative_part)
|
||
if not os.path.exists(real_img_path):
|
||
raise FileNotFoundError(f"图片不存在: {real_img_path}")
|
||
img_path = real_img_path
|
||
debug(f"自动修复图片路径成功: {img_path}")
|
||
# 统一模式处理 limit 参数,为了对接dify和coze
|
||
raw_limit = params_kw.get('limit') or (
|
||
params_kw.get('retrieval_setting', {}).get('top_k')
|
||
if isinstance(params_kw.get('retrieval_setting'), dict)
|
||
else None
|
||
)
|
||
|
||
# 标准化为整数值
|
||
if raw_limit is None:
|
||
limit = 5 # 两个来源都不存在时使用默认值
|
||
elif isinstance(raw_limit, (int, float)):
|
||
limit = int(raw_limit) # 数值类型直接转换
|
||
elif isinstance(raw_limit, str):
|
||
try:
|
||
# 字符串转换为整数
|
||
limit = int(raw_limit)
|
||
except (TypeError, ValueError):
|
||
limit = 5 # 转换失败使用默认值
|
||
else:
|
||
limit = 5 # 其他意外类型使用默认值
|
||
debug(f"limit: {limit}")
|
||
raw_fiids = params_kw.get('fiids') or params_kw.get('knowledge_id') #
|
||
|
||
# 标准化为列表格式
|
||
if raw_fiids is None:
|
||
fiids = [] # 两个参数都不存在
|
||
elif isinstance(raw_fiids, list):
|
||
fiids = [str(item).strip() for item in raw_fiids] # 已经是列表
|
||
elif isinstance(raw_fiids, str):
|
||
# fiids = [f.strip() for f in raw_fiids.split(',') if f.strip()]
|
||
try:
|
||
# 尝试解析 JSON 字符串
|
||
parsed = json.loads(raw_fiids)
|
||
if isinstance(parsed, list):
|
||
fiids = [str(item).strip() for item in parsed] # JSON 数组转为字符串列表
|
||
else:
|
||
# 处理逗号分隔的字符串或单个 ID 字符串
|
||
fiids = [f.strip() for f in raw_fiids.split(',') if f.strip()]
|
||
except json.JSONDecodeError:
|
||
# 如果不是合法 JSON，按逗号分隔
|
||
fiids = [f.strip() for f in raw_fiids.split(',') if f.strip()]
|
||
elif isinstance(raw_fiids, (int, float)):
|
||
fiids = [str(int(raw_fiids))] # 数值类型转为字符串列表
|
||
else:
|
||
fiids = [] # 其他意外类型
|
||
|
||
debug(f"fiids: {fiids}")
|
||
|
||
# 验证 fiids的orgid与orgid = await f()是否一致
|
||
await _validate_fiids_orgid(fiids, orgid, kw)
|
||
|
||
service_params = await get_service_params(orgid)
|
||
if not service_params:
|
||
raise ValueError("无法获取服务参数")
|
||
# 获取嵌入模式
|
||
embedding_mode = await get_embedding_mode(orgid)
|
||
debug(f"检测到 embedding_mode = {embedding_mode}（0=文本, 1=多模态）")
|
||
|
||
# 情况1：query 和 image 都为空 → 报错
|
||
if not query and not img_path:
|
||
raise ValueError("查询文本和图片不能同时为空")
|
||
|
||
# 情况2：query 和 image 都存在 → 报错（你当前业务不允许同时传）
|
||
if query and img_path:
|
||
raise ValueError("查询文本和图片只能二选一，不能同时提交")
|
||
|
||
# 3. 只有图片 → 以图搜图 走纯多模态分支
|
||
if img_path and not query:
|
||
try:
|
||
debug("检测到纯图片查询，执行以图搜图")
|
||
rag_ops = RagOperations()
|
||
|
||
timings = {}
|
||
start_time = time.time()
|
||
|
||
# 直接生成图片向量
|
||
img_vector = await rag_ops.generate_image_vector(
|
||
request, img_path, service_params, userid, timings, embedding_mode
|
||
)
|
||
|
||
# 向量搜索（多取 50 条再截断，和文本分支保持一致）
|
||
search_results = await rag_ops.vector_search(
|
||
request, img_vector, orgid, fiids, limit + 50, service_params, userid, timings
|
||
)
|
||
|
||
timings["total_time"] = time.time() - start_time
|
||
|
||
# 可选：搜索完后删除图片，省磁盘（看你需求）
|
||
# try:
|
||
# os.remove(img_path)
|
||
# except:
|
||
# pass
|
||
|
||
final_results = []
|
||
for item in search_results[:limit]:
|
||
final_results.append({
|
||
"text": item["text"],
|
||
"distance": item["distance"]
|
||
})
|
||
|
||
return {
|
||
"results": final_results,
|
||
"timings": timings
|
||
}
|
||
except Exception as e:
|
||
error(f"融合搜索失败: {str(e)}, 堆栈: {traceback.format_exc()}")
|
||
return {
|
||
"records": [],
|
||
"timings": {"total_time": time.time() - start_time if 'start_time' in locals() else 0},
|
||
"error": str(e)
|
||
}
|
||
|
||
if not img_path and query:
|
||
try:
|
||
timings = {}
|
||
start_time = time.time()
|
||
rag_ops = RagOperations()
|
||
|
||
query_entities = await rag_ops.extract_entities(request, query, service_params, userid, timings)
|
||
all_triplets = await rag_ops.match_triplets(request, query, query_entities, orgid, fiids, service_params,
|
||
userid, timings)
|
||
combined_text = _combine_query_with_triplets(query, all_triplets)
|
||
query_vector = await rag_ops.generate_query_vector(request, combined_text, service_params, userid, timings, embedding_mode)
|
||
search_results = await rag_ops.vector_search(request, query_vector, orgid, fiids, limit + 50, service_params,
|
||
userid, timings)
|
||
|
||
use_rerank = True
|
||
if use_rerank and search_results:
|
||
final_results = await rag_ops.rerank_results(request, combined_text, search_results, limit, service_params,
|
||
userid, timings)
|
||
debug(f"final_results: {final_results}")
|
||
else:
|
||
final_results = [{k: v for k, v in r.items() if k != 'rerank_score'} for r in search_results]
|
||
|
||
formatted_results = rag_ops.format_search_results(final_results, limit)
|
||
timings["total_time"] = time.time() - start_time
|
||
debug(f"融合搜索完成，返回 {len(formatted_results)} 条结果，总耗时: {timings['total_time']:.3f} 秒")
|
||
|
||
return {
|
||
"records": formatted_results,
|
||
"timings": timings
|
||
}
|
||
except Exception as e:
|
||
error(f"融合搜索失败: {str(e)}, 堆栈: {traceback.format_exc()}")
|
||
return {
|
||
"records": [],
|
||
"timings": {"total_time": time.time() - start_time if 'start_time' in locals() else 0},
|
||
"error": str(e)
|
||
}
|
||
|
||
# async def fusedsearch(request, params_kw, *params):
|
||
# """
|
||
# 融合搜索，调用服务化端点
|
||
#
|
||
# """
|
||
# kw = request._run_ns
|
||
# f = kw.get('get_userorgid')
|
||
# orgid = await f()
|
||
# debug(f"orgid: {orgid}，{f=}")
|
||
# f = kw.get('get_user')
|
||
# userid = await f()
|
||
# debug(f"params_kw: {params_kw}")
|
||
# # orgid = "04J6VbxLqB_9RPMcgOv_8"
|
||
# # userid = "04J6VbxLqB_9RPMcgOv_8"
|
||
# query = params_kw.get('query', '')
|
||
# # 统一模式处理 limit 参数,为了对接dify和coze
|
||
# raw_limit = params_kw.get('limit') or (
|
||
# params_kw.get('retrieval_setting', {}).get('top_k')
|
||
# if isinstance(params_kw.get('retrieval_setting'), dict)
|
||
# else None
|
||
# )
|
||
#
|
||
# # 标准化为整数值
|
||
# if raw_limit is None:
|
||
# limit = 5 # 两个来源都不存在时使用默认值
|
||
# elif isinstance(raw_limit, (int, float)):
|
||
# limit = int(raw_limit) # 数值类型直接转换
|
||
# elif isinstance(raw_limit, str):
|
||
# try:
|
||
# # 字符串转换为整数
|
||
# limit = int(raw_limit)
|
||
# except (TypeError, ValueError):
|
||
# limit = 5 # 转换失败使用默认值
|
||
# else:
|
||
# limit = 5 # 其他意外类型使用默认值
|
||
# debug(f"limit: {limit}")
|
||
# raw_fiids = params_kw.get('fiids') or params_kw.get('knowledge_id') #
|
||
#
|
||
# # 标准化为列表格式
|
||
# if raw_fiids is None:
|
||
# fiids = [] # 两个参数都不存在
|
||
# elif isinstance(raw_fiids, list):
|
||
# fiids = [str(item).strip() for item in raw_fiids] # 已经是列表
|
||
# elif isinstance(raw_fiids, str):
|
||
# # fiids = [f.strip() for f in raw_fiids.split(',') if f.strip()]
|
||
# try:
|
||
# # 尝试解析 JSON 字符串
|
||
# parsed = json.loads(raw_fiids)
|
||
# if isinstance(parsed, list):
|
||
# fiids = [str(item).strip() for item in parsed] # JSON 数组转为字符串列表
|
||
# else:
|
||
# # 处理逗号分隔的字符串或单个 ID 字符串
|
||
# fiids = [f.strip() for f in raw_fiids.split(',') if f.strip()]
|
||
# except json.JSONDecodeError:
|
||
# # 如果不是合法 JSON，按逗号分隔
|
||
# fiids = [f.strip() for f in raw_fiids.split(',') if f.strip()]
|
||
# elif isinstance(raw_fiids, (int, float)):
|
||
# fiids = [str(int(raw_fiids))] # 数值类型转为字符串列表
|
||
# else:
|
||
# fiids = [] # 其他意外类型
|
||
#
|
||
# debug(f"fiids: {fiids}")
|
||
#
|
||
# # 验证 fiids的orgid与orgid = await f()是否一致
|
||
# await _validate_fiids_orgid(fiids, orgid, kw)
|
||
#
|
||
# service_params = await get_service_params(orgid)
|
||
# if not service_params:
|
||
# raise ValueError("无法获取服务参数")
|
||
# # 获取嵌入模式
|
||
# embedding_mode = await get_embedding_mode(orgid)
|
||
# debug(f"检测到 embedding_mode = {embedding_mode}（0=文本, 1=多模态）")
|
||
#
|
||
# try:
|
||
# timings = {}
|
||
# start_time = time.time()
|
||
# rag_ops = RagOperations()
|
||
#
|
||
# query_entities = await rag_ops.extract_entities(request, query, service_params, userid, timings)
|
||
# all_triplets = await rag_ops.match_triplets(request, query, query_entities, orgid, fiids, service_params,
|
||
# userid, timings)
|
||
# combined_text = _combine_query_with_triplets(query, all_triplets)
|
||
# query_vector = await rag_ops.generate_query_vector(request, combined_text, service_params, userid, timings, embedding_mode)
|
||
# search_results = await rag_ops.vector_search(request, query_vector, orgid, fiids, limit + 50, service_params,
|
||
# userid, timings)
|
||
#
|
||
# use_rerank = False
|
||
# if use_rerank and search_results:
|
||
# final_results = await rag_ops.rerank_results(request, combined_text, search_results, limit, service_params,
|
||
# userid, timings)
|
||
# debug(f"final_results: {final_results}")
|
||
# else:
|
||
# final_results = [{k: v for k, v in r.items() if k != 'rerank_score'} for r in search_results]
|
||
#
|
||
# formatted_results = rag_ops.format_search_results(final_results, limit)
|
||
# timings["total_time"] = time.time() - start_time
|
||
# debug(f"融合搜索完成，返回 {len(formatted_results)} 条结果，总耗时: {timings['total_time']:.3f} 秒")
|
||
#
|
||
# return {
|
||
# "records": formatted_results,
|
||
# "timings": timings
|
||
# }
|
||
# except Exception as e:
|
||
# error(f"融合搜索失败: {str(e)}, 堆栈: {traceback.format_exc()}")
|
||
# return {
|
||
# "records": [],
|
||
# "timings": {"total_time": time.time() - start_time if 'start_time' in locals() else 0},
|
||
# "error": str(e)
|
||
# }
|
||
|
||
# async def text_insert(text: str, fiid: str, orgid: str, db_type: str):
|
||
async def textinsert(request, params_kw, *params):
|
||
kw = request._run_ns
|
||
f = kw.get('get_userorgid')
|
||
orgid = await f()
|
||
debug(f"orgid: {orgid}，{f=}")
|
||
f = kw.get('get_user')
|
||
userid = await f()
|
||
text = params_kw.get('text', '')
|
||
fiid = params_kw.get('fiid')
|
||
db_type = params_kw.get('db_type')
|
||
id = str(uuid.uuid4())
|
||
debug(f"params_kw: {params_kw}")
|
||
# orgid = "04J6VbxLqB_9RPMcgOv_8"
|
||
# userid = "04J6VbxLqB_9RPMcgOv_8"
|
||
# text = text
|
||
# fiid = fiid
|
||
# orgid = orgid
|
||
# db_type = db_type
|
||
# id = str(uuid.uuid4())
|
||
# debug(f'Inserting document: text={text},userid={orgid}, db_type={db_type}, knowledge_base_id={fiid}')
|
||
|
||
timings = {}
|
||
start_total = time.time()
|
||
|
||
result = {
|
||
"status": "error",
|
||
"userid": orgid,
|
||
"collection_name": f"ragdb_{db_type}",
|
||
"message": "",
|
||
"status_code": 400
|
||
}
|
||
|
||
try:
|
||
# 验证必填字段
|
||
if not orgid or not fiid or not text or not db_type:
|
||
raise ValueError("orgid、fiid、db_type 和 text 不能为空")
|
||
if len(orgid) > 32 or len(fiid) > 255:
|
||
raise ValueError("orgid 或 fiid 的长度超出限制")
|
||
|
||
# 获取服务参数
|
||
service_params = await get_service_params(orgid)
|
||
if not service_params:
|
||
raise ValueError("无法获取服务参数")
|
||
rag_ops = RagOperations()
|
||
# 生成嵌入向量
|
||
embedding = await rag_ops.generate_query_vector(request, text, service_params, userid, timings)
|
||
|
||
# 插入 Milvus
|
||
fields = {
|
||
"text": text,
|
||
"knowledge_base_id": fiid,
|
||
"userid": orgid,
|
||
"vector": embedding,
|
||
"document_id": id
|
||
}
|
||
chunks_data = await rag_ops.insert_to_vector_text(request, db_type, fields, service_params, userid, timings)
|
||
|
||
# 抽取三元组
|
||
document = Document(page_content=text)
|
||
chunks = [document]
|
||
triples = await rag_ops.extract_triples(
|
||
request, chunks, service_params, userid, timings
|
||
)
|
||
|
||
# 插入 Neo4j
|
||
await rag_ops.insert_to_graph_db(
|
||
request, triples, id, fiid, orgid, service_params, userid, timings
|
||
)
|
||
|
||
timings["total"] = time.time() - start_total
|
||
result.update({
|
||
"status": "success",
|
||
"unique_triples": triples,
|
||
"message": f"文本成功嵌入并处理三元组",
|
||
"status_code": 200
|
||
})
|
||
debug(f"总耗时: {timings['total']:.2f} 秒")
|
||
|
||
except Exception as e:
|
||
error(f"插入文档失败: {str(e)}, 堆栈: {traceback.format_exc()}")
|
||
timings["total"] = time.time() - start_total
|
||
result.update({
|
||
"message": f"插入文档失败: {str(e)}",
|
||
"timings": timings
|
||
})
|
||
raise ValueError(str(e)) from e
|
||
debug(f"最终结果是：{result}")
|
||
return result
|
||
|
||
async def _validate_fiids_orgid(fiids, orgid, kw):
|
||
"""验证 fiids 的 orgid 与当前用户 orgid 是否一致"""
|
||
if fiids:
|
||
db = DBPools()
|
||
dbname = kw.get('get_module_dbname')('rag')
|
||
sql_opts = """SELECT orgid FROM kdb WHERE id = ${id}$"""
|
||
try:
|
||
async with db.sqlorContext(dbname) as sor:
|
||
result = await sor.sqlExe(sql_opts, {"id": fiids[0]})
|
||
if not result:
|
||
raise ValueError(f"未找到 fiid={fiids[0]} 的记录")
|
||
kdb_orgid = result[0].get('orgid')
|
||
if kdb_orgid != orgid:
|
||
raise ValueError(f"orgid 不一致: kdb.orgid={kdb_orgid}, user orgid={orgid}")
|
||
except Exception as e:
|
||
error(f"orgid 验证失败: {str(e)}")
|
||
raise
|
||
|
||
async def _validate_fiids_orgid(fiids, orgid, kw):
|
||
"""验证 fiids 的 orgid 与当前用户 orgid 是否一致"""
|
||
if fiids:
|
||
db = DBPools()
|
||
dbname = kw.get('get_module_dbname')('rag')
|
||
sql_opts = """SELECT orgid FROM kdb WHERE id = ${id}$"""
|
||
try:
|
||
async with db.sqlorContext(dbname) as sor:
|
||
result = await sor.sqlExe(sql_opts, {"id": fiids[0]})
|
||
if not result:
|
||
raise ValueError(f"未找到 fiid={fiids[0]} 的记录")
|
||
kdb_orgid = result[0].get('orgid')
|
||
if kdb_orgid != orgid:
|
||
raise ValueError(f"orgid 不一致: kdb.orgid={kdb_orgid}, user orgid={orgid}")
|
||
except Exception as e:
|
||
error(f"orgid 验证失败: {str(e)}")
|
||
raise
|
||
|
||
|
||
def _combine_query_with_triplets(query, triplets):
|
||
"""拼接查询文本和三元组文本"""
|
||
triplet_texts = []
|
||
for triplet in triplets:
|
||
head = triplet.get('head', '')
|
||
type_ = triplet.get('type', '')
|
||
tail = triplet.get('tail', '')
|
||
if head and type_ and tail:
|
||
triplet_texts.append(f"{head} {type_} {tail}")
|
||
else:
|
||
debug(f"无效三元组: {triplet}")
|
||
|
||
combined_text = query
|
||
if triplet_texts:
|
||
combined_text += "".join(triplet_texts)
|
||
|
||
debug(f"拼接文本: {combined_text[:200]}... (总长度: {len(combined_text)}, 三元组数量: {len(triplet_texts)})")
|
||
return combined_text |