-
Notifications
You must be signed in to change notification settings - Fork 0
/
functions.py
72 lines (63 loc) · 2.03 KB
/
functions.py
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
71
72
import sqlite3
def get_database_schema():
conn = sqlite3.connect('database.db')
cursor = conn.cursor()
cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
tables = cursor.fetchall()
schema = ""
for table_name in tables:
table_name = table_name[0]
schema = f"Table: {table_name}\n"
cursor.execute(f"PRAGMA table_info({table_name})")
columns = cursor.fetchall()
for column in columns:
schema = f" Column: {column[1]} Type: {column[2]}\n"
conn.close()
return schema
def get_functions():
return [{
"name": "use_database",
"description": "This function runs a SQL query against a SQLite database. Use this function to insert, read, update, or delete data from the database. This is the current schema of the database:\n\n" get_database_schema(),
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The SQL query to run."
}
}
},
}, {
"name": "run_code",
"description": "This function runs arbitrary Python code. Use this function to do anything that you can do in Python.",
"parameters": {
"type": "object",
"properties": {
"code": {
"type": "string",
"description": "The Python code to run."
}
}
},
}]
def use_database(query: str):
try:
print("Using database", query)
conn = sqlite3.connect('database.db')
cursor = conn.cursor()
cursor.execute(query)
results = cursor.fetchall()
conn.close()
return results
except Exception as e:
return {
"error": e.args[0]
}
def run_code(code: str):
try:
print("Running code", code)
return exec(code)
except Exception as e:
return {
"error": e.args[0]
}