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("迁移完成")
|