KiloStar/pretor/api/agent.py

#  Copyright 2026 zhaoxi826
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.


from typing import Union
from pretor.utils.ray_hook import ray_actor_hook
from fastapi import APIRouter, Depends
from pydantic import BaseModel
from pretor.utils.access import Accessor, TokenData
from pretor.core.database.table.individual import AgentType
from fastapi import HTTPException
from typing import Optional, List, Dict
from pretor.utils.check_user.role_check import RoleChecker
from pretor.core.database.table.user import UserAuthority

agent_router = APIRouter(prefix="/api/v1/agent", tags=["agent"])


class AgentRegister(BaseModel):
    """AgentRegister 核心组件类。
    这是一个领域数据模型或功能封装类，承载了 AgentRegister 相关的内聚属性定义与状态维护。它的存在隔离了局部的业务复杂性，并对外提供了类型安全的访问接口。"""

    provider_title: str
    model_id: str
    individual_name: str
    tools: Optional[List[str]] = None


class AgentLocalRegister(BaseModel):
    """AgentLocalRegister 核心组件类。
    这是一个领域数据模型或功能封装类，承载了 AgentLocalRegister 相关的内聚属性定义与状态维护。它的存在隔离了局部的业务复杂性，并对外提供了类型安全的访问接口。"""

    path: str
    individual_name: str
    tools: Optional[List[str]] = None


@agent_router.get("")
async def get_system_nodes(
    _: TokenData = Depends(RoleChecker(allowed_roles=UserAuthority.USER)),
):
    """处理针对 get system nodes 相关的 HTTP API 请求。
    该接口负责解析前端传入的载荷数据，调用底层核心业务逻辑进行处理，并组装标准化的 JSON 响应。
    Args:        _ (TokenData): 参与 get system nodes 逻辑运算或数据构建的上下文依赖对象。
    Returns:        : 序列化后的标准网络响应模型（如包含业务状态码、成功标志及对应的数据载荷 Data）。"""
    postgres_database = ray_actor_hook("postgres_database").postgres_database
    configs = await postgres_database.get_all_system_node_configs.remote()
    return {"system_nodes": configs}


@agent_router.post("")
async def load_agent(
    agent_register: Union[AgentRegister, AgentLocalRegister],
    _: TokenData = Depends(RoleChecker(allowed_roles=UserAuthority.USER)),
):
    """处理针对 load agent 相关的 HTTP API 请求。
    该接口负责解析前端传入的载荷数据，调用底层核心业务逻辑进行处理，并组装标准化的 JSON 响应。
    Args:        agent_register (Union[AgentRegister, AgentLocalRegister]): 参与 load agent 逻辑运算或数据构建的上下文依赖对象。        _ (TokenData): 参与 load agent 逻辑运算或数据构建的上下文依赖对象。
    Returns:        : 序列化后的标准网络响应模型（如包含业务状态码、成功标志及对应的数据载荷 Data）。"""
    global_state_machine = ray_actor_hook("global_state_machine").global_state_machine
    postgres_database = ray_actor_hook("postgres_database").postgres_database

    if isinstance(agent_register, AgentLocalRegister):
        pass

    elif isinstance(agent_register, AgentRegister):
        try:
            # Persist configuration
            await postgres_database.upsert_system_node_config.remote(
                agent_register.individual_name,
                agent_register.provider_title,
                agent_register.model_id,
                agent_register.tools,
            )

            # Load agent into state machine
            match agent_register.individual_name:
                case "supervisory_node":
                    node = ray_actor_hook("supervisory_node").supervisory_node
                    await node.create_agent.remote(
                        global_state_machine,
                        agent_register.provider_title,
                        agent_register.model_id,
                        agent_register.tools,
                    )
                case "consciousness_node":
                    node = ray_actor_hook("consciousness_node").consciousness_node
                    await node.create_agent.remote(
                        global_state_machine,
                        agent_register.provider_title,
                        agent_register.model_id,
                        agent_register.tools,
                    )
                case "control_node":
                    node = ray_actor_hook("control_node").control_node
                    await node.create_agent.remote(
                        global_state_machine,
                        agent_register.provider_title,
                        agent_register.model_id,
                        agent_register.tools,
                    )
                case _:
                    pass
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"加载节点失败: {str(e)}")
    return {"message": "创建成功"}


class WorkerIndividualCreate(BaseModel):
    """WorkerIndividualCreate 核心组件类。
    这是一个具体的 Worker 智能体实体类，代表着具备特定人设、领域技能或长文本处理能力的数字员工。它可以被控制器动态拉起，并在安全沙箱内执行复杂的工作流指令与多步骤推理任务。"""

    agent_name: str
    agent_type: AgentType
    description: str
    provider_title: str
    model_id: str
    system_prompt: str
    output_template: dict
    bound_skill: Dict[str, List[str]]
    workspace: List[str]
    tools: Optional[List[str]] = None


class WorkerIndividualUpdate(BaseModel):
    """WorkerIndividualUpdate 核心组件类。
    这是一个具体的 Worker 智能体实体类，代表着具备特定人设、领域技能或长文本处理能力的数字员工。它可以被控制器动态拉起，并在安全沙箱内执行复杂的工作流指令与多步骤推理任务。"""

    agent_name: Optional[str] = None
    agent_type: Optional[AgentType] = None
    description: Optional[str] = None
    provider_title: Optional[str] = None
    model_id: Optional[str] = None
    system_prompt: Optional[str] = None
    output_template: Optional[dict] = None
    bound_skill: Optional[Dict[str, List[str]]] = None
    workspace: Optional[List[str]] = None
    tools: Optional[List[str]] = None


@agent_router.post("/worker")
async def create_worker_individual(
    worker_data: WorkerIndividualCreate,
    token_data: TokenData = Depends(RoleChecker(allowed_roles=UserAuthority.USER)),
):
    """处理针对 create worker individual 相关的 HTTP API 请求。
    该接口负责解析前端传入的载荷数据，调用底层核心业务逻辑进行处理，并组装标准化的 JSON 响应。
    Args:        worker_data (WorkerIndividualCreate): 从客户端传递过来或由上游组件生成的核心业务数据体，通常需要进一步的清洗和结构化解析。        token_data (TokenData): 从客户端传递过来或由上游组件生成的核心业务数据体，通常需要进一步的清洗和结构化解析。
    Returns:        : 序列化后的标准网络响应模型（如包含业务状态码、成功标志及对应的数据载荷 Data）。"""
    postgres_database = ray_actor_hook("postgres_database").postgres_database
    data_dict = worker_data.model_dump()
    data_dict["owner_id"] = token_data.user_id
    worker = await postgres_database.add_worker_individual.remote(**data_dict)
    return {"message": "success", "agent_id": worker.agent_id}


@agent_router.get("/worker")
async def get_worker_individual_list(
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    """处理针对 get worker individual list 相关的 HTTP API 请求。
    该接口负责解析前端传入的载荷数据，调用底层核心业务逻辑进行处理，并组装标准化的 JSON 响应。
    Args:        token_data (TokenData): 从客户端传递过来或由上游组件生成的核心业务数据体，通常需要进一步的清洗和结构化解析。
    Returns:        : 序列化后的标准网络响应模型（如包含业务状态码、成功标志及对应的数据载荷 Data）。"""
    postgres_database = ray_actor_hook("postgres_database").postgres_database
    workers = await postgres_database.get_worker_individual_list.remote(
        owner_id=token_data.user_id
    )
    return {"workers": workers}


@agent_router.get("/worker/{agent_id}")
async def get_worker_individual(
    agent_id: str, token_data: TokenData = Depends(Accessor.get_current_user)
):
    """处理针对 get worker individual 相关的 HTTP API 请求。
    该接口负责解析前端传入的载荷数据，调用底层核心业务逻辑进行处理，并组装标准化的 JSON 响应。
    Args:        agent_id (str): 目标对象的唯一全局标识符 (UUID/ULID)，用于在数据库表或缓存结构中精准匹配该 agent 实例。        token_data (TokenData): 从客户端传递过来或由上游组件生成的核心业务数据体，通常需要进一步的清洗和结构化解析。
    Returns:        : 序列化后的标准网络响应模型（如包含业务状态码、成功标志及对应的数据载荷 Data）。"""
    postgres_database = ray_actor_hook("postgres_database").postgres_database
    worker = await postgres_database.get_worker_individual.remote(agent_id=agent_id)
    if not worker:
        raise HTTPException(status_code=404, detail="Agent not found")
    if worker.owner_id != token_data.user_id:
        raise HTTPException(
            status_code=403, detail="Forbidden: You do not own this agent"
        )
    return worker


@agent_router.put("/worker/{agent_id}")
async def update_worker_individual(
    agent_id: str,
    worker_data: WorkerIndividualUpdate,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    """处理针对 update worker individual 相关的 HTTP API 请求。
    该接口负责解析前端传入的载荷数据，调用底层核心业务逻辑进行处理，并组装标准化的 JSON 响应。
    Args:        agent_id (str): 目标对象的唯一全局标识符 (UUID/ULID)，用于在数据库表或缓存结构中精准匹配该 agent 实例。        worker_data (WorkerIndividualUpdate): 从客户端传递过来或由上游组件生成的核心业务数据体，通常需要进一步的清洗和结构化解析。        token_data (TokenData): 从客户端传递过来或由上游组件生成的核心业务数据体，通常需要进一步的清洗和结构化解析。
    Returns:        : 序列化后的标准网络响应模型（如包含业务状态码、成功标志及对应的数据载荷 Data）。"""
    postgres_database = ray_actor_hook("postgres_database").postgres_database
    worker = await postgres_database.get_worker_individual.remote(agent_id=agent_id)
    if not worker:
        raise HTTPException(status_code=404, detail="Agent not found")
    if worker.owner_id != token_data.user_id:
        raise HTTPException(
            status_code=403, detail="Forbidden: You do not own this agent"
        )

    update_data = worker_data.model_dump(exclude_unset=True)
    updated_worker = await postgres_database.update_worker_individual.remote(
        agent_id=agent_id, **update_data
    )

    global_state_machine = ray_actor_hook("global_state_machine").global_state_machine
    try:
        await global_state_machine.remove_individual.remote(agent_id)
    except Exception:
        pass

    return {"message": "success", "worker": updated_worker}


@agent_router.post("/worker/{agent_id}/reload")
async def reload_worker_individual(
    agent_id: str, token_data: TokenData = Depends(Accessor.get_current_user)
):
    """处理针对 reload worker individual 相关的 HTTP API 请求。
    该接口负责解析前端传入的载荷数据，调用底层核心业务逻辑进行处理，并组装标准化的 JSON 响应。
    Args:        agent_id (str): 目标对象的唯一全局标识符 (UUID/ULID)，用于在数据库表或缓存结构中精准匹配该 agent 实例。        token_data (TokenData): 从客户端传递过来或由上游组件生成的核心业务数据体，通常需要进一步的清洗和结构化解析。
    Returns:        : 序列化后的标准网络响应模型（如包含业务状态码、成功标志及对应的数据载荷 Data）。"""
    postgres_database = ray_actor_hook("postgres_database").postgres_database
    worker = await postgres_database.get_worker_individual.remote(agent_id=agent_id)
    if not worker:
        raise HTTPException(status_code=404, detail="Agent not found")
    if worker.owner_id != token_data.user_id:
        raise HTTPException(
            status_code=403, detail="Forbidden: You do not own this agent"
        )

    global_state_machine = ray_actor_hook("global_state_machine").global_state_machine
    await global_state_machine.remove_individual.remote(agent_id)

    return {"message": "Worker will be reloaded on next use"}


@agent_router.delete("/worker/{agent_id}")
async def delete_worker_individual(
    agent_id: str, token_data: TokenData = Depends(Accessor.get_current_user)
):
    """处理针对 delete worker individual 相关的 HTTP API 请求。
    该接口负责解析前端传入的载荷数据，调用底层核心业务逻辑进行处理，并组装标准化的 JSON 响应。
    Args:        agent_id (str): 目标对象的唯一全局标识符 (UUID/ULID)，用于在数据库表或缓存结构中精准匹配该 agent 实例。        token_data (TokenData): 从客户端传递过来或由上游组件生成的核心业务数据体，通常需要进一步的清洗和结构化解析。
    Returns:        : 序列化后的标准网络响应模型（如包含业务状态码、成功标志及对应的数据载荷 Data）。"""
    postgres_database = ray_actor_hook("postgres_database").postgres_database
    worker = await postgres_database.get_worker_individual.remote(agent_id=agent_id)
    if not worker:
        raise HTTPException(status_code=404, detail="Agent not found")
    if worker.owner_id != token_data.user_id:
        raise HTTPException(
            status_code=403, detail="Forbidden: You do not own this agent"
        )
    await postgres_database.delete_worker_individual.remote(agent_id=agent_id)
    return {"message": "success"}