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	"gopkg.in/dnaeon/go-vcr.v4/pkg/recorder"
13)
14
15type builderFunc func(r *recorder.Recorder) (ai.LanguageModel, error)
16
17type builderPair struct {
18	name    string
19	builder builderFunc
20}
21
22var languageModelBuilders = []builderPair{
23	{"openai-gpt-4o", builderOpenaiGpt4o},
24	{"openai-gpt-4o-mini", builderOpenaiGpt4oMini},
25	{"anthropic-claude-sonnet", builderAnthropicClaudeSonnet4},
26	{"google-gemini-2.5-flash", builderGoogleGemini25Flash},
27	{"google-gemini-2.5-pro", builderGoogleGemini25Pro},
28}
29
30func builderOpenaiGpt4o(r *recorder.Recorder) (ai.LanguageModel, error) {
31	provider := openai.New(
32		openai.WithAPIKey(os.Getenv("OPENAI_API_KEY")),
33		openai.WithHTTPClient(&http.Client{Transport: r}),
34	)
35	return provider.LanguageModel("gpt-4o")
36}
37
38func builderOpenaiGpt4oMini(r *recorder.Recorder) (ai.LanguageModel, error) {
39	provider := openai.New(
40		openai.WithAPIKey(os.Getenv("OPENAI_API_KEY")),
41		openai.WithHTTPClient(&http.Client{Transport: r}),
42	)
43	return provider.LanguageModel("gpt-4o-mini")
44}
45
46func builderAnthropicClaudeSonnet4(r *recorder.Recorder) (ai.LanguageModel, error) {
47	provider := anthropic.New(
48		anthropic.WithAPIKey(os.Getenv("ANTHROPIC_API_KEY")),
49		anthropic.WithHTTPClient(&http.Client{Transport: r}),
50	)
51	return provider.LanguageModel("claude-sonnet-4-20250514")
52}
53
54func builderGoogleGemini25Flash(r *recorder.Recorder) (ai.LanguageModel, error) {
55	provider := google.New(
56		google.WithAPIKey(cmp.Or(os.Getenv("GEMINI_API_KEY"), "(missing)")),
57		google.WithHTTPClient(&http.Client{Transport: r}),
58	)
59	return provider.LanguageModel("gemini-2.5-flash")
60}
61
62func builderGoogleGemini25Pro(r *recorder.Recorder) (ai.LanguageModel, error) {
63	provider := google.New(
64		google.WithAPIKey(cmp.Or(os.Getenv("GEMINI_API_KEY"), "(missing)")),
65		google.WithHTTPClient(&http.Client{Transport: r}),
66	)
67	return provider.LanguageModel("gemini-2.5-flash")
68}