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: Agent
as the function will eventually be attached and executed as part of theAgent
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":
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
:
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:
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:
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:
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:
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.
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:
agent_state = client.create_agent(
tools=[custom_tool.name, "send_message", "pause_heartbeats", "archival_memory_insert", "archival_memory_search"],
include_base_tools=False
)
Updated about 1 month ago