|
|
@@ -0,0 +1,391 @@
|
|
|
+from pymilvus import Collection, connections, utility, FieldSchema, CollectionSchema, DataType
|
|
|
+import pymysql.cursors
|
|
|
+import time
|
|
|
+import requests
|
|
|
+
|
|
|
+# import pymysql
|
|
|
+from tqdm import tqdm
|
|
|
+from openai import OpenAI
|
|
|
+from typing import List, Dict, Any, Optional
|
|
|
+from cosine_similarity_example import cosine_similarity, cosine_similarity_numpy
|
|
|
+from datetime import datetime, timedelta
|
|
|
+################################连接milvus数据库 A
|
|
|
+# 配置信息
|
|
|
+MILVUS_CONFIG = {
|
|
|
+ "host": "c-981be0ee7225467b.milvus.aliyuncs.com",
|
|
|
+ "user": "root",
|
|
|
+ "password": "Piaoquan@2025",
|
|
|
+ "port": "19530",
|
|
|
+}
|
|
|
+print("正在连接 Milvus 数据库...")
|
|
|
+connections.connect("default", **MILVUS_CONFIG)
|
|
|
+################################连接milvus数据库 B
|
|
|
+
|
|
|
+def iterate_milvus_collection(collection, batch_size=1000):
|
|
|
+ """遍历Milvus集合的函数"""
|
|
|
+ # 获取集合中的实体总数
|
|
|
+ total_entities = collection.num_entities
|
|
|
+ print(f"集合中总共有 {total_entities} 条记录")
|
|
|
+
|
|
|
+ # 分页遍历
|
|
|
+ offset = 0
|
|
|
+ count = 0
|
|
|
+ batch_size =1
|
|
|
+
|
|
|
+ while offset < total_entities:
|
|
|
+ # 使用query方法进行分页查询
|
|
|
+ # expr="" 表示查询所有记录
|
|
|
+ # limit控制每页返回的记录数
|
|
|
+ # offset控制从第几条记录开始查询
|
|
|
+ results = collection.query(
|
|
|
+ expr="",
|
|
|
+ limit=batch_size,
|
|
|
+ offset=offset,
|
|
|
+ output_fields=["id", "doc_id", "chunk_id", "vector_text","vector_summary","vector_questions"] # 指定要返回的字段
|
|
|
+ )
|
|
|
+
|
|
|
+ # 处理当前批次的结果
|
|
|
+ for item in results:
|
|
|
+ count += 1
|
|
|
+ print(f"记录 {count}: id={item['id']}, doc_id={item['doc_id']}, chunk_id={item['chunk_id']}")
|
|
|
+ # print("item['vector_summary'] is ", item['vector_summary'])
|
|
|
+ # print("item['vector_questions'] is ", item['vector_questions'])
|
|
|
+ # 这里可以根据需要处理每条记录
|
|
|
+ # exit()
|
|
|
+ # 更新偏移量
|
|
|
+ offset += len(results)
|
|
|
+
|
|
|
+ # 如果本批次结果为空，说明已经遍历完毕
|
|
|
+ if len(results) == 0:
|
|
|
+ break
|
|
|
+
|
|
|
+################################连接mysql数据库 A
|
|
|
+MYSQL_CONFIG = {
|
|
|
+ "host": "rm-bp13g3ra2f59q49xs.mysql.rds.aliyuncs.com",
|
|
|
+ "user": "wqsd",
|
|
|
+ "password": "wqsd@2025",
|
|
|
+ "port": 3306,
|
|
|
+ "db": "rag",
|
|
|
+ "charset": "utf8mb4",
|
|
|
+}
|
|
|
+class DatabaseManager:
|
|
|
+ def __init__(self):
|
|
|
+ self.connection = None
|
|
|
+
|
|
|
+ def connect(self):
|
|
|
+ self.connection = pymysql.connect(
|
|
|
+ **MYSQL_CONFIG,
|
|
|
+ cursorclass=pymysql.cursors.DictCursor
|
|
|
+ )
|
|
|
+ print("✅ MySQL connected")
|
|
|
+
|
|
|
+ def close(self):
|
|
|
+ if self.connection:
|
|
|
+ self.connection.close()
|
|
|
+ print("MySQL closed")
|
|
|
+
|
|
|
+ def fetch(self, query, params=None):
|
|
|
+ with self.connection.cursor() as cursor:
|
|
|
+ cursor.execute(query, params or ())
|
|
|
+ return cursor.fetchall()
|
|
|
+
|
|
|
+ def fetch_one(self, query, params=None):
|
|
|
+ with self.connection.cursor() as cursor:
|
|
|
+ # print("query is ", query)
|
|
|
+ cursor.execute(query)
|
|
|
+ return cursor.fetchone()
|
|
|
+ def execute_ddl(self, query=None):
|
|
|
+ if query is None:
|
|
|
+ return False
|
|
|
+ with self.connection.cursor() as cursor:
|
|
|
+ cursor.execute(query)
|
|
|
+ self.connection.commit()
|
|
|
+ return True
|
|
|
+################################
|
|
|
+
|
|
|
+###########
|
|
|
+mysql_manager = DatabaseManager()
|
|
|
+mysql_manager.connect()
|
|
|
+###########
|
|
|
+
|
|
|
+test_query="set autocommit = 1"
|
|
|
+res = mysql_manager.fetch_one(test_query)
|
|
|
+################################连接mysql数据库 B
|
|
|
+
|
|
|
+################################连接Embedding service A
|
|
|
+VLLM_SERVER_URL = "http://192.168.100.31:8000/v1/embeddings"
|
|
|
+DEFAULT_MODEL = "/models/Qwen3-Embedding-4B"
|
|
|
+
|
|
|
+def get_basic_embedding(text: str, model=DEFAULT_MODEL):
|
|
|
+ """通过HTTP调用在线embedding服务"""
|
|
|
+ headers = {
|
|
|
+ "Content-Type": "application/json"
|
|
|
+ }
|
|
|
+ data = {
|
|
|
+ "model": model,
|
|
|
+ "input": text
|
|
|
+ }
|
|
|
+
|
|
|
+ response = requests.post(
|
|
|
+ VLLM_SERVER_URL,
|
|
|
+ headers=headers,
|
|
|
+ json=data
|
|
|
+ )
|
|
|
+
|
|
|
+ if response.status_code == 200:
|
|
|
+ result = response.json()
|
|
|
+ return result["data"][0]["embedding"]
|
|
|
+ else:
|
|
|
+ raise Exception(f"Failed to get embedding: {response.status_code} {response.text}")
|
|
|
+################################连接Embedding service B
|
|
|
+
|
|
|
+################################连接deepseek v3 A
|
|
|
+DEEPSEEK_API_KEY = "sk-cfd2df92c8864ab999d66a615ee812c5"
|
|
|
+DEEPSEEK_MODEL = {
|
|
|
+ "DeepSeek-R1": "deepseek-reasoner",
|
|
|
+ "DeepSeek-V3": "deepseek-chat",
|
|
|
+}
|
|
|
+
|
|
|
+def get_deepseek_completion(
|
|
|
+ model: str,
|
|
|
+ prompt: str,
|
|
|
+ output_type: str = "text",
|
|
|
+ tool_calls: bool = False,
|
|
|
+ tools: List[Dict] = None
|
|
|
+) -> Optional[Dict | List | str]:
|
|
|
+ """调用DeepSeek API获取回答"""
|
|
|
+ messages = [{"role": "user", "content": prompt}]
|
|
|
+ kwargs = {
|
|
|
+ "model": DEEPSEEK_MODEL.get(model, "deepseek-chat"),
|
|
|
+ "messages": messages,
|
|
|
+ }
|
|
|
+ # 添加工具调用参数
|
|
|
+ if tool_calls and tools:
|
|
|
+ kwargs["tools"] = tools
|
|
|
+ kwargs["tool_choice"] = "auto"
|
|
|
+
|
|
|
+ # 创建OpenAI客户端连接DeepSeek API
|
|
|
+ client = OpenAI(api_key=DEEPSEEK_API_KEY, base_url="https://api.deepseek.com")
|
|
|
+
|
|
|
+ # 设置JSON输出格式
|
|
|
+ if output_type == "json":
|
|
|
+ kwargs["response_format"] = {"type": "json_object"}
|
|
|
+
|
|
|
+ try:
|
|
|
+ response = client.chat.completions.create(**kwargs)
|
|
|
+ choice = response.choices[0]
|
|
|
+
|
|
|
+ if output_type == "text":
|
|
|
+ return choice.message.content # 只返回文本
|
|
|
+ elif output_type == "json":
|
|
|
+ import json
|
|
|
+ return json.loads(choice.message.content)
|
|
|
+ else:
|
|
|
+ raise ValueError(f"Invalid output_type: {output_type}")
|
|
|
+
|
|
|
+ except Exception as e:
|
|
|
+ print(f"[ERROR] DeepSeek API调用失败: {e}")
|
|
|
+ return None
|
|
|
+################################连接deepseek v3 B
|
|
|
+
|
|
|
+#######################################################
|
|
|
+#######################################################开发
|
|
|
+#######################################################
|
|
|
+
|
|
|
+# 方法1: 使用utility.describe_collection()查看集合信息
|
|
|
+collection_name = "chunk_multi_embeddings_v2_trsc"
|
|
|
+print(f"\n📋 正在查看集合 '{collection_name}' 的信息...")
|
|
|
+
|
|
|
+# 正确的方式：创建Collection对象
|
|
|
+milvus_client = Collection(name=collection_name)
|
|
|
+print(milvus_client)
|
|
|
+milvus_client.load()
|
|
|
+total = milvus_client.num_entities
|
|
|
+print("total is ", total)
|
|
|
+
|
|
|
+############################遍历mysql中的content_chunks_trsc,获取doc_id和chunk_id,条件为updated_at时间戳小于当天的前一天
|
|
|
+
|
|
|
+# query = """
|
|
|
+# select doc_id, chunk_id from content_chunks_trsc where updated_at<NOW()- INTERVAL 1 DAY;
|
|
|
+# """
|
|
|
+# results = mysql_manager.fetch(query=query)
|
|
|
+# print(f"获取到{len(results)}条需要处理的记录")
|
|
|
+# print(results[0])
|
|
|
+
|
|
|
+#############################遍历milvus中的collection chunk_multi_embeddings_v2_trsc 找出所有doc_id，chunk_id 在results 中的记录
|
|
|
+res_id = milvus_client.query(
|
|
|
+ expr="",
|
|
|
+ output_fields=['id'],
|
|
|
+ limit=10000,
|
|
|
+ consistency_level="Strong"
|
|
|
+ )
|
|
|
+print("res_id length is ", len(res_id)) # 6198
|
|
|
+
|
|
|
+total = len(res_id)
|
|
|
+# assert len(res_id) == total
|
|
|
+
|
|
|
+# exit()
|
|
|
+# batch_size=1000
|
|
|
+processed_count = 0
|
|
|
+
|
|
|
+for i in tqdm(range(0,total)):
|
|
|
+ try:
|
|
|
+ ####update selelct 数据
|
|
|
+ # print("expr is ", f"id == {res_id[i]['id']}")
|
|
|
+ # 从错误信息得知id字段是Int64类型，需要作为整数进行比较，不能用引号包裹
|
|
|
+ results = milvus_client.query(
|
|
|
+ expr=f"id == {res_id[i]['id']}", # 使用双等号，直接使用整数值
|
|
|
+ limit=1,
|
|
|
+ output_fields=['*'],
|
|
|
+ consistency_level="Strong"
|
|
|
+ )
|
|
|
+ if not results:
|
|
|
+ print(f"11111111########id {res_id[i]['id']} 没有查询到结果")
|
|
|
+ continue
|
|
|
+ ####找外健
|
|
|
+ temp_doc_id = results[0]['doc_id']
|
|
|
+ temp_chunk_id = results[0]['chunk_id']
|
|
|
+ temp_vector_text = results[0]['vector_text']
|
|
|
+ temp_vector_summary = results[0]['vector_summary']
|
|
|
+ temp_vector_questions = results[0]['vector_questions']
|
|
|
+
|
|
|
+ ####访问mysql，获取updated_at, 如果 updated_at<NOW()- INTERVAL 1 DAY 返回false，则已经处理过，pass，否则正常处理
|
|
|
+ query = f"""
|
|
|
+ select text,updated_at from content_chunks_trsc where doc_id = '{temp_doc_id}' and chunk_id = '{temp_chunk_id}';
|
|
|
+ """
|
|
|
+ results = mysql_manager.fetch_one(query=query)
|
|
|
+ if not results:
|
|
|
+ print(f"222222222#######doc_id {temp_doc_id} chunk_id {temp_chunk_id} 没有查询到结果")
|
|
|
+ continue
|
|
|
+ temp_updated_at = results['updated_at']
|
|
|
+ if temp_updated_at > datetime.now() - timedelta(days=1):
|
|
|
+ print(f"doc_id {temp_doc_id} chunk_id {temp_chunk_id} 已经处理过，updated_at {temp_updated_at}")
|
|
|
+ continue
|
|
|
+ #############
|
|
|
+ # exit()
|
|
|
+ insert_data = []
|
|
|
+
|
|
|
+ ####找mysql表text
|
|
|
+ query = f"""
|
|
|
+ select text from content_chunks_trsc where doc_id = '{temp_doc_id}' and chunk_id = '{temp_chunk_id}';
|
|
|
+ """
|
|
|
+ results = mysql_manager.fetch_one(query=query)
|
|
|
+
|
|
|
+ if not results:
|
|
|
+ print(f"3333333*******doc_id {temp_doc_id} chunk_id {temp_chunk_id} 没有查询到text结果")
|
|
|
+ continue
|
|
|
+ temp_text = results['text']
|
|
|
+
|
|
|
+ # ####访问deepseek
|
|
|
+ # """
|
|
|
+ # 你是一个专业的文本转写助手，负责将用户输入的文本进行准确的转写。请确保转写结果与原始文本信息保持一致，避免添加任何额外的解释，注释。\
|
|
|
+ # 请务必保持语言准确，精炼。以下是信息
|
|
|
+ # """
|
|
|
+ # ####grok
|
|
|
+ # system_pre_prompt = f'''
|
|
|
+ # 你是一个专业的文本精炼专家。请将以下原文改写成更简练、精确的版本。要求：
|
|
|
+ # 1. 逐句检查原文，保留所有信息、事实、数据、关系和含义，不丢失任何细节。
|
|
|
+ # 2. 去除冗余词语、重复表达和不必要的描述，使语言更紧凑和精准。
|
|
|
+ # 3. 输出纯文本，不要添加任何格式化元素，如标题、列表、编号、粗体、换行符、额外解释或总结语句。只输出改写后的文本。
|
|
|
+ # 4. 如果原文有专业术语，保持原样。
|
|
|
+
|
|
|
+ # 原文：{temp_text}
|
|
|
+ # '''
|
|
|
+
|
|
|
+ # trans_answer_grok = get_deepseek_completion(
|
|
|
+ # model="DeepSeek-V3", # 使用DeepSeek V3模型
|
|
|
+ # prompt=system_pre_prompt, # 使用query_text_1作为prompt
|
|
|
+ # output_type="text" # 返回文本格式
|
|
|
+ # )
|
|
|
+ # print("trans_answer grok is ", trans_answer_grok)
|
|
|
+
|
|
|
+ print("text is", temp_text)
|
|
|
+
|
|
|
+ ###chatgpt
|
|
|
+ system_pre_prompt = f'''
|
|
|
+ 请将以下文本转写为更简练、精确的表达方式。转写时需要满足以下要求：
|
|
|
+ 不丢失信息：转写后的内容必须包含原文中的所有关键信息，任何核心的事实、细节和数据都不能丢失。
|
|
|
+ 简洁清晰：去除冗余的词汇和不必要的修饰，确保语言简洁明了，信息传达精准。
|
|
|
+ 无额外格式：转写后的文本应为纯文本，不添加任何额外的格式、列表、标点或其他符号，保持原有的结构。
|
|
|
+ 请根据这些要求，转写以下段落：{temp_text}
|
|
|
+ '''
|
|
|
+ ###chatgpt
|
|
|
+
|
|
|
+ trans_answer_gpt = get_deepseek_completion(
|
|
|
+ model="DeepSeek-V3", # 使用DeepSeek V3模型
|
|
|
+ prompt=system_pre_prompt, # 使用query_text_1作为prompt
|
|
|
+ output_type="text" # 返回文本格式
|
|
|
+ )
|
|
|
+ # print("text is", temp_text)
|
|
|
+ print("trans_answer_chatgpt is ", trans_answer_gpt)
|
|
|
+
|
|
|
+ # ###compare
|
|
|
+ # system_pre_prompt = f'''
|
|
|
+ # 我有两个经过转写后的文本，请你根据以下标准来比较这两个文本，判断哪个文本更好：
|
|
|
+ # 1. 信息完整性：哪个文本更好地保留了原文的关键信息？如果有任何重要信息缺失，请指出。
|
|
|
+ # 2. 简洁性：哪个文本更加简洁明了，去除了冗余和不必要的部分？
|
|
|
+ # 3. 清晰度：哪个文本更容易理解，表达更清晰？
|
|
|
+ # 4. 精准度：哪个文本更准确地传达了原文的意思，没有歧义或误解？
|
|
|
+ # 请基于这些标准分别对两个文本进行评估，并给出理由说明哪个文本更优。
|
|
|
+ # 文本1: {trans_answer_grok}
|
|
|
+ # 文本2: {trans_answer_gpt}
|
|
|
+ # '''
|
|
|
+ # ###compare
|
|
|
+ # trans_answer_compare = get_deepseek_completion(
|
|
|
+ # model="DeepSeek-V3", # 使用DeepSeek V3模型
|
|
|
+ # prompt=system_pre_prompt, # 使用query_text_1作为prompt
|
|
|
+ # output_type="text" # 返回文本格式
|
|
|
+ # )
|
|
|
+ # print("trans_answer_compare is ", trans_answer_compare)
|
|
|
+
|
|
|
+ # exit()
|
|
|
+ # ####访问embedding model 获取转写后的text
|
|
|
+
|
|
|
+ temp_embedding = get_basic_embedding(text=temp_text, model=DEFAULT_MODEL)
|
|
|
+ temp_embedding_trsc = get_basic_embedding(text=trans_answer_gpt, model=DEFAULT_MODEL)
|
|
|
+
|
|
|
+ similarity_score = cosine_similarity(temp_embedding, temp_embedding_trsc)
|
|
|
+
|
|
|
+ print(f"\n 重写后问题与重写前问题的余弦相似度: {cosine_similarity(temp_embedding, temp_embedding_trsc)}")
|
|
|
+
|
|
|
+ # 更新mysql表，将转写后的文本、相似度分数和当前时间戳更新到对应记录中
|
|
|
+
|
|
|
+ # test_query="select @@autocommit"
|
|
|
+ # res = mysql_manager.fetch_one(test_query)
|
|
|
+ # print("autocommit is ", res['@@autocommit'])
|
|
|
+ # exit()
|
|
|
+ update_query = f"""
|
|
|
+ update content_chunks_trsc set simillarity = {similarity_score}, text = '{trans_answer_gpt}', updated_at = NOW() where doc_id = '{temp_doc_id}' and chunk_id = '{temp_chunk_id}';
|
|
|
+ """
|
|
|
+ mysql_manager.fetch_one(update_query)
|
|
|
+
|
|
|
+ ####更新milvus 的collection，将转写后的embedding替代原来的embedding
|
|
|
+ # 更新milvus的collection chunk_multi_embeddings_v2_trsc，使用转写过的text计算出的embedding temp_embedding_trsc 更新其vector_text
|
|
|
+ # Milvus的upsert方法要求数据格式为字段名到值列表的映射
|
|
|
+ # 注意：必须包含所有非空且无默认值的字段
|
|
|
+ upsert_expr = f"id == {res_id[i]['id']}"
|
|
|
+ # 正确格式：包含所有必需字段
|
|
|
+ upsert_data = {
|
|
|
+ "id": res_id[i]['id'], # 注意：upsert操作需要将值放入列表中
|
|
|
+ "doc_id": temp_doc_id, # 添加必需的doc_id字段
|
|
|
+ "chunk_id": temp_chunk_id, # 添加必需的chunk_id字段
|
|
|
+ "vector_text": temp_embedding_trsc, # vector_text字段也需要放入列表中
|
|
|
+ "vector_summary":temp_vector_summary,
|
|
|
+ "vector_questions":temp_vector_questions
|
|
|
+ }
|
|
|
+ # 执行更新操作
|
|
|
+ milvus_client.upsert(upsert_data, expr=upsert_expr) # 注意参数顺序：先数据后表达式
|
|
|
+
|
|
|
+ processed_count += 1
|
|
|
+ ##################
|
|
|
+ if processed_count % 200 == 0:
|
|
|
+ milvus_client.flush()
|
|
|
+ except Exception as e:
|
|
|
+ print(f"处理第{i}条数据时出现异常: {str(e)}")
|
|
|
+ # 跳过当前循环，继续处理下一条数据
|
|
|
+ continue
|
|
|
+milvus_client.flush()
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
lookathis@163.com