101 lines
3.4 KiB
Python
101 lines
3.4 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""迁移脚本：添加 token 消耗字段到对话记录表。"""
|
||
import logging
|
||
from sqlalchemy import text
|
||
from insurance.db.compat import db
|
||
|
||
|
||
def migrate():
|
||
"""添加 model_id、message_tokens、answer_tokens 字段到 insurance_chat_records 表。"""
|
||
try:
|
||
from sqlalchemy import inspect
|
||
|
||
inspector = inspect(db.engine)
|
||
columns = [col['name'] for col in inspector.get_columns('insurance_chat_records')]
|
||
|
||
dialect = db.engine.dialect.name
|
||
|
||
# 添加 model_id 字段
|
||
if 'model_id' not in columns:
|
||
if dialect == 'sqlite':
|
||
db.session.execute(text(
|
||
"ALTER TABLE insurance_chat_records "
|
||
"ADD COLUMN model_id VARCHAR(100)"
|
||
))
|
||
else:
|
||
db.session.execute(text(
|
||
"ALTER TABLE insurance_chat_records "
|
||
"ADD COLUMN model_id VARCHAR(100) NULL "
|
||
"COMMENT '模型 ID'"
|
||
))
|
||
logging.info("成功添加 model_id 字段")
|
||
|
||
# 添加 message_tokens 字段（用户输入 token 数）
|
||
if 'message_tokens' not in columns:
|
||
if dialect == 'sqlite':
|
||
db.session.execute(text(
|
||
"ALTER TABLE insurance_chat_records "
|
||
"ADD COLUMN message_tokens INTEGER DEFAULT 0"
|
||
))
|
||
else:
|
||
db.session.execute(text(
|
||
"ALTER TABLE insurance_chat_records "
|
||
"ADD COLUMN message_tokens INTEGER DEFAULT 0 "
|
||
"COMMENT '用户输入 token 数'"
|
||
))
|
||
logging.info("成功添加 message_tokens 字段")
|
||
|
||
# 添加 answer_tokens 字段（助手回复 token 数）
|
||
if 'answer_tokens' not in columns:
|
||
if dialect == 'sqlite':
|
||
db.session.execute(text(
|
||
"ALTER TABLE insurance_chat_records "
|
||
"ADD COLUMN answer_tokens INTEGER DEFAULT 0"
|
||
))
|
||
else:
|
||
db.session.execute(text(
|
||
"ALTER TABLE insurance_chat_records "
|
||
"ADD COLUMN answer_tokens INTEGER DEFAULT 0 "
|
||
"COMMENT '助手回复 token 数'"
|
||
))
|
||
logging.info("成功添加 answer_tokens 字段")
|
||
|
||
db.session.commit()
|
||
logging.info("迁移完成：token 字段添加成功")
|
||
|
||
except Exception as e:
|
||
logging.error("添加 token 字段失败: %s", e)
|
||
db.session.rollback()
|
||
raise
|
||
|
||
|
||
def downgrade():
|
||
"""删除 token 相关字段。"""
|
||
try:
|
||
from sqlalchemy import inspect
|
||
|
||
inspector = inspect(db.engine)
|
||
columns = [col['name'] for col in inspector.get_columns('insurance_chat_records')]
|
||
|
||
for field in ['model_id', 'message_tokens', 'answer_tokens']:
|
||
if field in columns:
|
||
db.session.execute(text(f"ALTER TABLE insurance_chat_records DROP COLUMN {field}"))
|
||
logging.info(f"成功删除 {field} 字段")
|
||
|
||
db.session.commit()
|
||
logging.info("回滚完成：token 字段删除成功")
|
||
|
||
except Exception as e:
|
||
logging.error("删除 token 字段失败: %s", e)
|
||
db.session.rollback()
|
||
raise
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import sys
|
||
sys.path.insert(0, ".")
|
||
from app import app
|
||
with app.app_context():
|
||
migrate()
|
||
print("迁移完成")
|