A demo LLM agent with math-related tools

demo_tools_agent.py

from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain_openai import ChatOpenAI
from langchain.tools import tool
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from dotenv import load_dotenv
import os

load_dotenv()
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")

@tool
def add(a: int, b: int) -> int:
    """
    Add two numbers together.
    
    Parameters:
    a (int): The first number to add
    b (int): The second number to add
    
    Returns:
    (int): The sum of a + b added together
    """
    return a + b

@tool
def multiply(a: int, b: int) -> int:
    """
    Multiply two numbers together.
    
    Parameters:
    a (int): The first number to multiply
    b (int): The second number to multiply
    
    Returns:
    (int): The product of a * b multiplied together
    """
    return a * b

@tool
def divide(a: int, b: int) -> int:
    """
    Divides a by b and returns a whole integer
    
    Parameters:
    a (int): The dividend to be used in the division process
    b (int): The divisor to be used in the division process
    
    Returns:
    (int): The quotient result of the division process
    """
    return a / b

@tool
def square(a: int) -> int:
    """
    Calculates the square of a number.
    
    Parameters:
    a (int): The number to square
    
    Returns:
    (int): The square of the number passed in as an argument
    """
    return a * a

@tool
def findAreaOfCircle(r: int) -> float:
  """
  Finds the area of a circle using a provided radius.
  
  Parameters:
  r (int): The radius of a circle
  
  Returns:
  (int): The area of the circle
  """
  return 3.14 * r ** 2

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", """You are a mathematical assistant. Use your tools to answer questions.
         If you do not have a tool to answer the question, say so.

        Return only the answers. e.g
        Human: What is 1 + 1?
        AI: 2
        """),
        MessagesPlaceholder("chat_history", optional=True),
        ("human", "{input}"),
        MessagesPlaceholder("agent_scratchpad"),
    ]
)

llm = ChatOpenAI(model="gpt-4o", temperature=0)
toolkit = [add, multiply, square, divide, findAreaOfCircle]
agent = create_openai_tools_agent(llm, toolkit, prompt)
agent_executor = AgentExecutor(agent=agent, tools=toolkit, verbose=False)

result = agent_executor.invoke({"input": "Find the area of a circle with a radius of 5"})
print(result['output'])