Databricksのai_query関数を外部モデル・カスタムモデルを指定して利用する

カスタムモデル周りはハマってしまった。。。

# 導入

以下の記事で、Databricksの```ai_query```関数がバッチ推論やサービングエンドポイントを利用できるようになったのを知りました。

https://qiita.com/taka_yayoi/items/00c736d656e0cf8434bb

同様のことをやろうとするとPythonでSparkのUDFを作ったりして・・・ということをしないといけなかったのでかなり便利です。SQL単体でできるし。

元記事では[プロビジョン済みスループットのモデルサービングエンドポイント](https://docs.databricks.com/ja/machine-learning/foundation-models/deploy-prov-throughput-foundation-model-apis.html)を利用されていましたが、[ドキュメント](https://docs.databricks.com/ja/sql/language-manual/functions/ai_query.html)を読むにOpenAIなどの外部モデルやMLflowのカスタムpyfuncモデルのエンドポイントも利用できるようなのでそちらを試してみました。

結構ハマりポイントが多かったので、備忘録も兼ねて記事に残します。

# Step1. 機能有効化

qi_queryからの外部モデルやカスタムモデルの利用については、「プレビュー」メニューから以下の項目を有効化しておく必要があります。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/cd263664-b016-f7cc-2540-d186c56f2356.png)

当初オフになっているのに気づかず、```ai_query```実行時に403エラーが出てハマってました。
上記記事でもしっかり解説されてますので、ちゃんと読みましょう。

# Step2. 外部モデルを利用してみる

外部モデルのサービングエンドポイントを利用して、```ai_query```を実行します。

まずは試験用のエンドポイントをサービングメニューから作成。
以下のような感じで、Azure上のOpenAI APIを利用するエンドポイントを作成しました。
（画像には出ていませんが、gpt-4o-miniを利用しています)

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/a68e08eb-a705-cf21-e71c-6b9c3a4c834e.png)

こちらのエンドポイントを利用して```qi_query```を実行してみます。
Databricksのサンプルデータとして提供されている```samples.tpch.part```テーブルを利用して商品名を日本語に翻訳してみました。

```sql
WITH tmp AS (
  SELECT p_name FROM samples.tpch.part LIMIT 10
)

SELECT
  p_name,
  ai_query(
    'test-external-openai-endpoint',
    CONCAT(
      '次の商品名・商品説明を日本語に翻訳してください。翻訳結果のみ出力してください。\n\n### 商品説明：',
      p_name
    )
  )
  AS translated
FROM
  tmp
;
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/1978fd4e-362f-9a34-2551-9a598803b581.png)

内容の妥当感はともかく、ちゃんと処理された結果が出力されました。

# Step3. カスタムChatModelのモデルを利用してみる

次にMLflowのカスタムモデルを利用してみます。
MLflowの[カスタムChatModel](https://qiita.com/isanakamishiro2/items/184f7ee8a1d00cc01fea)を使ったサービングエンドポイントならそのまま使えるかなと思い、以下の記事で作成したエンドポイントを流用します。

https://qiita.com/isanakamishiro2/items/885219c81aee1e2aea04

エンドポイント名だけを変えて実行。

```sql
WITH tmp AS (
  SELECT p_name FROM samples.tpch.part LIMIT 10
)

SELECT
  p_name,
  ai_query(
    'gemma-2-2b-jpn-it-endpoint',
    CONCAT(
      '次の商品名・商品説明を日本語に翻訳してください。翻訳結果のみ出力してください。\n\n### 商品説明：',
      p_name
    )
  )
  AS translated
FROM
  tmp
;
```

結果、ダメでした。

```:エラー
[AI_FUNCTION_MODEL_SCHEMA_PARSE_ERROR] Failed to parse the schema for the serving endpoint "gemma-2-2b-jpn-it-endpoint": Failed to convert the schema object returned from the model endpoint into a valid SQL data type, response JSON was: \"{"openapi":"3.1.0","info":{"title":"gemma-2-2b-jpn-it-endpoint","version":"2"},"servers":[{"url":"https://tokyo.cloud.databricks.com/serving-endpoints/gemma-2-2b-jpn-it-endpoint"}],"paths":{"/served-models/gemma-2-2b-jpn-it-2/invocations":{"post":{"requestBody":{"content":{"application/json":{"schema":{"oneOf":[{"type":"object","properties":{"dataframe_split":{"type":"object","properties":{"index":{"type":"array","items":{"type":"integer"}},"columns":{"description":"required fields: messages","type":"array","items":{"type":"string","enum":["messages","temperature","max_tokens","stop","n","stream","top_p","top_k","frequency_penalty","presence_penalty","tools","metadata"]}},"data":{"type":"array","items":{"type":"array","prefixItems":[{"type":"array","items":{"required":["role"],"type":"object","properties":{"name":{"type":"string"},"role":{"type":"string"},"tool_calls":{"type":"array","items":{"required":["function","id","type"],"type":"object","properties":{"function":{"required":["arguments","name"],"type":"object","properties":{"arguments":{"type":"string"},"name":{"type":"string"}}},"id":{"type":"string"},"type":{"type":"string"}}}},"refusal":{"type":"string"},"content":{"type":"string"},"tool_call_id":{"type":"string"}}}},{"type":"number","format":"double"},{"type":"integer","format":"int64"},{"type":"array","items":{"type":"string"}},{"type":"integer","format":"int64"},{"type":"boolean"},{"type":"number","format":"double"},{"type":"integer","format":"int64"},{"type":"number","format":"double"},{"type":"number","format":"double"},{"type":"array","items":{"required":["type"],"type":"object","properties":{"function":{"required":["name","parameters"],"type":"object","properties":{"description":{"type":"string"},"name":{"type":"string"},"parameters":{"required":["properties"],"type":"object","properties":{"additionalProperties":{"type":"boolean"},"properties":{"type":"object"},"required":{"type":"array","items":{"type":"string"}},"type":{"type":"string"}}},"strict":{"type":"boolean"}}},"type":{"type":"string"}}}},{"type":"object"}]}}}}}},{"type":"object","properties":{"dataframe_records":{"type":"array","items":{"required":["messages"],"type":"object","properties":{"messages":{"type":"array","items":{"required":["role"],"type":"object","properties":{"name":{"type":"string"},"role":{"type":"string"},"tool_calls":{"type":"array","items":{"required":["function","id","type"],"type":"object","properties":{"function":{"required":["arguments","name"],"type":"object","properties":{"arguments":{"type":"string"},"name":{"type":"string"}}},"id":{"type":"string"},"type":{"type":"string"}}}},"refusal":{"type":"string"},"content":{"type":"string"},"tool_call_id":{"type":"string"}}}},"frequency_penalty":{"type":"number","format":"double"},"n":{"type":"integer","format":"int64"},"max_tokens":{"type":"integer","format":"int64"},"top_p":{"type":"number","format":"double"},"presence_penalty":{"type":"number","format":"double"},"temperature":{"type":"number","format":"double"},"tools":{"type":"array","items":{"required":["type"],"type":"object","properties":{"function":{"required":["name","parameters"],"type":"object","properties":{"description":{"type":"string"},"name":{"type":"string"},"parameters":{"required":["properties"],"type":"object","properties":{"additionalProperties":{"type":"boolean"},"properties":{"type":"object"},"required":{"type":"array","items":{"type":"string"}},"type":{"type":"string"}}},"strict":{"type":"boolean"}}},"type":{"type":"string"}}}},"stop":{"type":"array","items":{"type":"string"}},"stream":{"type":"boolean"},"metadata":{"type":"object"},"top_k":{"type":"integer","format":"int64"}}}}}}]}}}},"responses":{"200":{"description":"Successful operation","content":{"application/json":{"schema":{"type":"object","properties":{"predictions":{"type":"array","items":{"type":"object","properties":{"model":{"type":"string"},"choices":{"type":"array","items":{"required":["finish_reason","index","message"],"type":"object","properties":{"finish_reason":{"type":"string"},"index":{"type":"integer","format":"int64"},"message":{"required":["role"],"type":"object","properties":{"name":{"type":"string"},"role":{"type":"string"},"tool_calls":{"type":"array","items":{"required":["function","id","type"],"type":"object","properties":{"function":{"required":["arguments","name"],"type":"object","properties":{"arguments":{"type":"string"},"name":{"type":"string"}}},"id":{"type":"string"},"type":{"type":"string"}}}},"refusal":{"type":"string"},"content":{"type":"string"},"tool_call_id":{"type":"string"}}}}}},"usage":{"required":["completion_tokens","prompt_tokens","total_tokens"],"type":"object","properties":{"completion_tokens":{"type":"integer","format":"int64"},"prompt_tokens":{"type":"integer","format":"int64"},"total_tokens":{"type":"integer","format":"int64"}}},"object":{"type":"string"},"id":{"type":"string"},"metadata":{"type":"object"},"created":{"type":"integer","format":"int64"}}}}}}}}}}}}}}\".
Set the `returnType` parameter manually in the AI_QUERY function to override schema resolution. SQLSTATE: 2203G
```

エラーメッセージを見るに、```returnType```のスキーマ指定が必須の模様。

エラー内容に基づいて```returnType```を指定、その後は入力スキーマの指定でもエラーが出るなど試行錯誤し、最終的には以下のSQLの実行まで行いました。
ただし、こちらでもエラー。

```sql
WITH tmp AS (
  SELECT
    p_name
  FROM
    samples.tpch.part
  LIMIT
    10
)
SELECT
  p_name,
  ai_query(
    'gemma-2-2b-jpn-it-endpoint',
    request => named_struct(
      'messages',
      array(
        named_struct(
          'role',
          'user',
          'content',
          CONCAT(
            '次の商品名・商品説明を日本語に翻訳してください。翻訳結果のみ出力してください。\n\n### 商品説明：',
            p_name
          )
        )
      )
    ),
    returnType => 'STRUCT<id:STRING, object:STRING, created:LONG, model:STRING, choices:ARRAY<STRUCT<finish_reason:STRING, index:LONG, message:STRUCT<content:STRING, role:STRING>>>, usage:STRUCT<completion_tokens:LONG, prompt_tokens:LONG, total_tokens:LONG>>'
  ) AS translated
FROM
  tmp;
```

```:エラー
[REMOTE_FUNCTION_HTTP_RESULT_PARSE_ERROR] Failed to evaluate the ai_query SQL function due to inability to parse the JSON result from the remote HTTP response; the error message is 'Can not parse the response from model serving endpoint. ai_query expect the model prediction output to be arrays, and the array size must be 1. However, the received predictions value is: {"model":"training.llm.gemma-2-2b-jpn-it","choices":[{"index":0,"message":{"role":"assistant","content":"ミステリーのような輝く青緑色。 \\n\\n\\n\\n<end_of_turn>"},"finish_reason":"stop"}],"usage":{"prompt_tokens":39,"completion_tokens":13,"total_tokens":52},"object":"chat.completion","id":"2","created":1728808851}.'. Check API documentation: https://docs.databricks.com/en/machine-learning/model-serving/index.html. Please fix the problem indicated in the error message and retry the query again. SQLSTATE: 22032
```

どうも、カスタムモデルだと出力は要素が一つのリスト(ARRAY)の結果である必要がある模様。
ChatModelだと出力が単一の辞書型になるので、ここで一旦断念。

通常のPyfuncモデルだと出力がリストになるので、現時点ではそこを想定した仕様になっているのだと思います。

# Step4. カスタムChatModelのモデルを利用してみる(リトライ)

カスタムChatモデルをそのまま利用できないなら、**カスタムモデルエンドポイントを外部モデルのエンドポイントとして使えば** いけるんじゃないか？と思ったのでトライ。

外部モデルのプロバイダとしてDatabricks Model Servingを指定し、上記のカスタムモデル```gemma-2-2b-jpn-it-endpoint```を利用するように外部モデルエンドポイントを作成。
そう、カスタムChatModelで作成したモデルは外部モデルのエンドポイントとして利用できるようです。（今回初めて知った）

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/4c48c8dd-93ca-e4b0-db96-0f05f9790ccf.png)

作成したエンドポイントを使って```ai_query```を実行。

```sql
WITH tmp AS (
  SELECT p_name FROM samples.tpch.part LIMIT 10
)

SELECT
  p_name,
  ai_query(
    'test-external-model-endpoint',
    CONCAT(
      '次の商品名・商品説明を日本語に翻訳してください。翻訳結果のみ出力してください。\n\n### 商品説明：',
      p_name
    )
  )
  AS translated
FROM
  tmp
;
```

通常の外部モデルエンドポイントと同様に動作します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/1e9dfa1c-0be8-1ab3-c248-cd9c1649d4b2.png)

こちらも結果内容はともかく、動きました！

# まとめ

```ai_query```を外部モデル/カスタムモデルを使って試してみました。
SQLで手軽にLLMのパワーを発揮できるのは非常に良いですね。カスタムモデルも利用できるので、自作のエージェントを用いたバッチ実行など応用範囲も広そうです。

エージェント関連はもっと勉強して、いろんなエージェントを今後試作してみたいと思っています。