chat_optimized_memory.py

from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.store.memory import InMemoryStore
from langmem import create_manage_memory_tool, create_search_memory_tool
from langgraph.config import get_store

checkpointer = InMemorySaver()
store = InMemoryStore(
    index={
        "dims": 1536,
        "embed": "openai:text-embedding-3-small"
    }
)

namespace = ("agent_memories",)
memory_tools = [
    create_manage_memory_tool(namespace),
    create_search_memory_tool(namespace)
]
def prompt(state):
    # Search over memories based on the messages
    store_obj = get_store()
    items = store_obj.search(namespace, query=state["messages"][-1].content)
    print("Items")
    print(items)
    memories = "\n\n".join(str(item) for item in items)
    system_msg = {"role": "system", "content": f"## Memories:\n\n{memories}"}
    return [system_msg] + state["messages"]

agentObj = create_react_agent("openai:gpt-4o", prompt=prompt, tools=memory_tools, store=store,
                              checkpointer=checkpointer)


def chat(agent, txt, thread_id):
    result_state = agent.invoke({"messages": [{"role": "user", "content": txt}]},
                                config={"configurable": {"thread_id": thread_id}})
    return result_state["messages"][-1].content