'commit'
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
import time
|
||||
import warnings
|
||||
import hashlib # 导入哈希库
|
||||
|
||||
from elasticsearch import Elasticsearch
|
||||
from langchain_core.documents import Document
|
||||
|
||||
@@ -1,238 +0,0 @@
|
||||
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)
|
||||
Binary file not shown.
Reference in New Issue
Block a user