Text-to-SQL

让 Claude 将自然语言转换为 SQL 查询

🌐 查看英文原文 | 源码 Notebook

---

使用 Claude 进行 Text to SQL

简介

Text to SQL 是一种自然语言处理任务，它将可读文本查询转换为结构化 SQL 查询。这使得用户可以使用自然语言与数据库交互。

Claude 可以理解上下文，解释复杂查询，并生成准确的 SQL 语句。本指南侧重于使用 Claude 构建一个健壮的 Text to SQL 系统。

为什么 Text to SQL 有用？

Text to SQL 有几个重要的原因：

1. 易用性：非技术用户可以在不了解 SQL 语法的情况下查询数据库，这使得组织内部的数据访问更加容易。

2. 效率：数据分析师和科学家可以使用自然语言快速原型设计查询。

3. 集成：它使得应用程序和聊天机器人中的数据库交互界面更加直观。

4. 复杂查询生成：LLM 可以生成涉及多个连接、子查询和聚合的复杂 SQL 查询，这对于人类来说可能需要花费大量时间编写。

本指南涵盖的内容

本指南将指导您使用 LLM 构建 Text to SQL 系统。我们将涵盖：

1. 设置测试 SQLite 数据库
2. 有效的 Text to SQL 转换提示
3. RAG（检索增强生成）以处理更复杂的数据库系统
4. Claude 输出的自我改进和迭代
5. 评估

在本指南结束时，您将了解如何使用 Claude 实施和改进 Text to SQL 任务，并拥有将这些技术应用于您自己的项目的框架。

目录

1. 简介
2. 设置
3. 创建基本的 Text to SQL 提示
4. 通过示例改进提示
5. 使用思维链提示
6. 实现针对复杂数据库模式的RAG
7. 实现查询自我改进
8. 评估
9. 进一步探索与下一步

环境配置

让我们设置我们的环境，并创建一个包含两个表（employees 和 departments）的测试 SQLite 数据库。我们将在整个指南中使用这个数据库。

# 创建数据库和表
import sqlite3

# 连接到SQLite数据库
# 如果文件不存在，会自动在当前目录创建:
conn = sqlite3.connect('test.db')
cursor = conn.cursor()

# 创建表
cursor.execute('''
CREATE TABLE IF NOT EXISTS employees (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    department_id INTEGER,
    FOREIGN KEY(department_id) REFERENCES departments(id)
)
''')

cursor.execute('''
CREATE TABLE IF NOT EXISTS departments (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL
)
''')

# 提交事务:
conn.commit()

# 关闭Cursor:
cursor.close()

# 关闭Connection:
conn.close()
``````python
%pip install -q anthropic pandas voyageai matplotlib seaborn

Note: you may need to restart the kernel to use updated packages.

import os
import sqlite3

import pandas as pd
from anthropic import Anthropic
from IPython.display import display

# Set your Claude API key
os.environ["ANTHROPIC_API_KEY"] = "YOUR_ANTHROPIC_API_KEY"
os.environ["VOYAGE_API_KEY"] = "YOUR_VOYAGE_API_KEY"

# Initialize the Anthropic client
client = Anthropic()
MODEL_NAME = "claude-sonnet-4-6"

# Filepath to the SQLite database
DATABASE_PATH = "data/data.db"

创建测试数据库

import random
from datetime import datetime, timedelta

if not os.path.exists(DATABASE_PATH):
    print("Database does not exist. Creating and populating...")
    # Create a new SQLite database and tables
    with sqlite3.connect(DATABASE_PATH) as conn:
        cursor = conn.cursor()

        cursor.executescript("""
        CREATE TABLE IF NOT EXISTS departments (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            location TEXT
        );
        CREATE TABLE IF NOT EXISTS employees (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            age INTEGER,
            department_id INTEGER,
            salary REAL,
            hire_date DATE,
            FOREIGN KEY (department_id) REFERENCES departments (id)
        );
        """)

        # Insert sample data
        cursor.executemany(
            "INSERT INTO departments VALUES (?,?,?)",
            [
                (1, "HR", "New York"),
                (2, "Engineering", "San Francisco"),
                (3, "Marketing", "Chicago"),
                (4, "Sales", "Los Angeles"),
                (5, "Finance", "Boston"),
                (6, "Customer Support", "Dallas"),
                (7, "Research", "Seattle"),
                (8, "Legal", "Washington D.C."),
                (9, "Product", "Austin"),
                (10, "Operations", "Denver"),
            ],
        )

        first_names = [
            "John",
            "Jane",
            "Bob",
            "Alice",
            "Charlie",
            "Diana",
            "Edward",
            "Fiona",
            "George",
            "Hannah",
            "Ian",
            "Julia",
            "Kevin",
            "Laura",
            "Michael",
            "Nora",
            "Oliver",
            "Patricia",
            "Quentin",
            "Rachel",
            "Steve",
            "Tina",
            "Ulysses",
            "Victoria",
            "William",
            "Xena",
            "Yannick",
            "Zoe",
        ]
        last_names = [
            "Smith",
            "Johnson",
            "Williams",
            "Jones",
            "Brown",
            "Davis",
            "Miller",
            "Wilson",
            "Moore",
            "Taylor",
            "Anderson",
            "Thomas",
            "Jackson",
            "White",
            "Harris",
            "Martin",
            "Thompson",
            "Garcia",
            "Martinez",
            "Robinson",
            "Clark",
            "Rodriguez",
            "Lewis",
            "Lee",
            "Walker",
            "Hall",
            "Allen",
            "Young",
            "King",
        ]

        employees_data = []
        for i in range(1, 201):  # Generate 200 employees
            name = f"{random.choice(first_names)} {random.choice(last_names)}"
            age = random.randint(22, 65)
            department_id = random.randint(1, 10)
            salary = round(random.uniform(40000, 200000), 2)
            hire_date = (datetime.now() - timedelta(days=random.randint(0, 3650))).strftime(
                "%Y-%m-%d"
            )
            employees_data.append((i, name, age, department_id, salary, hire_date))

        cursor.executemany("INSERT INTO employees VALUES (?,?,?,?,?,?)", employees_data)

    print("Database created and populated successfully.")
else:
    print("Database already exists. Skipping creation and population.")

# Display table contents
with sqlite3.connect(DATABASE_PATH) as conn:
    for table in ["departments", "employees"]:
        df = pd.read_sql_query(f"SELECT * FROM {table}", conn)
        print(f"\n{table.capitalize()} table:")
        display(df)

Database already exists. Skipping creation and population.

Departments table:

   id              name         location
0   1                HR         New York
1   2       Engineering    San Francisco
2   3         Marketing          Chicago
3   4             Sales      Los Angeles
4   5           Finance           Boston
5   6  Customer Support           Dallas
6   7          Research          Seattle
7   8             Legal  Washington D.C.
8   9           Product           Austin
9  10        Operations           Denver

Employees table:

      id             name  age  department_id     salary   hire_date
0      1    Michael Allen   57              9  151012.98  2016-02-04
1      2        Nora Hall   23              8  186548.83  2018-01-27
2      3  Patricia Miller   49              5   43540.04  2020-06-07
3      4   Alice Martinez   48              7  131993.17  2021-01-21
4      5  Patricia Walker   59              5  167151.15  2020-05-24
..   ...              ...  ...            ...        ...         ...
195  196     Hannah Clark   31             10  195944.00  2017-11-08
196  197      Alice Davis   46              5  145584.16  2022-02-13
197  198     Charlie Hall   37              1   53690.40  2024-06-18
198  199     Alice Garcia   50              5   92372.26  2024-02-01
199  200      Laura Young   25              9   64738.56  2015-08-02

[200 rows x 6 columns]
```## 创建基础 Text to SQL 提示

现在我们已经设置了数据库，让我们创建一个基本的 Text to SQL 转换提示。一个好的提示应该包括：

1. 清晰的指令，说明我们希望模型执行的操作
2. 用户的查询
3. 数据库的模式，这样 Claude 就知道如何将用户的查询进行翻译```python
def get_schema_info(db_path):
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    schema_info = []

    # Get all tables
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
    tables = cursor.fetchall()

    for (table_name,) in tables:
        # Get columns for this table
        cursor.execute(f"PRAGMA table_info({table_name})")
        columns = cursor.fetchall()

        table_info = f"Table: {table_name}\n"
        table_info += "\n".join(f"  - {col[1]} ({col[2]})" for col in columns)
        schema_info.append(table_info)

    conn.close()
    return "\n\n".join(schema_info)


# Get the schema info
schema = get_schema_info(DATABASE_PATH)
print(schema)

Table: departments
  - id (INTEGER)
  - name (TEXT)
  - location (TEXT)

Table: employees
  - id (INTEGER)
  - name (TEXT)
  - age (INTEGER)
  - department_id (INTEGER)
  - salary (REAL)
  - hire_date (DATE)

现在我们有了模式信息，让我们创建一个基础提示：

def generate_prompt(schema, query):
    return f"""
        You are an AI assistant that converts natural language queries into SQL.
        Given the following SQL database schema:

        <schema>
        {schema}
        </schema>

        Convert the following natural language query into SQL:
        <query>
        {query}
        </query>

        Provide only the SQL query in your response, enclosed within <sql> tags.
    """


# Test the prompt
user_query = "What are the names of all employees in the Engineering department?"
prompt = generate_prompt(schema, user_query)
print(prompt)

        You are an AI assistant that converts natural language queries into SQL.
        Given the following SQL database schema:

        <schema>
        Table: departments
  - id (INTEGER)
  - name (TEXT)
  - location (TEXT)

Table: employees
  - id (INTEGER)
  - name (TEXT)
  - age (INTEGER)
  - department_id (INTEGER)
  - salary (REAL)
  - hire_date (DATE)
        </schema>

        Convert the following natural language query into SQL:
        <query>
        What are the names of all employees in the Engineering department?
        </query>

        Provide only the SQL query in your response, enclosed within <sql> tags.

现在让我们使用此提示配合 Claude API 来生成 SQL：

def generate_sql(prompt):
    response = client.messages.create(
        model=MODEL_NAME,
        max_tokens=1000,
        temperature=0,
        messages=[{"role": "user", "content": prompt}],
    )
    return response.content[0].text.strip()


# Generate SQL
result = generate_sql(prompt)
sql = result.split("<sql>")[1].split("</sql>")[0].strip()
print("Generated SQL:")
print(sql)

Generated SQL:
SELECT e.name 
FROM employees e
JOIN departments d ON e.department_id = d.id
WHERE d.name = 'Engineering';

让我们通过在数据库上运行来测试生成的 SQL：

def run_sql(sql):
    conn = sqlite3.connect(DATABASE_PATH)
    result = pd.read_sql_query(sql, conn)
    conn.close()
    return result


result = run_sql(sql)
print("Query result:")
display(result)

Query result:

                name
0         Nora Lewis
1       Oliver White
2      William Clark
3     Tina Rodriguez
4       Diana Taylor
5        Nora Taylor
6       Steve Taylor
7      Oliver Martin
8     Tina Rodriguez
9   Quentin Anderson
10      Julia Miller
11     Jane Thompson
12       Julia Clark
13       Diana White
14       Xena Garcia
15      Ulysses Hall
16       Diana Brown
17   Charlie Johnson
18     Michael Clark
19    Yannick Harris
```## 通过示例改进提示

我们的基本提示已经有效，但我们可以通过包含示例来使其更加高效。这种技术称为少样本学习（few-shot learning），通过提供具体的输入-输出对示例，帮助模型更好地理解任务。

让我们修改我们的 `generate_prompt` 函数，以包含一些示例：

```python
def generate_prompt(prompt_text, examples):
    # ...

确保示例与提示文本相关，以便模型能够更好地学习。 :::```python def generate_prompt_with_examples(schema, query): examples = ““” Example 1: List all employees in the HR department.</

SELECT e.name FROM employees e JOIN departments d ON e.department_id = d.id WHERE d.name = ‘HR’;

    Example 2:
    User: What is the average salary of employees in the Engineering department?
    SQL: SELECT AVG(e.salary) FROM employees e JOIN departments d ON e.department_id = d.id WHERE d.name = 'Engineering';

    Example 3:
    User: Who is the oldest employee?
    SQL: SELECT name, age FROM employees ORDER BY age DESC LIMIT 1;
"""

return f"""
    You are an AI assistant that converts natural language queries into SQL.
    Given the following SQL database schema:

    <schema>
    {schema}
    </schema>

    Here are some examples of natural language queries and their corresponding SQL:

    <examples>
    {examples}
    </examples>

    Now, convert the following natural language query into SQL:
    <query>
    {query}
    </query>

    Provide only the SQL query in your response, enclosed within <sql> tags.
"""

Test the new prompt

user_query = “What are the names and salaries of employees in the Marketing department?” prompt = generate_prompt_with_examples(schema, user_query) print(prompt)

    You are an AI assistant that converts natural language queries into SQL.
    Given the following SQL database schema:

    <schema>
    Table: departments

• id (INTEGER)
• name (TEXT)
• location (TEXT)

Table: employees - id (INTEGER) - name (TEXT) - age (INTEGER) - department_id (INTEGER) - salary (REAL) - hire_date (DATE)

    Here are some examples of natural language queries and their corresponding SQL:

    <examples>
    
    Example 1:
    <query>List all employees in the HR department.</<query>
    <output>SELECT e.name FROM employees e JOIN departments d ON e.department_id = d.id WHERE d.name = 'HR';</output>

    Example 2:
    User: What is the average salary of employees in the Engineering department?
    SQL: SELECT AVG(e.salary) FROM employees e JOIN departments d ON e.department_id = d.id WHERE d.name = 'Engineering';

    Example 3:
    User: Who is the oldest employee?
    SQL: SELECT name, age FROM employees ORDER BY age DESC LIMIT 1;

    </examples>

    Now, convert the following natural language query into SQL:
    <query>
    What are the names and salaries of employees in the Marketing department?
    </query>

    Provide only the SQL query in your response, enclosed within <sql> tags.

现在让我们使用改进后的提示来生成 SQL：

```python
# Generate SQL using the improved prompt
result = generate_sql(prompt)
sql = result.split("<sql>")[1].split("</sql>")[0].strip()
print("Generated SQL:")
print(sql)

# Run the generated SQL
result = run_sql(sql)
print("\nQuery result:")
display(result)

Generated SQL:
SELECT e.name, e.salary 
FROM employees e 
JOIN departments d ON e.department_id = d.id 
WHERE d.name = 'Marketing';

Query result:

                name     salary
0          Zoe Jones  123511.58
1         Jane Smith  120291.41
2         John Young  179126.29
3       Charlie King   43323.02
4         Jane White   65134.81
5         Bob Harris   44083.34
6      Tina Robinson  131015.71
7       Steve Thomas  191563.64
8         Laura Hall  118691.73
9        Fiona Young  167114.79
10     Bob Rodriguez   64961.43
11       Diana Young  123255.78
12       John Harris  118778.51
13     Edward Taylor  112959.56
14  Michael Thompson  136840.04
15    William Taylor   49565.18
16        Diana King  154917.02
17        John Davis   46914.45
18      Julia Garcia   46486.44
19       Nora Wilson  153063.56
```通过在我们的提示中包含示例，我们使模型更好地理解了如何构建其响应的结构。这可能导致更准确和一致的SQL生成，特别是对于更复杂的查询。

**注意**：另一种可能值得探索的提示技术是在提示中包含数据库模式之外的一些真实数据行。这可能使Claude对数据结构和内容有更多的背景了解。

在下一节中，我们将探讨如何处理更复杂的查询，并使用思维链提示来改进模型的推理过程。

## 使用思维链提示

思维链提示鼓励模型将复杂问题分解为步骤。对于文本到SQL任务，这有助于处理需要多个操作或仔细考虑数据库模式的更复杂查询。

让我们修改我们的提示，以使用XML标签结合思维链推理：

```xml
<chain-of-thought>
  <step>首先，我们需要连接到数据库。</step>
  <step>然后，我们可以执行以下查询：</step>
  <step>SELECT * FROM table_name WHERE condition;</step>
</chain-of-thought>
``````python
def generate_prompt_with_cot(schema, query):
    examples = """
    <example>
    <query>List all employees in the HR department.</query>
    <thought_process>
    1. We need to join the employees and departments tables.
    2. We'll match employees.department_id with departments.id.
    3. We'll filter for the HR department.
    4. We only need to return the employee names.
    </thought_process>
    <sql>SELECT e.name FROM employees e JOIN departments d ON e.department_id = d.id WHERE d.name = 'HR';</sql>
    </example>

    <example>
    <query>What is the average salary of employees hired in 2022?</query>
    <thought_process>
    1. We need to work with the employees table.
    2. We need to filter for employees hired in 2022.
    3. We'll use the YEAR function to extract the year from the hire_date.
    4. We'll calculate the average of the salary column for the filtered rows.
    </thought_process>
    <sql>SELECT AVG(salary) FROM employees WHERE YEAR(hire_date) = 2022;</sql>
    </example>
    """

    return f"""You are an AI assistant that converts natural language queries into SQL.
    Given the following SQL database schema:

    <schema>
    {schema}
    </schema>

    Here are some examples of natural language queries, thought processes, and their corresponding SQL:

    <examples>
    {examples}
    </examples>

    Now, convert the following natural language query into SQL:
    <query>
    {query}
    </query>

    Within <thought_process> tags, explain your thought process for creating the SQL query.
    Then, within <sql> tags, provide your output SQL query.
    """


# Test the new prompt
user_query = "What are the names and hire dates of employees in the Engineering department, ordered by their salary?"
prompt = generate_prompt_with_cot(schema, user_query)
print(prompt)

You are an AI assistant that converts natural language queries into SQL.
    Given the following SQL database schema:

    <schema>
    Table: departments
  - id (INTEGER)
  - name (TEXT)
  - location (TEXT)

Table: employees
  - id (INTEGER)
  - name (TEXT)
  - age (INTEGER)
  - department_id (INTEGER)
  - salary (REAL)
  - hire_date (DATE)
    </schema>

    Here are some examples of natural language queries, thought processes, and their corresponding SQL:

    <examples>
    
    <example>
    <query>List all employees in the HR department.</query>
    <thought_process>
    1. We need to join the employees and departments tables.
    2. We'll match employees.department_id with departments.id.
    3. We'll filter for the HR department.
    4. We only need to return the employee names.
    </thought_process>
    <sql>SELECT e.name FROM employees e JOIN departments d ON e.department_id = d.id WHERE d.name = 'HR';</sql>
    </example>

    <example>
    <query>What is the average salary of employees hired in 2022?</query>
    <thought_process>
    1. We need to work with the employees table.
    2. We need to filter for employees hired in 2022.
    3. We'll use the YEAR function to extract the year from the hire_date.
    4. We'll calculate the average of the salary column for the filtered rows.
    </thought_process>
    <sql>SELECT AVG(salary) FROM employees WHERE YEAR(hire_date) = 2022;</sql>
    </example>
    
    </examples>

    Now, convert the following natural language query into SQL:
    <query>
    What are the names and hire dates of employees in the Engineering department, ordered by their salary?
    </query>

    Within <thought_process> tags, explain your thought process for creating the SQL query.
    Then, within <sql> tags, provide your output SQL query.

现在让我们使用带有 XML 标签的思维链提示来生成 SQL：

# Generate SQL using the chain-of-thought prompt
result = generate_sql(prompt)
print("Raw response from Claude:")
print(result)

# Extract thought process and SQL query
thought_process = result.split("<thought_process>")[1].split("</thought_process>")[0].strip()
sql = result.split("<sql>")[1].split("</sql>")[0].strip()

print("\nThought Process:")
print(thought_process)

print("\nGenerated SQL:")
print(sql)

# Run the generated SQL
query_result = run_sql(sql)
print("\nQuery result:")
display(query_result)

Raw response from Claude:
<thought_process>
1. We need to get information about employees in the Engineering department, so we'll need to join employees and departments tables
2. We'll match employees.department_id with departments.id
3. We'll filter for the Engineering department using departments.name
4. We need to select:
   - employee names (from employees table)
   - hire dates (from employees table)
5. The results should be ordered by salary
6. We'll use ORDER BY for the salary sorting
</thought_process>

<sql>
SELECT e.name, e.hire_date 
FROM employees e 
JOIN departments d ON e.department_id = d.id 
WHERE d.name = 'Engineering' 
ORDER BY e.salary;
</sql>

Thought Process:
1. We need to get information about employees in the Engineering department, so we'll need to join employees and departments tables
2. We'll match employees.department_id with departments.id
3. We'll filter for the Engineering department using departments.name
4. We need to select:
   - employee names (from employees table)
   - hire dates (from employees table)
5. The results should be ordered by salary
6. We'll use ORDER BY for the salary sorting

Generated SQL:
SELECT e.name, e.hire_date 
FROM employees e 
JOIN departments d ON e.department_id = d.id 
WHERE d.name = 'Engineering' 
ORDER BY e.salary;

Query result:

                name   hire_date
0     Tina Rodriguez  2019-09-16
1     Yannick Harris  2020-12-17
2        Nora Taylor  2016-08-24
3   Quentin Anderson  2024-02-12
4    Charlie Johnson  2019-01-16
5      Oliver Martin  2016-07-14
6        Julia Clark  2015-09-24
7       Diana Taylor  2023-07-16
8        Xena Garcia  2016-08-03
9      Jane Thompson  2016-06-27
10    Tina Rodriguez  2020-05-04
11     William Clark  2016-07-01
12       Diana White  2018-03-24
13      Oliver White  2017-03-27
14        Nora Lewis  2016-03-01
15      Julia Miller  2019-03-09
16       Diana Brown  2022-09-12
17      Ulysses Hall  2015-05-22
18     Michael Clark  2022-05-11
19      Steve Taylor  2020-01-26
```## 为复杂数据库模式实现 RAG

随着数据库的规模和复杂性不断增加，在每个提示中提供完整的模式变得不切实际。检索增强生成（RAG）可以通过根据用户的查询动态检索最相关的模式信息来帮助管理这种复杂性。

首先，让我们构建一个简单的 VectorDB 类，该类利用 VoyageAI 创建的嵌入模型：

```python
class VectorDB:
    def __init__(self, embedding_model):
        self.embedding_model = embedding_model

    def get_embedding(self, text):
        # 获取文本的嵌入表示
        return self.embedding_model.encode(text)

请注意，这里的 embedding_model 应该是一个已经训练好的嵌入模型实例。

接下来，我们将实现一个简单的 RAG 系统来处理查询并检索相关模式信息：

class RAG:
    def __init__(self, vector_db):
        self.vector_db = vector_db

    def get_relevant_schema(self, query):
        # 获取与查询最相关的模式信息
        query_embedding = self.vector_db.get_embedding(query)
        # ... 检索逻辑 ...
        return relevant_schema

通过这种方式，我们可以有效地管理大型和复杂的数据库模式，同时为用户提供更加精确和高效的查询结果。```python import json import os import pickle import sqlite3

import numpy as np import voyageai

class VectorDB: def init(self, db_path=“./data/vector_db.pkl”): self.client = voyageai.Client(api_key=os.getenv(“VOYAGE_API_KEY”)) self.db_path = db_path self.load_db()

def load_db(self):
    if os.path.exists(self.db_path):
        with open(self.db_path, "rb") as file:
            data = pickle.load(file)
        self.embeddings, self.metadata, self.query_cache = (
            data["embeddings"],
            data["metadata"],
            json.loads(data["query_cache"]),
        )
    else:
        self.embeddings, self.metadata, self.query_cache = [], [], {}

def load_data(self, data):
    if not self.embeddings:
        texts = [item["text"] for item in data]
        self.embeddings = [
            emb
            for batch in range(0, len(texts), 128)
            for emb in self.client.embed(
                texts[batch : batch + 128], model="voyage-2"
            ).embeddings
        ]
        self.metadata = [item["metadata"] for item in data]  # Store only the inner metadata
        self.save_db()

def search(self, query, k=5, similarity_threshold=0.3):
    if query not in self.query_cache:
        self.query_cache[query] = self.client.embed([query], model="voyage-2").embeddings[0]
        self.save_db()

    similarities = np.dot(self.embeddings, self.query_cache[query])
    top_indices = np.argsort(similarities)[::-1]

    return [
        {"metadata": self.metadata[i], "similarity": similarities[i]}
        for i in top_indices
        if similarities[i] >= similarity_threshold
    ][:k]

def save_db(self):
    with open(self.db_path, "wb") as file:
        pickle.dump(
            {
                "embeddings": self.embeddings,
                "metadata": self.metadata,
                "query_cache": json.dumps(self.query_cache),
            },
            file,
        )

Initialize and load schema data

vectordb = VectorDB() if not vectordb.embeddings: with sqlite3.connect(DATABASE_PATH) as conn: cursor = conn.cursor() cursor.execute(“SELECT name FROM sqlite_master WHERE type=‘table’;”) schema_data = [ { “text”: f”Table: {table[0]}, Column: {col[1]}, Type: {col[2]}“,”metadata”: {“table”: table[0], “column”: col[1], “type”: col[2]}, } for table in cursor.fetchall() for col in cursor.execute(f”PRAGMA table_info({table[0]})“).fetchall() ] vectordb.load_data(schema_data)

Test the search functionality

test_query = “What is the average salary of employees in each department?” results = vectordb.search(test_query) print(“Search results:”) for result in results: print(f”Similarity: {result[‘similarity’]}, Metadata: {result[‘metadata’]}“)

Search results: Similarity: 0.7318002364429477, Metadata: {‘table’: ‘employees’, ‘column’: ‘salary’, ‘type’: ‘REAL’} Similarity: 0.7284569547956667, Metadata: {‘table’: ‘employees’, ‘column’: ‘department_id’, ‘type’: ‘INTEGER’} Similarity: 0.6810496067975434, Metadata: {‘table’: ‘departments’, ‘column’: ‘name’, ‘type’: ‘TEXT’} Similarity: 0.6697669330753087, Metadata: {‘table’: ‘employees’, ‘column’: ‘name’, ‘type’: ‘TEXT’} Similarity: 0.6666317064533499, Metadata: {‘table’: ‘departments’, ‘column’: ‘location’, ‘type’: ‘TEXT’}

现在，让我们更新提示生成函数以使用 RAG：

```python
def generate_prompt_with_rag(query):
    relevant_schema = vectordb.search(query, k=10, similarity_threshold=0.3)
    schema_info = "\n".join(
        [
            f"Table: {item['metadata']['table']}, Column: {item['metadata']['column']}, Type: {item['metadata']['type']}"
            for item in relevant_schema
        ]
    )

    return generate_prompt_with_cot(schema_info, query)


# Test the RAG-based prompt
user_query = "What is the average salary of employees in each department?"
prompt = generate_prompt_with_rag(user_query)
print("Generated prompt:")
print(prompt)

# Generate and execute SQL
result = generate_sql(prompt)
print("\nGenerated result:")
print(result)

# Extract and run the SQL query
sql = result.split("<sql>")[1].split("</sql>")[0].strip()
print("\nExtracted SQL:")
print(sql)

query_result = run_sql(sql)
print("\nQuery result:")
display(query_result)

Generated prompt:
You are an AI assistant that converts natural language queries into SQL.
    Given the following SQL database schema:

    <schema>
    Table: employees, Column: salary, Type: REAL
Table: employees, Column: department_id, Type: INTEGER
Table: departments, Column: name, Type: TEXT
Table: employees, Column: name, Type: TEXT
Table: departments, Column: location, Type: TEXT
Table: employees, Column: id, Type: INTEGER
Table: departments, Column: id, Type: INTEGER
Table: employees, Column: age, Type: INTEGER
Table: employees, Column: hire_date, Type: DATE
    </schema>

    Here are some examples of natural language queries, thought processes, and their corresponding SQL:

    <examples>
    
    <example>
    <query>List all employees in the HR department.</query>
    <thought_process>
    1. We need to join the employees and departments tables.
    2. We'll match employees.department_id with departments.id.
    3. We'll filter for the HR department.
    4. We only need to return the employee names.
    </thought_process>
    <sql>SELECT e.name FROM employees e JOIN departments d ON e.department_id = d.id WHERE d.name = 'HR';</sql>
    </example>

    <example>
    <query>What is the average salary of employees hired in 2022?</query>
    <thought_process>
    1. We need to work with the employees table.
    2. We need to filter for employees hired in 2022.
    3. We'll use the YEAR function to extract the year from the hire_date.
    4. We'll calculate the average of the salary column for the filtered rows.
    </thought_process>
    <sql>SELECT AVG(salary) FROM employees WHERE YEAR(hire_date) = 2022;</sql>
    </example>
    
    </examples>

    Now, convert the following natural language query into SQL:
    <query>
    What is the average salary of employees in each department?
    </query>

    Within <thought_process> tags, explain your thought process for creating the SQL query.
    Then, within <sql> tags, provide your output SQL query.
    

Generated result:
<thought_process>
1. We need to get salary information from the employees table and department names from the departments table
2. We'll need to JOIN these tables using department_id and id
3. We want to group the results by department since we need averages per department
4. We need to show both the department name and the average salary
5. We'll use AVG() function to calculate the average salary
6. GROUP BY will be used on the department name to get per-department averages
</thought_process>

<sql>
SELECT d.name, AVG(e.salary) as average_salary
FROM employees e
JOIN departments d ON e.department_id = d.id
GROUP BY d.name;
</sql>

Extracted SQL:
SELECT d.name, AVG(e.salary) as average_salary
FROM employees e
JOIN departments d ON e.department_id = d.id
GROUP BY d.name;

Query result:

               name  average_salary
0  Customer Support   140937.064615
1       Engineering   131710.790500
2           Finance   134658.368800
3                HR   106710.123750
4             Legal   127712.166667
5         Marketing   109579.914500
6        Operations   141913.497619
7           Product    99374.619167
8          Research   115082.171667
9             Sales   136531.167826
```## 实现查询自我改进

Here, we'll build a self-improvement loop with Claude. This lets Claude execute the SQL it generates, analyze the results or errors, and improve the query if necessary.

This technique helps with:

1. Error Handling: It can catch and respond to SQL syntax errors or other execution issues.
2. Iterative Refinement: The model can learn from its mistakes and thought processes, and improve its output.
3. Performance: It increases the likelihood of generating a valid and executable SQL query.

In practice, you might want to adjust the `max_attempts` value based on your specific use case and performance requirements.

Let's start by creating a function that tries to execute the SQL and provides feedback:

```python
def execute_sql(sql):
    # Your code to execute SQL and provide feedback
    pass

请注意，以下代码示例仅供参考，具体实现可能需要根据您的环境和需求进行调整。 :::```python def execute_sql_with_feedback(sql): try: result = run_sql(sql) return True, result, “Query executed successfully.” except Exception as e: return False, None, str(e)

def generate_prompt_with_self_improvement(query, max_attempts=3): feedback = None sql = None

for attempt in range(max_attempts):
    if attempt == 0:
        prompt = generate_prompt_with_rag(query)
    else:
        prompt = f"""
        The previous SQL query resulted in this error: {feedback}
        Analyze the error and provide an improved SQL query.
        Original query: {sql}
        Explain your changes in <thought_process> tags and provide the corrected SQL in <sql> tags.
        """

    response = generate_sql(prompt)
    sql = response.split("<sql>")[1].split("</sql>")[0].strip()

    print(f"\nAttempt {attempt + 1}:")

    success, result, feedback = execute_sql_with_feedback(sql)
    if success:
        print("SQL executed successfully!")
        return sql, result
    else:
        print("SQL failed to execute")

print("Maximum attempts reached. Could not generate a valid SQL query.")
return None, None

Test the self-improving SQL generation

user_query = “For each department, show the ratio of the highest paid employee’s salary to the lowest paid employee’s salary, but only for departments where this ratio is greater than 3” final_sql, result = generate_prompt_with_self_improvement(user_query)

if final_sql: print(“SQL query:”) print(final_sql) print(“result:”) display(result) else: print(“Failed to generate a valid SQL query.”)

Attempt 1: SQL executed successfully!

Final SQL query: SELECT d.name, MAX(e.salary) / MIN(e.salary) as salary_ratio FROM departments d JOIN employees e ON d.id = e.department_id GROUP BY d.id, d.name HAVING MAX(e.salary) / MIN(e.salary) > 3;

Query result:

           name  salary_ratio

0 HR 4.542060 1 Engineering 3.106824 2 Marketing 4.421752 3 Sales 4.750024 4 Finance 4.474439 5 Customer Support 4.520068 6 Research 3.420039 7 Legal 4.888909 8 Product 3.687899 9 Operations 4.357528

## 评估

评估文本到SQL系统并不总是直截了当的。一个SQL查询可能写得正确，但给出的答案不正确，或者它可能工作，但不是获取结果的最佳方式。有些查询简单易生成，而有些则非常复杂。

在这里，我们将探讨使用[Promptfoo](https://promptfoo.dev)，一个开源的LLM评估工具包，来构建评估框架。要开始，请查看`./evaluation/README.md`。

我们框架中的测试难度各异，并在多个维度上评估系统：

1. **模型类型**：我们评估了Claude 3 Haiku和Claude 3.5 Sonnet。
2. **语法**：我们检查生成的SQL是否语法正确，并且可以无错误地执行。这是任何文本到SQL系统的基本要求。
3. **查询语义**：我们验证生成的SQL是否正确捕捉了自然语言查询的意图——例如，正确的表选择、适当的连接、以及聚合和分组的适当使用。
4. **结果正确性**：我们在测试数据库上执行生成的SQL，并将结果与预期结果进行比较。这确保了查询不仅看起来正确，而且产生了正确的数据。
5. **处理复杂查询**：我们测试系统处理越来越复杂查询的能力，包括多表连接、子查询、窗口函数和复杂的聚合。

### 结果

我们将使用以下代码来可视化我们的结果。您可以看到，Claude 3.5 Sonnet在采用我们介绍的技术时，更一致地通过了我们的断言，但它可能比Claude 3 Haiku产生略微更高的延迟、成本或token使用。

您可以使用此类评估来根据您的用例做出适当的权衡。

```python
# 示例代码
def evaluate_system(query):
    # 评估系统的函数
    pass

# 使用示例
evaluate_system("SELECT * FROM users WHERE age > 18")

请注意，上述代码仅为示例，实际评估逻辑可能更复杂。 :::```python import matplotlib.pyplot as plt import numpy as np

Data

models = [ “3H Basic”, “3H Few-Shot”, “3H CoT+Few-Shot”, “3H RAG+Few-Shot+CoT”, “3.5S Basic”, “3.5S Few-Shot”, “3.5S CoT+Few-Shot”, “3.5S RAG+Few-Shot+CoT”,] pass_rates = [83.3, 77.8, 88.9, 88.9, 88.9, 94.4, 100, 100] avg_latencies = [1561, 1758, 2187, 2564, 1887, 2297, 3900, 4614] avg_tokens = [383, 572, 765, 1001, 309, 496, 765, 984] costs = [0.0023, 0.0027, 0.0034, 0.0044, 0.020, 0.024, 0.041, 0.050]

Normalize the data

max_latency = max(avg_latencies) max_tokens = max(avg_tokens) max_cost = max(costs) norm_latencies = [latency / max_latency * 100 for latency in avg_latencies] norm_tokens = [t / max_tokens * 100 for t in avg_tokens] norm_costs = [c / max_cost * 100 for c in costs]

Set up the plot

x = np.arange(len(models)) width = 0.2

fig, ax = plt.subplots(figsize=(15, 8))

Create the bars

ax.bar(x - 1.5 * width, pass_rates, width, label=“Pass Rate (%)”, color=“green”) ax.bar(x - 0.5 * width, norm_latencies, width, label=“Avg. Latency (normalized)”, color=“blue”) ax.bar(x + 0.5 * width, norm_tokens, width, label=“Avg. Tokens (normalized)”, color=“orange”) ax.bar(x + 1.5 * width, norm_costs, width, label=“Cost (normalized)”, color=“red”)

Customize the plot

ax.set_ylabel(“Percentage / Normalized Value”) ax.set_title(“Comparison of Model Performance”) ax.set_xticks(x) ax.set_xticklabels(models, rotation=45, ha=“right”) ax.legend()

Add value labels on the bars

def add_labels(rects): for rect in rects: height = rect.get_height() ax.annotate( f”{height:.1f}“, xy=(rect.get_x() + rect.get_width() / 2, height), xytext=(0, 3), # 3 points vertical offset textcoords=”offset points”, ha=“center”, va=“bottom”, rotation=90, )

add_labels(ax.containers[0]) # Only add labels for pass rates

plt.tight_layout() plt.show()

<Figure size 1500x800 with 1 Axes> ```### 运行 Promptfoo

您可以使用命令 promptfoo eval 运行评估。您可以通过命令 promptfoo view 渲染评估结果。以下是结果的预览：

进一步探索和后续步骤

本指南涵盖了使用 Claude 构建文本到 SQL 系统的基础。以下是一些探索方向，可以帮助您改进解决方案：

优化检索性能

随着数据库的增长，确保您的 RAG 系统能够找到最相关和最新的信息非常重要：

1. 近期使用过滤器：通过只包括在设定时间段内被查询了一定次数的表来关注活跃维护和新鲜数据。

2. 查询频率排名：通过根据生产环境中表被查询的频率对 RAG 结果进行排名，优先考虑更常用的数据。

3. 定期更新：设置一个系统，在模式更改时更新您的向量数据库，使其与数据库保持最新。

4. 上下文感知嵌入：尝试嵌入表关系和用法模式，以及它们的名称，以改善搜索相关性。

向提示添加更多上下文

在提示中提供有关您的数据结构和内容的信息，除了数据库模式之外，可以帮助 Claude 生成更好的查询：

1. 数据样本：在您的提示中包括每个相关表的几行实际数据。例如：<data_sample> Sample data for employees table: id | name | age | department_id | salary | hire_date 1 | John Doe | 35 | 2 | 75000.0 | 2020-01-15 2 | Jane Smith | 28 | 3 | 65000.0 | 2021-03-01 </data_sample>2. 列统计信息：为每个列添加有用的信息，例如：

   • 空值百分比
   • 数字列的最小值和最大值
   • 类别列中最常见的值

2. 数据质量备注：提及任何已知的问题或特定列或表的特殊情况。

3. 数据目录信息：如果您使用dbt等工具来管理数据，请包括数据目录中的相关细节：

   • dbt .yml文件中的表概览
   • 列说明，包括计算或业务规则
   • 表关系和数据血缘
   • 数据质量检查和预期
   • 时间敏感查询的更新频率

4. 业务背景：添加有关数据在您的组织中的使用情况、常见分析类型或从数据中得出的重要业务指标的信息。

---

Tip

中文读者提示：本章节的代码和输出为英文原文。如需理解具体实现细节，可参考上方中文导读和代码注释。如有疑问，欢迎在 GitHub Issues 讨论。