'commit'

dsSchoolBuddy/ElasticSearch/T2_Vector.py

@@ -1,5 +1,5 @@
 # pip install pydantic requests
-from Util.VectorUtil import text_to_vector_db, query_vector_db
+from ElasticSearch.Utils.VectorUtil import text_to_vector_db, query_vector_db
 
 
 def main():

dsSchoolBuddy/ElasticSearch/T7_XiangLiangQuery.py

@@ -1,76 +1,14 @@
 import logging
 import warnings
-import json
-import requests
 from typing import List, Tuple, Dict
 
-from elasticsearch import Elasticsearch
-
-from Config import Config
-from Config.Config import ES_CONFIG, EMBED_MODEL_NAME, EMBED_BASE_URL, EMBED_API_KEY, RERANK_MODEL, RERANK_BASE_URL, RERANK_BINDING_API_KEY
-from langchain_openai import OpenAIEmbeddings
-from pydantic import SecretStr
+from Config.Config import ES_CONFIG
+from ElasticSearch.Utils.EsSearchUtil import EsSearchUtil
 
 # 初始化日志
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
-# 抑制HTTPS相关警告
-warnings.filterwarnings('ignore', message='Connecting to .* using TLS with verify_certs=False is insecure')
-warnings.filterwarnings('ignore', message='Unverified HTTPS request is being made to host')
-
-
-def text_to_embedding(text: str) -> List[float]:
-    """
-    将文本转换为嵌入向量
-    """
-    embeddings = OpenAIEmbeddings(
-        model=EMBED_MODEL_NAME,
-        base_url=EMBED_BASE_URL,
-        api_key=SecretStr(EMBED_API_KEY)
-    )
-    return embeddings.embed_query(text)
-
-
-def rerank_results(query: str, results: List[Dict]) -> List[Tuple[Dict, float]]:
-    """
-    对搜索结果进行重排
-    """
-    if len(results) <= 1:
-        return [(doc, 1.0) for doc in results]
-
-    # 准备重排请求数据
-    rerank_data = {
-        "model": RERANK_MODEL,
-        "query": query,
-        "documents": [doc['_source']['user_input'] for doc in results],
-        "top_n": len(results)
-    }
-
-    # 调用SiliconFlow API进行重排
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {RERANK_BINDING_API_KEY}"
-    }
-
-    try:
-        response = requests.post(RERANK_BASE_URL, headers=headers, data=json.dumps(rerank_data))
-        response.raise_for_status()
-        rerank_result = response.json()
-
-        # 处理重排结果
-        reranked_docs_with_scores = []
-        if "results" in rerank_result:
-            for item in rerank_result["results"]:
-                doc_idx = item.get("index")
-                score = item.get("relevance_score", 0.0)
-                if 0 <= doc_idx < len(results):
-                    reranked_docs_with_scores.append((results[doc_idx], score))
-        return reranked_docs_with_scores
-    except Exception as e:
-        logger.error(f"重排失败: {str(e)}")
-        return [(doc, 1.0) for doc in results]
-
 
 def merge_results(keyword_results: List[Tuple[Dict, float]], vector_results: List[Tuple[Dict, float]]) -> List[Tuple[Dict, float, str]]:
     """
@@ -97,11 +35,7 @@ def merge_results(keyword_results: List[Tuple[Dict, float]], vector_results: Lis
 
 if __name__ == "__main__":
     # 初始化EsSearchUtil
-    esClient = Elasticsearch(
-        hosts=Config.ES_CONFIG['hosts'],
-        basic_auth=Config.ES_CONFIG['basic_auth'],
-        verify_certs=False
-    )
+    search_util = EsSearchUtil(ES_CONFIG)
 
     # 获取用户输入
     user_query = input("请输入查询语句（例如：高性能的混凝土）: ")
@@ -114,83 +48,34 @@ if __name__ == "__main__":
     print(f"\n=== 开始执行查询 ===")
     print(f"原始查询文本: {user_query}")
 
-    # 执行搜索
-    es_conn = esClient.es_pool.get_connection()
     try:
         # 1. 向量搜索
         print("\n=== 向量搜索阶段 ===")
         print("1. 文本向量化处理中...")
-        query_embedding = text_to_embedding(user_query)
+        query_embedding = search_util.get_query_embedding(user_query)
         print(f"2. 生成的查询向量维度: {len(query_embedding)}")
         print(f"3. 前3维向量值: {query_embedding[:3]}")
 
         print("4. 正在执行Elasticsearch向量搜索...")
-        vector_results = es_conn.search(
-            index=ES_CONFIG['index_name'],
-            body={
-                "query": {
-                    "script_score": {
-                        "query": {
-                            "bool": {
-                                "should": [
-                                    {
-                                        "terms": {
-                                            "tags.tags": query_tags
-                                        }
-                                    }
-                                ] if query_tags else {"match_all": {}},
-                                "minimum_should_match": 1 if query_tags else 0
-                            }
-                        },
-                        "script": {
-                            "source": "double score = cosineSimilarity(params.query_vector, 'embedding'); return score >= 0 ? score : 0",
-                            "params": {"query_vector": query_embedding}
-                        }
-                    }
-                },
-                "size": 5
-            }
-        )
+        vector_results = search_util.search_by_vector(query_embedding, k=5)
         vector_hits = vector_results['hits']['hits']
         print(f"5. 向量搜索结果数量: {len(vector_hits)}")
 
         # 向量结果重排
         print("6. 正在进行向量结果重排...")
-        reranked_vector_results = rerank_results(user_query, vector_hits)
+        reranked_vector_results = search_util.rerank_results(user_query, vector_hits)
         print(f"7. 重排后向量结果数量: {len(reranked_vector_results)}")
 
         # 2. 关键字搜索
         print("\n=== 关键字搜索阶段 ===")
         print("1. 正在执行Elasticsearch关键字搜索...")
-        keyword_results = es_conn.search(
-            index=ES_CONFIG['index_name'],
-            body={
-                "query": {
-                    "bool": {
-                        "must": [
-                            {
-                                "match": {
-                                    "user_input": user_query
-                                }
-                            }
-                        ] + ([
-                            {
-                                "terms": {
-                                    "tags.tags": query_tags
-                                }
-                            }
-                        ] if query_tags else [])
-                    }
-                },
-                "size": 5
-            }
-        )
+        keyword_results = search_util.text_search(user_query, size=5)
         keyword_hits = keyword_results['hits']['hits']
         print(f"2. 关键字搜索结果数量: {len(keyword_hits)}")
 
         # 3. 合并结果
         print("\n=== 合并搜索结果 ===")
-        # 为关键字结果添加默认分数1.0
+        # 为关键字结果添加分数
         keyword_results_with_scores = [(doc, doc['_score']) for doc in keyword_hits]
         merged_results = merge_results(keyword_results_with_scores, reranked_vector_results)
         print(f"合并后唯一结果数量: {len(merged_results)}")
@@ -205,5 +90,3 @@ if __name__ == "__main__":
     except Exception as e:
         logger.error(f"搜索过程中发生错误: {str(e)}")
         print(f"搜索失败: {str(e)}")
-    finally:
-        esClient.es_pool.release_connection(es_conn)

dsSchoolBuddy/ElasticSearch/Utils/EsSearchUtil.py

@@ -13,7 +13,7 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_openai import OpenAIEmbeddings
 from pydantic import SecretStr
 from Config import Config
-
+from typing import List, Tuple, Dict
 # 初始化日志
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -285,72 +285,29 @@ class EsSearchUtil:
         query_embedding = embeddings.embed_query(query)
         return query_embedding
 
-    def search_by_vector(self, query_embedding: list, k: int = 10) -> list:
+    def rerank_results(self, query: str, results: List[Dict]) -> List[Tuple[Dict, float]]:
         """
-        在Elasticsearch中按向量搜索
-
-        参数:
-            query_embedding: 查询向量
-            k: 返回结果数量
-
-        返回:
-            list: 搜索结果
-        """
-        # 从连接池获取连接
-        conn = self.es_pool.get_connection()
-
-        try:
-            # 构建向量查询DSL
-            query = {
-                "query": {
-                    "script_score": {
-                        "query": {"match_all": {}},
-                        "script": {
-                            "source": "cosineSimilarity(params.query_vector, 'embedding') + 1.0",
-                            "params": {
-                                "query_vector": query_embedding
-                            }
-                        }
-                    }
-                },
-                "size": k
-            }
-
-            # 执行查询
-            response = conn.search(index=self.es_config['index_name'], body=query)
-            return response['hits']['hits']
-        except Exception as e:
-            logger.error(f"向量查询失败: {e}")
-            print(f"向量查询失败: {e}")
-            return []
-        finally:
-            # 释放连接回连接池
-            self.es_pool.release_connection(conn)
-
-    def rerank_results(self, query: str, results: list) -> list:
-        """
-        使用重排模型对结果进行排序
+        对搜索结果进行重排
 
         参数:
             query: 查询文本
-            results: 初始搜索结果
+            results: 搜索结果列表
 
         返回:
-            list: 重排后的结果
+            list: 重排后的结果列表，每个元素是(文档, 分数)元组
         """
         if len(results) <= 1:
-            # 结果太少，无需重排
-            return [(result, 1.0) for result in results]
+            return [(doc, 1.0) for doc in results]
 
         # 准备重排请求数据
         rerank_data = {
             "model": Config.RERANK_MODEL,
             "query": query,
-            "documents": [result['_source']['user_input'] for result in results],
+            "documents": [doc['_source']['user_input'] for doc in results],
             "top_n": len(results)
         }
 
-        # 调用重排API
+        # 调用API进行重排
         headers = {
             "Content-Type": "application/json",
             "Authorization": f"Bearer {Config.RERANK_BINDING_API_KEY}"
@@ -361,45 +318,78 @@ class EsSearchUtil:
             response.raise_for_status()
             rerank_result = response.json()
 
-            # 检查响应结构
-            if 'results' not in rerank_result:
-                logger.error(f"重排API响应结构不正确，缺少'results'字段: {rerank_result}")
-                print(f"重排API响应结构不正确，缺少'results'字段")
-                return [(result, 1.0) for result in results]
+            # 处理重排结果
+            reranked_docs_with_scores = []
+            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)
 
-            # 构建重排后的结果列表
-            reranked_pairs = []
-            for item in rerank_result['results']:
-                # 尝试获取文档索引，优先使用'index'字段，其次是'document'字段
-                doc_idx = item.get('index', item.get('document', -1))
-                if doc_idx == -1:
-                    logger.error(f"重排结果项缺少有效索引字段: {item}")
-                    print(f"重排结果项结构不正确")
-                    continue
-
-                # 尝试获取分数，优先使用'relevance_score'字段，其次是'score'字段
-                score = item.get('relevance_score', item.get('score', 1.0))
-
-                # 检查索引是否有效
-                if 0 <= doc_idx < len(results):
-                    reranked_pairs.append((results[doc_idx], score))
-                else:
-                    logger.error(f"文档索引{doc_idx}超出范围")
-                    print(f"文档索引超出范围")
+                    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}")
 
             # 如果没有有效的重排结果，返回原始结果
-            if not reranked_pairs:
-                logger.warning("没有有效的重排结果，返回原始结果")
-                return [(result, 1.0) for result in results]
+            if not reranked_docs_with_scores:
+                logger.warning("没有获取到有效的重排结果，返回原始结果")
+                return [(doc, 1.0) for doc in results]
 
-            # 按分数降序排序
-            reranked_pairs.sort(key=lambda x: x[1], reverse=True)
-            return reranked_pairs
+            return reranked_docs_with_scores
         except Exception as e:
             logger.error(f"重排失败: {str(e)}")
-            print(f"重排失败: {e}")
-            # 重排失败时返回原始结果
-            return [(result, 1.0) for result in results]
+            return [(doc, 1.0) for doc in results]
+
+    def search_by_vector(self, query_embedding: list, k: int = 10) -> dict:
+        """
+        在Elasticsearch中按向量搜索
+
+        参数:
+            query_embedding: 查询向量
+            k: 返回结果数量
+
+        返回:
+            dict: 搜索结果
+        """
+        # 从连接池获取连接
+        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
+            }
+
+            # 执行查询
+            response = conn.search(index=self.es_config['index_name'], body=query)
+            return response
+        except Exception as e:
+            logger.error(f"向量搜索失败: {str(e)}")
+            raise
+        finally:
+            # 释放连接回连接池
+            self.es_pool.release_connection(conn)
 
     def display_results(self, results: list, show_score: bool = True) -> None:
         """