llm prompt --schema X option and model.prompt(..., schema=) parameter (#777)

Refs #776

* Implemented new llm prompt --schema and model.prompt(schema=)
* Log schema to responses.schema_id and schemas table
* Include schema in llm logs Markdown output
* Test for schema=pydantic_model
* Initial --schema CLI documentation
* Python docs for schema=
* Advanced plugin docs on schemas

Author: simonw

docs/changelog.md

@@ -411,7 +411,7 @@ There's also a new {ref}`llm.Collection <embeddings-python-collections>` class f
 - The output format for `llm logs` has changed. Previously it was JSON - it's now a much more readable Markdown format suitable for pasting into other documents. [#160](https://github.com/simonw/llm/issues/160)
   - The new `llm logs --json` option can be used to get the old JSON format.
   - Pass `llm logs --conversation ID` or `--cid ID` to see the full logs for a specific conversation.
-- You can now combine piped input and a prompt in a single command: `cat script.py | llm 'explain this code'`. This works even for models that do not support {ref}`system prompts <system-prompts>`. [#153](https://github.com/simonw/llm/issues/153)
+- You can now combine piped input and a prompt in a single command: `cat script.py | llm 'explain this code'`. This works even for models that do not support {ref}`system prompts <usage-system-prompts>`. [#153](https://github.com/simonw/llm/issues/153)
 - Additional {ref}`openai-compatible-models` can now be configured with custom HTTP headers. This enables platforms such as [openrouter.ai](https://openrouter.ai/) to be used with LLM, which can provide Claude access even without an Anthropic API key.
 - Keys set in `keys.json` are now used in preference to environment variables. [#158](https://github.com/simonw/llm/issues/158)
 - The documentation now includes a {ref}`plugin directory <plugin-directory>` listing all available plugins for LLM. [#173](https://github.com/simonw/llm/issues/173)

docs/help.md

@@ -117,6 +117,7 @@ Options:
   --at, --attachment-type <TEXT TEXT>...
                                   Attachment with explicit mimetype
   -o, --option <TEXT TEXT>...     key/value options for the model
+  --schema TEXT                   JSON schema to use for output
   -t, --template TEXT             Template to use
   -p, --param <TEXT TEXT>...      Parameters for template
   --no-stream                     Do not stream output

docs/logging.md

@@ -157,7 +157,7 @@ Example output:
 (logs-conversation)=
 ### Logs for a conversation
 
-To view the logs for the most recent {ref}`conversation <conversation>` you have had with a model, use `-c`:
+To view the logs for the most recent {ref}`conversation <usage-conversation>` you have had with a model, use `-c`:
 
 ```bash
 llm logs -c
@@ -209,7 +209,7 @@ def cleanup_sql(sql):
     return first_line + '(\n  ' + ',\n  '.join(columns) + '\n);'
 
 cog.out("```sql\n")
-for table in ("conversations", "responses", "responses_fts", "attachments", "prompt_attachments"):
+for table in ("conversations", "schemas", "responses", "responses_fts", "attachments", "prompt_attachments"):
     schema = db[table].schema
     cog.out(format(cleanup_sql(schema)))
     cog.out("\n")
@@ -221,7 +221,11 @@ CREATE TABLE [conversations] (
   [name] TEXT,
   [model] TEXT
 );
-CREATE TABLE [responses] (
+CREATE TABLE [schemas] (
+  [id] TEXT PRIMARY KEY,
+  [content] TEXT
+);
+CREATE TABLE "responses" (
   [id] TEXT PRIMARY KEY,
   [model] TEXT,
   [prompt] TEXT,
@@ -235,7 +239,8 @@ CREATE TABLE [responses] (
   [datetime_utc] TEXT,
   [input_tokens] INTEGER,
   [output_tokens] INTEGER,
-  [token_details] TEXT
+  [token_details] TEXT,
+  [schema_id] TEXT REFERENCES [schemas]([id])
 );
 CREATE VIRTUAL TABLE [responses_fts] USING FTS5 (
   [prompt],

docs/plugins/advanced-model-plugins.md

@@ -90,6 +90,21 @@ def register_models(register):
     )
 ```
 
+(advanced-model-plugins-schemas)=
+
+## Supporting schemas
+
+If your model supports {ref}`structured output <python-api-schemas>` against a defined JSON schema you can implement support by first adding `supports_schema = True` to the class:
+
+```python
+class MyModel(llm.KeyModel):
+    ...
+    support_schema = True
+```
+And then adding code to your `.execute()` method that checks for `prompt.schema` and, if it is present, uses that to prompt the model. `prompt.schema` will always be a Python dictionary, even if the user passed in a Pydantic model class.
+
+Check the [llm-gemini](https://github.com/simonw/llm-gemini) and [llm-anthropic](https://github.com/simonw/llm-anthropic) plugins for example of this pattern in action.
+
 (advanced-model-plugins-attachments)=
 
 ## Attachments for multi-modal models

docs/python-api.md

@@ -83,6 +83,41 @@ if "image/jpeg" in model.attachment_types:
     ...
 ```
 
+(python-api-schemas)=
+
+### Schemas
+
+As with {ref}`the CLI tool <usage-schemas>` some models support passing a JSON schema should be used for the resulting response.
+
+You can pass this to the `prompt(schema=)` parameter as either a Python dictionary or a [Pydantic](https://docs.pydantic.dev/) `BaseModel` subclass:
+
+```python
+import llm, json
+from pydantic import BaseModel
+
+class Dog(BaseModel):
+    name: str
+    age: int
+
+model = llm.get_model("gpt-4o-mini")
+response = model.prompt("Describe a nice dog", schema=Dog)
+dog = json.loads(response.text())
+print(dog)
+# {"name":"Buddy","age":3}
+```
+You can also pass a schema directly, like this:
+```python
+response = model.prompt("Describe a nice dog", schema={
+    "properties": {
+        "name": {"title": "Name", "type": "string"},
+        "age": {"title": "Age", "type": "integer"},
+    },
+    "required": ["name", "age"],
+    "title": "Dog",
+    "type": "object",
+})
+```
+
 (python-api-model-options)=
 
 ### Model options

docs/usage.md

@@ -38,7 +38,7 @@ Will run a prompt of:
 ```
 <contents of myscript.py> explain this code
 ```
-For models that support them, {ref}`system prompts <system-prompts>` are a better tool for this kind of prompting.
+For models that support them, {ref}`system prompts <usage-system-prompts>` are a better tool for this kind of prompting.
 
 Some models support options. You can pass these using `-o/--option name value` - for example, to set the temperature to 1.5 run this:
 
@@ -88,7 +88,7 @@ LLM will attempt to automatically detect the content type of the image. If this
 cat myfile | llm "describe this image" --at - image/jpeg
 ```
 
-(system-prompts)=
+(usage-system-prompts)=
 ### System prompts
 
 You can use `-s/--system '...'` to set a system prompt.
@@ -122,7 +122,41 @@ cat llm/utils.py | llm -t pytest
 ```
 See {ref}`prompt templates <prompt-templates>` for more.
 
-(conversation)=
+(usage-schemas)=
+### Schemas
+
+Some models include the ability to return JSON that matches a provided [JSON schema](https://json-schema.org/). Models from OpenAI, Anthropic and Google Gemini all include this capability.
+
+LLM has alpha functionality for specifying a schema to use for the response to a prompt.
+
+Create the schema as a JSON string, then pass that to the `--schema` option. For example:
+
+```bash
+llm --schema '{
+  "type": "object",
+  "properties": {
+    "dogs": {
+      "type": "array",
+      "items": {
+        "type": "object",
+        "properties": {
+          "name": {
+            "type": "string"
+          },
+          "bio": {
+            "type": "string"
+          }
+        }
+      }
+    }
+  }
+}' -m gpt-4o-mini 'invent two dogs'
+```
+The JSON returned from the model should match that schema.
+
+Be warned that different models may support different dialects of the JSON schema specification.
+
+(usage-conversation)=
 ### Continuing a conversation
 
 By default, the tool will start a new conversation each time you run it.

llm/cli.py

@@ -114,16 +114,19 @@ def attachment_types_callback(ctx, param, values):
     return collected
 
 
-def _validate_metadata_json(ctx, param, value):
-    if value is None:
-        return value
-    try:
-        obj = json.loads(value)
-        if not isinstance(obj, dict):
-            raise click.BadParameter("Metadata must be a JSON object")
-        return obj
-    except json.JSONDecodeError:
-        raise click.BadParameter("Metadata must be valid JSON")
+def json_validator(object_name):
+    def validator(ctx, param, value):
+        if value is None:
+            return value
+        try:
+            obj = json.loads(value)
+            if not isinstance(obj, dict):
+                raise click.BadParameter(f"{object_name} must be a JSON object")
+            return obj
+        except json.JSONDecodeError:
+            raise click.BadParameter(f"{object_name} must be valid JSON")
+
+    return validator
 
 
 @click.group(
@@ -184,6 +187,9 @@ def cli():
     multiple=True,
     help="key/value options for the model",
 )
+@click.option(
+    "--schema", callback=json_validator("schema"), help="JSON schema to use for output"
+)
 @click.option("-t", "--template", help="Template to use")
 @click.option(
     "-p",
@@ -228,6 +234,7 @@ def prompt(
     attachments,
     attachment_types,
     options,
+    schema,
     template,
     param,
     no_stream,
@@ -429,6 +436,7 @@ async def inner():
                         prompt,
                         attachments=resolved_attachments,
                         system=system,
+                        schema=schema,
                         **kwargs,
                     )
                     async for chunk in response:
@@ -440,6 +448,7 @@ async def inner():
                         prompt,
                         attachments=resolved_attachments,
                         system=system,
+                        schema=schema,
                         **kwargs,
                     )
                     text = await response.text()
@@ -456,6 +465,7 @@ async def inner():
                 prompt,
                 attachments=resolved_attachments,
                 system=system,
+                schema=schema,
                 **kwargs,
             )
             if should_stream:
@@ -829,13 +839,15 @@ def logs_turn_off():
     responses.output_tokens,
     responses.token_details,
     conversations.name as conversation_name,
-    conversations.model as conversation_model"""
+    conversations.model as conversation_model,
+    schemas.content as schema_json"""
 
 LOGS_SQL = """
 select
 {columns}
 from
     responses
+left join schemas on responses.schema_id = schemas.id
 left join conversations on responses.conversation_id = conversations.id{extra_where}
 order by responses.id desc{limit}
 """
@@ -844,6 +856,7 @@ def logs_turn_off():
 {columns}
 from
     responses
+left join schemas on responses.schema_id = schemas.id
 left join conversations on responses.conversation_id = conversations.id
 join responses_fts on responses_fts.rowid = responses.rowid
 where responses_fts match :query{extra_where}
@@ -1117,6 +1130,12 @@ def logs_list(
                 if row["system"] is not None:
                     click.echo("\n## System:\n\n{}".format(row["system"]))
                 current_system = row["system"]
+            if row["schema_json"]:
+                click.echo(
+                    "\n## Schema:\n\n```json\n{}\n```".format(
+                        json.dumps(row["schema_json"], indent=2)
+                    )
+                )
             attachments = attachments_by_id.get(row["id"])
             if attachments:
                 click.echo("\n### Attachments\n")
@@ -1141,7 +1160,15 @@ def logs_list(
                             )
                         )
 
-            click.echo("\n## Response:\n\n{}\n".format(row["response"]))
+            # If a schema was provided and the row is valid JSON, pretty print and syntax highlight it
+            response = row["response"]
+            if row["schema_json"]:
+                try:
+                    parsed = json.loads(response)
+                    response = "```json\n{}\n```".format(json.dumps(parsed, indent=2))
+                except ValueError:
+                    pass
+            click.echo("\n## Response:\n\n{}\n".format(response))
             if usage:
                 token_usage = token_usage_string(
                     row["input_tokens"],
@@ -1510,7 +1537,7 @@ def uninstall(packages, yes):
 @click.option(
     "--metadata",
     help="JSON object metadata to store",
-    callback=_validate_metadata_json,
+    callback=json_validator("metadata"),
 )
 @click.option(
     "format_",

llm/default_plugins/openai_models.py

@@ -366,6 +366,8 @@ def _attachment(attachment):
 
 
 class _Shared:
+    supports_schema = True
+
     def __init__(
         self,
         model_id,
@@ -504,6 +506,11 @@ def build_kwargs(self, prompt, stream):
             kwargs["max_tokens"] = self.default_max_tokens
         if json_object:
             kwargs["response_format"] = {"type": "json_object"}
+        if prompt.schema:
+            kwargs["response_format"] = {
+                "type": "json_schema",
+                "json_schema": {"name": "output", "schema": prompt.schema},
+            }
         if stream:
             kwargs["stream_options"] = {"include_usage": True}
         return kwargs

llm/migrations.py

@@ -237,3 +237,21 @@ def m013_usage(db):
     db["responses"].add_column("input_tokens", int)
     db["responses"].add_column("output_tokens", int)
     db["responses"].add_column("token_details", str)
+
+
+@migration
+def m014_schemas(db):
+    db["schemas"].create(
+        {
+            "id": str,
+            "content": str,
+        },
+        pk="id",
+    )
+    db["responses"].add_column("schema_id", str, fk="schemas", fk_col="id")
+    # Clean up SQL create table indentation
+    db["responses"].transform()
+    # These changes may have dropped the FTS configuration, fix that
+    db["responses"].enable_fts(
+        ["prompt", "response"], create_triggers=True, replace=True
+    )