This commit is contained in:
2025-08-19 09:21:32 +08:00
parent 11d4b33170
commit 305bd177bb
3 changed files with 22 additions and 21 deletions

View File

@@ -1,8 +1,7 @@
from Config import Config from Config import Config
from ElasticSearch.Utils.EsSearchUtil import EsSearchUtil, disableWarning from ElasticSearch.Utils.EsSearchUtil import EsSearchUtil, disableWarning
# 禁用警告
disableWarning()
# 创建EsSearchUtil实例 # 创建EsSearchUtil实例
search_util = EsSearchUtil(Config.ES_CONFIG) search_util = EsSearchUtil(Config.ES_CONFIG)

View File

@@ -18,9 +18,6 @@ def disableWarning():
warnings.filterwarnings('ignore', message='Connecting to .* using TLS with verify_certs=False is insecure') 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') warnings.filterwarnings('ignore', message='Unverified HTTPS request is being made to host')
# 初始化配置
disableWarning()
# 初始化日志 # 初始化日志
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO) logger.setLevel(logging.INFO)
@@ -31,6 +28,8 @@ class EsSearchUtil:
初始化Elasticsearch搜索工具 初始化Elasticsearch搜索工具
:param es_config: Elasticsearch配置字典包含hosts, username, password, index_name等 :param es_config: Elasticsearch配置字典包含hosts, username, password, index_name等
""" """
# 禁用警告
disableWarning()
self.es_config = es_config self.es_config = es_config
# 初始化连接池 # 初始化连接池
@@ -154,16 +153,15 @@ def insert_long_text_to_es(long_text: str, tags: list = None) -> bool:
bool: 插入是否成功 bool: 插入是否成功
""" """
try: try:
# 1. 初始化Elasticsearch连接 # 1. 创建EsSearchUtil实例以使用连接
es = Elasticsearch( search_util = EsSearchUtil(Config.ES_CONFIG)
hosts=Config.ES_CONFIG['hosts'],
basic_auth=Config.ES_CONFIG['basic_auth'], # 2. 从连接池获取连接
verify_certs=False conn = search_util.es_pool.get_connection()
)
# 2. 检查索引是否存在,不存在则创建 # 3. 检查索引是否存在,不存在则创建
index_name = Config.ES_CONFIG['index_name'] index_name = Config.ES_CONFIG['index_name']
if not es.indices.exists(index=index_name): if not conn.indices.exists(index=index_name):
# 定义mapping结构 # 定义mapping结构
mapping = { mapping = {
"mappings": { "mappings": {
@@ -186,33 +184,33 @@ def insert_long_text_to_es(long_text: str, tags: list = None) -> bool:
} }
} }
} }
es.indices.create(index=index_name, body=mapping) conn.indices.create(index=index_name, body=mapping)
print(f"索引 '{index_name}' 创建成功") print(f"索引 '{index_name}' 创建成功")
# 3. 切割文本 # 4. 切割文本
text_chunks = split_text_into_chunks(long_text) text_chunks = split_text_into_chunks(long_text)
# 4. 准备标签 # 5. 准备标签
if tags is None: if tags is None:
tags = ["general_text"] tags = ["general_text"]
# 5. 获取当前时间 # 6. 获取当前时间
timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
# 6. 创建嵌入模型 # 7. 创建嵌入模型
embeddings = OpenAIEmbeddings( embeddings = OpenAIEmbeddings(
model=Config.EMBED_MODEL_NAME, model=Config.EMBED_MODEL_NAME,
base_url=Config.EMBED_BASE_URL, base_url=Config.EMBED_BASE_URL,
api_key=SecretStr(Config.EMBED_API_KEY) api_key=SecretStr(Config.EMBED_API_KEY)
) )
# 7. 为每个文本块生成向量并插入 # 8. 为每个文本块生成向量并插入
for i, chunk in enumerate(text_chunks): for i, chunk in enumerate(text_chunks):
# 生成文本块的哈希值作为文档ID # 生成文本块的哈希值作为文档ID
doc_id = hashlib.md5(chunk.encode('utf-8')).hexdigest() doc_id = hashlib.md5(chunk.encode('utf-8')).hexdigest()
# 检查文档是否已存在 # 检查文档是否已存在
if es.exists(index=index_name, id=doc_id): if conn.exists(index=index_name, id=doc_id):
print(f"文档块 {i+1} 已存在,跳过插入: {doc_id}") print(f"文档块 {i+1} 已存在,跳过插入: {doc_id}")
continue continue
@@ -228,13 +226,17 @@ def insert_long_text_to_es(long_text: str, tags: list = None) -> bool:
} }
# 插入数据到Elasticsearch # 插入数据到Elasticsearch
es.index(index=index_name, id=doc_id, document=doc) conn.index(index=index_name, id=doc_id, document=doc)
print(f"文档块 {i+1} 插入成功: {doc_id}") print(f"文档块 {i+1} 插入成功: {doc_id}")
return True return True
except Exception as e: except Exception as e:
print(f"插入数据失败: {e}") print(f"插入数据失败: {e}")
return False return False
finally:
# 确保释放连接回连接池
if 'conn' in locals() and 'search_util' in locals():
search_util.es_pool.release_connection(conn)