openvino-genai-node
Version:
OpenVINO™ GenAI pipelines for using from Node.js environment
102 lines • 4.08 kB
JavaScript
import util from 'node:util';
import addon from '../addon.js';
import { StreamingStatus } from '../utils.js';
export class LLMPipeline {
constructor(modelPath, device) {
this.modelPath = null;
this.device = null;
this.pipeline = null;
this.isInitialized = false;
this.isChatStarted = false;
this.modelPath = modelPath;
this.device = device;
}
async init() {
if (this.isInitialized)
throw new Error('LLMPipeline is already initialized');
this.pipeline = new addon.LLMPipeline();
const initPromise = util.promisify(this.pipeline.init.bind(this.pipeline));
const result = await initPromise(this.modelPath, this.device);
this.isInitialized = true;
return result;
}
async startChat() {
if (this.isChatStarted)
throw new Error('Chat is already started');
const startChatPromise = util.promisify(this.pipeline.startChat.bind(this.pipeline));
const result = await startChatPromise();
this.isChatStarted = true;
return result;
}
async finishChat() {
if (!this.isChatStarted)
throw new Error('Chat is not started');
const finishChatPromise = util.promisify(this.pipeline.finishChat.bind(this.pipeline));
const result = await finishChatPromise();
this.isChatStarted = false;
return result;
}
stream(prompt, generationConfig = {}) {
if (!this.isInitialized)
throw new Error('Pipeline is not initialized');
if (typeof prompt !== 'string')
throw new Error('Prompt must be a string');
if (typeof generationConfig !== 'object')
throw new Error('Options must be an object');
const queue = [];
let resolvePromise;
// Callback function that C++ will call when a chunk is ready
function chunkOutput(isDone, subword) {
if (resolvePromise) {
// Fulfill pending request
resolvePromise({ value: subword, done: isDone });
resolvePromise = null; // Reset promise resolver
}
else {
// Add data to queue if no pending promise
queue.push({ isDone, subword });
}
}
this.pipeline.generate(prompt, chunkOutput, generationConfig);
return {
async next() {
// If there is data in the queue, return it
// Otherwise, return a promise that will resolve when data is available
const data = queue.shift();
if (data !== undefined) {
const { isDone, subword } = data;
return { value: subword, done: isDone };
}
return new Promise((resolve) => (resolvePromise = resolve));
},
[Symbol.asyncIterator]() { return this; },
};
}
async generate(prompt, generationConfig = {}, callback) {
if (typeof prompt !== 'string'
&& !(Array.isArray(prompt)
&& prompt.every(item => typeof item === 'string')))
throw new Error('Prompt must be a string or string[]');
if (typeof generationConfig !== 'object')
throw new Error('Options must be an object');
if (callback !== undefined && typeof callback !== 'function')
throw new Error('Callback must be a function');
const options = {};
if (!callback) {
options['disableStreamer'] = true;
}
return new Promise((resolve) => {
const chunkOutput = (isDone, subword) => {
if (isDone) {
resolve(subword);
}
else if (callback) {
return callback(subword);
}
return StreamingStatus.RUNNING;
};
this.pipeline.generate(prompt, chunkOutput, generationConfig, options);
});
}
}
//# sourceMappingURL=llmPipeline.js.map