openai-compat-endpoint

// src/generator.ts
import { globalConfigSchematics, createConfigSchema } from "./schema";
import { type Chat, type GeneratorController, createConfigSchematics } from "@lmstudio/sdk";
import OpenAI from "openai";
import {
  type ChatCompletionMessageParam,
  type ChatCompletionMessageToolCall,
  type ChatCompletionTool,
  type ChatCompletionToolMessageParam,
} from "openai/resources/index";
import { getModels } from "./models";

/* -------------------------------------------------------------------------- */
/*                              Global Vars                                   */
/* -------------------------------------------------------------------------- */

const MAX_REQUESTS = 25;

// before 1st using of function
function getFormattedTime(): string {
  const now = new Date();
  const dateStr = now.toLocaleDateString('ru-RU');
  const timeStr = now.toLocaleTimeString('ru-RU', { hour12: false });
  return `${dateStr}, ${timeStr.split(':').slice(0, 2).join(':')} (UTC+3)`;
}

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

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

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

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

/** Convert internal chat history to the format expected by OpenAI. */

function toOpenAIMessages(history: Chat): ChatCompletionMessageParam[] {
  const messages: ChatCompletionMessageParam[] = [];

  for (const message of history) {
    switch (message.getRole()) {
      case "system":
        messages.push({ role: "system", content: message.getText() });
        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;
}

/** Convert LM Studio tool definitions to OpenAI function-tool descriptors. */
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;
}

/* -------------------------------------------------------------------------- */
/*                            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) {
    //console.info("Received chunk:", JSON.stringify(chunk));
    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 randomId = Math.random().toString(36).slice(-6);
 console.log(`[GENERATOR VERSION CHECK] Random ID: ${randomId}`);
 console.log(`[DEBUG] [ENTER] generate() PID=${randomId}`);

  console.log('[Generator] Creating runtime config schema (string type for model)');
  const runtimeConfigSchema = createConfigSchematics()
    .field("model", "string", { displayName: "Model" }, "auto")
    .field("customModel", "string", { displayName: "Custom Model ID" }, "")
    .field("onlyFreeModels", "boolean", { displayName: "Only Free Models" }, true)
    .build();
    
  console.log('[Generator] Getting plugin config');
  const config = ctl.getPluginConfig(runtimeConfigSchema as any) as any;
  const globalConfig = ctl.getGlobalPluginConfig(globalConfigSchematics as any) as any;
  console.log('[Generator] Config received');

  let requestCounter = 1;
  const allMessages = Array.from(history);

  for (let i = allMessages.length - 1; i >= 0; i--) {
    const msg = allMessages[i];
    if (msg.getRole() === "assistant") {
      const content = msg.getText();
      if (content) {
        const match = content.match(/✅ Request #(\d+)/);
        if (match) {
          requestCounter = parseInt(match[1], 10) + 1;
          break;
        }
      }
    }
  }

  console.log('[PLUGIN] History length:', Array.from(history).length);
  const lastMsg = Array.from(history).slice(-1)[0];
  if (lastMsg) console.log('[PLUGIN] Last role:', lastMsg.getRole(), 'preview:', lastMsg.getText()?.substring(0, 50));
  console.log('[PLUGIN] Parsed counter:', requestCounter);



  const apiKey = globalConfig.get("apiKey");
  const baseUrl = globalConfig.get("baseUrl");
  const modelFromSelect = config.get("model");
  const customModel = config.get("customModel") ?? "";
  const onlyFree = config.get("onlyFreeModels") ?? true;

  console.log('[Generator] Config values:');
  console.log('  - model (select):', modelFromSelect);
  console.log('  - customModel:', customModel);
  console.log('  - onlyFreeModels:', onlyFree);

  let model: string;
  if (customModel && customModel.trim()) {
    model = customModel.trim();
    console.info("✅ Using model from customModel:", model);
  } else if (modelFromSelect && modelFromSelect !== "auto") {
    model = modelFromSelect;
    console.info("✅ Using model from select:", model);
  } else {
    console.log('[Generator] Using auto, fetching models...');
    const freeModels = await getModels(ctl, onlyFree);
    model = freeModels[0] || "x-ai/grok-4.1-fast:free";
    if (freeModels.length === 0) {
      console.warn('⚠️ No models available in cache. Using fallback model:', model);
      console.warn('⚠️ Restart the plugin to load models from API.');
    }
    console.info("✅ Using auto model:", model);
  }
  console.log('[Generator] ✅ Final model selected:', model);
  model = model.trim();

  const openai = createOpenAI(globalConfig);
  const messages = toOpenAIMessages(history);
  const tools = toOpenAITools(ctl);

  const timeStr = getFormattedTime();

try {
  const stream = await openai.chat.completions.create({
    model,
    messages,
    tools,
    stream: true
  });

  wireAbort(ctl, stream);
  await consumeStream(stream, ctl);

  const timeStr = getFormattedTime();
  ctl.fragmentGenerated(`\n✅ Request #${requestCounter}/${MAX_REQUESTS} at ${timeStr}\n`);

} catch (error: unknown) {
  let msg = "❌ Generation failed.";
  if (typeof error === "object" && error !== null) {
    if ("status" in error && (error as any).status === 429) {
      msg = `❌ 429 Rate Limit Exceeded. You've used ${requestCounter}/${MAX_REQUESTS} free requests. Try again later or add your own API key.`;
      console.log('[PLUGIN] API key Limit exceeded/ Error: ', (error as any).status);
    } else if ("message" in error && typeof (error as any).message === "string") {
      const message = (error as any).message;
      if (message.includes("API Key")) {
        msg = "❌ Invalid or missing API key.";
      } else {
        msg = `❌ API error: ${message}`;
      }
    }
  }

  ctl.fragmentGenerated(`${msg}\n`);

  return;
}

}