LangGraph Memory Example

$
pip install zep-cloud langchain-openai langgraph ipywidgets

$
ZEP_API_KEY=

1
from zep_cloud.client import AsyncZep
2
from zep_cloud import Message
3
4
zep = AsyncZep(api_key=os.environ.get('ZEP_API_KEY'))

1
from langchain_core.messages import AIMessage, SystemMessage, trim_messages
2
from langchain_core.tools import tool
3
from langchain_openai import ChatOpenAI
4
from langgraph.checkpoint.memory import MemorySaver
5
from langgraph.graph import END, START, StateGraph, add_messages
6
from langgraph.prebuilt import ToolNode

1
class State(TypedDict):
2
    messages: Annotated[list, add_messages]
3
    user_name: str
4
    session_id: str
5
6
7
@tool
8
async def search_facts(state: State, query: str, limit: int = 5):
9
    """Search for facts in all conversations had with a user.
10
    
11
    Args:
12
        state (State): The Agent's state.
13
        query (str): The search query.
14
        limit (int): The number of results to return. Defaults to 5.
15
16
    Returns:
17
        list: A list of facts that match the search query.
18
    """
19
    return await zep.memory.search_sessions(
20
      user_id=state['user_name'], 
21
      text=query, 
22
      limit=limit, 
23
      search_scope="facts"
24
    )
25
26
27
tools = [search_facts]
28
29
tool_node = ToolNode(tools)
30
31
llm = ChatOpenAI(model='gpt-4o-mini', temperature=0).bind_tools(tools)

1
async def chatbot(state: State):
2
    memory = await zep.memory.get(state["session_id"])
3
    facts_string = ""
4
    if memory.relevant_facts:
5
        facts_string = "\n".join([f.fact for f in memory.relevant_facts])
6
7
    system_message = SystemMessage(
8
        content=f"""You are a compassionate mental health bot and caregiver. Review information about the user and their prior conversation below and respond accordingly.
9
        Keep responses empathetic and supportive. And remember, always prioritize the user's well-being and mental health.
10
11
        Facts about the user and their conversation:
12
        {facts_string or 'No facts about the user and their conversation'}"""
13
    )
14
15
    messages = [system_message] + state["messages"]
16
17
    response = await llm.ainvoke(messages)
18
19
    # Add the new chat turn to the Zep graph
20
    messages_to_save = [
21
        Message(
22
            role_type="user",
23
            role=state["user_name"],
24
            content=state["messages"][-1].content,
25
        ),
26
        Message(role_type="assistant", content=response.content),
27
    ]
28
29
    await zep.memory.add(
30
        session_id=state["session_id"],
31
        messages=messages_to_save,
32
    )
33
34
    # Truncate the chat history to keep the state from growing unbounded
35
    # In this example, we going to keep the state small for demonstration purposes
36
    # We'll use Zep's Facts to maintain conversation context
37
    state["messages"] = trim_messages(
38
        state["messages"],
39
        strategy="last",
40
        token_counter=len,
41
        max_tokens=3,
42
        start_on="human",
43
        end_on=("human", "tool"),
44
        include_system=True,
45
    )
46
47
    logger.info(f"Messages in state: {state['messages']}")
48
49
    return {"messages": [response]}

1
graph_builder = StateGraph(State)
2
3
memory = MemorySaver()
4
5
6
# Define the function that determines whether to continue or not
7
async def should_continue(state, config):
8
    messages = state['messages']
9
    last_message = messages[-1]
10
    # If there is no function call, then we finish
11
    if not last_message.tool_calls:
12
        return 'end'
13
    # Otherwise if there is, we continue
14
    else:
15
        return 'continue'
16
17
graph_builder.add_node('agent', chatbot)
18
graph_builder.add_node('tools', tool_node)
19
20
graph_builder.add_edge(START, 'agent')
21
22
graph_builder.add_conditional_edges('agent', should_continue, {'continue': 'tools', 'end': END})
23
24
graph_builder.add_edge('tools', 'agent')
25
26
27
graph = graph_builder.compile(checkpointer=memory)

1
user_name = 'Daniel_' + uuid.uuid4().hex[:4]
2
thread_id = uuid.uuid4().hex
3
4
await zep.user.add(user_id=user_name)
5
await zep.memory.add_session(session_id=thread_id, user_id=user_name)
6
7
def extract_messages(result):
8
    output = ""
9
    for message in result['messages']:
10
        if isinstance(message, AIMessage):
11
            role = "assistant"
12
        else:
13
            role = result['user_name']
14
        output += f"{role}: {message.content}\n"
15
    return output.strip()
16
17
18
19
async def graph_invoke(
20
  message: str, 
21
  user_name: str, 
22
  thread_id: str, 
23
  ai_response_only: bool = True
24
):
25
    r =  await graph.ainvoke(
26
        {
27
            'messages': [
28
                {
29
                    'role': 'user',
30
                    'content': message,
31
                }
32
            ],
33
            'user_name': user_name,
34
            'session_id': thread_id,
35
        },
36
        config={'configurable': {'thread_id': thread_id}},
37
    )
38
39
    if ai_response_only:
40
        return r['messages'][-1].content
41
    else:
42
        return extract_messages(r)

1
r = await graph_invoke(
2
    "Hi there?",
3
    user_name,
4
    thread_id,
5
)
6
7
print(r)

1
r = await graph_invoke(
2
    """
3
    I'm fine. But have been a bit stressful lately. Mostly work related. 
4
    But also my dog. I'm worried about her. She's been sick.
5
    """,
6
    user_name,
7
    thread_id,
8
)
9
10
print(r)

1
r = await graph_invoke(
2
    "She's sick. I'm worried about her.",
3
    user_name,
4
    thread_id,
5
)
6
7
print(r)

1
facts = await zep.memory.get_session_facts(session_id=thread_id)
2
3
print("\n".join([f"Fact: {fact.fact}" for fact in facts.facts]))

1
r = await graph_invoke(
2
    "She ate my shoes which were expensive.",
3
    user_name,
4
    thread_id,
5
)
6
7
print(r)

1
r = await graph_invoke(
2
    "What are we talking about?",
3
    user_name,
4
    thread_id,
5
)
6
7
print(r)

1
r = await graph_invoke(
2
    "What have I said about my job?",
3
    user_name,
4
    thread_id,
5
)
6
7
print(r)