pengxiao
4bb50ae9c1
Extract images from PDFs using pymupdf, generate descriptions via Qwen3-VL-8B, store in ChromaDB alongside text chunks, and render images in chat answers. Includes image proxy rewrite, force re-process endpoint, and VLM API timeout. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
219 lines
8.3 KiB
Python
219 lines
8.3 KiB
Python
"""
|
|
文档处理服务（LangChain 1.0 + 多模态图片处理）
|
|
"""
|
|
import os
|
|
import asyncio
|
|
from pathlib import Path
|
|
from typing import List, Dict, Any, Optional
|
|
from sqlalchemy.orm import Session
|
|
from langchain_core.documents import Document as LangChainDocument
|
|
|
|
from ..models.document import Document, DocumentChunk
|
|
from ..rag.vector_store import get_vector_store
|
|
from ..rag.document_loaders import DocumentLoaderFactory, PDFImageExtractor
|
|
from ..rag.text_splitters import get_text_splitter
|
|
from ..llm.siliconflow import get_llm_client
|
|
|
|
IMAGES_DIR = Path(__file__).parent.parent.parent.parent / "data" / "images"
|
|
|
|
|
|
class DocumentService:
|
|
"""文档处理服务"""
|
|
|
|
def __init__(self, db: Session):
|
|
self.db = db
|
|
self.vector_store = get_vector_store()
|
|
|
|
async def process_document(self, document_id: int) -> bool:
|
|
"""处理文档（使用LangChain 1.0 + 多模态图片处理）"""
|
|
try:
|
|
document = self.db.query(Document).filter(Document.id == document_id).first()
|
|
if not document:
|
|
return False
|
|
|
|
# 1. 使用LangChain加载文档（文本）
|
|
documents = DocumentLoaderFactory.load_document(
|
|
file_path=document.file_path,
|
|
file_type=document.file_type,
|
|
metadata={
|
|
"document_id": document.id,
|
|
"knowledge_base_id": document.knowledge_base_id,
|
|
"title": document.title,
|
|
"filename": document.filename
|
|
}
|
|
)
|
|
|
|
# 2. 使用中文优化的文本分割器
|
|
text_splitter = get_text_splitter(chunk_size=1000, chunk_overlap=200)
|
|
splits = text_splitter.split_documents(documents)
|
|
|
|
# 3. PDF图片提取和描述（仅PDF文件）
|
|
image_chunks = []
|
|
if document.file_type == ".pdf":
|
|
image_chunks = await self._process_pdf_images(
|
|
file_path=document.file_path,
|
|
document_id=document.id,
|
|
knowledge_base_id=document.knowledge_base_id,
|
|
title=document.title,
|
|
filename=document.filename
|
|
)
|
|
|
|
# 4. 将文本块和图片描述合并添加到向量存储
|
|
all_splits = splits + image_chunks
|
|
success = self.vector_store.add_documents(all_splits)
|
|
|
|
if success:
|
|
document.is_processed = True
|
|
self.db.commit()
|
|
print(f"[DocumentService] 文档 {document.filename} 处理完成: "
|
|
f"{len(splits)} 个文本块, {len(image_chunks)} 个图片描述块")
|
|
return True
|
|
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"处理文档失败: {str(e)}")
|
|
self.db.rollback()
|
|
return False
|
|
|
|
async def _process_pdf_images(
|
|
self,
|
|
file_path: str,
|
|
document_id: int,
|
|
knowledge_base_id: int,
|
|
title: str,
|
|
filename: str
|
|
) -> List[LangChainDocument]:
|
|
"""提取PDF图片并用VLM生成描述"""
|
|
image_chunks = []
|
|
try:
|
|
# 创建图片输出目录
|
|
img_output_dir = IMAGES_DIR / str(document_id)
|
|
|
|
# 提取图片
|
|
images = PDFImageExtractor.extract_images(str(file_path), str(img_output_dir))
|
|
|
|
if not images:
|
|
print(f"[DocumentService] 未发现可提取的图片: {filename}")
|
|
return []
|
|
|
|
print(f"[DocumentService] 提取到 {len(images)} 张图片, 开始VLM描述生成...")
|
|
|
|
# 批量调用VLM生成描述
|
|
llm_client = get_llm_client()
|
|
for idx, img in enumerate(images):
|
|
try:
|
|
description = await llm_client.describe_image(
|
|
img["path"],
|
|
img.get("context_text", "")
|
|
)
|
|
|
|
if description:
|
|
# 相对路径用于URL访问
|
|
rel_path = f"{document_id}/{img['filename']}"
|
|
image_url = f"/images/{rel_path}"
|
|
|
|
# 构建图片描述文本块（URL写入内容，LLM可直接引用）
|
|
chunk_content = (
|
|
f"[图片描述 - 第{img['page']}页]\n"
|
|
f"图片URL: {image_url}\n"
|
|
f"图片内容：{description}"
|
|
)
|
|
|
|
chunk = LangChainDocument(
|
|
page_content=chunk_content,
|
|
metadata={
|
|
"document_id": document_id,
|
|
"knowledge_base_id": knowledge_base_id,
|
|
"title": title,
|
|
"filename": filename,
|
|
"source_type": "image",
|
|
"image_path": str(img["path"]),
|
|
"image_url": image_url,
|
|
"page": img["page"],
|
|
}
|
|
)
|
|
image_chunks.append(chunk)
|
|
print(f"[DocumentService] 图片描述成功 {idx+1}/{len(images)}: {img['filename']}")
|
|
else:
|
|
print(f"[DocumentService] 图片描述为空 {idx+1}/{len(images)}: {img['filename']}")
|
|
|
|
except Exception as e:
|
|
print(f"[DocumentService] 图片处理失败 {img['filename']}: {e}")
|
|
continue
|
|
|
|
except Exception as e:
|
|
print(f"[DocumentService] PDF图片处理失败: {e}")
|
|
|
|
return image_chunks
|
|
|
|
def search_documents(self, query: str, knowledge_base_ids: Optional[List[int]] = None, limit: int = 5) -> List[Dict[str, Any]]:
|
|
"""搜索文档（保留原有接口兼容性）"""
|
|
try:
|
|
# 构建过滤条件
|
|
filter_dict = None
|
|
if knowledge_base_ids:
|
|
filter_dict = {"knowledge_base_id": {"$in": knowledge_base_ids}}
|
|
|
|
# 使用LangChain Chroma进行搜索
|
|
results = self.vector_store.similarity_search_with_score(
|
|
query=query,
|
|
k=limit,
|
|
filter=filter_dict
|
|
)
|
|
|
|
# 格式化结果
|
|
search_results = []
|
|
for doc, distance in results:
|
|
metadata = doc.metadata if hasattr(doc, 'metadata') else {}
|
|
score = self._convert_distance_to_score(distance)
|
|
|
|
search_results.append({
|
|
"content": doc.page_content,
|
|
"metadata": metadata,
|
|
"score": score,
|
|
"distance": distance
|
|
})
|
|
|
|
return search_results
|
|
|
|
except Exception as e:
|
|
print(f"搜索文档失败: {str(e)}")
|
|
return []
|
|
|
|
def _convert_distance_to_score(self, distance: float) -> float:
|
|
"""将ChromaDB的distance转换为0-1范围的相似度分数"""
|
|
import math
|
|
|
|
# 内积距离（负值）
|
|
if distance < 0:
|
|
return (1 + distance) / 2
|
|
|
|
# 大距离使用对数缩放
|
|
if distance > 100:
|
|
return 1 / (1 + math.log(distance))
|
|
|
|
# 标准距离转换
|
|
return 1 / (1 + distance)
|
|
|
|
def get_document_chunks(self, document_id: int) -> List[DocumentChunk]:
|
|
"""获取文档的所有块"""
|
|
return self.db.query(DocumentChunk).filter(
|
|
DocumentChunk.document_id == document_id
|
|
).order_by(DocumentChunk.chunk_index).all()
|
|
|
|
def delete_document_chunks(self, document_id: int) -> bool:
|
|
"""删除文档的所有块（数据库 + 向量存储）"""
|
|
try:
|
|
# 删除向量存储中的文档数据
|
|
self.vector_store.delete_by_document_id(document_id)
|
|
# 删除数据库中的chunk记录
|
|
self.db.query(DocumentChunk).filter(
|
|
DocumentChunk.document_id == document_id
|
|
).delete()
|
|
self.db.commit()
|
|
return True
|
|
except Exception as e:
|
|
print(f"删除文档块失败: {str(e)}")
|
|
self.db.rollback()
|
|
return False |