1113 lines
47 KiB
Python
1113 lines
47 KiB
Python
import os
|
||
import re
|
||
import time
|
||
import math
|
||
import numpy as np
|
||
from datetime import datetime
|
||
from typing import List, Dict, Any, Optional
|
||
from langchain_core.documents import Document
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
|
||
from appPublic.log import debug, error, info
|
||
from filetxt.loader import fileloader, File2Text
|
||
from rag.uapi_service import APIService
|
||
from rag.service_opts import get_service_params
|
||
from rag.transaction_manager import TransactionManager, OperationType
|
||
from pdf2image import convert_from_path
|
||
import pytesseract
|
||
import base64
|
||
from pathlib import Path
|
||
|
||
class RagOperations:
|
||
"""RAG 操作类，提供所有通用的 RAG 操作"""
|
||
|
||
def __init__(self):
|
||
self.api_service = APIService()
|
||
|
||
async def load_and_chunk_document(self, realpath: str, timings: Dict,
|
||
transaction_mgr: TransactionManager = None) -> List[Document]:
|
||
"""加载文件并进行文本分片"""
|
||
debug(f"加载文件: {realpath}")
|
||
start_load = time.time()
|
||
|
||
# 检查文件格式支持
|
||
supported_formats = File2Text.supported_types()
|
||
debug(f"支持的文件格式:{supported_formats}")
|
||
ext = realpath.rsplit('.', 1)[1].lower() if '.' in realpath else ''
|
||
if ext not in supported_formats:
|
||
raise ValueError(f"不支持的文件格式: {ext}, 支持的格式: {', '.join(supported_formats)}")
|
||
|
||
text = fileloader(realpath)
|
||
if ext == 'pdf':
|
||
debug(f"pdf原生提取结果是：{text}")
|
||
if not text or len(text.strip()) == 0: # 更严格的空值检查
|
||
debug(f"pdf原生提取失败，尝试扫描件提取")
|
||
ocr_text = self.pdf_to_text(realpath)
|
||
debug(f"pdf扫描件抽取的文本内容是：{ocr_text}")
|
||
text = ocr_text # 只在原生提取失败时使用OCR结果
|
||
|
||
# 只在有文本内容时进行清洗
|
||
if text and len(text.strip()) > 0:
|
||
# 或者保留更多有用字符
|
||
text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s.;,\n/]', '', text)
|
||
else:
|
||
error(f"文件 {realpath} 无法提取任何文本内容")
|
||
text = "" # 确保为空字符串
|
||
|
||
timings["load_file"] = time.time() - start_load
|
||
debug(f"加载文件耗时: {timings['load_file']:.2f} 秒, 文本长度: {len(text)}")
|
||
|
||
# # 加载文件内容
|
||
# text = fileloader(realpath)
|
||
# debug(f"pdf原生提取结果是：{text}")
|
||
# if len(text) == 0:
|
||
# debug(f"pdf原生提取失败，尝试扫描件提取")
|
||
# text = self.pdf_to_text(realpath)
|
||
# debug(f"pdf扫描件抽取的文本内容是：{text}")
|
||
# text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s.;,\n/]', '', text)
|
||
# timings["load_file"] = time.time() - start_load
|
||
# debug(f"加载文件耗时: {timings['load_file']:.2f} 秒, 文本长度: {len(text)}")
|
||
#
|
||
# if not text or not text.strip():
|
||
# raise ValueError(f"文件 {realpath} 加载为空")
|
||
|
||
# 分片处理
|
||
document = Document(page_content=text)
|
||
text_splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=500,
|
||
chunk_overlap=100,
|
||
length_function=len
|
||
)
|
||
debug("开始分片文件内容")
|
||
start_split = time.time()
|
||
chunks = text_splitter.split_documents([document])
|
||
timings["split_text"] = time.time() - start_split
|
||
debug(f"文本分片耗时: {timings['split_text']:.2f} 秒, 分片数量: {len(chunks)}")
|
||
|
||
if not chunks:
|
||
raise ValueError(f"文件 {realpath} 未生成任何文档块")
|
||
|
||
# 记录事务操作
|
||
if transaction_mgr:
|
||
transaction_mgr.add_operation(
|
||
OperationType.FILE_LOAD,
|
||
{'realpath': realpath, 'chunks_count': len(chunks)}
|
||
)
|
||
|
||
return chunks
|
||
|
||
def pdf_to_text(
|
||
self,
|
||
pdf_path: str,
|
||
output_txt: Optional[str] = None,
|
||
dpi: int = 300,
|
||
lang: str = 'chi_sim+chi_tra+eng'
|
||
) -> str:
|
||
"""
|
||
将扫描版 PDF 转为文字（你原来的代码，一行调用版）
|
||
|
||
参数:
|
||
pdf_path: PDF 文件路径（字符串）
|
||
output_txt: 如果提供，会自动保存到这个 txt 文件（可选）
|
||
dpi: 图片分辨率，默认 300（越高越清晰）
|
||
lang: 语言包，默认中文简体+繁体+英文
|
||
|
||
返回:
|
||
提取出的完整文字（字符串）
|
||
"""
|
||
# PDF 转图片
|
||
images = convert_from_path(pdf_path, dpi=dpi)
|
||
|
||
# OCR 识别
|
||
text = ''
|
||
for img in images:
|
||
text += pytesseract.image_to_string(img, lang=lang) + '\n'
|
||
|
||
# 可选：自动保存到文件
|
||
if output_txt:
|
||
with open(output_txt, 'w', encoding='utf-8') as f:
|
||
f.write(text)
|
||
|
||
return text
|
||
|
||
async def generate_embeddings(self, request, chunks: List[Document], service_params: Dict,
|
||
userid: str, timings: Dict,
|
||
transaction_mgr: TransactionManager = None) -> List[List[float]]:
|
||
"""生成嵌入向量"""
|
||
debug("调用嵌入服务生成向量")
|
||
start_embedding = time.time()
|
||
texts = [chunk.page_content for chunk in chunks]
|
||
embeddings = []
|
||
|
||
# 批量处理嵌入
|
||
for i in range(0, len(texts), 10):
|
||
batch_texts = texts[i:i + 10]
|
||
batch_embeddings = await self.api_service.get_embeddings(
|
||
request=request,
|
||
texts=batch_texts,
|
||
upappid=service_params['embedding'],
|
||
apiname="BAAI/bge-m3",
|
||
user=userid
|
||
)
|
||
embeddings.extend(batch_embeddings)
|
||
|
||
if not embeddings or not all(len(vec) == 1024 for vec in embeddings):
|
||
raise ValueError("所有嵌入向量必须是长度为 1024 的浮点数列表")
|
||
|
||
timings["generate_embeddings"] = time.time() - start_embedding
|
||
debug(f"生成嵌入向量耗时: {timings['generate_embeddings']:.2f} 秒, 嵌入数量: {len(embeddings)}")
|
||
|
||
# 记录事务操作
|
||
if transaction_mgr:
|
||
transaction_mgr.add_operation(
|
||
OperationType.EMBEDDING,
|
||
{'embeddings_count': len(embeddings)}
|
||
)
|
||
|
||
return embeddings
|
||
|
||
async def generate_multi_embeddings(self, request, inputs: List[Dict], service_params: Dict,
|
||
userid: str, timings: Dict,
|
||
transaction_mgr: TransactionManager = None) -> Dict[str, Dict]:
|
||
"""调用多模态嵌入服务（CLIP）"""
|
||
debug("调用多模态嵌入服务")
|
||
start = time.time()
|
||
|
||
result = await self.api_service.get_multi_embeddings(
|
||
request=request,
|
||
inputs=inputs,
|
||
upappid=service_params['embedding'],
|
||
apiname="black/clip",
|
||
user=userid
|
||
)
|
||
debug(f"多模态返回结果是{result}")
|
||
timings["multi_embedding"] = time.time() - start
|
||
debug(f"多模态嵌入耗时: {timings['multi_embedding']:.2f}秒，处理 {len(result)} 条")
|
||
|
||
# ==================== 新增：错误检查 + 过滤 ====================
|
||
valid_results = {}
|
||
error_count = 0
|
||
error_examples = []
|
||
|
||
for key, info in result.items():
|
||
if info.get("type") == "error":
|
||
error_count += 1
|
||
if len(error_examples) < 3: # 只记录前3个
|
||
error_examples.append(f"{key} → {info['error']}")
|
||
# 直接丢弃错误条目
|
||
continue
|
||
valid_results[key] = info
|
||
|
||
if error_count > 0:
|
||
error(f"多模态嵌入失败 {error_count} 条！示例：{'; '.join(error_examples)}")
|
||
raise RuntimeError(f"多模态嵌入有{error_count} 条失败")
|
||
else:
|
||
debug("多模态嵌入全部成功！")
|
||
|
||
if transaction_mgr:
|
||
transaction_mgr.add_operation(
|
||
OperationType.EMBEDDING,
|
||
{'count': len(result)}
|
||
)
|
||
|
||
return result
|
||
|
||
# async def force_l2_normalize(self, vector: List[float]) -> List[float]:
|
||
# """万无一失的 L2 归一化"""
|
||
# arr = np.array(vector, dtype=np.float32)
|
||
# norm = np.linalg.norm(arr)
|
||
# if norm == 0:
|
||
# return vector # 全零向量无法归一化
|
||
# return (arr / norm).tolist()
|
||
|
||
# 统一插入向量库
|
||
async def insert_all_vectors(
|
||
self,
|
||
request,
|
||
text_chunks: List[Document],
|
||
realpath: str,
|
||
orgid: str,
|
||
fiid: str,
|
||
document_id: str,
|
||
service_params: Dict,
|
||
userid: str,
|
||
db_type: str,
|
||
timings: Dict,
|
||
img_paths: List[str] = None,
|
||
text_embeddings: List[List[float]] = None,
|
||
multi_results: Dict = None,
|
||
transaction_mgr: TransactionManager = None
|
||
) -> Dict[str, int]:
|
||
"""
|
||
统一插入函数：支持两种模式
|
||
1. 纯文本模式：text_embeddings 有值
|
||
2. 多模态模式：multi_results 有值（来自 generate_multi_embeddings）
|
||
"""
|
||
img_paths = img_paths or []
|
||
all_chunks = []
|
||
start = time.time()
|
||
filename = os.path.basename(realpath)
|
||
upload_time = datetime.now().isoformat()
|
||
|
||
# ==================== 1. 纯文本模式（BGE） ====================
|
||
if text_embeddings is not None:
|
||
debug(f"【纯文本模式】插入 {len(text_embeddings)} 条文本向量")
|
||
for i, chunk in enumerate(text_chunks):
|
||
all_chunks.append({
|
||
"userid": orgid,
|
||
"knowledge_base_id": fiid,
|
||
"text": chunk.page_content,
|
||
"vector": text_embeddings[i],
|
||
"document_id": document_id,
|
||
"filename": filename,
|
||
"file_path": realpath,
|
||
"upload_time": upload_time,
|
||
"file_type": "text",
|
||
})
|
||
|
||
# ==================== 2. 多模态模式（CLIP 混排） ====================
|
||
if multi_results is not None:
|
||
debug(f"【多模态模式】解析 {len(multi_results)} 条 CLIP 结果")
|
||
# 遍历 multi_results
|
||
for raw_key, info in multi_results.items():
|
||
typ = info["type"]
|
||
# vector = info["vector"]
|
||
# debug(f"从后端传回来的向量数据是：{vector}")
|
||
# emb = await self.force_l2_normalize(info["vector"])
|
||
# debug(f"归一化后的向量数据是：{emb}")
|
||
# --- 文本 ---
|
||
if typ == "text":
|
||
# raw_key 就是原文
|
||
all_chunks.append({
|
||
"userid": orgid,
|
||
"knowledge_base_id": fiid,
|
||
"text": raw_key,
|
||
"vector": info["vector"],
|
||
"document_id": document_id,
|
||
"filename": filename,
|
||
"file_path": realpath,
|
||
"upload_time": upload_time,
|
||
"file_type": "text",
|
||
})
|
||
continue
|
||
|
||
# --- 图像 ---
|
||
if typ == "image":
|
||
img_path = info.get("path") or raw_key
|
||
img_name = os.path.basename(img_path)
|
||
|
||
# 整图向量
|
||
if "vector" in info:
|
||
all_chunks.append({
|
||
"userid": orgid,
|
||
"knowledge_base_id": fiid,
|
||
"text": f"[Image: {img_path}]图片来源于文件{realpath}",
|
||
"vector": info["vector"],
|
||
"document_id": document_id,
|
||
"filename": img_name,
|
||
"file_path": realpath,
|
||
"upload_time": upload_time,
|
||
"file_type": "image",
|
||
})
|
||
|
||
# 人脸向量
|
||
face_vecs = info.get("face_vecs", [])
|
||
face_count = len(face_vecs)
|
||
# if face_count > 0:
|
||
# for f_idx, fvec in enumerate(face_vecs):
|
||
# debug(f"人脸向量维度是：{len(fvec)}")
|
||
# all_chunks.append({
|
||
# "userid": orgid,
|
||
# "knowledge_base_id": fiid,
|
||
# "text": f"[Face {f_idx + 1}/{face_count} in {img_name}]人脸来源于{realpath}的{img_path}图片",
|
||
# "vector": fvec,
|
||
# "document_id": document_id,
|
||
# "filename": img_name,
|
||
# "file_path": realpath,
|
||
# "upload_time": upload_time,
|
||
# "file_type": "face",
|
||
# })
|
||
continue
|
||
|
||
# --- 视频 ---
|
||
if typ == "video":
|
||
video_path = info.get("path") or raw_key
|
||
video_name = os.path.basename(video_path)
|
||
|
||
if "vector" in info:
|
||
all_chunks.append({
|
||
"userid": orgid,
|
||
"knowledge_base_id": fiid,
|
||
"text": f"[Video: {video_name}]",
|
||
"vector": info["vector"],
|
||
"document_id": document_id,
|
||
"filename": video_path,
|
||
"file_path": realpath,
|
||
"upload_time": upload_time,
|
||
"file_type": "video",
|
||
})
|
||
|
||
# 视频人脸
|
||
face_vecs = info.get("face_vecs", [])
|
||
face_count = len(face_vecs)
|
||
# if face_count > 0 :
|
||
# for f_idx, fvec in enumerate(face_vecs):
|
||
# all_chunks.append({
|
||
# "userid": orgid,
|
||
# "knowledge_base_id": fiid,
|
||
# "text": f"[Face {f_idx + 1}/{face_count} in video {video_name}]来源于{video_path}",
|
||
# "vector": fvec,
|
||
# "document_id": document_id,
|
||
# "filename": video_path,
|
||
# "file_path": realpath,
|
||
# "upload_time": upload_time,
|
||
# "file_type": "face",
|
||
# })
|
||
continue
|
||
|
||
# --- 音频 ---
|
||
if typ == "audio":
|
||
audio_path = info.get("path") or raw_key
|
||
audio_name = os.path.basename(audio_path)
|
||
|
||
if "vector" in info:
|
||
all_chunks.append({
|
||
"userid": orgid,
|
||
"knowledge_base_id": fiid,
|
||
"text": f"[Audio: {audio_name}]",
|
||
"vector": info["vector"],
|
||
"document_id": document_id,
|
||
"filename": audio_path,
|
||
"file_path": realpath,
|
||
"upload_time": upload_time,
|
||
"file_type": "audio",
|
||
})
|
||
continue
|
||
|
||
# --- 未知类型 ---
|
||
debug(f"未知类型跳过: {typ} → {raw_key}")
|
||
|
||
# ==================== 3. 批量插入 Milvus ====================
|
||
if not all_chunks:
|
||
debug("无向量需要插入")
|
||
return {"text": 0, "image": 0, "face": 0}
|
||
|
||
for i in range(0, len(all_chunks), 10):
|
||
batch = all_chunks[i:i + 10]
|
||
result = await self.api_service.milvus_insert_document(
|
||
request=request,
|
||
chunks=batch,
|
||
upappid=service_params['vdb'],
|
||
apiname="milvus/insertdocument",
|
||
user=userid,
|
||
db_type=db_type
|
||
)
|
||
if result.get("status") != "success":
|
||
raise ValueError(f"Milvus 插入失败: {result.get('message')}")
|
||
|
||
# ==================== 4. 统一回滚（只登记一次） ====================
|
||
if transaction_mgr and all_chunks:
|
||
async def rollback_all(data, context):
|
||
try:
|
||
await self.delete_from_vector_db(
|
||
request=context['request'],
|
||
orgid=data['orgid'],
|
||
realpath=data['realpath'],
|
||
fiid=data['fiid'],
|
||
id=data['document_id'],
|
||
service_params=context['service_params'],
|
||
userid=context['userid'],
|
||
db_type=data['db_type']
|
||
)
|
||
return f"已回滚 document_id={data['document_id']} 的所有向量"
|
||
except Exception as e:
|
||
error(f"统一回滚失败: {e}")
|
||
raise
|
||
|
||
transaction_mgr.add_operation(
|
||
OperationType.VDB_INSERT,
|
||
{
|
||
'orgid': orgid,
|
||
'realpath': realpath,
|
||
'fiid': fiid,
|
||
'id': document_id,
|
||
'db_type': db_type
|
||
},
|
||
rollback_func=rollback_all
|
||
)
|
||
|
||
# ==================== 5. 统计返回 ====================
|
||
stats = {
|
||
"text": len([c for c in all_chunks if c["file_type"] == "text"]),
|
||
"image": len([c for c in all_chunks if c["file_type"] == "image"]),
|
||
"face": len([c for c in all_chunks if c["file_type"] == "face"])
|
||
}
|
||
|
||
timings["insert_all"] = time.time() - start
|
||
debug(
|
||
f"统一插入完成: 文本 {stats['text']}, 图像 {stats['image']}, 人脸 {stats['face']}, 耗时 {timings['insert_all']:.2f}s")
|
||
return stats
|
||
# async def insert_to_vector_db(self, request, chunks: List[Document], embeddings: List[List[float]],
|
||
# realpath: str, orgid: str, fiid: str, id: str, service_params: Dict,
|
||
# userid: str, db_type: str, timings: Dict,
|
||
# transaction_mgr: TransactionManager = None):
|
||
# """插入向量数据库"""
|
||
# debug(f"准备数据并调用插入文件端点: {realpath}")
|
||
# filename = os.path.basename(realpath).rsplit('.', 1)[0]
|
||
# ext = realpath.rsplit('.', 1)[1].lower() if '.' in realpath else ''
|
||
# upload_time = datetime.now().isoformat()
|
||
#
|
||
# chunks_data = [
|
||
# {
|
||
# "userid": orgid,
|
||
# "knowledge_base_id": fiid,
|
||
# "text": chunk.page_content,
|
||
# "vector": embeddings[i],
|
||
# "document_id": id,
|
||
# "filename": filename + '.' + ext,
|
||
# "file_path": realpath,
|
||
# "upload_time": upload_time,
|
||
# "file_type": ext,
|
||
# }
|
||
# for i, chunk in enumerate(chunks)
|
||
# ]
|
||
#
|
||
# start_milvus = time.time()
|
||
# for i in range(0, len(chunks_data), 10):
|
||
# batch_chunks = chunks_data[i:i + 10]
|
||
# debug(f"传入的数据是：{batch_chunks}")
|
||
# result = await self.api_service.milvus_insert_document(
|
||
# request=request,
|
||
# chunks=batch_chunks,
|
||
# db_type=db_type,
|
||
# upappid=service_params['vdb'],
|
||
# apiname="milvus/insertdocument",
|
||
# user=userid
|
||
# )
|
||
# if result.get("status") != "success":
|
||
# raise ValueError(result.get("message", "Milvus 插入失败"))
|
||
#
|
||
# timings["insert_milvus"] = time.time() - start_milvus
|
||
# debug(f"Milvus 插入耗时: {timings['insert_milvus']:.2f} 秒")
|
||
#
|
||
# # 记录事务操作，包含回滚函数
|
||
# if transaction_mgr:
|
||
# async def rollback_vdb_insert(data, context):
|
||
# try:
|
||
# # 防御性检查
|
||
# required_context = ['request', 'service_params', 'userid']
|
||
# missing_context = [k for k in required_context if k not in context or context[k] is None]
|
||
# if missing_context:
|
||
# raise ValueError(f"回滚上下文缺少字段: {', '.join(missing_context)}")
|
||
#
|
||
# required_data = ['orgid', 'realpath', 'fiid', 'id', 'db_type']
|
||
# missing_data = [k for k in required_data if k not in data or data[k] is None]
|
||
# if missing_data:
|
||
# raise ValueError(f"VDB_INSERT 数据缺少字段: {', '.join(missing_data)}")
|
||
#
|
||
# await self.delete_from_vector_db(
|
||
# context['request'], data['orgid'], data['realpath'],
|
||
# data['fiid'], data['id'], context['service_params'],
|
||
# context['userid'], data['db_type']
|
||
# )
|
||
# return f"已回滚向量数据库插入: {data['id']}"
|
||
# except Exception as e:
|
||
# error(f"回滚向量数据库失败: document_id={data.get('id', '未知')}, 错误: {str(e)}")
|
||
# raise
|
||
#
|
||
# transaction_mgr.add_operation(
|
||
# OperationType.VDB_INSERT,
|
||
# {
|
||
# 'orgid': orgid, 'realpath': realpath, 'fiid': fiid,
|
||
# 'id': id, 'db_type': db_type
|
||
# },
|
||
# rollback_func=rollback_vdb_insert
|
||
# )
|
||
#
|
||
# return chunks_data
|
||
#
|
||
# async def insert_image_vectors(
|
||
# self,
|
||
# request,
|
||
# multi_results: Dict[str, Dict],
|
||
# realpath: str,
|
||
# orgid: str,
|
||
# fiid: str,
|
||
# document_id: str,
|
||
# service_params: Dict,
|
||
# userid: str,
|
||
# db_type: str,
|
||
# timings: Dict,
|
||
# transaction_mgr: TransactionManager = None
|
||
# ) -> tuple[int, int]:
|
||
#
|
||
# start = time.time()
|
||
# image_chunks = []
|
||
# face_chunks = []
|
||
#
|
||
# for img_path, info in multi_results.items():
|
||
# # img_name = os.path.basename(img_path)
|
||
#
|
||
# # 1. 插入整张图
|
||
# if info.get("type") in ["image", "video"] and "vector" in info:
|
||
# image_chunks.append({
|
||
# "userid": orgid,
|
||
# "knowledge_base_id": fiid,
|
||
# "text": f"[Image: {img_path}]",
|
||
# "vector": info["vector"],
|
||
# "document_id": document_id,
|
||
# "filename": os.path.basename(realpath),
|
||
# "file_path": realpath,
|
||
# "upload_time": datetime.now().isoformat(),
|
||
# "file_type": "image"
|
||
# })
|
||
#
|
||
# # 2. 插入每张人脸
|
||
# face_vecs = info.get("face_vecs")
|
||
# face_count = info.get("face_count", 0)
|
||
#
|
||
# if face_count > 0 and face_vecs and len(face_vecs) == face_count:
|
||
# for idx, face_vec in enumerate(face_vecs):
|
||
# face_chunks.append({
|
||
# "userid": orgid,
|
||
# "knowledge_base_id": fiid,
|
||
# "text": f"[Face {idx + 1}/{face_count} in {img_path}]",
|
||
# "vector": face_vec,
|
||
# "document_id": document_id,
|
||
# "filename": os.path.basename(realpath),
|
||
# "file_path": realpath,
|
||
# "upload_time": datetime.now().isoformat(),
|
||
# "file_type": "face",
|
||
# })
|
||
#
|
||
# if image_chunks:
|
||
# for i in range(0, len(image_chunks), 10):
|
||
# await self.api_service.milvus_insert_document(
|
||
# request=request,
|
||
# chunks=image_chunks[i:i + 10],
|
||
# upappid=service_params['vdb'],
|
||
# apiname="milvus/insertdocument",
|
||
# user=userid,
|
||
# db_type=db_type
|
||
# )
|
||
#
|
||
# if face_chunks:
|
||
# for i in range(0, len(face_chunks), 10):
|
||
# await self.api_service.milvus_insert_document(
|
||
# request=request,
|
||
# chunks=face_chunks[i:i + 10],
|
||
# upappid=service_params['vdb'],
|
||
# apiname="milvus/insertdocument",
|
||
# user=userid,
|
||
# db_type=db_type
|
||
# )
|
||
# timings["insert_images"] = time.time() - start
|
||
# image_count = len(image_chunks)
|
||
# face_count = len(face_chunks)
|
||
#
|
||
# debug(f"多模态插入完成: 图像 {image_count} 条, 人脸 {face_count} 条")
|
||
#
|
||
# if transaction_mgr and (image_count + face_count > 0):
|
||
# transaction_mgr.add_operation(
|
||
# OperationType.IMAGE_VECTORS_INSERT,
|
||
# {"images": image_count, "faces": face_count, "document_id": document_id}
|
||
# )
|
||
#
|
||
# # 记录事务操作，包含回滚函数
|
||
# if transaction_mgr:
|
||
# async def rollback_multimodal(data, context):
|
||
# try:
|
||
# # 防御性检查
|
||
# required_context = ['request', 'service_params', 'userid']
|
||
# missing_context = [k for k in required_context if k not in context or context[k] is None]
|
||
# if missing_context:
|
||
# raise ValueError(f"回滚上下文缺少字段: {', '.join(missing_context)}")
|
||
#
|
||
# required_data = ['orgid', 'realpath', 'fiid', 'id', 'db_type']
|
||
# missing_data = [k for k in required_data if k not in data or data[k] is None]
|
||
# if missing_data:
|
||
# raise ValueError(f"多模态回滚数据缺少字段: {', '.join(missing_data)}")
|
||
#
|
||
# await self.delete_from_vector_db(
|
||
# context['request'], data['orgid'], data['realpath'],
|
||
# data['fiid'], data['id'], context['service_params'],
|
||
# context['userid'], data['db_type']
|
||
# )
|
||
# return f"已回滚多模态向量: {data['id']}"
|
||
# except Exception as e:
|
||
# error(f"多模态回滚向量数据库失败: document_id={data.get('id', '未知')}, 错误: {str(e)}")
|
||
# raise
|
||
#
|
||
# transaction_mgr.add_operation(
|
||
# OperationType.VDB_INSERT,
|
||
# {
|
||
# 'orgid': orgid, 'realpath': realpath, 'fiid': fiid,
|
||
# 'id': id, 'db_type': db_type
|
||
# },
|
||
# rollback_func=rollback_multimodal
|
||
# )
|
||
#
|
||
# return image_count, face_count
|
||
|
||
async def insert_to_vector_text(self, request,
|
||
db_type: str, fields: Dict, service_params: Dict, userid: str, timings: Dict) -> List[Dict]:
|
||
"""插入单一纯文本到向量数据库，支持动态 schema"""
|
||
|
||
chunk_data = {}
|
||
debug("准备单一纯文本数据并调用插入端点")
|
||
start = time.time()
|
||
for key, value in fields.items():
|
||
chunk_data[key] = value
|
||
|
||
chunks_data = [chunk_data]
|
||
debug(f"向量库插入传入的数据是：{chunks_data}")
|
||
|
||
# 调用 Milvus 插入
|
||
result = await self.api_service.milvus_insert_document(
|
||
request=request,
|
||
chunks=chunks_data,
|
||
upappid=service_params['vdb'],
|
||
apiname="milvus/insertdocument",
|
||
user=userid,
|
||
db_type=db_type
|
||
)
|
||
if result.get("status") != "success":
|
||
raise ValueError(result.get("message", "Milvus 插入失败"))
|
||
|
||
debug(f"成功插入纯文本到集合 {result.get('collection_name')}")
|
||
timings["textinsert"] = time.time() - start
|
||
debug(f"插入纯文本耗时: {timings['textinsert']:.2f} 秒")
|
||
|
||
return chunks_data
|
||
|
||
async def extract_triples(self, request, chunks: List[Document], service_params: Dict,
|
||
userid: str, timings: Dict,
|
||
transaction_mgr: TransactionManager = None) -> List[Dict]:
|
||
"""抽取三元组"""
|
||
debug("调用三元组抽取服务")
|
||
start_triples = time.time()
|
||
chunk_texts = [doc.page_content for doc in chunks]
|
||
triples = []
|
||
|
||
for i, chunk in enumerate(chunk_texts):
|
||
result = await self.api_service.extract_triples(
|
||
request=request,
|
||
text=chunk,
|
||
upappid=service_params['triples'],
|
||
apiname="Babelscape/mrebel-large",
|
||
user=userid
|
||
)
|
||
if isinstance(result, list):
|
||
triples.extend(result)
|
||
debug(f"分片 {i + 1} 抽取到 {len(result)} 个三元组")
|
||
else:
|
||
error(f"分片 {i + 1} 处理失败: {str(result)}")
|
||
|
||
# 去重和优化三元组
|
||
unique_triples = self._deduplicate_triples(triples)
|
||
|
||
timings["extract_triples"] = time.time() - start_triples
|
||
debug(f"三元组抽取耗时: {timings['extract_triples']:.2f} 秒, 抽取到 {len(unique_triples)} 个三元组")
|
||
|
||
# 记录事务操作
|
||
if transaction_mgr:
|
||
transaction_mgr.add_operation(
|
||
OperationType.TRIPLES_EXTRACT,
|
||
{'triples_count': len(unique_triples)}
|
||
)
|
||
|
||
return unique_triples
|
||
|
||
async def insert_to_graph_db(self, request, triples: List[Dict], id: str, fiid: str,
|
||
orgid: str, service_params: Dict, userid: str, timings: Dict,
|
||
transaction_mgr: TransactionManager = None):
|
||
|
||
"""插入图数据库"""
|
||
debug(f"插入 {len(triples)} 个三元组到 Neo4j")
|
||
start_neo4j = time.time()
|
||
|
||
if triples:
|
||
for i in range(0, len(triples), 30):
|
||
batch_triples = triples[i:i + 30]
|
||
neo4j_result = await self.api_service.neo4j_insert_triples(
|
||
request=request,
|
||
triples=batch_triples,
|
||
document_id=id,
|
||
knowledge_base_id=fiid,
|
||
userid=orgid,
|
||
upappid=service_params['gdb'],
|
||
apiname="neo4j/inserttriples",
|
||
user=userid
|
||
)
|
||
if neo4j_result.get("status") != "success":
|
||
raise ValueError(f"Neo4j 三元组插入失败: {neo4j_result.get('message', '未知错误')}")
|
||
info(f"文件三元组成功插入 Neo4j: {neo4j_result.get('message')}")
|
||
|
||
timings["insert_neo4j"] = time.time() - start_neo4j
|
||
debug(f"Neo4j 插入耗时: {timings['insert_neo4j']:.2f} 秒")
|
||
else:
|
||
debug("未抽取到三元组")
|
||
timings["insert_neo4j"] = 0.0
|
||
|
||
# 记录事务操作，包含回滚函数
|
||
if transaction_mgr:
|
||
async def rollback_gdb_insert(data, context):
|
||
await self.delete_from_graph_db(
|
||
context['request'], data['id'],
|
||
context['service_params'], context['userid']
|
||
)
|
||
return f"已回滚图数据库插入: {data['id']}"
|
||
|
||
transaction_mgr.add_operation(
|
||
OperationType.GDB_INSERT,
|
||
{'id': id, 'triples_count': len(triples)},
|
||
rollback_func=rollback_gdb_insert
|
||
)
|
||
|
||
async def delete_from_vector_db(self, request, orgid: str, realpath: str, fiid: str,
|
||
id: str, service_params: Dict, userid: str, db_type: str):
|
||
"""从向量数据库删除文档"""
|
||
debug(f"调用删除文件端点: userid={orgid}, file_path={realpath}, knowledge_base_id={fiid}, document_id={id}")
|
||
milvus_result = await self.api_service.milvus_delete_document(
|
||
request=request,
|
||
userid=orgid,
|
||
file_path=realpath,
|
||
knowledge_base_id=fiid,
|
||
document_id=id,
|
||
db_type=db_type,
|
||
upappid=service_params['vdb'],
|
||
apiname="milvus/deletedocument",
|
||
user=userid
|
||
)
|
||
if milvus_result.get("status") != "success":
|
||
raise ValueError(milvus_result.get("message", "Milvus 删除失败"))
|
||
|
||
async def delete_from_graph_db(self, request, id: str, service_params: Dict, userid: str):
|
||
"""从图数据库删除文档"""
|
||
debug(f"调用 Neo4j 删除文档端点: document_id={id}")
|
||
neo4j_result = await self.api_service.neo4j_delete_document(
|
||
request=request,
|
||
document_id=id,
|
||
upappid=service_params['gdb'],
|
||
apiname="neo4j/deletedocument",
|
||
user=userid
|
||
)
|
||
if neo4j_result.get("status") != "success":
|
||
raise ValueError(neo4j_result.get("message", "Neo4j 删除失败"))
|
||
nodes_deleted = neo4j_result.get("nodes_deleted", 0)
|
||
rels_deleted = neo4j_result.get("rels_deleted", 0)
|
||
info(f"成功删除 document_id={id} 的 {nodes_deleted} 个 Neo4j 节点和 {rels_deleted} 个关系")
|
||
return nodes_deleted, rels_deleted
|
||
|
||
async def extract_entities(self, request, query: str, service_params: Dict, userid: str,
|
||
timings: Dict) -> List[str]:
|
||
"""提取实体"""
|
||
debug(f"提取查询实体: {query}")
|
||
start_extract = time.time()
|
||
entities = await self.api_service.extract_entities(
|
||
request=request,
|
||
query=query,
|
||
upappid=service_params['entities'],
|
||
apiname="LTP/small",
|
||
user=userid
|
||
)
|
||
timings["entity_extraction"] = time.time() - start_extract
|
||
debug(f"提取实体: {entities}, 耗时: {timings['entity_extraction']:.3f} 秒")
|
||
return entities
|
||
|
||
async def match_triplets(self, request, query: str, entities: List[str], orgid: str,
|
||
fiids: List[str], service_params: Dict, userid: str,
|
||
timings: Dict) -> List[Dict]:
|
||
"""匹配三元组"""
|
||
debug("开始三元组匹配")
|
||
start_triplet = time.time()
|
||
all_triplets = []
|
||
|
||
for kb_id in fiids:
|
||
debug(f"调用 Neo4j 三元组匹配: knowledge_base_id={kb_id}")
|
||
try:
|
||
neo4j_result = await self.api_service.neo4j_match_triplets(
|
||
request=request,
|
||
query=query,
|
||
query_entities=entities,
|
||
userid=orgid,
|
||
knowledge_base_id=kb_id,
|
||
upappid=service_params['gdb'],
|
||
apiname="neo4j/matchtriplets",
|
||
user=userid
|
||
)
|
||
if neo4j_result.get("status") == "success":
|
||
triplets = neo4j_result.get("triplets", [])
|
||
all_triplets.extend(triplets)
|
||
debug(f"知识库 {kb_id} 匹配到 {len(triplets)} 个三元组")
|
||
else:
|
||
error(f"Neo4j 三元组匹配失败: knowledge_base_id={kb_id}, 错误: {neo4j_result.get('message', '未知错误')}")
|
||
except Exception as e:
|
||
error(f"Neo4j 三元组匹配失败: knowledge_base_id={kb_id}, 错误: {str(e)}")
|
||
continue
|
||
|
||
timings["triplet_matching"] = time.time() - start_triplet
|
||
debug(f"三元组匹配总耗时: {timings['triplet_matching']:.3f} 秒")
|
||
return all_triplets
|
||
|
||
async def generate_query_vector(
|
||
self,
|
||
request,
|
||
text: str,
|
||
service_params: Dict,
|
||
userid: str,
|
||
timings: Dict,
|
||
embedding_mode: int = 0
|
||
) -> List[float]:
|
||
"""生成查询向量（支持文本/多模态）"""
|
||
debug(f"生成查询向量: mode={embedding_mode}, text='{text[:100]}...'")
|
||
start_vector = time.time()
|
||
|
||
if embedding_mode == 0:
|
||
# === 模式 0：纯文本嵌入（BAAI/bge-m3）===
|
||
debug("使用 BAAI/bge-m3 文本嵌入")
|
||
vectors = await self.api_service.get_embeddings(
|
||
request=request,
|
||
texts=[text],
|
||
upappid=service_params['embedding'],
|
||
apiname="BAAI/bge-m3",
|
||
user=userid
|
||
)
|
||
if not vectors or not isinstance(vectors, list) or len(vectors) == 0:
|
||
raise ValueError("bge-m3 返回空结果")
|
||
query_vector = vectors[0]
|
||
if len(query_vector) != 1024:
|
||
raise ValueError(f"bge-m3 返回向量维度错误: {len(query_vector)}")
|
||
|
||
elif embedding_mode == 1:
|
||
# === 模式 1：多模态嵌入（black/clip）===
|
||
debug("使用 black/clip 多模态嵌入")
|
||
inputs = [{"type": "text", "content": text}]
|
||
|
||
result = await self.api_service.get_multi_embeddings(
|
||
request=request,
|
||
inputs=inputs,
|
||
upappid=service_params['embedding'],
|
||
apiname="black/clip",
|
||
user=userid
|
||
)
|
||
|
||
query_vector = None
|
||
for key, info in result.items():
|
||
if info.get("type") == "error":
|
||
debug(f"CLIP 返回错误跳过: {info['error']}")
|
||
continue
|
||
if "vector" in info and isinstance(info["vector"], list) and len(info["vector"]) == 1024:
|
||
query_vector = info["vector"]
|
||
debug(f"成功获取 CLIP 向量（来自 {info['type']}）")
|
||
break
|
||
|
||
if query_vector is None:
|
||
raise ValueError("black/clip 未返回任何有效 1024 维向量")
|
||
|
||
else:
|
||
raise ValueError(f"不支持的 embedding_mode: {embedding_mode}")
|
||
|
||
# 最终统一校验
|
||
if not isinstance(query_vector, list) or len(query_vector) != 1024:
|
||
raise ValueError(f"查询向量必须是长度为 1024 的浮点数列表，实际: {len(query_vector)}")
|
||
|
||
timings["vector_generation"] = time.time() - start_vector
|
||
debug(f"生成查询向量成功，耗时: {timings['vector_generation']:.3f} 秒，模式: {embedding_mode}")
|
||
return query_vector
|
||
|
||
async def generate_image_vector(
|
||
self,
|
||
request,
|
||
img_path: str,
|
||
service_params: Dict,
|
||
userid: str,
|
||
timings: Dict,
|
||
embedding_mode: int = 0
|
||
) -> List[float]:
|
||
"""生成查询向量（支持文本/多模态）"""
|
||
debug(f"生成查询向量: mode={embedding_mode}, image={img_path}")
|
||
start_vector = time.time()
|
||
|
||
if embedding_mode == 0:
|
||
raise ValueError(f"纯文本没有这个功能，请重新选择服务")
|
||
|
||
elif embedding_mode == 1:
|
||
# === 模式 1：多模态嵌入（black/clip）===
|
||
debug("使用 black/clip 多模态嵌入")
|
||
inputs = []
|
||
try:
|
||
ext = Path(img_path).suffix.lower()
|
||
if ext not in {".png", ".jpg", ".jpeg", ".webp", ".bmp"}:
|
||
ext = ".jpg"
|
||
|
||
mime_map = {
|
||
".png": "image/png",
|
||
".jpg": "image/jpeg",
|
||
".jpeg": "image/jpeg",
|
||
".webp": "image/webp",
|
||
".bmp": "image/bmp"
|
||
}
|
||
mime_type = mime_map.get(ext, "image/jpeg")
|
||
with open(img_path, "rb") as f:
|
||
b64 = base64.b64encode(f.read()).decode("utf-8")
|
||
data_uri = f"data:{mime_type};base64,{b64}"
|
||
|
||
inputs.append({
|
||
"type": "image",
|
||
"data": data_uri
|
||
})
|
||
debug(f"已添加图像（{mime_type}, {len(b64) / 1024:.1f}KB）: {Path(img_path).name}")
|
||
|
||
except Exception as e:
|
||
debug(f"图像处理失败，跳过: {img_path} → {e}")
|
||
|
||
result = await self.api_service.get_multi_embeddings(
|
||
request=request,
|
||
inputs=inputs,
|
||
upappid=service_params['embedding'],
|
||
apiname="black/clip",
|
||
user=userid
|
||
)
|
||
|
||
image_vector = None
|
||
for key, info in result.items():
|
||
if info.get("type") == "error":
|
||
debug(f"CLIP 返回错误跳过: {info['error']}")
|
||
continue
|
||
if "vector" in info and isinstance(info["vector"], list) and len(info["vector"]) == 1024:
|
||
image_vector = info["vector"]
|
||
debug(f"成功获取 CLIP 向量（来自 {info['type']}）")
|
||
break
|
||
|
||
if image_vector is None:
|
||
raise ValueError("black/clip 未返回任何有效 1024 维向量")
|
||
|
||
else:
|
||
raise ValueError(f"不支持的 embedding_mode: {embedding_mode}")
|
||
|
||
# 最终统一校验
|
||
if not isinstance(image_vector, list) or len(image_vector) != 1024:
|
||
raise ValueError(f"查询向量必须是长度为 1024 的浮点数列表，实际: {len(image_vector)}")
|
||
|
||
timings["vector_generation"] = time.time() - start_vector
|
||
debug(f"生成查询向量成功，耗时: {timings['vector_generation']:.3f} 秒，模式: {embedding_mode}")
|
||
return image_vector
|
||
|
||
async def vector_search(self, request, query_vector: List[float], orgid: str,
|
||
fiids: List[str], limit: int, service_params: Dict, userid: str,
|
||
timings: Dict) -> List[Dict]:
|
||
"""向量搜索"""
|
||
debug("开始向量搜索")
|
||
start_search = time.time()
|
||
result = await self.api_service.milvus_search_query(
|
||
request=request,
|
||
query_vector=query_vector,
|
||
userid=orgid,
|
||
knowledge_base_ids=fiids,
|
||
limit=limit,
|
||
offset=0,
|
||
upappid=service_params['vdb'],
|
||
apiname="mlvus/searchquery",
|
||
user=userid
|
||
)
|
||
|
||
if result.get("status") != "success":
|
||
raise ValueError(f"向量搜索失败: {result.get('message', '未知错误')}")
|
||
|
||
search_results = result.get("results", [])
|
||
timings["vector_search"] = time.time() - start_search
|
||
debug(f"向量搜索耗时: {timings['vector_search']:.3f} 秒")
|
||
debug(f"从向量数据中搜索到{len(search_results)}条数据")
|
||
return search_results
|
||
|
||
async def rerank_results(self, request, query: str, results: List[Dict], top_n: int,
|
||
service_params: Dict, userid: str, timings: Dict) -> List[Dict]:
|
||
"""重排序结果"""
|
||
debug("开始重排序")
|
||
start_rerank = time.time()
|
||
reranked_results = await self.api_service.rerank_results(
|
||
request=request,
|
||
query=query,
|
||
results=results,
|
||
top_n=top_n,
|
||
upappid=service_params['reranker'],
|
||
apiname="BAAI/bge-reranker-v2-m3",
|
||
user=userid
|
||
)
|
||
reranked_results = sorted(reranked_results, key=lambda x: x.get('rerank_score', 0), reverse=True)
|
||
timings["reranking"] = time.time() - start_rerank
|
||
debug(f"重排序耗时: {timings['reranking']:.3f} 秒")
|
||
debug(f"重排序分数分布: {[round(r.get('rerank_score', 0), 3) for r in reranked_results]}")
|
||
return reranked_results
|
||
|
||
def _deduplicate_triples(self, triples: List[Dict]) -> List[Dict]:
|
||
"""去重和优化三元组"""
|
||
unique_triples = []
|
||
seen = set()
|
||
|
||
for t in triples:
|
||
identifier = (t['head'].lower(), t['tail'].lower(), t['type'].lower())
|
||
if identifier not in seen:
|
||
seen.add(identifier)
|
||
unique_triples.append(t)
|
||
else:
|
||
# 如果发现更具体的类型，则替换
|
||
for existing in unique_triples:
|
||
if (existing['head'].lower() == t['head'].lower() and
|
||
existing['tail'].lower() == t['tail'].lower() and
|
||
len(t['type']) > len(existing['type'])):
|
||
unique_triples.remove(existing)
|
||
unique_triples.append(t)
|
||
debug(f"替换三元组为更具体类型: {t}")
|
||
break
|
||
|
||
return unique_triples
|
||
|
||
def format_search_results(self, results: List[Dict], limit: int) -> List[Dict]:
|
||
formatted = []
|
||
for res in results[:limit]:
|
||
# # 优先 rerank，其次用向量相似度（直接用，不要反）
|
||
# if res.get('rerank_score') is not None:
|
||
# score = res.get('rerank_score')
|
||
# else:
|
||
# score = res.get('distance', 0.0)
|
||
distance = res.get('distance', 0.0)
|
||
rerank_score = res.get('rerank_score', 0.0)
|
||
formatted.append({
|
||
"content": res.get('text', ''),
|
||
"title": res.get('metadata', {}).get('filename', 'Untitled'),
|
||
"metadata": {
|
||
"document_id": res.get('metadata', {}).get('document_id', ''),
|
||
"distance": distance,
|
||
"rerank_score": rerank_score,
|
||
}
|
||
})
|
||
return formatted
|
||
|
||
# async def save_uploaded_photo(self, image_file: FileStorage, orgid: str) -> str:
|
||
# """
|
||
# 把前端上传的图片保存到 /home/wangmeihua/kyrag/data/photo 目录下
|
||
# 返回保存后的绝对路径（字符串），供 generate_img_vector 使用
|
||
# """
|
||
# if not image_file or not hasattr(image_file, "filename"):
|
||
# raise ValueError("无效的图片上传对象")
|
||
#
|
||
# # 为了安全，按 orgid 分目录存放（避免不同公司文件混在一起）
|
||
# org_dir = UPLOAD_PHOTO_DIR / orgid
|
||
# org_dir.mkdir(parents=True, exist_ok=True)
|
||
#
|
||
# # 生成唯一文件名，保留原始后缀
|
||
# suffix = Path(image_file.filename).suffix.lower()
|
||
# if not suffix or suffix not in {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif"}:
|
||
# suffix = ".jpg"
|
||
#
|
||
# unique_name = f"{uuid.uuid4().hex}{suffix}"
|
||
# save_path = org_dir / unique_name
|
||
#
|
||
# # 真正落盘
|
||
# image_file.save(str(save_path))
|
||
# debug(f"图片已保存: {save_path} (原始名: {image_file.filename})")
|
||
#
|
||
# # 返回字符串路径，generate_img_vector 直接收 str 就行
|
||
# return str(save_path) |