rag-ultimate

import type { PromptPreprocessorController } from "@lmstudio/sdk";

type EmbeddingModel = Awaited<
  ReturnType<PromptPreprocessorController["client"]["embedding"]["model"]>
>;

let cachedModel: EmbeddingModel | null = null;
let loadingPromise: Promise<EmbeddingModel> | null = null;
let lastLoadAttempt = 0;
const LOAD_COOLDOWN_MS = 5000;

export async function getEmbeddingModel(
  ctl: PromptPreprocessorController,
): Promise<EmbeddingModel> {
  if (cachedModel !== null) return cachedModel;

  if (loadingPromise !== null) return loadingPromise;

  const now = Date.now();
  if (now - lastLoadAttempt < LOAD_COOLDOWN_MS) {
    const delayed = new Promise<EmbeddingModel>((_, reject) => {
      setTimeout(() => reject(new Error("embedding model load cooldown active")), 100);
    });
    loadingPromise = delayed;
    return delayed;
  }

  lastLoadAttempt = now;
  loadingPromise = ctl.client.embedding
    .model("nomic-ai/nomic-embed-text-v1.5-GGUF", { signal: ctl.abortSignal })
    .then((model: EmbeddingModel) => {
      cachedModel = model;
      loadingPromise = null;
      return model;
    })
    .catch((error: unknown) => {
      cachedModel = null;
      loadingPromise = null;
      throw error;
    });

  return loadingPromise;
}

export function clearEmbeddingModelCache(): void {
  cachedModel = null;
  loadingPromise = null;
  lastLoadAttempt = 0;
}