document_processor.py

"""
文档处理和向量化模块
负责文档加载、文本分块、向量化和向量数据库初始化
"""

try:
    from langchain_text_splitters import RecursiveCharacterTextSplitter
except ImportError:
    from langchain.text_splitter import RecursiveCharacterTextSplitter

from langchain_community.document_loaders import WebBaseLoader

# 尝试导入 langchain_milvus，如果失败则回退到 langchain_community 并应用补丁
try:
    from langchain_milvus import MilvusVectorStore as Milvus
    print("✅ 使用 langchain-milvus (新版)")
except ImportError:
    try:
        from langchain_community.vectorstores import Milvus
        print("⚠️ 使用 langchain_community.vectorstores.Milvus (旧版)")
        
        # Monkeypatch: 修复旧版 LangChain 对 Milvus Lite 本地文件路径的校验问题
        # 旧版 _create_connection_alias 强制要求 URI 以 http/https 开头
        def _patched_create_connection_alias(self, connection_args):
            uri = connection_args.get("uri")
            # 为本地文件生成唯一的 alias
            if uri:
                import hashlib
                return hashlib.md5(uri.encode()).hexdigest()
            return "default"
            
        # 应用补丁
        Milvus._create_connection_alias = _patched_create_connection_alias
        print("🔧 已应用 Milvus Lite 路径校验补丁")
    except ImportError:
        pass

from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.retrievers import BM25Retriever

from config import (
    KNOWLEDGE_BASE_URLS,
    CHUNK_SIZE,
    CHUNK_OVERLAP,
    COLLECTION_NAME,
    EMBEDDING_MODEL,
    # 混合检索配置
    ENABLE_HYBRID_SEARCH,
    HYBRID_SEARCH_WEIGHTS,
    KEYWORD_SEARCH_K,
    BM25_K1,
    BM25_B,
    # 向量库配置
    VECTOR_STORE_TYPE,
    MILVUS_HOST,
    MILVUS_PORT,
    MILVUS_USER,
    MILVUS_PASSWORD,
    MILVUS_URI,
    MILVUS_INDEX_TYPE,
    MILVUS_INDEX_PARAMS,
    MILVUS_SEARCH_PARAMS,
    # 查询扩展配置
    ENABLE_QUERY_EXPANSION,
    QUERY_EXPANSION_MODEL,
    QUERY_EXPANSION_PROMPT,
    MAX_EXPANDED_QUERIES,
    # 多模态配置
    ENABLE_MULTIMODAL,
    MULTIMODAL_IMAGE_MODEL,
    SUPPORTED_IMAGE_FORMATS,
    IMAGE_EMBEDDING_DIM,
    MULTIMODAL_WEIGHTS
)
from reranker import create_reranker

# 多模态支持相关导入
import base64
import io
from PIL import Image
import numpy as np
from typing import List, Dict, Any, Optional, Union

try:
    from langchain_core.documents import Document
except ImportError:
    try:
        from langchain_core.documents import Document
    except ImportError:
        from langchain.schema import Document


class CustomEnsembleRetriever:
    """自定义集成检索器，结合向量检索和BM25检索"""
    
    def __init__(self, retrievers, weights):
        self.retrievers = retrievers
        self.weights = weights
        
    def invoke(self, query):
        """执行检索并合并结果"""
        # 获取各检索器的结果
        all_results = []
        for i, retriever in enumerate(self.retrievers):
            results = retriever.invoke(query)
            for doc in results:
                # 添加检索器索引和权重信息
                doc.metadata["retriever_index"] = i
                doc.metadata["retriever_weight"] = self.weights[i]
                all_results.append(doc)
        
        return self._process_results(all_results)
    
    async def ainvoke(self, query):
        """异步执行检索并合并结果"""
        import asyncio
        
        # 并发获取各检索器的结果
        # 注意：假设所有 retriever 都支持 ainvoke
        tasks = [retriever.ainvoke(query) for retriever in self.retrievers]
        results_list = await asyncio.gather(*tasks)
        
        all_results = []
        for i, results in enumerate(results_list):
            for doc in results:
                # 添加检索器索引和权重信息
                doc.metadata["retriever_index"] = i
                doc.metadata["retriever_weight"] = self.weights[i]
                all_results.append(doc)
                
        return self._process_results(all_results)

    def _process_results(self, all_results):
        """排序和去重处理"""
        # 根据权重排序并去重
        # 简单实现：先按检索器索引排序，再按权重排序
        all_results.sort(key=lambda x: (x.metadata["retriever_index"], -x.metadata["retriever_weight"]))
        
        # 去重（基于文档内容）
        unique_results = []
        seen_content = set()
        for doc in all_results:
            content = doc.page_content
            if content not in seen_content:
                seen_content.add(content)
                unique_results.append(doc)
                
        return unique_results


class DocumentProcessor:
    """文档处理器类，负责文档加载、处理和向量化"""
    
    def __init__(self):
        self.text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
            chunk_size=CHUNK_SIZE, 
            chunk_overlap=CHUNK_OVERLAP
        )
        
        # Try to initialize embeddings with error handling
        try:
            import torch
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
            print(f"✅ 检测到设备: {device}")
            if device == 'cuda':
                print(f"   GPU型号: {torch.cuda.get_device_name(0)}")
                print(f"   GPU内存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f}GB")
            
            self.embeddings = HuggingFaceEmbeddings(
                model_name="sentence-transformers/all-MiniLM-L6-v2",  # 轻量级嵌入模型
                model_kwargs={'device': device},  # 自动选择GPU或CPU
                encode_kwargs={'normalize_embeddings': True}  # 标准化嵌入向量
            )
            print(f"✅ HuggingFace嵌入模型初始化成功 (设备: {device})")
        except Exception as e:
            print(f"⚠️ HuggingFace嵌入初始化失败: {e}")
            print("正在尝试备用嵌入方案...")
            # Fallback to OpenAI embeddings or other alternatives
            from langchain_community.embeddings import FakeEmbeddings
            self.embeddings = FakeEmbeddings(size=384)  # For testing purposes
            print("✅ 使用测试嵌入模型")
            
        self.vectorstore = None
        self.retriever = None
        self.bm25_retriever = None  # BM25检索器
        self.ensemble_retriever = None  # 集成检索器
        
        # 初始化重排器
        self.reranker = None
        self._setup_reranker()
        
        # 初始化多模态支持
        self.image_embeddings_model = None
        self._setup_multimodal()
        
        # 初始化查询扩展
        self.query_expansion_model = None
        self._setup_query_expansion()
    
    def _setup_reranker(self):
        """
        设置重排器
        使用 CrossEncoder 提升重排准确率
        """
        try:
            # 使用 CrossEncoder 重排器 (准确率最高) ⭐
            print("🔧 正在初始化 CrossEncoder 重排器...")
            self.reranker = create_reranker(
                'crossencoder',
                model_name='cross-encoder/ms-marco-MiniLM-L-6-v2',  # 轻量级模型
                max_length=512
            )
            print("✅ CrossEncoder 重排器初始化成功")
        except Exception as e:
            print(f"⚠️ CrossEncoder 初始化失败: {e}")
            print("🔄 尝试回退到混合重排器...")
            try:
                # 回退到混合重排器
                self.reranker = create_reranker('hybrid', self.embeddings)
                print("✅ 混合重排器初始化成功")
            except Exception as e2:
                print(f"⚠️ 重排器初始化完全失败: {e2}")
                print("⚠️ 将使用基础检索，不进行重排")
    
    def _setup_multimodal(self):
        """设置多模态支持"""
        if not ENABLE_MULTIMODAL:
            print("⚠️ 多模态支持已禁用")
            return
            
        try:
            print("🔧 正在初始化多模态支持...")
            from transformers import CLIPProcessor, CLIPModel
            import torch
            
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
            self.image_embeddings_model = CLIPModel.from_pretrained(MULTIMODAL_IMAGE_MODEL).to(device)
            self.image_processor = CLIPProcessor.from_pretrained(MULTIMODAL_IMAGE_MODEL)
            print(f"✅ 多模态支持初始化成功 (设备: {device})")
        except Exception as e:
            print(f"⚠️ 多模态支持初始化失败: {e}")
            print("⚠️ 将仅使用文本检索")
            self.image_embeddings_model = None
    
    def _setup_query_expansion(self):
        """设置查询扩展"""
        if not ENABLE_QUERY_EXPANSION:
            print("⚠️ 查询扩展已禁用")
            return
            
        try:
            print("🔧 正在初始化查询扩展...")
            from langchain_community.llms import Ollama
            
            self.query_expansion_model = Ollama(model=QUERY_EXPANSION_MODEL)
            print(f"✅ 查询扩展初始化成功 (模型: {QUERY_EXPANSION_MODEL})")
        except Exception as e:
            print(f"⚠️ 查询扩展初始化失败: {e}")
            print("⚠️ 将不使用查询扩展")
            self.query_expansion_model = None
    
    def load_documents(self, urls=None):
        """从URL加载文档"""
        if urls is None:
            urls = KNOWLEDGE_BASE_URLS
        
        print(f"正在加载 {len(urls)} 个URL的文档...")
        docs = [WebBaseLoader(url).load() for url in urls]
        docs_list = [item for sublist in docs for item in sublist]
        print(f"成功加载 {len(docs_list)} 个文档")
        return docs_list
    
    def split_documents(self, docs):
        """将文档分割成块"""
        print("正在分割文档...")
        doc_splits = self.text_splitter.split_documents(docs)
        print(f"文档分割完成，共 {len(doc_splits)} 个文档块")
        return doc_splits
    
    def initialize_vectorstore(self):
        """初始化向量数据库连接"""
        if self.vectorstore:
            return

        print("正在连接向量数据库...")
        
        # 强制使用 Milvus
        try:
            # 准备连接参数
            connection_args = {}
            is_local_file = False
            
            # 优先使用 URI
            if MILVUS_URI and len(MILVUS_URI.strip()) > 0:
                is_local_file = not (MILVUS_URI.startswith("http://") or MILVUS_URI.startswith("https://"))
                
                real_uri = MILVUS_URI
                if is_local_file:
                    import os
                    # Milvus Lite requires absolute path in some versions/environments
                    if not os.path.isabs(real_uri):
                        real_uri = os.path.abspath(real_uri)
                        print(f"📂 将相对路径转换为绝对路径: {real_uri}")
                    
                    # 确保父目录存在
                    parent_dir = os.path.dirname(real_uri)
                    if parent_dir and not os.path.exists(parent_dir):
                        print(f"📂 创建 Milvus 存储目录: {parent_dir}")
                        os.makedirs(parent_dir, exist_ok=True)
                
                mode_name = "Lite (Local File)" if is_local_file else "Cloud (HTTP)"
                print(f"🔄 正在连接 Milvus {mode_name} ({real_uri})...")
                connection_args["uri"] = real_uri
                
                if not is_local_file and MILVUS_PASSWORD:
                        connection_args["token"] = MILVUS_PASSWORD
            else:
                print(f"🔄 正在连接 Milvus Server ({MILVUS_HOST}:{MILVUS_PORT})...")
                connection_args = {
                    "host": MILVUS_HOST,
                    "port": MILVUS_PORT,
                    "user": MILVUS_USER,
                    "password": MILVUS_PASSWORD
                }

            # 显式建立全局连接 (修复 ConnectionNotExistException)
            try:
                from pymilvus import connections, utility
                print(f"🔌 尝试建立 pymilvus 全局连接 (Alias: default)...")
                # 移除旧连接（如果存在）以防参数变更
                if connections.has_connection("default"):
                    connections.disconnect("default")
                
                connections.connect(alias="default", **connection_args)
                print("✅ pymilvus 全局连接建立成功")
                
                # 检查集合是否存在 (提前检查，避免 LangChain 内部出错)
                if utility.has_collection(COLLECTION_NAME, using="default"):
                    print(f"✅ 集合 {COLLECTION_NAME} 已存在")
                else:
                    print(f"ℹ️ 集合 {COLLECTION_NAME} 不存在，将由 Milvus 类自动创建")
                    
            except ImportError:
                print("⚠️ 未找到 pymilvus 库，跳过显式连接")
            except Exception as e:
                print(f"⚠️ 显式连接尝试失败: {e}")
                # 继续尝试，也许 LangChain 内部能处理

            # 确定索引类型
            # Milvus Lite (本地模式) 仅支持 FLAT, IVF_FLAT, AUTOINDEX，不支持 HNSW
            final_index_type = MILVUS_INDEX_TYPE
            final_index_params = MILVUS_INDEX_PARAMS
            
            if is_local_file and MILVUS_INDEX_TYPE == "HNSW":
                print("⚠️ 检测到 Milvus Lite (本地模式)，HNSW 索引不受支持，自动切换为 AUTOINDEX")
                final_index_type = "AUTOINDEX"
                final_index_params = {} # AUTOINDEX 不需要复杂参数

            # 初始化 Milvus 连接 (不删除旧数据)
            # 注意：由于我们已经手动建立了全局连接 'default'，
            # 这里我们将 connection_args 简化为仅指向该 alias，
            # 避免 LangChain 再次尝试连接或因参数问题覆盖连接。
            self.vectorstore = Milvus(
                embedding_function=self.embeddings,
                collection_name=COLLECTION_NAME,
                connection_args={"alias": "default"}, # ✅ 复用已建立的连接
                index_params={
                    "metric_type": "L2",
                    "index_type": final_index_type,
                    "params": final_index_params
                },
                search_params={
                    "metric_type": "L2", 
                    "params": MILVUS_SEARCH_PARAMS
                },
                drop_old=False,  # ✅ 持久化关键：不删除旧索引
                auto_id=True
            )
            print("✅ Milvus 向量数据库连接成功")
        except ImportError:
            print("❌ 未安装 pymilvus，请运行: pip install pymilvus")
            raise
        except Exception as e:
            print(f"❌ Milvus 连接失败: {e}")
            raise

        # 配置检索器
        retriever_kwargs = {}
        # if ENABLE_MULTIMODAL:
            # retriever_kwargs["expr"] = "data_type == 'text'"
        self.retriever = self.vectorstore.as_retriever(search_kwargs=retriever_kwargs)

    def check_existing_urls(self, urls: List[str]) -> set:
        """检查哪些URL已经存在于向量库中"""
        if not self.vectorstore:
            return set()
            
        existing = set()
        print("正在检查已存在的文档...")
        try:
            # 尝试通过检索来检查
            # 注意：这里假设 source 字段可以作为过滤条件
            for url in urls:
                # 使用 similarity_search 但带有严格过滤，且只取1条
                # 这里的 query 没关系，主要看 filter
                try:
                    # 注意：Milvus 的 expr 语法
                    expr = f'source == "{url}"'
                    res = self.vectorstore.similarity_search(
                        "test", 
                        k=1, 
                        expr=expr
                    )
                    if res:
                        existing.add(url)
                except Exception as e:
                    # 如果失败，可能是 schema 问题，尝试 metadata 字段
                    try:
                        expr = f'metadata["source"] == "{url}"'
                        res = self.vectorstore.similarity_search(
                            "test", 
                            k=1, 
                            expr=expr
                        )
                        if res:
                            existing.add(url)
                    except:
                        pass
                        
            print(f"✅ 发现 {len(existing)} 个已存在的 URL")
        except Exception as e:
            print(f"⚠️ 检查现有URL失败: {e}")
            
        return existing

    def add_documents_to_vectorstore(self, doc_splits):
        """添加文档到向量库"""
        if not doc_splits:
            return

        print(f"正在添加 {len(doc_splits)} 个文档块到向量数据库...")
        if not self.vectorstore:
            self.initialize_vectorstore()
            
        # 添加元数据
        for doc in doc_splits:
            if 'source_type' not in doc.metadata:
                source = doc.metadata.get('source', '')
                if any(fmt in source.lower() for fmt in SUPPORTED_IMAGE_FORMATS):
                    doc.metadata['data_type'] = 'image'
                else:
                    doc.metadata['data_type'] = 'text'

        self.vectorstore.add_documents(doc_splits)
        print("✅ 文档添加完成")
        
    def create_vectorstore(self, doc_splits, persist_directory=None):
        """(已弃用) 兼容旧接口，但使用新逻辑"""
        print("⚠️ create_vectorstore 已弃用，请使用 initialize_vectorstore 和 add_documents_to_vectorstore")
        self.initialize_vectorstore()
        if doc_splits:
            self.add_documents_to_vectorstore(doc_splits)
        return self.vectorstore, self.retriever

    def get_all_documents_from_vectorstore(self, limit: Optional[int] = None) -> List[Document]:
        """从已持久化的向量数据库读取所有文档内容并构造 Document 列表"""
        if not self.vectorstore:
            return []
        try:
            data = self.vectorstore._collection.get(include=["documents", "metadatas"])  # type: ignore
            docs_raw = data.get("documents") or []
            metas = data.get("metadatas") or []
            docs: List[Document] = []
            for i, content in enumerate(docs_raw):
                if content:
                    meta = metas[i] if i < len(metas) else {}
                    docs.append(Document(page_content=content, metadata=meta))
            if limit:
                return docs[:limit]
            return docs
        except Exception as e:
            print(f"⚠️ 读取向量库文档失败: {e}")
            return []
    
    def setup_knowledge_base(self, urls=None, enable_graphrag=False):
        """设置完整的知识库（加载、分割、向量化）
        
        Args:
            urls: 文档URL列表
            enable_graphrag: 是否启用GraphRAG索引
            
        Returns:
            vectorstore, retriever, doc_splits
        """
        if urls is None:
            urls = KNOWLEDGE_BASE_URLS
            
        # 1. 初始化向量库连接
        self.initialize_vectorstore()
        
        # 2. 检查已存在的 URL (去重)
        existing_urls = self.check_existing_urls(urls)
        new_urls = [url for url in urls if url not in existing_urls]
        
        doc_splits = []
        if new_urls:
            print(f"🔄 发现 {len(new_urls)} 个新 URL，开始处理...")
            docs = self.load_documents(new_urls)
            doc_splits = self.split_documents(docs)
            self.add_documents_to_vectorstore(doc_splits)
        else:
            print("✅ 所有 URL 已存在，跳过文档加载和向量化")
            
        # 3. 初始化混合检索 (BM25)
        if ENABLE_HYBRID_SEARCH:
            print("正在初始化混合检索 (BM25)...")
            try:
                bm25_docs = []
                # 如果有旧数据且这次没有加载全部数据，必须从 DB 加载所有文档以重建 BM25
                # 注意：如果只有新文档，BM25 只会包含新文档，这是不对的。
                # 只要有 existing_urls，说明库里有旧数据。
                if len(existing_urls) > 0:
                    print("🔄 正在从向量库加载所有文档以重建 BM25 索引...")
                    # 注意：这里假设内存够大
                    all_docs = self.get_all_documents_from_vectorstore()
                    bm25_docs = all_docs
                else:
                    # 全新构建
                    bm25_docs = doc_splits
                
                if bm25_docs:
                    self.bm25_retriever = BM25Retriever.from_documents(
                        bm25_docs, 
                        k=KEYWORD_SEARCH_K,
                        k1=BM25_K1,
                        b=BM25_B
                    )
                    
                    self.ensemble_retriever = CustomEnsembleRetriever(
                        retrievers=[self.retriever, self.bm25_retriever],
                        weights=[HYBRID_SEARCH_WEIGHTS["vector"], HYBRID_SEARCH_WEIGHTS["keyword"]]
                    )
                    print("✅ 混合检索初始化成功")
                else:
                    print("⚠️ 没有文档用于初始化 BM25")
            except Exception as e:
                print(f"⚠️ 混合检索初始化失败: {e}")
                self.ensemble_retriever = None
        
        # 返回 doc_splits用于GraphRAG索引 (注意：这里只返回了新增的)
        return self.vectorstore, self.retriever, doc_splits
    
    async def async_expand_query(self, query: str) -> List[str]:
        """异步扩展查询"""
        if not self.query_expansion_model:
            return [query]
            
        try:
            # 使用LLM生成扩展查询
            prompt = QUERY_EXPANSION_PROMPT.format(query=query)
            expanded_queries_text = await self.query_expansion_model.ainvoke(prompt)
            
            # 解析扩展查询
            expanded_queries = [query]  # 包含原始查询
            for line in expanded_queries_text.strip().split('\n'):
                line = line.strip()
                if line and not line.startswith('#') and not line.startswith('//'):
                    # 移除可能的编号前缀
                    if line[0].isdigit() and '.' in line[:5]:
                        line = line.split('.', 1)[1].strip()
                    expanded_queries.append(line)
            
            # 限制扩展查询数量
            return expanded_queries[:MAX_EXPANDED_QUERIES + 1]
        except Exception as e:
            print(f"⚠️ 异步查询扩展失败: {e}")
            return [query]

    async def async_hybrid_retrieve(self, query: str, top_k: int = 5, filter_type: str = "text") -> List:
        """异步混合检索
        
        Args:
            filter_type: 数据类型过滤，"text" (默认), "image", 或 "all" (不过滤)
        """
        # 构建搜索参数
        search_kwargs = {}
        if filter_type != "all" and ENABLE_MULTIMODAL:
            search_kwargs["expr"] = f"data_type == '{filter_type}'"

        if not ENABLE_HYBRID_SEARCH or not self.ensemble_retriever:
            # 纯向量检索，直接支持 search_kwargs
            if self.vectorstore:
                return await self.vectorstore.asimilarity_search(query, k=top_k, **search_kwargs)
            return await self.retriever.ainvoke(query)
            
        try:
            # 混合检索
            # 注意：目前 CustomEnsembleRetriever 的 invoke/ainvoke 尚未透传 search_kwargs
            # 为了让混合检索也享受到过滤优化，我们需要修改 CustomEnsembleRetriever 或者在这里处理
            # 鉴于 CustomEnsembleRetriever 比较简单，我们假设它主要用于文本
            # 如果需要严格过滤，最好在 vectorstore 层面处理
            
            # 临时方案：如果是混合检索且需要过滤，我们可能需要传递给 retriever
            # 但标准 retriever 接口不支持动态传参。
            # 策略：如果 filter_type 是 text (默认)，且我们在 init 时已经设置了默认不严格过滤，
            # 这里其实无法动态改变 retriever 的行为，除非我们重新生成一个 retriever 或者修改 retriever.search_kwargs
            
            # 动态修改 retriever 的 search_kwargs (这是 LangChain retriever 的特性)
            if filter_type != "all" and ENABLE_MULTIMODAL:
                self.retriever.search_kwargs["expr"] = f"data_type == '{filter_type}'"
            else:
                self.retriever.search_kwargs.pop("expr", None)
                
            results = await self.ensemble_retriever.ainvoke(query)
            return results[:top_k]
        except Exception as e:
            print(f"⚠️ 异步混合检索失败: {e}")
            print("回退到向量检索")
            if self.vectorstore:
                return await self.vectorstore.asimilarity_search(query, k=top_k, **search_kwargs)
            return await self.retriever.ainvoke(query)

    async def async_enhanced_retrieve(self, query: str, top_k: int = 5, rerank_candidates: int = 20, 
                         image_paths: List[str] = None, use_query_expansion: bool = None):
        """异步增强检索"""
        import asyncio
        
        # 确定是否使用查询扩展
        if use_query_expansion is None:
            use_query_expansion = ENABLE_QUERY_EXPANSION
            
        # 如果启用查询扩展，生成扩展查询
        if use_query_expansion:
            expanded_queries = await self.async_expand_query(query)
            print(f"查询扩展: {len(expanded_queries)} 个查询")
        else:
            expanded_queries = [query]
            
        # 多模态检索（暂时保持同步，使用线程池）
        if image_paths and ENABLE_MULTIMODAL:
            loop = asyncio.get_running_loop()
            return await loop.run_in_executor(None, self.multimodal_retrieve, query, image_paths, top_k)
            
        # 混合检索或向量检索
        all_candidate_docs = []
        
        # 决定过滤策略
        # 默认情况下，如果只是文本查询，为了性能优化，我们只检索文本数据
        # 如果提供了图像，或者用户显式要求，可以放开限制
        filter_type = "text" # 默认只搜文本，实现百万级数据的性能优化
        if image_paths:
            filter_type = "all" # 跨模态时搜所有
            
        # 构建过滤表达式 (仅用于直接调用 vectorstore 的情况，async_hybrid_retrieve 内部已处理)
        search_kwargs = {}
        if filter_type != "all" and ENABLE_MULTIMODAL:
             search_kwargs["expr"] = f"data_type == '{filter_type}'"

        async def retrieve_single(q):
            if ENABLE_HYBRID_SEARCH:
                # 使用支持动态过滤的 hybrid retrieve
                 docs = await self.async_hybrid_retrieve(q, rerank_candidates, filter_type=filter_type)
            else:
                # 使用带有过滤条件的检索
                if self.vectorstore:
                    docs = await self.vectorstore.asimilarity_search(
                        q, 
                        k=rerank_candidates,
                        **search_kwargs # 传入 expr
                    )
                else:
                    # Fallback
                    docs = await self.retriever.ainvoke(q)
                
                if len(docs) > rerank_candidates:
                    docs = docs[:rerank_candidates]
            return docs

        # 并发执行所有查询的检索
        results = await asyncio.gather(*[retrieve_single(q) for q in expanded_queries])
        
        for docs in results:
            all_candidate_docs.extend(docs)
            
        # 去重（基于文档内容）
        unique_docs = []
        seen_content = set()
        for doc in all_candidate_docs:
            content = doc.page_content
            if content not in seen_content:
                seen_content.add(content)
                unique_docs.append(doc)
                
        print(f"检索获得 {len(unique_docs)} 个候选文档")
        
        # 重排（如果重排器可用）
        # 注意：重排通常是计算密集型，建议放入线程池
        if self.reranker and len(unique_docs) > top_k:
            try:
                loop = asyncio.get_running_loop()
                # rerank 方法内部可能也比较耗时
                reranked_results = await loop.run_in_executor(
                    None, 
                    self.reranker.rerank, 
                    query, unique_docs, top_k
                )
                final_docs = [doc for doc, score in reranked_results]
                scores = [score for doc, score in reranked_results]
                
                print(f"重排后返回 {len(final_docs)} 个文档")
                print(f"重排分数范围: {min(scores):.4f} - {max(scores):.4f}")
                
                return final_docs
            except Exception as e:
                print(f"⚠️ 重排失败: {e}，使用原始检索结果")
                return unique_docs[:top_k]
        else:
            return unique_docs[:top_k]
    
    def expand_query(self, query: str) -> List[str]:
        """扩展查询，生成相关查询"""
        if not self.query_expansion_model:
            return [query]
            
        try:
            # 使用LLM生成扩展查询
            prompt = QUERY_EXPANSION_PROMPT.format(query=query)
            expanded_queries_text = self.query_expansion_model.invoke(prompt)
            
            # 解析扩展查询
            expanded_queries = [query]  # 包含原始查询
            for line in expanded_queries_text.strip().split('\n'):
                line = line.strip()
                if line and not line.startswith('#') and not line.startswith('//'):
                    # 移除可能的编号前缀
                    if line[0].isdigit() and '.' in line[:5]:
                        line = line.split('.', 1)[1].strip()
                    expanded_queries.append(line)
            
            # 限制扩展查询数量
            return expanded_queries[:MAX_EXPANDED_QUERIES + 1]  # +1 因为包含原始查询
        except Exception as e:
            print(f"⚠️ 查询扩展失败: {e}")
            return [query]
    
    def encode_image(self, image_path: str) -> np.ndarray:
        """编码图像为嵌入向量"""
        if not self.image_embeddings_model:
            raise ValueError("多模态支持未初始化")
            
        try:
            # 加载并处理图像
            image = Image.open(image_path).convert('RGB')
            inputs = self.image_processor(images=image, return_tensors="pt")
            
            # 获取图像嵌入
            with torch.no_grad():
                image_features = self.image_embeddings_model.get_image_features(**inputs)
                # 标准化嵌入向量
                image_features = image_features / image_features.norm(p=2, dim=-1, keepdim=True)
                
            return image_features.cpu().numpy().flatten()
        except Exception as e:
            print(f"⚠️ 图像编码失败: {e}")
            raise
    
    def multimodal_retrieve(self, query: str, image_paths: List[str] = None, top_k: int = 5) -> List:
        """多模态检索，结合文本和图像"""
        if not ENABLE_MULTIMODAL or not self.image_embeddings_model:
            # 如果多模态未启用，回退到文本检索
            return self.hybrid_retrieve(query, top_k) if ENABLE_HYBRID_SEARCH else self.retriever.invoke(query)[:top_k]
        
        # 1. 文本查询 (Text-to-Text & Text-to-Image)
        # 如果提供了文本查询，我们希望它能检索到文本和相关图像
        # 此时不应该限制 data_type，或者应该显式包含两者
        
        # 如果没有提供图像，这可能是一个纯文本查询，但也可能想搜图
        # 这里我们让 self.retriever (或 hybrid) 负责所有模态的检索
        # (前提是它们都在同一个向量空间，CLIP 可以做到这一点)
        text_docs = []
        if query:
             text_docs = self.hybrid_retrieve(query, top_k) if ENABLE_HYBRID_SEARCH else self.retriever.invoke(query)[:top_k]
        
        # 如果没有提供图像输入，直接返回文本查询的结果
        if not image_paths:
            return text_docs
            
        try:
            # 2. 图像查询 (Image-to-Text & Image-to-Image)
            image_results = []
            for image_path in image_paths:
                # 检查文件格式
                file_ext = image_path.split('.')[-1].lower()
                if file_ext not in SUPPORTED_IMAGE_FORMATS:
                    print(f"⚠️ 不支持的图像格式: {file_ext}")
                    continue
                    
                # 编码图像
                image_embedding = self.encode_image(image_path)
                
                # 使用图像嵌入进行检索
                if self.vectorstore:
                    # 图像可以检索文本描述，也可以检索相似图像
                    # 这里我们不做限制，检索所有类型
                    img_docs = self.vectorstore.similarity_search_by_vector(
                        embedding=image_embedding,
                        k=top_k
                    )
                    image_results.extend(img_docs)
                
            # 合并文本查询结果和图像查询结果
            # 简单合并并去重
            all_docs = text_docs + image_results
            
            # 去重
            unique_docs = []
            seen_content = set()
            for doc in all_docs:
                content = doc.page_content
                if content not in seen_content:
                    seen_content.add(content)
                    unique_docs.append(doc)
            
            final_docs = unique_docs[:top_k]
            
            print(f"✅ 多模态检索完成，返回 {len(final_docs)} 个结果")
            return final_docs
        except Exception as e:
            print(f"⚠️ 多模态检索失败: {e}")
            print("回退到文本检索")
            return text_docs
    
    def hybrid_retrieve(self, query: str, top_k: int = 5) -> List:
        """混合检索，结合向量检索和关键词检索"""
        if not ENABLE_HYBRID_SEARCH or not self.ensemble_retriever:
            # 如果混合检索未启用，回退到向量检索
            return self.retriever.invoke(query)[:top_k]
            
        try:
            # 使用集成检索器进行混合检索
            results = self.ensemble_retriever.invoke(query)
            return results[:top_k]
        except Exception as e:
            print(f"⚠️ 混合检索失败: {e}")
            print("回退到向量检索")
            return self.retriever.invoke(query)[:top_k]
    
    def enhanced_retrieve(self, query: str, top_k: int = 5, rerank_candidates: int = 20, 
                         image_paths: List[str] = None, use_query_expansion: bool = None):
        """增强检索：先检索更多候选，然后重排，支持查询扩展和多模态
        
        Args:
            query: 查询字符串
            top_k: 返回的文档数量
            rerank_candidates: 重排前的候选文档数量
            image_paths: 图像路径列表，用于多模态检索
            use_query_expansion: 是否使用查询扩展，None表示使用配置默认值
        """
        # 确定是否使用查询扩展
        if use_query_expansion is None:
            use_query_expansion = ENABLE_QUERY_EXPANSION
            
        # 如果启用查询扩展，生成扩展查询
        if use_query_expansion:
            expanded_queries = self.expand_query(query)
            print(f"查询扩展: {len(expanded_queries)} 个查询")
        else:
            expanded_queries = [query]
            
        # 多模态检索（如果提供了图像）
        if image_paths and ENABLE_MULTIMODAL:
            return self.multimodal_retrieve(query, image_paths, top_k)
            
        # 混合检索或向量检索
        all_candidate_docs = []
        for expanded_query in expanded_queries:
            if ENABLE_HYBRID_SEARCH:
                # 使用混合检索
                docs = self.hybrid_retrieve(expanded_query, rerank_candidates)
            else:
                # 使用向量检索
                docs = self.retriever.invoke(expanded_query)
                if len(docs) > rerank_candidates:
                    docs = docs[:rerank_candidates]
            
            all_candidate_docs.extend(docs)
            
        # 去重（基于文档内容）
        unique_docs = []
        seen_content = set()
        for doc in all_candidate_docs:
            content = doc.page_content
            if content not in seen_content:
                seen_content.add(content)
                unique_docs.append(doc)
                
        print(f"检索获得 {len(unique_docs)} 个候选文档")
        
        # 重排（如果重排器可用）
        if self.reranker and len(unique_docs) > top_k:
            try:
                reranked_results = self.reranker.rerank(query, unique_docs, top_k)
                final_docs = [doc for doc, score in reranked_results]
                scores = [score for doc, score in reranked_results]
                
                print(f"重排后返回 {len(final_docs)} 个文档")
                print(f"重排分数范围: {min(scores):.4f} - {max(scores):.4f}")
                
                return final_docs
            except Exception as e:
                print(f"⚠️ 重排失败: {e}，使用原始检索结果")
                return unique_docs[:top_k]
        else:
            # 不重排或候选数量不足
            return unique_docs[:top_k]
    
    def compare_retrieval_methods(self, query: str, top_k: int = 5, image_paths: List[str] = None):
        """比较不同检索方法的效果"""
        if not self.retriever:
            return {}
        
        results = {
            'query': query,
            'image_paths': image_paths
        }
        
        # 原始检索 (使用 invoke 替代 get_relevant_documents)
        original_docs = self.retriever.invoke(query)[:top_k]
        results['vector_retrieval'] = {
            'count': len(original_docs),
            'documents': [{
                'content': doc.page_content[:200] + '...' if len(doc.page_content) > 200 else doc.page_content,
                'metadata': getattr(doc, 'metadata', {})
            } for doc in original_docs]
        }
        
        # 混合检索（如果启用）
        if ENABLE_HYBRID_SEARCH and self.ensemble_retriever:
            hybrid_docs = self.hybrid_retrieve(query, top_k)
            results['hybrid_retrieval'] = {
                'count': len(hybrid_docs),
                'documents': [{
                    'content': doc.page_content[:200] + '...' if len(doc.page_content) > 200 else doc.page_content,
                    'metadata': getattr(doc, 'metadata', {})
                } for doc in hybrid_docs]
            }
        
        # 查询扩展检索（如果启用）
        if ENABLE_QUERY_EXPANSION and self.query_expansion_model:
            expanded_docs = self.enhanced_retrieve(query, top_k, use_query_expansion=True)
            results['expanded_query_retrieval'] = {
                'count': len(expanded_docs),
                'documents': [{
                    'content': doc.page_content[:200] + '...' if len(doc.page_content) > 200 else doc.page_content,
                    'metadata': getattr(doc, 'metadata', {})
                } for doc in expanded_docs]
            }
        
        # 多模态检索（如果启用且有图像）
        if ENABLE_MULTIMODAL and image_paths:
            multimodal_docs = self.multimodal_retrieve(query, image_paths, top_k)
            results['multimodal_retrieval'] = {
                'count': len(multimodal_docs),
                'documents': [{
                    'content': doc.page_content[:200] + '...' if len(doc.page_content) > 200 else doc.page_content,
                    'metadata': getattr(doc, 'metadata', {})
                } for doc in multimodal_docs]
            }
        
        # 增强检索（带重排）
        enhanced_docs = self.enhanced_retrieve(query, top_k)
        results['enhanced_retrieval'] = {
            'count': len(enhanced_docs),
            'documents': [{
                'content': doc.page_content[:200] + '...' if len(doc.page_content) > 200 else doc.page_content,
                'metadata': getattr(doc, 'metadata', {})
            } for doc in enhanced_docs]
        }
        
        # 添加配置信息
        results['configuration'] = {
            'hybrid_search_enabled': ENABLE_HYBRID_SEARCH,
            'query_expansion_enabled': ENABLE_QUERY_EXPANSION,
            'multimodal_enabled': ENABLE_MULTIMODAL,
            'reranker_used': self.reranker is not None,
            'hybrid_weights': HYBRID_SEARCH_WEIGHTS if ENABLE_HYBRID_SEARCH else None,
            'multimodal_weights': MULTIMODAL_WEIGHTS if ENABLE_MULTIMODAL else None
        }
        
        return results

    def format_docs(self, docs):
        """格式化文档用于生成"""
        return "\n\n".join(doc.page_content for doc in docs)


def initialize_document_processor():
    """初始化文档处理器并设置知识库"""
    print("🚀 初始化文档处理器 (Milvus 版)...")
    processor = DocumentProcessor()
    
    # 直接设置知识库
    # Milvus 的连接和索引逻辑在 DocumentProcessor.create_vectorstore 中处理
    vectorstore, retriever, doc_splits = processor.setup_knowledge_base()
    
    return processor, vectorstore, retriever, doc_splits