This commit is contained in:
yumoqing 2025-10-09 17:30:32 +08:00
parent b9ec5336fa
commit 02b5413f4d

View File

@ -11,133 +11,133 @@ from appPublic.registerfunction import RegisterFunction
from sqlor.dbpools import DBPools from sqlor.dbpools import DBPools
async def _make_connection_request(action: str, params: dict = None) -> dict: async def _make_connection_request(action: str, params: dict = None) -> dict:
""" """
通用函数调用 MilvusConnection 的服务化接口 通用函数调用 MilvusConnection 的服务化接口
参数: 参数:
action (str): 操作类型例如 'initialize''insert_document' action (str): 操作类型例如 'initialize''insert_document'
params (dict): 操作参数默认为 None params (dict): 操作参数默认为 None
返回: 返回:
dict: 服务响应包含 statusmessagecollection_name 等字段 dict: 服务响应包含 statusmessagecollection_name 等字段
""" """
if params is None: if params is None:
params = {} params = {}
url = f"http://localhost:8888/v1/{action}" url = f"http://localhost:8888/v1/{action}"
try: try:
debug(f"发起 {action} 请求: params={params}") debug(f"发起 {action} 请求: params={params}")
async with ClientSession(timeout=ClientTimeout(total=10)) as session: async with ClientSession(timeout=ClientTimeout(total=10)) as session:
async with session.post( async with session.post(
url, url,
headers={"Content-Type": "application/json"}, headers={"Content-Type": "application/json"},
json=params json=params
) as response: ) as response:
response_text = await response.text() response_text = await response.text()
debug(f"收到 {action} 响应: status={response.status}, content={response_text}") debug(f"收到 {action} 响应: status={response.status}, content={response_text}")
if response.status != 200: if response.status != 200:
error(f"{action} 请求失败: 状态码={response.status}, 响应={response_text}") error(f"{action} 请求失败: 状态码={response.status}, 响应={response_text}")
return { return {
"status": "error", "status": "error",
"message": f"请求失败: 状态码 {response.status}", "message": f"请求失败: 状态码 {response.status}",
"collection_name": params.get("db_type", "ragdb"), "collection_name": params.get("db_type", "ragdb"),
"document_id": "", "document_id": "",
"status_code": response.status "status_code": response.status
} }
result = await response.json() result = await response.json()
info(f"{action} 请求成功: 结果={result}") info(f"{action} 请求成功: 结果={result}")
return result return result
except Exception as e: except Exception as e:
error(f"{action} 请求异常: 错误={str(e)}") error(f"{action} 请求异常: 错误={str(e)}")
return { return {
"status": "error", "status": "error",
"message": f"服务器错误: {str(e)}", "message": f"服务器错误: {str(e)}",
"collection_name": params.get("db_type", "ragdb"), "collection_name": params.get("db_type", "ragdb"),
"document_id": "", "document_id": "",
"status_code": 500 "status_code": 500
} }
async def create_collection(db_type: str = "") -> dict: async def create_collection(db_type: str = "") -> dict:
"""创建 Milvus 集合""" """创建 Milvus 集合"""
return await _make_connection_request("create_collection", {"db_type": db_type}) return await _make_connection_request("create_collection", {"db_type": db_type})
async def delete_collection(db_type: str = "") -> dict: async def delete_collection(db_type: str = "") -> dict:
"""删除 Milvus 集合""" """删除 Milvus 集合"""
return await _make_connection_request("delete_collection", {"db_type": db_type}) return await _make_connection_request("delete_collection", {"db_type": db_type})
async def insert_document(file_path: str, userid: str, knowledge_base_id: str, db_type: str = "") -> dict: async def insert_document(file_path: str, userid: str, knowledge_base_id: str, db_type: str = "") -> dict:
"""插入文档到 Milvus 并抽取三元组到 Neo4j""" """插入文档到 Milvus 并抽取三元组到 Neo4j"""
params = { params = {
"file_path": file_path, "file_path": file_path,
"userid": userid, "userid": userid,
"knowledge_base_id": knowledge_base_id, "knowledge_base_id": knowledge_base_id,
"db_type": db_type "db_type": db_type
} }
return await _make_connection_request("insert_document", params) return await _make_connection_request("insert_document", params)
async def delete_document(userid: str, filename: str, knowledge_base_id: str, db_type: str = "") -> dict: async def delete_document(userid: str, filename: str, knowledge_base_id: str, db_type: str = "") -> dict:
"""删除指定文档的 Milvus 和 Neo4j 记录""" """删除指定文档的 Milvus 和 Neo4j 记录"""
params = { params = {
"userid": userid, "userid": userid,
"filename": filename, "filename": filename,
"knowledge_base_id": knowledge_base_id, "knowledge_base_id": knowledge_base_id,
"db_type": db_type "db_type": db_type
} }
return await _make_connection_request("delete_document", params) return await _make_connection_request("delete_document", params)
async def delete_knowledge_base(userid: str, knowledge_base_id: str, db_type: str = "") -> dict: async def delete_knowledge_base(userid: str, knowledge_base_id: str, db_type: str = "") -> dict:
"""删除整个知识库的 Milvus 和 Neo4j 记录""" """删除整个知识库的 Milvus 和 Neo4j 记录"""
params = { params = {
"userid": userid, "userid": userid,
"knowledge_base_id": knowledge_base_id, "knowledge_base_id": knowledge_base_id,
"db_type": db_type "db_type": db_type
} }
return await _make_connection_request("delete_knowledge_base", params) return await _make_connection_request("delete_knowledge_base", params)
async def search_query(query: str, userid: str, knowledge_base_ids: list, limit: int = 5, offset: int = 0, async def search_query(query: str, userid: str, knowledge_base_ids: list, limit: int = 5, offset: int = 0,
use_rerank: bool = True, db_type: str = "") -> dict: use_rerank: bool = True, db_type: str = "") -> dict:
"""执行纯向量搜索""" """执行纯向量搜索"""
params = { params = {
"query": query, "query": query,
"userid": userid, "userid": userid,
"knowledge_base_ids": knowledge_base_ids, "knowledge_base_ids": knowledge_base_ids,
"limit": limit, "limit": limit,
"offset": offset, "offset": offset,
"use_rerank": use_rerank, "use_rerank": use_rerank,
"db_type": db_type "db_type": db_type
} }
return await _make_connection_request("search_query", params) return await _make_connection_request("search_query", params)
async def fused_search(query: str, userid: str, knowledge_base_ids: list, limit: int = 5, offset: int = 0, async def fused_search(query: str, userid: str, knowledge_base_ids: list, limit: int = 5, offset: int = 0,
use_rerank: bool = True, db_type: str = "") -> dict: use_rerank: bool = True, db_type: str = "") -> dict:
"""执行融合搜索(向量 + 三元组)""" """执行融合搜索(向量 + 三元组)"""
params = { params = {
"query": query, "query": query,
"userid": userid, "userid": userid,
"knowledge_base_ids": knowledge_base_ids, "knowledge_base_ids": knowledge_base_ids,
"limit": limit, "limit": limit,
"offset": offset, "offset": offset,
"use_rerank": use_rerank, "use_rerank": use_rerank,
"db_type": db_type "db_type": db_type
} }
return await _make_connection_request("fused_search", params) return await _make_connection_request("fused_search", params)
async def list_user_files(userid: str, db_type: str = "") -> dict: async def list_user_files(userid: str, db_type: str = "") -> dict:
"""列出用户的所有知识库及其文件""" """列出用户的所有知识库及其文件"""
params = { params = {
"userid": userid, "userid": userid,
"db_type": db_type "db_type": db_type
} }
return await _make_connection_request("list_user_files", params) return await _make_connection_request("list_user_files", params)
async def list_all_knowledge_bases(db_type: str = "") -> dict: async def list_all_knowledge_bases(db_type: str = "") -> dict:
"""列出所有用户的知识库及其文件""" """列出所有用户的知识库及其文件"""
return await _make_connection_request("list_all_knowledge_bases", {"db_type": db_type}) return await _make_connection_request("list_all_knowledge_bases", {"db_type": db_type})
async def docs() -> dict: async def docs() -> dict:
"""列出所有用户的知识库及其文件""" """列出所有用户的知识库及其文件"""
return await _make_connection_request("docs", {}) return await _make_connection_request("docs", {})
async def get_user_kdbs(request): async def get_user_kdbs(request):
env = request._run_ns env = request._run_ns
@ -151,23 +151,23 @@ async def get_user_kdbs(request):
return recs return recs
def load_rag(): def load_rag():
""" """
初始化 ServerEnv绑定 MilvusConnection 的所有功能 初始化 ServerEnv绑定 MilvusConnection 的所有功能
""" """
env = ServerEnv() env = ServerEnv()
env.create_collection = create_collection env.create_collection = create_collection
env.delete_collection = delete_collection env.delete_collection = delete_collection
env.insert_document = insert_document env.insert_document = insert_document
env.delete_document = delete_document env.delete_document = delete_document
env.delete_knowledge_base = delete_knowledge_base env.delete_knowledge_base = delete_knowledge_base
env.search_query = search_query env.search_query = search_query
env.fused_search = fused_search env.fused_search = fused_search
env.list_user_files = list_user_files env.list_user_files = list_user_files
env.list_all_knowledge_bases = list_all_knowledge_bases env.list_all_knowledge_bases = list_all_knowledge_bases
env.docs = docs env.docs = docs
env.RagFileMgr = RagFileMgr env.RagFileMgr = RagFileMgr
env.set_program = set_program env.set_program = set_program
env.get_rag_programs = get_rag_programs env.get_rag_programs = get_rag_programs
env.get_ragllms = get_ragllms env.get_ragllms = get_ragllms
env.get_user_kdbs = get_user_kdbs env.get_user_kdbs = get_user_kdbs