https://ollama.ac.cn/library/sqlcoder
https://blog.csdn.net/hzether/article/details/143816042
import ollama
import sqlite3
import json
from contextlib import closingdef generate_and_execute_sql(question: str, db_path: str) -> dict:# 1. 生成 SQL 查詢語句prompt = f"""### Instructions:Convert Chinese question to SQL query. Follow these rules strictly:1. ONLY return a valid SELECT SQL query2. Use EXACT table names from the mapping below3. DO NOT use any table that's not in the mapping### Examples:Question: 所有訂單記錄SQL: SELECT * FROM orders ORDER BY id;### Database Schema:{get_schema(db_path)}### Question:{question}### SQL Query:"""print(f"輸入: {prompt}")response = ollama.chat(model='sqlcoder:latest',messages=[{'role': 'user', 'content': prompt}])sql_query = response['message']['content'].strip()print(f"生成的SQL: {sql_query}") # 調試日志# 2. 執行 SQL 查詢try:with closing(sqlite3.connect(db_path)) as conn:conn.row_factory = sqlite3.Row # 設置為行工廠cursor = conn.cursor()cursor.execute(sql_query)# 3. 獲取結果并轉為 JSONrows = cursor.fetchall()result = [dict(row) for row in rows]return {"status": "success","sql": sql_query,"data": result}except Exception as e:return {"status": "error","sql": sql_query,"message": str(e)}def get_schema(db_path: str) -> str:"""獲取數據庫的 schema 信息"""with closing(sqlite3.connect(db_path)) as conn:cursor = conn.cursor()cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")tables = cursor.fetchall()schema_info = []for table in tables:table_name = table[0]cursor.execute(f"PRAGMA table_info({table_name});")columns = cursor.fetchall()col_details = [f"{col[1]} ({col[2]})" for col in columns]schema_info.append(f"表 {table_name}: {', '.join(col_details)}")return "\n".join(schema_info)# 使用示例
if __name__ == "__main__":# 配置參數DB_PATH = "data.db" # SQLite 數據庫文件路徑QUESTION = "查詢銷售額超過10000的訂單信息" # 用戶問題# 執行查詢result = generate_and_execute_sql(QUESTION, DB_PATH)print(f"返回數據")print(json.dumps(result, indent=2, ensure_ascii=False))
D:\ProgramData\anaconda3\python.exe F:/mark/sqlauto/main.py
輸入: ### Instructions:Convert Chinese question to SQL query. Follow these rules strictly:1. ONLY return a valid SELECT SQL query2. Use EXACT table names from the mapping below3. DO NOT use any table that's not in the mapping### Examples:Question: 所有訂單記錄SQL: SELECT * FROM orders ORDER BY id;### Database Schema:表 sqlite_sequence: name (), seq ()
表 orders: order_id (INTEGER), customer_name (TEXT), order_date (DATE), total_amount (REAL), status (TEXT)### Question:查詢銷售額超過10000的訂單信息### SQL Query:生成的SQL: SELECT * FROM orders WHERE CAST(total_amount AS integer) > 10000;
返回數據
{"status": "success","sql": "SELECT * FROM orders WHERE CAST(total_amount AS integer) > 10000;","data": [{"order_id": 2,"customer_name": "李四","order_date": "2023-09-20","total_amount": 12000.5,"status": "shipped"},{"order_id": 4,"customer_name": "趙六","order_date": "2023-10-18","total_amount": 21000.0,"status": "delivered"},{"order_id": 6,"customer_name": "孫八","order_date": "2023-11-15","total_amount": 15500.0,"status": "delivered"},{"order_id": 8,"customer_name": "吳十","order_date": "2023-12-10","total_amount": 18900.75,"status": "delivered"}]
}Process finished with exit code 0