jk-rag-python/rag/db copy.py

from rag.vector_db.milvus_vector import HybridRetriever
from response_info import generate_message, generate_response
from utils.get_logger import setup_logger
from datetime import datetime
from uuid import uuid1
import mysql.connector
from mysql.connector import pooling, Error, errors
import threading
from concurrent.futures import ThreadPoolExecutor, TimeoutError
from config import milvus_uri, mysql_config

logger = setup_logger(__name__)
# uri = "http://localhost:19530"

# if 'POOL' not in globals():
#     try:
#         POOL = pooling.MySQLConnectionPool(
#             pool_name="mysql_pool",
#             pool_size=10,
#             **mysql_config
#         )
#         logger.info("MySQL 连接池初始化成功")
#     except Error as e:
#         logger.info(f"初始化 MySQL 连接池失败: {e}")
#         POOL = None


class MilvusOperate:

    def __init__(self, collection_name: str = "default", embedding_name:str = "e5"):
        self.collection = collection_name
        self.hybrid_retriever = HybridRetriever(uri=milvus_uri, embedding_name=embedding_name, collection_name=collection_name)
        self.mysql_client = MysqlOperate()

    def _has_collection(self):
        is_collection = self.hybrid_retriever.has_collection()
        return is_collection
    
    def _create_collection(self):
        if self._has_collection():
            resp = {"code": 400, "message": "数据库已存在"}
        else:
            create_result = self.hybrid_retriever.build_collection()
            resp = generate_message(create_result)
        return resp

    
    def _delete_collection(self):
        delete_result = self.hybrid_retriever.delete_collection(self.collection)
        resp = generate_message(delete_result)
        return resp

    
    def _put_by_id(self, slice_json):
        slice_id = slice_json.get("slice_id", None)
        slice_text = slice_json.get("slice_text", None)
        update_result, chunk_len = self.hybrid_retriever.update_data(chunk_id=slice_id, chunk=slice_text)
        if update_result.endswith("success"):
            # 如果成功，更新mysql中知识库总长度和文档长度
            update_json = {}
            update_json["knowledge_id"] = slice_json.get("knowledge_id")
            update_json["doc_id"] = slice_json.get("document_id")
            update_json["chunk_len"] = chunk_len
            update_json["operate"] = "update"
            update_json["chunk_id"] = slice_id
            update_json["chunk_text"] = slice_text
            update_flag, update_str = self.mysql_client.update_total_doc_len(update_json)
        else:
            update_flag = False
            
        if not update_flag:
            update_result = "update_error"    
        
        resp = generate_message(update_result)
        return resp
    
    def _insert_slice(self, slice_json):
        slice_id = str(uuid1())
        knowledge_id = slice_json.get("knowledge_id")
        doc_id = slice_json.get("document_id")
        slice_text = slice_json.get("slice_text", None)
        doc_name = slice_json.get("doc_name")
        chunk_len = len(slice_text)
        metadata = {
            "content": slice_text,
            "doc_id": doc_id,
            "chunk_id": slice_id,
            "metadata": {"source": doc_name, "chunk_len": chunk_len},
            "Chapter": slice_json.get("Chapter",""),
            "Father_Chapter": slice_json.get("Father_Chapter",""),
        }
        insert_flag, insert_str = self.hybrid_retriever.insert_data(slice_text, metadata)
        if insert_flag:
            # 如果成功，更新mysql中知识库总长度和文档长度
            update_json = {}
            update_json["knowledge_id"] = slice_json.get("knowledge_id")
            update_json["doc_id"] = slice_json.get("document_id")
            update_json["chunk_len"] = chunk_len
            update_json["operate"] = "insert"
            update_json["chunk_id"] = slice_id
            update_json["chunk_text"] = slice_text
            update_json["slice_index"] = slice_json.get("slice_index")
            update_flag, update_str = self.mysql_client.update_total_doc_len(update_json)
        else:
            logger.error(f"插入向量库出错：{insert_str}")
            update_flag = False
            update_str = "向量库写入出错"
            # pass
            
        if not update_flag:
            logger.error(f"新增切片中mysql数据库出错：{update_str}")
            # insert_result = "insert_error"
            success = "insert_error"
        else:
            # insert_result = "insert_success"
            success = "insert_success"
        
        # resp = generate_message(insert_result)
        resp = {"status": success, "slice_id":slice_id}
        return resp
    

    def _delete_by_chunk_id(self, chunk_id, knowledge_id, document_id):
        logger.info(f"删除的切片id：{chunk_id}")
        delete_result, delete_chunk_len = self.hybrid_retriever.delete_by_chunk_id(chunk_id=chunk_id)
        if delete_result.endswith("success"):
            chunk_len = delete_chunk_len[0]
            update_json = {
                "knowledge_id": knowledge_id,
                "doc_id": document_id,
                "chunk_len": -chunk_len,
                "operate": "delete",
                "chunk_id": chunk_id
            }
            update_flag, update_str = self.mysql_client.update_total_doc_len(update_json)
        else:
            logger.error("根据chunk id删除向量库失败")
            update_flag = False
            update_str = "根据chunk id删除失败"
        
        if not update_flag:
            logger.error(update_str)
            delete_result = "delete_error"

        
        resp = generate_message(delete_result)
        return resp
    
    def _delete_by_doc_id(self, doc_id: str = None):
        logger.info(f"删除数据的id：{doc_id}")
        delete_result = self.hybrid_retriever.delete_by_doc_id(doc_id=doc_id)
        resp = generate_message(delete_result)
        return resp
    

    def _search_by_chunk_id(self, chunk_id):
        if self._has_collection():
            query_result = self.hybrid_retriever.query_chunk_id(chunk_id=chunk_id)
        else:
            query_result = []
        logger.info(f"根据切片查询到的信息：{query_result}")
        resp = generate_response(query_result)

        return resp
    
    def _search_by_chunk_id_list(self, chunk_id_list):
        if self._has_collection():
            query_result = self.hybrid_retriever.query_chunk_id_list(chunk_id_list)
        else:
            query_result = []
        logger.info(f"召回的切片列表查询切片信息：{query_result}")
        
        chunk_content_list = []
        for chunk_dict in query_result:
            chunk_content = chunk_dict.get("content")
            chunk_content_list.append(chunk_content)

        return chunk_content_list
    
    
    def _search_by_key_word(self, search_json):
        if self._has_collection():
            doc_id = search_json.get("document_id", None)
            text = search_json.get("text", None)
            page_num = search_json.get("pageNum", 1)
            page_size = search_json.get("pageSize", 10)
            page_num = search_json.get("pageNum")  # 根据传过来的id处理对应知识库
            query_result = self.hybrid_retriever.query_filter(doc_id=doc_id, filter_field=text)
        else:
            query_result = []
        resp = generate_response(query_result,page_num,page_size)

        return resp
    
    def _insert_data(self, docs):
        insert_flag = ""
        insert_info = ""
        for doc in docs:
            chunk = doc.get("content")
            insert_flag, insert_info = self.hybrid_retriever.insert_data(chunk, doc)
            if not insert_flag:
                break
        resp = insert_flag if insert_flag else "insert_error"
        return resp, insert_info
    
    def _batch_insert_data(self, docs, text_lists):
        insert_flag, insert_info = self.hybrid_retriever.batch_insert_data(text_lists, docs)

        resp = insert_flag
        return resp, insert_info

    def _search(self, query, k, mode):
        search_result = self.hybrid_retriever.search(query, k, mode)
        return search_result
    
    def _query_by_scalar_field(self, doc_id: str, field_name: str, field_value: str):
        """
        根据标量字段查询数据
        
        参数:
            doc_id: 文档ID
            field_name: 字段名（如 Father_Chapter）
            field_value: 字段值
        返回:
            查询结果列表
        """
        return self.hybrid_retriever.query_by_scalar_field(doc_id, field_name, field_value)
    
    def _copy_docs_to_new_collection(self, new_collection_name, doc_ids, embedding_name="e5"):
        """
        将指定的文档数据复制到新集合或现有集合
        使用雪花算法生成新的 doc_id 和 chunk_id
        
        参数:
            new_collection_name: 目标集合名称（也是新的 knowledge_id）
            doc_ids: 要复制的文档ID列表
            embedding_name: 向量模型名称
        
        返回:
            响应字典，包含 doc_id_mapping 映射关系
        """
        try:
            # 1. 从源集合查询数据
            logger.info(f"从集合 {self.collection} 查询文档: {doc_ids}")
            query_results = self.hybrid_retriever.query_by_doc_ids(doc_ids)
            
            if not query_results:
                return {"code": 404, "message": "未找到匹配的文档数据", "doc_id_mapping": {}, "chunk_id_mapping": {}}
            
            logger.info(f"查询到 {len(query_results)} 条数据")
            
            # 2. 检查目标集合是否存在，不存在则创建
            target_milvus_client = MilvusOperate(collection_name=new_collection_name, embedding_name=embedding_name)
            
            collection_exists = target_milvus_client._has_collection()
            if not collection_exists:
                logger.info(f"创建新集合: {new_collection_name}")
                create_result = target_milvus_client._create_collection()
                if create_result.get("code") != 200:
                    create_result["doc_id_mapping"] = {}
                    return create_result
            else:
                logger.info(f"集合 {new_collection_name} 已存在，直接插入数据")
            
            # 3. 为每个源 doc_id 生成新的 doc_id（使用雪花算法）
            doc_id_mapping = {}  # {old_doc_id: new_doc_id}
            for old_doc_id in doc_ids:
                doc_id_mapping[old_doc_id] = generate_snowflake_id()
            
            # 4. 准备插入数据（移除pk字段，使用新的 doc_id 和 chunk_id）
            insert_data = []
            chunk_id_mapping = {}  # {old_chunk_id: new_chunk_id}
            for item in query_results:
                old_doc_id = item.get("doc_id")
                old_chunk_id = item.get("chunk_id")
                new_doc_id = doc_id_mapping.get(old_doc_id, generate_snowflake_id())
                new_chunk_id = generate_snowflake_id()  # 每个切片生成新的 chunk_id
                chunk_id_mapping[old_chunk_id] = new_chunk_id  # 记录映射
                
                new_item = {
                    "content": item.get("content"),
                    "dense_vector": item.get("dense_vector"),
                    "doc_id": new_doc_id,
                    "chunk_id": new_chunk_id,
                    "Father_Chapter": item.get("Father_Chapter"),
                    "Chapter": item.get("Chapter"),
                    "metadata": item.get("metadata")
                }
                insert_data.append(new_item)
            
            # 5. 批量插入数据到目标集合
            logger.info(f"开始向集合 {new_collection_name} 插入 {len(insert_data)} 条数据")
            try:
                target_milvus_client.hybrid_retriever.client.insert(
                    collection_name=new_collection_name,
                    data=insert_data
                )
                logger.info(f"成功向集合插入数据")
                
                return {
                    "code": 200,
                    "message": "复制成功",
                    "data": {
                        "source_collection": self.collection,
                        "target_collection": new_collection_name,
                        "doc_ids": doc_ids,
                        "total_records": len(insert_data),
                        "collection_existed": collection_exists
                    },
                    "doc_id_mapping": doc_id_mapping,
                    "chunk_id_mapping": chunk_id_mapping
                }
            except Exception as e:
                logger.error(f"插入数据到集合失败: {e}")
                # 插入失败且是新创建的集合时删除
                if not collection_exists:
                    target_milvus_client._delete_collection()
                return {"code": 500, "message": f"插入数据失败: {str(e)}", "doc_id_mapping": {}, "chunk_id_mapping": {}}
                
        except Exception as e:
            logger.error(f"复制文档到新集合失败: {e}")
            return {"code": 500, "message": f"复制失败: {str(e)}", "doc_id_mapping": {}, "chunk_id_mapping": {}}


# class MysqlOperate:

#     def get_connection(self):
#         """
#         从连接池中获取一个连接
#         :return: 数据库连接对象
#         """
#         # try:
#             # with ThreadPoolExecutor() as executor:
#             #     future = executor.submit(POOL.get_connection)
#             #     connection = future.result(timeout=5.0)  # 设置超时时间为5秒
#             # logger.info("成功从连接池获取连接")
#             # return connection, "success"
#         # except TimeoutError:
#         #     logger.error("获取mysql数据库连接池超时")
#         #     return None, "mysql获取连接池超时"
#         # except errors.InterfaceError as e:
#         #     logger.error(f"MySQL 接口异常：{e}")
#         #     return None, "mysql接口异常"
#         # except errors.OperationalError as e:
#         #     logger.error(f"MySQL 操作错误：{e}")
#         #     return None, "mysql 操作错误"
#         # except Error as e:
#         #     logger.error(f"无法从连接池获取连接: {e}")
#         #     return None, str(e)
#         connection = None
#         event = threading.Event()

#         def target():
#             nonlocal connection
#             try:
#                 connection = POOL.get_connection()
#             finally:
#                 event.set()

#         thread = threading.Thread(target=target)
#         thread.start()
#         event.wait(timeout=5)

#         if thread.is_alive():
#             # 超时处理
#             logger.error("获取连接超时")
#             return None, "获取连接超时"
#         else:
#             if connection:
#                 return connection, "success"
#             else:
#                 logger.error("获取连接失败")
#                 return None, "获取连接失败"

#     def insert_to_slice(self, docs, knowledge_id, doc_id):
#         """
#         插入数据到切片信息表中 slice_info
#         """
#         connection = None
#         cursor = None
#         date_now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
#         values = []
#         connection, cennction_info = self.get_connection()
#         if not connection:
#             return False, cennction_info
        
#         for chunk in docs:
#             slice_id = chunk.get("chunk_id")
#             slice_text = chunk.get("content")
#             chunk_index = chunk.get("metadata").get("chunk_index")
#             values.append((slice_id, knowledge_id, doc_id, slice_text, date_now, chunk_index))
#         try:
#             cursor = connection.cursor()
#             # insert_sql = """
#             #     INSERT INTO slice_info (
#             #         slice_id,
#             #         knowledge_id,
#             #         document_id,
#             #         slice_text,
#             #         create_time,
#             #         slice_index
#             #     ) VALUES (%s, %s, %s, %s, %s,%s)
#             #     """
            
#             # 容错“for key 'UK_ID_TYPE_KEY'”
#             insert_sql = """
#             INSERT INTO slice_info (
#                 slice_id,
#                 knowledge_id,
#                 document_id,
#                 slice_text,
#                 create_time,
#                 slice_index
#             ) VALUES (%s, %s, %s, %s, %s, %s)
#             ON DUPLICATE KEY UPDATE
#                 slice_text = VALUES(slice_text),
#                 create_time = VALUES(create_time),
#                 slice_index = VALUES(slice_index)
#             """
            
#             cursor.executemany(insert_sql, values)
#             connection.commit()
#             logger.info(f"批量插入切片数据成功。")
#             return True, "success"

#         except Error as e:
#             logger.error(f"数据库操作出错：{e}")
#             connection.rollback()
#             return False, str(e)
#         finally:
#             # if cursor:
#             cursor.close()
#             # if connection and connection.is_connected():
#             connection.close()

#     def delete_to_slice(self, doc_id):
#         """
#         删除 slice_info库中切片信息
#         """
#         connection = None
#         cursor = None
#         connection, connection_info = self.get_connection()
#         if not connection:
#             return False, connection_info
#         try:
#             cursor = connection.cursor()
#             delete_sql = f"DELETE FROM slice_info WHERE document_id = %s"
#             cursor.execute(delete_sql, (doc_id,))
#             connection.commit()
#             logger.info(f"删除数据成功")
#             return True, "success"

#         except Error as e:
#             logger.error(f"根据{doc_id}删除数据失败：{e}")
#             connection.rollback()
#             return False, str(e)
#         finally:
#             # if cursor:
#             cursor.close()
#             # if connection and connection.is_connected():
#             connection.close()

#     def insert_to_image_url(self, image_dict, knowledge_id, doc_id):
#         """
#         批量插入数据到指定表
#         """
#         connection = None
#         cursor = None
#         connection, connection_info = self.get_connection()
#         if not connection:
#             return False, connection_info
        
#         date_now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
#         values = []
#         for img_key, img_value in image_dict.items():
#             origin_text = img_key
#             media_url = img_value
#             values.append((knowledge_id, doc_id, origin_text, "image", media_url, date_now))
#         try:
#             cursor = connection.cursor()
#             # insert_sql = """
#             #     INSERT INTO bm_media_replacement (
#             #         knowledge_id,
#             #         document_id,
#             #         origin_text,
#             #         media_type,
#             #         media_url,
#             #         create_time
#             #     ) VALUES (%s, %s, %s, %s, %s, %s)
#             #     """
            
#             # 容错“for key 'UK_ID_TYPE_KEY'”
#             insert_sql = """
#                 INSERT INTO bm_media_replacement (
#                     knowledge_id,
#                     document_id,
#                     origin_text,
#                     media_type,
#                     media_url,
#                     create_time
#                 ) VALUES (%s, %s, %s, %s, %s, %s)
#                 ON DUPLICATE KEY UPDATE
#                     origin_text = VALUES(origin_text),
#                     media_type = VALUES(media_type),
#                     media_url = VALUES(media_url),
#                     create_time = VALUES(create_time)
#                 """

#             cursor.executemany(insert_sql, values)
#             connection.commit()
#             logger.info(f"插入到bm_media_replacement表成功")
#             return True, "success"
#         except Error as e:
#             logger.error(f"数据库操作出错：{e}")
#             connection.rollback()
#             return False, str(e)
#         finally:
#             # if cursor:
#             cursor.close()
#             # if connection and connection.is_connected():
#             connection.close()

#     def delete_image_url(self, doc_id):
#         """
#         根据doc id删除bm_media_replacement中的数据
#         """
#         connection = None
#         cursor = None
#         connection, connection_info = self.get_connection()
#         if not connection:
#             return False, connection_info
        
#         try:
#             cursor = connection.cursor()
#             delete_sql = f"DELETE FROM bm_media_replacement WHERE document_id = %s"
#             cursor.execute(delete_sql, (doc_id,))
#             connection.commit()
#             logger.info(f"根据{doc_id} 删除bm_media_replacement表中数据成功")
#             return True, "success"
#         except Error as e:
#             logger.error(f"根据{doc_id}删除 bm_media_replacement 数据库操作出错：{e}")
#             connection.rollback()
#             return False, str(e)
#         finally:
#             # if cursor:
#             cursor.close()
#             # if connection and connection.is_connected():
#             connection.close()

#     def update_total_doc_len(self, update_json):
#         """
#         更新长度表和文档长度表，删除slice info表, 插入slice info 切片信息
#         """
#         knowledge_id = update_json.get("knowledge_id")
#         doc_id = update_json.get("doc_id")
#         chunk_len = update_json.get("chunk_len")
#         operate = update_json.get("operate")
#         chunk_id = update_json.get("chunk_id")
#         chunk_text = update_json.get("chunk_text")
#         connection = None
#         cursor = None
#         connection, connection_info = self.get_connection()
#         if not connection:
#             return False, connection_info
#         try:
#             cursor = connection.cursor()
#             query_doc_word_num_sql = f"select word_num,slice_total from bm_document where document_id = %s"
#             query_knowledge_word_num_sql = f"select word_num from bm_knowledge where knowledge_id = %s"
#             cursor.execute(query_doc_word_num_sql, (doc_id,))
#             doc_result = cursor.fetchone()
#             logger.info(f"查询到的文档长度信息：{doc_result}")
#             cursor.execute(query_knowledge_word_num_sql, (knowledge_id, ))
#             knowledge_result = cursor.fetchone()
#             logger.info(f"查询到的知识库总长度信息：{knowledge_result}")
#             if not doc_result:
#                 new_word_num = 0
#                 slice_total = 0
#             else:
#                 old_word_num = doc_result[0]
#                 slice_total = doc_result[1]
#                 new_word_num = old_word_num + chunk_len
#                 slice_total -= 1 if slice_total else 0

#             if not knowledge_result:
#                 new_knowledge_word_num = 0
#             else:
#                 old_knowledge_word_num = knowledge_result[0]
#                 new_knowledge_word_num = old_knowledge_word_num + chunk_len

#             if operate == "update":
#                 update_sql = f"UPDATE bm_document SET word_num = %s WHERE document_id = %s"
#                 cursor.execute(update_sql, (new_word_num, doc_id))

#                 date_now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
#                 update_slice_sql = f"UPDATE slice_info SET slice_text = %s, update_time = %s WHERE slice_id = %s"
#                 cursor.execute(update_slice_sql, (chunk_text, date_now, chunk_id))
#             elif operate == "insert":
#                 query_slice_info_index_sql = f"select MAX(slice_index) from slice_info where document_id = %s"
#                 cursor.execute(query_slice_info_index_sql, (doc_id,))
#                 chunk_index_result = cursor.fetchone()[0]
#                 # logger.info(chunk_index_result)
#                 if chunk_index_result:
#                     chunk_max_index = int(chunk_index_result)
#                 else:
#                     chunk_max_index = 0

#                 update_sql = f"UPDATE bm_document SET word_num = %s WHERE document_id = %s"
#                 cursor.execute(update_sql, (new_word_num, doc_id))

#                 date_now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
#                 insert_slice_sql = "INSERT INTO slice_info (slice_id,knowledge_id,document_id,slice_text,create_time, slice_index) VALUES (%s, %s, %s, %s, %s, %s)"
#                 cursor.execute(insert_slice_sql, (chunk_id, knowledge_id, doc_id, chunk_text, date_now, chunk_max_index+1))
#             else:
#                 update_sql = f"UPDATE bm_document SET word_num = %s, slice_total = %s WHERE document_id = %s"
#                 cursor.execute(update_sql, (new_word_num, slice_total, doc_id))

#                 # 删除切片id对应的切片  
#                 delete_slice_sql = f"DELETE FROM slice_info where slice_id = %s"
#                 cursor.execute(delete_slice_sql, (chunk_id, ))

#             update_knowledge_sql = f"UPDATE bm_knowledge SET word_num = %s WHERE knowledge_id = %s"
#             cursor.execute(update_knowledge_sql, (new_knowledge_word_num, knowledge_id))

#             connection.commit()
#             logger.info("bm_document和bm_knowledge数据更新成功")
#             return True, "success"
#         except Error as e:
#             logger.error(f"数据库操作出错：{e}")
#             connection.rollback()
#             return False, str(e)
#         finally:
#             # if cursor:
#             cursor.close()
#             # if connection and connection.is_connected():
#             connection.close()


TABLE_NAME = "bm_document"
STATUS_FIELD = "update_by"
TASK_ID_FIELD = "document_id"
USER_ID_FIELD = "remark"


import time
# ========== 全局初始化连接池（自动检测 + 超时保护） ==========
class SafeMySQLPool:
    def __init__(self, pool_size=10, conn_timeout=10, idle_timeout=60, **mysql_config):
        mysql_config.setdefault("connect_timeout", conn_timeout)
        mysql_config.setdefault("pool_reset_session", True)
        self._pool = pooling.MySQLConnectionPool(
            pool_name="safe_mysql_pool",
            pool_size=pool_size,
            **mysql_config
        )
        # 使用 RLock 可重入锁，避免 _auto_reclaimer 调用 close 时死锁
        self._lock = threading.RLock()
        self._active_conns = {}  # {id(conn): (conn, last_used_time)}
        self._idle_timeout = idle_timeout
        self._stop_event = threading.Event()
        threading.Thread(target=self._auto_reclaimer, daemon=True).start()

    def get_connection(self, timeout=10):
        """安全获取连接（带超时检测与追踪）"""
        start = time.time()
        while True:
            try:
                conn = self._pool.get_connection()
                conn.ping(reconnect=True, attempts=3, delay=2)
                with self._lock:
                    self._active_conns[id(conn)] = (conn, time.time())
                return self._wrap_connection(conn)
            except errors.PoolError:
                if time.time() - start > timeout:
                    raise TimeoutError(f"获取 MySQL 连接超时（超过 {timeout}s）")
                time.sleep(0.3)

    def _wrap_connection(self, conn):
        """包装连接对象以监控关闭事件"""
        pool = self
        orig_close = conn.close

        def safe_close():
            try:
                orig_close()
            finally:
                with pool._lock:
                    pool._active_conns.pop(id(conn), None)
        conn.close = safe_close
        return conn

    def _auto_reclaimer(self):
        """后台线程自动回收超时未关闭连接"""
        while not self._stop_event.is_set():
            time.sleep(5)
            now = time.time()
            with self._lock:
                to_remove = []
                for cid, (conn, last_used) in list(self._active_conns.items()):
                    if now - last_used > self._idle_timeout:
                        try:
                            conn.close()
                            logger.warning(f"[回收] 已回收超时未关闭连接 (idle={int(now - last_used)}s)")
                        except Exception as e:
                            logger.error(f"[回收] 回收连接失败: {e}")
                        to_remove.append(cid)
                for cid in to_remove:
                    self._active_conns.pop(cid, None)

    def close_all(self):
        """停止守护线程并关闭所有连接"""
        self._stop_event.set()
        with self._lock:
            for conn, _ in self._active_conns.values():
                try:
                    conn.close()
                except:
                    pass
            self._active_conns.clear()


# ========== 初始化连接池 ==========
if "POOL" not in globals():
    try:
        POOL = SafeMySQLPool(pool_size=10, idle_timeout=60, **mysql_config)
        logger.info("MySQL 连接池初始化成功")
    except Error as e:
        logger.error(f"MySQL 连接池初始化失败: {e}")
        POOL = None



"""
雪花算法获取唯一id
"""

import time
import threading

class Snowflake:
    def __init__(self, datacenter_id=1, worker_id=1):
        self.worker_id_bits = 5
        self.datacenter_id_bits = 5
        self.sequence_bits = 12

        self.max_worker_id = -1 ^ (-1 << self.worker_id_bits)
        self.max_datacenter_id = -1 ^ (-1 << self.datacenter_id_bits)

        self.worker_id = worker_id
        self.datacenter_id = datacenter_id
        self.sequence = 0

        self.worker_id_shift = self.sequence_bits
        self.datacenter_id_shift = self.sequence_bits + self.worker_id_bits
        self.timestamp_left_shift = self.sequence_bits + self.worker_id_bits + self.datacenter_id_bits

        self.twepoch = 1288834974657
        self.last_timestamp = -1

        self.lock = threading.Lock()

    def _timestamp(self):
        return int(time.time() * 1000)

    def _wait_next_ms(self, last):
        ts = self._timestamp()
        while ts <= last:
            ts = self._timestamp()
        return ts

    def generate_id(self):
        with self.lock:
            timestamp = self._timestamp()

            if timestamp < self.last_timestamp:
                raise Exception("Clock moved backwards, refusing to generate id")

            if timestamp == self.last_timestamp:
                self.sequence = (self.sequence + 1) & ((1 << self.sequence_bits) - 1)
                if self.sequence == 0:
                    timestamp = self._wait_next_ms(timestamp)
            else:
                self.sequence = 0

            self.last_timestamp = timestamp

            snowflake_id = (
                ((timestamp - self.twepoch) << self.timestamp_left_shift) |
                (self.datacenter_id << self.datacenter_id_shift) |
                (self.worker_id << self.worker_id_shift) |
                self.sequence
            )

            # 转成固定 20 位数字
            return str(snowflake_id).zfill(20)


# ========== 全局雪花算法 ID 生成器 ==========
snowflake_id_generator = Snowflake(datacenter_id=1, worker_id=1)

def generate_snowflake_id():
    """
    全局方法：使用雪花算法生成唯一ID
    返回20位数字字符串
    """
    return snowflake_id_generator.generate_id()


# ========== MysqlOperate 类 ==========
class MysqlOperate:
    def get_connection(self):
        """安全获取连接"""
        if not POOL:
            return None, "连接池未初始化"
        try:
            connection = POOL.get_connection(timeout=5)
            return connection, "success"
        except TimeoutError as e:
            logger.error(str(e))
            return None, "获取连接超时"
        except Error as e:
            logger.error(f"MySQL 获取连接失败: {e}")
            return None, str(e)

    def _execute_many(self, sql, values, success_msg, err_msg):
        """通用批量执行模板"""
        connection, info = self.get_connection()
        if not connection:
            return False, info
        cursor = None
        try:
            cursor = connection.cursor()
            cursor.executemany(sql, values)
            connection.commit()
            logger.info(success_msg)
            return True, "success"
        except Error as e:
            connection.rollback()
            logger.error(f"{err_msg}: {e}")
            return False, str(e)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()

    def insert_to_slice(self, docs, knowledge_id, doc_id):
        """批量插入切片信息（同时存储 slice_text 和 old_slice_text）"""
        date_now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        values = [
            (
                chunk.get("chunk_id"),
                knowledge_id,
                doc_id,
                chunk.get("content"),
                chunk.get("content"),  # old_slice_text 存储原始文本
                date_now,
                chunk.get("metadata", {}).get("chunk_index"),
                chunk.get("Chapter"),
                chunk.get("Father_Chapter", "")
            )
            for chunk in docs
        ]
 
        sql = """
        INSERT INTO slice_info (
            slice_id, knowledge_id, document_id, slice_text, old_slice_text, create_time, slice_index, section, parent_section
        ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
        ON DUPLICATE KEY UPDATE
           slice_text = VALUES(slice_text),
           create_time = VALUES(create_time),
           slice_index = VALUES(slice_index),
           section = VALUES(section),
           parent_section = VALUES(parent_section)
        """
        return self._execute_many(sql, values, "批量插入切片数据成功", "插入 slice_info 出错")

    def delete_to_slice(self, doc_id):
        """删除切片"""
        connection, info = self.get_connection()
        if not connection:
            return False, info
        cursor = None
        try:
            cursor = connection.cursor()
            cursor.execute("DELETE FROM slice_info WHERE document_id = %s", (doc_id,))
            connection.commit()
            logger.info(f"删除 slice_info 数据成功")
            return True, "success"
        except Error as e:
            connection.rollback()
            logger.error(f"删除 slice_info 出错: {e}")
            return False, str(e)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()

    def insert_to_image_url(self, image_dict, knowledge_id, doc_id):
        """插入图片映射表"""
        date_now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        values = [
            (knowledge_id, doc_id, k, "image", v, date_now)
            for k, v in image_dict.items()
        ]
        sql = """
        INSERT INTO bm_media_replacement (
            knowledge_id, document_id, origin_text, media_type, media_url, create_time
        ) VALUES (%s, %s, %s, %s, %s, %s)
        ON DUPLICATE KEY UPDATE
            origin_text = VALUES(origin_text),
            media_type = VALUES(media_type),
            media_url = VALUES(media_url),
            create_time = VALUES(create_time)
        """
        return self._execute_many(sql, values, "插入 bm_media_replacement 成功", "插入 bm_media_replacement 出错")

    def delete_image_url(self, doc_id):
        """删除图片映射"""
        connection, info = self.get_connection()
        if not connection:
            return False, info
        cursor = None
        try:
            cursor = connection.cursor()
            cursor.execute("DELETE FROM bm_media_replacement WHERE document_id = %s", (doc_id,))
            connection.commit()
            logger.info(f"删除 bm_media_replacement 成功")
            return True, "success"
        except Error as e:
            connection.rollback()
            logger.error(f"删除 bm_media_replacement 出错: {e}")
            return False, str(e)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()

    def update_task_status_start(self, task_id):
        """更新任务状态为开始(1)"""
        connection, info = self.get_connection()
        if not connection:
            return False, info
        cursor = None
        try:
            cursor = connection.cursor()
            sql = f"UPDATE {TABLE_NAME} SET {STATUS_FIELD} = %s WHERE {TASK_ID_FIELD} = %s"
            cursor.execute(sql, (1, task_id))
            connection.commit()
            logger.info(f"任务 {task_id} 状态更新为开始(1)")
            return True, "success"
        except Error as e:
            connection.rollback()
            logger.error(f"更新任务状态为开始失败: {e}")
            return False, str(e)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()
    
    def update_task_status_complete(self, task_id, user_id):
        """更新任务状态为完成(2)并更新用户ID"""
        connection, info = self.get_connection()
        if not connection:
            return False, info
        cursor = None
        try:
            cursor = connection.cursor()
            sql = f"UPDATE {TABLE_NAME} SET {STATUS_FIELD} = %s, {USER_ID_FIELD} = %s WHERE {TASK_ID_FIELD} = %s"
            cursor.execute(sql, (2, user_id, task_id))
            connection.commit()
            logger.info(f"任务 {task_id} 状态更新为完成(2)，用户ID: {user_id}")
            return True, "success"
        except Error as e:
            connection.rollback()
            logger.error(f"更新任务状态为完成失败: {e}")
            return False, str(e)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()
    
    def update_task_status_error(self, task_id):
        """更新任务状态为错误(0)"""
        connection, info = self.get_connection()
        if not connection:
            return False, info
        cursor = None
        try:
            cursor = connection.cursor()
            sql = f"UPDATE {TABLE_NAME} SET {STATUS_FIELD} = %s WHERE {TASK_ID_FIELD} = %s"
            cursor.execute(sql, (0, task_id))
            connection.commit()
            logger.info(f"任务 {task_id} 状态更新为错误(0)")
            return True, "success"
        except Error as e:
            connection.rollback()
            logger.error(f"更新任务状态为错误失败: {e}")
            return False, str(e)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()
    
    def delete_document(self, task_id):
        """删除 bm_document 表中的记录（取消任务前清理）"""
        connection, info = self.get_connection()
        if not connection:
            return False, info
        cursor = None
        try:
            cursor = connection.cursor()
            sql = f"DELETE FROM {TABLE_NAME} WHERE {TASK_ID_FIELD} = %s"
            cursor.execute(sql, (task_id,))
            affected_rows = cursor.rowcount
            connection.commit()
            logger.info(f"删除 bm_document 记录成功: task_id={task_id}, 影响行数={affected_rows}")
            return True, affected_rows
        except Error as e:
            connection.rollback()
            logger.error(f"删除 bm_document 记录失败: {e}")
            return False, str(e)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()
    
    def query_parent_generation_enabled(self, doc_ids: list):
        """
        查询文档的 parent_generation_enabled 字段
        返回 parent_generation_enabled=1 的 doc_id 集合
        """
        if not doc_ids:
            return set()
        connection, info = self.get_connection()
        if not connection:
            return set()
        cursor = None
        try:
            cursor = connection.cursor()
            placeholders = ", ".join(["%s"] * len(doc_ids))
            sql = f"SELECT document_id FROM bm_document WHERE document_id IN ({placeholders}) AND parent_generation_enabled = 1"
            cursor.execute(sql, tuple(doc_ids))
            results = cursor.fetchall()
            return {row[0] for row in results}
        except Error as e:
            logger.error(f"查询 parent_generation_enabled 失败: {e}")
            return set()
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()
    
    def insert_oss_record(self, file_name, url, tenant_id="000000", file_extension=""):
        """插入 OSS 记录到 sys_oss 表"""
        connection, info = self.get_connection()
        if not connection:
            return False, info
        cursor = None
        try:
            # 生成类似 "a2922173479520702464" 的 oss_id 20位
            # import random
            # oss_id = f"a{int(time.time() * 1000)}{random.randint(1000, 9999)}"
            oss_id = Snowflake().generate_id()
            create_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            
            cursor = connection.cursor()
            sql = """
            INSERT INTO sys_oss (oss_id, tenant_id, file_name, original_name, url, create_time, file_suffix)
            VALUES (%s, %s, %s, %s, %s, %s, %s)
            """
            cursor.execute(sql, (oss_id, tenant_id, file_name, file_name, url, create_time, file_extension))
            connection.commit()
            logger.info(f"OSS 记录插入成功: {oss_id}")
            return True, oss_id
        except Error as e:
            connection.rollback()
            logger.error(f"插入 OSS 记录失败: {e}")
            return False, str(e)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()
    
    def update_total_doc_len(self, update_json):
        """
        更新长度表和文档长度表，删除/更新/插入 slice_info 表
        """
        knowledge_id = update_json.get("knowledge_id")
        doc_id = update_json.get("doc_id")
        chunk_len = update_json.get("chunk_len")
        operate = update_json.get("operate")
        chunk_id = update_json.get("chunk_id")
        chunk_text = update_json.get("chunk_text", "")
        slice_index = update_json.get("slice_index")
        
        connection, info = self.get_connection()
        if not connection:
            return False, info
        
        cursor = None
        try:
            cursor = connection.cursor()
            
            # 查询文档当前信息
            cursor.execute(
                "SELECT word_num, slice_total FROM bm_document WHERE document_id = %s", 
                (doc_id,)
            )
            doc_result = cursor.fetchone()
            logger.info(f"查询到的文档长度信息：{doc_result}")
            
            # 查询知识库当前信息
            cursor.execute(
                "SELECT word_num FROM bm_knowledge WHERE knowledge_id = %s", 
                (knowledge_id,)
            )
            knowledge_result = cursor.fetchone()
            logger.info(f"查询到的知识库总长度信息：{knowledge_result}")
            
            # 计算新的文档长度
            if not doc_result:
                new_word_num = chunk_len if chunk_len > 0 else 0
                slice_total = 0
            else:
                old_word_num = doc_result[0] or 0
                slice_total = doc_result[1] or 0
                new_word_num = old_word_num + chunk_len
                if operate == "delete":
                    slice_total = max(0, slice_total - 1)
            
            # 计算新的知识库长度
            if not knowledge_result:
                new_knowledge_word_num = chunk_len if chunk_len > 0 else 0
            else:
                old_knowledge_word_num = knowledge_result[0] or 0
                new_knowledge_word_num = old_knowledge_word_num + chunk_len
            
            date_now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            
            # 根据操作类型执行不同的 SQL
            if operate == "update":
                # # 更新文档长度
                # cursor.execute(
                #     "UPDATE bm_document SET word_num = %s WHERE document_id = %s",
                #     (new_word_num, doc_id)
                # )
                # 更新切片内容
                cursor.execute(
                    "UPDATE slice_info SET slice_text = %s, update_time = %s WHERE slice_id = %s",
                    (chunk_text, date_now, chunk_id)
                )
                
            elif operate == "insert":
                # 获取最大切片索引
                # cursor.execute(
                #     "SELECT MAX(slice_index) FROM slice_info WHERE document_id = %s",
                #     (doc_id,)
                # )
                # chunk_index_result = cursor.fetchone()[0]
                # chunk_max_index = int(chunk_index_result) if chunk_index_result else 0
                
                # # 更新文档长度
                # cursor.execute(
                #     "UPDATE bm_document SET word_num = %s WHERE document_id = %s",
                #     (new_word_num, doc_id)
                # )
                # 插入新切片
                cursor.execute(
                    """INSERT INTO slice_info 
                       (slice_id, knowledge_id, document_id, slice_text, create_time, slice_index, old_slice_text) 
                       VALUES (%s, %s, %s, %s, %s, %s, %s)""",
                    (chunk_id, knowledge_id, doc_id, chunk_text, date_now, int(slice_index), chunk_text)
                )
                
            elif operate == "delete":
                # # 更新文档长度和切片总数
                # cursor.execute(
                #     "UPDATE bm_document SET word_num = %s, slice_total = %s WHERE document_id = %s",
                #     (new_word_num, slice_total, doc_id)
                # )
                # 删除切片
                # cursor.execute(
                #     "DELETE FROM slice_info WHERE slice_id = %s",
                #     (chunk_id,)
                # )
                cursor.execute(
                    "UPDATE slice_info SET del_flag = 1 WHERE slice_id = %s",
                    (chunk_id,)
                )
            
            # # 更新知识库总长度
            # cursor.execute(
            #     "UPDATE bm_knowledge SET word_num = %s WHERE knowledge_id = %s",
            #     (new_knowledge_word_num, knowledge_id)
            # )
            
            connection.commit()
            logger.info("bm_document 和 bm_knowledge 数据更新成功")
            return True, "success"
            
        except Error as e:
            connection.rollback()
            logger.error(f"update_total_doc_len 数据库操作出错：{e}")
            return False, str(e)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()

    def query_slice_info_by_doc_ids(self, knowledge_id: str, doc_ids: list):
        """
        根据 knowledge_id 和 doc_ids 列表查询 slice_info 表中的数据
        
        参数:
            knowledge_id: 知识库ID
            doc_ids: 文档ID列表
        返回:
            (success, data_or_error): 成功时返回数据列表，失败时返回错误信息
        """
        if not doc_ids:
            return True, []
        
        connection, info = self.get_connection()
        if not connection:
            return False, info
        cursor = None
        try:
            cursor = connection.cursor(dictionary=True)
            placeholders = ','.join(['%s'] * len(doc_ids))
            sql = f"""
            SELECT slice_id, knowledge_id, document_id, slice_text, old_slice_text, 
                   create_time, update_time, slice_index
            FROM slice_info 
            WHERE knowledge_id = %s AND document_id IN ({placeholders})
            """
            cursor.execute(sql, (knowledge_id, *doc_ids))
            results = cursor.fetchall()
            logger.info(f"查询 slice_info 成功，共 {len(results)} 条记录")
            return True, results
        except Error as e:
            logger.error(f"查询 slice_info 出错: {e}")
            return False, str(e)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()
    
    def query_media_replacement_by_doc_ids(self, knowledge_id: str, doc_ids: list):
        """
        根据 knowledge_id 和 doc_ids 列表查询 bm_media_replacement 表中的数据
        
        参数:
            knowledge_id: 知识库ID
            doc_ids: 文档ID列表
        返回:
            (success, data_or_error): 成功时返回数据列表，失败时返回错误信息
        """
        if not doc_ids:
            return True, []
        
        connection, info = self.get_connection()
        if not connection:
            return False, info
        cursor = None
        try:
            cursor = connection.cursor(dictionary=True)
            placeholders = ','.join(['%s'] * len(doc_ids))
            sql = f"""
            SELECT knowledge_id, document_id, origin_text, media_type, media_url, create_time
            FROM bm_media_replacement 
            WHERE knowledge_id = %s AND document_id IN ({placeholders})
            """
            cursor.execute(sql, (knowledge_id, *doc_ids))
            results = cursor.fetchall()
            logger.info(f"查询 bm_media_replacement 成功，共 {len(results)} 条记录")
            return True, results
        except Error as e:
            logger.error(f"查询 bm_media_replacement 出错: {e}")
            return False, str(e)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()

    def query_bm_document_by_doc_ids(self, knowledge_id: str, doc_ids: list):
        """
        根据 knowledge_id 和 doc_ids 列表查询 bm_document 表中的数据
        
        参数:
            knowledge_id: 知识库ID
            doc_ids: 文档ID列表
        返回:
            (success, data_or_error): 成功时返回数据列表，失败时返回错误信息
        """
        if not doc_ids:
            return True, []
        
        connection, info = self.get_connection()
        if not connection:
            return False, info
        cursor = None
        try:
            cursor = connection.cursor(dictionary=True)
            placeholders = ','.join(['%s'] * len(doc_ids))
            sql = f"""
            SELECT * FROM bm_document 
            WHERE knowledge_id = %s AND document_id IN ({placeholders})
            """
            cursor.execute(sql, (knowledge_id, *doc_ids))
            results = cursor.fetchall()
            logger.info(f"查询 bm_document 成功，共 {len(results)} 条记录")
            return True, results
        except Error as e:
            logger.error(f"查询 bm_document 出错: {e}")
            return False, str(e)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()

    def copy_docs_metadata_to_new_knowledge(self, source_knowledge_id: str, source_doc_ids: list, new_knowledge_id: str, doc_id_mapping: dict = None, chunk_id_mapping: dict = None, tenant_id: str = "000000"):
        """
        复制文档元数据到新知识库（bm_document、slice_info 和 bm_media_replacement）
        
        参数:
            source_knowledge_id: 源知识库ID
            source_doc_ids: 源文档ID列表
            new_knowledge_id: 新的知识库ID
            doc_id_mapping: 旧文档ID到新文档ID的映射 {old_doc_id: new_doc_id}
                           如果不提供，将为每个文档生成新的雪花算法ID
            chunk_id_mapping: 旧切片ID到新切片ID的映射 {old_chunk_id: new_chunk_id}
                             用于 slice_id，保持与向量库一致
            tenant_id: 租户ID（由前端传入）
        
        返回:
            (success, result_dict): result_dict 包含复制结果详情
        """
        if not source_doc_ids:
            return True, {"message": "无需复制的文档", "doc_id_mapping": {}}
        
        # 如果没有提供映射，为每个源文档生成新的 doc_id
        if doc_id_mapping is None:
            doc_id_mapping = {}
            for old_doc_id in source_doc_ids:
                doc_id_mapping[old_doc_id] = generate_snowflake_id()
        
        # 1. 查询源 slice_info 数据
        success, slice_data = self.query_slice_info_by_doc_ids(source_knowledge_id, source_doc_ids)
        if not success:
            return False, {"error": f"查询切片数据失败: {slice_data}"}
        
        # 2. 查询源 bm_media_replacement 数据
        success, media_data = self.query_media_replacement_by_doc_ids(source_knowledge_id, source_doc_ids)
        if not success:
            return False, {"error": f"查询媒体映射失败: {media_data}"}
        
        # 3. 查询源 bm_document 数据
        success, doc_data = self.query_bm_document_by_doc_ids(source_knowledge_id, source_doc_ids)
        if not success:
            return False, {"error": f"查询文档数据失败: {doc_data}"}
        
        connection, info = self.get_connection()
        if not connection:
            return False, {"error": info}
        
        cursor = None
        try:
            cursor = connection.cursor()
            date_now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            
            slice_count = 0
            media_count = 0
            doc_count = 0
            
            # 4. 插入 bm_document
            if doc_data:
                # 定义所有字段及特殊处理
                all_fields = [
                    'document_id', 'knowledge_id', 'custom_separator', 'sentence_size', 'length',
                    'word_num', 'slice_total', 'name', 'url', 'parse_image', 'tenant_id',
                    'create_dept', 'create_by', 'create_time', 'update_by', 'update_time',
                    'remark', 'parsing_type', 'oss_id', 'status', 'mark_oss_id', 'mark_url',
                    'ref_document_id', 'qa_checked', 'related_questions_enabled',
                    'summary_generation_enabled', 'parent_generation_enabled', 'suffix', 'pdf_url'
                ]
                doc_values = []
                for row in doc_data:
                    old_doc_id = row['document_id']
                    new_doc_id = doc_id_mapping.get(old_doc_id, generate_snowflake_id())
                    # 特殊处理字段
                    special = {
                        'document_id': new_doc_id,
                        'knowledge_id': new_knowledge_id,
                        'tenant_id': tenant_id,
                        'create_time': date_now,
                        'ref_document_id': old_doc_id
                    }
                    doc_values.append(tuple(special.get(f, row.get(f)) for f in all_fields))
                
                placeholders = ','.join(['%s'] * len(all_fields))
                doc_sql = f"INSERT INTO bm_document ({','.join(all_fields)}) VALUES ({placeholders})"
                cursor.executemany(doc_sql, doc_values)
                doc_count = len(doc_values)
                logger.info(f"插入文档数据成功: {doc_count} 条")
            
            # 5. 插入 slice_info（使用 old_slice_text 存储原始文本）
            if slice_data:
                slice_values = []
                
                for row in slice_data:
                    old_doc_id = row['document_id']
                    old_slice_id = row['slice_id']
                    new_doc_id = doc_id_mapping.get(old_doc_id, generate_snowflake_id())
                    # 使用向量库中的 chunk_id 作为 slice_id（保持一致）
                    new_slice_id = chunk_id_mapping.get(old_slice_id, generate_snowflake_id()) if chunk_id_mapping else generate_snowflake_id()
                    
                    # old_slice_text 存储原始文本
                    old_slice_text = row.get('old_slice_text') or row['slice_text']
                    
                    slice_values.append((
                        new_slice_id,
                        new_knowledge_id,
                        new_doc_id,
                        row['slice_text'],
                        old_slice_text,
                        date_now,
                        row.get('slice_index')
                    ))
                
                slice_sql = """
                INSERT INTO slice_info (
                    slice_id, knowledge_id, document_id, slice_text, old_slice_text, create_time, slice_index
                ) VALUES (%s, %s, %s, %s, %s, %s, %s)
                """
                
                cursor.executemany(slice_sql, slice_values)
                slice_count = len(slice_values)
                logger.info(f"插入切片数据成功: {slice_count} 条")
            
            # 6. 插入 bm_media_replacement
            if media_data:
                media_values = []
                for row in media_data:
                    old_doc_id = row['document_id']
                    new_doc_id = doc_id_mapping.get(old_doc_id, generate_snowflake_id())
                    
                    media_values.append((
                        new_knowledge_id,
                        new_doc_id,
                        row['origin_text'],
                        row['media_type'],
                        row['media_url'],
                        date_now
                    ))
                
                media_sql = """
                INSERT INTO bm_media_replacement (
                    knowledge_id, document_id, origin_text, media_type, media_url, create_time
                ) VALUES (%s, %s, %s, %s, %s, %s)
                """
                cursor.executemany(media_sql, media_values)
                media_count = len(media_values)
                logger.info(f"插入媒体映射数据成功: {media_count} 条")
            
            connection.commit()
            
            return True, {
                "message": "复制元数据成功",
                "doc_id_mapping": doc_id_mapping,
                "doc_count": doc_count,
                "slice_count": slice_count,
                "media_count": media_count
            }
            
        except Error as e:
            connection.rollback()
            logger.error(f"复制文档元数据失败: {e}")
            return False, {"error": str(e)}
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()

    def query_slice_by_id(self, knowledge_id: str, slice_id: str):
        """
        根据 knowledge_id 和 slice_id 查询单个切片数据
        """
        connection, info = self.get_connection()
        if not connection:
            return False, info
        cursor = None
        try:
            cursor = connection.cursor(dictionary=True)
            sql = """
            SELECT slice_id, knowledge_id, document_id, slice_text, old_slice_text, slice_index
            FROM slice_info WHERE knowledge_id = %s AND slice_id = %s
            """
            cursor.execute(sql, (knowledge_id, slice_id))
            results = cursor.fetchall()
            return True, results
        except Error as e:
            logger.error(f"查询 slice_info 出错: {e}")
            return False, str(e)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()

    def query_slice_by_knowledge_and_doc(self, knowledge_id: str, document_id: str):
        """
        根据 knowledge_id 和 document_id 查询 slice_info 表中的切片数据
        """
        connection, info = self.get_connection()
        if not connection:
            return False, info
        cursor = None
        try:
            cursor = connection.cursor(dictionary=True)
            sql = """
            SELECT slice_id, knowledge_id, document_id, slice_text, old_slice_text, slice_index
            FROM slice_info 
            WHERE knowledge_id = %s AND document_id = %s
            """
            cursor.execute(sql, (knowledge_id, document_id))
            results = cursor.fetchall()
            logger.info(f"查询 slice_info 成功，共 {len(results)} 条记录")
            return True, results
        except Error as e:
            logger.error(f"查询 slice_info 出错: {e}")
            return False, str(e)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()

    def update_slice_llm_fields(self, knowledge_id: str, slice_id: str, qa: str = None, question: str = None, summary: str = None):
        """
        更新 slice_info 表中的 qa、question、summary 字段，值为 None 的字段不更新
        基于 knowledge_id 和 slice_id 的包含关系进行更新
        """
        connection, info = self.get_connection()
        if not connection:
            return False, info
        cursor = None
        try:
            # 动态构建 SET 子句，只更新非 None 字段
            set_parts = []
            params = []
            if qa is not None:
                set_parts.append("qa = %s")
                params.append(qa)
            if question is not None:
                set_parts.append("question = %s")
                params.append(question)
            if summary is not None:
                set_parts.append("summary = %s")
                params.append(summary)
            
            if not set_parts:
                return True, "no fields to update"
            
            date_now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            set_parts.append("update_time = %s")
            params.append(date_now)
            params.append(knowledge_id)
            params.append(slice_id)
            
            cursor = connection.cursor()
            sql = f"UPDATE slice_info SET {', '.join(set_parts)} WHERE knowledge_id = %s AND slice_id = %s"
            cursor.execute(sql, tuple(params))
            connection.commit()
            logger.info(f"更新切片 {slice_id} 的LLM字段成功")
            return True, "success"
        except Error as e:
            connection.rollback()
            logger.error(f"更新 slice_info LLM字段出错: {e}")
            return False, str(e)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()