rag

import { type ChatMessage, type PromptPreprocessorController } from "@lmstudio/sdk";
import { pluginConfigSchematics } from "./config";
import { getEmbedFn } from "./embedder";
import { getIndex, queryChunks } from "./vectorStore";
import { join } from "path";

export async function promptPreprocessor(
  ctl: PromptPreprocessorController,
  userMessage: ChatMessage,
): Promise<string | ChatMessage> {
  const history = await ctl.pullHistory();
  if (history.length !== 0) return userMessage;

  const cfg = ctl.getPluginConfig(pluginConfigSchematics);

  if (!cfg.get("autoInject")) return userMessage;

  const query = userMessage.getText().trim();
  if (!query) return userMessage;

  try {
    const dp = cfg.get("dataPath").trim() || join(process.env.HOME ?? "~", "rag-data");
    const embId = cfg.get("embeddingModelIdentifier").trim();
    const topK = cfg.get("topK");

    const embed = await getEmbedFn(ctl.client, embId);
    const [queryVec] = await embed([query]);

    const idx = getIndex(dp, "default");
    const results = await queryChunks(idx, queryVec, topK);
    if (results.length === 0) return userMessage;

    const contextBlock = results
      .map((r, i) => `[${i + 1}] (${r.metadata.fileName})\n${r.metadata.text}`)
      .join("\n\n---\n\n");

    const injection = `[Retrieved context from your document index:\n\n${contextBlock}\n]`;
    return `${injection}\n\n${query}`;
  } catch {
    // No embedding model loaded or index empty — pass through silently
    return userMessage;
  }
}