probo/internal/cmd/genmodels/main.go

// 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 main

import (
	"encoding/json"
	"fmt"
	"go/format"
	"net/http"
	"os"
	"strings"
	"time"
)

const openRouterURL = "https://openrouter.ai/api/v1/models"

var includedProviders = map[string]struct{}{
	"anthropic":  {},
	"openai":     {},
	"google":     {},
	"x-ai":       {},
	"perplexity": {},
	"amazon":     {},
	"meta-llama": {},
	"mistralai":  {},
	"deepseek":   {},
}

var paramFieldMap = map[string]string{
	"temperature":           "Temperature",
	"top_p":                 "TopP",
	"top_k":                 "TopK",
	"frequency_penalty":     "FrequencyPenalty",
	"presence_penalty":      "PresencePenalty",
	"stop":                  "Stop",
	"seed":                  "Seed",
	"max_tokens":            "MaxTokens",
	"max_completion_tokens": "MaxTokens",
	"tool_choice":           "ToolChoice",
	"parallel_tool_calls":   "ParallelToolCalls",
	"response_format":       "ResponseFormat",
	"structured_outputs":    "StructuredOutputs",
	"reasoning":             "Reasoning",
	"include_reasoning":     "Reasoning",
}

type (
	openRouterResponse struct {
		Data []openRouterModel `json:"data"`
	}

	openRouterModel struct {
		ID          string `json:"id"`
		Name        string `json:"name"`
		ContextLen  int    `json:"context_length"`
		TopProvider struct {
			MaxCompletionTokens int `json:"max_completion_tokens"`
		} `json:"top_provider"`
		SupportedParams []string `json:"supported_parameters"`
	}
)

const iscHeader = `// 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.

`

func main() {
	client := &http.Client{Timeout: 30 * time.Second}

	resp, err := client.Get(openRouterURL)
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot fetch models: %v\n", err)
		os.Exit(1)
	}
	defer func() { _ = resp.Body.Close() }()

	if resp.StatusCode != http.StatusOK {
		fmt.Fprintf(os.Stderr, "cannot fetch models: unexpected status %s\n", resp.Status)
		os.Exit(1)
	}

	var result openRouterResponse
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		fmt.Fprintf(os.Stderr, "cannot decode response: %v\n", err)
		os.Exit(1)
	}

	var buf strings.Builder
	fmt.Fprintf(
		&buf,
		`%s// Code generated by genmodels; DO NOT EDIT.
// Source: %s
// Generated: %s

package llm

var generatedModels = map[string]ModelDefinition{
`,
		iscHeader,
		openRouterURL,
		time.Now().UTC().Format(time.RFC3339),
	)

	var count int
	for _, m := range result.Data {
		provider := providerFromID(m.ID)
		if _, ok := includedProviders[provider]; !ok {
			continue
		}

		fmt.Fprintf(
			&buf,
			`	%q: {
		Name:            %q,
		ContextLength:   %d,
		MaxOutputTokens: %d,
		Supports: SupportedParameters{
%s		},
	},
`,
			m.ID,
			m.Name,
			m.ContextLen,
			m.TopProvider.MaxCompletionTokens,
			buildSupports(m.SupportedParams),
		)
		count++
	}

	buf.WriteString("}\n")

	src, err := format.Source([]byte(buf.String()))
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot format generated code: %v\n", err)
		os.Exit(1)
	}

	if err := os.WriteFile("registry_gen.go", src, 0644); err != nil {
		fmt.Fprintf(os.Stderr, "cannot write file: %v\n", err)
		os.Exit(1)
	}

	fmt.Fprintf(os.Stderr, "genmodels: wrote %d models to registry_gen.go\n", count)
}

func providerFromID(id string) string {
	provider, _, ok := strings.Cut(id, "/")
	if ok {
		return provider
	}
	return ""
}

func buildSupports(params []string) string {
	fields := make(map[string]bool)
	for _, p := range params {
		if field, ok := paramFieldMap[p]; ok {
			fields[field] = true
		}
	}

	var buf strings.Builder
	fieldNames := []string{
		"Temperature", "TopP", "TopK", "FrequencyPenalty", "PresencePenalty",
		"Stop", "Seed", "MaxTokens", "ToolChoice", "ParallelToolCalls",
		"ResponseFormat", "StructuredOutputs", "Reasoning",
	}
	for _, f := range fieldNames {
		fmt.Fprintf(&buf, "			%s: %t,\n", f, fields[f])
	}
	return buf.String()
}