chore: Refactor mocking completions in tests

Author: dividedmind

packages/navie/test/fixture.ts

@@ -98,23 +98,3 @@ export function predictsSummary(): (messages: Message[]) => void {
     return Promise.resolve([]);
   });
 }
-
-export function mockAIResponse(completionWithRetry: jest.Mock, responses: string[]): void {
-  completionWithRetry.mockResolvedValueOnce(
-    responses.map((response, index) => ({
-      id: 'cmpl-3Z5z9J5Z5Z5Z5Z5Z5Z5Z5Z5Z5Z5',
-      choices: [
-        {
-          delta: {
-            content: response,
-          },
-          index,
-          finish_reason: index === responses.length - 1 ? 'stop' : null,
-        },
-      ],
-      created: 1635989729,
-      model: 'gpt-3.5',
-      object: 'chat.completion.chunk',
-    }))
-  );
-}

packages/navie/test/services/classification-service.spec.ts

@@ -1,39 +1,28 @@
-import { ChatOpenAI } from '@langchain/openai';
-
 import InteractionHistory from '../../src/interaction-history';
 import ClassificationService from '../../src/services/classification-service';
-import { mockAIResponse } from '../fixture';
-import OpenAICompletionService from '../../src/services/openai-completion-service';
-import Trajectory from '../../src/lib/trajectory';
-import { TrajectoryEvent } from '../../dist/lib/trajectory';
-import MessageTokenReducerService from '../../src/services/message-token-reducer-service';
 
-jest.mock('@langchain/openai');
-const completionWithRetry = jest.mocked(ChatOpenAI.prototype.completionWithRetry);
+import MockCompletionService from './mock-completion-service';
 
 describe('ClassificationService', () => {
   let interactionHistory: InteractionHistory;
-  let trajectory: Trajectory;
   let service: ClassificationService;
+  const completion = new MockCompletionService();
+  const completeSpy = jest.spyOn(completion, 'complete');
 
   beforeEach(() => {
     interactionHistory = new InteractionHistory();
     interactionHistory.on('event', (event) => console.log(event.message));
-    trajectory = new Trajectory();
-    service = new ClassificationService(
-      interactionHistory,
-      new OpenAICompletionService('gpt-4', 0.5, trajectory, new MessageTokenReducerService())
-    );
+    service = new ClassificationService(interactionHistory, completion);
   });
-  afterEach(() => jest.resetAllMocks());
+  afterEach(() => jest.restoreAllMocks());
 
   describe('when LLM responds', () => {
     const classification = `
 - architecture: high
 - troubleshoot: medium
 `;
 
-    beforeEach(() => mockAIResponse(completionWithRetry, [classification]));
+    beforeEach(() => completion.mock(classification));
 
     it('returns the response', async () => {
       const response = await service.classifyQuestion('user management');
@@ -47,7 +36,7 @@ describe('ClassificationService', () => {
           weight: 'medium',
         },
       ]);
-      expect(completionWithRetry).toHaveBeenCalledTimes(1);
+      expect(completeSpy).toHaveBeenCalledTimes(1);
     });
 
     it('emits classification event', async () => {
@@ -60,16 +49,5 @@ describe('ClassificationService', () => {
         classification: ['architecture=high', 'troubleshoot=medium'],
       });
     });
-
-    it('emits trajectory events', async () => {
-      const trajectoryEvents = new Array<TrajectoryEvent>();
-
-      trajectory.on('event', (event) => trajectoryEvents.push(event));
-
-      await service.classifyQuestion('user management');
-
-      expect(trajectoryEvents.map((e) => e.message.role)).toEqual(['system', 'user', 'assistant']);
-      expect(trajectoryEvents.map((e) => e.type)).toEqual(['sent', 'sent', 'received']);
-    });
   });
 });

packages/navie/test/services/compute-update-service.spec.ts

@@ -1,31 +1,20 @@
-import { ChatOpenAI } from '@langchain/openai';
-
 import InteractionHistory from '../../src/interaction-history';
-import ClassificationService from '../../src/services/classification-service';
-import { mockAIResponse } from '../fixture';
-import OpenAICompletionService from '../../src/services/openai-completion-service';
 import ComputeUpdateService from '../../src/services/compute-update-service';
-import Trajectory from '../../src/lib/trajectory';
-import MessageTokenReducerService from '../../src/services/message-token-reducer-service';
 
-jest.mock('@langchain/openai');
-const completionWithRetry = jest.mocked(ChatOpenAI.prototype.completionWithRetry);
+import MockCompletionService from './mock-completion-service';
 
 describe('ComputeUpdateService', () => {
   let interactionHistory: InteractionHistory;
-  let trajectory: Trajectory;
   let service: ComputeUpdateService;
+  const completion = new MockCompletionService();
+  const complete = jest.spyOn(completion, 'complete');
 
   beforeEach(() => {
     interactionHistory = new InteractionHistory();
     interactionHistory.on('event', (event) => console.log(event.message));
-    trajectory = new Trajectory();
-    service = new ComputeUpdateService(
-      interactionHistory,
-      new OpenAICompletionService('gpt-4', 0.5, trajectory, new MessageTokenReducerService())
-    );
+    service = new ComputeUpdateService(interactionHistory, completion);
   });
-  afterEach(() => jest.resetAllMocks());
+  afterEach(() => jest.restoreAllMocks());
 
   describe('when LLM responds', () => {
     const existingContent = `class User < ApplicationRecord
@@ -50,12 +39,12 @@ end
 </change>
     `;
 
-    beforeEach(() => mockAIResponse(completionWithRetry, [changeStr]));
+    beforeEach(() => completion.mock(changeStr));
 
     it('computes the update', async () => {
       const response = await service.computeUpdate(existingContent, newContent);
       expect(response).toStrictEqual(change);
-      expect(completionWithRetry).toHaveBeenCalledTimes(1);
+      expect(complete).toHaveBeenCalledTimes(1);
     });
   });
 });

packages/navie/test/services/mock-completion-service.ts

@@ -0,0 +1,55 @@
+import assert from 'node:assert';
+
+import { ZodType } from 'zod';
+
+import type { Message } from '../../src';
+import CompletionService, { type Completion, Usage } from '../../src/services/completion-service';
+
+export default class MockCompletionService implements CompletionService {
+  // eslint-disable-next-line @typescript-eslint/require-await
+  async *complete(messages: readonly Message[]): Completion {
+    const completion = this.completion(messages);
+    for (const c of completion) {
+      yield c;
+    }
+    return new Usage();
+  }
+
+  /**
+   * The mock completion function. This function can be used to mock the completion result.
+   * By default, it returns a hardcoded string split on spaces. It's a normal Jest mock, so you can manipulate the result
+   * further using eg. the `mockReturnValue` function. mock() method is provided as a shorthand.
+   * @param messages The messages to complete.
+   * @returns The mocked completion.
+   */
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  completion = jest.fn(function (this: MockCompletionService, messages: readonly Message[]) {
+    return ['Example ', 'response ', 'from ', 'the ', 'LLM', '.'];
+  });
+
+  /**
+   * A shorthand to mock the completion result.
+   * @param response The response to give to callers. This can be text or a JSON object.
+   * @returns The mock function that can be used to manipulate the results further.
+   */
+  mock(...response: string[]): typeof this.completion;
+  mock(response: string): typeof this.completion;
+  mock(response: unknown): typeof this.completion;
+  mock(...response: unknown[]): typeof this.completion {
+    if (response.length > 1 && response.every((x) => typeof x === 'string'))
+      return this.completion.mockReturnValue(response as string[]);
+    assert(response.length === 1, 'Only one response is supported');
+    const cpl = typeof response[0] === 'string' ? response[0] : JSON.stringify(response[0]);
+    return this.completion.mockReturnValue(cpl.split(/(?= )/));
+  }
+
+  // eslint-disable-next-line @typescript-eslint/require-await
+  async json<T>(messages: Message[], schema: ZodType<T>): Promise<T | undefined> {
+    const completion = this.completion(messages).join('');
+    return schema.parse(JSON.parse(completion));
+  }
+
+  modelName = 'mock-model';
+  miniModelName = 'mock-mini-model';
+  temperature = 0.7;
+}

packages/navie/test/services/vector-terms-service.spec.ts

@@ -1,38 +1,24 @@
-/* eslint-disable @typescript-eslint/no-unsafe-return */
-/* eslint-disable @typescript-eslint/no-unsafe-assignment */
-
-import { ChatOpenAI } from '@langchain/openai';
-
 import VectorTermsService from '../../src/services/vector-terms-service';
 import InteractionHistory from '../../src/interaction-history';
-import { mockAIResponse } from '../fixture';
-import OpenAICompletionService from '../../src/services/openai-completion-service';
-import Trajectory from '../../src/lib/trajectory';
-import MessageTokenReducerService from '../../src/services/message-token-reducer-service';
-
-jest.mock('@langchain/openai');
-const completionWithRetry = jest.mocked(ChatOpenAI.prototype.completionWithRetry);
+import MockCompletionService from './mock-completion-service';
 
 describe('VectorTermsService', () => {
   let interactionHistory: InteractionHistory;
   let service: VectorTermsService;
-  let trajectory: Trajectory;
+  const completion = new MockCompletionService();
+  const complete = jest.spyOn(completion, 'complete');
 
   beforeEach(() => {
     interactionHistory = new InteractionHistory();
     interactionHistory.on('event', (event) => console.log(event.message));
-    trajectory = new Trajectory();
-    service = new VectorTermsService(
-      interactionHistory,
-      new OpenAICompletionService('gpt-4', 0.5, trajectory, new MessageTokenReducerService())
-    );
+    service = new VectorTermsService(interactionHistory, completion);
   });
-  afterEach(() => jest.resetAllMocks());
+  afterEach(() => jest.restoreAllMocks());
 
   describe('when LLM suggested terms', () => {
     describe('is a valid JSON object', () => {
       it('is recorded in the interaction history', async () => {
-        mockAIResponse(completionWithRetry, [`{"terms": ["user", "management"]}`]);
+        completion.mock(`{"terms": ["user", "management"]}`);
         await service.suggestTerms('user management');
         expect(interactionHistory.events.map((e) => ({ ...e }))).toEqual([
           {
@@ -42,73 +28,73 @@ describe('VectorTermsService', () => {
         ]);
       });
       it('should return the terms', async () => {
-        mockAIResponse(completionWithRetry, [`{"terms": ["user", "management"]}`]);
+        completion.mock(`{"terms": ["user", "management"]}`);
         const terms = await service.suggestTerms('user management');
         expect(terms).toEqual(['user', 'management']);
-        expect(completionWithRetry).toHaveBeenCalledTimes(1);
+        expect(complete).toHaveBeenCalledTimes(1);
       });
       it('removes very short terms', async () => {
-        mockAIResponse(completionWithRetry, [`["user", "management", "a"]`]);
+        completion.mock(`["user", "management", "a"]`);
         const terms = await service.suggestTerms('user management');
         expect(terms).toEqual(['user', 'management', 'a']);
-        expect(completionWithRetry).toHaveBeenCalledTimes(1);
+        expect(complete).toHaveBeenCalledTimes(1);
       });
       it('converts underscore_words to distinct words', async () => {
-        mockAIResponse(completionWithRetry, [`["user_management"]`]);
+        completion.mock(`["user_management"]`);
         const terms = await service.suggestTerms('user management');
         expect(terms).toEqual(['user_management']);
-        expect(completionWithRetry).toHaveBeenCalledTimes(1);
+        expect(complete).toHaveBeenCalledTimes(1);
       });
     });
 
     describe('are a valid JSON list', () => {
       it('should return the terms', async () => {
-        mockAIResponse(completionWithRetry, ['["user", "management"]']);
+        completion.mock('["user", "management"]');
         const terms = await service.suggestTerms('user management');
         expect(terms).toEqual(['user', 'management']);
       });
     });
 
     describe('are valid JSON wrapped in fences', () => {
       it('should return the terms', async () => {
-        mockAIResponse(completionWithRetry, ['```json\n', '["user", "management"]\n', '```\n']);
+        completion.mock('```json\n["user", "management"]\n```\n');
         const terms = await service.suggestTerms('user management');
         expect(terms).toEqual(['user', 'management']);
       });
     });
 
     describe('is YAML', () => {
       it('parses the terms', async () => {
-        mockAIResponse(completionWithRetry, ['response_key:\n', '  - user\n', '  - management\n']);
+        completion.mock('response_key:\n', '  - user\n', '  - management\n');
         const terms = await service.suggestTerms('user management');
         expect(terms).toEqual(['response_key:', '-', 'user', 'management']);
       });
     });
 
     describe('is prefixed by "Terms:"', () => {
       it('is accepted and processed', async () => {
-        mockAIResponse(completionWithRetry, ['Terms: ["user", "management"]']);
+        completion.mock('Terms: ["user", "management"]');
         const terms = await service.suggestTerms('user management');
         expect(terms).toEqual(['user', 'management']);
-        expect(completionWithRetry).toHaveBeenCalledTimes(1);
+        expect(complete).toHaveBeenCalledTimes(1);
       });
     });
 
     describe('includes terms with "+" prefix', () => {
       it('is accepted and processed', async () => {
-        mockAIResponse(completionWithRetry, ['Terms: +user management']);
+        completion.mock('Terms: +user management');
         const terms = await service.suggestTerms('user management');
         expect(terms).toEqual(['+user', 'management']);
-        expect(completionWithRetry).toHaveBeenCalledTimes(1);
+        expect(complete).toHaveBeenCalledTimes(1);
       });
     });
 
-    describe('is list-ish ', () => {
+    describe('is list-ish', () => {
       it('is accepted and processed', async () => {
-        mockAIResponse(completionWithRetry, ['-user -mgmt']);
+        completion.mock('-user -mgmt');
         const terms = await service.suggestTerms('user management');
         expect(terms).toEqual(['-user', '-mgmt']);
-        expect(completionWithRetry).toHaveBeenCalledTimes(1);
+        expect(complete).toHaveBeenCalledTimes(1);
       });
     });
   });