-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy paths03e03.py
119 lines (99 loc) · 3.61 KB
/
s03e03.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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
# Standard library imports
import requests
# Local imports
from utilities.common import AIDevsClient, OpenAIClient
from utilities.config import AI_DEVS_API_KEY, S03E03_TASK_URL, S03E03_REPORT_URL
# =========================================================
# Configuration
# =========================================================
SUBMIT_URL = S03E03_REPORT_URL
TASK_NAME = "database"
API_URL = S03E03_TASK_URL
def get_table_info():
"""Get database table structures."""
# First get list of tables
payload = {
"task": TASK_NAME,
"apikey": AI_DEVS_API_KEY,
"query": "show tables"
}
response = requests.post(API_URL, json=payload).json()
# Extract table names dynamically by getting the first key from the first dictionary
if 'reply' in response and response['reply']:
table_key = list(response['reply'][0].keys())[0] # Gets the first (and only) key dynamically
tables = [table[table_key] for table in response['reply']]
print("Tables:", tables)
# Then get structure for each table
table_structures = {}
for table in tables:
payload["query"] = f"show create table {table}"
structure = requests.post(API_URL, json=payload).json()
print(f"\nStructure for {table}:", structure)
if 'reply' in structure:
table_structures[table] = structure
return table_structures
return None
def generate_sql_query(table_info, client_openai):
"""Generate SQL query using OpenAI."""
messages = [
{
"role": "system",
"content": "You are an SQL expert. Generate a SQL query based on the table structures provided. Return only the query without any formatting."
},
{
"role": "user",
"content": f"""
Table structures:
{table_info}
Question: Which active datacenter IDs are managed by users who are currently on vacation (inactive)?
"""
}
]
return client_openai.get_completion(
messages=messages,
model="gpt-4o-mini",
temperature=0.1
)
def test_query(sql_query):
"""Test the generated SQL query."""
test_payload = {
"task": TASK_NAME,
"apikey": AI_DEVS_API_KEY,
"query": sql_query
}
return requests.post(API_URL, json=test_payload).json()
def submit_results(dc_ids, client_aidevs):
"""Submit the final results."""
return client_aidevs.submit_answer(
answer={
"task": TASK_NAME,
"apikey": AI_DEVS_API_KEY,
"answer": dc_ids
},
submit_url=SUBMIT_URL
)
def main():
# Initialize clients
client_aidevs = AIDevsClient()
client_openai = OpenAIClient()
# Step 1: Get table structures
table_info = get_table_info()
if not table_info:
print("Failed to get table information")
return
# Step 2: Generate and test SQL query
sql_query = generate_sql_query(table_info, client_openai)
print("Generated SQL query:", sql_query)
# Step 3: Test the query
test_response = test_query(sql_query)
print("\nQuery test response:", test_response)
# Step 4: Process and submit results
if 'reply' in test_response:
dc_ids = [int(row['dc_id']) for row in test_response['reply']]
print("DC IDs:", dc_ids)
response = submit_results(dc_ids, client_aidevs)
print("Submission response:", response)
else:
print("Failed to get valid response from query test")
if __name__ == "__main__":
main()