more engine tests

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 {

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' } },
+    ]},
+  ])
+})