Skip to content

Latest commit

 

History

History
163 lines (119 loc) · 6.45 KB

tool_guide.md

File metadata and controls

163 lines (119 loc) · 6.45 KB
Raw

Custom Tools

Follow these instructions to create your own custom tools.

For information on how to add authentication to your tool, please refer to the Tool Auth guide.

If you'd like to contribute a custom tool to the main repository, please make sure to add it to the community folder . Make sure you've enabled the INSTALL_COMMUNITY_DEPS build arg in the docker-compose.yml file by setting it to true.

If you've forked the repo and want to customize Toolkit to your liking, then feel free to add it anywhere.

Step 1: Choose a Tool to Implement

A tool is anything that can retrieve data to augment the chat experience with RAG.

You can easily add a tool from Langchain or LlamaIndex by importing from their Python SDKs, or write your own custom implementation.

Step 2: Add your Tool code implementation

Add your tool implementation in the community folder here.

In a fork, feel free to add it directly to the backend folder.

If you need to install a new library to run your tool, execute the following command and run make dev again.

poetry add <MODULE> --group community

Tool example

Below is an example of adding a tool using the Langchain SDK.

To implement a tool, you will need to inherit from the BaseTool abstract class. You'll need to implement the following methods:

  • is_available: Should return True/False depending on if the tool is available during runtime. Often this will involve checking for example an API key is set, if it is required.
  • get_tool_definition: Should return the config for your tool. See existing implementations for details.
  • call: The main method that is called during chat. This should return a list of dictionary results. For best model performance, each dictionary entry should be serialized to contain at minimum a text field. Optionally, a url and title fields can be added to include more metadata for citations.

Let's take a look at the community-implemented ArxivRetriever:

from typing import Any, Dict, List

from langchain_community.utilities import ArxivAPIWrapper

from backend.schemas.tool import ToolCategory, ToolDefinition
from backend.tools.base import BaseTool


class ArxivRetriever(BaseTool):
    ID = "arxiv"

    def __init__(self):
        self.client = ArxivAPIWrapper()

    @classmethod
    def is_available(cls) -> bool:
        return True

    @classmethod
    def get_tool_definition(cls) -> ToolDefinition:
        return ToolDefinition(
            name=cls.ID,
            display_name="Arxiv",
            implementation=cls,
            parameter_definitions={
                "query": {
                    "description": "Query for retrieval.",
                    "type": "str",
                    "required": True,
                }
            },
            is_visible=False,
            is_available=cls.is_available(),
            error_message=cls.generate_error_message(),
            category=ToolCategory.DataLoader,
            description="Retrieves documents from Arxiv.",
        )

    async def call(self, parameters: dict, **kwargs: Any) -> List[Dict[str, Any]]:
        query = parameters.get("query", "")
        result = self.client.run(query)
        return [{"text": result}]

Important: the call() method is where the tool call itself happens, i.e: this is where you want to perform your business logic (API call, data ETL, hardcoded response, etc).

The return value of this tool MUST BE a list of dictionaries that have the text (required), url (optional), title (optional) fields as keys.

For example:

return [{"text": "The fox is blue", "url": "wikipedia.org/foxes", "title": "Color of foxes"}, {..}, {..}]

Custom Parameters

You can set custom parameters that you can use with the parameter_definitions key when specifying your tool config.

These parameters are automatically generated by Command during the chat process. For example, if you define

parameter_definitions={
    "query": {
        "description": "Query for retrieval.",
        "type": "str",
        "required": True,
    }
},

for some search API, Command will generate an appropriate query to the tool based on your prompt.

Step 3: Importing your Tool

Next, add your tool class to the init file by locating it in src/community/tools/__init__.py. Import your tool here, then add it to the __all__ list.

Finally, you will need to add your tool definition to the config file. Locate it in src/community/config/tools.py, and import your tool at the top with from backend.tools import ...

Finally, to enable your tool, add your tool class as an enum value in the Tool enum. For example, My_Tool = MyToolClass.

Step 4: Test Your Tool!

Now, when you run the toolkit, all the visible tools, including the one you just added, should be available!

  • Run make dev
  • Open http://localhost:4000/
  • Toggle your tool and send a message that triggers it
  • Appreciate a grounded response with something ✨you created from scratch✨!

Remember, you can also access your tools via the API.

  • List tools:
curl --location --request GET 'http://localhost:8000/v1/tools' \
--header 'User-Id: me' \
--header 'Content-Type: application/json' \
--data '{}'
  • Chat turns with tools:
curl --location 'http://localhost:8000/v1/chat-stream' \
--header 'User-Id: me' \
--header 'Content-Type: application/json' \
--data '{
    "message": "Tell me about the aya model",
    "tools": [{"name": "Arxiv"}]
}
'

Step 5 (extra): Add Unit tests

If you would like to go above and beyond, it would be helpful to add some unit tests to ensure that your tool is working as expected. Create a file here and add a few test cases.