Upload tool

app.py

@@ -0,0 +1,4 @@
+from transformers import launch_gradio_demo
+from postgresql_tool import PostgreSQLTool
+
+launch_gradio_demo(PostgreSQLTool)

postgresql_tool.py

@@ -0,0 +1,54 @@
+from transformers.tools import Tool
+from psycopg2 import connect
+
+
+class PostgreSQLTool(Tool):
+    name = "postgres_database_tool"
+    description = (
+        "This tool is used to query a PostgreSQL database with a SQL request. "
+        "The tool is already connected to the database. "
+        "Example: postgres_tool('SELECT field FROM my_table;')"
+        "It takes a SQL request as argument and returns the result of the query. "
+    )
+
+    inputs = ["text"]
+    outputs = ["text"]
+
+    debug = False
+
+    database = None
+    cursor = None
+
+    def __init__(self, debug: bool = False):
+        super().__init__()
+        self.debug = debug
+
+    def connect(
+        self, host: str, database: str, user: str, password: str, port: int = 5432
+    ):
+        # Connect to the database and create a cursor
+        self.database = connect(
+            database=database, host=host, user=user, password=password, port=port
+        )
+        self.cursor = self.database.cursor()
+
+    def disconnect(self):
+        # Close the connection to the database
+        self.database.close()
+
+    def __call__(self, query: str):
+        if self.debug:
+            print(f"[POSTGRESQL_TOOL] Executing: {query}")
+
+        try:
+            # Execute the query
+            self.cursor.execute(query)
+        except Exception as e:
+            if self.debug:
+                print(f"[POSTGRESQL_TOOL] Query failed: {e}")
+
+            # Return the error message
+            return "[POSTGRESQL_TOOL] Query failed: " + str(e)
+
+        # Return the result of the query
+        return self.cursor.fetchall()

requirements.txt

@@ -0,0 +1,2 @@
+transformers
+psycopg2

tool_config.json

@@ -0,0 +1,5 @@
+{
+  "description": "This tool is used to query a PostgreSQL database with a SQL request. The tool is already connected to the database. Example: postgres_tool('SELECT field FROM my_table;')It takes a SQL request as argument and returns the result of the query. ",
+  "name": "postgres_database_tool",
+  "tool_class": "postgresql_tool.PostgreSQLTool"
+}