1package providertests
 2
 3import (
 4	"net/http"
 5	"os"
 6	"testing"
 7
 8	"charm.land/fantasy"
 9	"charm.land/fantasy/providers/bedrock"
10	"gopkg.in/dnaeon/go-vcr.v4/pkg/recorder"
11)
12
13func TestBedrockCommon(t *testing.T) {
14	testCommon(t, []builderPair{
15		{"bedrock-anthropic-claude-3-sonnet", builderBedrockClaude3Sonnet, nil, nil},
16		{"bedrock-anthropic-claude-3-opus", builderBedrockClaude3Opus, nil, nil},
17		{"bedrock-anthropic-claude-3-haiku", builderBedrockClaude3Haiku, nil, nil},
18	})
19}
20
21func TestBedrockBasicAuth(t *testing.T) {
22	testSimple(t, builderPair{"bedrock-anthropic-claude-3-sonnet", buildersBedrockBasicAuth, nil, nil})
23}
24
25func builderBedrockClaude3Sonnet(r *recorder.Recorder) (fantasy.LanguageModel, error) {
26	provider, err := bedrock.New(
27		bedrock.WithHTTPClient(&http.Client{Transport: r}),
28		bedrock.WithSkipAuth(!r.IsRecording()),
29	)
30	if err != nil {
31		return nil, err
32	}
33	return provider.LanguageModel("us.anthropic.claude-3-sonnet-20240229-v1:0")
34}
35
36func builderBedrockClaude3Opus(r *recorder.Recorder) (fantasy.LanguageModel, error) {
37	provider, err := bedrock.New(
38		bedrock.WithHTTPClient(&http.Client{Transport: r}),
39		bedrock.WithSkipAuth(!r.IsRecording()),
40	)
41	if err != nil {
42		return nil, err
43	}
44	return provider.LanguageModel("us.anthropic.claude-3-opus-20240229-v1:0")
45}
46
47func builderBedrockClaude3Haiku(r *recorder.Recorder) (fantasy.LanguageModel, error) {
48	provider, err := bedrock.New(
49		bedrock.WithHTTPClient(&http.Client{Transport: r}),
50		bedrock.WithSkipAuth(!r.IsRecording()),
51	)
52	if err != nil {
53		return nil, err
54	}
55	return provider.LanguageModel("us.anthropic.claude-3-haiku-20240307-v1:0")
56}
57
58func buildersBedrockBasicAuth(r *recorder.Recorder) (fantasy.LanguageModel, error) {
59	provider, err := bedrock.New(
60		bedrock.WithHTTPClient(&http.Client{Transport: r}),
61		bedrock.WithAPIKey(os.Getenv("FANTASY_BEDROCK_API_KEY")),
62		bedrock.WithSkipAuth(true),
63	)
64	if err != nil {
65		return nil, err
66	}
67	return provider.LanguageModel("us.anthropic.claude-3-sonnet-20240229-v1:0")
68}