system-prompt-autoselector

import os from 'os';
import path from 'path';
import fs from 'fs/promises';
import { Chat, PromptPreprocessorController } from "@lmstudio/sdk";
import { SUB_PROMPT, MAX_TOKENS, PRESET_PATH } from '../const';

type Preset = {
  name: string;
  prompt: string;
}

export async function getSystemPrompts(): Promise< Array<Preset>> {
  const presets: Array<Preset> = [];
  
  const dir = path.join(os.homedir(), PRESET_PATH);
  let presetList;

  try {
    presetList = await fs.readdir(dir);
  } catch {
    console.error(`[preprocessor] cannot read dir: ${dir}`);
    return presets;
  }
  
  for(const fname of presetList) {
    if (!fname.endsWith('.json')) continue;
    
    try {
      const text = await fs.readFile(path.join(dir, fname), { encoding: 'utf8' });
      const json: any = JSON.parse(text);

      const field = json.operation?.fields?.find(f => f.key === 'llm.prediction.systemPrompt');
      if (!field) {
        console.warn(`[preprocessor] no systemPrompt field in preset: ${fname}`);
        continue;
      }
      presets.push({
        name: json.name,
        prompt: field.value
      });
    } catch {
      console.error(`[preprocessor] cannot load system promt: ${fname}`)
    }
  }

  return presets;
}

export function userMessageCount(chat: Chat): number {
  let userMessages = 0;
  chat.map( msg => {
    if( msg.getRole() === 'user' ) userMessages ++;
  });

  return userMessages;
}

export async function getNewSystemPrompt(userMessage: string, ctl: PromptPreprocessorController): Promise<{ prompt: string; preset: string } | null> {
  const loadedModels = await ctl.client.llm.listLoaded();
  if (loadedModels.length === 0) {
    console.warn("[preprocessor] no models, skip pre-call");
    return null;
  }

  const model = loadedModels[0];
  console.info(`[preprocessor] Pre-call using model: ${model.identifier}`);

  let rules = ``;
  const prompts = await getSystemPrompts();

  for(const p of prompts) rules += `${p.name}\n`;

  const prePrompt = SUB_PROMPT.replace('{{USER_MESSAGE}}', userMessage).replace('{{PRESETS}}', rules);

  let result;
  try {
    result = await model.respond(prePrompt, {
      maxTokens: MAX_TOKENS,
      signal: ctl.abortSignal
    });
  } catch (err) {
    console.error(`[preprocessor] model respond failed: ${err}`);
    return null;
  }
  
  const name = result.nonReasoningContent ? result.nonReasoningContent.trim() : result.content.trim();
  console.info(`[preprocessor] use prompt name: ${name}`);

  const newSys = prompts.find( p => p.name === name );

  if(!newSys) return null;

  return { prompt: newSys.prompt, preset: name };
}