100 lines
3.4 KiB
Go
100 lines
3.4 KiB
Go
package providers
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"path/filepath"
|
|
"testing"
|
|
|
|
"github.com/crypto/byteclaw/internal/config"
|
|
"github.com/crypto/byteclaw/internal/llm"
|
|
"github.com/crypto/byteclaw/internal/store"
|
|
)
|
|
|
|
type fake struct{ base, key, model string }
|
|
|
|
func (fake) Chat(context.Context, llm.Request, llm.DeltaFunc) (llm.Response, error) {
|
|
return llm.Response{}, nil
|
|
}
|
|
|
|
func newMgr(t *testing.T, env map[string]string) (*Manager, *store.Store) {
|
|
st, err := store.Open(filepath.Join(t.TempDir(), "p.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Cleanup(func() { st.Close() })
|
|
cfg := config.LLM{Provider: "openai", BaseURL: "https://api.openai.com/v1", Model: "gpt-4o-mini", APIKey: "sk-openai",
|
|
Providers: map[string]config.ProviderOverride{"gonkarouter": {APIKey: "sk-gonka"}}}
|
|
m := New(cfg, func(k string) string { return env[k] }, st)
|
|
m.New = func(b, k, mo string) llm.Provider { return fake{b, k, mo} }
|
|
return m, st
|
|
}
|
|
|
|
func TestListAndSwitch(t *testing.T) {
|
|
m, st := newMgr(t, map[string]string{"DAHL_API_KEY": "dahl_k"})
|
|
got := map[string]Info{}
|
|
for _, i := range m.List() {
|
|
got[i.Name] = i
|
|
}
|
|
if !got["openai"].Active || !got["dahl"].Configured || !got["gonkarouter"].Configured || got["openrouter"].Configured || !got["ollama"].Configured {
|
|
t.Fatalf("%+v", got)
|
|
}
|
|
if _, err := m.SetActive(context.Background(), "openrouter", ""); !errors.Is(err, ErrNotConfigured) {
|
|
t.Fatalf("want not configured, got %v", err)
|
|
}
|
|
if _, err := m.SetActive(context.Background(), "nope", ""); !errors.Is(err, ErrUnknown) {
|
|
t.Fatal(err)
|
|
}
|
|
a, err := m.SetActive(context.Background(), "Dahl", "")
|
|
if err != nil || a.Provider != "dahl" || a.Model != "MiniMaxAI/MiniMax-M2.7" {
|
|
t.Fatalf("%+v %v", a, err)
|
|
}
|
|
p, _ := m.Build("dahl", "zai-org/GLM-5.3-Flash")
|
|
if f := p.(fake); f.key != "dahl_k" || f.base != "https://inference.dahl.global/v1" || f.model != "zai-org/GLM-5.3-Flash" {
|
|
t.Fatalf("%+v", f)
|
|
}
|
|
// persisted across restarts
|
|
m2 := New(config.LLM{Provider: "openai", BaseURL: "x", Model: "y"}, func(k string) string { return map[string]string{"DAHL_API_KEY": "dahl_k"}[k] }, st)
|
|
if m2.Active().Provider != "dahl" {
|
|
t.Fatalf("not persisted: %+v", m2.Active())
|
|
}
|
|
}
|
|
|
|
func TestSessionOverride(t *testing.T) {
|
|
m, st := newMgr(t, nil)
|
|
ctx := context.Background()
|
|
ss, _ := st.CreateSession(ctx, "")
|
|
st.SetSessionProvider(ctx, ss.ID, "gonkarouter", "MiniMaxAI/MiniMax-M2.7")
|
|
p, a, err := m.ForSession(ctx, ss.ID)
|
|
if err != nil || a.Provider != "gonkarouter" || p.(fake).key != "sk-gonka" {
|
|
t.Fatalf("%+v %v", a, err)
|
|
}
|
|
other, _ := st.CreateSession(ctx, "")
|
|
if _, a, _ := m.ForSession(ctx, other.ID); a.Provider != "openai" {
|
|
t.Fatalf("%+v", a)
|
|
}
|
|
}
|
|
|
|
func TestModelsUsesServerSideKey(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/v1/models" || r.Header.Get("Authorization") != "Bearer sk-gonka" {
|
|
http.Error(w, "no", 401)
|
|
return
|
|
}
|
|
fmt.Fprint(w, `{"data":[{"id":"b-model"},{"id":"a-model"}]}`)
|
|
}))
|
|
defer srv.Close()
|
|
m, _ := newMgr(t, nil)
|
|
m.entries["gonkarouter"].baseURL = srv.URL + "/v1"
|
|
ids, err := m.Models(context.Background(), "gonkarouter")
|
|
if err != nil || len(ids) != 2 || ids[0] != "a-model" {
|
|
t.Fatalf("%v %v", ids, err)
|
|
}
|
|
m.entries["gonkarouter"].apiKey = "wrong"
|
|
if _, err := m.Models(context.Background(), "gonkarouter"); err == nil {
|
|
t.Fatal("expected error")
|
|
}
|
|
}
|