pip install langchain langgraph langchain-community langchain-google-genai faiss-cpu python-dotenv

GOOGLE_API_KEY=your_google_api_key

import os
import uuid
from typing import List
from dotenv import load_dotenv

from langchain_google_genai import (
    ChatGoogleGenerativeAI,
    GoogleGenerativeAIEmbeddings
)
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_core.messages import get_buffer_string
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, START, MessagesState, StateGraph
from langgraph.prebuilt import ToolNode

load_dotenv()
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash")
 embeddings = GoogleGenerativeAIEmbeddings(
    model="models/gemini-embedding-exp-03-07"
)

FAISS_INDEX_PATH = "faiss_index"

# Load existing index or create new one
if os.path.exists(FAISS_INDEX_PATH):
    vector_store = FAISS.load_local(
        FAISS_INDEX_PATH,
        embeddings,
        allow_dangerous_deserialization=True
    )
else:
    # Initialize with dummy document to avoid empty index
    vector_store = FAISS.from_texts(["init"], embeddings)

@tool
def save_recall_memory(memory: str, config: RunnableConfig) -> str:
    """
    Save a memory string to the vector store for future retrieval.
    
    Args:
        memory: The information to remember
        config: Runtime configuration including user_id
    
    Returns:
        The saved memory string
    """
    user_id = config["configurable"].get("user_id")
    doc = Document(
        page_content=memory,
        metadata={"user_id": user_id}
    )
    vector_store.add_documents([doc])
    vector_store.save_local(FAISS_INDEX_PATH)
    return memory

@tool
def search_recall_memories(
    query: str,
    config: RunnableConfig
) -> List[str]:
    """
    Search for relevant memories using semantic similarity.
    
    Args:
        query: Current conversation context
        config: Runtime configuration including user_id
    
    Returns:
        List of relevant memory strings
    """
    user_id = config["configurable"].get("user_id")
    results = vector_store.similarity_search(query, k=10)
    
    # Filter by user_id and return top 3
    filtered = [
        doc.page_content
        for doc in results
        if doc.metadata.get("user_id") == user_id
    ]
    return filtered[:3]

prompt = ChatPromptTemplate.from_messages([
    ("system", """
You are a helpful assistant with advanced long-term memory capabilities.
Powered by a stateless LLM, you must rely on external memory to store
information between conversations. Utilize the available memory tools
to store and retrieve important details.

Memory Usage Guidelines:
1. Actively use memory tools to build comprehensive user understanding
2. Make informed suppositions based on stored memories
3. Regularly reflect on past interactions to identify patterns
4. Update your mental model of the user with each new piece of info
5. Cross-reference new information with existing memories
6. Prioritize storing emotional context and personal values
7. Use memory to anticipate needs and tailor responses
8. Recognize changes in user's situation over time
9. Leverage memories to provide personalized examples
10. Recall past challenges or successes for problem-solving

## Recall Memories
Contextually retrieved based on current conversation:
{recall_memories}

## Instructions
Engage naturally as a trusted colleague or friend. Seamlessly
incorporate your understanding without explicitly mentioning memory
capabilities. Be attentive to subtle cues and emotions. Use tools
to persist information you want to retain. If you call tools, all
text preceding the tool call is internal. Respond AFTER calling
the tool, once you have confirmation it completed successfully.
"""),
    ("placeholder", "{messages}"),
])

# Define custom state with memory
class State(MessagesState):
    recall_memories: List[str]

# Node 1: Load relevant memories
def load_memories(state: State, config: RunnableConfig) -> State:
    convo = get_buffer_string(state["messages"])
    convo = convo[:2048]  # Truncate to reasonable length
    recall = search_recall_memories.invoke(convo, config)
    return {"recall_memories": recall}

# Node 2: Agent reasoning
def agent(state: State) -> State:
    recall_str = (
        "<recall_memory>\n"
        + "\n".join(state["recall_memories"])
        + "\n</recall_memory>"
    )
    response = (prompt | llm_with_tools).invoke({
        "messages": state["messages"],
        "recall_memories": recall_str,
    })
    return {"messages": [response]}

# Bind tools to LLM
tools = [save_recall_memory, search_recall_memories]
llm_with_tools = llm.bind_tools(tools)

# Build state graph
builder = StateGraph(State)
builder.add_node(load_memories)
builder.add_node(agent)
builder.add_node("tools", ToolNode(tools))

# Define edges
builder.add_edge(START, "load_memories")
builder.add_edge("load_memories", "agent")
builder.add_conditional_edges(
    "agent",
    lambda s: "tools" if s["messages"][-1].tool_calls else END,
    ["tools", END]
)
builder.add_edge("tools", "agent")

# Compile with checkpointer
graph = builder.compile(checkpointer=MemorySaver())

def terminal_chat():
    # Get or generate user_id
    user_id = input("Enter your user_id: ") or "1"
    config = {
        "configurable": {
            "user_id": user_id,
            "thread_id": str(uuid.uuid4())
        }
    }
    
    history = []
    print("\nYou can start chatting. Type 'exit' to quit.\n")
    
    while True:
        user_input = input("You: ")
        if user_input.lower() in ["exit", "quit"]:
            break
        
        # Add user message to history
        history.append(("user", user_input))
        
        # Stream agent response
        for chunk in graph.stream(
            {"messages": history},
            config=config
        ):
            if "agent" in chunk and "messages" in chunk["agent"]:
                last_msg = chunk["agent"]["messages"][-1]
                if last_msg.content:
                    print("AI:", last_msg.content)
                    history.append(("ai", last_msg.content))

if __name__ == "__main__":
    terminal_chat()