diff --git a/migrations/add_single_agent_mode.sql b/migrations/add_single_agent_mode.sql new file mode 100644 index 0000000..f86ac74 --- /dev/null +++ b/migrations/add_single_agent_mode.sql @@ -0,0 +1,15 @@ +-- ============================================================ +-- Single Agent Mode Migration +-- 为用户表添加单智能体 bot_id 字段 +-- ============================================================ + +-- 添加 single_agent_bot_id 字段 +ALTER TABLE agent_user +ADD COLUMN IF NOT EXISTS single_agent_bot_id UUID; + +-- 创建索引 +CREATE INDEX IF NOT EXISTS idx_agent_user_single_agent_bot +ON agent_user(single_agent_bot_id) WHERE single_agent_bot_id IS NOT NULL; + +-- 添加注释 +COMMENT ON COLUMN agent_user.single_agent_bot_id IS '单智能体模式下用户的专属智能体ID'; diff --git a/routes/bot_manager.py b/routes/bot_manager.py index cc529d2..5fb3266 100644 --- a/routes/bot_manager.py +++ b/routes/bot_manager.py @@ -17,6 +17,7 @@ from pydantic import BaseModel from agent.db_pool_manager import get_db_pool_manager from utils.fastapi_utils import extract_api_key_from_auth from utils.new_api_proxy import get_new_api_proxy +from utils.settings import SINGLE_AGENT_MODE, TEMPLATE_BOT_ID, TEMPLATE_BOT_NAME logger = logging.getLogger('app') @@ -55,6 +56,105 @@ def copy_skills_folder(source_bot_id: str, target_bot_id: str) -> bool: logger.error(f"Failed to copy skills folder: {e}") return False + +async def get_or_create_single_agent_bot(user_id: str, pool): + """ + 获取或创建用户的单智能体 bot + + Args: + user_id: 用户 ID + pool: 数据库连接池 + + Returns: + dict: { enabled: bool, bot_name: str, bot_id: str } 或 { enabled: False } + """ + if not SINGLE_AGENT_MODE: + return {"enabled": False} + + async with pool.connection() as conn: + async with conn.cursor() as cursor: + # 检查用户是否已有 single_agent_bot_id + await cursor.execute(""" + SELECT single_agent_bot_id + FROM agent_user + WHERE id = %s + """, (user_id,)) + row = await cursor.fetchone() + + user_bot_id = row[0] if row and row[0] else None + + # 如果用户没有 bot_id，自动复制模板 + if not user_bot_id: + if not TEMPLATE_BOT_ID: + logger.warning("TEMPLATE_BOT_ID not configured") + return {"enabled": False} + + try: + # 获取模板 bot + await cursor.execute(""" + SELECT id, name, settings + FROM agent_bots + WHERE id = %s AND is_published = TRUE + """, (TEMPLATE_BOT_ID,)) + template = await cursor.fetchone() + + if not template: + logger.warning(f"Template bot {TEMPLATE_BOT_ID} not found") + return {"enabled": False} + + template_id, template_name, template_settings = template + settings = template_settings if template_settings else {} + + # 创建新 bot + new_bot_id = str(uuid.uuid4()) + + # 复制设置（不复制知识库） + new_settings = { + 'model_id': settings.get('model_id'), + 'language': settings.get('language', 'zh'), + 'avatar_url': settings.get('avatar_url'), + 'description': settings.get('description'), + 'suggestions': settings.get('suggestions'), + 'system_prompt': settings.get('system_prompt'), + 'enable_memori': settings.get('enable_memori', False), + 'enable_thinking': settings.get('enable_thinking', False), + 'tool_response': settings.get('tool_response', False), + 'skills': settings.get('skills'), + } + + # 插入新 bot + await cursor.execute(""" + INSERT INTO agent_bots (name, bot_id, owner_id, settings, copied_from) + VALUES (%s, %s, %s, %s, %s) + RETURNING id + """, (TEMPLATE_BOT_NAME, new_bot_id, user_id, json.dumps(new_settings), template_id)) + new_row = await cursor.fetchone() + user_bot_id = new_row[0] + + # 更新用户的 single_agent_bot_id + await cursor.execute(""" + UPDATE agent_user + SET single_agent_bot_id = %s + WHERE id = %s + """, (user_bot_id, user_id)) + + # 复制 skills 文件夹 + copy_skills_folder(str(template_id), str(new_bot_id)) + + await conn.commit() + + logger.info(f"Created single agent bot {user_bot_id} for user {user_id}") + + except Exception as e: + logger.error(f"Failed to copy single agent bot: {e}") + return {"enabled": False} + + return { + "enabled": True, + "bot_name": TEMPLATE_BOT_NAME, + "bot_id": str(user_bot_id) + } + # ============== Admin 配置 ============== ADMIN_USERNAME = "admin" ADMIN_PASSWORD = "Admin123" # 生产环境应使用环境变量 @@ -335,6 +435,7 @@ class UserLoginResponse(BaseModel): email: Optional[str] = None is_admin: bool = False expires_at: str + single_agent: Optional[dict] = None # 单智能体模式配置 class UserVerifyResponse(BaseModel): @@ -2195,13 +2296,17 @@ async def user_login(request: UserLoginRequest): await conn.commit() + # 获取单智能体配置 + single_agent_config = await get_or_create_single_agent_bot(str(user_id), pool) + return UserLoginResponse( token=token, user_id=str(user_id), username=username, email=email, is_admin=is_admin or False, - expires_at=expires_at.isoformat() + expires_at=expires_at.isoformat(), + single_agent=single_agent_config ) diff --git a/utils/settings.py b/utils/settings.py index 3423b4f..810c026 100644 --- a/utils/settings.py +++ b/utils/settings.py @@ -113,3 +113,10 @@ NEW_API_TIMEOUT = int(os.getenv("NEW_API_TIMEOUT", "30")) # New API 管理员密钥（用于同步用户等管理操作，可选） NEW_API_ADMIN_KEY = os.getenv("NEW_API_ADMIN_KEY", "") + +# ============================================================ +# Single Agent Mode Configuration +# ============================================================ +SINGLE_AGENT_MODE = os.getenv("SINGLE_AGENT_MODE", "false") == "true" +TEMPLATE_BOT_ID = os.getenv("TEMPLATE_BOT_ID", "") +TEMPLATE_BOT_NAME = os.getenv("TEMPLATE_BOT_NAME", "智能助手")