289 lines
9.7 KiB
Go
289 lines
9.7 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"io/fs"
|
|
"strings"
|
|
"testing"
|
|
"testing/fstest"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/promptkit"
|
|
)
|
|
|
|
func TestAssetRegistryCombinesPromptAndSchemaSources(t *testing.T) {
|
|
registry := NewAssetRegistry()
|
|
mustRegisterPromptFS(t, registry, fstest.MapFS{
|
|
"prompts/test.yaml": {Data: []byte(validPromptYAML("schemas/out.json"))},
|
|
"prompts/messages/user.tmpl": {Data: []byte(`Input: {{ input "transcript" }}`)},
|
|
"prompts/messages/task.tmpl": {Data: []byte("Return JSON.")},
|
|
"schemas/ignored/schema.json": {Data: []byte(`{"type":"object"}`)},
|
|
}, "prompts")
|
|
mustRegisterSchemaFS(t, registry, fstest.MapFS{
|
|
"root/schemas/out.json": {Data: []byte(`{"type":"object"}`)},
|
|
}, "root")
|
|
|
|
engine := newPromptKitAssetTestEngine(t, registry)
|
|
prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{
|
|
PromptID: "asset.test",
|
|
ProfileID: "asset-test-profile",
|
|
Inputs: map[string]promptkit.ArtifactRef{
|
|
"transcript": promptkit.Inline(`{"ok":true}`),
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Prepare() error = %v, want nil", err)
|
|
}
|
|
if got := len(prepared.Messages); got != 2 {
|
|
t.Fatalf("message count = %d, want 2", got)
|
|
}
|
|
if prepared.OutputContract.SchemaPath != "schemas/out.json" {
|
|
t.Fatalf("schema path = %q, want schemas/out.json", prepared.OutputContract.SchemaPath)
|
|
}
|
|
}
|
|
|
|
func TestAssetRegistryPrepareFailsForMissingPromptAsset(t *testing.T) {
|
|
registry := NewAssetRegistry()
|
|
mustRegisterPromptFS(t, registry, fstest.MapFS{
|
|
"test.yaml": {Data: []byte(validPromptYAML("out.json"))},
|
|
}, ".")
|
|
mustRegisterSchemaFS(t, registry, fstest.MapFS{
|
|
"out.json": {Data: []byte(`{"type":"object"}`)},
|
|
}, ".")
|
|
|
|
engine := newPromptKitAssetTestEngine(t, registry)
|
|
_, err := engine.Prepare(context.Background(), promptkit.RunRequest{
|
|
PromptID: "asset.test",
|
|
ProfileID: "asset-test-profile",
|
|
Inputs: map[string]promptkit.ArtifactRef{
|
|
"transcript": promptkit.Inline(`{"ok":true}`),
|
|
},
|
|
})
|
|
if err == nil || !strings.Contains(err.Error(), "content_file") {
|
|
t.Fatalf("Prepare() error = %v, want missing content_file error", err)
|
|
}
|
|
}
|
|
|
|
func TestAssetRegistryPrepareFailsForMissingSchemaAsset(t *testing.T) {
|
|
registry := NewAssetRegistry()
|
|
mustRegisterPromptFS(t, registry, fstest.MapFS{
|
|
"test.yaml": {Data: []byte(validPromptYAML("missing.json"))},
|
|
"messages/user.tmpl": {Data: []byte(`Input: {{ input "transcript" }}`)},
|
|
"messages/task.tmpl": {Data: []byte("Return JSON.")},
|
|
}, ".")
|
|
mustRegisterSchemaFS(t, registry, fstest.MapFS{
|
|
"present.json": {Data: []byte(`{"type":"object"}`)},
|
|
}, ".")
|
|
|
|
engine := newPromptKitAssetTestEngine(t, registry)
|
|
_, err := engine.Prepare(context.Background(), promptkit.RunRequest{
|
|
PromptID: "asset.test",
|
|
ProfileID: "asset-test-profile",
|
|
Inputs: map[string]promptkit.ArtifactRef{
|
|
"transcript": promptkit.Inline(`{"ok":true}`),
|
|
},
|
|
})
|
|
if err == nil || !strings.Contains(err.Error(), "missing.json") {
|
|
t.Fatalf("Prepare() error = %v, want missing schema error", err)
|
|
}
|
|
}
|
|
|
|
func TestAssetRegistryRejectsDuplicateAssetPaths(t *testing.T) {
|
|
registry := NewAssetRegistry()
|
|
mustRegisterPromptFS(t, registry, fstest.MapFS{"one/prompt.yaml": {Data: []byte("id: one")}}, "one")
|
|
mustRegisterPromptFS(t, registry, fstest.MapFS{"two/prompt.yaml": {Data: []byte("id: two")}}, "two")
|
|
|
|
_, err := registry.PromptFS()
|
|
if err == nil || !strings.Contains(err.Error(), "duplicate asset path") {
|
|
t.Fatalf("PromptFS() error = %v, want duplicate path error", err)
|
|
}
|
|
}
|
|
|
|
func TestAssetRegistryCombinesFallbackProfileSources(t *testing.T) {
|
|
registry := NewAssetRegistry()
|
|
if err := registry.RegisterFallbackProfileFS(fstest.MapFS{
|
|
"first/profiles/one.yaml": {Data: []byte("id: one\nmodel: first\n")},
|
|
}, "first/profiles"); err != nil {
|
|
t.Fatalf("RegisterFallbackProfileFS() error = %v, want nil", err)
|
|
}
|
|
if err := registry.RegisterFallbackProfileFS(fstest.MapFS{
|
|
"second/two.yaml": {Data: []byte("id: two\nmodel: second\n")},
|
|
}, "second"); err != nil {
|
|
t.Fatalf("RegisterFallbackProfileFS() error = %v, want nil", err)
|
|
}
|
|
|
|
fallbackFS, err := registry.FallbackProfileFS()
|
|
if err != nil {
|
|
t.Fatalf("FallbackProfileFS() error = %v, want nil", err)
|
|
}
|
|
for _, name := range []string{"one.yaml", "two.yaml"} {
|
|
if _, err := fs.ReadFile(fallbackFS, name); err != nil {
|
|
t.Fatalf("FallbackProfileFS().ReadFile(%q) error = %v, want nil", name, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestAssetRegistryRejectsInvalidFallbackProfileRoot(t *testing.T) {
|
|
registry := NewAssetRegistry()
|
|
err := registry.RegisterFallbackProfileFS(fstest.MapFS{}, "../profiles")
|
|
if err == nil || !strings.Contains(err.Error(), "invalid path") {
|
|
t.Fatalf("RegisterFallbackProfileFS() error = %v, want invalid root error", err)
|
|
}
|
|
}
|
|
|
|
func TestAssetRegistryRejectsUnreadableFallbackProfileAssets(t *testing.T) {
|
|
registry := NewAssetRegistry()
|
|
if err := registry.RegisterFallbackProfileFS(unreadableAssetFS{}, "."); err != nil {
|
|
t.Fatalf("RegisterFallbackProfileFS() error = %v, want nil", err)
|
|
}
|
|
_, err := registry.FallbackProfileFS()
|
|
if err == nil || !strings.Contains(err.Error(), "permission denied") {
|
|
t.Fatalf("FallbackProfileFS() error = %v, want unreadable asset error", err)
|
|
}
|
|
}
|
|
|
|
func TestAssetRegistryRejectsDuplicateFallbackProfilePaths(t *testing.T) {
|
|
registry := NewAssetRegistry()
|
|
if err := registry.RegisterFallbackProfileFS(fstest.MapFS{"first/profile.yaml": {Data: []byte("id: first\n")}}, "first"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := registry.RegisterFallbackProfileFS(fstest.MapFS{"second/profile.yaml": {Data: []byte("id: second\n")}}, "second"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
_, err := registry.FallbackProfileFS()
|
|
if err == nil || !strings.Contains(err.Error(), "duplicate asset path") {
|
|
t.Fatalf("FallbackProfileFS() error = %v, want duplicate path error", err)
|
|
}
|
|
}
|
|
|
|
func TestAssetRegistryFallbackProfileDigestTracksContentWithoutLeakingIt(t *testing.T) {
|
|
digestFor := func(content string) string {
|
|
t.Helper()
|
|
registry := NewAssetRegistry()
|
|
if err := registry.RegisterFallbackProfileFS(fstest.MapFS{
|
|
"profiles/profile.yaml": {Data: []byte(content)},
|
|
}, "profiles"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
digest, err := registry.FallbackProfileDigest()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return digest
|
|
}
|
|
|
|
first := digestFor("id: fallback\nmodel: model-one\n")
|
|
second := digestFor("id: fallback\nmodel: model-two\n")
|
|
if first == second {
|
|
t.Fatalf("fallback profile digests = %q and %q, want content change", first, second)
|
|
}
|
|
if !strings.HasPrefix(first, "sha256:") || strings.Contains(first, "model-one") || strings.Contains(first, "profile.yaml") {
|
|
t.Fatalf("fallback profile digest leaked source details: %q", first)
|
|
}
|
|
}
|
|
|
|
func TestAssetRegistryCombinesNamespacedPromptSources(t *testing.T) {
|
|
registry := NewAssetRegistry()
|
|
mustRegisterPromptFS(t, registry, fstest.MapFS{
|
|
"dnd.spells/dnd.spells.yaml": {Data: []byte(validPromptYAML("schema.json"))},
|
|
"dnd.spells/task.md": {Data: []byte("spell task")},
|
|
"dnd.spells/instructions.md": {Data: []byte("spell instructions")},
|
|
}, ".")
|
|
mustRegisterPromptFS(t, registry, fstest.MapFS{
|
|
"dnd.scenes/dnd.scenes.yaml": {Data: []byte(validPromptYAML("schema.json"))},
|
|
"dnd.scenes/task.md": {Data: []byte("scene task")},
|
|
"dnd.scenes/instructions.md": {Data: []byte("scene instructions")},
|
|
}, ".")
|
|
|
|
fsys, err := registry.PromptFS()
|
|
if err != nil {
|
|
t.Fatalf("PromptFS() error = %v, want nil", err)
|
|
}
|
|
for _, name := range []string{
|
|
"dnd.spells/dnd.spells.yaml",
|
|
"dnd.spells/task.md",
|
|
"dnd.spells/instructions.md",
|
|
"dnd.scenes/dnd.scenes.yaml",
|
|
"dnd.scenes/task.md",
|
|
"dnd.scenes/instructions.md",
|
|
} {
|
|
if _, err := fsys.Open(name); err != nil {
|
|
t.Fatalf("PromptFS().Open(%q) error = %v, want nil", name, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestHashAssetsOmitsRawAssetContent(t *testing.T) {
|
|
hash, err := HashAssets([]AssetHashPart{{
|
|
FS: fstest.MapFS{"prompt.md": {Data: []byte("secret prompt text")}},
|
|
Path: "prompt.md",
|
|
}})
|
|
if err != nil {
|
|
t.Fatalf("HashAssets() error = %v, want nil", err)
|
|
}
|
|
if !strings.HasPrefix(hash, "sha256:") {
|
|
t.Fatalf("hash = %q, want sha256-prefixed value", hash)
|
|
}
|
|
if strings.Contains(hash, "secret prompt text") {
|
|
t.Fatalf("hash leaked asset content")
|
|
}
|
|
}
|
|
|
|
func newPromptKitAssetTestEngine(t *testing.T, registry *AssetRegistry) *promptkit.Engine {
|
|
t.Helper()
|
|
options, err := registry.PromptKitOptions()
|
|
if err != nil {
|
|
t.Fatalf("PromptKitOptions() error = %v, want nil", err)
|
|
}
|
|
options = append(options, promptkit.WithProfiles(promptkit.OpenAICompatibleProfile(promptkit.OpenAICompatibleProfileConfig{
|
|
ID: "asset-test-profile",
|
|
Endpoint: "http://127.0.0.1:1/v1",
|
|
Model: "asset-test-model",
|
|
})))
|
|
engine, err := promptkit.NewEngine(promptkit.Config{Timeout: time.Second}, options...)
|
|
if err != nil {
|
|
t.Fatalf("NewEngine() error = %v, want nil", err)
|
|
}
|
|
return engine
|
|
}
|
|
|
|
func mustRegisterPromptFS(t *testing.T, registry *AssetRegistry, fsys fstest.MapFS, root string) {
|
|
t.Helper()
|
|
if err := registry.RegisterPromptFS(fsys, root); err != nil {
|
|
t.Fatalf("RegisterPromptFS() error = %v, want nil", err)
|
|
}
|
|
}
|
|
|
|
func mustRegisterSchemaFS(t *testing.T, registry *AssetRegistry, fsys fstest.MapFS, root string) {
|
|
t.Helper()
|
|
if err := registry.RegisterSchemaFS(fsys, root); err != nil {
|
|
t.Fatalf("RegisterSchemaFS() error = %v, want nil", err)
|
|
}
|
|
}
|
|
|
|
func validPromptYAML(schemaPath string) string {
|
|
return `id: asset.test
|
|
version: "v1"
|
|
inputs:
|
|
- name: transcript
|
|
required: true
|
|
content_type: application/json
|
|
messages:
|
|
- role: user
|
|
content_file: ./messages/user.tmpl
|
|
- role: user
|
|
content_file: ./messages/task.tmpl
|
|
output:
|
|
format: json
|
|
validation_mode: json_schema
|
|
schema_path: ` + schemaPath + `
|
|
repair_attempts: 0
|
|
`
|
|
}
|
|
|
|
type unreadableAssetFS struct{}
|
|
|
|
func (unreadableAssetFS) Open(name string) (fs.File, error) {
|
|
return nil, &fs.PathError{Op: "open", Path: name, Err: fs.ErrPermission}
|
|
}
|