ADD(Agent): Can now specify how you want you tool to be found

docs/examples/agent/multi_document_agents-v1.ipynb

@@ -491,6 +491,59 @@
     "agents_dict, extra_info_dict = await build_agents(docs)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "7364a596",
+   "metadata": {},
+   "source": [
+    "# Custom Tool Retrieval Functionality\n",
+    "\n",
+    "In this section, we introduce a feature that allows users to define their own custom tool retrieval function. This is particularly beneficial for those who have encountered challenges in consistently obtaining real tools from OpenAI.\n",
+    "\n",
+    "The custom tool retrieval function provides the flexibility to alter the way tools are retrieved based on your specific needs. This not only gives you more control over the tool retrieval process but also ensures that you achieve the desired results.\n",
+    "\n",
+    "To utilize this functionality, you simply need to define your custom tool retrieval function and integrate it into the existing code. This allows you to customize the tool retrieval process to your specific requirements, ensuring a more efficient and effective workflow.\n",
+    "\n",
+    "By enabling users to define the way tools are retrieved, we aim to offer a more customizable and user-friendly experience. We believe that this feature will significantly enhance your ability to work with tools and improve the overall performance of your Jupyter Notebook."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "207394e7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define a function to retrieve the tools, here we will use a closest match method\n",
+    "\n",
+    "from typing import List\n",
+    "from llama_index.tools import BaseTool\n",
+    "from fuzzywuzzy import process\n",
+    "\n",
+    "\n",
+    "\n",
+    "def get_function_by_name(tools: List[BaseTool], name: str) -> BaseTool:\n",
+    "    \"\"\"Get function by name.\"\"\"\n",
+    "    name_to_tool = {tool.metadata.name: tool for tool in tools}\n",
+    "    closest_match, score = process.extractOne(name, name_to_tool.keys())\n",
+    "    if closest_match not in name_to_tool:\n",
+    "        raise ValueError(f\"Tool with name {name} not found\")\n",
+    "    return name_to_tool[closest_match]\n",
+    "\n",
+    "# Now when creating the agent we can use the function to retrieve the tools\n",
+    "\n",
+    "agent = OpenAIAgent.from_tools(\n",
+    "        query_engine_tools,\n",
+    "        llm=function_llm,\n",
+    "        verbose=True,\n",
+    "        system_prompt=f\"\"\"\\\n",
+    "You are a specialized agent designed to answer queries about the `{file_base}.html` part of the LlamaIndex docs.\n",
+    "You must ALWAYS use at least one of the tools provided when answering a question; do NOT rely on prior knowledge.\\\n",
+    "\"\"\",\n",
+    "        get_tool_by=get_function_by_name,\n",
+    "    )"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "899ca55b-0c02-429b-a765-8e4f806d503f",

llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/base.py

@@ -13,6 +13,7 @@
     List,
     Optional,
     Type,
+    Callable,
 )
 
 from llama_index.agent.openai.step import OpenAIAgentWorker
@@ -53,6 +54,7 @@ def __init__(
         default_tool_choice: str = "auto",
         callback_manager: Optional[CallbackManager] = None,
         tool_retriever: Optional[ObjectRetriever[BaseTool]] = None,
+        get_tool_by: Callable[[List[BaseTool], str], BaseTool] = get_function_by_name,
     ) -> None:
         """Init params."""
         callback_manager = callback_manager or llm.callback_manager
@@ -64,6 +66,7 @@ def __init__(
             max_function_calls=max_function_calls,
             callback_manager=callback_manager,
             prefix_messages=prefix_messages,
+            get_tool_by=get_tool_by,
         )
         super().__init__(
             step_engine,
@@ -88,6 +91,7 @@ def from_tools(
         callback_manager: Optional[CallbackManager] = None,
         system_prompt: Optional[str] = None,
         prefix_messages: Optional[List[ChatMessage]] = None,
+        get_tool_by: Callable[[List[BaseTool], str], BaseTool] = get_function_by_name,
         **kwargs: Any,
     ) -> "OpenAIAgent":
         """Create an OpenAIAgent from a list of tools.
@@ -133,4 +137,5 @@ def from_tools(
             max_function_calls=max_function_calls,
             callback_manager=callback_manager,
             default_tool_choice=default_tool_choice,
+            get_tool_by=get_tool_by,
         )

llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/openai_assistant_agent.py

@@ -3,7 +3,7 @@
 import json
 import logging
 import time
-from typing import Any, Dict, List, Optional, Tuple, Union, cast
+from typing import Any, Dict, List, Optional, Tuple, Union, cast, Callable
 
 from llama_index.agent.openai.utils import get_function_by_name
 from llama_index.core.agent.types import BaseAgent
@@ -59,7 +59,10 @@ def from_openai_thread_messages(thread_messages: List[Any]) -> List[ChatMessage]
 
 
 def call_function(
-    tools: List[BaseTool], fn_obj: Any, verbose: bool = False
+    tools: List[BaseTool],
+    fn_obj: Any,
+    verbose: bool = False,
+    get_tool_by: Callable[[List[BaseTool], str], BaseTool] = get_function_by_name,
 ) -> Tuple[ChatMessage, ToolOutput]:
     """Call a function and return the output as a string."""
     from openai.types.beta.threads.required_action_function_tool_call import Function
@@ -71,7 +74,7 @@ def call_function(
     if verbose:
         print("=== Calling Function ===")
         print(f"Calling function: {name} with args: {arguments_str}")
-    tool = get_function_by_name(tools, name)
+    tool = get_tool_by(tools, name)
     argument_dict = json.loads(arguments_str)
     output = tool(**argument_dict)
     if verbose:
@@ -90,7 +93,10 @@ def call_function(
 
 
 async def acall_function(
-    tools: List[BaseTool], fn_obj: Any, verbose: bool = False
+    tools: List[BaseTool],
+    fn_obj: Any,
+    verbose: bool = False,
+    get_tool_by: Callable[[List[BaseTool], str], BaseTool] = get_function_by_name,
 ) -> Tuple[ChatMessage, ToolOutput]:
     """Call an async function and return the output as a string."""
     from openai.types.beta.threads.required_action_function_tool_call import Function
@@ -102,7 +108,7 @@ async def acall_function(
     if verbose:
         print("=== Calling Function ===")
         print(f"Calling function: {name} with args: {arguments_str}")
-    tool = get_function_by_name(tools, name)
+    tool = get_tool_by(tools, name)
     argument_dict = json.loads(arguments_str)
     async_tool = adapt_to_async_tool(tool)
     output = await async_tool.acall(**argument_dict)
@@ -152,6 +158,7 @@ def __init__(
         run_retrieve_sleep_time: float = 0.1,
         file_dict: Dict[str, str] = {},
         verbose: bool = False,
+        get_tool_by: function = get_function_by_name,
     ) -> None:
         """Init params."""
         from openai import OpenAI
@@ -168,6 +175,7 @@ def __init__(
         self._run_retrieve_sleep_time = run_retrieve_sleep_time
         self._verbose = verbose
         self.file_dict = file_dict
+        self._get_tool_fn = get_tool_by
 
         self.callback_manager = callback_manager or CallbackManager([])
 
@@ -187,6 +195,7 @@ def from_new(
         verbose: bool = False,
         file_ids: Optional[List[str]] = None,
         api_key: Optional[str] = None,
+        get_tool_by: function = get_function_by_name,
     ) -> "OpenAIAssistantAgent":
         """From new assistant.
 
@@ -204,6 +213,7 @@ def from_new(
             verbose: verbose
             file_ids: list of file ids
             api_key: OpenAI API key
+            get_tool_by: Function to get tool by name, format: get_tool_by(tools, name)
 
         """
         from openai import OpenAI
@@ -244,6 +254,7 @@ def from_new(
             file_dict=file_dict,
             run_retrieve_sleep_time=run_retrieve_sleep_time,
             verbose=verbose,
+            get_tool_by=get_tool_by,
         )
 
     @classmethod
@@ -257,6 +268,7 @@ def from_existing(
         callback_manager: Optional[CallbackManager] = None,
         api_key: Optional[str] = None,
         verbose: bool = False,
+        get_tool_by: function = get_function_by_name,
     ) -> "OpenAIAssistantAgent":
         """From existing assistant id.
 
@@ -269,6 +281,7 @@ def from_existing(
             callback_manager: callback manager
             api_key: OpenAI API key
             verbose: verbose
+            get_tool_by: Function to get tool by name, format: get_tool_by(tools, name)
 
         """
         from openai import OpenAI
@@ -289,6 +302,7 @@ def from_existing(
             instructions_prefix=instructions_prefix,
             run_retrieve_sleep_time=run_retrieve_sleep_time,
             verbose=verbose,
+            get_tool_by=get_tool_by,
         )
 
     @property
@@ -350,7 +364,12 @@ def _run_function_calling(self, run: Any) -> List[ToolOutput]:
         tool_output_objs: List[ToolOutput] = []
         for tool_call in tool_calls:
             fn_obj = tool_call.function
-            _, tool_output = call_function(self._tools, fn_obj, verbose=self._verbose)
+            _, tool_output = call_function(
+                self._tools,
+                fn_obj,
+                verbose=self._verbose,
+                get_tool_by=self._get_tool_fn,
+            )
             tool_output_dicts.append(
                 {"tool_call_id": tool_call.id, "output": str(tool_output)}
             )
@@ -373,7 +392,10 @@ async def _arun_function_calling(self, run: Any) -> List[ToolOutput]:
         for tool_call in tool_calls:
             fn_obj = tool_call.function
             _, tool_output = await acall_function(
-                self._tools, fn_obj, verbose=self._verbose
+                self._tools,
+                fn_obj,
+                verbose=self._verbose,
+                get_tool_by=self._get_tool_fn,
             )
             tool_output_dicts.append(
                 {"tool_call_id": tool_call.id, "output": str(tool_output)}

llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/step.py

@@ -5,7 +5,7 @@
 import logging
 import uuid
 from threading import Thread
-from typing import Any, Dict, List, Optional, Tuple, Union, cast, get_args
+from typing import Any, Dict, List, Optional, Tuple, Union, cast, get_args, Callable
 
 from llama_index.agent.openai.utils import resolve_tool_choice
 from llama_index.core.agent.types import (
@@ -80,6 +80,7 @@ def call_function(
     tools: List[BaseTool],
     tool_call: OpenAIToolCall,
     verbose: bool = False,
+    get_tool_by: Callable[[List[BaseTool], str], BaseTool] = get_function_by_name,
 ) -> Tuple[ChatMessage, ToolOutput]:
     """Call a function and return the output as a string."""
     # validations to get passed mypy
@@ -95,7 +96,7 @@ def call_function(
     if verbose:
         print("=== Calling Function ===")
         print(f"Calling function: {name} with args: {arguments_str}")
-    tool = get_function_by_name(tools, name)
+    tool = get_tool_by(tools, name)
     argument_dict = json.loads(arguments_str)
 
     # Call tool
@@ -118,7 +119,10 @@ def call_function(
 
 
 async def acall_function(
-    tools: List[BaseTool], tool_call: OpenAIToolCall, verbose: bool = False
+    tools: List[BaseTool],
+    tool_call: OpenAIToolCall,
+    verbose: bool = False,
+    get_tool_by: Callable[[List[BaseTool], str], BaseTool] = get_function_by_name,
 ) -> Tuple[ChatMessage, ToolOutput]:
     """Call a function and return the output as a string."""
     # validations to get passed mypy
@@ -134,7 +138,7 @@ async def acall_function(
     if verbose:
         print("=== Calling Function ===")
         print(f"Calling function: {name} with args: {arguments_str}")
-    tool = get_function_by_name(tools, name)
+    tool = get_tool_by(tools, name)
     async_tool = adapt_to_async_tool(tool)
     argument_dict = json.loads(arguments_str)
     output = await async_tool.acall(**argument_dict)
@@ -166,12 +170,14 @@ def __init__(
         max_function_calls: int = DEFAULT_MAX_FUNCTION_CALLS,
         callback_manager: Optional[CallbackManager] = None,
         tool_retriever: Optional[ObjectRetriever[BaseTool]] = None,
+        get_tool_by: Callable[[List[BaseTool], str], BaseTool] = get_function_by_name,
     ):
         self._llm = llm
         self._verbose = verbose
         self._max_function_calls = max_function_calls
         self.prefix_messages = prefix_messages
         self.callback_manager = callback_manager or self._llm.callback_manager
+        self._get_tools_fn = get_tool_by
 
         if len(tools) > 0 and tool_retriever is not None:
             raise ValueError("Cannot specify both tools and tool_retriever")
@@ -195,6 +201,7 @@ def from_tools(
         callback_manager: Optional[CallbackManager] = None,
         system_prompt: Optional[str] = None,
         prefix_messages: Optional[List[ChatMessage]] = None,
+        get_tool_by: Callable[[List[BaseTool], str], BaseTool] = get_function_by_name,
         **kwargs: Any,
     ) -> "OpenAIAgentWorker":
         """Create an OpenAIAgent from a list of tools.
@@ -235,6 +242,7 @@ def from_tools(
             verbose=verbose,
             max_function_calls=max_function_calls,
             callback_manager=callback_manager,
+            get_tool_by=get_tool_by,
         )
 
     def get_all_messages(self, task: Task) -> List[ChatMessage]:
@@ -353,13 +361,13 @@ def _call_function(
             CBEventType.FUNCTION_CALL,
             payload={
                 EventPayload.FUNCTION_CALL: function_call.arguments,
-                EventPayload.TOOL: get_function_by_name(
+                EventPayload.TOOL: self._get_tools_fn(
                     tools, function_call.name
                 ).metadata,
             },
         ) as event:
             function_message, tool_output = call_function(
-                tools, tool_call, verbose=self._verbose
+                tools, tool_call, verbose=self._verbose, get_tool_by=self._get_tools_fn
             )
             event.on_end(payload={EventPayload.FUNCTION_OUTPUT: str(tool_output)})
         sources.append(tool_output)
@@ -382,13 +390,13 @@ async def _acall_function(
             CBEventType.FUNCTION_CALL,
             payload={
                 EventPayload.FUNCTION_CALL: function_call.arguments,
-                EventPayload.TOOL: get_function_by_name(
+                EventPayload.TOOL: self._get_tools_fn(
                     tools, function_call.name
                 ).metadata,
             },
         ) as event:
             function_message, tool_output = await acall_function(
-                tools, tool_call, verbose=self._verbose
+                tools, tool_call, verbose=self._verbose, get_tool_by=self._get_tools_fn
             )
             event.on_end(payload={EventPayload.FUNCTION_OUTPUT: str(tool_output)})
         sources.append(tool_output)