rag-v2-improved

import {
  text,
  type ChatMessage,
  type FileHandle,
  type PromptPreprocessorController,
} from "@lmstudio/sdk";

import { configSchematics } from "./config";
import { MemoryService } from "./memory/MemoryService";
import { buildDocumentContext } from "./templates/document_context";

async function retrieveContext(
  ctl: PromptPreprocessorController,
  prompt: string,
  files: FileHandle[],
) {
  const pluginConfig = ctl.getPluginConfig(configSchematics);

  const retrievalLimit = pluginConfig.get("retrievalLimit");

  const embeddingModel = await ctl.client.embedding.model(
    "nomic-ai/nomic-embed-text-v1.5-GGUF",
    { signal: ctl.abortSignal }
  );

  return await ctl.client.files.retrieve(prompt, files, {
    embeddingModel,
    limit: retrievalLimit,
    signal: ctl.abortSignal,
  });
}

export async function preprocess(
  ctl: PromptPreprocessorController,
  userMessage: ChatMessage,
) {
  const userPrompt = userMessage.getText();

  const history = await ctl.pullHistory();
  history.append(userMessage);

  const files = history
    .getAllFiles(ctl.client)
    .filter(file => file.type !== "image");

  if (files.length === 0) {
    return userMessage;
  }

  const retrieval = await retrieveContext(
    ctl,
    userPrompt,
    files,
  );

  const chunks = retrieval.entries
    .map(entry => entry.content)
    .join("\n\n");

  const citations = retrieval.entries
    .map((entry, index) => `[${index + 1}] ${entry.fileName}`)
    .join("\n");

  const memory = MemoryService.retrieve("default");

  const finalPrompt = buildDocumentContext(
    citations,
    `${memory}\n\n${chunks}`,
    userPrompt,
  );

  MemoryService.save(
    "default",
    `Processed query: ${userPrompt.slice(0, 120)}`
  );

  return text(finalPrompt);
}