amyeroberts HF staff commited on
Commit
dd2bd16
1 Parent(s): 4ecadab

Upload tool

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