From a396950f7934e82a9968703bb3ce9ab7ab62f7f8 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sat, 17 Jun 2023 17:42:13 +0100
Subject: [PATCH] Initial plugin framework and register_commands(cli) hook,
 refs #49

---
 docs/plugins.md       | 57 +++++++++++++++++++++++++++++++++++++++++++
 llm/__init__.py       |  2 ++
 llm/cli.py            | 10 ++++++++
 llm/hookspecs.py      | 10 ++++++++
 llm/plugins.py        | 26 ++++++++++++++++++++
 setup.py              |  1 +
 tests/conftest.py     |  6 +++++
 tests/test_plugins.py | 38 +++++++++++++++++++++++++++++
 8 files changed, 150 insertions(+)
 create mode 100644 docs/plugins.md
 create mode 100644 llm/hookspecs.py
 create mode 100644 llm/plugins.py
 create mode 100644 tests/test_plugins.py

diff --git a/docs/plugins.md b/docs/plugins.md
new file mode 100644
index 00000000..e9dfebba
--- /dev/null
+++ b/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`.
\ No newline at end of file
diff --git a/llm/__init__.py b/llm/__init__.py
index 194d9cf4..944a3d96 100644
--- a/llm/__init__.py
+++ b/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):
diff --git a/llm/cli.py b/llm/cli.py
index 485fde49..cde5a4d7 100644
--- a/llm/cli.py
+++ b/llm/cli.py
@@ -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
@@ -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)
diff --git a/llm/hookspecs.py b/llm/hookspecs.py
new file mode 100644
index 00000000..88422de7
--- /dev/null
+++ b/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 ...'"""
diff --git a/llm/plugins.py b/llm/plugins.py
new file mode 100644
index 00000000..97e3b825
--- /dev/null
+++ b/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
diff --git a/setup.py b/setup.py
index 0508b730..b3cdee67 100644
--- a/setup.py
+++ b/setup.py
@@ -39,6 +39,7 @@ def get_long_description():
         "sqlite-utils",
         "pydantic",
         "PyYAML",
+        "pluggy",
     ],
     extras_require={"test": ["pytest", "requests-mock"]},
     python_requires=">=3.7",
diff --git a/tests/conftest.py b/tests/conftest.py
index 844c872e..78d3afa5 100644
--- a/tests/conftest.py
+++ b/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"
diff --git a/tests/test_plugins.py b/tests/test_plugins.py
new file mode 100644
index 00000000..029608c8
--- /dev/null
+++ b/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() == []