import requests
import json

def get_access_token():
    """
    使用 API Key，Secret Key 获取access_token，替换下列示例中的应用API Key、应用Secret Key
    """
        
    url = "https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=[应用API Key]&client_secret=[应用Secret Key]"
    
    payload = json.dumps("")
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }
    
    response = requests.request("POST", url, headers=headers, data=payload)
    return response.json().get("access_token")


def main():
       
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/completions/sqlcoder_7b?access_token=" + get_access_token()
    
    payload = json.dumps({
         "prompt": "### Task\nGenerate a SQL query to answer the following question:\n`What is our total revenue by product in the last week?`\n\n### Database Schema\nThis query will run on a database whose schema is represented in this string:\nCREATE TABLE products (\n  product_id INTEGER PRIMARY KEY, -- Unique ID for each product\n  name VARCHAR(50), -- Name of the product\n  price DECIMAL(10,2), -- Price of each unit of the product\n  quantity INTEGER  -- Current quantity in stock\n);\n\nCREATE TABLE customers (\n   customer_id INTEGER PRIMARY KEY, -- Unique ID for each customer\n   name VARCHAR(50), -- Name of the customer\n   address VARCHAR(100) -- Mailing address of the customer\n);\n\nCREATE TABLE salespeople (\n  salesperson_id INTEGER PRIMARY KEY, -- Unique ID for each salesperson\n  name VARCHAR(50), -- Name of the salesperson\n  region VARCHAR(50) -- Geographic sales region\n);\n\nCREATE TABLE sales (\n  sale_id INTEGER PRIMARY KEY, -- Unique ID for each sale\n  product_id INTEGER, -- ID of product sold\n  customer_id INTEGER,  -- ID of customer who made purchase\n  salesperson_id INTEGER, -- ID of salesperson who made the sale\n  sale_date DATE, -- Date the sale occurred\n  quantity INTEGER -- Quantity of product sold\n);\n\nCREATE TABLE product_suppliers (\n  supplier_id INTEGER PRIMARY KEY, -- Unique ID for each supplier\n  product_id INTEGER, -- Product ID supplied\n  supply_price DECIMAL(10,2) -- Unit price charged by supplier\n);\n\n-- sales.product_id can be joined with products.product_id\n-- sales.customer_id can be joined with customers.customer_id\n-- sales.salesperson_id can be joined with salespeople.salesperson_id\n-- product_suppliers.product_id can be joined with products.product_id\n\n### SQL\nGiven the database schema, here is the SQL query that answers `What is our total revenue by product in the last week?`:\n```sql\n"
    })
    headers = {
        'Content-Type': 'application/json'
    }
    
    response = requests.request("POST", url, headers=headers, data=payload)
    
    print(response.text)
    

if __name__ == '__main__':
    main()