How can I setup memory for my Multi Agent System in Langgraph ?

[iamgauravpant]

Hello everyone , I'm working on a project that is built on Langgraph's multi-agent system ( Hierarchical architecture ) .

My multi-agent system is derived from here : https://langchain-ai.github.io/langgraph/tutorials/multi_agent/hierarchical_agent_teams

I want to add memory to the system now but haven't been successful so far . Every doc that I refer to is using memory on a single agent .

Has anyone tried setting up memory for multi agent system using Langgraph yet ? Can I get some references ?

Here is my Github repo link : https://github.com/MindzKonnectedAI/SkillSync/tree/gaurav

I've also attached my full code below :

app.py

import streamlit as st
from dotenv import load_dotenv
import os
from typing import Annotated, List
from langgraph.graph import END, StateGraph, START
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
import operator
from langchain_openai.chat_models import ChatOpenAI
from typing_extensions import TypedDict
import sql_agent_team_supervisor 
import github_team_supervisor

import csv_to_sql
import create_image_func
import create_team_supervisor_func

import os

# Load environment variables from .env file
load_dotenv()

# Access the environment variables
openai_api_key = os.getenv("OPENAI_API_KEY")
tavily_api_key = os.getenv("TAVILY_API_KEY")
langchain_api_key = os.getenv("LANGCHAIN_API_KEY")
langchain_tracking_v2 = os.getenv("LANGCHAIN_TRACING_V2")
langchain_endpoint = os.getenv("LANGCHAIN_ENDPOINT")
langchain_project = os.getenv("LANGCHAIN_PROJECT")

# Set environment variables if needed
os.environ["OPENAI_API_KEY"] = openai_api_key
os.environ["TAVILY_API_KEY"] = tavily_api_key
os.environ["LANGCHAIN_API_KEY"] = langchain_api_key
os.environ["LANGCHAIN_TRACING_V2"] = langchain_tracking_v2
os.environ["LANGCHAIN_ENDPOINT"] = langchain_endpoint
os.environ["LANGCHAIN_PROJECT"] = langchain_project

### Statefully manage chat history ###
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# Define the agent node function
def agent_node(state, agent, name):
    result = agent.invoke(state)
    return {"messages": [HumanMessage(content=result["messages"][-1].content, name=name)]}

# Initialize the agents
# research_chain = research_team_supervisor_agent.research_team_supervisor_agent(agent_node)
# authoring_chain = document_team_supervisor.document_team_supervisor(agent_node)
github_chain = github_team_supervisor.github_team_supervisor(agent_node)
sql_chain = sql_agent_team_supervisor.sql_agent_team_supervisor()
llm = ChatOpenAI(model="gpt-4o-mini")

# # Define the supervisor node
# supervisor_node = create_team_supervisor_func.create_team_supervisor_func(
#     llm,
#     "You are a supervisor tasked with managing a conversation between the following teams: {team_members}. "
#     "Given the following user request, respond with the worker to act next. "
#     "Carefully analyze the query and then respond with the worker to act next as for a user request , you can only call one worker "
#     "Each worker will perform a task and respond with their results and status. When finished, respond with FINISH.",
#     ["SqlTeam", "GithubTeam"],
# )

# Define the supervisor node
supervisor_node = create_team_supervisor_func.create_team_supervisor_func(
    llm,
    "You are a supervisor tasked with managing a conversation between the following teams: {team_members}. "
    "Given the following user request, respond with the worker to act next. For a user request , you can only call one worker , so carefully analyze the user request and then assign the worker. Each worker will perform a "
    "task and respond with their results and status. When finished, respond with FINISH."
    "When giving results back to the user , also mention the worker that was used.",
    ["SqlTeam", "GithubTeam"],
)

# Define the top-level state
class State(TypedDict):
    messages: Annotated[List[BaseMessage], operator.add]
    next: str

# Helper functions
def get_last_message(state: State) -> str:
    return state["messages"][-1].content

def join_graph(response: dict):
    return {"messages": [response["messages"][-1]]}

# Define the graph
super_graph = StateGraph(State)
# super_graph.add_node("ResearchTeam", get_last_message | research_chain | join_graph)
# super_graph.add_node("PaperWritingTeam", get_last_message | authoring_chain | join_graph)
super_graph.add_node("GithubTeam", get_last_message | github_chain | join_graph)
super_graph.add_node("SqlTeam", get_last_message | sql_chain | join_graph)
super_graph.add_node("super_supervisor", supervisor_node)

# super_graph.add_edge("ResearchTeam", "super_supervisor")
# super_graph.add_edge("PaperWritingTeam", "super_supervisor")
super_graph.add_edge("GithubTeam", "super_supervisor")
super_graph.add_edge("SqlTeam", "super_supervisor")
super_graph.add_conditional_edges(
    "super_supervisor",
    lambda x: x["next"],
    {
        "SqlTeam": "SqlTeam",
        "GithubTeam":"GithubTeam",
        "FINISH": END,
    },
)
super_graph.add_edge(START, "super_supervisor")
super_graph = super_graph.compile()

# Streamlit UI
st.title("Multi-Agent Supervisor System")

# Define folder paths
csv_folder = "csv"

# Ensure the csv and db directories exist
os.makedirs(csv_folder, exist_ok=True)

# File uploader widget
st.sidebar.title('File Upload and Processing')
uploaded_file = st.sidebar.file_uploader("Upload your CSV file", type=["csv"])
agent_name = st.sidebar.radio(
    "Search using",
    ["SQL", "Github"]
)

def get_agent_name(agent_name_here):
    if(agent_name_here=="SQL"):
        return "SQLTeam Agent"
    else:
        return "GithubTeam Agent"

if uploaded_file is not None:
    st.sidebar.write("Processing the uploaded file...")

    # Define the path to save the CSV in the csv folder
    csv_path = os.path.join(csv_folder, uploaded_file.name)

    # Save the uploaded file to the csv folder
    with open(csv_path, "wb") as file:
        file.write(uploaded_file.getbuffer())

    st.sidebar.success(f"CSV file saved successfully in {csv_folder} as {uploaded_file.name}")

    # Call the function to save CSV data into the database
    csv_to_sql.save_csv_to_sql(csv_path)

# Conversation History
for message in st.session_state.chat_history:
    if isinstance(message,HumanMessage):
        with st.chat_message("Human"):
            st.markdown(message.content)
    else:
        with st.chat_message("AI"):
            st.markdown(message.content)

# user_input = st.text_input("Enter your query:", "Total number of users in SQL Database?")
prompt = st.chat_input("Enter your query")
if prompt is not None and prompt != "":
    with st.chat_message("Human"):
        st.markdown(prompt)
    
    st.session_state.chat_history.append(HumanMessage(prompt))
    with st.chat_message("AI"):
        # Add a spinner with "Processing your query..." text
        # with st.spinner("Processing your query..."):
            # Generate the graph image and save it to the temporary file
        create_image_func.create_graph_image(super_graph, "super_graph")
        final_prompt = prompt +" using "+ get_agent_name(agent_name)
        print("the final prompt to go to supervisor :",final_prompt)
        res = super_graph.invoke(input={"messages": [HumanMessage(content=final_prompt)]})
        print("AI response :",res["messages"])
        aiRes = res["messages"][-1].content
        st.write(aiRes)            
        st.session_state.chat_history.append(AIMessage(aiRes))

sql_agent_team_supervisor.py

from dotenv import load_dotenv
import os
import create_image_func
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_community.utilities import SQLDatabase
from typing import Any
from langchain_core.messages import ToolMessage
from langchain_core.runnables import RunnableLambda, RunnableWithFallbacks
from langgraph.prebuilt import ToolNode
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langchain_core.utils.function_calling import convert_to_openai_function
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from typing import Annotated, Literal
from langchain_core.messages import AIMessage
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_openai import ChatOpenAI
from typing_extensions import TypedDict
from langgraph.graph import END, StateGraph, START
from langgraph.graph.message import AnyMessage, add_messages
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_community.utilities import SQLDatabase
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.messages import ToolMessage
from langchain_core.runnables import RunnableLambda
from langgraph.graph import END, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition

# Load environment variables from .env file
load_dotenv()

# Access the variables
openai_api_key = os.getenv("OPENAI_API_KEY")
langchain_api_key = os.getenv("LANGCHAIN_API_KEY")
langchain_tracking_v2 = os.getenv("LANGCHAIN_TRACING_V2")
langchain_endpoint = os.getenv("LANGCHAIN_ENDPOINT")
langchain_project = os.getenv("LANGCHAIN_PROJECT1")

# Set them as environment variables if needed
os.environ["OPENAI_API_KEY"] = openai_api_key
os.environ["LANGCHAIN_API_KEY"] = langchain_api_key
os.environ["LANGCHAIN_TRACING_V2"] = langchain_tracking_v2
os.environ["LANGCHAIN_ENDPOINT"] = langchain_endpoint
os.environ["LANGCHAIN_PROJECT"] = langchain_project

db_folder = "./db"

# Define the path to the SQLite database file
db_path = os.path.join(db_folder, "employee.db")

db = SQLDatabase.from_uri(f"sqlite:///{db_path}")

print(db.dialect)
print(db.get_usable_table_names())

llm=ChatOpenAI(model="gpt-4o-mini")
# SQL toolkit
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
tools = toolkit.get_tools()

# Query checking
query_check_system = """You are a SQL expert with a strong attention to detail.
Double check the SQLite query for common mistakes, including:
- Using NOT IN with NULL values
- Using UNION when UNION ALL should have been used
- Using BETWEEN for exclusive ranges
- Data type mismatch in predicates
- Properly quoting identifiers
- Using the correct number of arguments for functions
- Casting to the correct data type
- Using the proper columns for joins

If there are any of the above mistakes, rewrite the query. If there are no mistakes, just reproduce the original query.

Execute the correct query with the appropriate tool."""
query_check_prompt = ChatPromptTemplate.from_messages([("system", query_check_system),("user", "{query}")])
query_check = query_check_prompt | llm

@tool
def check_query_tool(query: str) -> str:
    """
    Use this tool to double check if your query is correct before executing it.
    """
    return query_check.invoke({"query": query}).content

# Query result checking
query_result_check_system = """You are grading the result of a SQL query from a DB.
- Check that the result is not empty.
- If it is empty, instruct the system to re-try!"""
query_result_check_prompt = ChatPromptTemplate.from_messages([("system", query_result_check_system),("user", "{query_result}")])
query_result_check = query_result_check_prompt | llm

@tool
def check_result(query_result: str) -> str:
    """
    Use this tool to check the query result from the database to confirm it is not empty and is relevant.
    """
    return query_result_check.invoke({"query_result": query_result}).content

tools.append(check_query_tool)
tools.append(check_result)

# Define the state for the agent
class State(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]
    team_members: str
    next: str

# Assistant
class Assistant:

    def __init__(self, runnable: Runnable):
        self.runnable = runnable

    def __call__(self, state: State, config: RunnableConfig):
        while True:
            # Append to state
            state = {**state}
            # Invoke the tool-calling LLM
            result = self.runnable.invoke(state)
            # If it is a tool call -> response is valid
            # If it has meaninful text -> response is valid
            # Otherwise, we re-prompt it b/c response is not meaninful
            if not result.tool_calls and (
                not result.content
                or isinstance(result.content, list)
                and not result.content[0].get("text")
            ):
                messages = state["messages"] + [("user", "Respond with a real output.")]
                state = {**state, "messages": messages}
            else:
                break
        return {"messages": result}

# Assistant runnable
query_gen_system = """
ROLE:
You are an agent designed to interact with a SQL database. You have access to tools for interacting with the database.
GOAL:
Given an input question, create a syntactically correct SQLite query to run, then look at the results of the query and return the answer.
INSTRUCTIONS:
- Only use the below tools for the following operations.
- Only use the information returned by the below tools to construct your final answer.
- To start you should ALWAYS look at the tables in the database to see what you can query. Do NOT skip this step.
- Then you should query the schema of the most relevant tables.
- Write your query based upon the schema of the tables. You MUST double check your query before executing it.
- You can order the results by a relevant column to return the most interesting examples in the database.
- Never query for all the columns from a specific table, only ask for the relevant columns given the question.
- ALWAYS prefer to give response in table format when returning the it to the user
- If you get an error while executing a query, rewrite the query and try again.
- If the query returns a result, use check_result tool to check the query result.
- If the query result result is empty, think about the table schema, rewrite the query, and try again.
- DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database."""

query_gen_prompt = ChatPromptTemplate.from_messages([("system", query_gen_system),("placeholder", "{messages}")])
assistant_runnable = query_gen_prompt | llm.bind_tools(tools)

def create_tool_node_with_fallback(tools: list) -> dict:
    return ToolNode(tools).with_fallbacks(
        [RunnableLambda(handle_tool_error)], exception_key="error"
    )

def handle_tool_error(state) -> dict:
    error = state.get("error")
    tool_calls = state["messages"][-1].tool_calls
    return {
        "messages": [
            ToolMessage(
                content=f"Error: {repr(error)}\n please fix your mistakes.",
                tool_call_id=tc["id"],
            )
            for tc in tool_calls
        ]
    }

def enter_chain(message: str):
    print("message", message)
    results = {
        "messages": [HumanMessage(content=message)],
    }
    return results

def sql_agent_team_supervisor() -> str:

    # Graph
    builder = StateGraph(State)

    # Define nodes: these do the work
    builder.add_node("assistant", Assistant(assistant_runnable))
    builder.add_node("tools", create_tool_node_with_fallback(tools))

    # Define edges: these determine how the control flow moves
    builder.set_entry_point("assistant")
    builder.add_conditional_edges(
        "assistant",
        # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools
        # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END
        tools_condition,
        # "tools" calls one of our tools. END causes the graph to terminate (and respond to the user)
        {"tools": "tools", END: END},
    )
    builder.add_edge("tools", "assistant")
    chain = builder.compile()
    sql_chain = enter_chain | chain
    create_image_func.create_graph_image(chain, "sql_graph_image3")
    
    return sql_chain

github_team_supervisor.py

import github_agent
from typing import Annotated, List
from langchain_core.messages import BaseMessage, HumanMessage
import operator
from typing_extensions import TypedDict
from langchain_openai.chat_models import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from create_team_supervisor_func
from langgraph.graph import END, StateGraph, START
import create_image_func
from typing import Annotated, Literal

# GithubTeamState graph state
class GithubTeamState(TypedDict):
    # A message is added after each team member finishes
    messages: Annotated[List[BaseMessage], operator.add]
    # The team members are tracked so they are aware of
    # the others' skill-sets
    team_members: List[str]
    # Used to route work. The supervisor calls a function
    # that will update this every time it makes a decision
    next: str


# The following functions interoperate between the top level graph state
# and the state of the research sub-graph
# this makes it so that the states of each graph don't get intermixed
def enter_chain(message: str):
    results = {
        "messages": [HumanMessage(content=message)],
    }
    return results

# Define a conditional edge to decide whether to continue or end the workflow
def condition(
    state: GithubTeamState,
) -> Literal["supervisor", "query_param_generator"]:
    messages = state["messages"]
    print("messages inside condition conditional edge :",messages)
    last_message = messages[-1]
    print("last_message here in conditional edge ", last_message)
    if last_message.content.startswith("Error:"):
        return "query_param_generator"
    else:
        return "supervisor"

def github_team_supervisor(agent_node)-> str:

    llm = ChatOpenAI(model="gpt-4o-mini")

    query_param_node = github_agent.query_param_generator_node(agent_node)

    fetch_node = github_agent.fetch_users_node(agent_node)

    def supervisor_agent(state):

        github_supervisor_agent = create_team_supervisor_func(
            llm,
            "You are a supervisor tasked with managing a conversation between the"
            " following workers:  query_param_generator, fetch_users. Given the following user request,"
            " always respond with the query_param_generator worker as it will be called first and take its output as parameter to fetch_users worker"
            " After calling fetch_users , respond with FINISH.",
            ["query_param_generator", "fetch_users"],
        )
        return github_supervisor_agent

    github_graph = StateGraph(GithubTeamState)
    github_graph.add_node("query_param_generator", query_param_node)
    github_graph.add_node("fetch_users", fetch_node)
    github_graph.add_node("supervisor", supervisor_agent)

    # Define the control flow
    github_graph.add_edge(START,"query_param_generator")
    github_graph.add_edge("query_param_generator", "fetch_users")
    github_graph.add_conditional_edges(
        "fetch_users",
        condition
    )
    github_graph.add_edge("supervisor",END)
    # github_graph.add_edge(START, "supervisor")
    chain = github_graph.compile()

    github_chain = enter_chain | chain
    create_image_func.create_graph_image(chain, "github_graph_image")
    return github_chain

github_agent.py

from langchain_core.tools import tool
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
import requests
import functools
from langgraph.prebuilt import create_react_agent
from langchain_core.messages import ToolMessage


@tool("query_param_generator")
def query_param_generator(query:str):
    """
    Takes a natural language query as input and returns the appropriate query parameters based on the rules defined in system_message
    """
    pass

@tool("fetch_users")
def fetch_users(query_param: str, per_page: int = 10):
    """Returns a list of Github users based on query parameters and results per page."""

    try:
        # Ensure the per_page value is within the allowed range
        if per_page > 100:
            per_page = 100
        elif per_page < 1:
            per_page = 1

        # Log query parameters
        print("query_param :", query_param)
        print("per_page :", per_page)
        print("the url :",f"https://api.github.com/search/users?q={query_param}&per_page={per_page}")
        
        # Make the API request with query parameters and per_page limit
        res = requests.get(
            f"https://api.github.com/search/users?q={query_param}&per_page={per_page}"
        )
        print("fetch_users github user search api response :",res.json())
        print("status code of res :",res.status_code)
        return res.json()
    except:
        print("inside exception !!!c")
        return {
            "messages": [
                ToolMessage(
                    content=f"Error: please fix your mistakes.",
                    tool_call_id=["id"],
                )
            ]
        }

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

def query_param_generator_node(agent_node):
    # give prompt to agent
    query_gen_system = """
    Strictly follow all the rules below to generate query parameters for GitHub User Search . Ensure all rules are followed to generate accurate search queries.
    Only return the query , no extra information .
    Prefix every query you generate with 'type:user' 
    # GitHub User Search - Query Parameter Generation Rules

    The following guidelines define how to construct query parameters for the GitHub User Search. Ensure all rules are followed to generate accurate search queries.

    # Overview

    1. Search Scope: Applies to public personal GitHub accounts (not organizations).
    2. Query Components: Queries can include:
    3. Keywords: For general information like usernames, names, emails, and bios.
    4. Qualifiers: For searching specific fields.
    5. Sort Parameters: Optional but can be added for ordering results.
    6. Case Sensitivity: Keywords are case-insensitive.
    7. Result Limit: The search returns the first 1000 results, sorted by best match (by default).

    # Qualifiers & Usage

    Each qualifier targets a specific field in the GitHub user data. These cannot be mixed with regular keywords.

    user:NAME: Matches exact usernames.
    Example: user:braingain

    in:login: Searches within usernames (non-exact matches allowed).
    Example: braingain in:login

    in:email: Searches within users' email addresses.
    Example: irina in:email

    in:name: Searches within users' full names.
    Example: Irina in:name

    fullname:NAME: Similar to in:name, searches users' full names.
    Example: fullname:john smith

    location:NAME: Searches users based on location.
    Example: location:Boston

    language:NAME: Finds users based on the primary language of their public repositories.
    Example: language:python

    repos:n: Searches users by the number of public repositories.
    Example: repos:>1000

    followers:n: Searches users by the number of followers.
    Example: followers:>1000

    created:DATE: Finds users by their GitHub account creation date.
    Example: created:>2020-01-01

    is:sponsorable: Finds users with a GitHub Sponsors profile.
    Example: is:sponsorable

    sort:: Sorts users based on specific attributes.
    Example: repos:>10000 sort:followers

    # Boolean Operators
    You can combine keywords and qualifiers using Boolean operators to refine the search. Follow these rules:

    AND (implied): Combining two different qualifiers or a qualifier and a keyword automatically implies AND.
    Example: location:"San Francisco" language:python (Finds users in San Francisco who primarily use Python).

    OR: Explicitly use OR between keywords or in: qualifiers only. For other qualifiers, using the same qualifier twice implies OR.
    Example: "front-end developer" OR "ui developer"
    Example: location:"new jersey" location:"new york" (Finds users in either New Jersey or New York).

    NOT (-): Use the minus sign (-) to exclude certain terms or qualifiers.
    Example: location:iceland -location:Reykjavik (Finds users in Iceland but not Reykjavik).

    # Key Limitations & Constraints

    Character Limit: Queries must not exceed 256 characters.

    No Parentheses: Do not use parentheses in queries.

    AND/OR/NOT Limits: You cannot use more than five AND, OR, or NOT operators in a single query.
    For example: location:"silicon valley" -language:java -language:c++ -language:python -language:javascript -language:html is valid (5 negations).
    Special Notes on Combining Operators

    AND is implied for combining qualifiers and keywords but cannot be used explicitly with certain fields like location, language, etc.

    OR cannot be used explicitly between different qualifiers.
    Example: fullname:irina user:braingain is interpreted as AND, while fullname:irina OR user:braingain is invalid.

    You cannot combine keywords and qualifiers in OR statements.
    Example: language:java OR "java developer" is invalid, while language:java "java developer" is interpreted as AND.
    
    # Output format

    Always return ONLY the query parameters you generated WITHOUT any extra text . 
    Below are examples of some output that will be considered correct : 

    1. type:user AI Engineer location:India created:<2019-01-01 repos:>50
    2. type:user Blockchain developer location:Ukraine sort:followers
    3. type:user Full Stack Developer location:USA created:>2020-01-01 
    4. type:user UX Engineer location:Delhi 
    """
    query_param_generator_agent = create_react_agent(llm,tools=[query_param_generator],state_modifier=query_gen_system)
    query_param_node = functools.partial(agent_node, agent=query_param_generator_agent, name="query_param_generator")
    return query_param_node


def fetch_users_node(agent_node):

    fetch_users_agent = create_react_agent(llm, tools=[fetch_users])
    fetch_node = functools.partial(agent_node, agent=fetch_users_agent, name="fetch_users")
    return fetch_node

csv_to_sql.py

import streamlit as st
import pandas as pd
from sqlalchemy import create_engine
import os

# Define folder paths
csv_folder = "csv"
db_folder = "db"

# Ensure the csv and db directories exist
os.makedirs(csv_folder, exist_ok=True)
os.makedirs(db_folder, exist_ok=True)

def save_csv_to_sql(csv_path):
    """
    Saves the content of a CSV file into an SQLite database.

    Args:
        csv_path (str): Path to the CSV file to be saved into the database.
        db_folder (str): Folder where the SQLite database will be stored.
    """
    # Construct the database path inside the db folder
    db_path = os.path.join(db_folder, "employee.db")

    # Load DataFrame from CSV
    try:
        df = pd.read_csv(csv_path, encoding='utf-8')
    except UnicodeDecodeError:
        df = pd.read_csv(csv_path, encoding='latin1')

    # Create SQLite engine
    engine = create_engine(f"sqlite:///{db_path}")

    # Load DataFrame into SQLite database
    df.to_sql("employee", engine, if_exists='replace', index=False)

    st.sidebar.success(f"Database created and data loaded successfully into {db_path}")

create_team_supervisor_func.py

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser
from langchain_openai.chat_models import ChatOpenAI

def create_team_supervisor_func(llm: ChatOpenAI, system_prompt, members) -> str:
    """An LLM-based router."""
    # print("state", state)
    options = ["FINISH"] + members
    function_def = {
        "name": "route",
        "description": "Select the next role.",
        "parameters": {
            "title": "routeSchema",
            "type": "object",
            "properties": {
                "next": {
                    "title": "Next",
                    "anyOf": [
                        {"enum": options},
                    ],
                },
            },
            "required": ["next"],
        },
    }
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system_prompt),
            MessagesPlaceholder(variable_name="messages"),
            (
                "system",
                "Given the conversation above, who should act next?"
                " Or should we FINISH? Select one of: {options}",
            ),
        ]
    ).partial(options=str(options), team_members=", ".join(members))
    return (
        prompt
        | llm.bind_functions(functions=[function_def], function_call="route")
        | JsonOutputFunctionsParser()
    )
    ```
    
    
 create_image_func.py
------------------------

from IPython.display import Image, display
import os

def create_graph_image(chain, image_name="graph_image", folder_path="images"):
"""
Saves and displays an image generated by the chain's graph in a specified folder.

Args:
    chain: The chain object with the graph generation capability.
    image_name (str): The base name for the image file without extension.
                      Default is 'graph_image'.
    folder_path (str): The directory path where the image should be saved.
                       Default is 'images'.
"""
# Ensure the folder exists
if not os.path.exists(folder_path):
    os.makedirs(folder_path)

# Create a dynamic file path with the specified folder and name
image_path = os.path.join(folder_path, f"{image_name}.png")

try:
    # Assuming `draw_mermaid_png()` returns binary image data
    graph_image_data = chain.get_graph(xray=True).draw_mermaid_png()

    # Save the image data to a file
    with open(image_path, "wb") as file:
        file.write(graph_image_data)

    # Display the saved image
    display(Image(filename=image_path))

    print(f"Image saved and displayed: {image_path}")

except AttributeError as e:
    print(f"Error: The chain object does not have the required methods. {e}")

except Exception as e:
    print(f"An error occurred: {e}")

[Arokamals]

You can add memory while you create your graph.
Here's how you can do that:

from langgraph.checkpoint.memory import MemorySaver

memory = MemorySaver()

Next define the graph.

from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition

class State(TypedDict):
messages: Annotated[list, add_messages]

graph_builder = StateGraph(State)
.
.
.
.
Finally, compile the graph with the provided checkpointer.

graph = graph_builder.compile(checkpointer=memory)

[iamgauravpant]

Hi @Arokamals , this way of adding memory to hierarchical agent system does not work as intended .

This method only works for single agent systems .

Have you tried adding memory to multi-agent systems ?

[shaibt]

I have this working by providing a memory instance as the parent graph checkpointer when it is compiled (like the example provided by @Arokamals) and then using the same memory instance in all subgraphs as well.
pay attention that create_react_agent also accepts a checkpointer param so even when creating a prebuilt subgraph/subagent you can provide the memory object.

[MitchellHarry16]

just a @Arokamals try to add a adding memory to multi-agent systems ?