amyeroberts HF staff commited on
Commit
cd1bfbe
1 Parent(s): 1da1bc5

Upload tool

Browse files
Files changed (4) hide show
  1. app.py +4 -0
  2. debug_25192_upload_a_tool.py +52 -0
  3. requirements.txt +2 -0
  4. tool_config.json +5 -0
app.py ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ from transformers import launch_gradio_demo
2
+ from debug_25192_upload_a_tool import PostgreSQLTool
3
+
4
+ launch_gradio_demo(PostgreSQLTool)
debug_25192_upload_a_tool.py ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from psycopg2 import connect
2
+ from transformers.tools import Tool
3
+
4
+
5
+ class PostgreSQLTool(Tool):
6
+ name = "postgres_database_tool"
7
+ description = (
8
+ "This tool is used to query a PostgreSQL database with a SQL request. "
9
+ "The tool is already connected to the database. "
10
+ "Example: postgres_tool('SELECT field FROM my_table;')"
11
+ "It takes a SQL request as argument and returns the result of the query. "
12
+ )
13
+
14
+ inputs = ["text"]
15
+ outputs = ["text"]
16
+
17
+ debug = False
18
+
19
+ database = None
20
+ cursor = None
21
+
22
+ def __init__(self, debug: bool = False, **kwargs):
23
+ super().__init__(**kwargs)
24
+ self.debug = debug
25
+
26
+ def connect(self, host: str, database: str, user: str, password: str, port: int = 5432):
27
+ # Connect to the database and create a cursor
28
+ self.database = connect(
29
+ database=database, host=host, user=user, password=password, port=port)
30
+ self.cursor = self.database.cursor()
31
+
32
+ def disconnect(self):
33
+ # Close the connection to the database
34
+ self.database.close()
35
+
36
+ def __call__(self, query: str):
37
+ if self.debug:
38
+ print(f"[POSTGRESQL_TOOL] Executing: {query}")
39
+
40
+ try:
41
+ # Execute the query
42
+ self.cursor.execute(query)
43
+ except Exception as e:
44
+
45
+ if self.debug:
46
+ print(f"[POSTGRESQL_TOOL] Query failed: {e}")
47
+
48
+ # Return the error message
49
+ return "[POSTGRESQL_TOOL] Query failed: " + str(e)
50
+
51
+ # Return the result of the query
52
+ return self.cursor.fetchall()
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ psycopg2
2
+ transformers
tool_config.json ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ {
2
+ "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. ",
3
+ "name": "postgres_database_tool",
4
+ "tool_class": "debug_25192_upload_a_tool.PostgreSQLTool"
5
+ }