Adding new slash command /tools to use LLMs with your own custom tools

[srdas]

This draft PR proposes a new feature enabling the use of custom tools (functions) with LLMs. Lots of work to do here and see the end of this note for all the design decisions to make.

In many situations LLMs will handle complex mathematical formulas quite well and return correct answers, but this is often not the case. Even for textual repsonses, using custom functions can constrain responses to formats and content that is more accurate and acceptable.

This PR adds a new a slash command /tools that directs the LLM to use functions from a tools library that you provide. This is a single file titled mytools.py which may be stored in the default directory, that is, the one from which Jupyter is started. We provide an example of the tools file here, containing just three functions. Make sure to add the @tool decorator to each function and to import all packages that are not already installed within each function. The functions below are common financial formulas that are widely in use and you may expect that an LLM would be trained on these. While this is accurate, we will see that the LLM is unable to accurately execute the math in these formulas.

# mytools.py

@tool
def BlackMertonScholes_Call(S: float, # current stock price
                            K: float, # exercise price of the option
                            T: float, # option maturity in years
                            d: float, # annualized dividend rate
                            r: float, # annualized risk free interest rate
                            v: float, # stock volatility
                           ):
    """Black-Scholes-Merton option pricing model for call options"""
    from scipy.stats import norm
    d1 = (np.log(S/K) + (r-d+0.5*v**2)*T)/(v*np.sqrt(T))
    d2 = d1 - v*np.sqrt(T)
    call_option_price = S*np.exp(-d*T)*norm.cdf(d1) - K*np.exp(-r*T)*norm.cdf(d2)
    return call_option_price

@tool
def BlackMertonScholes_Put(S: float, # current stock price
                            K: float, # exercise price of the option
                            T: float, # option maturity in years
                            d: float, # annualized dividend rate
                            r: float, # annualized risk free interest rate
                            v: float, # stock volatility
                           ):
    """Black-Scholes-Merton option pricing model for put options"""
    from scipy.stats import norm
    d1 = (np.log(S/K) + (r-d+0.5*v**2)*T)/(v*np.sqrt(T))
    d2 = d1 - v*np.sqrt(T)
    put_option_price =  K*np.exp(-r*T)*norm.cdf(-d2) - S*np.exp(-d*T)*norm.cdf(-d1)
    return put_option_price

@tool
def calculate_monthly_payment(principal, annual_interest_rate, loan_term_years):
    """
    Calculate the monthly mortgage payment.
    Args:
        principal (float): The principal amount of the loan.
        annual_interest_rate (float): The annual interest rate as a decimal (e.g., 0.06 for 6%).
        loan_term_years (int): The loan term in years.
    Returns:
        float: The monthly mortgage payment.
    """
    import math
    # Convert annual interest rate to monthly interest rate
    monthly_interest_rate = annual_interest_rate / 12
    # Calculate the number of monthly payments
    num_payments = loan_term_years * 12
    # Calculate the monthly payment using the annuity formula
    monthly_payment = (principal * monthly_interest_rate) / (1 - math.pow(1 + monthly_interest_rate, -num_payments))
    return monthly_payment

Each function contains the @tool decorator and the required imports. Note also the comment string that describes what each tool does. This will help direct the LLM to relevant tool. Providing sufficient guiding comments in the function is helpful in the form of comment strings, variable annotations, and expolicit argument comments, example of which are shown in the code above. For example, default values in comments will be used by the LLM if the user forgets to provide them (for example, see the explicit mention of a 6% interest rate in calculate_monthly_payment function above).

When the /tools command is used, Jupyter AI will bind the custom tools file to the LLM currently in use and build a LangGraph (https://langchain-ai.github.io/langgraph/). It will use this graph to respond to the query and use the appropriate tools, if available. In the examples below, we use Claude Haiku as the base model:

As an example, submit this query in the chat interface without using tools: "What is the price of a put option where the stock price is 100, the exercise price is 101, the time to maturity is 1 year, the risk free rate is 3%, the dividend rate is zero, and the stock volatility is 20%?" The correct answer to this query is $6.93. However, though the LLM returns the correct formula, it computes the answer incorrectly:

Next, use the /tools command with the same query to get the correct answer:

While this is the simplest implementation and the first use of LangGraph in Jupyter AI, there are several improvements and additions to consider:

1. Add a special icon for /tools
2. Make the location of the mytools.py file flexible, or pass it as a parameter in the slash command? DONE
3. Pass tools file location on startup, or in chat? Only allow one tools file, or handle multiple tool files. DONE
4. Make handling the different providers (Chat*) and models (model_id) part of the UI? Removed
5. To integrate with chat history and memory or not? Will using history confuse the tool and return worse answers?
6. Format output from the tools? For example, suppress the problem with % sign messing up output?
7. Show full exchange or only the answer? That is, all Human and AI messages or only the final AI response?
8. Improve error handling DONE
9. Documentation. Currently this is sufficient for users, but maybe add documentation for developers to use LangGraph in a more extended manner. Long term: Using the more advanced features of LangGraph, Agents, Multi-agentic workflows, etc. DONE
10. What's the best way to add tools to magics?

[krassowski]

I feel like using an entrypoint approach would be more scalable than reading from user directory. The kernel may live on a different computer than the server so the server will often not have access to read from the mytools.py that user created.

[srdas]

@dlqqq Updated this PR to include:

1. Multiple tool files within the tools folder (.jupyter/jupyter-ai/tools/) as suggested by @ellisonbg
2. /tools -t <tool_file_name> <query> for use with a single tool file (suggested by @dlqqq)
3. /tools <query> to use all tool files (suggested by @dlqqq)
4. /tools -l to list all tool file names (suggested by @dlqqq)
5. Updated extraction of tool names using ast parsing (suggested by @krassowski)
6. Updated tools.md for full documentation on using tools.
   Main concerns remaining are the use of exec() and eval() as security risks.

[dlqqq]

@srdas Thank you for working on this! I've left some points of feedback below. No rush on addressing this, since we're still using this PR to experiment & test.

High-level callouts (subject to future edits):

• (low priority) Can you add typing to the methods & attributes in tools.py? The lack of typing is causing the typing tests (done via mypy) to fail in CI.

[srdas]

@dlqqq Note: While the PR is working for Bedrock Chat models, it needs to be tested for other providers. Additional classes may need to be added to each partner provider module or new modules may be needed. I will work on this next.

1. Tested on Ollama models by changing the provider from Ollama to ChatOllama and adding the required import: from langchain_ollama import ChatOllama -- works as expected with tools and without tools.