wehub-resource-sync
420 lines
17 KiB
Python
420 lines
17 KiB
Python
from fastapi import APIRouter
|
|
from fastapi import HTTPException, Depends
|
|
from fastapi_jwt_auth import AuthJWT
|
|
from fastapi_sqlalchemy import db
|
|
from pydantic import BaseModel
|
|
from sqlalchemy import desc
|
|
import ast
|
|
|
|
from pytz import timezone
|
|
from sqlalchemy import func, or_
|
|
from superagi.models.agent import Agent
|
|
from superagi.models.agent_execution_config import AgentExecutionConfiguration
|
|
from superagi.models.agent_config import AgentConfiguration
|
|
from superagi.models.agent_schedule import AgentSchedule
|
|
from superagi.models.agent_template import AgentTemplate
|
|
from superagi.models.project import Project
|
|
from superagi.models.workflows.agent_workflow import AgentWorkflow
|
|
from superagi.models.agent_execution import AgentExecution
|
|
from superagi.models.tool import Tool
|
|
from superagi.controllers.types.agent_schedule import AgentScheduleInput
|
|
from superagi.controllers.types.agent_with_config import AgentConfigInput
|
|
from superagi.controllers.types.agent_with_config_schedule import AgentConfigSchedule
|
|
from jsonmerge import merge
|
|
from datetime import datetime
|
|
import json
|
|
|
|
from superagi.models.toolkit import Toolkit
|
|
from superagi.models.knowledges import Knowledges
|
|
|
|
from sqlalchemy import func
|
|
# from superagi.types.db import AgentOut, AgentIn
|
|
from superagi.helper.auth import check_auth
|
|
from superagi.apm.event_handler import EventHandler
|
|
from superagi.models.workflows.iteration_workflow import IterationWorkflow
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
class AgentOut(BaseModel):
|
|
id: int
|
|
name: str
|
|
project_id: int
|
|
description: str
|
|
created_at: datetime
|
|
updated_at: datetime
|
|
|
|
class Config:
|
|
orm_mode = True
|
|
|
|
|
|
class AgentIn(BaseModel):
|
|
name: str
|
|
project_id: int
|
|
description: str
|
|
|
|
class Config:
|
|
orm_mode = True
|
|
|
|
|
|
@router.post("/create", status_code=201)
|
|
def create_agent_with_config(agent_with_config: AgentConfigInput,
|
|
Authorize: AuthJWT = Depends(check_auth)):
|
|
"""
|
|
Create a new agent with configurations.
|
|
|
|
Args:
|
|
agent_with_config (AgentConfigInput): Data for creating a new agent with configurations.
|
|
- name (str): Name of the agent.
|
|
- project_id (int): Identifier of the associated project.
|
|
- description (str): Description of the agent.
|
|
- goal (List[str]): List of goals for the agent.
|
|
- constraints (List[str]): List of constraints for the agent.
|
|
- tools (List[int]): List of tool identifiers associated with the agent.
|
|
- exit (str): Exit condition for the agent.
|
|
- iteration_interval (int): Interval between iterations for the agent.
|
|
- model (str): Model information for the agent.
|
|
- permission_type (str): Permission type for the agent.
|
|
- LTM_DB (str): LTM database for the agent.
|
|
- max_iterations (int): Maximum number of iterations for the agent.
|
|
- user_timezone (string): Timezone of the user
|
|
|
|
Returns:
|
|
dict: Dictionary containing the created agent's ID, execution ID, name, and content type.
|
|
|
|
Raises:
|
|
HTTPException (status_code=404): If the associated project or any of the tools is not found.
|
|
"""
|
|
|
|
project = db.session.query(Project).get(agent_with_config.project_id)
|
|
if not project:
|
|
raise HTTPException(status_code=404, detail="Project not found")
|
|
|
|
invalid_tools = Tool.get_invalid_tools(agent_with_config.tools, db.session)
|
|
if len(invalid_tools) > 0: # If the returned value is not True (then it is an invalid tool_id)
|
|
raise HTTPException(status_code=404,
|
|
|
|
detail=f"Tool with IDs {str(invalid_tools)} does not exist. 404 Not Found.")
|
|
|
|
agent_toolkit_tools = Toolkit.fetch_tool_ids_from_toolkit(session=db.session,
|
|
toolkit_ids=agent_with_config.toolkits)
|
|
agent_with_config.tools.extend(agent_toolkit_tools)
|
|
db_agent = Agent.create_agent_with_config(db, agent_with_config)
|
|
|
|
start_step = AgentWorkflow.fetch_trigger_step_id(db.session, db_agent.agent_workflow_id)
|
|
iteration_step_id = IterationWorkflow.fetch_trigger_step_id(db.session,
|
|
start_step.action_reference_id).id if start_step.action_type == "ITERATION_WORKFLOW" else -1
|
|
|
|
# Creating an execution with RUNNING status
|
|
execution = AgentExecution(status='CREATED', last_execution_time=datetime.now(), agent_id=db_agent.id,
|
|
name="New Run", current_agent_step_id=start_step.id, iteration_workflow_step_id=iteration_step_id)
|
|
|
|
agent_execution_configs = {
|
|
"goal": agent_with_config.goal,
|
|
"instruction": agent_with_config.instruction,
|
|
"constraints": agent_with_config.constraints,
|
|
"toolkits": agent_with_config.toolkits,
|
|
"exit": agent_with_config.exit,
|
|
"tools": agent_with_config.tools,
|
|
"iteration_interval": agent_with_config.iteration_interval,
|
|
"model": agent_with_config.model,
|
|
"permission_type": agent_with_config.permission_type,
|
|
"LTM_DB": agent_with_config.LTM_DB,
|
|
"max_iterations": agent_with_config.max_iterations,
|
|
"user_timezone": agent_with_config.user_timezone,
|
|
"knowledge": agent_with_config.knowledge
|
|
}
|
|
db.session.add(execution)
|
|
db.session.commit()
|
|
db.session.flush()
|
|
AgentExecutionConfiguration.add_or_update_agent_execution_config(session=db.session, execution=execution,
|
|
agent_execution_configs=agent_execution_configs)
|
|
|
|
agent = db.session.query(Agent).filter(Agent.id == db_agent.id, ).first()
|
|
organisation = agent.get_agent_organisation(db.session)
|
|
|
|
EventHandler(session=db.session).create_event('run_created',
|
|
{'agent_execution_id': execution.id,
|
|
'agent_execution_name': execution.name},
|
|
db_agent.id,
|
|
organisation.id if organisation else 0),
|
|
|
|
if agent_with_config.knowledge:
|
|
knowledge_name = db.session.query(Knowledges.name).filter(Knowledges.id == agent_with_config.knowledge).first()[0]
|
|
EventHandler(session=db.session).create_event('knowledge_picked',
|
|
{'knowledge_name': knowledge_name,
|
|
'agent_execution_id': execution.id},
|
|
db_agent.id,
|
|
organisation.id if organisation else 0)
|
|
|
|
EventHandler(session=db.session).create_event('agent_created',
|
|
{'agent_name': agent_with_config.name,
|
|
'model': agent_with_config.model},
|
|
db_agent.id,
|
|
organisation.id if organisation else 0)
|
|
|
|
db.session.commit()
|
|
|
|
return {
|
|
"id": db_agent.id,
|
|
"execution_id": execution.id,
|
|
"name": db_agent.name,
|
|
"contentType": "Agents"
|
|
}
|
|
|
|
|
|
|
|
@router.post("/schedule", status_code=201)
|
|
def create_and_schedule_agent(agent_config_schedule: AgentConfigSchedule,
|
|
Authorize: AuthJWT = Depends(check_auth)):
|
|
"""
|
|
Create a new agent with configurations and scheduling.
|
|
|
|
Args:
|
|
agent_with_config_schedule (AgentConfigSchedule): Data for creating a new agent with configurations and scheduling.
|
|
|
|
Returns:
|
|
dict: Dictionary containing the created agent's ID, name, content type and schedule ID of the agent.
|
|
|
|
Raises:
|
|
HTTPException (status_code=500): If the associated agent fails to get scheduled.
|
|
"""
|
|
|
|
project = db.session.query(Project).get(agent_config_schedule.agent_config.project_id)
|
|
if not project:
|
|
raise HTTPException(status_code=404, detail="Project not found")
|
|
agent_config = agent_config_schedule.agent_config
|
|
invalid_tools = Tool.get_invalid_tools(agent_config.tools, db.session)
|
|
if len(invalid_tools) > 0: # If the returned value is not True (then it is an invalid tool_id)
|
|
raise HTTPException(status_code=404,
|
|
|
|
detail=f"Tool with IDs {str(invalid_tools)} does not exist. 404 Not Found.")
|
|
|
|
agent_toolkit_tools = Toolkit.fetch_tool_ids_from_toolkit(session=db.session,
|
|
toolkit_ids=agent_config.toolkits)
|
|
agent_config.tools.extend(agent_toolkit_tools)
|
|
db_agent = Agent.create_agent_with_config(db, agent_config)
|
|
|
|
# Update the agent_id of schedule before scheduling the agent
|
|
agent_schedule = agent_config_schedule.schedule
|
|
|
|
# Create a new agent schedule
|
|
agent_schedule = AgentSchedule(
|
|
agent_id=db_agent.id,
|
|
start_time=agent_schedule.start_time,
|
|
next_scheduled_time=agent_schedule.start_time,
|
|
recurrence_interval=agent_schedule.recurrence_interval,
|
|
expiry_date=agent_schedule.expiry_date,
|
|
expiry_runs=agent_schedule.expiry_runs,
|
|
current_runs=0,
|
|
status="SCHEDULED"
|
|
)
|
|
|
|
agent_schedule.agent_id = db_agent.id
|
|
db.session.add(agent_schedule)
|
|
db.session.commit()
|
|
|
|
if agent_schedule.id is None:
|
|
raise HTTPException(status_code=500, detail="Failed to schedule agent")
|
|
|
|
agent = db.session.query(Agent).filter(Agent.id == db_agent.id, ).first()
|
|
organisation = agent.get_agent_organisation(db.session)
|
|
|
|
EventHandler(session=db.session).create_event('agent_created', {'agent_name': agent_config.name,
|
|
'model': agent_config.model}, db_agent.id,
|
|
organisation.id if organisation else 0)
|
|
|
|
db.session.commit()
|
|
|
|
return {
|
|
"id": db_agent.id,
|
|
"name": db_agent.name,
|
|
"contentType": "Agents",
|
|
"schedule_id": agent_schedule.id
|
|
}
|
|
|
|
|
|
|
|
@router.post("/stop/schedule", status_code=200)
|
|
def stop_schedule(agent_id: int, Authorize: AuthJWT = Depends(check_auth)):
|
|
"""
|
|
Stopping the scheduling for a given agent.
|
|
|
|
Args:
|
|
agent_id (int): Identifier of the Agent
|
|
Authorize (AuthJWT, optional): Authorization dependency. Defaults to Depends(check_auth).
|
|
|
|
Raises:
|
|
HTTPException (status_code=404): If the agent schedule is not found.
|
|
"""
|
|
|
|
agent_to_delete = db.session.query(AgentSchedule).filter(AgentSchedule.agent_id == agent_id,
|
|
AgentSchedule.status == "SCHEDULED").first()
|
|
if not agent_to_delete:
|
|
raise HTTPException(status_code=404, detail="Schedule not found")
|
|
agent_to_delete.status = "STOPPED"
|
|
db.session.commit()
|
|
|
|
|
|
@router.put("/edit/schedule", status_code=200)
|
|
def edit_schedule(schedule: AgentScheduleInput,
|
|
Authorize: AuthJWT = Depends(check_auth)):
|
|
"""
|
|
Edit the scheduling for a given agent.
|
|
|
|
Args:
|
|
agent_id (int): Identifier of the Agent
|
|
schedule (AgentSchedule): New schedule data
|
|
Authorize (AuthJWT, optional): Authorization dependency. Defaults to Depends(check_auth).
|
|
|
|
Raises:
|
|
HTTPException (status_code=404): If the agent schedule is not found.
|
|
"""
|
|
|
|
agent_to_edit = db.session.query(AgentSchedule).filter(AgentSchedule.agent_id == schedule.agent_id, AgentSchedule.status == "SCHEDULED").first()
|
|
|
|
if not agent_to_edit:
|
|
raise HTTPException(status_code=404, detail="Schedule not found")
|
|
|
|
# Update agent schedule with new data
|
|
agent_to_edit.start_time = schedule.start_time
|
|
agent_to_edit.next_scheduled_time = schedule.start_time
|
|
agent_to_edit.recurrence_interval = schedule.recurrence_interval
|
|
agent_to_edit.expiry_date = schedule.expiry_date
|
|
agent_to_edit.expiry_runs = schedule.expiry_runs
|
|
|
|
db.session.commit()
|
|
|
|
|
|
@router.get("/get/schedule_data/{agent_id}")
|
|
def get_schedule_data(agent_id: int, Authorize: AuthJWT = Depends(check_auth)):
|
|
"""
|
|
Get the scheduling data for a given agent.
|
|
|
|
Args:
|
|
agent_id (int): Identifier of the Agent
|
|
|
|
Raises:
|
|
HTTPException (status_code=404): If the agent schedule is not found.
|
|
|
|
Returns:
|
|
current_datetime (DateTime): Current Date and Time.
|
|
recurrence_interval (String): Time interval for recurring schedule run.
|
|
expiry_date (DateTime): The date and time when the agent is scheduled to stop runs.
|
|
expiry_runs (Integer): The number of runs before the agent expires.
|
|
"""
|
|
agent = db.session.query(AgentSchedule).filter(AgentSchedule.agent_id == agent_id,
|
|
AgentSchedule.status == "SCHEDULED").first()
|
|
|
|
if not agent:
|
|
raise HTTPException(status_code=404, detail="Agent Schedule not found")
|
|
|
|
user_timezone = db.session.query(AgentConfiguration).filter(AgentConfiguration.key == "user_timezone",
|
|
AgentConfiguration.agent_id == agent_id).first()
|
|
|
|
if user_timezone and user_timezone.value != "None":
|
|
tzone = timezone(user_timezone.value)
|
|
else:
|
|
tzone = timezone('GMT')
|
|
|
|
current_datetime = datetime.now(tzone).strftime("%d/%m/%Y %I:%M %p")
|
|
|
|
return {
|
|
"current_datetime": current_datetime,
|
|
"start_date": agent.start_time.astimezone(tzone).strftime("%d %b %Y"),
|
|
"start_time": agent.start_time.astimezone(tzone).strftime("%I:%M %p"),
|
|
"recurrence_interval": agent.recurrence_interval if agent.recurrence_interval else None,
|
|
"expiry_date": agent.expiry_date.astimezone(tzone).strftime("%d/%m/%Y") if agent.expiry_date else None,
|
|
"expiry_runs": agent.expiry_runs if agent.expiry_runs != -1 else None
|
|
}
|
|
|
|
|
|
@router.get("/get/project/{project_id}")
|
|
def get_agents_by_project_id(project_id: int,
|
|
Authorize: AuthJWT = Depends(check_auth)):
|
|
"""
|
|
Get all agents by project ID.
|
|
|
|
Args:
|
|
project_id (int): Identifier of the project.
|
|
Authorize (AuthJWT, optional): Authorization dependency. Defaults to Depends(check_auth).
|
|
|
|
Returns:
|
|
list: List of agents associated with the project, including their status and scheduling information.
|
|
|
|
Raises:
|
|
HTTPException (status_code=404): If the project is not found.
|
|
"""
|
|
|
|
# Checking for project
|
|
project = db.session.query(Project).get(project_id)
|
|
if not project:
|
|
raise HTTPException(status_code=404, detail="Project not found")
|
|
|
|
agents = db.session.query(Agent).filter(Agent.project_id == project_id, or_(or_(Agent.is_deleted == False, Agent.is_deleted is None), Agent.is_deleted is None)).all()
|
|
|
|
new_agents, new_agents_sorted = [], []
|
|
for agent in agents:
|
|
agent_dict = vars(agent)
|
|
|
|
agent_id = agent.id
|
|
|
|
# Query the AgentExecution table using the agent ID
|
|
executions = db.session.query(AgentExecution).filter_by(agent_id=agent_id).all()
|
|
is_running = False
|
|
for execution in executions:
|
|
if execution.status == "RUNNING":
|
|
is_running = True
|
|
break
|
|
# Check if the agent is scheduled
|
|
is_scheduled = db.session.query(AgentSchedule).filter_by(agent_id=agent_id, status="SCHEDULED").first() is not None
|
|
|
|
|
|
new_agent = {
|
|
**agent_dict,
|
|
'is_running': is_running,
|
|
'is_scheduled': is_scheduled
|
|
}
|
|
new_agents.append(new_agent)
|
|
new_agents_sorted = sorted(new_agents, key=lambda agent: agent['is_running'] == True, reverse=True)
|
|
return new_agents_sorted
|
|
|
|
|
|
@router.put("/delete/{agent_id}", status_code=200)
|
|
def delete_agent(agent_id: int, Authorize: AuthJWT = Depends(check_auth)):
|
|
"""
|
|
Delete an existing Agent
|
|
- Updates the is_deleted flag: Executes a soft delete
|
|
- AgentExecutions are updated to: "TERMINATED" if agentexecution is created, All the agent executions are updated
|
|
- AgentExecutionPermission is set to: "REJECTED" if agentexecutionpersmision is created
|
|
|
|
Args:
|
|
agent_id (int): Identifier of the Agent to delete
|
|
|
|
Returns:
|
|
A dictionary containing a "success" key with the value True to indicate a successful delete.
|
|
|
|
Raises:
|
|
HTTPException (Status Code=404): If the Agent or associated Project is not found or deleted already.
|
|
"""
|
|
|
|
db_agent = db.session.query(Agent).filter(Agent.id == agent_id).first()
|
|
db_agent_executions = db.session.query(AgentExecution).filter(AgentExecution.agent_id == agent_id).all()
|
|
db_agent_schedule = db.session.query(AgentSchedule).filter(AgentSchedule.agent_id == agent_id, AgentSchedule.status == "SCHEDULED").first()
|
|
|
|
if not db_agent or db_agent.is_deleted:
|
|
raise HTTPException(status_code=404, detail="agent not found")
|
|
|
|
# Deletion Procedure
|
|
db_agent.is_deleted = True
|
|
if db_agent_executions:
|
|
# Updating all the RUNNING executions to TERMINATED
|
|
for db_agent_execution in db_agent_executions:
|
|
db_agent_execution.status = "TERMINATED"
|
|
|
|
if db_agent_schedule:
|
|
# Updating the schedule status to STOPPED
|
|
db_agent_schedule.status = "STOPPED"
|
|
|
|
db.session.commit()
|