Skip to content

Commit

Permalink
Merge pull request #998 from rounak610/edit_agent_feature_1
Browse files Browse the repository at this point in the history 
Edit agent feature 1
  • Loading branch information
I’m committed Aug 7, 2023
2 parents 3266d8c + 989ce4c commit 6c9afdd
Show file tree
Hide file tree
Showing 10 changed files with 371 additions and 121 deletions.
96 changes: 31 additions & 65 deletions superagi/controllers/agent.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
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_
Expand Down Expand Up @@ -187,6 +189,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 @@ -204,21 +207,37 @@ def create_agent_with_config(agent_with_config: AgentConfigInput,

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 @@ -233,6 +252,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 @@ -256,6 +276,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 @@ -302,6 +323,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 @@ -338,8 +360,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 @@ -418,7 +440,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 @@ -434,8 +456,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 @@ -447,62 +469,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 @@ -539,4 +505,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()
88 changes: 83 additions & 5 deletions superagi/controllers/agent_execution.py
View file
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 @@ -19,8 +20,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 @@ -41,7 +43,6 @@ class AgentExecutionOut(BaseModel):
class Config:
orm_mode = True


class AgentExecutionIn(BaseModel):
status: Optional[str]
name: Optional[str]
Expand All @@ -54,9 +55,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 @@ -88,10 +90,25 @@ def create_agent_execution(agent_execution: AgentExecutionIn,
num_of_tokens=0,
current_agent_step_id=start_step.id,
iteration_workflow_step_id=iteration_step_id)

agent_execution_configs = {
"goal": agent_execution.goal,
"instruction": agent_execution.instruction
}

agent_configs = db.session.query(AgentConfiguration).filter(AgentConfiguration.agent_id == agent_execution.agent_id).all()
keys_to_exclude = ["goal", "instruction"]
for agent_config in agent_configs:
if agent_config.key not in keys_to_exclude:
if agent_config.key == "toolkits":
toolkits = [int(item) for item in agent_config.value.strip('{}').split(',')]
agent_execution_configs[agent_config.key] = toolkits
elif agent_config.key == "constraints":
constraints = [item.strip('"') for item in agent_config.value.strip('{}').split(',')]
agent_execution_configs[agent_config.key] = constraints
else:
agent_execution_configs[agent_config.key] = agent_config.value

db.session.add(db_agent_execution)
db.session.commit()
db.session.flush()
Expand All @@ -107,6 +124,67 @@ 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.
"""
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")

#Update the agent configurations table with the data of the latest agent execution
AgentConfiguration.update_agent_configurations_table(session=db.session, agent_id=agent_execution.agent_id, updated_details=agent_execution)

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 @@ -282,4 +360,4 @@ def get_agent_by_latest_execution(project_id: int,
"id": agent.id,
"status": isRunning,
"contentType": "Agents"
}
}
Loading

0 comments on commit 6c9afdd

Please sign in to comment.