dsProject/dsRag/Util/EsSearchUtil.py

import logging
import os
from logging.handlers import RotatingFileHandler
import jieba
from gensim.models import KeyedVectors

from Config.Config import MODEL_LIMIT, MODEL_PATH, ES_CONFIG
from ElasticSearch.Utils.ElasticsearchConnectionPool import ElasticsearchConnectionPool

# 初始化日志
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
# 确保日志目录存在
os.makedirs('Logs', exist_ok=True)
handler = RotatingFileHandler('Logs/start.log', maxBytes=1024 * 1024, backupCount=5)
handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
logger.addHandler(handler)

class EsSearchUtil:
    def __init__(self, es_config):
        """
                初始化Elasticsearch搜索工具
                :param es_config: Elasticsearch配置字典，包含hosts, username, password, index_name等
                """
        self.es_config = es_config

        # 初始化连接池
        self.es_pool = ElasticsearchConnectionPool(
            hosts=es_config['hosts'],
            basic_auth=es_config['basic_auth'],
            verify_certs=es_config.get('verify_certs', False),
            max_connections=50
        )

        # 保留直接连接用于兼容
        from elasticsearch import Elasticsearch
        self.es = Elasticsearch(
            hosts=es_config['hosts'],
            basic_auth=es_config['basic_auth'],
            verify_certs=es_config.get('verify_certs', False)
        )

        # 确保es_conn属性存在以兼容旧代码
        self.es_conn = self.es

        # 确保es_conn属性存在以兼容旧代码
        self.es_conn = self.es


        # 加载预训练模型
        self.model = KeyedVectors.load_word2vec_format(MODEL_PATH, binary=False, limit=MODEL_LIMIT)
        logger.info(f"模型加载成功，词向量维度: {self.model.vector_size}")
        
        # 初始化Elasticsearch连接
        self.es = Elasticsearch(
            hosts=es_config['hosts'],
            basic_auth=es_config['basic_auth'],
            verify_certs=False
        )
        self.index_name = es_config['index_name']
        
    def text_to_embedding(self, text):
        # 使用已加载的模型
        # 对文本分词并计算平均向量
        words = jieba.lcut(text)
        vectors = [self.model[word] for word in words if word in self.model]
        
        if not vectors:
            return [0.0] * self.model.vector_size
            
        # 计算平均向量
        avg_vector = [sum(dim)/len(vectors) for dim in zip(*vectors)]
        return avg_vector

    def vector_search(self, query, size=10):
        query_embedding = self.text_to_embedding(query)
        script_query = {
            "script_score": {
                "query": {"match_all": {}},
                "script": {
                    "source": "double score = cosineSimilarity(params.query_vector, 'embedding'); return score >= 0 ? score : 0",
                    "params": {"query_vector": query_embedding}
                }
            }
        }
        
        return self.es_conn.search(
            index=self.es_config['index_name'],
            query=script_query,
            size=size
        )

    def text_search(self, query, size=10):
        return self.es_conn.search(
            index=self.es_config['index_name'],
            query={"match": {"user_input": query}},
            size=size
        )

    def hybrid_search(self, query, size=10):
        """
        执行混合搜索（向量搜索+文本搜索）
        :param query: 搜索查询文本
        :param size: 返回结果数量
        :return: 包含两种搜索结果的字典
        """
        vector_results = self.vector_search(query, size)
        text_results = self.text_search(query, size)
        
        return {
            'vector_results': vector_results,
            'text_results': text_results
        }
        
    def search(self, query, search_type='hybrid', size=10):
        """
        统一搜索接口
        :param query: 搜索查询文本
        :param search_type: 搜索类型（'vector', 'text' 或 'hybrid'）
        :param size: 返回结果数量
        :return: 搜索结果
        """
        if search_type == 'vector':
            return self.vector_search(query, size)
        elif search_type == 'text':
            return self.text_search(query, size)
        else:
            return self.hybrid_search(query, size)

    def queryByEs(query, query_tags, logger):
        # 获取EsSearchUtil实例
        es_search_util = EsSearchUtil(ES_CONFIG)

        # 执行混合搜索
        es_conn = es_search_util.es_pool.get_connection()
        try:
            # 向量搜索
            logger.info(f"\n=== 开始执行查询 ===")
            logger.info(f"原始查询文本: {query}")
            logger.info(f"查询标签: {query_tags}")

            logger.info("\n=== 向量搜索阶段 ===")
            logger.info("1. 文本分词和向量化处理中...")
            query_embedding = es_search_util.text_to_embedding(query)
            logger.info(f"2. 生成的查询向量维度: {len(query_embedding)}")
            logger.info(f"3. 前3维向量值: {query_embedding[:3]}")

            logger.info("4. 正在执行Elasticsearch向量搜索...")
            vector_results = es_conn.search(
                index=ES_CONFIG['index_name'],
                body={
                    "query": {
                        "script_score": {
                            "query": {
                                "bool": {
                                    "should": [
                                        {
                                            "terms": {
                                                "tags.tags": query_tags
                                            }
                                        }
                                    ],
                                    "minimum_should_match": 1
                                }
                            },
                            "script": {
                                "source": "double score = cosineSimilarity(params.query_vector, 'embedding'); return score >= 0 ? score : 0",
                                "params": {"query_vector": query_embedding}
                            }
                        }
                    },
                    "size": 3
                }
            )
            # 处理一下，判断是否到达阀值
            filtered_vector_hits = []
            vector_int = 0
            for hit in vector_results['hits']['hits']:
                if hit['_score'] > 0.8:  # 阀值0.8
                    # 新增语义相关性检查
                    if all(word in hit['_source']['user_input'] for word in jieba.lcut(query)):
                        logger.info(f"   {vector_int + 1}. 文档ID: {hit['_id']}, 相似度分数: {hit['_score']:.2f}")
                        logger.info(f"      内容: {hit['_source']['user_input']}")
                        filtered_vector_hits.append(hit)
                        vector_int += 1

            # 更新vector_results只包含通过过滤的文档
            vector_results['hits']['hits'] = filtered_vector_hits
            logger.info(f"5. 向量搜索结果数量(过滤后): {vector_int}")

            # 文本精确搜索
            logger.info("\n=== 文本精确搜索阶段 ===")
            logger.info("1. 正在执行Elasticsearch文本精确搜索...")
            text_results = es_conn.search(
                index=ES_CONFIG['index_name'],
                body={
                    "query": {
                        "bool": {
                            "must": [
                                {
                                    "match": {
                                        "user_input": query
                                    }
                                },
                                {
                                    "terms": {
                                        "tags.tags": query_tags
                                    }
                                }
                            ]
                        }
                    },
                    "size": 3
                }
            )
            logger.info(f"2. 文本搜索结果数量: {len(text_results['hits']['hits'])}")

            # 合并vector和text结果
            all_sources = [hit['_source'] for hit in vector_results['hits']['hits']] + \
                          [hit['_source'] for hit in text_results['hits']['hits']]
            
            # 去重处理
            unique_sources = []
            seen_user_inputs = set()
            
            for source in all_sources:
                if source['user_input'] not in seen_user_inputs:
                    seen_user_inputs.add(source['user_input'])
                    unique_sources.append(source)
            
            logger.info(f"合并后去重结果数量: {len(unique_sources)}条")

            search_results = {
                "text_results": unique_sources
            }
            return search_results
        finally:
            es_search_util.es_pool.release_connection(es_conn)