2024-11-09 14:32:24 -08:00
|
|
|
from typing import Dict
|
2024-11-26 15:39:36 -08:00
|
|
|
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException
|
|
|
|
|
|
|
|
from ...database import DatabaseManager # Add this import
|
2024-11-09 14:32:24 -08:00
|
|
|
from ...datamodel import Agent, Model, Tool
|
2024-11-26 15:39:36 -08:00
|
|
|
from ..deps import get_db
|
2024-11-09 14:32:24 -08:00
|
|
|
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
|
|
|
|
|
|
@router.get("/")
|
2024-11-26 15:39:36 -08:00
|
|
|
async def list_agents(user_id: str, db: DatabaseManager = Depends(get_db)) -> Dict:
|
2024-11-09 14:32:24 -08:00
|
|
|
"""List all agents for a user"""
|
|
|
|
response = db.get(Agent, filters={"user_id": user_id})
|
2024-11-26 15:39:36 -08:00
|
|
|
return {"status": True, "data": response.data}
|
2024-11-09 14:32:24 -08:00
|
|
|
|
|
|
|
|
|
|
|
@router.get("/{agent_id}")
|
2024-11-26 15:39:36 -08:00
|
|
|
async def get_agent(agent_id: int, user_id: str, db: DatabaseManager = Depends(get_db)) -> Dict:
|
2024-11-09 14:32:24 -08:00
|
|
|
"""Get a specific agent"""
|
2024-11-26 15:39:36 -08:00
|
|
|
response = db.get(Agent, filters={"id": agent_id, "user_id": user_id})
|
2024-11-09 14:32:24 -08:00
|
|
|
if not response.status or not response.data:
|
|
|
|
raise HTTPException(status_code=404, detail="Agent not found")
|
2024-11-26 15:39:36 -08:00
|
|
|
return {"status": True, "data": response.data[0]}
|
2024-11-09 14:32:24 -08:00
|
|
|
|
|
|
|
|
|
|
|
@router.post("/")
|
2024-11-26 15:39:36 -08:00
|
|
|
async def create_agent(agent: Agent, db: DatabaseManager = Depends(get_db)) -> Dict:
|
2024-11-09 14:32:24 -08:00
|
|
|
"""Create a new agent"""
|
|
|
|
response = db.upsert(agent)
|
|
|
|
if not response.status:
|
|
|
|
raise HTTPException(status_code=400, detail=response.message)
|
2024-11-26 15:39:36 -08:00
|
|
|
return {"status": True, "data": response.data}
|
2024-11-09 14:32:24 -08:00
|
|
|
|
|
|
|
|
|
|
|
@router.delete("/{agent_id}")
|
2024-11-26 15:39:36 -08:00
|
|
|
async def delete_agent(agent_id: int, user_id: str, db: DatabaseManager = Depends(get_db)) -> Dict:
|
2024-11-09 14:32:24 -08:00
|
|
|
"""Delete an agent"""
|
2024-11-26 15:39:36 -08:00
|
|
|
db.delete(filters={"id": agent_id, "user_id": user_id}, model_class=Agent)
|
|
|
|
return {"status": True, "message": "Agent deleted successfully"}
|
|
|
|
|
2024-11-09 14:32:24 -08:00
|
|
|
|
|
|
|
# Agent-Model link endpoints
|
|
|
|
|
|
|
|
|
|
|
|
@router.post("/{agent_id}/models/{model_id}")
|
2024-11-26 15:39:36 -08:00
|
|
|
async def link_agent_model(agent_id: int, model_id: int, db: DatabaseManager = Depends(get_db)) -> Dict:
|
2024-11-09 14:32:24 -08:00
|
|
|
"""Link a model to an agent"""
|
2024-11-26 15:39:36 -08:00
|
|
|
db.link(link_type="agent_model", primary_id=agent_id, secondary_id=model_id)
|
|
|
|
return {"status": True, "message": "Model linked to agent successfully"}
|