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

docs/source/users/index.md

@@ -607,10 +607,16 @@ contents of the failing cell.
     class="screenshot" style="max-width:65%" />
 
 
+### Using custom tools in chat
+
+In order to use your own custom tools in chat, create a tools file named `mytools.py` in the default directory. Then use the `/tools` command with your query. For details on how to build your custom tools file and usage of the `/tools` command, refer to [Using your custom tools library in the chat interface](tools.md).
+
+
 ### Additional chat commands
 
 To start a new conversation, use the `/clear` command. This will clear the chat panel and reset the model's memory.
 
+
 ## The `%ai` and `%%ai` magic commands
 
 Jupyter AI can also be used in notebooks via Jupyter AI magics. This section

docs/source/users/tools.md

@@ -0,0 +1,175 @@
+# Using your custom tools library in the chat interface
+
+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.
+
+Jupyter AI includes a slash command `/tools` that directs the LLM to use functions from a tools library that you provide. You can have multiple tool files stored in the subdirectory `.jupyter/jupyter-ai/tools/`.
+
+You can use a single tool file from this subdirectory. In this case the usage of this slash command is as follows:
+
+```
+/tools -t <tools_file_name> <query>
+```
+
+For example, we may try using a tools file called `arithmetic.py`. Note that since the file has to be placed in `.jupyter/jupyter-ai/tools/`, only file name is needed in the command.
+
+```
+/tools -t arithmetic.py What is the sum of 1 and 2?
+```
+
+The contents of the example file `arithmetic.py` are very simple:
+
+```
+@tool
+def multiply(first_number: float, second_number: float):
+    """Multiplies two numbers together."""
+    return first_number * second_number
+
+@tool
+def add(first_number: float, second_number: float):
+    """Adds two numbers together."""
+    return first_number + second_number
+```
+
+The result is shown below:
+
+<img src="../_static/tools_simple_example.png"
+    width="75%"
+    alt='Use of the arithmetic tools file.'
+    class="screenshot" />
+
+
+We provide another example of the tools file here, called `finance.py`, containing just three functions. Make sure to add the `@tool` decorator to each function and to import all packages that are required 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.
+
+```
+@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
+    import numpy as np
+    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
+    import numpy as np
+    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 explicit 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.
+
+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:
+
+<img src="../_static/tools_wrong_answer.png"
+    width="75%"
+    alt='Incorrect use of the Black-Merton-Scholes formula for pricing a put option.'
+    class="screenshot" />
+
+Next, use the `/tools` command with the same query to get the correct answer:
+
+<img src="../_static/tools_correct_answer.png"
+    width="75%"
+    alt='Incorrect use of the Black-Merton-Scholes formula for pricing a put option.'
+    class="screenshot" />
+
+As a third example, we prepare a tools file called `filetools.py` that has the following code:
+```
+@tool
+def list_files(dir_path):
+    """Returns a list of files in the directory path"""
+    import os
+    file_list = os.listdir(dir_path)
+    file_list = [os.path.join(dir_path, f) for f in file_list]
+    print('Preparing file list')
+    return file_list
+
+
+@tool
+def rename_file_extensions(file_list, old_extension, new_extension):
+    """Renames file extensions from old_extension to new_extension"""
+    import os
+    for filename in file_list:
+        if filename.endswith(old_extension):
+            new_filename = filename.replace(old_extension, new_extension)
+            os.rename(filename, new_filename)
+            print(f'Renamed {filename} to {new_filename}')
+    return
+```
+
+There is a function `list_files` that collects all the file names in a target folder and another function `rename_file_extensions` that changes file names with specific extensions to new ones. Note that the second function needs to invoke the first function as an intermediate step.
+
+We test the first function as shown below:
+
+<img src="../_static/tools_file_list.png"
+    width="75%"
+    alt='Get the list of all filenames in a target folder.'
+    class="screenshot" />
+
+Next, test the second function:
+
+<img src="../_static/tools_rename_files.png"
+    width="100%"
+    alt='Change file extensions for selected filenames in a target folder.'
+    class="screenshot" />
+
+We see that the file extensions have been changed. The logs show that the LLM has used both tools, getting the file list using the first function and then passing it to the second function to change the file extensions.
+
+You can try the other tools in this example or build your own custom tools file to experiment with this feature.
+
+If you do not want to use any specific tool file, that is, use all the tool files together, the command is simply:
+
+```
+/tools <query>
+```
+
+To list all the tool file names:
+
+```
+/tools -l
+```
+
+and the response is
+
+```
+The available tools files are: ['arithmetic.py', 'finance.py', 'filetools.py']
+```

packages/jupyter-ai-magics/pyproject.toml

@@ -29,6 +29,7 @@ dependencies = [
     "typing_extensions>=4.5.0",
     "click~=8.0",
     "jsonpath-ng>=1.5.3,<2",
+    "langgraph",
 ]
 
 [project.optional-dependencies]

packages/jupyter-ai/jupyter_ai/chat_handlers/__init__.py

@@ -7,3 +7,4 @@
 from .generate import GenerateChatHandler
 from .help import HelpChatHandler
 from .learn import LearnChatHandler
+from .tools import ToolsChatHandler