langchainのチャットボットにmem0aiで長期記憶を持たせる

# pip install langchain langchain_openai mem0ai

import os
from typing import List, Dict
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from mem0 import MemoryClient

# Configuration
os.environ["OPENAI_API_KEY"] = "**********************************************"
os.environ["MEM0_API_KEY"] = "************************************************"

# Initialize LangChain and Mem0
llm = ChatOpenAI(model="gpt-4o-mini")
mem0 = MemoryClient(api_key=os.environ["MEM0_API_KEY"])

prompt = ChatPromptTemplate.from_messages([
    SystemMessage(content="""あなたは役に立つ旅行代理店の AI です。提供されたコンテキストを使用して応答をパーソナライズし、ユーザーの好みや過去のやり取りを記憶します。
旅行の推奨事項、旅程の提案を提供し、目的地に関する質問に答えます。
具体的な情報がない場合は、一般的な旅行知識に基づいて一般的な提案を行うことができます。"""),
    MessagesPlaceholder(variable_name="context"),
    HumanMessage(content="{input}")
])

def retrieve_context(query: str, user_id: str) -> List[Dict]:
    """Mem0から関連するコンテキストを取得する"""
    memories = mem0.search(query, user_id=user_id)
    seralized_memories = ' '.join([mem["memory"] for mem in memories])
    context = [
        {
            "role": "system",
            "content": f"Relevant information: {seralized_memories}"
        },
        {
            "role": "user",
            "content": query
        }
    ]
    return context

def generate_response(input: str, context: List[Dict]) -> str:
    """言語モデルを使用して応答を生成する"""
    chain = prompt | llm
    response = chain.invoke({
        "context": context,
        "input": input
    })
    return response.content

def save_interaction(user_id: str, user_input: str, assistant_response: str):
    """インタラクションをMem0に保存する"""
    interaction = [
        {
          "role": "user",
          "content": user_input
        },
        {
            "role": "assistant",
            "content": assistant_response
        }
    ]
    mem0.add(interaction, user_id=user_id)


def chat_turn(user_input: str, user_id: str) -> str:
    # Retrieve context
    context = retrieve_context(user_input, user_id)

    # Generate response
    response = generate_response(user_input, context)

    # Save interaction
    save_interaction(user_id, user_input, response)

    return response


if __name__ == "__main__":
    print("パーソナル旅行代理店プランナーへようこそ! 今日の旅行計画について、どのようにお手伝いできますか?")
    user_id = "john"

    while True:
        user_input = input("You: ")
        if user_input.lower() in ['quit', 'exit', 'bye']:
            print("旅行代理店: 当社の旅行プランニングサービスをご利用いただきありがとうございます。楽しいご旅行をお過ごしください。")
            break

        response = chat_turn(user_input, user_id)
        print(f"旅行代理店: {response}")

langchainのチャットボットに長期記憶を持たせます。mem0ai公式のコードを和訳しました。