diff --git a/packages/core/src/ai/embedding-worker.ts b/packages/core/src/ai/embedding-worker.ts index e62f99bf..d1be3c9d 100644 --- a/packages/core/src/ai/embedding-worker.ts +++ b/packages/core/src/ai/embedding-worker.ts @@ -19,6 +19,25 @@ let pipeline: any = null; let currentModelId: string | null = null; +const MODEL_LOAD_MAX_ATTEMPTS = 3; +const MODEL_LOAD_RETRY_BASE_MS = 500; +const DEFAULT_REMOTE_HOSTS = ["https://hf-mirror.com/", "https://huggingface.co/"]; + +function getRemoteHosts(): string[] { + const configured = (globalThis as { __READANY_HF_REMOTE_HOST__?: unknown }) + .__READANY_HF_REMOTE_HOST__; + const hosts = + typeof configured === "string" ? [configured, ...DEFAULT_REMOTE_HOSTS] : DEFAULT_REMOTE_HOSTS; + return [...new Set(hosts.map((host) => `${host.trim().replace(/\/+$/, "")}/`).filter(Boolean))]; +} + +function describeError(error: unknown): { type: string; message: string } { + if (error instanceof Error) return { type: error.name || "Error", message: error.message }; + return { type: typeof error, message: String(error) }; +} + +const wait = (milliseconds: number) => new Promise((resolve) => setTimeout(resolve, milliseconds)); + self.onmessage = async (e: MessageEvent) => { const msg = e.data; @@ -34,43 +53,77 @@ self.onmessage = async (e: MessageEvent) => { }; async function handleLoad(modelId: string, hfModelId: string) { - try { - // Reuse if same model already loaded - if (pipeline && currentModelId === modelId) { - self.postMessage({ type: "load:done" }); - return; - } + let lastError: unknown; + const remoteHosts = getRemoteHosts(); - // Dispose previous pipeline - if (pipeline) { - try { - await pipeline.dispose?.(); - } catch { - /* ignore */ + for (let attempt = 1; attempt <= MODEL_LOAD_MAX_ATTEMPTS; attempt++) { + const remoteHost = remoteHosts[(attempt - 1) % remoteHosts.length]; + try { + // Reuse if same model already loaded + if (pipeline && currentModelId === modelId) { + self.postMessage({ type: "load:done" }); + return; } - pipeline = null; - currentModelId = null; - } - - const { pipeline: createPipeline, env } = await import("@huggingface/transformers"); - env.allowLocalModels = false; - pipeline = await createPipeline("feature-extraction", hfModelId, { - progress_callback: (p: any) => { - if (p.status === "progress") { - self.postMessage({ type: "load:progress", progress: Math.round(p.progress ?? 0) }); + // Dispose previous pipeline + if (pipeline) { + try { + await pipeline.dispose?.(); + } catch { + /* ignore */ } - }, - }); + pipeline = null; + currentModelId = null; + } - currentModelId = modelId; - self.postMessage({ type: "load:done" }); - } catch (err) { - self.postMessage({ - type: "load:error", - error: err instanceof Error ? err.message : String(err), - }); + const { pipeline: createPipeline, env } = await import("@huggingface/transformers"); + env.allowLocalModels = false; + env.remoteHost = remoteHost; + + console.info("[EmbeddingWorker] loading model", { + modelId, + hfModelId, + attempt, + remoteHost, + }); + + pipeline = await createPipeline("feature-extraction", hfModelId, { + progress_callback: (p: any) => { + if (p.status === "progress") { + self.postMessage({ type: "load:progress", progress: Math.round(p.progress ?? 0) }); + } + }, + }); + + currentModelId = modelId; + self.postMessage({ type: "load:done" }); + return; + } catch (err) { + lastError = err; + const details = describeError(err); + console.warn("[EmbeddingWorker] model load failed", { + modelId, + hfModelId, + attempt, + remoteHost, + errorType: details.type, + error: details.message, + }); + + // A partial model download can poison subsequent attempts. Clear only this + // model before retrying, then switch to the next configured host. + if (attempt < MODEL_LOAD_MAX_ATTEMPTS) { + await clearModelCacheEntries(hfModelId); + await wait(MODEL_LOAD_RETRY_BASE_MS * 2 ** (attempt - 1)); + } + } } + + const details = describeError(lastError); + self.postMessage({ + type: "load:error", + error: `模型下载失败(${details.type}):${details.message}`, + }); } async function handleEmbed(requestId: string, texts: string[]) { @@ -136,19 +189,7 @@ async function handleClearCache(hfModelId: string) { } // Transformers.js uses the Cache API with cache name "transformers-cache" - const cacheNames = await caches.keys(); - let deletedCount = 0; - for (const cacheName of cacheNames) { - const cache = await caches.open(cacheName); - const keys = await cache.keys(); - for (const key of keys) { - // Match URLs containing the HuggingFace model ID - if (key.url.includes(hfModelId)) { - await cache.delete(key); - deletedCount++; - } - } - } + const deletedCount = await clearModelCacheEntries(hfModelId); self.postMessage({ type: "clearCache:done", deletedCount }); } catch (err) { @@ -158,3 +199,18 @@ async function handleClearCache(hfModelId: string) { }); } } + +async function clearModelCacheEntries(hfModelId: string): Promise { + if (typeof caches === "undefined") return 0; + const cacheNames = await caches.keys(); + let deletedCount = 0; + for (const cacheName of cacheNames) { + const cache = await caches.open(cacheName); + for (const key of await cache.keys()) { + if (key.url.includes(hfModelId) && (await cache.delete(key))) deletedCount++; + } + } + if (deletedCount > 0) + console.info("[EmbeddingWorker] cleared cached model files", { hfModelId, deletedCount }); + return deletedCount; +}