LangGraph Memory Example

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
    first_name: str
4
    last_name: str
5
    session_id: str
6
    user_name: str
7
8
9
@tool
10
async def search_facts(state: State, query: str, limit: int = 5) -> list[str]:
11
    """Search for facts in all conversations had with a user.
12
    
13
    Args:
14
        state (State): The Agent's state.
15
        query (str): The search query.
16
        limit (int): The number of results to return. Defaults to 5.
17
18
    Returns:
19
        list: A list of facts that match the search query.
20
    """
21
    edges = await zep.graph.search(
22
        user_id=state["user_name"], text=query, limit=limit, search_scope="edges"
23
    )
24
    return [edge.fact for edge in edges]
25
26
27
@tool
28
async def search_nodes(state: State, query: str, limit: int = 5) -> list[str]:
29
    """Search for nodes in all conversations had with a user.
30
    
31
    Args:
32
        state (State): The Agent's state.
33
        query (str): The search query.
34
        limit (int): The number of results to return. Defaults to 5.
35
36
    Returns:
37
        list: A list of node summaries for nodes that match the search query.
38
    """
39
    nodes = await zep.graph.search(
40
        user_id=state["user_name"], text=query, limit=limit, search_scope="nodes"
41
    )
42
    return [node.summary for node in nodes]
43
44
45
tools = [search_facts, search_nodes]
46
47
tool_node = ToolNode(tools)
48
49
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
4
    system_message = SystemMessage(
5
        content=f"""You are a compassionate mental health bot and caregiver. Review information about the user and their prior conversation below and respond accordingly.
6
        Keep responses empathetic and supportive. And remember, always prioritize the user's well-being and mental health.
7
8
        {memory.context}"""
9
    )
10
11
    messages = [system_message] + state["messages"]
12
13
    response = await llm.ainvoke(messages)
14
15
    # Add the new chat turn to the Zep graph
16
    messages_to_save = [
17
        Message(
18
            role_type="user",
19
            role=state["first_name"] + " " + state["last_name"],
20
            content=state["messages"][-1].content,
21
        ),
22
        Message(role_type="assistant", content=response.content),
23
    ]
24
25
    await zep.memory.add(
26
        session_id=state["session_id"],
27
        messages=messages_to_save,
28
    )
29
30
    # Truncate the chat history to keep the state from growing unbounded
31
    # In this example, we going to keep the state small for demonstration purposes
32
    # We'll use Zep's Facts to maintain conversation context
33
    state["messages"] = trim_messages(
34
        state["messages"],
35
        strategy="last",
36
        token_counter=len,
37
        max_tokens=3,
38
        start_on="human",
39
        end_on=("human", "tool"),
40
        include_system=True,
41
    )
42
43
    logger.info(f"Messages in state: {state['messages']}")
44
45
    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
18
graph_builder.add_node("agent", chatbot)
19
graph_builder.add_node("tools", tool_node)
20
21
graph_builder.add_edge(START, "agent")
22
23
graph_builder.add_conditional_edges("agent", should_continue, {"continue": "tools", "end": END})
24
25
graph_builder.add_edge("tools", "agent")
26
27
28
graph = graph_builder.compile(checkpointer=memory)

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

1
r = await graph_invoke(
2
    "Hi there?",
3
    first_name,
4
    last_name,
5
    thread_id,
6
)
7
8
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.
5
    """,
6
    first_name,
7
    last_name,
8
    thread_id,
9
)
10
11
print(r)

1
memory = await zep.memory.get(session_id=thread_id)
2
3
print(memory.context)

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

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

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