LangGraph: Multi-agent Network

 

https://www.pexels.com/photo/blonde-model-against-binary-code-4389463/
https://www.pexels.com/photo/blonde-model-against-binary-code-4389463/

This is the last series of LangGraph multi-agent blogs. We started with

  1. LangGraph - tools where we use conditional edges for routing messages to different tool nodes, then
  2. LangGraph: The Supervisor Agent where we leverage a supervisor node to route messages to different nodes.
Here, we have the multi-agent network where messages are routed to different nodes to complete a task. We mocked a banking scenario.

Get my investment and saving account balance. And sum them up.

This involves getting the bank account ID, looking up the investment and saving account balance, and sum them. We have the network looks like this
  1. Customer Agent is able to look up bank account ID
  2. Investment Agent is able to retrieve the investment account balance
  3. Saving Agent is able to retrieve the saving account balance
Let's take a look at the tools. [code]. Pretty straightforward and self-explanatory.
import logging
import random

from langchain_core.tools import tool

from langgraph_network.hosting import container

logger = container[logging.Logger]


@tool
def get_bank_account_id() -> str:
    """Get the bank account ID."""
    id = random.choice(["123", "456", "789"])
    logger.info(f"[TOOL] Bank account ID: {id}")
    return id


@tool
def get_investment_account_balance(bank_account_id: str) -> float:
    """Get the investment balance of an account.

    :param bank_account_id: The account ID.
    :return: The investment balance.
    """
    balance = {
        "123": 100000.10,
        "456": 200000.20,
        "789": 300000.30,
    }[bank_account_id]
    logger.info(f"[TOOL] Investment balance: {balance}")
    return balance


@tool
def get_saving_account_balance(bank_account_id: str) -> float:
    """Get the saving balance of an account.

    :param bank_account_id: The account ID.
    :return: The saving balance.
    """
    balance = {
        "123": 10000.10,
        "456": 20000.20,
        "789": 30000.30,
    }[bank_account_id]
    logger.info(f"[TOOL] Saving balance: {balance}")
    return balance
Next, let's take a look at the agent nodes [code]. They are also self-explanatory. We use the create_react_agent helper function to create these executable agents. These agents used the tools that are mentioned above.
from langgraph.prebuilt import create_react_agent

from langgraph_network.hosting import container
from langgraph_network.protocols.i_azure_openai_service import (
    IAzureOpenAIService,
)
from langgraph_network.tools.system_prompt import make_system_prompt
from langgraph_network.tools.tools import (
    get_bank_account_id,
    get_investment_account_balance,
    get_saving_account_balance,
)

llm_model = container[IAzureOpenAIService].get_model()

customer_agent = create_react_agent(
    llm_model,
    tools=[
        get_bank_account_id,
    ],
    state_modifier=make_system_prompt("You can provide bank account ID."),
)

investment_account_agent = create_react_agent(
    llm_model,
    [get_investment_account_balance],
    state_modifier=make_system_prompt("You can provide investment account balance."),
)

saving_account_agent = create_react_agent(
    llm_model,
    [get_saving_account_balance],
    state_modifier=make_system_prompt("You can provide saving account balance."),
)
Lastly, we pull everything together. [code
  1. The constructor creates the graph.
  2. In the graph, we have 
    1. an edge from START node to customer agent node.
    2. We have customer agent node, investment account agent node and saving account agent node.
  3. Message will be routed to customer agent node initially because we need it to get the bank account id, then message is routed to investment account agent node and then saving account agent node.
import asyncio
import logging

from langchain_core.messages import BaseMessage, HumanMessage
from langgraph.graph import END, START, MessagesState, StateGraph
from langgraph.graph.graph import CompiledGraph
from langgraph.graph.state import CompiledStateGraph
from langgraph.types import Command

from langgraph_network.hosting import container
from langgraph_network.protocols.i_azure_openai_service import (
    IAzureOpenAIService,
)
from langgraph_network.tools.agents import (
    customer_agent,
    investment_account_agent,
    saving_account_agent,
)

CUSTOMER_AGENT = "customer_agent"
INVESTMENT_ACCOUNT_ASSISTANT = "investment_account_assistant"
SAVING_ACCOUNT_ASSISTANT = "saving_account_assistant"


def get_next_node(last_message: BaseMessage, goto: str):
    if "FINAL ANSWER" in last_message.content:
        # Any agent decided the work is done
        return END
    return goto


class Graph:
    def __init__(self):
        self.graph = self.create_graph()
        self.logger = container[logging.Logger]

    def create_node(self, node_name: str, agent: CompiledGraph, next_node_name: str):
        """
        Create a node for the graph.

        :param node_name: The name of the node.
        :param agent: The agent to invoke.
        :param next_node_name: The name of the next node.
        """

        def wrapper(
            state: MessagesState,
        ) -> Command:  # type: ignore
            result = agent.invoke(state)
            goto = get_next_node(result["messages"][-1], next_node_name)
            result["messages"][-1] = HumanMessage(
                content=result["messages"][-1].content, name=node_name
            )
            return Command(
                update={"messages": result["messages"]},
                goto=goto,
            )

        return wrapper

    def create_graph(self) -> CompiledStateGraph:
        """
        Create the graph for the workflow.

        Create the nodes and edge for the workflow.
        """
        workflow = StateGraph(MessagesState)

        workflow.add_node(
            CUSTOMER_AGENT,
            self.create_node(
                CUSTOMER_AGENT, customer_agent, INVESTMENT_ACCOUNT_ASSISTANT
            ),
        )
        workflow.add_node(
            INVESTMENT_ACCOUNT_ASSISTANT,
            self.create_node(
                INVESTMENT_ACCOUNT_ASSISTANT,
                investment_account_agent,
                SAVING_ACCOUNT_ASSISTANT,
            ),
        )
        workflow.add_node(
            SAVING_ACCOUNT_ASSISTANT,
            self.create_node(
                SAVING_ACCOUNT_ASSISTANT, saving_account_agent, CUSTOMER_AGENT
            ),
        )

        # above a create 3 nodes and here we create the edge
        workflow.add_edge(START, CUSTOMER_AGENT)
        return workflow.compile()

    async def astream(self):
        async for state in self.graph.astream(
            {
                "messages": [
                    HumanMessage(
                        content="Get my investment and saving account balance. "
                        "And sum them up.",
                    )
                ]
            },
            {"recursion_limit": 10},
        ):
            self.logger.info(f"[STREAMING]: {next(iter(state.keys()))}")
            for message in state.values():
                for msg in message.values():
                    for m in msg:
                        if m.content:
                            self.logger.info(f"[STREAMING]:     {m.type}: {m.content}")


if __name__ == "__main__":
    graph = Graph()
    asyncio.run(graph.astream())
This is the traces. We can see that the agents are erroring out because it does not have the tools to perform certain tasks and route the message to the next node.
> customer_agent
>     human: Get my investment and saving account balance. And sum them up.
>     tool: 123
>     ai: I have obtained your bank account ID as 123. I will now proceed to
             get the balances of your investment and savings accounts.
>     tool: Error: get_investment_account_balance is not a valid tool, try one
            of [get_bank_account_id].
>     tool: Error: get_savings_account_balance is not a valid tool, try one of
            [get_bank_account_id].
>     human: It seems that I don't have the necessary tools to retrieve your
                investment and savings account balances. Another assistant with
                the appropriate tools will assist you further.
>
> investment_account_assistant
>     human: Get my investment and saving account balance. And sum them up.
>     tool: 123
>     ai: I have obtained your bank account ID as 123. I will now proceed to
             get the balances of your investment and savings accounts.
>     tool: Error: get_investment_account_balance is not a valid tool, try
            one of [get_bank_account_id].
>     tool: Error: get_savings_account_balance is not a valid tool, try one
            of [get_bank_account_id].
>     human: It seems that I don't have the necessary tools to retrieve your
             investment and savings account balances. Another assistant with
             the appropriate tools will assist you further.
>     ai: I will assist you in retrieving your savings account balance.
          Please hold on.
>     tool: 100000.1
>     human: I have retrieved your investment account balance, which is
             $100,000.10. Now, let's wait for the savings account balance
             to be provided by another assistant so we can sum them up.
>
> saving_account_assistant
>     human: Get my investment and saving account balance. And sum them up.
>     tool: 123
>     ai: I have obtained your bank account ID as 123. I will now proceed to
          get the balances of your investment and savings accounts.
>     tool: Error: get_investment_account_balance is not a valid tool,
            try one of [get_bank_account_id].
>     tool: Error: get_savings_account_balance is not a valid tool,
            try one of [get_bank_account_id].
>     human: It seems that I don't have the necessary tools to retrieve your
             investment and savings account balances. Another assistant with the
             appropriate tools will assist you further.
>     ai: I will assist you in retrieving your savings account balance.
          Please hold on.
>     tool: 100000.1
>     human: I have retrieved your investment account balance, which is
             $100,000.10. Now, let's wait for the savings account balance to be
             provided by another assistant so we can sum them up.
>     tool: 10000.1
>     human: Your savings account balance is $10,000.10. 
>
> Now, let's sum up your investment and savings account balances:
>
> - Investment Account Balance: $100,000.10
> - Savings Account Balance: $10,000.10
>
> Total Balance = $100,000.10 + $10,000.10 = $110,000.20
>
> FINAL ANSWER: Your total balance, combining both investment and
  savings accounts, is $110,000.20.
It is nice to see that LLM is able to reason things and pick the tools to complete the task.



Comments