Databricksのai_query関数の基礎から高度な使い方

`ai_query`のリファレンスで紹介されていた使い方をウォークスルーします。

https://docs.databricks.com/aws/ja/sql/language-manual/functions/ai_query

# ai_query関数とは

Databricksで大量のテキストデータにLLMを適用する際に使える汎用型の関数です。モデルサービングエンドポイントとプロンプトを指定することで、SQLの文脈で要約や翻訳などLLMタスクを柔軟にこなすことができます。SQLの中で使えるので、`CREATE TABLE AS SELECT`(CTAS)を通じて、大量のテキストデータに対してバッチのLLM推論処理を高速に行うことができます。

https://docs.databricks.com/aws/ja/large-language-models/ai-query-batch-inference

基本的な構文はこちらです。`endpoint`はLLMがデプロイされているモデルサービングポイント名、`request`はプロンプトを含むLLMへの指示を示すテキストです。

```sql
ai_query(endpoint, request)
```

# 準備

LLMがデプロイされたモデルサービングエンドポイントが必要です。今回は`system.ai.llama_v3_3_70b_instruct`をデプロイした[プロビジョン済みスループット](https://docs.databricks.com/aws/ja/machine-learning/foundation-model-apis/deploy-prov-throughput-foundation-model-apis)のモデルサービングエンドポイントを使用しています。

![Screenshot 2025-04-21 at 19.43.05.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/d939b3d0-bde4-4fe5-890f-9d686ff47bc4.png)


# 基本的な使い方

ある意味一番シンプルな例です。

```sql
SELECT ai_query(
    'taka-llama-pt',
    '100文字でDatabricks SQLを説明して'
  ) AS summary
```

![Screenshot 2025-04-21 at 19.45.33.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/ca2c7957-58f0-4bc9-94a9-0dd0b866cb28.png)

# モデルパラメータの指定

引数`modelParameters`でモデルのトークン数やtemperatureのようなパラメータを指定できます。

```sql
SELECT contents, ai_query(
    "taka-llama-pt",
    "以下の文を要約して: " || contents,
    modelParameters => named_struct('max_tokens', 100, 'temperature', 0.7)
) AS summary
FROM takaakiyayoi_catalog.doc_maintenance.databricks_docs LIMIT 1;
```
![Screenshot 2025-04-21 at 19.46.41.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/7ff66306-8a6c-4dc1-ac55-fa1bb35236aa.png)

# 出力の構造化による出力スキーマの強制

これは強力な機能です。通常ですとLLMのレスポンスは構造化されていないテキストですが、戻り値の構造を指定することで、それに合わせたレスポンスを得ることができます。後段でデータを加工したい場合に便利です・

`responseFormat`で戻り値の構造を指定します。

```sql
SELECT ai_query(
    "taka-llama-pt",
    "以下の文からマニュアルの詳細情報を抽出して: " || contents,
    responseFormat => 'STRUCT<doc_extraction:STRUCT<title:STRING, products:ARRAY<STRING>, summary:STRING, keywords:ARRAY<STRING>>>'
) AS extracted_info
FROM takaakiyayoi_catalog.doc_maintenance.databricks_docs LIMIT 2;
```

![Screenshot 2025-04-21 at 19.48.30.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/a7072df5-61e1-4a36-ab73-dd69390ec67e.png)

JSONとして出力されていますが文字列なので、JSONとしてパースします。

```py
SELECT parse_json(ai_query(
    "taka-llama-pt",
    "以下の文からマニュアルの詳細情報を抽出して: " || contents,
    responseFormat => 'STRUCT<doc_extraction:STRUCT<title:STRING, products:ARRAY<STRING>, summary:STRING, keywords:ARRAY<STRING>>>'
)) AS extracted_info
FROM takaakiyayoi_catalog.doc_maintenance.databricks_docs LIMIT 2;
```
![Screenshot 2025-04-21 at 19.49.19.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/b2c02738-902a-445f-b606-622fb3692e9f.png)

# UDFでのai_queryの使用

関数の中に`ai_query`を含めることができますので、処理をモジュール化できます。

```sql
CREATE FUNCTION takaakiyayoi_catalog.doc_maintenance.create_title(text STRING)
  RETURNS STRING
  RETURN ai_query(
    'taka-llama-pt',
    CONCAT('次の記事から20文字程度のタイトルを作成して: \n', text));
```

```sql
SELECT
    title,
    takaakiyayoi_catalog.doc_maintenance.create_title(contents) AS generated_title
  FROM takaakiyayoi_catalog.doc_maintenance.databricks_docs LIMIT 5;
```
|title|generated_title|
|---|---|
|Databricks管理の概要 \| Databricks Documentation|Databricks管理の概要と管理者の役割|
|外部Databricksサービスへの認証を有効にする \| Databricks Documentation|外部Databricksサービスへの認証を有効にする|
|個人用アクセス トークンの監視と取り消し \| Databricks Documentation|パーソナルアクセストークンの監視と取り消し|
|ステップ 4: ログ配信設定を作成する \| Databricks Documentation|Databricksのログ配信設定を作成する方法|
|ステップ 2: 監査ログ配信の認証情報を設定する \| Databricks Documentation|監査ログ配信の認証情報の設定方法|

# CTASによる処理結果の保存

これまでの処理では、結果が保存されていませんでした。結果を保存するにはCREATE TABLE AS SELECT(CTAS)を使うのが便利です。

```sql
CREATE TABLE takaakiyayoi_catalog.doc_maintenance.generated_titles AS
SELECT
  title,
  takaakiyayoi_catalog.doc_maintenance.create_title(contents) AS generated_title
FROM
  takaakiyayoi_catalog.doc_maintenance.databricks_docs
LIMIT 50;
```

テーブルに結果が保存されました。

![Screenshot 2025-04-21 at 19.55.42.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/109afff4-4d88-4794-a664-f45fc7a5eedd.png)


### はじめてのDatabricks

[はじめてのDatabricks](https://qiita.com/taka_yayoi/items/8dc72d083edb879a5e5d)

### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)