99 lines
3.0 KiB
TypeScript
99 lines
3.0 KiB
TypeScript
import { workerData, parentPort } from 'node:worker_threads';
|
|
import Database from 'better-sqlite3';
|
|
import { EmbeddingService } from '$lib/server/embeddings/embedding.service.js';
|
|
import { createProviderFromProfile } from '$lib/server/embeddings/registry.js';
|
|
import { EmbeddingProfileMapper } from '$lib/server/mappers/embedding-profile.mapper.js';
|
|
import { EmbeddingProfileEntity, type EmbeddingProfileEntityProps } from '$lib/server/models/embedding-profile.js';
|
|
import type { EmbedWorkerRequest, EmbedWorkerResponse, WorkerInitData } from './worker-types.js';
|
|
|
|
const { dbPath, embeddingProfileId } = workerData as WorkerInitData;
|
|
|
|
if (!embeddingProfileId) {
|
|
parentPort!.postMessage({
|
|
type: 'embed-failed',
|
|
jobId: 'init',
|
|
error: 'embeddingProfileId is required in workerData'
|
|
} satisfies EmbedWorkerResponse);
|
|
process.exit(1);
|
|
}
|
|
|
|
const db = new Database(dbPath);
|
|
db.pragma('journal_mode = WAL');
|
|
db.pragma('foreign_keys = ON');
|
|
db.pragma('busy_timeout = 5000');
|
|
db.pragma('synchronous = NORMAL');
|
|
db.pragma('cache_size = -65536');
|
|
db.pragma('temp_store = MEMORY');
|
|
db.pragma('mmap_size = 268435456');
|
|
db.pragma('wal_autocheckpoint = 1000');
|
|
|
|
// Load the embedding profile from DB
|
|
const rawProfile = db.prepare('SELECT * FROM embedding_profiles WHERE id = ?').get(embeddingProfileId);
|
|
|
|
if (!rawProfile) {
|
|
db.close();
|
|
parentPort!.postMessage({
|
|
type: 'embed-failed',
|
|
jobId: 'init',
|
|
error: `Embedding profile ${embeddingProfileId} not found`
|
|
} satisfies EmbedWorkerResponse);
|
|
process.exit(1);
|
|
}
|
|
|
|
const profileEntity = new EmbeddingProfileEntity(rawProfile as EmbeddingProfileEntityProps);
|
|
const profile = EmbeddingProfileMapper.fromEntity(profileEntity);
|
|
|
|
// Create provider and embedding service
|
|
const provider = createProviderFromProfile(profile);
|
|
const embeddingService = new EmbeddingService(db, provider, embeddingProfileId);
|
|
|
|
// Signal ready after service initialization
|
|
parentPort!.postMessage({
|
|
type: 'ready'
|
|
} satisfies EmbedWorkerResponse);
|
|
|
|
parentPort!.on('message', async (msg: EmbedWorkerRequest) => {
|
|
if (msg.type === 'shutdown') {
|
|
db.close();
|
|
process.exit(0);
|
|
}
|
|
|
|
if (msg.type === 'embed') {
|
|
try {
|
|
const snippetIds = embeddingService.findSnippetIdsMissingEmbeddings(
|
|
msg.repositoryId,
|
|
msg.versionId
|
|
);
|
|
|
|
await embeddingService.embedSnippets(snippetIds, (done: number, total: number) => {
|
|
parentPort!.postMessage({
|
|
type: 'embed-progress',
|
|
jobId: msg.jobId,
|
|
done,
|
|
total
|
|
} satisfies EmbedWorkerResponse);
|
|
});
|
|
|
|
parentPort!.postMessage({
|
|
type: 'embed-done',
|
|
jobId: msg.jobId
|
|
} satisfies EmbedWorkerResponse);
|
|
} catch (err) {
|
|
parentPort!.postMessage({
|
|
type: 'embed-failed',
|
|
jobId: msg.jobId,
|
|
error: err instanceof Error ? err.message : String(err)
|
|
} satisfies EmbedWorkerResponse);
|
|
}
|
|
}
|
|
});
|
|
|
|
process.on('uncaughtException', (err) => {
|
|
parentPort!.postMessage({
|
|
type: 'embed-failed',
|
|
jobId: 'uncaught',
|
|
error: err instanceof Error ? err.message : String(err)
|
|
} satisfies EmbedWorkerResponse);
|
|
process.exit(1);
|
|
});
|