user-docs

/**
 * Shared Retriever singleton.
 *
 * Pattern: lazy on-demand init, onStatus forwarded directly to ctx.status()
 * by the caller.
 *
 * Lifecycle:
 *   1. find_doc calls ensureRetriever(ctl, onStatus) on every tool invocation.
 *   2. First call: runs venv setup + initialize + waitForIndexing, streaming
 *      all status through onStatus.
 *   3. Subsequent calls: _instance already set → returns immediately.
 */

import path from "path";
import { type ToolsProviderController } from "@lmstudio/sdk";
import { globalConfigSchematics } from "../config.js";
import { Retriever } from "./retriever.js";
import { ensureRagVenv } from "../utils/ragVenvSetup.js";
import { ragDebug } from "../utils/ragLogger.js";
import { defaultLmStudioHome } from "../sources/lmStudioConversationMarkdown.js";

let _instance: Retriever | null = null;
let _initPromise: Promise<Retriever> | null = null;
let _updatePromise: Promise<void> | null = null;

function buildEffectiveDirs(gcfg: { get: (k: string) => any }): string[] {
  const notesDir: string = gcfg.get("notesDirectory") ?? "";
  const configuredDirs: string[] = gcfg.get("contentDirectories") ?? [];
  const remoteSources: string[] = gcfg.get("remoteSources") ?? [];
  const set = new Set<string>();
  if (notesDir) set.add(notesDir);
  for (const dir of configuredDirs) set.add(dir);
  for (const source of remoteSources) {
    const watchPath = conversationWatchPathFromSource(source);
    if (watchPath) set.add(watchPath);
  }
  return Array.from(set);
}

function dirsEqual(a: string[], b: string[]): boolean {
  if (a.length !== b.length) return false;
  const setA = new Set(a);
  return b.every((d) => setA.has(d));
}

function chatIdFromWorkingDirectory(ctl: ToolsProviderController): string | null {
  try {
    const workingDir = ctl.getWorkingDirectory();
    if (typeof workingDir !== "string" || !workingDir.trim()) return null;
    const chatId = path.basename(workingDir);
    return /^\d{13}$/.test(chatId) ? chatId : null;
  } catch {
    return null;
  }
}

function conversationWatchPathFromSource(source: string): string | null {
  const trimmed = source.trim();
  if (/^lmstudio-conversations?:\/\/(?:\d{13})?$/i.test(trimmed)) {
    const chatId = trimmed.match(/(\d{13})$/)?.[1];
    return chatId
      ? path.join(defaultLmStudioHome(), "conversations", `${chatId}.conversation.json`)
      : path.join(defaultLmStudioHome(), "conversations");
  }
  if (path.isAbsolute(trimmed) && path.basename(trimmed) === "conversations") return trimmed;
  if (path.isAbsolute(trimmed) && /(?:^|\/)(\d{13})\.conversation\.json$/i.test(trimmed)) return trimmed;
  return null;
}

export function getRetriever(): Retriever {
  if (!_instance) throw new Error("[RAG] Retriever not initialised yet.");
  return _instance;
}

export function isRetrieverReady(): boolean {
  return _instance !== null;
}

export async function ensureRetriever(
  ctl: ToolsProviderController,
  onStatus: (msg: string) => void
): Promise<Retriever> {
  if (_instance) {
    // Wait for any in-progress directory update to finish.
    if (_updatePromise) await _updatePromise;

    // Check whether effective content directories have changed since init.
    const getter: any = (ctl as any).getGlobalPluginConfig || (ctl as any).getGlobalConfig;
    if (getter) {
      const gcfg = getter.call(ctl, globalConfigSchematics);
      const newDirs = buildEffectiveDirs(gcfg);
      const oldDirs = _instance.getContentDirectories();
      if (!dirsEqual(oldDirs, newDirs)) {
        ragDebug("RAG", `Content directories changed (${oldDirs.length} → ${newDirs.length}), updating index…`);
        _updatePromise = _instance.updateContentDirectories(newDirs).finally(() => {
          _updatePromise = null;
        });
        await _updatePromise;
      }
    }
    return _instance;
  }

  // Another concurrent call already started init — join it.
  if (_initPromise) return _initPromise;

  _initPromise = (async () => {
    const getter: any =
      (ctl as any).getGlobalPluginConfig || (ctl as any).getGlobalConfig;
    if (!getter) throw new Error("[RAG] No global config accessor on ToolsProviderController");
    const gcfg = getter.call(ctl, globalConfigSchematics);

    // Step 1: Python venv
    await ensureRagVenv((msg) => {
      onStatus(msg);
      ragDebug("RAG setup", msg);
    });

    // Step 2: Create retriever
    // notesDirectory is always included in the index — it is the primary write
    // target of this plugin and must always be searchable, regardless of whether
    // it overlaps with contentDirectories.
    const dbPath = path.join(process.cwd(), ".rag-data", "vectors.db");
    const remoteSources: string[] = gcfg.get("remoteSources") ?? [];
    const contentDirectories = buildEffectiveDirs(gcfg);
    const activeChatId = chatIdFromWorkingDirectory(ctl);
    const embeddingBaseUrl = String(
      gcfg.get("embeddingBaseUrl") || "http://127.0.0.1:1234/v1"
    ).replace(/\/v1\/?$/, "");
    const retriever = new Retriever({
      dbPath,
      contentDirectories,
      embeddingModel:     gcfg.get("embeddingModel"),
      embeddingBaseUrl,
      embeddingApiKey:    gcfg.get("embeddingApiKey"),
      chunkSize:          gcfg.get("chunkSize"),
      chunkOverlap:       gcfg.get("chunkOverlap"),
      embeddingBatchSize: gcfg.get("embeddingBatchSize"),
      embeddingRequestTimeoutMs: gcfg.get("embeddingRequestTimeoutMs"),
      retrievalLimit:     gcfg.get("retrievalLimit"),
      watchFiles:         gcfg.get("watchFiles"),
      hybridSearchWeight: gcfg.get("hybridSearchWeight"),
      languageBiasStrength: gcfg.get("languageBiasStrength"),
      remoteSources,
      remoteFetchTimeoutMs: gcfg.get("remoteFetchTimeoutMs"),
      remoteMaxBytes: gcfg.get("remoteMaxBytes"),
      remoteMaxPages: gcfg.get("remoteMaxPages"),
      githubToken: gcfg.get("githubToken"),
      huggingFaceToken: gcfg.get("huggingFaceToken"),
      activeChatId: activeChatId ?? undefined,
    });

    retriever.setStatusCallback((text) => {
      onStatus(text);
      ragDebug("RAG", text);
    });

    // Step 3: Initialize (connects to embedding model)
    onStatus("Connecting to embedding model\u2026");
    await retriever.initialize();

    // Step 4: Wait for change events that arrived while initialization was running.
    await retriever.waitForIndexing(null, (filesRemaining, currentFile) => {
      const name = currentFile ? path.basename(currentFile) : "";
      if (filesRemaining > 0) {
        onStatus(`Indexing: ${name} (${filesRemaining} file${filesRemaining !== 1 ? "s" : ""} remaining)`);
        ragDebug("RAG", `Indexing: ${name} (${filesRemaining} left)`);
      }
    });

    retriever.setStatusCallback(null);

    const stats = retriever.getStats();
    const readyMsg = `Document index ready \u2014 ${stats.documentCount} doc${stats.documentCount !== 1 ? "s" : ""}, ${stats.chunkCount} chunk${stats.chunkCount !== 1 ? "s" : ""}`;
    onStatus(readyMsg);
    ragDebug("RAG", `Ready: ${stats.documentCount} docs, ${stats.chunkCount} chunks`);

    _instance = retriever;
    return retriever;
  })();

  try {
    return await _initPromise;
  } catch (err) {
    _initPromise = null; // allow retry on next call
    throw err;
  }
}