添加rag服务
This commit is contained in:
parent
8cf8b975f7
commit
08fac45422
27
rag/base_connection.py
Normal file
27
rag/base_connection.py
Normal file
@ -0,0 +1,27 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict
|
||||
from appPublic.log import debug, error, info, exception
|
||||
|
||||
connection_pathMap = {}
|
||||
|
||||
def connection_register(connection_key, Klass):
|
||||
"""为给定的连接键注册一个连接类"""
|
||||
global connection_pathMap
|
||||
connection_pathMap[connection_key] = Klass
|
||||
info(f"Registered {connection_key} with class {Klass}")
|
||||
|
||||
def get_connection_class(connection_path):
|
||||
"""根据连接路径查找对应的连接类"""
|
||||
global connection_pathMap
|
||||
debug(f"connection_pathMap: {connection_pathMap}")
|
||||
klass = connection_pathMap.get(connection_path)
|
||||
if klass is None:
|
||||
error(f"{connection_path} has not mapping to a connection class")
|
||||
raise Exception(f"{connection_path} has not mapping to a connection class")
|
||||
return klass
|
||||
|
||||
class BaseConnection(ABC):
|
||||
@abstractmethod
|
||||
async def handle_connection(self, action: str, params: Dict = None) -> Dict:
|
||||
"""处理数据库操作，根据 action 执行创建集合等"""
|
||||
pass
|
||||
649
rag/connection.py
Normal file
649
rag/connection.py
Normal file
@ -0,0 +1,649 @@
|
||||
import llmengine.milvus_connection
|
||||
from traceback import format_exc
|
||||
import argparse
|
||||
from aiohttp import web
|
||||
from llmengine.base_connection import get_connection_class
|
||||
from appPublic.registerfunction import RegisterFunction
|
||||
from appPublic.log import debug, error, info
|
||||
from ahserver.serverenv import ServerEnv
|
||||
from ahserver.webapp import webserver
|
||||
import os
|
||||
import json
|
||||
|
||||
helptext = """Milvus Connection Service API (using pymilvus Collection API):
|
||||
|
||||
1. Create Collection Endpoint:
|
||||
path: /v1/createcollection
|
||||
method: POST
|
||||
headers: {"Content-Type": "application/json"}
|
||||
data: {
|
||||
"db_type": "textdb" // 可选，若不提供则使用默认集合 ragdb
|
||||
}
|
||||
response:
|
||||
- Success: HTTP 200, {"status": "success", "collection_name": "ragdb" or "ragdb_textdb", "message": "集合 ragdb 或 ragdb_textdb 创建成功"}
|
||||
- Error: HTTP 400, {"status": "error", "collection_name": "ragdb" or "ragdb_textdb", "message": "<error message>"}
|
||||
|
||||
2. Delete Collection Endpoint:
|
||||
path: /v1/deletecollection
|
||||
method: POST
|
||||
headers: {"Content-Type": "application/json"}
|
||||
data: {
|
||||
"db_type": "textdb" // 可选，若不提供则删除默认集合 ragdb
|
||||
}
|
||||
response:
|
||||
- Success: HTTP 200, {"status": "success", "collection_name": "ragdb" or "ragdb_textdb", "message": "集合 ragdb 或 ragdb_textdb 删除成功"}
|
||||
- Success (collection does not exist): HTTP 200, {"status": "success", "collection_name": "ragdb" or "ragdb_textdb", "message": "集合 ragdb 或 ragdb_textdb 不存在，无需删除"}
|
||||
- Error: HTTP 400, {"status": "error", "collection_name": "ragdb" or "ragdb_textdb", "message": "<error message>"}
|
||||
|
||||
3. Insert File Endpoint:
|
||||
path: /v1/insertfile
|
||||
method: POST
|
||||
headers: {"Content-Type": "application/json"}
|
||||
data: {
|
||||
"file_path": "/path/to/file.txt", // 必填，文件路径
|
||||
"userid": "user123", // 必填，用户 ID
|
||||
"db_type": "textdb", // 可选，若不提供则使用默认集合 ragdb
|
||||
"knowledge_base_id": "kb123" // 必填，知识库 ID
|
||||
}
|
||||
response:
|
||||
- Success: HTTP 200, {"status": "success", "document_id": "<uuid>", "collection_name": "ragdb" or "ragdb_textdb", "message": "文件 <file_path> 成功嵌入并处理三元组", "status_code": 200}
|
||||
- Success (triples failed): HTTP 200, {"status": "success", "document_id": "<uuid>", "collection_name": "ragdb" or "ragdb_textdb", "message": "文件 <file_path> 成功嵌入，但三元组处理失败: <error>", "status_code": 200}
|
||||
- Error: HTTP 400, {"status": "error", "document_id": "", "collection_name": "ragdb" or "ragdb_textdb", "message": "<error message>", "status_code": 400}
|
||||
|
||||
4. Delete Document Endpoint:
|
||||
path: /v1/deletefile
|
||||
method: POST
|
||||
headers: {"Content-Type": "application/json"}
|
||||
data: {
|
||||
"userid": "user123", // 必填，用户 ID
|
||||
"filename": "file.txt", // 必填，文件名
|
||||
"db_type": "textdb", // 可选，若不提供则使用默认集合 ragdb
|
||||
"knowledge_base_id": "kb123" // 必填，知识库 ID
|
||||
}
|
||||
response:
|
||||
- Success: HTTP 200, {"status": "success", "document_id": "<uuid1,uuid2>", "collection_name": "ragdb" or "ragdb_textdb", "message": "成功删除 <count> 条 Milvus 记录，<nodes> 个 Neo4j 节点，<rels> 个 Neo4j 关系，userid=<userid>, filename=<filename>", "status_code": 200}
|
||||
- Success (no records): HTTP 200, {"status": "success", "document_id": "", "collection_name": "ragdb" or "ragdb_textdb", "message": "没有找到 userid=<userid>, filename=<filename>, knowledge_base_id=<knowledge_base_id> 的记录，无需删除", "status_code": 200}
|
||||
- Success (collection missing): HTTP 200, {"status": "success", "document_id": "", "collection_name": "ragdb" or "ragdb_textdb", "message": "集合 <collection_name> 不存在，无需删除", "status_code": 200}
|
||||
- Error: HTTP 400, {"status": "error", "document_id": "", "collection_name": "ragdb" or "ragdb_textdb", "message": "<error message>", "status_code": 400}
|
||||
|
||||
5. Fused Search Query Endpoint:
|
||||
path: /v1/fusedsearchquery
|
||||
method: POST
|
||||
headers: {"Content-Type": "application/json"}
|
||||
data: {
|
||||
"query": "苹果公司在北京开设新店",
|
||||
"userid": "user1",
|
||||
"db_type": "textdb", // 可选，若不提供则使用默认集合 ragdb
|
||||
"knowledge_base_ids": ["kb123"],
|
||||
"limit": 5,
|
||||
"offset": 0,
|
||||
"use_rerank": true
|
||||
}
|
||||
response:
|
||||
- Success: HTTP 200, {
|
||||
"status": "success",
|
||||
"results": [
|
||||
{
|
||||
"text": "<完整文本内容>",
|
||||
"distance": 0.95,
|
||||
"source": "fused_query_with_triplets",
|
||||
"rerank_score": 0.92, // 若 use_rerank=true
|
||||
"metadata": {
|
||||
"userid": "user1",
|
||||
"document_id": "<uuid>",
|
||||
"filename": "file.txt",
|
||||
"file_path": "/path/to/file.txt",
|
||||
"upload_time": "<iso_timestamp>",
|
||||
"file_type": "txt"
|
||||
}
|
||||
},
|
||||
...
|
||||
],
|
||||
"timing": {
|
||||
"collection_load": <float>, // 集合加载耗时（秒）
|
||||
"entity_extraction": <float>, // 实体提取耗时（秒）
|
||||
"triplet_matching": <float>, // 三元组匹配耗时（秒）
|
||||
"triplet_text_combine": <float>, // 拼接三元组文本耗时（秒）
|
||||
"embedding_generation": <float>, // 嵌入向量生成耗时（秒）
|
||||
"vector_search": <float>, // 向量搜索耗时（秒）
|
||||
"deduplication": <float>, // 去重耗时（秒）
|
||||
"reranking": <float>, // 重排序耗时（秒，若 use_rerank=true）
|
||||
"total_time": <float> // 总耗时（秒）
|
||||
},
|
||||
"collection_name": "ragdb" or "ragdb_textdb"
|
||||
}
|
||||
- Error: HTTP 400, {
|
||||
"status": "error",
|
||||
"message": "<error message>",
|
||||
"collection_name": "ragdb" or "ragdb_textdb"
|
||||
}
|
||||
6. Search Query Endpoint:
|
||||
path: /v1/searchquery
|
||||
method: POST
|
||||
headers: {"Content-Type": "application/json"}
|
||||
data: {
|
||||
"query": "知识图谱的知识融合是什么？",
|
||||
"userid": "user1",
|
||||
"db_type": "textdb", // 可选，若不提供则使用默认集合 ragdb
|
||||
"knowledge_base_ids": ["kb123"],
|
||||
"limit": 5,
|
||||
"offset": 0,
|
||||
"use_rerank": true
|
||||
}
|
||||
response:
|
||||
- Success: HTTP 200, {
|
||||
"status": "success",
|
||||
"results": [
|
||||
{
|
||||
"text": "<完整文本内容>",
|
||||
"distance": 0.95,
|
||||
"source": "vector_query",
|
||||
"rerank_score": 0.92, // 若 use_rerank=true
|
||||
"metadata": {
|
||||
"userid": "user1",
|
||||
"document_id": "<uuid>",
|
||||
"filename": "file.txt",
|
||||
"file_path": "/path/to/file.txt",
|
||||
"upload_time": "<iso_timestamp>",
|
||||
"file_type": "txt"
|
||||
}
|
||||
},
|
||||
...
|
||||
],
|
||||
"timing": {
|
||||
"collection_load": <float>, // 集合加载耗时（秒）
|
||||
"embedding_generation": <float>, // 嵌入向量生成耗时（秒）
|
||||
"vector_search": <float>, // 向量搜索耗时（秒）
|
||||
"deduplication": <float>, // 去重耗时（秒）
|
||||
"reranking": <float>, // 重排序耗时（秒，若 use_rerank=true）
|
||||
"total_time": <float> // 总耗时（秒）
|
||||
},
|
||||
"collection_name": "ragdb" or "ragdb_textdb"
|
||||
}
|
||||
- Error: HTTP 400, {
|
||||
"status": "error",
|
||||
"message": "<error message>",
|
||||
"collection_name": "ragdb" or "ragdb_textdb"
|
||||
}
|
||||
|
||||
7. List User Files Endpoint:
|
||||
path: /v1/listuserfiles
|
||||
method: POST
|
||||
headers: {"Content-Type": "application/json"}
|
||||
data: {
|
||||
"userid": "user1",
|
||||
"db_type": "textdb" // 可选，若不提供则使用默认集合 ragdb
|
||||
}
|
||||
response:
|
||||
- Success: HTTP 200, {
|
||||
"status": "success",
|
||||
"files_by_knowledge_base": {
|
||||
"kb123": [
|
||||
{
|
||||
"document_id": "<uuid>",
|
||||
"filename": "file1.txt",
|
||||
"file_path": "/path/to/file1.txt",
|
||||
"upload_time": "<iso_timestamp>",
|
||||
"file_type": "txt",
|
||||
"knowledge_base_id": "kb123"
|
||||
},
|
||||
...
|
||||
],
|
||||
"kb456": [
|
||||
{
|
||||
"document_id": "<uuid>",
|
||||
"filename": "file2.pdf",
|
||||
"file_path": "/path/to/file2.pdf",
|
||||
"upload_time": "<iso_timestamp>",
|
||||
"file_type": "pdf",
|
||||
"knowledge_base_id": "kb456"
|
||||
},
|
||||
...
|
||||
]
|
||||
},
|
||||
"collection_name": "ragdb" or "ragdb_textdb"
|
||||
}
|
||||
- Error: HTTP 400, {
|
||||
"status": "error",
|
||||
"message": "<error message>",
|
||||
"collection_name": "ragdb" or "ragdb_textdb"
|
||||
}
|
||||
8. Connection Endpoint (for compatibility):
|
||||
path: /v1/connection
|
||||
method: POST
|
||||
headers: {"Content-Type": "application/json"}
|
||||
data: {
|
||||
"action": "<initialize|get_params|create_collection|delete_collection|insert_document|delete_document|fused_search|search_query|list_user_files>",
|
||||
"params": {...}
|
||||
}
|
||||
response:
|
||||
- Success: HTTP 200, {"status": "success", ...}
|
||||
- Error: HTTP 400, {"status": "error", "message": "<error message>"}
|
||||
|
||||
9. Docs Endpoint:
|
||||
path: /docs
|
||||
method: GET
|
||||
response: This help text
|
||||
|
||||
10. Delete Knowledge Base Endpoint:
|
||||
path: /v1/deleteknowledgebase
|
||||
method: POST
|
||||
headers: {"Content-Type": "application/json"}
|
||||
data: {
|
||||
"userid": "user123", // 必填，用户 ID
|
||||
"knowledge_base_id": "kb123",// 必填，知识库 ID
|
||||
"db_type": "textdb" // 可选，若不提供则使用默认集合 ragdb
|
||||
}
|
||||
response:
|
||||
- Success: HTTP 200, {"status": "success", "document_id": "<uuid1,uuid2>", "filename": "<filename1,filename2>", "collection_name": "ragdb" or "ragdb_textdb", "message": "成功删除 <count> 条 Milvus 记录，<nodes> 个 Neo4j 节点，<rels> 个 Neo4j 关系，userid=<userid>, knowledge_base_id=<knowledge_base_id>", "status_code": 200}
|
||||
- Success (no records): HTTP 200, {"status": "success", "document_id": "", "filename": "", "collection_name": "ragdb" or "ragdb_textdb", "message": "没有找到 userid=<userid>, knowledge_base_id=<knowledge_base_id> 的记录，无需删除", "status_code": 200}
|
||||
- Success (collection missing): HTTP 200, {"status": "success", "document_id": "", "filename": "", "collection_name": "ragdb" or "ragdb_textdb", "message": "集合 <collection_name> 不存在，无需删除", "status_code": 200}
|
||||
- Error: HTTP 400, {"status": "error", "document_id": "", "filename": "", "collection_name": "ragdb" or "ragdb_textdb", "message": "<error message>", "status_code": 400}
|
||||
|
||||
10. List All Knowledge Bases Endpoint:
|
||||
path: /v1/listallknowledgebases
|
||||
method: POST
|
||||
headers: {"Content-Type": "application/json"}
|
||||
data: {
|
||||
"db_type": "textdb" // 可选，若不提供则使用默认集合 ragdb
|
||||
}
|
||||
response:
|
||||
- Success: HTTP 200, {
|
||||
"status": "success",
|
||||
"users_knowledge_bases": {
|
||||
"user1": {
|
||||
"kb123": [
|
||||
{
|
||||
"document_id": "<uuid>",
|
||||
"filename": "file1.txt",
|
||||
"file_path": "/path/to/file1.txt",
|
||||
"upload_time": "<iso_timestamp>",
|
||||
"file_type": "txt",
|
||||
"knowledge_base_id": "kb123"
|
||||
},
|
||||
...
|
||||
],
|
||||
"kb456": [
|
||||
{
|
||||
"document_id": "<uuid>",
|
||||
"filename": "file2.pdf",
|
||||
"file_path": "/path/to/file2.pdf",
|
||||
"upload_time": "<iso_timestamp>",
|
||||
"file_type": "pdf",
|
||||
"knowledge_base_id": "kb456"
|
||||
},
|
||||
...
|
||||
]
|
||||
},
|
||||
"user2": {...}
|
||||
},
|
||||
"collection_name": "ragdb" or "ragdb_textdb",
|
||||
"message": "成功列出 <count> 个用户的知识库和文件",
|
||||
"status_code": 200
|
||||
}
|
||||
- Error: HTTP 400, {
|
||||
"status": "error",
|
||||
"users_knowledge_bases": {},
|
||||
"collection_name": "ragdb" or "ragdb_textdb",
|
||||
"message": "<error message>",
|
||||
"status_code": 400
|
||||
}
|
||||
"""
|
||||
|
||||
def init():
|
||||
rf = RegisterFunction()
|
||||
rf.register('createcollection', create_collection)
|
||||
rf.register('deletecollection', delete_collection)
|
||||
rf.register('insertfile', insert_file)
|
||||
rf.register('deletefile', delete_file)
|
||||
rf.register('deleteknowledgebase', delete_knowledge_base)
|
||||
rf.register('fusedsearchquery', fused_search_query)
|
||||
rf.register('searchquery', search_query)
|
||||
rf.register('listuserfiles', list_user_files)
|
||||
rf.register('listallknowledgebases', list_all_knowledge_bases)
|
||||
rf.register('connection', handle_connection)
|
||||
rf.register('docs', docs)
|
||||
|
||||
async def docs(request, params_kw, *params, **kw):
|
||||
return web.Response(text=helptext, content_type='text/plain')
|
||||
|
||||
async def not_implemented(request, params_kw, *params, **kw):
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"message": "功能尚未实现"
|
||||
}, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=501)
|
||||
|
||||
async def create_collection(request, params_kw, *params, **kw):
|
||||
debug(f'{params_kw=}')
|
||||
se = ServerEnv()
|
||||
engine = se.engine
|
||||
db_type = params_kw.get('db_type', '')
|
||||
collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
|
||||
try:
|
||||
result = await engine.handle_connection("create_collection", {"db_type": db_type})
|
||||
debug(f'{result=}')
|
||||
return web.json_response(result, dumps=lambda obj: json.dumps(obj, ensure_ascii=False))
|
||||
except Exception as e:
|
||||
error(f'创建集合失败: {str(e)}')
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"collection_name": collection_name,
|
||||
"message": str(e)
|
||||
}, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
|
||||
|
||||
async def delete_collection(request, params_kw, *params, **kw):
|
||||
debug(f'{params_kw=}')
|
||||
se = ServerEnv()
|
||||
engine = se.engine
|
||||
db_type = params_kw.get('db_type', '')
|
||||
collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
|
||||
try:
|
||||
result = await engine.handle_connection("delete_collection", {"db_type": db_type})
|
||||
debug(f'{result=}')
|
||||
return web.json_response(result, dumps=lambda obj: json.dumps(obj, ensure_ascii=False))
|
||||
except Exception as e:
|
||||
error(f'删除集合失败: {str(e)}')
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"collection_name": collection_name,
|
||||
"message": str(e)
|
||||
}, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
|
||||
|
||||
async def insert_file(request, params_kw, *params, **kw):
|
||||
debug(f'Received params: {params_kw=}')
|
||||
se = ServerEnv()
|
||||
engine = se.engine
|
||||
file_path = params_kw.get('file_path', '')
|
||||
userid = params_kw.get('userid', '')
|
||||
db_type = params_kw.get('db_type', '')
|
||||
knowledge_base_id = params_kw.get('knowledge_base_id', '')
|
||||
collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
|
||||
try:
|
||||
required_fields = ['file_path', 'userid', 'knowledge_base_id']
|
||||
missing_fields = [field for field in required_fields if field not in params_kw or not params_kw[field]]
|
||||
if missing_fields:
|
||||
raise ValueError(f"缺少必填字段: {', '.join(missing_fields)}")
|
||||
|
||||
debug(
|
||||
f'Calling insert_document with: file_path={file_path}, userid={userid}, db_type={db_type}, knowledge_base_id={knowledge_base_id}')
|
||||
result = await engine.handle_connection("insert_document", {
|
||||
"file_path": file_path,
|
||||
"userid": userid,
|
||||
"db_type": db_type,
|
||||
"knowledge_base_id": knowledge_base_id
|
||||
})
|
||||
debug(f'Insert result: {result=}')
|
||||
status = 200 if result.get("status") == "success" else 400
|
||||
return web.json_response(result, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=status)
|
||||
except Exception as e:
|
||||
error(f'插入文件失败: {str(e)}')
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"collection_name": collection_name,
|
||||
"document_id": "",
|
||||
"message": str(e)
|
||||
}, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
|
||||
|
||||
async def delete_file(request, params_kw, *params, **kw):
|
||||
debug(f'Received delete_file params: {params_kw=}')
|
||||
se = ServerEnv()
|
||||
engine = se.engine
|
||||
userid = params_kw.get('userid', '')
|
||||
filename = params_kw.get('filename', '')
|
||||
db_type = params_kw.get('db_type', '')
|
||||
knowledge_base_id = params_kw.get('knowledge_base_id', '')
|
||||
collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
|
||||
try:
|
||||
required_fields = ['userid', 'filename', 'knowledge_base_id']
|
||||
missing_fields = [field for field in required_fields if field not in params_kw or not params_kw[field]]
|
||||
if missing_fields:
|
||||
raise ValueError(f"缺少必填字段: {', '.join(missing_fields)}")
|
||||
|
||||
debug(f'Calling delete_document with: userid={userid}, filename={filename}, db_type={db_type}, knowledge_base_id={knowledge_base_id}')
|
||||
result = await engine.handle_connection("delete_document", {
|
||||
"userid": userid,
|
||||
"filename": filename,
|
||||
"knowledge_base_id": knowledge_base_id,
|
||||
"db_type": db_type
|
||||
})
|
||||
debug(f'Delete result: {result=}')
|
||||
status = 200 if result.get("status") == "success" else 400
|
||||
return web.json_response(result, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=status)
|
||||
except Exception as e:
|
||||
error(f'删除文件失败: {str(e)}')
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"collection_name": collection_name,
|
||||
"document_id": "",
|
||||
"message": str(e),
|
||||
"status_code": 400
|
||||
}, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
|
||||
|
||||
async def delete_knowledge_base(request, params_kw, *params, **kw):
|
||||
debug(f'Received delete_knowledge_base params: {params_kw=}')
|
||||
se = ServerEnv()
|
||||
engine = se.engine
|
||||
userid = params_kw.get('userid', '')
|
||||
knowledge_base_id = params_kw.get('knowledge_base_id', '')
|
||||
db_type = params_kw.get('db_type', '')
|
||||
collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
|
||||
try:
|
||||
required_fields = ['userid', 'knowledge_base_id']
|
||||
missing_fields = [field for field in required_fields if field not in params_kw or not params_kw[field]]
|
||||
if missing_fields:
|
||||
raise ValueError(f"缺少必填字段: {', '.join(missing_fields)}")
|
||||
|
||||
debug(
|
||||
f'Calling delete_knowledge_base with: userid={userid}, knowledge_base_id={knowledge_base_id}, db_type={db_type}')
|
||||
result = await engine.handle_connection("delete_knowledge_base", {
|
||||
"userid": userid,
|
||||
"knowledge_base_id": knowledge_base_id,
|
||||
"db_type": db_type
|
||||
})
|
||||
debug(f'Delete knowledge base result: {result=}')
|
||||
status = 200 if result.get("status") == "success" else 400
|
||||
return web.json_response(result, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=status)
|
||||
except Exception as e:
|
||||
error(f'删除知识库失败: {str(e)}')
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"collection_name": collection_name,
|
||||
"document_id": "",
|
||||
"filename": "",
|
||||
"message": str(e),
|
||||
"status_code": 400
|
||||
}, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
|
||||
|
||||
async def search_query(request, params_kw, *params, **kw):
|
||||
debug(f'Received search_query params: {params_kw=}')
|
||||
se = ServerEnv()
|
||||
engine = se.engine
|
||||
query = params_kw.get('query')
|
||||
userid = params_kw.get('userid')
|
||||
db_type = params_kw.get('db_type', '')
|
||||
knowledge_base_ids = params_kw.get('knowledge_base_ids')
|
||||
limit = params_kw.get('limit', 5)
|
||||
offset = params_kw.get('offset', 0)
|
||||
use_rerank = params_kw.get('use_rerank', True)
|
||||
collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
|
||||
try:
|
||||
if not all([query, userid, knowledge_base_ids]):
|
||||
debug(f'query, userid 或 knowledge_base_ids 未提供')
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"message": "query, userid 或 knowledge_base_ids 未提供",
|
||||
"collection_name": collection_name
|
||||
}, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
|
||||
result = await engine.handle_connection("search_query", {
|
||||
"query": query,
|
||||
"userid": userid,
|
||||
"knowledge_base_ids": knowledge_base_ids,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
"use_rerank": use_rerank,
|
||||
"db_type": db_type
|
||||
})
|
||||
debug(f'Search result: {result=}')
|
||||
response = {
|
||||
"status": "success",
|
||||
"results": result.get("results", []),
|
||||
"timing": result.get("timing", {}),
|
||||
"collection_name": collection_name
|
||||
}
|
||||
return web.json_response(response, dumps=lambda obj: json.dumps(obj, ensure_ascii=False))
|
||||
except Exception as e:
|
||||
error(f'纯向量搜索失败: {str(e)}')
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"message": str(e),
|
||||
"collection_name": collection_name
|
||||
}, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
|
||||
|
||||
async def fused_search_query(request, params_kw, *params, **kw):
|
||||
debug(f'Received fused_search_query params: {params_kw=}')
|
||||
se = ServerEnv()
|
||||
engine = se.engine
|
||||
query = params_kw.get('query')
|
||||
userid = params_kw.get('userid')
|
||||
db_type = params_kw.get('db_type', '')
|
||||
knowledge_base_ids = params_kw.get('knowledge_base_ids')
|
||||
limit = params_kw.get('limit', 5)
|
||||
offset = params_kw.get('offset', 0)
|
||||
use_rerank = params_kw.get('use_rerank', True)
|
||||
collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
|
||||
try:
|
||||
if not all([query, userid, knowledge_base_ids]):
|
||||
debug(f'query, userid 或 knowledge_base_ids 未提供')
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"message": "query, userid 或 knowledge_base_ids 未提供",
|
||||
"collection_name": collection_name
|
||||
}, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
|
||||
result = await engine.handle_connection("fused_search", {
|
||||
"query": query,
|
||||
"userid": userid,
|
||||
"knowledge_base_ids": knowledge_base_ids,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
"use_rerank": use_rerank,
|
||||
"db_type": db_type
|
||||
})
|
||||
debug(f'Fused search result: {result=}')
|
||||
response = {
|
||||
"status": "success",
|
||||
"results": result.get("results", []),
|
||||
"timing": result.get("timing", {}),
|
||||
"collection_name": collection_name
|
||||
}
|
||||
return web.json_response(response, dumps=lambda obj: json.dumps(obj, ensure_ascii=False))
|
||||
except Exception as e:
|
||||
error(f'融合搜索失败: {str(e)}')
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"message": str(e),
|
||||
"collection_name": collection_name
|
||||
}, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
|
||||
|
||||
async def list_user_files(request, params_kw, *params, **kw):
|
||||
debug(f'{params_kw=}')
|
||||
se = ServerEnv()
|
||||
engine = se.engine
|
||||
userid = params_kw.get('userid')
|
||||
db_type = params_kw.get('db_type', '')
|
||||
collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
|
||||
try:
|
||||
if not userid:
|
||||
debug(f'userid 未提供')
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"message": "userid 未提供",
|
||||
"collection_name": collection_name
|
||||
}, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
|
||||
result = await engine.handle_connection("list_user_files", {
|
||||
"userid": userid,
|
||||
"db_type": db_type
|
||||
})
|
||||
debug(f'{result=}')
|
||||
response = {
|
||||
"status": "success",
|
||||
"files_by_knowledge_base": result,
|
||||
"collection_name": collection_name
|
||||
}
|
||||
return web.json_response(response, dumps=lambda obj: json.dumps(obj, ensure_ascii=False))
|
||||
except Exception as e:
|
||||
error(f'列出用户文件失败: {str(e)}')
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"message": str(e),
|
||||
"collection_name": collection_name
|
||||
}, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
|
||||
|
||||
async def list_all_knowledge_bases(request, params_kw, *params, **kw):
|
||||
debug(f'{params_kw=}')
|
||||
se = ServerEnv()
|
||||
engine = se.engine
|
||||
db_type = params_kw.get('db_type', '')
|
||||
collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
|
||||
try:
|
||||
result = await engine.handle_connection("list_all_knowledge_bases", {
|
||||
"db_type": db_type
|
||||
})
|
||||
debug(f'{result=}')
|
||||
response = {
|
||||
"status": result.get("status", "success"),
|
||||
"users_knowledge_bases": result.get("users_knowledge_bases", {}),
|
||||
"collection_name": collection_name,
|
||||
"message": result.get("message", ""),
|
||||
"status_code": result.get("status_code", 200)
|
||||
}
|
||||
return web.json_response(response, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=response["status_code"])
|
||||
except Exception as e:
|
||||
error(f'列出所有用户知识库失败: {str(e)}')
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"users_knowledge_bases": {},
|
||||
"collection_name": collection_name,
|
||||
"message": str(e),
|
||||
"status_code": 400
|
||||
}, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
|
||||
|
||||
async def handle_connection(request, params_kw, *params, **kw):
|
||||
debug(f'{params_kw=}')
|
||||
se = ServerEnv()
|
||||
engine = se.engine
|
||||
try:
|
||||
data = await request.json()
|
||||
action = data.get('action')
|
||||
if not action:
|
||||
debug(f'action 未提供')
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"message": "action 参数未提供"
|
||||
}, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
|
||||
result = await engine.handle_connection(action, data.get('params', {}))
|
||||
debug(f'{result=}')
|
||||
return web.json_response(result, dumps=lambda obj: json.dumps(obj, ensure_ascii=False))
|
||||
except Exception as e:
|
||||
error(f'处理连接操作失败: {str(e)}')
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"message": str(e)
|
||||
}, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(prog="Milvus Connection Service")
|
||||
parser.add_argument('-w', '--workdir')
|
||||
parser.add_argument('-p', '--port', default='8888')
|
||||
parser.add_argument('connection_path')
|
||||
args = parser.parse_args()
|
||||
debug(f"Arguments: {args}")
|
||||
Klass = get_connection_class(args.connection_path)
|
||||
se = ServerEnv()
|
||||
se.engine = Klass()
|
||||
workdir = args.workdir or os.getcwd()
|
||||
port = args.port
|
||||
debug(f'{args=}')
|
||||
webserver(init, workdir, port)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
1073
rag/milvus_connection.py
Normal file
1073
rag/milvus_connection.py
Normal file
File diff suppressed because it is too large
Load Diff
Loading…
x
Reference in New Issue
Block a user