openrouter_test.go

 1package providertests
 2
 3import (
 4	"net/http"
 5	"os"
 6	"testing"
 7
 8	"github.com/charmbracelet/fantasy/ai"
 9	"github.com/charmbracelet/fantasy/openrouter"
10	"gopkg.in/dnaeon/go-vcr.v4/pkg/recorder"
11)
12
13type openrouterModel struct {
14	name        string
15	builderFunc builderFunc
16	providers   []string
17}
18
19func TestOpenRouterCommon(t *testing.T) {
20	var pairs []builderPair
21	models := []openrouterModel{
22		kimiK2(),
23		grokCodeFast1(),
24	}
25
26	for _, model := range models {
27		for _, provider := range model.providers {
28			pairs = append(
29				pairs,
30				builderPair{
31					model.name + "_" + provider,
32					model.builderFunc,
33					ai.ProviderOptions{
34						openrouter.Name: &openrouter.ProviderOptions{
35							Provider: &openrouter.Provider{
36								Only: []string{provider},
37							},
38						},
39					},
40				})
41		}
42	}
43
44	testCommon(t, pairs)
45}
46
47func kimiK2() openrouterModel {
48	return openrouterModel{
49		name: "kimi-k2",
50		builderFunc: func(r *recorder.Recorder) (ai.LanguageModel, error) {
51			provider := openrouter.New(
52				openrouter.WithAPIKey(os.Getenv("OPENROUTER_API_KEY")),
53				openrouter.WithHTTPClient(&http.Client{Transport: r}),
54			)
55			return provider.LanguageModel("moonshotai/kimi-k2-0905")
56		},
57		providers: []string{
58			"chutes",
59			"deepinfra",
60			"siliconflow",
61			"fireworks",
62			"moonshotai",
63			"novita",
64			"baseten",
65			"together",
66			"groq",
67			"moonshotai/turbo",
68			"wandb",
69		},
70	}
71}
72
73func grokCodeFast1() openrouterModel {
74	return openrouterModel{
75		name: "grok-code-fast-1",
76		builderFunc: func(r *recorder.Recorder) (ai.LanguageModel, error) {
77			provider := openrouter.New(
78				openrouter.WithAPIKey(os.Getenv("OPENROUTER_API_KEY")),
79				openrouter.WithHTTPClient(&http.Client{Transport: r}),
80			)
81			return provider.LanguageModel("x-ai/grok-code-fast-1")
82		},
83		providers: []string{
84			"xai",
85		},
86	}
87}