knowledge_base_router.py 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582
  1. import json
  2. import os
  3. import io
  4. import logging
  5. import urllib.parse
  6. import time
  7. import glob
  8. import shutil
  9. import subprocess
  10. from typing import List, Optional
  11. from datetime import datetime
  12. from fastapi import APIRouter, FastAPI, Depends, HTTPException, UploadFile, File, Form
  13. from fastapi.middleware.cors import CORSMiddleware
  14. from fastapi.responses import StreamingResponse
  15. from fastapi.staticfiles import StaticFiles
  16. from fastapi.openapi.docs import (
  17. get_redoc_html,
  18. get_swagger_ui_html,
  19. get_swagger_ui_oauth2_redirect_html,
  20. )
  21. from sqlalchemy import create_engine, and_
  22. from sqlalchemy.orm import Session, sessionmaker
  23. from sqlalchemy.ext.declarative import declarative_base
  24. from agent.models.db.graph import DbUserDataRelation as UserDataRelation
  25. from pydantic import BaseModel, ConfigDict, Field, field_serializer
  26. from agent.libs.auth import SessionValues, verify_session_id
  27. from agent.libs.user_data_relation import UserDataRelationBusiness
  28. from agent.models.web.knowledge_base import Base, KnowledgeBase, KnowledgeFile
  29. from agent.utils import DatabaseUtils, MinioUtils, FileUtils
  30. from config.site import settings
  31. # 响应模型
  32. class ResponseModel(BaseModel):
  33. code: int
  34. message: str
  35. data: Optional[dict | list | bool | None]
  36. class KnowledgeBaseResponse(BaseModel):
  37. model_config = ConfigDict(from_attributes=True)
  38. id: int
  39. name: str
  40. description: Optional[str] = None
  41. tags: Optional[str] = None
  42. creator: Optional[str] = None
  43. user_name: Optional[str] = None # 新增字段
  44. file_count: int = 0
  45. created_at: datetime = Field(default_factory=datetime.utcnow)
  46. updated_at: datetime = Field(default_factory=datetime.utcnow)
  47. @field_serializer('created_at', 'updated_at')
  48. def serialize_datetime(self, dt: datetime) -> str:
  49. return dt.strftime('%Y-%m-%d')
  50. @field_serializer('tags')
  51. def serialize_tags(self, tags: str) -> Optional[List[str]]:
  52. if tags:
  53. return json.loads(tags)
  54. return None
  55. class KnowledgeFileResponse(BaseModel):
  56. model_config = ConfigDict(from_attributes=True)
  57. id: int
  58. knowledge_base_id: int
  59. file_name: str
  60. file_size: float
  61. file_type: str
  62. minio_url: str
  63. status: bool = False
  64. user_name: Optional[str] = None # 用户名
  65. version: Optional[str] = None
  66. author: Optional[str] = None
  67. year: Optional[int] = None
  68. page_count: Optional[int] = None
  69. creator: Optional[str] = None
  70. knowledge_type: Optional[str] = None
  71. created_at: datetime = Field(default_factory=datetime.utcnow)
  72. updated_at: datetime = Field(default_factory=datetime.utcnow)
  73. @field_serializer('created_at', 'updated_at')
  74. def serialize_datetime(self, dt: datetime) -> str:
  75. return dt.strftime('%Y-%m-%d %H:%M')
  76. # 配置日志
  77. logging.basicConfig(
  78. level=logging.INFO,
  79. format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
  80. )
  81. logger = logging.getLogger(__name__)
  82. # 创建数据库引擎
  83. engine = create_engine(settings.DATABASE_URL)
  84. SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
  85. # 创建数据库表
  86. Base.metadata.create_all(bind=engine)
  87. router = APIRouter(tags=["knowledge base interface"])
  88. # logger = logging.getLogger(__name__)
  89. # config = SiteConfig()
  90. # 初始化MinIO工具类
  91. minio_utils = MinioUtils()
  92. # 全局异常处理
  93. # @router.exception_handler(Exception)
  94. # async def global_exception_handler(request, exc):
  95. # logger.error(f"全局异常: {exc}", exc_info=True)
  96. # return {
  97. # "code": 500,
  98. # "message": "服务器内部错误",
  99. # "data": None
  100. # }
  101. # 依赖项:获取数据库会话
  102. def get_db():
  103. db = SessionLocal()
  104. try:
  105. yield db
  106. finally:
  107. db.close()
  108. # 请求模型
  109. class KnowledgeBaseCreate(BaseModel):
  110. name: str
  111. description: Optional[str] = None
  112. tags: Optional[List[str]] = Field(default_factory=list)
  113. class KnowledgeBaseUpdate(BaseModel):
  114. name: str
  115. description: Optional[str] = None
  116. tags: Optional[List[str]] = Field(default_factory=list)
  117. class FileUpdate(BaseModel):
  118. id: int
  119. file_name: Optional[str] = None
  120. version: Optional[str] = None
  121. author: Optional[str] = None
  122. year: Optional[int] = None
  123. page_count: Optional[int] = None
  124. creator: Optional[str] = None
  125. knowledge_type: Optional[str] = None
  126. class BatchFileUpdate(BaseModel):
  127. files: List[FileUpdate]
  128. # 使用utils.py中的FileUtils类进行文件转换
  129. @router.post("/knowledge-base/", response_model=ResponseModel)
  130. def create_knowledge_base(kb: KnowledgeBaseCreate, db: Session = Depends(get_db),
  131. sess:SessionValues = Depends(verify_session_id)):
  132. # 1. 从session获取user_id
  133. user_id = sess.user_id
  134. user_name = sess.username
  135. tags = json.dumps(kb.tags, ensure_ascii=False)
  136. # 2. 创建知识库
  137. kb_data = DatabaseUtils.create_knowledge_base(db, kb.name, user_id, kb.description, tags)
  138. # 3. 创建用户数据关联
  139. relation_business = UserDataRelationBusiness(db)
  140. relation = relation_business.create_relation(
  141. user_id=user_id,
  142. data_category='KnowledgeBase',
  143. data_id=kb_data.id,
  144. user_name=user_name,
  145. role_id=None,
  146. role_name=None
  147. )
  148. return ResponseModel(
  149. code=200,
  150. message="创建成功",
  151. data=KnowledgeBaseResponse.model_validate(kb_data).model_dump()
  152. )
  153. @router.put("/knowledge-base/{kb_id}", response_model=ResponseModel)
  154. def update_knowledge_base(kb_id: int, kb: KnowledgeBaseUpdate, db: Session = Depends(get_db)):
  155. tags = json.dumps(kb.tags, ensure_ascii=False)
  156. kb_data = DatabaseUtils.update_knowledge_base(db, kb_id, kb.name, kb.description, tags)
  157. return ResponseModel(
  158. code=200,
  159. message="更新成功",
  160. data=KnowledgeBaseResponse.model_validate(kb_data).model_dump()
  161. )
  162. @router.delete("/knowledge-base/{kb_id}", response_model=ResponseModel)
  163. def delete_knowledge_base(kb_id: int, db: Session = Depends(get_db)):
  164. result = DatabaseUtils.delete_knowledge_base(db, kb_id)
  165. return ResponseModel(
  166. code=200,
  167. message="删除成功",
  168. data=result
  169. )
  170. @router.get("/knowledge-base/{kb_id}", response_model=ResponseModel)
  171. def get_knowledge_base(kb_id: int, db: Session = Depends(get_db)):
  172. kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id, KnowledgeBase.is_deleted == 0).first()
  173. if not kb:
  174. raise HTTPException(status_code=404, detail="知识库不存在")
  175. kb_data = KnowledgeBaseResponse.model_validate(kb).model_dump()
  176. return ResponseModel(
  177. code=200,
  178. message="查询成功",
  179. data=kb_data
  180. )
  181. @router.get("/knowledge-base/", response_model=ResponseModel)
  182. def list_knowledge_bases(pageNo: int = 1, pageSize: int = 10, name: Optional[str] = None,
  183. db: Session = Depends(get_db)):
  184. if pageNo < 1:
  185. raise HTTPException(status_code=400, detail="页码必须大于等于1")
  186. if pageSize < 1:
  187. raise HTTPException(status_code=400, detail="每页条数必须大于等于1")
  188. skip = (pageNo - 1) * pageSize
  189. kb_list, total = DatabaseUtils.get_knowledge_bases(db, skip, pageSize, name)
  190. return ResponseModel(
  191. code=200,
  192. message="查询成功",
  193. data={
  194. "list": [KnowledgeBaseResponse.model_validate(kb).model_dump() for kb in kb_list],
  195. "total": total
  196. }
  197. )
  198. @router.get("/knowledge-base/name/{name}", response_model=ResponseModel)
  199. def get_knowledge_base_by_name(name: str, db: Session = Depends(get_db)):
  200. kb = DatabaseUtils.get_knowledge_base_by_name(db, name)
  201. if not kb:
  202. raise HTTPException(status_code=404, detail="知识库不存在")
  203. return ResponseModel(
  204. code=200,
  205. message="查询成功",
  206. data=KnowledgeBaseResponse.model_validate(kb).model_dump()
  207. )
  208. @router.post("/knowledge-base/{kb_id}/files/", response_model=ResponseModel)
  209. async def upload_files(
  210. kb_id: int,
  211. files: List[UploadFile] = File(...),
  212. db: Session = Depends(get_db),
  213. sess: SessionValues = Depends(verify_session_id) # 添加session依赖
  214. ):
  215. # 验证知识库是否存在
  216. kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id, KnowledgeBase.is_deleted == 0).first()
  217. if not kb:
  218. raise HTTPException(status_code=404, detail="知识库不存在")
  219. # 获取当前用户信息
  220. user_id = sess.user_id
  221. user_name = sess.username
  222. # 验证文件数量
  223. if len(files) > settings.MAX_FILE_COUNT:
  224. raise HTTPException(status_code=400, detail=f"单次上传文件数量不能超过{settings.MAX_FILE_COUNT}个")
  225. # 导入所需模块
  226. import tempfile
  227. uploaded_files = []
  228. for file in files:
  229. # 获取文件扩展名
  230. file_ext = os.path.splitext(file.filename)[1].lower().lstrip('.')
  231. original_filename = file.filename
  232. converted_content = None
  233. # 读取文件内容
  234. content = await file.read()
  235. # 处理需要转换的文件格式
  236. if file_ext in ["doc", "ppt"]:
  237. # 创建临时目录用于文件转换
  238. with tempfile.TemporaryDirectory() as temp_dir:
  239. # 创建临时文件
  240. temp_input_path = os.path.join(temp_dir, original_filename)
  241. with open(temp_input_path, "wb") as temp_file:
  242. temp_file.write(content)
  243. # 确定目标格式
  244. target_format = "docx" if file_ext == "doc" else "pptx"
  245. # 使用FileUtils中的文件转换方法
  246. converted_file_path = FileUtils.convert_office_file(temp_input_path, temp_dir, target_format)
  247. if converted_file_path and os.path.exists(converted_file_path):
  248. # 读取转换后的文件内容
  249. with open(converted_file_path, "rb") as converted_file:
  250. converted_content = converted_file.read()
  251. # 更新文件名和扩展名
  252. file_ext = target_format
  253. file.filename = os.path.splitext(original_filename)[0] + f".{target_format}"
  254. else:
  255. # 转换失败,使用原始文件
  256. raise HTTPException(status_code=500, detail=f"文件格式转换失败: {original_filename}")
  257. # 检查文件格式是否支持
  258. if file_ext not in settings.ALLOWED_EXTENSIONS:
  259. raise HTTPException(status_code=400, detail=f"不支持的文件格式:{file_ext}")
  260. # 使用转换后的内容或原始内容
  261. file_content = converted_content if converted_content else content
  262. file_size = len(file_content) / (1024 * 1024) # 转换为MB
  263. # 验证文件大小
  264. max_size = settings.ALLOWED_EXTENSIONS[file_ext]["max_size"]
  265. if file_size > max_size:
  266. raise HTTPException(status_code=400, detail=f"{original_filename}超过最大允许大小{max_size}MB")
  267. # 上传到MinIO
  268. minio_url = minio_utils.upload_file(file_content, file.filename, file.content_type)
  269. # 从文件名识别知识类型
  270. knowledge_type = None
  271. if '指南' in file.filename:
  272. knowledge_type = '指南'
  273. elif '教材' in file.filename:
  274. knowledge_type = '教材'
  275. # 创建文件记录
  276. db_file = KnowledgeFile(
  277. knowledge_base_id=kb_id,
  278. file_name=file.filename,
  279. file_size=file_size,
  280. file_type=file_ext,
  281. minio_url=minio_url,
  282. creator=user_id,
  283. knowledge_type=knowledge_type
  284. )
  285. db.add(db_file)
  286. uploaded_files.append(db_file)
  287. # 创建用户数据关联
  288. relation_business = UserDataRelationBusiness(db)
  289. relation = relation_business.create_relation(
  290. user_id=user_id,
  291. data_category='KnowledgeFile',
  292. data_id=db_file.id,
  293. user_name=user_name,
  294. role_id=None,
  295. role_name=None
  296. )
  297. # 更新知识库文件计数
  298. DatabaseUtils.increment_file_count(db, kb_id)
  299. db.commit()
  300. return ResponseModel(
  301. code=200,
  302. message="上传成功",
  303. data=[KnowledgeFileResponse.model_validate(file).model_dump() for file in uploaded_files]
  304. )
  305. @router.get("/knowledge-base/{kb_id}/files/", response_model=ResponseModel)
  306. def list_files(kb_id: int, pageNo: int = 1, pageSize: int = 10, file_name: Optional[str] = None,
  307. db: Session = Depends(get_db)):
  308. if pageNo < 1:
  309. raise HTTPException(status_code=400, detail="页码必须大于等于1")
  310. if pageSize < 1:
  311. raise HTTPException(status_code=400, detail="每页条数必须大于等于1")
  312. skip = (pageNo - 1) * pageSize
  313. query = db.query(KnowledgeFile,UserDataRelation.user_name).\
  314. outerjoin(UserDataRelation,
  315. and_(
  316. UserDataRelation.data_id == KnowledgeFile.id,
  317. UserDataRelation.data_category == 'KnowledgeFile'
  318. )).\
  319. filter(
  320. KnowledgeFile.knowledge_base_id == kb_id,
  321. KnowledgeFile.is_deleted == 0
  322. ).order_by(KnowledgeFile.status.desc())
  323. if file_name:
  324. query = query.filter(KnowledgeFile.file_name.ilike(f"%{file_name}%"))
  325. total = query.count()
  326. files = query.offset(skip).limit(pageSize).all()
  327. return ResponseModel(
  328. code=200,
  329. message="查询成功",
  330. data={
  331. "list": [KnowledgeFileResponse.model_validate(
  332. {
  333. **file[0].__dict__,
  334. "user_name": file[1]
  335. }
  336. ).model_dump() for file in files],
  337. "total": total
  338. }
  339. )
  340. @router.get("/knowledge-base/{kb_id}/files/search/", response_model=ResponseModel)
  341. def search_files(kb_id: int, file_name: str, db: Session = Depends(get_db)):
  342. files = db.query(KnowledgeFile, UserDataRelation.user_name).\
  343. join(UserDataRelation,
  344. and_(
  345. UserDataRelation.data_id == KnowledgeFile.id,
  346. UserDataRelation.data_category == 'KnowledgeFile'
  347. )).\
  348. filter(
  349. KnowledgeFile.knowledge_base_id == kb_id,
  350. KnowledgeFile.file_name.ilike(f"%{file_name}%"),
  351. KnowledgeFile.is_deleted == 0
  352. ).all()
  353. result = []
  354. for file, user_name in files:
  355. file_response = KnowledgeFileResponse.model_validate(file)
  356. file_response.user_name = user_name
  357. result.append(file_response.model_dump())
  358. return ResponseModel(
  359. code=200,
  360. message="查询成功",
  361. data=result
  362. )
  363. @router.get("/files/{file_id}/download")
  364. def download_file(file_id: int, db: Session = Depends(get_db)):
  365. # 获取文件信息
  366. file = db.query(KnowledgeFile).filter(
  367. KnowledgeFile.id == file_id,
  368. KnowledgeFile.is_deleted == 0
  369. ).first()
  370. if not file:
  371. raise HTTPException(status_code=404, detail="文件不存在")
  372. # 从MinIO下载文件
  373. object_name = file.minio_url.split("/")[-1]
  374. file_content = minio_utils.download_file(object_name)
  375. # 创建文件流
  376. file_stream = io.BytesIO(file_content)
  377. # 对文件名进行URL编码
  378. encoded_filename = urllib.parse.quote(file.file_name)
  379. return StreamingResponse(
  380. file_stream,
  381. media_type="application/octet-stream",
  382. headers={
  383. "Content-Disposition": f"attachment; filename*=UTF-8''{encoded_filename}"
  384. }
  385. )
  386. @router.delete("/files/{file_id}", response_model=dict)
  387. def delete_file(file_id: int, db: Session = Depends(get_db)):
  388. # 获取文件信息
  389. file = db.query(KnowledgeFile).filter(
  390. KnowledgeFile.id == file_id,
  391. KnowledgeFile.is_deleted == 0
  392. ).first()
  393. if not file:
  394. raise HTTPException(status_code=404, detail="文件不存在")
  395. # 从MinIO删除文件
  396. object_name = file.minio_url.split("/")[-1]
  397. minio_utils.delete_file(object_name)
  398. # 标记文件为已删除
  399. file.is_deleted = 1
  400. file.updated_at = datetime.utcnow()
  401. # 更新知识库文件计数
  402. DatabaseUtils.decrement_file_count(db, file.knowledge_base_id)
  403. db.commit()
  404. return {
  405. "code": 200,
  406. "message": "删除成功",
  407. "data": True
  408. }
  409. @router.get("/files/{file_id}/changeStatus", response_model=dict)
  410. def change_file_status(file_id: int, status: bool, db: Session = Depends(get_db)):
  411. # 获取文件信息
  412. file = db.query(KnowledgeFile).filter(
  413. KnowledgeFile.id == file_id,
  414. KnowledgeFile.is_deleted == 0
  415. ).first()
  416. if not file:
  417. raise HTTPException(status_code=404, detail="文件不存在")
  418. # 标记文件停用状态
  419. if status:
  420. file.status = 1
  421. else:
  422. file.status = 0
  423. file.updated_at = datetime.utcnow()
  424. db.commit()
  425. return {
  426. "code": 200,
  427. "message": "修改成功",
  428. "data": True
  429. }
  430. @router.put("/files/batch-update", response_model=ResponseModel)
  431. def batch_update_files(update_data: BatchFileUpdate, db: Session = Depends(get_db)):
  432. updated_files = []
  433. for file_update in update_data.files:
  434. # 获取文件信息
  435. file = db.query(KnowledgeFile).filter(
  436. KnowledgeFile.id == file_update.id,
  437. KnowledgeFile.is_deleted == 0
  438. ).first()
  439. if not file:
  440. raise HTTPException(status_code=404, detail=f"文件ID {file_update.id} 不存在")
  441. # 如果需要更新文件名,同时更新MinIO中的文件
  442. if file_update.file_name and file_update.file_name != file.file_name:
  443. old_object_name = file.minio_url.split("/")[-1]
  444. new_object_name = f"{datetime.utcnow().strftime('%Y%m%d%H%M%S')}_{file_update.file_name}"
  445. # 从MinIO下载文件
  446. file_content = minio_utils.download_file(old_object_name)
  447. # 上传到MinIO新的位置
  448. new_minio_url = minio_utils.upload_file(
  449. file_content,
  450. file_update.file_name,
  451. file.file_type
  452. )
  453. # 删除旧文件
  454. minio_utils.delete_file(old_object_name)
  455. # 更新数据库中的文件名和MinIO URL
  456. file.file_name = file_update.file_name
  457. file.minio_url = new_minio_url
  458. # 更新其他字段
  459. if file_update.version is not None:
  460. file.version = file_update.version
  461. if file_update.author is not None:
  462. file.author = file_update.author
  463. if file_update.year is not None:
  464. file.year = file_update.year
  465. if file_update.page_count is not None:
  466. file.page_count = file_update.page_count
  467. if file_update.creator is not None:
  468. file.creator = file_update.creator
  469. if file_update.knowledge_type is not None:
  470. file.knowledge_type = file_update.knowledge_type
  471. file.updated_at = datetime.utcnow()
  472. updated_files.append(file)
  473. db.commit()
  474. return ResponseModel(
  475. code=200,
  476. message="更新成功",
  477. data=[KnowledgeFileResponse.model_validate(file).model_dump() for file in updated_files]
  478. )
  479. knowledge_base_router = router