114 lines
4.9 KiB
Go
114 lines
4.9 KiB
Go
package promptkit_test
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"reflect"
|
|
"strings"
|
|
"testing"
|
|
|
|
"go.maximumdirect.net/promptkit"
|
|
)
|
|
|
|
func TestExchangeCaptureAndPreparedContinuation(t *testing.T) {
|
|
client := &fakeLLMClient{response: &promptkit.GenerateResponse{Content: "answer"}}
|
|
engine := newSimpleEngine(t, "question", promptkit.WithLLMClient(client))
|
|
ctx := context.Background()
|
|
plain, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt"})
|
|
if err != nil || plain.Exchange != nil {
|
|
t.Fatalf("plain run = %#v, %v", plain, err)
|
|
}
|
|
initial, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
|
|
if err != nil || initial.Exchange == nil {
|
|
t.Fatalf("captured run = %#v, %v", initial, err)
|
|
}
|
|
want := []promptkit.RenderedMessage{{Role: promptkit.RoleUser, Content: "question"}, {Role: promptkit.RoleAssistant, Content: "answer"}}
|
|
if details := initial.Exchange.Details(); details.RunID != initial.RunID || !reflect.DeepEqual(details.Messages, want) {
|
|
t.Fatalf("details = %#v", details)
|
|
}
|
|
encoded, err := json.Marshal(initial)
|
|
if err != nil || strings.Contains(string(encoded), "exchange") || strings.Contains(string(encoded), "question") {
|
|
t.Fatalf("result JSON = %s, %v", encoded, err)
|
|
}
|
|
handle, err := engine.PrepareExecution(ctx, promptkit.RunRequest{PromptID: "prompt", ContinueFrom: initial.Exchange})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if handle.Details().ParentRunID != initial.RunID {
|
|
t.Fatal("prepared ancestry missing")
|
|
}
|
|
continued, err := engine.RunPrepared(ctx, handle)
|
|
if err != nil || continued.Exchange == nil || continued.ParentRunID != initial.RunID {
|
|
t.Fatalf("continued = %#v, %v", continued, err)
|
|
}
|
|
want = append(want, promptkit.RenderedMessage{Role: promptkit.RoleUser, Content: "question"}, promptkit.RenderedMessage{Role: promptkit.RoleAssistant, Content: "answer"})
|
|
if details := continued.Exchange.Details(); details.ParentRunID != initial.RunID || !reflect.DeepEqual(details.Messages, want) {
|
|
t.Fatalf("continued details = %#v", details)
|
|
}
|
|
}
|
|
|
|
func TestExchangeCapturesInvalidFinalOutput(t *testing.T) {
|
|
client := &fakeLLMClient{responses: []*promptkit.GenerateResponse{{Content: ""}, {Content: " "}}}
|
|
engine := newSimpleEngine(t, "question", promptkit.WithLLMClient(client))
|
|
result, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true,
|
|
Validation: &promptkit.OutputContract{Format: promptkit.FormatText, ValidationMode: promptkit.ValidationBasic, RepairAttempts: 1}})
|
|
if err != nil || result == nil || result.Validation.Status != promptkit.ValidationFailed || result.Exchange == nil {
|
|
t.Fatalf("invalid result = %#v, %v", result, err)
|
|
}
|
|
if messages := result.Exchange.Details().Messages; len(messages) != 2 || messages[1].Content != " " {
|
|
t.Fatalf("history = %#v", messages)
|
|
}
|
|
}
|
|
|
|
func TestExchangeFailuresLeaveSourceReusable(t *testing.T) {
|
|
client := &fakeLLMClient{response: &promptkit.GenerateResponse{Content: "answer"}}
|
|
engine := newSimpleEngine(t, "question", promptkit.WithLLMClient(client))
|
|
ctx := context.Background()
|
|
initial, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
original := initial.Exchange.Details()
|
|
request := promptkit.RunRequest{PromptID: "prompt", ContinueFrom: initial.Exchange}
|
|
handle, err := engine.PrepareExecution(ctx, request)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
handle.Discard()
|
|
if result, err := engine.RunPrepared(ctx, handle); result != nil || !errors.Is(err, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("discarded = %#v, %v", result, err)
|
|
}
|
|
client.err = errors.New("provider failure")
|
|
if result, err := engine.Run(ctx, request); result != nil || !errors.Is(err, promptkit.ErrLLMGenerate) {
|
|
t.Fatalf("failure = %#v, %v", result, err)
|
|
}
|
|
client.err = nil
|
|
result, err := engine.Run(ctx, request)
|
|
if err != nil || result.Exchange == nil || !reflect.DeepEqual(initial.Exchange.Details(), original) {
|
|
t.Fatalf("reused source = %#v, %v", result, err)
|
|
}
|
|
}
|
|
|
|
func TestEmptyExchangeRejectedBeforePreparation(t *testing.T) {
|
|
client := &fakeLLMClient{response: &promptkit.GenerateResponse{Content: "answer"}}
|
|
engine := newSimpleEngine(t, "question", promptkit.WithLLMClient(client))
|
|
request := promptkit.RunRequest{PromptID: "missing", ContinueFrom: &promptkit.Exchange{}}
|
|
ctx := context.Background()
|
|
if _, err := engine.Prepare(ctx, request); !errors.Is(err, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("prepare error = %v", err)
|
|
}
|
|
if err := engine.ValidateRequest(ctx, request); !errors.Is(err, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("validate error = %v", err)
|
|
}
|
|
if _, err := engine.PrepareExecution(ctx, request); !errors.Is(err, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("prepare execution error = %v", err)
|
|
}
|
|
if _, err := engine.Run(ctx, request); !errors.Is(err, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("run error = %v", err)
|
|
}
|
|
if len(client.requests) != 0 {
|
|
t.Fatal("invalid exchange reached the model")
|
|
}
|
|
}
|