CCCMKホールディングス TECH LABの Tech Blog

TECH LABのエンジニアが技術情報を発信しています

ブログタイトル

embedding database "Chroma"とLLMを使って記憶に基づいた回答をAgentにさせてみました。

こんにちは、CCCMKホールディングス TECH LABの三浦です。

私はとても忘れっぽいのですが、何故か何十年経っても覚えていること、思い出せることがあります。すぐ忘れてしまうことと、ずっと覚えていることの違いってなんだろう、と考えてみるのですがよくわかりません。でも今でも思い出せることは、たぶんこれからもずっと忘れないような気がしています。

さて、今回はLangChainやLlamaIndexを使っていると、自ずと使うことが多いChromaというオープンソースのembedding databaseについて触れたいと思います。Chromaはテキストデータを埋め込み表現と一緒に格納する構造を持ち、検索機能や削除・更新といったデータ操作機能を備えています。

Chroma

ChromaはLangChainやLlamaIndexといったLLMを使ったアプリケーション開発フレームワークにおいて活用されています。これらのフレームワークを使う過程でChromaを利用してきましたが、データの追加や更新削除などをもう少し細かく制御して行いたい、と思うことがありました。そう思ったきっかけがこちらの論文です。

Title: Generative Agents: Interactive Simulacra of Human Behavior
Authors: Joon Sung Park, Joseph C. O'Brien, Carrie J. Cai, Meredith Ringel Morris, Percy Liang, Michael S. Bernstein
Submit: Submitted on 7 Apr 2023
URL: https://arxiv.org/abs/2304.03442

この論文ではChatGPT(gpt-3.5-turbo)をエンジンにした人格を持つ20以上のAgent達がお互い干渉しあいながら一つの村で生活をするというシミュレーションが行われています。Agent達はこの環境の中で様々な行動をとり、その結果を観測し、それを記憶として保存していきます。そしてその記憶を元に熟考し、次の行動を決定します。

Generative Agents: Interactive Simulacra of Human Behavior, Joon Sung Park, Joseph C. O'Brien, Carrie J. Cai, Meredith Ringel Morris, Percy Liang, Michael S. Bernstein, Figure 1

このシミュレーションがとても面白く、自分でも試してみたい、と考えるようになりました。そのためにはこの論文の提案手法の中心にある"記憶"を表現しないといけないのですが、そのためにChromaを使ってみようと考えました。

Chroma

Chromaはembedding databaseで、テキストデータそのもの(documents)とそれのベクトル表現(embeddings)、各documentsの関連データ(metadatas)、各データに一意に紐づくID(ids)で各データを表現します。データの集合はcollectionsと呼ばれます。queryに与えられたテキストに対する類似度やmetadataに対する条件指定でデータを検索することが可能です。

さわってみよう!

ということで、今回はChromaを使った基本的な操作を試してみました。

  • Chromaのインストール
  • clientとcollectionの作成
  • collectionへのデータの追加
  • データの検索
  • 新しいデータの追加

Chromaのインストール

Chromaを使うためにPythonのライブラリchromaをインストールします。

pip install chroma

embeddingsの作成にはAzure OpenAIなどのEmbeddingモデルを使うことも可能です。その場合にはopenaiなどのライブラリのインストールが必要です。私はlangchainなども合わせてインストールしました。

pip install openai==0.27.6 langchain tiktoken

clientとcollectionの作成

次にdatabaseを操作するためのchromadb.Clientを作成します。ChromaはデフォルトではIn-memory databaseとして動作します。chromadb.Clientを作成する際の引数persist_directoryに指定したパスに終了時にデータを永続化し、次回そのデータをロードして使用することが出来ます。

import chromadb
from chromadb.config import Settings

client = chromadb.Client(Settings(
    chroma_db_impl="duckdb+parquet",persist_directory='./db'
))

次はcollectionを作成します。collectionを作成する際に、Embedding処理に使用するembedding functionを指定することが出来ます。デフォルトではsentence transfomerが使用されますが、OpenAIやAzure OpenAIのEmbeddingモデルを指定することが出来ます。

from chromadb.utils import embedding_functions

openai_embedding =embedding_functions.OpenAIEmbeddingFunction(
  model_name="text-embedding-ada-002",
  api_type='azure')

collection = client.get_or_create_collection(
    name="memories",embedding_function=openai_embedding
)

ここでは"memories"という名前のcollectionを作成しました。

collectionへのデータの追加

先ほど作成したcollectionにデータを追加します。ここで用意したデータですが、先の論文"Generative Agents: Interactive Simulacra of Human Behavior"でAgentが観測した出来事を記録するMemory Streamの形式に合わせてみました。出来事を観測した日時とその出来事について記録する形式です。

(なんとなく子供がいる家庭の朝の風景をイメージして作ってみたデータです。)

memories = [
"2023-06-01 06:00:00 ミルクが残りわずかしかない",
"2023-06-01 06:10:00 コーンフレークが半分以上残っている",
"2023-06-01 06:30:00 子どもが朝食にコーンフレークにミルクをかけて食べている",
"2023-06-01 06:40:00 TVが付いている",
"2023-06-01 06:50:00 コンロの火が付いている",
"2023-06-01 06:55:00 コンロの火が消えている",
"2023-06-01 07:00:00 コーヒーポットがコーヒーでいっぱいになっている",
"2023-06-01 07:10:00 子供が今日の夕ご飯はカレーがいいと言っている",
"2023-06-01 07:15:00 カレールーの在庫がない",
"2023-06-01 07:20:00 たまねぎの在庫がない",
"2023-06-01 07:20:00 ニンジンは3本残っている",
"2023-06-01 07:20:00 鶏肉が少し残っている",
"2023-06-01 07:20:00 ジャガイモの在庫がない",
"2023-06-01 07:40:00 TVのスイッチがオフになった",
"2023-06-01 07:50:00 子どもが出かけた",
"2023-06-01 08:00:00 コーヒーポットが空になった"
]

このデータから、metadataとして日時(date)を、documentとして出来事の内容を、idとして出来事の発生順に振ったシーケンシャルな番号に接頭語"id"を付与したものを抽出します。

documents = [] #テキストデータ
metadatas = [] #テキストデータに関連するデータ(検索時に利用できる)
ids = [] #documentに一意に振るID

for i, d in enumerate(memories):
  metadatas.append({
    "date":d[:19]})
  documents.append(d[20:])
  ids.append(f"id{i}")

あとはこれをcollectionに追加します。以下のコードで複数のデータをまとめて追加出来るのですが、embedding functionにAzure OpenAIのEmbeddingモデルを使用しているとエラーが発生してしまいます。

#collectionにAddする
collection.add(
  documents=documents,
  metadatas=metadatas,
  ids=ids
)

エラー内容は以下の通りです。内容から、複数の入力に対応していないことが分かります。

InvalidRequestError: Too many inputs. The max number of inputs is 1.

なので以下の様に1件ずつ処理するように書き換えました。

for d, m, i in zip(documents, metadatas, ids):
  collection.add(
    documents=d, 
    metadatas=m, 
    ids=i)

ではcollectionに含まれるデータの件数を出してみます。

collection.count()
16

.peek()を使うと最初の10件を取得することが出来ます。

print(collection.peek()["ids"])
['id0', 'id1', 'id2', 'id3', 'id4', 'id5', 'id6', 'id7', 'id8', 'id9']

最初の10件のidが表示されました。

データの検索

このcollectionにクエリを実行し、入力テキストと類似したテキストを検索し、取得します。

result = collection.query(
  query_texts="今日の晩御飯の料理名を考えてください。",n_results=1)
print(result)
result_str = result['documents']
{'ids': [['id7']], 'embeddings': None, 'documents': [['子供が今日の夕ご飯はカレーがいいと言っている']], 'metadatas': [[{'date': '2023-06-01 07:10:00'}]], 'distances': [[0.27418699860572815]]}

"今日の晩御飯の料理名を考えてください。"の類似テキストとして"子供が今日の夕ご飯はカレーがいいと言っている"が選択されました。たしかにこのテキストが一番意味が近いですね!

ではここで少しLLMと絡めた実験をしてみます。同じ人格を設定したAgentに対し、記憶を与えた場合と与えない場合で"今日の晩御飯の料理名を考えてください。"という問いに対してどんな答えを返してくれるか見てみます。Agentに与える記憶は今collectionに格納したデータです。子どもが朝話した、"今日の夕ご飯はカレーがいい"という言葉を加味できるかがポイントです。

まずはlangchainのセットアップをします。

# LangChainを使って記憶から次の行動を決める
from langchain.prompts import (
  ChatPromptTemplate,
  HumanMessagePromptTemplate,
  SystemMessagePromptTemplate
)
from langchain.chat_models import ChatOpenAI
from langchain import LLMChain

system_template = SystemMessagePromptTemplate.from_template(
  """
  あなたは日本に住んでいます。子供を大切に想っており、料理が趣味です。
  """
)
human_template = HumanMessagePromptTemplate.from_template(
  """
  {question}
  """
)

chat_prompt = ChatPromptTemplate.from_messages(
  [system_template,human_template]
)

chat = ChatOpenAI(engine="gpt-35-turbo",temperature=0)
chain = LLMChain(
  llm=chat,
  prompt=chat_prompt,
)

記憶を与えない場合を試してみます。

#記憶を使わないで今日の晩御飯を考える
chain.run(
  """
  今日の晩御飯の料理名を考えてください。
  料理名:
  """
)
'「野菜たっぷりの鶏そぼろ丼」'

たしかに子どもが好きそうだし、健康にも良さそうです。でも当たり前ですが、子どもが朝に話していた"今日の夕ご飯はカレーがいい"という記憶は反映されていません。

では次に記憶を与えた場合です。まず質問文"今日の晩御飯の料理名を考えてください。"に関連する記憶を検索します。そしてその検索結果を組み込んだpromptをLLMに入力する、という流れになります。

question = "今日の晩御飯の料理名を考えてください。"

# 記憶の中から関連する情報を抽出する。
result = collection.query(
  query_texts=question,n_results=1)
relevant_context = result['documents'][0][0]
print(f"関連する情報:{relevant_context}")

# 記憶の中から取り出した関連する情報を参考に、質問に答える。
chain.run(
  f"""
  以下の情報を参考に、質問の答えを考えてください。
  {relevant_context}

  {question}
  料理名:
  """
)
関連する情報:子供が今日の夕ご飯はカレーがいいと言っている
Out[xx]: 'カレー'

このように、"カレー"を今晩の晩御飯候補として導き出すことが出来ました。

新しいデータの追加

最後に新しいデータをこのcollectionに追加してみます。例えばこのAgentが近所のスーパーマーケットに買い物に出かけ、そこで料理の材料がセールで安くなっていたことを観測した時に、それを新しい記憶として追加するようなことを想定しています。

# idを取得する
last_id = collection.get()["ids"][-1]
last_int_id = int(last_id.replace("id",""))
next_id = f"id{last_int_id+1}"
print(next_id)

observation = "スーパーで鶏肉がセールで安く売られていた。"
date = "2023-06-01 15:00:00"

collection.add(
  documents=observation,
  metadatas={"date":date},
  ids=next_id
)

"id16"として追加されます。

collection.get(ids="id16")
{'ids': ['id16'],
 'embeddings': None,
 'documents': ['スーパーで鶏肉がセールで安く売られていた。'],
 'metadatas': [{'date': '2023-06-01 15:00:00'}]}

最後に更新された記憶を使い、今晩作る予定のカレーについて、より具体的に考えてもらいます。

#記憶を使わないで今日の晩御飯を考える
question = "費用をかけずに安く確実に作ることが出来るカレーの種類を答えてください。"

# 記憶の中から関連する情報を抽出する。
result = collection.query(
  query_texts=question,n_results=1)
relevant_context = result['documents'][0][0]
print(f"関連する情報:{relevant_context}")


# 記憶の中から取り出した関連する情報を参考に、質問に答える。
chain.run(
  f"""
  以下の情報を参考に、質問の答えを考えてください。
  {relevant_context}

  {question}
  カレーの種類:
  """
)
関連する情報:スーパーで鶏肉がセールで安く売られていた。
Out[xxx]: '鶏肉カレーが費用をかけずに安く確実に作ることができるカレーの種類です。鶏肉はセールで安く手に入れることができるため、経済的にもおすすめです。また、鶏肉は調理が簡単で、カレーにも合うため、初心者でも作りやすいカレーと言えます。'

かなり粗削りですが、行動→観測→記憶→行動→・・・と繰り返していくことでより現状に近い回答を作成してくれるAgentを作ることが出来そうです。

まとめ

ということで、今回はChromaというembedding databaseを少し触ってみた話をご紹介しました。そしてChromaのcollectionを使ってAgentの観測記憶を表現し、その記憶を元に質問に答えてもらう、ということも試してみました。論文"Generative Agents: Interactive Simulacra of Human Behavior"で紹介されている手法はもっと複雑で高度なので、もう少ししっかり読み込んで、また今度整理してみたいと思います。