kjarni-wasm 0.1.9 → 0.2.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "kjarni-wasm",
-  "version": "0.1.9",
+  "version": "0.2.0",
   "description": "Local AI inference in the browser: embeddings, classification, reranking, semantic search and chat, running in a Web Worker.",
   "type": "module",
   "main": "./dist/index.js",

package/pkg/kjarni_wasm.d.ts

@@ -1,19 +1,6 @@
 /* tslint:disable */
 /* eslint-disable */
 
-/**
- * A local language model running in the browser.
- *
- * Text generation is synchronous here, and deliberately so. The engine's streaming
- * path builds a background task with `tokio::task::spawn_blocking`, and wasm has no
- * blocking pool for that to run on — it compiles but would fail at runtime. Calling
- * the generation loop directly and draining the channel afterwards avoids the
- * problem entirely.
- *
- * The cost is that a call blocks until the whole response is ready, which for a few
- * hundred tokens is seconds, not milliseconds. Run it in a Web Worker; on the main
- * thread it will freeze the tab.
- */
 export class WasmChat {
     private constructor();
     free(): void;
@@ -26,6 +13,20 @@ export class WasmChat {
      * The model's context window, in tokens.
      */
     context_size(): number;
+    /**
+     * Reuse the KV cache across turns, so a conversation prefills only its newest
+     * message instead of replaying the whole transcript.
+     *
+     * Measured natively on Qwen2.5-0.5B with a long system prompt: turn one is
+     * unchanged, turns two and three go from about 2.2s to roughly 0.2s, for
+     * identical answers. A browser tab is one conversation, so there is no sharing
+     * hazard here.
+     *
+     * `tokens` is the longest transcript that can be reused, and what the memory
+     * cost scales with; pass 0 for [`DEFAULT_WASM_PREFIX_CACHE_TOKENS`]. Longer
+     * conversations keep working, they simply stop reusing once they outgrow it.
+     */
+    enablePrefixCache(tokens: number): void;
     /**
      * Generate a completion for `prompt`.
      */
@@ -45,6 +46,10 @@ export class WasmChat {
      * itself is in `generate_stream_core`, which builds and is tested natively.
      */
     generate_stream(prompt: string, max_new_tokens: number, temperature: number, on_token: Function): string;
+    /**
+     * Whether prefix reuse is on.
+     */
+    hasPrefixCache(): boolean;
     /**
      * Number of messages held, system prompt included.
      */
@@ -215,8 +220,10 @@ export interface InitOutput {
     readonly set_debug_logging: (a: number) => void;
     readonly wasmchat_clear_history: (a: number) => void;
     readonly wasmchat_context_size: (a: number) => number;
+    readonly wasmchat_enablePrefixCache: (a: number, b: number, c: number) => void;
     readonly wasmchat_generate: (a: number, b: number, c: number, d: number, e: number, f: number) => void;
     readonly wasmchat_generate_stream: (a: number, b: number, c: number, d: number, e: number, f: number, g: number) => void;
+    readonly wasmchat_hasPrefixCache: (a: number) => number;
     readonly wasmchat_history_len: (a: number) => number;
     readonly wasmchat_load: (a: number, b: number, c: number, d: number, e: number) => void;
     readonly wasmchat_send_stream: (a: number, b: number, c: number, d: number, e: number, f: number, g: number) => void;
@@ -248,8 +255,8 @@ export interface InitOutput {
     readonly wasmsearch_search_keywords: (a: number, b: number, c: number, d: number, e: number) => void;
     readonly wasmsearch_search_semantic: (a: number, b: number, c: number, d: number, e: number) => void;
     readonly wasmsearch_update_file: (a: number, b: number, c: number, d: number, e: number, f: number) => void;
-    readonly __wasm_bindgen_func_elem_5516: (a: number, b: number, c: number, d: number) => void;
-    readonly __wasm_bindgen_func_elem_5525: (a: number, b: number, c: number, d: number) => void;
+    readonly __wasm_bindgen_func_elem_5674: (a: number, b: number, c: number, d: number) => void;
+    readonly __wasm_bindgen_func_elem_5694: (a: number, b: number, c: number, d: number) => void;
     readonly __wbindgen_export: (a: number, b: number) => number;
     readonly __wbindgen_export2: (a: number, b: number, c: number, d: number) => number;
     readonly __wbindgen_export3: (a: number) => void;

package/pkg/kjarni_wasm.js

@@ -1,18 +1,5 @@
 /* @ts-self-types="./kjarni_wasm.d.ts" */
 
-/**
- * A local language model running in the browser.
- *
- * Text generation is synchronous here, and deliberately so. The engine's streaming
- * path builds a background task with `tokio::task::spawn_blocking`, and wasm has no
- * blocking pool for that to run on — it compiles but would fail at runtime. Calling
- * the generation loop directly and draining the channel afterwards avoids the
- * problem entirely.
- *
- * The cost is that a call blocks until the whole response is ready, which for a few
- * hundred tokens is seconds, not milliseconds. Run it in a Web Worker; on the main
- * thread it will freeze the tab.
- */
 export class WasmChat {
     static __wrap(ptr) {
         const obj = Object.create(WasmChat.prototype);
@@ -44,6 +31,33 @@ export class WasmChat {
         const ret = wasm.wasmchat_context_size(this.__wbg_ptr);
         return ret >>> 0;
     }
+    /**
+     * Reuse the KV cache across turns, so a conversation prefills only its newest
+     * message instead of replaying the whole transcript.
+     *
+     * Measured natively on Qwen2.5-0.5B with a long system prompt: turn one is
+     * unchanged, turns two and three go from about 2.2s to roughly 0.2s, for
+     * identical answers. A browser tab is one conversation, so there is no sharing
+     * hazard here.
+     *
+     * `tokens` is the longest transcript that can be reused, and what the memory
+     * cost scales with; pass 0 for [`DEFAULT_WASM_PREFIX_CACHE_TOKENS`]. Longer
+     * conversations keep working, they simply stop reusing once they outgrow it.
+     * @param {number} tokens
+     */
+    enablePrefixCache(tokens) {
+        try {
+            const retptr = wasm.__wbindgen_add_to_stack_pointer(-16);
+            wasm.wasmchat_enablePrefixCache(retptr, this.__wbg_ptr, tokens);
+            var r0 = getDataViewMemory0().getInt32(retptr + 4 * 0, true);
+            var r1 = getDataViewMemory0().getInt32(retptr + 4 * 1, true);
+            if (r1) {
+                throw takeObject(r0);
+            }
+        } finally {
+            wasm.__wbindgen_add_to_stack_pointer(16);
+        }
+    }
     /**
      * Generate a completion for `prompt`.
      * @param {string} prompt
@@ -123,6 +137,14 @@ export class WasmChat {
             wasm.__wbindgen_export4(deferred3_0, deferred3_1, 1);
         }
     }
+    /**
+     * Whether prefix reuse is on.
+     * @returns {boolean}
+     */
+    hasPrefixCache() {
+        const ret = wasm.wasmchat_hasPrefixCache(this.__wbg_ptr);
+        return ret !== 0;
+    }
     /**
      * Number of messages held, system prompt included.
      * @returns {number}
@@ -1092,7 +1114,7 @@ function __wbg_get_imports() {
                     const a = state0.a;
                     state0.a = 0;
                     try {
-                        return __wasm_bindgen_func_elem_5525(a, state0.b, arg0, arg1);
+                        return __wasm_bindgen_func_elem_5694(a, state0.b, arg0, arg1);
                     } finally {
                         state0.a = a;
                     }
@@ -1206,8 +1228,8 @@ function __wbg_get_imports() {
             return addHeapObject(ret);
         },
         __wbindgen_generic_0000000000000001: function(arg0, arg1) {
-            // Cast intrinsic for `Closure(Closure { owned: true, function: Function { arguments: [Externref], shim_idx: 2060, ret: Result(Unit), inner_ret: Some(Result(Unit)) }, mutable: true }) -> Externref`.
-            const ret = makeMutClosure(arg0, arg1, __wasm_bindgen_func_elem_5516);
+            // Cast intrinsic for `Closure(Closure { owned: true, function: Function { arguments: [Externref], shim_idx: 2082, ret: Result(Unit), inner_ret: Some(Result(Unit)) }, mutable: true }) -> Externref`.
+            const ret = makeMutClosure(arg0, arg1, __wasm_bindgen_func_elem_5674);
             return addHeapObject(ret);
         },
         __wbindgen_generic_0000000000000002: function(arg0) {
@@ -1244,14 +1266,14 @@ function __wbg_get_imports() {
     };
 }
 
-function __wasm_bindgen_func_elem_5525(arg0, arg1, arg2, arg3) {
-    wasm.__wasm_bindgen_func_elem_5525(arg0, arg1, addHeapObject(arg2), addHeapObject(arg3));
+function __wasm_bindgen_func_elem_5694(arg0, arg1, arg2, arg3) {
+    wasm.__wasm_bindgen_func_elem_5694(arg0, arg1, addHeapObject(arg2), addHeapObject(arg3));
 }
 
-function __wasm_bindgen_func_elem_5516(arg0, arg1, arg2) {
+function __wasm_bindgen_func_elem_5674(arg0, arg1, arg2) {
     try {
         const retptr = wasm.__wbindgen_add_to_stack_pointer(-16);
-        wasm.__wasm_bindgen_func_elem_5516(retptr, arg0, arg1, addHeapObject(arg2));
+        wasm.__wasm_bindgen_func_elem_5674(retptr, arg0, arg1, addHeapObject(arg2));
         var r0 = getDataViewMemory0().getInt32(retptr + 4 * 0, true);
         var r1 = getDataViewMemory0().getInt32(retptr + 4 * 1, true);
         if (r1) {

package/pkg/kjarni_wasm_bg.wasm.d.ts

@@ -12,8 +12,10 @@ export const init: () => void;
 export const set_debug_logging: (a: number) => void;
 export const wasmchat_clear_history: (a: number) => void;
 export const wasmchat_context_size: (a: number) => number;
+export const wasmchat_enablePrefixCache: (a: number, b: number, c: number) => void;
 export const wasmchat_generate: (a: number, b: number, c: number, d: number, e: number, f: number) => void;
 export const wasmchat_generate_stream: (a: number, b: number, c: number, d: number, e: number, f: number, g: number) => void;
+export const wasmchat_hasPrefixCache: (a: number) => number;
 export const wasmchat_history_len: (a: number) => number;
 export const wasmchat_load: (a: number, b: number, c: number, d: number, e: number) => void;
 export const wasmchat_send_stream: (a: number, b: number, c: number, d: number, e: number, f: number, g: number) => void;
@@ -45,8 +47,8 @@ export const wasmsearch_search: (a: number, b: number, c: number, d: number, e:
 export const wasmsearch_search_keywords: (a: number, b: number, c: number, d: number, e: number) => void;
 export const wasmsearch_search_semantic: (a: number, b: number, c: number, d: number, e: number) => void;
 export const wasmsearch_update_file: (a: number, b: number, c: number, d: number, e: number, f: number) => void;
-export const __wasm_bindgen_func_elem_5516: (a: number, b: number, c: number, d: number) => void;
-export const __wasm_bindgen_func_elem_5525: (a: number, b: number, c: number, d: number) => void;
+export const __wasm_bindgen_func_elem_5674: (a: number, b: number, c: number, d: number) => void;
+export const __wasm_bindgen_func_elem_5694: (a: number, b: number, c: number, d: number) => void;
 export const __wbindgen_export: (a: number, b: number) => number;
 export const __wbindgen_export2: (a: number, b: number, c: number, d: number) => number;
 export const __wbindgen_export3: (a: number) => void;

package/pkg/package.json

@@ -5,7 +5,7 @@
     "Ólafur Aron Jóhannsson <olafurjohannss@gmail.com>"
   ],
   "description": "Local AI inference in the browser: embeddings, classification, reranking, semantic search and chat, compiled to WebAssembly.",
-  "version": "0.1.9",
+  "version": "0.2.0",
   "license": "MIT",
   "repository": {
     "type": "git",