Skip to content

Commit

Permalink
Initial plugin framework and register_commands(cli) hook, refs #49
Browse files Browse the repository at this point in the history
  • Loading branch information
simonw committed Jun 17, 2023
1 parent 8859199 commit a396950
Show file tree
Hide file tree
Showing 8 changed files with 150 additions and 0 deletions.
57 changes: 57 additions & 0 deletions docs/plugins.md
@@ -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`.
2 changes: 2 additions & 0 deletions llm/__init__.py
@@ -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):
Expand Down
10 changes: 10 additions & 0 deletions llm/cli.py
Expand Up @@ -4,6 +4,7 @@
import json
from llm import Template
from .migrations import migrate
from .plugins import pm, get_plugins
import openai
import os
import pathlib
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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 changes: 10 additions & 0 deletions llm/hookspecs.py
@@ -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 changes: 26 additions & 0 deletions llm/plugins.py
@@ -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
1 change: 1 addition & 0 deletions setup.py
Expand Up @@ -39,6 +39,7 @@ def get_long_description():
"sqlite-utils",
"pydantic",
"PyYAML",
"pluggy",
],
extras_require={"test": ["pytest", "requests-mock"]},
python_requires=">=3.7",
Expand Down
6 changes: 6 additions & 0 deletions tests/conftest.py
@@ -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"
Expand Down
38 changes: 38 additions & 0 deletions tests/test_plugins.py
@@ -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() == []

0 comments on commit a396950

Please sign in to comment.