Use PromptKit prepared execution for structured model calls
This commit is contained in:
257
internal/framework/llm/adapter_test.go
Normal file
257
internal/framework/llm/adapter_test.go
Normal file
@@ -0,0 +1,257 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
ck "gitea.maximumdirect.net/eric/ConfigKit"
|
||||
"gitea.maximumdirect.net/eric/audita/internal/core/config"
|
||||
"gitea.maximumdirect.net/eric/audita/internal/framework/responseschema"
|
||||
"gitea.maximumdirect.net/eric/audita/internal/framework/structuredoutput"
|
||||
"gitea.maximumdirect.net/eric/audita/internal/prompts"
|
||||
pk "gitea.maximumdirect.net/eric/promptkit"
|
||||
)
|
||||
|
||||
type reply struct {
|
||||
status int
|
||||
content string
|
||||
}
|
||||
|
||||
func adapterFixture(t *testing.T, replies []reply, retries, repairs int, hook func(*http.Request)) (*Runtime, *int, string) {
|
||||
t.Helper()
|
||||
count := 0
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
count++
|
||||
if hook != nil {
|
||||
hook(r)
|
||||
}
|
||||
step := replies[len(replies)-1]
|
||||
if count <= len(replies) {
|
||||
step = replies[count-1]
|
||||
}
|
||||
if step.status == 0 {
|
||||
connection, _, hijackErr := w.(http.Hijacker).Hijack()
|
||||
if hijackErr != nil {
|
||||
t.Error(hijackErr)
|
||||
return
|
||||
}
|
||||
_ = connection.Close()
|
||||
return
|
||||
}
|
||||
if step.status != http.StatusOK {
|
||||
w.WriteHeader(step.status)
|
||||
_, _ = fmt.Fprint(w, `{"error":{"message":"provider error"}}`)
|
||||
return
|
||||
}
|
||||
if step.content == "MALFORMED_ENVELOPE" {
|
||||
_, _ = fmt.Fprint(w, `{"choices":[]}`)
|
||||
return
|
||||
}
|
||||
encoded, _ := json.Marshal(step.content)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = fmt.Fprintf(w, `{"model":"returned-model","choices":[{"message":{"content":%s}}],"usage":{"prompt_tokens":2,"completion_tokens":3,"total_tokens":5}}`, encoded)
|
||||
}))
|
||||
t.Cleanup(server.Close)
|
||||
t.Setenv("ADAPTER_KEY", "adapter-secret")
|
||||
base := fmt.Sprintf("version: 2\nbackends:\n offline:\n endpoint: %s/v1\n credential: ADAPTER_KEY\nllm:\n proposal_profile: base\n transient_retries: %d\n repair_attempts: %d\n", server.URL, retries, repairs)
|
||||
dir := profileLayout(t, base, "")
|
||||
main := filepath.Join(dir, "config.yml")
|
||||
profile := "id: base\nbackend: offline\nmodel: full-model\n"
|
||||
if err := os.WriteFile(filepath.Join(dir, "profiles", "base.yaml"), []byte(profile), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
loaded, err := config.Load(ck.File(main), "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
runtime, err := NewRuntime(context.Background(), loaded.Config, main, server.Client())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return runtime, &count, dir
|
||||
}
|
||||
|
||||
func adapterRequest(schema responseschema.Key) structuredoutput.Request {
|
||||
id := prompts.PromptIDModuleGlossaryProposal
|
||||
vars := map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}
|
||||
if schema == responseschema.ValidatorDecisionSetKey {
|
||||
id = prompts.PromptIDValidatorEditorialReview
|
||||
vars = map[string]string{"TranscriptDescriptionBlock": "", "PayloadJSON": "[]"}
|
||||
}
|
||||
return structuredoutput.Request{PromptID: id, PromptVersion: prompts.VersionV1, Variables: vars, Schema: schema, Context: structuredoutput.DiagnosticContext{ModuleKey: "glossary"}}
|
||||
}
|
||||
|
||||
func TestPreparedAdapterRepairsAndClassifies(t *testing.T) {
|
||||
for _, schema := range []responseschema.Key{responseschema.CorrectionSetKey, responseschema.ValidatorDecisionSetKey} {
|
||||
valid := `{"corrections":[]}`
|
||||
if schema == responseschema.ValidatorDecisionSetKey {
|
||||
valid = `{"validations":[]}`
|
||||
}
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
replies []reply
|
||||
retries, repairs, attempts, repairCalls, transientRetries int
|
||||
malformed bool
|
||||
cause error
|
||||
}{
|
||||
{"success", []reply{{200, valid}}, 0, 0, 1, 0, 0, false, nil},
|
||||
{"repair succeeds", []reply{{200, `{"wrong":true}`}, {200, valid}}, 0, 1, 2, 1, 0, false, nil},
|
||||
{"repair exhausted", []reply{{200, `{"wrong":true}`}, {200, `{"wrong":true}`}}, 0, 1, 2, 1, 0, true, nil},
|
||||
{"429 recovery", []reply{{429, ""}, {200, valid}}, 1, 0, 2, 0, 1, false, nil},
|
||||
{"5xx recovery", []reply{{503, ""}, {200, valid}}, 1, 0, 2, 0, 1, false, nil},
|
||||
{"zero retry", []reply{{503, ""}}, 0, 0, 1, 0, 0, false, pk.ErrLLMGenerate},
|
||||
{"positive retry exhausted", []reply{{503, ""}, {503, ""}}, 1, 0, 2, 0, 1, false, pk.ErrRetriesExhausted},
|
||||
{"transport failure", []reply{{0, ""}}, 0, 0, 1, 0, 0, false, pk.ErrTransport},
|
||||
{"malformed envelope", []reply{{200, "MALFORMED_ENVELOPE"}}, 0, 0, 1, 0, 0, false, pk.ErrMalformedResponse},
|
||||
{"retry and repair", []reply{{503, ""}, {200, `{"wrong":true}`}, {503, ""}, {200, valid}}, 1, 1, 4, 1, 2, false, nil},
|
||||
} {
|
||||
t.Run(string(schema)+"/"+tc.name, func(t *testing.T) {
|
||||
runtime, count, _ := adapterFixture(t, tc.replies, tc.retries, tc.repairs, nil)
|
||||
role := "proposal"
|
||||
if schema == responseschema.ValidatorDecisionSetKey {
|
||||
role = "validation"
|
||||
}
|
||||
var events []pk.ExecutionEvent
|
||||
var records []structuredoutput.CallRecord
|
||||
client, err := runtime.NewRoleClient(role, CallHooks{OnEvent: func(e pk.ExecutionEvent) { events = append(events, e) }, OnCall: func(r structuredoutput.CallRecord) { records = append(records, r) }})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var out map[string]any
|
||||
response, err := client.CompleteStructured(context.Background(), adapterRequest(schema), &out)
|
||||
if *count != tc.attempts {
|
||||
t.Fatalf("HTTP attempts=%d, want %d", *count, tc.attempts)
|
||||
}
|
||||
if tc.malformed != structuredoutput.IsMalformedError(err) {
|
||||
t.Fatalf("malformed classification: %v", err)
|
||||
}
|
||||
if tc.cause != nil && !errors.Is(err, tc.cause) {
|
||||
t.Fatalf("cause %v missing from %v", tc.cause, err)
|
||||
}
|
||||
if tc.cause == nil && !tc.malformed && err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tc.name == "zero retry" && errors.Is(err, pk.ErrRetriesExhausted) {
|
||||
t.Fatal("zero retry reported exhaustion")
|
||||
}
|
||||
if response.Prepared == nil || response.Prompt.DefinitionHash == "" || response.Prompt.RenderedHash == "" || response.Prompt.ProfileID != "base" || response.Prompt.BackendID != "offline" {
|
||||
t.Fatalf("prepared metadata missing: %+v", response)
|
||||
}
|
||||
if len(records) != 2 || records[0].Finished.IsZero() == false || records[1].Finished.IsZero() || records[1].Context.Role != role {
|
||||
t.Fatalf("logical records: %+v", records)
|
||||
}
|
||||
if tc.malformed && records[1].FinalContent == "" {
|
||||
t.Fatal("invalid final content lost")
|
||||
}
|
||||
if len(events) == 0 || events[len(events)-1].Kind != pk.ExecutionFinished || int(events[len(events)-1].RepairCalls) != tc.repairCalls || int(events[len(events)-1].TransientRetries) != tc.transientRetries {
|
||||
t.Fatalf("execution counts: %+v", events)
|
||||
}
|
||||
if err == nil && (out == nil || string(response.Content) != valid || response.TotalTokens != 5*(tc.repairCalls+1) || response.Model != "full-model") {
|
||||
t.Fatalf("decoded response: %+v, %v", response, out)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPreparedAdapterMetadataAndTargetGuard(t *testing.T) {
|
||||
requestSeen := false
|
||||
var posted struct {
|
||||
Messages []structuredoutput.Message `json:"messages"`
|
||||
}
|
||||
runtime, count, _ := adapterFixture(t, []reply{{200, `{"corrections":[]}`}}, 0, 0, func(r *http.Request) {
|
||||
requestSeen = true
|
||||
if r.Header.Get("Authorization") != "Bearer adapter-secret" {
|
||||
t.Error("backend credential not used")
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&posted); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
})
|
||||
client, err := runtime.NewRoleClient("proposal", CallHooks{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
req := adapterRequest(responseschema.CorrectionSetKey)
|
||||
req.Variables["HardeningPolicy"] = "attacker-controlled"
|
||||
var out map[string]any
|
||||
response, err := client.CompleteStructured(context.Background(), req, &out)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !requestSeen || *count != 1 || response.Prepared == nil || len(response.Prepared.Messages) == 0 {
|
||||
t.Fatal("prepared call missing")
|
||||
}
|
||||
if len(posted.Messages) != len(response.Prepared.Messages) {
|
||||
t.Fatal("prepared and dispatched message counts differ")
|
||||
}
|
||||
for i := range posted.Messages {
|
||||
if posted.Messages[i] != response.Prepared.Messages[i] {
|
||||
t.Fatal("prepared messages differ from dispatched request")
|
||||
}
|
||||
}
|
||||
found := false
|
||||
for _, msg := range response.Prepared.Messages {
|
||||
if strings.Contains(msg.Content, prompts.HardeningText()) {
|
||||
found = true
|
||||
}
|
||||
if strings.Contains(msg.Content, "attacker-controlled") {
|
||||
t.Fatal("reserved hardening variable was overridden")
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatal("hardening policy missing")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPreparedAdapterRejectsTargetDriftAndPreparationErrors(t *testing.T) {
|
||||
runtime, count, dir := adapterFixture(t, []reply{{200, `{"corrections":[]}`}}, 0, 0, nil)
|
||||
client, err := runtime.NewRoleClient("proposal", CallHooks{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
req := adapterRequest(responseschema.CorrectionSetKey)
|
||||
req.PromptID = "missing.prompt"
|
||||
var out map[string]any
|
||||
if _, err := client.CompleteStructured(context.Background(), req, &out); err == nil || *count != 0 {
|
||||
t.Fatalf("preparation failure dispatched: %v, %d", err, *count)
|
||||
}
|
||||
// PrepareExecution reloads the profile, so a changed target must be blocked.
|
||||
if err := os.WriteFile(filepath.Join(dir, "profiles", "base.yaml"), []byte("id: base\nbackend: offline\nendpoint: http://127.0.0.1:1/v1\nmodel: full-model\n"), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
req = adapterRequest(responseschema.CorrectionSetKey)
|
||||
_, err = client.CompleteStructured(context.Background(), req, &out)
|
||||
if err == nil || !strings.Contains(err.Error(), "target changed") || *count != 0 {
|
||||
t.Fatalf("target drift dispatched: %v, %d", err, *count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPreparedAdapterCancellationStopsRetryWait(t *testing.T) {
|
||||
runtime, count, _ := adapterFixture(t, []reply{{503, ""}}, 2, 0, nil)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
var once sync.Once
|
||||
client, err := runtime.NewRoleClient("proposal", CallHooks{OnEvent: func(e pk.ExecutionEvent) {
|
||||
if e.Kind == pk.ExecutionAttempt && e.Dispatched {
|
||||
once.Do(cancel)
|
||||
}
|
||||
}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var out map[string]any
|
||||
_, err = client.CompleteStructured(ctx, adapterRequest(responseschema.CorrectionSetKey), &out)
|
||||
if !errors.Is(err, context.Canceled) || *count != 1 {
|
||||
t.Fatalf("cancellation: %v, HTTP calls=%d", err, *count)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user