Pretor/pretor/api/agent.py

139 lines
5.8 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
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(Accessor.get_current_user)):
global_state_machine = ray_actor_hook("global_state_machine")
if isinstance(agent_register, AgentLocalRegister):
pass
elif isinstance(agent_register, AgentRegister):
match agent_register.individual_title:
case "supervisory_node":
node = ray_actor_hook("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")
node.create_agent.remote(global_state_machine,agent_register.provider_title,agent_register.model_id)
case "control_node":
node = ray_actor_hook("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(Accessor.get_current_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.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)):
postgres_database = ray_actor_hook("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)):
postgres_database = ray_actor_hook("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)):
postgres_database = ray_actor_hook("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)
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.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"}