koishi-plugin-chatluna-rmkv-adapter/lib/index.mjs

rwkv adapter for chatluna

var __defProp = Object.defineProperty;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __name = (target, value) => __defProp(target, "name", { value, configurable: true });
var __commonJS = (cb, mod) => function __require() {
  return mod || (0, cb[__getOwnPropNames(cb)[0]])((mod = { exports: {} }).exports, mod), mod.exports;
};

// src/locales/zh-CN.schema.yml
var require_zh_CN_schema = __commonJS({
  "src/locales/zh-CN.schema.yml"(exports, module) {
    module.exports = { $inner: [{}, { $desc: "请求选项", apiKeys: { $inner: ["RWKV Runner 的 API Key", "RWKV Runner 的 API 地址"], $desc: "RWKV Runner API 的 API Key 和请求地址列表。" } }, { $desc: "模型配置", maxTokens: "输入的最大上下文 Token（16~4096，必须是 16 的倍数）。注意：仅当您使用的模型最大 Token 为 8000 及以上时，才建议设置超过 2000 token。", temperature: "回复的随机性程度，数值越高，回复越随机。", presencePenalty: "重复惩罚系数，数值越高，越不易重复出现已出现过至少一次的 Token（范围：-2~2，步长：0.1）。", frequencyPenalty: "频率惩罚系数，数值越高，越不易重复出现次数较多的 Token（范围：-2~2，步长：0.1）。" }] };
  }
});

// src/locales/en-US.schema.yml
var require_en_US_schema = __commonJS({
  "src/locales/en-US.schema.yml"(exports, module) {
    module.exports = { $inner: [{}, { $desc: "API Configuration", apiKeys: { $inner: ["RWKV Runner API Key", "RWKV Runner API Endpoint"], $desc: "RWKV Runner API Keys and endpoints" } }, { $desc: "Model Parameters", maxTokens: "Max input tokens (16-4096, multiple of 16). Note: >2000 for 8k+ token models", temperature: "Sampling temperature (higher values increase randomness)", presencePenalty: "Token presence penalty (-2 to 2, step 0.1, discourages token repetition)", frequencyPenalty: "Token frequency penalty (-2 to 2, step 0.1, reduces frequent token repetition)" }] };
  }
});

// src/index.ts
import { ChatLunaPlugin } from "koishi-plugin-chatluna/services/chat";
import { Schema } from "koishi";

// src/client.ts
import { PlatformModelAndEmbeddingsClient } from "koishi-plugin-chatluna/llm-core/platform/client";
import {
  ChatLunaChatModel,
  ChatLunaEmbeddings
} from "koishi-plugin-chatluna/llm-core/platform/model";
import {
  ModelType
} from "koishi-plugin-chatluna/llm-core/platform/types";
import {
  ChatLunaError as ChatLunaError2,
  ChatLunaErrorCode as ChatLunaErrorCode2
} from "koishi-plugin-chatluna/utils/error";

// src/requester.ts
import {
  ModelRequester
} from "koishi-plugin-chatluna/llm-core/platform/api";
import { ChatGenerationChunk } from "@langchain/core/outputs";
import {
  ChatLunaError,
  ChatLunaErrorCode
} from "koishi-plugin-chatluna/utils/error";
import { sseIterable } from "koishi-plugin-chatluna/utils/sse";

// src/utils.ts
import {
  AIMessageChunk,
  ChatMessageChunk,
  FunctionMessageChunk,
  HumanMessageChunk,
  SystemMessageChunk,
  ToolMessageChunk
} from "@langchain/core/messages";
import { zodToJsonSchema } from "zod-to-json-schema";
function langchainMessageToOpenAIMessage(messages) {
  const result = [];
  for (const rawMessage of messages) {
    const role = messageTypeToOpenAIRole(rawMessage._getType());
    const msg = {
      content: rawMessage.content || null,
      name: role === "assistant" || role === "tool" ? rawMessage.name : void 0,
      role,
      //  function_call: rawMessage.additional_kwargs.function_call,
      tool_calls: rawMessage.additional_kwargs.tool_calls,
      tool_call_id: rawMessage.tool_call_id
    };
    if (msg.tool_calls == null) {
      delete msg.tool_calls;
    }
    if (msg.tool_call_id == null) {
      delete msg.tool_call_id;
    }
    if (msg.tool_calls) {
      for (const toolCall of msg.tool_calls) {
        const tool = toolCall.function;
        if (!tool.arguments) {
          continue;
        }
        tool.arguments = JSON.stringify(JSON.parse(tool.arguments));
      }
    }
    result.push(msg);
  }
  return result;
}
__name(langchainMessageToOpenAIMessage, "langchainMessageToOpenAIMessage");
function messageTypeToOpenAIRole(type) {
  switch (type) {
    case "system":
      return "system";
    case "ai":
      return "assistant";
    case "human":
      return "user";
    case "function":
      return "function";
    case "tool":
      return "tool";
    default:
      throw new Error(`Unknown message type: ${type}`);
  }
}
__name(messageTypeToOpenAIRole, "messageTypeToOpenAIRole");
function formatToolsToOpenAITools(tools) {
  if (tools.length < 1) {
    return void 0;
  }
  return tools.map(formatToolToOpenAITool);
}
__name(formatToolsToOpenAITools, "formatToolsToOpenAITools");
function formatToolToOpenAITool(tool) {
  return {
    type: "function",
    function: {
      name: tool.name,
      description: tool.description,
      // any?
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      parameters: zodToJsonSchema(tool.schema)
    }
  };
}
__name(formatToolToOpenAITool, "formatToolToOpenAITool");
function convertDeltaToMessageChunk(delta, defaultRole) {
  const role = ((delta.role?.length ?? 0) > 0 ? delta.role : defaultRole).toLowerCase();
  const content = delta.content ?? "";
  let additional_kwargs;
  if (delta.function_call) {
    additional_kwargs = {
      function_call: delta.function_call
    };
  } else if (delta.tool_calls) {
    additional_kwargs = {
      tool_calls: delta.tool_calls
    };
  } else {
    additional_kwargs = {};
  }
  if (role === "user") {
    return new HumanMessageChunk({ content });
  } else if (role === "assistant") {
    const toolCallChunks = [];
    if (Array.isArray(delta.tool_calls)) {
      for (const rawToolCall of delta.tool_calls) {
        toolCallChunks.push({
          name: rawToolCall.function?.name,
          args: rawToolCall.function?.arguments,
          id: rawToolCall.id,
          index: rawToolCall.index
        });
      }
    }
    return new AIMessageChunk({
      content,
      tool_call_chunks: toolCallChunks,
      additional_kwargs
    });
  } else if (role === "system") {
    return new SystemMessageChunk({ content });
  } else if (role === "function") {
    return new FunctionMessageChunk({
      content,
      additional_kwargs,
      name: delta.name
    });
  } else if (role === "tool") {
    return new ToolMessageChunk({
      content,
      additional_kwargs,
      tool_call_id: delta.tool_call_id
    });
  } else {
    return new ChatMessageChunk({ content, role });
  }
}
__name(convertDeltaToMessageChunk, "convertDeltaToMessageChunk");

// src/requester.ts
var RWKVRequester = class extends ModelRequester {
  constructor(_config, _plugin) {
    super();
    this._config = _config;
    this._plugin = _plugin;
  }
  static {
    __name(this, "RWKVRequester");
  }
  async *completionStream(params) {
    try {
      const response = await this._post(
        "chat/completions",
        {
          model: params.model,
          messages: langchainMessageToOpenAIMessage(params.input),
          tools: params.tools != null ? formatToolsToOpenAITools(params.tools) : void 0,
          stop: params.stop != null ? params.stop : void 0,
          // remove max_tokens
          max_tokens: params.model.includes("vision") ? void 0 : params.maxTokens,
          temperature: params.temperature,
          presence_penalty: params.presencePenalty,
          frequency_penalty: params.frequencyPenalty,
          n: params.n,
          top_p: params.topP,
          user: params.user ?? "user",
          stream: true,
          logit_bias: params.logitBias
        },
        {
          signal: params.signal
        }
      );
      const iterator = sseIterable(response);
      let defaultRole = "assistant";
      let errorCount = 0;
      for await (const event of iterator) {
        const chunk = event.data;
        if (chunk === "[DONE]") {
          return;
        }
        try {
          const data = JSON.parse(chunk);
          if (data.error) {
            throw new ChatLunaError(
              ChatLunaErrorCode.API_REQUEST_FAILED,
              new Error(
                "error when calling rwkv completion, Result: " + chunk
              )
            );
          }
          const choice = data.choices?.[0];
          if (!choice) {
            continue;
          }
          const { delta } = choice;
          const messageChunk = convertDeltaToMessageChunk(
            delta,
            defaultRole
          );
          defaultRole = (delta.role?.length ?? 0) > 0 ? delta.role : defaultRole;
          const generationChunk = new ChatGenerationChunk({
            message: messageChunk,
            text: messageChunk.content
          });
          yield generationChunk;
        } catch (e) {
          if (errorCount > 5) {
            throw new ChatLunaError(
              ChatLunaErrorCode.API_REQUEST_FAILED,
              e
            );
          } else {
            errorCount++;
            continue;
          }
        }
      }
    } catch (e) {
      if (e instanceof ChatLunaError) {
        throw e;
      } else {
        throw new ChatLunaError(ChatLunaErrorCode.API_REQUEST_FAILED, e);
      }
    }
  }
  async embeddings(params) {
    let data;
    try {
      const response = await this._post("embeddings", {
        input: params.input,
        model: "rwkv"
      });
      data = await response.text();
      data = JSON.parse(data);
      if (data.data && data.data.length > 0) {
        return data.data.map(
          (it) => it.embedding
        );
      }
      throw new Error();
    } catch (e) {
      const error = new Error(
        "error when calling rwkv embeddings, Result: " + JSON.stringify(data)
      );
      error.stack = e.stack;
      error.cause = e.cause;
      throw new ChatLunaError(ChatLunaErrorCode.API_REQUEST_FAILED, error);
    }
  }
  async getModels() {
    let data;
    try {
      const response = await this._get("models");
      data = await response.text();
      data = JSON.parse(data);
      return data.data.map((model) => model.id);
    } catch (e) {
      const error = new Error(
        "error when listing rwkv models, Result: " + JSON.stringify(data)
      );
      error.stack = e.stack;
      error.cause = e.cause;
      throw error;
    }
  }
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  _post(url, data, params = {}) {
    const requestUrl = this._concatUrl(url);
    const body = JSON.stringify(data);
    return this._plugin.fetch(requestUrl, {
      body,
      headers: this._buildHeaders(),
      method: "POST",
      ...params
    });
  }
  _get(url) {
    const requestUrl = this._concatUrl(url);
    return this._plugin.fetch(requestUrl, {
      method: "GET",
      headers: this._buildHeaders()
    });
  }
  _buildHeaders() {
    return {
      Authorization: `Bearer ${this._config.apiKey}`,
      "Content-Type": "application/json"
    };
  }
  _concatUrl(url) {
    const apiEndPoint = this._config.apiEndpoint;
    if (!apiEndPoint.match(/\/v1\/?$/)) {
      if (apiEndPoint.endsWith("/")) {
        return apiEndPoint + "v1/" + url;
      }
      return apiEndPoint + "/v1/" + url;
    }
    if (apiEndPoint.endsWith("/")) {
      return apiEndPoint + url;
    }
    return apiEndPoint + "/" + url;
  }
  async init() {
  }
  async dispose() {
  }
};

// src/client.ts
var RWKVClient = class extends PlatformModelAndEmbeddingsClient {
  constructor(ctx, _config, clientConfig, plugin) {
    super(ctx, clientConfig);
    this._config = _config;
    this._requester = new RWKVRequester(clientConfig, plugin);
  }
  static {
    __name(this, "RWKVClient");
  }
  platform = "rwkv";
  _requester;
  _models;
  async init() {
    await this.getModels();
  }
  async refreshModels() {
    try {
      const rawModels = await this._requester.getModels();
      return rawModels.map((model) => {
        return {
          name: model,
          type: ModelType.llm,
          supportMode: ["all"]
        };
      }).concat([
        {
          name: "rwkv-embeddings",
          type: ModelType.embeddings,
          supportMode: ["all"]
        }
      ]);
    } catch (e) {
      throw new ChatLunaError2(ChatLunaErrorCode2.MODEL_INIT_ERROR, e);
    }
  }
  async getModels() {
    if (this._models) {
      return Object.values(this._models);
    }
    const models = await this.refreshModels();
    this._models = {};
    for (const model of models) {
      this._models[model.name] = model;
    }
    return models;
  }
  _createModel(model) {
    const info = this._models[model];
    if (info == null) {
      throw new ChatLunaError2(ChatLunaErrorCode2.MODEL_NOT_FOUND);
    }
    if (info.type === ModelType.llm) {
      return new ChatLunaChatModel({
        modelInfo: info,
        requester: this._requester,
        model,
        maxTokenLimit: this._config.maxTokens,
        frequencyPenalty: this._config.frequencyPenalty,
        presencePenalty: this._config.presencePenalty,
        timeout: this._config.timeout,
        temperature: this._config.temperature,
        maxRetries: this._config.maxRetries,
        llmType: "rwkv"
      });
    }
    return new ChatLunaEmbeddings({
      client: this._requester,
      maxRetries: this._config.maxRetries
    });
  }
};

// src/index.ts
function apply(ctx, config) {
  const plugin = new ChatLunaPlugin(ctx, config, "rwkv");
  ctx.on("ready", async () => {
    plugin.registerToService();
    await plugin.parseConfig((config2) => {
      return config2.apiKeys.map(([apiKey, apiEndpoint]) => {
        return {
          apiKey,
          apiEndpoint,
          platform: "rwkv",
          chatLimit: config2.chatTimeLimit,
          timeout: config2.timeout,
          maxRetries: config2.maxRetries,
          concurrentMaxSize: config2.chatConcurrentMaxSize
        };
      });
    });
    plugin.registerClient(
      (_, clientConfig) => new RWKVClient(ctx, config, clientConfig, plugin)
    );
    await plugin.initClients();
  });
}
__name(apply, "apply");
var Config = Schema.intersect([
  ChatLunaPlugin.Config,
  Schema.object({
    apiKeys: Schema.array(
      Schema.tuple([
        Schema.string().role("secret"),
        Schema.string().default("http://127.0.0.1:8000")
      ])
    ).default([["", "http://127.0.0.1:8000"]])
  }),
  Schema.object({
    maxTokens: Schema.number().min(16).max(4096).step(16).default(4096),
    temperature: Schema.percent().min(0).max(2).step(0.1).default(0.8),
    presencePenalty: Schema.number().min(-2).max(2).step(0.1).default(0.2),
    frequencyPenalty: Schema.number().min(-2).max(2).step(0.1).default(0.2)
  })
]).i18n({
  "zh-CN": require_zh_CN_schema(),
  "en-US": require_en_US_schema()
});
var inject = ["chatluna"];
var name = "chatluna-rmkv-adapter";
export {
  Config,
  apply,
  inject,
  name
};