ollama - sqlcoder模型：面向提示词编程（根据用户信息生成sql语句并执行返回结果）

https://ollama.ac.cn/library/sqlcoder

https://blog.csdn.net/hzether/article/details/143816042

import ollama
import sqlite3
import json
from contextlib import closing

def 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 query
        2. Use EXACT table names from the mapping below
        3. 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. 获取结果并转为 JSON
            rows = 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 query
        2. Use EXACT table names from the mapping below
        3. 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