Unleashing the Power of Multiple Agents with LangGraph

Introduction

Welcome back to this ongoing series on building an AI-powered stock analysis tool! I took a significant leap forward in the previous article by converting the AI stock agent from AgentExecutor to LangGraph. I explored the advantages of using LangGraph for complex workflows and state management, setting the stage for more advanced enhancements to this tool.

In this article, I’ll dive into the exciting world of multiple agents and discover how they can revolutionize the stock analysis tool. Expanding the LangGraph implementation to incorporate specialized agents for various tasks can unlock new levels of efficiency and effectiveness in the AI-driven trading process. So, let’s get started and explore the power of multiple agents with LangGraph!

TL;DR? -> Code

The Need for Multiple Agents

While the previous implementation of the AI stock analysis tool using LangGraph showed promising results, it still relied on a single agent to handle all tasks. This approach has limitations, as a single agent may struggle to excel at every aspect of the complex stock analysis process. One can overcome these limitations by introducing multiple agents specializing in a specific area and creating a more robust and efficient analysis pipeline.

Using multiple agents offers several key benefits. First, it allows us to break down the stock analysis process into smaller, more manageable tasks. Each agent can focus on a specific aspect, such as scanning the market, performing fundamental analysis, conducting chart analysis, or managing risk. This specialization allows agents to focus on their specific areas, resulting in more accurate and insightful outcomes.

Furthermore, multiple agents can collaborate, sharing information and building upon each other’s findings. This collaboration can lead to a more comprehensive understanding of the market and individual stocks, ultimately enhancing the decision-making process for traders. By leveraging each agent’s strengths, one can create a powerful AI-powered stock analysis tool that adapts to the ever-changing market conditions and provides users with timely and actionable insights.

Implementing Multiple Agents with LangGraph

To implement multiple agents with LangGraph, I need to make key updates to the code structure. Let’s examine the changes required to accommodate this new approach.

Python

# agent.py

from typing import Annotated, TypedDict, Optional, Literal, Callable

from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_anthropic import ChatAnthropic
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.messages import ToolMessage
from langgraph.graph import END, StateGraph
from langgraph.prebuilt.tool_node import tools_condition
from langgraph.graph.message import AnyMessage, add_messages
from dotenv import load_dotenv

from app.tools.stock_stats import (
    get_gainers,
    get_losers,
    get_stock_price_history,
    get_stock_quantstats,
    get_stock_ratios,
    get_key_metrics,
    get_stock_sector_info,
    get_valuation_multiples,
    get_stock_universe,
)
from app.tools.stock_sentiment import get_news_sentiment
from app.tools.stock_relative_strength import get_relative_strength
from app.tools.stock_charts import get_stock_chart_analysis
from app.tools.risk_management import (
    calculate_r_multiples,
    calculate_technical_stops,
    calculate_position_size,
)
from app.tools.utils import create_tool_node_with_fallback
from app.chains.templates import *

load_dotenv()


def update_dialog_stack(left: list[str], right: Optional[str]) -> list[str]:
    """Push or pop the state."""
    if right is None:
        return left
    if right == "pop":
        return left[:-1]
    return left + [right]


class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]
    dialog_state: Annotated[
        list[
            Literal[
                "assistant",
                "full_scan",
                "full_analysis",
                "chart_analysis",
                "determine_risk",
                "gainers_losers",
            ]
        ],
        update_dialog_stack,
    ]


class ToFullScanAssistant(BaseModel):
    """Transfers work to a specialized assistant to handle a stock universe scan."""

    symbol: str = Field(description="The symbol of the stock to analyze.")


class ToFullAnalysisAssistant(BaseModel):
    """Transfers work to a specialized assistant to handle a full analysis of a stock."""

    symbol: str = Field(description="The symbol of the stock to analyze.")


class ToChartAnalysisAssistant(BaseModel):
    """Transfers work to a specialized assistant to handle a chart analysis of a stock."""

    symbol: str = Field(description="The symbol of the stock to analyze.")


class ToRiskManagementAssistant(BaseModel):
    """Transfers work to a specialized assistant to handle risk management."""

    symbol: str = Field(description="The symbol of the stock to analyze.")
    request: str = Field(
        description="Any necessary followup questions the risk assistant should clarify before proceeding."
    )


class ToGainersLosersAssistant(BaseModel):
    """Transfers work to a specialized assistant to handle fetching gainers and losers."""

    request: str = Field(
        description="Any necessary followup questions the gainers/losers assistant should clarify before proceeding."
    )


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

    def __call__(self, state: AgentState, config: RunnableConfig):
        while True:
            result = self.runnable.invoke(state)

            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}
                messages = state["messages"] + [("user", "Respond with a real output.")]
                state = {**state, "messages": messages}
            else:
                break
        return {"messages": result}


def create_full_scan_agent(llm: Runnable) -> Assistant:
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", SCAN_TEMPLATE),
            ("placeholder", "{messages}"),
        ]
    )
    scan_tools = [get_stock_universe, get_stock_price_history, get_stock_quantstats]
    runnable = prompt | llm.bind_tools(scan_tools)
    return Assistant(runnable)


def create_full_analysis_agent(llm: Runnable) -> Assistant:
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", FULL_ANALYSIS_TEMPLATE),
            ("placeholder", "{messages}"),
        ]
    )
    analysis_tools = [
        get_stock_price_history,
        get_key_metrics,
        get_stock_ratios,
        get_stock_sector_info,
        get_valuation_multiples,
        get_news_sentiment,
        get_relative_strength,
        get_stock_quantstats,
    ]
    runnable = prompt | llm.bind_tools(analysis_tools)
    return Assistant(runnable)


def create_chart_analysis_agent(llm: Runnable) -> Assistant:
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", CHART_ANALYSIS_TEMPLATE),
            ("placeholder", "{messages}"),
        ]
    )
    chart_tools = [get_stock_chart_analysis]
    runnable = prompt | llm.bind_tools(chart_tools)
    return Assistant(runnable)


def create_risk_management_agent(llm: Runnable) -> Assistant:
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", RISK_TEMPLATE),
            ("placeholder", "{messages}"),
        ]
    )
    risk_tools = [
        calculate_technical_stops,
        calculate_r_multiples,
        calculate_position_size,
    ]
    runnable = prompt | llm.bind_tools(risk_tools)
    return Assistant(runnable)


def create_gainers_losers_agent(llm: Runnable) -> Assistant:
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", GAINERS_LOSERS_TEMPLATE),
            ("placeholder", "{messages}"),
        ]
    )
    runnable = prompt | llm.bind_tools([get_gainers, get_losers])
    return Assistant(runnable)


def create_primary_assistant(llm: Runnable) -> Assistant:
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", BASE_TEMPLATE),
            ("placeholder", "{messages}"),
        ]
    )
    tools = [TavilySearchResults(max_results=1)]
    runnable = prompt | llm.bind_tools(
        tools
        + [
            ToFullScanAssistant,
            ToFullAnalysisAssistant,
            ToChartAnalysisAssistant,
            ToRiskManagementAssistant,
            ToGainersLosersAssistant,
        ]
    )
    return Assistant(runnable)


def create_entry_node(assistant_name: str, new_dialog_state: str) -> Callable:
    def entry_node(state: AgentState) -> dict:
        tool_call_id = state["messages"][-1].tool_calls[0]["id"]
        return {
            "messages": [
                ToolMessage(
                    content=f"The assistant is now the {assistant_name}. Reflect on the above conversation between the host assistant and the user."
                    f" The user's intent is unsatisfied. Use the provided tools to assist the user. Remember, you are {assistant_name},"
                    " and actions are not complete until after you have successfully invoked the appropriate tool."
                    " Do not mention who you are - just act as the proxy for the assistant.",
                    tool_call_id=tool_call_id,
                )
            ],
            "dialog_state": new_dialog_state,
        }

    return entry_node


def should_continue(state: AgentState) -> Literal[
    "scan_stocks_tools",
    "analyze_stocks_tools",
    "chart_analysis_tools",
    "risk_management_tools",
    "gainers_losers_tools",
    "__end__",
]:
    route = tools_condition(state)
    if route == END:
        return END

    tool_calls = state["messages"][-1].tool_calls
    if tool_calls:
        tool_name = tool_calls[0]["name"]
        if tool_name in [
            "get_stock_price_history",
            "get_key_metrics",
            "get_stock_ratios",
            "get_stock_sector_info",
            "get_valuation_multiples",
            "get_news_sentiment",
            "get_relative_strength",
            "get_stock_quantstats",
        ]:
            return "analyze_stocks_tools"
        elif tool_name == "get_stock_universe":
            return "scan_stocks_tools"
        elif tool_name == "get_stock_chart_analysis":
            return "chart_analysis_tools"
        elif tool_name in [
            "calculate_technical_stops",
            "calculate_r_multiples",
            "calculate_position_size",
        ]:
            return "risk_management_tools"
        elif tool_name in ["get_gainers", "get_losers"]:
            return "gainers_losers_tools"
    return END


def route_primary_assistant(state: AgentState) -> Literal[
    "primary_assistant_tools",
    "enter_scan_stocks",
    "enter_analyze_stocks",
    "enter_chart_analysis",
    "enter_risk_management",
    "enter_gainers_losers",
    "__end__",
]:
    route = tools_condition(state)
    if route == END:
        return END

    tool_calls = state["messages"][-1].tool_calls
    if tool_calls:
        if tool_calls[0]["name"] == ToFullScanAssistant.__name__:
            return "enter_scan_stocks"
        elif tool_calls[0]["name"] == ToFullAnalysisAssistant.__name__:
            return "enter_analyze_stocks"
        elif tool_calls[0]["name"] == ToChartAnalysisAssistant.__name__:
            return "enter_chart_analysis"
        elif tool_calls[0]["name"] == ToRiskManagementAssistant.__name__:
            return "enter_risk_management"
        elif tool_calls[0]["name"] == ToGainersLosersAssistant.__name__:
            return "enter_gainers_losers"
        return "primary_assistant_tools"
    raise ValueError("Invalid route")


def create_anthropic_agent_graph() -> StateGraph:
    llm = ChatAnthropic(temperature=0, model_name="claude-3-opus-20240229")

    builder = StateGraph(AgentState)

    # Scan Assistant
    builder.add_node(
        "enter_scan_stocks",
        create_entry_node("Stock Scan Assistant", "scan_stocks"),
    )
    builder.add_node("scan_stocks", create_full_scan_agent(llm))
    builder.add_edge("enter_scan_stocks", "scan_stocks")
    builder.add_node(
        "scan_stocks_tools",
        create_tool_node_with_fallback(
            [get_stock_universe, get_stock_price_history, get_stock_quantstats]
        ),
    )
    builder.add_edge("scan_stocks_tools", "scan_stocks")
    builder.add_conditional_edges("scan_stocks", should_continue)

    # Analysis Assistant
    builder.add_node(
        "enter_analyze_stocks",
        create_entry_node("Stock Analysis Assistant", "analyze_stocks"),
    )
    builder.add_node("analyze_stocks", create_full_analysis_agent(llm))
    builder.add_edge("enter_analyze_stocks", "analyze_stocks")
    builder.add_node(
        "analyze_stocks_tools",
        create_tool_node_with_fallback(
            [
                get_stock_price_history,
                get_key_metrics,
                get_stock_ratios,
                get_stock_sector_info,
                get_valuation_multiples,
                get_news_sentiment,
                get_relative_strength,
                get_stock_quantstats,
            ]
        ),
    )
    builder.add_edge("analyze_stocks_tools", "analyze_stocks")
    builder.add_conditional_edges("analyze_stocks", should_continue)

    # Chart Assistant
    builder.add_node(
        "enter_chart_analysis",
        create_entry_node("Stock Chart Analysis Assistant", "chart_analysis"),
    )
    builder.add_node("chart_analysis", create_chart_analysis_agent(llm))
    builder.add_edge("enter_chart_analysis", "chart_analysis")
    builder.add_node(
        "chart_analysis_tools",
        create_tool_node_with_fallback([get_stock_chart_analysis]),
    )
    builder.add_edge("chart_analysis_tools", "chart_analysis")
    builder.add_conditional_edges("chart_analysis", should_continue)

    # Risk Management Assistant
    builder.add_node(
        "enter_risk_management",
        create_entry_node("Stock Risk Management Assistant", "risk_management"),
    )
    builder.add_node("risk_management", create_risk_management_agent(llm))
    builder.add_edge("enter_risk_management", "risk_management")
    builder.add_node(
        "risk_management_tools",
        create_tool_node_with_fallback(
            [
                calculate_technical_stops,
                calculate_r_multiples,
                calculate_position_size,
            ]
        ),
    )
    builder.add_edge("risk_management_tools", "risk_management")
    builder.add_conditional_edges("risk_management", should_continue)

    # Gainers/Losers Assistant
    builder.add_node(
        "enter_gainers_losers",
        create_entry_node("Stock Gainers/Losers Assistant", "gainers_losers"),
    )
    builder.add_node("gainers_losers", create_gainers_losers_agent(llm))
    builder.add_edge("enter_gainers_losers", "gainers_losers")
    builder.add_node(
        "gainers_losers_tools",
        create_tool_node_with_fallback([get_gainers, get_losers]),
    )
    builder.add_edge("gainers_losers_tools", "gainers_losers")
    builder.add_conditional_edges("gainers_losers", should_continue)

    # Primary Assistant
    builder.add_node("primary_assistant", create_primary_assistant(llm))
    builder.add_node(
        "primary_assistant_tools",
        create_tool_node_with_fallback([TavilySearchResults(max_results=1)]),
    )
    builder.add_conditional_edges(
        "primary_assistant",
        route_primary_assistant,
        {
            "enter_scan_stocks": "enter_scan_stocks",
            "enter_analyze_stocks": "enter_analyze_stocks",
            "enter_chart_analysis": "enter_chart_analysis",
            "enter_risk_management": "enter_risk_management",
            "enter_gainers_losers": "enter_gainers_losers",
            "primary_assistant_tools": "primary_assistant_tools",
            END: END,
        },
    )
    builder.add_edge("primary_assistant_tools", "primary_assistant")
    builder.set_entry_point("primary_assistant")

    graph = builder.compile()
    return graph

Alright, that’s a big paste. Let’s break it down.

A. Overview of the updated code structure

New agent classes: I’ve introduced new agent classes for each specialized task, such as FullScanAssistant, FullAnalysisAssistant, ChartAnalysisAssistant, RiskManagementAssistant, and GainersLosersAssistant. These classes encapsulate the specific functionality and tools required for each agent.
Updated AgentState and message types: I’ve expanded the AgentState to include additional properties, such as dialog_state keeping track of the current agent in the conversation. I’ve also introduced new message types, like ToFullScanAssistant, ToFullAnalysisAssistant, etc., to facilitate communication between agents.

B. Creating specialized agents

Prompt templates for each agent: I’ve defined tailored prompt templates for each agent, ensuring they have the necessary context and instructions to perform their specific tasks effectively. These templates are stored in the templates.py file for easy reference and maintenance.
Binding tools to each agent: Each agent is equipped with specific tools relevant to its specialized task. For example, the FullScanAssistant has access to tools like get_stock_universe, get_stock_price_history, and get_stock_quantstats, while the RiskManagementAssistant utilizes tools such as calculate_technical_stops, calculate_r_multiples, and calculate_position_size.

C. Designing the agent graph

Adding nodes for each agent: In the create_anthropic_agent_graph function, I’ve added nodes for each specialized agent, such as scan_stocks, analyze_stocks, chart_analysis, risk_management, and gainers_losers. These nodes represent the different states in the conversation flow.
Defining entry points and transitions between agents: I’ve created entry point nodes like enter_scan_stocks, enter_analyze_stocks, etc., which serve as the starting points for each specialized agent. These nodes are connected to their respective agent nodes, ensuring a smooth transition between agents.
Conditional edges and routing logic: I’ve implemented conditional edges and routing logic using functions like should_continue and route_primary_assistant. These functions determine the next agent to be invoked based on the user’s input and the current state of the conversation.

By structuring the code this way, the groundwork has been established for a versatile and scalable multi-agent system capable of effectively managing various facets of the stock analysis process.

Walkthrough of the Enhanced AI Stock Analysis Tool

Let’s explore how the enhanced AI stock analysis tool, powered by multiple agents, works in practice.

The primary assistant recognizes the need for a full stock analysis and transfers the request to the FullScanAssistant. This specialized agent fetches top stocks from the universe scan, the latest price history for each stock, key metrics, ratios, sector information, valuation multiples, news sentiment, relative strength, and fundamental data, compiling a comprehensive report for the user.

The primary assistant can also route this request to theChartAnalysisAssistant, which generates an in-depth chart analysis. The agents work together seamlessly, providing a holistic view of the stock.

The modular approach of using multiple agents allows for greater flexibility and adaptability. As market conditions change or new analysis techniques emerge, one can easily add or modify agents to keep the AI-powered stock analysis tool up-to-date and effective.

Conclusion

In this article, I’ve explored the potential of using multiple agents with LangGraph to enhance the AI-powered stock analysis tool. By creating specialized agents for tasks like full stock analysis and chart analysis, one can unlock a new level of efficiency and effectiveness in their trading process.

Using multiple agents allows us to break down the complex stock analysis process into manageable components. Each agent contributes unique expertise, collaborating seamlessly to provide users with a comprehensive and reliable assessment of the stocks they’re interested in.

The true power of the multi-agent system lies in its flexibility and adaptability. As market conditions evolve and new analysis techniques emerge, the tool can adapt by adding or modifying agents. This ensures that the AI-powered stock analysis tool remains a valuable resource for traders, regardless of market changes.

I have one more article on cloud deployment for this series in the works. I’m also working on articles about fine-tuning open-source LLMs. Axolotl fine-tuning Llama-3 for finance, anyone? But for now, embrace the power of multiple agents and LangGraph, and join me on this journey as I continue swapping symbols.

Swapping Symbols