LibreChat/api/app/clients/specs/FakeClient.js

const crypto = require('crypto');
const BaseClient = require('../BaseClient');
const { maxTokensMap } = require('../../../utils');

class FakeClient extends BaseClient {
  constructor(apiKey, options = {}) {
    super(apiKey, options);
    this.sender = 'AI Assistant';
    this.setOptions(options);
  }
  setOptions(options) {
    if (this.options && !this.options.replaceOptions) {
      this.options.modelOptions = {
        ...this.options.modelOptions,
        ...options.modelOptions,
      };
      delete options.modelOptions;
      this.options = {
        ...this.options,
        ...options,
      };
    } else {
      this.options = options;
    }

    if (this.options.openaiApiKey) {
      this.apiKey = this.options.openaiApiKey;
    }

    const modelOptions = this.options.modelOptions || {};
    if (!this.modelOptions) {
      this.modelOptions = {
        ...modelOptions,
        model: modelOptions.model || 'gpt-3.5-turbo',
        temperature: typeof modelOptions.temperature === 'undefined' ? 0.8 : modelOptions.temperature,
        top_p: typeof modelOptions.top_p === 'undefined' ? 1 : modelOptions.top_p,
        presence_penalty: typeof modelOptions.presence_penalty === 'undefined' ? 1 : modelOptions.presence_penalty,
        stop: modelOptions.stop,
      };
    }

    this.maxContextTokens = maxTokensMap[this.modelOptions.model] ?? 4097;
  }
  getCompletion() {}
  buildMessages() {}
  getTokenCount(str) {
    return str.length;
  }
  getTokenCountForMessage(message) {
    return message?.content?.length || message.length;
  }
}

const initializeFakeClient = (apiKey, options, fakeMessages) => {
  let TestClient = new FakeClient(apiKey);
  TestClient.options = options;
  TestClient.abortController = { abort: jest.fn() };
  TestClient.saveMessageToDatabase = jest.fn();
  TestClient.loadHistory = jest
    .fn()
    .mockImplementation((conversationId, parentMessageId = null) => {
      if (!conversationId) {
        TestClient.currentMessages = [];
        return Promise.resolve([]);
      }

      const orderedMessages = TestClient.constructor.getMessagesForConversation(
        fakeMessages,
        parentMessageId
      );

      TestClient.currentMessages = orderedMessages;
      return Promise.resolve(orderedMessages);
    });

  TestClient.getSaveOptions = jest.fn().mockImplementation(() => {
    return {};
  });

  TestClient.getBuildMessagesOptions = jest.fn().mockImplementation(() => {
    return {};
  });

  TestClient.sendCompletion = jest.fn(async () => {
    return 'Mock response text';
  });

  TestClient.sendMessage = jest.fn().mockImplementation(async (message, opts = {}) => {
    if (opts && typeof opts === 'object') {
      TestClient.setOptions(opts);
    }

    const user = opts.user || null;
    const conversationId = opts.conversationId || crypto.randomUUID();
    const parentMessageId = opts.parentMessageId || '00000000-0000-0000-0000-000000000000';
    const userMessageId = opts.overrideParentMessageId || crypto.randomUUID();
    const saveOptions = TestClient.getSaveOptions();

    this.pastMessages = await TestClient.loadHistory(
      conversationId,
      TestClient.options?.parentMessageId
    );

    const userMessage = {
      text: message,
      sender: TestClient.sender,
      isCreatedByUser: true,
      messageId: userMessageId,
      parentMessageId,
      conversationId
    };

    const response = {
      sender: TestClient.sender,
      text: 'Hello, User!',
      isCreatedByUser: false,
      messageId: crypto.randomUUID(),
      parentMessageId: userMessage.messageId,
      conversationId
    };

    fakeMessages.push(userMessage);
    fakeMessages.push(response);

    if (typeof opts.getIds === 'function') {
      opts.getIds({
        userMessage,
        conversationId,
        responseMessageId: response.messageId
      });
    }

    if (typeof opts.onStart === 'function') {
      opts.onStart(userMessage);
    }

    let { prompt: payload, tokenCountMap } = await TestClient.buildMessages(
      this.currentMessages,
      userMessage.messageId,
      TestClient.getBuildMessagesOptions(opts),
    );

    if (tokenCountMap) {
      payload = payload.map((message, i) => {
        const { tokenCount, ...messageWithoutTokenCount } = message;
        // userMessage is always the last one in the payload
        if (i === payload.length - 1) {
          userMessage.tokenCount = message.tokenCount;
          console.debug(`Token count for user message: ${tokenCount}`, `Instruction Tokens: ${tokenCountMap.instructions || 'N/A'}`);
        }
        return messageWithoutTokenCount;
      });
      TestClient.handleTokenCountMap(tokenCountMap);
    }

    await TestClient.saveMessageToDatabase(userMessage, saveOptions, user);
    response.text = await TestClient.sendCompletion(payload, opts);
    if (tokenCountMap && TestClient.getTokenCountForResponse) {
      response.tokenCount = TestClient.getTokenCountForResponse(response);
    }
    await TestClient.saveMessageToDatabase(response, saveOptions, user);
    return response;
  });

  TestClient.buildMessages = jest.fn(async (messages, parentMessageId) => {
    const orderedMessages = TestClient.constructor.getMessagesForConversation(messages, parentMessageId);
    const formattedMessages = orderedMessages.map((message) => {
      let { role: _role, sender, text } = message;
      const role = _role ?? sender;
      const content = text ?? '';
      return {
        role: role?.toLowerCase() === 'user' ? 'user' : 'assistant',
        content,
      };
    });
    return {
      prompt: formattedMessages,
      tokenCountMap: null, // Simplified for the mock
    };
  });

  return TestClient;
}

module.exports = { FakeClient, initializeFakeClient };