Initial plugin framework and register_commands(cli) hook, refs #49

This commit is contained in:
Simon Willison 2023-06-17 17:42:13 +01:00
parent 88591998e4
commit a396950f79
8 changed files with 150 additions and 0 deletions

57
docs/plugins.md Normal file
View file

@ -0,0 +1,57 @@
# Plugins
LLM plugins can provide extra features to the tool.
## Installing plugins
Plugins can be installed by running `pip install` in the same virtual environment as `llm` itself:
```bash
pip install llm-hello-world
```
The [llm-hello-world](https://github.com/simonw/llm-hello-world) plugin is the current best example of how to build and package a plugin.
## Listing installed plugins
Run `llm plugins` to list installed plugins:
```bash
llm plugins
```
```json
[
{
"name": "llm-hello-world",
"hooks": [
"register_commands"
],
"version": "0.1"
}
]
```
## Plugin hooks
Plugins use **plugin hooks** to customize LLM's behavior. These hooks are powered by the [Pluggy plugin system](https://pluggy.readthedocs.io/).
Each plugin can implement one or more hooks using the @hookimpl decorator against one of the hook function names described on this page.
LLM imitates the Datasette plugin system. The [Datasette plugin documentation](https://docs.datasette.io/en/stable/writing_plugins.html) describes how plugins work.
### register_commands(cli)
This hook adds new commands to the `llm` CLI tool - for example `llm extra-command`.
This example plugin adds a new `hello-world` command that prints "Hello world!":
```python
from llm import hookimpl
import click
@hookimpl
def register_commands(cli):
@cli.command(name="hello-world")
def hello_world():
"Print hello world"
click.echo("Hello world!")
```
This new command will be added to `llm --help` and can be run using `llm hello-world`.

View file

@ -1,6 +1,8 @@
from pydantic import BaseModel
import string
from typing import Optional
from .hookspecs import hookimpl # noqa
from .hookspecs import hookspec # noqa
class Template(BaseModel):

View file

@ -4,6 +4,7 @@ import datetime
import json
from llm import Template
from .migrations import migrate
from .plugins import pm, get_plugins
import openai
import os
import pathlib
@ -307,6 +308,12 @@ def templates_list():
click.echo(display_truncated(text))
@cli.command(name="plugins")
def plugins_list():
"List installed plugins"
click.echo(json.dumps(get_plugins(), indent=2))
def display_truncated(text):
console_width = shutil.get_terminal_size()[0]
if len(text) > console_width:
@ -468,3 +475,6 @@ def get_history(chat_id):
"id = ? or chat_id = ?", [chat_id, chat_id], order_by="id"
)
return chat_id, rows
pm.hook.register_commands(cli=cli)

10
llm/hookspecs.py Normal file
View file

@ -0,0 +1,10 @@
from pluggy import HookimplMarker
from pluggy import HookspecMarker
hookspec = HookspecMarker("llm")
hookimpl = HookimplMarker("llm")
@hookspec
def register_commands(cli):
"""Register additional CLI commands, e.g. 'llm mycommand ...'"""

26
llm/plugins.py Normal file
View file

@ -0,0 +1,26 @@
import pluggy
import sys
from . import hookspecs
pm = pluggy.PluginManager("llm")
pm.add_hookspecs(hookspecs)
if not hasattr(sys, "_called_from_test"):
# Only load plugins if not running tests
pm.load_setuptools_entrypoints("llm")
def get_plugins():
plugins = []
plugin_to_distinfo = dict(pm.list_plugin_distinfo())
for plugin in pm.get_plugins():
plugin_info = {
"name": plugin.__name__,
"hooks": [h.name for h in pm.get_hookcallers(plugin)],
}
distinfo = plugin_to_distinfo.get(plugin)
if distinfo:
plugin_info["version"] = distinfo.version
plugin_info["name"] = distinfo.project_name
plugins.append(plugin_info)
return plugins

View file

@ -39,6 +39,7 @@ setup(
"sqlite-utils",
"pydantic",
"PyYAML",
"pluggy",
],
extras_require={"test": ["pytest", "requests-mock"]},
python_requires=">=3.7",

View file

@ -1,6 +1,12 @@
import pytest
def pytest_configure(config):
import sys
sys._called_from_test = True
@pytest.fixture
def log_path(tmpdir):
return tmpdir / "logs.db"

38
tests/test_plugins.py Normal file
View file

@ -0,0 +1,38 @@
from click.testing import CliRunner
import click
import importlib
from llm import cli, hookimpl, plugins
import pytest
def test_register_commands():
importlib.reload(cli)
assert plugins.get_plugins() == []
class HelloWorldPlugin:
__name__ = "HelloWorldPlugin"
@hookimpl
def register_commands(self, cli):
@cli.command(name="hello-world")
def hello_world():
"Print hello world"
click.echo("Hello world!")
try:
plugins.pm.register(HelloWorldPlugin(), name="HelloWorldPlugin")
importlib.reload(cli)
assert plugins.get_plugins() == [
{"name": "HelloWorldPlugin", "hooks": ["register_commands"]}
]
runner = CliRunner()
result = runner.invoke(cli.cli, ["hello-world"])
assert result.exit_code == 0
assert result.output == "Hello world!\n"
finally:
plugins.pm.unregister(name="HelloWorldPlugin")
importlib.reload(cli)
assert plugins.get_plugins() == []