Updated

BigPapaBear2217 · Jul 3, 2023 · b232c67 · b232c67
1 parent 6f614c3
commit b232c67
Show file tree

Hide file tree

Showing 11 changed files with 142 additions and 82 deletions.
diff --git a/docker-compose.yaml b/docker-compose.yaml
@@ -45,6 +45,8 @@ services:
       - superagi_postgres_data:/var/lib/postgresql/data/
     networks:
       - super_network
+    ports:
+      - "5432:5432"
 
   proxy:
     image: nginx:stable-alpine

diff --git a/main.py b/main.py
@@ -269,7 +269,7 @@ def check_toolkit_registration():
 
     build_single_step_agent()
     build_task_based_agents()
-    check_toolkit_registration()
+    # check_toolkit_registration()
     session.close()
 
 

diff --git a/migrations/versions/864d7872b598_added_agent_execution_config.py b/migrations/versions/864d7872b598_added_agent_execution_config.py
@@ -0,0 +1,46 @@
+"""added_agent_execution_config
+
+Revision ID: 864d7872b598
+Revises: c5c19944c90c
+Create Date: 2023-07-03 11:50:04.038668
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = '864d7872b598'
+down_revision = 'c5c19944c90c'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('agent_execution_configs',
+    sa.Column('created_at', sa.DateTime(), nullable=True),
+    sa.Column('updated_at', sa.DateTime(), nullable=True),
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('agent_execution_id', sa.Integer(), nullable=True),
+    sa.Column('key', sa.String(), nullable=True),
+    sa.Column('value', sa.Text(), nullable=True),
+    sa.PrimaryKeyConstraint('id')
+    )
+    op.drop_table('tool_statistics')
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('tool_statistics',
+    sa.Column('created_at', postgresql.TIMESTAMP(), autoincrement=False, nullable=True),
+    sa.Column('updated_at', postgresql.TIMESTAMP(), autoincrement=False, nullable=True),
+    sa.Column('id', sa.INTEGER(), autoincrement=True, nullable=False),
+    sa.Column('toolkit_id', sa.INTEGER(), autoincrement=False, nullable=True),
+    sa.Column('key', sa.VARCHAR(), autoincrement=False, nullable=True),
+    sa.Column('value', sa.VARCHAR(), autoincrement=False, nullable=True),
+    sa.PrimaryKeyConstraint('id', name='tool_statistics_pkey')
+    )
+    op.drop_table('agent_execution_configs')
+    # ### end Alembic commands ###
diff --git a/superagi/controllers/agent.py b/superagi/controllers/agent.py
@@ -1,28 +1,26 @@
-from fastapi_sqlalchemy import db
-from fastapi import HTTPException, Depends, Request
-from fastapi_jwt_auth import AuthJWT
+import json
+from datetime import datetime
 
-from superagi.controllers.types.agent_execution_config_request import AgentExecutionConfigRequest
-from superagi.models.agent import Agent
-from superagi.models.agent_execution_config import AgentExecutionConfiguration
-from superagi.models.agent_template import AgentTemplate
-from superagi.models.agent_template_config import AgentTemplateConfig
-from superagi.models.project import Project
 from fastapi import APIRouter
+from fastapi import HTTPException, Depends
+from fastapi_jwt_auth import AuthJWT
+from fastapi_sqlalchemy import db
+from jsonmerge import merge
 from pydantic_sqlalchemy import sqlalchemy_to_pydantic
+from sqlalchemy import func
 
-from superagi.models.agent_workflow import AgentWorkflow
-from superagi.models.types.agent_with_config import AgentWithConfig
+# from superagi.controllers.types.agent_create_request import AgentCreateRequest
+from superagi.helper.auth import check_auth
+from superagi.models.agent import Agent
 from superagi.models.agent_config import AgentConfiguration
 from superagi.models.agent_execution import AgentExecution
-from superagi.models.agent_execution_feed import AgentExecutionFeed
+from superagi.models.agent_execution_config import AgentExecutionConfiguration
+from superagi.models.agent_template import AgentTemplate
+from superagi.models.agent_workflow import AgentWorkflow
+from superagi.models.project import Project
 from superagi.models.tool import Tool
-from jsonmerge import merge
+from superagi.models.types.agent_with_config import AgentWithConfig
 from superagi.worker import execute_agent
-from datetime import datetime
-import json
-from sqlalchemy import func
-from superagi.helper.auth import check_auth, get_user_organisation
 
 router = APIRouter()
 
@@ -122,7 +120,6 @@ def update_agent(agent_id: int, agent: sqlalchemy_to_pydantic(Agent, exclude=["i
 
 @router.post("/create", status_code=201)
 def create_agent_with_config(agent_with_config: AgentWithConfig,
-                             agent_execution_config_request: AgentExecutionConfigRequest,
                              Authorize: AuthJWT = Depends(check_auth)):
     """
     Create a new agent with configurations.
@@ -170,11 +167,16 @@ def create_agent_with_config(agent_with_config: AgentWithConfig,
     # Creating an execution with RUNNING status
     execution = AgentExecution(status='RUNNING', last_execution_time=datetime.now(), agent_id=db_agent.id,
                                name="New Run", current_step_id=start_step_id)
-    AgentExecutionConfiguration.add_or_update_agent_execution_config(session=db.session, execution=execution,
-                                                                     agent_execution_config_request=agent_execution_config_request)
 
+    agent_execution_configs = {
+        "goal": agent_with_config.goal,
+        "instruction": agent_with_config.instruction
+    }
     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)
     execute_agent.delay(execution.id, datetime.now())
 
     return {

diff --git a/superagi/controllers/agent_execution.py b/superagi/controllers/agent_execution.py
@@ -3,6 +3,7 @@
 from fastapi import HTTPException, Depends
 from fastapi_jwt_auth import AuthJWT
 
+from superagi.controllers.types.agent_create_request import AgentExecutionConfigRequest
 from superagi.helper.time_helper import get_time_difference
 from superagi.models.agent_workflow import AgentWorkflow
 from superagi.worker import execute_agent
@@ -17,8 +18,8 @@
 
 
 # CRUD Operations
-@router.post("/add", response_model=sqlalchemy_to_pydantic(AgentExecution), status_code=201)
-def create_agent_execution(agent_execution: sqlalchemy_to_pydantic(AgentExecution, exclude=["id"]),
+@router.post("/add", response_model=AgentExecutionConfigRequest, status_code=201)
+def create_agent_execution(agent_execution: AgentExecutionConfigRequest,
                            Authorize: AuthJWT = Depends(check_auth)):
     """
     Create a new agent execution/run.

diff --git a/superagi/controllers/agent_execution_config.py b/superagi/controllers/agent_execution_config.py
@@ -1,49 +1,45 @@
+
 from fastapi import APIRouter
 from fastapi import HTTPException, Depends
 from fastapi_jwt_auth import AuthJWT
 from fastapi_sqlalchemy import db
+from pydantic.fields import List
 from pydantic_sqlalchemy import sqlalchemy_to_pydantic
 
 from superagi.helper.auth import check_auth
 from superagi.models.agent_execution_config import AgentExecutionConfiguration
-from superagi.controllers.types.agent_execution_config_request import AgentExecutionConfigRequest
-router = APIRouter()
-
-
-@router.put("/update", response_model=sqlalchemy_to_pydantic(AgentExecutionConfiguration))
-def update_agent(agent_execution_config: AgentExecutionConfigRequest,
-                 Authorize: AuthJWT = Depends(check_auth)):
-    """
-        Update a particular agent execution configuration value for the given agent_id and agent_config key.
-
-        Args:
-            agent_execution_config (AgentExecutionConfiguration): The updated agent configuration data.
+from superagi.controllers.types.agent_create_request import AgentExecutionConfigRequest
 
-        Returns:
-            AgentConfiguration: The updated agent execution configuration.
 
-        Raises:
-            HTTPException (Status Code=404): If the agent configuration is not found.
-    """
+router = APIRouter()
 
-    db_agent_execution_config = db.session.query(AgentExecutionConfiguration).filter(
-        AgentExecutionConfiguration.key == agent_execution_config.key,
-        AgentExecutionConfiguration.agent_id == agent_execution_config.agent_id).first()
-    if not db_agent_execution_config:
-        raise HTTPException(status_code=404, detail="Agent Configuration not found")
 
-    db_agent_execution_config.key = agent_execution_config.key
-    if isinstance(agent_execution_config.value, list):
-        db_agent_execution_config.value = str(agent_execution_config.value)
-    else:
-        db_agent_execution_config.value = agent_execution_config.value
-    db.session.commit()
-    db.session.flush()
-    return db_agent_execution_config
+# @router.put("/update", response_model=AgentExecutionConfigRequest)
+# def update_agent_execution_config(agent_execution_configs: AgentExecutionConfigRequest,
+#                  Authorize: AuthJWT = Depends(check_auth)):
+#     # db_agent_execution_config = db.session.query(AgentExecutionConfiguration).filter(
+#     #     AgentExecutionConfiguration.key == agent_execution_config.key,
+#     #     AgentExecutionConfiguration.agent_id == agent_execution_config.agent_id).first()
+#     # if not db_agent_execution_config:
+#     #     raise HTTPException(status_code=404, detail="Agent Configuration not found")
+#     #
+#     # db_agent_execution_config.key = agent_execution_config.key
+#     # if isinstance(agent_execution_config.value, list):
+#     #     db_agent_execution_config.value = str(agent_execution_config.value)
+#     # else:
+#     #     db_agent_execution_config.value = agent_execution_config.value
+#     # db.session.commit()
+#     # db.session.flush()
+#     # return db_agent_execution_config
+#     execution = db.session.query(AgentExecutionConfiguration).filter(AgentExecutionConfiguration.id == agent_execution_configs.execution_id).first()
+#     if execution is None:
+#         raise HTTPException(status_code=404,detail='Execution not found')
+#     AgentExecutionConfiguration.add_or_update_agent_execution_config(session=db.session, execution=execution,
+#                                                                      agent_execution_configs=agent_execution_configs)
 
 
 @router.get("/details/{agent_execution_id}")
-def get_agent_configuration(agent_execution_id: int,
+def get_agent_execution_configuration(agent_execution_id: int,
                             Authorize: AuthJWT = Depends(check_auth)):
     """
     Get the agent execution configuration using the agent execution ID.

diff --git a/superagi/controllers/types/agent_create_request.py b/superagi/controllers/types/agent_create_request.py
@@ -0,0 +1,11 @@
+from typing import List, Optional
+
+from pydantic import BaseModel
+from pydantic_sqlalchemy import sqlalchemy_to_pydantic
+from superagi.models.agent_execution import AgentExecution
+
+
+class AgentExecutionConfigRequest(BaseModel):
+    AgentExecution: Optional[sqlalchemy_to_pydantic(AgentExecution, exclude=["id"])]
+    goal: List[str]
+    instruction: List[str]
diff --git a/superagi/controllers/types/agent_execution_config_request.py b/superagi/controllers/types/agent_execution_config_request.py
diff --git a/superagi/jobs/agent_executor.py b/superagi/jobs/agent_executor.py
@@ -29,6 +29,7 @@
 from superagi.vector_store.embedding.openai import OpenAiEmbedding
 from superagi.vector_store.vector_factory import VectorFactory
 import yaml
+
 # from superagi.helper.tool_helper import get_tool_config_by_key
 
 engine = connect_db()
@@ -49,6 +50,7 @@ def get_tool_config(self, key: str):
             return tool_config.value
         return super().get_tool_config(key=key)
 
+
 class AgentExecutor:
     @staticmethod
     def validate_filename(filename):
@@ -66,7 +68,7 @@ def validate_filename(filename):
         return filename
 
     @staticmethod
-    def create_object(tool,session):
+    def create_object(tool, session):
         """
         Create an object of a agent usable tool dynamically.
 
@@ -158,7 +160,7 @@ def execute_next_action(self, agent_execution_id):
         ]
 
         parsed_config = Agent.fetch_configuration(session, agent.id)
-        parsed_execution_config = AgentExecutionConfiguration.fetch_configuration(session, agent_execution.id)
+        parsed_execution_config = AgentExecutionConfiguration.fetch_configuration(session, agent_execution)
         max_iterations = (parsed_config["max_iterations"])
         total_calls = agent_execution.num_of_calls
 
@@ -172,7 +174,7 @@ def execute_next_action(self, agent_execution_id):
         parsed_config["agent_execution_id"] = agent_execution.id
 
         model_api_key = AgentExecutor.get_model_api_key_from_execution(agent_execution, session)
-
+        memory = None
         try:
             if parsed_config["LTM_DB"] == "Pinecone":
                 memory = VectorFactory.get_vector_storage("PineCone", "super-agent-index1",
@@ -186,10 +188,10 @@ def execute_next_action(self, agent_execution_id):
 
         user_tools = session.query(Tool).filter(Tool.id.in_(parsed_config["tools"])).all()
         for tool in user_tools:
-            tool = AgentExecutor.create_object(tool,session)
+            tool = AgentExecutor.create_object(tool, session)
             tools.append(tool)
 
-        tools = self.set_default_params_tools(tools, parsed_config, agent_execution.agent_id,
+        tools = self.set_default_params_tools(tools, parsed_config, parsed_execution_config, agent_execution.agent_id,
                                               model_api_key=model_api_key, session=session)
 
         spawned_agent = SuperAgi(ai_name=parsed_config["name"], ai_role=parsed_config["description"],
@@ -226,7 +228,7 @@ def execute_next_action(self, agent_execution_id):
         session.close()
         engine.dispose()
 
-    def set_default_params_tools(self, tools, parsed_config, parsed_execution_config,agent_id, model_api_key, session):
+    def set_default_params_tools(self, tools, parsed_config, parsed_execution_config, agent_id, model_api_key, session):
         """
         Set the default parameters for the tools.
 

diff --git a/superagi/models/agent_execution_config.py b/superagi/models/agent_execution_config.py
@@ -29,40 +29,37 @@ def __repr__(self):
             str: String representation of the AgentTemplateConfig.
         """
 
-        return f"AgentExecutionConfig(id={self.id}, agent_execution_id='{self.agent_template_id}', " \
+        return f"AgentExecutionConfig(id={self.id}, agent_execution_id='{self.agent_execution_id}', " \
                f"key='{self.key}', value='{self.value}')"
 
     @classmethod
-    def add_or_update_agent_execution_config(cls, session, execution, agent_execution_config_request):
-        agent_config_values = {
-            "goal": agent_execution_config_request.goal,
-            "instruction": agent_execution_config_request.instruction
-        }
-
+    def add_or_update_agent_execution_config(cls, session, execution, agent_execution_configs):
+        print("EXECUTION : ",execution)
         agent_execution_configurations = [
             AgentExecutionConfiguration(agent_execution_id=execution.id, key=key, value=str(value))
-            for key, value in agent_config_values.items()
+            for key, value in agent_execution_configs.items()
         ]
-
-        for key, value in agent_execution_configurations.items():
+        for agent_execution in agent_execution_configurations:
             agent_execution_config = (
                 session.query(AgentExecutionConfiguration)
                 .filter(
                     AgentExecutionConfiguration.agent_execution_id == execution.id,
-                    AgentExecutionConfiguration.key == key
+                    AgentExecutionConfiguration.key == agent_execution.key
                 )
                 .first()
             )
 
             if agent_execution_config:
-                agent_execution_config.value = str(value)
+                agent_execution_config.value = str(agent_execution.value)
             else:
                 agent_execution_config = AgentExecutionConfiguration(
                     agent_execution_id=execution.id,
-                    key=key,
-                    value=str(value)
+                    key=agent_execution.key,
+                    value=str(agent_execution.value)
                 )
+                print("CONFIG : ",agent_execution_config)
                 session.add(agent_execution_config)
+            session.commit()
 
     @classmethod
     def fetch_configuration(cls, session, execution):
@@ -77,13 +74,15 @@ def fetch_configuration(cls, session, execution):
             dict: Parsed agent configuration.
 
         """
-
+        print("fetching-------------------")
+        print(execution)
         agent_configurations = session.query(AgentExecutionConfiguration).filter_by(
             agent_execution_id=execution.id).all()
         parsed_config = {
             "goal": [],
             "instruction": [],
         }
+        print("PARSED CONFIG : ",parsed_config)
         if not agent_configurations:
             return parsed_config
         for item in agent_configurations:

diff --git a/superagi/models/types/agent_execution_config.py b/superagi/models/types/agent_execution_config.py
@@ -0,0 +1,8 @@
+# from typing import List
+#
+# from pydantic import BaseModel
+#
+#
+# class AgentExecutionConfig(BaseModel):
+#     goal: List[str]
+#     instruction: List[str]