huanghai/dsProject
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

'commit'

Browse Source
This commit is contained in:
HuangHai
2025-08-19 10:41:58 +08:00
parent 4e47f5eff5
commit 328882505e
8 changed files with 294 additions and 222 deletions
Download Patch File Download Diff File Expand all files Collapse all files

262
dsSchoolBuddy/ElasticSearch/Utils/EsSearchUtil.py
View File

@@ -185,33 +185,33 @@ class EsSearchUtil:
# 2. 从连接池获取连接
conn = search_util.es_pool.get_connection()
# 3. 检查索引是否存在，不存在则创建
# # 3. 检查索引是否存在，不存在则创建
index_name = Config.ES_CONFIG['index_name']
if not conn.indices.exists(index=index_name):
# 定义mapping结构
mapping = {
"mappings": {
"properties": {
"embedding": {
"type": "dense_vector",
"dims": 1024, # 根据实际embedding维度调整
"index": True,
"similarity": "l2_norm"
},
"user_input": {"type": "text"},
"tags": {
"type": "object",
"properties": {
"tags": {"type": "keyword"},
"full_content": {"type": "text"}
}
},
"timestamp": {"type": "date"}
}
}
}
conn.indices.create(index=index_name, body=mapping)
print(f"索引 '{index_name}' 创建成功")
# if not conn.indices.exists(index=index_name):
# # 定义mapping结构
# mapping = {
# "mappings": {
# "properties": {
# "embedding": {
# "type": "dense_vector",
# "dims": Config.EMBED_DIM, # 根据实际embedding维度调整
# "index": True,
# "similarity": "l2_norm"
# },
# "user_input": {"type": "text"},
# "tags": {
# "type": "object",
# "properties": {
# "tags": {"type": "keyword"},
# "full_content": {"type": "text"}
# }
# },
# "timestamp": {"type": "date"}
# }
# }
# }
# conn.indices.create(index=index_name, body=mapping)
# print(f"索引 '{index_name}' 创建成功")
# 4. 切割文本
text_chunks = self.split_text_into_chunks(long_text)
@@ -285,108 +285,128 @@ class EsSearchUtil:
query_embedding = embeddings.embed_query(query)
return query_embedding
def rerank_results(self, query: str, results: List[Dict]) -> List[Tuple[Dict, float]]:
def rerank_results(self, query: str, results: list) -> list:
"""
对搜索结果进行重排
使用重排模型对搜索结果进行重排
参数:
query: 查询文本
results: 搜索结果列表
返回:
list: 重排后的结果列表，每个元素是(文档, 分数)元组
list: 重排后的结果列表，每个元素是(文档对象, 分数)的元组
"""
if len(results) <= 1:
return [(doc, 1.0) for doc in results]
# 准备重排请求数据
rerank_data = {
"model": Config.RERANK_MODEL,
"query": query,
"documents": [doc['_source']['user_input'] for doc in results],
"top_n": len(results)
}
# 调用API进行重排
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {Config.RERANK_BINDING_API_KEY}"
}
if not results:
print("警告: 没有搜索结果可供重排")
return []
try:
# 准备重排请求数据
# 确保doc是字典并包含'_source'和'user_input'字段
documents = []
valid_results = []
for i, doc in enumerate(results):
if isinstance(doc, dict) and '_source' in doc and 'user_input' in doc['_source']:
documents.append(doc['_source']['user_input'])
valid_results.append(doc)
else:
print(f"警告: 结果项 {i} 格式不正确，跳过该结果")
print(f"结果项内容: {doc}")
if not documents:
print("警告: 没有有效的文档可供重排")
# 返回原始结果，但转换为(结果, 分数)的元组格式
return [(doc, doc.get('_score', 0.0)) for doc in results]
rerank_data = {
"model": Config.RERANK_MODEL,
"query": query,
"documents": documents,
"top_n": len(documents)
}
# 调用重排API
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {Config.RERANK_BINDING_API_KEY}"
}
response = requests.post(Config.RERANK_BASE_URL, headers=headers, data=json.dumps(rerank_data))
response.raise_for_status()
response.raise_for_status() # 检查请求是否成功
rerank_result = response.json()
# 处理重排结果
reranked_docs_with_scores = []
reranked_results = []
if "results" in rerank_result:
for item in rerank_result["results"]:
# 尝试获取index和relevance_score字段
doc_idx = item.get("index")
score = item.get("relevance_score", 0.0)
# 如果找不到，尝试fallback到document和score字段
if doc_idx is None:
doc_idx = item.get("document")
if score == 0.0:
score = item.get("score", 0.0)
if 0 <= doc_idx < len(valid_results):
result = valid_results[doc_idx]
reranked_results.append((result, score))
else:
print("警告: 无法识别重排API响应格式")
# 返回原始结果，但转换为(结果, 分数)的元组格式
reranked_results = [(doc, doc.get('_score', 0.0)) for doc in valid_results]
if doc_idx is not None and 0 <= doc_idx < len(results):
reranked_docs_with_scores.append((results[doc_idx], score))
logger.debug(f"重排结果: 文档索引={doc_idx}, 分数={score}")
else:
logger.warning(f"重排结果项索引无效: {doc_idx}")
print(f"重排后结果数量：{len(reranked_results)}")
return reranked_results
# 如果没有有效的重排结果，返回原始结果
if not reranked_docs_with_scores:
logger.warning("没有获取到有效的重排结果，返回原始结果")
return [(doc, 1.0) for doc in results]
return reranked_docs_with_scores
except Exception as e:
logger.error(f"重排失败: {str(e)}")
return [(doc, 1.0) for doc in results]
print(f"重排失败: {e}")
print("将使用原始搜索结果")
# 返回原始结果，但转换为(结果, 分数)的元组格式
return [(doc, doc.get('_score', 0.0)) for doc in results]
def search_by_vector(self, query_embedding: list, k: int = 10) -> dict:
def search_by_vector(self, query_embedding: list, k: int = 10) -> list:
"""
在Elasticsearch中按向量搜索
根据向量进行相似性搜索
参数:
query_embedding: 查询向量
k: 返回结果数量
k: 返回的结果数量
返回:
dict: 搜索结果
list: 搜索结果列表
"""
# 从连接池获取连接
conn = self.es_pool.get_connection()
try:
# 构建向量搜索查询
query = {
"query": {
"script_score": {
"query": {
"bool": {
"should": [],
"minimum_should_match": 0
}
},
"script": {
"source": "double score = cosineSimilarity(params.query_vector, 'embedding'); return score >= 0 ? score : 0",
"params": {"query_vector": query_embedding}
}
}
},
"size": k
}
# 从连接池获取连接
conn = self.es_pool.get_connection()
index_name = Config.ES_CONFIG['index_name']
# 执行向量搜索
response = conn.search(
index=index_name,
body={
"query": {
"script_score": {
"query": {"match_all": {}},
"script": {
"source": "cosineSimilarity(params.query_vector, 'embedding') + 1.0",
"params": {
"query_vector": query_embedding
}
}
}
},
"size": k
}
)
# 提取结果
# 确保我们提取的是 hits.hits 部分
if 'hits' in response and 'hits' in response['hits']:
results = response['hits']['hits']
print(f"向量搜索结果数量: {len(results)}")
return results
else:
print("警告: 向量搜索响应格式不正确")
print(f"响应内容: {response}")
return []
# 执行查询
response = conn.search(index=self.es_config['index_name'], body=query)
return response
except Exception as e:
logger.error(f"向量搜索失败: {str(e)}")
raise
print(f"向量搜索失败: {e}")
return []
finally:
# 释放连接回连接池
self.es_pool.release_connection(conn)
@@ -404,11 +424,53 @@ class EsSearchUtil:
return
print(f"找到 {len(results)} 条结果:\n")
for i, (result, score) in enumerate(results, 1):
for i, item in enumerate(results, 1):
print(f"结果 {i}:")
print(f"内容: {result['_source']['user_input']}")
if show_score:
print(f"分数: {score:.4f}")
try:
# 检查item是否为元组格式 (result, score)
if isinstance(item, tuple):
if len(item) >= 2:
result, score = item[0], item[1]
else:
result, score = item[0], 0.0
else:
# 如果不是元组，假设item就是result
result = item
score = result.get('_score', 0.0)
# 确保result是字典类型
if not isinstance(result, dict):
print(f"警告: 结果项 {i} 不是字典类型，跳过显示")
print(f"结果项内容: {result}")
print("---")
continue
# 尝试获取user_input内容
if '_source' in result and 'user_input' in result['_source']:
content = result['_source']['user_input']
print(f"内容: {content}")
elif 'user_input' in result:
content = result['user_input']
print(f"内容: {content}")
else:
print(f"警告: 结果项 {i} 缺少'user_input'字段")
print(f"结果项内容: {result}")
print("---")
continue
# 显示分数
if show_score:
print(f"分数: {score:.4f}")
# 如果有标签信息，也显示出来
if '_source' in result and 'tags' in result['_source']:
tags = result['_source']['tags']
if isinstance(tags, dict) and 'tags' in tags:
print(f"标签: {tags['tags']}")
except Exception as e:
print(f"处理结果项 {i} 时出错: {str(e)}")
print(f"结果项内容: {item}")
print("---")
def merge_results(self, keyword_results: List[Tuple[Dict, float]], vector_results: List[Tuple[Dict, float]]) -> List[Tuple[Dict, float, str]]: