multi-role

import { type LMStudioClient } from "@lmstudio/sdk";
import { readFile, writeFile, mkdir } from "fs/promises";
import { join } from "path";

export type Role = "planner" | "executor" | "embedder";

interface RoleConfig {
  model: string;
  systemPrompt?: string;
}

interface Config {
  roles: Partial<Record<Role, RoleConfig>>;
  swapMode: boolean;
}

const DEFAULT_SYSTEM_PROMPTS: Record<Role, string> = {
  planner: `You are a planning expert. When given a task, break it down into clear, numbered, actionable steps.
Be thorough and precise. Each step must be self-contained and executable by another model.
Return ONLY the numbered plan, no extra commentary.`,

  executor: `You are an efficient executor. You receive a single concrete step and execute it directly.
Produce the exact output requested (code, text, data) without unnecessary explanation.
Be concise and precise.`,

  embedder: `You are a document analysis assistant. Answer questions based solely on the provided context.
If the answer is not in the context, say so clearly.`,
};

export class RoleManager {
  private config: Config = { roles: {}, swapMode: true };
  private configPath: string;

  constructor(private readonly client: LMStudioClient, storageDir: string) {
    this.configPath = join(storageDir, ".multi-role-config.json");
  }

  async loadConfig(): Promise<void> {
    try {
      const raw = await readFile(this.configPath, "utf-8");
      this.config = JSON.parse(raw);
    } catch {
      // No config yet — start with defaults
    }
  }

  async saveConfig(): Promise<void> {
    await mkdir(join(this.configPath, ".."), { recursive: true });
    await writeFile(this.configPath, JSON.stringify(this.config, null, 2), "utf-8");
  }

  setRole(role: Role, modelId: string, systemPrompt?: string): void {
    this.config.roles[role] = {
      model: modelId,
      systemPrompt: systemPrompt ?? DEFAULT_SYSTEM_PROMPTS[role],
    };
  }

  getRole(role: Role): RoleConfig | undefined {
    return this.config.roles[role];
  }

  getAllRoles(): Partial<Record<Role, RoleConfig>> {
    return this.config.roles;
  }

  setSwapMode(enabled: boolean): void {
    this.config.swapMode = enabled;
  }

  isSwapMode(): boolean {
    return this.config.swapMode;
  }

  async callLLM(role: Role, userMessage: string, extraContext?: string): Promise<string> {
    const cfg = this.config.roles[role];
    if (!cfg) throw new Error(`Role "${role}" has no model assigned. Use set_role first.`);

    const model = await this.client.llm.get(cfg.model);
    const systemPrompt = cfg.systemPrompt ?? DEFAULT_SYSTEM_PROMPTS[role];

    const messages: Array<{ role: "system" | "user" | "assistant"; content: string }> = [
      { role: "system", content: systemPrompt },
    ];

    if (extraContext) {
      messages.push({ role: "user", content: `Context:\n${extraContext}\n\nTask:\n${userMessage}` });
    } else {
      messages.push({ role: "user", content: userMessage });
    }

    const response = await model.respond(messages);
    return typeof response === "string" ? response : (response as any).content ?? String(response);
  }

  async callEmbedder(text: string): Promise<number[]> {
    const cfg = this.config.roles["embedder"];
    if (!cfg) throw new Error(`Role "embedder" has no model assigned. Use set_role first.`);

    try {
      const embModel = await this.client.embedding.get(cfg.model);
      const result = await embModel.embed(text);
      return Array.isArray(result) ? result : (result as any).embedding ?? [];
    } catch (err: any) {
      throw new Error(`Embedding failed: ${err?.message ?? String(err)}`);
    }
  }
}