From 45d9a9e43b3327285704ea8f6edb2925463230ff Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Sibiril?=
 <81782+aureliensibiril@users.noreply.github.com>
Date: Mon, 13 Apr 2026 16:05:31 +0200
Subject: [PATCH] Add model registry types and lookup
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

ModelDefinition, SupportedParameters, and Registry types with
multi-key lookup supporting canonical, bare, and normalized
model IDs. NewRegistry constructor accepts model definitions
for testability; DefaultRegistry caches the generated data.

Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>
---
 pkg/llm/registry.go | 112 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 112 insertions(+)
 create mode 100644 pkg/llm/registry.go

diff --git a/pkg/llm/registry.go b/pkg/llm/registry.go
new file mode 100644
index 000000000..60d48deb0
--- /dev/null
+++ b/pkg/llm/registry.go
@@ -0,0 +1,112 @@
+// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package llm
+
+import (
+	"strings"
+	"sync"
+)
+
+//go:generate go run go.probo.inc/probo/internal/cmd/genmodels
+
+type (
+	// ModelDefinition describes a model's capabilities and limits.
+	ModelDefinition struct {
+		ID              string
+		Name            string
+		Provider        string
+		ContextLength   int
+		MaxOutputTokens int
+		Supports        SupportedParameters
+	}
+
+	// SupportedParameters tracks which request parameters a model accepts.
+	SupportedParameters struct {
+		Temperature       bool
+		TopP              bool
+		TopK              bool
+		FrequencyPenalty  bool
+		PresencePenalty   bool
+		Stop              bool
+		Seed              bool
+		MaxTokens         bool
+		ToolChoice        bool
+		ParallelToolCalls bool
+		ResponseFormat    bool
+		StructuredOutputs bool
+		Reasoning         bool
+	}
+
+	// Registry provides model capability lookups.
+	Registry struct {
+		byID map[string]*ModelDefinition
+	}
+)
+
+var (
+	defaultRegistry     *Registry
+	defaultRegistryOnce sync.Once
+)
+
+// NewRegistry builds a registry from the given model definitions.
+func NewRegistry(models []ModelDefinition) *Registry {
+	r := &Registry{byID: make(map[string]*ModelDefinition, len(models)*3)}
+	for i := range models {
+		m := &models[i]
+		r.index(m)
+	}
+	return r
+}
+
+// DefaultRegistry returns the cached registry built from generated model data.
+func DefaultRegistry() *Registry {
+	defaultRegistryOnce.Do(func() {
+		defaultRegistry = NewRegistry(generatedModels)
+	})
+	return defaultRegistry
+}
+
+// Lookup finds a model by ID. It accepts both provider-prefixed IDs
+// ("anthropic/claude-opus-4.6") and bare provider IDs ("claude-opus-4-6",
+// "gpt-5.4"). Returns false if the model is not in the registry.
+func (r *Registry) Lookup(modelID string) (ModelDefinition, bool) {
+	if m, ok := r.byID[modelID]; ok {
+		return *m, true
+	}
+	if m, ok := r.byID[normalizeModelID(modelID)]; ok {
+		return *m, true
+	}
+	return ModelDefinition{}, false
+}
+
+func (r *Registry) index(m *ModelDefinition) {
+	r.byID[m.ID] = m
+
+	if idx := strings.IndexByte(m.ID, '/'); idx >= 0 {
+		r.byID[m.ID[idx+1:]] = m
+	}
+
+	normalized := normalizeModelID(m.ID)
+	if normalized != m.ID {
+		r.byID[normalized] = m
+	}
+}
+
+func normalizeModelID(id string) string {
+	if idx := strings.IndexByte(id, '/'); idx >= 0 {
+		id = id[idx+1:]
+	}
+	return strings.ReplaceAll(id, ".", "-")
+}