feat: implement binding for runQueryStream

developer239 · Jul 27, 2024 · e18cb8c · e18cb8c
1 parent fa70521
commit e18cb8c
Show file tree

Hide file tree

Showing 6 changed files with 60 additions and 28 deletions.
diff --git a/cpp/src/LlamaCPPBinding.cpp b/cpp/src/LlamaCPPBinding.cpp
@@ -9,6 +9,7 @@ Napi::Object LlamaCPPBinding::Init(Napi::Env env, Napi::Object exports) {
     Napi::Function func = DefineClass(env, "LlamaCPP", {
         InstanceMethod("initialize", &LlamaCPPBinding::Initialize),
         InstanceMethod("runQuery", &LlamaCPPBinding::RunQuery),
+        InstanceMethod("runQueryStream", &LlamaCPPBinding::RunQueryStream)
     });
 
     constructor = Napi::Persistent(func);
@@ -58,3 +59,24 @@ Napi::Value LlamaCPPBinding::RunQuery(const Napi::CallbackInfo& info) {
     std::string response = llama_->RunQuery(prompt, max_tokens);
     return Napi::String::New(env, response);
 }
+
+Napi::Value LlamaCPPBinding::RunQueryStream(const Napi::CallbackInfo& info) {
+    Napi::Env env = info.Env();
+    if (info.Length() < 2 || !info[0].IsString() || !info[1].IsFunction()) {
+        Napi::TypeError::New(env, "String and function expected").ThrowAsJavaScriptException();
+        return env.Null();
+    }
+
+    std::string prompt = info[0].As<Napi::String>().Utf8Value();
+    Napi::Function callback = info[1].As<Napi::Function>();
+    size_t max_tokens = 1000;
+    if (info.Length() > 2 && info[2].IsNumber()) {
+        max_tokens = info[2].As<Napi::Number>().Uint32Value();
+    }
+
+    llama_->RunQueryStream(prompt, max_tokens, [&env, &callback](const std::string& token) {
+        callback.Call(env.Global(), {Napi::String::New(env, token)});
+    });
+
+    return env.Undefined();
+}
diff --git a/cpp/src/LlamaCPPBinding.h b/cpp/src/LlamaCPPBinding.h
@@ -13,6 +13,7 @@ class LlamaCPPBinding : public Napi::ObjectWrap<LlamaCPPBinding> {
 
     Napi::Value Initialize(const Napi::CallbackInfo& info);
     Napi::Value RunQuery(const Napi::CallbackInfo& info);
+    Napi::Value RunQueryStream(const Napi::CallbackInfo& info);
 
     std::unique_ptr<LlamaWrapper> llama_;
 };
diff --git a/example/index.js b/example/index.js
@@ -8,6 +8,23 @@ async function main() {
     console.log("\nRunning a simple query:");
     const response = llama.runQuery("Tell me a short story.", 100);
     console.log("Response:", response);
+
+    console.log("\nRunning a streaming query:");
+    let streamingResponse = "";
+    llama.runQueryStream(
+        "List 5 interesting facts about space.",
+        (token) => {
+            process.stdout.write(token);
+            streamingResponse += token;
+        },
+        200
+    );
+
+    // Wait for the streaming to finish
+    await new Promise(resolve => setTimeout(resolve, 5000));
+
+    console.log("\n\nFull streaming response:");
+    console.log(streamingResponse);
 }
 
 main().catch(console.error);
diff --git a/index.d.ts b/index.d.ts
@@ -2,4 +2,5 @@ export class Llama {
   constructor();
   initialize(modelPath: string, contextSize?: number): boolean;
   runQuery(prompt: string, maxTokens?: number): string;
+  runQueryStream(prompt: string, callback: (token: string) => void, maxTokens?: number): void;
 }
diff --git a/package.json b/package.json
@@ -14,7 +14,7 @@
   },
   "dependencies": {
     "node-addon-api": "8.1.0",
-    "cmake-js": "7.2.1"
+    "cmake-js": "7.3.0"
   },
   "scripts": {
     "install": "cmake-js compile -d cpp print-configure",

diff --git a/yarn.lock b/yarn.lock
@@ -706,7 +706,7 @@ asynckit@^0.4.0:
   resolved "https://registry.yarnpkg.com/asynckit/-/asynckit-0.4.0.tgz#c79ed97f7f34cb8f2ba1bc9790bcc366474b4b79"
   integrity sha512-Oei9OH4tRh0YqU3GxhX79dM/mwVgvbZJaSNaRk+bshkj0S5cfHcgYakreBjrHwatXKbz+IoIdYLxrKim2MjW0Q==
 
-axios@^1.3.2:
+axios@^1.6.5:
   version "1.7.2"
   resolved "https://registry.yarnpkg.com/axios/-/axios-1.7.2.tgz#b625db8a7051fbea61c35a3cbb3a1daa7b9c7621"
   integrity sha512-2A8QhOMrbomlDuiLeK9XibIBzuHeRcqqNOHp0Cyp5EoJ1IFDh+XZH3A6BkXtv0K4gFGCI0Y4BM7B1wOEi0Rmgw==
@@ -917,24 +917,24 @@ clone@^1.0.2:
   resolved "https://registry.yarnpkg.com/clone/-/clone-1.0.4.tgz#da309cc263df15994c688ca902179ca3c7cd7c7e"
   integrity sha512-JQHZ2QMW6l3aH/j6xCqQThY/9OH4D/9ls34cgkUBiEeocRTU04tHfKPBsUK1PqZCUQM7GiA0IIXJSuXHI64Kbg==
 
-cmake-js@7.2.1:
-  version "7.2.1"
-  resolved "https://registry.yarnpkg.com/cmake-js/-/cmake-js-7.2.1.tgz#757c0d39994121b084bab96290baf115ee7712cd"
-  integrity sha512-AdPSz9cSIJWdKvm0aJgVu3X8i0U3mNTswJkSHzZISqmYVjZk7Td4oDFg0mCBA383wO+9pG5Ix7pEP1CZH9x2BA==
+cmake-js@7.3.0:
+  version "7.3.0"
+  resolved "https://registry.yarnpkg.com/cmake-js/-/cmake-js-7.3.0.tgz#6fd6234b7aeec4545c1c806f9e3f7ffacd9798b2"
+  integrity sha512-dXs2zq9WxrV87bpJ+WbnGKv8WUBXDw8blNiwNHoRe/it+ptscxhQHKB1SJXa1w+kocLMeP28Tk4/eTCezg4o+w==
   dependencies:
-    axios "^1.3.2"
+    axios "^1.6.5"
     debug "^4"
-    fs-extra "^10.1.0"
+    fs-extra "^11.2.0"
     lodash.isplainobject "^4.0.6"
     memory-stream "^1.0.0"
-    node-api-headers "^0.0.2"
+    node-api-headers "^1.1.0"
     npmlog "^6.0.2"
     rc "^1.2.7"
-    semver "^7.3.8"
-    tar "^6.1.11"
+    semver "^7.5.4"
+    tar "^6.2.0"
     url-join "^4.0.1"
     which "^2.0.2"
-    yargs "^17.6.0"
+    yargs "^17.7.2"
 
 cmd-shim@^6.0.0:
   version "6.0.3"
@@ -1395,16 +1395,7 @@ from2@^2.3.0:
     inherits "^2.0.1"
     readable-stream "^2.0.0"
 
-fs-extra@^10.1.0:
-  version "10.1.0"
-  resolved "https://registry.yarnpkg.com/fs-extra/-/fs-extra-10.1.0.tgz#02873cfbc4084dde127eaa5f9905eef2325d1abf"
-  integrity sha512-oRXApq54ETRj4eMiFzGnHWGy+zo5raudjuxN0b8H7s/RU2oW0Wvsx9O0ACRN/kRq9E8Vu/ReskGB5o3ji+FzHQ==
-  dependencies:
-    graceful-fs "^4.2.0"
-    jsonfile "^6.0.1"
-    universalify "^2.0.0"
-
-fs-extra@^11.0.0:
+fs-extra@^11.0.0, fs-extra@^11.2.0:
   version "11.2.0"
   resolved "https://registry.yarnpkg.com/fs-extra/-/fs-extra-11.2.0.tgz#e70e17dfad64232287d01929399e0ea7c86b0e5b"
   integrity sha512-PmDi3uwK5nFuXh7XDTlVnS17xJS7vW36is2+w3xcv8SVxiB4NyATf4ctkVY5bkSjX0Y4nbvZCq1/EjtEyr9ktw==
@@ -2503,10 +2494,10 @@ node-addon-api@8.1.0:
   resolved "https://registry.yarnpkg.com/node-addon-api/-/node-addon-api-8.1.0.tgz#55a573685dd4bd053f189cffa4e6332d2b1f1645"
   integrity sha512-yBY+qqWSv3dWKGODD6OGE6GnTX7Q2r+4+DfpqxHSHh8x0B4EKP9+wVGLS6U/AM1vxSNNmUEuIV5EGhYwPpfOwQ==
 
-node-api-headers@^0.0.2:
-  version "0.0.2"
-  resolved "https://registry.yarnpkg.com/node-api-headers/-/node-api-headers-0.0.2.tgz#31f4c6c2750b63e598128e76a60aefca6d76ac5d"
-  integrity sha512-YsjmaKGPDkmhoNKIpkChtCsPVaRE0a274IdERKnuc/E8K1UJdBZ4/mvI006OijlQZHCfpRNOH3dfHQs92se8gg==
+node-api-headers@^1.1.0:
+  version "1.2.0"
+  resolved "https://registry.yarnpkg.com/node-api-headers/-/node-api-headers-1.2.0.tgz#b717cd420aec79031f8dc83a50eb0a8bdf24c70d"
+  integrity sha512-L9AiEkBfgupC0D/LsudLPOhzy/EdObsp+FHyL1zSK0kKv5FDA9rJMoRz8xd+ojxzlqfg0tTZm2h8ot2nS7bgRA==
 
 node-emoji@^1.11.0:
   version "1.11.0"
@@ -3377,7 +3368,7 @@ semver@^6.0.0:
   resolved "https://registry.yarnpkg.com/semver/-/semver-6.3.1.tgz#556d2ef8689146e46dcea4bfdd095f3434dffcb4"
   integrity sha512-BR7VvDCVHO+q2xBEWskxS6DJE1qRnb7DxzUrogb71CWoSficBxYsiAGd+Kl0mmq/MprG9yArRkyrQxTO6XjMzA==
 
-semver@^7.1.1, semver@^7.1.2, semver@^7.3.2, semver@^7.3.4, semver@^7.3.5, semver@^7.3.7, semver@^7.3.8, semver@^7.5.3, semver@^7.6.0:
+semver@^7.1.1, semver@^7.1.2, semver@^7.3.2, semver@^7.3.4, semver@^7.3.5, semver@^7.3.7, semver@^7.5.3, semver@^7.5.4, semver@^7.6.0:
   version "7.6.3"
   resolved "https://registry.yarnpkg.com/semver/-/semver-7.6.3.tgz#980f7b5550bc175fb4dc09403085627f9eb33143"
   integrity sha512-oVekP1cKtI+CTDvHWYFUcMtsK/00wmAEfyqKfNdARm8u1wNVhSgaX7A8d4UuIlUI5e84iEwOhs7ZPYRmzU9U6A==
@@ -3958,7 +3949,7 @@ yargs-parser@^21.1.1:
   resolved "https://registry.yarnpkg.com/yargs-parser/-/yargs-parser-21.1.1.tgz#9096bceebf990d21bb31fa9516e0ede294a77d35"
   integrity sha512-tVpsJW7DdjecAiFpbIB1e3qxIQsE6NoPc5/eTdrbbIC4h0LVsWhnoa3g+m2HclBIujHzsxZ4VJVA+GUuc2/LBw==
 
-yargs@^17.5.1, yargs@^17.6.0:
+yargs@^17.5.1, yargs@^17.7.2:
   version "17.7.2"
   resolved "https://registry.yarnpkg.com/yargs/-/yargs-17.7.2.tgz#991df39aca675a192b816e1e0363f9d75d2aa269"
   integrity sha512-7dSzzRQ++CKnNI/krKnYRV7JKKPUXMEh61soaHKg9mrWEhzFWhFnxPxGl+69cD1Ou63C13NUPCnmIcrvqCuM6w==