'commit'

dsSchoolBuddy/ElasticSearch/T3_InsertData.py

@@ -1,75 +1,144 @@
-import warnings
-
-from Config import Config
-from Config.Config import *
-from elasticsearch import Elasticsearch
-from gensim.models import KeyedVectors
-import jieba
 import os
 import time
+import warnings
+
+from elasticsearch import Elasticsearch
+from langchain_core.documents import Document
+
+from Config import Config
+from Util.VectorUtil import text_to_vector_db  # 导入向量化工具函数
 
 # 抑制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')
 
-# 1. 加载预训练的 Word2Vec 模型
-model_path = MODEL_PATH
-model = KeyedVectors.load_word2vec_format(model_path, binary=False, limit=MODEL_LIMIT)
-print(f"模型加载成功，词向量维度: {model.vector_size}")
 
+def insert_long_text_to_es(long_text: str, tags: list = None) -> bool:
+    """
+    将长文本向量化并插入到Elasticsearch
+    
+    参数:
+        long_text: 要插入的长文本
+        tags: 可选的标签列表
+    
+    返回:
+        bool: 插入是否成功
+    """
+    try:
+        # 1. 初始化Elasticsearch连接
+        es = Elasticsearch(
+            hosts=Config.ES_CONFIG['hosts'],
+            basic_auth=Config.ES_CONFIG['basic_auth'],
+            verify_certs=False
+        )
 
-# 功能：将文本转换为嵌入向量
+        # 2. 检查索引是否存在，不存在则创建
-def text_to_embedding(text):
+        index_name = Config.ES_CONFIG['index_name']
-    words = jieba.lcut(text)
+        if not es.indices.exists(index=index_name):
-    embeddings = [model[word] for word in words if word in model]
+            # 定义mapping结构
-    if embeddings:
+            mapping = {
-        return sum(embeddings) / len(embeddings)
+                "mappings": {
-    return [0.0] * model.vector_size
+                    "properties": {
-
+                        "embedding": {
-
+                            "type": "dense_vector",
-# 2. 初始化Elasticsearch连接
+                            "dims": 1024,  # 根据实际embedding维度调整
-es = Elasticsearch(
+                            "index": True,
-    hosts=Config.ES_CONFIG['hosts'],
+                            "similarity": "l2_norm"
-    basic_auth=Config.ES_CONFIG['basic_auth'],
+                        },
-    verify_certs=False
+                        "user_input": {"type": "text"},
-)
+                        "tags": {
-
+                            "type": "object",
-# 3. 处理processed_chunks目录下的所有文件
+                            "properties": {
-txt_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'Txt')
+                                "tags": {"type": "keyword"},
-
+                                "full_content": {"type": "text"}
-for filename in os.listdir(txt_dir):
+                            }
-    if filename.endswith('.txt'):
+                        },
-        filepath = os.path.join(txt_dir, filename)
+                        "timestamp": {"type": "date"}
-        with open(filepath, 'r', encoding='utf-8') as f:
+                    }
-            # 只读取第一行作为向量计算
+                }
-            first_line = f.readline().strip()
-            # 读取全部内容用于后续查询
-            full_content = first_line + '\n' + f.read()
-
-            if not first_line:
-                print(f"跳过空文件: {filename}")
-                continue
-
-            print(f"正在处理文件: {filename}")
-
-            # 4. 获取当前时间和会话ID
-            timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
-            # 需要进行标记的标签
-            x = filename.split("_")
-            selectedTags = [x[0] + "_" + x[1]]
-            tags = {"tags": selectedTags, "full_content": full_content}  # 添加完整内容
-
-            # 5. 将第一行文本转换为嵌入向量
-            embedding = text_to_embedding(first_line)
-
-            # 6. 插入数据到Elasticsearch
-            doc = {
-                'tags': tags,
-                'user_input': first_line,
-                'timestamp': timestamp,
-                'embedding': embedding
             }
-            es.index(index=ES_CONFIG['index_name'], document=doc)
+            es.indices.create(index=index_name, body=mapping)
-            print(f"文件 {filename} 数据插入成功")
+            print(f"索引 '{index_name}' 创建成功")
 
-print("所有文件处理完成")
+        # 3. 使用VectorUtil处理文本
+        # 这里我们创建一个临时的Document对象
+        docs = [Document(page_content=long_text, metadata={"source": "user_provided_text"})]
+
+        # 4. 获取当前时间
+        timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
+
+        # 5. 准备标签
+        if tags is None:
+            tags = ["general_text"]
+        tags_dict = {"tags": tags, "full_content": long_text}
+
+        # 6. 使用VectorUtil中的功能获取嵌入向量
+        # 注意：我们需要修改text_to_vector_db函数或创建新函数来获取单个文档的嵌入
+        # 这里为了演示，我们直接调用text_to_vector_db并获取第一个文档的嵌入
+        vector_store, _, _ = text_to_vector_db(long_text)
+
+        # 7. 提取嵌入向量
+        # 注意：这里假设我们只处理一个文档
+        embedding = vector_store._embeddings.embed_documents([long_text])[0]
+
+        # 8. 插入数据到Elasticsearch
+        doc = {
+            'tags': tags_dict,
+            'user_input': long_text[:500],  # 取前500个字符作为摘要
+            'timestamp': timestamp,
+            'embedding': embedding
+        }
+        es.index(index=index_name, document=doc)
+        print(f"长文本数据插入成功")
+        return True
+    except Exception as e:
+        print(f"插入数据失败: {e}")
+        return False
+
+
+def process_text_directory(txt_dir: str) -> None:
+    """
+    处理指定目录下的所有文本文件，将其向量化并插入到Elasticsearch
+    
+    参数:
+        txt_dir: 包含文本文件的目录路径
+    """
+    for filename in os.listdir(txt_dir):
+        if filename.endswith('.txt'):
+            filepath = os.path.join(txt_dir, filename)
+            with open(filepath, 'r', encoding='utf-8') as f:
+                full_content = f.read()
+
+                if not full_content:
+                    print(f"跳过空文件: {filename}")
+                    continue
+
+                print(f"正在处理文件: {filename}")
+
+                # 提取标签
+                x = filename.split("_")
+                if len(x) >= 2:
+                    selected_tags = [x[0] + "_" + x[1]]
+                else:
+                    selected_tags = ["uncategorized"]
+
+                # 插入文本
+                insert_long_text_to_es(full_content, selected_tags)
+
+
+def main():
+    # 示例1：插入单个长文本
+    long_text = """混凝土是一种广泛使用的建筑材料，由水泥、砂、石子和水混合而成。它具有高强度、耐久性和良好的可塑性，被广泛应用于建筑、桥梁、道路等土木工程领域。
+
+混凝土的历史可以追溯到古罗马时期，当时人们使用火山灰、石灰和碎石混合制成类似混凝土的材料。现代混凝土技术始于19世纪，随着波特兰水泥的发明而得到快速发展。
+
+混凝土的性能取决于其配合比，包括水灰比、砂率等参数。水灰比是影响混凝土强度的关键因素，较小的水灰比通常会产生更高强度的混凝土。"""
+    insert_long_text_to_es(long_text, tags=["construction", "materials"])
+
+    # 示例2：处理目录中的所有文本文件
+    # txt_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'Txt')
+    # process_text_directory(txt_dir)
+
+
+if __name__ == "__main__":
+    main()