#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 整合上传功能 1. 上传PDF到MinIO并获取document 2. 处理TXT的chunk 3. 处理chunk中的图片链接:上传图片到MinIO,删除chunk中的图片链接 4. 上传chunk并获取chunk_id 5. 根据chunk_id更新Elasticsearch中的img_id """ import os import re from pathlib import Path import sys # 添加当前目录到Python路径 sys.path.append(os.path.dirname(os.path.abspath(__file__))) from ragflow_sdk import RAGFlow from minio import Minio from minio.error import S3Error from elasticsearch import Elasticsearch import chunk_operations from markdown_image2minio import get_minio_client, upload_file_to_minio # 配置 MINIO_BUCKET = "pdf-documents" ELASTICSEARCH_HOST = "192.168.107.165" ELASTICSEARCH_PORT = 1200 # RAGFlow配置 API_KEY = "ragflow-MyMjM2ODE2NThlMTExZjBiMzJlNzY5Mj" BASE_URL = "http://127.0.0.1:8099" class IntegratedUploader: def __init__(self, tenant_id, dataset_id): self.tenant_id = tenant_id self.dataset_id = dataset_id self.ragflow = RAGFlow(api_key=API_KEY, base_url=BASE_URL) self.minio_client = get_minio_client() self.es = Elasticsearch( [{'host': ELASTICSEARCH_HOST, 'port': ELASTICSEARCH_PORT, 'scheme': 'http'}], basic_auth=('elastic', 'infini_rag_flow') ) # 确保MinIO bucket存在 self._ensure_bucket_exists() def _ensure_bucket_exists(self): """确保MinIO bucket存在""" try: if not self.minio_client.bucket_exists(MINIO_BUCKET): self.minio_client.make_bucket(MINIO_BUCKET) print(f"Bucket '{MINIO_BUCKET}' created") except S3Error as e: print(f"MinIO bucket error: {e}") def upload_pdf_to_minio(self, pdf_path): """上传PDF到MinIO并返回URL""" if not os.path.exists(pdf_path): raise FileNotFoundError(f"PDF文件不存在: {pdf_path}") filename = os.path.basename(pdf_path) object_name = f"{self.dataset_id}/{filename}" try: upload_file_to_minio(self.minio_client, MINIO_BUCKET, object_name, pdf_path) url = f"http://127.0.0.1:9000/{MINIO_BUCKET}/{object_name}" print(f"PDF已上传到MinIO: {url}") return url except Exception as e: print(f"上传PDF失败: {e}") raise def get_or_create_document(self, pdf_path, display_name=None): """获取或创建文档对象""" if display_name is None: display_name = os.path.basename(pdf_path) try: datasets = self.ragflow.list_datasets() dataset = None for ds in datasets: if ds.id == self.dataset_id: dataset = ds break if not dataset: raise ValueError(f"数据集 {self.dataset_id} 不存在") # 检查文档是否已存在 try: documents = dataset.list_documents(name=display_name) if documents: print(f"文档已存在: {display_name}") return documents[0] except Exception: pass # 上传PDF文档 print(f"上传文档: {display_name}") with open(pdf_path, "rb") as f: blob = f.read() dataset.upload_documents([{"display_name": display_name, "blob": blob}]) documents = dataset.list_documents(name=display_name) return documents[0] except Exception as e: print(f"获取/创建文档失败: {e}") raise def extract_images_from_chunk(self, content): """从chunk内容中提取图片链接""" img_pattern = r'!\[.*?\]\((.*?)\)' return re.findall(img_pattern, content) def remove_images_from_content(self, content): """从内容中移除图片链接""" # 移除markdown图片语法 ![alt](url) content = re.sub(r'!\[.*?\]\(.*?\)', '', content) # 清理多余的空行 content = re.sub(r'\n\s*\n\s*\n', '\n\n', content) return content.strip() def upload_chunk_images(self, images, base_path, chunk_index): """上传chunk中的图片到MinIO""" uploaded_images = [] img_ids = [] for idx, img_path in enumerate(images): try: # 处理相对路径 if not os.path.isabs(img_path): img_abs_path = os.path.join(os.path.dirname(base_path), img_path) else: img_abs_path = img_path if not os.path.exists(img_abs_path): print(f"图片文件不存在: {img_abs_path}") continue # 生成图片ID和对象名称 img_id = f"{self.dataset_id}-{chunk_index:04d}-{idx:03d}" filename = os.path.basename(img_abs_path) ext = os.path.splitext(filename)[1] or ".jpg" object_name = f"{self.dataset_id}/images/chunk_{chunk_index:04d}/img_{idx:03d}{ext}" # 上传到MinIO upload_file_to_minio(self.minio_client, MINIO_BUCKET, object_name, img_abs_path) url = f"http://127.0.0.1:9000/{MINIO_BUCKET}/{object_name}" uploaded_images.append({ 'original_path': img_path, 'img_id': img_id, 'url': url, 'object_name': object_name }) img_ids.append(img_id) print(f"图片已上传: {img_path} -> {url}") except Exception as e: print(f"上传图片失败 {img_path}: {e}") continue return uploaded_images, img_ids def process_txt_chunks(self, txt_path, document): """处理TXT文件中的chunks""" if not os.path.exists(txt_path): raise FileNotFoundError(f"TXT文件不存在: {txt_path}") try: with open(txt_path, 'r', encoding='utf-8') as f: content = f.read() except Exception as e: print(f"读取TXT文件失败: {e}") raise # 按段落分割内容 chunks = content.split('\n\n') processed_chunks = [] for chunk_index, chunk_content in enumerate(chunks): if not chunk_content.strip(): continue print(f"\n处理第 {chunk_index + 1} 个chunk...") print(f"原始内容长度: {len(chunk_content)} 字符") # 提取图片 images = self.extract_images_from_chunk(chunk_content) print(f"找到 {len(images)} 张图片") # 上传图片并获取img_ids uploaded_images, img_ids = self.upload_chunk_images( images, txt_path, chunk_index ) # 移除图片链接后的内容 clean_content = self.remove_images_from_content(chunk_content) print(f"清理后内容长度: {len(clean_content)} 字符") if clean_content.strip(): try: # 上传chunk chunk = document.add_chunk(content=clean_content) chunk_id = chunk.id print(f"Chunk已添加,ID: {chunk_id}") # 如果有图片,更新Elasticsearch中的img_id if img_ids: img_id_str = ",".join(img_ids) result = chunk_operations.update_img_id_in_elasticsearch( self.tenant_id, document.id, chunk_id, img_id_str ) if result["code"] == 0: print(f"Elasticsearch img_id已更新: {img_id_str}") else: print(f"更新Elasticsearch失败: {result['message']}") processed_chunks.append({ 'chunk_id': chunk_id, 'content': clean_content, 'images': uploaded_images, 'img_ids': img_ids }) except Exception as e: print(f"添加chunk失败: {e}") continue else: print("跳过空chunk") return processed_chunks def upload_pdf_and_process_txt(self, pdf_path, txt_path): """完整的上传和处理流程""" print("=== 开始整合上传流程 ===") try: # 1. 上传PDF到MinIO pdf_url = self.upload_pdf_to_minio(pdf_path) print(f"PDF URL: {pdf_url}") # 2. 获取或创建文档 document = self.get_or_create_document(pdf_path) print(f"文档ID: {document.id}") # 3. 处理TXT chunks processed_chunks = self.process_txt_chunks(txt_path, document) print(f"处理完成,共 {len(processed_chunks)} 个chunks") return { 'success': True, 'document_id': document.id, 'pdf_url': pdf_url, 'chunks': processed_chunks } except Exception as e: print(f"整合上传失败: {e}") return { 'success': False, 'error': str(e) } def main(): """主函数 - 示例用法""" # 配置参数 tenant_id = "d669205e57a211f0b9e7324e7f243034" dataset_id = "10345832587311f0919f3a2728512a4b" # 文件路径 pdf_path = r"G:\11\22\规范\example.pdf" txt_path = r"G:\11\22\规范\example.txt" # 创建上传器实例 uploader = IntegratedUploader(tenant_id, dataset_id) # 执行上传和处理 result = uploader.upload_pdf_and_process_txt(pdf_path, txt_path) if result['success']: print("\n=== 上传成功 ===") print(f"文档ID: {result['document_id']}") print(f"PDF URL: {result['pdf_url']}") print(f"处理chunks数: {len(result['chunks'])}") for chunk in result['chunks']: print(f" - Chunk ID: {chunk['chunk_id']}, 图片数: {len(chunk['images'])}") else: print(f"\n上传失败: {result['error']}") if __name__ == "__main__": main()