382 lines
16 KiB
Python
382 lines
16 KiB
Python
import os
|
||
import asyncio
|
||
import traceback
|
||
from fastapi import APIRouter, File, UploadFile, HTTPException, Depends, Body, Form, Query
|
||
|
||
from src.utils import logger, hashstr
|
||
from src import executor, config, knowledge_base, graph_base
|
||
from server.utils.auth_middleware import get_admin_user
|
||
from server.models.user_model import User
|
||
|
||
data = APIRouter(prefix="/data")
|
||
|
||
|
||
@data.get("/")
|
||
async def get_databases(current_user: User = Depends(get_admin_user)):
|
||
try:
|
||
database = knowledge_base.get_databases()
|
||
except Exception as e:
|
||
logger.error(f"获取数据库列表失败 {e}, {traceback.format_exc()}")
|
||
return {"message": f"获取数据库列表失败 {e}", "databases": []}
|
||
return database
|
||
|
||
@data.get("/kb-types")
|
||
async def get_knowledge_base_types(current_user: User = Depends(get_admin_user)):
|
||
"""获取支持的知识库类型"""
|
||
try:
|
||
kb_types = knowledge_base.get_supported_kb_types()
|
||
return {"kb_types": kb_types, "message": "success"}
|
||
except Exception as e:
|
||
logger.error(f"获取知识库类型失败 {e}, {traceback.format_exc()}")
|
||
return {"message": f"获取知识库类型失败 {e}", "kb_types": {}}
|
||
|
||
@data.get("/stats")
|
||
async def get_knowledge_base_statistics(current_user: User = Depends(get_admin_user)):
|
||
"""获取知识库统计信息"""
|
||
try:
|
||
stats = knowledge_base.get_statistics()
|
||
return {"stats": stats, "message": "success"}
|
||
except Exception as e:
|
||
logger.error(f"获取知识库统计失败 {e}, {traceback.format_exc()}")
|
||
return {"message": f"获取知识库统计失败 {e}", "stats": {}}
|
||
|
||
@data.get("/query-params/{db_id}")
|
||
async def get_knowledge_base_query_params(db_id: str, current_user: User = Depends(get_admin_user)):
|
||
"""获取知识库类型特定的查询参数"""
|
||
try:
|
||
# 获取数据库信息
|
||
db_info = knowledge_base.get_database_info(db_id)
|
||
if not db_info:
|
||
raise HTTPException(status_code=404, detail="Database not found")
|
||
|
||
kb_type = db_info.get("kb_type", "lightrag")
|
||
|
||
# 根据知识库类型返回不同的查询参数
|
||
if kb_type == "lightrag":
|
||
params = {
|
||
"type": "lightrag",
|
||
"options": [
|
||
{
|
||
"key": "mode",
|
||
"label": "检索模式",
|
||
"type": "select",
|
||
"default": "mix",
|
||
"options": [
|
||
{"value": "local", "label": "Local", "description": "上下文相关信息"},
|
||
{"value": "global", "label": "Global", "description": "全局知识"},
|
||
{"value": "hybrid", "label": "Hybrid", "description": "本地和全局混合"},
|
||
{"value": "naive", "label": "Naive", "description": "基本搜索"},
|
||
{"value": "mix", "label": "Mix", "description": "知识图谱和向量检索混合"},
|
||
]
|
||
},
|
||
{
|
||
"key": "only_need_context",
|
||
"label": "只使用上下文",
|
||
"type": "boolean",
|
||
"default": True,
|
||
"description": "只返回上下文，不生成回答"
|
||
},
|
||
{
|
||
"key": "only_need_prompt",
|
||
"label": "只使用提示",
|
||
"type": "boolean",
|
||
"default": False,
|
||
"description": "只返回提示，不进行检索"
|
||
},
|
||
{
|
||
"key": "top_k",
|
||
"label": "TopK",
|
||
"type": "number",
|
||
"default": 10,
|
||
"min": 1,
|
||
"max": 100,
|
||
"description": "返回的最大结果数量"
|
||
}
|
||
]
|
||
}
|
||
elif kb_type == "chroma":
|
||
params = {
|
||
"type": "chroma",
|
||
"options": [
|
||
{
|
||
"key": "top_k",
|
||
"label": "TopK",
|
||
"type": "number",
|
||
"default": 10,
|
||
"min": 1,
|
||
"max": 100,
|
||
"description": "返回的最大结果数量"
|
||
},
|
||
{
|
||
"key": "similarity_threshold",
|
||
"label": "相似度阈值",
|
||
"type": "number",
|
||
"default": 0.0,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.1,
|
||
"description": "过滤相似度低于此值的结果"
|
||
},
|
||
{
|
||
"key": "include_distances",
|
||
"label": "显示相似度",
|
||
"type": "boolean",
|
||
"default": True,
|
||
"description": "在结果中显示相似度分数"
|
||
}
|
||
]
|
||
}
|
||
elif kb_type == "milvus":
|
||
params = {
|
||
"type": "milvus",
|
||
"options": [
|
||
{
|
||
"key": "top_k",
|
||
"label": "TopK",
|
||
"type": "number",
|
||
"default": 10,
|
||
"min": 1,
|
||
"max": 100,
|
||
"description": "返回的最大结果数量"
|
||
},
|
||
{
|
||
"key": "similarity_threshold",
|
||
"label": "相似度阈值",
|
||
"type": "number",
|
||
"default": 0.0,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.1,
|
||
"description": "过滤相似度低于此值的结果"
|
||
},
|
||
{
|
||
"key": "include_distances",
|
||
"label": "显示相似度",
|
||
"type": "boolean",
|
||
"default": True,
|
||
"description": "在结果中显示相似度分数"
|
||
},
|
||
{
|
||
"key": "metric_type",
|
||
"label": "距离度量类型",
|
||
"type": "select",
|
||
"default": "COSINE",
|
||
"options": [
|
||
{"value": "COSINE", "label": "余弦相似度", "description": "适合文本语义相似度"},
|
||
{"value": "L2", "label": "欧几里得距离", "description": "适合数值型数据"},
|
||
{"value": "IP", "label": "内积", "description": "适合标准化向量"}
|
||
],
|
||
"description": "向量相似度计算方法"
|
||
}
|
||
]
|
||
}
|
||
else:
|
||
# 未知类型，返回基本参数
|
||
params = {
|
||
"type": "unknown",
|
||
"options": [
|
||
{
|
||
"key": "top_k",
|
||
"label": "TopK",
|
||
"type": "number",
|
||
"default": 10,
|
||
"min": 1,
|
||
"max": 100,
|
||
"description": "返回的最大结果数量"
|
||
}
|
||
]
|
||
}
|
||
|
||
return {"params": params, "message": "success"}
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取知识库查询参数失败 {e}, {traceback.format_exc()}")
|
||
return {"message": f"获取知识库查询参数失败 {e}", "params": {}}
|
||
|
||
@data.post("/")
|
||
async def create_database(
|
||
database_name: str = Body(...),
|
||
description: str = Body(...),
|
||
embed_model_name: str = Body(...),
|
||
kb_type: str = Body("lightrag"), # 新增：知识库类型参数，默认为lightrag
|
||
current_user: User = Depends(get_admin_user)
|
||
):
|
||
logger.debug(f"Create database {database_name} with kb_type {kb_type}")
|
||
try:
|
||
embed_info = config.embed_model_names[embed_model_name]
|
||
database_info = knowledge_base.create_database(
|
||
database_name,
|
||
description,
|
||
kb_type=kb_type, # 传递知识库类型
|
||
embed_info=embed_info
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"创建数据库失败 {e}, {traceback.format_exc()}")
|
||
return {"message": f"创建数据库失败 {e}", "status": "failed"}
|
||
return database_info
|
||
|
||
@data.delete("/")
|
||
async def delete_database(db_id, current_user: User = Depends(get_admin_user)):
|
||
logger.debug(f"Delete database {db_id}")
|
||
knowledge_base.delete_database(db_id)
|
||
return {"message": "删除成功"}
|
||
|
||
@data.post("/query-test")
|
||
async def query_test(query: str = Body(...), meta: dict = Body(...), current_user: User = Depends(get_admin_user)):
|
||
logger.debug(f"Query test in {meta}: {query}")
|
||
result = await knowledge_base.aquery(query, **meta)
|
||
return result
|
||
|
||
@data.post("/add-files")
|
||
async def add_files(db_id: str = Body(...), items: list[str] = Body(...), params: dict = Body(...), current_user: User = Depends(get_admin_user)):
|
||
logger.debug(f"Add files/urls for db_id {db_id}: {items} {params=}")
|
||
|
||
# 从 params 中获取 content_type，默认为 'file'
|
||
content_type = params.get('content_type', 'file')
|
||
|
||
try:
|
||
# 使用统一的 add_content 方法
|
||
processed_items = await knowledge_base.add_content(db_id, items, params=params)
|
||
|
||
item_type = "URLs" if content_type == 'url' else "files"
|
||
processed_failed_count = len([_p for _p in processed_items if _p['status'] == 'failed'])
|
||
processed_info = f"Processed {len(processed_items)} {item_type}, {processed_failed_count} {item_type} failed"
|
||
return {"message": processed_info, "items": processed_items, "status": "success"}
|
||
except Exception as e:
|
||
logger.error(f"Failed to process {content_type}s: {e}, {traceback.format_exc()}")
|
||
return {"message": f"Failed to process {content_type}s: {e}", "status": "failed"}
|
||
|
||
@data.post("/file-to-chunk")
|
||
async def file_to_chunk(db_id: str = Body(...), files: list[str] = Body(...), params: dict = Body(...), current_user: User = Depends(get_admin_user)):
|
||
logger.debug(f"File to chunk for db_id {db_id}: {files} {params=} (deprecated, use /add-files)")
|
||
# 兼容性路由，转发到新的统一接口
|
||
params['content_type'] = 'file'
|
||
return await add_files(db_id, files, params, current_user)
|
||
|
||
@data.post("/url-to-chunk")
|
||
async def url_to_chunk(db_id: str = Body(...), urls: list[str] = Body(...), params: dict = Body(...), current_user: User = Depends(get_admin_user)):
|
||
logger.debug(f"Url to chunk for db_id {db_id}: {urls} {params=} (deprecated, use /add-files)")
|
||
# 兼容性路由，转发到新的统一接口
|
||
params['content_type'] = 'url'
|
||
return await add_files(db_id, urls, params, current_user)
|
||
|
||
@data.post("/add-by-file")
|
||
async def create_document_by_file(db_id: str = Body(...), files: list[str] = Body(...), current_user: User = Depends(get_admin_user)):
|
||
raise ValueError("This method is deprecated. Use /add-files instead.")
|
||
|
||
@data.post("/add-by-chunks")
|
||
async def add_by_chunks(db_id: str = Body(...), file_chunks: dict = Body(...), current_user: User = Depends(get_admin_user)):
|
||
raise ValueError("This method is deprecated. Use /add-files instead.")
|
||
|
||
@data.get("/info")
|
||
async def get_database_info(db_id: str, current_user: User = Depends(get_admin_user)):
|
||
# logger.debug(f"Get database {db_id} info")
|
||
database = knowledge_base.get_database_info(db_id)
|
||
if database is None:
|
||
raise HTTPException(status_code=404, detail="Database not found")
|
||
return database
|
||
|
||
@data.delete("/document")
|
||
async def delete_document(db_id: str = Body(...), file_id: str = Body(...), current_user: User = Depends(get_admin_user)):
|
||
logger.debug(f"DELETE document {file_id} info in {db_id}")
|
||
await knowledge_base.delete_file(db_id, file_id)
|
||
return {"message": "删除成功"}
|
||
|
||
@data.get("/document")
|
||
async def get_document_info(db_id: str, file_id: str, current_user: User = Depends(get_admin_user)):
|
||
logger.debug(f"GET document {file_id} info in {db_id}")
|
||
|
||
try:
|
||
info = await knowledge_base.get_file_info(db_id, file_id)
|
||
except Exception as e:
|
||
logger.error(f"Failed to get file info, {e}, {db_id=}, {file_id=}, {traceback.format_exc()}")
|
||
info = {"message": "Failed to get file info", "status": "failed"}
|
||
|
||
return info
|
||
|
||
@data.post("/upload")
|
||
async def upload_file(
|
||
file: UploadFile = File(...),
|
||
db_id: str | None = Query(None),
|
||
current_user: User = Depends(get_admin_user)
|
||
):
|
||
if not file.filename:
|
||
raise HTTPException(status_code=400, detail="No selected file")
|
||
|
||
# 根据db_id获取上传路径，如果db_id为None则使用默认路径
|
||
if db_id:
|
||
upload_dir = knowledge_base.get_db_upload_path(db_id)
|
||
else:
|
||
upload_dir = os.path.join(config.save_dir, "database", "uploads")
|
||
|
||
basename, ext = os.path.splitext(file.filename)
|
||
filename = f"{basename}_{hashstr(basename, 4, with_salt=True)}{ext}".lower()
|
||
file_path = os.path.join(upload_dir, filename)
|
||
os.makedirs(upload_dir, exist_ok=True)
|
||
|
||
with open(file_path, "wb") as buffer:
|
||
buffer.write(await file.read())
|
||
|
||
return {"message": "File successfully uploaded", "file_path": file_path, "db_id": db_id}
|
||
|
||
@data.get("/graph")
|
||
async def get_graph_info(current_user: User = Depends(get_admin_user)):
|
||
graph_info = graph_base.get_graph_info()
|
||
if graph_info is None:
|
||
raise HTTPException(status_code=400, detail="图数据库获取出错")
|
||
return graph_info
|
||
|
||
@data.post("/graph/index-nodes")
|
||
async def index_nodes(data: dict = Body(default={}), current_user: User = Depends(get_admin_user)):
|
||
if not graph_base.is_running():
|
||
raise HTTPException(status_code=400, detail="图数据库未启动")
|
||
|
||
# 获取参数或使用默认值
|
||
kgdb_name = data.get('kgdb_name', 'neo4j')
|
||
|
||
# 调用GraphDatabase的add_embedding_to_nodes方法
|
||
count = graph_base.add_embedding_to_nodes(kgdb_name=kgdb_name)
|
||
|
||
return {"status": "success", "message": f"已成功为{count}个节点添加嵌入向量", "indexed_count": count}
|
||
|
||
@data.get("/graph/node")
|
||
async def get_graph_node(entity_name: str, current_user: User = Depends(get_admin_user)):
|
||
result = graph_base.query_node(entity_name=entity_name)
|
||
return {"result": graph_base.format_query_result_to_graph(result), "message": "success"}
|
||
|
||
@data.get("/graph/nodes")
|
||
async def get_graph_nodes(kgdb_name: str, num: int, current_user: User = Depends(get_admin_user)):
|
||
|
||
logger.debug(f"Get graph nodes in {kgdb_name} with {num} nodes")
|
||
result = graph_base.get_sample_nodes(kgdb_name, num)
|
||
return {"result": graph_base.format_general_results(result), "message": "success"}
|
||
|
||
@data.post("/graph/add-by-jsonl")
|
||
async def add_graph_entity(file_path: str = Body(...), kgdb_name: str | None = Body(None), current_user: User = Depends(get_admin_user)):
|
||
|
||
if not file_path.endswith('.jsonl'):
|
||
return {"message": "文件格式错误，请上传jsonl文件", "status": "failed"}
|
||
|
||
try:
|
||
await graph_base.jsonl_file_add_entity(file_path, kgdb_name)
|
||
return {"message": "实体添加成功", "status": "success"}
|
||
except Exception as e:
|
||
logger.error(f"添加实体失败: {e}, {traceback.format_exc()}")
|
||
return {"message": f"添加实体失败: {e}", "status": "failed"}
|
||
|
||
@data.post("/update")
|
||
async def update_database_info(
|
||
db_id: str = Body(...),
|
||
name: str = Body(...),
|
||
description: str = Body(...),
|
||
current_user: User = Depends(get_admin_user)
|
||
):
|
||
logger.debug(f"Update database {db_id} info: {name}, {description}")
|
||
try:
|
||
database = knowledge_base.update_database(db_id, name, description)
|
||
return {"message": "更新成功", "database": database}
|
||
except Exception as e:
|
||
logger.error(f"更新数据库失败 {e}, {traceback.format_exc()}")
|
||
raise HTTPException(status_code=400, detail=f"更新数据库失败: {e}")
|
||
|