adding init db alemic and fixing acitvity feed

dataaccess-recovery · May 28, 2023 · 2d5dc55 · 2d5dc55
1 parent 79c1583
commit 2d5dc55
Show file tree

Hide file tree

Showing 12 changed files with 211 additions and 49 deletions.
diff --git a/.gitignore b/.gitignore
@@ -15,4 +15,6 @@ venv
 gui/node_modules
 gui/.next
 .DS_Store
-.DS_Store?
+.DS_Store?
+workspace/output
+workspace/input
diff --git a/alembic/versions/115a710c4685_init_db.py b/alembic/versions/115a710c4685_init_db.py
@@ -0,0 +1,130 @@
+"""init_db
+
+Revision ID: 115a710c4685
+Revises: 
+Create Date: 2023-05-28 14:35:55.544155
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '115a710c4685'
+down_revision = None
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        'organisations',
+        sa.Column('id', sa.Integer()),
+        sa.Column('name', sa.String(length=256), nullable=False),
+        sa.Column('description', sa.VARCHAR(1000)),
+        sa.Column('created_at', sa.DateTime()),
+        sa.Column('updated_at', sa.DateTime()),
+        sa.PrimaryKeyConstraint('id')
+    )
+
+    op.create_table(
+        'projects',
+        sa.Column('id', sa.Integer()),
+        sa.Column('name', sa.String(length=256), nullable=False),
+        sa.Column('description', sa.VARCHAR(1000)),
+        sa.Column('organisation_id', sa.Integer(), nullable=False),
+        sa.Column('created_at', sa.DateTime()),
+        sa.Column('updated_at', sa.DateTime()),
+        sa.PrimaryKeyConstraint('id')
+    )
+
+    op.create_table(
+        'users',
+        sa.Column('id', sa.Integer()),
+        sa.Column('name', sa.VARCHAR(250)),
+        sa.Column('email', sa.VARCHAR(250), nullable=False),
+        sa.Column('password', sa.VARCHAR(250), nullable=False),
+
+        sa.Column('organisation_id', sa.Integer(), nullable=False),
+        sa.Column('created_at', sa.DateTime()),
+        sa.Column('updated_at', sa.DateTime()),
+        sa.PrimaryKeyConstraint('id'),
+        sa.UniqueConstraint('email')
+    )
+
+    op.create_table(
+        'agents',
+        sa.Column('id', sa.Integer()),
+        sa.Column('project_id', sa.Integer()),
+        sa.Column('name', sa.VARCHAR(250)),
+        sa.Column('description', sa.VARCHAR(1000), nullable=False),
+
+        sa.Column('created_at', sa.DateTime()),
+        sa.Column('updated_at', sa.DateTime()),
+        sa.PrimaryKeyConstraint('id'),
+    )
+
+    op.create_table(
+        'agent_configurations',
+        sa.Column('id', sa.Integer()),
+        sa.Column('agent_id', sa.Integer()),
+        sa.Column('key', sa.VARCHAR(250)),
+        sa.Column('value', sa.VARCHAR(3000)),
+
+        sa.Column('created_at', sa.DateTime()),
+        sa.Column('updated_at', sa.DateTime()),
+        sa.PrimaryKeyConstraint('id'),
+    )
+
+    op.create_table(
+        'agent_executions',
+        sa.Column('id', sa.Integer()),
+        sa.Column('agent_id', sa.Integer()),
+        sa.Column('status', sa.VARCHAR(250)),
+        sa.Column('last_execution_time', sa.DateTime()),
+
+        sa.Column('created_at', sa.DateTime()),
+        sa.Column('updated_at', sa.DateTime()),
+        sa.PrimaryKeyConstraint('id'),
+    )
+
+    op.create_table(
+        'agent_execution_feeds',
+        sa.Column('id', sa.Integer()),
+        sa.Column('agent_id', sa.Integer()),
+        sa.Column('agent_execution_id', sa.Integer()),
+        sa.Column('feed', sa.TEXT()),
+        sa.Column('role', sa.VARCHAR(250)),
+        sa.Column('created_at', sa.DateTime()),
+        sa.Column('updated_at', sa.DateTime()),
+        sa.PrimaryKeyConstraint('id'),
+    )
+
+    op.create_table(
+        'budgets',
+        sa.Column('id', sa.Integer()),
+        sa.Column('agent_id', sa.Integer()),
+        sa.Column('cycle', sa.String()),
+        sa.Column('budget', sa.Float()),
+        sa.Column('created_at', sa.DateTime()),
+        sa.Column('updated_at', sa.DateTime()),
+        sa.PrimaryKeyConstraint('id'),
+    )
+
+    op.create_table('tools',
+        sa.Column('created_at', sa.DateTime()),
+        sa.Column('updated_at', sa.DateTime()),
+        sa.Column('id', sa.INTEGER(), autoincrement=True, nullable=False),
+        sa.Column('name', sa.VARCHAR(), autoincrement=False, nullable=True),
+        sa.Column('folder_name', sa.VARCHAR(), autoincrement=False, nullable=True),
+        sa.Column('class_name', sa.VARCHAR(), autoincrement=False, nullable=True),
+        sa.Column('file_name', sa.VARCHAR(), autoincrement=False, nullable=True),
+        sa.PrimaryKeyConstraint('id', name='tools_pkey')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    pass
+    # ### end Alembic commands ###
diff --git a/config_template.yaml b/config_template.yaml
@@ -31,4 +31,4 @@ DB_PASSWORD:
 MAX_TOOL_TOKEN_LIMIT: 600
 
 # redis details
-REDIS_URL: localhost:6379
+REDIS_URL: "localhost:6379"
diff --git a/gui/pages/Content/Agents/ActivityFeed.js b/gui/pages/Content/Agents/ActivityFeed.js
@@ -19,6 +19,15 @@ export default function ActivityFeed({selectedRunId, selectedRunStatus}) {
     return () => clearInterval(interval);
   }, []);
 
+  useEffect(() => {
+    const selectedRunId1 = selectedRunId;
+    const interval = window.setInterval(function(){
+      fetchFeeds(selectedRunId1);
+    }, 10000);
+
+    return () => clearInterval(interval);
+  }, [selectedRunId]);
+
   function checkEmptyText(text) {
     return text.replace(/\s/g, '') !== ''
   }
@@ -33,6 +42,16 @@ export default function ActivityFeed({selectedRunId, selectedRunStatus}) {
       });
   }, [selectedRunId])
 
+  function fetchFeeds(selectedRunId) {
+    getExecutionFeeds(selectedRunId)
+      .then((response) => {
+        setFeeds(response.data);
+      })
+      .catch((error) => {
+        console.error('Error fetching execution feeds:', error);
+      });
+  }
+
   return (<>
     <Head>
       {/* eslint-disable-next-line @next/next/no-page-custom-font */}

diff --git a/main.py b/main.py
@@ -110,10 +110,10 @@ def authjwt_exception_handler(request: Request, exc: AuthJWTException):
 organisation = session.query(Organisation).filter_by(id=1).first()
 
 if not organisation or organisation is None:
-    default_organization = Organisation(id=1, name='Default Organization',
+    organisation = Organisation(id=1, name='Default Organization',
                                         description='This is the default organization')
     print("Org create.....")
-    session.add(default_organization)
+    session.add(organisation)
     session.commit()
 
 project_name = "Default Project"

diff --git a/superagi/agent/super_agi.py b/superagi/agent/super_agi.py
@@ -138,10 +138,10 @@ def execute(self, goals: List[str]):
 
         current_tokens = TokenCounter.count_message_tokens(messages, self.llm.get_model())
 
-        spinner = Halo(text='Thinking...', spinner='dots')
-        spinner.start()
+        # spinner = Halo(text='Thinking...', spinner='dots')
+        # spinner.start()
         response = self.llm.chat_completion(messages, token_limit - current_tokens)
-        spinner.stop()
+        # spinner.stop()
         print("\n")
 
         if response['content'] is None:

diff --git a/superagi/controllers/agent_execution.py b/superagi/controllers/agent_execution.py
@@ -64,15 +64,14 @@ def update_agent_execution(agent_execution_id: int,
         raise HTTPException(status_code=400, detail="Invalid Request")
     db_agent_execution.status = agent_execution.status
     db_agent_execution.last_execution_time = datetime.now()
-    db_agent_execution.name = agent_execution.name
     db.session.commit()
 
     if db_agent_execution.status == "RUNNING":
         print("DB EXEC : ")
         print(db_agent_execution)
         print("JSON:")
         print(db_agent_execution.to_json())
-        execute_agent.delay(db_agent_execution.id)
+        execute_agent.delay(db_agent_execution.id, datetime.now())
         # AgentExecutor.create_execute_agent_task(db_agent_execution.id)
 
     return db_agent_execution

diff --git a/superagi/controllers/agent_execution_feed.py b/superagi/controllers/agent_execution_feed.py
@@ -1,3 +1,5 @@
+import json
+
 from fastapi_sqlalchemy import db
 from fastapi import HTTPException, Depends, Request
 from fastapi_jwt_auth import AuthJWT
@@ -58,4 +60,33 @@ def update_agent_execution_feed(agent_execution_feed_id: int,
 @router.get("/get/execution/{agent_execution_id}")
 def get_agent_execution_feed(agent_execution_id: int, Authorize: AuthJWT = Depends()):
     feeds = db.session.query(AgentExecutionFeed).filter_by(agent_execution_id=agent_execution_id).all()
-    return feeds
+    # parse json
+    final_feeds = []
+    for feed in feeds:
+        final_feeds.append(parse_feed(feed))
+    return final_feeds
+
+
+def parse_feed(feed):
+    if feed.role == "assistant":
+        try:
+            parsed = json.loads(feed.feed, strict=False)
+            format_prefix_yellow = "\033[93m\033[1m"
+            format_suffix_yellow = "\033[0m\033[0m"
+            format_prefix_green = "\033[92m\033[1m"
+            format_suffix_green = "\033[0m\033[0m"
+            final_output = format_prefix_yellow + "Thoughts: " + format_suffix_yellow + parsed["thoughts"][
+                "reasoning"] + "<br>"
+            final_output += format_prefix_yellow + "Plan: " + format_suffix_yellow + parsed["thoughts"]["plan"] + "<br>"
+            final_output += format_prefix_yellow + "Criticism: " + format_suffix_yellow + parsed["thoughts"][
+                "criticism"] + "<br>"
+            final_output += format_prefix_green + "Action : " + format_suffix_green + "<br>"
+            final_output += format_prefix_yellow + "Tool: " + format_suffix_yellow + parsed["command"]["name"] + "<br>"
+
+            return {"role": "assistant", "feed": final_output}
+        except Exception:
+            return feed
+    if feed.role == "assistant":
+        return feed
+
+    return feed
diff --git a/superagi/controllers/user.py b/superagi/controllers/user.py
@@ -14,7 +14,7 @@ def create_user(user: sqlalchemy_to_pydantic(User, exclude=["id"]), Authorize: A
     db_user = db.session.query(User).filter(User.email == user.email).first()
     if db_user:
         return db_user
-    db_user = User(name=user.name, email=user.email, password=user.password, organisation=user.organisation)
+    db_user = User(name=user.name, email=user.email, password=user.password, organisation_id=user.organisation_id)
     db.session.add(db_user)
     db.session.commit()
     print("User created", db_user)

diff --git a/superagi/jobs/agent_executor.py b/superagi/jobs/agent_executor.py
@@ -1,6 +1,8 @@
 # from superagi.models.types.agent_with_config import AgentWithConfig
 import importlib
 import json
+from datetime import datetime
+from time import time
 
 from celery import Celery
 from sqlalchemy.orm import sessionmaker
@@ -31,17 +33,6 @@
 Session = sessionmaker(bind=engine)
 
 class AgentExecutor:
-    app = None
-    @classmethod
-    def create_execute_agent_task(cls, agent_execution_id: int):
-        superagi.worker.execute_agent.apply_async(agent_execution_id, 10)
-        # if cls.app is None:
-        #     cls.app = Celery("superagi", include=["superagi.worker"], imports=["superagi.worker"])
-        #     cls.app.conf.broker_url = "redis://localhost:6379"  # 'redis://' + redis_url
-        #     cls.app.conf.result_backend = "redis://localhost:6379"  # 'redis://' + redis_url
-        #     cls.app.autodiscover_tasks(['superagi.worker'])
-        # cls.app.send_task("superagi.worker.execute_agent", args=[agent_execution_id])
-
     @staticmethod
     def validate_filename(filename):
         if filename.endswith(".py"):
@@ -106,9 +97,7 @@ def execute_next_action(self, agent_execution_id):
             return
         else:
             print("Starting next job for agent execution id: ", agent_execution_id)
-            # AgentExecutor.create_execute_agent_task(agent_execution_id)
-
-            # worker.execute_agent.delay(agent_execution_id)
+            superagi.worker.execute_agent.delay(agent_execution_id, datetime.now())
 
     def fetch_agent_configuration(self, session, agent, agent_execution):
         agent_configurations = session.query(AgentConfiguration).filter_by(agent_id=agent_execution.agent_id).all()

diff --git a/superagi/worker.py b/superagi/worker.py
@@ -7,16 +7,13 @@
 
 redis_url = get_config('REDIS_URL')
 
-CELERY_IMPORTS = ('superagi.worker')
 app = Celery("superagi", include=["superagi.worker"], imports=["superagi.worker"])
-app.conf.broker_url = "redis://localhost:6379" #'redis://' + redis_url
-app.conf.result_backend = "redis://localhost:6379" #'redis://' + redis_url
+app.conf.broker_url = "redis://" + redis_url + "/0"
+app.conf.result_backend = "redis://" + redis_url + "/0"
 app.conf.worker_concurrency = 10
-app.autodiscover_tasks(['superagi.worker'])
 
-
-@app.task(bind=True)
-def execute_agent(agent_execution_id: int):
+@app.task(name="execute_agent")
+def execute_agent(agent_execution_id: int, time):
     """Execute an agent step in background."""
-    print(agent_execution_id)
+    print("Execute agent:" + str(time) + "," + str(agent_execution_id))
     AgentExecutor().execute_next_action(agent_execution_id=agent_execution_id)
diff --git a/test.py b/test.py
@@ -1,21 +1,16 @@
-from superagi.tools.file.write_file import WriteFileTool
-from superagi.tools.file.read_file import ReadFileTool
-from superagi.tools.google_search.google_search import GoogleSearchSchema, GoogleSearchTool
-from superagi.tools.google_serp_search.google_serp_search import GoogleSerpTool
-from superagi.models.agent_config import AgentConfiguration
-from sqlalchemy import Column, Integer, String, Text, DateTime, ForeignKey
-from sqlalchemy.orm import relationship
-from superagi.models.base_model import DBBaseModel
-from superagi.models.organisation import Organisation
-from superagi.models.project import Project
+import argparse
+from datetime import datetime
+from time import time
+
+from sqlalchemy.orm import sessionmaker
+
+from superagi.worker import execute_agent
 from superagi.models.agent import Agent
+from superagi.models.agent_config import AgentConfiguration
 from superagi.models.agent_execution import AgentExecution
-from datetime import datetime
 from superagi.models.db import connectDB
-from sqlalchemy.orm import sessionmaker, query
-from celery_app import test_fucntion
-
-import argparse
+from superagi.models.organisation import Organisation
+from superagi.models.project import Project
 
 parser = argparse.ArgumentParser(description='Create a new agent.')
 parser.add_argument('--name', type=str, help='Agent name for the script.')
@@ -109,6 +104,6 @@ def run_superagi_cli(agent_name=None,agent_description=None,agent_goals=None):
     print("Final Execution")
     print(execution)
 
-    test_fucntion.delay(execution.to_json())
+    execute_agent.delay(execution.id, datetime.now())
 
 run_superagi_cli(agent_name=agent_name,agent_description=agent_description,agent_goals=agent_goals)