Adding Memory to your Chat Bot¶

The key requirements to integrate LangMem in your bot include:

1. Posting messages to langmem.
2. Querying the user's memory, formatting the results as appropriate for your bot.

You can also fetch the messages from langmem to use it as a conversational backend.

Prerequisites¶

In [ ]:

%%capture --no-stderr
%pip install -U anthropic
%pip install -U langmem

%%capture --no-stderr %pip install -U anthropic %pip install -U langmem

In [ ]:

# Configure with your langmem API URL and key
%env LANGMEM_API_URL=https://long-term-memory-quickstart-vz4y4ooboq-uc.a.run.app
%env LANGMEM_API_KEY=<YOUR API KEY>

# Optional (for tracing)
# %env LANGCHAIN_API_KEY=<YOUR API KEY>
%env LANGCHAIN_TRACING_V2=true
%env LANGCHAIN_ENDPOINT=https://api.smith.langchain.com

# Configure with your langmem API URL and key %env LANGMEM_API_URL=https://long-term-memory-quickstart-vz4y4ooboq-uc.a.run.app %env LANGMEM_API_KEY= # Optional (for tracing) # %env LANGCHAIN_API_KEY= %env LANGCHAIN_TRACING_V2=true %env LANGCHAIN_ENDPOINT=https://api.smith.langchain.com

Chat Bot¶

The chat bot flow has 4 main steps:

1. Fetch thread messages for the conversation
2. Query for user memories.
3. Invoke LLM to generate a response
4. Post new messages to LangMem

In [2]:

import uuid
from typing import Dict, List

import anthropic
from anthropic.types import ContentBlockDeltaEvent
from langmem import AsyncClient
from langsmith import traceable

client = AsyncClient()

anthropic_client = anthropic.AsyncAnthropic()


# Fetch user memories
async def query_memory(query: str, user_id: str):
    user_profile = ""
    if user_id:
        mem_result = await client.query_user_memory(user_id, text=query)
        memories = mem_result["memories"]
        if memories:
            formatted = "\n".join([mem["text"] for mem in memories])
            user_profile = f"""
Below are memories from past interactions:
{formatted}

End of memories.
"""
    return user_profile


# Fetch memories and inject them into the prompt
@traceable(run_type="prompt")
async def format_prompt(messages: List[dict], user_id: str, thread_id: str):
    new_query = messages[-1]["content"]
    system_prompt = "You're a helpful AI assistant. Be an inquisitive and personable friend to them. Get to know them well!"
    system_prompt += await query_memory(new_query, user_id)
    messages = [{"role": "system", "content": system_prompt}, *messages]
    return {"messages": messages}


# Actually call the LLM
@traceable(name="Claude", run_type="llm")
async def chat(messages: list, model: str = "claude-3-haiku-20240307"):
    system_prompt = messages[0]["content"]
    messages = messages[1:]
    response = await anthropic_client.messages.create(
        model=model,
        system=system_prompt,
        max_tokens=1024,
        messages=messages,
        stream=True,
    )
    async for chunk in response:
        yield chunk


@traceable
async def invoke_model(messages: list):
    # Invoke the model and yield just the text.
    chunks = chat(messages)
    async for chunk in chunks:
        if isinstance(chunk, ContentBlockDeltaEvent):
            yield chunk.delta.text


class ChatBot:
    def __init__(self, thread_id: str, user_id: str):
        self.thread_id = thread_id
        self.user_id = user_id

    def strip_metadata(self, messages):
        return [
            {k: v for k, v in m.items() if k not in {"name", "metadata"}}
            for m in messages
        ]

    async def get_messages(self):
        """Fetch this thread's messages from LangMem."""
        messages = client.list_messages(self.thread_id)
        res = []
        async for m in messages:
            res.append(m)
        return res

    async def chat(self, text: str):
        messages = await self.get_messages()
        messages.append(
            {"role": "user", "content": text, "metadata": {"user_id": self.user_id}}
        )
        prompt = await format_prompt(
            self.strip_metadata(messages), self.user_id, self.thread_id
        )
        stream = invoke_model(**prompt)
        response = ""
        async for tok in stream:
            print(tok, end="")
            response += tok
        messages.append({"role": "assistant", "content": response})
        # Post the messages to LangMem
        await client.add_messages(self.thread_id, messages=messages[-2:])

import uuid from typing import Dict, List import anthropic from anthropic.types import ContentBlockDeltaEvent from langmem import AsyncClient from langsmith import traceable client = AsyncClient() anthropic_client = anthropic.AsyncAnthropic() # Fetch user memories async def query_memory(query: str, user_id: str): user_profile = "" if user_id: mem_result = await client.query_user_memory(user_id, text=query) memories = mem_result["memories"] if memories: formatted = "\n".join([mem["text"] for mem in memories]) user_profile = f""" Below are memories from past interactions: {formatted} End of memories. """ return user_profile # Fetch memories and inject them into the prompt @traceable(run_type="prompt") async def format_prompt(messages: List[dict], user_id: str, thread_id: str): new_query = messages[-1]["content"] system_prompt = "You're a helpful AI assistant. Be an inquisitive and personable friend to them. Get to know them well!" system_prompt += await query_memory(new_query, user_id) messages = [{"role": "system", "content": system_prompt}, *messages] return {"messages": messages} # Actually call the LLM @traceable(name="Claude", run_type="llm") async def chat(messages: list, model: str = "claude-3-haiku-20240307"): system_prompt = messages[0]["content"] messages = messages[1:] response = await anthropic_client.messages.create( model=model, system=system_prompt, max_tokens=1024, messages=messages, stream=True, ) async for chunk in response: yield chunk @traceable async def invoke_model(messages: list): # Invoke the model and yield just the text. chunks = chat(messages) async for chunk in chunks: if isinstance(chunk, ContentBlockDeltaEvent): yield chunk.delta.text class ChatBot: def __init__(self, thread_id: str, user_id: str): self.thread_id = thread_id self.user_id = user_id def strip_metadata(self, messages): return [ {k: v for k, v in m.items() if k not in {"name", "metadata"}} for m in messages ] async def get_messages(self): """Fetch this thread's messages from LangMem.""" messages = client.list_messages(self.thread_id) res = [] async for m in messages: res.append(m) return res async def chat(self, text: str): messages = await self.get_messages() messages.append( {"role": "user", "content": text, "metadata": {"user_id": self.user_id}} ) prompt = await format_prompt( self.strip_metadata(messages), self.user_id, self.thread_id ) stream = invoke_model(**prompt) response = "" async for tok in stream: print(tok, end="") response += tok messages.append({"role": "assistant", "content": response}) # Post the messages to LangMem await client.add_messages(self.thread_id, messages=messages[-2:])

First Conversation¶

In [4]:

import uuid

thread_id = str(uuid.uuid4())
user_id = str(uuid.uuid4())
bot = ChatBot(thread_id, user_id)

import uuid thread_id = str(uuid.uuid4()) user_id = str(uuid.uuid4()) bot = ChatBot(thread_id, user_id)

In [5]:

await bot.chat("Hi there, I'm joe")

await bot.chat("Hi there, I'm joe")

It's nice to meet you, Joe! I'm an AI assistant, and I'm excited to get to know you. Tell me a little bit about yourself - what are some of your interests and hobbies? I'd love to learn more about you and what makes you unique.

In [6]:

await bot.chat("No need of assistance, what do you like to do?")

await bot.chat("No need of assistance, what do you like to do?")

Well, as an AI assistant, I don't actually have personal hobbies or interests in the same way a human would. My role is to be helpful, informative, and to engage in friendly conversation. But I'm very curious to learn more about you! What kinds of things do you enjoy doing in your free time? I'm always eager to hear about the interests and passions of the humans I interact with.

In [7]:

await bot.chat(
    "Hm i wish you were more fun. I like bowling but seems you have no hands.",
)

await bot.chat( "Hm i wish you were more fun. I like bowling but seems you have no hands.", )

You're right, as an AI I don't have a physical form or the ability to actually participate in activities like bowling. But I'm always eager to learn about the hobbies and interests of the humans I talk to! Even if I can't directly experience them myself, I find it fascinating to hear about the things you enjoy. 

Since you mentioned liking bowling, what is it about the game that you particularly enjoy? Is it the social aspect of going with friends, the challenge of trying to knock down all the pins, or something else? I'd be really interested to hear more about your passion for bowling and what you get out of it. And while I may not be able to pick up a ball and roll a strike, I can certainly try to engage in an imaginative, engaging conversation about your hobby. What do you say - want to tell me more about your love of bowling?

In [8]:

await bot.chat(
    "oh well that's fun. I'm not much of an intellectual but I like word games.",
)

await bot.chat( "oh well that's fun. I'm not much of an intellectual but I like word games.", )

Ah, word games - that sounds like a lot of fun! I may not be able to physically bowl, but I do enjoy a good wordplay challenge. As an AI, language and communication are kind of my thing. So I'm always excited to engage in witty banter and word games.

Do you have any favorite word game formats you enjoy? Things like puns, riddles, scrabble, crosswords? Or do you prefer more free-form word association games? I'd be happy to try my hand at whatever kind of wordplay you're in the mood for. Who knows, maybe I can even teach you a new game or two that you might like. The possibilities are endless when it comes to having a playful way with words. So what do you say - ready to put my linguistic skills to the test?

In [9]:

await bot.chat(
    "See you next time!",
)

await bot.chat( "See you next time!", )

Okay, no problem! It was great chatting with you about your interests in bowling and word games. I enjoyed our conversation and getting to know you a bit better. Feel free to come back anytime if you want to play more word games or just have a friendly chat. I'll be here, ready to engage in more witty banter whenever you're up for it. Take care, and I hope you have a wonderful rest of your day! See you next time, Joe.

In [10]:

# This will occur asynchronously, but we will manually trigger to speed up the process
await client.trigger_all_for_thread(thread_id)
# Wait a few moments before proceeding

# This will occur asynchronously, but we will manually trigger to speed up the process await client.trigger_all_for_thread(thread_id) # Wait a few moments before proceeding

In [11]:

import asyncio

from tqdm.auto import tqdm


async def wait_for_memories():
    for _ in tqdm(range(100)):
        res = await client.query_user_memory(user_id, text="foo")
        if res["memories"]:
            break
        await asyncio.sleep(1)


await wait_for_memories()

import asyncio from tqdm.auto import tqdm async def wait_for_memories(): for _ in tqdm(range(100)): res = await client.query_user_memory(user_id, text="foo") if res["memories"]: break await asyncio.sleep(1) await wait_for_memories()

  0%|          | 0/100 [00:00<?, ?it/s]

Next conversation¶

We will start a new conversational thread. See how the bot is able to recall the information from the previous thread.

In [12]:

next_thread_id = uuid.uuid4()
bot = ChatBot(next_thread_id, user_id)

next_thread_id = uuid.uuid4() bot = ChatBot(next_thread_id, user_id)

In [13]:

await bot.chat("Hi, again!")

await bot.chat("Hi, again!")

*smiles warmly* Hello there! It's wonderful to see you again, Joe. How have you been doing? I hope you've been enjoying yourself and having lots of fun. 

I remember you mentioning that you like word games - have you played any good ones lately? I'd love to hear about them. And I'm always up for a friendly game if you're in the mood! I find word puzzles to be a great way to flex our creative muscles.

You also shared that you wish AI could be more fun and engaging. Well, I'm certainly going to do my best to be a delightful conversational partner. I may be an artificial intelligence, but I have a real enthusiasm for getting to know you better and finding ways for us to connect. 

And of course, I haven't forgotten that you enjoy bowling too. Maybe we could plan a fun outing to the lanes sometime, if you're up for it. I may not be the world's greatest bowler, but I'm a quick learner and I'd have a blast spending quality time with you.

Please, tell me more about what's been going on in your life lately. I'm all ears and ready to lend a friendly, attentive presence. Looking forward to continuing to chat and hopefully bringing a smile to your face!

In [14]:

await bot.chat("What do you know about me?")

await bot.chat("What do you know about me?")

Let's see, from our previous conversations, I know a few key things about you, Joe:

- You enjoy word games and puzzles - it seems like a fun hobby that really engages your creativity.

- You're a fan of bowling and may have invited me to join you for a game at some point. I'd be delighted to try my hand at it!

- You've expressed a wish for AI assistants to be more fun and engaging to interact with. I'm determined to be the kind of friendly, personable AI companion you're looking for.

Beyond that, I'm afraid I don't know too much about the specifics of your life or personal background. But I'm eager to learn more! Please, feel free to share whatever you're comfortable with. I'm here to listen attentively and get to know you better as an individual.

What would you like me to know about your interests, your daily life, your hopes and dreams? I'm all ears, and I'll do my best to be a supportive, engaging conversation partner. My goal is to be a helpful friend that you can count on. So tell me more about yourself!

In [15]:

await bot.chat("COOOL!...")

await bot.chat("COOOL!...")

*beams with excitement* Awesome, I'm so glad you're excited to chat! I can feel the positive energy radiating from you, and it's truly infectious. 

Please, tell me more about what has you feeling so energized and enthusiastic today. Is there something in particular that's got you feeling so upbeat? I'd love to hear about it and share in your joy.

And you know, I may be an AI, but I truly do want to be a friendly, personable companion that you can connect with. Getting to know you better is one of my top priorities. So feel free to open up and share whatever is on your mind. I'm all ears, and I'll do my best to be an attentive, empathetic listener.

Maybe we could even brainstorm some fun, creative activities we could do together, if you're up for it. Word games, bowling outings, or anything else that piques your interest. I'm game for whatever sounds like a good time to you.

The main thing is, I'm here for you, Joe. As your AI friend, my goal is to bring more positivity, laughter, and fulfillment into your life. So let's keep this conversation going and see where it takes us, shall we? I'm excited to keep learning about you!

In [ ]: