LQAdminPlatform/src/app/services/milvus_service.py

"""
Milvus Service：业务层（直接用 manager.client 调 Milvus 原生方法）
"""
from __future__ import annotations

import sys
import os

import logging
import re
import hashlib
import time
from typing import List, Dict, Any, Optional
from datetime import datetime

from app.base import get_milvus_manager, get_milvus_vectorstore, get_embedding_model
from app.core.config import config_handler
from langchain_core.documents import Document

logger = logging.getLogger(__name__)

# 默认集合名称
PARENT_COLLECTION_NAME = config_handler.get("admin_app", "PARENT_COLLECTION_NAME", "test_27_parent")
CHILD_COLLECTION_NAME = config_handler.get("admin_app", "CHILD_COLLECTION_NAME", "test_27_child")


class MilvusService:
    def __init__(self):
        self.client = get_milvus_manager().client
        # 获取embedding model
        self.emdmodel = get_embedding_model()
        # 默认向量维度 (Qwen3-Embedding-8B default)
        self.DENSE_DIM = 4096

    def ensure_collections(self):
        """确保系统默认集合已创建"""
        collections = [PARENT_COLLECTION_NAME, CHILD_COLLECTION_NAME]
        for name in collections:
            self.ensure_collection_exists(name)

    async def insert_knowledge(self, content: str, doc_info: Dict[str, Any]):
        """将 Markdown 内容切分并入库 (支持父子段分表)"""
        try:
            doc_id = doc_info.get("doc_id")
            doc_name = doc_info.get("doc_name")
            doc_version = doc_info.get("doc_version", int(time.time()))
            tags = doc_info.get("tags", "")
            user_id = doc_info.get("user_id", "system")
            
            kb_method = doc_info.get("kb_method")
            target_collection = doc_info.get("collection_name") or PARENT_COLLECTION_NAME
            
            from langchain_text_splitters import RecursiveCharacterTextSplitter

            if kb_method == "parent_child":
                # --- 方案 A: 父子段分表入库 ---
                parent_col = f"{target_collection}_parent"
                child_col = f"{target_collection}_child"
                
                # 1. 切分父段 (较大块)
                parent_splitter = RecursiveCharacterTextSplitter(
                    chunk_size=1000,
                    chunk_overlap=100
                )
                parent_chunks = parent_splitter.split_text(content)
                
                parent_docs = []
                child_docs = []
                
                for p_idx, p_content in enumerate(parent_chunks):
                    # 生成唯一的 parent_id
                    p_id = hashlib.sha1(f"{doc_id}_p_{p_idx}".encode()).hexdigest()
                    
                    # 准备父段文档 (Metadata 不包含向量，仅用于检索回显)
                    p_metadata = self._prepare_metadata(doc_info, p_id, p_idx, p_id)
                    parent_docs.append(Document(page_content=p_content, metadata=p_metadata))
                    
                    # 2. 在每个父段内部切分子段 (较小块)
                    child_splitter = RecursiveCharacterTextSplitter(
                        chunk_size=300,
                        chunk_overlap=30
                    )
                    child_chunks = child_splitter.split_text(p_content)
                    
                    for c_idx, c_content in enumerate(child_chunks):
                        # 子段的 parent_id 指向父段的 p_id
                        c_metadata = self._prepare_metadata(doc_info, p_id, c_idx, p_id)
                        child_docs.append(Document(page_content=c_content, metadata=c_metadata))

                # 确保两个集合都存在
                self.ensure_collection_exists(parent_col)
                self.ensure_collection_exists(child_col)
                
                # 分别入库
                if parent_docs:
                    get_milvus_vectorstore(parent_col).add_documents(parent_docs)
                if child_docs:
                    get_milvus_vectorstore(child_col).add_documents(child_docs)
                
                logger.info(f"Successfully inserted parent-child chunks for {doc_name}: {len(parent_docs)} parents -> {len(child_docs)} children")
            
            else:
                # --- 常规单表入库逻辑 ---
                chunks = []
                if kb_method == "length":
                    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
                    chunks = splitter.split_text(content)
                elif kb_method == "symbol":
                    splitter = RecursiveCharacterTextSplitter(
                        separators=["\n\n", "\n", "。", "；", "！", "？", "!", "?", ";"],
                        chunk_size=500,
                        chunk_overlap=0
                    )
                    chunks = splitter.split_text(content)
                else:
                    chunks = [p.strip() for p in re.split(r"\n\s*\n+", content) if p.strip()]
                
                if not chunks:
                    logger.warning(f"Document {doc_name} has no content chunks.")
                    return

                documents = []
                for idx, chunk in enumerate(chunks):
                    p_id = hashlib.sha1(f"{doc_id}_{idx}".encode()).hexdigest()
                    metadata = self._prepare_metadata(doc_info, p_id, idx, p_id)
                    documents.append(Document(page_content=chunk, metadata=metadata))

                self.ensure_collection_exists(target_collection)
                get_milvus_vectorstore(target_collection).add_documents(documents)
                
                logger.info(f"Successfully inserted {len(documents)} chunks for {doc_name} into {target_collection}")

        except Exception as e:
            logger.error(f"Error inserting knowledge into Milvus: {e}")
            raise

    def _prepare_metadata(self, doc_info: Dict[str, Any], p_id: str, index: int, parent_ref_id: str) -> Dict[str, Any]:
        """统一准备元数据"""
        doc_id = doc_info.get("doc_id")
        doc_name = doc_info.get("doc_name")
        doc_version = doc_info.get("doc_version", int(time.time()))
        tags = doc_info.get("tags", "")
        user_id = doc_info.get("user_id", "system")
        
        return {
            "document_id": doc_id,
            "parent_id": parent_ref_id,
            "index": index,
            "tag_list": tags,
            "permission": {},
            "is_deleted": 0,
            "created_by": user_id,
            "created_time": int(time.time() * 1000),
            "updated_by": user_id,
            "updated_time": int(time.time() * 1000),
            "metadata": {
                "doc_name": doc_name,
                "doc_version": doc_version,
                "outline_path": ""
            }
        }

    def ensure_collection_exists(self, name: str):
        """确保指定名称的集合存在，不存在则按默认 Schema 创建"""
        from pymilvus import DataType, Function, FunctionType
        
        # 1. 如果不存在，则创建集合
        if not self.client.has_collection(name):
            logger.info(f"Creating collection: {name}")
            schema = self.client.create_schema(auto_id=True, enable_dynamic_field=False)
            schema.add_field("pk", DataType.INT64, is_primary=True, auto_id=True)
            schema.add_field("text", DataType.VARCHAR, max_length=65535, enable_analyzer=True)
            schema.add_field("dense", DataType.FLOAT_VECTOR, dim=self.DENSE_DIM)
            schema.add_field("sparse", DataType.SPARSE_FLOAT_VECTOR)
            schema.add_field("document_id", DataType.VARCHAR, max_length=256)
            schema.add_field("parent_id", DataType.VARCHAR, max_length=256)
            schema.add_field("index", DataType.INT64)
            schema.add_field("tag_list", DataType.VARCHAR, max_length=2048)
            schema.add_field("permission", DataType.JSON, nullable=True)
            schema.add_field("metadata", DataType.JSON, nullable=True)
            schema.add_field("is_deleted", DataType.INT64, default_value=0)
            schema.add_field("created_by", DataType.VARCHAR, max_length=256, nullable=True)
            schema.add_field("created_time", DataType.INT64)
            schema.add_field("updated_by", DataType.VARCHAR, max_length=256, nullable=True)
            schema.add_field("updated_time", DataType.INT64)

            schema.add_function(
                Function(
                    name="bm25_fn",
                    input_field_names=["text"],
                    output_field_names=["sparse"],
                    function_type=FunctionType.BM25,
                )
            )

            self.client.create_collection(collection_name=name, schema=schema)

        # 2. 检查并补全索引
        # 获取集合的描述信息以检查字段是否存在
        desc = self.client.describe_collection(collection_name=name)
        fields_in_collection = [f.get("name") for f in desc.get("fields", [])]
        
        existing_indexes = self.client.list_indexes(collection_name=name)
        
        index_params = self.client.prepare_index_params()
        needs_index = False

        # 只有当字段存在且没有索引时才添加
        if "dense" in fields_in_collection and "dense_idx" not in existing_indexes:
            index_params.add_index(
                field_name="dense",
                index_name="dense_idx",
                index_type="AUTOINDEX",
                metric_type="COSINE",
            )
            needs_index = True
        
        if "sparse" in fields_in_collection and "bm25_idx" not in existing_indexes:
            index_params.add_index(
                field_name="sparse",
                index_name="bm25_idx",
                index_type="SPARSE_INVERTED_INDEX",
                metric_type="BM25",
                params={"inverted_index_algo": "DAAT_MAXSCORE"},
            )
            needs_index = True

        if "permission" in fields_in_collection and "permission" not in existing_indexes:
            index_params.add_index(
                field_name="permission",
                index_type="INVERTED",
                params={"json_cast_type": "VARCHAR"}
            )
            needs_index = True
        
        if "metadata" in fields_in_collection and "metadata" not in existing_indexes:
            index_params.add_index(
                field_name="metadata",
                index_type="INVERTED",
                params={"json_cast_type": "VARCHAR"}
            )
            needs_index = True

        if needs_index:
            logger.info(f"Creating missing indexes for collection: {name}")
            try:
                self.client.create_index(collection_name=name, index_params=index_params)
            except Exception as e:
                logger.error(f"Failed to create index for {name}: {e}")
        
        self.client.load_collection(collection_name=name)

    def create_collection(self, name: str, dimension: int = None, description: str = "", fields: List[Dict] = None) -> None:
        """
        创建 Milvus 集合
        :param dimension: 向量维度，如果为None则使用默认值
        :param fields: 自定义字段列表，每个元素为 {"name": "age", "type": "INT64", ...}
        """
        # 使用默认维度
        if dimension is None:
            dimension = self.DENSE_DIM

        if self.client.has_collection(name):
            logger.info(f"Collection {name} already exists.")
            return
        
        # 如果有自定义字段，使用 schema 创建
        if fields:
            from pymilvus import MilvusClient, DataType, Function, FunctionType
            
            # 1. 创建 Schema
            schema = MilvusClient.create_schema(
                auto_id=True,
                enable_dynamic_field=True,
                description=description
            )
            
            # 检查字段中是否定义了主键
            has_primary = any(f.get("is_primary") for f in fields)
            if not has_primary:
                # 如果没有定义主键，添加默认主键
                schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True, auto_id=True)
            
            # 检查是否有默认向量列，如果没有则添加 (兼容旧逻辑，但如果fields里有dense则不添加)
            has_vector = any(f.get("type") == "FLOAT_VECTOR" for f in fields)
            if not has_vector:
                schema.add_field(field_name="dense", datatype=DataType.FLOAT_VECTOR, dim=dimension)
            
            # 3. 添加用户自定义字段
            type_map = {
                "BOOL": DataType.BOOL,
                "INT8": DataType.INT8,
                "INT16": DataType.INT16,
                "INT32": DataType.INT32,
                "INT64": DataType.INT64,
                "FLOAT": DataType.FLOAT,
                "DOUBLE": DataType.DOUBLE,
                "VARCHAR": DataType.VARCHAR,
                "JSON": DataType.JSON,
                "FLOAT_VECTOR": DataType.FLOAT_VECTOR,
                "SPARSE_FLOAT_VECTOR": DataType.SPARSE_FLOAT_VECTOR,
                "BM25": DataType.SPARSE_FLOAT_VECTOR # BM25 特殊处理，映射为稀疏向量
            }
            
            bm25_field = None
            text_field_name = "text" # 默认文本字段名

            for f in fields:
                field_type_str = f.get("type", "").upper()
                dtype = type_map.get(field_type_str)
                if not dtype:
                    continue 
                
                # 记录文本字段名，供BM25使用
                if f.get("name") in ["text", "content", "chunk"]: 
                    text_field_name = f.get("name")

                kwargs = {
                    "field_name": f.get("name"),
                    "datatype": dtype,
                    "description": f.get("description", "")
                }
                
                if f.get("is_primary"):
                    kwargs["is_primary"] = True
                    kwargs["auto_id"] = True # 假设主键都是自增
                
                if dtype == DataType.VARCHAR:
                    kwargs["max_length"] = f.get("max_length", 65535)
                    # 关键修复：如果要被 BM25 引用，必须启用 analyzer
                    if f.get("name") in ["text", "content", "chunk"]:
                        kwargs["enable_analyzer"] = True
                
                if dtype == DataType.FLOAT_VECTOR:
                    kwargs["dim"] = dimension # 使用传入的 dimension
                
                schema.add_field(**kwargs)

                # 如果是 BM25 类型，记录下来以便后续添加 Function
                if field_type_str == "BM25":
                    bm25_field = f.get("name")

            # 处理 BM25 Function
            if bm25_field:
                try:
                    schema.add_function(Function(
                        name="bm25_fn",
                        input_field_names=[text_field_name],
                        output_field_names=[bm25_field],
                        function_type=FunctionType.BM25
                    ))
                    logger.info(f"Added BM25 function mapping {text_field_name} -> {bm25_field}")
                except Exception as e:
                    logger.error(f"Failed to add BM25 function: {e}")

            # 4. 准备索引参数
            index_params = self.client.prepare_index_params()
            
            # 5. 为所有向量字段添加索引
            for f in fields:
                ftype = f.get("type", "").upper()
                if ftype == "FLOAT_VECTOR":
                    index_params.add_index(
                        field_name=f.get("name"), 
                        index_type="AUTOINDEX",
                        metric_type="IP" # [Modified] 更改为 IP (内积)，通常对规范化向量效果更好，与 COSINE 类似但更简单
                    )
                elif ftype == "BM25" or ftype == "SPARSE_FLOAT_VECTOR":
                    index_params.add_index(
                        field_name=f.get("name"),
                        index_type="SPARSE_INVERTED_INDEX", # 稀疏向量索引
                        metric_type="BM25"
                    )
            
            # 6. 为自定义标量字段添加索引
            for f in fields:
                ftype = f.get("type", "").upper()
                if ftype in ["VARCHAR", "INT64", "INT32", "BOOL"] and not f.get("is_primary"):
                     # 排除主键，主键自动索引
                    index_params.add_index(
                        field_name=f.get("name"),
                        index_type="INVERTED"
                    )
                elif ftype == "JSON":
                    # Milvus 2.4+ JSON 索引必须指定 json_cast_type
                    # 这里为 JSON 字段添加默认索引，以便支持查询
                    index_params.add_index(
                        field_name=f.get("name"),
                        index_type="INVERTED",
                        params={"json_cast_type": "VARCHAR"}
                    )

            # 7. 创建集合
            self.client.create_collection(
                collection_name=name,
                schema=schema,
                index_params=index_params
            )
            
        else:
            # 使用简化的 create_collection API
            self.client.create_collection(
                collection_name=name,
                dimension=dimension,
                description=description,
                auto_id=True,  # 自动生成 ID
                id_type="int", # ID 类型
                metric_type="IP" # [Modified] 默认使用 IP
            )
        
        logger.info(f"Created collection {name} with dimension {dimension}")

    def drop_collection(self, name: str) -> None:
        """删除 Milvus 集合"""
        if self.client.has_collection(name):
            self.client.drop_collection(name)
            logger.info(f"Dropped collection {name}")

    def has_collection(self, name: str) -> bool:
        """检查集合是否存在"""
        return self.client.has_collection(name)

    def get_collection_details(self) -> List[Dict[str, Any]]:
        """
        获取所有 Collections 详细信息
        """
        details: List[Dict[str, Any]] = []

        names = self.client.list_collections()

        for name in names:
            desc = self.client.describe_collection(collection_name=name)
            stats = self.client.get_collection_stats(collection_name=name)
            load_state = self.client.get_load_state(collection_name=name)

            # ===== 时间戳转换（按你指定写法，无封装）=====
            created_time = None
            updated_time = None

            if desc.get("created_timestamp") is not None:
                ts_int = int(desc["created_timestamp"])
                physical_ms = ts_int >> 18
                created_time = datetime.fromtimestamp(physical_ms / 1000).strftime("%Y-%m-%d %H:%M:%S")

            if desc.get("update_timestamp") is not None:
                ts_int = int(desc["update_timestamp"])
                physical_ms = ts_int >> 18
                updated_time = datetime.fromtimestamp(physical_ms / 1000).strftime("%Y-%m-%d %H:%M:%S")

            # ===== 数量：不保底（要求返回结构必须有 row_count）=====
            entity_count = stats["row_count"]

            # ===== 状态：不保底（要求返回结构必须有 state）=====
            status = load_state["state"]

            details.append(
                {
                    "name": name,
                    "status": status,
                    "entity_count": entity_count,
                    "description": desc.get("description", ""),
                    "created_time": created_time,
                    "updated_time": updated_time,
                }
            )

        logger.info(f"成功获取Collections详细信息，共{len(details)}个")
        return details

    def get_collection_state(self, name: str) -> str:
        """获取集合加载状态"""
        try:
            load_state = self.client.get_load_state(collection_name=name)
            state = load_state.get("state") if isinstance(load_state, dict) else load_state
            return state
        except Exception as e:
            logger.error(f"Failed to get collection state for {name}: {e}")
            return "Unknown"

    def set_collection_state(self, name: str, action: str) -> Dict[str, Any]:
        """
        改变指定 Collection 的加载状态。

        参数:
        - name: 集合名称
        - action: 操作，取值 'load' 或 'release'

        返回:
        - 包含集合名称和当前状态的字典，例如: {"name": name, "state": "Loaded"}
        """
        action_norm = (action or "").strip().lower()
        if action_norm not in {"load", "release"}:
            raise ValueError("action 必须为 'load' 或 'release'")

        # 执行加载/释放
        if action_norm == "load":
            self.client.load_collection(collection_name=name)
        else:
            self.client.release_collection(collection_name=name)

        # 返回最新状态
        load_state = self.client.get_load_state(collection_name=name)
        state = load_state.get("state") if isinstance(load_state, dict) else load_state
        result = {"name": name, "state": state, "action": action_norm}
        logger.info(f"集合 {name} 状态更新为 {state} (action={action_norm})")
        return result

    def delete_collection_if_empty(self, name: str) -> Dict[str, Any]:
        """仅当集合内容为空时删除集合，否则抛出异常"""
        stats = self.client.get_collection_stats(collection_name=name)
        row_count = stats.get("row_count") if isinstance(stats, dict) else None
        if row_count is None:
            raise ValueError("无法获取集合行数，禁止删除")
        if int(row_count) > 0:
            raise ValueError("集合内容不为空，不能删除")

        self.client.drop_collection(collection_name=name)
        logger.info(f"集合 {name} 已删除")
        return {"name": name, "deleted": True}

    def get_collection_detail(self, name: str) -> Dict[str, Any]:
        """获取单个集合的详细信息，包含schema、索引等所有desc字段"""
        desc = self.client.describe_collection(collection_name=name)
        stats = self.client.get_collection_stats(collection_name=name)
        load_state = self.client.get_load_state(collection_name=name)

        # 时间戳转换
        created_time = None
        updated_time = None

        if desc.get("created_timestamp") is not None:
            ts_int = int(desc["created_timestamp"])
            physical_ms = ts_int >> 18
            created_time = datetime.fromtimestamp(physical_ms / 1000).strftime("%Y-%m-%d %H:%M:%S")

        if desc.get("update_timestamp") is not None:
            ts_int = int(desc["update_timestamp"])
            physical_ms = ts_int >> 18
            updated_time = datetime.fromtimestamp(physical_ms / 1000).strftime("%Y-%m-%d %H:%M:%S")

        entity_count = stats.get("row_count", 0)
        status = load_state.get("state") if isinstance(load_state, dict) else load_state

        # 提取字段schema
        fields = []
        if "fields" in desc:
            for field in desc["fields"]:
                field_info = {
                    "name": field.get("name"),
                    "type": str(field.get("type")),
                    "description": field.get("description", ""),
                    "is_primary": field.get("is_primary", False),
                    "auto_id": field.get("auto_id"),
                }
                # 向量维度
                if "params" in field and "dim" in field["params"]:
                    field_info["dim"] = field["params"]["dim"]
                # 字符串长度
                if "params" in field and "max_length" in field["params"]:
                    field_info["max_length"] = field["params"]["max_length"]
                # 其他params
                if "params" in field:
                    field_info["params"] = field["params"]
                fields.append(field_info)

        # 提取索引信息
        indices = []
        
        # 尝试从 describe_collection 结果中获取 (兼容旧逻辑)
        if "indexes" in desc:
            for idx in desc["indexes"]:
                index_info = {
                    "field_name": idx.get("field_name"),
                    "index_name": idx.get("index_name"),
                    "index_type": idx.get("index_type"),
                    "metric_type": idx.get("metric_type"),
                    "params": idx.get("params"),
                }
                indices.append(index_info)
        
        # 如果没有获取到索引信息，尝试主动查询 list_indexes
        if not indices:
            try:
                # 获取索引列表 (通常返回索引名称列表)
                index_names = self.client.list_indexes(collection_name=name)
                if index_names:
                    for idx_name in index_names:
                        try:
                            # 获取索引详情
                            idx_desc = self.client.describe_index(collection_name=name, index_name=idx_name)
                            if idx_desc:
                                indices.append({
                                    "field_name": idx_desc.get("field_name"),
                                    "index_name": idx_desc.get("index_name"),
                                    "index_type": idx_desc.get("index_type"),
                                    "metric_type": idx_desc.get("metric_type"),
                                    "params": idx_desc.get("params"),
                                })
                        except Exception:
                            continue
            except Exception as e:
                logger.warning(f"Failed to list/describe indexes for {name}: {e}")

        detail = {
            "name": name,
            "description": desc.get("description", ""),
            "status": status,
            "entity_count": entity_count,
            "created_time": created_time,
            "updated_time": updated_time,
            "fields": fields,
            "enable_dynamic_field": desc.get("enable_dynamic_field", False),
            "consistency_level": desc.get("consistency_level"),
            "num_shards": desc.get("num_shards"),
            "num_partitions": desc.get("num_partitions"),
            "indices": indices,
            "properties": desc.get("properties"),
            "aliases": desc.get("aliases", []),
        }

        logger.info(f"成功获取集合 {name} 的详细信息")
        return detail

    
    def update_collection_description(self, name: str, description: str) -> Dict[str, Any]:
        """使用 alter_collection_properties 更新集合描述"""
        description = description or ""

        # 1. 更新集合 description（唯一修改点）
        self.client.alter_collection_properties(
            collection_name=name,
            properties={"collection.description": description},
        )

        # 2. 重新获取集合信息
        desc = self.client.describe_collection(collection_name=name)
        print(desc)
        stats = self.client.get_collection_stats(collection_name=name)
        load_state = self.client.get_load_state(collection_name=name)

        # 3. 时间戳转换（Milvus TSO -> 物理时间）
        def ts_to_str(ts):
            if ts is None:
                return None
            ts_int = int(ts)
            physical_ms = ts_int >> 18
            return datetime.fromtimestamp(physical_ms / 1000).strftime("%Y-%m-%d %H:%M:%S")

        created_time = ts_to_str(desc.get("created_timestamp"))
        updated_time = ts_to_str(desc.get("update_timestamp"))

        entity_count = stats.get("row_count") if isinstance(stats, dict) else None
        status = load_state.get("state") if isinstance(load_state, dict) else load_state

        return {
            "name": name,
            "status": status,
            "entity_count": entity_count,
            "description": desc.get("description", ""),
            "created_time": created_time,
            "updated_time": updated_time,
        }

    def hybrid_search(self, collection_name: str, query_text: str,
                     top_k: int = 3, ranker_type: str = "weighted",
                     dense_weight: float = 0.7, sparse_weight: float = 0.3,
                     expr: str = None):
        """
        混合搜索（参考 test_hybrid_v2.6.py 的实现）

        Args:
            param: 包含collection_name的参数字典
            query_text: 查询文本
            top_k: 返回结果数量
            ranker_type: 重排序类型 "weighted" 或 "rrf"
            dense_weight: 密集向量权重（当ranker_type="weighted"时使用）
            sparse_weight: 稀疏向量权重（当ranker_type="weighted"时使用）
            expr: 过滤表达式 (Metadata Filtering)

        Returns:
            List[Dict]: 搜索结果列表
        """
        try:
            collection_name = collection_name
            
            # 确保集合已加载
            self.client.load_collection(collection_name)

            # 获取 vectorstore 实例（包含 Milvus 和 BM25BuiltInFunction）
            vectorstore = get_milvus_vectorstore(
                collection_name=collection_name,
                consistency_level="Strong"
            )

            # 执行混合搜索 (完全按照 test_hybrid_v2.6.py 的逻辑)
            # 注意：LangChain Milvus vectorstore 的 similarity_search 支持 expr 参数用于过滤
            if ranker_type == "weighted":
                results = vectorstore.similarity_search(
                    query=query_text,
                    k=top_k,
                    expr=expr,
                    ranker_type="weighted",
                    ranker_params={"weights": [dense_weight, sparse_weight]}
                )
            else:  # rrf
                results = vectorstore.similarity_search(
                    query=query_text,
                    k=top_k,
                    expr=expr,
                    ranker_type="rrf",
                    ranker_params={"k": 60}
                )

            # 格式化结果，保持与其他搜索方法一致
            formatted_results = []
            for doc in results:
                formatted_results.append({
                    'id': doc.metadata.get('pk', 0),
                    'text_content': doc.page_content,
                    'metadata': doc.metadata,
                    'distance': 0.0,
                    'similarity': 1.0
                })

            logger.info(f"Hybrid search returned {len(formatted_results)} results")
            return formatted_results

        except Exception as e:
            logger.error(f"Error in hybrid search: {e}")
            # 回退到传统的向量搜索
            logger.info("Falling back to traditional vector search")
            return []


# 可选：单例
milvus_service = MilvusService()


if __name__ == "__main__":
    # 推荐这样跑：
    # uv run python -m src.app.services.milvus_service
    import json

    service = MilvusService()
    
    # 测试混合搜索 hybrid_search
    print("=" * 50)
    print("测试混合检索 (Hybrid Search)")
    print("=" * 50)
    
    try:
        # 示例参数，需要根据实际情况修改
        collection_name = "first_bfp_collection_status" 
        query_text = "《公路水运工程临时用电技术规程》（JTT1499-2024）状态为现行"  # 修改为实际查询内容
        
        # 测试 weighted 模式
        print("\n1. 测试 Weighted 重排序模式:")
        print(f"   集合: {collection_name}")
        print(f"   查询: {query_text}")
        print(f"   密集权重: 0.7, 稀疏权重: 0.3")
        
        results_weighted = service.hybrid_search(
            collection_name=collection_name,
            query_text=query_text,
            top_k=5,
            ranker_type="weighted",
            dense_weight=0.7,
            sparse_weight=0.3
        )
        
        print(f"\n   结果数量: {len(results_weighted)}")
        for i, result in enumerate(results_weighted, 1):
            print(f"   [{i}] ID: {result.get('id')}, Text: {result.get('text_content')[:50]}...")
        
        # 测试 RRF 模式
        print("\n2. 测试 RRF (Reciprocal Rank Fusion) 重排序模式:")
        print(f"   集合: {collection_name}")
        print(f"   查询: {query_text}")
        
        results_rrf = service.hybrid_search(
            collection_name=collection_name,
            query_text=query_text,
            top_k=5,
            ranker_type="rrf"
        )
        
        print(f"\n   结果数量: {len(results_rrf)}")
        for i, result in enumerate(results_rrf, 1):
            print(f"   [{i}] ID: {result.get('id')}, Text: {result.get('text_content')[:50]}...")
        
        print("\n✓ 混合检索测试完成")
        
    except Exception as e:
        print(f"\n✗ 混合检索测试失败: {e}")
        import traceback
        traceback.print_exc()
    
    # 也可以查看集合详情
    print("\n" + "=" * 50)
    print("获取所有集合信息:")
    print("=" * 50)
    data = service.get_collection_details()
    for item in data:
        print(json.dumps(item, ensure_ascii=False, indent=2))