langchain agents custom function calling

How To Use Custom Functions With LangChain Agents

In this post, we’re going to take a look at how we can setup LangChain agents with custom function calling. Basically, we’ll make our custom tools and integrate them with our agent.

In order to demonstrate this process, we’ll make an agent that’ll be able to fetch a stock price from a certain ticker, and make calculations with it.

Custom function calling shows remarkable potential for development of next generation AI applications.

Prerequisites

Firstly, we need to import all the necessary libraries and tools for this project. Among these is also yfinance, which we’ll use to retrieve the stock price data.

from datetime import datetime, timedelta
import os
import json
import yfinance as yf
from typing import List
from langchain.tools import BaseTool
from typing import Optional, Type
from pydantic import BaseModel, Field
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from langchain.chat_models import ChatOpenAI

Next thing we need to do is get the API key for OpenAI and set it as an environment variable. This will allow us to use LangChain library.

ROOT = os.path.dirname(__file__)
def get_token(token_name):
    
    auth_file = open(os.path.join(ROOT, 'auth.json'))
    auth_data = json.load(auth_file)
    token = auth_data[token_name]
    return token

os.environ['OPENAI_API_KEY'] = get_token('openai-token')

Custom functions

In the heart of this tutorial the custom functions, which our agents is going to call to get the results it’ll need to output the correct answer.

First function, we’ll need, is to retrieve stock price information of a ticker, our agent will pass through an argument.

def get_stock_price(symbol):
    ticker = yf.Ticker(symbol)
    todays_data = ticker.history(period='1d')
    return round(todays_data['Close'][0], 2)

Next, we’ll define a function that will calculate the percentage of how much the price changed in last number of days. We’ll also set an argument for the number of days.

def get_price_change_percent(symbol, days_ago):
    ticker = yf.Ticker(symbol)

    # Get today's date
    end_date = datetime.now()

    # Get the date N days ago
    start_date = end_date - timedelta(days=days_ago)

    # Convert dates to string format that yfinance can accept
    start_date = start_date.strftime('%Y-%m-%d')
    end_date = end_date.strftime('%Y-%m-%d')

    # Get the historical data
    historical_data = ticker.history(start=start_date, end=end_date)

    # Get the closing price N days ago and today's closing price
    old_price = historical_data['Close'].iloc[0]
    new_price = historical_data['Close'].iloc[-1]

    # Calculate the percentage change
    percent_change = ((new_price - old_price) / old_price) * 100

    return round(percent_change, 2)

And for the last custom tool, we’ll define a couple of functions that will calculate which ticker is the best performing in the last x amount of days.

def calculate_performance(symbol, days_ago):
    ticker = yf.Ticker(symbol)
    end_date = datetime.now()
    start_date = end_date - timedelta(days=days_ago)
    start_date = start_date.strftime('%Y-%m-%d')
    end_date = end_date.strftime('%Y-%m-%d')
    historical_data = ticker.history(start=start_date, end=end_date)
    old_price = historical_data['Close'].iloc[0]
    new_price = historical_data['Close'].iloc[-1]
    percent_change = ((new_price - old_price) / old_price) * 100
    return round(percent_change, 2)

def get_best_performing(stocks, days_ago):
    best_stock = None
    best_performance = None
    for stock in stocks:
        try:
            performance = calculate_performance(stock, days_ago)
            if best_performance is None or performance > best_performance:
                best_stock = stock
                best_performance = performance
        except Exception as e:
            print(f"Could not calculate performance for {stock}: {e}")
    return best_stock, best_performance

Defining custom tools for LangChain agents

This is the part where we connect our custom functions with our agent. First, let’s take a look at the custom tool for retrieving price data.

class StockPriceCheckInput(BaseModel):
    """Input for Stock price check."""

    stockticker: str = Field(..., description="Ticker symbol for stock or index")

class StockPriceTool(BaseTool):
    name = "get_stock_ticker_price"
    description = "Useful for when you need to find out the price of stock. You should input the stock ticker used on the yfinance API"

    def _run(self, stockticker: str):
        # print("i'm running")
        price_response = get_stock_price(stockticker)

        return price_response

    def _arun(self, stockticker: str):
        raise NotImplementedError("This tool does not support async")

    args_schema: Optional[Type[BaseModel]] = StockPriceCheckInput

It is essential we provide descriptions for the arguments of our custom function and tool itself. This is because our agent will use it to know what it’s good for and when to use it.

We can do the same thing for the other 2 tools as well.

class StockChangePercentageCheckInput(BaseModel):
    """Input for Stock ticker check. for percentage check"""

    stockticker: str = Field(..., description="Ticker symbol for stock or index")
    days_ago: int = Field(..., description="Int number of days to look back")

class StockPercentageChangeTool(BaseTool):
    name = "get_price_change_percent"
    description = "Useful for when you need to find out the percentage change in a stock's value. You should input the stock ticker used on the yfinance API and also input the number of days to check the change over"

    def _run(self, stockticker: str, days_ago: int):
        price_change_response = get_price_change_percent(stockticker, days_ago)

        return price_change_response

    def _arun(self, stockticker: str, days_ago: int):
        raise NotImplementedError("This tool does not support async")

    args_schema: Optional[Type[BaseModel]] = StockChangePercentageCheckInput


# the best performing

class StockBestPerformingInput(BaseModel):
    """Input for Stock ticker check. for percentage check"""

    stocktickers: List[str] = Field(..., description="Ticker symbols for stocks or indices")
    days_ago: int = Field(..., description="Int number of days to look back")

class StockGetBestPerformingTool(BaseTool):
    name = "get_best_performing"
    description = "Useful for when you need to get the performance of multiple stocks over a period. You should input a list of stock tickers used on the yfinance API and also input the number of days to check the change over"

    def _run(self, stocktickers: List[str], days_ago: int):
        price_change_response = get_best_performing(stocktickers, days_ago)

        return price_change_response

    def _arun(self, stockticker: List[str], days_ago: int):
        raise NotImplementedError("This tool does not support async")

    args_schema: Optional[Type[BaseModel]] = StockBestPerformingInput

Setting up components of LangChain agents

For the last part of this tutorial, we need to define all the core components of our agent and put them all together. Furthermore, it’s important we set our AgentType to OPENAI_FUNCTIONS, which is the type that’s able to perform custom function calling.

tools = [StockPriceTool(),StockPercentageChangeTool(), StockGetBestPerformingTool()]
llm = ChatOpenAI(temperature=0, model="gpt-3.5-turbo-0613")
open_ai_agent = initialize_agent(tools,
                        llm,
                        agent=AgentType.OPENAI_FUNCTIONS,
                        verbose=True)

price = open_ai_agent.run("Has google's stock gone up over the past 90 days?")
print(price)

And the following is what we get from it.

langchain agents custom function calling

Remarkable, as you can see, we didn’t even have to specifically define from which ticker (GOOGL in this case) it needs to get the data.

Conclusion

To conclude, we built an agent, that’s able to use custom functions to give us accurate results, by automatically recognizing what arguments it needs to pass through them. I learned a lot when working on this project and I hope it helps you as well.

Also feel free to check out my other LangChain tutorials.

Share this article:

Related posts

Discussion(0)