
Snowflake CortexのCOMPLETEでSQL文のみを取得できるか試してみた
データ事業本部の鈴木です。
Snowflake CortexのCOMPLETEでSQL文のみを取得できるか試してみたのでご共有します。
やりたいこと
最近は自然言語からSQLの作成をするのに凝っています。ネイティブの機能としてはCortex Analystがあり、Cortex AnalystはSQLを単体で返すことができます。
自分でSQL生成を行う場合は既存の研究や実装例を参考に、Completeを使ってLLMに生成させることになりますが、意外とSQLだけを返させるのが難しかったため、いくつか実現パターンを試してみました。
SQLだけ返してくれないと、生成したSQL文を使ったデータ処理の実行が自動でできません。具体的なイメージとしては、例えばStreamlitから自前のSQL生成処理実行を実行しても、回答内のSQL部分を手で抜き出してSQLワークシートから実行し...というようにちょっとテンポ感が悪いアプリになってしまいます。SQL文だけ確実に返してくれるのであれば、SnowparkからSQLを実行して、結果をデータフレームとして表示するもよし、作図用のライブラリでプロットを表示してもよしで、使いやすいアプリになりますね。
やってみた
選択肢としては以下の2つを試してみました。
- COMPLETEでシステムプロンプトでSQL文のみ返すように指示する
- COMPLETE Structured Outputsを使う
肌感としては後者を試すのがよいと思いました。
1. システムプロンプトでSQL文のみ返すように指示する
COMPLETEではシステムプロンプトの指定や過去の会話履歴を含めたテキスト生成ができます。
今回はシステムプロンプトにSQL文のみ出力させるよう指定することにします。
何も指定しないとコードブロックで回答が生成されるため、SQL文だけを生成するようにも指示しました。
import pandas as pd
from snowflake.cortex import complete, CompleteOptions
messages = []
messages.append({"role": "system", "content": "SQL文のみを回答してください。```sqlのようなコードブロックは付けないでください。"})
messages.append({"role": "user", "content": "Snowflakeの構文のSQLでなにか複雑なものを考えてください。"})
answer = complete(model='mistral-large2', prompt=messages)
以下のような回答が生成されました。
SELECT
employee_id,
first_name,
last_name,
department,
salary,
CASE
WHEN salary > 100000 THEN 'High'
WHEN salary BETWEEN 50000 AND 100000 THEN 'Medium'
ELSE 'Low'
END AS salary_level,
RANK() OVER (PARTITION BY department ORDER BY salary DESC) AS rank_in_department,
LAG(salary, 1) OVER (PARTITION BY department ORDER BY salary DESC) AS previous_salary,
LEAD(salary, 1) OVER (PARTITION BY department ORDER BY salary DESC) AS next_salary
FROM
employees
WHERE
hire_date > '2020-01-01'
ORDER BY
department, salary DESC;
2. COMPLETE Structured Outputsを使う
COMPLETEではStructured Outputsを使い、スキーマを指定することができます。このとき、キーをSQL文だと分かるものにすることで、SQL文を確実に取得できるようにします。
以下のように回答を生成しました。
import pandas as pd
from snowflake.cortex import complete, CompleteOptions
response_format = {
"type": "json",
"schema": {
"type": "object",
"properties": {
"sql": {
"type": "string",
"description": "SQL文"
},
"text": {
"type": "string",
"description": "SQL文の解説"
},
}
},
"required": ["sql", "text"],
}
prompt = [{
"role": "user",
"content": "snowflakeの構文のSQLでなにか複雑なものを考えてください。SQL文および解説文をください。"
}]
options = CompleteOptions(
max_tokens=4096,
guardrails=False,
response_format=response_format
)
result = complete(
model="mistral-large2",
prompt=prompt,
options=options,
)
output = "".join(result)
print(output)
以下の回答が生成されました。
{
"sql": "WITH RECURSIVE cte AS (\n SELECT 1 AS n\n UNION ALL\n SELECT n + 1\n FROM cte\n WHERE n < 10\n)\nSELECT n, n * 2 AS double_n\nFROM cte;",
"text": "このSQL文は、再帰的な共通テーブル式(CTE)を使用して、1から10までの数字を生成し、それぞれの数字の2倍を計算します。\n\n1. `WITH RECURSIVE cte AS (...)`: 再帰的なCTEを定義します。\n2. `SELECT 1 AS n`: 初期値として1を設定します。\n3. `UNION ALL`: 初期値と再帰的な部分を結合します。\n4. `SELECT n + 1 FROM cte WHERE n < 10`: 再帰的に次の数字を生成し、nが10未満の場合に繰り返します。\n5. `SELECT n, n * 2 AS double_n FROM cte`: 生成された数字とその2倍を選択します。\n\nこのクエリは、再帰的なCTEを使用して数列を生成し、その数列に対して計算を行う方法を示しています。"
}
response_format
は、descriptionも指定することができました。スキーマ情報だけだとLLMの出力を制御しにくい場合もあると思うので、上手く組み合わせられるのがよいですね。
response_format
はpydanticで定義することも可能です。詳しくはドキュメントを確認ください。
終わりに
Snowflake CortexのCOMPLETEでSQL文のみを取得できるか試してみた結果でした。
システムプロンプトかCOMPLETE Structured Outputsで、何かしらプログラム的にパース可能な形式でSQL文のみを生成することはできましたが、後者の方が自然そうでした。