bugfix
This commit is contained in:
parent
b9ec5336fa
commit
02b5413f4d
240
rag/init.py
240
rag/init.py
@ -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: 服务响应，包含 status、message、collection_name 等字段。
|
dict: 服务响应，包含 status、message、collection_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
|
||||||
|
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user