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)