#332: Long-Term Memory in LangGraph

Last week we added short-term memory to our LangGraph application. That works great as long as we stay in the same session. But when we want to keep the memory around between sessions, we need a different approach.

In this post we create our hand-written approach for a long-term memory solution. Do this only to understand what is going on and not to use it in production. For that purpose, we can use pre-build solutions that we explore next week.

Persist the InMemoryStore

To keep the memory around, we can take the InMemoryStore of LangGraph and dump it into a JSON file. That way we can see what is inside the store and load it back into the memory when we run our script once more:

import json
import sys
import uuid
from pathlib import Path
from typing import TypedDict

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import START, END, StateGraph
from langgraph.store.base import BaseStore
from langgraph.store.memory import InMemoryStore

sys.stdout.reconfigure(encoding="utf-8")

MEMORY_FILE = Path("memory.json")
DEFAULT_USER_ID = "demo-user"


def load_store() -> InMemoryStore:
    """Build a fresh InMemoryStore and rehydrate it from memory.json if present."""
    store = InMemoryStore()
    if MEMORY_FILE.exists():
        for record in json.loads(MEMORY_FILE.read_text(encoding="utf-8")):
            store.put(tuple(record["namespace"]), record["key"], record["value"])
    return store


def dump_store(store: InMemoryStore) -> None:
    """Serialize every item in the store to memory.json."""
    records = [
        {"namespace": list(ns), "key": item.key, "value": item.value}
        for ns in store.list_namespaces()
        for item in store.search(ns)
    ]
    MEMORY_FILE.write_text(json.dumps(records, indent=2), encoding="utf-8")

Build the workflow

Our minimalistic workflow uses a State object and a greeter() method. The method gives us a customised greeting if it knows us and a generic one if we did not yet tell who we are:

class State(TypedDict):
    user_id: str
    greeting: str
    known: bool


def greeter(state: State, *, store: BaseStore) -> dict:
    record = store.get(("users", state["user_id"]), "profile")
    if record is not None:
        return {
            "greeting": f"Hi {record.value['name']}, welcome back!",
            "known": True,
        }
    return {"greeting": "Hello user, what is your name?", "known": False}


store = load_store()

workflow = StateGraph(State)

workflow.add_node("greeter", greeter)

workflow.add_edge(START, "greeter")
workflow.add_edge("greeter", END)

# Use the configured memory object
graph = workflow.compile(checkpointer=InMemorySaver(), store=store)

Interact with the workflow

This time we generate a unique thread_id for each session and use a user_id to keep track of the user. As in a real use case, our sessions will come and go, but the knowledge about our user should stay around. To achieve that, we must not forget to persist the store after we let the user answered our questions:

def main() -> None:
    user_id = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_USER_ID
    config = {"configurable": {"thread_id": str(uuid.uuid4())}}

    result = graph.invoke({"user_id": user_id}, config)
    print(result["greeting"])

    if not result["known"]:
        name = input("> ").strip()
        if name:
            store.put(("users", user_id), "profile", {"name": name})
            dump_store(store)
            print(f"Nice to meet you, {name}. Run me again to see me remember.")


if __name__ == "__main__":
    main()

We can run our script and pass our user id as a parameter:

$ python .\long_term_memory.py AA
Hello user, what is your name?
> Alice
Nice to meet you, Alice. Run me again to see me remember.

$ python .\long_term_memory.py AA
Hi Alice, welcome back!

$ python .\long_term_memory.py BB
Hello user, what is your name?
> Bob
Nice to meet you, Bob. Run me again to see me remember.

$ python .\long_term_memory.py BB
Hi Bob, welcome back!

What got persisted?

After our experiment with the two users, our memory.json file looks like this:

[
  {
    "namespace": [
      "users",
      "AA"
    ],
    "key": "profile",
    "value": {
      "name": "Alice"
    }
  },
  {
    "namespace": [
      "users",
      "BB"
    ],
    "key": "profile",
    "value": {
      "name": "Bob"
    }
  }
]

We can see the two names we specified and that are now tracked for the specific user id.

#332: Long-Term Memory in LangGraph

Persist the InMemoryStore

Build the workflow

Interact with the workflow

What got persisted?

Next