feat: improve grammar support (#215)

* feat: improve grammar support
* feat: improve JSON schema grammar

Author: giladgd

.config/typedoc.json

@@ -23,6 +23,6 @@
     "propertiesFormat": "list",
     "enumMembersFormat": "table",
     "typeDeclarationFormat": "list",
-    "hideInPageTOC": true,
+    "sort": ["source-order"],
     "docsRoot": "../docs"
 }

.releaserc.ts

@@ -40,7 +40,9 @@ export default {
             }
         }],
         "@semantic-release/npm",
-        "@semantic-release/github",
+        ["@semantic-release/github", {
+            "discussionCategoryName": "Releases"
+        }],
         ["@semantic-release/exec", {
             "publishCmd": "echo \"${nextRelease.version}\" > .semanticRelease.npmPackage.deployedVersion.txt"
         }]

.vitepress/config.ts

@@ -137,7 +137,10 @@ export default defineConfig({
                         ...(await fs.readJSON(path.join(__dirname, "..", "tsconfig.json"))).compilerOptions,
                         moduleResolution: undefined,
                         paths: {
-                            "node-llama-cpp": [path.resolve(__dirname, "..", "src", "index.ts")]
+                            "node-llama-cpp": [
+                                path.resolve(__dirname, "..", "dist", "index.d.ts"),
+                                path.resolve(__dirname, "..", "src", "index.ts")
+                            ]
                         },
                         typeRoots: [
                             path.resolve(__dirname, "..", "node_modules"),
@@ -157,7 +160,7 @@ export default defineConfig({
         },
         nav: [
             {text: "Guide", link: "/guide/", activeMatch: "/guide/"},
-            {text: "API Reference", link: "/api/classes/LlamaModel", activeMatch: "/api/"},
+            {text: "API Reference", link: "/api/functions/getLlama", activeMatch: "/api/"},
             {
                 text: packageVersion,
                 items: [{

package-lock.json

@@ -62,7 +62,7 @@
     "dev:build": "npm run build && node ./dist/cli/cli.js build --noUsageExample",
     "clean": "rm -rf ./node_modules ./dist ./tsconfig.tsbuildinfo ./test/.models ./docs/api ./docs/api-overrides",
     "docs:generateTypedoc": "typedoc && rimraf ./docs/api/index.md ./docs/api/globals.md ./docs/api/functions/LlamaText.md && npm run docs:generateTypedoc:overrides",
-    "docs:generateTypedoc:overrides": "typedoc --entryPoints ./src/apiDocsOverrides.ts --out ./docs/api-overrides && copyfiles --flat \"./docs/api-overrides/classes/*.md\" ./docs/api/classes && rimraf ./docs/api-overrides",
+    "docs:generateTypedoc:overrides": "typedoc --entryPoints ./src/apiDocsOverrides.ts --out ./docs/api-overrides && copyfiles --flat \"./docs/api-overrides/classes/LlamaText.md\" ./docs/api/classes && rimraf ./docs/api-overrides",
     "docs:dev": "npm run docs:generateTypedoc && vitepress dev",
     "docs:build": "npm run docs:generateTypedoc && vitepress build",
     "docs:preview": "npm run docs:generateTypedoc && vitepress preview"
@@ -139,10 +139,10 @@
     "rimraf": "^5.0.1",
     "semantic-release": "^22.0.8",
     "tslib": "^2.6.1",
-    "typedoc": "^0.25.3",
+    "typedoc": "^0.25.13",
     "typedoc-plugin-markdown": "^4.0.0-next.55",
-    "typedoc-plugin-mdn-links": "^3.1.19",
-    "typedoc-vitepress-theme": "^1.0.0-next.10",
+    "typedoc-plugin-mdn-links": "^3.1.24",
+    "typedoc-vitepress-theme": "1.0.0-next.10",
     "typescript": "^5.2.2",
     "vite-node": "^1.4.0",
     "vitepress": "^1.1.4",

package.json

@@ -1,5 +1,24 @@
 /** @internal */
-import {_LlamaText} from "./utils/LlamaText.js";
+import {Tokenizer} from "./types.js";
+/** @internal */
+import {
+    _LlamaText,
+    type LlamaTextJSON,
+    type LlamaTextJSONValue,
+    type LlamaTextSpecialTokenJSON,
+    type LlamaTextSpecialTokensTextJSON,
+    type LlamaTextValue,
+    type LlamaTextInputValue
+} from "./utils/LlamaText.js";
 
 /** @internal */
-export {_LlamaText as LlamaText};
+export {
+    _LlamaText as LlamaText,
+    type Tokenizer,
+    type LlamaTextJSON,
+    type LlamaTextJSONValue,
+    type LlamaTextSpecialTokensTextJSON,
+    type LlamaTextSpecialTokenJSON,
+    type LlamaTextValue,
+    type LlamaTextInputValue
+};

src/apiDocsOverrides.ts

@@ -3,6 +3,9 @@ import {DisposedError, EventRelay, withLock} from "lifecycle-utils";
 import {getConsoleLogPrefix} from "../utils/getConsoleLogPrefix.js";
 import {LlamaModel, LlamaModelOptions} from "../evaluator/LlamaModel.js";
 import {DisposeGuard} from "../utils/DisposeGuard.js";
+import {GbnfJsonSchema} from "../utils/gbnfJson/types.js";
+import {LlamaJsonSchemaGrammar} from "../evaluator/LlamaJsonSchemaGrammar.js";
+import {LlamaGrammar, LlamaGrammarOptions} from "../evaluator/LlamaGrammar.js";
 import {BindingModule} from "./AddonTypes.js";
 import {BuildGpu, BuildMetadataFile, LlamaLocks, LlamaLogLevel} from "./types.js";
 import {MemoryOrchestrator, MemoryReservation} from "./utils/MemoryOrchestrator.js";
@@ -229,6 +232,18 @@ export class Llama {
         });
     }
 
+    public async createGrammarForJsonSchema<const T extends Readonly<GbnfJsonSchema>>(schema: T) {
+        return new LlamaJsonSchemaGrammar<T>(this, schema);
+    }
+
+    public async getGrammarFor(type: Parameters<typeof LlamaGrammar.getFor>[1]) {
+        return await LlamaGrammar.getFor(this, type);
+    }
+
+    public async createGrammar(options: LlamaGrammarOptions) {
+        return new LlamaGrammar(this, options);
+    }
+
     /** @internal */
     public async _init() {
         await this._bindings.init();

src/bindings/Llama.ts

@@ -169,7 +169,7 @@ const defaultBuildOption: Exclude<LlamaOptions["build"], undefined> = runningInE
     : "auto";
 
 /**
- * Get a llama.cpp binding.
+ * Get a `llama.cpp` binding.
  * Defaults to prefer a prebuilt binary, and fallbacks to building from source if a prebuilt binary is not found.
  * Pass `"lastCliBuild"` to default to use the last successful build created using the `download` or `build` CLI commands if one exists.
  */

src/bindings/getLlama.ts

@@ -657,7 +657,7 @@ export class LlamaChat {
                         pendingTokens.push(...streamRegulator.popFreeChunkTokens());
 
                         const triggeredStops = functionSyntaxStartDetector.getTriggeredStops();
-                        const partiallyFreeTokens = streamRegulator.getPartiallyFreeChunk();
+                        const partiallyFreeTokens = streamRegulator.getPartiallyFreeChunk(model.tokenizer);
 
                         const queuedTokensBeforeStopTrigger = getQueuedTokensBeforeStopTrigger(
                             triggeredStops,
@@ -775,10 +775,15 @@ export class LlamaChat {
                     if (stopGenerationDetector.hasTriggeredStops || customStopGenerationTriggersDetector.hasTriggeredStops ||
                         model.isEogToken(token)
                     ) {
+                        stopGenerationDetector.clearInProgressStops();
+                        customStopGenerationTriggersDetector.clearInProgressStops();
+                        pendingTokens.push(...streamRegulator.popFreeChunkTokens());
+
                         const triggeredStops = stopGenerationDetector.hasTriggeredStops
                             ? stopGenerationDetector.getTriggeredStops()
                             : customStopGenerationTriggersDetector.getTriggeredStops();
-                        const partiallyFreeTokens = streamRegulator.getPartiallyFreeChunk();
+
+                        const partiallyFreeTokens = streamRegulator.getPartiallyFreeChunk(model.tokenizer);
 
                         const queuedTokensBeforeStopTrigger = getQueuedTokensBeforeStopTrigger(
                             triggeredStops,

src/evaluator/LlamaChat/LlamaChat.ts

@@ -20,8 +20,7 @@ export class FunctionCallGrammar<const Functions extends ChatModelFunctions> ext
     public constructor(llama: Llama, functions: Functions, chatWrapper: ChatWrapper, initialFunctionCallEngaged: boolean) {
         const grammar = getGbnfGrammarForFunctionCalls(functions, chatWrapper, initialFunctionCallEngaged);
 
-        super({
-            llama,
+        super(llama, {
             grammar,
             stopGenerationTriggers: [LlamaText(chatWrapper.settings.functions.call.suffix, "\n".repeat(4))],
             trimWhitespaceSuffix: true