lint: 代码格式化
This commit is contained in:
Wenjie Zhang
parent
66c550c18c
commit
aebe63ce1c
@ -27,7 +27,7 @@ class DBManager:
|
||||
|
||||
# 首先创建基本表结构
|
||||
self.create_tables()
|
||||
|
||||
|
||||
# 然后检查并执行数据库迁移
|
||||
self.run_migrations()
|
||||
|
||||
@ -41,7 +41,7 @@ class DBManager:
|
||||
# 确保所有表都会被创建
|
||||
Base.metadata.create_all(self.engine)
|
||||
logger.info("Database tables created/checked")
|
||||
|
||||
|
||||
def run_migrations(self):
|
||||
"""运行数据库迁移"""
|
||||
try:
|
||||
|
||||
@ -16,7 +16,7 @@ class User(Base):
|
||||
role = Column(String, nullable=False, default="user") # 角色: superadmin, admin, user
|
||||
created_at = Column(DateTime, default=func.now())
|
||||
last_login = Column(DateTime, nullable=True)
|
||||
|
||||
|
||||
# 登录失败限制相关字段
|
||||
login_failed_count = Column(Integer, nullable=False, default=0) # 登录失败次数
|
||||
last_failed_login = Column(DateTime, nullable=True) # 最后一次登录失败时间
|
||||
@ -39,44 +39,46 @@ class User(Base):
|
||||
if include_password:
|
||||
result["password_hash"] = self.password_hash
|
||||
return result
|
||||
|
||||
|
||||
def is_login_locked(self):
|
||||
"""检查用户是否处于登录锁定状态"""
|
||||
if self.login_locked_until is None:
|
||||
return False
|
||||
from datetime import datetime
|
||||
|
||||
return datetime.now() < self.login_locked_until
|
||||
|
||||
|
||||
def get_remaining_lock_time(self):
|
||||
"""获取剩余锁定时间（秒）"""
|
||||
if not self.is_login_locked():
|
||||
return 0
|
||||
from datetime import datetime
|
||||
|
||||
return int((self.login_locked_until - datetime.now()).total_seconds())
|
||||
|
||||
|
||||
def calculate_lock_duration(self):
|
||||
"""根据失败次数计算锁定时长（秒）"""
|
||||
if self.login_failed_count < 10:
|
||||
return 0
|
||||
|
||||
|
||||
# 从第10次失败开始，等待时间从1秒开始，每次翻倍
|
||||
wait_seconds = 2 ** (self.login_failed_count - 10)
|
||||
|
||||
|
||||
# 最大锁定时间：365天
|
||||
max_seconds = 365 * 24 * 60 * 60
|
||||
return min(wait_seconds, max_seconds)
|
||||
|
||||
|
||||
def increment_failed_login(self):
|
||||
"""增加登录失败次数并设置锁定时间"""
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
|
||||
self.login_failed_count += 1
|
||||
self.last_failed_login = datetime.now()
|
||||
|
||||
|
||||
lock_duration = self.calculate_lock_duration()
|
||||
if lock_duration > 0:
|
||||
self.login_locked_until = datetime.now() + timedelta(seconds=lock_duration)
|
||||
|
||||
|
||||
def reset_failed_login(self):
|
||||
"""重置登录失败相关字段"""
|
||||
self.login_failed_count = 0
|
||||
|
||||
@ -72,17 +72,14 @@ async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends(
|
||||
detail="用户名或密码错误",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
|
||||
# 检查用户是否处于登录锁定状态
|
||||
if user.is_login_locked():
|
||||
remaining_time = user.get_remaining_lock_time()
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_423_LOCKED,
|
||||
detail=f"登录被锁定，请等待 {remaining_time} 秒后再试",
|
||||
headers={
|
||||
"WWW-Authenticate": "Bearer",
|
||||
"X-Lock-Remaining": str(remaining_time)
|
||||
},
|
||||
headers={"WWW-Authenticate": "Bearer", "X-Lock-Remaining": str(remaining_time)},
|
||||
)
|
||||
|
||||
# 验证密码
|
||||
@ -90,20 +87,17 @@ async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends(
|
||||
# 密码错误，增加失败次数
|
||||
user.increment_failed_login()
|
||||
db.commit()
|
||||
|
||||
|
||||
# 记录失败操作
|
||||
log_operation(db, user.id if user else None, "登录失败", f"密码错误，失败次数: {user.login_failed_count}")
|
||||
|
||||
|
||||
# 检查是否需要锁定
|
||||
if user.is_login_locked():
|
||||
remaining_time = user.get_remaining_lock_time()
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_423_LOCKED,
|
||||
detail=f"由于多次登录失败，账户已被锁定 {remaining_time} 秒",
|
||||
headers={
|
||||
"WWW-Authenticate": "Bearer",
|
||||
"X-Lock-Remaining": str(remaining_time)
|
||||
},
|
||||
headers={"WWW-Authenticate": "Bearer", "X-Lock-Remaining": str(remaining_time)},
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
|
||||
@ -15,7 +15,8 @@ from server.models.thread_model import Thread
|
||||
from server.models.user_model import User
|
||||
from server.routers.auth_router import get_admin_user
|
||||
from server.utils.auth_middleware import get_db, get_required_user
|
||||
from src import config, executor
|
||||
from src import executor
|
||||
from src import config as conf
|
||||
from src.agents import agent_manager
|
||||
from src.agents.common.tools import gen_tool_info, get_buildin_tools
|
||||
from src.models import select_model
|
||||
@ -32,7 +33,7 @@ chat = APIRouter(prefix="/chat", tags=["chat"])
|
||||
async def get_default_agent(current_user: User = Depends(get_required_user)):
|
||||
"""获取默认智能体ID（需要登录）"""
|
||||
try:
|
||||
default_agent_id = config.default_agent_id
|
||||
default_agent_id = conf.default_agent_id
|
||||
# 如果没有设置默认智能体，尝试获取第一个可用的智能体
|
||||
if not default_agent_id:
|
||||
agents = await agent_manager.get_agents_info()
|
||||
@ -61,9 +62,9 @@ async def set_default_agent(request_data: dict = Body(...), current_user=Depends
|
||||
raise HTTPException(status_code=404, detail=f"智能体 {agent_id} 不存在")
|
||||
|
||||
# 设置默认智能体ID
|
||||
config.default_agent_id = agent_id
|
||||
conf.default_agent_id = agent_id
|
||||
# 保存配置
|
||||
config.save()
|
||||
conf.save()
|
||||
|
||||
return {"success": True, "default_agent_id": agent_id}
|
||||
except HTTPException as he:
|
||||
@ -187,9 +188,9 @@ async def get_chat_models(model_provider: str, current_user: User = Depends(get_
|
||||
@chat.post("/models/update")
|
||||
async def update_chat_models(model_provider: str, model_names: list[str], current_user=Depends(get_admin_user)):
|
||||
"""更新指定模型提供商的模型列表 (仅管理员)"""
|
||||
config.model_names[model_provider]["models"] = model_names
|
||||
config._save_models_to_file()
|
||||
return {"models": config.model_names[model_provider]["models"]}
|
||||
conf.model_names[model_provider]["models"] = model_names
|
||||
conf._save_models_to_file()
|
||||
return {"models": conf.model_names[model_provider]["models"]}
|
||||
|
||||
|
||||
@chat.get("/tools")
|
||||
|
||||
@ -1,39 +1,39 @@
|
||||
"""
|
||||
数据库迁移系统
|
||||
"""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import sqlite3
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
from src.utils import logger
|
||||
|
||||
|
||||
class DatabaseMigrator:
|
||||
"""数据库迁移器"""
|
||||
|
||||
|
||||
def __init__(self, db_path: str):
|
||||
self.db_path = db_path
|
||||
self.backup_dir = os.path.join(os.path.dirname(db_path), "backups")
|
||||
self.migration_version_key = "migration_version"
|
||||
|
||||
|
||||
def ensure_backup_dir(self):
|
||||
"""确保备份目录存在"""
|
||||
Path(self.backup_dir).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
def backup_database(self) -> str:
|
||||
"""备份数据库文件"""
|
||||
if not os.path.exists(self.db_path):
|
||||
logger.info("数据库文件不存在，无需备份")
|
||||
return ""
|
||||
|
||||
|
||||
self.ensure_backup_dir()
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
backup_filename = f"server_backup_{timestamp}.db"
|
||||
backup_path = os.path.join(self.backup_dir, backup_filename)
|
||||
|
||||
|
||||
try:
|
||||
shutil.copy2(self.db_path, backup_path)
|
||||
logger.info(f"数据库已备份到: {backup_path}")
|
||||
@ -41,26 +41,25 @@ class DatabaseMigrator:
|
||||
except Exception as e:
|
||||
logger.error(f"数据库备份失败: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def get_current_version(self) -> int:
|
||||
"""获取当前数据库版本"""
|
||||
if not os.path.exists(self.db_path):
|
||||
return 0
|
||||
|
||||
|
||||
try:
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
|
||||
# 检查版本表是否存在
|
||||
cursor.execute("""
|
||||
SELECT name FROM sqlite_master
|
||||
WHERE type='table' AND name='migration_versions'
|
||||
""")
|
||||
|
||||
SELECT name FROM sqlite_master
|
||||
WHERE type='table' AND name='migration_versions' """)
|
||||
|
||||
if not cursor.fetchone():
|
||||
# 版本表不存在，检查是否为旧版本数据库
|
||||
cursor.execute("""
|
||||
SELECT name FROM sqlite_master
|
||||
SELECT name FROM sqlite_master
|
||||
WHERE type='table' AND name='users'
|
||||
""")
|
||||
if cursor.fetchone():
|
||||
@ -69,25 +68,25 @@ class DatabaseMigrator:
|
||||
else:
|
||||
# 全新数据库
|
||||
return 0
|
||||
|
||||
|
||||
# 获取当前版本
|
||||
cursor.execute("SELECT version FROM migration_versions ORDER BY version DESC LIMIT 1")
|
||||
result = cursor.fetchone()
|
||||
return result[0] if result else 0
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"获取数据库版本失败: {e}")
|
||||
return 0
|
||||
finally:
|
||||
if 'conn' in locals():
|
||||
if "conn" in locals():
|
||||
conn.close()
|
||||
|
||||
|
||||
def set_version(self, version: int):
|
||||
"""设置数据库版本"""
|
||||
try:
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
|
||||
# 创建版本表
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS migration_versions (
|
||||
@ -97,93 +96,98 @@ class DatabaseMigrator:
|
||||
description TEXT
|
||||
)
|
||||
""")
|
||||
|
||||
|
||||
# 插入版本记录
|
||||
cursor.execute("""
|
||||
INSERT INTO migration_versions (version, description)
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO migration_versions (version, description)
|
||||
VALUES (?, ?)
|
||||
""", (version, f"Migration to version {version}"))
|
||||
|
||||
""",
|
||||
(version, f"Migration to version {version}"),
|
||||
)
|
||||
|
||||
conn.commit()
|
||||
logger.info(f"数据库版本设置为: {version}")
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"设置数据库版本失败: {e}")
|
||||
raise
|
||||
finally:
|
||||
if 'conn' in locals():
|
||||
if "conn" in locals():
|
||||
conn.close()
|
||||
|
||||
def execute_migration(self, version: int, description: str, sql_commands: List[str]):
|
||||
|
||||
def execute_migration(self, version: int, description: str, sql_commands: list[str]):
|
||||
"""执行迁移"""
|
||||
logger.info(f"执行迁移 v{version}: {description}")
|
||||
|
||||
|
||||
try:
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
|
||||
# 执行迁移SQL命令
|
||||
for sql in sql_commands:
|
||||
if sql.strip(): # 跳过空命令
|
||||
logger.info(f"执行SQL: {sql}")
|
||||
cursor.execute(sql)
|
||||
|
||||
|
||||
conn.commit()
|
||||
logger.info(f"迁移 v{version} 执行成功")
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"迁移 v{version} 执行失败: {e}")
|
||||
raise
|
||||
finally:
|
||||
if 'conn' in locals():
|
||||
if "conn" in locals():
|
||||
conn.close()
|
||||
|
||||
|
||||
def check_column_exists(self, table_name: str, column_name: str) -> bool:
|
||||
"""检查列是否存在"""
|
||||
try:
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
|
||||
cursor.execute(f"PRAGMA table_info({table_name})")
|
||||
columns = [column[1] for column in cursor.fetchall()]
|
||||
return column_name in columns
|
||||
|
||||
|
||||
except Exception:
|
||||
return False
|
||||
finally:
|
||||
if 'conn' in locals():
|
||||
if "conn" in locals():
|
||||
conn.close()
|
||||
|
||||
|
||||
def run_migrations(self):
|
||||
"""运行所有待执行的迁移"""
|
||||
current_version = self.get_current_version()
|
||||
latest_version = self.get_latest_migration_version()
|
||||
|
||||
|
||||
# 如果数据库已存在但没有版本表，创建版本表并设置为最新版本
|
||||
if current_version == 0 and latest_version > 0 and os.path.exists(self.db_path):
|
||||
# 检查users表是否已有新字段，如果有，说明是通过SQLAlchemy创建的
|
||||
if (self.check_column_exists('users', 'login_failed_count') and
|
||||
self.check_column_exists('users', 'last_failed_login') and
|
||||
self.check_column_exists('users', 'login_locked_until')):
|
||||
if (
|
||||
self.check_column_exists("users", "login_failed_count")
|
||||
and self.check_column_exists("users", "last_failed_login")
|
||||
and self.check_column_exists("users", "login_locked_until")
|
||||
):
|
||||
# 字段已存在，直接设置为最新版本
|
||||
logger.info(f"检测到现有数据库已包含最新字段，设置版本为 v{latest_version}")
|
||||
self.set_version(latest_version)
|
||||
return
|
||||
|
||||
|
||||
if current_version >= latest_version:
|
||||
logger.info(f"数据库已是最新版本 v{current_version}")
|
||||
return
|
||||
|
||||
|
||||
logger.info(f"开始数据库迁移: v{current_version} -> v{latest_version}")
|
||||
|
||||
|
||||
# 备份数据库
|
||||
backup_path = self.backup_database()
|
||||
|
||||
|
||||
try:
|
||||
# 执行迁移
|
||||
migrations = self.get_migrations()
|
||||
has_executed_migrations = False
|
||||
|
||||
|
||||
for version, description, sql_commands in migrations:
|
||||
if version > current_version:
|
||||
if sql_commands: # 只有当有SQL命令时才执行迁移
|
||||
@ -191,15 +195,15 @@ class DatabaseMigrator:
|
||||
has_executed_migrations = True
|
||||
else:
|
||||
logger.info(f"迁移 v{version}: {description} - 无需执行，字段已存在")
|
||||
|
||||
|
||||
# 无论是否有SQL命令，都设置版本
|
||||
self.set_version(version)
|
||||
|
||||
|
||||
if has_executed_migrations:
|
||||
logger.info("数据库迁移完成")
|
||||
else:
|
||||
logger.info("数据库结构已是最新，仅更新版本记录")
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"数据库迁移失败: {e}")
|
||||
if backup_path and os.path.exists(backup_path):
|
||||
@ -210,39 +214,39 @@ class DatabaseMigrator:
|
||||
except Exception as restore_error:
|
||||
logger.error(f"数据库恢复失败: {restore_error}")
|
||||
raise
|
||||
|
||||
|
||||
def get_latest_migration_version(self) -> int:
|
||||
"""获取最新迁移版本号"""
|
||||
# 这里返回硬编码的最新版本号，不依赖迁移定义
|
||||
# 因为迁移定义可能为空（字段已存在）
|
||||
return 1 # 当前最新版本是 v1
|
||||
|
||||
def get_migrations(self) -> List[Tuple[int, str, List[str]]]:
|
||||
|
||||
def get_migrations(self) -> list[tuple[int, str, list[str]]]:
|
||||
"""获取所有迁移定义
|
||||
返回格式: [(version, description, [sql_commands])]
|
||||
"""
|
||||
migrations = []
|
||||
|
||||
|
||||
# 迁移 v1: 为 users 表添加登录失败限制字段
|
||||
# 使用条件检查来避免重复添加字段
|
||||
v1_commands = []
|
||||
|
||||
|
||||
# 检查并添加 login_failed_count 字段
|
||||
if not self.check_column_exists('users', 'login_failed_count'):
|
||||
if not self.check_column_exists("users", "login_failed_count"):
|
||||
v1_commands.append("ALTER TABLE users ADD COLUMN login_failed_count INTEGER NOT NULL DEFAULT 0")
|
||||
|
||||
|
||||
# 检查并添加 last_failed_login 字段
|
||||
if not self.check_column_exists('users', 'last_failed_login'):
|
||||
if not self.check_column_exists("users", "last_failed_login"):
|
||||
v1_commands.append("ALTER TABLE users ADD COLUMN last_failed_login DATETIME")
|
||||
|
||||
|
||||
# 检查并添加 login_locked_until 字段
|
||||
if not self.check_column_exists('users', 'login_locked_until'):
|
||||
if not self.check_column_exists("users", "login_locked_until"):
|
||||
v1_commands.append("ALTER TABLE users ADD COLUMN login_locked_until DATETIME")
|
||||
|
||||
|
||||
# 如果有命令需要执行，才添加迁移
|
||||
if v1_commands:
|
||||
migrations.append((1, "为用户表添加登录失败限制字段", v1_commands))
|
||||
|
||||
|
||||
# 未来的迁移可以在这里添加
|
||||
# migrations.append((
|
||||
# 2,
|
||||
@ -252,17 +256,17 @@ class DatabaseMigrator:
|
||||
# "ALTER TABLE existing_table ADD COLUMN new_field ..."
|
||||
# ]
|
||||
# ))
|
||||
|
||||
|
||||
return migrations
|
||||
|
||||
|
||||
def check_and_migrate(db_path: str):
|
||||
"""检查并执行数据库迁移"""
|
||||
migrator = DatabaseMigrator(db_path)
|
||||
|
||||
|
||||
try:
|
||||
migrator.run_migrations()
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"数据库迁移过程中发生错误: {e}")
|
||||
return False
|
||||
return False
|
||||
|
||||
@ -1,5 +1,3 @@
|
||||
import os
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv("src/.env", override=True)
|
||||
|
||||
@ -1,3 +1,3 @@
|
||||
from .app import config
|
||||
|
||||
__all__ = ["config"]
|
||||
__all__ = ["config"]
|
||||
|
||||
@ -228,4 +228,4 @@ class Config(SimpleConfig):
|
||||
return value
|
||||
|
||||
|
||||
config = Config()
|
||||
config = Config()
|
||||
|
||||
@ -21,4 +21,3 @@ knowledge_base = KnowledgeBaseManager(work_dir)
|
||||
graph_base = GraphDatabase()
|
||||
|
||||
__all__ = ["GraphDatabase", "knowledge_base", "graph_base"]
|
||||
|
||||
|
||||
@ -1,9 +0,0 @@
|
||||
ChatbotAgent:
|
||||
name: 食品安全知识助手
|
||||
icon: /agents-icon/chatbot.png
|
||||
examples:
|
||||
- 冰箱里不同食材应该怎么存放？
|
||||
- 有机食品和普通食品有什么区别？
|
||||
- 外卖食品安全需要注意什么？
|
||||
- 如何判断食品是否已经变质？
|
||||
- 食品添加剂对健康有害吗？
|
||||
Loading…
Reference in New Issue
Block a user