code cleanup

superagi/tools/eva_db/db_create_connection_tool.py

@@ -1,90 +1,89 @@
-import json
 import evadb
 from typing import Type, Optional
 
 from pydantic import BaseModel, Field
 
-from superagi.helper.google_search import GoogleSearchWrap
-from superagi.helper.token_counter import TokenCounter
 from superagi.llms.base_llm import BaseLlm
 from superagi.tools.base_tool import BaseTool
 
 
-
 class DbCreateConnectionSchema(BaseModel):
     db_engine: str = Field(
         ...,
-        description = "The name of underlying database engine to connect to",
+        description="The name of underlying database engine to connect to",
     )
     username: str = Field(
         ...,
-        description = "The username for connecting to the database",
+        description="The username for connecting to the database",
     )
     host: str = Field(
         ...,
-        description = "The hostname for connecting to the database",
+        description="The hostname for connecting to the database",
     )
     port: str = Field(
         ...,
-        description = "The port for connecting to the database",
+        description="The port for connecting to the database",
     )
     database_name: str = Field(
         ...,
-        description = "The name of the database in the underlying database to connect to",
+        description="The database name in the underlying engine to connect to",
     )
 
 
 class DbCreateConnectionTool(BaseTool):
     """
-    Creates a connection entry in EvaDB for the underlying database engine to use it to execute future queries on the underlying database.
+    Register a connection entry within EvaDB, enabling the underlying database engine to execute future queries effectively.
 
     Attributes:
         name : The name of the tool.
         description : The description of the tool.
         args_schema : The args schema.
     """
+
     llm: Optional[BaseLlm] = None
     name = "DbCreateConnectionTool"
-    description = (
-        """
-        A tool for creating connection to the required database in the corresponding database engine using the credentials provided by the user.
+    description = """
+        A tool for establishing connections to the designated database within the corresponding db engine, utilizing user-provided credentials.
         """
-    )
     args_schema: Type[DbCreateConnectionSchema] = DbCreateConnectionSchema
 
     class Config:
         arbitrary_types_allowed = True
 
-    def _execute(self, db_engine: str, username: str, host: str, port: str, database_name: str) -> tuple:
+    def _execute(
+        self, db_engine: str, username: str, host: str, port: str, database_name: str
+    ) -> tuple:
         """
         Execute the EvaDB create connection tool.
 
         Args:
-            db_engine: The name of underlying database engine to connect to
+            db_engine: The name of db engine to connect to eg. postgres, sqlite
             username: The username for connecting to the database
             password: The password for connecting to the database
             host: The hostname for connecting to the database
             port: The port for connecting to the database
-            database_name: The name of the database in the underlying database to connect to
+            database_name: The database name in the underlying engine to connect to
 
         Returns:
             The response of create connection query
         """
-        # print("####^^^^####")
         cursor = evadb.connect().cursor()
-        db_connection_name = db_engine + '__' + database_name
-        create_db_query = """
-            CREATE DATABASE IF NOT EXISTS {connection_name} WITH ENGINE = '{db_engine}', PARAMETERS = {{
+        db_connection_name = db_engine + "__" + database_name
+
+        # todo: discuss a better way to get the password from user
+        password = ""
+        create_db_query = f"""
+            CREATE DATABASE IF NOT EXISTS {db_connection_name} 
+            WITH ENGINE = "{db_engine}", 
+            PARAMETERS = {{
                 "user": "{username}",
-                "host": "{hostname}",
+                "host": "{host}",
                 "port": "{port}",
-                "database": "{db_name}",
+                "database": "{database_name}",
                 "password": "{password}"
             }};
-        """.format(connection_name = db_connection_name, db_engine = db_engine, username = username, 
-        hostname = host, port = port, db_name = database_name, password = "ada@Gatech")
-        print("create_db_query", create_db_query)
-        #cursor.query(f"DROP DATABASE IF EXISTS {db_connection_name}").df()
-        ret = cursor.query(create_db_query).df()
+        """
+        print(f"DbCreateConnectionTool running query {create_db_query}")
+        cursor.query(create_db_query).df()
         cursor.close()
-        return "Sucessfully created the desired connection to underlying databse engine with name {connection}. This connection can now be used to query requested database.".format(connection = db_connection_name)
+        return f"Sucessfully created the desired connection to underlying database engine with name {db_connection_name}. This connection can now be used to query requested database."

superagi/tools/eva_db/db_query_connection_tool.py

@@ -1,42 +1,36 @@
-import json
 import evadb
 from typing import Type, Optional
 
 from pydantic import BaseModel, Field
 
-from superagi.helper.google_search import GoogleSearchWrap
-from superagi.helper.token_counter import TokenCounter
 from superagi.llms.base_llm import BaseLlm
 from superagi.tools.base_tool import BaseTool
 
 
-
 class DbQueryConnectionSchema(BaseModel):
     db_connection: str = Field(
         ...,
-        description="The name of the connection on which the provided connection should be executed. This connection is created using create connection tool",
+        description="The name of the connection on which the provided query should be executed. This connection is created using create connection tool",
     )
     query: str = Field(
         ...,
-        description="The query to be executed on the undelying database engine though EvaDB",
+        description="The query to be executed on the connection through EvaDB",
     )
 
 
-
 class DbQueryConnectionTool(BaseTool):
     """
-    Executes the provided query on the db engine attached in the provided db connection through EvaDB.
+    Executes the provided query on the provided db connection using EvaDB.
 
     Attributes:
         name : The name of the tool.
         description : The description of the tool.
         args_schema : The args schema.
     """
+
     llm: Optional[BaseLlm] = None
     name = "DbQueryConnectionTool"
-    description = (
-        "A tool for executing queries on the undelying database thorugh EvaDB"
-    )
+    description = "A tool for executing queries on the provided connection using EvaDB"
     args_schema: Type[DbQueryConnectionSchema] = DbQueryConnectionSchema
 
     class Config:
@@ -48,23 +42,26 @@ def _execute(self, db_connection: str, query: str) -> tuple:
 
         Args:
             db_connection: the connection on which the query should be executed
-            query : the query which should be executed on the underlying database through the provided connection
+            query : the query which should be executed on the provided connection
 
         Returns:
             Response of executed query.
         """
-        # print("####^^^^###")
         db_query = """
             USE {db_connection} {{
                 {query}
             }};
-        """.format(db_connection = db_connection, query = query.strip(';'))
+        """.format(
+            db_connection=db_connection, query=query.strip(";")
+        )
 
-        print("FINAL_QUERY: ", db_query)
+        print("DbQueryConnectionTool running {db_query} ....")
         cursor = evadb.connect().cursor()
-        ret = cursor.query(db_query).df()
-        print(ret)
+        response = cursor.query(db_query).df().to_string()
+        print(f"DbQueryConnectionTool query response {response}")
         cursor.close()
 
-        return ("Successfully executed the provided query", "result: " + ret.to_string())
-
+        return (
+            "Successfully executed the provided query",
+            "result: " + response,
+        )