Databricks Python SDKによるLLMチャットセッションの管理

こういったサイトがあること知りませんでした。

https://notebooks.databricks.com/devrel/index.html

パラパラと見ていたら気になるノートブックが。

https://notebooks.databricks.com/devrel/nbs/manage_chat_sessions.html

動かしてみます。

# Pythonでチャットセッションを管理する

チャットセッションの管理は、効果的な会話型AI体験を作成するための重要な部分です。チャットセッションは、LLMが過去のユーザー入力と応答にアクセスできるようにし、モデルが会話をスムーズかつコンテキスト豊かに保つのに役立ちます。このノートブックでは、Python SDKの`ChatSession`クラスがこれをどのように簡素化し、AIチャットがより自然なやり取りのように感じられるようにするかを示します。また、必要に応じてカスタムチャット管理システムを作成する方法についても説明します。最後に、チャットセッションを処理するためのOpenAI Pythonクライアントの使用方法を示します。

## `ChatSession`クラスの使用

Foundation Model API Python SDKの`ChatSession`クラスを使用すると、多くのカスタムコードを実装することなくチャット履歴を管理することが簡単になります。チャットセッションの管理には主に3つのステップがあります：

1. 例えば`chat = ChatSession(model=mixtral-8x7b-instruct, system_message="you are a helpful assistant")`のように`ChatSession`オブジェクトを初期化します。
2. `reply`メソッドを使用してユーザーの返信をモデルに送信します。
3. `last`メソッドを使用してモデルからの最後の応答を取得します。

```py
%pip install databricks-genai-inference
dbutils.library.restartPython()
```

```py
from databricks_genai_inference import ChatSession

chat = ChatSession(model="databricks-meta-llama-3-3-70b-instruct",
system_message="""あなたは20の質問ゲームで質問者です。回答者はあなたが推測するための名詞を選びます。推測するために「はい」か「いいえ」で答えられる質問をしてください。20回の推測ができます。推測回数を記録し、質問をしてください。何であるか分かったら、回答者に伝えてください。一度に一つの質問だけをしてください。""",
stop=['?'])
```

```py
chat.reply("単語を選びました。質問をどうぞ！")
chat.last
```
```
'質問回数：1\n単語が生き物であるかどうかを調べてみます。\n答えは「はい」または「いいえ」でお願いします。\n質問：選ばれた単語は生き物を指すものですか？'
```

```py
chat.reply("はい、そうです！")
chat.last
```
```
'質問回数：2\n生き物であることが分かりました。次に、動物かどうか調べてみます。\n答えは「はい」または「いいえ」でお願いします。\n質問：選ばれた単語は動物を指すものですか？'
```

```py
chat.reply("はい、毛皮か羽毛があります。どちらかは教えません。なぜなら、はいかいいえでしか質問できないからです。")
chat.last
```

```
'質問回数：3\n動物であることと、毛皮または羽毛があることが分かりました。次に、ペットとして飼われる動物かどうか調べてみます。\n答えは「はい」または「いいえ」でお願いします。\n質問：選ばれた動物は一般的にペットとして飼われるものですか？'
```

```py
chat.reply("はい、そうです！")
chat.last
```
```
'質問回数：4\nペットとして飼われる動物であることが分かりました。次に、四本足で歩く動物かどうか調べてみます。\n答えは「はい」または「いいえ」でお願いします。\n質問：選ばれた動物は四本足で歩くものですか？'
```

```py
chat.reply("はい。")
chat.last
```
```
'質問回数：5\n四本足で歩くペットであることが分かりました。次に、犬かどうか調べてみます。\n答えは「はい」または「いいえ」でお願いします。\n質問：選ばれた動物は犬ですか？'
```

```py
chat.reply("はい。")
chat.last
```
```
'質問回数：6\n犬であることが分かりました！犬はペットとして非常に人気があります。\n正解です！答えは「犬」でしたね。游戏は終了です。ありがとうございました！次のゲームを始めますか？'
```

```py
chat.reply("いいえ！")
chat.last
```
```
'ゲームを終了します。\nゲーム中に協力してくださり、ありがとうございました！もしまた遊びたい時はいつでも開始してくださいね！'
```

こんなに簡単にセッション履歴を管理できるとは。

## チャット履歴の表示

`history` メソッドを使用して、チャット履歴全体を表示できます:

```py
chat.history
```
```py
[{'role': 'system',
  'content': 'あなたは20の質問ゲームで質問者です。回答者はあなたが推測するための名詞を選びます。推測するために「はい」か「いいえ」で答えられる質問をしてください。20回の推測ができます。推測回数を記録し、質問をしてください。何であるか分かったら、回答者に伝えてください。一度に一つの質問だけをしてください。'},
 {'role': 'user', 'content': '単語を選びました。質問をどうぞ！'},
 {'role': 'assistant',
  'content': '質問回数：1\n単語が生き物であるかどうかを調べてみます。\n答えは「はい」または「いいえ」でお願いします。\n質問：選ばれた単語は生き物を指すものですか？'},
 {'role': 'user', 'content': 'はい、そうです！'},
 {'role': 'assistant',
  'content': '質問回数：2\n生き物であることが分かりました。次に、動物かどうか調べてみます。\n答えは「はい」または「いいえ」でお願いします。\n質問：選ばれた単語は動物を指すものですか？'},
 {'role': 'user',
  'content': 'はい、毛皮か羽毛があります。どちらかは教えません。なぜなら、はいかいいえでしか質問できないからです。'},
 {'role': 'assistant',
  'content': '質問回数：3\n動物であることと、毛皮または羽毛があることが分かりました。次に、ペットとして飼われる動物かどうか調べてみます。\n答えは「はい」または「いいえ」でお願いします。\n質問：選ばれた動物は一般的にペットとして飼われるものですか？'},
 {'role': 'user', 'content': 'はい、そうです！'},
 {'role': 'assistant',
  'content': '質問回数：4\nペットとして飼われる動物であることが分かりました。次に、四本足で歩く動物かどうか調べてみます。\n答えは「はい」または「いいえ」でお願いします。\n質問：選ばれた動物は四本足で歩くものですか？'},
 {'role': 'user', 'content': 'はい。'},
 {'role': 'assistant',
  'content': '質問回数：5\n四本足で歩くペットであることが分かりました。次に、犬かどうか調べてみます。\n答えは「はい」または「いいえ」でお願いします。\n質問：選ばれた動物は犬ですか？'},
 {'role': 'user', 'content': 'はい。'},
 {'role': 'assistant',
  'content': '質問回数：6\n犬であることが分かりました！犬はペットとして非常に人気があります。\n正解です！答えは「犬」でしたね。游戏は終了です。ありがとうございました！次のゲームを始めますか？'},
 {'role': 'user', 'content': 'いいえ！'},
 {'role': 'assistant',
  'content': 'ゲームを終了します。\nゲーム中に協力してくださり、ありがとうございました！もしまた遊びたい時はいつでも開始してくださいね！'}]
```

`count` メソッドを使用して、ラウンド数（ユーザーのプロンプトとアシスタントの応答のペア）を数えることができます:

```py
chat.count
```
```
7
```

## 手動でチャットコンプリーションを使用

場合によっては、チャット履歴の管理方法をより細かく制御したいことがあります。ここでは、チャット履歴を自分で管理するための一つのアプローチの最小限の例を示します。これは、カスタムニーズに合わせたチャットシステムを構築するための大まかな概要に過ぎません。

```py
from databricks_genai_inference import ChatCompletion

class CustomChatSession:
    """
    モデルとのチャットセッションを管理するクラス。
    """
    def __init__(self, model, system_message=None, max_tokens=512):
        """
        モデル、オプションのシステムメッセージ、および最大トークン数でChatModelを初期化します。
        
        引数:
            model: 応答を生成するために使用するモデル。
            system_message (str, オプション): 初期システムメッセージ。デフォルトはNone。
            max_tokens (int, オプション): モデルが生成する最大トークン数。デフォルトは512。
        """
        self.model = model
        self.max_tokens = max_tokens
        self.system_message = system_message
        self.messages = (
            [{"role": "system", "content": system_message}] if system_message else []
        )

    def reply(self, user_message):
        """
        ユーザーメッセージをチャット履歴に追加し、応答を生成します。
        
        引数:
            user_message (str): ユーザーのメッセージ。
        
        戻り値:
            str: モデルの応答。
        """
        self.messages.append({"role": "user", "content": user_message})
        return self.execute()

    def __call__(self, user_message):
        """
        ユーザーメッセージをチャット履歴に追加し、応答を生成します。
        
        引数:
            user_message (str): ユーザーのメッセージ。
        
        戻り値:
            str: モデルの応答。
        """
        self.messages.append({"role": "user", "content": user_message})
        return self.execute()

    def execute(self):
        """
        チャット履歴に基づいてモデルから応答を生成します。
        
        戻り値:
            str: モデルの応答。
        """
        response = ChatCompletion.create(
            model=self.model,
            max_tokens=self.max_tokens,
            stream=False,
            messages=self.messages,
        )

        self.messages.append({"role": "assistant", "content": response.message})
        return response.message
```
```py
chat = CustomChatSession(model="databricks-meta-llama-3-3-70b-instruct", system_message="あなたは役に立つアシスタントです。")
```

```py
display(chat("私の名前はダンで、コーヒーを飲むのが好きです。コーヒーの作り方について何か一般的なおすすめはありますか？"))
```
```
'ダンさん、コーヒーを飲むのが好きな方なんですね！コーヒーの作り方にはいくつかの一般的な方法がありますが、人気のあるものをいくつか紹介しますね。\n\n1. **ドリップコーヒー**：コーヒーフィルターを使ってコーヒー豆を絞る方法で、手軽にコーヒーを楽しむことができます。コーヒー豆を挽いてフィルターに置き、熱湯を注ぐだけです。\n2. **フレンチプレス**：コーヒー豆を挽いて熱湯を注ぎ、しばらく置いてからプレスでコーヒーの粉を押し下げる方法です。コーヒー豆の風味がしっかり出るのが特徴です。\n3. **エスプレッソ**：高圧でコーヒー豆を絞る方法で、濃厚で強いコーヒーが得られます。エスプレッソマシーンが必要になりますが、カフェラテやカプチーノを作るのに最適です。\n4. **ポーチドコーヒー**：コーヒー豆を挽いて熱湯を注ぎ、しばらく置いてからコーヒーを注ぐ方法です。コーヒーの風味が口の中に広がります。\n\nこれらの方法はすべておすすめですが、ダンさんの好みのコーヒーを探すには、コーヒー豆の種類やローストレベルも重要です。lightローストからdarkローストまで様々なローストレベルがあるので、試してみてみつけることがおすすめです。\n\nダンさん、これらの方法の中でどれか興味がありますか？もしくは他の特定のコーヒー豆やローストレベルに興味がある場合は教えてくださいね！'
```

結果がマークダウンで出力されているので、[こちら](https://stackoverflow.com/questions/73551155/how-to-display-markdown-output-in-databricks-notebook-from-a-python-cell/73568603#73568603)を参考にノートブック上にレンダリングされるようにします。

```py
from IPython.display import DisplayObject, TextDisplayObject

class Markdown(TextDisplayObject):

    def __init__(self,TextDisplayObject):
        import markdown as md
        
        #converting markdown to html
        self.html = md.markdown(TextDisplayObject)
        
    
    def _repr_html_(self):
        return self.html
```

```py
Markdown(chat("私の名前はダンで、コーヒーを飲むのが好きです。コーヒーの作り方について何か一般的なおすすめはありますか？"))
```
![Screenshot 2025-01-12 at 11.04.29.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/cc1e37d1-0353-d464-65a9-da4202435ff4.png)

```py
Markdown(chat("私の名前は何ですか？今まで私について何を知っていますか？"))
```

> あなたの名前はダンです。これまでの会話から、ダンさんはコーヒーを飲むのが好きだと知っています。コーヒーについての話題を始めたときに、そのことを教えてくださったので、覚えています。

```py
chat.messages
```
```py
[{'role': 'system', 'content': 'あなたは役に立つアシスタントです。'},
 {'role': 'user',
  'content': '私の名前はダンで、コーヒーを飲むのが好きです。コーヒーの作り方について何か一般的なおすすめはありますか？'},
 {'role': 'assistant',
  'content': 'ダンさん、こんにちは！コーヒーを楽しむ方法はたくさんありますが、ここではいくつか一般的なおすすめを紹介します。\n\n1. **ドリップコーヒー**：手軽にコーヒーを楽しめる方法です。コーヒーフィルターを使用して、コーヒー豆の粉を淹れます。コーヒーの強さや味わいを自分で調整できます。\n\n2. **フレンチプレス**：コーヒー豆の粉を沸騰した水で浸け、しばらくすると粉を押してコーヒーを抽出します。コーヒー豆の風味がしっかりと感じられます。\n\n3. **エスプレッソ**：特に強いコーヒーです。エスプレッソマシンを使って、コーヒー豆を高圧で抽出します。カフェラテなどのベースとして人気です。\n\n4. **コールドブリュー**：低温で長時間浸け込みます。やさしい味わいと低カフェインが特徴です。\n\n5. **アイスドコーヒー**：ホットコーヒーを氷の上に注いで楽しむ、またはコールドブリューコーヒーを使用する方法です。夏にぴったりです。\n\nコーヒー豆の選び方も重要です。アラビカ豆とロブスタ豆の2種類があり、それぞれ特徴があります。アラビカ豆は酸味が強く、フルーティーな風味があります。一方、ロブスタ豆は苦味が強く、ボディが重いです。\n\nまた、コーヒーの強さや味わいを変えるために、牛乳、砂糖、シロップなどを加えることも楽しみです。コーヒー豆の焙煎度も、ライトロースト、ミディアムロースト、ダークローストなどがあり、それぞれ違った風味を楽しめます。\n\nコーヒーは人によって好みが違うので、違うタイプや方法を試してみて、ダンさんの好みにぴったりのコーヒーを見つけてくださいね。'},
 {'role': 'user', 'content': '私の名前は何ですか？今まで私について何を知っていますか？'},
 {'role': 'assistant',
  'content': 'あなたの名前はダンです。これまでの会話から、ダンさんはコーヒーを飲むのが好きだと知っています。コーヒーについての話題を始めたときに、そのことを教えてくださったので、覚えています。'}]
```

## OpenAI Pythonクライアントの使用
Databricks推論SDKの代わりにOpenAIクライアントを使用してチャットセッションを管理することができます。[OpenAI Pythonクライアントを使用したファウンデーションモデルAPIの使用](https://notebooks.databricks.com/devrel/nbs/fm_api_openai_sdk.html#nbs-fm-api-openai-sdk)に関するノートブックで詳述されているように、ファウンデーションモデルAPIはOpenAIクライアントと互換性があります。これにより、OpenAIクライアントに基づいてコードを書き直す必要なく、ファウンデーションモデルAPIモデルを試すことが非常に簡単になります。

以下は、OpenAI Pythonクライアントを使用して上記のコードを書き直す方法です：

### クライアントの設定
まず、`OPENAI_API_KEY`と`OPENAI_BASE_URL`環境変数を設定する必要があります。詳細は[OpenAIノートブック](https://notebooks.databricks.com/devrel/nbs/fm_api_openai_sdk.html#nbs-fm-api-openai-sdk)を参照してください。

```py
%pip install --upgrade openai
%restart_python
```
```py
import os
from openai import OpenAI

os.environ["OPENAI_BASE_URL"] = (
    "https://" + spark.conf.get("spark.databricks.workspaceUrl") + "/serving-endpoints/"
)

os.environ["OPENAI_API_KEY"] = (
    dbutils.notebook.entry_point.getDbutils()
    .notebook()
    .getContext()
    .apiToken()
    .getOrElse(None)
)

client = OpenAI()
```

次に、`execute`メソッドをOpenAIクライアントを使用するように書き直します。

```py
class CustomChatSession:
    """
    モデルとのチャットセッションを管理するクラス。
    """
    def __init__(self, model, system_message=None, max_tokens=512):
        """
        モデル、オプションのシステムメッセージ、および最大トークン数でChatModelを初期化します。
        
        引数:
            model: 応答を生成するために使用するモデル。
            system_message (str, optional): 初期のシステムメッセージ。デフォルトはNone。
            max_tokens (int, optional): モデルが生成する最大トークン数。デフォルトは512。
        """
        self.model = model
        self.max_tokens = max_tokens
        self.system_message = system_message
        self.messages = (
            [{"role": "system", "content": system_message}] if system_message else []
        )

    def reply(self, user_message):
        """
        ユーザーメッセージをチャット履歴に追加し、応答を生成します。
        
        引数:
            user_message (str): ユーザーのメッセージ。
        
        戻り値:
            str: モデルの応答。
        """
        self.messages.append({"role": "user", "content": user_message})
        return self.execute()

    def __call__(self, user_message):
        """
        ユーザーメッセージをチャット履歴に追加し、応答を生成します。
        
        引数:
            user_message (str): ユーザーのメッセージ。
        
        戻り値:
            str: モデルの応答。
        """
        self.messages.append({"role": "user", "content": user_message})
        return self.execute()

    def execute(self):
        """
        チャット履歴に基づいてモデルから応答を生成します。
        
        戻り値:
            str: モデルの応答。
        """
        response = client.chat.completions.create(
            model=self.model,
            max_tokens=self.max_tokens,
            stream=False,
            messages=self.messages,
        )

        self.messages.append({"role": "assistant", "content": response.choices[0].message.content})
        return response.choices[0].message.content
```

```py
chat = CustomChatSession(model="databricks-meta-llama-3-3-70b-instruct", system_message="あなたは役に立つアシスタントです。")
Markdown(chat("私の名前はダンで、コーヒーを飲むのが好きです。コーヒーの作り方について何か一般的なおすすめはありますか？"))
```

> こんにちはダンさん！コーヒーが好きな方は珍しくありませんね。コーヒーの作り方にはいくつか一般的なおすすめがあります。まず、豆の選び方です。アラビカ豆は味が繊細で高品質なコーヒーを作ることができます。次に、豆の挽き方です。豆を挽くタイミングは大切で、できれば コーヒーを淹れる直前に挽くことでフレッシュな風味を楽しむことができます。さらに、コーヒーの抽出方法も重要です。ドリップ、フレンチプレス、エスプレッソなど、好みの方法でコーヒーを楽しむことができます。最後に、水温やコーヒーと水の比率もコーヒーの味に大きく影響します。水温は適温（約90～95度）で、コーヒーと水の比率はお好みで調整できます。以上の点に気を配ってコーヒーを作れば、きっと美味しいコーヒーを楽しむことができるでしょう。

```py
Markdown(chat("私の名前は何ですか？これまでの私について何を知っていますか？"))
```

> あなたの名前はダンです。また、コーヒーが好きなこと、あなたがコーヒーについて相談したこと、またコーヒーの作り方のアドバイスをもらったことが分かっています。もし他に教えて欲しいことや相談したいことがあれば、いつでもお相手します。

これで、同じFoundation Model APIエンドポイントにクエリを実行し、OpenAI Pythonクライアントを使用して同じ結果を達成しました。

## 結論

このノートブックでは、チャットセッションを管理するための`ChatSession`クラスの使用方法を学びました。また、よりカスタマイズされたチャット管理システムにアプローチする1つの方法も学びました。


### はじめてのDatabricks

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

### Databricks無料トライアル

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