109 lines
4.9 KiB
Go
109 lines
4.9 KiB
Go
package promptkit_test
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/santhosh-tekuri/jsonschema/v6"
|
|
"go.maximumdirect.net/configkit"
|
|
"go.maximumdirect.net/promptkit"
|
|
)
|
|
|
|
type rejectWireSchemaLoad struct{}
|
|
|
|
func (rejectWireSchemaLoad) Load(uri string) (any, error) {
|
|
return nil, fmt.Errorf("unexpected external resource %s", uri)
|
|
}
|
|
|
|
func TestSchemaFilePreparationAndNativeSubmission(t *testing.T) {
|
|
firstDir, secondDir := t.TempDir(), t.TempDir()
|
|
first, second := filepath.Join(firstDir, "schema.json"), filepath.Join(secondDir, "schema.json")
|
|
if err := os.WriteFile(first, []byte(`{"type":"string"}`), 0600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := os.WriteFile(second, []byte(`{"$defs":{"value":{"type":"integer"}},"$ref":"#/$defs/value"}`), 0600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
input := configkit.Input{Source: configkit.Mapping("schema test", map[string]any{
|
|
"profiles": []any{map[string]any{"id": "profile", "endpoint": "https://example.invalid/v1", "model": "model"}},
|
|
"prompts": []any{map[string]any{"id": "prompt", "default_profile": "profile", "messages": []any{map[string]any{"role": "user", "content": "hello"}}, "output": map[string]any{"format": "json", "validation_mode": "json_schema", "schema_path": "schema.json"}}},
|
|
})}
|
|
var payload map[string]any
|
|
client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
|
|
if err := json.NewDecoder(request.Body).Decode(&payload); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"3"}}]}`))}, nil
|
|
})}
|
|
engine, err := promptkit.NewEngine(promptkit.Config{HTTPClient: client, Sources: promptkit.ConfigurationSources{Main: &input}}, promptkit.WithSchemaFile(first), promptkit.WithSchemaFile(second))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := engine.ValidateRequest(context.Background(), promptkit.RunRequest{PromptID: "prompt"}); err != nil {
|
|
t.Fatalf("validate request: %v", err)
|
|
}
|
|
prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
|
|
if err != nil || prepared.Details().StructuredOutput == nil {
|
|
t.Fatalf("prepared = %+v, %v", prepared, err)
|
|
}
|
|
if err := os.Remove(second); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
result, err := engine.RunPrepared(context.Background(), prepared)
|
|
if err != nil || result.Validation.Status != promptkit.ValidationPassed {
|
|
t.Fatalf("run = %+v, %v", result, err)
|
|
}
|
|
encoded, err := json.Marshal(payload["response_format"])
|
|
if err != nil || !strings.Contains(string(encoded), `"$defs"`) || !strings.Contains(string(encoded), `#/$defs/value`) {
|
|
t.Fatalf("wire schema = %s, %v", encoded, err)
|
|
}
|
|
compiler := jsonschema.NewCompiler()
|
|
compiler.UseLoader(rejectWireSchemaLoad{})
|
|
const submission = "https://provider.invalid/request-schema"
|
|
document := payload["response_format"].(map[string]any)["json_schema"].(map[string]any)["schema"]
|
|
if err := compiler.AddResource(submission, document); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := compiler.Compile(submission); err != nil {
|
|
t.Fatalf("outbound schema needs local source context: %v", err)
|
|
}
|
|
if _, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prompt"}); !errors.Is(err, promptkit.ErrValidation) {
|
|
t.Fatalf("missing root error = %v", err)
|
|
}
|
|
}
|
|
|
|
func TestSchemaFileRejectsInvalidSourcesAndExternalReferences(t *testing.T) {
|
|
dir := t.TempDir()
|
|
for _, name := range []string{filepath.Join(dir, "missing.json"), dir} {
|
|
if _, err := promptkit.NewEngine(promptkit.Config{}, promptkit.WithSchemaFile(name)); !errors.Is(err, promptkit.ErrInvalidConfig) {
|
|
t.Fatalf("source %q = %v", name, err)
|
|
}
|
|
}
|
|
root := filepath.Join(dir, "schema.json")
|
|
if err := os.WriteFile(root, []byte(`{"$defs":{"unused":{"$ref":"other.json"}}}`), 0600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
input := configkit.Input{Source: configkit.Mapping("schema test", map[string]any{
|
|
"profiles": []any{map[string]any{"id": "profile", "endpoint": "https://example.invalid/v1", "model": "model"}},
|
|
"prompts": []any{map[string]any{"id": "prompt", "default_profile": "profile", "messages": []any{map[string]any{"role": "user", "content": "hello"}}, "output": map[string]any{"format": "json", "validation_mode": "json_schema", "schema_path": "schema.json"}}},
|
|
})}
|
|
calls := 0
|
|
client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { calls++; return nil, errors.New("unexpected dispatch") })}
|
|
engine, err := promptkit.NewEngine(promptkit.Config{HTTPClient: client, Sources: promptkit.ConfigurationSources{Main: &input}}, promptkit.WithSchemaFile(root))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
_, err = engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
|
|
if !errors.Is(err, promptkit.ErrValidation) || calls != 0 {
|
|
t.Fatalf("external reference = %v, calls = %d", err, calls)
|
|
}
|
|
}
|