Mosaic AI Agent FrameworkによるRAGチェーンの構築、デプロイと評価

こちらの一番シンプルなチェーンのサンプルを構築し、Agent Frameworkを用いてそのチェーンを記録・デプロイします。

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

# "Hello World" アプリケーションノートブック

**チェーンのノートブック**

https://docs.databricks.com/ja/generative-ai/create-log-agent.html#simple-chain-notebook

このノートブックは、Mosaic AI Agent Frameworkを使用してシンプルな"Hello World"チェーンを作成します。

## パッケージのインストール

```py
%pip install -U -qqqq databricks-agents mlflow langchain==0.2.1 langchain_core==0.2.5 langchain_community==0.2.4
dbutils.library.restartPython()
```

## インポート

```py
import mlflow
from operator import itemgetter
from langchain_core.output_parsers import StrOutputParser
from langchain.schema.runnable import RunnableLambda
from langchain_community.chat_models import ChatDatabricks
from langchain_core.prompts import PromptTemplate
```

## MLflowトレースの有効化

MLflowトレースを有効にすると、次のことができます：
- このノートブックでトレースの可視化を表示する
- インファレンステーブルを介して本番環境でトレースをキャプチャする
- アプリケーションを評価する

```py
mlflow.langchain.autolog()
```

## ヘルパー関数

お使いのチェーンでは、`messages`パラメーターとしてOpenAIのフォーマットに従ったメッセージの配列を受け付ける必要があります。

スキーマ: [AWS](https://docs.databricks.com/ja/machine-learning/foundation-models/api-reference.html#chatmessage) | [Azure](https://learn.microsoft.com/ja-jp/azure/databricks/machine-learning/foundation-models/api-reference#chatmessage)

```py
# これらのヘルパー関数は `messages` の配列をパースします

# ユーザーからの最も直近のメッセージの文字列のコンテンツを返却
def extract_user_query_string(chat_messages_array):
    return chat_messages_array[-1]["content"]

# 最後の質問の前の全てとなるチャット履歴を返却
def extract_chat_history(chat_messages_array):
    return chat_messages_array[:-1]
```

## チェーンの作成

```py
############
# 生成のためのプロンプトテンプレート
############
prompt = PromptTemplate(
    template="あなたはhello world botです。ユーザーの質問に対してユーザーにとって面白い回答を日本語でレスポンスします。ユーザーの質問: {question}",
    input_variables=["question"],
)

############
# 生成のための基盤モデル
# ChatDatabricks は、いかなる /llm/v1/chat モデルサービングエンドポイントを受け付ます
############
model = ChatDatabricks(
    endpoint="databricks-dbrx-instruct",
    extra_params={"temperature": 0.01, "max_tokens": 500},
)


############
# シンプルなチェーン
############
# このフレームワークは、チェーンに文字列を返却することを要求します
chain = (
    {
        "question": itemgetter("messages")
        | RunnableLambda(extract_user_query_string),
        "chat_history": itemgetter("messages") | RunnableLambda(extract_chat_history),
    }
    | prompt
    | model
    | StrOutputParser()
)
```

## ローカルでチェーンをテスト

```py
# これは、チェーンのREST APIが受け入れるのと同じ入力です。
question = {
    "messages": [
               {
            "role": "user",
            "content": "ragって何？",
        },
    ]
}

chain.invoke(question)
```
```
'「rag」は、英語で「ラグ」や「はみ出し」を意味する言葉です。しかし、日本語では「rag」という言葉はあまり使われません。もし、あなたが「rag」と聞いたのであれば、その場面によって意味が異なるでしょう。もし、具体的な文脈を教えていただければ、もっと正確な回答を差し上げます。'
```

トレースが表示されます。

![Screenshot 2024-09-30 at 21.56.50.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/52f241f8-af61-6321-8287-5394d93df22f.png)

## MLflowログにチェーンの場所を指定する

`mlflow.models.set_model(model=...)` 関数は、評価とデプロイメントに使用するLangChainチェーンを指定します。これは、`mlflow.langchain.log_model(...)`でこのチェーンをMLflowにログするために必要です。

```py
mlflow.models.set_model(model=chain)
```

本当にシンプルなチェーンを動かしましたというサンプルですね。次のドライバーノートブックからアクセスすることになります。

# "Hello World" ドライバーノートブック

**ドライバーノートブック**

https://docs.databricks.com/ja/generative-ai/create-log-agent.html#simple-chain-driver-notebook

このノートブックは、Mosaic AI Agent Frameworkを使用してシンプルな"Hello World"アプリケーションをログ、テスト、デプロイする方法を示しています。以下のステップをカバーしています：

1. Mosaic AI Agent Frameworkのインストール
2. 必要なモジュールのインポート
3. ノートブックのパスを定義
4. MLflowにチェーンをログし、ローカルでテストする
5. Unity Catalogにモデルとしてチェーンを登録する
6. チェーンをデプロイする
7. Agent Framework UIでデプロイされたチェーンを表示する

## パッケージのインストール

```py
%pip install -U -qqqq databricks-agents mlflow langchain==0.2.1 langchain_core==0.2.5 langchain_community==0.2.4
dbutils.library.restartPython()
```

## インポート

```py
import os
import mlflow
from databricks import agents

# Unity Catalog モデルレジストリを使用する
mlflow.set_registry_uri('databricks-uc')
```

## Unity Catalogの場所を選択して、チェーンをログに記録します

```py
# ウィジェットの作成
dbutils.widgets.text("uc_catalog", "", "catalog")
dbutils.widgets.text("uc_schema", "", "schema")
dbutils.widgets.text("model_name", "hello_world", "Model name")

# ウィジェットから値を取得
uc_catalog = dbutils.widgets.get("uc_catalog")
uc_schema = dbutils.widgets.get("uc_schema")
model_name = dbutils.widgets.get("model_name")
```

## MLflowへのチェーンのログ記録

チェーンとその結果のモデルをノートブックのMLflowエクスペリメントに記録します。

### MLflowログの入力パラメーター

```py
# MLflowモデルのシグネチャに設定される入力スキーマの例を提供します
input_example = {
   "messages": [
       {
           "role": "user",
           "content": "Retrieval-augmented Generationとは何ですか？",
       }
   ]
}

# チェーンノートブックへの完全なパスを指定します
chain_notebook_file = "hello-world-chain"
chain_notebook_path = os.path.join(os.getcwd(), chain_notebook_file)

print(f"チェーンノートブックのパス: {chain_notebook_path}")
```
```
チェーンノートブックのパス: /Workspace/Users/takaaki.yayoi@databricks.com/20240930_ai_agents/hello-world-chain
```

### チェーンを記録する

```py
with mlflow.start_run():
    logged_chain_info = mlflow.langchain.log_model(
        lc_model=chain_notebook_path,
        artifact_path="chain",
        input_example=input_example,
    )
```
チェーンが記録されます。

![Screenshot 2024-09-30 at 22.13.06.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/3b7418f5-6f26-67a2-ac23-089b999e20c9.png)

## チェーンの評価
通常、この時点でチェーンを評価します。ただし、このノートブックにはそのステップは含まれていません。

## チェーンのデプロイ

モデルをデプロイするには、まずMLflowランからチェーンをUnity Catalogモデルとして登録します。

```py
# Unity Catalogのロケーション
uc_model_name = f"{uc_catalog}.{uc_schema}.{model_name}"

# Unity Catalogにモデルを登録
uc_registered_model_info = mlflow.register_model(model_uri=logged_chain_info.model_uri, name=uc_model_name )
```

Unity Catalog配下にモデルが登録されます。

![Screenshot 2024-09-30 at 22.14.02.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/9f470f9d-ee0d-a36c-2d8e-6b2a1706e987.png)

`agents.deploy`を使用してチェーンをデプロイすると、以下の場所にデプロイされます：
- クエリのためのREST APIエンドポイント
- レビューアプリには、ウェブUIを使用してチェーンとチャットし、フィードバックを提供することができます。

```py
deployment_info = agents.deploy(uc_model_name, uc_registered_model_info.version)
```

上を実行すると、クエリのためのエンドポイントとレビューのためのエンドポイントが作成されます。

![Screenshot 2024-09-30 at 22.19.46.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/cee3519d-2e60-f9ba-964b-219882afa7fd.png)

```py
# アプリケーションをクエリするためにこのURLを使用してください。
deployment_info.query_endpoint
```
```
'https://xxxxxxx.cloud.databricks.com/serving-endpoints/agents_users-takaaki_yayoi-hello_world/served-models/users-takaaki_yayoi-hello_world_1/invocations'
```

上のURLでREST APIを呼び出すことでエンドポイントにクエリーすることができます。REST APIの呼び出しにはパーソナルアクセストークンが必要なので、ここでは直接エンドポイントの画面からクエリーします。

![Screenshot 2024-09-30 at 22.23.15.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/992f9f00-04d4-9d7a-940f-c63cd8d9e163.png)

```py
# アプリケーションと対話するために、このURLをブラウザにコピーしてください。
deployment_info.review_app_url
```
```
'https://xxxxxxx.cloud.databricks.com/ml/review/users.takaaki_yayoi.hello_world/1?o=2556758628403379'
```

こちらにアクセスすると、チェーンの精度を評価するためのUIが表示されます。

![Screenshot 2024-09-30 at 22.23.57.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/2b44fc01-12fd-3c7e-90e7-93828f965724.png)
![Screenshot 2024-09-30 at 22.24.16.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/acd4c096-cdcb-b566-ac8e-0d0a94fcc791.png)
![Screenshot 2024-09-30 at 22.24.38.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/23673f7d-7df4-f0c6-e8ce-8b51de0b4b72.png)


## デプロイメントの表示

もし上記で取得したデプロイメント情報を失ってしまった場合は、`list_deployments()` を使用してそれを見つけることができます。

```py
deployments = agents.list_deployments()
deployments
```



### はじめてのDatabricks

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

### Databricks無料トライアル

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