91 lines
3.7 KiB
Go
91 lines
3.7 KiB
Go
|
package promptkit_test
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"errors"
|
||
|
|
"io"
|
||
|
|
"net/http"
|
||
|
|
"reflect"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"go.maximumdirect.net/configkit"
|
||
|
|
"go.maximumdirect.net/promptkit"
|
||
|
|
)
|
||
|
|
|
||
|
|
func bundledPromptConfig(client *http.Client) promptkit.Config {
|
||
|
|
input := configkit.Input{Source: configkit.Mapping("prompt", map[string]any{"prompts": []any{map[string]any{
|
||
|
|
"id": "prompt", "messages": []any{map[string]any{"role": "user", "content": "hello"}}, "output": map[string]any{"format": "text", "validation_mode": "none"},
|
||
|
|
}}})}
|
||
|
|
return promptkit.Config{HTTPClient: client, Sources: promptkit.ConfigurationSources{Main: &input}}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestBundledProfilePreparationAndExecution(t *testing.T) {
|
||
|
|
client := &fakeLLMClient{response: &promptkit.GenerateResponse{Content: "ok"}}
|
||
|
|
engine, err := promptkit.NewEngine(bundledPromptConfig(nil), promptkit.WithLLMClient(client))
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if _, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prompt"}); !errors.Is(err, promptkit.ErrProfileRequired) {
|
||
|
|
t.Fatalf("missing profile = %v", err)
|
||
|
|
}
|
||
|
|
temperature := 0.3
|
||
|
|
request := promptkit.RunRequest{PromptID: "prompt", ProfileID: "gpt-luna-latest", Execution: &promptkit.ExecutionTargetOverride{Temperature: &temperature}}
|
||
|
|
prepared, err := engine.PrepareExecution(context.Background(), request)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
details := prepared.Details()
|
||
|
|
if details.SelectedProfileID != "gpt-luna-latest" || details.EffectiveModelParams.Model != "~openai/gpt-luna-latest" || details.EffectiveModelParams.BackendID != "openrouter" || details.EffectiveModelParams.Temperature != temperature {
|
||
|
|
t.Fatalf("prepared target = %+v", details.EffectiveModelParams)
|
||
|
|
}
|
||
|
|
if _, err := engine.Run(context.Background(), request); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if len(client.requests) != 1 || client.requests[0].Target.Model != "~openai/gpt-luna-latest" || client.requests[0].Target.BackendID != "openrouter" {
|
||
|
|
t.Fatalf("requests = %+v", client.requests)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestBundledProfileHTTPPayloads(t *testing.T) {
|
||
|
|
t.Setenv("OPENROUTER_API_KEY", "synthetic-key")
|
||
|
|
for _, test := range []struct{ id, model, tier string }{
|
||
|
|
{"auto-low", "openrouter/auto", "low"}, {"auto-medium", "openrouter/auto", "medium"}, {"auto-high", "openrouter/auto", "high"}, {"auto-xhigh", "openrouter/auto", "xhigh"}, {"auto-max", "openrouter/auto", "max"},
|
||
|
|
{"auto", "openrouter/auto", ""}, {"free", "openrouter/free", ""}, {"gpt-luna-latest", "~openai/gpt-luna-latest", ""},
|
||
|
|
} {
|
||
|
|
t.Run(test.id, func(t *testing.T) {
|
||
|
|
var body map[string]any
|
||
|
|
transport := roundTripFunc(func(request *http.Request) (*http.Response, error) {
|
||
|
|
if request.Header.Get("Authorization") != "Bearer synthetic-key" {
|
||
|
|
t.Errorf("authorization = %q", request.Header.Get("Authorization"))
|
||
|
|
}
|
||
|
|
if err := json.NewDecoder(request.Body).Decode(&body); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"ok"}}]}`))}, nil
|
||
|
|
})
|
||
|
|
engine, err := promptkit.NewEngine(bundledPromptConfig(&http.Client{Transport: transport}))
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if _, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", ProfileID: test.id}); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if body["model"] != test.model || body["extra_params"] != nil {
|
||
|
|
t.Fatalf("request = %+v", body)
|
||
|
|
}
|
||
|
|
if test.tier == "" {
|
||
|
|
if _, exists := body["plugins"]; exists {
|
||
|
|
t.Fatalf("unexpected plugins = %+v", body)
|
||
|
|
}
|
||
|
|
return
|
||
|
|
}
|
||
|
|
want := []any{map[string]any{"id": "auto-router", "cost_tier": test.tier}}
|
||
|
|
if !reflect.DeepEqual(body["plugins"], want) {
|
||
|
|
t.Fatalf("plugins = %#v, want %#v", body["plugins"], want)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|