diff --git a/routes/bot_manager.py b/routes/bot_manager.py index d33f788..43e2e2f 100644 --- a/routes/bot_manager.py +++ b/routes/bot_manager.py @@ -346,6 +346,10 @@ async def get_parent_user_id(user_id: str) -> Optional[str]: Returns: Optional[str]: 主账号ID,如果不是子账号则返回自身ID """ + # masterkey 用户无需查询数据库 + if user_id == "__masterkey__": + return user_id + pool = get_db_pool_manager().pool async with pool.connection() as conn: async with conn.cursor() as cursor: @@ -1703,37 +1707,38 @@ async def get_bot_settings(bot_uuid: str, authorization: Optional[str] = Header( # 获取用户可用的模型列表 models_list = [] try: - async with pool.connection() as conn: - async with conn.cursor() as cursor: - # 获取用户的 new_api_session 和 new_api_user_id(子账号使用主账号的) - effective_user_id = await get_parent_user_id(user_id) - await cursor.execute(""" - SELECT new_api_session, new_api_user_id - FROM agent_user - WHERE id = %s - """, (effective_user_id,)) - row = await cursor.fetchone() + if user_id != "__masterkey__": + async with pool.connection() as conn: + async with conn.cursor() as cursor: + # 获取用户的 new_api_session 和 new_api_user_id(子账号使用主账号的) + effective_user_id = await get_parent_user_id(user_id) + await cursor.execute(""" + SELECT new_api_session, new_api_user_id + FROM agent_user + WHERE id = %s + """, (effective_user_id,)) + row = await cursor.fetchone() - if row and row[0] and row[1]: - new_api_session = row[0] - new_api_user_id = row[1] + if row and row[0] and row[1]: + new_api_session = row[0] + new_api_user_id = row[1] - # 调用 New API 获取模型列表 - proxy = get_new_api_proxy() - cookies = {"session": new_api_session} + # 调用 New API 获取模型列表 + proxy = get_new_api_proxy() + cookies = {"session": new_api_session} - result = await proxy.get_user_models(cookies, new_api_user_id) + result = await proxy.get_user_models(cookies, new_api_user_id) - if result.get("success"): - data = result.get("data", []) - if isinstance(data, list): - for model in data: - models_list.append(NewAPIModelResponse( - id=model if isinstance(model, str) else model.get("id", ""), - object="model", - created=None, - owned_by="system" - )) + if result.get("success"): + data = result.get("data", []) + if isinstance(data, list): + for model in data: + models_list.append(NewAPIModelResponse( + id=model if isinstance(model, str) else model.get("id", ""), + object="model", + created=None, + owned_by="system" + )) except Exception as e: # 获取模型列表失败不影响主流程,返回空列表 print(f"Failed to fetch models: {e}") @@ -1756,9 +1761,12 @@ async def get_bot_settings(bot_uuid: str, authorization: Optional[str] = Header( bot_id, bot_name, settings_json, updated_at, is_published, copied_from, owner_id = row settings = settings_json if settings_json else {} - # 判断当前用户是否是所有者(子账号使用主账号ID判断) - effective_user_id = await get_parent_user_id(user_id) - is_owner = (str(owner_id) == str(effective_user_id)) + # 判断当前用户是否是所有者(子账号使用主账号ID判断,masterkey 视为所有者) + if user_id == "__masterkey__": + is_owner = True + else: + effective_user_id = await get_parent_user_id(user_id) + is_owner = (str(owner_id) == str(effective_user_id)) # 检查 model_id 是否在可用模型列表中 model_id = settings.get('model_id') @@ -1800,13 +1808,14 @@ async def get_bot_settings(bot_uuid: str, authorization: Optional[str] = Header( # 处理 dataset_ids # 单智能体模式:加载用户的所有知识库(子账号使用主账号的知识库) # 普通模式:从 settings 读取 - if SINGLE_AGENT_MODE: + if SINGLE_AGENT_MODE and user_id != "__masterkey__": + effective_user_id_for_ds = await get_parent_user_id(user_id) if user_id != "__masterkey__" else user_id await cursor.execute(""" SELECT dataset_id FROM user_datasets WHERE user_id = %s ORDER BY created_at DESC - """, (effective_user_id,)) + """, (effective_user_id_for_ds,)) user_datasets = await cursor.fetchall() dataset_ids = [row[0] for row in user_datasets] if user_datasets else [] else: