385 lines
14 KiB
Go
385 lines
14 KiB
Go
package promptkitadapter
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
promptkit "gitea.maximumdirect.net/eric/promptkit"
|
|
"gitea.maximumdirect.net/eric/weatherreporter/internal/promptexec"
|
|
)
|
|
|
|
type fakeClient struct {
|
|
mu sync.Mutex
|
|
response *promptkit.GenerateResponse
|
|
err error
|
|
calls int
|
|
requests []promptkit.GenerateRequest
|
|
block bool
|
|
}
|
|
|
|
type recordingReader struct {
|
|
ref promptkit.ArtifactRef
|
|
}
|
|
|
|
func (reader *recordingReader) Read(_ context.Context, ref promptkit.ArtifactRef) (*promptkit.Artifact, error) {
|
|
reader.ref = ref
|
|
return &promptkit.Artifact{
|
|
Name: "data_package",
|
|
ContentType: "application/yaml",
|
|
Body: []byte(ref.Body),
|
|
URI: ref.URI,
|
|
Hash: "input-hash",
|
|
}, nil
|
|
}
|
|
|
|
func (client *fakeClient) Generate(ctx context.Context, request promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
|
client.mu.Lock()
|
|
client.calls++
|
|
client.requests = append(client.requests, request)
|
|
block := client.block
|
|
response := client.response
|
|
err := client.err
|
|
client.mu.Unlock()
|
|
if block {
|
|
<-ctx.Done()
|
|
return nil, ctx.Err()
|
|
}
|
|
return response, err
|
|
}
|
|
|
|
func (client *fakeClient) callCount() int {
|
|
client.mu.Lock()
|
|
defer client.mu.Unlock()
|
|
return client.calls
|
|
}
|
|
|
|
func (client *fakeClient) request() promptkit.GenerateRequest {
|
|
client.mu.Lock()
|
|
defer client.mu.Unlock()
|
|
return client.requests[0]
|
|
}
|
|
|
|
func TestInspectPromptAndProfile(t *testing.T) {
|
|
adapter := newTestAdapter(t, &fakeClient{})
|
|
inspection, err := adapter.InspectPrompt(context.Background(), "weather.daily_generated_text", "1.0.0")
|
|
if err != nil {
|
|
t.Fatalf("InspectPrompt() error = %v", err)
|
|
}
|
|
if inspection.PromptID != "weather.daily_generated_text" || inspection.PromptVersion != "1.0.0" || inspection.DefaultProfileID != "gemini-flash-latest" {
|
|
t.Fatalf("inspection = %#v", inspection)
|
|
}
|
|
if len(inspection.Inputs) != 1 || inspection.Inputs[0].Name != "data_package" || !inspection.Inputs[0].Required || inspection.Inputs[0].ContentType != "application/yaml" {
|
|
t.Fatalf("inputs = %#v", inspection.Inputs)
|
|
}
|
|
if inspection.Output.Format != "json" || inspection.Output.ValidationMode != "json_schema" || inspection.Output.SchemaPath != "daily.generated_text.schema.json" {
|
|
t.Fatalf("output = %#v", inspection.Output)
|
|
}
|
|
|
|
profile, err := adapter.InspectProfile(context.Background(), "test-profile")
|
|
if err != nil {
|
|
t.Fatalf("InspectProfile() error = %v", err)
|
|
}
|
|
if profile.ProfileID != "test-profile" || profile.BackendID != "" || profile.ModelName != "test-model" || profile.CredentialRequired {
|
|
t.Fatalf("profile = %#v", profile)
|
|
}
|
|
if strings.Contains(fmt.Sprintf("%#v", profile), "https://profile.example") {
|
|
t.Fatalf("profile leaks endpoint: %#v", profile)
|
|
}
|
|
|
|
builtin, err := adapter.InspectProfile(context.Background(), "gemini-flash-latest")
|
|
if err != nil {
|
|
t.Fatalf("InspectProfile(builtin) error = %v", err)
|
|
}
|
|
if builtin.ProfileID != "gemini-flash-latest" || builtin.ModelName == "" {
|
|
t.Fatalf("builtin profile = %#v", builtin)
|
|
}
|
|
}
|
|
|
|
func TestExecuteUsesPreparedInlineDataPackage(t *testing.T) {
|
|
client := &fakeClient{response: validResponse()}
|
|
adapter := newTestAdapter(t, client)
|
|
request := testExecuteRequest()
|
|
callbackCalls := 0
|
|
result, err := adapter.Execute(context.Background(), request, func(preparation promptexec.Preparation, debug *promptexec.PreparationDebug) error {
|
|
callbackCalls++
|
|
if preparation.PromptID != request.PromptID || preparation.PromptVersion != request.PromptVersion || preparation.DataPackagePath != request.DataPackagePath || preparation.ModelName != "test-model" {
|
|
t.Fatalf("preparation = %#v", preparation)
|
|
}
|
|
if debug != nil {
|
|
t.Fatalf("debug = %#v, want nil", debug)
|
|
}
|
|
if client.callCount() != 0 {
|
|
t.Fatal("provider called before preparation callback")
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Execute() error = %v", err)
|
|
}
|
|
if callbackCalls != 1 || client.callCount() != 1 {
|
|
t.Fatalf("callback/provider calls = %d/%d, want 1/1", callbackCalls, client.callCount())
|
|
}
|
|
if result == nil || result.Validation.Status != promptexec.ValidationPassed || string(result.RawOutput) != client.response.Content || result.DataPackagePath != request.DataPackagePath {
|
|
t.Fatalf("result = %#v", result)
|
|
}
|
|
if result.Debug != nil {
|
|
t.Fatalf("debug = %#v, want nil", result.Debug)
|
|
}
|
|
providerRequest := client.request()
|
|
if providerRequest.Target.Model != "test-model" || providerRequest.Target.Endpoint != "https://profile.example/v1" {
|
|
t.Fatalf("provider target = %#v", providerRequest.Target)
|
|
}
|
|
if len(providerRequest.Prompt.Messages) == 0 || !strings.Contains(providerRequest.Prompt.Messages[2].Content, string(request.DataPackage)) {
|
|
t.Fatalf("rendered messages do not contain exact data package: %#v", providerRequest.Prompt.Messages)
|
|
}
|
|
}
|
|
|
|
func TestExecuteUsesExactInlineDataPackageProvenance(t *testing.T) {
|
|
client := &fakeClient{response: validResponse()}
|
|
reader := &recordingReader{}
|
|
adapter := newTestAdapterWithOptions(t, client, promptkit.WithArtifactReader(reader))
|
|
request := testExecuteRequest()
|
|
if _, err := adapter.Execute(context.Background(), request, nil); err != nil {
|
|
t.Fatalf("Execute() error = %v", err)
|
|
}
|
|
if reader.ref.Type != promptkit.ArtifactRefInline || reader.ref.URI != request.DataPackagePath || reader.ref.Body != string(request.DataPackage) {
|
|
t.Fatalf("artifact ref = %#v, want exact inline data package provenance", reader.ref)
|
|
}
|
|
}
|
|
|
|
func TestExecuteCapturesSensitiveDebugOnlyWhenRequested(t *testing.T) {
|
|
client := &fakeClient{response: validResponse()}
|
|
adapter := newTestAdapter(t, client)
|
|
request := testExecuteRequest()
|
|
request.CaptureDebug = true
|
|
var preparationDebug *promptexec.PreparationDebug
|
|
result, err := adapter.Execute(context.Background(), request, func(preparation promptexec.Preparation, debug *promptexec.PreparationDebug) error {
|
|
preparationDebug = debug
|
|
if strings.Contains(fmt.Sprintf("%#v", preparation), "https://profile.example") || strings.Contains(fmt.Sprintf("%#v", preparation), string(request.DataPackage)) {
|
|
t.Fatalf("safe preparation leaks sensitive content: %#v", preparation)
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Execute() error = %v", err)
|
|
}
|
|
if preparationDebug == nil || preparationDebug.Endpoint != "https://profile.example/v1" || len(preparationDebug.RenderedMessages) == 0 || len(preparationDebug.StructuredSchema) == 0 || len(preparationDebug.ParametersJSON) == 0 {
|
|
t.Fatalf("preparation debug = %#v", preparationDebug)
|
|
}
|
|
if result.Debug == nil || string(result.Debug.RawOutput) != client.response.Content {
|
|
t.Fatalf("execution debug = %#v", result.Debug)
|
|
}
|
|
}
|
|
|
|
func TestExecuteCallbackFailurePreventsGeneration(t *testing.T) {
|
|
client := &fakeClient{response: validResponse()}
|
|
adapter := newTestAdapter(t, client)
|
|
callbackError := errors.New("save preparation")
|
|
result, err := adapter.Execute(context.Background(), testExecuteRequest(), func(promptexec.Preparation, *promptexec.PreparationDebug) error {
|
|
return callbackError
|
|
})
|
|
if result != nil || !errors.Is(err, callbackError) || client.callCount() != 0 {
|
|
t.Fatalf("result/error/provider calls = %#v/%v/%d", result, err, client.callCount())
|
|
}
|
|
}
|
|
|
|
func TestExecuteReturnsCompletedValidationRejection(t *testing.T) {
|
|
client := &fakeClient{response: &promptkit.GenerateResponse{Content: `{"summary":42}`, Usage: promptkit.TokenUsage{TotalTokens: 5}}}
|
|
adapter := newTestAdapter(t, client)
|
|
result, err := adapter.Execute(context.Background(), testExecuteRequest(), nil)
|
|
if err != nil {
|
|
t.Fatalf("Execute() error = %v", err)
|
|
}
|
|
if result == nil || result.Validation.Status != promptexec.ValidationFailed || len(result.Validation.Diagnostics) == 0 || string(result.RawOutput) != client.response.Content {
|
|
t.Fatalf("result = %#v", result)
|
|
}
|
|
}
|
|
|
|
func TestExecuteClassifiesOperationalFailures(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
client *fakeClient
|
|
context func() (context.Context, context.CancelFunc)
|
|
category promptexec.ErrorCategory
|
|
}{
|
|
{
|
|
name: "generation",
|
|
client: &fakeClient{err: errors.New("provider response body")},
|
|
context: func() (context.Context, context.CancelFunc) {
|
|
return context.WithCancel(context.Background())
|
|
},
|
|
category: promptexec.Generation,
|
|
},
|
|
{
|
|
name: "canceled",
|
|
client: &fakeClient{block: true},
|
|
context: func() (context.Context, context.CancelFunc) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
return ctx, func() {}
|
|
},
|
|
category: promptexec.Canceled,
|
|
},
|
|
{
|
|
name: "deadline",
|
|
client: &fakeClient{block: true},
|
|
context: func() (context.Context, context.CancelFunc) {
|
|
return context.WithTimeout(context.Background(), time.Nanosecond)
|
|
},
|
|
category: promptexec.DeadlineExceeded,
|
|
},
|
|
}
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
adapter := newTestAdapter(t, test.client)
|
|
ctx, cancel := test.context()
|
|
defer cancel()
|
|
result, err := adapter.Execute(ctx, testExecuteRequest(), nil)
|
|
if result != nil || err == nil || promptexec.CategoryOf(err) != test.category {
|
|
t.Fatalf("result/error/category = %#v/%v/%q, want %q", result, err, promptexec.CategoryOf(err), test.category)
|
|
}
|
|
if strings.Contains(err.Error(), "provider response body") {
|
|
t.Fatalf("error leaks provider detail: %v", err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestClassifyPromptkitErrors(t *testing.T) {
|
|
tests := []struct {
|
|
err error
|
|
category promptexec.ErrorCategory
|
|
}{
|
|
{promptkit.ErrInvalidConfig, promptexec.InvalidConfiguration},
|
|
{promptkit.ErrInvalidRequest, promptexec.InvalidRequest},
|
|
{promptkit.ErrPromptNotFound, promptexec.PromptNotFound},
|
|
{promptkit.ErrPromptLoad, promptexec.PromptLoad},
|
|
{promptkit.ErrProfileNotFound, promptexec.ProfileNotFound},
|
|
{promptkit.ErrProfileLoad, promptexec.ProfileLoad},
|
|
{promptkit.ErrAPIKeyEnvMissing, promptexec.MissingCredential},
|
|
{promptkit.ErrArtifactLoad, promptexec.ArtifactLoad},
|
|
{promptkit.ErrPromptRender, promptexec.PromptRender},
|
|
{promptkit.ErrLLMGenerate, promptexec.Generation},
|
|
{promptkit.ErrValidation, promptexec.OperationalValidation},
|
|
{&promptkit.CapacityError{BackendID: "local"}, promptexec.Capacity},
|
|
}
|
|
for _, test := range tests {
|
|
t.Run(string(test.category), func(t *testing.T) {
|
|
got := classifyError(test.err)
|
|
if promptexec.CategoryOf(got) != test.category {
|
|
t.Fatalf("category = %q, want %q", promptexec.CategoryOf(got), test.category)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestNewValidatesConfiguration(t *testing.T) {
|
|
if _, err := New(Config{ProfileDirectory: "profiles", ProfileFile: "profile.yml"}); promptexec.CategoryOf(err) != promptexec.InvalidConfiguration {
|
|
t.Fatalf("profile source error = %v", err)
|
|
}
|
|
if _, err := New(Config{LocalConcurrencyLimit: 1}); promptexec.CategoryOf(err) != promptexec.InvalidConfiguration {
|
|
t.Fatalf("local concurrency error = %v", err)
|
|
}
|
|
if _, err := New(Config{LocalEndpoint: "not a URL"}); promptexec.CategoryOf(err) != promptexec.InvalidConfiguration {
|
|
t.Fatalf("local endpoint error = %v", err)
|
|
}
|
|
}
|
|
|
|
func TestLocalBackendAndMissingCredentialBehavior(t *testing.T) {
|
|
profiles := testProfileDirectory(t, `id: local-profile
|
|
backend: local
|
|
model: local-model
|
|
`)
|
|
adapter, err := newAdapterForTest(Config{
|
|
ProfileDirectory: profiles,
|
|
LocalEndpoint: "https://local.example/v1",
|
|
LocalConcurrencyLimit: 1,
|
|
}, &fakeClient{})
|
|
if err != nil {
|
|
t.Fatalf("newAdapterForTest(local) error = %v", err)
|
|
}
|
|
profile, err := adapter.InspectProfile(context.Background(), "local-profile")
|
|
if err != nil || profile.BackendID != promptkit.BackendLocal || profile.ModelName != "local-model" {
|
|
t.Fatalf("local profile/error = %#v/%v", profile, err)
|
|
}
|
|
if got := classifyError(&promptkit.CapacityError{BackendID: promptkit.BackendLocal}); promptexec.CategoryOf(got) != promptexec.Capacity {
|
|
t.Fatalf("capacity classification = %v", got)
|
|
}
|
|
|
|
credentialProfiles := testProfileDirectory(t, `id: credential-profile
|
|
endpoint: https://profile.example/v1
|
|
model: test-model
|
|
api_key_env: WEATHERREPORTER_TEST_MISSING_KEY
|
|
`)
|
|
client := &fakeClient{response: validResponse()}
|
|
credentialAdapter, err := newAdapterForTest(Config{ProfileDirectory: credentialProfiles}, client)
|
|
if err != nil {
|
|
t.Fatalf("newAdapterForTest(credential) error = %v", err)
|
|
}
|
|
credentialProfile, err := credentialAdapter.InspectProfile(context.Background(), "credential-profile")
|
|
if err != nil || credentialProfile.CredentialRequired || credentialProfile.APIKeyEnv != "WEATHERREPORTER_TEST_MISSING_KEY" {
|
|
t.Fatalf("credential profile/error = %#v/%v", credentialProfile, err)
|
|
}
|
|
request := testExecuteRequest()
|
|
request.ProfileID = "credential-profile"
|
|
result, err := credentialAdapter.Execute(context.Background(), request, nil)
|
|
if result != nil || promptexec.CategoryOf(err) != promptexec.MissingCredential || client.callCount() != 0 {
|
|
t.Fatalf("credential result/category/calls = %#v/%q/%d", result, promptexec.CategoryOf(err), client.callCount())
|
|
}
|
|
}
|
|
|
|
func newTestAdapter(t *testing.T, client promptkit.LLMClient) *Adapter {
|
|
return newTestAdapterWithOptions(t, client)
|
|
}
|
|
|
|
func newTestAdapterWithOptions(t *testing.T, client promptkit.LLMClient, options ...promptkit.Option) *Adapter {
|
|
t.Helper()
|
|
profiles := testProfileDirectory(t, `id: test-profile
|
|
endpoint: https://profile.example/v1
|
|
model: test-model
|
|
temperature: 0.2
|
|
max_tokens: 300
|
|
top_p: 1
|
|
timeout_seconds: 30
|
|
`)
|
|
options = append(options, promptkit.WithLLMClient(client))
|
|
adapter, err := newAdapter(Config{ProfileDirectory: profiles, Timeout: time.Second}, options...)
|
|
if err != nil {
|
|
t.Fatalf("newAdapter() error = %v", err)
|
|
}
|
|
return adapter
|
|
}
|
|
|
|
func testProfileDirectory(t *testing.T, profile string) string {
|
|
t.Helper()
|
|
profiles := t.TempDir()
|
|
if err := os.WriteFile(filepath.Join(profiles, "profile.yml"), []byte(profile), 0o600); err != nil {
|
|
t.Fatalf("write profile: %v", err)
|
|
}
|
|
return profiles
|
|
}
|
|
|
|
func testExecuteRequest() promptexec.ExecuteRequest {
|
|
return promptexec.ExecuteRequest{
|
|
PromptID: "weather.daily_generated_text",
|
|
PromptVersion: "1.0.0",
|
|
ProfileID: "test-profile",
|
|
DataPackage: []byte("report:\n id: daily\nbriefing: {}\n"),
|
|
DataPackagePath: "data-packages/daily/data_package.yaml",
|
|
}
|
|
}
|
|
|
|
func validResponse() *promptkit.GenerateResponse {
|
|
return &promptkit.GenerateResponse{
|
|
Content: `{"summary":"A quiet day is expected.","forecast_discussion":["High pressure keeps conditions settled."],"confidence":"High."}`,
|
|
Usage: promptkit.TokenUsage{PromptTokens: 12, CompletionTokens: 8, TotalTokens: 20},
|
|
}
|
|
}
|