Moved iter_prompt from Response to Model, moved a lot of other stuff

- Moved a whole bunch of things from llm/cli.py into llm/__init__.py
- Switched plugin listings to use importlib.metadata to avoid deprecation warning
- iter_prompt() is now a method on Model, not on Response

llm/__init__.py

@@ -1,9 +1,25 @@
 from .hookspecs import hookimpl
-from .models import LogMessage, Model, Options, Prompt, Response, OptionsError
+from .models import (
+    LogMessage,
+    Model,
+    ModelWithAliases,
+    Options,
+    Prompt,
+    Response,
+    OptionsError,
+)
 from .templates import Template
+from .plugins import pm
+import click
+from typing import Dict, List
+import json
+import os
+import pathlib
 
 __all__ = [
     "hookimpl",
+    "get_model",
+    "get_key",
     "LogMessage",
     "Model",
     "Options",
@@ -12,3 +28,76 @@
     "Response",
     "Template",
 ]
+
+
+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
+
+
+def get_models_with_aliases() -> List["ModelWithAliases"]:
+    model_aliases = []
+
+    def register(model, aliases=None):
+        model_aliases.append(ModelWithAliases(model, aliases or set()))
+
+    pm.hook.register_models(register=register)
+    return model_aliases
+
+
+def get_model_aliases() -> Dict[str, Model]:
+    model_aliases = {}
+    for model_with_aliases in get_models_with_aliases():
+        for alias in model_with_aliases.aliases:
+            model_aliases[alias] = model_with_aliases.model
+        model_aliases[model_with_aliases.model.model_id] = model_with_aliases.model
+    return model_aliases
+
+
+class UnknownModelError(KeyError):
+    pass
+
+
+def get_model(name):
+    aliases = get_model_aliases()
+    try:
+        return aliases[name]
+    except KeyError:
+        raise UnknownModelError(name)
+
+
+def get_key(key_arg, default_key, env_var=None):
+    keys = load_keys()
+    if key_arg in keys:
+        return keys[key_arg]
+    if key_arg:
+        return key_arg
+    if env_var and os.environ.get(env_var):
+        return os.environ[env_var]
+    return keys.get(default_key)
+
+
+def load_keys():
+    path = user_dir() / "keys.json"
+    if path.exists():
+        return json.loads(path.read_text())
+    else:
+        return {}
+
+
+def user_dir():
+    llm_user_path = os.environ.get("LLM_USER_PATH")
+    if llm_user_path:
+        return pathlib.Path(llm_user_path)
+    return pathlib.Path(click.get_app_dir("io.datasette.llm"))

llm/cli.py

@@ -1,16 +1,18 @@
 import click
 from click_default_group import DefaultGroup
 import json
-from llm import Template
-from .migrations import migrate
-from .plugins import (
-    pm,
+from llm import (
+    Template,
+    get_key,
     get_plugins,
     get_model,
     get_model_aliases,
     get_models_with_aliases,
+    user_dir,
 )
-import os
+
+from .migrations import migrate
+from .plugins import pm
 import pathlib
 import pydantic
 from runpy import run_module
@@ -313,11 +315,7 @@ def keys():
 @keys.command(name="path")
 def keys_path_command():
     "Output the path to the keys.json file"
-    click.echo(keys_path())
-
-
-def keys_path():
-    return user_dir() / "keys.json"
+    click.echo(user_dir() / "keys.json")
 
 
 @keys.command(name="set")
@@ -334,7 +332,7 @@ def set_(name, value):
         Enter key: ...
     """
     default = {"// Note": "This file stores secret API credentials. Do not share!"}
-    path = pathlib.Path(keys_path())
+    path = user_dir() / "keys.json"
     path.parent.mkdir(parents=True, exist_ok=True)
     if not path.exists():
         path.write_text(json.dumps(default))
@@ -548,32 +546,6 @@ def _truncate_string(s, max_length=100):
     return s
 
 
-def get_key(key_arg, default_key, env_var=None):
-    keys = load_keys()
-    if key_arg in keys:
-        return keys[key_arg]
-    if key_arg:
-        return key_arg
-    if env_var and os.environ.get(env_var):
-        return os.environ[env_var]
-    return keys.get(default_key)
-
-
-def load_keys():
-    path = pathlib.Path(keys_path())
-    if path.exists():
-        return json.loads(path.read_text())
-    else:
-        return {}
-
-
-def user_dir():
-    llm_user_path = os.environ.get("LLM_USER_PATH")
-    if llm_user_path:
-        return pathlib.Path(llm_user_path)
-    return pathlib.Path(click.get_app_dir("io.datasette.llm"))
-
-
 def get_default_model():
     path = user_dir() / "default_model.txt"
     if path.exists():

llm/default_plugins/openai_models.py

@@ -105,43 +105,43 @@ def __init__(self, prompt, model, stream, key):
             super().__init__(prompt, model, stream)
             self.key = key
 
-        def iter_prompt(self, prompt):
-            messages = []
-            if prompt.system:
-                messages.append({"role": "system", "content": prompt.system})
-            messages.append({"role": "user", "content": prompt.prompt})
-            openai.api_key = self.key
-            self._prompt_json = {"messages": messages}
-            if self.stream:
-                completion = openai.ChatCompletion.create(
-                    model=prompt.model.model_id,
-                    messages=messages,
-                    stream=True,
-                    **not_nulls(prompt.options),
-                )
-                chunks = []
-                for chunk in completion:
-                    chunks.append(chunk)
-                    content = chunk["choices"][0].get("delta", {}).get("content")
-                    if content is not None:
-                        yield content
-                self._response_json = combine_chunks(chunks)
-            else:
-                response = openai.ChatCompletion.create(
-                    model=prompt.model.model_id,
-                    messages=messages,
-                    stream=False,
-                )
-                self._response_json = response.to_dict_recursive()
-                yield response.choices[0].message.content
-
     def __init__(self, model_id, key=None):
         self.model_id = model_id
         self.key = key
 
     def __str__(self):
         return "OpenAI Chat: {}".format(self.model_id)
 
+    def iter_prompt(self, prompt, stream, response):
+        messages = []
+        if prompt.system:
+            messages.append({"role": "system", "content": prompt.system})
+        messages.append({"role": "user", "content": prompt.prompt})
+        openai.api_key = self.key
+        response._prompt_json = {"messages": messages}
+        if stream:
+            completion = openai.ChatCompletion.create(
+                model=prompt.model.model_id,
+                messages=messages,
+                stream=True,
+                **not_nulls(prompt.options),
+            )
+            chunks = []
+            for chunk in completion:
+                chunks.append(chunk)
+                content = chunk["choices"][0].get("delta", {}).get("content")
+                if content is not None:
+                    yield content
+            response._response_json = combine_chunks(chunks)
+        else:
+            completion = openai.ChatCompletion.create(
+                model=prompt.model.model_id,
+                messages=messages,
+                stream=False,
+            )
+            response._response_json = completion.to_dict_recursive()
+            yield completion.choices[0].message.content
+
 
 def not_nulls(data) -> dict:
     return {key: value for key, value in data if value is not None}

llm/models.py

@@ -2,7 +2,7 @@
 import datetime
 from .errors import NeedsKeyException
 import time
-from typing import cast, Any, Callable, Dict, Iterator, List, Optional, Set
+from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Union
 from abc import ABC, abstractmethod
 import os
 from pydantic import ConfigDict, BaseModel
@@ -70,17 +70,14 @@ def __iter__(self) -> Iterator[str]:
         self._start_utcnow = datetime.datetime.utcnow()
         if self._done:
             return self._chunks
-        for chunk in self.iter_prompt(self.prompt):
+        for chunk in self.model.iter_prompt(
+            self.prompt, stream=self.stream, response=self
+        ):
             yield chunk
             self._chunks.append(chunk)
         self._end = time.monotonic()
         self._done = True
 
-    @abstractmethod
-    def iter_prompt(self, prompt: Prompt) -> Iterator[str]:
-        "Execute prompt and yield chunks of text, or yield a single big chunk"
-        pass
-
     def _force(self):
         if not self._done:
             list(self)
@@ -160,6 +157,16 @@ def get_key(self):
             message += " or set the {} environment variable".format(self.key_env_var)
         raise NeedsKeyException(message)
 
+    @abstractmethod
+    def iter_prompt(
+        self, prompt: Prompt, stream: bool, response: Response
+    ) -> Iterator[str]:
+        """
+        Execute a prompt and yield chunks of text, or yield a single big chunk.
+        Any additional useful information about the execution should be assigned to the response.
+        """
+        pass
+
     def prompt(
         self,
         prompt: Optional[str],
@@ -172,12 +179,36 @@ def prompt(
             stream=stream,
         )
 
+    def chain(
+        self,
+        prompt: Union[Prompt, str],
+        system: Optional[str] = None,
+        stream: bool = True,
+        proceed: Optional[Callable] = None,
+        **options
+    ):
+        if proceed is None:
+
+            def proceed(response):
+                return None
+
+        while True:
+            if isinstance(prompt, str):
+                prompt = Prompt(
+                    prompt, model=self, system=system, options=self.Options(**options)
+                )
+            response = self.execute(
+                prompt,
+                stream=stream,
+            )
+            yield response
+            next_prompt = proceed(response)
+            if not next_prompt:
+                break
+            prompt = next_prompt
+
     def execute(self, prompt: Prompt, stream: bool = True) -> Response:
-        r = cast(Callable, getattr(self, "Response"))
-        kwargs = {}
-        if self.needs_key:
-            kwargs["key"] = self.get_key()
-        return r(prompt, self, stream, **kwargs)
+        return Response(prompt, self, stream)
 
     def __str__(self) -> str:
         return "{}: {}".format(self.__class__.__name__, self.model_id)