Skip to content

Commit

Permalink
feat: allow to load tools from external modules
Browse files Browse the repository at this point in the history
  • Loading branch information
jrmi committed Dec 21, 2024
1 parent 1934fcf commit 860157b
Show file tree
Hide file tree
Showing 22 changed files with 419 additions and 136 deletions.
5 changes: 5 additions & 0 deletions docs/config.rst
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,11 @@ Here is an example:
#MODEL = "local/<model-name>"
#OPENAI_BASE_URL = "http://localhost:11434/v1"
# Uncomment to change tool configuration
#TOOL_FORMAT = "markdown" # Select the tool formal. One of `markdown`, `xml`, `tool`
#TOOL_ALLOW_LIST = "save,append,patch,python" # Comma separated list of allowed tools
#TOOL_MODULES = "gptme.tools,custom.tools" # List of python comma separated python module path
The ``prompt`` section contains options for the prompt.

The ``env`` section contains environment variables that gptme will fall back to if they are not set in the shell environment. This is useful for setting the default model and API keys for :doc:`providers`.
Expand Down
83 changes: 83 additions & 0 deletions docs/custom_tool.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
Creating a Custom Tool for gptme
=================================

Introduction
------------
In gptme, a custom tool allows you to extend the functionality of the assistant by
defining new tools that can be executed.
This guide will walk you through the process of creating and registering a custom tool.

Creating a Custom Tool
-----------------------
To create a custom tool, you need to define a new instance of the `ToolSpec` class.
This class requires several parameters:

- **name**: The name of the tool.
- **desc**: A description of what the tool does.
- **instructions**: Instructions on how to use the tool.
- **examples**: Example usage of the tool.
- **execute**: A function that defines the tool's behavior when executed.
- **block_types**: The block types to detects.
- **parameters**: A list of parameters that the tool accepts.

Here is a basic example of defining a custom tool:

.. code-block:: python
import random
from gptme.tools import ToolSpec, Parameter, ToolUse
from gptme.message import Message
def execute(code, args, kwargs, confirm):
if code is None and kwargs is not None:
code = kwargs.get('side_count')
yield Message('system', f"Result: {random.randint(1,code)}")
def examples(tool_format):
return f"""
> User: Throw a dice and give me the result.
> Assistant:
{ToolUse("dice", [], "6").to_output(tool_format)}
> System: 3
> assistant: The result is 3
""".strip()
tool = ToolSpec(
name="dice",
desc="A dice simulator.",
instructions="This tool generate a random integer value like a dice.",
examples=examples,
execute=execute,
block_types=["dice"],
parameters=[
Parameter(
name="side_count",
type="integer",
description="The number of faces of the dice to throw.",
required=True,
),
],
)
Registering the Tool
---------------------
To ensure your tool is available for use, you can specify the module in the `TOOL_MODULES` env variable or
setting in your :doc:`project configuration file <config>`, which will automatically load your custom tools.

.. code-block:: toml
TOOL_MODULES = "gptme.tools,path.to.your.custom_tool_module"
Don't remove the `gptme.tools` package unless you know exactly what you are doing.

Ensure your module is in the Python path by either installing it (e.g., with `pip install .`) or
by temporarily modifying the `PYTHONPATH` environment variable. For example:

.. code-block:: bash
export PYTHONPATH=$PYTHONPATH:/path/to/your/module
This lets Python locate your module during development and testing without requiring installation.
1 change: 1 addition & 0 deletions docs/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ See the `README <https://github.com/ErikBjare/gptme/blob/master/README.md>`_ fil
evals
bot
finetuning
custom_tool
arewetiny
timeline
alternatives
Expand Down
28 changes: 20 additions & 8 deletions gptme/chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,13 @@
import re
import sys
import termios
from typing import cast
import urllib.parse
from collections.abc import Generator
from pathlib import Path

from .commands import action_descriptions, execute_cmd
from .config import get_config
from .constants import PROMPT_USER
from .init import init
from .llm import reply
Expand All @@ -19,11 +21,12 @@
from .tools import (
ToolFormat,
ToolUse,
execute_msg,
has_tool,
loaded_tools,
get_tools,
execute_msg,
ConfirmFunc,
set_tool_format,
)
from .tools.base import ConfirmFunc
from .tools.browser import read_url
from .util import console, path_with_tilde, print_bell
from .util.ask_execute import ask_execute
Expand All @@ -46,7 +49,7 @@ def chat(
show_hidden: bool = False,
workspace: Path | None = None,
tool_allowlist: list[str] | None = None,
tool_format: ToolFormat = "markdown",
tool_format: ToolFormat | None = None,
) -> None:
"""
Run the chat loop.
Expand All @@ -71,6 +74,15 @@ def chat(
console.log(f"Using logdir {path_with_tilde(logdir)}")
manager = LogManager.load(logdir, initial_msgs=initial_msgs, create=True)

config = get_config()
tool_format_with_default: ToolFormat = tool_format or cast(
ToolFormat, config.get_env("TOOL_FORMAT", "markdown")
)

# By defining the tool_format at the last moment we ensure we can use the
# configuration for subagent
set_tool_format(tool_format_with_default)

# change to workspace directory
# use if exists, create if @log, or use given path
log_workspace = logdir / "workspace"
Expand Down Expand Up @@ -128,7 +140,7 @@ def confirm_func(msg) -> bool:
manager.log,
stream,
confirm_func,
tool_format,
tool_format_with_default,
workspace,
)
)
Expand Down Expand Up @@ -177,7 +189,7 @@ def confirm_func(msg) -> bool:
# ask for input if no prompt, generate reply, and run tools
clear_interruptible() # Ensure we're not interruptible during user input
for msg in step(
manager.log, stream, confirm_func, tool_format, workspace
manager.log, stream, confirm_func, tool_format_with_default, workspace
): # pragma: no cover
manager.append(msg)
# run any user-commands, if msg is from user
Expand All @@ -189,7 +201,7 @@ def step(
log: Log | list[Message],
stream: bool,
confirm: ConfirmFunc,
tool_format: ToolFormat = "markdown",
tool_format: ToolFormat,
workspace: Path | None = None,
) -> Generator[Message, None, None]:
"""Runs a single pass of the chat."""
Expand Down Expand Up @@ -222,7 +234,7 @@ def step(

tools = None
if tool_format == "tool":
tools = [t for t in loaded_tools if t.is_runnable()]
tools = [t for t in get_tools() if t.is_runnable()]

# generate response
msg_response = reply(msgs, get_model().model, stream, tools)
Expand Down
12 changes: 3 additions & 9 deletions gptme/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,9 @@
import click
from pick import pick

from gptme.config import get_config

from .chat import chat
from .config import get_config
from .commands import _gen_help
from .constants import MULTIPROMPT_SEPARATOR
from .dirs import get_logs_dir
Expand All @@ -22,12 +22,7 @@
from .logmanager import ConversationMeta, get_user_conversations
from .message import Message
from .prompts import get_prompt
from .tools import (
ToolFormat,
ToolSpec,
init_tools,
set_tool_format,
)
from .tools import ToolFormat, init_tools, get_available_tools
from .util import epoch_to_age
from .util.generate_name import generate_name
from .util.interrupt import handle_keyboard_interrupt, set_interruptible
Expand All @@ -39,7 +34,7 @@
script_path = Path(os.path.realpath(__file__))
commands_help = "\n".join(_gen_help(incl_langtags=False))
available_tool_names = ", ".join(
sorted([tool.name for tool in ToolSpec.get_tools().values() if tool.available])
sorted([tool.name for tool in get_available_tools() if tool.available])
)


Expand Down Expand Up @@ -190,7 +185,6 @@ def main(
selected_tool_format: ToolFormat = (
tool_format or config.get_env("TOOL_FORMAT") or "markdown" # type: ignore
)
set_tool_format(selected_tool_format)

# early init tools to generate system prompt
init_tools(frozenset(tool_allowlist) if tool_allowlist else None)
Expand Down
7 changes: 3 additions & 4 deletions gptme/commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,7 @@
print_msg,
toml_to_msgs,
)
from .tools import ToolUse, execute_msg, loaded_tools
from .tools.base import ConfirmFunc, get_tool_format
from .tools import ToolUse, execute_msg, get_tools, ConfirmFunc, get_tool_format
from .util.export import export_chat_to_html
from .util.useredit import edit_text_with_editor
from .util.cost import log_costs
Expand Down Expand Up @@ -138,7 +137,7 @@ def handle_cmd(
case "tools":
manager.undo(1, quiet=True)
print("Available tools:")
for tool in loaded_tools:
for tool in get_tools():
print(
f"""
# {tool.name}
Expand Down Expand Up @@ -220,7 +219,7 @@ def _gen_help(incl_langtags: bool = True) -> Generator[str, None, None]:
yield " /python print('hello')"
yield ""
yield "Supported langtags:"
for tool in loaded_tools:
for tool in get_tools():
if tool.block_types:
yield f" - {tool.block_types[0]}" + (
f" (alias: {', '.join(tool.block_types[1:])})"
Expand Down
4 changes: 2 additions & 2 deletions gptme/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,11 @@ class Config:
env: dict

def get_env(self, key: str, default: str | None = None) -> str | None:
"""Gets an enviromnent variable, checks the config file if it's not set in the environment."""
"""Gets an environment variable, checks the config file if it's not set in the environment."""
return os.environ.get(key) or self.env.get(key) or default

def get_env_required(self, key: str) -> str:
"""Gets an enviromnent variable, checks the config file if it's not set in the environment."""
"""Gets an environment variable, checks the config file if it's not set in the environment."""
if val := os.environ.get(key) or self.env.get(key):
return val
raise KeyError( # pragma: no cover
Expand Down
2 changes: 1 addition & 1 deletion gptme/llm/llm_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from ..config import Config
from ..constants import TEMPERATURE, TOP_P
from ..message import Message, msgs2dicts
from ..tools.base import Parameter, ToolSpec, ToolUse
from ..tools import Parameter, ToolSpec, ToolUse
from .models import ModelMeta, Provider, get_model

if TYPE_CHECKING:
Expand Down
6 changes: 3 additions & 3 deletions gptme/prompts.py
Original file line number Diff line number Diff line change
Expand Up @@ -199,14 +199,14 @@ def prompt_tools(
examples: bool = True, tool_format: ToolFormat = "markdown"
) -> Generator[Message, None, None]:
"""Generate the tools overview prompt."""
from .tools import loaded_tools # fmt: skip
from .tools import get_tools # fmt: skip

assert loaded_tools, "No tools loaded"
assert get_tools(), "No tools loaded"

use_tool = tool_format == "tool"

prompt = "# Tools aliases" if use_tool else "# Tools Overview"
for tool in loaded_tools:
for tool in get_tools():
if not use_tool or not tool.is_runnable():
prompt += tool.get_tool_prompt(examples, tool_format)

Expand Down
3 changes: 1 addition & 2 deletions gptme/server/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,7 @@
from ..llm.models import get_model
from ..logmanager import LogManager, get_user_conversations, prepare_messages
from ..message import Message
from ..tools import execute_msg
from ..tools.base import ToolUse
from ..tools import ToolUse, execute_msg

logger = logging.getLogger(__name__)

Expand Down
Loading

0 comments on commit 860157b

Please sign in to comment.