232 lines
8.7 KiB
Go
232 lines
8.7 KiB
Go
package promptkit_test
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
"testing/fstest"
|
|
|
|
"gitea.maximumdirect.net/eric/ConfigKit"
|
|
"gitea.maximumdirect.net/eric/promptkit"
|
|
)
|
|
|
|
func TestProfileEnvironmentLookupAcceptsNonShellNames(t *testing.T) {
|
|
const name = "promptkit.test-api-key"
|
|
for _, test := range []struct {
|
|
name, value, direct, want string
|
|
unset, clear bool
|
|
}{
|
|
{name: "set", value: " synthetic-key ", want: "Bearer synthetic-key"},
|
|
{name: "empty"},
|
|
{name: "whitespace", value: " "},
|
|
{name: "unset", unset: true},
|
|
{name: "direct precedence", value: "synthetic-key", direct: "direct-key", want: "Bearer direct-key"},
|
|
{name: "explicit clear", value: "synthetic-key", clear: true},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
t.Setenv(name, test.value)
|
|
if test.unset {
|
|
if err := os.Unsetenv(name); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
cfg := simpleConfig("hello")
|
|
cfg.Sources.Underlays = []configkit.Underlay{{Inputs: []configkit.Input{*cfg.Sources.Main}}}
|
|
lowerName := name
|
|
cfg.Sources.Underlays = append(cfg.Sources.Underlays, configkit.Underlay{Inputs: []configkit.Input{{
|
|
Source: promptkit.ProfileSource("default credentials", map[string]promptkit.ProfileConfig{
|
|
"profile": {APIKeyEnv: &lowerName},
|
|
}),
|
|
}}})
|
|
env := " " + name + " "
|
|
if test.clear {
|
|
env = ""
|
|
}
|
|
cfg.Sources.Main = &configkit.Input{Source: promptkit.ProfileSource("credentials", map[string]promptkit.ProfileConfig{
|
|
"profile": {APIKeyEnv: &env},
|
|
})}
|
|
var calls int
|
|
var auth string
|
|
cfg.HTTPClient = &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
|
|
calls++
|
|
auth = request.Header.Get("Authorization")
|
|
return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"ok"}}]}`))}, nil
|
|
})}
|
|
engine, err := promptkit.NewEngine(cfg)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
_, err = engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", APIKey: test.direct})
|
|
if err != nil || calls != 1 || auth != test.want {
|
|
t.Fatalf("run: err=%v calls=%d authorization=%q, want %q", err, calls, auth, test.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestPublicConstructionWithLocalBackendAndProgrammaticProfile(t *testing.T) {
|
|
backendID, model := "workstation", "offline-model"
|
|
prompt := configkit.Input{Source: configkit.Mapping("prompt", map[string]any{
|
|
"prompts": map[string]any{"hello": map[string]any{"1": map[string]any{
|
|
"default_profile": "offline",
|
|
"messages": []any{map[string]any{"role": "user", "content": "Hello"}},
|
|
"output": map[string]any{"format": "text", "validation_mode": "none"},
|
|
}}},
|
|
})}
|
|
client := &fakeLLMClient{response: &promptkit.GenerateResponse{Content: "reply"}}
|
|
engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{
|
|
Underlays: []configkit.Underlay{{Inputs: []configkit.Input{
|
|
{Source: promptkit.BackendSource("local backend", map[string]promptkit.BackendConfig{
|
|
backendID: promptkit.LocalBackend("http://localhost:8000/v1", 2),
|
|
})},
|
|
{Source: promptkit.ProfileSource("profile", map[string]promptkit.ProfileConfig{
|
|
"offline": {Backend: &backendID, Model: &model},
|
|
})},
|
|
}}},
|
|
Main: &prompt,
|
|
}}, promptkit.WithLLMClient(client))
|
|
if err != nil {
|
|
t.Fatalf("construct engine: %v", err)
|
|
}
|
|
result, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "hello"})
|
|
if err != nil {
|
|
t.Fatalf("run prompt: %v", err)
|
|
}
|
|
if result.RawOutput != "reply" || result.SelectedBackendID != backendID || result.EffectiveModelParams.Model != model {
|
|
t.Fatalf("unexpected result: %+v", result)
|
|
}
|
|
if len(client.requests) != 1 || client.requests[0].Target.Endpoint != "http://localhost:8000/v1" {
|
|
t.Fatalf("unexpected model request: %+v", client.requests)
|
|
}
|
|
}
|
|
|
|
func TestPublicConstructionCombinesEmbeddedOperatorAndProgrammaticInputs(t *testing.T) {
|
|
embedded, err := promptkit.FilesystemInputs("embedded defaults", fstest.MapFS{
|
|
"config/default.yaml": &fstest.MapFile{Data: []byte(`
|
|
profiles:
|
|
application:
|
|
endpoint: https://example.invalid/v1
|
|
model: default-model
|
|
prompts:
|
|
hello:
|
|
"1":
|
|
default_profile: application
|
|
messages:
|
|
- role: user
|
|
content: Hello
|
|
output:
|
|
format: text
|
|
validation_mode: none
|
|
`)},
|
|
}, "config")
|
|
if err != nil {
|
|
t.Fatalf("select embedded defaults: %v", err)
|
|
}
|
|
operatorDirectory := t.TempDir()
|
|
if err := os.WriteFile(filepath.Join(operatorDirectory, "override.yaml"), []byte("profiles:\n application:\n model: operator-model\n"), 0o644); err != nil {
|
|
t.Fatalf("write operator input: %v", err)
|
|
}
|
|
operator, err := promptkit.DirectoryInputs(operatorDirectory)
|
|
if err != nil {
|
|
t.Fatalf("select operator inputs: %v", err)
|
|
}
|
|
programmaticModel := "programmatic-model"
|
|
programmatic := configkit.Input{Source: promptkit.ProfileSource("programmatic settings", map[string]promptkit.ProfileConfig{
|
|
"application": {Model: &programmaticModel},
|
|
})}
|
|
for _, test := range []struct {
|
|
name string
|
|
apply bool
|
|
want string
|
|
}{
|
|
{name: "embedded", want: "default-model"},
|
|
{name: "operator", apply: true, want: "operator-model"},
|
|
{name: "programmatic", apply: true, want: "programmatic-model"},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
underlays := []configkit.Underlay{{Inputs: embedded}}
|
|
if test.apply {
|
|
underlays = append(underlays, configkit.Underlay{Inputs: operator})
|
|
}
|
|
if test.name == "programmatic" {
|
|
underlays = append(underlays, configkit.Underlay{Inputs: []configkit.Input{programmatic}})
|
|
}
|
|
engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{Underlays: underlays}})
|
|
if err != nil {
|
|
t.Fatalf("construct engine: %v", err)
|
|
}
|
|
prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{PromptID: "hello"})
|
|
if err != nil {
|
|
t.Fatalf("prepare prompt: %v", err)
|
|
}
|
|
if prepared.EffectiveModelParams.Model != test.want {
|
|
t.Fatalf("model = %q, want %q", prepared.EffectiveModelParams.Model, test.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestPublicConstructionPreservesConfigKitLoadError(t *testing.T) {
|
|
engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{Main: &configkit.Input{
|
|
Source: configkit.Mapping("invalid", map[string]any{"profiles": map[string]any{"broken": map[string]any{"unknown": true}}}),
|
|
}}})
|
|
var loadErr *configkit.LoadError
|
|
if engine != nil || !errors.Is(err, promptkit.ErrInvalidConfig) || !errors.As(err, &loadErr) {
|
|
t.Fatalf("construction = (%#v, %v), want invalid config preserving LoadError", engine, err)
|
|
}
|
|
}
|
|
|
|
func TestPublicConstructionRejectsSuppliedInvalidMainInput(t *testing.T) {
|
|
main := configkit.Input{}
|
|
engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{Main: &main}})
|
|
if engine != nil || !errors.Is(err, promptkit.ErrInvalidConfig) {
|
|
t.Fatalf("construction = (%#v, %v), want invalid supplied main input", engine, err)
|
|
}
|
|
}
|
|
|
|
func TestPublicConstructionCanOverlayMaintainedBackend(t *testing.T) {
|
|
endpoint := "http://local-openrouter.example/v1"
|
|
engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{
|
|
Underlays: []configkit.Underlay{{Inputs: []configkit.Input{{Source: promptkit.BackendSource("operator backend", map[string]promptkit.BackendConfig{
|
|
promptkit.BackendOpenRouter: {Endpoint: &endpoint},
|
|
})}}}},
|
|
}})
|
|
if err != nil {
|
|
t.Fatalf("construct engine with backend overlay: %v", err)
|
|
}
|
|
inspection, err := engine.InspectProfile(context.Background(), "mistral-small-3")
|
|
if err != nil {
|
|
t.Fatalf("inspect catalog profile: %v", err)
|
|
}
|
|
if inspection.EffectiveModelParams.Endpoint != endpoint {
|
|
t.Fatalf("effective endpoint = %q, want %q", inspection.EffectiveModelParams.Endpoint, endpoint)
|
|
}
|
|
}
|
|
|
|
func TestPreparationRejectsAmbiguousPromptVersion(t *testing.T) {
|
|
definition := func(content string) map[string]any {
|
|
return map[string]any{
|
|
"messages": []any{map[string]any{"role": "user", "content": content}},
|
|
"output": map[string]any{"format": "text", "validation_mode": "none"},
|
|
}
|
|
}
|
|
engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{
|
|
Main: &configkit.Input{Source: configkit.Mapping("versions", map[string]any{
|
|
"prompts": map[string]any{"hello": map[string]any{
|
|
"1": definition("first"), "2": definition("second"),
|
|
}},
|
|
})},
|
|
}})
|
|
if err != nil {
|
|
t.Fatalf("construct engine: %v", err)
|
|
}
|
|
if result, err := engine.Prepare(context.Background(), promptkit.RunRequest{PromptID: "hello"}); result != nil || !errors.Is(err, promptkit.ErrPromptLoad) {
|
|
t.Fatalf("ambiguous preparation = (%#v, %v), want ErrPromptLoad", result, err)
|
|
}
|
|
}
|