diff --git a/src/services/engine.ts b/src/services/engine.ts index a2cfa82..12665fb 100644 --- a/src/services/engine.ts +++ b/src/services/engine.ts @@ -130,7 +130,7 @@ export default class LlmEngine { } - buildPayload(thread: Message[], model: string): LLmCompletionPayload[] { + buildPayload(thread: Message[] | string, model: string): LLmCompletionPayload[] { if (typeof thread === 'string') { return [{ role: 'user', content: thread }] } else { diff --git a/tests/unit/engine.test.ts b/tests/unit/engine.test.ts index 521006d..6bef444 100644 --- a/tests/unit/engine.test.ts +++ b/tests/unit/engine.test.ts @@ -1,17 +1,36 @@ -import { beforeEach, expect, test } from 'vitest' -import { isEngineReady, igniteEngine, hasVisionModels, isVisionModel } from '../../src/services/llm' +import { vi, beforeEach, expect, test } from 'vitest' +import { isEngineReady, igniteEngine, hasVisionModels, isVisionModel, loadOpenAIModels } from '../../src/services/llm' import { store } from '../../src/services/store' import defaults from '../../defaults/settings.json' +import Message from '../../src/models/message' import OpenAI from '../../src/services/openai' import Ollama from '../../src/services/ollama' import MistralAI from '../../src/services/mistralai' import Anthropic from '../../src/services/anthropic' import Google from '../../src/services/google' import Groq from '../../src/services/groq' +import { Model } from '../../src/types/config.d' +import { text } from 'stream/consumers' const model = [{ id: 'llava:latest', name: 'llava:latest', meta: {} }] +vi.mock('openai', async() => { + const OpenAI = vi.fn() + OpenAI.prototype.apiKey = '123' + OpenAI.prototype.models = { + list: vi.fn(() => { + return { data: [ + { id: 'gpt-model2', name: 'model2' }, + { id: 'gpt-model1', name: 'model1' }, + { id: 'dall-e-model2', name: 'model2' }, + { id: 'dall-e-model1', name: 'model1' }, + ] } + }) + } + return { default: OpenAI } +}) + beforeEach(() => { store.config = defaults }) @@ -98,3 +117,78 @@ test('Is Vision Model', async () => { expect(isVisionModel('openai', 'gpt-4-turbo')).toBe(true) expect(isVisionModel('openai', 'gpt-vision')).toBe(true) }) + +test('Get Chat Models', async () => { + await loadOpenAIModels() + const openai = new OpenAI(store.config) + expect(openai.getChatModel()).toBe('gpt-model1') + expect(openai.getChatModels().map((m: Model) => { return { id: m.id, name: m.name }})).toStrictEqual([ + { id: 'gpt-model1', name: 'gpt-model1' }, + { id: 'gpt-model2', name: 'gpt-model2' }, + ]) +}) + +test('Find Models', async () => { + const models = [ + { id: 'gpt-model1', name: 'gpt-model1', meta: {} }, + { id: 'gpt-model2', name: 'gpt-model2', meta: {} }, + ] + const openai = new OpenAI(store.config) + expect(openai.findModel(models, ['gpt-model'])).toBeNull() + expect(openai.findModel(models, ['gpt-vision*'])).toBeNull() + expect(openai.findModel(models, ['*']).id).toBe('gpt-model1') + expect(openai.findModel(models, ['gpt-model2']).id).toBe('gpt-model2') + expect(openai.findModel(models, ['gpt-model*']).id).toBe('gpt-model1') + expect(openai.findModel(models, ['gpt-vision', '*gpt*2*']).id).toBe('gpt-model2') +}) + +test('Build payload no attachment', async () => { + const openai = new OpenAI(store.config) + expect(openai.buildPayload([], 'gpt-model1')).toStrictEqual([]) + expect(openai.buildPayload('content', 'gpt-model1')).toStrictEqual([{ role: 'user', content: 'content' }]) + expect(openai.buildPayload([ + new Message('system', { role: 'system', type: 'text', content: 'instructions' }), + new Message('user', { role: 'user', type: 'text', content: 'prompt1' }), + new Message('assistant', { role: 'assistant', type: 'image', content: 'response1' }), + new Message('user', { role: 'user', type: 'text', content: 'prompt2' }), + new Message('assistant', { role: 'assistant', type: 'text', content: 'response2' }), + ], 'gpt-model1')).toStrictEqual([ + { role: 'system', content: 'instructions' }, + { role: 'user', content: 'prompt1' }, + { role: 'user', content: 'prompt2' }, + { role: 'assistant', content: 'response2'} + ]) +}) + +test('Build payload with text attachment', async () => { + const openai = new OpenAI(store.config) + const messages = [ + new Message('system', { role: 'system', type: 'text', content: 'instructions' }), + new Message('user', { role: 'user', type: 'text', content: 'prompt1' }), + ] + messages[1].attachFile({ format: 'txt', contents: 'attachment', downloaded: true, url: '' }) + expect(openai.buildPayload(messages, 'gpt-model1')).toStrictEqual([ + { role: 'system', content: 'instructions' }, + { role: 'user', content: 'prompt1\n\nattachment' }, + ]) +}) + +test('Build payload with image attachment', async () => { + const openai = new OpenAI(store.config) + const messages = [ + new Message('system', { role: 'system', type: 'text', content: 'instructions' }), + new Message('user', { role: 'user', type: 'text', content: 'prompt1' }), + ] + messages[1].attachFile({ format: 'png', contents: 'attachment', downloaded: true, url: '' }) + expect(openai.buildPayload(messages, 'gpt-model1')).toStrictEqual([ + { role: 'system', content: 'instructions' }, + { role: 'user', content: 'prompt1' }, + ]) + expect(openai.buildPayload(messages, 'gpt-4-vision')).toStrictEqual([ + { role: 'system', content: 'instructions' }, + { role: 'user', content: [ + { type: 'text', text: 'prompt1' }, + { 'type': 'image_url', 'image_url': { 'url': 'data:image/jpeg;base64,attachment' } }, + ]}, + ]) +})