Skip to content

Commit 3992661

Browse files
committed
feat: allow to load tools from external modules
1 parent 7abeb1e commit 3992661

23 files changed

+422
-139
lines changed

docs/config.rst

+5
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,11 @@ Here is an example:
4242
#MODEL = "local/<model-name>"
4343
#OPENAI_BASE_URL = "http://localhost:11434/v1"
4444
45+
# Uncomment to change tool configuration
46+
#TOOL_FORMAT = "markdown" # Select the tool formal. One of `markdown`, `xml`, `tool`
47+
#TOOL_ALLOW_LIST = "save,append,patch,python" # Comma separated list of allowed tools
48+
#TOOL_MODULES = "gptme.tools,custom.tools" # List of python comma separated python module path
49+
4550
The ``prompt`` section contains options for the prompt.
4651

4752
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`.

docs/custom_tool.rst

+83
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
Creating a Custom Tool for gptme
2+
=================================
3+
4+
Introduction
5+
------------
6+
In gptme, a custom tool allows you to extend the functionality of the assistant by
7+
defining new tools that can be executed.
8+
This guide will walk you through the process of creating and registering a custom tool.
9+
10+
Creating a Custom Tool
11+
-----------------------
12+
To create a custom tool, you need to define a new instance of the `ToolSpec` class.
13+
This class requires several parameters:
14+
15+
- **name**: The name of the tool.
16+
- **desc**: A description of what the tool does.
17+
- **instructions**: Instructions on how to use the tool.
18+
- **examples**: Example usage of the tool.
19+
- **execute**: A function that defines the tool's behavior when executed.
20+
- **block_types**: The block types to detects.
21+
- **parameters**: A list of parameters that the tool accepts.
22+
23+
Here is a basic example of defining a custom tool:
24+
25+
.. code-block:: python
26+
27+
import random
28+
from gptme.tools import ToolSpec, Parameter, ToolUse
29+
from gptme.message import Message
30+
31+
def execute(code, args, kwargs, confirm):
32+
33+
if code is None and kwargs is not None:
34+
code = kwargs.get('side_count')
35+
36+
yield Message('system', f"Result: {random.randint(1,code)}")
37+
38+
def examples(tool_format):
39+
return f"""
40+
> User: Throw a dice and give me the result.
41+
> Assistant:
42+
{ToolUse("dice", [], "6").to_output(tool_format)}
43+
> System: 3
44+
> assistant: The result is 3
45+
""".strip()
46+
47+
tool = ToolSpec(
48+
name="dice",
49+
desc="A dice simulator.",
50+
instructions="This tool generate a random integer value like a dice.",
51+
examples=examples,
52+
execute=execute,
53+
block_types=["dice"],
54+
parameters=[
55+
Parameter(
56+
name="side_count",
57+
type="integer",
58+
description="The number of faces of the dice to throw.",
59+
required=True,
60+
),
61+
],
62+
)
63+
64+
Registering the Tool
65+
---------------------
66+
To ensure your tool is available for use, you can specify the module in the `TOOL_MODULES` env variable or
67+
setting in your :doc:`project configuration file <config>`, which will automatically load your custom tools.
68+
69+
.. code-block:: toml
70+
71+
TOOL_MODULES = "gptme.tools,path.to.your.custom_tool_module"
72+
73+
Don't remove the `gptme.tools` package unless you know exactly what you are doing.
74+
75+
Ensure your module is in the Python path by either installing it (e.g., with `pip install .`) or
76+
by temporarily modifying the `PYTHONPATH` environment variable. For example:
77+
78+
.. code-block:: bash
79+
80+
export PYTHONPATH=$PYTHONPATH:/path/to/your/module
81+
82+
83+
This lets Python locate your module during development and testing without requiring installation.

docs/index.rst

+1
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ See the `README <https://github.com/ErikBjare/gptme/blob/master/README.md>`_ fil
3737
evals
3838
bot
3939
finetuning
40+
custom_tool
4041
arewetiny
4142
timeline
4243
alternatives

gptme/chat.py

+20-8
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,13 @@
44
import re
55
import sys
66
import termios
7+
from typing import cast
78
import urllib.parse
89
from collections.abc import Generator
910
from pathlib import Path
1011

1112
from .commands import action_descriptions, execute_cmd
13+
from .config import get_config
1214
from .constants import PROMPT_USER
1315
from .init import init
1416
from .llm import reply
@@ -19,11 +21,12 @@
1921
from .tools import (
2022
ToolFormat,
2123
ToolUse,
22-
execute_msg,
2324
has_tool,
24-
loaded_tools,
25+
get_tools,
26+
execute_msg,
27+
ConfirmFunc,
28+
set_tool_format,
2529
)
26-
from .tools.base import ConfirmFunc
2730
from .tools.browser import read_url
2831
from .util import console, path_with_tilde, print_bell
2932
from .util.ask_execute import ask_execute
@@ -46,7 +49,7 @@ def chat(
4649
show_hidden: bool = False,
4750
workspace: Path | None = None,
4851
tool_allowlist: list[str] | None = None,
49-
tool_format: ToolFormat = "markdown",
52+
tool_format: ToolFormat | None = None,
5053
) -> None:
5154
"""
5255
Run the chat loop.
@@ -71,6 +74,15 @@ def chat(
7174
console.log(f"Using logdir {path_with_tilde(logdir)}")
7275
manager = LogManager.load(logdir, initial_msgs=initial_msgs, create=True)
7376

77+
config = get_config()
78+
tool_format_with_default: ToolFormat = tool_format or cast(
79+
ToolFormat, config.get_env("TOOL_FORMAT", "markdown")
80+
)
81+
82+
# By defining the tool_format at the last moment we ensure we can use the
83+
# configuration for subagent
84+
set_tool_format(tool_format_with_default)
85+
7486
# change to workspace directory
7587
# use if exists, create if @log, or use given path
7688
log_workspace = logdir / "workspace"
@@ -128,7 +140,7 @@ def confirm_func(msg) -> bool:
128140
manager.log,
129141
stream,
130142
confirm_func,
131-
tool_format,
143+
tool_format_with_default,
132144
workspace,
133145
)
134146
)
@@ -177,7 +189,7 @@ def confirm_func(msg) -> bool:
177189
# ask for input if no prompt, generate reply, and run tools
178190
clear_interruptible() # Ensure we're not interruptible during user input
179191
for msg in step(
180-
manager.log, stream, confirm_func, tool_format, workspace
192+
manager.log, stream, confirm_func, tool_format_with_default, workspace
181193
): # pragma: no cover
182194
manager.append(msg)
183195
# run any user-commands, if msg is from user
@@ -189,7 +201,7 @@ def step(
189201
log: Log | list[Message],
190202
stream: bool,
191203
confirm: ConfirmFunc,
192-
tool_format: ToolFormat = "markdown",
204+
tool_format: ToolFormat,
193205
workspace: Path | None = None,
194206
) -> Generator[Message, None, None]:
195207
"""Runs a single pass of the chat."""
@@ -222,7 +234,7 @@ def step(
222234

223235
tools = None
224236
if tool_format == "tool":
225-
tools = [t for t in loaded_tools if t.is_runnable()]
237+
tools = [t for t in get_tools() if t.is_runnable()]
226238

227239
# generate response
228240
msg_response = reply(msgs, get_model().model, stream, tools)

gptme/cli.py

+4-10
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,9 @@
1111
import click
1212
from pick import pick
1313

14-
from gptme.config import get_config
1514

1615
from .chat import chat
16+
from .config import get_config
1717
from .commands import _gen_help
1818
from .constants import MULTIPROMPT_SEPARATOR
1919
from .dirs import get_logs_dir
@@ -22,12 +22,7 @@
2222
from .logmanager import ConversationMeta, get_user_conversations
2323
from .message import Message
2424
from .prompts import get_prompt
25-
from .tools import (
26-
ToolFormat,
27-
ToolSpec,
28-
init_tools,
29-
set_tool_format,
30-
)
25+
from .tools import ToolFormat, init_tools, get_available_tools
3126
from .util import epoch_to_age
3227
from .util.generate_name import generate_name
3328
from .util.interrupt import handle_keyboard_interrupt, set_interruptible
@@ -39,7 +34,7 @@
3934
script_path = Path(os.path.realpath(__file__))
4035
commands_help = "\n".join(_gen_help(incl_langtags=False))
4136
available_tool_names = ", ".join(
42-
sorted([tool.name for tool in ToolSpec.get_tools().values() if tool.available])
37+
sorted([tool.name for tool in get_available_tools() if tool.available])
4338
)
4439

4540

@@ -187,8 +182,7 @@ def main(
187182

188183
config = get_config()
189184

190-
tool_format = tool_format or config.get_env("TOOL_FORMAT") or "markdown"
191-
set_tool_format(tool_format)
185+
tool_format = tool_format or config.get_env("TOOL_FORMAT", "markdown")
192186

193187
# early init tools to generate system prompt
194188
init_tools(frozenset(tool_allowlist) if tool_allowlist else None)

gptme/commands.py

+3-4
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,7 @@
1515
print_msg,
1616
toml_to_msgs,
1717
)
18-
from .tools import ToolUse, execute_msg, loaded_tools
19-
from .tools.base import ConfirmFunc, get_tool_format
18+
from .tools import ToolUse, execute_msg, get_tools, ConfirmFunc, get_tool_format
2019
from .util.export import export_chat_to_html
2120
from .util.useredit import edit_text_with_editor
2221
from .util.cost import log_costs
@@ -138,7 +137,7 @@ def handle_cmd(
138137
case "tools":
139138
manager.undo(1, quiet=True)
140139
print("Available tools:")
141-
for tool in loaded_tools:
140+
for tool in get_tools():
142141
print(
143142
f"""
144143
# {tool.name}
@@ -220,7 +219,7 @@ def _gen_help(incl_langtags: bool = True) -> Generator[str, None, None]:
220219
yield " /python print('hello')"
221220
yield ""
222221
yield "Supported langtags:"
223-
for tool in loaded_tools:
222+
for tool in get_tools():
224223
if tool.block_types:
225224
yield f" - {tool.block_types[0]}" + (
226225
f" (alias: {', '.join(tool.block_types[1:])})"

gptme/config.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,11 @@ class Config:
1919
env: dict
2020

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

2525
def get_env_required(self, key: str) -> str:
26-
"""Gets an enviromnent variable, checks the config file if it's not set in the environment."""
26+
"""Gets an environment variable, checks the config file if it's not set in the environment."""
2727
if val := os.environ.get(key) or self.env.get(key):
2828
return val
2929
raise KeyError( # pragma: no cover

gptme/llm/llm_anthropic.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212

1313
from ..constants import TEMPERATURE, TOP_P
1414
from ..message import Message, msgs2dicts
15-
from ..tools.base import Parameter, ToolSpec
15+
from ..tools import Parameter, ToolSpec
1616

1717
if TYPE_CHECKING:
1818
# noreorder

gptme/llm/llm_openai.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from ..config import Config
88
from ..constants import TEMPERATURE, TOP_P
99
from ..message import Message, msgs2dicts
10-
from ..tools.base import Parameter, ToolSpec
10+
from ..tools import Parameter, ToolSpec
1111
from .models import Provider, get_model
1212

1313
if TYPE_CHECKING:

gptme/prompts.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -199,14 +199,14 @@ def prompt_tools(
199199
examples: bool = True, tool_format: ToolFormat = "markdown"
200200
) -> Generator[Message, None, None]:
201201
"""Generate the tools overview prompt."""
202-
from .tools import loaded_tools # fmt: skip
202+
from .tools import get_tools # fmt: skip
203203

204-
assert loaded_tools, "No tools loaded"
204+
assert get_tools(), "No tools loaded"
205205

206206
use_tool = tool_format == "tool"
207207

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

gptme/server/api.py

+1-2
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,7 @@
2424
from ..llm.models import get_model
2525
from ..logmanager import LogManager, get_user_conversations, prepare_messages
2626
from ..message import Message
27-
from ..tools import execute_msg
28-
from ..tools.base import ToolUse
27+
from ..tools import ToolUse, execute_msg
2928

3029
logger = logging.getLogger(__name__)
3130

0 commit comments

Comments
 (0)