diff --git a/docs/source/_static/tools_correct_answer.png b/docs/source/_static/tools_correct_answer.png new file mode 100644 index 000000000..816b96df5 Binary files /dev/null and b/docs/source/_static/tools_correct_answer.png differ diff --git a/docs/source/_static/tools_file_list.png b/docs/source/_static/tools_file_list.png new file mode 100644 index 000000000..2972ed463 Binary files /dev/null and b/docs/source/_static/tools_file_list.png differ diff --git a/docs/source/_static/tools_rename_files.png b/docs/source/_static/tools_rename_files.png new file mode 100644 index 000000000..8bf6afc18 Binary files /dev/null and b/docs/source/_static/tools_rename_files.png differ diff --git a/docs/source/_static/tools_simple_example.png b/docs/source/_static/tools_simple_example.png new file mode 100644 index 000000000..e7aa096e4 Binary files /dev/null and b/docs/source/_static/tools_simple_example.png differ diff --git a/docs/source/_static/tools_wrong_answer.png b/docs/source/_static/tools_wrong_answer.png new file mode 100644 index 000000000..25372aab1 Binary files /dev/null and b/docs/source/_static/tools_wrong_answer.png differ diff --git a/docs/source/users/index.md b/docs/source/users/index.md index 918170187..f9e3300a3 100644 --- a/docs/source/users/index.md +++ b/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 diff --git a/docs/source/users/tools.md b/docs/source/users/tools.md new file mode 100644 index 000000000..9d26155cb --- /dev/null +++ b/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 +``` + +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: + +Use of the arithmetic tools file. + + +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: + +Incorrect use of the Black-Merton-Scholes formula for pricing a put option. + +Next, use the `/tools` command with the same query to get the correct answer: + +Incorrect use of the Black-Merton-Scholes formula for pricing a put option. + +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: + +Get the list of all filenames in a target folder. + +Next, test the second function: + +Change file extensions for selected filenames in a target folder. + +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 +``` + +To list all the tool file names: + +``` +/tools -l +``` + +and the response is + +``` +The available tools files are: ['arithmetic.py', 'finance.py', 'filetools.py'] +``` diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml index 91ef5699f..42057ce1b 100644 --- a/packages/jupyter-ai-magics/pyproject.toml +++ b/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] diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/__init__.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/__init__.py index a8fe9eb50..974365718 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/__init__.py +++ b/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 diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/tools.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/tools.py new file mode 100644 index 000000000..83ade590e --- /dev/null +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/tools.py @@ -0,0 +1,291 @@ +import argparse +import ast +import os +from pathlib import Path +from typing import Dict, Literal, Type + +import numpy as np +from jupyter_ai.models import HumanChatMessage +from jupyter_ai_magics.providers import BaseProvider +from jupyter_core.paths import jupyter_config_dir +from langchain_core.prompts import PromptTemplate +from langchain_core.runnables import ConfigurableFieldSpec +from langchain_core.runnables.history import RunnableWithMessageHistory +from langchain_core.tools import tool +from langgraph.graph import MessagesState, StateGraph +from langgraph.prebuilt import ToolNode + +from .base import BaseChatHandler, SlashCommandRoutingType + +TOOLS_DIR = os.path.join(jupyter_config_dir(), "jupyter-ai", "tools") + +PROMPT_TEMPLATE = """Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question. + +Chat History: +{chat_history} +Follow Up Input: {question} +Standalone question: +Format the answer to be as pretty as possible. +""" +CONDENSE_PROMPT = PromptTemplate.from_template(PROMPT_TEMPLATE) + + +class ExceptionNoToolsFile(Exception): + """Missing tools file""" + + pass + + +class ExceptionModelDoesTakeTools(Exception): + """Model is not a chat model that takes tools""" + + pass + + +class ExceptionModelNotAuthorized(Exception): + """Authentication failed for model authorization""" + + pass + + +class ExceptionNotChatModel(Exception): + """Not a chat model""" + + pass + + +class ToolsChatHandler(BaseChatHandler): + """Processes messages prefixed with /tools. This actor will + bind a .py collection of tools to the LLM and + build a computational graph to direct queries to tools + that apply to the prompt. If there is no appropriate tool, + the LLM will default to a standard chat response from the LLM + without using tools. + """ + + id = "tools" + name = "Use tools with LLM" + help = "Ask a question that uses your custom tools" + routing_type = SlashCommandRoutingType(slash_id="tools") + + uses_llm = True + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.parser.prog = "/tools" + self.parser.add_argument( + "-t", + "--tools", + action="store", + default=None, + type=str, + help="Uses tools in the given file name", + ) + self.parser.add_argument( + "-l", + "--list", + action="store_true", + help="Lists available files in tools directory.", + ) + + self.parser.add_argument("query", nargs=argparse.REMAINDER) + self.tools_file_path = None + + def setup_llm(self, provider: Type[BaseProvider], provider_params: Dict[str, str]): + """Sets up the LLM before creating the LLM Chain""" + unified_parameters = { + "verbose": True, + **provider_params, + **(self.get_model_parameters(provider, provider_params)), + } + llm = provider(**unified_parameters) + self.llm = llm + return llm + + # https://python.langchain.com/v0.2/docs/integrations/platforms/' + def create_llm_chain( + self, provider: Type[BaseProvider], provider_params: Dict[str, str] + ): + """Uses the LLM set up to create the LLM Chain""" + llm = self.setup_llm(provider, provider_params) + prompt_template = llm.get_chat_prompt_template() + self.llm = llm + + runnable = prompt_template | llm # type:ignore + if not llm.manages_history: + runnable = RunnableWithMessageHistory( + runnable=runnable, # type:ignore[arg-type] + get_session_history=self.get_llm_chat_memory, + input_messages_key="input", + history_messages_key="history", + history_factory_config=[ + ConfigurableFieldSpec( + id="last_human_msg", + annotation=HumanChatMessage, + ), + ], + ) + self.llm_chain = runnable + + def get_tool_files(self) -> list: + """ + Gets required tool files from TOOLS_DIR: `.jupyter/jupyter-ai/tools/` + which is the directory in which all tool files are placed. + """ + try: + if os.path.isfile(self.tools_file_path): + file_paths = [self.tools_file_path] + elif os.path.isdir(self.tools_file_path): + file_paths = [] + for filename in os.listdir(self.tools_file_path): + file_paths.append(os.path.join(self.tools_file_path, filename)) + return file_paths + except UnboundLocalError as e: + raise ExceptionNoToolsFile() + + def use_llm_with_tools(self, query: str) -> str: + """ + LangGraph documentation : https://langchain-ai.github.io/langgraph/tutorials/introduction/ + The code below: + 1. Extracts the function names in the custom tools file + 2. Adds the tools to the Tool Node + 3. Binds the Tool Node to the LLM + 4. Sets up a basic LangGraph with nodes and edges + 5. Compiles the graph into a runnable app + 6. This function is then called with a prompt + Every time a query is submitted the langgraph is rebuilt in case the tools file has been changed. + """ + + def call_tool(state: MessagesState) -> Dict[str, list]: + """Calls the requisite tool in the LangGraph""" + messages = state["messages"] + response = self.model_with_tools.invoke(messages) + return {"messages": [response]} + + def conditional_continue(state: MessagesState) -> Literal["tools", "__end__"]: + """ + Branches from tool back to the agent or ends the + computation on the langgraph + """ + messages = state["messages"] + last_message = messages[-1] + # if last_message.tool_calls: + if hasattr(last_message, "tool_calls") and last_message.tool_calls: + return "tools" + return "__end__" + + def get_tools(file_paths: list): + """ + Gets all tool objects from the tool files. + Returns tool objects of functions that have the `@tool` decorator. + Ignores all code in tool files that does not relate to tool functions. + """ + if len(file_paths) > 0: + tools = [] # tool objects + for file_path in file_paths: + try: # For each tool file, collect tool list and function source code + with open(file_path) as file: + content = file.read() + tree = ast.parse(content) # Build AST tree + for node in ast.walk( + tree + ): # Get the nodes with @tool decorator + if isinstance(node, ast.FunctionDef): + for decorator in node.decorator_list: + if ( + isinstance(decorator, ast.Name) + and decorator.id == "tool" + ): + exec( + ast.unparse(node) + ) # dynamically execute the tool function (object in memory) + tools.append( + eval(node.name) + ) # adds function to tool objects list + except FileNotFoundError: + raise ExceptionNoToolsFile() + return tools # a list of function objects + else: + self.reply("No available tool files.") + + # Get tool file(s), then tools within tool files, and create tool node from tools + tool_files = self.get_tool_files() # Get all tool files (python modules) + tools = get_tools(tool_files) # get tool objects + tool_node = ToolNode(tools) # create a LangGraph node with tool objects + + # Bind tools to LLM + # Check if the LLM class takes tools else advise user accordingly. + # Can be extended to include temperature parameter + self.llm = self.setup_llm( + self.config_manager.lm_provider, self.config_manager.lm_provider_params + ) + assert self.llm + if not self.llm.is_chat_provider: + raise ExceptionNotChatModel() + try: + self.model_with_tools = self.llm.bind_tools( # type:ignore[attr-defined] + tools + ) + except AttributeError: + raise ExceptionModelDoesTakeTools() + except Exception: + raise ExceptionModelNotAuthorized() + + # Initialize graph + agentic_workflow = StateGraph(MessagesState) + # Define the agent and tool nodes we will cycle between + agentic_workflow.add_node("agent", call_tool) + agentic_workflow.add_node("tools", tool_node) + # Add edges to the graph + agentic_workflow.add_edge("__start__", "agent") + agentic_workflow.add_conditional_edges("agent", conditional_continue) + agentic_workflow.add_edge("tools", "agent") + # Compile graph + app = agentic_workflow.compile() + + # Run query + res = app.invoke({"messages": query}) + return res["messages"][-1].content + + async def process_message(self, message: HumanChatMessage): + args = self.parse_args(message) + if args is None: + return + + if args.list: + tool_files = os.listdir(os.path.join(Path.home(), TOOLS_DIR)) + self.reply(f"The available tools files are: {tool_files}") + return + elif args.tools: + self.tools_file_path = os.path.join(Path.home(), TOOLS_DIR, args.tools) + else: + self.tools_file_path = os.path.join(Path.home(), TOOLS_DIR) + + query = " ".join(args.query) + if not query: + self.reply(f"{self.parser.format_usage()}", message) + return + + # self.get_llm_chain() + + try: + with self.pending("Using LLM with tools ..."): + response = self.use_llm_with_tools(query) + self.reply(response, message) + except ExceptionNoToolsFile as e: + self.reply(f"Tools file not found at {self.tools_file_path}.") + self.log.error(e) + except ExceptionModelDoesTakeTools as e: + self.reply(f"Not a chat model that takes tools.") + self.log.error(e) + except ExceptionModelNotAuthorized as e: + self.reply( + f"API failed. Model not authorized or provider package not installed." + ) + self.log.error(e) + except ExceptionNotChatModel as e: + self.reply(f"Not a chat model, cannot be used with tools.") + self.log.error(e) + except Exception as e: + self.log.error(e) diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index 08c8c5a47..81b11342f 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -21,6 +21,7 @@ GenerateChatHandler, HelpChatHandler, LearnChatHandler, + ToolsChatHandler, ) from .completions.handlers import DefaultInlineCompletionHandler from .config_manager import ConfigManager @@ -381,6 +382,7 @@ def _init_chat_handlers(self): export_chat_handler = ExportChatHandler(**chat_handler_kwargs) fix_chat_handler = FixChatHandler(**chat_handler_kwargs) + tools_chat_handler = ToolsChatHandler(**chat_handler_kwargs) chat_handlers["default"] = default_chat_handler chat_handlers["/ask"] = ask_chat_handler @@ -389,6 +391,7 @@ def _init_chat_handlers(self): chat_handlers["/learn"] = learn_chat_handler chat_handlers["/export"] = export_chat_handler chat_handlers["/fix"] = fix_chat_handler + chat_handlers["/tools"] = tools_chat_handler slash_command_pattern = r"^[a-zA-Z0-9_]+$" for chat_handler_ep in chat_handler_eps: