feat: add JSON schema grammar support (#68)

* feat: add JSON schema grammar support
* test: `LlamaJsonSchemaGrammar`
* feat: add JSON schema grammar support to the `chat` command
* feat: add `promptWithMeta` function to `LlamaChatSession`
* docs: document JSON schema grammar support
* ci: remove redundant step
* chore: eslint config fix

.eslintrc.json

@@ -5,7 +5,7 @@
         "browser": false,
         "es6": true
     },
-    "ignorePatterns": ["/dist", "/llama", "/docs"],
+    "ignorePatterns": ["/dist", "/llama", "/docs-site"],
     "extends": [
         "eslint:recommended"
     ],
@@ -40,6 +40,11 @@
             "@typescript-eslint/semi": ["warn", "always"],
             "@typescript-eslint/no-inferrable-types": ["off"]
         }
+    }, {
+        "files": ["test/**/**.ts"],
+        "rules": {
+            "max-len": ["off"]
+        }
     }],
     "plugins": [
         "@typescript-eslint",

.github/workflows/build.yml

@@ -89,10 +89,6 @@ jobs:
           name: llama.cpp
           path: llama/llama.cpp
 
-      - uses: actions/setup-python@v4
-        with:
-          python-version: "3.10"
-
       - name: Install dependencies on windows
         if: startsWith(matrix.config.os, 'windows')
         run: |
@@ -204,6 +200,40 @@ jobs:
           name: "bins-${{ matrix.config.artifact }}"
           path: "llamaBins/*"
 
+  standalone-tests:
+    name: Standalone tests
+    runs-on: ubuntu-22.04
+    needs:
+      - build
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-node@v3
+        with:
+          node-version: "20"
+
+      - name: Download build artifact
+        uses: actions/download-artifact@v3
+        with:
+          name: build
+          path: dist
+
+      - name: Download llama.cpp artifact
+        uses: actions/download-artifact@v3
+        with:
+          name: llama.cpp
+          path: llama/llama.cpp
+
+      - name: Install dependencies on ubuntu
+        run: |
+          sudo apt-get update
+          sudo apt-get install ninja-build cmake 
+
+      - name: Install modules
+        run: npm ci
+
+      - name: Run standalone tests
+        run: npm run test:standalone
+
   release:
     name: Release
     if: github.ref == 'refs/heads/master'

.vitepress/config.ts

@@ -73,7 +73,14 @@ export default defineConfig({
         }
     },
     transformPageData(pageData) {
-        if (pageData.filePath.startsWith("api/") || pageData.filePath.startsWith("guide/cli/")) {
+        if (pageData.filePath.startsWith("api/")) {
+            pageData.frontmatter.editLink = false;
+            pageData.frontmatter.lastUpdated = false;
+            pageData.frontmatter ||= {}
+            pageData.frontmatter.outline = [2, 3];
+        }
+
+        if (pageData.filePath.startsWith("guide/cli/")) {
             pageData.frontmatter.editLink = false;
             pageData.frontmatter.lastUpdated = false;
         }
@@ -181,13 +188,20 @@ function getApiReferenceSidebar(): typeof typedocSidebar {
 }
 
 function orderApiReferenceSidebar(sidebar: typeof typedocSidebar): typeof typedocSidebar {
+    orderClasses(sidebar);
+    orderTypes(sidebar);
+
+    return sidebar;
+}
+
+function orderClasses(sidebar: typeof typedocSidebar) {
     const baseChatPromptWrapper = "ChatPromptWrapper";
     const chatPromptWrapperItems: DefaultTheme.SidebarItem[] = [];
 
     const classes = sidebar.find((item) => item.text === "Classes");
 
     if (classes == null || !(classes.items instanceof Array))
-        return sidebar;
+        return;
 
     (classes.items as DefaultTheme.SidebarItem[]).unshift({
         text: "Chat wrappers",
@@ -222,7 +236,54 @@ function orderApiReferenceSidebar(sidebar: typeof typedocSidebar): typeof typedo
 
         return aIndex - bIndex;
     });
-
-    return sidebar;
 }
 
+function orderTypes(sidebar: typeof typedocSidebar) {
+    const types = sidebar.find((item) => item.text === "Types");
+
+    if (types == null || !(types.items instanceof Array))
+        return;
+
+    function groupGbnfJsonSchema() {
+        if (types == null || !(types.items instanceof Array))
+            return;
+
+        const gbnfJsonSchemaItemTitle = "GbnfJsonSchema";
+        const gbnfItemsPrefix = "GbnfJson";
+        const gbnfJsonSchemaItems: DefaultTheme.SidebarItem[] = [];
+
+        const gbnfJsonSchemaItem = types.items
+            .find((item) => item.text === gbnfJsonSchemaItemTitle) as DefaultTheme.SidebarItem | null;
+
+        if (gbnfJsonSchemaItem == null)
+            return;
+
+        gbnfJsonSchemaItem.collapsed = true;
+        gbnfJsonSchemaItem.items = gbnfJsonSchemaItems;
+
+        for (const item of types.items.slice()) {
+            if (item.text === gbnfJsonSchemaItemTitle || !item.text.startsWith(gbnfItemsPrefix))
+                continue;
+
+            types.items.splice(types.items.indexOf(item), 1);
+            gbnfJsonSchemaItems.push(item);
+        }
+    }
+
+    function moveCollapseItemsToTheEnd() {
+        if (types == null || !(types.items instanceof Array))
+            return;
+
+        types.items.sort((a, b) => {
+            if (a.collapsed && !b.collapsed)
+                return 1;
+            if (!a.collapsed && b.collapsed)
+                return -1;
+
+            return 0;
+        });
+    }
+
+    groupGbnfJsonSchema();
+    moveCollapseItemsToTheEnd();
+}

README.md

@@ -22,7 +22,7 @@
 * Chat with a model using a chat wrapper
 * Use the CLI to chat with a model without writing any code
 * Up-to-date with the latest version of `llama.cpp`. Download and compile the latest release with a single CLI command.
-* Force a model to generate output in a parseable format, like JSON
+* Force a model to generate output in a parseable format, like JSON, or even force it to follow a specific JSON schema
 
 ## [Documentation](https://withcatai.github.io/node-llama-cpp/)
 * [Getting started guide](https://withcatai.github.io/node-llama-cpp/guide/)

docs/guide/chat-session.md

@@ -212,3 +212,48 @@ const a2 = await session.prompt(q2, {
 console.log("AI: " + a2);
 console.log(JSON.parse(a2));
 ```
+
+## JSON response with schema
+To learn more about the JSON schema grammar, see the [grammar guide](./grammar.md#using-a-json-schema-grammar).
+```typescript
+import {fileURLToPath} from "url";
+import path from "path";
+import {
+    LlamaModel, LlamaJsonSchemaGrammar, LlamaContext, LlamaChatSession
+} from "node-llama-cpp";
+
+const __dirname = path.dirname(fileURLToPath(import.meta.url));
+
+const model = new LlamaModel({
+    modelPath: path.join(__dirname, "models", "codellama-13b.Q3_K_M.gguf")
+})
+const grammar = new LlamaJsonSchemaGrammar({
+    "type": "object",
+    "properties": {
+        "responseMessage": {
+            "type": "string"
+        },
+        "requestPositivityScoreFromOneToTen": {
+            "type": "number"
+        }
+    }
+} as const);
+const context = new LlamaContext({model});
+const session = new LlamaChatSession({context});
+
+
+const q1 = 'How are you doing?';
+console.log("User: " + q1);
+
+const a1 = await session.prompt(q1, {
+    grammar,
+    maxTokens: context.getContextSize()
+});
+console.log("AI: " + a1);
+
+const parsedA1 = grammar.parse(a1);
+console.log(
+    parsedA1.responseMessage,
+    parsedA1.requestPositivityScoreFromOneToTen
+);
+```

docs/guide/cli/build.md

@@ -1,3 +1,6 @@
+---
+outline: deep
+---
 # `build` command
 
 <script setup lang="ts">

docs/guide/cli/chat.md

@@ -1,3 +1,6 @@
+---
+outline: deep
+---
 # `chat` command
 
 <script setup lang="ts">

docs/guide/cli/clear.md

@@ -1,3 +1,6 @@
+---
+outline: deep
+---
 # `clear` command
 
 <script setup lang="ts">

docs/guide/cli/cli.data.ts

@@ -1,9 +1,9 @@
+import {CommandModule} from "yargs";
 import {getCommandHtmlDoc} from "../../../.vitepress/utils/getCommandHtmlDoc.js";
 import {BuildCommand} from "../../../src/cli/commands/BuildCommand.js";
 import {ChatCommand} from "../../../src/cli/commands/ChatCommand.js";
 import {DownloadCommand} from "../../../src/cli/commands/DownloadCommand.js";
 import {ClearCommand} from "../../../src/cli/commands/ClearCommand.js";
-import {CommandModule} from "yargs";
 import {htmlEscape} from "../../../.vitepress/utils/htmlEscape.js";
 import {cliBinName, npxRunPrefix} from "../../../src/config.js";
 import {buildHtmlHeading} from "../../../.vitepress/utils/buildHtmlHeading.js";
@@ -28,7 +28,7 @@ export default {
             clear: await getCommandHtmlDoc(ClearCommand)
         };
     }
-}
+};
 
 function buildIndexTable(commands: [pageLink: string, command: CommandModule<any, any>][], cliName: string = cliBinName) {
     let res = "";

docs/guide/cli/download.md

@@ -1,3 +1,6 @@
+---
+outline: deep
+---
 # `download` command
 
 <script setup lang="ts">

docs/guide/cli/index.md

@@ -1,3 +1,6 @@
+---
+outline: deep
+---
 # CLI
 
 <script setup lang="ts">

docs/guide/contributing.md

@@ -1,3 +1,6 @@
+---
+outline: deep
+---
 # Opening a PR on `node-llama-cpp`
 This document describes the guidelines of how to open a PR on the `node-llama-cpp` project.
 

docs/guide/development.md

@@ -1,3 +1,6 @@
+---
+outline: deep
+---
 # Developing `node-llama-cpp`
 This document describes how to set up your development environment to contribute to `node-llama-cpp`.
 

docs/guide/grammar.md

@@ -57,6 +57,56 @@ console.log("AI: " + a2);
 console.log(JSON.parse(a2));
 ```
 
+## Using a JSON schema grammar
+The [`LlamaJsonSchemaGrammar`](/api/classes/LlamaJsonSchemaGrammar) class uses a GBNF grammar that's generated based on the [JSON schema](https://json-schema.org/learn/getting-started-step-by-step) you provide.
+
+It only supports [a small subset of the JSON schema spec](/api/type-aliases/GbnfJsonSchema), but it's enough to generate useful JSON objects using a text generation model.
+
+To see what subset of the JSON schema spec is supported, see the [`GbnfJsonSchema` type](/api/type-aliases/GbnfJsonSchema).
+
+```typescript
+import {fileURLToPath} from "url";
+import path from "path";
+import {
+    LlamaModel, LlamaJsonSchemaGrammar, LlamaContext, LlamaChatSession
+} from "node-llama-cpp";
+
+const __dirname = path.dirname(fileURLToPath(import.meta.url));
+
+const model = new LlamaModel({
+    modelPath: path.join(__dirname, "models", "codellama-13b.Q3_K_M.gguf")
+})
+const grammar = new LlamaJsonSchemaGrammar({
+    "type": "object",
+    "properties": {
+        "responseMessage": {
+            "type": "string"
+        },
+        "requestPositivityScoreFromOneToTen": {
+            "type": "number"
+        }
+    }
+} as const);
+const context = new LlamaContext({model});
+const session = new LlamaChatSession({context});
+
+
+const q1 = 'How are you doing?';
+console.log("User: " + q1);
+
+const a1 = await session.prompt(q1, {
+    grammar,
+    maxTokens: context.getContextSize()
+});
+console.log("AI: " + a1);
+
+const parsedA1 = grammar.parse(a1);
+console.log(
+    parsedA1.responseMessage,
+    parsedA1.requestPositivityScoreFromOneToTen
+);
+```
+
 ## Creating your own grammar
 To create your own grammar, read the [GBNF guide](https://github.com/ggerganov/llama.cpp/blob/f5fe98d11bdf9e7797bcfb05c0c3601ffc4b9d26/grammars/README.md) to create a GBNF grammar file.