debug_25192_upload_a_tool.py

from psycopg2 import connect
from transformers.tools import Tool


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, **kwargs):
        super().__init__(**kwargs)
        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()