AIEC-RAG---/AIEC-RAG/retriver/langgraph/es_vector_retriever.py

"""
ES向量检索器
用于直接与ES向量库进行向量匹配检索
"""

import os
import sys
from typing import List, Dict, Any, Optional
from langchain_core.documents import Document

# 添加路径
project_root = os.path.join(os.path.dirname(__file__), '..', '..')
sys.path.append(project_root)

from retriver.langgraph.dashscope_embedding import DashScopeEmbeddingModel
from elasticsearch_vectorization.es_client_wrapper import ESClientWrapper
from elasticsearch_vectorization.config import ElasticsearchConfig


class ESVectorRetriever:
    """ES向量检索器，用于直接进行向量相似度匹配"""
    
    def __init__(
        self,
        keyword: str,
        top_k: int = 3,
        oneapi_key: Optional[str] = None,
        oneapi_base_url: Optional[str] = None,
        embed_model_name: Optional[str] = None
    ):
        """
        初始化ES向量检索器
        
        Args:
            keyword: ES索引关键词
            top_k: 返回的文档数量
            oneapi_key: OneAPI密钥
            oneapi_base_url: OneAPI基础URL  
            embed_model_name: 嵌入模型名称
        """
        self.keyword = keyword
        self.top_k = top_k
        
        # 初始化嵌入模型
        self.embedding_model = DashScopeEmbeddingModel(
            api_key=oneapi_key,
            model_name=embed_model_name
        )
        
        # 初始化ES客户端
        self.es_client = ESClientWrapper()
        
        # 设置索引名称
        self.passages_index = ElasticsearchConfig.get_passages_index_name(keyword)
        
        print(f"ES向量检索器初始化完成，目标索引: {self.passages_index}")
    
    def retrieve(self, query: str) -> List[Document]:
        """
        检索相关文档
        
        Args:
            query: 查询文本
            
        Returns:
            检索到的文档列表
        """
        try:
            # 生成查询向量
            query_embedding = self.embedding_model.encode([query], normalize_embeddings=True)[0]
            
            # 确保向量是列表格式
            if hasattr(query_embedding, 'tolist'):
                query_vector = query_embedding.tolist()
            else:
                query_vector = list(query_embedding)
            
            # 执行向量搜索
            search_result = self.es_client.vector_search(
                index_name=self.passages_index,
                vector=query_vector,
                field="embedding",
                size=self.top_k
            )
            
            # 解析搜索结果
            documents = []
            hits = search_result.get("hits", {}).get("hits", [])
            
            for hit in hits:
                source = hit["_source"]
                score = hit["_score"]
                
                # 创建Document对象
                doc = Document(
                    page_content=source.get("content", ""),
                    metadata={
                        "passage_id": source.get("passage_id", ""),
                        "file_id": source.get("file_id", ""),
                        "evidence": source.get("evidence", ""),
                        "score": score,
                        "source": "es_vector_search"
                    }
                )
                documents.append(doc)
            
            print(f"ES向量检索完成，找到 {len(documents)} 个相关文档")
            return documents
            
        except Exception as e:
            print(f"ES向量检索失败: {e}")
            return []
    
    def test_connection(self) -> bool:
        """测试ES连接"""
        try:
            return self.es_client.ping()
        except:
            return False
    
    def get_index_stats(self) -> Dict[str, Any]:
        """获取索引统计信息"""
        try:
            query = {"match_all": {}}
            result = self.es_client.search(self.passages_index, query, size=0)
            total = result.get("hits", {}).get("total", 0)
            
            # 兼容不同ES版本的total格式
            if isinstance(total, dict):
                count = total.get("value", 0)
            else:
                count = total
                
            return {
                "index_name": self.passages_index,
                "document_count": count,
                "top_k": self.top_k
            }
        except Exception as e:
            print(f"获取索引统计信息失败: {e}")
            return {
                "index_name": self.passages_index,
                "document_count": 0,
                "top_k": self.top_k,
                "error": str(e)
            }


def create_es_vector_retriever(
    keyword: str,
    top_k: int = 3,
    **kwargs
) -> ESVectorRetriever:
    """
    创建ES向量检索器的便捷函数
    
    Args:
        keyword: ES索引关键词
        top_k: 返回的文档数量
        **kwargs: 其他参数
        
    Returns:
        ES向量检索器实例
    """
    return ESVectorRetriever(
        keyword=keyword,
        top_k=top_k,
        **kwargs
    )