kit/internal/models/generate_models.go

//go:build ignore

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"text/template"
	"time"
)

// ModelInfo represents information about a specific model.
// This struct is used during code generation to parse model data
// from the models.dev API and generate the static Go code.
type ModelInfo struct {
	// ID is the unique identifier for the model
	ID string `json:"id"`
	// Name is the human-readable name of the model
	Name string `json:"name"`
	// Attachment indicates whether the model supports file attachments
	Attachment bool `json:"attachment"`
	// Reasoning indicates whether this is a reasoning/chain-of-thought model
	Reasoning bool `json:"reasoning"`
	// Temperature indicates whether the model supports temperature parameter
	Temperature bool `json:"temperature"`
	// Cost contains the pricing information for the model
	Cost Cost `json:"cost"`
	// Limit contains the context and output token limits
	Limit Limit `json:"limit"`
}

// Cost represents the pricing information for a model.
// Used during code generation to parse pricing data from models.dev.
type Cost struct {
	// Input is the cost per million input tokens
	Input float64 `json:"input"`
	// Output is the cost per million output tokens
	Output float64 `json:"output"`
	// CacheRead is the cost per million cached read tokens (optional)
	CacheRead *float64 `json:"cache_read,omitempty"`
	// CacheWrite is the cost per million cached write tokens (optional)
	CacheWrite *float64 `json:"cache_write,omitempty"`
}

// Limit represents the context and output limits for a model.
// Used during code generation to parse token limit data from models.dev.
type Limit struct {
	// Context is the maximum number of input tokens
	Context int `json:"context"`
	// Output is the maximum number of output tokens
	Output int `json:"output"`
}

// ProviderInfo represents information about a model provider.
// Used during code generation to parse provider data from models.dev
// and generate the static provider registry.
type ProviderInfo struct {
	// ID is the unique identifier for the provider
	ID string `json:"id"`
	// Env lists the environment variables for API credentials
	Env []string `json:"env"`
	// NPM is the NPM package name (for reference)
	NPM string `json:"npm"`
	// Name is the human-readable provider name
	Name string `json:"name"`
	// Models maps model IDs to their information
	Models map[string]ModelInfo `json:"models"`
}

const codeTemplate = `// Code generated by go generate; DO NOT EDIT.
// Generated at: {{.Timestamp}}

package models

// ModelInfo represents information about a specific model
type ModelInfo struct {
	ID          string
	Name        string
	Attachment  bool
	Reasoning   bool
	Temperature bool
	Cost        Cost
	Limit       Limit
}

// Cost represents the pricing information for a model
type Cost struct {
	Input      float64
	Output     float64
	CacheRead  *float64
	CacheWrite *float64
}

// Limit represents the context and output limits for a model
type Limit struct {
	Context int
	Output  int
}

// ProviderInfo represents information about a model provider
type ProviderInfo struct {
	ID     string
	Env    []string
	NPM    string
	Name   string
	Models map[string]ModelInfo
}

// GetModelsData returns the static models data from models.dev
func GetModelsData() map[string]ProviderInfo {
	return map[string]ProviderInfo{
{{- range $providerID, $provider := .Providers}}
		"{{$providerID}}": {
			ID:   "{{$provider.ID}}",
			Env:  []string{ {{- range $i, $env := $provider.Env}}{{if $i}}, {{end}}"{{$env}}"{{end}} },
			NPM:  "{{$provider.NPM}}",
			Name: "{{$provider.Name}}",
			Models: map[string]ModelInfo{
{{- range $modelID, $model := $provider.Models}}
				"{{$modelID}}": {
					ID:          "{{$model.ID}}",
					Name:        "{{$model.Name}}",
					Attachment:  {{$model.Attachment}},
					Reasoning:   {{$model.Reasoning}},
					Temperature: {{$model.Temperature}},
					Cost: Cost{
						Input:  {{$model.Cost.Input}},
						Output: {{$model.Cost.Output}},
						{{- if $model.Cost.CacheRead}}
						CacheRead: &[]float64{{"{"}}{{$model.Cost.CacheRead}}{{"}"}}[0],
						{{- else}}
						CacheRead: nil,
						{{- end}}
						{{- if $model.Cost.CacheWrite}}
						CacheWrite: &[]float64{{"{"}}{{$model.Cost.CacheWrite}}{{"}"}}[0],
						{{- else}}
						CacheWrite: nil,
						{{- end}}
					},
					Limit: Limit{
						Context: {{$model.Limit.Context}},
						Output:  {{$model.Limit.Output}},
					},
				},
{{- end}}
			},
		},
{{- end}}
	}
}
`

func main() {
	fmt.Println("Fetching models data from models.dev...")

	// Fetch data from API
	resp, err := http.Get("https://models.dev/api.json")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error fetching data: %v\n", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		fmt.Fprintf(os.Stderr, "API returned status %d\n", resp.StatusCode)
		os.Exit(1)
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error reading response: %v\n", err)
		os.Exit(1)
	}

	// Parse JSON
	var providers map[string]ProviderInfo
	if err := json.Unmarshal(body, &providers); err != nil {
		fmt.Fprintf(os.Stderr, "Error parsing JSON: %v\n", err)
		os.Exit(1)
	}

	// Fix Google provider environment variables to match our implementation
	if googleProvider, exists := providers["google"]; exists {
		googleProvider.Env = []string{"GOOGLE_API_KEY", "GEMINI_API_KEY", "GOOGLE_GENERATIVE_AI_API_KEY"}
		providers["google"] = googleProvider
	}

	// Generate Go code
	tmpl, err := template.New("models").Parse(codeTemplate)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error parsing template: %v\n", err)
		os.Exit(1)
	}

	// Create output file
	file, err := os.Create("models_data.go")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error creating output file: %v\n", err)
		os.Exit(1)
	}
	defer file.Close()

	// Execute template
	data := struct {
		Providers map[string]ProviderInfo
		Timestamp string
	}{
		Providers: providers,
		Timestamp: time.Now().Format(time.RFC3339),
	}

	if err := tmpl.Execute(file, data); err != nil {
		fmt.Fprintf(os.Stderr, "Error executing template: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Generated models_data.go with %d providers\n", len(providers))

	// Print summary
	for providerID, provider := range providers {
		fmt.Printf("  %s: %d models\n", providerID, len(provider.Models))
	}
}