236 lines
8.6 KiB
Go
236 lines
8.6 KiB
Go
package promptkit_test
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"testing"
|
|
|
|
"gitea.maximumdirect.net/eric/promptkit"
|
|
)
|
|
|
|
func TestBuiltInGenerationError(t *testing.T) {
|
|
const (
|
|
codeMarker = "provider-code-marker"
|
|
typeMarker = "provider-type-marker"
|
|
messageMarker = "provider-message-marker"
|
|
)
|
|
engine := newBuiltInGenerationErrorEngine(t, http.StatusUnprocessableEntity,
|
|
`{"error":{"code":"`+codeMarker+`","type":"`+typeMarker+`","message":"`+messageMarker+`"}}`)
|
|
|
|
result, err := engine.Run(context.Background(), generationErrorRunRequest())
|
|
if result != nil {
|
|
t.Fatalf("Run result = %#v, want nil", result)
|
|
}
|
|
assertGenerationError(t, err, http.StatusUnprocessableEntity, codeMarker, typeMarker, messageMarker)
|
|
|
|
preparedEngine := newBuiltInGenerationErrorEngine(t, http.StatusServiceUnavailable, `{"error":{}}`)
|
|
prepared, err := preparedEngine.PrepareExecution(context.Background(), generationErrorRunRequest())
|
|
if err != nil {
|
|
t.Fatalf("PrepareExecution: %v", err)
|
|
}
|
|
result, err = preparedEngine.RunPrepared(context.Background(), prepared)
|
|
if result != nil {
|
|
t.Fatalf("RunPrepared result = %#v, want nil", result)
|
|
}
|
|
assertGenerationError(t, err, http.StatusServiceUnavailable, "", "", "")
|
|
}
|
|
|
|
func TestBuiltInGenerationErrorWithAppendedMessages(t *testing.T) {
|
|
const messageMarker = "combined-request-provider-marker"
|
|
engine := newBuiltInGenerationErrorEngine(t, http.StatusUnprocessableEntity,
|
|
`{"error":{"message":"`+messageMarker+`"}}`)
|
|
request := generationErrorRunRequest()
|
|
request.AppendedMessages = []promptkit.RenderedMessage{
|
|
{Role: promptkit.RoleAssistant, Content: "previous response"},
|
|
{Role: promptkit.RoleUser, Content: "consumer correction"},
|
|
}
|
|
|
|
result, err := engine.Run(context.Background(), request)
|
|
if result != nil {
|
|
t.Fatalf("Run result = %#v, want nil", result)
|
|
}
|
|
assertGenerationError(t, err, http.StatusUnprocessableEntity, "", "", messageMarker)
|
|
}
|
|
|
|
func TestBuiltInRepairGenerationError(t *testing.T) {
|
|
const (
|
|
codeMarker = "repair-code-marker"
|
|
typeMarker = "repair-type-marker"
|
|
messageMarker = "repair-message-marker"
|
|
)
|
|
calls := 0
|
|
config := contractConfig(frameworkSchemaDir)
|
|
config.HTTPClient = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
calls++
|
|
if calls == 1 {
|
|
body := `{"choices":[{"message":{"content":"not-json"}}]}`
|
|
return &http.Response{StatusCode: http.StatusOK, ContentLength: int64(len(body)), Body: io.NopCloser(strings.NewReader(body))}, nil
|
|
}
|
|
body := `{"error":{"code":"` + codeMarker + `","type":"` + typeMarker + `","message":"` + messageMarker + `"}}`
|
|
return &http.Response{StatusCode: http.StatusUnprocessableEntity, ContentLength: int64(len(body)), Body: io.NopCloser(strings.NewReader(body))}, nil
|
|
})}
|
|
engine, err := promptkit.NewEngine(config)
|
|
if err != nil {
|
|
t.Fatalf("NewEngine: %v", err)
|
|
}
|
|
req := generationErrorRunRequest()
|
|
req.Validation = &promptkit.OutputContract{
|
|
Format: promptkit.FormatJSON,
|
|
ValidationMode: promptkit.ValidationJSON,
|
|
RepairAttempts: 1,
|
|
}
|
|
|
|
result, err := engine.Run(context.Background(), req)
|
|
if result != nil {
|
|
t.Fatalf("Run result = %#v, want nil", result)
|
|
}
|
|
if calls != 2 {
|
|
t.Fatalf("provider calls = %d, want 2", calls)
|
|
}
|
|
assertGenerationError(t, err, http.StatusUnprocessableEntity, codeMarker, typeMarker, messageMarker)
|
|
}
|
|
|
|
func TestBuiltInGenerationFailureCategories(t *testing.T) {
|
|
transportCause := errors.New("private transport diagnostic")
|
|
config := contractConfig(frameworkSchemaDir)
|
|
config.HTTPClient = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
return nil, transportCause
|
|
})}
|
|
engine, err := promptkit.NewEngine(config)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
result, err := engine.Run(context.Background(), generationErrorRunRequest())
|
|
if result != nil || !errors.Is(err, promptkit.ErrLLMGenerate) || !errors.Is(err, promptkit.ErrTransport) || !errors.Is(err, transportCause) {
|
|
t.Fatalf("transport result=%#v, error=%v", result, err)
|
|
}
|
|
if errors.Is(err, promptkit.ErrMalformedResponse) || strings.Contains(fmt.Sprintf("%#v", err), transportCause.Error()) {
|
|
t.Fatalf("transport category or formatting = %v", err)
|
|
}
|
|
|
|
for _, tc := range []struct {
|
|
name string
|
|
body string
|
|
}{
|
|
{name: "invalid framing", body: `{"choices":[]}{}`},
|
|
{name: "missing choices", body: `{"choices":[]}`},
|
|
{name: "missing content", body: `{"choices":[{"message":{}}]}`},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
engine := newBuiltInGenerationErrorEngine(t, http.StatusOK, tc.body)
|
|
result, err := engine.Run(context.Background(), generationErrorRunRequest())
|
|
if result != nil || !errors.Is(err, promptkit.ErrLLMGenerate) || !errors.Is(err, promptkit.ErrMalformedResponse) || errors.Is(err, promptkit.ErrTransport) {
|
|
t.Fatalf("malformed result=%#v, error=%v", result, err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestBuiltInGenerationPreservesContextIdentity(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name string
|
|
cause error
|
|
}{
|
|
{name: "canceled", cause: context.Canceled},
|
|
{name: "deadline", cause: context.DeadlineExceeded},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
config := contractConfig(frameworkSchemaDir)
|
|
config.HTTPClient = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
return nil, tc.cause
|
|
})}
|
|
engine, err := promptkit.NewEngine(config)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
result, err := engine.Run(context.Background(), generationErrorRunRequest())
|
|
if result != nil || !errors.Is(err, tc.cause) || !errors.Is(err, promptkit.ErrLLMGenerate) || !errors.Is(err, promptkit.ErrTransport) {
|
|
t.Fatalf("context result=%#v, error=%v", result, err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestInjectedGenerationFailureCategories(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name string
|
|
client *fakeLLMClient
|
|
malformed bool
|
|
}{
|
|
{name: "nil response", client: &fakeLLMClient{}, malformed: true},
|
|
{name: "arbitrary error", client: &fakeLLMClient{err: errors.New("injected failure")}},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
engine := newContractEngineWithOptions(t, frameworkSchemaDir, promptkit.WithLLMClient(tc.client))
|
|
result, err := engine.Run(context.Background(), generationErrorRunRequest())
|
|
if result != nil || !errors.Is(err, promptkit.ErrLLMGenerate) || errors.Is(err, promptkit.ErrTransport) || errors.Is(err, promptkit.ErrMalformedResponse) != tc.malformed {
|
|
t.Fatalf("injected result=%#v, error=%v", result, err)
|
|
}
|
|
if tc.client.err != nil && !errors.Is(err, tc.client.err) {
|
|
t.Fatalf("injected cause lost: %v", err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func assertGenerationError(t *testing.T, err error, statusCode int, code, providerType, message string) {
|
|
t.Helper()
|
|
|
|
if !errors.Is(err, promptkit.ErrLLMGenerate) {
|
|
t.Fatalf("errors.Is(%v, ErrLLMGenerate) = false", err)
|
|
}
|
|
var generationErr *promptkit.GenerationError
|
|
if !errors.As(err, &generationErr) || generationErr == nil {
|
|
t.Fatalf("error = %T, want *GenerationError", err)
|
|
}
|
|
if generationErr.StatusCode() != statusCode || generationErr.ProviderCode() != code || generationErr.ProviderType() != providerType || generationErr.ProviderMessage() != message {
|
|
t.Fatalf("GenerationError = %#v", generationErr)
|
|
}
|
|
|
|
wantFormatted := fmt.Sprintf("failed to generate output: provider returned HTTP status %d", statusCode)
|
|
for _, rendered := range []string{fmt.Sprintf("%v", generationErr), fmt.Sprintf("%+v", generationErr), fmt.Sprintf("%#v", generationErr)} {
|
|
if rendered != wantFormatted {
|
|
t.Fatalf("formatted error = %q, want %q", rendered, wantFormatted)
|
|
}
|
|
for _, marker := range []string{code, providerType, message} {
|
|
if marker != "" && strings.Contains(rendered, marker) {
|
|
t.Fatalf("formatted error exposed provider marker %q: %q", marker, rendered)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func newBuiltInGenerationErrorEngine(t *testing.T, statusCode int, body string) *promptkit.Engine {
|
|
t.Helper()
|
|
|
|
config := contractConfig(frameworkSchemaDir)
|
|
noRetries := 0
|
|
config.TransientRetries = &noRetries
|
|
config.HTTPClient = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
return &http.Response{
|
|
StatusCode: statusCode,
|
|
ContentLength: int64(len(body)),
|
|
Body: io.NopCloser(strings.NewReader(body)),
|
|
}, nil
|
|
})}
|
|
engine, err := promptkit.NewEngine(config)
|
|
if err != nil {
|
|
t.Fatalf("NewEngine: %v", err)
|
|
}
|
|
return engine
|
|
}
|
|
|
|
func generationErrorRunRequest() promptkit.RunRequest {
|
|
return promptkit.RunRequest{
|
|
PromptID: frameworkMarkdownSummaryPromptID,
|
|
Inputs: map[string]promptkit.ArtifactRef{
|
|
"transcript": promptkit.Inline("Rin opens the gate."),
|
|
"glossary": promptkit.Inline("gate: A guarded passage."),
|
|
},
|
|
}
|
|
}
|