// Copyright (c) 2026 Probo Inc . // // Permission is hereby granted, free of charge, to any person obtaining a copy // of this software and associated documentation files (the "Software"), to deal // in the Software without restriction, including without limitation the rights // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell // copies of the Software, and to permit persons to whom the Software is // furnished to do so, subject to the following conditions: // // The above copyright notice and this permission notice shall be included in // all copies or substantial portions of the Software. // // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE // SOFTWARE. package llm_test import ( "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.probo.inc/probo/pkg/llm" ) func TestRegistry_Lookup(t *testing.T) { t.Parallel() r := llm.NewRegistry(map[string]llm.ModelDefinition{ "acme/test-model-1.5": { Name: "Acme: Test Model 1.5", ContextLength: 8192, MaxOutputTokens: 4096, Supports: llm.SupportedParameters{ Temperature: true, Reasoning: true, }, }, }) t.Run( "by full id", func(t *testing.T) { t.Parallel() m, ok := r.Lookup("acme/test-model-1.5") require.True(t, ok) assert.Equal(t, "acme/test-model-1.5", m.ID) assert.Equal(t, "acme", m.Provider()) assert.Equal(t, 8192, m.ContextLength) assert.Equal(t, 4096, m.MaxOutputTokens) }, ) t.Run( "by bare name with dots", func(t *testing.T) { t.Parallel() m, ok := r.Lookup("test-model-1.5") require.True(t, ok) assert.Equal(t, "acme", m.Provider()) }, ) t.Run( "by bare name with dashes", func(t *testing.T) { t.Parallel() m, ok := r.Lookup("test-model-1-5") require.True(t, ok) assert.Equal(t, "acme", m.Provider()) }, ) t.Run( "dated snapshot falls back to base model", func(t *testing.T) { t.Parallel() r := llm.NewRegistry(map[string]llm.ModelDefinition{ "openai/gpt-5-nano": { Name: "OpenAI: GPT-5 Nano", }, }) m, ok := r.Lookup("gpt-5-nano-2025-08-07") require.True(t, ok) assert.Equal(t, "openai/gpt-5-nano", m.ID) }, ) t.Run( "unknown model returns false", func(t *testing.T) { t.Parallel() _, ok := r.Lookup("nonexistent-model-42") assert.False(t, ok) }, ) t.Run( "empty string returns false", func(t *testing.T) { t.Parallel() _, ok := r.Lookup("") assert.False(t, ok) }, ) t.Run( "provider prefix only returns false", func(t *testing.T) { t.Parallel() _, ok := r.Lookup("acme/") assert.False(t, ok) }, ) } func TestRegistry_Capabilities(t *testing.T) { t.Parallel() r := llm.NewRegistry(map[string]llm.ModelDefinition{ "acme/reasoning-model": { Name: "Acme: Reasoning Model", Supports: llm.SupportedParameters{ Reasoning: true, Seed: true, }, }, "acme/chat-model": { Name: "Acme: Chat Model", Supports: llm.SupportedParameters{ Temperature: true, TopP: true, TopK: true, FrequencyPenalty: true, }, }, }) t.Run( "reasoning model has reasoning and seed but not temperature", func(t *testing.T) { t.Parallel() m, ok := r.Lookup("acme/reasoning-model") require.True(t, ok) assert.True(t, m.Supports.Reasoning) assert.True(t, m.Supports.Seed) assert.False(t, m.Supports.Temperature) assert.False(t, m.Supports.TopP) }, ) t.Run( "chat model has temperature and penalties but not reasoning", func(t *testing.T) { t.Parallel() m, ok := r.Lookup("acme/chat-model") require.True(t, ok) assert.True(t, m.Supports.Temperature) assert.True(t, m.Supports.TopP) assert.True(t, m.Supports.TopK) assert.True(t, m.Supports.FrequencyPenalty) assert.False(t, m.Supports.Reasoning) }, ) }