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}