AI Playgroundによるエージェントシステムのプロトタイピング

よく理解できていなかったので、こちらの[AI Playgroundにおけるツール呼び出しエージェントのプロトタイプ](https://docs.databricks.com/ja/generative-ai/create-log-agent.html#prototype-tool-calling-agents-in-ai-playground)の機能を試してみます。

# どのような機能か

[AI Playground](https://docs.databricks.com/ja/large-language-models/ai-playground.html)は、様々なLLMに対するプロンプトエンジニアリングを容易に行うことができる環境です。最近、こちらにエージェントシステムをプロトタイピングする機能が追加されました。

エージェントが呼び出すツールをUnity Catalog配下の関数として定義することで、エージェントシステムから関数を呼び出して、複雑な挙動をさせることが可能になります。

マニュアルはこちら。

https://docs.databricks.com/ja/generative-ai/create-log-agent.html

全体的な流れは以下のようになります。

1. エージェントシステムの設計(どのようなツールを活用してどのような挙動をさせたいのか)
1. データの準備
1. ツールの実装
1. AI Playgroundでのプロトタイピング
1. エージェントノートブック、ドライバーノートブックのエクスポート
1. ドライバーノートブックを実行してエージェントシステムを構築

今回はエクスポート以降には踏み込みません。別の記事で説明する予定です。

# エージェントシステムの設計

RAG(Retrieval Augmented Generation)もエージェントシステムの一種ですが、同じことをやっても詮無いので今回はクッキーの売り上げデータを取り扱うエージェントシステムを組んでみます。今年のData + AIサミットでデモされたものです。

私自身、エージェントシステムは勉強中ですが、以下のようなツールを用いたエージェントが今後増えてくるのではないかと思っています。

- LLMに含まれていないデータに基づく挙動をするエージェント(RAGもこの一種です)
- 外部APIを呼び出し、その結果に基づいた挙動をするエージェント

![Screenshot 2024-10-20 at 11.52.19.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/04647b8d-fdbe-e355-4efb-2623a9cb5de0.png)

# データの準備

ツールからアクセスするデータを格納するテーブルを作成します。フランチャイズ店舗の情報を持つ`franchises`と、トランザクションデータを保持する`transactions`です。

![Screenshot 2024-10-20 at 11.54.23.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/d6f431de-16bf-a5a9-82cb-0fac5b6b0099.png)
![Screenshot 2024-10-20 at 11.54.57.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/24935136-34cc-ee5b-ea62-570d6eff705a.png)

# ツールの実装

Databricksでは、[SQLあるいはPythonで関数を定義](https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-create-sql-function.html)することができますので、柔軟にロジックを組むことができます。

ここでは、SQLを用いた関数を定義します。上で準備したテーブルにアクセスするクエリーを記述します。

**都市名からフランチャイズ情報を取得する関数**

```sql:franchise_by_city
%sql
CREATE
OR REPLACE FUNCTION takaakiyayoi_catalog.cookie.franchise_by_city (
  city_name STRING COMMENT 'フランチャイズが存在する都市',
  country_name STRING COMMENT 'フランチャイズが存在する国'
) RETURNS TABLE READS SQL DATA SQL SECURITY DEFINER COMMENT '国、都市にあるフランチャイズを取得' RETURN
SELECT franchiseID, name, size from takaakiyayoi_catalog.cookie.franchises where city=city_name and country=CASE 
    WHEN country_name IN ('USA', 'America', 'United States') THEN 'US' 
    ELSE country_name 
    END order by size desc
```

**特定のフランチャイズの売上データを取得する関数**

```sql:franchise_sales
%sql
CREATE
OR REPLACE FUNCTION takaakiyayoi_catalog.cookie.franchise_sales (
  franchise_id INT COMMENT 'bakehouseフランチャイズのID',
  start_date STRING COMMENT '販売期間の開始日'
) RETURNS TABLE READS SQL DATA SQL SECURITY DEFINER COMMENT '指定されたフランチャイズで販売された総量と売上合計を取得' RETURN
SELECT SUM(totalPrice) AS total_sales, SUM(quantity) AS total_quantity, product 
FROM takaakiyayoi_catalog.cookie.transactions 
WHERE franchiseID = franchise_id AND dateTime > start_date GROUP BY product
```

これで、データとツールの準備ができました。

![Screenshot 2024-10-20 at 11.57.20.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/ff602152-1194-fc07-f1dd-be5138cab014.png)

:::note info
**注意**
AI Playgroundで稼働するツールはサーバレスコンピュートを使用します。本記事の執筆時点では、サーバレスコンピュートからインターネットへのアクセスはできないので、外部APIの呼び出しはサポートされていません。通常のクラスターにアタッチしたノートブックからツールを実行した際には外部APIの呼び出しは成功します。
:::


# AI Playgroundでのプロトタイピング

AI Playgroundでツールをサポートしているモデルを選択します。

![Screenshot 2024-10-20 at 11.58.25.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/98383543-139a-9c6b-2115-918206f07ddc.png)

上のステップで作成したツール(関数)を追加します。

![Screenshot 2024-10-20 at 11.58.51.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/35273ac5-94c4-6011-afb0-2b2b9c3e16e6.png)
![Screenshot 2024-10-20 at 11.59.24.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/b2ce4c5c-a4d8-c8c2-fab8-530aa67346e2.png)

これでエージェントの準備ができたので問い合わせをしてみます。

> Tokyoで人気のクッキー

すると、Tokyoのフランチャイズを特定するために`franchise_by_city`を呼び出します。

![Screenshot 2024-10-20 at 12.00.55.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/b815a694-7c8e-a121-de7f-70d3bf4b4afa.png)

続いて、Tokyoのフランチャイズの売上データを取得するために、上のステップで取得したフランチャイズIDを用いて`franchise_sales`を呼び出します。

![Screenshot 2024-10-20 at 12.02.12.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/c36a9db1-7d7f-dd52-97d3-92df7d8f25f5.png)

結果として以下を返してくれました。

> Based on the response, the top-selling cookie in Tokyo is "Outback Oatmeal" with a total sales of 81 and a total quantity of 27.

モデルを切り替えたり、システムプロンプトを追加することもできます。

![Screenshot 2024-10-20 at 12.05.54.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/a0e7c0cf-deda-5942-f6b0-66667a92a075.png)
![Screenshot 2024-10-20 at 12.05.59.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/edb1bfb4-f688-991c-fdba-8e80ca66ed72.png)

ユーザーとのインタラクション、エージェントに期待する挙動をイメージしつつ、ツールの実装、LLMとの連携を組み上げていく流れを容易に試行錯誤することができます。是非ご活用ください！

プロトタイピングの次はプロダクションですが、Databricksではそちらもカバーしています。プロダクションへの移行は次回説明します。

続きはこちらです。

https://qiita.com/taka_yayoi/items/df9cb37679b1b1a39911

### はじめてのDatabricks

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

### Databricks無料トライアル

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