Edit agent feature

superagi/controllers/agent.py

@@ -1,15 +1,17 @@
 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
-from sqlalchemy import func
+from sqlalchemy import func, or_
 from superagi.models.agent_execution_permission import AgentExecutionPermission
 from superagi.worker import execute_agent
 from superagi.helper.auth import check_auth
@@ -109,11 +111,9 @@ def get_agent(agent_id: int,
             HTTPException (Status Code=404): If the Agent is not found or deleted.
     """
 
-    if (
-        db_agent := db.session.query(Agent)
-        .filter(Agent.id == agent_id, Agent.is_deleted == False)
-        .first()
-    ):
+    if (db_agent := db.session.query(Agent)
+            .filter(Agent.id == agent_id, or_(Agent.is_deleted == False, Agent.is_deleted is None))
+            .first()):
         return db_agent
     else:
         raise HTTPException(status_code=404, detail="agent not found")
@@ -141,7 +141,7 @@ def update_agent(agent_id: int, agent: AgentIn,
             HTTPException (Status Code=404): If the Agent or associated Project is not found.
     """
 
-    db_agent = db.session.query(Agent).filter(Agent.id == agent_id, Agent.is_deleted == False).first()
+    db_agent = db.session.query(Agent).filter(Agent.id == agent_id, or_(Agent.is_deleted == False, Agent.is_deleted is None)).first()
     if not db_agent:
         raise HTTPException(status_code=404, detail="agent not found")
 
@@ -190,35 +190,52 @@ def create_agent_with_config(agent_with_config: AgentConfigInput,
     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.")
+        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_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
     }
     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, 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)
+    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),
+    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())
 
@@ -231,6 +248,7 @@ def create_agent_with_config(agent_with_config: AgentConfigInput,
         "contentType": "Agents"
     }
 
+
 @router.post("/schedule", status_code=201)
 def create_and_schedule_agent(agent_config_schedule: AgentConfigSchedule,
                               Authorize: AuthJWT = Depends(check_auth)):
@@ -253,7 +271,8 @@ def create_and_schedule_agent(agent_config_schedule: AgentConfigSchedule,
     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.")
+        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)
@@ -289,6 +308,7 @@ def create_and_schedule_agent(agent_config_schedule: AgentConfigSchedule,
         "schedule_id": agent_schedule.id
     }
 
+
 @router.post("/stop/schedule", status_code=200)
 def stop_schedule(agent_id: int, Authorize: AuthJWT = Depends(check_auth)):
     """
@@ -303,7 +323,7 @@ def stop_schedule(agent_id: int, Authorize: AuthJWT = Depends(check_auth)):
     """
 
     agent_to_delete = db.session.query(AgentSchedule).filter(AgentSchedule.agent_id == agent_id,
-                                                              AgentSchedule.status == "SCHEDULED").first()
+                                                             AgentSchedule.status == "SCHEDULED").first()
     if not agent_to_delete:
         raise HTTPException(status_code=404, detail="Schedule not found")
     agent_to_delete.status = "STOPPED"
@@ -326,7 +346,7 @@ def edit_schedule(schedule: AgentScheduleInput,
     """
 
     agent_to_edit = db.session.query(AgentSchedule).filter(AgentSchedule.agent_id == schedule.agent_id,
-                                                            AgentSchedule.status == "SCHEDULED").first()
+                                                           AgentSchedule.status == "SCHEDULED").first()
     if not agent_to_edit:
         raise HTTPException(status_code=404, detail="Schedule not found")
 
@@ -358,7 +378,7 @@ def get_schedule_data(agent_id: int, Authorize: AuthJWT = Depends(check_auth)):
         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()
+                                                   AgentSchedule.status == "SCHEDULED").first()
 
     if not agent:
         raise HTTPException(status_code=404, detail="Agent Schedule not found")
@@ -371,7 +391,6 @@ def get_schedule_data(agent_id: int, Authorize: AuthJWT = Depends(check_auth)):
     else:
         tzone = timezone('GMT')
 
-
     current_datetime = datetime.now(tzone).strftime("%d/%m/%Y %I:%M %p")
 
     return {
@@ -406,7 +425,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, Agent.is_deleted == False).all()
+    agents = db.session.query(Agent).filter(Agent.project_id == project_id, or_(Agent.is_deleted == False, Agent.is_deleted is None)).all()
 
     new_agents, new_agents_sorted = [], []
     for agent in agents:
@@ -423,7 +442,7 @@ def get_agents_by_project_id(project_id: int,
                 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
+                                                                 status="SCHEDULED").first() is not None
 
         new_agent = {
             **agent_dict,
@@ -434,58 +453,7 @@ def get_agents_by_project_id(project_id: int,
         new_agents_sorted = sorted(new_agents, key=lambda agent: agent['is_running'] == True, reverse=True)
     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, Agent.is_deleted == False).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()
-
-    # Construct the JSON response
-    response = {result.key: result.value for result in results}
-    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", '[]')),
-                                "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)
+@router.put("/delete/{agent_id}", status_code=200)
 def delete_agent(agent_id: int, Authorize: AuthJWT = Depends(check_auth)):
     """
         Delete an existing Agent
@@ -502,18 +470,18 @@ def delete_agent(agent_id: int, Authorize: AuthJWT = Depends(check_auth)):
         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_executions = db.session.query(AgentExecution).filter(AgentExecution.agent_id == agent_id).all()
+
     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"
-    
+
     db.session.commit()

superagi/controllers/agent_execution.py

@@ -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
@@ -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()
 
@@ -40,7 +42,6 @@ class AgentExecutionOut(BaseModel):
     class Config:
         orm_mode = True
 
-
 class AgentExecutionIn(BaseModel):
     status: Optional[str]
     name: Optional[str]
@@ -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,
@@ -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
+    }
+    #agent_execution_configs = agent_execution.dict()
+    #print(agent_execution_configs)
+    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,