Compare commits

..

2 Commits

Author SHA1 Message Date
pengxiao a6987ff996 修正删除文档不能同时删除图片的bug 2026-06-02 19:31:28 +08:00
pengxiao 3250e87b30 fix: knowledge base name collision, orphan dirs, and document download
- Name uniqueness check scoped to user domain (system KBs no longer block user KBs)
- Delete KB now cleans up its directory with shutil.rmtree
- User KB uploads saved to knowledge_bases/{kb_name}/ subdirectory
- New /documents/{id}/download endpoint with auth-aware FileResponse
- Frontend uses fetch+token instead of direct window.open for downloads

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-06-02 11:30:23 +08:00
9 changed files with 2883 additions and 2743 deletions
+5
View File
@@ -43,6 +43,7 @@ dependencies = [
"docx2txt>=0.9",
"pypdf>=6.12.0",
"pymupdf>=1.27.2.3",
"opencv-python-headless>=4.13.0.92",
]
[project.optional-dependencies]
@@ -61,6 +62,10 @@ packages = ["src"]
[tool.uv]
required-version = ">=0.6.15"
[[tool.uv.index]]
url = "https://mirrors.aliyun.com/pypi/simple/"
default = true
[tool.ruff]
line-length = 88
indent-width = 4
+44 -112
View File
@@ -48,6 +48,18 @@ class DocumentStats(BaseModel):
file_types: dict
def _check_document_access(document: Document, user) -> None:
"""检查用户是否有权访问该文档(自己的文档,或系统知识库且为admin)"""
if document.user_id == user.id:
return
if document.knowledge_base and document.knowledge_base.is_system and user.is_superuser:
return
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="文档不存在"
)
@router.post("/upload", response_model=DocumentUploadResponse, deprecated=True)
async def upload_document(
file: UploadFile = File(...),
@@ -72,29 +84,20 @@ async def get_documents(
):
"""获取文档列表"""
try:
# 获取用户ID
from ..models.user import User
user = db.query(User).filter(User.username == current_user).first()
if not user:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="用户不存在"
)
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="用户不存在")
# 获取用户的文档
documents = db.query(Document).filter(
Document.user_id == user.id
).offset(skip).limit(limit).all()
return [
DocumentResponse(
id=doc.id,
filename=doc.filename,
title=doc.title,
file_size=doc.file_size,
file_type=doc.file_type,
is_processed=doc.is_processed,
is_public=doc.is_public,
id=doc.id, filename=doc.filename, title=doc.title,
file_size=doc.file_size, file_type=doc.file_type,
is_processed=doc.is_processed, is_public=doc.is_public,
created_at=doc.created_at.isoformat()
)
for doc in documents
@@ -103,10 +106,7 @@ async def get_documents(
except HTTPException:
raise
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"获取文档列表失败: {str(e)}"
)
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"获取文档列表失败: {str(e)}")
@router.get("/{document_id}", response_model=DocumentResponse)
@@ -117,45 +117,27 @@ async def get_document(
):
"""获取单个文档信息"""
try:
# 获取用户ID
from ..models.user import User
user = db.query(User).filter(User.username == current_user).first()
if not user:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="用户不存在"
)
# 获取文档
document = db.query(Document).filter(
Document.id == document_id,
Document.user_id == user.id
).first()
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="用户不存在")
document = db.query(Document).filter(Document.id == document_id).first()
if not document:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="文档不存在"
)
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="文档不存在")
_check_document_access(document, user)
return DocumentResponse(
id=document.id,
filename=document.filename,
title=document.title,
file_size=document.file_size,
file_type=document.file_type,
is_processed=document.is_processed,
is_public=document.is_public,
id=document.id, filename=document.filename, title=document.title,
file_size=document.file_size, file_type=document.file_type,
is_processed=document.is_processed, is_public=document.is_public,
created_at=document.created_at.isoformat()
)
except HTTPException:
raise
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"获取文档信息失败: {str(e)}"
)
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"获取文档信息失败: {str(e)}")
@router.delete("/{document_id}")
@@ -166,36 +148,25 @@ async def delete_document(
):
"""删除文档"""
try:
# 获取用户ID
from ..models.user import User
user = db.query(User).filter(User.username == current_user).first()
if not user:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="用户不存在"
)
# 获取文档
document = db.query(Document).filter(
Document.id == document_id,
Document.user_id == user.id
).first()
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="用户不存在")
document = db.query(Document).filter(Document.id == document_id).first()
if not document:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="文档不存在"
)
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="文档不存在")
_check_document_access(document, user)
# 1. 删除向量数据和文档块
# 1. 删除向量数据和文档块
try:
logger.info(f"开始删除文档向量数据: {document.filename} (ID: {document.id})")
logger.info(f"开始删除文档向量数据: {document.filename} (ID={document.id})")
document_service = DocumentService(db)
vector_deleted = document_service.delete_document_chunks(document.id)
vector_deleted = document_service.delete_document_chunks(document.id, knowledge_base_id=document.knowledge_base_id)
if vector_deleted:
logger.info(f"成功删除文档向量数据: {document.filename}")
else:
logger.warning(f"警告:删除文档向量数据失败: {document.filename}")
logger.warning(f"删除文档向量数据失败: {document.filename}")
except Exception as e:
logger.error(f"删除向量数据时发生错误: {str(e)}", exc_info=True)
@@ -219,10 +190,7 @@ async def delete_document(
except HTTPException:
raise
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"删除文档失败: {str(e)}"
)
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"删除文档失败: {str(e)}")
@router.post("/{document_id}/process")
@@ -234,54 +202,35 @@ async def process_document(
):
"""处理文档(向量化),force=true 强制重新处理"""
try:
# 获取用户ID
from ..models.user import User
user = db.query(User).filter(User.username == current_user).first()
if not user:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="用户不存在"
)
# 获取文档
document = db.query(Document).filter(
Document.id == document_id,
Document.user_id == user.id
).first()
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="用户不存在")
document = db.query(Document).filter(Document.id == document_id).first()
if not document:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="文档不存在"
)
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="文档不存在")
_check_document_access(document, user)
if document.is_processed and not force:
return {"message": "文档已经处理过了,使用 force=true 强制重新处理"}
# 强制重新处理时,先删除已有的向量数据
if force and document.is_processed:
document_service = DocumentService(db)
document_service.delete_document_chunks(document.id)
document_service.delete_document_chunks(document.id, knowledge_base_id=document.knowledge_base_id)
# 处理文档(直接 await 异步方法)
document_service = DocumentService(db)
success = await document_service.process_document(document.id)
if success:
return {"message": "文档处理成功"}
else:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="文档处理失败"
)
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="文档处理失败")
except HTTPException:
raise
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"处理文档失败: {str(e)}"
)
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"处理文档失败: {str(e)}")
@router.get("/stats/overview", response_model=DocumentStats)
@@ -291,27 +240,19 @@ async def get_document_stats(
):
"""获取文档统计信息"""
try:
# 获取用户ID
from ..models.user import User
user = db.query(User).filter(User.username == current_user).first()
if not user:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="用户不存在"
)
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="用户不存在")
# 统计信息
total_documents = db.query(Document).filter(Document.user_id == user.id).count()
processed_documents = db.query(Document).filter(
Document.user_id == user.id,
Document.is_processed == True
Document.user_id == user.id, Document.is_processed == True
).count()
# 计算总大小
documents = db.query(Document).filter(Document.user_id == user.id).all()
total_size = sum(doc.file_size for doc in documents)
# 文件类型统计
file_types = {}
for doc in documents:
file_type = doc.file_type
@@ -327,13 +268,4 @@ async def get_document_stats(
except HTTPException:
raise
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"获取统计信息失败: {str(e)}"
)
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"获取统计信息失败: {str(e)}")
+73 -8
View File
@@ -3,10 +3,12 @@
"""
import logging
import os
import shutil
import uuid
from pathlib import Path
from typing import List, Optional, Dict, Any
from fastapi import APIRouter, Depends, HTTPException, status, UploadFile, File
from fastapi.responses import FileResponse
from sqlalchemy.orm import Session
from sqlalchemy import or_, and_
from pydantic import BaseModel
@@ -131,9 +133,19 @@ async def create_knowledge_base(
detail="用户不存在"
)
# 检查知识库名称是否已存在
# 管理员创建的知识库标记为系统知识库
is_system = user.is_superuser
# 检查知识库名称是否在同域内已存在(系统知识库与用户知识库互不冲突)
if is_system:
existing_kb = db.query(KnowledgeBase).filter(
KnowledgeBase.name == data.name
KnowledgeBase.name == data.name,
KnowledgeBase.is_system == True
).first()
else:
existing_kb = db.query(KnowledgeBase).filter(
KnowledgeBase.name == data.name,
KnowledgeBase.user_id == user.id
).first()
if existing_kb:
@@ -142,9 +154,6 @@ async def create_knowledge_base(
detail="知识库名称已存在"
)
# 管理员创建的知识库标记为系统知识库
is_system = user.is_superuser
# 创建知识库
knowledge_base = KnowledgeBase(
name=data.name,
@@ -381,7 +390,7 @@ async def delete_knowledge_base(
for document in documents:
try:
# 删除向量数据
vector_deleted = document_service.delete_document_chunks(document.id)
vector_deleted = document_service.delete_document_chunks(document.id, knowledge_base_id=document.knowledge_base_id)
if vector_deleted:
logger.info(f"成功删除文档向量数据: {document.filename}")
else:
@@ -404,6 +413,22 @@ async def delete_knowledge_base(
logger.info(f"文档清理完成: 成功 {success_count} 个,失败 {error_count}")
# 清理知识库目录
if knowledge_base.is_system:
kb_dir = Path(settings.knowledge_base_dir) / knowledge_base.name
else:
kb_dir = Path(settings.upload_dir) / user.username / "knowledge_bases" / knowledge_base.name
if kb_dir.exists():
shutil.rmtree(kb_dir)
logger.info(f"已删除知识库目录: {kb_dir}")
# 清理知识库对应的图片目录
from ..services.document_service import IMAGES_DIR
kb_img_dir = IMAGES_DIR / str(knowledge_base.id)
if kb_img_dir.exists():
shutil.rmtree(kb_img_dir)
logger.info(f"已删除知识库图片目录: {kb_img_dir}")
# 3. 删除知识库(级联删除文档记录)
db.delete(knowledge_base)
db.commit()
@@ -498,8 +523,8 @@ async def upload_document_to_knowledge_base(
save_dir = Path(settings.knowledge_base_dir) / knowledge_base.name
source_type = "knowledge_base"
else:
# 用户知识库:保存到 uploads/{username}/
save_dir = Path(settings.upload_dir) / user.username
# 用户知识库:保存到 uploads/{username}/knowledge_bases/{kb_name}/
save_dir = Path(settings.upload_dir) / user.username / "knowledge_bases" / knowledge_base.name
source_type = "upload"
save_dir.mkdir(parents=True, exist_ok=True)
@@ -740,6 +765,46 @@ async def reindex_document(
)
@router.get("/documents/{document_id}/download")
async def download_document(
document_id: int,
current_user: str = Depends(get_current_user),
db: Session = Depends(get_db)
):
"""下载/查看知识库文档"""
try:
user = db.query(User).filter(User.username == current_user).first()
if not user:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="用户不存在")
document = db.query(Document).filter(Document.id == document_id).first()
if not document:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="文档不存在")
# 权限检查:系统知识库所有人可访问,用户知识库仅限所有者
if not document.knowledge_base:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="知识库不存在")
if not document.knowledge_base.is_system and document.knowledge_base.user_id != user.id:
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="无权访问此文档")
if not os.path.exists(document.file_path):
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="文件不存在")
return FileResponse(
document.file_path,
filename=document.original_filename,
media_type="application/octet-stream"
)
except HTTPException:
raise
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"下载文档失败: {str(e)}"
)
@router.delete("/documents/{document_id}", response_model=ReindexResponse)
async def delete_knowledge_base_document(
document_id: int,
+16 -6
View File
@@ -178,14 +178,24 @@ class SiliconFlowLLM:
import asyncio
import time
# 读取图片并编码为base64
with open(image_path, "rb") as f:
image_data = base64.b64encode(f.read()).decode("utf-8")
# 检测图片格式
# 读取图片并编码为base64,不支持的格式先转为PNG
ext = os.path.splitext(image_path)[1].lower()
mime_map = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".gif": "image/gif", ".webp": "image/webp"}
mime_type = mime_map.get(ext, "image/png")
if ext not in mime_map:
from PIL import Image
import io
img = Image.open(image_path)
if img.mode in ("CMYK", "P"):
img = img.convert("RGB")
buf = io.BytesIO()
img.save(buf, format="PNG")
image_data = base64.b64encode(buf.getvalue()).decode("utf-8")
mime_type = "image/png"
else:
with open(image_path, "rb") as f:
image_data = base64.b64encode(f.read()).decode("utf-8")
mime_type = mime_map[ext]
prompt = IMAGE_DESCRIPTION_PROMPT.format(context_text=context_text[:600])
+53 -6
View File
@@ -2,9 +2,11 @@
LangChain 1.0 文档加载器封装 + PDF图片提取
"""
import logging
import io
from typing import List, Optional, Dict
from pathlib import Path
import fitz # pymupdf
from PIL import Image
from langchain_community.document_loaders import (
PyPDFLoader,
Docx2txtLoader,
@@ -21,18 +23,45 @@ class PDFImageExtractor:
@staticmethod
def extract_images(file_path: str, output_dir: str) -> List[dict]:
"""提取PDF中所有图片,返回图片元数据列表"""
"""提取PDF中所有图片,返回图片元数据列表
对正常页面提取内嵌图片,对瓦片式页面(>50个小图片)渲染整页截图。
"""
Path(output_dir).mkdir(parents=True, exist_ok=True)
images = []
doc = fitz.open(file_path)
TILE_THRESHOLD = 50
vlm_formats = {"png", "jpg", "jpeg", "webp", "gif"}
dpi = 150 # 页面渲染 DPI
for page_num in range(len(doc)):
page = doc[page_num]
# 获取页面文本作为图片上下文
page_text = page.get_text("text")
# 提取页面内嵌图片
image_list = page.get_images(full=True)
if len(image_list) > TILE_THRESHOLD:
# 瓦片式页面:渲染整页为一张完整图片
try:
pix = page.get_pixmap(dpi=dpi)
img_bytes = pix.tobytes("png")
filename = f"page{page_num+1}_full.png"
output_path = Path(output_dir) / filename
output_path.write_bytes(img_bytes)
images.append({
"path": str(output_path),
"filename": filename,
"page": page_num + 1,
"context_text": page_text[:600].strip(),
"size": len(img_bytes),
})
logger.info(f"[ImageExtractor] 瓦片页面渲染为整图 page={page_num+1} size={len(img_bytes)}")
except Exception as e:
logger.warning(f"[ImageExtractor] 页面渲染失败 page={page_num+1}: {e}")
continue
# 正常页面:提取内嵌图片
for img_idx, img_info in enumerate(image_list):
xref = img_info[0]
try:
@@ -40,16 +69,34 @@ class PDFImageExtractor:
image_bytes = base_image["image"]
ext = base_image["ext"]
# 过滤太小的图片(图标装饰元素等)
# 过滤太小的碎片/图标/装饰
if len(image_bytes) < 2048:
continue
w = base_image.get("width", 0)
h = base_image.get("height", 0)
if w < 150 or h < 150:
continue
# 图片周围文本(取该页文字前后各300字作为上下文)
# 图片周围文本
context_start = max(0, page_text.find(
page_text[:len(page_text)//2]) if len(page_text) > 600
else 0)
context_text = page_text[context_start:context_start+600].strip()
# 不支持的格式(jpx/jpeg2000等)转为 PNG
if ext.lower() not in vlm_formats:
try:
img = Image.open(io.BytesIO(image_bytes))
if img.mode in ("CMYK", "P"):
img = img.convert("RGB")
buf = io.BytesIO()
img.save(buf, format="PNG")
image_bytes = buf.getvalue()
ext = "png"
except Exception as e:
logger.warning(f"[ImageExtractor] 格式转换失败 page={page_num+1} img={img_idx}: {e}")
continue
filename = f"page{page_num+1}_img{img_idx+1}.{ext}"
output_path = Path(output_dir) / filename
output_path.write_bytes(image_bytes)
@@ -62,7 +109,7 @@ class PDFImageExtractor:
"size": len(image_bytes),
})
except Exception as e:
logger.warning(f"提取图片失败 page={page_num+1} img={img_idx}: {e}")
logger.warning(f"[ImageExtractor] 提取图片失败 page={page_num+1} img={img_idx}: {e}")
continue
doc.close()
+25 -5
View File
@@ -2,6 +2,7 @@
文档处理服务(LangChain 1.0 + 多模态图片处理)
"""
import os
import shutil
import asyncio
import logging
from pathlib import Path
@@ -91,8 +92,8 @@ class DocumentService:
"""提取PDF图片并用VLM生成描述"""
image_chunks = []
try:
# 创建图片输出目录
img_output_dir = IMAGES_DIR / str(document_id)
# 创建图片输出目录: images/{knowledge_base_id}/{document_id}/
img_output_dir = IMAGES_DIR / str(knowledge_base_id) / str(document_id)
# 提取图片
images = PDFImageExtractor.extract_images(str(file_path), str(img_output_dir))
@@ -115,7 +116,7 @@ class DocumentService:
img.get("context_text", "")
)
if description:
rel_path = f"{document_id}/{img['filename']}"
rel_path = f"{knowledge_base_id}/{document_id}/{img['filename']}"
image_url = f"/images/{rel_path}"
chunk_content = (
f"[图片描述 - 第{img['page']}页]\n"
@@ -193,8 +194,8 @@ class DocumentService:
DocumentChunk.document_id == document_id
).order_by(DocumentChunk.chunk_index).all()
def delete_document_chunks(self, document_id: int) -> bool:
"""删除文档的所有块(数据库 + 向量存储)"""
def delete_document_chunks(self, document_id: int, knowledge_base_id: int = None) -> bool:
"""删除文档的所有块(数据库 + 向量存储 + 图片目录"""
try:
# 删除向量存储中的文档数据
self.vector_store.delete_by_document_id(document_id)
@@ -203,6 +204,25 @@ class DocumentService:
DocumentChunk.document_id == document_id
).delete()
self.db.commit()
# 清理提取的图片目录: images/{kb_id}/{doc_id}/
if knowledge_base_id:
img_dir = IMAGES_DIR / str(knowledge_base_id) / str(document_id)
else:
# 兼容旧数据:尝试查找
for kb_subdir in IMAGES_DIR.iterdir():
if not kb_subdir.is_dir():
continue
candidate = kb_subdir / str(document_id)
if candidate.exists():
img_dir = candidate
break
else:
img_dir = None
if img_dir and img_dir.exists():
shutil.rmtree(img_dir)
logger.info(f"已清理图片目录: {img_dir}")
return True
except Exception as e:
logger.error(f"删除文档块失败: {str(e)}")
+20 -1
View File
@@ -398,15 +398,34 @@ class KnowledgeBaseService:
if not document:
return {"success": False, "message": "文档不存在"}
# 删除向量数据
try:
from ..rag.vector_store import get_vector_store
vector_store = get_vector_store()
vector_store.delete_by_document_id(document_id)
except Exception as e:
logger.warning(f"删除向量数据失败: {e}")
# 删除文档块
self.db.query(DocumentChunk).filter(DocumentChunk.document_id == document_id).delete()
# 删除图片目录
try:
from .document_service import IMAGES_DIR
kb_id = document.knowledge_base_id
img_dir = IMAGES_DIR / str(kb_id) / str(document_id)
if img_dir.exists():
import shutil
shutil.rmtree(img_dir)
logger.info(f"已清理图片目录: {img_dir}")
except Exception as e:
logger.warning(f"清理图片目录失败: {e}")
# 删除文档记录
self.db.delete(document)
self.db.commit()
return {"success": True, "message": "文档删除成功"}
except Exception as e:
self.db.rollback()
return {"success": False, "message": str(e)}
+2595 -2575
View File
File diff suppressed because it is too large Load Diff
+27 -5
View File
@@ -155,11 +155,33 @@ export default function KnowledgeBaseDetailPage() {
}
};
const handleViewDocument = (doc: Document) => {
// Use the file_path from the document to construct the download URL
// Since backend serves static files from /uploads, we can use the file_path directly
const fileUrl = `/api/uploads/${doc.filename}`;
window.open(fileUrl, '_blank');
const handleViewDocument = async (doc: Document) => {
try {
const token = localStorage.getItem("auth_token");
const res = await fetch(`/api/knowledge-bases/documents/${doc.id}/download`, {
headers: { Authorization: `Bearer ${token}` },
});
if (!res.ok) throw new Error("下载失败");
const blob = await res.blob();
// 从 Content-Disposition 提取文件名,或用 doc 信息拼接
const disposition = res.headers.get("Content-Disposition");
let filename = `${doc.title}${doc.file_type}`;
if (disposition) {
const match = disposition.match(/filename\*?=(?:UTF-8'')?(.+)/i);
if (match) filename = decodeURIComponent(match[1].replace(/["']/g, ""));
}
// 触发浏览器下载
const url = URL.createObjectURL(blob);
const a = document.createElement("a");
a.href = url;
a.download = filename;
a.click();
URL.revokeObjectURL(url);
} catch {
window.open(`/api/knowledge-bases/documents/${doc.id}/download`, '_blank');
}
};
const filteredDocuments = knowledgeBase?.documents.filter(doc =>