From f88aca74f294dc923c68b08de7f555a2841edeb5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9C=B1=E6=BD=AE?= Date: Thu, 29 Jan 2026 15:59:42 +0800 Subject: [PATCH] add table prefix --- routes/bot_manager.py | 267 ++++++++++++++++++++++++++++++++++-------- 1 file changed, 218 insertions(+), 49 deletions(-) diff --git a/routes/bot_manager.py b/routes/bot_manager.py index 221e5ce..30bfcae 100644 --- a/routes/bot_manager.py +++ b/routes/bot_manager.py @@ -4,7 +4,9 @@ Bot Manager API 路由 """ import logging import uuid -from datetime import datetime +import hashlib +import secrets +from datetime import datetime, timedelta from typing import Optional, List from fastapi import APIRouter, HTTPException, Header from pydantic import BaseModel @@ -16,9 +18,47 @@ logger = logging.getLogger('app') router = APIRouter() +# ============== Admin 配置 ============== +ADMIN_USERNAME = "admin" +ADMIN_PASSWORD = "Admin123" # 生产环境应使用环境变量 +TOKEN_EXPIRE_HOURS = 24 + # ============== 认证函数 ============== +async def verify_admin_auth(authorization: Optional[str]) -> tuple[bool, Optional[str]]: + """ + 验证管理员认证 + + Args: + authorization: Authorization header 值 + + Returns: + tuple[bool, Optional[str]]: (是否有效, 用户名) + """ + provided_token = extract_api_key_from_auth(authorization) + if not provided_token: + return False, None + + pool = get_db_pool_manager().pool + + async with pool.connection() as conn: + async with conn.cursor() as cursor: + # 检查 token 是否有效且未过期 + await cursor.execute(""" + SELECT username, expires_at + FROM agent_admin_tokens + WHERE token = %s + AND expires_at > NOW() + """, (provided_token,)) + row = await cursor.fetchone() + + if not row: + return False, None + + return True, row[0] + + def verify_auth(authorization: Optional[str]) -> None: """ 验证请求认证 @@ -39,6 +79,26 @@ def verify_auth(authorization: Optional[str]) -> None: # ============== Pydantic Models ============== +# --- Admin 登录相关 --- +class AdminLoginRequest(BaseModel): + """管理员登录请求""" + username: str + password: str + + +class AdminLoginResponse(BaseModel): + """管理员登录响应""" + token: str + username: str + expires_at: str + + +class AdminVerifyResponse(BaseModel): + """管理员验证响应""" + valid: bool + username: Optional[str] = None + + # --- 模型相关 --- class ModelCreate(BaseModel): """创建模型请求""" @@ -77,7 +137,6 @@ class ModelResponse(BaseModel): class BotCreate(BaseModel): """创建 Bot 请求""" name: str - bot_id: str class BotUpdate(BaseModel): @@ -196,9 +255,23 @@ async def init_bot_manager_tables(): # SQL 表创建语句 tables_sql = [ + # admin_tokens 表（用于存储登录 token） + """ + CREATE TABLE IF NOT EXISTS agent_admin_tokens ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + username VARCHAR(255) NOT NULL, + token VARCHAR(255) NOT NULL UNIQUE, + expires_at TIMESTAMP WITH TIME ZONE NOT NULL, + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() + ) + """, + # admin_tokens 索引 + "CREATE INDEX IF NOT EXISTS idx_agent_admin_tokens_token ON agent_admin_tokens(token)", + "CREATE INDEX IF NOT EXISTS idx_agent_admin_tokens_expires ON agent_admin_tokens(expires_at)", + # models 表 """ - CREATE TABLE IF NOT EXISTS models ( + CREATE TABLE IF NOT EXISTS agent_models ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), name VARCHAR(255) NOT NULL, provider VARCHAR(100) NOT NULL, @@ -211,11 +284,11 @@ async def init_bot_manager_tables(): ) """, # models 索引 - "CREATE INDEX IF NOT EXISTS idx_models_is_default ON models(is_default)", + "CREATE INDEX IF NOT EXISTS idx_agent_models_is_default ON agent_models(is_default)", # bots 表 """ - CREATE TABLE IF NOT EXISTS bots ( + CREATE TABLE IF NOT EXISTS agent_bots ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), name VARCHAR(255) NOT NULL, bot_id VARCHAR(255) NOT NULL UNIQUE, @@ -224,13 +297,13 @@ async def init_bot_manager_tables(): ) """, # bots 索引 - "CREATE INDEX IF NOT EXISTS idx_bots_bot_id ON bots(bot_id)", + "CREATE INDEX IF NOT EXISTS idx_agent_bots_bot_id ON agent_bots(bot_id)", # bot_settings 表 """ - CREATE TABLE IF NOT EXISTS bot_settings ( - bot_id UUID PRIMARY KEY REFERENCES bots(id) ON DELETE CASCADE, - model_id UUID REFERENCES models(id) ON DELETE SET NULL, + CREATE TABLE IF NOT EXISTS agent_bot_settings ( + bot_id UUID PRIMARY KEY REFERENCES agent_bots(id) ON DELETE CASCADE, + model_id UUID REFERENCES agent_models(id) ON DELETE SET NULL, language VARCHAR(10) DEFAULT 'zh', robot_type VARCHAR(50), dataset_ids TEXT, @@ -245,9 +318,9 @@ async def init_bot_manager_tables(): # mcp_servers 表 """ - CREATE TABLE IF NOT EXISTS mcp_servers ( + CREATE TABLE IF NOT EXISTS agent_mcp_servers ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), - bot_id UUID REFERENCES bots(id) ON DELETE CASCADE, + bot_id UUID REFERENCES agent_bots(id) ON DELETE CASCADE, name VARCHAR(255) NOT NULL, type VARCHAR(50) NOT NULL, config JSONB NOT NULL, @@ -257,22 +330,22 @@ async def init_bot_manager_tables(): ) """, # mcp_servers 索引 - "CREATE INDEX IF NOT EXISTS idx_mcp_servers_bot_id ON mcp_servers(bot_id)", - "CREATE INDEX IF NOT EXISTS idx_mcp_servers_enabled ON mcp_servers(enabled)", + "CREATE INDEX IF NOT EXISTS idx_agent_mcp_servers_bot_id ON agent_mcp_servers(bot_id)", + "CREATE INDEX IF NOT EXISTS idx_agent_mcp_servers_enabled ON agent_mcp_servers(enabled)", # chat_sessions 表 """ - CREATE TABLE IF NOT EXISTS chat_sessions ( + CREATE TABLE IF NOT EXISTS agent_chat_sessions ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), - bot_id UUID REFERENCES bots(id) ON DELETE CASCADE, + bot_id UUID REFERENCES agent_bots(id) ON DELETE CASCADE, title VARCHAR(500), created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(), updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() ) """, # chat_sessions 索引 - "CREATE INDEX IF NOT EXISTS idx_chat_sessions_bot_id ON chat_sessions(bot_id)", - "CREATE INDEX IF NOT EXISTS idx_chat_sessions_created ON chat_sessions(created_at DESC)", + "CREATE INDEX IF NOT EXISTS idx_agent_chat_sessions_bot_id ON agent_chat_sessions(bot_id)", + "CREATE INDEX IF NOT EXISTS idx_agent_chat_sessions_created ON agent_chat_sessions(created_at DESC)", ] async with pool.connection() as conn: @@ -321,7 +394,7 @@ async def get_models(authorization: Optional[str] = Header(None)): async with conn.cursor() as cursor: await cursor.execute(""" SELECT id, name, provider, model, server, api_key, is_default, created_at, updated_at - FROM models + FROM agent_models ORDER BY is_default DESC, created_at DESC """) rows = await cursor.fetchall() @@ -362,13 +435,13 @@ async def create_model(request: ModelCreate, authorization: Optional[str] = Head if request.is_default: async with pool.connection() as conn: async with conn.cursor() as cursor: - await cursor.execute("UPDATE models SET is_default = FALSE WHERE is_default = TRUE") + await cursor.execute("UPDATE agent_models SET is_default = FALSE WHERE is_default = TRUE") await conn.commit() async with pool.connection() as conn: async with conn.cursor() as cursor: await cursor.execute(""" - INSERT INTO models (name, provider, model, server, api_key, is_default) + INSERT INTO agent_models (name, provider, model, server, api_key, is_default) VALUES (%s, %s, %s, %s, %s, %s) RETURNING id, created_at, updated_at """, ( @@ -449,13 +522,13 @@ async def update_model( if request.is_default is True: async with pool.connection() as conn: async with conn.cursor() as cursor: - await cursor.execute("UPDATE models SET is_default = FALSE WHERE is_default = TRUE") + await cursor.execute("UPDATE agent_models SET is_default = FALSE WHERE is_default = TRUE") await conn.commit() async with pool.connection() as conn: async with conn.cursor() as cursor: await cursor.execute(f""" - UPDATE models + UPDATE agent_models SET {', '.join(update_fields)} WHERE id = %s RETURNING id, name, provider, model, server, api_key, is_default, created_at, updated_at @@ -498,7 +571,7 @@ async def delete_model(model_id: str, authorization: Optional[str] = Header(None async with pool.connection() as conn: async with conn.cursor() as cursor: - await cursor.execute("DELETE FROM models WHERE id = %s RETURNING id", (model_id,)) + await cursor.execute("DELETE FROM agent_models WHERE id = %s RETURNING id", (model_id,)) row = await cursor.fetchone() if not row: @@ -528,17 +601,17 @@ async def set_default_model(model_id: str, authorization: Optional[str] = Header # 首先检查模型是否存在 async with pool.connection() as conn: async with conn.cursor() as cursor: - await cursor.execute("SELECT id FROM models WHERE id = %s", (model_id,)) + await cursor.execute("SELECT id FROM agent_models WHERE id = %s", (model_id,)) row = await cursor.fetchone() if not row: raise HTTPException(status_code=404, detail="Model not found") # 取消所有默认设置 - await cursor.execute("UPDATE models SET is_default = FALSE WHERE is_default = TRUE") + await cursor.execute("UPDATE agent_models SET is_default = FALSE WHERE is_default = TRUE") # 设置新的默认模型 - await cursor.execute("UPDATE models SET is_default = TRUE WHERE id = %s", (model_id,)) + await cursor.execute("UPDATE agent_models SET is_default = TRUE WHERE id = %s", (model_id,)) await conn.commit() @@ -566,7 +639,7 @@ async def get_bots(authorization: Optional[str] = Header(None)): async with conn.cursor() as cursor: await cursor.execute(""" SELECT id, name, bot_id, created_at, updated_at - FROM bots + FROM agent_bots ORDER BY created_at DESC """) rows = await cursor.fetchall() @@ -599,19 +672,22 @@ async def create_bot(request: BotCreate, authorization: Optional[str] = Header(N pool = get_db_pool_manager().pool + # 自动生成 bot_id + bot_id = str(uuid.uuid4()) + try: async with pool.connection() as conn: async with conn.cursor() as cursor: await cursor.execute(""" - INSERT INTO bots (name, bot_id) + INSERT INTO agent_bots (name, bot_id) VALUES (%s, %s) RETURNING id, created_at, updated_at - """, (request.name, request.bot_id)) + """, (request.name, bot_id)) row = await cursor.fetchone() # 创建对应的设置记录 await cursor.execute(""" - INSERT INTO bot_settings (bot_id, language) + INSERT INTO agent_bot_settings (bot_id, language) VALUES (%s, 'zh') """, (str(row[0]),)) @@ -620,7 +696,7 @@ async def create_bot(request: BotCreate, authorization: Optional[str] = Header(N return BotResponse( id=str(row[0]), name=request.name, - bot_id=request.bot_id, + bot_id=bot_id, created_at=datetime_to_str(row[1]), updated_at=datetime_to_str(row[2]) ) @@ -671,7 +747,7 @@ async def update_bot( async with pool.connection() as conn: async with conn.cursor() as cursor: await cursor.execute(f""" - UPDATE bots + UPDATE agent_bots SET {', '.join(update_fields)} WHERE id = %s RETURNING id, name, bot_id, created_at, updated_at @@ -710,7 +786,7 @@ async def delete_bot(bot_uuid: str, authorization: Optional[str] = Header(None)) async with pool.connection() as conn: async with conn.cursor() as cursor: - await cursor.execute("DELETE FROM bots WHERE id = %s RETURNING id", (bot_uuid,)) + await cursor.execute("DELETE FROM agent_bots WHERE id = %s RETURNING id", (bot_uuid,)) row = await cursor.fetchone() if not row: @@ -745,7 +821,7 @@ async def get_bot_settings(bot_uuid: str, authorization: Optional[str] = Header( SELECT bot_id, model_id, language, robot_type, dataset_ids, system_prompt, user_identifier, enable_memori, tool_response, skills, updated_at - FROM bot_settings + FROM agent_bot_settings WHERE bot_id = %s """, (bot_uuid,)) row = await cursor.fetchone() @@ -759,7 +835,7 @@ async def get_bot_settings(bot_uuid: str, authorization: Optional[str] = Header( if model_id: await cursor.execute(""" SELECT id, name, provider, model, server, api_key - FROM models WHERE id = %s + FROM agent_models WHERE id = %s """, (model_id,)) model_row = await cursor.fetchone() if model_row: @@ -821,7 +897,7 @@ async def update_bot_settings( if model_id_value: async with pool.connection() as conn: async with conn.cursor() as cursor: - await cursor.execute("SELECT id FROM models WHERE id = %s", (model_id_value,)) + await cursor.execute("SELECT id FROM agent_models WHERE id = %s", (model_id_value,)) if not await cursor.fetchone(): raise HTTPException(status_code=400, detail=f"Model with id '{request.model_id}' not found") # 使用 NULL 或占位符 @@ -864,7 +940,7 @@ async def update_bot_settings( async with pool.connection() as conn: async with conn.cursor() as cursor: # 检查设置是否存在 - await cursor.execute("SELECT bot_id FROM bot_settings WHERE bot_id = %s", (bot_uuid,)) + await cursor.execute("SELECT bot_id FROM agent_bot_settings WHERE bot_id = %s", (bot_uuid,)) if not await cursor.fetchone(): # 不存在则创建 # 需要分别处理带占位符和不带占位符的字段 @@ -935,13 +1011,13 @@ async def update_bot_settings( # 构建 SQL：混合使用占位符和 NULL values_clause = ", ".join(insert_placeholders) - sql = f"INSERT INTO bot_settings ({', '.join(insert_fields)}) VALUES ({values_clause})" + sql = f"INSERT INTO agent_bot_settings ({', '.join(insert_fields)}) VALUES ({values_clause})" await cursor.execute(sql, insert_values) else: # 存在则更新 await cursor.execute(f""" - UPDATE bot_settings + UPDATE agent_bot_settings SET {', '.join(update_fields)} WHERE bot_id = %s """, values) @@ -973,7 +1049,7 @@ async def get_bot_sessions(bot_uuid: str, authorization: Optional[str] = Header( async with conn.cursor() as cursor: await cursor.execute(""" SELECT id, bot_id, title, created_at, updated_at - FROM chat_sessions + FROM agent_chat_sessions WHERE bot_id = %s ORDER BY updated_at DESC """, (bot_uuid,)) @@ -1015,13 +1091,13 @@ async def create_session( # 验证 Bot 是否存在 async with pool.connection() as conn: async with conn.cursor() as cursor: - await cursor.execute("SELECT id FROM bots WHERE id = %s", (bot_uuid,)) + await cursor.execute("SELECT id FROM agent_bots WHERE id = %s", (bot_uuid,)) if not await cursor.fetchone(): raise HTTPException(status_code=404, detail="Bot not found") # 创建会话 await cursor.execute(""" - INSERT INTO chat_sessions (bot_id, title) + INSERT INTO agent_chat_sessions (bot_id, title) VALUES (%s, %s) RETURNING id, created_at, updated_at """, (bot_uuid, request.title)) @@ -1056,7 +1132,7 @@ async def delete_session(session_id: str, authorization: Optional[str] = Header( async with pool.connection() as conn: async with conn.cursor() as cursor: - await cursor.execute("DELETE FROM chat_sessions WHERE id = %s RETURNING id", (session_id,)) + await cursor.execute("DELETE FROM agent_chat_sessions WHERE id = %s RETURNING id", (session_id,)) row = await cursor.fetchone() if not row: @@ -1089,7 +1165,7 @@ async def get_mcp_servers(bot_uuid: str, authorization: Optional[str] = Header(N async with conn.cursor() as cursor: await cursor.execute(""" SELECT id, bot_id, name, type, config, enabled, created_at, updated_at - FROM mcp_servers + FROM agent_mcp_servers WHERE bot_id = %s ORDER BY created_at DESC """, (bot_uuid,)) @@ -1134,12 +1210,12 @@ async def update_mcp_servers( async with pool.connection() as conn: async with conn.cursor() as cursor: # 删除旧的 MCP 配置 - await cursor.execute("DELETE FROM mcp_servers WHERE bot_id = %s", (bot_uuid,)) + await cursor.execute("DELETE FROM agent_mcp_servers WHERE bot_id = %s", (bot_uuid,)) # 插入新的 MCP 配置 for server in servers: await cursor.execute(""" - INSERT INTO mcp_servers (bot_id, name, type, config, enabled) + INSERT INTO agent_mcp_servers (bot_id, name, type, config, enabled) VALUES (%s, %s, %s, %s, %s) """, ( bot_uuid, @@ -1181,13 +1257,13 @@ async def add_mcp_server( # 验证 Bot 是否存在 async with pool.connection() as conn: async with conn.cursor() as cursor: - await cursor.execute("SELECT id FROM bots WHERE id = %s", (bot_uuid,)) + await cursor.execute("SELECT id FROM agent_bots WHERE id = %s", (bot_uuid,)) if not await cursor.fetchone(): raise HTTPException(status_code=404, detail="Bot not found") # 创建 MCP 服务器 await cursor.execute(""" - INSERT INTO mcp_servers (bot_id, name, type, config, enabled) + INSERT INTO agent_mcp_servers (bot_id, name, type, config, enabled) VALUES (%s, %s, %s, %s, %s) RETURNING id, created_at, updated_at """, ( @@ -1237,7 +1313,7 @@ async def delete_mcp_server( async with pool.connection() as conn: async with conn.cursor() as cursor: await cursor.execute( - "DELETE FROM mcp_servers WHERE id = %s AND bot_id = %s RETURNING id", + "DELETE FROM agent_mcp_servers WHERE id = %s AND bot_id = %s RETURNING id", (mcp_id, bot_uuid) ) row = await cursor.fetchone() @@ -1248,3 +1324,96 @@ async def delete_mcp_server( await conn.commit() return SuccessResponse(success=True, message="MCP server deleted successfully") + + +# ============== Admin 登录 API ============== + +@router.post("/api/v1/admin/login", response_model=AdminLoginResponse) +async def admin_login(request: AdminLoginRequest): + """ + 管理员登录 + + Args: + request: 登录请求（用户名和密码） + + Returns: + AdminLoginResponse: 登录成功返回 token + """ + # 硬编码验证账号密码 + if request.username != ADMIN_USERNAME or request.password != ADMIN_PASSWORD: + raise HTTPException( + status_code=401, + detail="用户名或密码错误" + ) + + pool = get_db_pool_manager().pool + + # 生成 token + token = secrets.token_urlsafe(32) + expires_at = datetime.now() + timedelta(hours=TOKEN_EXPIRE_HOURS) + + async with pool.connection() as conn: + async with conn.cursor() as cursor: + # 清理该用户的旧 token + await cursor.execute("DELETE FROM agent_admin_tokens WHERE username = %s", (request.username,)) + + # 保存新 token + await cursor.execute(""" + INSERT INTO agent_admin_tokens (username, token, expires_at) + VALUES (%s, %s, %s) + """, (request.username, token, expires_at)) + + await conn.commit() + + return AdminLoginResponse( + token=token, + username=request.username, + expires_at=expires_at.isoformat() + ) + + +@router.post("/api/v1/admin/verify", response_model=AdminVerifyResponse) +async def admin_verify(authorization: Optional[str] = Header(None)): + """ + 验证管理员 token 是否有效 + + Args: + authorization: Bearer token + + Returns: + AdminVerifyResponse: 验证结果 + """ + valid, username = await verify_admin_auth(authorization) + + return AdminVerifyResponse( + valid=valid, + username=username + ) + + +@router.post("/api/v1/admin/logout", response_model=SuccessResponse) +async def admin_logout(authorization: Optional[str] = Header(None)): + """ + 管理员登出（删除 token） + + Args: + authorization: Bearer token + + Returns: + SuccessResponse: 登出结果 + """ + provided_token = extract_api_key_from_auth(authorization) + if not provided_token: + raise HTTPException( + status_code=401, + detail="Authorization header is required" + ) + + pool = get_db_pool_manager().pool + + async with pool.connection() as conn: + async with conn.cursor() as cursor: + await cursor.execute("DELETE FROM agent_admin_tokens WHERE token = %s", (provided_token,)) + await conn.commit() + + return SuccessResponse(success=True, message="Logged out successfully")