Pretor/pretor/api/agent.py

140 lines
6.1 KiB
Python

# 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):
provider_title: str
model_id: str
individual_name: str
class AgentLocalRegister(BaseModel):
path: str
individual_name: str
@agent_router.post("")
async def load_agent(agent_register: Union[AgentRegister, AgentLocalRegister],
_: TokenData = Depends(RoleChecker(allowed_roles=UserAuthority.USER))):
global_state_machine = ray_actor_hook("global_state_machine").global_state_machine
if isinstance(agent_register, AgentLocalRegister):
pass
elif isinstance(agent_register, AgentRegister):
match agent_register.individual_name:
case "supervisory_node":
node = ray_actor_hook("supervisory_node").supervisory_node
node.create_agent.remote(global_state_machine,agent_register.provider_title,agent_register.model_id)
case "consciousness_node":
node = ray_actor_hook("consciousness_node").consciousness_node
node.create_agent.remote(global_state_machine,agent_register.provider_title,agent_register.model_id)
case "control_node":
node = ray_actor_hook("control_node").control_node
node.create_agent.remote(global_state_machine,agent_register.provider_title,agent_register.model_id)
case _:
pass
return {"message": "创建成功"}
class WorkerIndividualCreate(BaseModel):
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]
class WorkerIndividualUpdate(BaseModel):
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
@agent_router.post("/worker")
async def create_worker_individual(worker_data: WorkerIndividualCreate,
token_data: TokenData = Depends(RoleChecker(allowed_roles=UserAuthority.USER))):
postgres_database = ray_actor_hook("postgres_database")
data_dict = worker_data.model_dump()
data_dict["owner_id"] = token_data.user_id
worker = await postgres_database.individual_database.remote("add_worker_individual", **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)):
postgres_database = ray_actor_hook("postgres_database")
workers = await postgres_database.individual_database.remote("get_worker_individual_list", 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)):
postgres_database = ray_actor_hook("postgres_database")
worker = await postgres_database.individual_database.remote("get_worker_individual", 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)):
postgres_database = ray_actor_hook("postgres_database")
worker = await postgres_database.individual_database.remote("get_worker_individual", 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.individual_database.remote("update_worker_individual", agent_id=agent_id, **update_data)
return {"message": "success", "worker": updated_worker}
@agent_router.delete("/worker/{agent_id}")
async def delete_worker_individual(agent_id: str,
token_data: TokenData = Depends(Accessor.get_current_user)):
postgres_database = ray_actor_hook("postgres_database")
worker = await postgres_database.individual_database.remote("get_worker_individual", 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.individual_database.remote("delete_worker_individual", agent_id=agent_id)
return {"message": "success"}