test_tool_2 / debug_25192_upload_a_tool.py
amyeroberts's picture
amyeroberts HF staff
Upload tool
cd1bfbe
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()