Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Edit agent feature #963

Merged
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
99 changes: 33 additions & 66 deletions superagi/controllers/agent.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
import json
from datetime import datetime

from typing import Union, List
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 jsonmerge import merge
from pytz import timezone
Expand Down Expand Up @@ -193,6 +195,7 @@ def create_agent_with_config(agent_with_config: AgentConfigInput,
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,
Expand All @@ -201,27 +204,44 @@ def create_agent_with_config(agent_with_config: AgentConfigInput,
db_agent = Agent.create_agent_with_config(db, agent_with_config)

start_step_id = AgentWorkflow.fetch_trigger_step_id(db.session, db_agent.agent_workflow_id)

# Creating an execution with RUNNING status
execution = AgentExecution(status='CREATED', last_execution_time=datetime.now(), agent_id=db_agent.id,
name="New Run", current_step_id=start_step_id)

agent_execution_configs = {
"goal": agent_with_config.goal,
"instruction": agent_with_config.instruction
"instruction": agent_with_config.instruction,
"agent_type": agent_with_config.agent_type,
"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()
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,

'agent_execution_name': execution.name}, 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)

# execute_agent.delay(execution.id, datetime.now())
Expand All @@ -236,6 +256,7 @@ def create_agent_with_config(agent_with_config: AgentConfigInput,
}



@router.post("/schedule", status_code=201)
def create_and_schedule_agent(agent_config_schedule: AgentConfigSchedule,
Authorize: AuthJWT = Depends(check_auth)):
Expand All @@ -259,6 +280,7 @@ def create_and_schedule_agent(agent_config_schedule: AgentConfigSchedule,
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,
Expand Down Expand Up @@ -305,6 +327,7 @@ def create_and_schedule_agent(agent_config_schedule: AgentConfigSchedule,
}



@router.post("/stop/schedule", status_code=200)
def stop_schedule(agent_id: int, Authorize: AuthJWT = Depends(check_auth)):
"""
Expand Down Expand Up @@ -341,8 +364,8 @@ def edit_schedule(schedule: AgentScheduleInput,
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()
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")

Expand Down Expand Up @@ -421,7 +444,7 @@ def get_agents_by_project_id(project_id: int,
if not project:
raise HTTPException(status_code=404, detail="Project not found")

agents = db.session.query(Agent).filter(Agent.project_id == project_id, or_(Agent.is_deleted == False, Agent.is_deleted is None)).all()
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:
Expand All @@ -437,8 +460,8 @@ def get_agents_by_project_id(project_id: int,
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
is_scheduled = db.session.query(AgentSchedule).filter_by(agent_id=agent_id, status="SCHEDULED").first() is not None


new_agent = {
**agent_dict,
Expand All @@ -450,62 +473,6 @@ def get_agents_by_project_id(project_id: int,
return new_agents_sorted


@router.get("/get/details/{agent_id}")
def get_agent_configuration(agent_id: int,
Authorize: AuthJWT = Depends(check_auth)):
"""
Get the agent configuration using the agent ID.

Args:
agent_id (int): Identifier of the agent.
Authorize (AuthJWT, optional): Authorization dependency. Defaults to Depends(check_auth).

Returns:
dict: Agent configuration including its details.

Raises:
HTTPException (status_code=404): If the agent is not found or deleted.
"""

# Define the agent_config keys to fetch
keys_to_fetch = AgentTemplate.main_keys()
agent = db.session.query(Agent).filter(agent_id == Agent.id,or_(Agent.is_deleted == False, Agent.is_deleted is None)).first()

if not agent:
raise HTTPException(status_code=404, detail="Agent not found")

# Query the AgentConfiguration table for the specified keys
results = db.session.query(AgentConfiguration).filter(AgentConfiguration.key.in_(keys_to_fetch),
AgentConfiguration.agent_id == agent_id).all()
total_calls = db.session.query(func.sum(AgentExecution.num_of_calls)).filter(
AgentExecution.agent_id == agent_id).scalar()
total_tokens = db.session.query(func.sum(AgentExecution.num_of_tokens)).filter(
AgentExecution.agent_id == agent_id).scalar()

name = ""
# Construct the JSON response
response = {result.key: result.value for result in results}
if 'knowledge' in response.keys() and response['knowledge'] != 'None':
knowledge = db.session.query(Knowledges).filter(Knowledges.id == response['knowledge']).first()
name = knowledge.name if knowledge is not None else ""
response = merge(response, {"name": agent.name, "description": agent.description,
# Query the AgentConfiguration table for the speci
"goal": eval(response["goal"]),
"instruction": eval(response.get("instruction", '[]')),
"knowledge_name": name,
"calls": total_calls,
"tokens": total_tokens,
"constraints": eval(response.get("constraints")),
"tools": [int(x) for x in json.loads(response["tools"])]})
tools = db.session.query(Tool).filter(Tool.id.in_(response["tools"])).all()
response["tools"] = tools

# Close the session
db.session.close()

return response


@router.put("/delete/{agent_id}", status_code=200)
def delete_agent(agent_id: int, Authorize: AuthJWT = Depends(check_auth)):
"""
Expand Down Expand Up @@ -542,4 +509,4 @@ def delete_agent(agent_id: int, Authorize: AuthJWT = Depends(check_auth)):
# Updating the schedule status to STOPPED
db_agent_schedule.status = "STOPPED"

db.session.commit()
db.session.commit()
74 changes: 69 additions & 5 deletions superagi/controllers/agent_execution.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
from datetime import datetime
from typing import Optional
from typing import Optional, Union, List

from fastapi_sqlalchemy import db
from fastapi import HTTPException, Depends
from fastapi_jwt_auth import AuthJWT
from pydantic import BaseModel
from pydantic.fields import List
from superagi.controllers.types.agent_execution_config import AgentRunIn

from superagi.helper.time_helper import get_time_difference
from superagi.models.agent_execution_config import AgentExecutionConfiguration
Expand All @@ -18,8 +19,9 @@
from sqlalchemy import desc
from superagi.helper.auth import check_auth
from superagi.controllers.types.agent_schedule import AgentScheduleInput
# from superagi.types.db import AgentExecutionOut, AgentExecutionIn
from superagi.apm.event_handler import EventHandler
from superagi.controllers.tool import ToolOut
from superagi.models.agent_config import AgentConfiguration

router = APIRouter()

Expand All @@ -40,7 +42,6 @@ class AgentExecutionOut(BaseModel):
class Config:
orm_mode = True


class AgentExecutionIn(BaseModel):
status: Optional[str]
name: Optional[str]
Expand All @@ -53,9 +54,10 @@ class AgentExecutionIn(BaseModel):
goal: Optional[List[str]]
instruction: Optional[List[str]]

class Config:
class config:
orm_mode = True


# CRUD Operations
@router.post("/add", response_model=AgentExecutionOut, status_code=201)
def create_agent_execution(agent_execution: AgentExecutionIn,
Expand Down Expand Up @@ -102,6 +104,68 @@ def create_agent_execution(agent_execution: AgentExecutionIn,

return db_agent_execution

@router.post("/add_run", status_code = 201)
def create_agent_run(agent_execution: AgentRunIn, Authorize: AuthJWT = Depends(check_auth)):

"""
Create a new agent run with all the information(goals, instructions, model, etc).

Args:
agent_execution (AgentExecution): The agent execution data.

Returns:
AgentExecution: The created agent execution.

Raises:
HTTPException (Status Code=404): If the agent is not found.
"""
#Update the agent configurations table with the data of the latest agent execution and returns -1 if agent id not found.
if((AgentConfiguration.update_agent_configurations_table(session=db.session, agent_id=agent_execution.agent_id, updated_details=agent_execution))==-1):
raise HTTPException(status_code=404, detail="Agent ID not found")

agent = db.session.query(Agent).filter(Agent.id == agent_execution.agent_id, Agent.is_deleted == False).first()
if not agent:
raise HTTPException(status_code = 404, detail = "Agent not found")

start_step_id = AgentWorkflow.fetch_trigger_step_id(db.session, agent.agent_workflow_id)

db_agent_execution = AgentExecution(status = "RUNNING", last_execution_time = datetime.now(),
agent_id = agent_execution.agent_id, name = agent_execution.name, num_of_calls = 0,
num_of_tokens = 0,
current_step_id = start_step_id)
agent_execution_configs = {
"goal": agent_execution.goal,
"instruction": agent_execution.instruction,
"agent_type": agent_execution.agent_type,
"constraints": agent_execution.constraints,
"toolkits": agent_execution.toolkits,
"exit": agent_execution.exit,
"tools": agent_execution.tools,
"iteration_interval": agent_execution.iteration_interval,
"model": agent_execution.model,
"permission_type": agent_execution.permission_type,
"LTM_DB": agent_execution.LTM_DB,
"max_iterations": agent_execution.max_iterations,
"user_timezone": agent_execution.user_timezone,
"knowledge": agent_execution.knowledge
}

db.session.add(db_agent_execution)
db.session.commit()
db.session.flush()

AgentExecutionConfiguration.add_or_update_agent_execution_config(session = db.session, execution = db_agent_execution,
agent_execution_configs = agent_execution_configs)

organisation = agent.get_agent_organisation(db.session)
EventHandler(session=db.session).create_event('run_created', {'agent_execution_id': db_agent_execution.id,'agent_execution_name':db_agent_execution.name},
agent_execution.agent_id, organisation.id if organisation else 0)

if db_agent_execution.status == "RUNNING":
execute_agent.delay(db_agent_execution.id, datetime.now())

return db_agent_execution


@router.post("/schedule", status_code=201)
def schedule_existing_agent(agent_schedule: AgentScheduleInput,
Expand Down Expand Up @@ -277,4 +341,4 @@ def get_agent_by_latest_execution(project_id: int,
"id": agent.id,
"status": isRunning,
"contentType": "Agents"
}
}
Loading