deepseek-endpoint

import { promises as fs } from "node:fs";
import path from "node:path";
import { type Chat, type ChatMessage, type FileHandle, type GeneratorController, type InferParsedConfig } from "@lmstudio/sdk";
import OpenAI from "openai";
import {
  type ChatCompletionContentPart,
  type ChatCompletionContentPartImage,
  type ChatCompletionContentPartText,
  type ChatCompletionMessageParam,
  type ChatCompletionMessageToolCall,
  type ChatCompletionTool,
  type ChatCompletionToolMessageParam,
  type ChatCompletionUserMessageParam,
} from "openai/resources/index";
import { configSchematics, globalConfigSchematics } from "./config";

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

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

type OpenAICompatAssistantMessage = {
  role: "assistant";
  content: string;
  tool_calls?: ChatCompletionMessageToolCall[];
  [key: string]: unknown;
};

type OpenAICompatMessageParam = ChatCompletionMessageParam | OpenAICompatAssistantMessage;

type OpenAICompatDelta = {
  content?: string | null;
  tool_calls?: Array<{
    index: number;
    id?: string;
    function?: { name?: string; arguments?: string };
  }>;
  [key: string]: unknown;
};

const imageMimeTypesByExtension: Record<string, string> = {
  ".bmp": "image/bmp",
  ".gif": "image/gif",
  ".jpeg": "image/jpeg",
  ".jpg": "image/jpeg",
  ".png": "image/png",
  ".webp": "image/webp",
};

const reasoningStartTag = "<think>";
const reasoningEndTag = "</think>";

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

/** Build a pre-configured OpenAI-compatible client. */
function createOpenAI(globalConfig: InferParsedConfig<typeof globalConfigSchematics>) {
  const baseURL = globalConfig.get("baseUrl") || "https://api.deepseek.com";
  const apiKey = globalConfig.get("apiKey");

  return new OpenAI({
    apiKey,
    baseURL,
  });
}

function detectImageMimeType(buffer: Buffer, fileName: string) {
  if (buffer.length >= 8 && buffer.subarray(0, 8).equals(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]))) {
    return "image/png";
  }

  if (buffer.length >= 3 && buffer[0] === 0xff && buffer[1] === 0xd8 && buffer[2] === 0xff) {
    return "image/jpeg";
  }

  if (buffer.length >= 6) {
    const header = buffer.subarray(0, 6).toString("ascii");
    if (header === "GIF87a" || header === "GIF89a") {
      return "image/gif";
    }
  }

  if (
    buffer.length >= 12 &&
    buffer.subarray(0, 4).toString("ascii") === "RIFF" &&
    buffer.subarray(8, 12).toString("ascii") === "WEBP"
  ) {
    return "image/webp";
  }

  if (buffer.length >= 2 && buffer.subarray(0, 2).toString("ascii") === "BM") {
    return "image/bmp";
  }

  return imageMimeTypesByExtension[path.extname(fileName).toLowerCase()] ?? "image/png";
}

function parseToolCallArguments(rawArguments: string) {
  if (rawArguments.trim().length === 0) {
    return {};
  }

  return JSON.parse(rawArguments);
}

function extractReasoningBlocks(text: string) {
  const reasoningParts: string[] = [];
  let visibleContent = "";
  let cursor = 0;

  while (cursor < text.length) {
    const startIndex = text.indexOf(reasoningStartTag, cursor);
    if (startIndex === -1) {
      visibleContent += text.slice(cursor);
      break;
    }

    visibleContent += text.slice(cursor, startIndex);

    const reasoningStartIndex = startIndex + reasoningStartTag.length;
    const endIndex = text.indexOf(reasoningEndTag, reasoningStartIndex);
    if (endIndex === -1) {
      reasoningParts.push(text.slice(reasoningStartIndex));
      cursor = text.length;
      break;
    }

    reasoningParts.push(text.slice(reasoningStartIndex, endIndex));
    cursor = endIndex + reasoningEndTag.length;
  }

  const reasoningContent = reasoningParts.join("");
  return {
    content: visibleContent,
    reasoningContent: reasoningContent.length > 0 ? reasoningContent : undefined,
  };
}

function normalizeReasoningDelta(value: unknown, preferredKey?: string | null): string | null {
  if (typeof value === "string") {
    return value;
  }

  if (Array.isArray(value)) {
    const combined = value
      .map(item => normalizeReasoningDelta(item, preferredKey))
      .filter((item): item is string => item !== null && item.length > 0)
      .join("");

    return combined.length > 0 ? combined : null;
  }

  if (value === null || typeof value !== "object") {
    return null;
  }

  const record = value as Record<string, unknown>;
  const candidateKeys = [preferredKey, "text", "content", "reasoning_content", "delta", "parts"]
    .filter((key): key is string => typeof key === "string" && key.length > 0);
  for (const key of candidateKeys) {
    const normalized = normalizeReasoningDelta(record[key], preferredKey);
    if (normalized !== null && normalized.length > 0) {
      return normalized;
    }
  }

  return null;
}

async function fileToImagePart(file: FileHandle): Promise<ChatCompletionContentPartImage> {
  const filePath = await file.getFilePath();
  const fileBytes = await fs.readFile(filePath);
  const mimeType = detectImageMimeType(fileBytes, file.name);

  return {
    type: "image_url",
    image_url: {
      url: `data:${mimeType};base64,${fileBytes.toString("base64")}`,
    },
  };
}

async function toOpenAIUserMessage(
  message: ChatMessage,
  ctl: GeneratorController,
): Promise<ChatCompletionUserMessageParam> {
  const text = message.getText();
  const imageFiles = message.getFiles(ctl.client).filter(file => file.isImage());

  if (imageFiles.length === 0) {
    return {
      role: "user",
      content: text,
    };
  }

  const parts: ChatCompletionContentPart[] = [];
  if (text.length > 0) {
    const textPart: ChatCompletionContentPartText = {
      type: "text",
      text,
    };
    parts.push(textPart);
  }

  const imageParts = await Promise.all(imageFiles.map(file => fileToImagePart(file)));
  parts.push(...imageParts);

  return {
    role: "user",
    content: parts,
  };
}

/** Convert internal chat history to the format expected by OpenAI. */
async function toOpenAIMessages(
  history: Chat,
  ctl: GeneratorController,
  requestReasoningKey: string | null,
): Promise<OpenAICompatMessageParam[]> {
  const messages: OpenAICompatMessageParam[] = [];

  for (const message of history) {
    switch (message.getRole()) {
      case "system":
        messages.push({ role: "system", content: message.getText() });
        break;

      case "user":
        messages.push(await toOpenAIUserMessage(message, ctl));
        break;

      case "assistant": {
        const assistantText = message.getText().replace(statsBlockRegex, "");
        const { content, reasoningContent } = requestReasoningKey === null
          ? { content: assistantText, reasoningContent: undefined }
          : extractReasoningBlocks(assistantText);
        const toolCalls: ChatCompletionMessageToolCall[] = message
          .getToolCallRequests()
          .map(toolCall => ({
            id: toolCall.id ?? "",
            type: "function",
            function: {
              name: toolCall.name,
              arguments: JSON.stringify(toolCall.arguments ?? {}),
            },
          }));

        const assistantMessage: OpenAICompatAssistantMessage = {
          role: "assistant",
          content,
          ...(toolCalls.length ? { tool_calls: toolCalls } : {}),
        };
        if (requestReasoningKey !== null && reasoningContent) {
          assistantMessage[requestReasoningKey] = reasoningContent;
        }

        messages.push(assistantMessage);
        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();
  });
}

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

const statsBlockRegex = /\n```\n[\s\S]*?\n```\s*$/;

function formatStats(elapsedMs: number, outputTokens: number, inputTokens: number, cacheHit: number): string {
  const elapsedS = (elapsedMs / 1000).toFixed(2);
  const speed = outputTokens > 0 && elapsedMs > 0
    ? (outputTokens / (elapsedMs / 1000)).toFixed(1)
    : "0.0";
  const total = inputTokens + outputTokens;
  return "```\n" +
    `⏱️${elapsedS}s ${speed}t/s  📦️${total} ↑${inputTokens} ↓${outputTokens} ♻${cacheHit}\n` +
    "```";
}

async function consumeStream(
  stream: AsyncIterable<any>,
  ctl: GeneratorController,
  responseReasoningKey: string | null,
  debug: boolean,
) {
  let current: ToolCallState | null = null;
  let reasoningOpen = false;
  let fullContent = "";
  let fullReasoning = "";
  let usageInputTokens = 0;
  let usageOutputTokens = 0;
  let usageCacheHit = 0;
  let hasExactUsage = false;
  const fullToolCalls: Array<{ id: string; name: string; arguments: string }> = [];

  function openReasoningBlock() {
    if (reasoningOpen) {
      return;
    }

    ctl.fragmentGenerated(reasoningStartTag, { reasoningType: "reasoningStartTag" });
    reasoningOpen = true;
  }

  function closeReasoningBlock() {
    if (!reasoningOpen) {
      return;
    }

    ctl.fragmentGenerated(reasoningEndTag, { reasoningType: "reasoningEndTag" });
    reasoningOpen = false;
  }

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

  for await (const chunk of stream) {
    if (debug) {
      console.info("[DEBUG] Raw chunk:", JSON.stringify(chunk));
    }
    if (chunk.usage) {
      hasExactUsage = true;
      usageInputTokens = chunk.usage.prompt_tokens ?? 0;
      usageOutputTokens = chunk.usage.completion_tokens ?? 0;
      usageCacheHit = chunk.usage.prompt_cache_hit_tokens ?? chunk.usage.prompt_tokens_details?.cached_tokens ?? 0;
    }
    const delta = chunk.choices?.[0]?.delta as OpenAICompatDelta | undefined;

    if (!delta) continue;

    /* Collect reasoning content for debug regardless of responseReasoningKey */
    if (debug && delta.reasoning_content) {
      fullReasoning += delta.reasoning_content;
    }

    const reasoningValue = responseReasoningKey === null ? null : delta[responseReasoningKey];
    const reasoningDelta = normalizeReasoningDelta(reasoningValue, responseReasoningKey);
    if (reasoningDelta !== null && reasoningDelta.length > 0) {
      fullReasoning += reasoningDelta;
      openReasoningBlock();
      ctl.fragmentGenerated(reasoningDelta, { reasoningType: "reasoning" });
    }

    /* Text streaming */
    if (delta.content) {
      fullContent += delta.content;
      closeReasoningBlock();
      ctl.fragmentGenerated(delta.content);
    }

    /* Tool-call streaming */
    if ((delta.tool_calls?.length ?? 0) > 0) {
      closeReasoningBlock();
    }

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

      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();
    }
  }

  maybeFlushCurrentToolCall();
  closeReasoningBlock();

  if (debug) {
    console.info("[DEBUG] === Full Response ===");
    console.info("[DEBUG] Reasoning content:", fullReasoning || "(none)");
    console.info("[DEBUG] Visible content:", fullContent || "(none)");
    console.info("[DEBUG] Tool calls:", fullToolCalls.length ? JSON.stringify(fullToolCalls) : "(none)");
    console.info("[DEBUG] ======================");
  }

  console.info("Generation completed.");

  return {
    outputTokens: hasExactUsage ? usageOutputTokens : null,
    inputTokens: hasExactUsage ? usageInputTokens : null,
    cacheHit: usageCacheHit,
    hasExactUsage,
  };
}

type ReasoningConfig = {
  requestReasoningKey: string | null;
  responseReasoningKey: string | null;
  thinking: { type: "enabled" } | undefined;
  reasoningEffort: string | undefined;
};

function resolveReasoningConfig(reasoningLevel: string): ReasoningConfig {
  switch (reasoningLevel) {
    case "off":
      return { requestReasoningKey: null, responseReasoningKey: null, thinking: { type: "disabled" }, reasoningEffort: undefined };
    case "max":
      return { requestReasoningKey: "reasoning_content", responseReasoningKey: "reasoning_content", thinking: { type: "enabled" }, reasoningEffort: "max" };
    case "high":
    default:
      return { requestReasoningKey: "reasoning_content", responseReasoningKey: "reasoning_content", thinking: { type: "enabled" }, reasoningEffort: "high" };
  }
}

export async function generate(ctl: GeneratorController, history: Chat) {
  const config = ctl.getPluginConfig(configSchematics);
  const model = config.get("model");
  const customModel = config.get("customModel");
  const effectiveModel = customModel || model;

  const globalConfig = ctl.getGlobalPluginConfig(globalConfigSchematics);
  const reasoningLevel = globalConfig.get("reasoningLevel");
  const statisticsMode = globalConfig.get("statistics");
  const showDebug = statisticsMode === "debug";
  const showStats = statisticsMode === "stats";
  const { requestReasoningKey, responseReasoningKey, thinking, reasoningEffort } = resolveReasoningConfig(reasoningLevel);

  /* 1. Setup client & payload */
  const openai = createOpenAI(globalConfig);
  const messages = await toOpenAIMessages(history, ctl, requestReasoningKey);
  const tools = toOpenAITools(ctl);

  /* 2. Kick off streaming completion */
  const params: any = { messages, tools, stream: true };
  if (effectiveModel) {
    params.model = effectiveModel;
  }
  if (thinking) {
    params.thinking = thinking;
  }
  if (reasoningEffort) {
    params.reasoning_effort = reasoningEffort;
  }
  if (showDebug) {
    console.info("[DEBUG] === Sent Request ===");
    console.info("[DEBUG] Model:", effectiveModel || "(not set)");
    console.info("[DEBUG] Params:", JSON.stringify({ ...params, messages: "(see below)" }));
    console.info("[DEBUG] Messages:", JSON.stringify(messages, null, 2));
    console.info("[DEBUG] ======================");
  }
  const startTime = performance.now();
  const stream = await openai.chat.completions.create(params);

  /* 3. Abort wiring & stream processing */
  wireAbort(ctl, stream as any);
  const streamResult = await consumeStream(stream as any, ctl, responseReasoningKey, showDebug);

  /* 4. Display statistics */
  if (showStats && streamResult) {
    const elapsedMs = Math.max(performance.now() - startTime, 1);
    const tIn = streamResult.inputTokens ?? 0;
    const tOut = streamResult.outputTokens ?? 0;
    const cache = streamResult.cacheHit;
    const statsText = formatStats(elapsedMs, tOut, tIn, cache);
    ctl.fragmentGenerated("\n\n" + statsText);
  }
}