Escape schema resources and reuse compiled plans
This commit is contained in:
230
engine_test.go
230
engine_test.go
@@ -6,14 +6,17 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/fs"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"testing/fstest"
|
||||
"time"
|
||||
@@ -2424,6 +2427,162 @@ func TestWithProfilesRejectsCyclicExtraParams(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestPublicSchemaGraphErrorsHaveSourceParity(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
files map[string]string
|
||||
valid bool
|
||||
}{
|
||||
{
|
||||
name: "invalid keyword",
|
||||
files: map[string]string{"root.json": `{"type":42}`},
|
||||
},
|
||||
{
|
||||
name: "malformed direct reference",
|
||||
files: map[string]string{
|
||||
"root.json": `{"$ref":"child.json"}`,
|
||||
"child.json": `{`,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "missing direct reference",
|
||||
files: map[string]string{"root.json": `{"$ref":"missing.json"}`},
|
||||
},
|
||||
{
|
||||
name: "malformed second-level reference",
|
||||
files: map[string]string{
|
||||
"root.json": `{"$ref":"child.json"}`,
|
||||
"child.json": `{"$ref":"grandchild.json"}`,
|
||||
"grandchild.json": `{`,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "missing second-level reference",
|
||||
files: map[string]string{
|
||||
"root.json": `{"$ref":"child.json"}`,
|
||||
"child.json": `{"$ref":"missing.json"}`,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "unsupported referenced dialect",
|
||||
files: map[string]string{
|
||||
"root.json": `{"$ref":"child.json"}`,
|
||||
"child.json": `{"$schema":"http://json-schema.org/draft-07/schema#","type":"integer"}`,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "escaping reference",
|
||||
files: map[string]string{
|
||||
"root.json": `{"$ref":"../outside.json"}`,
|
||||
"../outside.json": `{"type":"integer"}`,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "remote reference",
|
||||
files: map[string]string{"root.json": `{"$ref":"https://example.test/schema.json"}`},
|
||||
},
|
||||
{
|
||||
name: "valid multi-document graph",
|
||||
files: map[string]string{
|
||||
"root.json": `{
|
||||
"type":"object",
|
||||
"required":["value"],
|
||||
"properties":{"value":{"$ref":"child.json"}}
|
||||
}`,
|
||||
"child.json": `{"$ref":"nested/value.json"}`,
|
||||
"nested/value.json": `{"type":"integer","minimum":2}`,
|
||||
},
|
||||
valid: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, source := range []string{"operating system files", "fs.FS"} {
|
||||
t.Run(source, func(t *testing.T) {
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
client := &fakeLLMClient{response: &promptkit.GenerateResponse{Content: `{"value":3}`}}
|
||||
engine := newPublicSchemaGraphEngine(t, source, tc.files, client)
|
||||
result, err := engine.Run(context.Background(), promptkit.RunRequest{
|
||||
PromptID: "schema.graph.prompt",
|
||||
Inputs: map[string]promptkit.ArtifactRef{
|
||||
"transcript": promptkit.Inline("Rin opens the gate."),
|
||||
},
|
||||
})
|
||||
if tc.valid {
|
||||
if err != nil {
|
||||
t.Fatalf("run valid schema graph: %v", err)
|
||||
}
|
||||
if result.Validation.Status != promptkit.ValidationPassed || !result.Validation.IsValid {
|
||||
t.Fatalf("validation = %+v, want passed", result.Validation)
|
||||
}
|
||||
if len(client.requests) != 1 {
|
||||
t.Fatalf("generation calls = %d, want 1", len(client.requests))
|
||||
}
|
||||
return
|
||||
}
|
||||
if !errors.Is(err, promptkit.ErrValidation) {
|
||||
t.Fatalf("error = %v, want ErrValidation", err)
|
||||
}
|
||||
if result != nil {
|
||||
t.Fatalf("result = %+v, want nil", result)
|
||||
}
|
||||
if len(client.requests) != 0 {
|
||||
t.Fatalf("invalid schema reached generation: %+v", client.requests)
|
||||
}
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunReadsSchemaGraphOncePerOperation(t *testing.T) {
|
||||
source := &countingSchemaFS{
|
||||
FS: fstest.MapFS{
|
||||
"schemas/root.json": &fstest.MapFile{Data: []byte(`{
|
||||
"type":"object",
|
||||
"required":["value"],
|
||||
"properties":{"value":{"$ref":"child.json"}}
|
||||
}`)},
|
||||
"schemas/child.json": &fstest.MapFile{Data: []byte(`{"$ref":"nested/value.json"}`)},
|
||||
"schemas/nested/value.json": &fstest.MapFile{Data: []byte(`{"type":"integer","minimum":2}`)},
|
||||
},
|
||||
reads: make(map[string]int),
|
||||
}
|
||||
client := &fakeLLMClient{response: &promptkit.GenerateResponse{Content: `{"value":3}`}}
|
||||
engine, err := promptkit.NewEngine(promptkit.Config{},
|
||||
promptkit.WithPromptFS(publicStructuredPromptFS("schema.count.prompt", "root.json"), "prompts"),
|
||||
promptkit.WithSchemaFS(source, "schemas"),
|
||||
promptkit.WithProfiles(promptkit.Profile{
|
||||
ID: "contract-fast", Endpoint: "http://example.test/v1", Model: "schema-model",
|
||||
}),
|
||||
promptkit.WithLLMClient(client),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("construct engine: %v", err)
|
||||
}
|
||||
request := promptkit.RunRequest{
|
||||
PromptID: "schema.count.prompt",
|
||||
Inputs: map[string]promptkit.ArtifactRef{
|
||||
"transcript": promptkit.Inline("Rin opens the gate."),
|
||||
},
|
||||
}
|
||||
|
||||
for operation := 1; operation <= 2; operation++ {
|
||||
result, err := engine.Run(context.Background(), request)
|
||||
if err != nil {
|
||||
t.Fatalf("run operation %d: %v", operation, err)
|
||||
}
|
||||
if !result.Validation.IsValid {
|
||||
t.Fatalf("run operation %d validation = %+v", operation, result.Validation)
|
||||
}
|
||||
for _, name := range []string{"schemas/root.json", "schemas/child.json", "schemas/nested/value.json"} {
|
||||
if got := source.readCount(name); got != operation {
|
||||
t.Fatalf("after operation %d, reads for %q = %d, want %d", operation, name, got, operation)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPreparedStructuredOutputRetainsExactSchemaNumbers(t *testing.T) {
|
||||
const schema = `{
|
||||
"type": "number",
|
||||
@@ -2987,7 +3146,7 @@ messages:
|
||||
output:
|
||||
format: json
|
||||
validation_mode: json_schema
|
||||
schema_path: ` + schemaPath + `
|
||||
schema_path: ` + strconv.Quote(schemaPath) + `
|
||||
repair_attempts: 0
|
||||
`)},
|
||||
}
|
||||
@@ -3095,6 +3254,75 @@ func (r *recordingArtifactReader) Read(ctx context.Context, ref promptkit.Artifa
|
||||
return r.artifact, nil
|
||||
}
|
||||
|
||||
type countingSchemaFS struct {
|
||||
fs.FS
|
||||
mu sync.Mutex
|
||||
reads map[string]int
|
||||
}
|
||||
|
||||
func (f *countingSchemaFS) ReadFile(name string) ([]byte, error) {
|
||||
f.mu.Lock()
|
||||
f.reads[name]++
|
||||
f.mu.Unlock()
|
||||
return fs.ReadFile(f.FS, name)
|
||||
}
|
||||
|
||||
func (f *countingSchemaFS) readCount(name string) int {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
return f.reads[name]
|
||||
}
|
||||
|
||||
func newPublicSchemaGraphEngine(
|
||||
t *testing.T,
|
||||
sourceName string,
|
||||
files map[string]string,
|
||||
client promptkit.LLMClient,
|
||||
) *promptkit.Engine {
|
||||
t.Helper()
|
||||
options := []promptkit.Option{
|
||||
promptkit.WithPromptFS(publicStructuredPromptFS("schema.graph.prompt", "root.json"), "prompts"),
|
||||
promptkit.WithProfiles(promptkit.Profile{
|
||||
ID: "contract-fast", Endpoint: "http://example.test/v1", Model: "schema-model",
|
||||
}),
|
||||
promptkit.WithLLMClient(client),
|
||||
}
|
||||
config := promptkit.Config{}
|
||||
|
||||
switch sourceName {
|
||||
case "operating system files":
|
||||
workspace := t.TempDir()
|
||||
schemaRoot := filepath.Join(workspace, "schemas")
|
||||
if err := os.Mkdir(schemaRoot, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for name, body := range files {
|
||||
fileName := filepath.Clean(filepath.Join(schemaRoot, filepath.FromSlash(name)))
|
||||
if err := os.MkdirAll(filepath.Dir(fileName), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(fileName, []byte(body), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
config.SchemaDir = schemaRoot
|
||||
case "fs.FS":
|
||||
schemaFS := make(fstest.MapFS, len(files))
|
||||
for name, body := range files {
|
||||
schemaFS[path.Clean(path.Join("schemas", name))] = &fstest.MapFile{Data: []byte(body)}
|
||||
}
|
||||
options = append(options, promptkit.WithSchemaFS(schemaFS, "schemas"))
|
||||
default:
|
||||
t.Fatalf("unknown schema source %q", sourceName)
|
||||
}
|
||||
|
||||
engine, err := promptkit.NewEngine(config, options...)
|
||||
if err != nil {
|
||||
t.Fatalf("construct engine: %v", err)
|
||||
}
|
||||
return engine
|
||||
}
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
|
||||
Reference in New Issue
Block a user