@nomyx/assistant/src/tests/OpenAIProvider.test.ts

A powerful assistant library and cli for your AI projects. works with Vertex AI (Claude and Gemini)

import { OpenAIProvider } from '../ai/providers/OpenAIProvider/index';
import { OpenAIProviderConfig } from '../ai/providers/OpenAIProvider/types';
import { ILogger } from '../types/common';
import { ChatMessage, ChatOptions, ProviderResponse } from '../types/chat';
import { OpenAIProviderError } from '../ai/providers/OpenAIProvider/errors';

// Mock the external dependencies
jest.mock('../chat', () => ({
  chat: jest.fn(),
}));

jest.mock('../streaming', () => ({
  streamChat: jest.fn(),
}));

jest.mock('../embedding', () => ({
  embed: jest.fn(),
  embedBatch: jest.fn(),
}));

describe('OpenAIProvider', () => {
  let provider: OpenAIProvider;
  let mockLogger: ILogger;
  let mockConfig: OpenAIProviderConfig;

  beforeEach(() => {
    mockLogger = {
      debug: jest.fn(),
      info: jest.fn(),
      warn: jest.fn(),
      error: jest.fn(),
    };

    mockConfig = {
      type: 'openai',
      apiKey: 'test-api-key',
      model: 'gpt-4o',
    };

    provider = new OpenAIProvider(mockConfig, mockLogger);
  });

  describe('chat', () => {
    it('should return cached response if available', async () => {
      const messages: ChatMessage[] = [{ role: 'user', content: 'Hello' }];
      const options: ChatOptions = { maxTokens: 100 };
      const cachedResponse: ProviderResponse = { content: 'Cached response' };

      // @ts-ignore: Accessing private property for testing
      provider.cache.set({ messages, options }, cachedResponse);

      const result = await provider.chat(messages, options);

      expect(result).toEqual(cachedResponse);
      expect(mockLogger.debug).toHaveBeenCalledWith('Cache hit for chat request');
    });

    it('should process messages and return response when cache is empty', async () => {
      const messages: ChatMessage[] = [{ role: 'user', content: 'Hello' }];
      const options: ChatOptions = { maxTokens: 100 };
      const mockResponse: ProviderResponse = { content: 'API response' };

      const chatMock = jest.requireMock('../chat').chat;
      chatMock.mockResolvedValue(mockResponse);

      const result = await provider.chat(messages, options);

      expect(result).toEqual(mockResponse);
      expect(chatMock).toHaveBeenCalledWith(expect.any(Object), messages, options, undefined);
    });

    it('should handle errors and throw OpenAIProviderError', async () => {
      const messages: ChatMessage[] = [{ role: 'user', content: 'Hello' }];
      const options: ChatOptions = { maxTokens: 100 };

      const chatMock = jest.requireMock('../chat').chat;
      chatMock.mockRejectedValue(new Error('API Error'));

      await expect(provider.chat(messages, options)).rejects.toThrow(OpenAIProviderError);
      expect(mockLogger.error).toHaveBeenCalledWith('Error in OpenAIProvider chat', expect.any(Object));
    });
  });

  describe('streamChat', () => {
    it('should yield processed responses', async () => {
      const messages: ChatMessage[] = [{ role: 'user', content: 'Hello' }];
      const options: ChatOptions = { maxTokens: 100 };
      const mockResponses: ProviderResponse[] = [
        { content: 'Partial response 1' },
        { content: 'Partial response 2' },
      ];

      const streamChatMock = jest.requireMock('../streaming').streamChat;
      streamChatMock.mockImplementation(async function* () {
        yield* mockResponses;
      });

      const results: ProviderResponse[] = [];
      for await (const response of provider.streamChat(messages, options)) {
        results.push(response);
      }

      expect(results).toEqual(mockResponses);
      expect(streamChatMock).toHaveBeenCalledWith(expect.any(Object), messages, options, undefined);
    });

    it('should handle errors and throw OpenAIProviderError', async () => {
      const messages: ChatMessage[] = [{ role: 'user', content: 'Hello' }];
      const options: ChatOptions = { maxTokens: 100 };

      const streamChatMock = jest.requireMock('../streaming').streamChat;
      streamChatMock.mockImplementation(async function* () {
        throw new Error('Stream Error');
      });

      await expect(async () => {
        for await (const _ of provider.streamChat(messages, options)) {
          // Consume the iterator
        }
      }).rejects.toThrow(OpenAIProviderError);

      expect(mockLogger.error).toHaveBeenCalledWith('Error in OpenAIProvider streamChat', expect.any(Object));
    });
  });

  describe('embed', () => {
    it('should return embedding for given text', async () => {
      const text = 'Hello, world!';
      const mockEmbedding = [0.1, 0.2, 0.3];

      const embedMock = jest.requireMock('../embedding').embed;
      embedMock.mockResolvedValue(mockEmbedding);

      const result = await provider.embed(text);

      expect(result).toEqual(mockEmbedding);
      expect(embedMock).toHaveBeenCalledWith(expect.any(Object), text);
    });

    it('should handle errors and throw OpenAIProviderError', async () => {
      const text = 'Hello, world!';

      const embedMock = jest.requireMock('../embedding').embed;
      embedMock.mockRejectedValue(new Error('Embedding Error'));

      await expect(provider.embed(text)).rejects.toThrow(OpenAIProviderError);
      expect(mockLogger.error).toHaveBeenCalledWith('Error in OpenAIProvider embed', expect.any(Object));
    });
  });

  describe('embedBatch', () => {
    it('should return embeddings for given texts', async () => {
      const texts = ['Hello', 'world'];
      const mockEmbeddings = [[0.1, 0.2], [0.3, 0.4]];

      const embedBatchMock = jest.requireMock('../embedding').embedBatch;
      embedBatchMock.mockResolvedValue(mockEmbeddings);

      const result = await provider.embedBatch(texts);

      expect(result).toEqual(mockEmbeddings);
      expect(embedBatchMock).toHaveBeenCalledWith(expect.any(Object), texts);
    });

    it('should handle errors and throw OpenAIProviderError', async () => {
      const texts = ['Hello', 'world'];

      const embedBatchMock = jest.requireMock('../embedding').embedBatch;
      embedBatchMock.mockRejectedValue(new Error('Embedding Batch Error'));

      await expect(provider.embedBatch(texts)).rejects.toThrow(OpenAIProviderError);
      expect(mockLogger.error).toHaveBeenCalledWith('Error in OpenAIProvider embedBatch', expect.any(Object));
    });
  });

  describe('getCapabilities', () => {
    it('should return the correct capabilities', () => {
      const capabilities = provider.getCapabilities();

      expect(capabilities).toEqual({
        maxTokens: 4096,
        supportsFunctionCalling: true,
        supportsStreaming: true,
        supportedModels: ['gpt-4o', 'o1-preview'],
        maxSimultaneousCalls: 1,
        supportsSemanticCaching: false,
      });
    });
  });

  describe('clearCache', () => {
    it('should clear the cache and log a debug message', () => {
      provider.clearCache();

      expect(mockLogger.debug).toHaveBeenCalledWith('OpenAIProvider cache cleared');
    });
  });
});