168 lines
7.1 KiB
Python
168 lines
7.1 KiB
Python
import os
|
||
import json
|
||
import pymysql
|
||
from pymysql.cursors import DictCursor
|
||
import pymongo
|
||
import requests
|
||
from dotenv import load_dotenv
|
||
from typing import List, Dict, Generator
|
||
|
||
# 加载环境变量
|
||
load_dotenv()
|
||
|
||
class StudentDataProcessorLMStudio:
|
||
def __init__(self):
|
||
# 初始化 MySQL 连接（支持批量遍历）
|
||
self.mysql_conn = pymysql.connect(
|
||
host=os.getenv("MYSQL_HOST"),
|
||
port=int(os.getenv("MYSQL_PORT")),
|
||
user=os.getenv("MYSQL_USER"),
|
||
password=os.getenv("MYSQL_PASSWORD"),
|
||
database=os.getenv("MYSQL_DB"),
|
||
charset="utf8mb4"
|
||
)
|
||
self.mysql_cursor = self.mysql_conn.cursor()
|
||
|
||
# 初始化 MongoDB 连接
|
||
self.mongo_client = pymongo.MongoClient(os.getenv("MONGO_URI"))
|
||
self.mongo_db = self.mongo_client[os.getenv("MONGO_DB")]
|
||
self.mongo_collection = self.mongo_db[os.getenv("MONGO_COLLECTION")]
|
||
|
||
# LM Studio 配置
|
||
self.lm_studio_api_url = os.getenv("LM_STUDIO_API_URL")
|
||
self.lm_studio_model = os.getenv("LM_STUDIO_MODEL_NAME")
|
||
self.lm_studio_api_token = os.getenv("LM_STUDIO_API_TOKEN")
|
||
self.batch_size = int(os.getenv("BATCH_SIZE"))
|
||
|
||
def batch_read_student_data(self) -> Generator[List[Dict], None, None]:
|
||
"""分批遍历 MySQL 所有学生数据（避免一次性加载过多数据）"""
|
||
try:
|
||
# 先获取总条数，用于分批
|
||
count_query = "SELECT COUNT(*) as total FROM student WHERE is_deleted = 0"
|
||
self.mysql_cursor.execute(count_query)
|
||
total = self.mysql_cursor.fetchone()["total"]
|
||
print(f"MySQL 中共计 {total} 条学生数据，将按每批 {self.batch_size} 条处理")
|
||
|
||
# 分批读取
|
||
offset = 0
|
||
while offset < total:
|
||
query = f"""
|
||
SELECT id, name, age, gender, class, phone, email, address
|
||
FROM student
|
||
WHERE is_deleted = 0
|
||
LIMIT {self.batch_size} OFFSET {offset}
|
||
"""
|
||
self.mysql_cursor.execute(query)
|
||
batch_data = self.mysql_cursor.fetchall()
|
||
if not batch_data:
|
||
break
|
||
yield batch_data
|
||
offset += self.batch_size
|
||
print(f"已读取第 {offset//self.batch_size} 批数据（累计 {offset} 条）")
|
||
except Exception as e:
|
||
print(f"分批读取 MySQL 数据失败：{str(e)}")
|
||
raise
|
||
|
||
def process_batch_with_lmstudio(self, batch_data: List[Dict]) -> List[Dict]:
|
||
"""调用 LM Studio 本地模型整理单批数据"""
|
||
if not batch_data:
|
||
return []
|
||
|
||
# 构造适配本地模型的提示词（更简洁，避免本地模型处理复杂指令出错）
|
||
prompt = f"""
|
||
请严格按照以下要求整理学生数据，仅返回JSON数组，不要添加任何额外文字、解释或备注：
|
||
整理规则：
|
||
1. 标准化字段：
|
||
- 姓名：去除首尾空格，仅保留中文（如" 李四 "→"李四"）
|
||
- 年龄：转为整数，非数字则设为null
|
||
- 手机号：仅保留11位纯数字，不符合则设为null
|
||
- 邮箱：验证格式，无效则设为null
|
||
- 性别：非"男"/"女"则设为"未知"
|
||
- 班级：为空则设为"未分配"
|
||
2. 新增字段：
|
||
- student_id：格式为"2024"+补零后的id（如id=5→"202400005"，id=123→"202400123"）
|
||
- data_quality："完整"（所有字段非null）/"部分缺失"（1-2个字段null）/"严重缺失"（≥3个字段null）
|
||
原始数据：{json.dumps(batch_data, ensure_ascii=False, indent=2)}
|
||
"""
|
||
|
||
# 构造 LM Studio API 请求体（兼容 OpenAI 格式）
|
||
payload = {
|
||
"model": self.lm_studio_model,
|
||
"messages": [
|
||
{"role": "system", "content": "你是数据整理专家，仅返回JSON格式的处理结果，无其他内容"},
|
||
{"role": "user", "content": prompt}
|
||
],
|
||
"temperature": 0.1, # 低随机性保证结果稳定
|
||
"max_tokens": 2000 # 足够容纳批量数据的JSON
|
||
}
|
||
|
||
try:
|
||
# 调用 LM Studio 本地 API
|
||
response = requests.post(
|
||
self.lm_studio_api_url,
|
||
headers={"Authorization": f"Bearer {self.lm_studio_api_token}"},
|
||
json=payload,
|
||
timeout=60 # 本地模型处理稍慢，延长超时时间
|
||
)
|
||
response.raise_for_status() # 抛出HTTP错误
|
||
result = response.json()
|
||
|
||
# 解析本地模型返回的结果
|
||
processed_text = result["choices"][0]["message"]["content"].strip()
|
||
# 清理可能的多余字符（本地模型可能返回```json开头/结尾）
|
||
if processed_text.startswith("```json"):
|
||
processed_text = processed_text.replace("```json", "").replace("```", "").strip()
|
||
processed_data = json.loads(processed_text)
|
||
return processed_data
|
||
except json.JSONDecodeError:
|
||
print(f"本地模型返回非标准JSON，跳过该批：{processed_text[:200]}")
|
||
return []
|
||
except Exception as e:
|
||
print(f"调用LM Studio失败：{str(e)}")
|
||
return []
|
||
|
||
def save_to_mongodb(self, processed_data: List[Dict]):
|
||
"""将整理后的数据存入MongoDB（支持增量插入）"""
|
||
if not processed_data:
|
||
return
|
||
try:
|
||
# 批量插入，跳过重复id（避免重复存储）
|
||
for data in processed_data:
|
||
self.mongo_collection.update_one(
|
||
{"id": data["id"]}, # 以原始id为唯一键
|
||
{"$set": data}, # 存在则更新，不存在则插入
|
||
upsert=True
|
||
)
|
||
print(f"成功存入/更新 {len(processed_data)} 条数据到MongoDB")
|
||
except Exception as e:
|
||
print(f"存入MongoDB失败：{str(e)}")
|
||
|
||
def close_connections(self):
|
||
"""关闭数据库连接"""
|
||
if self.mysql_cursor:
|
||
self.mysql_cursor.close()
|
||
if self.mysql_conn:
|
||
self.mysql_conn.close()
|
||
if self.mongo_client:
|
||
self.mongo_client.close()
|
||
print("所有连接已关闭")
|
||
|
||
def run(self):
|
||
"""执行完整流程：分批读取→本地模型处理→存入MongoDB"""
|
||
try:
|
||
# 遍历所有批次数据
|
||
for batch in self.batch_read_student_data():
|
||
# 处理单批数据
|
||
processed_batch = self.process_batch_with_lmstudio(batch)
|
||
# 存入MongoDB
|
||
self.save_to_mongodb(processed_batch)
|
||
print("✅ 全量数据处理完成！")
|
||
except Exception as e:
|
||
print(f"程序执行异常：{str(e)}")
|
||
finally:
|
||
self.close_connections()
|
||
|
||
# 程序入口
|
||
if __name__ == "__main__":
|
||
processor = StudentDataProcessorLMStudio()
|
||
processor.run() |