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) {
176 throw new Error('Custom API URL is not set.');176 throw new Error('Custom API URL is not set.');
177 }177 }
178}178}
179
180/**
181 * Check if the WebLLM extension is installed and supported.
182 * @returns {boolean} Whether the extension is installed and supported
183 */
184export 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 */
223export 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 */
238export 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 */
252export 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}