Improve logging functionality with logger integration

hugbubby · Jun 16, 2023 · 1fbba0e · 1fbba0e
1 parent de81722
commit 1fbba0e
Show file tree

Hide file tree

Showing 34 changed files with 196 additions and 103 deletions.
diff --git a/cli2.py b/cli2.py
@@ -5,11 +5,12 @@
 import shutil
 from sys import platform
 from multiprocessing import Process
+from superagi.lib.logger import logger
 
 
 def check_command(command, message):
     if not shutil.which(command):
-        print(message)
+        logger.info(message)
         sys.exit(1)
 
 
@@ -18,7 +19,7 @@ def run_npm_commands(shell=False):
     try:
         subprocess.run(["npm", "install"], check=True, shell=shell)
     except subprocess.CalledProcessError:
-        print(f"Error during '{' '.join(sys.exc_info()[1].cmd)}'. Exiting.")
+        logger.error(f"Error during '{' '.join(sys.exc_info()[1].cmd)}'. Exiting.")
         sys.exit(1)
     os.chdir("..")
 
@@ -36,11 +37,11 @@ def run_server(shell=False,a_name=None,a_description=None,goals=None):
 
 
 def cleanup(api_process, ui_process, celery_process):
-    print("Shutting down processes...")
+    logger.info("Shutting down processes...")
     api_process.terminate()
     ui_process.terminate()
     celery_process.terminate()
-    print("Processes terminated. Exiting.")
+    logger.info("Processes terminated. Exiting.")
     sys.exit(1)
 
 

diff --git a/run_gui.py b/run_gui.py
@@ -3,18 +3,19 @@
 import subprocess
 from time import sleep
 import shutil
+from superagi.lib.logger import logger
 
 def check_command(command, message):
     if not shutil.which(command):
-        print(message)
+        logger.info(message)
         sys.exit(1)
 
 def run_npm_commands():
     os.chdir("gui")
     try:
         subprocess.run(["npm", "install"], check=True)
     except subprocess.CalledProcessError:
-        print(f"Error during '{' '.join(sys.exc_info()[1].cmd)}'. Exiting.")
+        logger.error(f"Error during '{' '.join(sys.exc_info()[1].cmd)}'. Exiting.")
         sys.exit(1)
     os.chdir("..")
 
@@ -26,10 +27,10 @@ def run_server():
     return api_process, ui_process
 
 def cleanup(api_process, ui_process):
-    print("Shutting down processes...")
+    logger.info("Shutting down processes...")
     api_process.terminate()
     ui_process.terminate()
-    print("Processes terminated. Exiting.")
+    logger.info("Processes terminated. Exiting.")
     sys.exit(1)
 
 if __name__ == "__main__":

diff --git a/superagi/agent/output_parser.py b/superagi/agent/output_parser.py
@@ -3,6 +3,7 @@
 from typing import Dict, NamedTuple, List
 import re
 from superagi.helper.json_cleaner import JsonCleaner
+from superagi.lib.logger import logger
 
 
 class AgentGPTAction(NamedTuple):
@@ -25,7 +26,7 @@ def parse(self, text: str) -> AgentGPTAction:
 class AgentOutputParser(BaseOutputParser):
     def parse(self, text: str) -> AgentGPTAction:
         try:
-            print(text)
+            logger.info(text)
             text = JsonCleaner.check_and_clean_json(text)
             parsed = json.loads(text, strict=False)
         except json.JSONDecodeError:
@@ -38,24 +39,25 @@ def parse(self, text: str) -> AgentGPTAction:
             format_suffix_yellow = "\033[0m\033[0m"
             format_prefix_green = "\033[92m\033[1m"
             format_suffix_green = "\033[0m\033[0m"
-            print(format_prefix_green + "Intelligence : " + format_suffix_green)
+            logger.info(format_prefix_green + "Intelligence : " + format_suffix_green)
             if "text" in parsed["thoughts"]:
-                print(format_prefix_yellow + "Thoughts: " + format_suffix_yellow + parsed["thoughts"]["text"] + "\n")
+                logger.info(format_prefix_yellow + "Thoughts: " + format_suffix_yellow + parsed["thoughts"]["text"] + "\n")
+
             if "reasoning" in parsed["thoughts"]:
-                print(format_prefix_yellow + "Reasoning: " + format_suffix_yellow + parsed["thoughts"]["reasoning"] + "\n")
+                logger.info(format_prefix_yellow + "Reasoning: " + format_suffix_yellow + parsed["thoughts"]["reasoning"] + "\n")
 
             if "plan" in parsed["thoughts"]:
-                print(format_prefix_yellow + "Plan: " + format_suffix_yellow + parsed["thoughts"]["plan"] + "\n")
+                logger.info(format_prefix_yellow + "Plan: " + format_suffix_yellow + parsed["thoughts"]["plan"] + "\n")
 
             if "criticism" in parsed["thoughts"]:
-                print(format_prefix_yellow + "Criticism: " + format_suffix_yellow + parsed["thoughts"]["criticism"] + "\n")
+                logger.info(format_prefix_yellow + "Criticism: " + format_suffix_yellow + parsed["thoughts"]["criticism"] + "\n")
 
-            print(format_prefix_green + "Action : " + format_suffix_green)
+            logger.info(format_prefix_green + "Action : " + format_suffix_green)
             # print(format_prefix_yellow + "Args: "+ format_suffix_yellow + parsed["tool"]["args"] + "\n")
             if parsed["tool"] is None or not parsed["tool"]:
                 return AgentGPTAction(name="", args="")
             if "name" in parsed["tool"]:
-                print(format_prefix_yellow + "Tool: " + format_suffix_yellow + parsed["tool"]["name"] + "\n")
+                logger.info(format_prefix_yellow + "Tool: " + format_suffix_yellow + parsed["tool"]["name"] + "\n")
             return AgentGPTAction(
                 name=parsed["tool"]["name"],
                 args=parsed["tool"]["args"],
@@ -78,7 +80,7 @@ def parse_tasks(self, text: str) -> AgentTasks:
                     error=f"Could not parse invalid json: {text}",
                 )
         try:
-            print("Tasks: ", parsed["tasks"])
+            logger.info("Tasks: ", parsed["tasks"])
             return AgentTasks(
                 tasks=parsed["tasks"]
             )

diff --git a/superagi/agent/super_agi.py b/superagi/agent/super_agi.py
@@ -32,6 +32,7 @@
 from superagi.models.resource import Resource
 from superagi.config.config import get_config
 import os
+from superagi.lib.logger import logger
 
 FINISH = "finish"
 WRITE_FILE = "Write File"
@@ -143,7 +144,7 @@ def execute(self, workflow_step: AgentWorkflowStep):
             #                                           agent_id=self.agent_config["agent_id"], feed=template_step.prompt,
             #                                           role="user")
 
-        print(prompt)
+        logger.info(prompt)
         if len(agent_feeds) <= 0:
             for message in messages:
                 agent_execution_feed = AgentExecutionFeed(agent_execution_id=self.agent_config["agent_execution_id"],
@@ -184,7 +185,7 @@ def execute(self, workflow_step: AgentWorkflowStep):
             for task in reversed(tasks):
                 task_queue.add_task(task)
             if len(tasks) > 0:
-                print("Tasks reprioritized in order: " + str(tasks))
+                logger.info("Tasks reprioritized in order: " + str(tasks))
             current_tasks = task_queue.get_tasks()
             if len(current_tasks) == 0:
                 final_response = {"result": "COMPLETE", "pending_task_count": 0}
@@ -195,7 +196,7 @@ def execute(self, workflow_step: AgentWorkflowStep):
             for task in reversed(tasks):
                 task_queue.add_task(task)
             if len(tasks) > 0:
-                print("Adding task to queue: " + str(tasks))
+                logger.info("Adding task to queue: " + str(tasks))
             for task in tasks:
                 agent_execution_feed = AgentExecutionFeed(agent_execution_id=self.agent_config["agent_execution_id"],
                                                           agent_id=self.agent_config["agent_id"],
@@ -215,7 +216,7 @@ def execute(self, workflow_step: AgentWorkflowStep):
                 final_response["result"] = "PENDING"
         session.commit()
 
-        print("Iteration completed moving to next iteration!")
+        logger.info("Iteration completed moving to next iteration!")
         session.close()
         return final_response
 
@@ -224,15 +225,15 @@ def handle_tool_response(self, assistant_reply):
         tools = {t.name: t for t in self.tools}
 
         if action.name == FINISH or action.name == "":
-            print("\nTask Finished :) \n")
+            logger.info("\nTask Finished :) \n")
             output = {"result": "COMPLETE", "retry": False}
             return output
         if action.name in tools:
             tool = tools[action.name]
             try:
                 observation = tool.execute(action.args)
-                print("Tool Observation : ")
-                print(observation)
+                logger.info("Tool Observation : ")
+                logger.info(observation)
 
             except ValidationError as e:
                 observation = (
@@ -255,7 +256,7 @@ def handle_tool_response(self, assistant_reply):
             )
             output = {"result": result, "retry": True}
 
-        print("Tool Response : " + str(output) + "\n")
+        logger.info("Tool Response : " + str(output) + "\n")
         return output
 
     def update_agent_execution_tokens(self, current_calls, total_tokens):

diff --git a/superagi/config/config.py b/superagi/config/config.py
@@ -2,6 +2,7 @@
 from pydantic import BaseSettings
 from pathlib import Path
 import yaml
+from superagi.lib.logger import logger
 
 CONFIG_FILE = "config.yaml"
 
@@ -21,7 +22,7 @@ def load_config(cls, config_file: str) -> dict:
                 config_data = {}
         else:
             # If config file doesn't exist, prompt for credentials and create new file
-            print("\033[91m\033[1m"
+            logger.info("\033[91m\033[1m"
         + "\nConfig file not found. Enter required keys and values."
         + "\033[0m\033[0m")
             config_data = {

diff --git a/superagi/controllers/config.py b/superagi/controllers/config.py
@@ -8,6 +8,7 @@
 from superagi.helper.auth import check_auth
 from fastapi_jwt_auth import AuthJWT
 from superagi.helper.encyption_helper import encrypt_data,decrypt_data
+from superagi.lib.logger import logger
 
 router = APIRouter()
 
@@ -40,10 +41,10 @@ def create_config(config: sqlalchemy_to_pydantic(Configuration, exclude=["id"]),
         db.session.flush()
         return existing_config
 
-    print("NEW CONFIG")
+    logger.info("NEW CONFIG")
     new_config = Configuration(organisation_id=organisation_id, key=config.key, value=config.value)
-    print(new_config)
-    print("ORGANISATION ID : ",organisation_id)
+    logger.info(new_config)
+    logger.info("ORGANISATION ID : ",organisation_id)
     db.session.add(new_config)
     db.session.commit()
     db.session.flush()

diff --git a/superagi/controllers/organisation.py b/superagi/controllers/organisation.py
@@ -8,6 +8,7 @@
 from superagi.helper.auth import check_auth
 from superagi.models.project import Project
 from superagi.models.user import User
+from superagi.lib.logger import logger
 
 router = APIRouter()
 
@@ -25,7 +26,7 @@ def create_organisation(organisation: sqlalchemy_to_pydantic(Organisation, exclu
     db.session.add(new_organisation)
     db.session.commit()
     db.session.flush()
-    print(new_organisation)
+    logger.info(new_organisation)
 
     return new_organisation
 

diff --git a/superagi/controllers/project.py b/superagi/controllers/project.py
@@ -6,6 +6,7 @@
 from fastapi import APIRouter
 from pydantic_sqlalchemy import sqlalchemy_to_pydantic
 from superagi.helper.auth import check_auth
+from superagi.lib.logger import logger
 
 router = APIRouter()
 
@@ -17,7 +18,7 @@ def create_project(project: sqlalchemy_to_pydantic(Project, exclude=["id"]),
 
     """Create a new project"""
 
-    print("Organisation_id : ", project.organisation_id)
+    logger.info("Organisation_id : ", project.organisation_id)
     organisation = db.session.query(Organisation).get(project.organisation_id)
 
     if not organisation:

diff --git a/superagi/controllers/resources.py b/superagi/controllers/resources.py
@@ -20,7 +20,7 @@
 from botocore.exceptions import NoCredentialsError
 import tempfile
 import requests
-
+from superagi.lib.logger import logger
 
 router = APIRouter()
 
@@ -63,7 +63,7 @@ async def upload(agent_id: int, file: UploadFile = File(...), name=Form(...), si
         path = 'input/'+file_name[0]+ '_'+str(datetime.datetime.now()).replace(' ','').replace('.','').replace(':','')+'.'+file_name[1]
         try:
             s3.upload_fileobj(file.file, bucket_name, path)
-            print("File uploaded successfully!")
+            logger.info("File uploaded successfully!")
         except NoCredentialsError:
             raise HTTPException(status_code=500, detail="AWS credentials not found. Check your configuration.")
 
@@ -72,7 +72,7 @@ async def upload(agent_id: int, file: UploadFile = File(...), name=Form(...), si
     db.session.add(resource)
     db.session.commit()
     db.session.flush()
-    print(resource)
+    logger.info(resource)
     return resource
 
 

diff --git a/superagi/controllers/user.py b/superagi/controllers/user.py
@@ -8,6 +8,7 @@
 from fastapi import APIRouter
 from pydantic_sqlalchemy import sqlalchemy_to_pydantic
 from superagi.helper.auth import check_auth
+from superagi.lib.logger import logger
 
 router = APIRouter()
 
@@ -25,7 +26,7 @@ def create_user(user: sqlalchemy_to_pydantic(User, exclude=["id"]),
     db.session.flush()
     organisation = Organisation.find_or_create_organisation(db.session, db_user)
     Project.find_or_create_default_project(db.session, organisation.id)
-    print("User created", db_user)
+    logger.info("User created", db_user)
     return db_user