builders_test.go

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