From cb60a0b0c53ff0c1392321659a761b375b7f8909 Mon Sep 17 00:00:00 2001 From: ztimson Date: Wed, 28 Jan 2026 22:17:39 -0500 Subject: [PATCH] Moved embeddings to worker to prevent blocking --- package.json | 2 +- src/embedder.ts | 11 +++++++++++ src/llm.ts | 32 +++++++++++++++++++++++--------- vite.config.ts | 11 +++++++++-- 4 files changed, 44 insertions(+), 12 deletions(-) create mode 100644 src/embedder.ts diff --git a/package.json b/package.json index 7876703..a3435c3 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@ztimson/ai-utils", - "version": "0.2.5", + "version": "0.2.6", "description": "AI Utility library", "author": "Zak Timson", "license": "MIT", diff --git a/src/embedder.ts b/src/embedder.ts new file mode 100644 index 0000000..9b2c852 --- /dev/null +++ b/src/embedder.ts @@ -0,0 +1,11 @@ +import { pipeline } from '@xenova/transformers'; +import { parentPort } from 'worker_threads'; + +let model: any; + +parentPort?.on('message', async ({ id, text }) => { + if(!model) model = await pipeline('feature-extraction', 'Xenova/all-MiniLM-L6-v2'); + const output = await model(text, { pooling: 'mean', normalize: true }); + const embedding = Array.from(output.data); + parentPort?.postMessage({ id, embedding }); +}); diff --git a/src/llm.ts b/src/llm.ts index d154295..8c9df29 100644 --- a/src/llm.ts +++ b/src/llm.ts @@ -1,4 +1,3 @@ -import {pipeline} from '@xenova/transformers'; import {JSONAttemptParse} from '@ztimson/utils'; import {Ai} from './ai.ts'; import {Anthropic} from './antrhopic.ts'; @@ -6,7 +5,9 @@ import {Ollama} from './ollama.ts'; import {OpenAi} from './open-ai.ts'; import {AbortablePromise, LLMProvider} from './provider.ts'; import {AiTool} from './tools.ts'; -import * as tf from '@tensorflow/tfjs'; +import {Worker} from 'worker_threads'; +import {fileURLToPath} from 'url'; +import {dirname, join} from 'path'; export type LLMMessage = { /** Message originator */ @@ -83,11 +84,22 @@ export type LLMRequest = { } export class LLM { - private embedModel: any; + private embedWorker: Worker | null = null; + private embedQueue = new Map void; reject: (error: any) => void }>(); + private embedId = 0; private providers: {[key: string]: LLMProvider} = {}; + constructor(public readonly ai: Ai) { - this.embedModel = pipeline('feature-extraction', 'Xenova/all-MiniLM-L6-v2'); + this.embedWorker = new Worker(join(dirname(fileURLToPath(import.meta.url)), 'embedder.js')); + this.embedWorker.on('message', ({ id, embedding }) => { + const pending = this.embedQueue.get(id); + if (pending) { + pending.resolve(embedding); + this.embedQueue.delete(id); + } + }); + if(ai.options.anthropic?.token) this.providers.anthropic = new Anthropic(this.ai, ai.options.anthropic.token, ai.options.anthropic.model); if(ai.options.ollama?.host) this.providers.ollama = new Ollama(this.ai, ai.options.ollama.host, ai.options.ollama.model); if(ai.options.openAi?.token) this.providers.openAi = new OpenAi(this.ai, ai.options.openAi.token, ai.options.openAi.model); @@ -159,10 +171,12 @@ export class LLM { }); }; - const embed = async (text: string): Promise => { - const model = await this.embedModel; - const output = await model(text, {pooling: 'mean', normalize: true}); - return Array.from(output.data); + const embed = (text: string): Promise => { + return new Promise((resolve, reject) => { + const id = this.embedId++; + this.embedQueue.set(id, { resolve, reject }); + this.embedWorker?.postMessage({ id, text }); + }); }; // Tokenize @@ -188,7 +202,7 @@ export class LLM { const cleanText = text.replace(/\s*\n\s*/g, '\n').trim(); if(cleanText) chunks.push(cleanText); start = end - overlapTokens; - if (start <= end - tokens.length + end) start = end; // Safety: prevent infinite loop + if(start <= end - tokens.length + end) start = end; } return Promise.all(chunks.map(async (text, index) => ({ diff --git a/vite.config.ts b/vite.config.ts index 467da37..ef51e05 100644 --- a/vite.config.ts +++ b/vite.config.ts @@ -1,12 +1,19 @@ import {defineConfig} from 'vite'; import dts from 'vite-plugin-dts'; +import {resolve} from 'path'; export default defineConfig({ build: { lib: { - entry: './src/index.ts', + entry: { + index: './src/index.ts', + embedder: './src/embedder.ts', + }, name: 'utils', - fileName: (format) => (format === 'es' ? 'index.mjs' : 'index.js'), + fileName: (format, entryName) => { + if (entryName === 'embedder') return 'embedder.js'; + return format === 'es' ? 'index.mjs' : 'index.js'; + }, }, ssr: true, emptyOutDir: true,