Spaces:
Sleeping
Sleeping
File size: 1,852 Bytes
dd2bd16 bf7dd81 dd2bd16 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
from transformers.tools import Tool
class Connection:
def __init__(self, database, host, user, password, port):
pass
def cursor(self):
return self
def execute(self, query):
return query
def close(self):
pass
def connect(*args, **kwargs):
return Connection(*args, **kwargs)
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()
|