132 lines
4.4 KiB
TypeScript
132 lines
4.4 KiB
TypeScript
import { beforeEach, describe, expect, test } from 'vitest'
|
|
import { getBestModels, getFastModels, Model, ModelPreferences, pickModel, RemoteModelProvider } from '../src/models'
|
|
|
|
import MODELS from './models.json'
|
|
import { writeFileSync } from 'node:fs'
|
|
import { getTestClient } from './client'
|
|
|
|
describe('Models', () => {
|
|
test.skip('should fetch models', async () => {
|
|
// Run me manually if you need to re-generate the models.json file
|
|
// Make sure to setup the environment variables
|
|
const provider = new RemoteModelProvider(getTestClient())
|
|
const models = await provider.fetchInstalledModels()
|
|
writeFileSync('./models.json', JSON.stringify(models, null, 2))
|
|
})
|
|
|
|
test('Models ranking (best)', () => {
|
|
const best = getBestModels(MODELS as Model[])
|
|
expect(best.slice(0, 10).map((x) => x.ref)).toEqual([
|
|
'openai:gpt-4o-2024-11-20',
|
|
'openai:gpt-4o-2024-08-06',
|
|
'google-ai:models/gemini-1.5-pro-002',
|
|
'anthropic:claude-3-5-sonnet-20240620',
|
|
'openai:gpt-4o-mini-2024-07-18',
|
|
'groq:llama-3.2-90b-vision-preview',
|
|
'groq:llama-3.3-70b-versatile',
|
|
'fireworks-ai:accounts/fireworks/models/llama-v3p1-405b-instruct',
|
|
'google-ai:models/gemini-1.5-flash-002',
|
|
'openai:o1-mini-2024-09-12',
|
|
])
|
|
})
|
|
|
|
test('Models ranking (fast)', () => {
|
|
const fast = getFastModels(MODELS as Model[])
|
|
expect(fast.slice(0, 10).map((x) => x.ref)).toEqual([
|
|
'openai:gpt-4o-mini-2024-07-18',
|
|
'google-ai:models/gemini-1.5-flash-002',
|
|
'google-ai:models/gemini-1.5-flash-8b-001',
|
|
'openai:gpt-4o-2024-11-20',
|
|
'openai:gpt-4o-2024-08-06',
|
|
'google-ai:models/gemini-1.5-pro-002',
|
|
'anthropic:claude-3-haiku-20240307',
|
|
'anthropic:claude-3-5-sonnet-20240620',
|
|
'groq:llama-3.2-90b-vision-preview',
|
|
'groq:llama-3.3-70b-versatile',
|
|
])
|
|
})
|
|
|
|
test('Models ranking (boosted)', () => {
|
|
const fast = getFastModels(MODELS as Model[], {
|
|
'groq:llama-3.3-70b-versatile': 10,
|
|
'openai:gpt-4o-mini-2024-07-18': -10,
|
|
'google-ai:': 20,
|
|
})
|
|
expect(fast.slice(0, 10).map((x) => x.ref)).toEqual([
|
|
'google-ai:models/gemini-1.5-flash-002',
|
|
'google-ai:models/gemini-1.5-flash-8b-001',
|
|
'google-ai:models/gemini-1.5-pro-002',
|
|
'groq:llama-3.3-70b-versatile',
|
|
'openai:gpt-4o-2024-11-20',
|
|
'openai:gpt-4o-2024-08-06',
|
|
'anthropic:claude-3-haiku-20240307',
|
|
'anthropic:claude-3-5-sonnet-20240620',
|
|
'groq:llama-3.2-90b-vision-preview',
|
|
'fireworks-ai:accounts/fireworks/models/llama-v3p1-405b-instruct',
|
|
])
|
|
})
|
|
|
|
test('Pick model throws if none provided', () => {
|
|
expect(() => pickModel([])).toThrow()
|
|
expect(() => pickModel([], [])).toThrow()
|
|
})
|
|
|
|
test('Pick model throws if all models down', () => {
|
|
expect(() =>
|
|
pickModel(
|
|
['a:b', 'b:c'],
|
|
[
|
|
{ ref: 'a:b', reason: 'down', startedAt: new Date().toISOString() },
|
|
{ ref: 'b:c', reason: 'down', startedAt: new Date().toISOString() },
|
|
]
|
|
)
|
|
).toThrow()
|
|
})
|
|
|
|
test('Pick model picks the first one if all are up', () => {
|
|
expect(pickModel(['a:b', 'b:c'])).toEqual('a:b')
|
|
})
|
|
|
|
test('Pick model picks fallback when first down', () => {
|
|
expect(pickModel(['a:b', 'b:c'], [{ ref: 'a:b', reason: 'down', startedAt: new Date().toISOString() }])).toEqual(
|
|
'b:c'
|
|
)
|
|
})
|
|
})
|
|
|
|
describe('Remote Model Provider', () => {
|
|
beforeEach(async () => {
|
|
const client = getTestClient()
|
|
const provider = new RemoteModelProvider(client)
|
|
await provider.deleteModelPreferences()
|
|
})
|
|
|
|
test('fetch models preferences', async () => {
|
|
const client = getTestClient()
|
|
const provider = new RemoteModelProvider(client)
|
|
const preferences = await provider.fetchModelPreferences()
|
|
expect(preferences).toEqual(null)
|
|
})
|
|
|
|
test('save file preferences', async () => {
|
|
const client = getTestClient()
|
|
const provider = new RemoteModelProvider(client)
|
|
|
|
const customPreferences = {
|
|
best: ['openai:gpt-4o-2024-11-20' as const],
|
|
fast: ['openai:gpt-4o-mini-2024-07-18' as const],
|
|
downtimes: [],
|
|
} satisfies ModelPreferences
|
|
|
|
await provider.saveModelPreferences(customPreferences)
|
|
|
|
const preferences = await provider.fetchModelPreferences()
|
|
|
|
expect(preferences).toEqual({
|
|
best: ['openai:gpt-4o-2024-11-20'],
|
|
downtimes: [],
|
|
fast: ['openai:gpt-4o-mini-2024-07-18'],
|
|
})
|
|
})
|
|
})
|