Databricksからインターネットに接続するエージェントツールの構築

存在は知っていたのですが、動かせていませんでした。

https://docs.databricks.com/aws/ja/generative-ai/agent-framework/external-connection-tools

> AIエージェントツールをSlack、Googleカレンダー、またはHTTPリクエストを使用してAPIを使用する任意のサービスなどの外部アプリケーションに接続する方法を学びます。 エージェントは、外部に接続されたツールを使用して、タスクの自動化、メッセージの送信、およびサードパーティプラットフォームからのデータの取得を行うことができます。

この手順では、Unity Catalogが管理する[接続](https://docs.databricks.com/aws/ja/query-federation/connections)として**HTTP接続**を作成し、認証情報と共にアクセス先をポイントするようにします。

ただ、こちらを読んだ時に「[エグレスコントロール](https://docs.databricks.com/aws/ja/security/network/serverless-network-security/network-policies)などでネットワークが疎通していればここで説明されている手順でなくても、インターネットに接続するツールは作れるのは？」と思いました。

例えば、以下のようなツールを作成してUntiy Catalogに登録します。ちなみに、https://icanhazdadjoke.com/ はお父さんのジョークをランダムに返してくれます。

```py
CREATE OR REPLACE FUNCTION takaakiyayoi_catalog.tools.tell_dad_joke() 
RETURNS STRING 
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}"
$$
```

こんな感じでPlaygroundで冗談を教えてもらうことができます。

![Screenshot 2025-03-10 at 14.56.18.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/2ca9f1d6-7b4e-442d-8249-d7c144ddf26b.png)

では、なぜ接続を用いてエージェントツールを構築するのでしょうか。[元のマニュアルに説明](https://docs.databricks.com/aws/ja/generative-ai/agent-framework/external-connection-tools#%E5%A4%96%E9%83%A8%E3%82%B5%E3%83%BC%E3%83%93%E3%82%B9%E3%81%B8%E3%81%AE%E6%8E%A5%E7%B6%9A%E3%82%92%E4%BD%9C%E6%88%90%E3%81%99%E3%82%8B)があります:

> Unity Catalog Connection を使用する利点は次のとおりです。
> - **安全な認証情報管理:** シークレットとトークンは Unity Catalog に安全に保管され、管理されるため、ユーザーに公開されることはありません。
> - **きめ細かなアクセス制御:** Unity Catalog では、 `USE_CONNECTION` と `MANAGE_CONNECTION` の特権を使用して、接続を使用または管理できるユーザーをきめ細かく制御できます。
> - **ホスト固有のトークンの強制:** トークンは、接続の作成時に指定された `host_name` に制限されるため、承認されていないホストでは使用できません。

つまり、認証情報やアクセス先、そして接続へのアクセス権を集中管理できるというところにメリットがあるということです。また、接続として登録されたアクセス先には[`http_request`](https://docs.databricks.com/aws/ja/sql/language-manual/functions/http_request)関数でSQLからもアクセスできるようになります。

[こちらのSlackのサンプルノートブック](https://docs.databricks.com/aws/ja/generative-ai/agent-framework/external-connection-tools#%E3%83%8E%E3%83%BC%E3%83%88%E3%83%96%E3%83%83%E3%82%AF%E3%81%AE%E4%BE%8B)を動かします。

# AIエージェントのSlackメッセージングツール

このノートブックでは、事前に定義されたSlackチャンネルにメッセージを送信するAIエージェントツールを作成する2つの方法を示します。

**方法1: Unity Catalog関数を使用してツールを作成**

1. **Unity Catalog接続オブジェクトを作成**: 接続オブジェクトは、Slackに接続するために必要な認証情報を保存します。
2. **前処理および後処理のUnity Catalog関数を定義**: リクエストを前処理してSlackメッセージをフォーマットし、Slack API呼び出し後のステータスを確認するためにレスポンスを処理します。
3. **Unity Catalog関数ツールを作成**: Slack APIにクエリを送信してメッセージを投稿するUnity Catalog関数を作成します。

**方法2: エージェントコードでツールを定義**

1. **Unity Catalog接続オブジェクトを作成**: 接続オブジェクトは、Slackに接続するために必要な認証情報を保存します。
2. **Pythonでツールを定義**: 入力テキストを処理し、Slack APIにクエリを送信し、返り値を処理するPython関数を定義します。

Unity Catalog関数ツールとエージェントコードで定義されたツールの違いについては、Databricksのドキュメントを参照してください ([AWS](https://docs.databricks.com/ja/generative-ai/agent-framework/agent-tool.html#unity-catalog-function-tools-vs-agent-code-tools) | [Azure](https://learn.microsoft.com/ja-jp/azure/databricks/generative-ai/agent-framework/agent-tool#unity-catalog-function-tools-vs-agent-code-tools))。

```py
%pip install -U databricks-sdk
%restart_python
```

## UC接続の作成

このコードはUC接続オブジェクトを作成します。接続オブジェクトは、外部関数に認証するためのbearer_tokenを保存します。Slackのトークンの取得方法はこちらをご覧ください。

https://zenn.dev/kou_pg_0131/articles/slack-api-post-message

```sql
%sql
DROP CONNECTION IF EXISTS test_sql_slack_connection;
CREATE CONNECTION test_sql_slack_connection
TYPE HTTP
OPTIONS (
  host 'https://slack.com',
  port '443',
  base_path '/',
  bearer_token 'xoxb-...'
);
```

これでSlackへの接続が作成されました。

![Screenshot 2025-03-10 at 15.10.26.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/b54c9a3a-17a2-48bf-b060-355f7066265f.png)

## 方法1: Unity Catalog関数を使用してツールを作成

### 前処理関数

外部API呼び出しの前に呼び出される前処理関数を定義します。前処理関数は、Slackチャンネルとテキストを含む構造体を受け取り、テキストに免責事項メッセージを追加して修正された構造体を返します。

```sql
%sql
CREATE OR REPLACE FUNCTION takaakiyayoi_catalog.agents.pre_process_slack_message(data Struct < channel: STRING, text: STRING >) 
RETURNS Struct < channel: STRING, text: STRING > 
LANGUAGE PYTHON 
AS $$

# このメッセージはAIエージェントからのものですという免責事項を追加
slack_text = "このメッセージはAIエージェントからです:\n" + data.text

# チャンネル名とクリーンなメッセージを希望の形式で結合
return {"channel": data.channel, "text": data.text}
$$
```

### 後処理関数

外部API呼び出し後に呼び出される後処理関数を定義します。後処理関数は、ステータスコードとテキストとしてレスポンスデータを受け取り、文字列を返すことを期待します。この場合、関数はステータスコードを確認し、成功または失敗のメッセージを返します。

```sql
%sql
CREATE OR REPLACE FUNCTION takaakiyayoi_catalog.agents.post_process_slack_message(
  response_data Struct < statusCode: INT, text: STRING >
) 
RETURNS STRING 
LANGUAGE PYTHON AS $$

# 処理結果
if response_data.statusCode == 200:
    return "Slackへのメッセージ送信に成功しました！"
else:
    return "リクエスト実行中にエラーに遭遇しました。呼び出しからのメッセージ: " + text
$$
```

### エージェントツールの定義

このコードは、テキストパラメータを受け取り、Slackチャンネルにメッセージを送信するUnity Catalog関数を定義します。関数は、接続オブジェクト、HTTPメソッド、リクエストボディ、および外部呼び出しのJSONペイロードを受け取る`http_request`を使用します。

この例では、SlackチャンネルIDをハードコーディングしています。

```sql
%sql
CREATE OR REPLACE FUNCTION takaakiyayoi_catalog.agents.slack_send_message(text STRING COMMENT 'メッセージ内容')
RETURNS STRING
COMMENT "事前定義されたSlackチャンネル（ID: 'C085QCA8UKF'、名前: 'slack-integration-test'）にメッセージを送信します。自動的に現在のユーザー名を先頭に追加します。入力: 送信するテキストメッセージ。戻り値: Slack APIの応答。使用例: SELECT main.default.slack_send_message('ここにメッセージを入力')."
RETURN (
  takaakiyayoi_catalog.agents.post_process_slack_message(
    http_request(
      conn => 'test_sql_slack_connection',
      method => 'POST',
      path => '/api/chat.postMessage',
      json =>
        to_json(
          takaakiyayoi_catalog.agents.pre_process_slack_message(
            named_struct('channel', "C085QCA8UKF", 'text', concat("On behalf of ", current_user(), ": ", text))
          )
        )
    )
  )
)
```

これでPlaygroundから動作確認することができます。

![Screenshot 2025-03-10 at 15.12.57.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/a01c365d-4674-4c96-ad41-cce3438ccbaa.png)

![Screenshot 2025-03-10 at 15.14.07.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/3872a936-8415-4b6a-849a-1dc711e59a02.png)

## 方法2: エージェントコードでツールを定義する

[Python SDK](https://databricks-sdk-py.readthedocs.io/en/latest/index.html)を用いることで、同じ挙動をするエージェントを作成することができます。

```py
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod

def send_slack_message(text: str):
  """
    Slackチャンネル 'C085QCA8UKF' にメッセージを送信します。メッセージには自動的に現在のユーザー名が追加されます。戻り値: Slack APIのレスポンス。

    引数:
        text: Slackチャンネル 'channel-name' に投稿するテキストメッセージ。
  """

  slack_text = "This message is from an AI Agent:\n" + text

  response = WorkspaceClient().serving_endpoints.http_request(
    conn="test_sql_slack_connection",
    method=ExternalFunctionRequestHttpMethod.POST,
    path="/api/chat.postMessage",
    json={"channel": "C085QCA8UKF", "text": slack_text},
  )

  print(response.text)
  print(response.status_code)

  if response.status_code == 200:
    return "Successfully sent message to Slack!"
  else:
    return "Encountered failure when executing request. Message from Call: " + text
```
```py
send_slack_message("Hello, world!")
```
![Screenshot 2025-03-10 at 15.15.26.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/4d30fb69-ef7d-4f48-99ee-1e64b4703c64.png)



### はじめてのDatabricks

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

### Databricks無料トライアル

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