102 lines
3.7 KiB
Python
102 lines
3.7 KiB
Python
# 调用智谱的api，修改每一个项目的分类
|
||
# 从db文件读取表，读取第二个项，标题，根据标题，提交到api，获取回复，返回，并更新到db文件
|
||
|
||
import sqlite3
|
||
import time
|
||
from loguru import logger
|
||
from zhipuai import ZhipuAI
|
||
|
||
# 配置日志
|
||
logger.add("db_modify_zhipu.log", rotation="10 MB", level="INFO")
|
||
|
||
# 初始化客户端
|
||
client = ZhipuAI(api_key="fad3d9f9a45f4d939f0e7a7133fa07bf.X4bOO053GAIPKLE5")
|
||
|
||
def get_simplified_category(title):
|
||
"""
|
||
调用智谱API获取简化的分类名称
|
||
"""
|
||
try:
|
||
# 创建聊天完成请求
|
||
response = client.chat.completions.create(
|
||
model="glm-4-flash",
|
||
messages=[
|
||
{
|
||
"role": "system",
|
||
"content": "你是一个专业的分类助手。请根据文章标题，提供一个3-6个汉字的简化分类名称，去除空格和特殊符号，更容易理解，并保持原意。"
|
||
},
|
||
{
|
||
"role": "user",
|
||
"content": f"对以下文字内容进行分类，返回结果为类别，如\"社会新闻\"，\"机器人\"，\"金融\"，\"历史\"，\"购物\"，\"新质生产力\"等等。目的：只返回2-6个汉字，不返回其它内容。内容：'{title}'"
|
||
}
|
||
],
|
||
temperature=0.7
|
||
)
|
||
|
||
# 提取回复内容
|
||
category = response.choices[0].message.content.strip()
|
||
logger.info(f"标题: {title[:30]}... -> 分类: {category}")
|
||
return category
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取分类失败: {str(e)}")
|
||
return None
|
||
|
||
def update_database_categories():
|
||
"""
|
||
更新数据库中的分类信息
|
||
"""
|
||
# 连接到数据库
|
||
conn = sqlite3.connect('tophub_data.db')
|
||
cursor = conn.cursor()
|
||
|
||
try:
|
||
# 获取所有记录
|
||
cursor.execute("SELECT id, title, category FROM articles")
|
||
records = cursor.fetchall()
|
||
|
||
logger.info(f"共找到 {len(records)} 条记录需要处理")
|
||
|
||
updated_count = 0
|
||
failed_count = 0
|
||
|
||
# 处理每条记录
|
||
for record in records:
|
||
record_id, title, current_category = record
|
||
|
||
# 跳过已经简化的分类（长度<=6且不包含特殊字符）
|
||
if current_category and len(current_category) <= 6 and not any(c in current_category for c in " ,.!?;:，。！？；："):
|
||
logger.info(f"跳过记录 {record_id}，分类已简化: {current_category}")
|
||
continue
|
||
|
||
logger.info(f"处理记录 {record_id}: {title[:30]}...")
|
||
|
||
# 获取新的分类
|
||
new_category = get_simplified_category(title)
|
||
|
||
if new_category:
|
||
# 更新数据库
|
||
cursor.execute("UPDATE articles SET category = ? WHERE id = ?", (new_category, record_id))
|
||
conn.commit()
|
||
updated_count += 1
|
||
logger.info(f"已更新记录 {record_id} 的分类为: {new_category}")
|
||
else:
|
||
failed_count += 1
|
||
logger.error(f"无法获取记录 {record_id} 的新分类")
|
||
|
||
# 添加延迟，避免API调用过于频繁
|
||
time.sleep(1)
|
||
|
||
logger.info(f"处理完成! 成功更新 {updated_count} 条记录，失败 {failed_count} 条记录")
|
||
|
||
except Exception as e:
|
||
logger.error(f"更新数据库时出错: {str(e)}")
|
||
conn.rollback()
|
||
finally:
|
||
conn.close()
|
||
|
||
if __name__ == "__main__":
|
||
logger.info("开始更新数据库分类...")
|
||
update_database_categories()
|
||
logger.info("程序执行完成")
|