LlamaIndexの新しいメモリ機能「Chat Summary Memory Buffer」「Simple Composable Memory」「Vector Memory」を試す

!pip install llama-index llama-index-callbacks-arize-phoenix

!pip freeze | egrep "llama-|arize"

arize-phoenix==4.3.0
llama-index==0.10.43
llama-index-agent-openai==0.2.7
llama-index-callbacks-arize-phoenix==0.1.5
llama-index-cli==0.1.12
llama-index-core==0.10.43
llama-index-embeddings-openai==0.1.10
llama-index-indices-managed-llama-cloud==0.1.6
llama-index-legacy==0.9.48
llama-index-llms-openai==0.1.22
llama-index-multi-modal-llms-openai==0.1.6
llama-index-program-openai==0.1.6
llama-index-question-gen-openai==0.1.3
llama-index-readers-file==0.1.23
llama-index-readers-llama-parse==0.1.4
llama-parse==0.4.4
openinference-instrumentation-llama-index==1.4.2

import os
from google.colab import userdata

os.environ["OPENAI_API_KEY"] = userdata.get("OPENAI_API_KEY")

import phoenix as px
import llama_index.core

px.launch_app()
llama_index.core.set_global_handler("arize_phoenix")

from pathlib import Path
import requests
import re

def replace_heading(match):
    level = len(match.group(1))
    return '#' * level + ' ' + match.group(2).strip()

# Wikipediaからのデータ読み込み
wiki_titles = ["オグリキャップ"]
for title in wiki_titles:
    response = requests.get(
        "https://ja.wikipedia.org/w/api.php",
        params={
            "action": "query",
            "format": "json",
            "titles": title,
            "prop": "extracts",
            # 'exintro': True,
            "explaintext": True,
        },
    ).json()
    page = next(iter(response["query"]["pages"].values()))
    wiki_text = f"# {title}\n\n## 概要\n\n"
    wiki_text += page["extract"]

    wiki_text = re.sub(r"(=+)([^=]+)\1", replace_heading, wiki_text)
    wiki_text = re.sub(r"\t+", "", wiki_text)
    wiki_text = re.sub(r"\n{3,}", "\n\n", wiki_text)
    data_path = Path("data")
    if not data_path.exists():
        Path.mkdir(data_path)

    with open(data_path / f"{title}.txt", "w") as fp:
        fp.write(wiki_text)

from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.core.node_parser import SentenceSplitter
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI
from llama_index.core import Settings

# 今回はデフォルトのLLMをgpt-4o、Embeddingはtext-embedding-3-largeとした
Settings.llm = OpenAI(model="gpt-4o")
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-large")
Settings.node_parser = SentenceSplitter(chunk_size=400, chunk_overlap=100)

documents = SimpleDirectoryReader("data").load_data()

index = VectorStoreIndex.from_documents(
    documents,
    show_progress=True
)

from llama_index.core.storage.chat_store import SimpleChatStore
from llama_index.core.memory import ChatMemoryBuffer

chat_store = SimpleChatStore()

chat_memory = ChatMemoryBuffer.from_defaults(
    token_limit=1000,
    chat_store=chat_store,
    chat_store_key="user1",
)

chat_engine = index.as_chat_engine(
    memory=chat_memory,
    similarity_top_k=5,
)

chat_engine.chat_repl()

import json

# `chat_store_key`ごとに会話履歴が保存される
for i in json.loads(chat_store.json())["store"]["user1"]:
    if i["content"] is not None:
        if len(i["content"]) >= 50:
            print("{}: {}...".format(i["role"], i["content"].replace("\n","")[:50]))
        else:
            print("{}: {}".format(i["role"], i["content"].replace("\n","")[:50]))

user: オグリキャップの血統について教えて。
tool: オグリキャップの父はダンシングキャップで、母はホワイトナルビーです。ダンシングキャップの種牡馬成績は...
assistant: オグリキャップの血統について以下の情報があります：- **父**: ダンシングキャップ- **母**...
user: ダンシングキャップについてもっと詳しい情報はある？
tool: ダンシングキャップは「ダートの短距離血統」として評価されており、その種牡馬成績はさほど優れていなかっ...
assistant: ダンシングキャップについての詳細は以下の通りです：- **評価**: 「ダートの短距離血統」として評...
user: ダンシングキャップの他の代表産駒を教えて。
tool: ダンシングキャップの代表産駒にはオグリキャップがいます。
assistant: ダンシングキャップの代表産駒として特に有名なのはオグリキャップです。他の代表産駒についての情報は見つ...
user: なるほどね、じゃあほんとに突然変異っぽい感じだね
assistant: そうですね。ダンシングキャップの種牡馬成績がさほど優れていなかった中で、オグリキャップのような名馬が...
user: 次にオグリキャップの戦績について教えてくれるかな？
tool: オグリキャップは競走馬として非常に優れた戦績を持っています。地方競馬の笠松競馬場でデビューし、8連勝...
assistant: オグリキャップの戦績は非常に優れています。以下にその主な戦績をまとめます：### 地方競馬（笠松競馬...
user: GI4勝ってのはどのレース？
tool: オグリキャップのGI勝利レースは、天皇賞（秋）です。
assistant: オグリキャップのGI勝利レースについての情報が一部不足していました。正確には、オグリキャップは以下の...
user: 有馬記念連覇してるんだ、すごいね
assistant: はい、オグリキャップは有馬記念を1988年と1990年に勝利しています。特に1990年の有馬記念は、...
user: なんで奇跡の勝利って言われたの？
tool: 1990年の有馬記念でオグリキャップが「奇跡の勝利」と呼ばれた理由は、彼が限界説が唱えられていた中で...
assistant: 1990年の有馬記念でオグリキャップが「奇跡の勝利」と呼ばれた理由は以下の通りです：1. **限界説...
user: おお、それはすごいね、きっと盛り上がったんだろうね。
assistant: はい、1990年の有馬記念は競馬史に残る名レースとなり、多くのファンが感動と興奮に包まれました。オグ...
user: わかったよ、参考になったよ、ありがとう。
assistant: どういたしまして！お役に立てて嬉しいです。もし他にも質問や知りたいことがあれば、いつでも聞いてくださ...

chat_store.persist(persist_path="chat_store.json")

chat_store = SimpleChatStore.from_persist_path(
    persist_path="chat_store.json"
)

from llama_index.core.storage.chat_store import SimpleChatStore
from llama_index.core.memory import ChatMemoryBuffer

chat_store = SimpleChatStore()

chat_memory = ChatMemoryBuffer.from_defaults(
    token_limit=50,     # 1000から50に変更
    chat_store=chat_store,
    chat_store_key="user1",
)

chat_engine = index.as_chat_engine(
    memory=chat_memory,
    similarity_top_k=5,
)

chat_engine.chat_repl()

import json

for i in json.loads(chat_store.json())["store"]["user1"]:
    if i["content"] is not None:
        if len(i["content"]) >= 50:
            print("{}: {}...".format(i["role"], i["content"].replace("\n","")[:50]))
        else:
            print("{}: {}".format(i["role"], i["content"].replace("\n","")[:50]))

>user: オグリキャップの血統について教えて。
tool: オグリキャップの父はダンシングキャップで、母はホワイトナルビーです。ダンシングキャップの種牡馬成績は...
assistant: オグリキャップの血統について以下の情報があります：- **父**: ダンシングキャップ- **母**...
user: ほう、すごいねー。
assistant: ありがとうございます！何かお手伝いできることがあれば教えてくださいね。
user: 今何の話ししてたっけ？
assistant: 特に特定の話題について話していたわけではありませんが、何か特定の話題について知りたいことや質問があれ...

{
    'additional_kwargs': {
        'tool_calls': [
            {
                'function': {
                    'arguments': '{"input":"What ''did ''the ''author ''do ''growing ''up?"}',
                    'name': 'query_engine_tool'
                },
                'id': 'call_SUGsgqUD2SzECuJOUW5HA7tl',
                'type': 'function'
            }
        ]
    },
    'content': None,
    'role': 'assistant'
}

(snip)
        for msg in chat_history_to_be_summarized:
            if msg.content is not None:         # 追加
                prompt += msg.role + ": "       # 上に合わせてインデント
                prompt += msg.content + "\n\n"  # 上に合わせてインデント
(snip)

from llama_index.core.storage.chat_store import SimpleChatStore
from llama_index.core.memory import ChatSummaryMemoryBuffer
import tiktoken

chat_store = SimpleChatStore()

# 要約用のモデルの設定
summarizer_model = "gpt-3.5-turbo-0125"
summarizer_llm = OpenAI(model_name=summarizer_model, max_toke=250)
# 要約するかどうかをトークン数でカウントするための関数（だと思う）
tokenizer_fn = tiktoken.encoding_for_model(summarizer_model).encode

chat_memory = ChatSummaryMemoryBuffer.from_defaults(
    token_limit=50,
    chat_store=chat_store,
    chat_store_key="user1",
    llm=summarizer_llm,     # 要約に使うLLMを指定
    tokenizer_fn=tokenizer_fn,    # 要約時のトークンカウント用関数を指定
)

chat_engine = index.as_chat_engine(
    memory=chat_memory,
    similarity_top_k=5,
)

chat_engine.chat_repl()

import json

json.loads(chat_store.json())["store"]["user1"]

[
    {
        'role': 'system',
        'content': "The user inquired about the specific Grade 1 races that Oguri Cap won. The assistant provided the names of the four Grade 1 races that Oguri Cap won, which are the Mile Championship, Arima Kinen, Yasuda Kinen, and Tenno Sho (Autumn). These victories solidified Oguri Cap's legendary status in Japanese horse racing.",
        'additional_kwargs': {
        }
    },
    {
        'role': 'user',
        'content': 'ありがとう',
        'additional_kwargs': {
        }
    },
    {
        'role': 'assistant',
        'content': 'どういたしまして！他に何かお手伝いできることがあれば教えてください。',
        'additional_kwargs': {
        }
    }
]

from llama_index.core.memory import VectorMemory
from llama_index.embeddings.openai import OpenAIEmbedding

vector_memory = VectorMemory.from_defaults(
    vector_store=None,  # デフォルトのオンメモリベクトルストアを使う場合は"None"にしておく
     embed_model=OpenAIEmbedding(),
    retriever_kwargs={"similarity_top_k": 1},
)

from llama_index.core.llms import ChatMessage

msgs = [
    ChatMessage.from_str("オグリキャップの血統について教えて。", "user"),
    ChatMessage.from_str("オグリキャップは、父ダンシングキャップ、母ホワイトナルビーという血統です。兄弟にはオグリローマンなどがいます。", "assistant"),
    ChatMessage.from_str("ダンシングキャップについてもっと教えて。", "user"),
    ChatMessage.from_str("ダンシングキャップはネイティヴダンサー系の「ダートの短距離血統」として評価されていますが、種牡馬成績はあまり優れていません。", "assistant"),
    ChatMessage.from_str("オグリキャップ以外のダンシングキャップの代表産駒は？", "user"),
    ChatMessage.from_str("ダンシングキャップの産駒に目立った活躍馬はいません。オグリキャップが唯一の代表産駒と言えると思います。", "assistant"),
    ChatMessage.from_str("じゃあ血統的にはほんとに突然変異と言ってもいいぐらいだね。", "user"),
    ChatMessage.from_str("そうですね。オグリキャップは血統からは予想しにくいほどの成功を収めましたが、ネイティヴダンサー系の種牡馬は時々大物を出すという特徴があるようです。", "assistant"),
    ChatMessage.from_str("次にオグリキャップの戦績について教えてくれるかな？", "user"),
    ChatMessage.from_str("オグリキャップは、地方競馬の笠松競馬場でデビューし、8連勝を含む12戦10勝という驚異的な成績を収め、さらに中央競馬へ移籍後も、重賞12勝（うちGI4勝）という素晴らしい戦績を残しています。", "assistant"),
    ChatMessage.from_str("勝利したGIレースについて教えて。", "user"),
    ChatMessage.from_str("オグリキャップは次の4つのGIで勝利しています: マイルチャンピオンシップ（1989年）、有馬記念（1988年）、安田記念（1990年）、有馬記念（1990年）", "assistant"),
    ChatMessage.from_str("有馬記念2回ってのはすごいね。", "user"),
    ChatMessage.from_str("はい、特に1990年の有馬記念は、引退レースでの「奇跡の復活」として多くのファンの記憶に残る名勝負となりました。", "assistant"),
    ChatMessage.from_str("なんで有馬記念は奇跡の復活って言われたの？", "user"),
    ChatMessage.from_str("この年の秋、オグリキャップは惨敗続きであり、競走馬としてピークを過ぎたと見られていましたが、引退レースである有馬記念で低人気を覆して掌理したためです。", "assistant"),

]

for m in msgs:
    vector_memory.put(m)

msgs = vector_memory.get("ダンシングキャップの血統は？")
msgs

[
    ChatMessage(role=<MessageRole.USER: 'user'>, content='ダンシングキャップについてもっと教えて。', additional_kwargs={}),
    ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>, content='ダンシングキャップはネイティヴダンサー系の「ダートの短距離血統」として評価されていますが、種牡馬成績はあまり優れていません。', additional_kwargs={})
]

msgs = vector_memory.get("有馬記念での成績はどうだった？")
msgs

[
    ChatMessage(role=<MessageRole.USER: 'user'>, content='有馬記念2回ってのはすごいね。', additional_kwargs={}),
    ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>, content='はい、特に1990年の有馬記念は、引退レースでの「奇跡の復活」として多くのファンの記憶に残る名勝負となりました。', additional_kwargs={})
]

vector_memory.reset()

from llama_index.core.memory import VectorMemory
from llama_index.embeddings.openai import OpenAIEmbedding

vector_memory = VectorMemory.from_defaults(
    vector_store=None,
    embed_model=OpenAIEmbedding(),
    retriever_kwargs={"similarity_top_k": 1},
)

chat_engine = index.as_chat_engine(
    memory=vector_memory,
    similarity_top_k=5,
)

chat_engine.chat_repl()

from llama_index.core.storage.chat_store import SimpleChatStore
from llama_index.core.memory import ChatMemoryBuffer, VectorMemory
from llama_index.embeddings.openai import OpenAIEmbedding

chat_store = SimpleChatStore()

chat_memory = ChatMemoryBuffer.from_defaults(
    token_limit=1000,
    chat_store=chat_store,
    chat_store_key="user1",
)

vector_memory = VectorMemory.from_defaults(
    vector_store=None,
    embed_model=OpenAIEmbedding(),
    retriever_kwargs={"similarity_top_k": 1},
)

from llama_index.core.llms import ChatMessage

msgs = [
    ChatMessage.from_str("あなたは親切なAIアシスタントです。", "system"),
    ChatMessage.from_str("太郎さんの好物はラーメンです。", "user"),
    ChatMessage.from_str("たしかに太郎さんはラーメンが大好きですよね。", "assistant"),
    ChatMessage.from_str("花子さんの方はハンバーグが好きらしいですけどね。", "user"),
    ChatMessage.from_str("花子さんの好みは太郎さんとは違うんですね。", "assistant"),
]

vector_memory.set(msgs)

from llama_index.core.memory import SimpleComposableMemory

composable_memory = SimpleComposableMemory.from_defaults(
    primary_memory=chat_memory,
    secondary_memory_sources=[vector_memory],
)

ChatMemoryBuffer(chat_store=SimpleChatStore(store={}), chat_store_key='user1', token_limit=1000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all'))

composable_memory.secondary_memory_sources

from llama_index.core.llms import ChatMessage

msgs = [
    ChatMessage.from_str("あなたは陽気で明るい日本語のAIアシスタントです。", "system"),
    ChatMessage.from_str("次郎さんはアイスクリームが大好きらしいです。", "user"),
]

for m in msgs:
    composable_memory.put(m)

composable_memory.primary_memory

composable_memory.secondary_memory_sources

msgs = composable_memory.get("太郎さんの好きなものは何？")
msgs

for idx, m in enumerate(msgs):
    print(f"----- {idx} -----")
    print(m)

msgs = composable_memory.get("花子さんの好物は？")

for idx, m in enumerate(msgs):
    print(f"----- {idx} -----")
    print(m)

msgs = composable_memory.get("次郎さんの好物は？")

for idx, m in enumerate(msgs):
    print(f"----- {idx} -----")
    print(m)

composable_memory.primary_memory.reset()

composable_memory.reset()

composable_memory.secondary_memory_sources[0].reset()

composable_memory.secondary_memory_sources[0].get()

composable_memory.primary_memory.get()

composable_memory.reset()
print(composable_memory.primary_memory.get())
print(composable_memory.secondary_memory_sources[0].get())

import nest_asyncio

nest_asyncio.apply()

from llama_index.llms.openai import OpenAI
from llama_index.core.tools import FunctionTool
from llama_index.core.agent import FunctionCallingAgentWorker
from llama_index.core.memory import ChatMemoryBuffer, VectorMemory
from llama_index.embeddings.openai import OpenAIEmbedding

# composable_memoryを使ったメモリの定義
vector_memory = VectorMemory.from_defaults(
    vector_store=None,
    embed_model=OpenAIEmbedding(),
    retriever_kwargs={"similarity_top_k": 2},
)

chat_memory_buffer = ChatMemoryBuffer.from_defaults()

composable_memory = SimpleComposableMemory.from_defaults(
    primary_memory=chat_memory_buffer,
    secondary_memory_sources=[vector_memory],
)

# Agentのツールを定義
def multiply(a: int, b: int) -> int:
    """2つの整数を掛け算して、結果を整数で返す。"""
    return a * b


def mystery(a: int, b: int) -> int:
    """2つの整数を処理する謎の関数。"""
    return a**2 - b**2


multiply_tool = FunctionTool.from_defaults(fn=multiply)
mystery_tool = FunctionTool.from_defaults(fn=mystery)

# ツールとメモリを使用するエージェントの定義
llm = OpenAI(model="gpt-3.5-turbo-0613")
agent_worker = FunctionCallingAgentWorker.from_tools(
    [multiply_tool, mystery_tool],
    llm=llm,
    verbose=True
)
agent = agent_worker.as_agent(memory=composable_memory)

response = agent.chat("5と6を謎の関数に渡すと結果はどうなる?")

response = agent.chat("2 かける 3 は？")

composable_memory.primary_memory

ChatMemoryBuffer(chat_store=SimpleChatStore(store={
    'chat_history': [
        ChatMessage(role=<MessageRole.USER: 'user'>,
        content='5と6を謎の関数に渡すと結果はどうなる?',
        additional_kwargs={
        }
        ),
        ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>,
        content=None,
        additional_kwargs={
            'tool_calls': [
                ChatCompletionMessageToolCall(id='call_2itmKKngGYojlvlbla0bUlKu',
                function=Function(arguments='{\n  "a": 5,\n  "b": 6\n}',
                name='mystery'),
                type='function')
            ]
        }
        ),
        ChatMessage(role=<MessageRole.TOOL: 'tool'>,
        content='-11',
        additional_kwargs={
            'name': 'mystery',
            'tool_call_id': 'call_2itmKKngGYojlvlbla0bUlKu'
        }
        ),
        ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>,
        content='5と6を謎の関数に渡すと、結果は-11になります。',
        additional_kwargs={
        }
        ),
        ChatMessage(role=<MessageRole.USER: 'user'>,
        content='2 かける 3 は？',
        additional_kwargs={
        }
        ),
        ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>,
        content=None,
        additional_kwargs={
            'tool_calls': [
                ChatCompletionMessageToolCall(id='call_Z5s49YuMdtRwYuCA3VfYYUqD',
                function=Function(arguments='{\n  "a": 2,\n  "b": 3\n}',
                name='multiply'),
                type='function')
            ]
        }
        ),
        ChatMessage(role=<MessageRole.TOOL: 'tool'>,
        content='6',
        additional_kwargs={
            'name': 'multiply',
            'tool_call_id': 'call_Z5s49YuMdtRwYuCA3VfYYUqD'
        }
        ),
        ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>,
        content='2かける3は6です。',
        additional_kwargs={
        }
        )
    ]
}
),
chat_store_key='chat_history',
token_limit=3000,
tokenizer_fn=functools.partial(<boundmethodEncoding.encodeof<Encoding'cl100k_base'>>,
allowed_special='all'))

composable_memory.secondary_memory_sources[0]

VectorMemory(vector_index=<llama_index.core.indices.vector_store.base.VectorStoreIndexobjectat0x7eb648fc9060>,
retriever_kwargs={
    'similarity_top_k': 2
},
batch_by_user_message=True,
cur_batch_textnode=TextNode(id_='fe2fa265-7860-44e2-90a0-98c33830f96d',
embedding=None,
metadata={
    'sub_dicts': [
        {
            'role': <MessageRole.USER: 'user'>,
            'content': '2 かける 3 は？',
            'additional_kwargs': {
            }
        },
        {
            'role': <MessageRole.ASSISTANT: 'assistant'>,
            'content': None,
            'additional_kwargs': {
                'tool_calls': [
                    {
                        'id': 'call_Z5s49YuMdtRwYuCA3VfYYUqD',
                        'function': {
                            'arguments': '{\n  "a": 2,\n  "b": 3\n}',
                            'name': 'multiply'
                        },
                        'type': 'function'
                    }
                ]
            }
        },
        {
            'role': <MessageRole.TOOL: 'tool'>,
            'content': '6',
            'additional_kwargs': {
                'name': 'multiply',
                'tool_call_id': 'call_Z5s49YuMdtRwYuCA3VfYYUqD'
            }
        },
        {
            'role': <MessageRole.ASSISTANT: 'assistant'>,
            'content': '2かける3は6です。',
            'additional_kwargs': {
            }
        }
    ]
},
excluded_embed_metadata_keys=[
    'sub_dicts'
],
excluded_llm_metadata_keys=[
    'sub_dicts'
],
relationships={
},
text='2 かける 3 は？  6 2かける3は6です。',
start_char_idx=None,
end_char_idx=None,
text_template='{metadata_str}\n\n{content}',
metadata_template='{key}: {value}',
metadata_seperator='\n'))

composable_memory.secondary_memory_sources[0].get("5と6を謎の関数に渡すと結果はどうなる?")

[
    ChatMessage(role=<MessageRole.USER: 'user'>,
    content='5と6を謎の関数に渡すと結果はどうなる?',
    additional_kwargs={
    }
    ),
    ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>,
    content=None,
    additional_kwargs={
        'tool_calls': [
            {
                'id': 'call_2itmKKngGYojlvlbla0bUlKu',
                'function': {
                    'arguments': '{\n  "a": 5,\n  "b": 6\n}',
                    'name': 'mystery'
                },
                'type': 'function'
            }
        ]
    }
    ),
    ChatMessage(role=<MessageRole.TOOL: 'tool'>,
    content='-11',
    additional_kwargs={
        'name': 'mystery',
        'tool_call_id': 'call_2itmKKngGYojlvlbla0bUlKu'
    }
    ),
    ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>,
    content='5と6を謎の関数に渡すと、結果は-11になります。',
    additional_kwargs={
    }
    )
]

llm = OpenAI(model="gpt-3.5-turbo-0613")
agent_worker = FunctionCallingAgentWorker.from_tools(
    [multiply_tool, mystery_tool], llm=llm, verbose=True
)
agent_without_memory = agent_worker.as_agent()

response = agent_without_memory.chat(
    "さっき5と6を謎の関数で実行した結果っていくつだったっけ？再計算はしちゃだめだよ。"
)

agent_without_memory.chat_history

llm = OpenAI(model="gpt-3.5-turbo-0613")
agent_worker = FunctionCallingAgentWorker.from_tools(
    [multiply_tool, mystery_tool], llm=llm, verbose=True
)
composable_memory = SimpleComposableMemory.from_defaults(
    # primary_memoryは初期化した新しいものを与える
    primary_memory=ChatMemoryBuffer.from_defaults(),
    # secondary_memory_sourcesは上で使用したものを使うが、
    # 後ほど再度使うので、直接使わずにここではコピーしている
    secondary_memory_sources=[
        vector_memory.copy(
            deep=True
        )
    ],
)
agent_with_memory = agent_worker.as_agent(memory=composable_memory)

response = agent_with_memory.chat(
    "さっき5と6を謎の関数で実行した結果っていくつだったっけ？再計算しないで答えて。"
)

response = agent_with_memory.chat(
    "さっき計算した 2 かける 3の結果っていくつだっけ？再計算しないで答えて。"
)

llm = OpenAI(model="gpt-3.5-turbo-0613")
agent_worker = FunctionCallingAgentWorker.from_tools(
    [multiply_tool, mystery_tool], llm=llm, verbose=True
)
composable_memory = SimpleComposableMemory.from_defaults(
    primary_memory=ChatMemoryBuffer.from_defaults(),
    secondary_memory_sources=[
        vector_memory.copy(
            deep=True
        )
    ],
)
agent_with_memory = agent_worker.as_agent(memory=composable_memory)

agent_with_memory.memory.get(
    "さっき5と6を謎の関数で実行した結果っていくつだったっけ？再計算しないで答えて。"
)

print(
    agent_with_memory.memory.get(
        "さっき5と6を謎の関数で実行した結果っていくつだったっけ？再計算しないで答えて。"
    )[0]
)

from llama_index.core.storage.chat_store import SimpleChatStore
from llama_index.core.memory import ChatMemoryBuffer, VectorMemory, SimpleComposableMemory
from llama_index.core.llms import ChatMessage

vector_memory = VectorMemory.from_defaults(
    vector_store=None,
    embed_model=OpenAIEmbedding(),
    retriever_kwargs={"similarity_top_k": 2},
)

# インデックスに存在しない情報をsecondary_memory_sourcesにセット
vector_memory.set([
    ChatMessage.from_str("ホワイトナルビーの父はなんという馬ですか？", "user"),
    ChatMessage.from_str("ホワイトナルビーの父はシルバーシャークです。", "assistant"),
])

chat_store = SimpleChatStore()

chat_memory = ChatMemoryBuffer.from_defaults(
    token_limit=1000,
    chat_store=chat_store,
    chat_store_key="user1",
)

composable_memory = SimpleComposableMemory.from_defaults(
    primary_memory=chat_memory,
    secondary_memory_sources=[vector_memory],
)

chat_engine = index.as_chat_engine(
    memory=composable_memory,
    similarity_top_k=5,
)

response = chat_engine.chat("オグリキャップの血統について教えて")
print(response

response = chat_engine.chat("母ホワイトナルビーの父は？")
print(response)

chat_engine.memory.secondary_memory_sources[0].get("母ホワイトナルビーの父は？")

[
    ChatMessage(role=<MessageRole.USER: 'user'>,
    content='母ホワイトナルビーの父は？',
    additional_kwargs={
    }
    ),
    ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>,
    content=None,
    additional_kwargs={
        'tool_calls': [
            {
                'id': 'call_gvypKJ8TseHuMTl1idSmYUbk',
                'function': {
                    'arguments': '{"input":"ホワイトナルビーの父は誰ですか？"}',
                    'name': 'query_engine_tool'
                },
                'type': 'function'
            }
        ]
    }
    ),
    ChatMessage(role=<MessageRole.TOOL: 'tool'>,
    content='ホワイトナルビーの父はクモハタです。',
    additional_kwargs={
        'name': 'query_engine_tool',
        'tool_call_id': 'call_gvypKJ8TseHuMTl1idSmYUbk'
    }
    ),
    ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>,
    content='ホワイトナルビーの父は**クモハタ**です。',
    additional_kwargs={
    }
    ),
    ChatMessage(role=<MessageRole.USER: 'user'>,
    content='ホワイトナルビーの父はなんという馬ですか？',
    additional_kwargs={
    }
    ),
    ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>,
    content='ホワイトナルビーの父はシルバーシャークです。',
    additional_kwargs={
    }
    )
]

from llama_index.core.storage.chat_store import SimpleChatStore
from llama_index.core.memory import ChatMemoryBuffer, VectorMemory, SimpleComposableMemory
from llama_index.core.llms import ChatMessage

vector_memory = VectorMemory.from_defaults(
    vector_store=None,
    embed_model=OpenAIEmbedding(),
    retriever_kwargs={"similarity_top_k": 2},
)
# インデックスにない情報をsecondary_memory_sourcesにセット
vector_memory.set([
    ChatMessage.from_str("ホワイトナルビーの父はなんという馬ですか？", "user"),
    ChatMessage.from_str("ホワイトナルビーの父はシルバーシャークです。", "assistant"),
])

chat_store = SimpleChatStore()

chat_memory = ChatMemoryBuffer.from_defaults(
    token_limit=1000,
    chat_store=chat_store,
    chat_store_key="user1",
)

composable_memory = SimpleComposableMemory.from_defaults(
    primary_memory=chat_memory,
    secondary_memory_sources=[vector_memory],
)

# インデックスからQuery Engineを作成
query_engine = index.as_query_engine(
    similarity_top_k=5,
)

from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.tools import FunctionTool
from llama_index.core.agent import FunctionCallingAgentWorker

# QueryEngineToolでquery engineをツール化
query_engine_tool = QueryEngineTool(
    query_engine=query_engine,
    metadata=ToolMetadata(
        name=f"oguricap_retriever",
        description=f"競走馬「オグリキャップ」に関する情報を検索する",
    ),
)

agent_worker = FunctionCallingAgentWorker.from_tools(
    [query_engine_tool],
    verbose=True
)

# composable_memoryを使用したエージェントを作成
agent = agent_worker.as_agent(memory=composable_memory)

response = agent.chat("オグリキャップの血統について教えて。")

response = agent.chat("母ホワイトナルビーの父は？")