完成三种构建方法
This commit is contained in:
46291
random_sample/cms_1k_simple.json
Normal file
46291
random_sample/cms_1k_simple.json
Normal file
File diff suppressed because it is too large
Load Diff
17
random_sample/core_tables.json
Normal file
17
random_sample/core_tables.json
Normal file
@@ -0,0 +1,17 @@
|
||||
[
|
||||
"catalog_product_entity", "catalog_product_entity_varchar", "catalog_product_entity_int",
|
||||
"catalog_product_entity_decimal", "catalog_product_entity_text", "catalog_category_entity",
|
||||
"catalog_category_product", "catalog_eav_attribute", "eav_attribute", "eav_attribute_option",
|
||||
"eav_attribute_option_value", "eav_entity_type", "cataloginventory_stock_item",
|
||||
"cataloginventory_stock", "customer_entity", "customer_address_entity", "customer_group",
|
||||
"customer_grid_flat", "sales_order", "sales_order_item", "sales_order_address",
|
||||
"sales_order_payment", "sales_order_grid", "sales_order_status", "sales_order_status_state",
|
||||
"sales_invoice", "sales_invoice_item", "sales_shipment", "sales_shipment_item",
|
||||
"sales_creditmemo_grid", "sales_bestsellers_aggregated_daily",
|
||||
"sales_bestsellers_aggregated_monthly", "sales_bestsellers_aggregated_yearly",
|
||||
"sales_order_aggregated_created", "review", "review_detail", "review_status", "rating",
|
||||
"rating_option", "rating_option_vote", "cms_page", "store", "store_group", "store_website",
|
||||
"directory_country", "directory_country_region", "directory_country_region_name",
|
||||
"search_query", "sequence_invoice_1", "sequence_order_1", "sequence_shipment_1",
|
||||
"sales_sequence_meta", "sales_sequence_profile", "catalog_category_entity_varchar"
|
||||
]
|
||||
364
random_sample/generate_tasks copy 2.py
Normal file
364
random_sample/generate_tasks copy 2.py
Normal file
@@ -0,0 +1,364 @@
|
||||
import os
|
||||
import random
|
||||
import json
|
||||
import mysql.connector
|
||||
from openai import OpenAI
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# --- Configuration ---
|
||||
load_dotenv()
|
||||
|
||||
MYSQL_CONFIG = {
|
||||
"host": "localhost",
|
||||
"port": "23306",
|
||||
"user": "mcpuser",
|
||||
"password": "StrongPass123!",
|
||||
"database": "magentodb"
|
||||
}
|
||||
|
||||
OPENAI_CONFIG = {
|
||||
"api_key": os.getenv("OPENAI_API_KEY"),
|
||||
"base_url": os.getenv("OPENAI_BASE_URL"),
|
||||
"model": "gpt-4o"
|
||||
}
|
||||
|
||||
# --- Prompt Template ---
|
||||
# This is a carefully engineered prompt to guide the LLM's output.
|
||||
PROMPT_TEMPLATE = """
|
||||
You are an expert database analyst and a creative test case designer for e-commerce web applications.
|
||||
Your goal is to generate realistic administrative tasks that can be solved by a Web Agent navigating an admin panel.
|
||||
|
||||
I will provide you with the following context:
|
||||
1. **Full Database Schema**: A list of `CREATE TABLE` statements for the core tables of a Magento e-commerce platform.
|
||||
2. **Sampled Data**: A JSON object containing 5 random rows of data from 5 randomly selected core tables. This data is REAL and should be used to inspire specific, answerable questions.
|
||||
|
||||
## Your Task
|
||||
|
||||
Based on the provided schema and sample data, create a JSON object containing a single key, "questions", which holds an array of up to 10 unique task objects.
|
||||
|
||||
### Requirements for Each Question:
|
||||
- **Web Agent Solvable**: The task must represent a realistic action an administrator would perform in a web UI (e.g., "Find all orders for customer X", "Update the stock for product Y", "Approve a pending review").
|
||||
- **Grounded in Data**: The questions should be specific, using names, IDs, or values from the provided **Sampled Data** to make them concrete.
|
||||
- **Utilize Schema**: You can formulate questions that require joining tables, even if not all tables were sampled. The full schema is your guide.
|
||||
|
||||
### Output Format
|
||||
The final output MUST be a single, valid JSON object. Do not include any other text, explanations, or markdown formatting like ```json.
|
||||
The JSON object must have one key: "questions", containing a JSON array of task objects.
|
||||
|
||||
Each object in the array must contain exactly three keys: `question`, `answer`, and `sql`.
|
||||
|
||||
- **`question`**: (string) A natural language description of the task for a web agent.
|
||||
- **`answer`**: (string, integer, float, or list) The precise and concise answer to the question, derived by running the SQL query against the database.
|
||||
- **`sql`**: (string) The exact, runnable MySQL query that was used to find the answer.
|
||||
|
||||
### Output Format Example
|
||||
```json
|
||||
{{
|
||||
"questions": [
|
||||
{{
|
||||
"question": "What is the email address for customer with ID 5?",
|
||||
"answer": "customer5@example.com",
|
||||
"sql": "SELECT email FROM customer_entity WHERE entity_id = 5;"
|
||||
}},
|
||||
{{
|
||||
"question": "Find the total quantity of item with SKU 'ABC-123' in the cart.",
|
||||
"answer": 3,
|
||||
"sql": "SELECT SUM(qty) FROM quote_item WHERE sku = 'ABC-123';"
|
||||
}}
|
||||
]
|
||||
}}
|
||||
```
|
||||
|
||||
---
|
||||
### Full Database Schema
|
||||
{schema_context}
|
||||
|
||||
---
|
||||
### Sampled Data
|
||||
Here is the sample data from randomly selected tables. Use this to make your questions specific.
|
||||
|
||||
{sampled_data_str}
|
||||
|
||||
---
|
||||
Now, generate the JSON object based on these instructions.
|
||||
"""
|
||||
|
||||
# This is a carefully engineered prompt to verify the LLM's own output.
|
||||
SEMANTIC_VERIFICATION_PROMPT_TEMPLATE = """
|
||||
You are a meticulous data verifier. Your task is to determine if a given "answer" is semantically correct and accurately supported by the "SQL query result".
|
||||
|
||||
I will provide you with a JSON object containing:
|
||||
1. `question`: The original question asked.
|
||||
2. `sql`: The SQL query used to find the answer.
|
||||
3. `answer`: The answer generated by a previous AI.
|
||||
4. `sql_result`: The actual data returned by executing the SQL query.
|
||||
|
||||
## Your Task
|
||||
Carefully analyze the `sql_result` and compare it to the `answer`. The match should be semantic, not just a simple substring match. For example, if the question is "How many products are in stock?", an answer of "5" should be verifiable from the SQL result which might be `[(5,)]`.
|
||||
|
||||
### Requirements:
|
||||
- Respond with a single JSON object.
|
||||
- Do not include any other text, explanations, or markdown formatting.
|
||||
- The JSON object must have exactly two keys:
|
||||
- `is_match`: (boolean) `true` if the `answer` is fully and accurately supported by the `sql_result`, otherwise `false`.
|
||||
- `reason`: (string) A brief explanation for your decision. If it's a mismatch, explain why (e.g., "The answer is 'John Doe' but the result contains 'Jane Doe'", "The answer is a count but the result is a list of names").
|
||||
|
||||
---
|
||||
### Verification Data
|
||||
{task_data_json}
|
||||
---
|
||||
|
||||
Now, provide your verification as a JSON object.
|
||||
"""
|
||||
|
||||
def get_db_connection():
|
||||
"""Establishes a connection to the MySQL database."""
|
||||
try:
|
||||
conn = mysql.connector.connect(**MYSQL_CONFIG)
|
||||
return conn
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Error connecting to MySQL: {err}")
|
||||
return None
|
||||
|
||||
def get_full_schema(cursor, tables):
|
||||
"""Fetches the CREATE TABLE statements for all core tables."""
|
||||
schema_parts = []
|
||||
for table_name in tables:
|
||||
try:
|
||||
cursor.execute(f"SHOW CREATE TABLE `{table_name}`")
|
||||
result = cursor.fetchone()
|
||||
if result:
|
||||
schema_parts.append(result[1]) # result[1] is the CREATE TABLE statement
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Warning: Could not get schema for table {table_name}: {err}")
|
||||
return "\n\n".join(schema_parts)
|
||||
|
||||
def get_random_tables_and_samples(cursor, tables, num_tables=5, num_samples=5):
|
||||
"""Selects random tables and samples random rows from them."""
|
||||
selected_tables = random.sample(tables, num_tables)
|
||||
sampled_data = {}
|
||||
|
||||
for table_name in selected_tables:
|
||||
try:
|
||||
# Use ORDER BY RAND() for random sampling. Can be slow on very large tables.
|
||||
query = f"SELECT * FROM `{table_name}` ORDER BY RAND() LIMIT {num_samples}"
|
||||
cursor.execute(query)
|
||||
|
||||
rows = cursor.fetchall()
|
||||
if not rows:
|
||||
sampled_data[table_name] = []
|
||||
continue
|
||||
|
||||
columns = [desc[0] for desc in cursor.description]
|
||||
|
||||
# Convert rows (tuples) to a list of dictionaries
|
||||
sampled_rows = []
|
||||
for row in rows:
|
||||
row_dict = {}
|
||||
for i, col_value in enumerate(row):
|
||||
# Handle bytes by decoding, fall back to string representation
|
||||
if isinstance(col_value, bytes):
|
||||
try:
|
||||
row_dict[columns[i]] = col_value.decode('utf-8')
|
||||
except UnicodeDecodeError:
|
||||
row_dict[columns[i]] = str(col_value)
|
||||
else:
|
||||
row_dict[columns[i]] = col_value
|
||||
sampled_rows.append(row_dict)
|
||||
|
||||
sampled_data[table_name] = sampled_rows
|
||||
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Warning: Could not sample data from table {table_name}: {err}")
|
||||
sampled_data[table_name] = f"Error: {err}"
|
||||
|
||||
return sampled_data
|
||||
|
||||
def generate_questions(client, schema_context, sampled_data):
|
||||
"""Generates questions by calling the OpenAI API."""
|
||||
if not client:
|
||||
raise ValueError("OpenAI client not provided.")
|
||||
|
||||
sampled_data_str = json.dumps(sampled_data, indent=2, default=str)
|
||||
|
||||
prompt = PROMPT_TEMPLATE.format(
|
||||
schema_context=schema_context,
|
||||
sampled_data_str=sampled_data_str
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.chat.completions.create(
|
||||
model=OPENAI_CONFIG["model"],
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant designed to output JSON."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0.7,
|
||||
response_format={"type": "json_object"},
|
||||
)
|
||||
content = response.choices[0].message.content
|
||||
data = json.loads(content)
|
||||
|
||||
# The prompt asks for {"questions": [...]}, so we extract the list.
|
||||
if isinstance(data, dict) and "questions" in data and isinstance(data["questions"], list):
|
||||
return data["questions"]
|
||||
elif isinstance(data, list):
|
||||
# Fallback in case the model returns a list directly
|
||||
print("Warning: Model returned a raw list instead of an object with a 'questions' key.")
|
||||
return data
|
||||
else:
|
||||
print(f"Warning: Failed to find a 'questions' list in the model's output. Got: {content}")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error calling OpenAI API or parsing JSON: {e}")
|
||||
return None
|
||||
|
||||
def semantic_validate_tasks(tasks, client):
|
||||
"""
|
||||
Uses an LLM to semantically validate if the task's answer matches the SQL result.
|
||||
"""
|
||||
if not tasks:
|
||||
return []
|
||||
|
||||
final_validated_tasks = []
|
||||
print("\nPerforming semantic validation with GPT-4o...")
|
||||
|
||||
for task in tasks:
|
||||
# Prepare data for the prompt, including the SQL result
|
||||
task_data_for_prompt = {
|
||||
"question": task["question"],
|
||||
"sql": task["sql"],
|
||||
"answer": task["answer"],
|
||||
"sql_result": task["sql_result"]
|
||||
}
|
||||
task_data_json = json.dumps(task_data_for_prompt, indent=2, default=str)
|
||||
|
||||
prompt = SEMANTIC_VERIFICATION_PROMPT_TEMPLATE.format(task_data_json=task_data_json)
|
||||
|
||||
try:
|
||||
print(f" - Verifying question: \"{task['question'][:80]}...\"")
|
||||
response = client.chat.completions.create(
|
||||
model=OPENAI_CONFIG["model"],
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant designed to output JSON."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0.0, # We want deterministic validation
|
||||
response_format={"type": "json_object"},
|
||||
)
|
||||
content = response.choices[0].message.content
|
||||
verification_result = json.loads(content)
|
||||
|
||||
if verification_result.get("is_match") is True:
|
||||
# Task is valid. Rename sql_result for the final output.
|
||||
print(f" - Validation PASSED.")
|
||||
task['sql_execute_result'] = task.pop('sql_result')
|
||||
final_validated_tasks.append(task)
|
||||
else:
|
||||
reason = verification_result.get('reason', 'No reason provided.')
|
||||
print(f" - Validation FAILED. Filtering task.")
|
||||
print(f" - Reason: {reason}")
|
||||
print(f" - Question: {task['question']}")
|
||||
print(f" - Expected Answer: {json.dumps(task['answer'], default=str)}")
|
||||
print(f" - SQL: {task['sql']}")
|
||||
sql_result_str = json.dumps(task['sql_result'], indent=2, default=str)
|
||||
print(f" - SQL Result: {sql_result_str}")
|
||||
|
||||
except Exception as e:
|
||||
print(f" - An error occurred during semantic validation for task, filtering it out: {e}")
|
||||
print(f" - Question: {task.get('question', 'N/A')}")
|
||||
print(f" - SQL: {task.get('sql', 'N/A')}")
|
||||
|
||||
return final_validated_tasks
|
||||
|
||||
def main():
|
||||
"""Main function to run the script."""
|
||||
# 1. Load the list of core tables
|
||||
try:
|
||||
with open('core_tables.json', 'r') as f:
|
||||
core_tables = json.load(f)
|
||||
except FileNotFoundError:
|
||||
print("Error: core_tables.json not found. Please create it.")
|
||||
return
|
||||
|
||||
# 2. Connect to the database
|
||||
conn = get_db_connection()
|
||||
if not conn:
|
||||
return
|
||||
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 3. Setup OpenAI Client
|
||||
if not OPENAI_CONFIG["api_key"]:
|
||||
print("Error: OPENAI_API_KEY environment variable not set.")
|
||||
return
|
||||
client = OpenAI(api_key=OPENAI_CONFIG["api_key"], base_url=OPENAI_CONFIG["base_url"])
|
||||
|
||||
try:
|
||||
# 4. Get full schema context
|
||||
print("Fetching full database schema...")
|
||||
schema_context = get_full_schema(cursor, core_tables)
|
||||
|
||||
# 5. Get random samples and print them
|
||||
print("Sampling data from 5 random tables...")
|
||||
sampled_data = get_random_tables_and_samples(cursor, core_tables, num_tables=5, num_samples=5)
|
||||
print(f"Sampled from tables: {list(sampled_data.keys())}")
|
||||
print("\n--- Sampled Data ---")
|
||||
print(json.dumps(sampled_data, indent=2, default=str))
|
||||
print("---------------------\n")
|
||||
|
||||
# 6. Generate questions using the LLM
|
||||
print("Generating questions with GPT-4o...")
|
||||
generated_tasks = generate_questions(client, schema_context, sampled_data)
|
||||
|
||||
# 7. Initial validation (SQL execution and substring check)
|
||||
pre_validated_tasks = []
|
||||
if generated_tasks:
|
||||
print("\nPerforming initial validation (SQL execution and substring match)...")
|
||||
for task in generated_tasks:
|
||||
if not isinstance(task, dict) or not all(k in task for k in ['sql', 'answer', 'question']):
|
||||
print(f"Filtering task due to malformed structure or missing keys: {task}")
|
||||
continue
|
||||
|
||||
try:
|
||||
cursor.execute(task['sql'])
|
||||
sql_result = cursor.fetchall()
|
||||
answer_str = str(task['answer'])
|
||||
result_str = str(sql_result)
|
||||
|
||||
if answer_str in result_str:
|
||||
task['sql_result'] = sql_result # Attach result for the next validation step
|
||||
pre_validated_tasks.append(task)
|
||||
else:
|
||||
print(f"Filtering task: Answer '{answer_str}' not found in SQL result.")
|
||||
print(f" - Question: {task['question']}")
|
||||
print(f" - SQL: {task['sql']}")
|
||||
print(f" - Result: {result_str[:250]}...")
|
||||
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Filtering task due to SQL error: {err}")
|
||||
print(f" - Question: {task['question']}")
|
||||
print(f" - SQL: {task['sql']}")
|
||||
except Exception as e:
|
||||
print(f"An unexpected error occurred during initial validation for task {task}: {e}")
|
||||
|
||||
# 8. Semantic validation using LLM
|
||||
validated_tasks = semantic_validate_tasks(pre_validated_tasks, client)
|
||||
|
||||
# 9. Print the final JSON output
|
||||
if validated_tasks:
|
||||
print("\n--- Final Validated Tasks ---")
|
||||
print(json.dumps(validated_tasks, indent=2, default=str))
|
||||
else:
|
||||
print("Failed to generate any valid tasks after all validation steps.")
|
||||
|
||||
finally:
|
||||
# 10. Close the database connection
|
||||
if conn.is_connected():
|
||||
cursor.close()
|
||||
conn.close()
|
||||
print("\nDatabase connection closed.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
370
random_sample/generate_tasks copy 3.py
Normal file
370
random_sample/generate_tasks copy 3.py
Normal file
@@ -0,0 +1,370 @@
|
||||
import os
|
||||
import random
|
||||
import json
|
||||
import mysql.connector
|
||||
from openai import OpenAI
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# --- Configuration ---
|
||||
load_dotenv()
|
||||
|
||||
MYSQL_CONFIG = {
|
||||
"host": "localhost",
|
||||
"port": "23306",
|
||||
"user": "mcpuser",
|
||||
"password": "StrongPass123!",
|
||||
"database": "magentodb"
|
||||
}
|
||||
|
||||
OPENAI_CONFIG = {
|
||||
"api_key": os.getenv("OPENAI_API_KEY"),
|
||||
"base_url": os.getenv("OPENAI_BASE_URL"),
|
||||
"model": "gpt-4o"
|
||||
}
|
||||
|
||||
# --- Prompt Template ---
|
||||
# This is a carefully engineered prompt to guide the LLM's output.
|
||||
PROMPT_TEMPLATE = """
|
||||
You are an expert database analyst and a creative test case designer for e-commerce web applications.
|
||||
Your goal is to generate realistic administrative tasks that can be solved by a Web Agent navigating an admin panel.
|
||||
|
||||
I will provide you with the following context:
|
||||
1. **Full Database Schema**: A list of `CREATE TABLE` statements for the core tables of a Magento e-commerce platform.
|
||||
2. **Sampled Data**: A JSON object containing 5 random rows of data from 5 randomly selected core tables. This data is REAL and should be used to inspire specific, answerable questions.
|
||||
|
||||
## Your Task
|
||||
|
||||
Based on the provided schema and sample data, create a JSON object containing a single key, "questions", which holds an array of up to 10 unique task objects.
|
||||
|
||||
### Requirements for Each Question:
|
||||
- **Web Agent Solvable**: The task must represent a realistic action an administrator would perform in a web UI (e.g., "Find all orders for customer X", "Update the stock for product Y", "Approve a pending review").
|
||||
- **Grounded in Data**: The questions should be specific, using names, IDs, or values from the provided **Sampled Data** to make them concrete.
|
||||
- **Utilize Schema**: You can formulate questions that require joining tables, even if not all tables were sampled. The full schema is your guide.
|
||||
|
||||
### Output Format
|
||||
The final output MUST be a single, valid JSON object. Do not include any other text, explanations, or markdown formatting like ```json.
|
||||
The JSON object must have one key: "questions", containing a JSON array of task objects.
|
||||
|
||||
Each object in the array must contain exactly three keys: `question`, `answer`, and `sql`.
|
||||
|
||||
- **`question`**: (string) A natural language description of the task for a web agent.
|
||||
- **`answer`**: (string, integer, float, or list) The precise and concise answer to the question, derived by running the SQL query against the database.
|
||||
- **`sql`**: (string) The exact, runnable MySQL query that was used to find the answer.
|
||||
|
||||
### Output Format Example
|
||||
```json
|
||||
{{
|
||||
"questions": [
|
||||
{{
|
||||
"question": "What is the email address for customer with ID 5?",
|
||||
"answer": "customer5@example.com",
|
||||
"sql": "SELECT email FROM customer_entity WHERE entity_id = 5;"
|
||||
}},
|
||||
{{
|
||||
"question": "Find the total quantity of item with SKU 'ABC-123' in the cart.",
|
||||
"answer": 3,
|
||||
"sql": "SELECT SUM(qty) FROM quote_item WHERE sku = 'ABC-123';"
|
||||
}}
|
||||
]
|
||||
}}
|
||||
```
|
||||
|
||||
---
|
||||
### Full Database Schema
|
||||
{schema_context}
|
||||
|
||||
---
|
||||
### Sampled Data
|
||||
Here is the sample data from randomly selected tables. Use this to make your questions specific.
|
||||
|
||||
{sampled_data_str}
|
||||
|
||||
---
|
||||
Now, generate the JSON object based on these instructions.
|
||||
"""
|
||||
|
||||
# This is a new prompt to evaluate results and generate a corrected answer.
|
||||
SEMANTIC_EVALUATION_PROMPT_TEMPLATE = """
|
||||
You are a precise data analyst. Your task is to evaluate if a SQL query's result adequately answers a given natural language question. If it does, you must formulate a concise, natural-language answer.
|
||||
|
||||
I will provide you with a JSON object containing:
|
||||
1. `question`: The original question asked.
|
||||
2. `sql`: The SQL query that was executed.
|
||||
3. `sql_result`: The actual data returned by executing the SQL query.
|
||||
|
||||
## Your Task
|
||||
1. **Analyze**: Determine if the `sql_result` contains the necessary information to definitively answer the `question`.
|
||||
2. **Respond**: Based on your analysis, generate a JSON object with one of two structures.
|
||||
|
||||
### Case 1: The question CAN be answered
|
||||
If the `sql_result` provides a clear answer, respond with:
|
||||
```json
|
||||
{{
|
||||
"can_answer": true,
|
||||
"new_answer": "..."
|
||||
}}
|
||||
```
|
||||
- `can_answer`: (boolean) Must be `true`.
|
||||
- `new_answer`: (string, integer, float, or list) A concise, human-readable answer derived *only* from the `sql_result`. For example, if the result is `[(52.00,)]`, the answer can be "52.00" or 52.00.
|
||||
|
||||
### Case 2: The question CANNOT be answered
|
||||
If the `sql_result` is empty, irrelevant, or insufficient to answer the question, respond with:
|
||||
```json
|
||||
{{
|
||||
"can_answer": false,
|
||||
"reason": "..."
|
||||
}}
|
||||
```
|
||||
- `can_answer`: (boolean) Must be `false`.
|
||||
- `reason`: (string) A brief explanation for why the question cannot be answered from the given data (e.g., "The query returned no results.", "The result contains internal IDs, not the requested customer names.").
|
||||
|
||||
---
|
||||
### Evaluation Data
|
||||
{task_data_json}
|
||||
---
|
||||
|
||||
Now, provide your evaluation as a JSON object.
|
||||
"""
|
||||
|
||||
def get_db_connection():
|
||||
"""Establishes a connection to the MySQL database."""
|
||||
try:
|
||||
conn = mysql.connector.connect(**MYSQL_CONFIG)
|
||||
return conn
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Error connecting to MySQL: {err}")
|
||||
return None
|
||||
|
||||
def get_full_schema(cursor, tables):
|
||||
"""Fetches the CREATE TABLE statements for all core tables."""
|
||||
schema_parts = []
|
||||
for table_name in tables:
|
||||
try:
|
||||
cursor.execute(f"SHOW CREATE TABLE `{table_name}`")
|
||||
result = cursor.fetchone()
|
||||
if result:
|
||||
schema_parts.append(result[1]) # result[1] is the CREATE TABLE statement
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Warning: Could not get schema for table {table_name}: {err}")
|
||||
return "\n\n".join(schema_parts)
|
||||
|
||||
def get_random_tables_and_samples(cursor, tables, num_tables=5, num_samples=5):
|
||||
"""Selects random tables and samples random rows from them."""
|
||||
selected_tables = random.sample(tables, num_tables)
|
||||
sampled_data = {}
|
||||
|
||||
for table_name in selected_tables:
|
||||
try:
|
||||
# Use ORDER BY RAND() for random sampling. Can be slow on very large tables.
|
||||
query = f"SELECT * FROM `{table_name}` ORDER BY RAND() LIMIT {num_samples}"
|
||||
cursor.execute(query)
|
||||
|
||||
rows = cursor.fetchall()
|
||||
if not rows:
|
||||
sampled_data[table_name] = []
|
||||
continue
|
||||
|
||||
columns = [desc[0] for desc in cursor.description]
|
||||
|
||||
# Convert rows (tuples) to a list of dictionaries
|
||||
sampled_rows = []
|
||||
for row in rows:
|
||||
row_dict = {}
|
||||
for i, col_value in enumerate(row):
|
||||
# Handle bytes by decoding, fall back to string representation
|
||||
if isinstance(col_value, bytes):
|
||||
try:
|
||||
row_dict[columns[i]] = col_value.decode('utf-8')
|
||||
except UnicodeDecodeError:
|
||||
row_dict[columns[i]] = str(col_value)
|
||||
else:
|
||||
row_dict[columns[i]] = col_value
|
||||
sampled_rows.append(row_dict)
|
||||
|
||||
sampled_data[table_name] = sampled_rows
|
||||
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Warning: Could not sample data from table {table_name}: {err}")
|
||||
sampled_data[table_name] = f"Error: {err}"
|
||||
|
||||
return sampled_data
|
||||
|
||||
def generate_questions(client, schema_context, sampled_data):
|
||||
"""Generates questions by calling the OpenAI API."""
|
||||
if not client:
|
||||
raise ValueError("OpenAI client not provided.")
|
||||
|
||||
sampled_data_str = json.dumps(sampled_data, indent=2, default=str)
|
||||
|
||||
prompt = PROMPT_TEMPLATE.format(
|
||||
schema_context=schema_context,
|
||||
sampled_data_str=sampled_data_str
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.chat.completions.create(
|
||||
model=OPENAI_CONFIG["model"],
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant designed to output JSON."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0.7,
|
||||
response_format={"type": "json_object"},
|
||||
)
|
||||
content = response.choices[0].message.content
|
||||
data = json.loads(content)
|
||||
|
||||
# The prompt asks for {"questions": [...]}, so we extract the list.
|
||||
if isinstance(data, dict) and "questions" in data and isinstance(data["questions"], list):
|
||||
return data["questions"]
|
||||
elif isinstance(data, list):
|
||||
# Fallback in case the model returns a list directly
|
||||
print("Warning: Model returned a raw list instead of an object with a 'questions' key.")
|
||||
return data
|
||||
else:
|
||||
print(f"Warning: Failed to find a 'questions' list in the model's output. Got: {content}")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error calling OpenAI API or parsing JSON: {e}")
|
||||
return None
|
||||
|
||||
def evaluate_and_refine_tasks(tasks, client):
|
||||
"""
|
||||
Uses an LLM to evaluate if a SQL result answers the question and refines the answer.
|
||||
"""
|
||||
if not tasks:
|
||||
return []
|
||||
|
||||
final_validated_tasks = []
|
||||
print("\nPerforming semantic evaluation and answer refinement with GPT-4o...")
|
||||
|
||||
for task in tasks:
|
||||
# Prepare data for the prompt, excluding the original 'answer'
|
||||
task_data_for_prompt = {
|
||||
"question": task["question"],
|
||||
"sql": task["sql"],
|
||||
"sql_result": task["sql_result"]
|
||||
}
|
||||
task_data_json = json.dumps(task_data_for_prompt, indent=2, default=str)
|
||||
|
||||
prompt = SEMANTIC_EVALUATION_PROMPT_TEMPLATE.format(task_data_json=task_data_json)
|
||||
|
||||
try:
|
||||
print(f" - Evaluating question: \"{task['question'][:80]}...\"")
|
||||
response = client.chat.completions.create(
|
||||
model=OPENAI_CONFIG["model"],
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant designed to output JSON."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0.0, # We want deterministic evaluation
|
||||
response_format={"type": "json_object"},
|
||||
)
|
||||
content = response.choices[0].message.content
|
||||
evaluation_result = json.loads(content)
|
||||
|
||||
if evaluation_result.get("can_answer") is True and "new_answer" in evaluation_result:
|
||||
# Task is valid. Update the answer with the refined one from the LLM.
|
||||
task['answer'] = evaluation_result['new_answer']
|
||||
task['sql_execute_result'] = task.pop('sql_result')
|
||||
final_validated_tasks.append(task)
|
||||
print(f" - Evaluation PASSED. New answer: {json.dumps(task['answer'])}")
|
||||
else:
|
||||
reason = evaluation_result.get('reason', 'No reason provided.')
|
||||
print(f" - Evaluation FAILED. Filtering task.")
|
||||
print(f" - Reason: {reason}")
|
||||
print(f" - Question: {task['question']}")
|
||||
print(f" - Original Answer: {json.dumps(task['answer'], default=str)}")
|
||||
print(f" - SQL: {task['sql']}")
|
||||
sql_result_str = json.dumps(task['sql_result'], indent=2, default=str)
|
||||
print(f" - SQL Result: {sql_result_str}")
|
||||
|
||||
except Exception as e:
|
||||
print(f" - An error occurred during semantic evaluation for task, filtering it out: {e}")
|
||||
print(f" - Question: {task.get('question', 'N/A')}")
|
||||
print(f" - SQL: {task.get('sql', 'N/A')}")
|
||||
|
||||
return final_validated_tasks
|
||||
|
||||
def main():
|
||||
"""Main function to run the script."""
|
||||
# 1. Load the list of core tables
|
||||
try:
|
||||
with open('core_tables.json', 'r') as f:
|
||||
core_tables = json.load(f)
|
||||
except FileNotFoundError:
|
||||
print("Error: core_tables.json not found. Please create it.")
|
||||
return
|
||||
|
||||
# 2. Connect to the database
|
||||
conn = get_db_connection()
|
||||
if not conn:
|
||||
return
|
||||
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 3. Setup OpenAI Client
|
||||
if not OPENAI_CONFIG["api_key"]:
|
||||
print("Error: OPENAI_API_KEY environment variable not set.")
|
||||
return
|
||||
client = OpenAI(api_key=OPENAI_CONFIG["api_key"], base_url=OPENAI_CONFIG["base_url"])
|
||||
|
||||
try:
|
||||
# 4. Get full schema context
|
||||
print("Fetching full database schema...")
|
||||
schema_context = get_full_schema(cursor, core_tables)
|
||||
|
||||
# 5. Get random samples and print them
|
||||
print("Sampling data from 5 random tables...")
|
||||
sampled_data = get_random_tables_and_samples(cursor, core_tables, num_tables=5, num_samples=5)
|
||||
print(f"Sampled from tables: {list(sampled_data.keys())}")
|
||||
print("\n--- Sampled Data ---")
|
||||
print(json.dumps(sampled_data, indent=2, default=str))
|
||||
print("---------------------\n")
|
||||
|
||||
# 6. Generate questions using the LLM
|
||||
print("Generating questions with GPT-4o...")
|
||||
generated_tasks = generate_questions(client, schema_context, sampled_data)
|
||||
|
||||
# 7. Execute SQL for all generated tasks
|
||||
tasks_for_evaluation = []
|
||||
if generated_tasks:
|
||||
print("\nExecuting SQL for generated tasks...")
|
||||
for task in generated_tasks:
|
||||
if not isinstance(task, dict) or not all(k in task for k in ['sql', 'answer', 'question']):
|
||||
print(f"Filtering task due to malformed structure or missing keys: {task}")
|
||||
continue
|
||||
|
||||
try:
|
||||
cursor.execute(task['sql'])
|
||||
sql_result = cursor.fetchall()
|
||||
task['sql_result'] = sql_result
|
||||
tasks_for_evaluation.append(task)
|
||||
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Filtering task due to SQL error: {err}")
|
||||
print(f" - Question: {task['question']}")
|
||||
print(f" - SQL: {task['sql']}")
|
||||
except Exception as e:
|
||||
print(f"An unexpected error occurred during SQL execution for task {task}: {e}")
|
||||
|
||||
# 8. Semantic evaluation and answer refinement
|
||||
validated_tasks = evaluate_and_refine_tasks(tasks_for_evaluation, client)
|
||||
|
||||
# 9. Print the final JSON output
|
||||
if validated_tasks:
|
||||
print("\n--- Final Validated Tasks ---")
|
||||
print(json.dumps(validated_tasks, indent=2, default=str))
|
||||
else:
|
||||
print("Failed to generate any valid tasks after all validation steps.")
|
||||
|
||||
finally:
|
||||
# 10. Close the database connection
|
||||
if conn.is_connected():
|
||||
cursor.close()
|
||||
conn.close()
|
||||
print("\nDatabase connection closed.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
408
random_sample/generate_tasks copy 4.py
Normal file
408
random_sample/generate_tasks copy 4.py
Normal file
@@ -0,0 +1,408 @@
|
||||
import os
|
||||
import random
|
||||
import json
|
||||
import mysql.connector
|
||||
import argparse
|
||||
from openai import OpenAI
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# --- Configuration ---
|
||||
load_dotenv()
|
||||
|
||||
MYSQL_CONFIG = {
|
||||
"host": "localhost",
|
||||
"port": "23306",
|
||||
"user": "mcpuser",
|
||||
"password": "StrongPass123!",
|
||||
"database": "magentodb"
|
||||
}
|
||||
|
||||
OPENAI_CONFIG = {
|
||||
"api_key": os.getenv("OPENAI_API_KEY"),
|
||||
"base_url": os.getenv("OPENAI_BASE_URL"),
|
||||
"model": "gpt-4o"
|
||||
}
|
||||
|
||||
# --- Prompt Template ---
|
||||
# This is a carefully engineered prompt to guide the LLM's output.
|
||||
PROMPT_TEMPLATE = """
|
||||
You are an expert database analyst and a creative test case designer for e-commerce web applications.
|
||||
Your goal is to generate realistic administrative tasks that can be solved by a Web Agent navigating an admin panel.
|
||||
|
||||
I will provide you with the following context:
|
||||
1. **Full Database Schema**: A list of `CREATE TABLE` statements for the core tables of a Magento e-commerce platform.
|
||||
2. **Sampled Data**: A JSON object containing 5 random rows of data from 5 randomly selected core tables. This data is REAL and should be used to inspire specific, answerable questions.
|
||||
|
||||
## Your Task
|
||||
|
||||
Based on the provided schema and sample data, create a JSON object containing a single key, "questions", which holds an array of up to 10 unique task objects.
|
||||
|
||||
### Requirements for Each Question:
|
||||
- **Web Agent Solvable**: The task must represent a realistic action an administrator would perform in a web UI (e.g., "Find all orders for customer X", "Update the stock for product Y", "Approve a pending review").
|
||||
- **Grounded in Data**: The questions should be specific, using names, IDs, or values from the provided **Sampled Data** to make them concrete.
|
||||
- **Utilize Schema**: You can formulate questions that require joining tables, even if not all tables were sampled. The full schema is your guide.
|
||||
|
||||
### Output Format
|
||||
The final output MUST be a single, valid JSON object. Do not include any other text, explanations, or markdown formatting like ```json.
|
||||
The JSON object must have one key: "questions", containing a JSON array of task objects.
|
||||
|
||||
Each object in the array must contain exactly three keys: `question`, `answer`, and `sql`.
|
||||
|
||||
- **`question`**: (string) A natural language description of the task for a web agent.
|
||||
- **`answer`**: (string, integer, float, or list) The precise and concise answer to the question, derived by running the SQL query against the database.
|
||||
- **`sql`**: (string) The exact, runnable MySQL query that was used to find the answer.
|
||||
|
||||
### Output Format Example
|
||||
```json
|
||||
{{
|
||||
"questions": [
|
||||
{{
|
||||
"question": "What is the email address for customer with ID 5?",
|
||||
"answer": "customer5@example.com",
|
||||
"sql": "SELECT email FROM customer_entity WHERE entity_id = 5;"
|
||||
}},
|
||||
{{
|
||||
"question": "Find the total quantity of item with SKU 'ABC-123' in the cart.",
|
||||
"answer": 3,
|
||||
"sql": "SELECT SUM(qty) FROM quote_item WHERE sku = 'ABC-123';"
|
||||
}}
|
||||
]
|
||||
}}
|
||||
```
|
||||
|
||||
---
|
||||
### Full Database Schema
|
||||
{schema_context}
|
||||
|
||||
---
|
||||
### Sampled Data
|
||||
Here is the sample data from randomly selected tables. Use this to make your questions specific.
|
||||
|
||||
{sampled_data_str}
|
||||
|
||||
---
|
||||
Now, generate the JSON object based on these instructions.
|
||||
"""
|
||||
|
||||
# This is a new prompt to evaluate results and generate a corrected answer.
|
||||
SEMANTIC_EVALUATION_PROMPT_TEMPLATE = """
|
||||
You are a precise data analyst. Your task is to evaluate if a SQL query's result adequately answers a given natural language question. If it does, you must formulate a concise, natural-language answer.
|
||||
|
||||
I will provide you with a JSON object containing:
|
||||
1. `question`: The original question asked.
|
||||
2. `sql`: The SQL query that was executed.
|
||||
3. `sql_result`: The actual data returned by executing the SQL query.
|
||||
|
||||
## Your Task
|
||||
1. **Analyze**: Determine if the `sql_result` contains the necessary information to definitively answer the `question`.
|
||||
2. **Respond**: Based on your analysis, generate a JSON object with one of two structures.
|
||||
|
||||
### Case 1: The question CAN be answered
|
||||
If the `sql_result` provides a clear answer, respond with:
|
||||
```json
|
||||
{{
|
||||
"can_answer": true,
|
||||
"new_answer": "..."
|
||||
}}
|
||||
```
|
||||
- `can_answer`: (boolean) Must be `true`.
|
||||
- `new_answer`: (string, integer, float, or list) A concise, human-readable answer derived *only* from the `sql_result`. For example, if the result is `[(52.00,)]`, the answer can be "52.00" or 52.00.
|
||||
|
||||
### Case 2: The question CANNOT be answered
|
||||
If the `sql_result` is empty, irrelevant, or insufficient to answer the question, respond with:
|
||||
```json
|
||||
{{
|
||||
"can_answer": false,
|
||||
"reason": "..."
|
||||
}}
|
||||
```
|
||||
- `can_answer`: (boolean) Must be `false`.
|
||||
- `reason`: (string) A brief explanation for why the question cannot be answered from the given data (e.g., "The query returned no results.", "The result contains internal IDs, not the requested customer names.").
|
||||
|
||||
---
|
||||
### Evaluation Data
|
||||
{task_data_json}
|
||||
---
|
||||
|
||||
Now, provide your evaluation as a JSON object.
|
||||
"""
|
||||
|
||||
def get_db_connection():
|
||||
"""Establishes a connection to the MySQL database."""
|
||||
try:
|
||||
conn = mysql.connector.connect(**MYSQL_CONFIG)
|
||||
return conn
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Error connecting to MySQL: {err}")
|
||||
return None
|
||||
|
||||
def get_full_schema(cursor, tables):
|
||||
"""Fetches the CREATE TABLE statements for all core tables."""
|
||||
schema_parts = []
|
||||
for table_name in tables:
|
||||
try:
|
||||
cursor.execute(f"SHOW CREATE TABLE `{table_name}`")
|
||||
result = cursor.fetchone()
|
||||
if result:
|
||||
schema_parts.append(result[1]) # result[1] is the CREATE TABLE statement
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Warning: Could not get schema for table {table_name}: {err}")
|
||||
return "\n\n".join(schema_parts)
|
||||
|
||||
def get_random_tables_and_samples(cursor, tables, num_tables=5, num_samples=5):
|
||||
"""Selects random tables and samples random rows from them."""
|
||||
selected_tables = random.sample(tables, num_tables)
|
||||
sampled_data = {}
|
||||
|
||||
for table_name in selected_tables:
|
||||
try:
|
||||
# Use ORDER BY RAND() for random sampling. Can be slow on very large tables.
|
||||
query = f"SELECT * FROM `{table_name}` ORDER BY RAND() LIMIT {num_samples}"
|
||||
cursor.execute(query)
|
||||
|
||||
rows = cursor.fetchall()
|
||||
if not rows:
|
||||
sampled_data[table_name] = []
|
||||
continue
|
||||
|
||||
columns = [desc[0] for desc in cursor.description]
|
||||
|
||||
# Convert rows (tuples) to a list of dictionaries
|
||||
sampled_rows = []
|
||||
for row in rows:
|
||||
row_dict = {}
|
||||
for i, col_value in enumerate(row):
|
||||
# Handle bytes by decoding, fall back to string representation
|
||||
if isinstance(col_value, bytes):
|
||||
try:
|
||||
row_dict[columns[i]] = col_value.decode('utf-8')
|
||||
except UnicodeDecodeError:
|
||||
row_dict[columns[i]] = str(col_value)
|
||||
else:
|
||||
row_dict[columns[i]] = col_value
|
||||
sampled_rows.append(row_dict)
|
||||
|
||||
sampled_data[table_name] = sampled_rows
|
||||
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Warning: Could not sample data from table {table_name}: {err}")
|
||||
sampled_data[table_name] = f"Error: {err}"
|
||||
|
||||
return sampled_data
|
||||
|
||||
def generate_questions(client, schema_context, sampled_data):
|
||||
"""Generates questions by calling the OpenAI API."""
|
||||
if not client:
|
||||
raise ValueError("OpenAI client not provided.")
|
||||
|
||||
sampled_data_str = json.dumps(sampled_data, indent=2, default=str)
|
||||
|
||||
prompt = PROMPT_TEMPLATE.format(
|
||||
schema_context=schema_context,
|
||||
sampled_data_str=sampled_data_str
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.chat.completions.create(
|
||||
model=OPENAI_CONFIG["model"],
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant designed to output JSON."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0.7,
|
||||
response_format={"type": "json_object"},
|
||||
)
|
||||
content = response.choices[0].message.content
|
||||
data = json.loads(content)
|
||||
|
||||
# The prompt asks for {"questions": [...]}, so we extract the list.
|
||||
if isinstance(data, dict) and "questions" in data and isinstance(data["questions"], list):
|
||||
return data["questions"]
|
||||
elif isinstance(data, list):
|
||||
# Fallback in case the model returns a list directly
|
||||
print("Warning: Model returned a raw list instead of an object with a 'questions' key.")
|
||||
return data
|
||||
else:
|
||||
print(f"Warning: Failed to find a 'questions' list in the model's output. Got: {content}")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error calling OpenAI API or parsing JSON: {e}")
|
||||
return None
|
||||
|
||||
def load_existing_tasks(filepath):
|
||||
"""Loads tasks from a JSON file if it exists."""
|
||||
if not os.path.exists(filepath):
|
||||
return []
|
||||
try:
|
||||
with open(filepath, 'r') as f:
|
||||
content = f.read()
|
||||
if not content: # Handle empty file
|
||||
return []
|
||||
return json.loads(content)
|
||||
except (json.JSONDecodeError, FileNotFoundError):
|
||||
print(f"Warning: Could not read or parse {filepath}. Starting with an empty list.")
|
||||
return []
|
||||
|
||||
def evaluate_and_refine_tasks(tasks, client):
|
||||
"""
|
||||
Uses an LLM to evaluate if a SQL result answers the question and refines the answer.
|
||||
"""
|
||||
if not tasks:
|
||||
return []
|
||||
|
||||
final_validated_tasks = []
|
||||
print("\nPerforming semantic evaluation and answer refinement with GPT-4o...")
|
||||
|
||||
for task in tasks:
|
||||
# Prepare data for the prompt, excluding the original 'answer'
|
||||
task_data_for_prompt = {
|
||||
"question": task["question"],
|
||||
"sql": task["sql"],
|
||||
"sql_result": task["sql_result"]
|
||||
}
|
||||
task_data_json = json.dumps(task_data_for_prompt, indent=2, default=str)
|
||||
|
||||
prompt = SEMANTIC_EVALUATION_PROMPT_TEMPLATE.format(task_data_json=task_data_json)
|
||||
|
||||
try:
|
||||
print(f" - Evaluating question: \"{task['question'][:80]}...\"")
|
||||
response = client.chat.completions.create(
|
||||
model=OPENAI_CONFIG["model"],
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant designed to output JSON."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0.0, # We want deterministic evaluation
|
||||
response_format={"type": "json_object"},
|
||||
)
|
||||
content = response.choices[0].message.content
|
||||
evaluation_result = json.loads(content)
|
||||
|
||||
if evaluation_result.get("can_answer") is True and "new_answer" in evaluation_result:
|
||||
# Task is valid. Update the answer with the refined one from the LLM.
|
||||
task['answer'] = evaluation_result['new_answer']
|
||||
task['sql_execute_result'] = task.pop('sql_result')
|
||||
final_validated_tasks.append(task)
|
||||
print(f" - Evaluation PASSED. New answer: {json.dumps(task['answer'])}")
|
||||
else:
|
||||
reason = evaluation_result.get('reason', 'No reason provided.')
|
||||
print(f" - Evaluation FAILED. Filtering task.")
|
||||
print(f" - Reason: {reason}")
|
||||
print(f" - Question: {task['question']}")
|
||||
print(f" - Original Answer: {json.dumps(task['answer'], default=str)}")
|
||||
print(f" - SQL: {task['sql']}")
|
||||
sql_result_str = json.dumps(task['sql_result'], indent=2, default=str)
|
||||
print(f" - SQL Result: {sql_result_str}")
|
||||
|
||||
except Exception as e:
|
||||
print(f" - An error occurred during semantic evaluation for task, filtering it out: {e}")
|
||||
print(f" - Question: {task.get('question', 'N/A')}")
|
||||
print(f" - SQL: {task.get('sql', 'N/A')}")
|
||||
|
||||
return final_validated_tasks
|
||||
|
||||
def main():
|
||||
"""Main function to run the script."""
|
||||
parser = argparse.ArgumentParser(description="Generate and validate e-commerce admin tasks.")
|
||||
parser.add_argument(
|
||||
"--target-count",
|
||||
type=int,
|
||||
required=True,
|
||||
help="The total number of questions to generate."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-file",
|
||||
type=str,
|
||||
default="generated_tasks.json",
|
||||
help="The file to save the generated tasks to (in JSON format)."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Load existing tasks from the output file
|
||||
all_tasks = load_existing_tasks(args.output_file)
|
||||
print(f"Found {len(all_tasks)} existing valid tasks in '{args.output_file}'.")
|
||||
|
||||
# Connect to DB and set up client
|
||||
conn = get_db_connection()
|
||||
if not conn:
|
||||
return
|
||||
cursor = conn.cursor()
|
||||
|
||||
if not OPENAI_CONFIG["api_key"]:
|
||||
print("Error: OPENAI_API_KEY environment variable not set.")
|
||||
return
|
||||
client = OpenAI(api_key=OPENAI_CONFIG["api_key"], base_url=OPENAI_CONFIG["base_url"])
|
||||
|
||||
try:
|
||||
# Load core tables and schema once
|
||||
try:
|
||||
with open('core_tables.json', 'r') as f:
|
||||
core_tables = json.load(f)
|
||||
except FileNotFoundError:
|
||||
print("Error: core_tables.json not found. Please create it.")
|
||||
return
|
||||
|
||||
print("Fetching full database schema...")
|
||||
schema_context = get_full_schema(cursor, core_tables)
|
||||
|
||||
# Start the generation loop
|
||||
round_num = 1
|
||||
while len(all_tasks) < args.target_count:
|
||||
print(f"\n--- Starting Generation Round {round_num} ---")
|
||||
print(f"Goal: {args.target_count} | Current: {len(all_tasks)} | Needed: {args.target_count - len(all_tasks)}")
|
||||
|
||||
# Get random samples for this round
|
||||
print("Sampling data from 5 random tables...")
|
||||
sampled_data = get_random_tables_and_samples(cursor, core_tables, num_tables=5, num_samples=5)
|
||||
|
||||
# Generate questions
|
||||
print("Generating questions with GPT-4o...")
|
||||
generated_tasks = generate_questions(client, schema_context, sampled_data)
|
||||
|
||||
# Execute SQL for generated tasks
|
||||
tasks_for_evaluation = []
|
||||
if generated_tasks:
|
||||
print("\nExecuting SQL for generated tasks...")
|
||||
for task in generated_tasks:
|
||||
if not isinstance(task, dict) or not all(k in task for k in ['sql', 'answer', 'question']):
|
||||
print(f"Filtering task due to malformed structure: {task}")
|
||||
continue
|
||||
try:
|
||||
cursor.execute(task['sql'])
|
||||
sql_result = cursor.fetchall()
|
||||
task['sql_result'] = sql_result
|
||||
tasks_for_evaluation.append(task)
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Filtering task due to SQL error: {err} on SQL: {task['sql']}")
|
||||
|
||||
# Perform semantic evaluation and get validated tasks
|
||||
validated_tasks = evaluate_and_refine_tasks(tasks_for_evaluation, client)
|
||||
|
||||
# Append new tasks and save to file
|
||||
if validated_tasks:
|
||||
all_tasks.extend(validated_tasks)
|
||||
with open(args.output_file, 'w') as f:
|
||||
json.dump(all_tasks, f, indent=2, default=str)
|
||||
|
||||
print("\n--- Round Summary ---")
|
||||
print(f"Generated {len(validated_tasks)} new valid tasks in this round.")
|
||||
print(f"Progress: {len(all_tasks)} / {args.target_count} tasks.")
|
||||
else:
|
||||
print("\n--- Round Summary ---")
|
||||
print("No new valid tasks were generated in this round. Retrying...")
|
||||
|
||||
round_num += 1
|
||||
|
||||
finally:
|
||||
# Close the database connection
|
||||
if conn.is_connected():
|
||||
cursor.close()
|
||||
conn.close()
|
||||
print("\nDatabase connection closed.")
|
||||
|
||||
print(f"\nTarget of {args.target_count} tasks reached. Final output saved to {args.output_file}.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
437
random_sample/generate_tasks copy 5.py
Normal file
437
random_sample/generate_tasks copy 5.py
Normal file
@@ -0,0 +1,437 @@
|
||||
import os
|
||||
import random
|
||||
import json
|
||||
import mysql.connector
|
||||
import argparse
|
||||
from openai import OpenAI
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# --- Configuration ---
|
||||
load_dotenv()
|
||||
|
||||
MYSQL_CONFIG = {
|
||||
"host": "localhost",
|
||||
"port": "23306",
|
||||
"user": "mcpuser",
|
||||
"password": "StrongPass123!",
|
||||
"database": "magentodb"
|
||||
}
|
||||
|
||||
OPENAI_CONFIG = {
|
||||
"api_key": os.getenv("OPENAI_API_KEY"),
|
||||
"base_url": os.getenv("OPENAI_BASE_URL"),
|
||||
"model": "gpt-4o"
|
||||
}
|
||||
|
||||
# --- Prompt Template ---
|
||||
# This is a carefully engineered prompt to guide the LLM's output.
|
||||
PROMPT_TEMPLATE = """
|
||||
You are an expert database analyst and a creative test case designer for e-commerce web applications.
|
||||
Your goal is to generate realistic administrative tasks that can be solved by a Web Agent navigating an admin panel.
|
||||
|
||||
I will provide you with the following context:
|
||||
1. **Full Database Schema**: A list of `CREATE TABLE` statements for the core tables of a Magento e-commerce platform.
|
||||
2. **Sampled Data**: A JSON object containing 5 random rows of data from 5 randomly selected core tables. This data is REAL and should be used to inspire specific, answerable questions.
|
||||
|
||||
## Your Task
|
||||
|
||||
Based on the provided schema and sample data, create a JSON object containing a single key, "questions", which holds an array of up to 10 unique task objects.
|
||||
|
||||
### Requirements for Each Question:
|
||||
- **Web Agent Solvable**: The task must represent a realistic action an administrator would perform in a web UI (e.g., "Find all orders for customer X", "Update the stock for product Y", "Approve a pending review").
|
||||
- **Grounded in Data**: The questions should be specific, using names, IDs, or values from the provided **Sampled Data** to make them concrete.
|
||||
- **Utilize Schema**: You can formulate questions that require joining tables, even if not all tables were sampled. The full schema is your guide.
|
||||
|
||||
### Output Format
|
||||
The final output MUST be a single, valid JSON object. Do not include any other text, explanations, or markdown formatting like ```json.
|
||||
The JSON object must have one key: "questions", containing a JSON array of task objects.
|
||||
|
||||
Each object in the array must contain exactly three keys: `question`, `answer`, and `sql`.
|
||||
|
||||
- **`question`**: (string) A natural language description of the task for a web agent.
|
||||
- **`answer`**: (string, integer, float, or list) The precise and concise answer to the question, derived by running the SQL query against the database.
|
||||
- **`sql`**: (string) The exact, runnable MySQL query that was used to find the answer.
|
||||
|
||||
### Output Format Example
|
||||
```json
|
||||
{{
|
||||
"questions": [
|
||||
{{
|
||||
"question": "What is the email address for customer with ID 5?",
|
||||
"answer": "customer5@example.com",
|
||||
"sql": "SELECT email FROM customer_entity WHERE entity_id = 5;"
|
||||
}},
|
||||
{{
|
||||
"question": "Find the total quantity of item with SKU 'ABC-123' in the cart.",
|
||||
"answer": 3,
|
||||
"sql": "SELECT SUM(qty) FROM quote_item WHERE sku = 'ABC-123';"
|
||||
}}
|
||||
]
|
||||
}}
|
||||
```
|
||||
|
||||
---
|
||||
### Full Database Schema
|
||||
{schema_context}
|
||||
|
||||
---
|
||||
### Sampled Data
|
||||
Here is the sample data from randomly selected tables. Use this to make your questions specific.
|
||||
|
||||
{sampled_data_str}
|
||||
|
||||
---
|
||||
Now, generate the JSON object based on these instructions.
|
||||
"""
|
||||
|
||||
# This is a new prompt to evaluate results and generate a corrected answer.
|
||||
SEMANTIC_EVALUATION_PROMPT_TEMPLATE = """
|
||||
You are a precise data analyst. Your task is to evaluate if a SQL query's result adequately answers a given natural language question. You will then either refine the answer, or completely rephrase the question if the result set is large.
|
||||
|
||||
I will provide you with a JSON object containing:
|
||||
1. `question`: The original question asked.
|
||||
2. `sql`: The SQL query that was executed.
|
||||
3. `sql_result`: The actual data returned by executing the SQL query.
|
||||
4. `row_count`: The number of rows in `sql_result`.
|
||||
|
||||
## Your Task
|
||||
Analyze the inputs and respond with a JSON object. You have three cases. The `new_answer` field MUST always be an array of strings.
|
||||
|
||||
### Case 1: Large Result Set (Question Transformation)
|
||||
If `row_count` is greater than 10 AND the original `question` does NOT already ask for a count (e.g., it is not phrased like "How many..."), you must transform the question.
|
||||
Respond with:
|
||||
```json
|
||||
{{
|
||||
"can_answer": true,
|
||||
"new_question": "How many items were found?",
|
||||
"new_answer": ["42"]
|
||||
}}
|
||||
```
|
||||
- `can_answer`: (boolean) Must be `true`.
|
||||
- `new_question`: (string) A rephrased question that asks for the quantity of items. For example, if the original question was "List all products", the new question should be "How many products were found?".
|
||||
- `new_answer`: (array of strings) An array containing the `row_count` as a single string element.
|
||||
|
||||
### Case 2: Standard Answer (No Transformation)
|
||||
If Case 1 does not apply, but the `sql_result` still provides a clear answer to the original `question`, respond with:
|
||||
```json
|
||||
{{
|
||||
"can_answer": true,
|
||||
"new_answer": ["value1", "value2", ...]
|
||||
}}
|
||||
```
|
||||
- `can_answer`: (boolean) Must be `true`.
|
||||
- `new_answer`: (array of strings) An array containing all the essential parts of the answer extracted from `sql_result`. Every value from the result set that contributes to the answer should be included as a string in the array. This ensures answer completeness.
|
||||
- **Example 1**: If `question` is "What is the status of order 123?" and `sql_result` is `[["processing"]]`, `new_answer` should be `["processing"]`.
|
||||
- **Example 2**: If `question` is "List emails for pending customers" and `sql_result` is `[["test@a.com"], ["test@b.com"]]`, `new_answer` should be `["test@a.com", "test@b.com"]`.
|
||||
- **Example 3**: If `question` is "Get product name and price for SKU 'XYZ'" and `sql_result` is `[["My Product", 19.99]]`, `new_answer` should be `["My Product", "19.99"]`.
|
||||
|
||||
### Case 3: The question CANNOT be answered
|
||||
If the `sql_result` is empty, irrelevant, or insufficient to answer the question, respond with:
|
||||
```json
|
||||
{{
|
||||
"can_answer": false,
|
||||
"reason": "..."
|
||||
}}
|
||||
```
|
||||
- `can_answer`: (boolean) Must be `false`.
|
||||
- `reason`: (string) A brief explanation for why the question cannot be answered.
|
||||
|
||||
---
|
||||
### Evaluation Data
|
||||
{task_data_json}
|
||||
---
|
||||
|
||||
Now, provide your evaluation as a JSON object.
|
||||
"""
|
||||
|
||||
def get_db_connection():
|
||||
"""Establishes a connection to the MySQL database."""
|
||||
try:
|
||||
conn = mysql.connector.connect(**MYSQL_CONFIG)
|
||||
return conn
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Error connecting to MySQL: {err}")
|
||||
return None
|
||||
|
||||
def get_full_schema(cursor, tables):
|
||||
"""Fetches the CREATE TABLE statements for all core tables."""
|
||||
schema_parts = []
|
||||
for table_name in tables:
|
||||
try:
|
||||
cursor.execute(f"SHOW CREATE TABLE `{table_name}`")
|
||||
result = cursor.fetchone()
|
||||
if result:
|
||||
schema_parts.append(result[1]) # result[1] is the CREATE TABLE statement
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Warning: Could not get schema for table {table_name}: {err}")
|
||||
return "\n\n".join(schema_parts)
|
||||
|
||||
def get_random_tables_and_samples(cursor, tables, num_tables=5, num_samples=5):
|
||||
"""Selects random tables and samples random rows from them."""
|
||||
selected_tables = random.sample(tables, num_tables)
|
||||
sampled_data = {}
|
||||
|
||||
for table_name in selected_tables:
|
||||
try:
|
||||
# Use ORDER BY RAND() for random sampling. Can be slow on very large tables.
|
||||
query = f"SELECT * FROM `{table_name}` ORDER BY RAND() LIMIT {num_samples}"
|
||||
cursor.execute(query)
|
||||
|
||||
rows = cursor.fetchall()
|
||||
if not rows:
|
||||
sampled_data[table_name] = []
|
||||
continue
|
||||
|
||||
columns = [desc[0] for desc in cursor.description]
|
||||
|
||||
# Convert rows (tuples) to a list of dictionaries
|
||||
sampled_rows = []
|
||||
for row in rows:
|
||||
row_dict = {}
|
||||
for i, col_value in enumerate(row):
|
||||
# Handle bytes by decoding, fall back to string representation
|
||||
if isinstance(col_value, bytes):
|
||||
try:
|
||||
row_dict[columns[i]] = col_value.decode('utf-8')
|
||||
except UnicodeDecodeError:
|
||||
row_dict[columns[i]] = str(col_value)
|
||||
else:
|
||||
row_dict[columns[i]] = col_value
|
||||
sampled_rows.append(row_dict)
|
||||
|
||||
sampled_data[table_name] = sampled_rows
|
||||
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Warning: Could not sample data from table {table_name}: {err}")
|
||||
sampled_data[table_name] = f"Error: {err}"
|
||||
|
||||
return sampled_data
|
||||
|
||||
def generate_questions(client, schema_context, sampled_data):
|
||||
"""Generates questions by calling the OpenAI API."""
|
||||
if not client:
|
||||
raise ValueError("OpenAI client not provided.")
|
||||
|
||||
sampled_data_str = json.dumps(sampled_data, indent=2, default=str)
|
||||
|
||||
prompt = PROMPT_TEMPLATE.format(
|
||||
schema_context=schema_context,
|
||||
sampled_data_str=sampled_data_str
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.chat.completions.create(
|
||||
model=OPENAI_CONFIG["model"],
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant designed to output JSON."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0.7,
|
||||
response_format={"type": "json_object"},
|
||||
)
|
||||
content = response.choices[0].message.content
|
||||
data = json.loads(content)
|
||||
|
||||
# The prompt asks for {"questions": [...]}, so we extract the list.
|
||||
if isinstance(data, dict) and "questions" in data and isinstance(data["questions"], list):
|
||||
return data["questions"]
|
||||
elif isinstance(data, list):
|
||||
# Fallback in case the model returns a list directly
|
||||
print("Warning: Model returned a raw list instead of an object with a 'questions' key.")
|
||||
return data
|
||||
else:
|
||||
print(f"Warning: Failed to find a 'questions' list in the model's output. Got: {content}")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error calling OpenAI API or parsing JSON: {e}")
|
||||
return None
|
||||
|
||||
def load_existing_tasks(filepath):
|
||||
"""Loads tasks from a JSON file if it exists."""
|
||||
if not os.path.exists(filepath):
|
||||
return []
|
||||
try:
|
||||
with open(filepath, 'r') as f:
|
||||
content = f.read()
|
||||
if not content: # Handle empty file
|
||||
return []
|
||||
return json.loads(content)
|
||||
except (json.JSONDecodeError, FileNotFoundError):
|
||||
print(f"Warning: Could not read or parse {filepath}. Starting with an empty list.")
|
||||
return []
|
||||
|
||||
def evaluate_and_refine_tasks(tasks, client):
|
||||
"""
|
||||
Uses an LLM to evaluate if a SQL result answers the question and refines the answer.
|
||||
"""
|
||||
if not tasks:
|
||||
return []
|
||||
|
||||
final_validated_tasks = []
|
||||
print("\nPerforming semantic evaluation and answer refinement with GPT-4o...")
|
||||
|
||||
for task in tasks:
|
||||
# Prepare data for the prompt, excluding the original 'answer'
|
||||
task_data_for_prompt = {
|
||||
"question": task["question"],
|
||||
"sql": task["sql"],
|
||||
"sql_result": task["sql_result"],
|
||||
"row_count": task["row_count"]
|
||||
}
|
||||
task_data_json = json.dumps(task_data_for_prompt, indent=2, default=str)
|
||||
|
||||
prompt = SEMANTIC_EVALUATION_PROMPT_TEMPLATE.format(task_data_json=task_data_json)
|
||||
|
||||
try:
|
||||
print(f" - Evaluating question: \"{task['question'][:80]}...\"")
|
||||
response = client.chat.completions.create(
|
||||
model=OPENAI_CONFIG["model"],
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant designed to output JSON."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0.0, # We want deterministic evaluation
|
||||
response_format={"type": "json_object"},
|
||||
)
|
||||
content = response.choices[0].message.content
|
||||
evaluation_result = json.loads(content)
|
||||
|
||||
if evaluation_result.get("can_answer") is True and "new_answer" in evaluation_result:
|
||||
# Task is valid. Update the answer with the refined one from the LLM.
|
||||
task['answer'] = evaluation_result['new_answer']
|
||||
|
||||
# If the LLM provides a new question, update it.
|
||||
if 'new_question' in evaluation_result:
|
||||
task['question'] = evaluation_result['new_question']
|
||||
print(f" - Question was rephrased: \"{task['question']}\"")
|
||||
|
||||
task['sql_execute_result'] = task.pop('sql_result')
|
||||
task.pop('row_count', None) # Clean up temp key
|
||||
final_validated_tasks.append(task)
|
||||
print(f" - Evaluation PASSED. New answer: {json.dumps(task['answer'])}")
|
||||
else:
|
||||
reason = evaluation_result.get('reason', 'No reason provided.')
|
||||
print(f" - Evaluation FAILED. Filtering task.")
|
||||
print(f" - Reason: {reason}")
|
||||
print(f" - Question: {task['question']}")
|
||||
print(f" - SQL: {task['sql']}")
|
||||
sql_result_str = json.dumps(task['sql_result'], indent=2, default=str)
|
||||
print(f" - SQL Result: {sql_result_str}")
|
||||
|
||||
except Exception as e:
|
||||
print(f" - An error occurred during semantic evaluation for task, filtering it out: {e}")
|
||||
print(f" - Question: {task.get('question', 'N/A')}")
|
||||
print(f" - SQL: {task.get('sql', 'N/A')}")
|
||||
|
||||
return final_validated_tasks
|
||||
|
||||
def main():
|
||||
"""Main function to run the script."""
|
||||
parser = argparse.ArgumentParser(description="Generate and validate e-commerce admin tasks.")
|
||||
parser.add_argument(
|
||||
"--target-count",
|
||||
type=int,
|
||||
required=True,
|
||||
help="The total number of questions to generate."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-file",
|
||||
type=str,
|
||||
default="generated_tasks.json",
|
||||
help="The file to save the generated tasks to (in JSON format)."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Load existing tasks from the output file
|
||||
all_tasks = load_existing_tasks(args.output_file)
|
||||
print(f"Found {len(all_tasks)} existing valid tasks in '{args.output_file}'.")
|
||||
|
||||
# Connect to DB and set up client
|
||||
conn = get_db_connection()
|
||||
if not conn:
|
||||
return
|
||||
cursor = conn.cursor()
|
||||
|
||||
if not OPENAI_CONFIG["api_key"]:
|
||||
print("Error: OPENAI_API_KEY environment variable not set.")
|
||||
return
|
||||
client = OpenAI(api_key=OPENAI_CONFIG["api_key"], base_url=OPENAI_CONFIG["base_url"])
|
||||
|
||||
try:
|
||||
# Load core tables and schema once
|
||||
try:
|
||||
with open('core_tables.json', 'r') as f:
|
||||
core_tables = json.load(f)
|
||||
except FileNotFoundError:
|
||||
print("Error: core_tables.json not found. Please create it.")
|
||||
return
|
||||
|
||||
print("Fetching full database schema...")
|
||||
schema_context = get_full_schema(cursor, core_tables)
|
||||
|
||||
# Start the generation loop
|
||||
round_num = 1
|
||||
while len(all_tasks) < args.target_count:
|
||||
print(f"\n--- Starting Generation Round {round_num} ---")
|
||||
print(f"Goal: {args.target_count} | Current: {len(all_tasks)} | Needed: {args.target_count - len(all_tasks)}")
|
||||
|
||||
# Get random samples for this round
|
||||
print("Sampling data from 5 random tables...")
|
||||
sampled_data = get_random_tables_and_samples(cursor, core_tables, num_tables=5, num_samples=5)
|
||||
|
||||
# Generate questions
|
||||
print("Generating questions with GPT-4o...")
|
||||
generated_tasks = generate_questions(client, schema_context, sampled_data)
|
||||
|
||||
# Execute SQL for generated tasks
|
||||
tasks_for_evaluation = []
|
||||
if generated_tasks:
|
||||
print("\nExecuting SQL for generated tasks...")
|
||||
for task in generated_tasks:
|
||||
if not isinstance(task, dict) or not all(k in task for k in ['sql', 'answer', 'question']):
|
||||
print(f"Filtering task due to malformed structure: {task}")
|
||||
continue
|
||||
try:
|
||||
cursor.execute(task['sql'])
|
||||
sql_result = cursor.fetchall()
|
||||
# Create a new dict for evaluation, excluding the original 'answer'.
|
||||
tasks_for_evaluation.append({
|
||||
'question': task['question'],
|
||||
'sql': task['sql'],
|
||||
'sql_result': sql_result,
|
||||
'row_count': len(sql_result)
|
||||
})
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Filtering task due to SQL error: {err} on SQL: {task['sql']}")
|
||||
|
||||
# Perform semantic evaluation and get validated tasks
|
||||
validated_tasks = evaluate_and_refine_tasks(tasks_for_evaluation, client)
|
||||
|
||||
# Append new tasks and save to file
|
||||
if validated_tasks:
|
||||
all_tasks.extend(validated_tasks)
|
||||
with open(args.output_file, 'w') as f:
|
||||
json.dump(all_tasks, f, indent=2, default=str)
|
||||
|
||||
print("\n--- Round Summary ---")
|
||||
print(f"Generated {len(validated_tasks)} new valid tasks in this round.")
|
||||
print(f"Progress: {len(all_tasks)} / {args.target_count} tasks.")
|
||||
else:
|
||||
print("\n--- Round Summary ---")
|
||||
print("No new valid tasks were generated in this round. Retrying...")
|
||||
|
||||
round_num += 1
|
||||
|
||||
finally:
|
||||
# Close the database connection
|
||||
if conn.is_connected():
|
||||
cursor.close()
|
||||
conn.close()
|
||||
print("\nDatabase connection closed.")
|
||||
|
||||
print(f"\nTarget of {args.target_count} tasks reached. Final output saved to {args.output_file}.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
246
random_sample/generate_tasks copy.py
Normal file
246
random_sample/generate_tasks copy.py
Normal file
@@ -0,0 +1,246 @@
|
||||
import os
|
||||
import random
|
||||
import json
|
||||
import mysql.connector
|
||||
from openai import OpenAI
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# --- Configuration ---
|
||||
load_dotenv()
|
||||
|
||||
MYSQL_CONFIG = {
|
||||
"host": "localhost",
|
||||
"port": "23306",
|
||||
"user": "mcpuser",
|
||||
"password": "StrongPass123!",
|
||||
"database": "magentodb"
|
||||
}
|
||||
|
||||
OPENAI_CONFIG = {
|
||||
"api_key": os.getenv("OPENAI_API_KEY"),
|
||||
"base_url": os.getenv("OPENAI_BASE_URL"),
|
||||
"model": "gpt-4o"
|
||||
}
|
||||
|
||||
# --- Prompt Template ---
|
||||
# This is a carefully engineered prompt to guide the LLM's output.
|
||||
PROMPT_TEMPLATE = """
|
||||
You are an expert database analyst and a creative test case designer for e-commerce web applications.
|
||||
Your goal is to generate realistic administrative tasks that can be solved by a Web Agent navigating an admin panel.
|
||||
|
||||
I will provide you with the following context:
|
||||
1. **Full Database Schema**: A list of `CREATE TABLE` statements for the core tables of a Magento e-commerce platform.
|
||||
2. **Sampled Data**: A JSON object containing 5 random rows of data from 5 randomly selected core tables. This data is REAL and should be used to inspire specific, answerable questions.
|
||||
|
||||
## Your Task
|
||||
|
||||
Based on the provided schema and sample data, create a JSON array of up to 10 unique questions.
|
||||
|
||||
### Requirements for Each Question:
|
||||
- **Web Agent Solvable**: The task must represent a realistic action an administrator would perform in a web UI (e.g., "Find all orders for customer X", "Update the stock for product Y", "Approve a pending review").
|
||||
- **Grounded in Data**: The questions should be specific, using names, IDs, or values from the provided **Sampled Data** to make them concrete.
|
||||
- **Utilize Schema**: You can formulate questions that require joining tables, even if not all tables were sampled. The full schema is your guide.
|
||||
|
||||
### Output Format
|
||||
The final output MUST be a single, valid JSON array of objects. Do not include any other text, explanations, or markdown formatting like ```json.
|
||||
|
||||
Each object in the array must contain exactly three keys: `question`, `answer`, and `sql`.
|
||||
|
||||
- **`question`**: (string) A natural language description of the task for a web agent.
|
||||
- **`answer`**: (string, integer, float, or list) The precise and concise answer to the question, derived by running the SQL query against the database.
|
||||
- **`sql`**: (string) The exact, runnable MySQL query that was used to find the answer.
|
||||
|
||||
---
|
||||
### Full Database Schema
|
||||
{schema_context}
|
||||
|
||||
---
|
||||
### Sampled Data
|
||||
Here is the sample data from randomly selected tables. Use this to make your questions specific.
|
||||
|
||||
{sampled_data_str}
|
||||
|
||||
---
|
||||
Now, generate the JSON array based on these instructions.
|
||||
"""
|
||||
|
||||
def get_db_connection():
|
||||
"""Establishes a connection to the MySQL database."""
|
||||
try:
|
||||
conn = mysql.connector.connect(**MYSQL_CONFIG)
|
||||
return conn
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Error connecting to MySQL: {err}")
|
||||
return None
|
||||
|
||||
def get_full_schema(cursor, tables):
|
||||
"""Fetches the CREATE TABLE statements for all core tables."""
|
||||
schema_parts = []
|
||||
for table_name in tables:
|
||||
try:
|
||||
cursor.execute(f"SHOW CREATE TABLE `{table_name}`")
|
||||
result = cursor.fetchone()
|
||||
if result:
|
||||
schema_parts.append(result[1]) # result[1] is the CREATE TABLE statement
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Warning: Could not get schema for table {table_name}: {err}")
|
||||
return "\n\n".join(schema_parts)
|
||||
|
||||
def get_random_tables_and_samples(cursor, tables, num_tables=5, num_samples=5):
|
||||
"""Selects random tables and samples random rows from them."""
|
||||
selected_tables = random.sample(tables, num_tables)
|
||||
sampled_data = {}
|
||||
|
||||
for table_name in selected_tables:
|
||||
try:
|
||||
# Use ORDER BY RAND() for random sampling. Can be slow on very large tables.
|
||||
query = f"SELECT * FROM `{table_name}` ORDER BY RAND() LIMIT {num_samples}"
|
||||
cursor.execute(query)
|
||||
|
||||
rows = cursor.fetchall()
|
||||
if not rows:
|
||||
sampled_data[table_name] = []
|
||||
continue
|
||||
|
||||
columns = [desc[0] for desc in cursor.description]
|
||||
|
||||
# Convert rows (tuples) to a list of dictionaries
|
||||
sampled_rows = []
|
||||
for row in rows:
|
||||
row_dict = {}
|
||||
for i, col_value in enumerate(row):
|
||||
# Handle bytes by decoding, fall back to string representation
|
||||
if isinstance(col_value, bytes):
|
||||
try:
|
||||
row_dict[columns[i]] = col_value.decode('utf-8')
|
||||
except UnicodeDecodeError:
|
||||
row_dict[columns[i]] = str(col_value)
|
||||
else:
|
||||
row_dict[columns[i]] = col_value
|
||||
sampled_rows.append(row_dict)
|
||||
|
||||
sampled_data[table_name] = sampled_rows
|
||||
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Warning: Could not sample data from table {table_name}: {err}")
|
||||
sampled_data[table_name] = f"Error: {err}"
|
||||
|
||||
return sampled_data
|
||||
|
||||
def generate_questions(schema_context, sampled_data):
|
||||
"""Generates questions by calling the OpenAI API."""
|
||||
if not OPENAI_CONFIG["api_key"]:
|
||||
raise ValueError("OPENAI_API_KEY environment variable not set.")
|
||||
|
||||
client = OpenAI(api_key=OPENAI_CONFIG["api_key"], base_url=OPENAI_CONFIG["base_url"])
|
||||
|
||||
sampled_data_str = json.dumps(sampled_data, indent=2, default=str)
|
||||
|
||||
prompt = PROMPT_TEMPLATE.format(
|
||||
schema_context=schema_context,
|
||||
sampled_data_str=sampled_data_str
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.chat.completions.create(
|
||||
model=OPENAI_CONFIG["model"],
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant designed to output JSON."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0.7,
|
||||
)
|
||||
content = response.choices[0].message.content
|
||||
return json.loads(content)
|
||||
except Exception as e:
|
||||
print(f"Error calling OpenAI API: {e}")
|
||||
return None
|
||||
|
||||
def main():
|
||||
"""Main function to run the script."""
|
||||
# 1. Load the list of core tables
|
||||
try:
|
||||
with open('core_tables.json', 'r') as f:
|
||||
core_tables = json.load(f)
|
||||
except FileNotFoundError:
|
||||
print("Error: core_tables.json not found. Please create it.")
|
||||
return
|
||||
|
||||
# 2. Connect to the database
|
||||
conn = get_db_connection()
|
||||
if not conn:
|
||||
return
|
||||
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
# 3. Get full schema context
|
||||
print("Fetching full database schema...")
|
||||
schema_context = get_full_schema(cursor, core_tables)
|
||||
|
||||
# 4. Get random samples and print them
|
||||
print("Sampling data from 5 random tables...")
|
||||
sampled_data = get_random_tables_and_samples(cursor, core_tables, num_tables=5, num_samples=5)
|
||||
print(f"Sampled from tables: {list(sampled_data.keys())}")
|
||||
print("\n--- Sampled Data ---")
|
||||
print(json.dumps(sampled_data, indent=2, default=str))
|
||||
print("---------------------\n")
|
||||
|
||||
# 5. Generate questions using the LLM
|
||||
print("Generating questions with GPT-4o...")
|
||||
generated_tasks = generate_questions(schema_context, sampled_data)
|
||||
|
||||
# 6. Validate and filter the generated tasks
|
||||
validated_tasks = []
|
||||
if generated_tasks:
|
||||
print("\nValidating generated tasks...")
|
||||
for task in generated_tasks:
|
||||
# Basic validation for task structure
|
||||
if not isinstance(task, dict) or not all(k in task for k in ['sql', 'answer', 'question']):
|
||||
print(f"Filtering task due to malformed structure or missing keys: {task}")
|
||||
continue
|
||||
|
||||
try:
|
||||
# Execute the SQL query from the task
|
||||
cursor.execute(task['sql'])
|
||||
sql_result = cursor.fetchall()
|
||||
|
||||
# Convert both answer and result to string for flexible substring matching
|
||||
answer_str = str(task['answer'])
|
||||
result_str = str(sql_result)
|
||||
|
||||
# If the answer exists in the result, the task is valid
|
||||
if answer_str in result_str:
|
||||
validated_tasks.append(task)
|
||||
else:
|
||||
# Log tasks that are filtered because the answer doesn't match
|
||||
print(f"Filtering task: Answer '{answer_str}' not found in SQL result.")
|
||||
print(f" - Question: {task['question']}")
|
||||
print(f" - SQL: {task['sql']}")
|
||||
# Showing a snippet of a large result is helpful for debugging
|
||||
print(f" - Result: {result_str[:250]}...")
|
||||
|
||||
except mysql.connector.Error as err:
|
||||
# Log tasks that are filtered due to SQL errors
|
||||
print(f"Filtering task due to SQL error: {err}")
|
||||
print(f" - Question: {task['question']}")
|
||||
print(f" - SQL: {task['sql']}")
|
||||
except Exception as e:
|
||||
print(f"An unexpected error occurred during validation for task {task}: {e}")
|
||||
|
||||
# 7. Print the final JSON output
|
||||
if validated_tasks:
|
||||
print("\n--- Generated and Validated Tasks ---")
|
||||
print(json.dumps(validated_tasks, indent=2))
|
||||
else:
|
||||
print("Failed to generate any valid tasks.")
|
||||
|
||||
finally:
|
||||
# 8. Close the database connection
|
||||
if conn.is_connected():
|
||||
cursor.close()
|
||||
conn.close()
|
||||
print("\nDatabase connection closed.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
445
random_sample/generate_tasks.py
Normal file
445
random_sample/generate_tasks.py
Normal file
@@ -0,0 +1,445 @@
|
||||
import os
|
||||
import random
|
||||
import json
|
||||
import mysql.connector
|
||||
import argparse
|
||||
from openai import OpenAI
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# --- Configuration ---
|
||||
load_dotenv()
|
||||
|
||||
MYSQL_CONFIG = {
|
||||
"host": "localhost",
|
||||
"port": "23306",
|
||||
"user": "mcpuser",
|
||||
"password": "StrongPass123!",
|
||||
"database": "magentodb"
|
||||
}
|
||||
|
||||
OPENAI_CONFIG = {
|
||||
"api_key": os.getenv("OPENAI_API_KEY"),
|
||||
"base_url": os.getenv("OPENAI_BASE_URL"),
|
||||
"model": "gpt-4o"
|
||||
}
|
||||
|
||||
# --- Prompt Template ---
|
||||
# This is a carefully engineered prompt to guide the LLM's output.
|
||||
PROMPT_TEMPLATE = """
|
||||
You are an expert database analyst and a creative test case designer for e-commerce web applications.
|
||||
Your goal is to generate realistic administrative tasks that can be solved by a Web Agent navigating an admin panel.
|
||||
|
||||
I will provide you with the following context:
|
||||
1. **Full Database Schema**: A list of `CREATE TABLE` statements for the core tables of a Magento e-commerce platform.
|
||||
2. **Sampled Data**: A JSON object containing 5 random rows of data from 5 randomly selected core tables. This data is REAL and should be used to inspire specific, answerable questions.
|
||||
|
||||
## Your Task
|
||||
|
||||
Based on the provided schema and sample data, create a JSON object containing a single key, "questions", which holds an array of up to 10 unique task objects.
|
||||
|
||||
### Requirements for Each Question:
|
||||
- **Web Agent Solvable**: The task must represent a realistic action an administrator would perform in a web UI (e.g., "Find all orders for customer X", "Update the stock for product Y", "Approve a pending review").
|
||||
- **Grounded in Data**: The questions should be specific, using names, IDs, or values from the provided **Sampled Data** to make them concrete.
|
||||
- **Utilize Schema**: You can formulate questions that require joining tables, even if not all tables were sampled. The full schema is your guide.
|
||||
|
||||
### Output Format
|
||||
The final output MUST be a single, valid JSON object. Do not include any other text, explanations, or markdown formatting like ```json.
|
||||
The JSON object must have one key: "questions", containing a JSON array of task objects.
|
||||
|
||||
Each object in the array must contain exactly three keys: `question`, `answer`, and `sql`.
|
||||
|
||||
- **`question`**: (string) A natural language description of the task for a web agent.
|
||||
- **`answer`**: (string, integer, float, or list) The precise and concise answer to the question, derived by running the SQL query against the database.
|
||||
- **`sql`**: (string) The exact, runnable MySQL query that was used to find the answer.
|
||||
|
||||
### Output Format Example
|
||||
```json
|
||||
{{
|
||||
"questions": [
|
||||
{{
|
||||
"question": "What is the email address for customer with ID 5?",
|
||||
"answer": "customer5@example.com",
|
||||
"sql": "SELECT email FROM customer_entity WHERE entity_id = 5;"
|
||||
}},
|
||||
{{
|
||||
"question": "Find the total quantity of item with SKU 'ABC-123' in the cart.",
|
||||
"answer": 3,
|
||||
"sql": "SELECT SUM(qty) FROM quote_item WHERE sku = 'ABC-123';"
|
||||
}}
|
||||
]
|
||||
}}
|
||||
```
|
||||
|
||||
---
|
||||
### Full Database Schema
|
||||
{schema_context}
|
||||
|
||||
---
|
||||
### Sampled Data
|
||||
Here is the sample data from randomly selected tables. Use this to make your questions specific.
|
||||
|
||||
{sampled_data_str}
|
||||
|
||||
---
|
||||
Now, generate the JSON object based on these instructions.
|
||||
"""
|
||||
|
||||
# This is a new prompt to evaluate results and generate a corrected answer.
|
||||
SEMANTIC_EVALUATION_PROMPT_TEMPLATE = """
|
||||
You are a precise data analyst. Your task is to evaluate if a SQL query's result adequately answers a given natural language question. You will then either refine the answer, or completely rephrase the question if the result set is large.
|
||||
|
||||
I will provide you with a JSON object containing:
|
||||
1. `question`: The original question asked.
|
||||
2. `sql`: The SQL query that was executed.
|
||||
3. `sql_result`: The actual data returned by executing the SQL query.
|
||||
4. `row_count`: The number of rows in `sql_result`.
|
||||
|
||||
## Your Task
|
||||
Analyze the inputs and respond with a JSON object. You have three cases. The `new_answer` field MUST always be an array of strings.
|
||||
|
||||
### Data Analysis and Refinement Rules
|
||||
1. **Analyze SQL and Question Intent**: Look at the SQL query (`SELECT`, `COUNT`, `DISTINCT`, etc.) and the natural language `question` to understand the user's goal. Is the goal to count things, list unique items, or retrieve specific related data points?
|
||||
2. **Handle Duplicates and Merge Data**:
|
||||
- **De-duplication**: If the `question` implies a list of unique items (e.g., "List the cities..." or "What are the unique order statuses?"), you MUST de-duplicate the values in `sql_result` to form the `new_answer`. For example, if `sql_result` is `[["pending"], ["shipped"], ["pending"]]`, the `new_answer` should be `["pending", "shipped"]`.
|
||||
- **Data Merging**: If the `sql_result` contains multiple rows related to the same entity (e.g., different attributes of one product), combine the relevant information into a concise `new_answer`. For instance, if the question is "What are the name and price of product 'XYZ'?" and `sql_result` is `[["Product XYZ", 99.99]]`, the `new_answer` is `["Product XYZ", "99.99"]`. If the result was `[["Product XYZ", "Red"], ["Product XYZ", "Blue"]]` for a question about colors, `new_answer` could be `["Red", "Blue"]`. Extract only the information that directly answers the question.
|
||||
|
||||
After applying these rules, select one of the three cases below for your response format.
|
||||
|
||||
### Case 1: Large Result Set (Question Transformation)
|
||||
If `row_count` is greater than 10 AND the original `question` does NOT already ask for a count (e.g., it is not phrased like "How many..."), you must transform the question.
|
||||
Respond with:
|
||||
```json
|
||||
{{
|
||||
"can_answer": true,
|
||||
"new_question": "How many items were found?",
|
||||
"new_answer": ["42"]
|
||||
}}
|
||||
```
|
||||
- `can_answer`: (boolean) Must be `true`.
|
||||
- `new_question`: (string) A rephrased question that asks for the quantity of items. For example, if the original question was "List all products", the new question should be "How many products were found?".
|
||||
- `new_answer`: (array of strings) An array containing the `row_count` as a single string element.
|
||||
|
||||
### Case 2: Standard Answer (No Transformation)
|
||||
If Case 1 does not apply, but the `sql_result` still provides a clear answer to the original `question` (after applying the refinement rules), respond with:
|
||||
```json
|
||||
{{
|
||||
"can_answer": true,
|
||||
"new_answer": ["value1", "value2", ...]
|
||||
}}
|
||||
```
|
||||
- `can_answer`: (boolean) Must be `true`.
|
||||
- `new_answer`: (array of strings) An array containing all the essential parts of the answer extracted and refined from `sql_result`. Every value from the result set that contributes to the answer should be included as a string in the array. This ensures answer completeness.
|
||||
- **Example 1**: If `question` is "What is the status of order 123?" and `sql_result` is `[["processing"]]`, `new_answer` should be `["processing"]`.
|
||||
- **Example 2**: If `question` is "List emails for pending customers" and `sql_result` is `[["test@a.com"], ["test@b.com"]]`, `new_answer` should be `["test@a.com", "test@b.com"]`.
|
||||
- **Example 3**: If `question` is "Get product name and price for SKU 'XYZ'" and `sql_result` is `[["My Product", 19.99]]`, `new_answer` should be `["My Product", "19.99"]`.
|
||||
|
||||
### Case 3: The question CANNOT be answered
|
||||
If the `sql_result` is empty, irrelevant, or insufficient to answer the question, respond with:
|
||||
```json
|
||||
{{
|
||||
"can_answer": false,
|
||||
"reason": "..."
|
||||
}}
|
||||
```
|
||||
- `can_answer`: (boolean) Must be `false`.
|
||||
- `reason`: (string) A brief explanation for why the question cannot be answered.
|
||||
|
||||
---
|
||||
### Evaluation Data
|
||||
{task_data_json}
|
||||
---
|
||||
|
||||
Now, provide your evaluation as a JSON object.
|
||||
"""
|
||||
|
||||
def get_db_connection():
|
||||
"""Establishes a connection to the MySQL database."""
|
||||
try:
|
||||
conn = mysql.connector.connect(**MYSQL_CONFIG)
|
||||
return conn
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Error connecting to MySQL: {err}")
|
||||
return None
|
||||
|
||||
def get_full_schema(cursor, tables):
|
||||
"""Fetches the CREATE TABLE statements for all core tables."""
|
||||
schema_parts = []
|
||||
for table_name in tables:
|
||||
try:
|
||||
cursor.execute(f"SHOW CREATE TABLE `{table_name}`")
|
||||
result = cursor.fetchone()
|
||||
if result:
|
||||
schema_parts.append(result[1]) # result[1] is the CREATE TABLE statement
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Warning: Could not get schema for table {table_name}: {err}")
|
||||
return "\n\n".join(schema_parts)
|
||||
|
||||
def get_random_tables_and_samples(cursor, tables, num_tables=5, num_samples=5):
|
||||
"""Selects random tables and samples random rows from them."""
|
||||
selected_tables = random.sample(tables, num_tables)
|
||||
sampled_data = {}
|
||||
|
||||
for table_name in selected_tables:
|
||||
try:
|
||||
# Use ORDER BY RAND() for random sampling. Can be slow on very large tables.
|
||||
query = f"SELECT * FROM `{table_name}` ORDER BY RAND() LIMIT {num_samples}"
|
||||
cursor.execute(query)
|
||||
|
||||
rows = cursor.fetchall()
|
||||
if not rows:
|
||||
sampled_data[table_name] = []
|
||||
continue
|
||||
|
||||
columns = [desc[0] for desc in cursor.description]
|
||||
|
||||
# Convert rows (tuples) to a list of dictionaries
|
||||
sampled_rows = []
|
||||
for row in rows:
|
||||
row_dict = {}
|
||||
for i, col_value in enumerate(row):
|
||||
# Handle bytes by decoding, fall back to string representation
|
||||
if isinstance(col_value, bytes):
|
||||
try:
|
||||
row_dict[columns[i]] = col_value.decode('utf-8')
|
||||
except UnicodeDecodeError:
|
||||
row_dict[columns[i]] = str(col_value)
|
||||
else:
|
||||
row_dict[columns[i]] = col_value
|
||||
sampled_rows.append(row_dict)
|
||||
|
||||
sampled_data[table_name] = sampled_rows
|
||||
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Warning: Could not sample data from table {table_name}: {err}")
|
||||
sampled_data[table_name] = f"Error: {err}"
|
||||
|
||||
return sampled_data
|
||||
|
||||
def generate_questions(client, schema_context, sampled_data):
|
||||
"""Generates questions by calling the OpenAI API."""
|
||||
if not client:
|
||||
raise ValueError("OpenAI client not provided.")
|
||||
|
||||
sampled_data_str = json.dumps(sampled_data, indent=2, default=str)
|
||||
|
||||
prompt = PROMPT_TEMPLATE.format(
|
||||
schema_context=schema_context,
|
||||
sampled_data_str=sampled_data_str
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.chat.completions.create(
|
||||
model=OPENAI_CONFIG["model"],
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant designed to output JSON."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0.7,
|
||||
response_format={"type": "json_object"},
|
||||
)
|
||||
content = response.choices[0].message.content
|
||||
data = json.loads(content)
|
||||
|
||||
# The prompt asks for {"questions": [...]}, so we extract the list.
|
||||
if isinstance(data, dict) and "questions" in data and isinstance(data["questions"], list):
|
||||
return data["questions"]
|
||||
elif isinstance(data, list):
|
||||
# Fallback in case the model returns a list directly
|
||||
print("Warning: Model returned a raw list instead of an object with a 'questions' key.")
|
||||
return data
|
||||
else:
|
||||
print(f"Warning: Failed to find a 'questions' list in the model's output. Got: {content}")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error calling OpenAI API or parsing JSON: {e}")
|
||||
return None
|
||||
|
||||
def load_existing_tasks(filepath):
|
||||
"""Loads tasks from a JSON file if it exists."""
|
||||
if not os.path.exists(filepath):
|
||||
return []
|
||||
try:
|
||||
with open(filepath, 'r') as f:
|
||||
content = f.read()
|
||||
if not content: # Handle empty file
|
||||
return []
|
||||
return json.loads(content)
|
||||
except (json.JSONDecodeError, FileNotFoundError):
|
||||
print(f"Warning: Could not read or parse {filepath}. Starting with an empty list.")
|
||||
return []
|
||||
|
||||
def evaluate_and_refine_tasks(tasks, client):
|
||||
"""
|
||||
Uses an LLM to evaluate if a SQL result answers the question and refines the answer.
|
||||
"""
|
||||
if not tasks:
|
||||
return []
|
||||
|
||||
final_validated_tasks = []
|
||||
print("\nPerforming semantic evaluation and answer refinement with GPT-4o...")
|
||||
|
||||
for task in tasks:
|
||||
# Prepare data for the prompt, excluding the original 'answer'
|
||||
task_data_for_prompt = {
|
||||
"question": task["question"],
|
||||
"sql": task["sql"],
|
||||
"sql_result": task["sql_result"],
|
||||
"row_count": task["row_count"]
|
||||
}
|
||||
task_data_json = json.dumps(task_data_for_prompt, indent=2, default=str)
|
||||
|
||||
prompt = SEMANTIC_EVALUATION_PROMPT_TEMPLATE.format(task_data_json=task_data_json)
|
||||
|
||||
try:
|
||||
print(f" - Evaluating question: \"{task['question'][:80]}...\"")
|
||||
response = client.chat.completions.create(
|
||||
model=OPENAI_CONFIG["model"],
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant designed to output JSON."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0.0, # We want deterministic evaluation
|
||||
response_format={"type": "json_object"},
|
||||
)
|
||||
content = response.choices[0].message.content
|
||||
evaluation_result = json.loads(content)
|
||||
|
||||
if evaluation_result.get("can_answer") is True and "new_answer" in evaluation_result:
|
||||
# Task is valid. Update the answer with the refined one from the LLM.
|
||||
task['answer'] = evaluation_result['new_answer']
|
||||
|
||||
# If the LLM provides a new question, update it.
|
||||
if 'new_question' in evaluation_result:
|
||||
task['question'] = evaluation_result['new_question']
|
||||
print(f" - Question was rephrased: \"{task['question']}\"")
|
||||
|
||||
task['sql_execute_result'] = task.pop('sql_result')
|
||||
task.pop('row_count', None) # Clean up temp key
|
||||
final_validated_tasks.append(task)
|
||||
print(f" - Evaluation PASSED. New answer: {json.dumps(task['answer'])}")
|
||||
else:
|
||||
reason = evaluation_result.get('reason', 'No reason provided.')
|
||||
print(f" - Evaluation FAILED. Filtering task.")
|
||||
print(f" - Reason: {reason}")
|
||||
print(f" - Question: {task['question']}")
|
||||
print(f" - SQL: {task['sql']}")
|
||||
sql_result_str = json.dumps(task['sql_result'], indent=2, default=str)
|
||||
print(f" - SQL Result: {sql_result_str}")
|
||||
|
||||
except Exception as e:
|
||||
print(f" - An error occurred during semantic evaluation for task, filtering it out: {e}")
|
||||
print(f" - Question: {task.get('question', 'N/A')}")
|
||||
print(f" - SQL: {task.get('sql', 'N/A')}")
|
||||
|
||||
return final_validated_tasks
|
||||
|
||||
def main():
|
||||
"""Main function to run the script."""
|
||||
parser = argparse.ArgumentParser(description="Generate and validate e-commerce admin tasks.")
|
||||
parser.add_argument(
|
||||
"--target-count",
|
||||
type=int,
|
||||
required=True,
|
||||
help="The total number of questions to generate."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-file",
|
||||
type=str,
|
||||
default="generated_tasks.json",
|
||||
help="The file to save the generated tasks to (in JSON format)."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Load existing tasks from the output file
|
||||
all_tasks = load_existing_tasks(args.output_file)
|
||||
print(f"Found {len(all_tasks)} existing valid tasks in '{args.output_file}'.")
|
||||
|
||||
# Connect to DB and set up client
|
||||
conn = get_db_connection()
|
||||
if not conn:
|
||||
return
|
||||
cursor = conn.cursor()
|
||||
|
||||
if not OPENAI_CONFIG["api_key"]:
|
||||
print("Error: OPENAI_API_KEY environment variable not set.")
|
||||
return
|
||||
client = OpenAI(api_key=OPENAI_CONFIG["api_key"], base_url=OPENAI_CONFIG["base_url"])
|
||||
|
||||
try:
|
||||
# Load core tables and schema once
|
||||
try:
|
||||
with open('core_tables.json', 'r') as f:
|
||||
core_tables = json.load(f)
|
||||
except FileNotFoundError:
|
||||
print("Error: core_tables.json not found. Please create it.")
|
||||
return
|
||||
|
||||
print("Fetching full database schema...")
|
||||
schema_context = get_full_schema(cursor, core_tables)
|
||||
|
||||
# Start the generation loop
|
||||
round_num = 1
|
||||
while len(all_tasks) < args.target_count:
|
||||
print(f"\n--- Starting Generation Round {round_num} ---")
|
||||
print(f"Goal: {args.target_count} | Current: {len(all_tasks)} | Needed: {args.target_count - len(all_tasks)}")
|
||||
|
||||
# Get random samples for this round
|
||||
print("Sampling data from 5 random tables...")
|
||||
sampled_data = get_random_tables_and_samples(cursor, core_tables, num_tables=5, num_samples=5)
|
||||
|
||||
# Generate questions
|
||||
print("Generating questions with GPT-4o...")
|
||||
generated_tasks = generate_questions(client, schema_context, sampled_data)
|
||||
|
||||
# Execute SQL for generated tasks
|
||||
tasks_for_evaluation = []
|
||||
if generated_tasks:
|
||||
print("\nExecuting SQL for generated tasks...")
|
||||
for task in generated_tasks:
|
||||
if not isinstance(task, dict) or not all(k in task for k in ['sql', 'answer', 'question']):
|
||||
print(f"Filtering task due to malformed structure: {task}")
|
||||
continue
|
||||
try:
|
||||
cursor.execute(task['sql'])
|
||||
sql_result = cursor.fetchall()
|
||||
# Create a new dict for evaluation, excluding the original 'answer'.
|
||||
tasks_for_evaluation.append({
|
||||
'question': task['question'],
|
||||
'sql': task['sql'],
|
||||
'sql_result': sql_result,
|
||||
'row_count': len(sql_result)
|
||||
})
|
||||
except mysql.connector.Error as err:
|
||||
print(f"Filtering task due to SQL error: {err} on SQL: {task['sql']}")
|
||||
|
||||
# Perform semantic evaluation and get validated tasks
|
||||
validated_tasks = evaluate_and_refine_tasks(tasks_for_evaluation, client)
|
||||
|
||||
# Append new tasks and save to file
|
||||
if validated_tasks:
|
||||
all_tasks.extend(validated_tasks)
|
||||
with open(args.output_file, 'w') as f:
|
||||
json.dump(all_tasks, f, indent=2, default=str)
|
||||
|
||||
print("\n--- Round Summary ---")
|
||||
print(f"Generated {len(validated_tasks)} new valid tasks in this round.")
|
||||
print(f"Progress: {len(all_tasks)} / {args.target_count} tasks.")
|
||||
else:
|
||||
print("\n--- Round Summary ---")
|
||||
print("No new valid tasks were generated in this round. Retrying...")
|
||||
|
||||
round_num += 1
|
||||
|
||||
finally:
|
||||
# Close the database connection
|
||||
if conn.is_connected():
|
||||
cursor.close()
|
||||
conn.close()
|
||||
print("\nDatabase connection closed.")
|
||||
|
||||
print(f"\nTarget of {args.target_count} tasks reached. Final output saved to {args.output_file}.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
46291
random_sample/generated_tasks.json
Normal file
46291
random_sample/generated_tasks.json
Normal file
File diff suppressed because it is too large
Load Diff
3539
random_sample/generated_tasks.json.bak
Normal file
3539
random_sample/generated_tasks.json.bak
Normal file
File diff suppressed because it is too large
Load Diff
8182
random_sample/generated_tasks.json.bak2
Normal file
8182
random_sample/generated_tasks.json.bak2
Normal file
File diff suppressed because it is too large
Load Diff
9329
random_sample/generated_tasks.json.bak3
Normal file
9329
random_sample/generated_tasks.json.bak3
Normal file
File diff suppressed because it is too large
Load Diff
123
random_sample/make_config_files.py
Normal file
123
random_sample/make_config_files.py
Normal file
@@ -0,0 +1,123 @@
|
||||
"""
|
||||
读取generated_tasks.json文件格式如下输入，将每个问题一一转换按照输出格式的的json，其中question作为intent, answer对应到must_include的数组，sql对应填充到reference_answer_raw_annotation，task_id从0开始编号，其他字段保留原始内容，
|
||||
输出文件保存为test_rlvr.raw.json
|
||||
|
||||
输入
|
||||
[
|
||||
{
|
||||
"question": "Does the attribute ID 120 use page builder?",
|
||||
"sql": "SELECT is_pagebuilder_enabled FROM catalog_eav_attribute WHERE attribute_id = 120;",
|
||||
"answer": [
|
||||
"No"
|
||||
],
|
||||
"sql_execute_result": [
|
||||
[
|
||||
0
|
||||
]
|
||||
]
|
||||
},
|
||||
...
|
||||
]
|
||||
|
||||
|
||||
输出
|
||||
[
|
||||
{
|
||||
"sites": [
|
||||
"shopping_admin"
|
||||
],
|
||||
"task_id": 0,
|
||||
"require_login": true,
|
||||
"storage_state": "./.auth/shopping_admin_state.json",
|
||||
"start_url": "__SHOPPING_ADMIN__",
|
||||
"geolocation": null,
|
||||
"intent_template": "",
|
||||
"instantiation_dict": {},
|
||||
"intent": "What are the top-3 best-selling product in Jan 2023",
|
||||
"require_reset": false,
|
||||
"eval": {
|
||||
"eval_types": [
|
||||
"string_match"
|
||||
],
|
||||
"reference_answers": {
|
||||
"must_include": [
|
||||
"Impulse Duffle",
|
||||
"Overnight Duffle",
|
||||
"Hawkeye Yoga Short-32-Blue"
|
||||
]
|
||||
},
|
||||
"reference_url": "",
|
||||
"program_html": [],
|
||||
"string_note": "",
|
||||
"reference_answer_raw_annotation": ""
|
||||
},
|
||||
"intent_template_id": 0,
|
||||
"old_task_id": 0
|
||||
},
|
||||
...
|
||||
]
|
||||
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
def generate_config():
|
||||
"""
|
||||
Reads tasks from 'generated_tasks.json', converts them to the WebArena
|
||||
format, and saves them to 'test_rlvr.raw.json'.
|
||||
"""
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
input_path = os.path.join(script_dir, 'generated_tasks.json')
|
||||
output_path = os.path.join(script_dir, 'test_rlvr.raw.json')
|
||||
|
||||
try:
|
||||
with open(input_path, 'r', encoding='utf-8') as f:
|
||||
input_data = json.load(f)
|
||||
except FileNotFoundError:
|
||||
print(f"Error: Input file '{input_path}' not found.")
|
||||
return
|
||||
except json.JSONDecodeError:
|
||||
print(f"Error: Failed to decode JSON from '{input_path}'.")
|
||||
return
|
||||
|
||||
output_data = []
|
||||
for i, task in enumerate(input_data):
|
||||
if not all(k in task for k in ['question', 'answer', 'sql']):
|
||||
print(f"Warning: Skipping task at index {i} due to missing keys.")
|
||||
continue
|
||||
|
||||
new_task = {
|
||||
"sites": ["shopping_admin"],
|
||||
"task_id": i,
|
||||
"require_login": True,
|
||||
"storage_state": "./.auth/shopping_admin_state.json",
|
||||
"start_url": "__SHOPPING_ADMIN__",
|
||||
"geolocation": None,
|
||||
"intent_template": "",
|
||||
"instantiation_dict": {},
|
||||
"intent": task['question'],
|
||||
"require_reset": False,
|
||||
"eval": {
|
||||
"eval_types": ["string_match"],
|
||||
"reference_answers": {
|
||||
"must_include": task['answer']
|
||||
},
|
||||
"reference_url": "",
|
||||
"program_html": [],
|
||||
"string_note": "",
|
||||
"reference_answer_raw_annotation": task['sql']
|
||||
},
|
||||
"intent_template_id": 0,
|
||||
"old_task_id": i
|
||||
}
|
||||
output_data.append(new_task)
|
||||
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(output_data, f, indent=4)
|
||||
|
||||
print(f"Successfully created '{output_path}' with {len(output_data)} tasks.")
|
||||
|
||||
if __name__ == '__main__':
|
||||
generate_config()
|
||||
|
||||
30260
random_sample/test_rlvr.raw.json
Normal file
30260
random_sample/test_rlvr.raw.json
Normal file
File diff suppressed because it is too large
Load Diff
30260
random_sample/webrlvr_cms1k_easy.raw.json
Normal file
30260
random_sample/webrlvr_cms1k_easy.raw.json
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user