Upload sql_command.py

sql_command.py

@@ -0,0 +1,50 @@
+'''
+
+'''
+##TODO:
+
+import sqlite3
+import pandas as pd
+
+## run the following function to set up a dedicated database using SQLite.
+def construct_db(data_file=None):
+    excel_file = "/Users/yunshi/Downloads/360Data/Data Center/Working-On Task/演讲与培训/2023ChatGPT/Coding/Text2SQL/模拟数据.csv"  # Replace with your actual file path
+    df = pd.read_csv(excel_file)
+    print('df:', df.head())
+    
+    conn = sqlite3.connect('myexcelDB.db')  # Replace 'mydatabase.db' with your desired name
+    # Create a cursor object to execute SQL commands
+    cursor = conn.cursor()
+
+    ##NOTE: Insert data from DataFrame into the table. 注意这里的if_exists选项，考虑是否要覆盖原始内容。这里如果需要指定index，那么就需要指定index_label，且index=False，否则会选择默认的index，这样就会产生duplicate错误。
+    df.to_sql('table01', conn, if_exists='replace', index=False, index_label="产品ID")
+    
+    return None
+
+
+def llm_query(sql_command):
+    # Connect to the database (or create it if it doesn't exist)
+    conn = sqlite3.connect('./myexcelDB.db')  # Replace 'myexcelDB.db' with your desired name
+
+    # Create a cursor object to execute SQL commands
+    cursor = conn.cursor()
+    ## SQL command
+    cursor.execute(sql_command)
+    query_result = cursor.fetchall()  # Fetch all rows as a list of tuples
+    # print('query_result:', query_result)
+    
+    ## 将列名也取出来。
+    column_names = [description[0] for description in cursor.description]
+    
+    query_df = pd.DataFrame(query_result, columns=column_names)
+    # query_df.set_index("产品ID", inplace=True)
+    print('query_df:', query_df)
+    
+    return query_df
+
+# llm_query("SELECT * FROM table01 WHERE 宽度 > 300") ## sample function call. 
+
+# construct_db()
+# sql = "SELECT 产品ID FROM table01 WHERE 长度 > 50"
+# res = llm_query(sql)
+# res