LangGraph Memory Example

1
import os
2
import uuid
3
import logging
4
from typing import Annotated, TypedDict
5
6
from zep_cloud.client import AsyncZep
7
from zep_cloud import Message
8
9
from langchain_core.messages import AIMessage, SystemMessage, HumanMessage, trim_messages
10
from langchain_core.tools import tool
11
from langchain_openai import ChatOpenAI
12
from langgraph.checkpoint.memory import MemorySaver
13
from langgraph.graph import END, START, StateGraph, add_messages
14
from langgraph.prebuilt import ToolNode
15
16
# Optional: Load environment variables from .env file
17
# from dotenv import load_dotenv
18
# load_dotenv()
19
20
# Set up logging
21
logging.basicConfig(level=logging.INFO)
22
logger = logging.getLogger(__name__)
23
24
# Initialize Zep client
25
zep = AsyncZep(api_key=os.environ.get('ZEP_API_KEY'))

1
class State(TypedDict):
2
    messages: Annotated[list, add_messages]
3
    first_name: str
4
    last_name: str
5
    thread_id: str
6
    user_name: str

1
def create_zep_tools(user_name: str):
2
    """Create Zep search tools configured for a specific user."""
3
    
4
    @tool
5
    async def search_facts(query: str, limit: int = 5) -> list[str]:
6
        """Search for facts in all conversations had with a user.
7
        
8
        Args:
9
            query (str): The search query.
10
            limit (int): The number of results to return. Defaults to 5.
11
12
        Returns:
13
            list: A list of facts that match the search query.
14
        """
15
        result = await zep.graph.search(
16
            user_id=user_name, query=query, limit=limit, scope="edges"
17
        )
18
        facts = [edge.fact for edge in result.edges or []]
19
        if not facts:
20
            return ["No facts found for the query."]
21
        return facts
22
23
    @tool
24
    async def search_nodes(query: str, limit: int = 5) -> list[str]:
25
        """Search for nodes in all conversations had with a user.
26
        
27
        Args:
28
            query (str): The search query.
29
            limit (int): The number of results to return. Defaults to 5.
30
31
        Returns:
32
            list: A list of node summaries for nodes that match the search query.
33
        """
34
        result = await zep.graph.search(
35
            user_id=user_name, query=query, limit=limit, scope="nodes"
36
        )
37
        summaries = [node.summary for node in result.nodes or []]
38
        if not summaries:
39
            return ["No nodes found for the query."]
40
        return summaries
41
    
42
    return [search_facts, search_nodes]
43
44
# We'll create the actual tools after we have a user_name
45
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

1
async def chatbot(state: State):
2
    memory = await zep.thread.get_user_context(state["thread_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="user",
19
            name=state["first_name"] + " " + state["last_name"],
20
            content=state["messages"][-1].content,
21
        ),
22
        Message(role="assistant", content=response.content),
23
    ]
24
25
    await zep.thread.add_messages(
26
        thread_id=state["thread_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
def create_agent(user_name: str):
2
    """Create a LangGraph agent configured for a specific user."""
3
    
4
    # Create tools configured for this user
5
    tools = create_zep_tools(user_name)
6
    tool_node = ToolNode(tools)
7
    llm_with_tools = ChatOpenAI(model="gpt-4o-mini", temperature=0).bind_tools(tools)
8
    
9
    # Update the chatbot function to use the configured LLM
10
    async def chatbot_with_tools(state: State):
11
        memory = await zep.thread.get_user_context(state["thread_id"])
12
13
        system_message = SystemMessage(
14
            content=f"""You are a compassionate mental health bot and caregiver. Review information about the user and their prior conversation below and respond accordingly.
15
            Keep responses empathetic and supportive. And remember, always prioritize the user's well-being and mental health.
16
17
            {memory.context}"""
18
        )
19
20
        messages = [system_message] + state["messages"]
21
22
        response = await llm_with_tools.ainvoke(messages)
23
24
        # Add the new chat turn to the Zep graph
25
        messages_to_save = [
26
            Message(
27
                role="user",
28
                name=state["first_name"] + " " + state["last_name"],
29
                content=state["messages"][-1].content,
30
            ),
31
            Message(role="assistant", content=response.content),
32
        ]
33
34
        await zep.thread.add_messages(
35
            thread_id=state["thread_id"],
36
            messages=messages_to_save,
37
        )
38
39
        # Truncate the chat history to keep the state from growing unbounded
40
        state["messages"] = trim_messages(
41
            state["messages"],
42
            strategy="last",
43
            token_counter=len,
44
            max_tokens=3,
45
            start_on="human",
46
            end_on=("human", "tool"),
47
            include_system=True,
48
        )
49
50
        logger.info(f"Messages in state: {state['messages']}")
51
52
        return {"messages": [response]}
53
    
54
    # Define the function that determines whether to continue or not
55
    async def should_continue(state, config):
56
        messages = state["messages"]
57
        last_message = messages[-1]
58
        # If there is no function call, then we finish
59
        if not last_message.tool_calls:
60
            return "end"
61
        # Otherwise if there is, we continue
62
        else:
63
            return "continue"
64
    
65
    # Build the graph
66
    graph_builder = StateGraph(State)
67
    memory = MemorySaver()
68
    
69
    graph_builder.add_node("agent", chatbot_with_tools)
70
    graph_builder.add_node("tools", tool_node)
71
    
72
    graph_builder.add_edge(START, "agent")
73
    graph_builder.add_conditional_edges("agent", should_continue, {"continue": "tools", "end": END})
74
    graph_builder.add_edge("tools", "agent")
75
    
76
    return 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
# Create user and thread in Zep
7
await zep.user.add(user_id=user_name, first_name=first_name, last_name=last_name)
8
await zep.thread.create(thread_id=thread_id, user_id=user_name)
9
10
# Create the agent configured for this user
11
graph = create_agent(user_name)
12
13
14
def extract_messages(result, user_name):
15
    output = ""
16
    for message in result["messages"]:
17
        if isinstance(message, AIMessage):
18
            name = "assistant"
19
        else:
20
            name = user_name
21
        output += f"{name}: {message.content}\n"
22
    return output.strip()
23
24
25
async def graph_invoke(
26
    message: str,
27
    first_name: str,
28
    last_name: str,
29
    user_name: str,
30
    thread_id: str,
31
    ai_response_only: bool = True,
32
):
33
    r = await graph.ainvoke(
34
        {
35
            "messages": [HumanMessage(content=message)],
36
            "first_name": first_name,
37
            "last_name": last_name,
38
            "thread_id": thread_id,
39
            "user_name": user_name,
40
        },
41
        config={"configurable": {"thread_id": thread_id}},
42
    )
43
44
    if ai_response_only:
45
        return r["messages"][-1].content
46
    else:
47
        return extract_messages(r, user_name)

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

1
memory = await zep.thread.get_user_context(thread_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
    user_name,
6
    thread_id,
7
)
8
9
print(r)

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

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