Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| a6987ff996 | |||
| 3250e87b30 |
@@ -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
@@ -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)}")
|
||||
|
||||
@@ -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,10 +133,20 @@ async def create_knowledge_base(
|
||||
detail="用户不存在"
|
||||
)
|
||||
|
||||
# 检查知识库名称是否已存在
|
||||
existing_kb = db.query(KnowledgeBase).filter(
|
||||
KnowledgeBase.name == data.name
|
||||
).first()
|
||||
# 管理员创建的知识库标记为系统知识库
|
||||
is_system = user.is_superuser
|
||||
|
||||
# 检查知识库名称是否在同域内已存在(系统知识库与用户知识库互不冲突)
|
||||
if is_system:
|
||||
existing_kb = db.query(KnowledgeBase).filter(
|
||||
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:
|
||||
raise HTTPException(
|
||||
@@ -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,
|
||||
|
||||
@@ -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])
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)}")
|
||||
|
||||
@@ -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)}
|
||||
|
||||
Generated
+2595
-2575
File diff suppressed because it is too large
Load Diff
@@ -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 =>
|
||||
|
||||
Reference in New Issue
Block a user