Ilker ManapandClaude Opus 5.5
fe310576e7
CI / test (push) Successful in 8m52s
An external model takes api_key_files, each a file and an optional weight, in place of api_key_file. Each key is an endpoint of its own (external:MODEL#1, #2), so the router's scoring, retries and breakers work per key. A key answering 429 is passed over while its breaker is open. A key answering 401 is passed over too, and external_key_refused is sent, naming the endpoint and never the key; with a single key the 401 is still the client's. A key whose file cannot be read is left out, and usage records which key answered by its endpoint. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
101 lines
3.7 KiB
Go
101 lines
3.7 KiB
Go
package control
|
|
|
|
import (
|
|
"context"
|
|
"log/slog"
|
|
"os"
|
|
"path/filepath"
|
|
"sync"
|
|
"testing"
|
|
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
|
)
|
|
|
|
type recordSink struct {
|
|
mu sync.Mutex
|
|
infos []endpoint.Info
|
|
times int
|
|
}
|
|
|
|
func (s *recordSink) SetEndpoints(_ context.Context, infos []endpoint.Info) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
s.infos, s.times = infos, s.times+1
|
|
}
|
|
|
|
type noDispatch struct{}
|
|
|
|
func (noDispatch) StartShard(context.Context, reconciler.ShardCommand) error { return nil }
|
|
func (noDispatch) StopShard(context.Context, string, string) error { return nil }
|
|
func (noDispatch) DrainShard(context.Context, string, string) error { return nil }
|
|
|
|
func TestAnExternalModelIsPublishedWithItsKeyAndNothingIsPlacedForIt(t *testing.T) {
|
|
dir := t.TempDir()
|
|
keyFile := filepath.Join(dir, "provider.key")
|
|
if err := os.WriteFile(keyFile, []byte("sk-provider-1\n"), 0o600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
models := []catalog.Model{
|
|
{Name: "cloud", Source: "external", Aliases: []string{"gpt"}, Scaling: catalog.Scaling{Min: 3, Max: 3},
|
|
External: &catalog.External{URL: "https://api.example.com/", Model: "gpt-4o-mini", APIKeyFile: keyFile}},
|
|
{Name: "gone", Source: "external",
|
|
External: &catalog.External{URL: "https://api.example.com", Model: "m", APIKeyFile: filepath.Join(dir, "missing")}},
|
|
}
|
|
cat, err := catalog.New(models)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
sink := &recordSink{}
|
|
c := New(DefaultConfig(), catalog.NewLive(cat), noDispatch{}, sink, slog.New(slog.DiscardHandler))
|
|
c.Publish(t.Context())
|
|
|
|
if len(sink.infos) != 1 {
|
|
t.Fatalf("published %+v; the model whose key is missing must not be offered", sink.infos)
|
|
}
|
|
e := sink.infos[0]
|
|
if e.ID != "external:cloud" || e.URL != "https://api.example.com" || e.APIKey != "sk-provider-1" ||
|
|
e.UpstreamModel != "gpt-4o-mini" || e.Engine != endpoint.EngineExternal || len(e.Models) != 2 {
|
|
t.Errorf("endpoint %+v", e)
|
|
}
|
|
// A key rotated on disk is published on the next pass.
|
|
_ = os.WriteFile(keyFile, []byte("sk-provider-2"), 0o600)
|
|
c.Publish(t.Context())
|
|
if sink.times != 2 || sink.infos[0].APIKey != "sk-provider-2" {
|
|
t.Errorf("the rotated key was not published: %d pushes, key %q", sink.times, sink.infos[0].APIKey)
|
|
}
|
|
// Scaling of 3 or not, nothing is wanted of the cluster for a model
|
|
// served elsewhere, and so it never reads as short or unplaceable.
|
|
for _, h := range c.Reconciler().ModelHealth() {
|
|
if h.Model == "cloud" && (h.Desired != 0 || h.Unplaceable) {
|
|
t.Errorf("an external model wants replicas: %+v", h)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Several keys for one provider: an endpoint each, numbered, with its
|
|
// weight; one whose file cannot be read is left out and the others stay.
|
|
func TestEachOfAProvidersKeysIsAnEndpointOfItsOwn(t *testing.T) {
|
|
dir := t.TempDir()
|
|
a, c := filepath.Join(dir, "a.key"), filepath.Join(dir, "c.key")
|
|
_ = os.WriteFile(a, []byte("sk-a\n"), 0o600)
|
|
_ = os.WriteFile(c, []byte("sk-c"), 0o600)
|
|
cat, err := catalog.New([]catalog.Model{{Name: "cloud", Source: "external",
|
|
External: &catalog.External{URL: "https://api.example.com", Model: "gpt-4o-mini",
|
|
APIKeyFiles: []catalog.ExternalKey{{File: a, Weight: 3}, {File: filepath.Join(dir, "gone")}, {File: c}}}}})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
sink := &recordSink{}
|
|
New(DefaultConfig(), catalog.NewLive(cat), noDispatch{}, sink, slog.New(slog.DiscardHandler)).Publish(t.Context())
|
|
got := map[string]endpoint.Info{}
|
|
for _, e := range sink.infos {
|
|
got[e.ID] = e
|
|
}
|
|
if len(got) != 2 || got["external:cloud#1"].APIKey != "sk-a" || got["external:cloud#1"].Weight != 3 ||
|
|
got["external:cloud#3"].APIKey != "sk-c" || got["external:cloud#3"].Weight != 1 {
|
|
t.Errorf("published %+v", sink.infos)
|
|
}
|
|
}
|