openai-compat-endpoint

import { type Chat, type GeneratorController } from "@lmstudio/sdk";
import { loadPresetByName } from "./presets";
import OpenAI from "openai";
import {
  type ChatCompletionMessageParam,
  type ChatCompletionMessageToolCall,
  type ChatCompletionTool,
  type ChatCompletionToolMessageParam,
} from "openai/resources/index";
import type { ChatCompletionCreateParamsStreaming } from "openai/resources/chat/completions";
import { configSchematics, globalConfigSchematics } from "./config";

/* -------------------------------------------------------------------------- */
/*                                   Types                                    */
/* -------------------------------------------------------------------------- */

type ToolCallState = {
  id: string;
  name: string | null;
  index: number;
  arguments: string;
};

/* -------------------------------------------------------------------------- */
/*                               Build helpers                              /* -------------------------------------------------------------------------- */

function createOpenAI(globalConfig: any) {
  const baseURL = globalConfig?.get("baseUrl") || "";
  const apiKey = globalConfig?.get("apiKey");
  return new OpenAI({
    apiKey,
    baseURL,
  });
}

/* -------------------------------------------------------------------------- */
/*                              Message conversion                           /* -------------------------------------------------------------------------- */

/** Convert internal chat history to the format expected by OpenAI.
    If overrideSystemPrompt is provided, the first system message's content is replaced;
    subsequent system messages are kept as-is. */

function toOpenAIMessages(
  history: Chat,
  overrideSystemPrompt: string | null = null,
): ChatCompletionMessageParam[] {
  const messages: ChatCompletionMessageParam[] = [];
  let systemOverrideUsed = false;

  for (const message of history) {
    switch (message.getRole()) {
      case "system": {
        const content =
          (!systemOverrideUsed && overrideSystemPrompt !== null)
            ? overrideSystemPrompt
            : message.getText();
        if (content) {
          messages.push({ role: "system", content });
          if (!systemOverrideUsed && overrideSystemPrompt !== null) {
            systemOverrideUsed = true;
          }
        }
        break;
      }

      case "user":
        messages.push({ role: "user", content: message.getText() });
        break;

      case "assistant": {
        const toolCalls: ChatCompletionMessageToolCall[] = message
          .getToolCallRequests()
          .map(toolCall => ({
            id: toolCall.id ?? "",
            type: "function",
            function: {
              name: toolCall.name,
              arguments: JSON.stringify(toolCall.arguments ?? {}),
            },
          }));

        messages.push({
          role: "assistant",
          content: message.getText(),
          ...(toolCalls.length ? { tool_calls: toolCalls } : {}),
        });
        break;
      }

      case "tool": {
        message.getToolCallResults().forEach(toolCallResult => {
          messages.push({
            role: "tool",
            tool_call_id: toolCallResult.toolCallId ?? "",
            content: toolCallResult.content,
          } as ChatCompletionToolMessageParam);
        });
        break;
      }
    }
  }

  return messages;
}

/* -------------------------------------------------------------------------- */
/*                              Stream-handling utils                         /* -------------------------------------------------------------------------- */

function wireAbort(ctl: GeneratorController, stream: { controller: AbortController }) {
  ctl.onAborted(() => {
    console.info("Generation aborted by user.");
    stream.controller.abort();
  });
}

async function consumeStream(stream: AsyncIterable<any>, ctl: GeneratorController) {
  let current: ToolCallState | null = null;

  function maybeFlushCurrentToolCall() {
    if (current === null || current.name === null) {
      return;
    }
    ctl.toolCallGenerationEnded({
      type: "function",
      name: current.name,
      arguments: JSON.parse(current.arguments),
      id: current.id,
    });
    current = null;
  }

  for await (const chunk of stream) {
    const delta = chunk.choices?.[0]?.delta as
      | {
          content?: string;
          tool_calls?: Array<{
            index: number;
            id?: string;
            function?: { name?: string; arguments?: string };
          }>;
        }
      | undefined;

    if (!delta) continue;

    /* Text streaming */
    if (delta.content) {
      ctl.fragmentGenerated(delta.content);
    }

    /* Tool-call streaming */
    for (const toolCall of delta.tool_calls ?? []) {
      if (toolCall.id !== undefined) {
        maybeFlushCurrentToolCall();
        current = { id: toolCall.id, name: null, index: toolCall.index, arguments: "" };
        ctl.toolCallGenerationStarted();
      }

      if (toolCall.function?.name && current) {
        current.name = toolCall.function.name;
        ctl.toolCallGenerationNameReceived(toolCall.function.name);
      }

      if (toolCall.function?.arguments && current) {
        current.arguments += toolCall.function.arguments;
        ctl.toolCallGenerationArgumentFragmentGenerated(toolCall.function.arguments);
      }
    }

    /* Finalize tool call */
    if (chunk.choices?.[0]?.finish_reason === "tool_calls" && current?.name) {
      maybeFlushCurrentToolCall();
    }
  }

  console.info("Generation completed.");
}

/* -------------------------------------------------------------------------- */
/*                                     API                                    /* -------------------------------------------------------------------------- */

export async function generate(ctl: GeneratorController, history: Chat) {
  const config = ctl.getPluginConfig(configSchematics as any) as any;
  const globalConfig = ctl.getGlobalPluginConfig(globalConfigSchematics as any) as any;
  const debug = Boolean(config.get("debug"));

  // Determine system prompt override from preset or custom text
  const selectedPreset = config.get("systemPromptPreset") as string;
  const customSystemPrompt = config.get("customSystemPrompt") as string;
  let overrideSystemPrompt: string | null = null;
  if (customSystemPrompt) {
    overrideSystemPrompt = customSystemPrompt;
  } else if (selectedPreset && selectedPreset !== "none") {
    const preset = loadPresetByName(selectedPreset);
    if (preset) {
      overrideSystemPrompt = preset.systemPrompt;
    }
  }

  if (debug) {
    console.log(`[DEBUG] System prompt override: ${overrideSystemPrompt ? `${overrideSystemPrompt.length} chars` : "none"}`);
  }

  let model = config.get("model");
  model = typeof model === "string" ? model.trim() : "";

  if (!model) {
    ctl.fragmentGenerated("⚠ Missing model. Set a model ID in plugin settings.\n");
    return;
  }

  const openai = createOpenAI(globalConfig);
  const messages = toOpenAIMessages(history, overrideSystemPrompt);
  const tools = toOpenAITools(ctl);
  const requestPayload: ChatCompletionCreateParamsStreaming = {
    model,
    messages,
    tools,
    stream: true,
  };

  try {
    if (debug) {
      console.info("[PLUGIN] Request payload:", JSON.stringify(requestPayload, null, 2));
    }

    const stream = await openai.chat.completions.create(requestPayload);
    wireAbort(ctl, stream as any);
    await consumeStream(stream, ctl);
  } catch (error: unknown) {
    const err = error && typeof error === "object" ? (error as any) : null;
    const status = err?.status ?? err?.response?.status ?? (typeof err?.code === "number" ? err.code : undefined);
    let msg = "⚠ Generation failed.";
    if (status === 429) {
      msg = "⚠ 429 Rate Limit Exceeded. Try again later or add your own API key.";
    } else if (typeof err?.message === "string") {
      msg = `⚠ API error: ${err.message}`;
    }
    ctl.fragmentGenerated(`${msg}\n`);
    return;
  }
}

function toOpenAITools(ctl: GeneratorController): ChatCompletionTool[] | undefined {
  const tools = ctl.getToolDefinitions().map<ChatCompletionTool>(t => ({
    type: "function",
    function: {
      name: t.function.name,
      description: t.function.description,
      parameters: t.function.parameters ?? {},
    },
  }));
  return tools.length ? tools : undefined;
}