Files
2026-07-13 13:34:48 +08:00

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