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
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings

from config import (
    KNOWLEDGE_BASE_URLS,
    CHUNK_SIZE,
    CHUNK_OVERLAP,
    COLLECTION_NAME,
    EMBEDDING_MODEL
)
from reranker import create_reranker


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.reranker = None
        self._setup_reranker()
    
    def _setup_reranker(self):
        """设置重排器"""
        try:
            # 使用混合重排器获得最佳效果
            self.reranker = create_reranker('hybrid', self.embeddings)
            print("✅ 重排器初始化成功")
        except Exception as e:
            print(f"⚠️ 重排器初始化失败: {e}")
            print("将使用基础检索，不进行重排")
    
    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 create_vectorstore(self, doc_splits):
        """创建向量数据库"""
        print("正在创建向量数据库...")
        self.vectorstore = Chroma.from_documents(
            documents=doc_splits,
            collection_name=COLLECTION_NAME,
            embedding=self.embeddings,
        )
        self.retriever = self.vectorstore.as_retriever()
        print("向量数据库创建完成")
        return self.vectorstore, self.retriever
    
    def setup_knowledge_base(self, urls=None, enable_graphrag=False):
        """设置完整的知识库（加载、分割、向量化）
        
        Args:
            urls: 文档URL列表
            enable_graphrag: 是否启用GraphRAG索引
            
        Returns:
            vectorstore, retriever, doc_splits
        """
        docs = self.load_documents(urls)
        doc_splits = self.split_documents(docs)
        vectorstore, retriever = self.create_vectorstore(doc_splits)
        
        # 返回doc_splits用于GraphRAG索引
        return vectorstore, retriever, doc_splits
    
    def enhanced_retrieve(self, query: str, top_k: int = 5, rerank_candidates: int = 20):
        """增强检索：先检索更多候选，然后重排"""
        if not self.retriever:
            print("⚠️ 检索器未初始化")
            return []
        
        # 1. 初始检索：获取更多候选文档
        initial_docs = self.retriever.get_relevant_documents(query)
        
        # 获取更多候选（如果可能）
        if hasattr(self.retriever, 'search_kwargs'):
            # 修改检索参数以获取更多结果
            original_k = self.retriever.search_kwargs.get('k', 4)
            self.retriever.search_kwargs['k'] = min(rerank_candidates, len(initial_docs))
            candidate_docs = self.retriever.get_relevant_documents(query)
            self.retriever.search_kwargs['k'] = original_k  # 恢复原设置
        else:
            candidate_docs = initial_docs
        
        print(f"初始检索获得 {len(candidate_docs)} 个候选文档")
        
        # 2. 重排（如果重排器可用）
        if self.reranker and len(candidate_docs) > top_k:
            try:
                reranked_results = self.reranker.rerank(query, candidate_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 candidate_docs[:top_k]
        else:
            # 不重排或候选数量不足
            return candidate_docs[:top_k]
    
    def compare_retrieval_methods(self, query: str, top_k: int = 5):
        """比较不同检索方法的效果"""
        if not self.retriever:
            return {}
        
        # 原始检索
        original_docs = self.retriever.get_relevant_documents(query)[:top_k]
        
        # 增强检索（带重排）
        enhanced_docs = self.enhanced_retrieve(query, top_k)
        
        return {
            'query': query,
            'original_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]
            },
            '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]
            },
            'reranker_used': self.reranker is not None
        }

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


def initialize_document_processor():
    """初始化文档处理器并设置知识库"""
    processor = DocumentProcessor()
    vectorstore, retriever, doc_splits = processor.setup_knowledge_base()
    return processor, vectorstore, retriever, doc_splits