修正删除文档不能同时删除图片的bug

This commit is contained in:
2026-06-02 19:31:28 +08:00
parent 3250e87b30
commit a6987ff996
8 changed files with 2788 additions and 2728 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
+61 -129
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,41 +84,29 @@ 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
]
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,39 +148,28 @@ 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="文档不存在"
)
# 1. 先删除向量数据和文档块
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="文档不存在")
_check_document_access(document, user)
# 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)
# 2. 删除物理文件
try:
if os.path.exists(document.file_path):
@@ -208,21 +179,18 @@ async def delete_document(
logger.info(f"物理文件不存在: {document.file_path}")
except Exception as e:
logger.error(f"删除物理文件时发生错误: {str(e)}")
# 3. 删除数据库记录
db.delete(document)
db.commit()
logger.info(f"成功删除文档数据库记录: {document.filename}")
return {"message": "文档删除成功"}
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,49 +240,32 @@ 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
file_types[file_type] = file_types.get(file_type, 0) + 1
return DocumentStats(
total_documents=total_documents,
processed_documents=processed_documents,
total_size=total_size,
file_types=file_types
)
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)}")
+8 -1
View File
@@ -390,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:
@@ -422,6 +422,13 @@ async def delete_knowledge_base(
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()
+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)}")
+25 -6
View File
@@ -394,19 +394,38 @@ class KnowledgeBaseService:
Document.source_type == "knowledge_base"
)
).first()
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": "文档删除成功"}
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