builders_test.go

  1package providertests
  2
  3import (
  4	"cmp"
  5	"net/http"
  6	"os"
  7
  8	"github.com/charmbracelet/fantasy/ai"
  9	"github.com/charmbracelet/fantasy/anthropic"
 10	"github.com/charmbracelet/fantasy/google"
 11	"github.com/charmbracelet/fantasy/openai"
 12	"github.com/charmbracelet/fantasy/openrouter"
 13	"gopkg.in/dnaeon/go-vcr.v4/pkg/recorder"
 14)
 15
 16type builderFunc func(r *recorder.Recorder) (ai.LanguageModel, error)
 17
 18type builderPair struct {
 19	name    string
 20	builder builderFunc
 21}
 22
 23var languageModelBuilders = []builderPair{
 24	{"openai-gpt-4o", builderOpenaiGpt4o},
 25	{"openai-gpt-4o-mini", builderOpenaiGpt4oMini},
 26	{"anthropic-claude-sonnet", builderAnthropicClaudeSonnet4},
 27	{"google-gemini-2.5-flash", builderGoogleGemini25Flash},
 28	{"google-gemini-2.5-pro", builderGoogleGemini25Pro},
 29	{"openrouter-kimi-k2", builderOpenrouterKimiK2},
 30}
 31
 32var thinkingLanguageModelBuilders = []builderPair{
 33	{"openai-gpt-5", builderOpenaiGpt5},
 34	{"anthropic-claude-sonnet", builderAnthropicClaudeSonnet4},
 35	{"google-gemini-2.5-pro", builderGoogleGemini25Pro},
 36	{"openrouter-glm-4.5", builderOpenrouterGLM45},
 37}
 38
 39func builderOpenaiGpt4o(r *recorder.Recorder) (ai.LanguageModel, error) {
 40	provider := openai.New(
 41		openai.WithAPIKey(os.Getenv("OPENAI_API_KEY")),
 42		openai.WithHTTPClient(&http.Client{Transport: r}),
 43	)
 44	return provider.LanguageModel("gpt-4o")
 45}
 46
 47func builderOpenaiGpt4oMini(r *recorder.Recorder) (ai.LanguageModel, error) {
 48	provider := openai.New(
 49		openai.WithAPIKey(os.Getenv("OPENAI_API_KEY")),
 50		openai.WithHTTPClient(&http.Client{Transport: r}),
 51	)
 52	return provider.LanguageModel("gpt-4o-mini")
 53}
 54
 55func builderOpenaiGpt5(r *recorder.Recorder) (ai.LanguageModel, error) {
 56	provider := openai.New(
 57		openai.WithAPIKey(os.Getenv("OPENAI_API_KEY")),
 58		openai.WithHTTPClient(&http.Client{Transport: r}),
 59	)
 60	return provider.LanguageModel("gpt-5")
 61}
 62
 63func builderAnthropicClaudeSonnet4(r *recorder.Recorder) (ai.LanguageModel, error) {
 64	provider := anthropic.New(
 65		anthropic.WithAPIKey(os.Getenv("ANTHROPIC_API_KEY")),
 66		anthropic.WithHTTPClient(&http.Client{Transport: r}),
 67	)
 68	return provider.LanguageModel("claude-sonnet-4-20250514")
 69}
 70
 71func builderGoogleGemini25Flash(r *recorder.Recorder) (ai.LanguageModel, error) {
 72	provider := google.New(
 73		google.WithAPIKey(cmp.Or(os.Getenv("GEMINI_API_KEY"), "(missing)")),
 74		google.WithHTTPClient(&http.Client{Transport: r}),
 75	)
 76	return provider.LanguageModel("gemini-2.5-flash")
 77}
 78
 79func builderGoogleGemini25Pro(r *recorder.Recorder) (ai.LanguageModel, error) {
 80	provider := google.New(
 81		google.WithAPIKey(cmp.Or(os.Getenv("GEMINI_API_KEY"), "(missing)")),
 82		google.WithHTTPClient(&http.Client{Transport: r}),
 83	)
 84	return provider.LanguageModel("gemini-2.5-pro")
 85}
 86
 87func builderOpenrouterKimiK2(r *recorder.Recorder) (ai.LanguageModel, error) {
 88	provider := openrouter.New(
 89		openrouter.WithAPIKey(os.Getenv("OPENROUTER_API_KEY")),
 90		openrouter.WithHTTPClient(&http.Client{Transport: r}),
 91	)
 92	return provider.LanguageModel("moonshotai/kimi-k2-0905")
 93}
 94
 95func builderOpenrouterGLM45(r *recorder.Recorder) (ai.LanguageModel, error) {
 96	provider := openrouter.New(
 97		openrouter.WithAPIKey(os.Getenv("OPENROUTER_API_KEY")),
 98		openrouter.WithHTTPClient(&http.Client{Transport: r}),
 99	)
100	return provider.LanguageModel("z-ai/glm-4.5")
101}