Spaces:
Sleeping
Sleeping
| 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() | |