178 lines
6.4 KiB
Go
178 lines
6.4 KiB
Go
package promptkit_test
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io/fs"
|
|
"os"
|
|
"path/filepath"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"testing/fstest"
|
|
|
|
"gitea.maximumdirect.net/eric/ConfigKit"
|
|
"gitea.maximumdirect.net/eric/promptkit"
|
|
)
|
|
|
|
const capturedPromptDocument = `
|
|
prompts:
|
|
captured:
|
|
"1":
|
|
default_profile: offline
|
|
session_id: '{{.session}}'
|
|
messages:
|
|
- role: user
|
|
content_file: body.tmpl
|
|
output:
|
|
format: text
|
|
validation_mode: none
|
|
`
|
|
|
|
func capturedPromptEngine(t *testing.T, source fs.FS) *promptkit.Engine {
|
|
t.Helper()
|
|
inputs, err := promptkit.FilesystemInputs("application prompts", source, "config")
|
|
if err != nil {
|
|
t.Fatalf("select prompt inputs: %v", err)
|
|
}
|
|
endpoint, model := "https://example.invalid/v1", "offline-model"
|
|
engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{
|
|
Underlays: []configkit.Underlay{
|
|
{Inputs: inputs},
|
|
{Inputs: []configkit.Input{{Source: promptkit.ProfileSource("offline profile", map[string]promptkit.ProfileConfig{
|
|
"offline": {Endpoint: &endpoint, Model: &model},
|
|
})}}},
|
|
},
|
|
}})
|
|
if err != nil {
|
|
t.Fatalf("construct engine: %v", err)
|
|
}
|
|
return engine
|
|
}
|
|
|
|
func TestEngineCapturesPromptContentUntilReconstruction(t *testing.T) {
|
|
files := fstest.MapFS{
|
|
"config/prompt.yaml": &fstest.MapFile{Data: []byte(capturedPromptDocument)},
|
|
"config/body.tmpl": &fstest.MapFile{Data: []byte("before {{.name}}")},
|
|
}
|
|
first := capturedPromptEngine(t, files)
|
|
files["config/body.tmpl"] = &fstest.MapFile{Data: []byte("after {{.name}}")}
|
|
request := promptkit.RunRequest{PromptID: "captured", Vars: map[string]string{"name": "Ada", "session": "session-1"}}
|
|
oldPrepared, err := first.Prepare(context.Background(), request)
|
|
if err != nil {
|
|
t.Fatalf("prepare with existing engine: %v", err)
|
|
}
|
|
second := capturedPromptEngine(t, files)
|
|
newPrepared, err := second.Prepare(context.Background(), request)
|
|
if err != nil {
|
|
t.Fatalf("prepare with rebuilt engine: %v", err)
|
|
}
|
|
if oldPrepared.Messages[0].Content != "before Ada" || newPrepared.Messages[0].Content != "after Ada" || oldPrepared.PromptHash == newPrepared.PromptHash {
|
|
t.Fatalf("captured content or hash did not follow engine lifetime: old=%+v new=%+v", oldPrepared, newPrepared)
|
|
}
|
|
}
|
|
|
|
type unreadableAfterConstructionFS struct {
|
|
fs.FS
|
|
closed atomic.Bool
|
|
reads atomic.Int64
|
|
}
|
|
|
|
func (source *unreadableAfterConstructionFS) Open(name string) (fs.File, error) {
|
|
if source.closed.Load() {
|
|
source.reads.Add(1)
|
|
return nil, fs.ErrPermission
|
|
}
|
|
return source.FS.Open(name)
|
|
}
|
|
|
|
func TestCapturedTemplatesBindInputsIndependentlyAcrossConcurrentPreparations(t *testing.T) {
|
|
source := &unreadableAfterConstructionFS{FS: fstest.MapFS{
|
|
"config/prompt.yaml": &fstest.MapFile{Data: []byte(capturedPromptDocument)},
|
|
"config/body.tmpl": &fstest.MapFile{Data: []byte(`{{.name}}: {{input "note"}}`)},
|
|
}}
|
|
engine := capturedPromptEngine(t, source)
|
|
source.closed.Store(true)
|
|
|
|
const calls = 24
|
|
var wait sync.WaitGroup
|
|
errs := make(chan error, calls)
|
|
for index := 0; index < calls; index++ {
|
|
wait.Add(1)
|
|
go func(index int) {
|
|
defer wait.Done()
|
|
name := fmt.Sprintf("caller-%d", index)
|
|
body := fmt.Sprintf("body-%d", index)
|
|
prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{
|
|
PromptID: "captured", SessionID: "direct-session",
|
|
Vars: map[string]string{"name": name},
|
|
Inputs: map[string]promptkit.ArtifactRef{"note": promptkit.Inline(body)},
|
|
})
|
|
if err == nil && (prepared.Messages[0].Content != name+": "+body || prepared.SessionID != "direct-session") {
|
|
err = fmt.Errorf("unexpected prepared content: %+v", prepared)
|
|
}
|
|
errs <- err
|
|
}(index)
|
|
}
|
|
wait.Wait()
|
|
close(errs)
|
|
for err := range errs {
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
if source.reads.Load() != 0 {
|
|
t.Fatalf("preparation reopened captured source %d times", source.reads.Load())
|
|
}
|
|
|
|
_, err := engine.Prepare(context.Background(), promptkit.RunRequest{
|
|
PromptID: "captured", Vars: map[string]string{"name": "Ada"},
|
|
Inputs: map[string]promptkit.ArtifactRef{"note": promptkit.Inline("body")},
|
|
})
|
|
if !errors.Is(err, promptkit.ErrPromptRender) {
|
|
t.Fatalf("missing session variable = %v, want runtime render error", err)
|
|
}
|
|
_, err = engine.Prepare(context.Background(), promptkit.RunRequest{
|
|
PromptID: "captured", SessionID: "direct-session", Vars: map[string]string{"name": "Ada"},
|
|
})
|
|
if !errors.Is(err, promptkit.ErrPromptRender) {
|
|
t.Fatalf("missing input helper artifact = %v, want runtime render error", err)
|
|
}
|
|
}
|
|
|
|
func TestPreparationStillReadsRequestArtifactsAndSchemasAfterConstruction(t *testing.T) {
|
|
endpoint, model := "https://example.invalid/v1", "offline-model"
|
|
schemas := fstest.MapFS{}
|
|
engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{
|
|
Underlays: []configkit.Underlay{{Inputs: []configkit.Input{{Source: promptkit.ProfileSource("profile", map[string]promptkit.ProfileConfig{
|
|
"offline": {Endpoint: &endpoint, Model: &model},
|
|
})}}}},
|
|
Main: &configkit.Input{Source: configkit.Mapping("structured prompt", map[string]any{
|
|
"prompts": map[string]any{"structured": map[string]any{"1": map[string]any{
|
|
"default_profile": "offline",
|
|
"inputs": []any{map[string]any{"name": "note", "required": true}},
|
|
"messages": []any{map[string]any{"role": "user", "content": `{{input "note"}}`}},
|
|
"output": map[string]any{"format": "json", "validation_mode": "json_schema", "schema_path": "schema.json"},
|
|
}}},
|
|
})},
|
|
}}, promptkit.WithSchemaFS(schemas, "."))
|
|
if err != nil {
|
|
t.Fatalf("construct engine before request resources exist: %v", err)
|
|
}
|
|
schemas["schema.json"] = &fstest.MapFile{Data: []byte(`{"$schema":"https://json-schema.org/draft/2020-12/schema","type":"object"}`)}
|
|
artifactPath := filepath.Join(t.TempDir(), "note.txt")
|
|
if err := os.WriteFile(artifactPath, []byte("request artifact"), 0o600); err != nil {
|
|
t.Fatalf("write request artifact: %v", err)
|
|
}
|
|
prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{PromptID: "structured", Inputs: map[string]promptkit.ArtifactRef{
|
|
"note": promptkit.File(artifactPath),
|
|
}})
|
|
if err != nil {
|
|
t.Fatalf("prepare with resources supplied after construction: %v", err)
|
|
}
|
|
if len(prepared.Messages) != 1 || prepared.Messages[0].Content != "request artifact" || prepared.StructuredOutput == nil || prepared.StructuredOutput.JSONSchema == nil {
|
|
t.Fatalf("runtime artifact or schema was not prepared: %+v", prepared)
|
|
}
|
|
}
|