builders_test.go

 1package providertests
 2
 3import (
 4	"net/http"
 5	"os"
 6
 7	"github.com/charmbracelet/ai/ai"
 8	"github.com/charmbracelet/ai/anthropic"
 9	"github.com/charmbracelet/ai/openai"
10	"gopkg.in/dnaeon/go-vcr.v4/pkg/recorder"
11)
12
13type builderFunc func(r *recorder.Recorder) (ai.LanguageModel, error)
14
15type builderPair struct {
16	name    string
17	builder builderFunc
18}
19
20var languageModelBuilders = []builderPair{
21	{"openai-gpt-4o", builderOpenaiGpt4o},
22	{"anthropic-claude-sonnet", builderAnthropicClaudeSonnet4},
23}
24
25func builderOpenaiGpt4o(r *recorder.Recorder) (ai.LanguageModel, error) {
26	provider := openai.New(
27		openai.WithAPIKey(os.Getenv("OPENAI_API_KEY")),
28		openai.WithHTTPClient(&http.Client{Transport: r}),
29	)
30	model, err := provider.LanguageModel("gpt-4o")
31	if err != nil {
32		return nil, err
33	}
34	return model, nil
35}
36
37func builderAnthropicClaudeSonnet4(r *recorder.Recorder) (ai.LanguageModel, error) {
38	provider := anthropic.New(
39		anthropic.WithAPIKey(os.Getenv("ANTHROPIC_API_KEY")),
40		anthropic.WithHTTPClient(&http.Client{Transport: r}),
41	)
42	model, err := provider.LanguageModel("claude-sonnet-4-20250514")
43	if err != nil {
44		return nil, err
45	}
46	return model, nil
47}