Tinz Twins

Tinz Twins

Technology is our passion.


Build a Multi-Agent Stock Market Analyst to Compare Stock Price Performance

6 minute read

Do you want to build an AI chatbot that can answer complex financial questions, like plotting a chart for stock price performance? If so, this article might be very interesting for you!

Llama sitting in front of a stock market terminal
Generated with Grok

Building a chatbot app for stock market questions can be complex and time-consuming, but Python packages like Chainlit and AG2 make it much easier.

In addition, more and more small open-source models are entering the market, enabling you to run LLM apps entirely on your local machine for free. That opens the door for developers worldwide to create amazing AI apps.

In this step-by-step guide, we’ll show you how to build an AI-powered stock market analyst that can compare year-to-date (YTD) stock price gains between different stocks.

Let’s build a local chatbot for the stock market!

Sneak Peak

Tech Stack

Our chatbot uses Meta’s Llama 3.1:8b as LLM and has tool use support. So you can easily add new functions or tools to the application.

For this, we use the following frameworks:

  • Chainlit as an intuitive user-interface
  • AG2 to implement the AI agents

We’ve written a comprehensive article about the AG2 framework. If you are not familiar with AG2, we recommend reading this article first.

Prerequisites

You will need the following prerequisites:

  • Python package manager of your choice (We use conda).
  • A code editor of your choice (We use Visual Studio Code).
  • Download Ollama and install Llama3.1. Make sure that it runs on your computer.
  • A computer with a GPU (We use a MacBook Pro M4 Max 48 GB).

ℹ️ We recommend a computer with at least 16GB RAM to run the examples in this guide efficiently.

Step-by-Step Guide

Step 1: Setup the development environment

  • Create a conda environment: A virtual environment keeps your main system clean.
conda create -n stock-price-performance python=3.11.9
conda activate stock-price-performance
  • Clone the GitHub repo:
git clone https://github.com/tinztwins/finllm-apps.git
  • Install requirements: Go to the folder compare-stock-price-performance and run the following command:
pip install -r requirements.txt
  • Make sure that Ollama is running on your computer:

Screenshot: Is Ollama running?

Step 2: Create the Chainlit App

  • Import required libraries: First, we import all necessary libraries.
import chainlit as cl
from autogen import ConversableAgent, register_function, GroupChat, GroupChatManager
import matplotlib.pyplot as plt
import yfinance
from typing import Annotated
  • Start a new chat session: Every Chainlit app follows a life cycle. The framework creates a new chat session when a user opens the Chainlit app. When a new chat session begins, Chainlit triggers the on_chat_start() function. In addition, we can store data in memory during the user session. So, we can persist the agent objects with the command cl.user_session.set("agent_name", agent_object).
# Agent code and Tool registration
# ...
cl.user_session.set("manager", manager)
cl.user_session.set("user_proxy_agent", user_proxy_agent)
  • Create Agents: In AG2, you can create a new agent object with ConversableAgent(...). You can pass this object several parameters, like name, system_message, and llm_config. We use Llama3.1:8b as a language model. It supports tool use and provides reasoning capabilities to the agents.
# Configure the API endpoint for your model
config_list = [
  {
    "model": "llama3.1:latest",
    "api_type": "ollama",
  }
]

# Define the financial analyst agent that suggests tool calls.
  financial_analyst = ConversableAgent(
    name="Financial Analyst",
    system_message="You are a helpful AI assistant."
    "You can help with plotting financial charts."
    "Don't plot a chart when no stock symbols are given."
    "Return 'TERMINATE' when the task is done.",
    llm_config={"config_list": config_list},
  )

  # This is the assistant agent.
  assistant_agent = ConversableAgent(
    name="Assistant",
    system_message="You are a helpful AI assistant."
    "You can help with general queries and questions about the stock market."
    "Add 'TERMINATE' to every answer.",
    code_execution_config=False,
    llm_config={"config_list": config_list},
    human_input_mode="NEVER",
  )

  # The user proxy agent can interact with the assistant agent and the financial analyst.
  user_proxy_agent = ConversableAgent(
    name="User",
    llm_config=False,
    is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
    human_input_mode="NEVER",
    max_consecutive_auto_reply=3,
  )
  • Registration of the function: Next, we register a Python function/tool with the agents so they can use it to answer user questions. In AG2, you can use the register_function(...). The following code shows the function and its registration.
def get_stock_prices(stock_symbols, 
                     start_date, 
                     end_date):
  stock_data = yfinance.download(
    stock_symbols, 
    start=start_date, 
    end=end_date
  )
  return stock_data.get("Close")

def plot_ytd_gains(stock_symbols: Annotated[list[str], "The stock symbols to get the prices for. (list[str])"], 
                   start_date: Annotated[str, "Format: YYYY-MM-DD"], 
                   end_date: Annotated[str, "Format: YYYY-MM-DD"]) -> str:
  stock_data_list = list()
  for i in stock_symbols:
    stock_price_data = get_stock_prices(i, start_date, end_date)
    stock_ytd_gain = (stock_price_data - stock_price_data.iloc[0]) / stock_price_data.iloc[0] * 100
    stock_data_list.append(stock_ytd_gain)
  
  fig, ax = plt.subplots(figsize=(10, 6))
  for df in stock_data_list:
    plt.plot(df, label=df.columns.values[0])
  plt.title('Stock Price Performance YTD')
  plt.xlabel('Date')
  plt.setp(ax.get_xticklabels(), rotation=45, ha='right')
  plt.ylabel('Gain in %')
  plt.legend()
  plt.grid(True)
  cl.user_session.set("figure", fig)

  return "The task is done."

# Register the function to the two agents.
register_function(
  plot_ytd_gains,
  caller=financial_analyst, 
  executor=user_proxy_agent,
  name="plot_ytd_gains",
  description="A function to plot the stock gain YTD for the given stock symbols.",  # A description of the tool.
)
  • Create a GroupChat: Next, we need to connect the three agents. That makes it possible for the agents to talk to each other. In this context, we can configure the conversation. The following Python code shows that only the user_proxy_agent can talk to the other agents. However, the other agents can only talk to the user_proxy_agent. The group chat manager manages the conversation.
# Create a GroupChat
groupchat = GroupChat(agents=[user_proxy_agent, 
                              assistant_agent, 
                              financial_analyst], 
                    messages=[], 
                    max_round=5,
                    allowed_or_disallowed_speaker_transitions={
        user_proxy_agent: [assistant_agent, 
                           financial_analyst],
        assistant_agent: [user_proxy_agent],
        financial_analyst: [user_proxy_agent],
    },
    speaker_transitions_type="allowed",)
manager = GroupChatManager(groupchat=groupchat, llm_config={"config_list": config_list})
  • New message from the user: The on_message(message: cl.Message) function is called when a user sends a new message. In this context, the user_proxy_agent initiates a new group chat. Then, the manager decides which agent can answer the user’s question best. The manager forwards general questions about the stock market to the assistant_agent and questions about creating a chart to the financial_analyst. If a plot is created, it will be displayed. If not, a text response will be provided.
@cl.on_message
async def on_message(message: cl.Message):
  manager = cl.user_session.get("manager")
  user_proxy_agent = cl.user_session.get("user_proxy_agent")

  groupchat_result = await cl.make_async(user_proxy_agent.initiate_chat)(manager, message=message.content)
  fig = cl.user_session.get("figure")

  if fig:
    elements = [
      cl.Pyplot(name="plot", figure=fig, display="inline"),
    ]
    await cl.Message(
      content="You can see the stock gain YTD plot below:",
      elements=elements,
    ).send()
    cl.user_session.set("figure", 0)
  else:
    msg = cl.Message(content="")
    for chunk in groupchat_result.summary:
      await msg.stream_token(chunk)
    await msg.send()

ℹ️ If you want to learn more about Chainlit, check out our introduction article on Chainlit.

Step 3: Run the App

  • Start the Chainlit App: Go to the project folder and run the chatbot app with the following command:
chainlit run app.py
  • Chatbot App: Open http://localhost:8000 in your web browser and enter a question about the stock market.

Conclusion

Great! You have built a local AI chatbot to plot a financial chart. The combination of Chainlit and AG2 provides a strong foundation for building more advanced AI applications.

You can use the code provided as a starting point for your next project. Feel free to share your project on social media and mention us so we can see what you’re working on.

Happy coding!

💡 Do you enjoy our content and want to read super-detailed articles about data science topics? If so, be sure to check out our premium offer!

You may also enjoy