770 lines
34 KiB
Python
770 lines
34 KiB
Python
"""文档处理服务，负责文档的分段、向量化和索引"""
|
||
|
||
import os
|
||
from typing import List, Dict, Any, Optional
|
||
from pathlib import Path
|
||
from fastapi import HTTPException
|
||
from requests import Session
|
||
from sqlalchemy import text
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
from langchain_community.document_loaders import (
|
||
TextLoader,
|
||
PyPDFLoader,
|
||
Docx2txtLoader,
|
||
UnstructuredMarkdownLoader
|
||
)
|
||
import pdfplumber
|
||
from langchain_core.documents import Document
|
||
from langchain_postgres import PGVector
|
||
from typing import List
|
||
|
||
|
||
from ..core.config import BaseSettings, get_settings
|
||
from ..models.knowledge_base import Document as DocumentModel
|
||
from ..db.database import get_session
|
||
from loguru import logger
|
||
|
||
settings = get_settings()
|
||
class DocumentProcessor:
|
||
"""文档处理器，负责文档的加载、分段和向量化"""
|
||
|
||
def __init__(self):
|
||
# 初始化语义分割器配置
|
||
self.embeddings = None
|
||
self.semantic_splitter_enabled = settings.file.semantic_splitter_enabled
|
||
self.text_splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=settings.file.chunk_size,
|
||
chunk_overlap=settings.file.chunk_overlap,
|
||
length_function=len,
|
||
separators=["\n\n", "\n", " ", ""]
|
||
)
|
||
|
||
async def initialize(self, session: Session = None):
|
||
# 初始化嵌入模型 - 根据配置选择提供商
|
||
await self._init_embeddings(session)
|
||
|
||
# 初始化连接池（仅对PGVector）
|
||
self.pgvector_pool = None
|
||
|
||
# PostgreSQL pgvector连接配置
|
||
# if settings.vector_db.type == "pgvector":
|
||
# # 新版本PGVector使用psycopg3连接字符串
|
||
# # 对密码进行URL编码以处理特殊字符（如@符号）
|
||
# encoded_password = quote(settings.vector_db.pgvector_password, safe="")
|
||
# self.connection_string = (
|
||
# f"postgresql+psycopg://{settings.vector_db.pgvector_user}:"
|
||
# f"{encoded_password}@"
|
||
# f"{settings.vector_db.pgvector_host}:"
|
||
# f"{settings.vector_db.pgvector_port}/"
|
||
# f"{settings.vector_db.pgvector_database}"
|
||
# )
|
||
# # 初始化连接池
|
||
# self.pgvector_pool = PGVectorConnectionPool()
|
||
# logger.info("新版本PGVector使用psycopg3连接字符串: %s", self.connection_string)
|
||
# else:
|
||
# 向量数据库存储路径（Chroma兼容）
|
||
vector_db_path = settings.vector_db.persist_directory
|
||
if not os.path.isabs(vector_db_path):
|
||
# 如果是相对路径，则基于项目根目录计算绝对路径
|
||
# 项目根目录是backend的父目录
|
||
backend_dir = Path(__file__).parent.parent.parent
|
||
vector_db_path = str(backend_dir / vector_db_path)
|
||
self.vector_db_path = vector_db_path
|
||
session.desc = f"初始化向量数据库 - 路径 = {self.vector_db_path}"
|
||
|
||
async def _init_embeddings(self, session: Optional[Any] = None):
|
||
"""初始化嵌入模型。"""
|
||
try:
|
||
if not self.embeddings:
|
||
# 使用llm_config_service获取嵌入配置
|
||
from .llm_config_service import LLMConfigService
|
||
llm_config_service = LLMConfigService()
|
||
|
||
# 获取嵌入配置
|
||
config = None
|
||
if session:
|
||
config = await llm_config_service.get_default_embedding_config(session)
|
||
if config and session:
|
||
session.desc = f"获取默认嵌入模型配置: {config}"
|
||
# # 转换配置格式
|
||
# config = {
|
||
# "provider": config.provider,
|
||
# "api_key": config.api_key,
|
||
# "model": config.model_name
|
||
# }
|
||
|
||
# 如果未找到配置，使用默认配置
|
||
if not config:
|
||
if session:
|
||
session.desc = f"ERROR: 未找到嵌入模型配置"
|
||
raise HTTPException(status_code=400, detail="未找到嵌入模型配置")
|
||
if session:
|
||
session.desc = f"获取嵌入模型配置 > 结果：{config}"
|
||
|
||
# 根据配置创建嵌入模型
|
||
if config.provider == "openai":
|
||
from langchain_openai import OpenAIEmbeddings
|
||
self.embeddings = OpenAIEmbeddings(
|
||
model=config.model_name or "text-embedding-3-small",
|
||
api_key=config.api_key
|
||
)
|
||
if session:
|
||
session.desc = f"创建嵌入模型 - OpenAIEmbeddings(model={config.model_name or 'text-embedding-3-small'})"
|
||
elif config.provider == "zhipu":
|
||
from .zhipu_embeddings import ZhipuOpenAIEmbeddings
|
||
self.embeddings = ZhipuOpenAIEmbeddings(
|
||
api_key=config.api_key,
|
||
base_url=config.base_url or "https://open.bigmodel.cn/api/paas/v4",
|
||
model=config.model_name or "embedding-3",
|
||
dimensions=settings.vector_db.embedding_dimension
|
||
)
|
||
if session:
|
||
session.desc = f"创建嵌入模型 - ZhipuOpenAIEmbeddings(model={config.model_name or 'embedding-3'}, base_url={config.base_url})"
|
||
elif config.provider == "ollama":
|
||
from langchain_ollama import OllamaEmbeddings
|
||
self.embeddings = OllamaEmbeddings(
|
||
model=config.model_name,
|
||
base_url=config.base_url
|
||
)
|
||
if session:
|
||
session.desc = f"创建嵌入模型 - OllamaEmbeddings({self.embeddings.base_url} - {self.embeddings.model})"
|
||
elif config.provider == "local":
|
||
from langchain_huggingface import HuggingFaceEmbeddings
|
||
self.embeddings = HuggingFaceEmbeddings(
|
||
model_name=config.model_name or "sentence-transformers/all-MiniLM-L6-v2"
|
||
)
|
||
if session:
|
||
session.desc = f"创建嵌入模型 - HuggingFaceEmbeddings(model={config.model_name or 'sentence-transformers/all-MiniLM-L6-v2'})"
|
||
else:
|
||
# 默认使用OpenAI
|
||
from langchain_openai import OpenAIEmbeddings
|
||
self.embeddings = OpenAIEmbeddings(
|
||
model=config.model_name or "text-embedding-3-small",
|
||
api_key=config.api_key
|
||
)
|
||
if session:
|
||
session.desc = f"ERROR: 未支持的嵌入提供者: {config.provider}，已使用默认的 OpenAIEmbeddings - 可能不正确或无效"
|
||
|
||
return self.embeddings
|
||
except Exception as e:
|
||
logger.error(f"初始化嵌入模型时出错: {e}")
|
||
raise
|
||
|
||
def load_document(self, session: Session, file_path: str) -> List[Document]:
|
||
"""根据文件类型加载文档"""
|
||
file_extension = Path(file_path).suffix.lower()
|
||
try:
|
||
if file_extension == '.txt':
|
||
session.desc = f"加载文档 - 文件路径: {file_path} - 类型: txt"
|
||
loader = TextLoader(file_path, encoding='utf-8')
|
||
documents = loader.load()
|
||
elif file_extension == '.pdf':
|
||
# 使用pdfplumber处理PDF文件，更稳定
|
||
session.desc = f"加载文档 - 文件路径: {file_path} - 类型: pdf"
|
||
from langchain_community.document_loaders import PyPDFLoader
|
||
loader = PyPDFLoader(file_path)
|
||
documents = loader.load()
|
||
# documents = self._load_pdf_with_pdfplumber(file_path)
|
||
elif file_extension == '.docx':
|
||
session.desc = f"加载文档 - 文件路径: {file_path} - 类型: docx"
|
||
loader = Docx2txtLoader(file_path)
|
||
documents = loader.load()
|
||
elif file_extension == '.md':
|
||
session.desc = f"加载文档 - 文件路径: {file_path} - 类型: md"
|
||
loader = UnstructuredMarkdownLoader(file_path)
|
||
documents = loader.load()
|
||
else:
|
||
raise ValueError(f"不支持的文件类型: {file_extension}")
|
||
|
||
session.desc = f"已载文档: {file_path}, 页数: {len(documents)}"
|
||
# if len(documents) > 0:
|
||
# session.desc = f"文档内容示例: {type(documents[0])} - {documents[0]}"
|
||
return documents
|
||
|
||
except Exception as e:
|
||
session.desc = f"ERROR: 加载文档失败 {file_path}: {str(e)}"
|
||
raise e
|
||
|
||
def _load_pdf_with_pdfplumber(self, file_path: str) -> List[Document]:
|
||
"""使用pdfplumber加载PDF文档"""
|
||
documents = []
|
||
try:
|
||
with pdfplumber.open(file_path) as pdf:
|
||
for page_num, page in enumerate(pdf.pages):
|
||
text = page.extract_text()
|
||
if text and text.strip(): # 只处理有文本内容的页面
|
||
doc = Document(
|
||
page_content=text,
|
||
metadata={
|
||
"source": file_path,
|
||
"page": page_num + 1
|
||
}
|
||
)
|
||
documents.append(doc)
|
||
return documents
|
||
except Exception as e:
|
||
logger.error(f"使用pdfplumber加载PDF失败 {file_path}: {str(e)}")
|
||
# 如果pdfplumber失败，回退到PyPDFLoader
|
||
try:
|
||
loader = PyPDFLoader(file_path)
|
||
return loader.load()
|
||
except Exception as fallback_e:
|
||
logger.error(f"PyPDFLoader回退也失败 {file_path}: {str(fallback_e)}")
|
||
raise fallback_e
|
||
|
||
def _merge_documents(self, documents: List[Document]) -> Document:
|
||
"""将多个文档合并成一个文档"""
|
||
merged_text = ""
|
||
merged_metadata = {}
|
||
|
||
for doc in documents:
|
||
if merged_text:
|
||
merged_text += "\n\n"
|
||
merged_text += doc.page_content
|
||
# 合并元数据
|
||
merged_metadata.update(doc.metadata)
|
||
|
||
return Document(page_content=merged_text, metadata=merged_metadata)
|
||
|
||
def _split_by_semantic_points(self, text: str, split_points: List[str]) -> List[str]:
|
||
"""根据语义分割点切分文本"""
|
||
chunks = []
|
||
current_pos = 0
|
||
|
||
# 按顺序查找每个分割点并切分文本
|
||
for point in split_points:
|
||
pos = text.find(point, current_pos)
|
||
if pos != -1:
|
||
# 添加当前位置到分割点位置的文本块
|
||
if pos > current_pos:
|
||
chunk = text[current_pos:pos].strip()
|
||
if chunk:
|
||
chunks.append(chunk)
|
||
current_pos = pos
|
||
|
||
# 添加最后一个文本块
|
||
if current_pos < len(text):
|
||
chunk = text[current_pos:].strip()
|
||
if chunk:
|
||
chunks.append(chunk)
|
||
|
||
return chunks
|
||
|
||
async def split_documents(self, session: Session, documents: List[Document]) -> List[Document]:
|
||
"""将文档分割成小块（含短段落合并和超长强制分割功能）"""
|
||
try:
|
||
chunks = self.text_splitter.split_documents(documents)
|
||
|
||
session.desc = f"文档分割完成，共生成 {len(chunks)} 个文档块"
|
||
if len(chunks) > 0:
|
||
session.desc = f"文档块内容示例: {type(chunks[0])} - {chunks[0]}"
|
||
return chunks
|
||
|
||
except Exception as e:
|
||
session.desc = f"ERROR: 文档分割失败: {str(e)}"
|
||
raise e
|
||
|
||
def _force_split_long_chunk(self, chunk: str) -> List[str]:
|
||
"""强制分割超长段落（超过1000字符）"""
|
||
max_length = 1000
|
||
chunks = []
|
||
|
||
# 先尝试按换行符分割
|
||
if '\n' in chunk:
|
||
lines = chunk.split('\n')
|
||
current_chunk = ""
|
||
for line in lines:
|
||
if len(current_chunk) + len(line) + 1 > max_length:
|
||
if current_chunk:
|
||
chunks.append(current_chunk)
|
||
current_chunk = line
|
||
else:
|
||
chunks.append(line[:max_length])
|
||
current_chunk = line[max_length:]
|
||
else:
|
||
if current_chunk:
|
||
current_chunk += "\n" + line
|
||
else:
|
||
current_chunk = line
|
||
if current_chunk:
|
||
chunks.append(current_chunk)
|
||
else:
|
||
# 没有换行符则直接按长度分割
|
||
chunks = [chunk[i:i + max_length] for i in range(0, len(chunk), max_length)]
|
||
|
||
return chunks
|
||
|
||
def create_vector_store(self, knowledge_base_id: int, documents: List[Document], document_id: int = None) -> str:
|
||
"""为知识库创建向量存储"""
|
||
try:
|
||
# if settings.vector_db.type == "pgvector":
|
||
# # 添加元数据
|
||
# for i, doc in enumerate(documents):
|
||
# doc.metadata.update({
|
||
# "knowledge_base_id": knowledge_base_id,
|
||
# "document_id": str(document_id) if document_id else "unknown",
|
||
# "chunk_id": f"{knowledge_base_id}_{document_id}_{i}",
|
||
# "chunk_index": i
|
||
# })
|
||
|
||
# # 创建PostgreSQL pgvector存储
|
||
# collection_name = f"{settings.vector_db.pgvector_table_name}_kb_{knowledge_base_id}"
|
||
|
||
# # 创建新版本PGVector实例
|
||
# vector_store = PGVector(
|
||
# connection=self.connection_string,
|
||
# embeddings=self.embeddings,
|
||
# collection_name=collection_name,
|
||
# use_jsonb=True # 使用JSONB存储元数据
|
||
# )
|
||
|
||
# # 手动添加文档
|
||
# vector_store.add_documents(documents)
|
||
|
||
# logger.info(f"PostgreSQL pgvector存储创建成功: {collection_name}")
|
||
# return collection_name
|
||
# else:
|
||
# Chroma兼容模式
|
||
from langchain_chroma import Chroma
|
||
kb_vector_path = os.path.join(self.vector_db_path, f"kb_{knowledge_base_id}")
|
||
|
||
# 添加元数据
|
||
for i, doc in enumerate(documents):
|
||
doc.metadata.update({
|
||
"knowledge_base_id": knowledge_base_id,
|
||
"document_id": str(document_id) if document_id else "unknown",
|
||
"chunk_id": f"{knowledge_base_id}_{document_id}_{i}",
|
||
"chunk_index": i
|
||
})
|
||
|
||
# 创建向量存储
|
||
vector_store = Chroma.from_documents(
|
||
documents=documents,
|
||
embedding=self.embeddings,
|
||
persist_directory=kb_vector_path
|
||
)
|
||
|
||
logger.info(f"向量存储创建成功: {kb_vector_path}")
|
||
return kb_vector_path
|
||
|
||
except Exception as e:
|
||
logger.error(f"创建向量存储失败: {str(e)}")
|
||
raise
|
||
|
||
def add_documents_to_vector_store(self, session: Session, knowledge_base_id: int, documents: List[Document], document_id: int = None) -> None:
|
||
"""向现有向量存储添加文档"""
|
||
if len(documents) == 0:
|
||
session.desc = f"WARNING: 文档列表为空，不执行添加操作"
|
||
return
|
||
from langchain_chroma import Chroma
|
||
|
||
kb_vector_path = os.path.join(self.vector_db_path, f"kb_{knowledge_base_id}")
|
||
session.desc = f"添加文档到向量存储: {kb_vector_path} - documents number: {len(documents)}"
|
||
# 检查向量存储是否存在
|
||
if not os.path.exists(kb_vector_path):
|
||
# 如果不存在，创建新的向量存储
|
||
session.desc = f"WARNING: 向量存储不存在，创建新的向量存储"
|
||
self.create_vector_store(knowledge_base_id, documents, document_id)
|
||
return
|
||
session.desc = f"添加文档到向量存储: exists"
|
||
# 添加元数据
|
||
for i, doc in enumerate(documents):
|
||
doc.metadata.update({
|
||
"knowledge_base_id": knowledge_base_id,
|
||
"document_id": str(document_id) if document_id else "unknown",
|
||
"chunk_id": f"{knowledge_base_id}_{document_id}_{i}",
|
||
"chunk_index": i
|
||
})
|
||
|
||
session.desc = f"添加文档到向量存储: enumerate"
|
||
# 加载现有向量存储
|
||
vector_store = Chroma(
|
||
persist_directory=kb_vector_path,
|
||
embedding_function=self.embeddings
|
||
)
|
||
|
||
session.desc = f"添加文档到向量存储: Chroma"
|
||
# 添加新文档
|
||
ids = vector_store.add_documents(documents)
|
||
session.desc = f"文档已添加到向量存储: {kb_vector_path} -> {len(ids)} IDS - \n{ids}"
|
||
|
||
async def process_document(self, session: Session, document_id: int, file_path: str, knowledge_base_id: int) -> Dict[str, Any]:
|
||
"""处理单个文档：加载、分段、向量化"""
|
||
try:
|
||
session.desc = f"处理文档 ID: {document_id} 文件路径: {file_path}"
|
||
|
||
# 1. 加载文档
|
||
documents = self.load_document(session, file_path)
|
||
|
||
# 2. 分割文档
|
||
chunks = await self.split_documents(session, documents)
|
||
|
||
# 3. 添加到向量存储
|
||
self.add_documents_to_vector_store(session, knowledge_base_id, chunks, document_id)
|
||
|
||
# 4. 更新文档状态
|
||
# Python 3.9 兼容：使用 async for 替代 anext
|
||
async for db_session in get_session():
|
||
try:
|
||
from sqlalchemy import select
|
||
document = await db_session.scalar(select(DocumentModel).where(DocumentModel.id == document_id))
|
||
|
||
if document:
|
||
document.is_processed = True
|
||
document.chunk_count = len(chunks)
|
||
await db_session.commit()
|
||
finally:
|
||
await db_session.close()
|
||
break # 只取第一个 session
|
||
|
||
result = {
|
||
"document_id": document_id,
|
||
"status": "success",
|
||
"chunks_count": len(chunks),
|
||
"message": "文档处理完成"
|
||
}
|
||
|
||
|
||
session.desc = f"文档处理完成: {result}"
|
||
return result
|
||
|
||
except Exception as e:
|
||
session.desc = f"ERROR: 文档处理失败 ID: {document_id}: {str(e)}"
|
||
|
||
# 更新文档状态为失败
|
||
try:
|
||
# Python 3.9 兼容：使用 async for 替代 anext
|
||
async for db_session in get_session():
|
||
try:
|
||
from sqlalchemy import select
|
||
document = await db_session.scalar(select(DocumentModel).where(DocumentModel.id == document_id))
|
||
if document:
|
||
document.is_processed = False
|
||
document.processing_error = str(e)
|
||
await db_session.commit()
|
||
finally:
|
||
await db_session.close()
|
||
break # 只取第一个 session
|
||
except Exception as db_error:
|
||
session.desc = f"ERROR: 更新文档状态失败: {str(db_error)}"
|
||
|
||
return {
|
||
"document_id": document_id,
|
||
"status": "failed",
|
||
"error": str(e),
|
||
"message": "文档处理失败"
|
||
}
|
||
|
||
def delete_document_from_vector_store(self, knowledge_base_id: int, document_id: int) -> None:
|
||
"""从向量存储中删除文档"""
|
||
try:
|
||
# Chroma兼容模式
|
||
from langchain_chroma import Chroma
|
||
kb_vector_path = os.path.join(self.vector_db_path, f"kb_{knowledge_base_id}")
|
||
|
||
if not os.path.exists(kb_vector_path):
|
||
logger.warning(f"向量存储不存在: {kb_vector_path}")
|
||
return
|
||
|
||
chunks = self.get_document_chunks(knowledge_base_id, document_id)
|
||
# 加载向量存储
|
||
vector_store = Chroma(
|
||
persist_directory=kb_vector_path,
|
||
embedding_function=self.embeddings
|
||
)
|
||
|
||
count_before = vector_store._collection.count()
|
||
count_after = count_before
|
||
|
||
if len(chunks) > 0:
|
||
where_filter = {"document_id": str(document_id)}
|
||
vector_store.delete(where=where_filter)
|
||
count_after = vector_store._collection.count()
|
||
|
||
# 注意：Chroma的删除功能可能需要特定的实现方式
|
||
logger.info(f"文档已从向量存储中删除: document_id={document_id}，删除前有 {count_before} 个向量，删除后有 {count_after} 个向量")
|
||
|
||
except Exception as e:
|
||
logger.error(f"从向量存储删除文档失败: {str(e)}")
|
||
raise
|
||
|
||
def get_document_chunks(self, knowledge_base_id: int, document_id: int) -> List[Dict[str, Any]]:
|
||
"""获取文档的所有分段内容
|
||
|
||
改进说明:
|
||
- 避免使用空查询进行相似性搜索，防止触发不必要的embedding API调用
|
||
- 优先使用直接SQL查询，提高性能
|
||
- 确保结果按chunk_index排序
|
||
"""
|
||
try:
|
||
return self._get_chunks_chroma(knowledge_base_id, document_id)
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取文档分段失败 document_id: {document_id}, kb_id: {knowledge_base_id}: {str(e)}")
|
||
return []
|
||
|
||
def _get_chunks_by_sql(self, knowledge_base_id: int, document_id: int) -> List[Dict[str, Any]]:
|
||
"""使用SQLAlchemy连接池查询获取文档分段（推荐方法）"""
|
||
try:
|
||
if not self.pgvector_pool:
|
||
logger.error("PGVector连接池未初始化")
|
||
return []
|
||
|
||
# 直接SQL查询，避免相似性搜索和embedding计算
|
||
query = """
|
||
SELECT
|
||
id,
|
||
document,
|
||
cmetadata
|
||
FROM langchain_pg_embedding
|
||
WHERE cmetadata->>'document_id' = :document_id
|
||
AND cmetadata->>'knowledge_base_id' = :knowledge_base_id
|
||
ORDER BY
|
||
CAST(cmetadata->>'chunk_index' AS INTEGER) ASC;
|
||
"""
|
||
|
||
# 使用连接池执行查询
|
||
session = self.pgvector_pool.get_session()
|
||
try:
|
||
result = session.execute(
|
||
text(query),
|
||
{
|
||
'document_id': str(document_id),
|
||
'knowledge_base_id': str(knowledge_base_id)
|
||
}
|
||
)
|
||
results = result.fetchall()
|
||
|
||
chunks = []
|
||
for row in results:
|
||
# SQLAlchemy结果行访问
|
||
metadata = row.cmetadata
|
||
chunk = {
|
||
"id": f"chunk_{document_id}_{metadata.get('chunk_index', 0)}",
|
||
"content": row.document,
|
||
"metadata": metadata,
|
||
"page_number": metadata.get("page"),
|
||
"chunk_index": metadata.get("chunk_index", 0),
|
||
"start_char": metadata.get("start_char"),
|
||
"end_char": metadata.get("end_char")
|
||
}
|
||
chunks.append(chunk)
|
||
|
||
logger.info(f"通过SQLAlchemy连接池查询获取到文档 {document_id} 的 {len(chunks)} 个分段")
|
||
return chunks
|
||
|
||
finally:
|
||
session.close()
|
||
|
||
except Exception as e:
|
||
logger.error(f"SQLAlchemy连接池查询失败: {e}")
|
||
return []
|
||
|
||
def _get_chunks_by_langchain_improved(self, knowledge_base_id: int, document_id: int, collection_name: str) -> List[Dict[str, Any]]:
|
||
"""改进的LangChain查询方法（回退方案）"""
|
||
try:
|
||
vector_store = PGVector(
|
||
connection=self.connection_string,
|
||
embeddings=self.embeddings,
|
||
collection_name=collection_name,
|
||
use_jsonb=True
|
||
)
|
||
|
||
# 使用有意义的查询而不是空查询，避免触发embedding API错误
|
||
# 先尝试获取少量结果来构造查询
|
||
try:
|
||
sample_results = vector_store.similarity_search(
|
||
query="文档内容", # 使用通用查询词而非空字符串
|
||
k=5,
|
||
filter={"document_id": {"$eq": str(document_id)}}
|
||
)
|
||
|
||
if sample_results:
|
||
# 使用第一个结果的内容片段作为查询
|
||
first_content = sample_results[0].page_content[:50]
|
||
results = vector_store.similarity_search(
|
||
query=first_content,
|
||
k=1000,
|
||
filter={"document_id": {"$eq": str(document_id)}}
|
||
)
|
||
else:
|
||
# 如果没有结果，尝试不使用filter的查询
|
||
results = vector_store.similarity_search(
|
||
query="文档",
|
||
k=1000
|
||
)
|
||
# 手动过滤结果
|
||
results = [doc for doc in results if doc.metadata.get("document_id") == str(document_id)]
|
||
|
||
except Exception as e:
|
||
logger.warning(f"改进的相似性搜索失败: {e}")
|
||
return []
|
||
|
||
chunks = []
|
||
for i, doc in enumerate(results):
|
||
chunk = {
|
||
"id": f"chunk_{document_id}_{i}",
|
||
"content": doc.page_content,
|
||
"metadata": doc.metadata,
|
||
"page_number": doc.metadata.get("page"),
|
||
"chunk_index": doc.metadata.get("chunk_index", i),
|
||
"start_char": doc.metadata.get("start_char"),
|
||
"end_char": doc.metadata.get("end_char")
|
||
}
|
||
chunks.append(chunk)
|
||
|
||
# 按chunk_index排序
|
||
chunks.sort(key=lambda x: x.get("chunk_index", 0))
|
||
|
||
logger.info(f"通过改进的LangChain方法获取到文档 {document_id} 的 {len(chunks)} 个分段")
|
||
return chunks
|
||
|
||
except Exception as e:
|
||
logger.error(f"LangChain改进方法失败: {e}")
|
||
return []
|
||
|
||
def _get_chunks_chroma(self, knowledge_base_id: int, document_id: int) -> List[Dict[str, Any]]:
|
||
"""Chroma存储的处理逻辑"""
|
||
from langchain_chroma import Chroma
|
||
# 构建向量数据库路径
|
||
vector_db_path = os.path.join(self.vector_db_path, f"kb_{knowledge_base_id}")
|
||
|
||
if not os.path.exists(vector_db_path):
|
||
logger.warning(f"向量数据库不存在: {vector_db_path}")
|
||
return []
|
||
|
||
# 加载向量数据库
|
||
vectorstore = Chroma(
|
||
persist_directory=vector_db_path,
|
||
embedding_function=self.embeddings
|
||
)
|
||
|
||
# 获取所有文档的元数据，筛选出指定文档的分段
|
||
collection = vectorstore._collection
|
||
all_docs = collection.get(include=["metadatas", "documents"])
|
||
all_ids_data = collection.get()
|
||
|
||
chunks = []
|
||
chunk_index = 0
|
||
|
||
for i, metadata in enumerate(all_docs["metadatas"]):
|
||
if metadata.get("document_id") == str(document_id):
|
||
chunk_content = all_docs["documents"][i]
|
||
vector_id = all_ids_data["ids"][i]
|
||
|
||
chunk = {
|
||
"id": f"chunk_{document_id}_{chunk_index}",
|
||
"content": chunk_content,
|
||
"metadata": metadata,
|
||
"page_number": metadata.get("page"),
|
||
"chunk_index": chunk_index,
|
||
"start_char": metadata.get("start_char"),
|
||
"end_char": metadata.get("end_char"),
|
||
"vector_id": vector_id
|
||
}
|
||
chunks.append(chunk)
|
||
chunk_index += 1
|
||
|
||
return chunks
|
||
|
||
def search_similar_documents(self, knowledge_base_id: int, query: str, k: int = 5) -> List[Dict[str, Any]]:
|
||
"""在知识库中搜索相似文档"""
|
||
try:
|
||
# if settings.vector_db.type == "pgvector":
|
||
# # PostgreSQL pgvector存储
|
||
# collection_name = f"{settings.vector_db.pgvector_table_name}_kb_{knowledge_base_id}"
|
||
|
||
# try:
|
||
# vector_store = PGVector(
|
||
# connection=self.connection_string,
|
||
# embeddings=self.embeddings,
|
||
# collection_name=collection_name,
|
||
# use_jsonb=True
|
||
# )
|
||
|
||
# # 执行相似性搜索
|
||
# results = vector_store.similarity_search_with_score(query, k=k)
|
||
|
||
# # 格式化结果
|
||
# formatted_results = []
|
||
# for doc, distance_score in results:
|
||
# # pgvector使用余弦距离，距离越小相似度越高
|
||
# # 将距离转换为0-1之间的相似度分数
|
||
# similarity_score = 1.0 / (1.0 + distance_score)
|
||
|
||
# formatted_results.append({
|
||
# "content": doc.page_content,
|
||
# "metadata": doc.metadata,
|
||
# "similarity_score": distance_score, # 保留原始距离分数
|
||
# "normalized_score": similarity_score, # 归一化相似度分数
|
||
# "source": doc.metadata.get('filename', 'unknown'),
|
||
# "document_id": doc.metadata.get('document_id', 'unknown'),
|
||
# "chunk_id": doc.metadata.get('chunk_id', 'unknown')
|
||
# })
|
||
|
||
# # 按相似度分数排序（距离越小越相似）
|
||
# formatted_results.sort(key=lambda x: x['similarity_score'])
|
||
|
||
# logger.info(f"PostgreSQL pgvector搜索完成，找到 {len(formatted_results)} 个相关文档")
|
||
# return formatted_results
|
||
|
||
# except Exception as e:
|
||
# logger.warning(f"PostgreSQL pgvector存储不存在: {collection_name}, {str(e)}")
|
||
# return []
|
||
# else:
|
||
# Chroma兼容模式
|
||
from langchain_chroma import Chroma
|
||
kb_vector_path = os.path.join(self.vector_db_path, f"kb_{knowledge_base_id}")
|
||
|
||
if not os.path.exists(kb_vector_path):
|
||
logger.warning(f"向量存储不存在: {kb_vector_path}")
|
||
return []
|
||
|
||
# 加载向量存储
|
||
vector_store = Chroma(
|
||
persist_directory=kb_vector_path,
|
||
embedding_function=self.embeddings
|
||
)
|
||
|
||
# 执行相似性搜索
|
||
results = vector_store.similarity_search_with_score(query, k=k)
|
||
|
||
# 格式化结果
|
||
formatted_results = []
|
||
for doc, distance_score in results:
|
||
# Chroma使用欧几里得距离，距离越小相似度越高
|
||
# 将距离转换为0-1之间的相似度分数
|
||
similarity_score = 1.0 / (1.0 + distance_score)
|
||
|
||
formatted_results.append({
|
||
"content": doc.page_content,
|
||
"metadata": doc.metadata,
|
||
"similarity_score": distance_score, # 保留原始距离分数
|
||
"normalized_score": similarity_score, # 归一化相似度分数
|
||
"source": doc.metadata.get('filename', 'unknown'),
|
||
"document_id": doc.metadata.get('document_id', 'unknown'),
|
||
"chunk_id": doc.metadata.get('chunk_id', 'unknown')
|
||
})
|
||
|
||
# 按相似度分数排序（距离越小越相似）
|
||
formatted_results.sort(key=lambda x: x['similarity_score'])
|
||
|
||
logger.info(f"搜索完成，找到 {len(formatted_results)} 个相关文档")
|
||
return formatted_results
|
||
|
||
except Exception as e:
|
||
logger.error(f"搜索文档失败: {str(e)}")
|
||
return [] # 返回空列表而不是抛出异常
|
||
|
||
# 全局文档处理器实例（延迟初始化）
|
||
document_processor = None
|
||
|
||
async def get_document_processor(session: Session = None):
|
||
"""获取文档处理器实例（延迟初始化）"""
|
||
global document_processor
|
||
if session:
|
||
session.desc = "获取文档处理器实例"
|
||
if document_processor is None:
|
||
document_processor = DocumentProcessor()
|
||
await document_processor.initialize(session)
|
||
return document_processor |