intellinode/wrappers/VLLMWrapper.js

Create AI agents using the latest models, including ChatGPT, Llama, Diffusion, Cohere, Gemini, and Hugging Face.

const FetchClient = require('../utils/FetchClient');
const connHelper = require('../utils/ConnHelper');

class VLLMWrapper {
  constructor(apiBaseUrl) {
    this.client = new FetchClient({
      baseURL: apiBaseUrl,
      headers: {
        'Content-Type': 'application/json'
      }
    });
  }

  async generateText(params) {
    const endpoint = '/v1/completions';
    try {
      const extraConfig = params.stream ? { responseType: 'stream' } : {};
      return await this.client.post(endpoint, params, extraConfig);
    } catch (error) {
      throw new Error(connHelper.getErrorMessage(error));
    }
  }

  async generateChatText(params) {
    const endpoint = '/v1/chat/completions';
    try {
      const extraConfig = params.stream ? { responseType: 'stream' } : {};
      return await this.client.post(endpoint, params, extraConfig);
    } catch (error) {
      throw new Error(connHelper.getErrorMessage(error));
    }
  }

  async getEmbeddings(texts) {
    const endpoint = '/embed';
    try {
      return await this.client.post(endpoint, { texts });
    } catch (error) {
      throw new Error(connHelper.getErrorMessage(error));
    }
  }
}

module.exports = VLLMWrapper;