Sync callback object with llama.rn

mybigday · Apr 24, 2024 · 23d9d61 · 23d9d61
1 parent a32bf54
commit 23d9d61
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 3 deletions.
diff --git a/lib/binding.ts b/lib/binding.ts
@@ -30,10 +30,14 @@ export type LlamaCompletionResult = {
   truncated: boolean
 }
 
+export type LlamaCompletionToken = {
+  token: string
+}
+
 export interface LlamaContext {
   new (options: LlamaModelOptions): LlamaContext
   getSystemInfo(): string
-  completion(options: LlamaCompletionOptions, onToken?: (token: string) => void): Promise<LlamaCompletionResult>
+  completion(options: LlamaCompletionOptions, callback?: (token: LlamaCompletionToken) => void): Promise<LlamaCompletionResult>
   stopCompletion(): void
   saveSession(path: string): Promise<void>
   loadSession(path: string): Promise<void>

diff --git a/src/addons.cpp b/src/addons.cpp
@@ -274,7 +274,9 @@ class LlamaCompletionWorker : public Napi::AsyncWorker,
         const char *c_token = strdup(token.c_str());
         _tsfn.BlockingCall(c_token, [](Napi::Env env, Napi::Function jsCallback,
                                        const char *value) {
-          jsCallback.Call({Napi::String::New(env, value)});
+          auto obj = Napi::Object::New(env);
+          obj.Set("token", Napi::String::New(env, value));
+          jsCallback.Call({obj});
         });
       }
       // is it an end of generation?

diff --git a/test/__snapshots__/index.test.ts.snap b/test/__snapshots__/index.test.ts.snap
@@ -8,3 +8,63 @@ exports[`work fine 1`] = `
   "truncated": false,
 }
 `;
+
+exports[`work fine: partial 1`] = `
+{
+  "token": " swo",
+}
+`;
+
+exports[`work fine: partial 2`] = `
+{
+  "token": "cha",
+}
+`;
+
+exports[`work fine: partial 3`] = `
+{
+  "token": "door",
+}
+`;
+
+exports[`work fine: partial 4`] = `
+{
+  "token": "ter",
+}
+`;
+
+exports[`work fine: partial 5`] = `
+{
+  "token": " scientific",
+}
+`;
+
+exports[`work fine: partial 6`] = `
+{
+  "token": " Windows",
+}
+`;
+
+exports[`work fine: partial 7`] = `
+{
+  "token": "Ca",
+}
+`;
+
+exports[`work fine: partial 8`] = `
+{
+  "token": " occupied",
+}
+`;
+
+exports[`work fine: partial 9`] = `
+{
+  "token": "rå",
+}
+`;
+
+exports[`work fine: partial 10`] = `
+{
+  "token": " alta",
+}
+`;
diff --git a/test/index.test.ts b/test/index.test.ts
@@ -10,6 +10,8 @@ it('work fine', async () => {
     max_length: 100,
     n_predict: 10,
     seed: 0,
-  }, () => {})
+  }, (data) => {
+    expect(data).toMatchSnapshot('partial')
+  })
   expect(result).toMatchSnapshot()
 })