Databricks Unity Catalogの関数: ユーザーとエージェントのためのツール

[Unlocking the Power of Mosaic AI: A Comprehensive Databricks Demo](https://community.databricks.com/t5/technical-blog/unity-catalog-functions-tools-for-users-and-agents/ba-p/98221)の翻訳です。

:::note warn
本書は著者が手動で翻訳したものであり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。
:::

誰でもビジネス問題を解決するために生成AIを活用したいと考えていますが、開発者はシンプルに大規模言語モデルをデプロイするだけでは、ゴールラインに到達するには十分ではないことに気づきました。ビジネス問題は複雑で、さまざまなタイプのタスクにシームレスに対応できるソリューションを必要とします。我々は、これらのソリューションを複合AIシステムと呼びます。

この記事では、Mosaic AIプラットフォームの能力を説明する注目すべきDatabricksのデモをウォークスルーします。このデモでは、ツールとしてのUnity Catalogの関数、すぐに利用できるAI関数、Mosaic AIエージェントフレームワークとエージェント評価をハイライトし、パワフルな複合AIシステムを構築するために、これらのコンポーネントがどのように協調するのかを説明します。これは、たくさんのように聞こえると思いますがその通りですが、以下の3ステップに従うことで、Databricksプラットフォームがこれら全てを簡単にするのかを説明します:

1. SQL、Pythonを用いたUnity Catalogの関数としてツールを定義
1. DatabricksのPlaygroundを通じたツールのテスト
1. あなたのソリューション全体をDatabricksノートブックに対するコードとしてエクスポートし、Mosaic AI Agent & Evaluation Frameworkと完全に連携し、モデルサービングにデプロイ可能

![](https://community.databricks.com/t5/image/serverpage/image-id/12777iFE4D907DEF008B2B/image-size/large?v=v2&px=999)

# ツールの定義

ツールは、テキストを生成する大規模言語モデルの能力を拡張します。文書の収集や、サポートチケットを作成するためのAPIコールなどのアクションを取るためにモデルを強化することで、AIがビジネスプロセスを実行できるようにする機能を追加する事が、競合に打ち勝つための唯一のパスとなります。[DSPy](https://dspy-docs.vercel.app/deep-dive/modules/react/?h=tool#react-tools)、[Langchain](https://python.langchain.com/docs/integrations/tools/databricks/)、[CrewAI](https://github.com/crewAIInc/crewAI)のようなツール呼び出しのためのライブラリやアプローチは多数存在しており、これらはDatabricksで稼働しますが、使いやすさと広範なDatabricksプラットフォームとの連携の観点から、MLflowからのUnity Catalog関数の活用にフォーカスします。

Unity Catalogの関数によって、ユーザーはSQLやPythonでカスタムロジックを定義できるようになります。これらの関数は、Unity Catalogの3階層名前空間に従うので、他のUnity Catalogのオブジェクト(テーブル、ボリューム、モデルなど)のようにセキュアに統治し、共有する事ができます。関数を実行するには、関数に対するEXECUTE権限と、スキーマとカタログに対するUSAGE権限が必要となります。このデモでは、我々のすべてのツールを格納するカタログとスキーマを作成しました。

# SQL関数: ビジネス洞察とAIを結びつける

AIの機能を活用するSQLで記述された3つのツールの紹介からスタートします: AI予測、類似検索、テキスト分類です。シンプルな方法で時系列の予測を行うために、[AI\_FORECAST](https://docs.databricks.com/en/sql/language-manual/functions/ai_forecast.html)を活用する事ができます。例えば、エンジンから送信される時間のタイムスタンプ、出力、IDの列を持つIOTデータを持つテーブルがあるとした場合、将来的なエネルギー出力を予測するために以下のコマンドを実行する事ができます:

```sql
SELECT *
FROM AI_FORECAST(
   TABLE(josh_melton.iot_schema.engine_output),
   horizon => date_add(current_date(), 1),
   time_col => 'hourly_timestamp',
   value_col => 'avg_output',
   group_col => 'engine_id'
)
```

しかし、これを他のユーザーやAIエージェントによりシームレスに使わせたいと考えた場合にはどうしたらいいのでしょうか？エージェントがいつこの関数を使うべきかを理解できるように、説明的な名称やコメントを用いるのように、[CREATE FUNCTION](https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-create-sql-function.html)でシンプルにラッピングする事ができます。

```sql
DROP FUNCTION IF EXISTS josh_melton.tools_demo.forecast_energy_output;
CREATE FUNCTION josh_melton.tools_demo.forecast_energy_output()
RETURNS TABLE (
   hourly_timestamp TIMESTAMP,
   engine_id STRING,
   avg_energy_forecast FLOAT,
   avg_energy_upper FLOAT,
   avg_energy_lower FLOAT
) COMMENT "Returns the forecasted energy output through the given date" RETURN
SELECT *
FROM AI_FORECAST(
   TABLE(josh_melton.iot_schema.engine_output),
   horizon => date_add(current_date(), 1),
   time_col => 'hourly_timestamp',
   value_col => 'avg_output',
   group_col => 'engine_id'
)
```

これで、エージェントが将来を予測できるように強化し、他のSQLユーザーが同じことをするにも2行を実行するだけで済むようになります:

```sql
SELECT *
FROM josh_melton.tools_demo.forecast_energy_output()
```

同様に、サポートチケットの緊急性を分類するために、[ai\_classify\(\)](https://docs.databricks.com/en/sql/language-manual/functions/ai_classify.html)を用いて別の関数を作成しましょう。`ai_classify()`関数は2つの引数を受け取ります: ラベル付されるべきコンテンツとラベルのリストです。

```sql
CREATE OR REPLACE FUNCTION josh_melton.tools_demo.ai_classify_urgency(ticket_text_input STRING)
RETURNS STRING
RETURN
SELECT ai_classify(ai_classify_urgency.ticket_text_input, ARRAY('very urgent', 'somewhat urgent', 'not urgent')) AS urgency
```

最後のSQLツールとして、類似検索関数を作成します。Dataricksでは、SQLを用いてベクトル検索インデックスにクエリーを行えるようにするSQLの[vector\_search\(\)](https://docs.databricks.com/en/sql/language-manual/functions/vector_search.html)関数を提供しています。これは3つの引数を受け取ります:

- index - 検索したいベクトル検索インデックス
- query - 検索文字列
- num_results(オプション) - 返却するレコード数

このデモでは、すでにベクトル検索インデックスを作成して、カスタマーサービスのチケットをロードしています。作成していない場合には、こちらの[ステップ](https://docs.databricks.com/ja/generative-ai/create-query-vector-search.html#create-a-vector-search-index)に従ってください。

入力としてクエリー文字列を受け取る類似検索関数の実装方法を以下に示します:

```sql
CREATE FUNCTION josh_melton.tools_demo.similarity_search (
 query STRING COMMENT "The string to search for similar tickets to" DEFAULT "Turbochargers malfunctioning due to overheating"
) RETURNS TABLE (
 ticket_number STRING,
 issue_description STRING
) COMMENT "Returns the support ticket issues related to the query" RETURN
SELECT ticket_number, issue_description
FROM VECTOR_SEARCH(index => "josh_melton.tools_demo.customer_service_tickets_index",
                 query => similarity_search.query,
                 num_results => 3)
```

# Python関数: 機能を拡張

Databricks Unity CatalogにおけるPython UDFは、従来のSQL以上に我々のデータ能力を拡張し、Pythonのパワーと柔軟性を活用する事ができます。次に、Pythonを用いた2つの関数、`tell_data_joke()`と`update_ticket()`を定義します。

最初のPythonツール`tell_data_joke()`では、外部APIを呼び出すHTTPコールを行うためにPythonのrequestsライブラリを活用します。この例では、ランダムにジョークを返却する*icanhazdadjoke.com* APIを呼び出します。これは、リアルタイムの情報を取得したり、他の適切なマイクロサービスとの連携が可能な、外部サービスとの連携能力をデモンストレーションするお手軽な例です。

```sql
CREATE FUNCTION josh_melton.tools_demo.tell_dad_joke()

RETURNS STRING COMMENT "Returns a dad joke"

LANGUAGE PYTHON

AS $$
import requests

url = "https://icanhazdadjoke.com/"
headers = {"Accept": "application/json"}
response = requests.get(url, headers=headers)

if response.status_code == 200:
   joke_data = response.json()
   return joke_data['joke']
else:
   return f"Failed to retrieve joke. Status code: {response.status_code}"
$$;
```

データのトークン化、データマスキング、複雑な計算のような様々なタスクを実行するために、我々のPython関数の中でAzure Databricksに含まれるすべての標準Pythonライブラリを活用する事ができます。二つ目の`update_ticket()`の例では、デモ目的なので、いくつかの文字列を変更したレスポンスをハードコードします。チケット番号と緊急度でこの関数が呼び出されると、変更がなされた旨を示す文字列を返却します。

```sql
CREATE FUNCTION josh_melton.tools_demo.update_ticket(
   ticket_number STRING COMMENT "The ticket number to update",
   urgency STRING COMMENT "The urgency or severity of the ticket"
)
RETURNS STRING COMMENT "Updates the ticket urgency status with the given ticket_number, or returns a status" 
LANGUAGE PYTHON
AS $$
return f"Ticket {ticket_number} updated with urgency '{urgency}'"
$$;
```

# エージェント: 全てをまとめ上げる

これで、すべてのツールが定義されたのでテストの時間となります！DatabricksのAI Playgroundは、LLMと我々のツールのインタラクションのための使いやすいUIを提供しています。我々の複合AIシステムを設定するために行なったことを、以下に示します:

1. **Databricks AI Playgroundにアクセス:** Databrikcsワークスペースの左側にある**機械学習**セクションにある**Playground**を選択するだけです。
1. **LLMを選択:** **ツールが有効**アイコンが表示されているツール呼び出し能力を持つLLMを選択する必要があります。このLLMは、我々の複合AIシステムの頭脳として動作します。
1. **ツールの追加:** シンプルに、ツールドロップダウンの矢印からツールの追加をクリックし、ホストされている関数を追加します。ここでは、`josh_melton.tools_demo.*`をタイプすることですべてのツールを追加する事ができます。

これで、追加したツールで拡張されたエージェントをテストする事ができます。特定のトピックに関係する問題について質問しようとすると、我々のモデルは質問に回答するために、類似検索関数からコンテキストを取得して最新の情報を活用します。

![](https://community.databricks.com/t5/image/serverpage/image-id/12775i8F2159BC40926A9B/image-size/large?v=v2&px=999)

同じように、2つ目のチケットの緊急度を分類するようにお願いすr事ができ、非常に緊急度が高い場合にはチケットのステータスを更新することもあり得ます。

汎用的なllama 70bモデルは単独で最もパワフルなLLMというわけではありませんが、「こちらがチケットの分類のツールなのでこのチケットを分類してください」とお願いしているユーザーが分類を行うためにそのツールを使って欲しいことや、高度に緊急なチケットはすぐに更新されるべきであるということを理解するために、PhDレベルのインテリジェンスを必要とはしません。小規模で高速、安価なモデルをツールで拡張することで、遅くて高価、汎用的な知識を持つモデルよりもさらなる価値を提供する事ができます。

# コードのアップデートとデプロイメント

これで、我々のエージェントを開発、テストしたので、他の人が使用できるようにデプロイしたいと思います。Playgroundの画面の上部にある「エクスポート」ボタンをクリックし、ドライバーノートブック(TODOコメントを参照ください)でカタログ、スキーマ、モデル名を追加し、「run all」をクリックします([詳細はこちら](https://docs.databricks.com/ja/generative-ai/agent-framework/create-agent.html#export-and-deploy-ai-playground-agents))。ベクトル検索インデックスのようなアセットへのアクセスが必要なツールが含まれている場合、それらを使用する権限をエージェントに与える必要があります(関数自体は自動で追加されます)。

```py
from mlflow.models.resources import DatabricksFunction, DatabricksServingEndpoint, DatabricksVectorSearchIndex

resources = [DatabricksServingEndpoint(endpoint_name=config.get("llm_endpoint")),
            DatabricksVectorSearchIndex(index_name="josh_melton.tools_demo.customer_service_tickets_index")]
```

ノートブックを実行すると、MLflowモデルとしてエージェントを記録、登録し、テストするためのフロントエンドインタフェースと共に[デプロイ](https://docs.databricks.com/ja/generative-ai/agent-framework/deploy-agent.html)します。デプロイメントのプロセスは約20分を要します。

# まとめ

DatabricksのAI Playgroundは、我々の複合AIシステムをテスト、洗練するための直感的なインタフェースを提供します。この環境によって、迅速なプロトタイピングやイテレーションが可能となり、開発者はデプロイメントの前にクイックに自分たちのAIソリューションを検証する事ができます。

ツールとしてUnity Catalogの関数を活用することで、どのように大規模言語モデルを単なるテキスト生成以上の能力に拡張できるのかを説明しました。これらのツールはSQLで実装しようがPythonでしようが、予測や分類、類似検索、外部APIとの連携のような固有のタスクをAIが実行できるようにします。

このデモを通じて、Mosaic AIプラットフォームはいくつかのキーエリアでの価値をデモンストレーションしました:

1. **柔軟性:** カスタムツールを定義することで、AIシステムを特定のビジネス要件やワークフローに適応させる事ができます。
1. **効率性:** ツールで拡張された小規模、高速なモデルは、多くの場合特定のタスクにおいて、大規模で汎用的なモデルを上回ります。
1. **コスト効率:** ターゲットされたツールを活用することで、高価な大規模言語モデルへの定常的なクエリーを削減する事ができます。
1. **スケーラビリティ:** ビジネス要件が進化したら、システムの能力を拡張するために新たなツールを追加する事ができます。

### はじめてのDatabricks

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

### Databricks無料トライアル

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