Add shared utilities for generating text with WebLLM

77ab694ea0326f381a97057e320723d530aac74e

Cohee <18619528+Cohee1207@users.noreply.github.com>

1 files changed, +83 -0Showing whitespace changes
public/scripts/extensions/shared.js+83 -0
@@ -176,3 +176,86 @@ function throwIfInvalidModel(useReverseProxy) {
176176 throw new Error('Custom API URL is not set.');
177177 }
178178}
179+
180+/**
181+ * Check if the WebLLM extension is installed and supported.
182+ * @returns {boolean} Whether the extension is installed and supported
183+ */
184+export function isWebLlmSupported() {
185+ if (!('gpu' in navigator)) {
186+ toastr.error('Your browser does not support the WebGPU API. Please use a different browser.', 'WebLLM', {
187+ preventDuplicates: true,
188+ timeOut: 0,
189+ extendedTimeOut: 0,
190+ });
191+ return false;
192+ }
193+
194+ if (!('llm' in SillyTavern)) {
195+ toastr.error('WebLLM extension is not installed. Click here to install it.', 'WebLLM', {
196+ timeOut: 0,
197+ extendedTimeOut: 0,
198+ preventDuplicates: true,
199+ onclick: () => {
200+ const button = document.getElementById('third_party_extension_button');
201+ if (button) {
202+ button.click();
203+ }
204+
205+ const input = document.querySelector('dialog textarea');
206+
207+ if (input instanceof HTMLTextAreaElement) {
208+ input.value = 'https://github.com/SillyTavern/Extension-WebLLM';
209+ }
210+ },
211+ });
212+ return false;
213+ }
214+
215+ return true;
216+}
217+
218+/**
219+ * Generates text in response to a chat prompt using WebLLM.
220+ * @param {any[]} messages Messages to use for generating
221+ * @returns {Promise<string>} Generated response
222+ */
223+export async function generateWebLlmChatPrompt(messages) {
224+ if (!isWebLlmSupported()) {
225+ throw new Error('WebLLM extension is not installed.');
226+ }
227+
228+ const engine = SillyTavern.llm;
229+ const response = await engine.generateChatPrompt(messages);
230+ return response;
231+}
232+
233+/**
234+ * Counts the number of tokens in the provided text using WebLLM's default model.
235+ * @param {string} text Text to count tokens in
236+ * @returns {Promise<number>} Number of tokens in the text
237+ */
238+export async function countWebLlmTokens(text) {
239+ if (!isWebLlmSupported()) {
240+ throw new Error('WebLLM extension is not installed.');
241+ }
242+
243+ const engine = SillyTavern.llm;
244+ const response = await engine.countTokens(text);
245+ return response;
246+}
247+
248+/**
249+ * Gets the size of the context in the WebLLM's default model.
250+ * @returns {Promise<number>} Size of the context in the WebLLM model
251+ */
252+export async function getWebLlmContextSize() {
253+ if (!isWebLlmSupported()) {
254+ throw new Error('WebLLM extension is not installed.');
255+ }
256+
257+ const engine = SillyTavern.llm;
258+ await engine.loadModel();
259+ const model = await engine.getCurrentModelInfo();
260+ return model?.context_size;
261+}