Adding custom tools to MemGPT

MemGPT allows you to create custom tools for agents to use. MemGPT will automatically generate an OpenAI JSON schema for the tool based on the function arguments and docstring.

Defining functions for custom tools

The format of the function definition for the tool must have the following:

• The first argument is self: Agentas the function will eventually be attached and executed as part of the Agent class (this is so that functions can modify agent state if needed, such as the agent's memories, though this is of course optional)
• Properly formatted docstrings with a description, arguments, and returns.
• (Optional) Examples of how the tool should be called

Below is an example of a function definition that is properly formatted and allows the agent to "roll a dice":

Python

from memgpt.agent import Agent 

def roll_d20(self: Agents) -> str:
    """
    Simulate the roll of a 20-sided die (d20).

    This function generates a random integer between 1 and 20, inclusive,
    which represents the outcome of a single roll of a d20.

    Returns:
        int: A random integer between 1 and 20, representing the die roll.

    Example:
        >>> roll_d20()
        15  # This is an example output and may vary each time the function is called.
    """
    import random 
    
    dice_role_outcome = random.randint(1, 20)
    output_string = f"You rolled a {dice_role_outcome}"
    return output_string

⚠️

Function requirements

The functions you write MUST have proper docstrings and type hints - this is because MemGPT will use these docstrings and types to automatically create a JSON schema that is used in the LLM prompt. Use the docstrings and types annotations from the example functions for guidance.

⚠️

Function output length

Your custom function should always return a string that is capped in length. If your string goes over the specified limit, it will be truncated internally. This is to prevent potential context overflows caused by uncapped string returns (for example, a rogue HTTP request that returns a string larger than the LLM context window).

If you return any type other than str (e.g. `dict``) in your custom functions, MemGPT will attempt to cast the result to a string (and truncate the result if it is too long). It is preferable to return strings - think of your function returning a natural language description of the outcome.

Functions using external packages

To import external packages, you need to make sure to include the import <package-name> line inside of the function definition, as it otherwise will not be loaded by MemGPT. We show an example of this below for a google_search function that imports serpapi:

Python

def google_search(agent: Agent, query: str): 
    """

    A tool to search google with the provided query, and return a list of relevant summaries 
    and URLs.

    Args:
        query (str): The search query.

    Returns:
        List[Tuple[str, str]]: A list of up to 5 tuples, each containing a summary of the search 
        result and the URL of the search result in the form (summary, URL)

    Example:
        >>> google_search("How can I make a french 75?")
        [
            (
                "To make a French 75 cocktail, combine 1½ oz. gin, ¾ oz. fresh lemon juice, 
                and ¾ oz. simple syrup in a cocktail shaker with ice. Shake vigorously, then 
                strain into a large flute. Top with 2 oz. Champagne and garnish with a long 
                spiral lemon twist. The recipe prefers gin, but cognac is also traditional. 
                Serve in Champagne flutes for the full effect.",
                "https://www.bonappetit.com/recipe/french-75-3"
            )
        ]
    """

    # imports must be inside the function
    import serpapi
    ... 

If you are using MemGPT with a MemGPT server, you must make sure that the packages are installed into the environment that your server is running. Otherwise, when the server executes the function, it will be missing the packages and be unable to import them when the tool is called.

Functions using environment variables & secrets

Similar to packages, environment variables you access from the function must be set in the same environment as the MemGPT server. If you are starting the server with the CLI, you can set the environment variables in the server command:

CLI

SERP_API_KEY=<SERP_API_KEY> memgpt server 

If you are running the docker compose up, then you can set the environment variables in the compose.yaml file (link).

If your environment variables are set in the server environment, you can access them with:

Python

def google_search(agent: Agent, query: str): 
    """
		...
    """

    # imports must be inside the function
    import serpapi
    import os 
    
    api_key = os.environ["SERPAPI_API_KEY"]

Creating tools

You can create a tool through the Python client by passing the function:

Python

from memgpt import create_client 

client = create_client() 
tool = client.create_tool(print_tool, update=False) # update=whether to override an existing function

Creating a tool will return a ToolModel object which contains the source code, JSON schema, tool_id, and tool name.

You can view existing tools with:

Python

tools = client.list_tools()

Adding tools to agents

Currently, tools can only be added to agents at creation time (support for adding tools to existing agents will be coming soon). You can create an agent by specifying the tool name when creating the agent.

Python

agent_state = client.create_agent(tools=[tool.name])

Overriding default tools

By default, MemGPT agents are created with the following tools (source code):

- "send_message"
- "pause_heartbeats"
- "conversation_search"
- "conversation_search_date"
- "archival_memory_insert"
- "archival_memory_search"

If you don't want your agent to use these tools, you can specify include_base_tools=False when creating an agent to only attach the tools you manually specify:

Python

agent_state = client.create_agent(
  tools=[custom_tool.name, "send_message", "pause_heartbeats", "archival_memory_insert", "archival_memory_search"],
  include_base_tools=False
)